299 lines
11 KiB
Go
299 lines
11 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)
|
|
}
|
|
|
|
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,
|
|
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)
|
|
}
|
|
}
|
|
|
|
// 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 string
|
|
}{
|
|
{"deepseek-auto", "heuristic", `["deepseek-chat","deepseek-reasoner"]`},
|
|
{"openai-auto", "heuristic", `["gpt-4o-mini","gpt-4o"]`},
|
|
{"gemini-auto", "heuristic", `["gemini-2.0-flash","gemini-2.5-pro"]`},
|
|
}
|
|
for _, g := range defaultGroups {
|
|
db.Exec(`INSERT OR IGNORE INTO model_groups (id, strategy, targets) VALUES (?, ?, ?)`,
|
|
g.id, g.strategy, g.targets)
|
|
}
|
|
|
|
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"`
|
|
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
|
UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
|
|
}
|