feat: implement AI model API proxy service
This commit is contained in:
84
backend/app/database.py
Normal file
84
backend/app/database.py
Normal file
@@ -0,0 +1,84 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user