107 lines
2.9 KiB
Go
107 lines
2.9 KiB
Go
package provider
|
|
|
|
import (
|
|
"fmt"
|
|
"sort"
|
|
|
|
"llm-gateway/internal/config"
|
|
)
|
|
|
|
// Route maps a model to a specific provider with pricing.
|
|
type Route struct {
|
|
Provider Provider
|
|
ProviderModel string
|
|
Priority int
|
|
InputPrice float64 // per 1M tokens
|
|
OutputPrice float64 // per 1M tokens
|
|
}
|
|
|
|
// Registry maps model names to provider routes.
|
|
type Registry struct {
|
|
routes map[string][]Route
|
|
order []string // preserves config order
|
|
}
|
|
|
|
func NewRegistry(cfg *config.Config) (*Registry, error) {
|
|
// Build providers
|
|
providers := make(map[string]Provider)
|
|
for _, pc := range cfg.Providers {
|
|
providers[pc.Name] = NewOpenAIProvider(pc.Name, pc.BaseURL, pc.APIKey, pc.Timeout)
|
|
}
|
|
|
|
// Build routes (preserving config order)
|
|
routes := make(map[string][]Route)
|
|
order := make([]string, 0, len(cfg.Models))
|
|
for _, mc := range cfg.Models {
|
|
var modelRoutes []Route
|
|
for _, rc := range mc.Routes {
|
|
p, ok := providers[rc.Provider]
|
|
if !ok {
|
|
return nil, fmt.Errorf("model %s: unknown provider %s", mc.Name, rc.Provider)
|
|
}
|
|
pc := cfg.ProviderByName(rc.Provider)
|
|
priority := pc.Priority
|
|
modelRoutes = append(modelRoutes, Route{
|
|
Provider: p,
|
|
ProviderModel: rc.Model,
|
|
Priority: priority,
|
|
InputPrice: rc.Pricing.Input,
|
|
OutputPrice: rc.Pricing.Output,
|
|
})
|
|
}
|
|
// Sort by priority (lower = higher priority)
|
|
sort.Slice(modelRoutes, func(i, j int) bool {
|
|
return modelRoutes[i].Priority < modelRoutes[j].Priority
|
|
})
|
|
routes[mc.Name] = modelRoutes
|
|
order = append(order, mc.Name)
|
|
}
|
|
|
|
return &Registry{routes: routes, order: order}, nil
|
|
}
|
|
|
|
// Lookup returns the routes for a model name.
|
|
func (r *Registry) Lookup(model string) ([]Route, bool) {
|
|
routes, ok := r.routes[model]
|
|
return routes, ok
|
|
}
|
|
|
|
// ModelNames returns all registered model names in config order.
|
|
func (r *Registry) ModelNames() []string {
|
|
return r.order
|
|
}
|
|
|
|
// RouteInfo exposes route details for dashboard display.
|
|
type RouteInfo struct {
|
|
ProviderName string `json:"provider_name"`
|
|
ProviderModel string `json:"provider_model"`
|
|
Priority int `json:"priority"`
|
|
InputPrice float64 `json:"input_price"`
|
|
OutputPrice float64 `json:"output_price"`
|
|
}
|
|
|
|
// ModelRouteInfo exposes a model and its routes for dashboard display.
|
|
type ModelRouteInfo struct {
|
|
Name string `json:"name"`
|
|
Routes []RouteInfo `json:"routes"`
|
|
}
|
|
|
|
// AllRoutes returns all models and their routes in config order.
|
|
func (r *Registry) AllRoutes() []ModelRouteInfo {
|
|
results := make([]ModelRouteInfo, 0, len(r.order))
|
|
for _, name := range r.order {
|
|
routes := r.routes[name]
|
|
info := ModelRouteInfo{Name: name}
|
|
for _, rt := range routes {
|
|
info.Routes = append(info.Routes, RouteInfo{
|
|
ProviderName: rt.Provider.Name(),
|
|
ProviderModel: rt.ProviderModel,
|
|
Priority: rt.Priority,
|
|
InputPrice: rt.InputPrice,
|
|
OutputPrice: rt.OutputPrice,
|
|
})
|
|
}
|
|
results = append(results, info)
|
|
}
|
|
return results
|
|
}
|