diff options
Diffstat (limited to 'internal/catalog/catalog.go')
| -rw-r--r-- | internal/catalog/catalog.go | 103 |
1 files changed, 103 insertions, 0 deletions
diff --git a/internal/catalog/catalog.go b/internal/catalog/catalog.go new file mode 100644 index 0000000..76e33b1 --- /dev/null +++ b/internal/catalog/catalog.go @@ -0,0 +1,103 @@ +package catalog + +import ( + "fmt" + "sort" + "strings" + "sync/atomic" + + "aigw/internal/config" + "aigw/internal/domain" +) + +type Catalog struct { + state atomic.Pointer[snapshot] +} + +type snapshot struct { + models map[string]domain.Model + list []domain.Model +} + +func New(cfg config.Config) *Catalog { + providers := make(map[string]domain.Provider, len(cfg.Providers)) + for _, provider := range cfg.Providers { + providers[provider.ID] = domain.Provider{ + ID: provider.ID, + Protocol: provider.Protocol, + BaseURL: strings.TrimRight(provider.BaseURL, "/"), + APIKey: provider.APIKey, + } + } + + models := make([]domain.Model, 0, len(cfg.Models)) + for _, modelCfg := range cfg.Models { + model := domain.Model{ID: modelCfg.ID, OwnedBy: modelCfg.OwnedBy} + for _, route := range modelCfg.Routes { + model.Routes = append(model.Routes, domain.Route{ + Provider: providers[route.Provider], + UpstreamModel: route.UpstreamModel, + Priority: route.Priority, + Weight: route.Weight, + }) + } + models = append(models, model) + } + return NewModels(models) +} + +func NewModels(models []domain.Model) *Catalog { + catalog := &Catalog{} + catalog.Replace(models) + return catalog +} + +func (c *Catalog) Replace(source []domain.Model) { + models := make(map[string]domain.Model, len(source)) + list := make([]domain.Model, 0, len(source)) + for _, sourceModel := range source { + model := sourceModel + model.Routes = append([]domain.Route(nil), sourceModel.Routes...) + models[model.ID] = model + list = append(list, model) + } + sort.Slice(list, func(i, j int) bool { return list[i].ID < list[j].ID }) + c.state.Store(&snapshot{models: models, list: list}) +} + +func (c *Catalog) Model(id string) (domain.Model, error) { + current := c.state.Load() + if current == nil { + return domain.Model{}, fmt.Errorf("model %q not found", id) + } + model, ok := current.models[id] + if !ok { + return domain.Model{}, fmt.Errorf("model %q not found", id) + } + return model, nil +} + +func (c *Catalog) Models(protocol domain.Protocol) []domain.Model { + current := c.state.Load() + if current == nil { + return nil + } + result := make([]domain.Model, 0, len(current.list)) + for _, model := range current.list { + for _, route := range model.Routes { + if route.Provider.Protocol == protocol { + result = append(result, model) + break + } + } + } + return result +} + +func (c *Catalog) Count() int { + current := c.state.Load() + if current == nil { + return 0 + } + return len(current.list) +} |
