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) } 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 }