mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-19 16:39:15 +00:00
868ac59692
Upstreams like justwoker/tabitoken sit behind Cloudflare geo/IP blocks and only respond through a proxy. A global env proxy is wrong (intranet sources trae/localzen must stay direct), so add an explicit per-source proxy_url that overrides http.ProxyFromEnvironment for just that source. Sources without proxy_url keep the existing env/direct behavior. Also accepts a User-Agent header per source (config already supported headers) so CF-fronted resellers can be reached. Verified e2e: tabitoken and justwoker now return tool_calls through 127.0.0.1:7890 (clash); trae stays direct. Full go test -tags luajit passes.
405 lines
14 KiB
Go
405 lines
14 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"
|
||
"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"`
|
||
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 密钥页管理）
|
||
}
|
||
|
||
// 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"),
|
||
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.
|
||
func UpsertSourceInYAML(path, name string, src Source) error {
|
||
if path == "" {
|
||
return fmt.Errorf("config path is empty")
|
||
}
|
||
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).
|
||
func (c *Config) Save() error {
|
||
if c.Path == "" {
|
||
return fmt.Errorf("config path is empty")
|
||
}
|
||
out, err := yaml.Marshal(c)
|
||
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"
|
||
}
|
||
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"`
|
||
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
|
||
}
|
||
|
||
// ModelScope is one allowed model for a key, or one AUTO scheduling slot,
|
||
// with an optional token quota and reset period. TokenQuota 0 = unlimited;
|
||
// Period "" = never resets; "hour"/"week"/"month" are fixed windows; "nhour"
|
||
// uses Hours as the window length in 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"`
|
||
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"`
|
||
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.Period = o.Period
|
||
m.Hours = o.Hours
|
||
return nil
|
||
}
|