hobokenchicken
73a82e6175
Upgrades the routing engine to support tag, token limit, multimodal, reasoning, and tool calling conditions. Adds unit tests for the new routing features.
332 lines
13 KiB
Go
332 lines
13 KiB
Go
package db
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"github.com/jmoiron/sqlx"
|
|
_ "modernc.org/sqlite"
|
|
"golang.org/x/crypto/bcrypt"
|
|
)
|
|
|
|
type DB struct {
|
|
*sqlx.DB
|
|
}
|
|
|
|
func Init(path string) (*DB, error) {
|
|
// Ensure directory exists
|
|
dir := filepath.Dir(path)
|
|
if dir != "." && dir != "/" {
|
|
if err := os.MkdirAll(dir, 0755); err != nil {
|
|
return nil, fmt.Errorf("failed to create database directory: %w", err)
|
|
}
|
|
}
|
|
|
|
// Connect to SQLite
|
|
dsn := fmt.Sprintf("file:%s?_pragma=foreign_keys(1)", path)
|
|
db, err := sqlx.Connect("sqlite", dsn)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to connect to database: %w", err)
|
|
}
|
|
|
|
// Enable Write-Ahead Logging (WAL) and set a busy timeout to handle concurrent access
|
|
if _, err := db.Exec("PRAGMA journal_mode=WAL;"); err != nil {
|
|
log.Printf("failed to enable WAL mode: %v", err)
|
|
}
|
|
if _, err := db.Exec("PRAGMA busy_timeout=5000;"); err != nil {
|
|
log.Printf("failed to set busy timeout: %v", err)
|
|
}
|
|
|
|
instance := &DB{db}
|
|
|
|
// Run migrations
|
|
if err := instance.RunMigrations(); err != nil {
|
|
return nil, fmt.Errorf("failed to run migrations: %w", err)
|
|
}
|
|
|
|
return instance, nil
|
|
}
|
|
|
|
func (db *DB) RunMigrations() error {
|
|
// Tables creation
|
|
queries := []string{
|
|
`CREATE TABLE IF NOT EXISTS clients (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
client_id TEXT UNIQUE NOT NULL,
|
|
name TEXT,
|
|
description TEXT,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
is_active BOOLEAN DEFAULT TRUE,
|
|
rate_limit_per_minute INTEGER DEFAULT 60,
|
|
total_requests INTEGER DEFAULT 0,
|
|
total_tokens INTEGER DEFAULT 0,
|
|
total_cost REAL DEFAULT 0.0
|
|
)`,
|
|
`CREATE TABLE IF NOT EXISTS llm_requests (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
client_id TEXT,
|
|
provider TEXT,
|
|
model TEXT,
|
|
prompt_tokens INTEGER,
|
|
completion_tokens INTEGER,
|
|
reasoning_tokens INTEGER DEFAULT 0,
|
|
total_tokens INTEGER,
|
|
cost REAL,
|
|
has_images BOOLEAN DEFAULT FALSE,
|
|
status TEXT DEFAULT 'success',
|
|
error_message TEXT,
|
|
duration_ms INTEGER,
|
|
request_body TEXT,
|
|
response_body TEXT,
|
|
cache_read_tokens INTEGER DEFAULT 0,
|
|
cache_write_tokens INTEGER DEFAULT 0,
|
|
FOREIGN KEY (client_id) REFERENCES clients(client_id) ON DELETE SET NULL
|
|
)`,
|
|
`CREATE TABLE IF NOT EXISTS provider_configs (
|
|
id TEXT PRIMARY KEY,
|
|
display_name TEXT NOT NULL,
|
|
enabled BOOLEAN DEFAULT TRUE,
|
|
base_url TEXT,
|
|
api_key TEXT,
|
|
credit_balance REAL DEFAULT 0.0,
|
|
low_credit_threshold REAL DEFAULT 5.0,
|
|
billing_mode TEXT,
|
|
api_key_encrypted BOOLEAN DEFAULT FALSE,
|
|
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
|
)`,
|
|
`CREATE TABLE IF NOT EXISTS model_configs (
|
|
id TEXT PRIMARY KEY,
|
|
provider_id TEXT NOT NULL,
|
|
display_name TEXT,
|
|
enabled BOOLEAN DEFAULT TRUE,
|
|
prompt_cost_per_m REAL,
|
|
completion_cost_per_m REAL,
|
|
cache_read_cost_per_m REAL,
|
|
cache_write_cost_per_m REAL,
|
|
mapping TEXT,
|
|
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (provider_id) REFERENCES provider_configs(id) ON DELETE CASCADE
|
|
)`,
|
|
`CREATE TABLE IF NOT EXISTS users (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
username TEXT UNIQUE NOT NULL,
|
|
password_hash TEXT NOT NULL,
|
|
display_name TEXT,
|
|
role TEXT DEFAULT 'admin',
|
|
must_change_password BOOLEAN DEFAULT FALSE,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
|
)`,
|
|
`CREATE TABLE IF NOT EXISTS client_tokens (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
client_id TEXT NOT NULL,
|
|
token TEXT NOT NULL UNIQUE,
|
|
name TEXT DEFAULT 'default',
|
|
is_active BOOLEAN DEFAULT TRUE,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
last_used_at DATETIME,
|
|
FOREIGN KEY (client_id) REFERENCES clients(client_id) ON DELETE CASCADE
|
|
)`,
|
|
`CREATE TABLE IF NOT EXISTS model_groups (
|
|
id TEXT PRIMARY KEY,
|
|
strategy TEXT NOT NULL DEFAULT 'heuristic',
|
|
selector_model TEXT,
|
|
targets TEXT NOT NULL DEFAULT '[]',
|
|
complexity_threshold INTEGER,
|
|
heuristic_rules TEXT,
|
|
logic_level INTEGER,
|
|
primary_use TEXT,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
|
)`,
|
|
}
|
|
|
|
for _, q := range queries {
|
|
if _, err := db.Exec(q); err != nil {
|
|
return fmt.Errorf("migration failed for query [%s]: %w", q, err)
|
|
}
|
|
}
|
|
|
|
// Add indices
|
|
indices := []string{
|
|
"CREATE INDEX IF NOT EXISTS idx_clients_client_id ON clients(client_id)",
|
|
"CREATE INDEX IF NOT EXISTS idx_clients_created_at ON clients(created_at)",
|
|
"CREATE INDEX IF NOT EXISTS idx_llm_requests_timestamp ON llm_requests(timestamp)",
|
|
"CREATE INDEX IF NOT EXISTS idx_llm_requests_client_id ON llm_requests(client_id)",
|
|
"CREATE INDEX IF NOT EXISTS idx_llm_requests_provider ON llm_requests(provider)",
|
|
"CREATE INDEX IF NOT EXISTS idx_llm_requests_status ON llm_requests(status)",
|
|
"CREATE UNIQUE INDEX IF NOT EXISTS idx_client_tokens_token ON client_tokens(token)",
|
|
"CREATE INDEX IF NOT EXISTS idx_client_tokens_client_id ON client_tokens(client_id)",
|
|
"CREATE INDEX IF NOT EXISTS idx_llm_requests_client_timestamp ON llm_requests(client_id, timestamp)",
|
|
"CREATE INDEX IF NOT EXISTS idx_llm_requests_provider_timestamp ON llm_requests(provider, timestamp)",
|
|
"CREATE INDEX IF NOT EXISTS idx_model_configs_provider_id ON model_configs(provider_id)",
|
|
}
|
|
|
|
for _, idx := range indices {
|
|
if _, err := db.Exec(idx); err != nil {
|
|
return fmt.Errorf("failed to create index [%s]: %w", idx, err)
|
|
}
|
|
}
|
|
|
|
// Add columns to existing model_groups tables (safe — SQLite ignores duplicates on error)
|
|
db.Exec("ALTER TABLE model_groups ADD COLUMN logic_level INTEGER")
|
|
db.Exec("ALTER TABLE model_groups ADD COLUMN primary_use TEXT")
|
|
|
|
// Default admin user
|
|
var count int
|
|
if err := db.Get(&count, "SELECT COUNT(*) FROM users"); err != nil {
|
|
return fmt.Errorf("failed to count users: %w", err)
|
|
}
|
|
|
|
if count == 0 {
|
|
hash, err := bcrypt.GenerateFromPassword([]byte("admin123"), 12)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to hash default password: %w", err)
|
|
}
|
|
_, err = db.Exec("INSERT INTO users (username, password_hash, role, must_change_password) VALUES ('admin', ?, 'admin', 1)", string(hash))
|
|
if err != nil {
|
|
return fmt.Errorf("failed to insert default admin: %w", err)
|
|
}
|
|
log.Println("Created default admin user with password 'admin123' (must change on first login)")
|
|
}
|
|
|
|
// Default client
|
|
_, err := db.Exec(`INSERT OR IGNORE INTO clients (client_id, name, description)
|
|
VALUES ('default', 'Default Client', 'Default client for anonymous requests')`)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to insert default client: %w", err)
|
|
}
|
|
|
|
// Seed default model groups
|
|
defaultGroups := []struct {
|
|
id, strategy, targets, selectorModel string
|
|
complexityThreshold, logicLevel *int
|
|
primaryUse *string
|
|
}{
|
|
{"deepseek-auto", "heuristic", `["deepseek-chat","deepseek-reasoner"]`, "", nil, nil, nil},
|
|
{"openai-auto", "heuristic", `["gpt-4o-mini","gpt-4o"]`, "", nil, nil, nil},
|
|
{"gemini-auto", "heuristic", `["gemini-2.0-flash","gemini-2.5-pro"]`, "", nil, nil, nil},
|
|
{"heavy-logic", "heuristic", `["grok-4.3","kimi-k2.6","deepseek-v4-pro"]`, "", nil, intPtr(9), strPtr("Complex Coding, Logic, Agents.")},
|
|
{"standard-pro", "heuristic", `["gpt-5.4-mini","gemini-3-flash-preview"]`, "", nil, intPtr(5), strPtr("General Assistant, Long Docs.")},
|
|
{"fast-flow", "heuristic", `["deepseek-v4-flash","gpt-5.4-nano"]`, "", nil, intPtr(2), strPtr("Classification, JSON, Basic Q&A.")},
|
|
{"dispatcher", "classifier", `["fast-flow","standard-pro","heavy-logic"]`, "gpt-5.4-nano", intPtr(10), nil, strPtr("Auto-dispatches to tier groups by complexity.")},
|
|
}
|
|
for _, g := range defaultGroups {
|
|
db.Exec(`INSERT OR IGNORE INTO model_groups (id, strategy, targets, selector_model, complexity_threshold, logic_level, primary_use) VALUES (?, ?, ?, ?, ?, ?, ?)`,
|
|
g.id, g.strategy, g.targets, nilStr(g.selectorModel), g.complexityThreshold, g.logicLevel, g.primaryUse)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// Data models for DB tables
|
|
|
|
type Client struct {
|
|
ID int `db:"id"`
|
|
ClientID string `db:"client_id"`
|
|
Name *string `db:"name"`
|
|
Description *string `db:"description"`
|
|
CreatedAt time.Time `db:"created_at"`
|
|
UpdatedAt time.Time `db:"updated_at"`
|
|
IsActive bool `db:"is_active"`
|
|
RateLimitPerMinute int `db:"rate_limit_per_minute"`
|
|
TotalRequests int `db:"total_requests"`
|
|
TotalTokens int `db:"total_tokens"`
|
|
TotalCost float64 `db:"total_cost"`
|
|
}
|
|
|
|
type LLMRequest struct {
|
|
ID int `db:"id"`
|
|
Timestamp time.Time `db:"timestamp"`
|
|
ClientID *string `db:"client_id"`
|
|
Provider *string `db:"provider"`
|
|
Model *string `db:"model"`
|
|
PromptTokens *int `db:"prompt_tokens"`
|
|
CompletionTokens *int `db:"completion_tokens"`
|
|
ReasoningTokens int `db:"reasoning_tokens"`
|
|
TotalTokens *int `db:"total_tokens"`
|
|
Cost *float64 `db:"cost"`
|
|
HasImages bool `db:"has_images"`
|
|
Status string `db:"status"`
|
|
ErrorMessage *string `db:"error_message"`
|
|
DurationMS *int `db:"duration_ms"`
|
|
RequestBody *string `db:"request_body"`
|
|
ResponseBody *string `db:"response_body"`
|
|
CacheReadTokens int `db:"cache_read_tokens"`
|
|
CacheWriteTokens int `db:"cache_write_tokens"`
|
|
}
|
|
|
|
type ProviderConfig struct {
|
|
ID string `db:"id"`
|
|
DisplayName string `db:"display_name"`
|
|
Enabled bool `db:"enabled"`
|
|
BaseURL *string `db:"base_url"`
|
|
APIKey *string `db:"api_key"`
|
|
CreditBalance float64 `db:"credit_balance"`
|
|
LowCreditThreshold float64 `db:"low_credit_threshold"`
|
|
BillingMode *string `db:"billing_mode"`
|
|
APIKeyEncrypted bool `db:"api_key_encrypted"`
|
|
UpdatedAt time.Time `db:"updated_at"`
|
|
}
|
|
|
|
type ModelConfig struct {
|
|
ID string `db:"id"`
|
|
ProviderID string `db:"provider_id"`
|
|
DisplayName *string `db:"display_name"`
|
|
Enabled bool `db:"enabled"`
|
|
PromptCostPerM *float64 `db:"prompt_cost_per_m"`
|
|
CompletionCostPerM *float64 `db:"completion_cost_per_m"`
|
|
CacheReadCostPerM *float64 `db:"cache_read_cost_per_m"`
|
|
CacheWriteCostPerM *float64 `db:"cache_write_cost_per_m"`
|
|
Mapping *string `db:"mapping"`
|
|
UpdatedAt time.Time `db:"updated_at"`
|
|
}
|
|
|
|
type User struct {
|
|
ID int `db:"id" json:"id"`
|
|
Username string `db:"username" json:"username"`
|
|
PasswordHash string `db:"password_hash" json:"-"`
|
|
DisplayName *string `db:"display_name" json:"display_name"`
|
|
Role string `db:"role" json:"role"`
|
|
MustChangePassword bool `db:"must_change_password" json:"must_change_password"`
|
|
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
|
}
|
|
|
|
type ClientToken struct {
|
|
ID int `db:"id"`
|
|
ClientID string `db:"client_id"`
|
|
Token string `db:"token"`
|
|
Name string `db:"name"`
|
|
IsActive bool `db:"is_active"`
|
|
CreatedAt time.Time `db:"created_at"`
|
|
LastUsedAt *time.Time `db:"last_used_at"`
|
|
}
|
|
|
|
type ModelGroup struct {
|
|
ID string `db:"id" json:"id"`
|
|
Strategy string `db:"strategy" json:"strategy"`
|
|
SelectorModel *string `db:"selector_model" json:"selector_model"`
|
|
Targets string `db:"targets" json:"targets"` // JSON array
|
|
ComplexityThreshold *int `db:"complexity_threshold" json:"complexity_threshold"`
|
|
HeuristicRules *string `db:"heuristic_rules" json:"heuristic_rules"`
|
|
LogicLevel *int `db:"logic_level" json:"logic_level"`
|
|
PrimaryUse *string `db:"primary_use" json:"primary_use"`
|
|
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
|
UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
|
|
}
|
|
|
|
func intPtr(v int) *int { return &v }
|
|
func strPtr(v string) *string { return &v }
|
|
|
|
// nilStr returns a *string for non-empty strings, nil for empty.
|
|
func nilStr(v string) *string {
|
|
if v == "" {
|
|
return nil
|
|
}
|
|
return &v
|
|
}
|