mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-19 16:39:15 +00:00
feat: ModelRouter — unified OpenAI-compatible multi-source LLM gateway
- Lua adapters per upstream (transform_request/response/stream_chunk, build_headers signing hooks) - AUTO priority routing with per-model kind (chat/image), explicit source/model routing - Per-source concurrency caps with queueing, exponential backoff, AUTO failover - OpenAI-compatible API: chat completions, SSE streaming, image generations, models - Gateway key auth, web UI for adapter/source management, runtime persistence - e2e test running the real binary against mocked upstreams
This commit is contained in:
root
124
internal/config/config.go
Normal file
124
internal/config/config.go
Normal file
@ -0,0 +1,124 @@
|
||||
// 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 model selection (higher number = preferred).
|
||||
type Model struct {
|
||||
ID string `yaml:"id"`
|
||||
Priority int `yaml:"priority"`
|
||||
Kind string `yaml:"kind"` // "chat" (default) | "image"
|
||||
Meta map[string]interface{} `yaml:"meta"`
|
||||
}
|
||||
|
||||
// Source describes a single upstream LLM provider.
|
||||
type Source struct {
|
||||
Name string `yaml:"name"`
|
||||
BaseURL string `yaml:"base_url"`
|
||||
APIKey string `yaml:"api_key"`
|
||||
Adapter string `yaml:"adapter"`
|
||||
Endpoint string `yaml:"endpoint"` // chat endpoint override
|
||||
ImageEndpoint string `yaml:"image_endpoint"` // image endpoint override
|
||||
Models []Model `yaml:"models"`
|
||||
Headers map[string]string `yaml:"headers"`
|
||||
Meta map[string]interface{} `yaml:"meta"`
|
||||
Temperature float64 `yaml:"temperature"`
|
||||
MaxTokens int `yaml:"max_tokens"`
|
||||
Timeout time.Duration `yaml:"timeout"`
|
||||
MaxConcurrent int `yaml:"max_concurrent"` // per-source inflight cap
|
||||
QueueTimeout time.Duration `yaml:"queue_timeout"` // wait for slot before failing
|
||||
}
|
||||
|
||||
// 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"`
|
||||
}
|
||||
122
internal/config/config_test.go
Normal file
122
internal/config/config_test.go
Normal file
@ -0,0 +1,122 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLoadAndApplyDefaults(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "cfg.yaml")
|
||||
content := `
|
||||
listen: 127.0.0.1:9999
|
||||
gateway_keys: [sk-1]
|
||||
default_model: AUTO
|
||||
adapter_dir: adapters
|
||||
runtime_file: runtime.json
|
||||
sources:
|
||||
- name: deepseek
|
||||
base_url: https://api.deepseek.com
|
||||
api_key: sk-d
|
||||
adapter: deepseek
|
||||
models:
|
||||
- id: deepseek-v4-flash
|
||||
priority: 100
|
||||
- name: ollama
|
||||
base_url: http://127.0.0.1:11434
|
||||
adapter: ollama
|
||||
endpoint: /api/chat
|
||||
models:
|
||||
- id: llama3
|
||||
priority: 50
|
||||
`
|
||||
if err := os.WriteFile(path, []byte(content), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := Load(path)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
if len(cfg.Sources) != 2 {
|
||||
t.Fatalf("sources = %d", len(cfg.Sources))
|
||||
}
|
||||
if cfg.Sources[1].Endpoint != "/api/chat" {
|
||||
t.Fatalf("endpoint = %q", cfg.Sources[1].Endpoint)
|
||||
}
|
||||
if cfg.Sources[1].Timeout == 0 {
|
||||
t.Fatal("default timeout not applied")
|
||||
}
|
||||
if cfg.Sources[1].MaxConcurrent == 0 {
|
||||
t.Fatal("default max_concurrent not applied")
|
||||
}
|
||||
if cfg.Sources[0].Models[0].Priority != 100 {
|
||||
t.Fatalf("priority = %d", cfg.Sources[0].Models[0].Priority)
|
||||
}
|
||||
if cfg.DefaultModel != "AUTO" {
|
||||
t.Fatalf("default model = %q", cfg.DefaultModel)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyDefaultsDuplicateSource(t *testing.T) {
|
||||
cfg := Config{Sources: []Source{
|
||||
{Name: "a", BaseURL: "http://x", Models: []Model{{ID: "m1"}}},
|
||||
{Name: "a", BaseURL: "http://y", Models: []Model{{ID: "m2"}}},
|
||||
}}
|
||||
if err := cfg.ApplyDefaults(); err == nil {
|
||||
t.Fatal("expected duplicate source error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyDefaultsNoSources(t *testing.T) {
|
||||
// empty source list is valid (sources may be added later via the Web UI)
|
||||
cfg := Config{}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if cfg.Listen != ":8080" || cfg.DefaultModel != "AUTO" {
|
||||
t.Fatalf("defaults not applied: %+v", cfg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyDefaultsDuplicateModel(t *testing.T) {
|
||||
cfg := Config{Sources: []Source{
|
||||
{Name: "a", BaseURL: "http://x", Models: []Model{{ID: "m1"}}},
|
||||
{Name: "b", BaseURL: "http://y", Models: []Model{{ID: "m1"}}},
|
||||
}}
|
||||
if err := cfg.ApplyDefaults(); err == nil {
|
||||
t.Fatal("expected duplicate model error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStoreUpsertRemove(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "runtime.json")
|
||||
s := NewStore(path)
|
||||
if err := s.Load(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.Upsert(Source{Name: "a", BaseURL: "http://a", Models: []Model{{ID: "m"}}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.Upsert(Source{Name: "b", BaseURL: "http://b", Models: []Model{{ID: "m2"}}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(s.List()) != 2 {
|
||||
t.Fatalf("list = %d", len(s.List()))
|
||||
}
|
||||
removed, err := s.Remove("a")
|
||||
if err != nil || !removed {
|
||||
t.Fatalf("remove: %v %v", removed, err)
|
||||
}
|
||||
if len(s.List()) != 1 {
|
||||
t.Fatalf("after remove list = %d", len(s.List()))
|
||||
}
|
||||
// reload from disk
|
||||
s2 := NewStore(path)
|
||||
if err := s2.Load(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(s2.List()) != 1 {
|
||||
t.Fatalf("reloaded list = %d", len(s2.List()))
|
||||
}
|
||||
}
|
||||
85
internal/config/store.go
Normal file
85
internal/config/store.go
Normal file
@ -0,0 +1,85 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Store persists web-UI editable runtime state (sources added/edited) to a
|
||||
// JSON file so edits survive restarts. Base YAML sources are merged underneath.
|
||||
type Store struct {
|
||||
mu sync.Mutex
|
||||
path string
|
||||
data RuntimeConfig
|
||||
}
|
||||
|
||||
func NewStore(path string) *Store {
|
||||
return &Store{path: path}
|
||||
}
|
||||
|
||||
// Load reads the runtime file (missing file = empty state).
|
||||
func (s *Store) Load() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
data, err := os.ReadFile(s.path)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
s.data = RuntimeConfig{}
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
return json.Unmarshal(data, &s.data)
|
||||
}
|
||||
|
||||
// List returns the runtime sources (those edited via web UI).
|
||||
func (s *Store) List() []Source {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
out := make([]Source, len(s.data.Sources))
|
||||
copy(out, s.data.Sources)
|
||||
return out
|
||||
}
|
||||
|
||||
// Upsert adds or replaces a runtime source and persists.
|
||||
func (s *Store) Upsert(src Source) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for i := range s.data.Sources {
|
||||
if s.data.Sources[i].Name == src.Name {
|
||||
s.data.Sources[i] = src
|
||||
return s.persistLocked()
|
||||
}
|
||||
}
|
||||
s.data.Sources = append(s.data.Sources, src)
|
||||
return s.persistLocked()
|
||||
}
|
||||
|
||||
// Remove deletes a runtime source and persists.
|
||||
func (s *Store) Remove(name string) (bool, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
kept := s.data.Sources[:0]
|
||||
removed := false
|
||||
for _, src := range s.data.Sources {
|
||||
if src.Name == name {
|
||||
removed = true
|
||||
continue
|
||||
}
|
||||
kept = append(kept, src)
|
||||
}
|
||||
if !removed {
|
||||
return false, nil
|
||||
}
|
||||
s.data.Sources = kept
|
||||
return true, s.persistLocked()
|
||||
}
|
||||
|
||||
func (s *Store) persistLocked() error {
|
||||
b, err := json.MarshalIndent(s.data, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return os.WriteFile(s.path, b, 0644)
|
||||
}
|
||||
Reference in New Issue
Block a user