refactor: replaced ModelMatcher interface with ModelCatalog struct to avoid having to do to nil checks using reflect

Author: danpiths

framework/modelcatalog/main.go

@@ -634,3 +634,17 @@ func (mc *ModelCatalog) Cleanup() error {
 
 	return nil
 }
+
+// NewTestCatalog creates a minimal ModelCatalog for testing purposes.
+// It does not start background sync workers or connect to external services.
+func NewTestCatalog(baseModelIndex map[string]string) *ModelCatalog {
+	if baseModelIndex == nil {
+		baseModelIndex = make(map[string]string)
+	}
+	return &ModelCatalog{
+		modelPool:      make(map[schemas.ModelProvider][]string),
+		baseModelIndex: baseModelIndex,
+		pricingData:    make(map[string]configstoreTables.TableModelPricing),
+		done:           make(chan struct{}),
+	}
+}

plugins/governance/model_provider_governance_test.go

@@ -1947,7 +1947,7 @@ func TestStore_CheckModelBudget_CrossProviderModelMatch(t *testing.T) {
 	budget := buildBudgetWithUsage("budget1", 100.0, 100.0, "1h") // At limit
 	modelConfig := buildModelConfig("mc1", "gpt-4o", nil, budget, nil)
 
-	mc := newMockModelMatcher(t)
+	mc := newTestModelCatalog(t)
 	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
 		ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig},
 		Budgets:      []configstoreTables.TableBudget{*budget},
@@ -1967,7 +1967,7 @@ func TestStore_CheckModelBudget_CrossProviderModelMatch_WithinLimit(t *testing.T
 	budget := buildBudget("budget1", 100.0, "1h")
 	modelConfig := buildModelConfig("mc1", "gpt-4o", nil, budget, nil)
 
-	mc := newMockModelMatcher(t)
+	mc := newTestModelCatalog(t)
 	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
 		ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig},
 		Budgets:      []configstoreTables.TableBudget{*budget},
@@ -1985,7 +1985,7 @@ func TestStore_CheckModelRateLimit_CrossProviderModelMatch(t *testing.T) {
 	rateLimit := buildRateLimitWithUsage("rl1", 10000, 10000, 1000, 0) // Token limit at max
 	modelConfig := buildModelConfig("mc1", "gpt-4o", nil, nil, rateLimit)
 
-	mc := newMockModelMatcher(t)
+	mc := newTestModelCatalog(t)
 	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
 		ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig},
 		RateLimits:   []configstoreTables.TableRateLimit{*rateLimit},
@@ -2005,7 +2005,7 @@ func TestStore_UpdateModelBudgetUsage_CrossProviderModelMatch(t *testing.T) {
 	budget := buildBudget("budget1", 100.0, "1h")
 	modelConfig := buildModelConfig("mc1", "gpt-4o", nil, budget, nil)
 
-	mc := newMockModelMatcher(t)
+	mc := newTestModelCatalog(t)
 	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
 		ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig},
 		Budgets:      []configstoreTables.TableBudget{*budget},
@@ -2033,7 +2033,7 @@ func TestStore_UpdateModelRateLimitUsage_CrossProviderModelMatch(t *testing.T) {
 	rateLimit := buildRateLimitWithUsage("rl1", 100, 0, 1000, 0) // Low token limit
 	modelConfig := buildModelConfig("mc1", "gpt-4o", nil, nil, rateLimit)
 
-	mc := newMockModelMatcher(t)
+	mc := newTestModelCatalog(t)
 	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
 		ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig},
 		RateLimits:   []configstoreTables.TableRateLimit{*rateLimit},
@@ -2059,7 +2059,7 @@ func TestStore_CheckModelBudget_ModelWithProvider_ExactMatchOnly(t *testing.T) {
 	providerStr := "openai"
 	modelConfig := buildModelConfig("mc1", "gpt-4o", &providerStr, budget, nil)
 
-	mc := newMockModelMatcher(t)
+	mc := newTestModelCatalog(t)
 	store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{
 		ModelConfigs: []configstoreTables.TableModelConfig{*modelConfig},
 		Budgets:      []configstoreTables.TableBudget{*budget},

plugins/governance/store.go

@@ -13,16 +13,10 @@ import (
 	"github.com/maximhq/bifrost/core/schemas"
 	"github.com/maximhq/bifrost/framework/configstore"
 	configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables"
+	"github.com/maximhq/bifrost/framework/modelcatalog"
 	"gorm.io/gorm"
 )
 
-// ModelMatcher provides cross-provider model name matching.
-// This is satisfied by *modelcatalog.ModelCatalog.
-type ModelMatcher interface {
-	GetBaseModelName(model string) string
-	IsSameModel(model1, model2 string) bool
-}
-
 // LocalGovernanceStore provides in-memory cache for governance data with fast, non-blocking access
 type LocalGovernanceStore struct {
 	// Core data maps using sync.Map for lock-free reads
@@ -50,8 +44,8 @@ type LocalGovernanceStore struct {
 	// Config store for refresh operations
 	configStore configstore.ConfigStore
 
-	// Model matcher for cross-provider model matching (optional)
-	modelMatcher ModelMatcher
+	// Model catalog for cross-provider model matching (optional)
+	modelCatalog *modelcatalog.ModelCatalog
 
 	// Logger
 	logger schemas.Logger
@@ -141,9 +135,9 @@ type GovernanceStore interface {
 }
 
 // NewLocalGovernanceStore creates a new in-memory governance store
-// The modelMatcher parameter is optional (can be nil) and enables cross-provider model matching
+// The modelCatalog parameter is optional (can be nil) and enables cross-provider model matching
 // for governance lookups (e.g., "openai/gpt-4o" matching config for "gpt-4o").
-func NewLocalGovernanceStore(ctx context.Context, logger schemas.Logger, configStore configstore.ConfigStore, governanceConfig *configstore.GovernanceConfig, modelMatcher ModelMatcher) (*LocalGovernanceStore, error) {
+func NewLocalGovernanceStore(ctx context.Context, logger schemas.Logger, configStore configstore.ConfigStore, governanceConfig *configstore.GovernanceConfig, modelCatalog *modelcatalog.ModelCatalog) (*LocalGovernanceStore, error) {
 	// Create singleton CEL environment once for all routing rule compilations
 	env, err := createCELEnvironment()
 	if err != nil {
@@ -154,7 +148,7 @@ func NewLocalGovernanceStore(ctx context.Context, logger schemas.Logger, configS
 		configStore:                    configStore,
 		logger:                         logger,
 		routingCELEnv:                  env,
-		modelMatcher:                   modelMatcher,
+		modelCatalog:                   modelCatalog,
 		LastDBUsagesBudgets:            make(map[string]float64),
 		LastDBUsagesRequestsRateLimits: make(map[string]int64),
 		LastDBUsagesTokensRateLimits:   make(map[string]int64),
@@ -599,8 +593,8 @@ func (gs *LocalGovernanceStore) CheckProviderRateLimit(ctx context.Context, requ
 // Returns the matching config and the display name for error messages.
 func (gs *LocalGovernanceStore) findModelOnlyConfig(model string) (*configstoreTables.TableModelConfig, string) {
 	// If modelMatcher is available, try normalized base model name first (cross-provider matching)
-	if gs.modelMatcher != nil {
-		baseName := gs.modelMatcher.GetBaseModelName(model)
+	if gs.modelCatalog != nil {
+		baseName := gs.modelCatalog.GetBaseModelName(model)
 		if baseName != model {
 			if value, exists := gs.modelConfigs.Load(baseName); exists && value != nil {
 				if mc, ok := value.(*configstoreTables.TableModelConfig); ok && mc != nil {
@@ -1787,8 +1781,8 @@ func (gs *LocalGovernanceStore) rebuildInMemoryStructures(ctx context.Context, c
 		} else {
 			// Global config (applies to all providers) - store under normalized model name
 			key := mc.ModelName
-			if gs.modelMatcher != nil {
-				key = gs.modelMatcher.GetBaseModelName(mc.ModelName)
+			if gs.modelCatalog != nil {
+				key = gs.modelCatalog.GetBaseModelName(mc.ModelName)
 			}
 			gs.modelConfigs.Store(key, mc)
 		}
@@ -2406,8 +2400,8 @@ func (gs *LocalGovernanceStore) UpdateModelConfigInMemory(mc *configstoreTables.
 		gs.modelConfigs.Store(key, &clone)
 	} else {
 		key := clone.ModelName
-		if gs.modelMatcher != nil {
-			key = gs.modelMatcher.GetBaseModelName(clone.ModelName)
+		if gs.modelCatalog != nil {
+			key = gs.modelCatalog.GetBaseModelName(clone.ModelName)
 		}
 		gs.modelConfigs.Store(key, &clone)
 	}

plugins/governance/test_utils.go

@@ -325,46 +325,13 @@ func fetchDatasheetBaseIndex() {
 	datasheetBaseIndex = index
 }
 
-// mockModelMatcher implements ModelMatcher for testing.
-// It fetches the default datasheet to build a base model index,
-// matching the production ModelCatalog.GetBaseModelName behavior:
-// 1. Direct lookup in base model index
-// 2. Strip provider prefix and retry lookup
-// 3. Fallback to algorithmic date/version stripping
-type mockModelMatcher struct {
-	baseModelIndex map[string]string
-}
-
-func (m *mockModelMatcher) GetBaseModelName(model string) string {
-	// Step 1: Direct lookup in base model index
-	if base, ok := m.baseModelIndex[model]; ok {
-		return base
-	}
-
-	// Step 2: Strip provider prefix and try again
-	_, baseName := schemas.ParseModelString(model, "")
-	if baseName != model {
-		if base, ok := m.baseModelIndex[baseName]; ok {
-			return base
-		}
-	}
-
-	// Step 3: Fallback to algorithmic date/version stripping
-	return schemas.BaseModelName(baseName)
-}
-
-func (m *mockModelMatcher) IsSameModel(model1, model2 string) bool {
-	if model1 == model2 {
-		return true
-	}
-	return m.GetBaseModelName(model1) == m.GetBaseModelName(model2)
-}
-
-func newMockModelMatcher(t *testing.T) ModelMatcher {
+// newTestModelCatalog creates a test ModelCatalog using the fetched datasheet base model index.
+// This provides proper nil-pointer semantics (unlike an interface wrapper).
+func newTestModelCatalog(t *testing.T) *modelcatalog.ModelCatalog {
 	t.Helper()
 	datasheetOnce.Do(fetchDatasheetBaseIndex)
 	if datasheetErr != nil {
-		t.Skipf("skipping: failed to fetch datasheet for mock model matcher: %v", datasheetErr)
+		t.Skipf("skipping: failed to fetch datasheet for test model catalog: %v", datasheetErr)
 	}
-	return &mockModelMatcher{baseModelIndex: datasheetBaseIndex}
+	return modelcatalog.NewTestCatalog(datasheetBaseIndex)
 }