llm-in-text/backend/audit_store.py

import json
import os
import threading
from typing import Any

try:
    import psycopg
except Exception:  # pragma: no cover
    psycopg = None


class BaseAuditStore:
    def record_api_request(self, payload: dict[str, Any]) -> None:
        return None

    def record_llm_call(self, payload: dict[str, Any]) -> None:
        return None

    def upsert_daily_usage(self, payload: dict[str, Any]) -> None:
        return None


class NullAuditStore(BaseAuditStore):
    pass


class PostgresAuditStore(BaseAuditStore):
    def __init__(self, database_url: str) -> None:
        if psycopg is None:
            raise RuntimeError("psycopg 未安装，无法使用 PostgreSQL 审计存储")
        self.database_url = database_url
        self._init_lock = threading.Lock()
        self._initialized = False

    def _connect(self):
        return psycopg.connect(self.database_url, autocommit=True)

    def _ensure_initialized(self) -> None:
        if self._initialized:
            return
        with self._init_lock:
            if self._initialized:
                return
            with self._connect() as conn:
                with conn.cursor() as cur:
                    cur.execute(
                        """
                        CREATE TABLE IF NOT EXISTS api_request_audit (
                            id BIGSERIAL PRIMARY KEY,
                            request_id TEXT NOT NULL,
                            session_hash TEXT NOT NULL,
                            ip_hash TEXT NOT NULL,
                            route TEXT NOT NULL,
                            method TEXT NOT NULL,
                            status_code INTEGER NOT NULL,
                            decision TEXT NOT NULL,
                            delay_ms INTEGER NOT NULL DEFAULT 0,
                            queue_ms INTEGER NOT NULL DEFAULT 0,
                            error_code TEXT NOT NULL DEFAULT '',
                            created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
                            metadata_json JSONB NOT NULL DEFAULT '{}'::jsonb
                        )
                        """
                    )
                    cur.execute(
                        """
                        CREATE TABLE IF NOT EXISTS llm_call_audit (
                            id BIGSERIAL PRIMARY KEY,
                            request_id TEXT NOT NULL,
                            session_hash TEXT NOT NULL,
                            ip_hash TEXT NOT NULL,
                            job_type TEXT NOT NULL,
                            model TEXT NOT NULL,
                            estimated_input_tokens INTEGER NOT NULL DEFAULT 0,
                            max_output_tokens INTEGER NOT NULL DEFAULT 0,
                            estimated_cost NUMERIC(18, 8) NOT NULL DEFAULT 0,
                            actual_output_chars INTEGER NOT NULL DEFAULT 0,
                            actual_cost NUMERIC(18, 8) NOT NULL DEFAULT 0,
                            status TEXT NOT NULL,
                            error_code TEXT NOT NULL DEFAULT '',
                            started_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
                            finished_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
                            metadata_json JSONB NOT NULL DEFAULT '{}'::jsonb
                        )
                        """
                    )
                    cur.execute(
                        """
                        CREATE TABLE IF NOT EXISTS risk_events (
                            id BIGSERIAL PRIMARY KEY,
                            session_hash TEXT NOT NULL,
                            ip_hash TEXT NOT NULL,
                            event_type TEXT NOT NULL,
                            severity TEXT NOT NULL,
                            reason TEXT NOT NULL,
                            created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
                            metadata_json JSONB NOT NULL DEFAULT '{}'::jsonb
                        )
                        """
                    )
                    cur.execute(
                        """
                        CREATE TABLE IF NOT EXISTS daily_budget_usage (
                            usage_day DATE NOT NULL,
                            scope TEXT NOT NULL,
                            scope_hash TEXT NOT NULL,
                            estimated_cost NUMERIC(18, 8) NOT NULL DEFAULT 0,
                            actual_cost NUMERIC(18, 8) NOT NULL DEFAULT 0,
                            request_count INTEGER NOT NULL DEFAULT 0,
                            updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
                            PRIMARY KEY (usage_day, scope, scope_hash)
                        )
                        """
                    )
            self._initialized = True

    def record_api_request(self, payload: dict[str, Any]) -> None:
        self._ensure_initialized()
        metadata = payload.get("metadata") or {}
        with self._connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO api_request_audit (
                        request_id, session_hash, ip_hash, route, method, status_code,
                        decision, delay_ms, queue_ms, error_code, metadata_json
                    )
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s::jsonb)
                    """,
                    (
                        payload["request_id"],
                        payload["session_hash"],
                        payload["ip_hash"],
                        payload["route"],
                        payload["method"],
                        int(payload["status_code"]),
                        payload["decision"],
                        int(payload.get("delay_ms", 0)),
                        int(payload.get("queue_ms", 0)),
                        payload.get("error_code", ""),
                        json.dumps(metadata, ensure_ascii=False),
                    ),
                )

    def record_llm_call(self, payload: dict[str, Any]) -> None:
        self._ensure_initialized()
        metadata = payload.get("metadata") or {}
        with self._connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO llm_call_audit (
                        request_id, session_hash, ip_hash, job_type, model,
                        estimated_input_tokens, max_output_tokens, estimated_cost,
                        actual_output_chars, actual_cost, status, error_code, metadata_json
                    )
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s::jsonb)
                    """,
                    (
                        payload["request_id"],
                        payload["session_hash"],
                        payload["ip_hash"],
                        payload["job_type"],
                        payload["model"],
                        int(payload.get("estimated_input_tokens", 0)),
                        int(payload.get("max_output_tokens", 0)),
                        float(payload.get("estimated_cost", 0.0)),
                        int(payload.get("actual_output_chars", 0)),
                        float(payload.get("actual_cost", 0.0)),
                        payload["status"],
                        payload.get("error_code", ""),
                        json.dumps(metadata, ensure_ascii=False),
                    ),
                )

    def upsert_daily_usage(self, payload: dict[str, Any]) -> None:
        self._ensure_initialized()
        with self._connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO daily_budget_usage (
                        usage_day, scope, scope_hash, estimated_cost, actual_cost, request_count
                    )
                    VALUES (%s, %s, %s, %s, %s, %s)
                    ON CONFLICT (usage_day, scope, scope_hash)
                    DO UPDATE SET
                        estimated_cost = daily_budget_usage.estimated_cost + EXCLUDED.estimated_cost,
                        actual_cost = daily_budget_usage.actual_cost + EXCLUDED.actual_cost,
                        request_count = daily_budget_usage.request_count + EXCLUDED.request_count,
                        updated_at = CURRENT_TIMESTAMP
                    """,
                    (
                        payload["usage_day"],
                        payload["scope"],
                        payload["scope_hash"],
                        float(payload.get("estimated_cost", 0.0)),
                        float(payload.get("actual_cost", 0.0)),
                        int(payload.get("request_count", 0)),
                    ),
                )


_audit_store: BaseAuditStore | None = None


def get_audit_store(database_url: str | None = None) -> BaseAuditStore:
    global _audit_store
    if _audit_store is not None:
        return _audit_store
    if database_url or os.getenv("DATABASE_URL"):
        _audit_store = PostgresAuditStore(database_url or os.getenv("DATABASE_URL", ""))
    else:
        _audit_store = NullAuditStore()
    return _audit_store


def reset_audit_store() -> None:
    global _audit_store
    _audit_store = None