85 lines
2.8 KiB
Python
85 lines
2.8 KiB
Python
"""Database connection and initialization."""
|
|
import aiosqlite
|
|
from pathlib import Path
|
|
from contextlib import asynccontextmanager
|
|
|
|
DB_PATH = Path(__file__).parent.parent / "data" / "zzrouter.db"
|
|
|
|
|
|
async def get_db() -> aiosqlite.Connection:
|
|
"""Get database connection."""
|
|
DB_PATH.parent.mkdir(parents=True, exist_ok=True)
|
|
conn = await aiosqlite.connect(DB_PATH)
|
|
conn.row_factory = aiosqlite.Row
|
|
return conn
|
|
|
|
|
|
@asynccontextmanager
|
|
async def db_connection():
|
|
"""Context manager for database connection."""
|
|
conn = await get_db()
|
|
try:
|
|
yield conn
|
|
finally:
|
|
await conn.close()
|
|
|
|
|
|
async def init_db():
|
|
"""Initialize database schema."""
|
|
async with db_connection() as db:
|
|
await db.executescript("""
|
|
CREATE TABLE IF NOT EXISTS client_keys (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
key TEXT UNIQUE NOT NULL,
|
|
name TEXT,
|
|
is_active BOOLEAN DEFAULT TRUE,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS providers (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
name TEXT UNIQUE NOT NULL,
|
|
base_url TEXT NOT NULL,
|
|
api_type TEXT NOT NULL,
|
|
is_active BOOLEAN DEFAULT TRUE
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS provider_keys (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
provider_id INTEGER NOT NULL REFERENCES providers(id),
|
|
key TEXT NOT NULL,
|
|
is_active BOOLEAN DEFAULT TRUE,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS request_logs (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
client_key_id INTEGER REFERENCES client_keys(id),
|
|
provider_id INTEGER REFERENCES providers(id),
|
|
model TEXT,
|
|
prompt_tokens INTEGER,
|
|
completion_tokens INTEGER,
|
|
latency_ms INTEGER,
|
|
success BOOLEAN,
|
|
error_message TEXT,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_request_logs_created_at ON request_logs(created_at);
|
|
CREATE INDEX IF NOT EXISTS idx_request_logs_client_key_id ON request_logs(client_key_id);
|
|
""")
|
|
await db.commit()
|
|
|
|
|
|
async def seed_default_providers():
|
|
"""Seed default providers if not exist."""
|
|
async with db_connection() as db:
|
|
await db.executemany(
|
|
"""INSERT OR IGNORE INTO providers (name, base_url, api_type) VALUES (?, ?, ?)""",
|
|
[
|
|
("openai", "https://api.openai.com/v1", "openai"),
|
|
("anthropic", "https://api.anthropic.com/v1", "anthropic"),
|
|
]
|
|
)
|
|
await db.commit()
|