package kb import ( "encoding/json" "os" "github.com/yaoapp/gou/graphrag" "github.com/yaoapp/gou/graphrag/graph/neo4j" "github.com/yaoapp/gou/graphrag/types" "github.com/yaoapp/gou/graphrag/vector/qdrant" "github.com/yaoapp/gou/store" "github.com/yaoapp/kun/log" ) // Config parses the Knowledge Base configuration // ParseConfigFromJSON parses config from JSON bytes func ParseConfigFromJSON(data []byte) (*Config, error) { var config Config if err := json.Unmarshal(data, &config); err != nil { return nil, err } return &config, nil } // ParseConfigFromFile parses config from JSON file func ParseConfigFromFile(filename string) (*Config, error) { data, err := os.ReadFile(filename) if err != nil { return nil, err } return ParseConfigFromJSON(data) } // ToJSON converts config to JSON bytes func (c *Config) ToJSON() ([]byte, error) { return json.MarshalIndent(c, "", " ") } // GraphRagConfig converts KB config to GraphRag config func (c *Config) GraphRagConfig() (*graphrag.Config, error) { config := &graphrag.Config{ Logger: log.StandardLogger(), System: "__yao_kb_system", // Default system collection name Vector: nil, Graph: nil, Store: nil, } // Configure Vector Store (required) vectorStore, err := c.createVectorStore() if err != nil { return nil, err } config.Vector = vectorStore // Configure Graph Store (optional) if c.Graph != nil { graphStore, err := c.createGraphStore() if err != nil { return nil, err } config.Graph = graphStore } // Configure Store storeName := c.getStoreName() kvStore, err := store.Get(storeName) if err != nil { return nil, err } config.Store = kvStore return config, nil } // getStoreName returns the store name, using default if not configured func (c *Config) getStoreName() string { if c.Store != "" { return c.Store } return "__yao.kb.store" } // createVectorStore creates a vector store from config func (c *Config) createVectorStore() (types.VectorStore, error) { switch c.Vector.Driver { case "qdrant": // Convert config to VectorStoreConfig vectorConfig, err := c.toVectorStoreConfig() if err != nil { return nil, err } return qdrant.NewStoreWithConfig(vectorConfig), nil default: return nil, nil } } // createGraphStore creates a graph store from config func (c *Config) createGraphStore() (types.GraphStore, error) { switch c.Graph.Driver { case "neo4j": // Convert config to GraphStoreConfig graphConfig, err := c.toGraphStoreConfig() if err != nil { return nil, err } return neo4j.NewStoreWithConfig(graphConfig), nil default: return nil, nil } } // toVectorStoreConfig converts the vector config to VectorStoreConfig func (c *Config) toVectorStoreConfig() (types.VectorStoreConfig, error) { // Convert map[string]interface{} to types.VectorStoreConfig via JSON jsonData, err := json.Marshal(c.Vector.Config) if err != nil { return types.VectorStoreConfig{}, err } var config types.VectorStoreConfig if err := json.Unmarshal(jsonData, &config); err != nil { return types.VectorStoreConfig{}, err } return config, nil } // toGraphStoreConfig converts the graph config to GraphStoreConfig func (c *Config) toGraphStoreConfig() (types.GraphStoreConfig, error) { // Convert map[string]interface{} to types.GraphStoreConfig via JSON jsonData, err := json.Marshal(c.Graph.Config) if err != nil { return types.GraphStoreConfig{}, err } var config types.GraphStoreConfig if err := json.Unmarshal(jsonData, &config); err != nil { return types.GraphStoreConfig{}, err } return config, nil } // UnmarshalJSON implements json.Unmarshaler interface func (c *Config) UnmarshalJSON(data []byte) error { // Use alias type to avoid infinite recursion raw := (*RawConfig)(c) if err := json.Unmarshal(data, raw); err != nil { return err } // Compute features after parsing c.Features = c.ComputeFeatures() return nil } // MarshalJSON implements json.Marshaler interface func (c *Config) MarshalJSON() ([]byte, error) { // Use alias type for standard JSON marshaling (Features field is ignored) raw := (*RawConfig)(c) return json.Marshal(raw) } // ComputeFeatures calculates available features based on current configuration func (c *Config) ComputeFeatures() Features { features := Features{} // Core features features.GraphDatabase = c.Graph != nil features.PDFProcessing = c.PDF != nil features.VideoProcessing = c.FFmpeg != nil // File format support (based on converters) converterMap := make(map[string]bool) for _, provider := range c.Converters { converterMap[provider.ID] = true } features.PlainText = true // Plain text is always supported as a basic feature features.OfficeDocuments = converterMap["__yao.office"] features.OCRProcessing = converterMap["__yao.ocr"] features.AudioTranscript = converterMap["__yao.whisper"] features.ImageAnalysis = converterMap["__yao.vision"] // Advanced features features.EntityExtraction = len(c.Extractors) > 0 features.WebFetching = len(c.Fetchers) > 0 features.CustomSearch = len(c.Searchers) > 0 features.ResultReranking = len(c.Rerankers) > 0 features.SegmentVoting = len(c.Votes) > 0 features.SegmentWeighting = len(c.Weights) > 0 features.SegmentScoring = len(c.Scores) > 0 return features }