yao/kb/config.go

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
}