89 lines
2.2 KiB
Go
89 lines
2.2 KiB
Go
package router
|
|
|
|
import (
|
|
"fmt"
|
|
"sync"
|
|
|
|
"github.com/edgeai/gateway/internal/config"
|
|
)
|
|
|
|
// LogicalModelMapping maps logical model names to actual model configurations.
|
|
type LogicalModelMapping struct {
|
|
mu sync.RWMutex
|
|
mapping map[string]*ModelTarget
|
|
}
|
|
|
|
// ModelTarget represents the resolved target for a logical model.
|
|
type ModelTarget struct {
|
|
LogicalModel string
|
|
ActualModel string
|
|
Provider string
|
|
Endpoint string
|
|
ContextWindow int
|
|
MaxOutputTokens int
|
|
MaxConcurrency int
|
|
CancelSupported bool
|
|
}
|
|
|
|
// NewLogicalModelMapping creates a mapping from config.
|
|
func NewLogicalModelMapping(cfg *config.Config) *LogicalModelMapping {
|
|
m := &LogicalModelMapping{mapping: make(map[string]*ModelTarget)}
|
|
for logical, mc := range cfg.Models {
|
|
m.mapping[logical] = &ModelTarget{
|
|
LogicalModel: logical,
|
|
ActualModel: mc.ActualModel,
|
|
Provider: mc.Provider,
|
|
Endpoint: mc.Endpoint,
|
|
ContextWindow: mc.ContextWindow,
|
|
MaxOutputTokens: mc.MaxOutputTokens,
|
|
MaxConcurrency: mc.MaxConcurrency,
|
|
CancelSupported: mc.CancelSupported,
|
|
}
|
|
}
|
|
return m
|
|
}
|
|
|
|
// Resolve returns the ModelTarget for a logical model name.
|
|
func (m *LogicalModelMapping) Resolve(logicalModel string) (*ModelTarget, error) {
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
|
|
target, ok := m.mapping[logicalModel]
|
|
if !ok {
|
|
return nil, fmt.Errorf("logical model not found: %s", logicalModel)
|
|
}
|
|
return target, nil
|
|
}
|
|
|
|
// List returns all logical model names.
|
|
func (m *LogicalModelMapping) List() []string {
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
|
|
names := make([]string, 0, len(m.mapping))
|
|
for n := range m.mapping {
|
|
names = append(names, n)
|
|
}
|
|
return names
|
|
}
|
|
|
|
// Update updates the mapping (for config hot-reload).
|
|
func (m *LogicalModelMapping) Update(cfg *config.Config) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
|
|
m.mapping = make(map[string]*ModelTarget)
|
|
for logical, mc := range cfg.Models {
|
|
m.mapping[logical] = &ModelTarget{
|
|
LogicalModel: logical,
|
|
ActualModel: mc.ActualModel,
|
|
Provider: mc.Provider,
|
|
Endpoint: mc.Endpoint,
|
|
ContextWindow: mc.ContextWindow,
|
|
MaxOutputTokens: mc.MaxOutputTokens,
|
|
MaxConcurrency: mc.MaxConcurrency,
|
|
CancelSupported: mc.CancelSupported,
|
|
}
|
|
}
|
|
}
|