ModelRouter/internal/config/config.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"
	"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"`
	MaxConcurrent int          `yaml:"max_concurrent"`            // global inflight cap, 0 = unlimited
	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 优先级页编辑）
	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"`
	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:"-"`
}

// Load reads and validates a config file.
func Load(path string) (*Config, error) {
	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
}

// 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"`
	DeletedSources  []string     `json:"deleted_sources,omitempty"`
	DeletedAdapters []string     `json:"deleted_adapters,omitempty"`
	Keys            []GWKey      `json:"keys,omitempty"`
	Auto            []ModelScope `json:"auto,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
}