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 (
	"fmt"
	"os"
	"time"

	"gopkg.in/yaml.v3"
)

// Config is the top-level gateway configuration.
type Config struct {
	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
	Sources       []Source `yaml:"sources"`
}

// 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"`
	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)
	}
	if err := cfg.ApplyDefaults(); err != nil {
		return nil, err
	}
	return &cfg, nil
}

// 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{}
	modelOwners := map[string]string{}
	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 = 120 * time.Second
		}
		if s.QueueTimeout == 0 {
			s.QueueTimeout = 60 * time.Second
		}
		if s.MaxConcurrent == 0 {
			s.MaxConcurrent = 8
		}
		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)
			}
			if owner, ok := modelOwners[m.ID]; ok {
				return fmt.Errorf("config: model %q defined by both %s and %s", m.ID, owner, s.Name)
			}
			modelOwners[m.ID] = s.Name
		}
	}
	return nil
}

// RuntimeConfig is the persisted web-UI editable slice (sources added/edited).
type RuntimeConfig struct {
	Sources []Source `json:"sources"`
}