Merge pull request #1093 from trheyi/main

Enhance knowledge base provider management with multi-language support

Author: trheyi

kb/kb.go

@@ -23,7 +23,8 @@ var Instance types.GraphRag = nil
 
 // KnowledgeBase is the Knowledge Base instance
 type KnowledgeBase struct {
-	Config *kbtypes.Config // Knowledge Base configuration
+	Config    *kbtypes.Config         // Knowledge Base configuration
+	Providers *kbtypes.ProviderConfig // Multi-language provider configurations
 	*graphrag.GraphRag
 }
 
@@ -52,6 +53,13 @@ func Load(appConfig config.Config) (*KnowledgeBase, error) {
 		return nil, err
 	}
 
+	// Load providers from directories
+	providers, err := kbtypes.LoadProviders("kb")
+	if err != nil {
+		return nil, err
+	}
+	config.Providers = providers
+
 	// Set global configurations for providers to use
 	kbtypes.SetGlobalPDF(config.PDF)
 	kbtypes.SetGlobalFFmpeg(config.FFmpeg)
@@ -69,7 +77,7 @@ func Load(appConfig config.Config) (*KnowledgeBase, error) {
 	}
 
 	// Set the instance
-	instance := &KnowledgeBase{Config: &config, GraphRag: graphRag}
+	instance := &KnowledgeBase{Config: &config, Providers: providers, GraphRag: graphRag}
 
 	// Set the instance to the global variable
 	Instance = instance
@@ -88,48 +96,13 @@ func GetProviders(typ string, ids []string, locale string) ([]kbtypes.Provider,
 		return nil, fmt.Errorf("knowledge base not initialized")
 	}
 
-	// Get the configuration
-	conf := knowledgeBase.Config
-	if conf == nil {
-		return nil, fmt.Errorf("configuration not found")
+	// Default locale to "en" if empty
+	if locale == "" {
+		locale = "en"
 	}
 
-	providers := []*kbtypes.Provider{}
-	switch typ {
-	case "chunking":
-		providers = conf.Chunkings
-
-	case "converter":
-		providers = conf.Converters
-
-	case "embedding":
-		providers = conf.Embeddings
-
-	case "extractor":
-		providers = conf.Extractors
-
-	case "fetcher":
-		providers = conf.Fetchers
-
-	case "searcher":
-		providers = conf.Searchers
-
-	case "reranker":
-		providers = conf.Rerankers
-
-	case "vote":
-		providers = conf.Votes
-
-	case "weight":
-		providers = conf.Weights
-
-	case "score":
-		providers = conf.Scores
-
-	default:
-		return nil, fmt.Errorf("invalid provider type: %s", typ)
-
-	}
+	// Get providers for the requested type and language
+	providers := knowledgeBase.Providers.GetProviders(typ, locale)
 
 	// Filter empty ids
 	filteredIds := []string{}
@@ -149,8 +122,13 @@ func GetProviders(typ string, ids []string, locale string) ([]kbtypes.Provider,
 	return filteredProviders, nil
 }
 
-// GetProvider returns a provider by id
+// GetProvider returns a provider by id with default language "en"
 func GetProvider(typ string, id string) (*kbtypes.Provider, error) {
+	return GetProviderWithLanguage(typ, id, "en")
+}
+
+// GetProviderWithLanguage returns a provider by id, type, and language
+func GetProviderWithLanguage(typ string, id string, locale string) (*kbtypes.Provider, error) {
 	if Instance == nil {
 		return nil, fmt.Errorf("knowledge base not initialized")
 	}
@@ -160,53 +138,10 @@ func GetProvider(typ string, id string) (*kbtypes.Provider, error) {
 		return nil, fmt.Errorf("knowledge base not initialized")
 	}
 
-	conf := knowledgeBase.Config
-	if conf == nil {
-		return nil, fmt.Errorf("configuration not found")
-	}
-
-	providers := []*kbtypes.Provider{}
-	switch typ {
-	case "chunking":
-		providers = conf.Chunkings
-
-	case "converter":
-		providers = conf.Converters
-
-	case "embedding":
-		providers = conf.Embeddings
-
-	case "extractor":
-		providers = conf.Extractors
-
-	case "fetcher":
-		providers = conf.Fetchers
-
-	case "searcher":
-		providers = conf.Searchers
-
-	case "reranker":
-		providers = conf.Rerankers
-
-	case "vote":
-		providers = conf.Votes
-
-	case "weight":
-		providers = conf.Weights
-
-	case "score":
-		providers = conf.Scores
-
-	default:
-		return nil, fmt.Errorf("invalid provider type: %s", typ)
-	}
-
-	// Find the provider by id
-	for _, provider := range providers {
-		if provider.ID == id {
-			return provider, nil
-		}
+	// Default locale to "en" if empty
+	if locale == "" {
+		locale = "en"
 	}
 
-	return nil, fmt.Errorf("provider %s not found", id)
+	return knowledgeBase.Providers.GetProvider(typ, id, locale)
 }

kb/kb_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/yaoapp/yao/config"
+	kbtypes "github.com/yaoapp/yao/kb/types"
 	"github.com/yaoapp/yao/test"
 )
 
@@ -12,8 +13,134 @@ func TestLoad(t *testing.T) {
 	test.Prepare(&testing.T{}, config.Conf)
 	defer test.Clean()
 
+	kb, err := Load(config.Conf)
+	if err != nil {
+		t.Fatalf("Failed to load knowledge base: %v", err)
+	}
+
+	// Test that providers are loaded
+	if kb != nil && kb.Providers != nil {
+		t.Logf("Knowledge base loaded successfully with providers")
+	}
+}
+
+func TestGetProviders(t *testing.T) {
+	// Setup
+	test.Prepare(&testing.T{}, config.Conf)
+	defer test.Clean()
+
+	_, err := Load(config.Conf)
+	if err != nil {
+		t.Fatalf("Failed to load knowledge base: %v", err)
+	}
+
+	// Test getting providers for different languages
+	testCases := []struct {
+		providerType string
+		locale       string
+		expectEmpty  bool
+	}{
+		{"chunking", "en", false},
+		{"embedding", "en", false},
+		{"chunking", "zh-cn", false},
+		{"embedding", "zh-cn", false},
+		{"chunking", "nonexistent", false}, // Should fallback to "en"
+	}
+
+	for _, tc := range testCases {
+		providers, err := GetProviders(tc.providerType, []string{}, tc.locale)
+		if err != nil {
+			t.Errorf("Failed to get %s providers for locale %s: %v", tc.providerType, tc.locale, err)
+			continue
+		}
+
+		if tc.expectEmpty && len(providers) > 0 {
+			t.Errorf("Expected empty providers for %s/%s, got %d", tc.providerType, tc.locale, len(providers))
+		} else if !tc.expectEmpty && len(providers) == 0 {
+			t.Logf("No providers found for %s/%s (this may be expected if no provider files exist)", tc.providerType, tc.locale)
+		} else {
+			t.Logf("Found %d providers for %s/%s", len(providers), tc.providerType, tc.locale)
+		}
+	}
+}
+
+func TestGetProviderWithLanguage(t *testing.T) {
+	// Setup
+	test.Prepare(&testing.T{}, config.Conf)
+	defer test.Clean()
+
 	_, err := Load(config.Conf)
 	if err != nil {
 		t.Fatalf("Failed to load knowledge base: %v", err)
 	}
+
+	// Test getting a specific provider with language
+	provider, err := GetProviderWithLanguage("chunking", "__yao.structured", "en")
+	if err != nil {
+		t.Logf("Provider __yao.structured not found for chunking/en: %v (this may be expected if provider files don't exist)", err)
+	} else {
+		t.Logf("Found provider: %s", provider.ID)
+	}
+
+	// Test language fallback
+	provider, err = GetProviderWithLanguage("chunking", "__yao.structured", "nonexistent")
+	if err != nil {
+		t.Logf("Provider __yao.structured not found with fallback: %v (this may be expected if provider files don't exist)", err)
+	} else {
+		t.Logf("Found provider with fallback: %s", provider.ID)
+	}
+}
+
+func TestLoadProviders(t *testing.T) {
+	// Setup
+	test.Prepare(&testing.T{}, config.Conf)
+	defer test.Clean()
+
+	// Test loading providers from a directory
+	providers, err := kbtypes.LoadProviders("kb")
+	if err != nil {
+		t.Fatalf("Failed to load providers: %v", err)
+	}
+
+	if providers == nil {
+		t.Fatal("Providers config is nil")
+	}
+
+	// Check if provider maps are initialized
+	if providers.Chunkings == nil {
+		t.Error("Chunkings map is nil")
+	}
+	if providers.Embeddings == nil {
+		t.Error("Embeddings map is nil")
+	}
+
+	t.Logf("Loaded providers successfully")
+}
+
+func TestProviderConfigGetProviders(t *testing.T) {
+	// Setup
+	test.Prepare(&testing.T{}, config.Conf)
+	defer test.Clean()
+
+	providers, err := kbtypes.LoadProviders("kb")
+	if err != nil {
+		t.Fatalf("Failed to load providers: %v", err)
+	}
+
+	// Test getting providers for different types and languages
+	testCases := []string{"chunking", "embedding", "converter", "extractor", "fetcher"}
+
+	for _, providerType := range testCases {
+		// Test with "en"
+		enProviders := providers.GetProviders(providerType, "en")
+		t.Logf("Found %d %s providers for 'en'", len(enProviders), providerType)
+
+		// Test with "zh-cn"
+		zhProviders := providers.GetProviders(providerType, "zh-cn")
+		t.Logf("Found %d %s providers for 'zh-cn'", len(zhProviders), providerType)
+
+		// Test with nonexistent language (should fallback to "en")
+		fallbackProviders := providers.GetProviders(providerType, "nonexistent")
+		t.Logf("Found %d %s providers for 'nonexistent' (fallback)", len(fallbackProviders), providerType)
+	}
 }

kb/types/config.go

@@ -272,20 +272,31 @@ func (c *Config) resolveAllEnvVars() error {
 
 // resolveProviderEnvVars resolves environment variables in provider configurations
 func (c *Config) resolveProviderEnvVars() error {
-	providerLists := [][]*Provider{
-		c.Chunkings, c.Embeddings, c.Converters, c.Extractors,
-		c.Fetchers, c.Searchers, c.Rerankers, c.Votes, c.Weights, c.Scores,
-	}
-
-	for _, providers := range providerLists {
-		for _, provider := range providers {
-			for _, option := range provider.Options {
-				if option.Properties != nil {
-					resolved, err := c.resolveEnvVars(option.Properties)
-					if err != nil {
-						return err
+	if c.Providers == nil {
+		return nil
+	}
+
+	// Resolve env vars for all provider types and languages
+	providerMaps := []map[string][]*Provider{
+		c.Providers.Chunkings, c.Providers.Embeddings, c.Providers.Converters, c.Providers.Extractors,
+		c.Providers.Fetchers, c.Providers.Searchers, c.Providers.Rerankers, c.Providers.Votes,
+		c.Providers.Weights, c.Providers.Scores,
+	}
+
+	for _, providerMap := range providerMaps {
+		if providerMap == nil {
+			continue
+		}
+		for _, providers := range providerMap {
+			for _, provider := range providers {
+				for _, option := range provider.Options {
+					if option.Properties != nil {
+						resolved, err := c.resolveEnvVars(option.Properties)
+						if err != nil {
+							return err
+						}
+						option.Properties = resolved
 					}
-					option.Properties = resolved
 				}
 			}
 		}
@@ -335,8 +346,13 @@ func (c *Config) ComputeFeatures() Features {
 
 	// File format support (based on converters)
 	converterMap := make(map[string]bool)
-	for _, provider := range c.Converters {
-		converterMap[provider.ID] = true
+	if c.Providers != nil && c.Providers.Converters != nil {
+		// Check all languages for converter availability
+		for _, providers := range c.Providers.Converters {
+			for _, provider := range providers {
+				converterMap[provider.ID] = true
+			}
+		}
 	}
 
 	features.PlainText = true // Plain text is always supported as a basic feature
@@ -346,13 +362,28 @@ func (c *Config) ComputeFeatures() Features {
 	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
+	if c.Providers != nil {
+		features.EntityExtraction = c.hasProvidersInAnyLanguage(c.Providers.Extractors)
+		features.WebFetching = c.hasProvidersInAnyLanguage(c.Providers.Fetchers)
+		features.CustomSearch = c.hasProvidersInAnyLanguage(c.Providers.Searchers)
+		features.ResultReranking = c.hasProvidersInAnyLanguage(c.Providers.Rerankers)
+		features.SegmentVoting = c.hasProvidersInAnyLanguage(c.Providers.Votes)
+		features.SegmentWeighting = c.hasProvidersInAnyLanguage(c.Providers.Weights)
+		features.SegmentScoring = c.hasProvidersInAnyLanguage(c.Providers.Scores)
+	}
 
 	return features
 }
+
+// hasProvidersInAnyLanguage checks if there are providers available in any language
+func (c *Config) hasProvidersInAnyLanguage(providerMap map[string][]*Provider) bool {
+	if providerMap == nil {
+		return false
+	}
+	for _, providers := range providerMap {
+		if len(providers) > 0 {
+			return true
+		}
+	}
+	return false
+}