turf_saas/ai_router/models.py

import logging
import os
import sqlite3
from datetime import datetime, timezone

logger = logging.getLogger("ai_router.models")

DB_PATH = os.environ.get("AI_ROUTER_DB", "/home/h3r7/turf_saas/ai_router.db")


def get_db():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn


def init_db():
    conn = get_db()
    try:
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS ai_providers (
                id              INTEGER PRIMARY KEY AUTOINCREMENT,
                name            TEXT NOT NULL,
                provider_type   TEXT NOT NULL CHECK(provider_type IN ('openai','anthropic','google','mistral')),
                api_key         TEXT NOT NULL DEFAULT '',
                base_url        TEXT DEFAULT '',
                config          TEXT DEFAULT '{}',
                priority        INTEGER NOT NULL DEFAULT 99,
                is_active       INTEGER NOT NULL DEFAULT 1,
                created_at      TEXT NOT NULL DEFAULT (datetime('now')),
                updated_at      TEXT NOT NULL DEFAULT (datetime('now'))
            );

            CREATE TABLE IF NOT EXISTS ai_model_mapping (
                id              INTEGER PRIMARY KEY AUTOINCREMENT,
                model_alias     TEXT NOT NULL UNIQUE,
                provider_id     INTEGER NOT NULL REFERENCES ai_providers(id),
                real_model_id   TEXT NOT NULL,
                cost_per_1k_tokens REAL NOT NULL DEFAULT 0,
                is_active       INTEGER NOT NULL DEFAULT 1,
                created_at      TEXT NOT NULL DEFAULT (datetime('now'))
            );

            CREATE TABLE IF NOT EXISTS ai_router_log (
                id              INTEGER PRIMARY KEY AUTOINCREMENT,
                request_id      TEXT NOT NULL,
                user_id         INTEGER,
                model_alias     TEXT NOT NULL,
                provider_used   TEXT NOT NULL,
                tokens_in       INTEGER NOT NULL DEFAULT 0,
                tokens_out      INTEGER NOT NULL DEFAULT 0,
                duration_ms     INTEGER NOT NULL DEFAULT 0,
                status          TEXT NOT NULL CHECK(status IN ('success','error')),
                error_message   TEXT DEFAULT '',
                created_at      TEXT NOT NULL DEFAULT (datetime('now'))
            );

            CREATE INDEX IF NOT EXISTS idx_ai_router_log_request_id ON ai_router_log(request_id);
            CREATE INDEX IF NOT EXISTS idx_ai_router_log_created_at ON ai_router_log(created_at);
            CREATE INDEX IF NOT EXISTS idx_ai_model_mapping_alias ON ai_model_mapping(model_alias);
        """)
        conn.commit()
        logger.info("AI Router database tables initialized")
    except Exception as e:
        logger.error(f"Failed to initialize AI Router DB: {e}")
    finally:
        conn.close()


def get_providers_from_db():
    conn = get_db()
    try:
        rows = conn.execute("""
            SELECT p.id, p.name, p.provider_type, p.api_key, p.base_url, p.config,
                   p.priority, p.is_active, m.model_alias, m.real_model_id, m.cost_per_1k_tokens
            FROM ai_providers p
            LEFT JOIN ai_model_mapping m ON m.provider_id = p.id
            WHERE p.is_active = 1
        """).fetchall()
        return [dict(r) for r in rows]
    except Exception as e:
        logger.warning(f"Could not query providers: {e}")
        return []
    finally:
        conn.close()


def log_router_attempt(request_id, user_id, model_alias, provider_used,
                       tokens_in, tokens_out, duration_ms, status,
                       error_message=""):
    conn = get_db()
    try:
        conn.execute(
            """INSERT INTO ai_router_log
               (request_id, user_id, model_alias, provider_used,
                tokens_in, tokens_out, duration_ms, status, error_message)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (request_id, user_id, model_alias, provider_used,
             tokens_in, tokens_out, duration_ms, status, error_message),
        )
        conn.commit()
    except Exception as e:
        logger.warning(f"Failed to log router attempt: {e}")
    finally:
        conn.close()


def upsert_provider(name, provider_type, api_key="", base_url="",
                    config=None, priority=99, is_active=1):
    conn = get_db()
    try:
        existing = conn.execute(
            "SELECT id FROM ai_providers WHERE name = ?", (name,)
        ).fetchone()
        if existing:
            conn.execute(
                """UPDATE ai_providers SET provider_type=?, api_key=?, base_url=?,
                   config=?, priority=?, is_active=?, updated_at=datetime('now')
                   WHERE name=?""",
                (provider_type, api_key, base_url,
                 config or "{}", priority, is_active, name),
            )
        else:
            conn.execute(
                """INSERT INTO ai_providers
                   (name, provider_type, api_key, base_url, config, priority, is_active)
                   VALUES (?, ?, ?, ?, ?, ?, ?)""",
                (name, provider_type, api_key, base_url,
                 config or "{}", priority, is_active),
            )
        conn.commit()
        return True
    except Exception as e:
        logger.error(f"Failed to upsert provider: {e}")
        return False
    finally:
        conn.close()


def upsert_model_mapping(model_alias, provider_id, real_model_id, cost_per_1k=0):
    conn = get_db()
    try:
        existing = conn.execute(
            "SELECT id FROM ai_model_mapping WHERE model_alias = ?", (model_alias,)
        ).fetchone()
        if existing:
            conn.execute(
                """UPDATE ai_model_mapping SET provider_id=?, real_model_id=?,
                   cost_per_1k_tokens=? WHERE model_alias=?""",
                (provider_id, real_model_id, cost_per_1k, model_alias),
            )
        else:
            conn.execute(
                """INSERT INTO ai_model_mapping
                   (model_alias, provider_id, real_model_id, cost_per_1k_tokens)
                   VALUES (?, ?, ?, ?)""",
                (model_alias, provider_id, real_model_id, cost_per_1k),
            )
        conn.commit()
        return True
    except Exception as e:
        logger.error(f"Failed to upsert model mapping: {e}")
        return False
    finally:
        conn.close()