mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-10-03 23:54:06 +00:00
此前 AUTO 链是全局单值(cfg.Auto + Core.AutoChain()),所有用户共用一条链。
管理员无法为某个用户单独指定调度链。
按 per-key 覆盖 + 全局兜底实现:
- config.GWKey 增加 Auto 字段与 HasOwnAuto()。未配置即继承全局链,
存量部署零改动,新密钥天然继承全局链。
- Core 把 buildAutoChain 的归一化逻辑抽成 chainForRules,全局链、
生图链、per-key 链共用同一套编译,避免两处漂移。
- Core 增加 keyAutoChains 缓存 + AutoChainFor(key)。请求路径读缓存不
加锁,与全局链查询一致。缓存整表原子替换,不会看到半成品。
- 请求侧 chat.go 改用 AutoChainFor(reqKey)。冷却仍在 Provider 上按
model+source 共享:两条链指向同一个 slot 时共用冷却,与今天单链行为
相同,也避免为 per-key 维度重构冷却而改变现有可观测语义。
- API:GET/PUT/DELETE /api/keys/{key}/auto(admin),GET
/api/keys/me/auto(任意角色,只能读自己的)。空 PUT 与 DELETE 等价于
"恢复继承",无法持久化一条会 503 的空链。写入时回报解析出的槽位数,
让管理员当场看到模型名写错,而不是等用户下次请求 503。
- 源变更时一并重编译 per-key 链,加源后无需重启即可生效。
判据 18 条,5 个变异全部被抓住:AutoChainFor 忽略 key、清空后不重建
缓存、源变更不重建、空 PUT 落盘成空链、me/auto 误要求 admin。
UI 判据做变异时发现漏放:只查函数定义存在,删掉按钮后仍通过。已改为
断言 keyCanvasHtml 内的调用点。
端到端实测(真实 HTTP + 两个 mock 上游):admin 与 bob 初始同为 m-fast,
给 bob 配 m-cheap 后两者分流,重启后仍分流,DELETE 后 bob 回到 m-fast。
Co-Authored-By: ModelRouter <noreply@modelrouter.dev>
648 lines
24 KiB
Go
648 lines
24 KiB
Go
// Package config loads the gateway YAML configuration plus a runtime overlay
|
||
// (web UI edits) and resolves them into sources with per-model priority.
|
||
package config
|
||
|
||
import (
|
||
"encoding/json"
|
||
"fmt"
|
||
"os"
|
||
"path/filepath"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"gopkg.in/yaml.v3"
|
||
)
|
||
|
||
// Config is the top-level gateway configuration.
|
||
type Config struct {
|
||
Path string `yaml:"-" json:"-"`
|
||
Listen string `yaml:"listen"`
|
||
GatewayKeys []string `yaml:"gateway_keys"`
|
||
DefaultModel string `yaml:"default_model"` // e.g. "AUTO" or a model id
|
||
AdapterDir string `yaml:"adapter_dir"`
|
||
PluginDir string `yaml:"plugin_dir,omitempty"` // Lua plugin directory (optional; hooks + WebUI extensions)
|
||
RuntimeFile string `yaml:"runtime_file"`
|
||
TLSCertFile string `yaml:"tls_cert_file,omitempty"` // PEM cert; when set together with tls_key_file, serve HTTPS
|
||
TLSKeyFile string `yaml:"tls_key_file,omitempty"` // PEM private key
|
||
PublicBaseURL string `yaml:"public_base_url,omitempty"` // external base for generated config snippets; default inferred from request
|
||
Sources []Source `yaml:"sources"`
|
||
Auto []ModelScope `yaml:"auto,omitempty"` // AUTO 调度链规则(WebUI 优先级页编辑,chat)
|
||
AutoImage []ModelScope `yaml:"auto_image,omitempty"` // AUTO 生图调度链规则(WebUI 优先级页·生图)
|
||
Keys []GWKey `yaml:"keys,omitempty"` // 网关密钥(WebUI 密钥页管理)
|
||
// BillingDSL declares per-URL pricing profiles for the billing plugin.
|
||
//
|
||
// Declarative, because hand-writing the prices JSON was error-prone in ways
|
||
// that looked successful: the first PUT went out without the required
|
||
// `prices` wrapper and silently REPLACED the accumulated totals with the
|
||
// price table; the second used a peak-window shape the plugin does not read,
|
||
// so peak traffic was billed at off-peak rates — no error, the bill just
|
||
// quietly halved. A schema-checked config file catches both at load.
|
||
BillingDSL *BillingDSL `yaml:"billing,omitempty"`
|
||
// box seals credentials (sources' api_key/headers, keys' key) at rest.
|
||
// In-memory values are always plaintext; only the bytes on disk are sealed.
|
||
// Wired by AttachSecretBox — Load leaves it nil so `-check` and tests stay
|
||
// filesystem-free.
|
||
box *SecretBox
|
||
}
|
||
|
||
// Defaults applied to any source (YAML or runtime) that leaves a field unset.
|
||
const (
|
||
DefaultSourceTimeout = 120 * time.Second
|
||
DefaultSourceQueueTimeout = 60 * time.Second
|
||
DefaultSourceConcurrency = 8
|
||
)
|
||
|
||
// Model is a single exposed model id bound to a source, with priority used by
|
||
// AUTO auto selection (higher number = preferred).
|
||
type Model struct {
|
||
ID string `yaml:"id" json:"id"`
|
||
Priority int `yaml:"priority" json:"priority"`
|
||
Kind string `yaml:"kind" json:"kind"` // "chat" (default) | "image"
|
||
Meta map[string]interface{} `yaml:"meta" json:"meta,omitempty"`
|
||
}
|
||
|
||
// Source describes a single upstream LLM provider.
|
||
type Source struct {
|
||
Name string `yaml:"name" json:"name"`
|
||
BaseURL string `yaml:"base_url" json:"base_url"`
|
||
APIKey string `yaml:"api_key" json:"api_key"`
|
||
APIKeyEnv string `yaml:"api_key_env,omitempty" json:"-"` // reference to an env var holding the key (overrides api_key)
|
||
Adapter string `yaml:"adapter" json:"adapter"`
|
||
Endpoint string `yaml:"endpoint" json:"endpoint,omitempty"` // chat endpoint override
|
||
ImageEndpoint string `yaml:"image_endpoint" json:"image_endpoint,omitempty"` // image endpoint override
|
||
Models []Model `yaml:"models" json:"models"`
|
||
Headers map[string]string `yaml:"headers" json:"headers,omitempty"`
|
||
// ProxyURL routes this source's HTTP(S) traffic through an explicit proxy
|
||
// (e.g. http://127.0.0.1:7890). Empty = direct connection. Needed for
|
||
// upstreams behind geo/IP blocks that only respond through a proxy.
|
||
ProxyURL string `yaml:"proxy_url" json:"proxy_url,omitempty"`
|
||
Meta map[string]interface{} `yaml:"meta" json:"meta,omitempty"`
|
||
Temperature float64 `yaml:"temperature" json:"temperature,omitempty"`
|
||
MaxTokens int `yaml:"max_tokens" json:"max_tokens,omitempty"`
|
||
Timeout time.Duration `yaml:"timeout" json:"-"`
|
||
MaxConcurrent int `yaml:"max_concurrent" json:"max_concurrent"`
|
||
QueueTimeout time.Duration `yaml:"queue_timeout" json:"-"`
|
||
RPM int `yaml:"rpm" json:"rpm,omitempty"` // optional requests-per-minute cap (0 = unlimited)
|
||
}
|
||
|
||
// Load reads and validates a config file. When the file does not exist yet a
|
||
// default config is generated at that path first (first-run bootstrap), so a
|
||
// fresh binary just works: `llmsproxy -config /path/to/config.yaml`.
|
||
func Load(path string) (*Config, error) {
|
||
if _, err := EnsureDefault(path); err != nil {
|
||
return nil, err
|
||
}
|
||
data, err := os.ReadFile(path)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
var cfg Config
|
||
if err := yaml.Unmarshal(data, &cfg); err != nil {
|
||
return nil, fmt.Errorf("parse config: %w", err)
|
||
}
|
||
cfg.Path = path
|
||
if err := cfg.ApplyDefaults(); err != nil {
|
||
return nil, err
|
||
}
|
||
return &cfg, nil
|
||
}
|
||
|
||
// EnsureDefault creates a default config file at path when it does not exist
|
||
// yet and returns whether it was created. The repo ships no config file
|
||
// (config files carry real keys); the binary generates one per install with a
|
||
// fresh random admin key. An existing file is never touched.
|
||
func EnsureDefault(path string) (bool, error) {
|
||
if _, err := os.Stat(path); err == nil {
|
||
return false, nil
|
||
} else if !os.IsNotExist(err) {
|
||
return false, err
|
||
}
|
||
if err := writeDefaultConfig(path); err != nil {
|
||
return false, err
|
||
}
|
||
return true, nil
|
||
}
|
||
|
||
// writeDefaultConfig writes a minimal, safe-by-default config: loopback-only
|
||
// listen and a fresh random admin key. No upstream sources are preconfigured
|
||
// (the repo ships none — add them via WebUI or config.yaml). adapter_dir /
|
||
// runtime_file live next to the config file so the binary works regardless of
|
||
// the working directory it is started from.
|
||
func writeDefaultConfig(path string) error {
|
||
key, err := NewGatewayKey()
|
||
if err != nil {
|
||
return fmt.Errorf("generate gateway key: %w", err)
|
||
}
|
||
dir := filepath.Dir(path)
|
||
abs, err := filepath.Abs(dir)
|
||
if err != nil {
|
||
abs = dir
|
||
}
|
||
cfg := Config{
|
||
Listen: "127.0.0.1:8080",
|
||
GatewayKeys: []string{key},
|
||
DefaultModel: "AUTO",
|
||
AdapterDir: filepath.Join(abs, "adapters"),
|
||
PluginDir: filepath.Join(abs, "plugins"),
|
||
RuntimeFile: filepath.Join(abs, "runtime.json"),
|
||
}
|
||
out, err := yaml.Marshal(&cfg)
|
||
if err != nil {
|
||
return fmt.Errorf("marshal default config: %w", err)
|
||
}
|
||
// The config holds the plaintext admin key — restrict permissions.
|
||
if err := os.MkdirAll(abs, 0755); err != nil {
|
||
return fmt.Errorf("mkdir config dir: %w", err)
|
||
}
|
||
return os.WriteFile(path, out, 0600)
|
||
}
|
||
|
||
// RemoveSourceFromYAML deletes the named source entry from the config file so
|
||
// the delete is a real one (no tombstone needed). Uses yaml.Node to preserve
|
||
// the rest of the file's comments and formatting.
|
||
func RemoveSourceFromYAML(path, name string) error {
|
||
data, err := os.ReadFile(path)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
var doc yaml.Node
|
||
if err := yaml.Unmarshal(data, &doc); err != nil {
|
||
return err
|
||
}
|
||
content := doc.Content
|
||
if len(content) == 0 {
|
||
return nil
|
||
}
|
||
root := content[0]
|
||
if root.Kind != yaml.MappingNode {
|
||
return nil
|
||
}
|
||
for i := 0; i+1 < len(root.Content); i += 2 {
|
||
key, val := root.Content[i], root.Content[i+1]
|
||
if key.Value != "sources" || val.Kind != yaml.SequenceNode {
|
||
continue
|
||
}
|
||
kept := val.Content[:0]
|
||
for _, item := range val.Content {
|
||
if item.Kind != yaml.MappingNode {
|
||
continue
|
||
}
|
||
found := false
|
||
for j := 0; j+1 < len(item.Content); j += 2 {
|
||
if item.Content[j].Value == "name" && item.Content[j+1].Value == name {
|
||
found = true
|
||
break
|
||
}
|
||
}
|
||
if !found {
|
||
kept = append(kept, item)
|
||
}
|
||
}
|
||
val.Content = kept
|
||
}
|
||
out, err := yaml.Marshal(&doc)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return os.WriteFile(path, out, 0644)
|
||
}
|
||
|
||
// UpsertSourceInYAML adds or updates a source entry in the YAML config file.
|
||
// Uses yaml.Node to preserve the rest of the file's comments and formatting.
|
||
// When box is non-nil the source's credentials are sealed before writing, so a
|
||
// newly added source never lands in the file as plaintext.
|
||
func UpsertSourceInYAML(path, name string, src Source, box *SecretBox) error {
|
||
if path == "" {
|
||
return fmt.Errorf("config path is empty")
|
||
}
|
||
if err := sealSource(&src, box); err != nil {
|
||
return err
|
||
}
|
||
data, err := os.ReadFile(path)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
var doc yaml.Node
|
||
if err := yaml.Unmarshal(data, &doc); err != nil {
|
||
return err
|
||
}
|
||
content := doc.Content
|
||
if len(content) == 0 {
|
||
return nil
|
||
}
|
||
root := content[0]
|
||
if root.Kind != yaml.MappingNode {
|
||
return nil
|
||
}
|
||
// Build yaml.Node for the new source value
|
||
srcRaw, _ := yaml.Marshal(src)
|
||
var srcNode yaml.Node
|
||
yaml.Unmarshal(srcRaw, &srcNode)
|
||
|
||
for i := 0; i+1 < len(root.Content); i += 2 {
|
||
key, val := root.Content[i], root.Content[i+1]
|
||
if key.Value != "sources" || val.Kind != yaml.SequenceNode {
|
||
continue
|
||
}
|
||
replaced := false
|
||
for _, item := range val.Content {
|
||
if item.Kind != yaml.MappingNode {
|
||
continue
|
||
}
|
||
for j := 0; j+1 < len(item.Content); j += 2 {
|
||
if item.Content[j].Value == "name" && item.Content[j+1].Value == name {
|
||
// Replace the existing entry with the new source
|
||
item.Content = srcNode.Content[0].Content
|
||
replaced = true
|
||
break
|
||
}
|
||
}
|
||
if replaced {
|
||
break
|
||
}
|
||
}
|
||
if !replaced {
|
||
val.Content = append(val.Content, srcNode.Content[0])
|
||
}
|
||
break
|
||
}
|
||
out, err := yaml.Marshal(&doc)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return os.WriteFile(path, out, 0644)
|
||
}
|
||
|
||
// Save writes the current config back to the YAML file (preserving comments
|
||
// via yaml.Node round-trip when possible, or full marshaling as fallback).
|
||
// Credentials are sealed on the way out and the in-memory copy is restored to
|
||
// plaintext afterwards, so callers keep working with usable values.
|
||
func (c *Config) Save() error {
|
||
if c.Path == "" {
|
||
return fmt.Errorf("config path is empty")
|
||
}
|
||
if err := c.sealInPlace(c.box); err != nil {
|
||
return err
|
||
}
|
||
out, err := yaml.Marshal(c)
|
||
c.unsealAfterWrite(c.box)
|
||
if err != nil {
|
||
return fmt.Errorf("marshal config: %w", err)
|
||
}
|
||
return os.WriteFile(c.Path, out, 0644)
|
||
}
|
||
|
||
// ApplyDefaults sets missing values and validates the config.
|
||
func (c *Config) ApplyDefaults() error {
|
||
if c.Listen == "" {
|
||
c.Listen = ":8080"
|
||
}
|
||
if c.AdapterDir == "" {
|
||
c.AdapterDir = "adapters"
|
||
}
|
||
// PluginDir is deliberately NOT defaulted: with no plugins a missing
|
||
// directory would be created and seeded on every start for nothing. An
|
||
// operator opts in by setting it (or by dropping files into it).
|
||
if c.RuntimeFile == "" {
|
||
c.RuntimeFile = "runtime.json"
|
||
}
|
||
if c.DefaultModel == "" {
|
||
c.DefaultModel = "AUTO"
|
||
}
|
||
seen := map[string]bool{}
|
||
for i := range c.Sources {
|
||
s := &c.Sources[i]
|
||
if s.Name == "" {
|
||
return fmt.Errorf("config: sources[%d] missing name", i)
|
||
}
|
||
if s.BaseURL == "" {
|
||
return fmt.Errorf("config: source %s missing base_url", s.Name)
|
||
}
|
||
if s.Adapter == "" {
|
||
s.Adapter = "openai"
|
||
}
|
||
if s.Timeout == 0 {
|
||
s.Timeout = DefaultSourceTimeout
|
||
}
|
||
if s.QueueTimeout == 0 {
|
||
s.QueueTimeout = DefaultSourceQueueTimeout
|
||
}
|
||
if s.MaxConcurrent == 0 {
|
||
s.MaxConcurrent = DefaultSourceConcurrency
|
||
}
|
||
if seen[s.Name] {
|
||
return fmt.Errorf("config: duplicate source name %q", s.Name)
|
||
}
|
||
seen[s.Name] = true
|
||
for j := range s.Models {
|
||
m := &s.Models[j]
|
||
if m.ID == "" {
|
||
return fmt.Errorf("config: source %s has a model without id", s.Name)
|
||
}
|
||
// Allow same model ID on multiple sources — disambiguation is via
|
||
// "source-model" / "source:model" / "source/model" pinning and
|
||
// the AUTO chain slots carry explicit (model, source) pairs.
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// RuntimeConfig is the legacy runtime file format (kept for migration only).
|
||
type RuntimeConfig struct {
|
||
Sources []Source `json:"sources,omitempty"`
|
||
SourceTemplates []SourceTemplate `json:"source_templates,omitempty"`
|
||
DeletedSources []string `json:"deleted_sources,omitempty"`
|
||
DeletedAdapters []string `json:"deleted_adapters,omitempty"`
|
||
PresetTemplates []string `json:"preset_templates,omitempty"` // preset names the user has seen (or deleted) — never re-seeded
|
||
Keys []GWKey `json:"keys,omitempty"`
|
||
Auto []ModelScope `json:"auto,omitempty"`
|
||
AutoImage []ModelScope `json:"auto_image,omitempty"`
|
||
}
|
||
|
||
// SourceTemplate stores reusable source configuration (everything except
|
||
// name and api_key) so the WebUI can spin up multiple key-bearing sources
|
||
// from one shared template.
|
||
type SourceTemplate struct {
|
||
Name string `json:"name"`
|
||
BaseURL string `json:"base_url"`
|
||
Adapter string `json:"adapter"`
|
||
Endpoint string `json:"endpoint,omitempty"`
|
||
ImageEndpoint string `json:"image_endpoint,omitempty"`
|
||
Models []Model `json:"models"`
|
||
Headers map[string]string `json:"headers,omitempty"`
|
||
// ProxyURL routes this source's HTTP(S) traffic through an explicit proxy.
|
||
ProxyURL string `json:"proxy_url,omitempty"`
|
||
Meta map[string]interface{} `json:"meta,omitempty"`
|
||
Temperature float64 `json:"temperature,omitempty"`
|
||
MaxTokens int `json:"max_tokens,omitempty"`
|
||
MaxConcurrent int `json:"max_concurrent,omitempty"`
|
||
RPM int `json:"rpm,omitempty"`
|
||
}
|
||
|
||
// GWKey is a gateway API key persisted in the config file. Role is "admin"
|
||
// (full management) or "user" (sees only its own key); Models is the allowed
|
||
// model scope with per-model token quota (0 = unlimited).
|
||
type GWKey struct {
|
||
Key string `yaml:"key" json:"key"`
|
||
Role string `yaml:"role" json:"role"`
|
||
Name string `yaml:"name,omitempty" json:"name,omitempty"`
|
||
// Auto is this key's own AUTO scheduling chain. It overrides the global
|
||
// cfg.Auto for requests authenticated with this key; nil (or an empty
|
||
// list) means "inherit the global chain", so keys created before per-key
|
||
// chains existed keep working untouched.
|
||
Auto []ModelScope `yaml:"auto,omitempty" json:"auto,omitempty"`
|
||
Models []ModelScope `yaml:"models,omitempty" json:"models,omitempty"`
|
||
Note string `yaml:"note,omitempty" json:"note,omitempty"`
|
||
CreatedAt int64 `yaml:"created_at,omitempty" json:"created_at,omitempty"`
|
||
Seed bool `yaml:"seed,omitempty" json:"seed,omitempty"` // true if migrated from config gateway_keys
|
||
}
|
||
|
||
// HasOwnAuto reports whether the key declares its own AUTO chain rather than
|
||
// inheriting the global one. An explicitly empty list still counts as "no own
|
||
// chain" so a user can clear an override back to the global default.
|
||
func (k GWKey) HasOwnAuto() bool { return len(k.Auto) > 0 }
|
||
|
||
// BillingDSL holds declarative per-URL pricing profiles for the billing plugin.
|
||
//
|
||
// Profiles are the "let the user pick" axis: the same upstream URL can be
|
||
// declared under multiple profiles, and the operator chooses which is active.
|
||
// Switching recomputes and re-injects the billing plugin's prices table.
|
||
type BillingDSL struct {
|
||
// Active is the id of the profile in effect. Empty => the first profile.
|
||
Active string `yaml:"active,omitempty" json:"active,omitempty"`
|
||
Profiles []BillingProfile `yaml:"profiles,omitempty" json:"profiles,omitempty"`
|
||
}
|
||
|
||
// BillingProfile is one named set of per-URL pricing rules.
|
||
type BillingProfile struct {
|
||
ID string `yaml:"id" json:"id"`
|
||
Label string `yaml:"label,omitempty" json:"label,omitempty"`
|
||
// Default applies when no rule matches: "free" or "unpriced".
|
||
Default string `yaml:"default,omitempty" json:"default,omitempty"`
|
||
Rules []BillingRule `yaml:"rules,omitempty" json:"rules,omitempty"`
|
||
}
|
||
|
||
// BillingRule declares the pricing for everything under one URL.
|
||
type BillingRule struct {
|
||
// URL matches a source's base_url. "*" = any unmatched URL.
|
||
URL string `yaml:"url" json:"url"`
|
||
// Mode is one of: free, token, subscription, unpriced.
|
||
Mode string `yaml:"mode" json:"mode"`
|
||
Currency string `yaml:"currency,omitempty" json:"currency,omitempty"`
|
||
// Subscription is the flat monthly cost (for mode=subscription).
|
||
Subscription float64 `yaml:"subscription,omitempty" json:"subscription,omitempty"`
|
||
// Models declares per-token prices when mode=token.
|
||
Models map[string]BillingToken `yaml:"models,omitempty" json:"models,omitempty"`
|
||
// Free is true when mode=free (kept for explicitness).
|
||
// Peak applies to ALL token-priced requests under this URL.
|
||
Peak *BillingPeak `yaml:"peak,omitempty" json:"peak,omitempty"`
|
||
}
|
||
|
||
// BillingToken is the per-model price.
|
||
type BillingToken struct {
|
||
Prompt string `yaml:"prompt" json:"prompt"` // USD per million tokens, e.g. "0.15"
|
||
Completion string `yaml:"completion" json:"completion"` // USD per million tokens
|
||
CacheDiscount *float64 `yaml:"cache_discount,omitempty" json:"cache_discount,omitempty"`
|
||
}
|
||
|
||
// BillingPeak declares peak-hour pricing.
|
||
type BillingPeak struct {
|
||
Multiplier float64 `yaml:"multiplier" json:"multiplier"`
|
||
// Weekdays: 1=Mon..5=Fri (Lua os.date !%w, 0=Sun)
|
||
Weekdays []int `yaml:"weekdays,omitempty" json:"weekdays,omitempty"`
|
||
// Hours: each pair is [lo,hi] inclusive
|
||
Hours [][2]int `yaml:"hours,omitempty" json:"hours,omitempty"`
|
||
}
|
||
|
||
// Validate rejects a BillingDSL that could not work as written.
|
||
func (b *BillingDSL) Validate() error {
|
||
if b == nil {
|
||
return nil
|
||
}
|
||
ids := map[string]bool{}
|
||
for i := range b.Profiles {
|
||
p := &b.Profiles[i]
|
||
if p.ID == "" {
|
||
return fmt.Errorf("billing profile [%d] has no id", i)
|
||
}
|
||
if ids[p.ID] {
|
||
return fmt.Errorf("billing profile id %q appears more than once", p.ID)
|
||
}
|
||
ids[p.ID] = true
|
||
if p.Default != "" && p.Default != "free" && p.Default != "unpriced" {
|
||
return fmt.Errorf("profile %q: default must be free or unpriced, got %q", p.ID, p.Default)
|
||
}
|
||
for j := range p.Rules {
|
||
r := &p.Rules[j]
|
||
if r.URL == "" {
|
||
return fmt.Errorf("profile %q rule [%d]: url is required", p.ID, j)
|
||
}
|
||
switch r.Mode {
|
||
case "free", "token", "subscription", "unpriced":
|
||
default:
|
||
return fmt.Errorf("profile %q rule [%d]: mode must be free/token/subscription/unpriced, got %q", p.ID, j, r.Mode)
|
||
}
|
||
if r.Mode == "token" {
|
||
for m, t := range r.Models {
|
||
if t.Prompt == "" || t.Completion == "" {
|
||
return fmt.Errorf("profile %q model %q: prompt and completion are required for token mode", p.ID, m)
|
||
}
|
||
if _, err := ParseFloatUSDPerM(t.Prompt); err != nil {
|
||
return fmt.Errorf("profile %q model %q prompt: %w", p.ID, m, err)
|
||
}
|
||
if _, err := ParseFloatUSDPerM(t.Completion); err != nil {
|
||
return fmt.Errorf("profile %q model %q completion: %w", p.ID, m, err)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
// Active must resolve to a profile id.
|
||
if b.Active != "" && !ids[b.Active] {
|
||
return fmt.Errorf("billing.active = %q but no profile has that id", b.Active)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// Resolve returns the profile with the given id; empty id (or an id that does
|
||
// not exist) falls back to the FIRST profile. A config with profiles but no
|
||
// active marker should mean "use the one I wrote first" rather than "price
|
||
// nothing", which is what a nil here would silently do.
|
||
func (b *BillingDSL) Resolve(id string) *BillingProfile {
|
||
if b == nil || len(b.Profiles) == 0 {
|
||
return nil
|
||
}
|
||
if id != "" {
|
||
for i := range b.Profiles {
|
||
if b.Profiles[i].ID == id {
|
||
return &b.Profiles[i]
|
||
}
|
||
}
|
||
}
|
||
return &b.Profiles[0]
|
||
}
|
||
|
||
// parseFloatUSDPerM parses a USD-per-million price string. Returns the
|
||
// per-single-token rate the plugin expects (price / 1e6).
|
||
//
|
||
// Kept as a string in the DSL rather than a float: "0.15" survives round-trips
|
||
// through YAML editors, and a malformed value must fail VALIDATION (with the
|
||
// file and field named) rather than parse as 0 and under-bill silently — which
|
||
// is exactly the failure mode that motivated the DSL.
|
||
func ParseFloatUSDPerM(s string) (float64, error) {
|
||
f, err := strconv.ParseFloat(strings.TrimSpace(s), 64)
|
||
if err != nil {
|
||
return 0, fmt.Errorf("not a number: %q", s)
|
||
}
|
||
if f < 0 {
|
||
return 0, fmt.Errorf("negative price: %q", s)
|
||
}
|
||
return f / 1e6, nil
|
||
}
|
||
|
||
// KeyQuota is retained only to carry a scope entry's caps through the admin
|
||
// API. Quotas are per model, never per key: there is deliberately no key-wide
|
||
// total, so exhausting one model's budget never blocks the others.
|
||
type KeyQuota struct {
|
||
TokenQuota int64 `json:"token_quota"`
|
||
ReqQuota int64 `json:"req_quota"`
|
||
Period string `json:"period"`
|
||
Hours int64 `json:"hours"`
|
||
}
|
||
|
||
// NormalizeRole defaults an empty role to "user", so a key can never end up in
|
||
// a state where no role means "neither admin nor user".
|
||
func NormalizeRole(role string) string {
|
||
if role == "admin" {
|
||
return "admin"
|
||
}
|
||
return "user"
|
||
}
|
||
|
||
// Validate rejects a quota configuration that could not work as written. A
|
||
// period is only meaningful when at least one cap is set, and a cap of zero
|
||
// means "unlimited" rather than "deny everything", so those are the only two
|
||
// things worth rejecting.
|
||
func (q KeyQuota) Validate() error {
|
||
if q.TokenQuota < 0 {
|
||
return fmt.Errorf("token_quota must be >= 0 (0 = unlimited)")
|
||
}
|
||
if q.ReqQuota < 0 {
|
||
return fmt.Errorf("req_quota must be >= 0 (0 = unlimited)")
|
||
}
|
||
if q.Hours < 0 {
|
||
return fmt.Errorf("hours must be >= 0")
|
||
}
|
||
if q.TokenQuota > 0 || q.ReqQuota > 0 {
|
||
if err := ValidatePeriod(q.Period, q.Hours); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// ValidatePeriod accepts the quota period vocabulary: "" (never resets),
|
||
// "hour", "week", "month", or "nhour" with hours >= 1. An unknown period is
|
||
// rejected rather than silently treated as "never resets", which would turn a
|
||
// typo into an all-time quota — the opposite of what the operator typed.
|
||
func ValidatePeriod(period string, hours int64) error {
|
||
switch period {
|
||
case "", "hour", "week", "month":
|
||
return nil
|
||
case "nhour":
|
||
if hours < 1 {
|
||
return fmt.Errorf("period %q needs hours >= 1", period)
|
||
}
|
||
return nil
|
||
}
|
||
return fmt.Errorf("period must be one of \"\", hour, week, month, nhour (got %q)", period)
|
||
}
|
||
|
||
// ModelScope is one allowed model for a key, or one AUTO scheduling slot.
|
||
// Its TokenQuota and ReqQuota cap THAT entry only, independently of every
|
||
// other entry on the same key: a model that runs out of budget stops being
|
||
// served while the key's other models keep working. TokenQuota 0 / ReqQuota 0
|
||
// = unlimited. Period "" = never resets; "hour"/"week"/"month" are fixed
|
||
// windows; "nhour" uses Hours.
|
||
type ModelScope struct {
|
||
Model string `yaml:"model" json:"model"`
|
||
Source string `yaml:"source,omitempty" json:"source,omitempty"` // optional: pin to one upstream source; "" = any source
|
||
Tier int `yaml:"tier,omitempty" json:"tier,omitempty"`
|
||
TokenQuota int64 `yaml:"token_quota" json:"token_quota"`
|
||
ReqQuota int64 `yaml:"req_quota,omitempty" json:"req_quota,omitempty"`
|
||
Period string `yaml:"period,omitempty" json:"period,omitempty"`
|
||
Hours int64 `yaml:"hours,omitempty" json:"hours,omitempty"`
|
||
}
|
||
|
||
// UnmarshalJSON accepts both the legacy "model-id" string form and the
|
||
// {"model":"...","token_quota":N} object form so old runtime files keep
|
||
// loading.
|
||
func (m *ModelScope) UnmarshalJSON(b []byte) error {
|
||
var s string
|
||
if err := json.Unmarshal(b, &s); err == nil {
|
||
m.Model = s
|
||
return nil
|
||
}
|
||
var o struct {
|
||
Model string `json:"model"`
|
||
Source string `json:"source"`
|
||
Tier int `json:"tier"`
|
||
TokenQuota int64 `json:"token_quota"`
|
||
ReqQuota int64 `json:"req_quota"`
|
||
Period string `json:"period"`
|
||
Hours int64 `json:"hours"`
|
||
}
|
||
if err := json.Unmarshal(b, &o); err != nil {
|
||
return err
|
||
}
|
||
m.Model = o.Model
|
||
m.Source = o.Source
|
||
m.Tier = o.Tier
|
||
m.TokenQuota = o.TokenQuota
|
||
m.ReqQuota = o.ReqQuota
|
||
m.Period = o.Period
|
||
m.Hours = o.Hours
|
||
return nil
|
||
}
|