1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
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)
}
|
