From 163bcf2fc555d56620f986a7040a93a1dc88429c Mon Sep 17 00:00:00 2001 From: roroghost17 Date: Tue, 28 Jul 2026 19:18:55 +0530 Subject: [PATCH] refactor: updates modelcatalog to memoise provider configs recomputing for allowed models / providers --- framework/modelcatalog/datasheet/store.go | 12 ++++ framework/modelcatalog/keyconfig/store.go | 14 ++++ framework/modelcatalog/live/store.go | 14 ++++ framework/modelcatalog/main.go | 7 ++ framework/modelcatalog/models.go | 82 ++++++++++++++++++++--- 5 files changed, 118 insertions(+), 11 deletions(-) diff --git a/framework/modelcatalog/datasheet/store.go b/framework/modelcatalog/datasheet/store.go index a915e705d46..e3df581322c 100644 --- a/framework/modelcatalog/datasheet/store.go +++ b/framework/modelcatalog/datasheet/store.go @@ -4,6 +4,7 @@ import ( "slices" "strings" "sync" + "sync/atomic" "time" bifrost "github.com/maximhq/bifrost/core" @@ -67,6 +68,10 @@ type Store struct { datasheetByProvider map[schemas.ModelProvider][]string // rebuilt every reload deprecatedByProvider map[schemas.ModelProvider][]string // rebuilt every reload + // writeGen counts membership rebuilds; the composer's model→provider memo + // stamps it to detect staleness. Atomic so readers skip mu. + writeGen atomic.Uint64 + // Overrides under their own mutex: writes here don't block pricing reads // (the hot CalculateCost path takes mu.RLock and overridesMu.RLock // independently and the orderings never invert). @@ -500,4 +505,11 @@ func (s *Store) rebuildDatasheetViewUnsafe() { slices.Sort(models) s.deprecatedByProvider[provider] = models } + + s.writeGen.Add(1) +} + +// WriteGen returns the monotonic count of membership rebuilds. +func (s *Store) WriteGen() uint64 { + return s.writeGen.Load() } diff --git a/framework/modelcatalog/keyconfig/store.go b/framework/modelcatalog/keyconfig/store.go index 8c29f8bd6f3..14465edd4fc 100644 --- a/framework/modelcatalog/keyconfig/store.go +++ b/framework/modelcatalog/keyconfig/store.go @@ -18,6 +18,7 @@ import ( "slices" "strings" "sync" + "sync/atomic" bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" @@ -60,6 +61,10 @@ type Store struct { mu sync.RWMutex entries map[schemas.ModelProvider]*providerState logger schemas.Logger + + // writeGen counts entry mutations; the composer's model→provider memo + // stamps it to detect staleness. Atomic so readers skip mu. + writeGen atomic.Uint64 } // New constructs an empty Store. @@ -88,6 +93,12 @@ func (s *Store) Replace(snapshot map[schemas.ModelProvider][]schemas.Key) { } } s.entries = next + s.writeGen.Add(1) +} + +// WriteGen returns the monotonic count of entry mutations. +func (s *Store) WriteGen() uint64 { + return s.writeGen.Load() } // SetProvider replaces the cached state for one provider. Call after a @@ -98,9 +109,11 @@ func (s *Store) SetProvider(provider schemas.ModelProvider, keys []schemas.Key) defer s.mu.Unlock() if st == nil { delete(s.entries, provider) + s.writeGen.Add(1) return } s.entries[provider] = st + s.writeGen.Add(1) } // RemoveProvider drops all state for the provider. Call on provider delete. @@ -108,6 +121,7 @@ func (s *Store) RemoveProvider(provider schemas.ModelProvider) { s.mu.Lock() defer s.mu.Unlock() delete(s.entries, provider) + s.writeGen.Add(1) } // EntriesFor returns all per-key entries for the provider (including diff --git a/framework/modelcatalog/live/store.go b/framework/modelcatalog/live/store.go index 783e8a67458..47dabb6aa77 100644 --- a/framework/modelcatalog/live/store.go +++ b/framework/modelcatalog/live/store.go @@ -18,6 +18,7 @@ package live import ( "slices" "sync" + "sync/atomic" bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" @@ -50,6 +51,11 @@ type Store struct { // look current again after the provider was re-added. gen map[schemas.ModelProvider]uint64 logger schemas.Logger + + // writeGen counts entry mutations across all providers; the composer's + // model→provider memo stamps it to detect staleness. Distinct from gen, + // which is per-provider and guards in-flight fetch races. + writeGen atomic.Uint64 } func New(logger schemas.Logger) *Store { @@ -73,6 +79,12 @@ func (s *Store) Upsert(provider schemas.ModelProvider, keyID string, unfiltered s.mu.Lock() s.entries[k] = Entry{Models: cp} s.mu.Unlock() + s.writeGen.Add(1) +} + +// WriteGen returns the monotonic count of entry mutations across all providers. +func (s *Store) WriteGen() uint64 { + return s.writeGen.Load() } // Generation returns the provider's current invalidation counter. Read it @@ -102,6 +114,7 @@ func (s *Store) UpsertIfCurrent(provider schemas.ModelProvider, keyID string, un return false } s.entries[k] = Entry{Models: cp} + s.writeGen.Add(1) return true } @@ -113,6 +126,7 @@ func (s *Store) UpsertIfCurrent(provider schemas.ModelProvider, keyID string, un // nothing cached yet, and that is the write this has to stop. func (s *Store) bumpLocked(provider schemas.ModelProvider) { s.gen[provider]++ + s.writeGen.Add(1) } // Invalidate drops both filtered and unfiltered entries for one key. Called diff --git a/framework/modelcatalog/main.go b/framework/modelcatalog/main.go index ea44998beaa..c105d53be5e 100644 --- a/framework/modelcatalog/main.go +++ b/framework/modelcatalog/main.go @@ -40,6 +40,12 @@ type ModelCatalog struct { // loadCapabilities fills a capability cache miss. loadCapabilities func(schemas.ModelProvider, string) (*schemas.ModelCapabilities, error) + // providerMemo caches GetProvidersForModel per model, stamped with + // catalogGeneration() at compute time; any store write invalidates every + // entry. Capped at providerMemoMaxEntries, flushed on overflow. + providerMemoMu sync.RWMutex + providerMemo map[string]providerMemoEntry + // MCP library sync configuration (protected by syncMu) mcpLibraryURL string mcpLibrarySyncInterval time.Duration @@ -114,6 +120,7 @@ func Init(ctx context.Context, config *Config, configStore configstore.ConfigSto live: live.New(logger), keyconf: keyconfig.New(logger), capabilities: lrucache.New[*schemas.ModelCapabilities](capabilityCacheSize), + providerMemo: make(map[string]providerMemoEntry), done: make(chan struct{}), } mc.syncCtx, mc.syncCancel = context.WithCancel(ctx) diff --git a/framework/modelcatalog/models.go b/framework/modelcatalog/models.go index 92e36b98a06..716e267aa71 100644 --- a/framework/modelcatalog/models.go +++ b/framework/modelcatalog/models.go @@ -150,11 +150,63 @@ func (mc *ModelCatalog) GetDistinctBaseModelNames() []string { return mc.datasheet.DistinctBaseModelNames() } +// providerMemoEntry is a memoized GetProvidersForModel result and the +// catalogGeneration it was computed under. +type providerMemoEntry struct { + gen uint64 + val []schemas.ModelProvider +} + +// providerMemoMaxEntries bounds the memo; model strings are client-controlled +// and would otherwise grow it without limit. On overflow the map is flushed +// and hot entries repopulate on demand. +const providerMemoMaxEntries = 4096 + +// catalogGeneration sums the three stores' write counters. Every write bumps +// exactly one counter by one, so the sum strictly increases and no two +// catalog states share a value. +func (mc *ModelCatalog) catalogGeneration() uint64 { + return mc.datasheet.WriteGen() + mc.keyconf.WriteGen() + mc.live.WriteGen() +} + // GetProvidersForModel returns every provider that can serve the model. -// Composes across stores and applies the cross-provider special cases -// (openrouter / vertex / groq-gpt / bedrock-claude) preserved verbatim from -// the pre-refactor implementation. +// Memoized per model (bounded, empty results uncached); any store write +// invalidates (see catalogGeneration). Returns a fresh clone the caller +// may mutate. func (mc *ModelCatalog) GetProvidersForModel(model string) []schemas.ModelProvider { + gen := mc.catalogGeneration() + + mc.providerMemoMu.RLock() + entry, ok := mc.providerMemo[model] + mc.providerMemoMu.RUnlock() + if ok && entry.gen == gen { + return slices.Clone(entry.val) + } + + val := mc.computeProvidersForModel(model) + // Unknown models resolve to nothing; skipping memoising + if len(val) == 0 { + return val + } + + // Stamp with the generation read before compute: a write during compute + // leaves the entry stale and the next call recomputes. + mc.providerMemoMu.Lock() + if mc.providerMemo == nil { + mc.providerMemo = make(map[string]providerMemoEntry) + } + if _, exists := mc.providerMemo[model]; !exists && len(mc.providerMemo) >= providerMemoMaxEntries { + mc.providerMemo = make(map[string]providerMemoEntry) + } + mc.providerMemo[model] = providerMemoEntry{gen: gen, val: val} + mc.providerMemoMu.Unlock() + + return slices.Clone(val) +} + +// computeProvidersForModel composes the uncached answer across stores, +// including the openrouter / vertex / groq-gpt / bedrock-claude special cases. +func (mc *ModelCatalog) computeProvidersForModel(model string) []schemas.ModelProvider { baseModel := mc.datasheet.BaseModelName(model) providers := make([]schemas.ModelProvider, 0) @@ -265,17 +317,25 @@ func (mc *ModelCatalog) IsModelAllowedForProvider(provider schemas.ModelProvider return false } + // Bare-name match needs no catalog access and covers most allowlists. + if slices.Contains(allowedModels, model) { + return true + } + + // Only provider-prefixed entries ("openai/gpt-4o") need the provider + // catalog; build it once, and only when one exists. + if !slices.ContainsFunc(allowedModels, func(m string) bool { return strings.Contains(m, "/") }) { + return false + } providerCatalogModels := mc.GetModelsForProvider(provider) for _, allowedModel := range allowedModels { - if allowedModel == model { - return true + if !strings.Contains(allowedModel, "/") { + continue } - if strings.Contains(allowedModel, "/") { - if slices.Contains(providerCatalogModels, allowedModel) { - _, modelPart := schemas.ParseModelString(allowedModel, "") - if modelPart == model { - return true - } + if slices.Contains(providerCatalogModels, allowedModel) { + _, modelPart := schemas.ParseModelString(allowedModel, "") + if modelPart == model { + return true } } }