PolyGateway/src/polygateway/telemetry/sqlite.py

"""SQLite 遥测后端(默认): WAL + 单持久连接 + to_thread 桥接。

蓝本 VT `adapters/telemetry.py`: 构造期建连接与表,失败降级为 no-op
(记录基础设施不得拖垮业务调用);`INSERT OR IGNORE` 幂等(call_id 主键);
写入经 threading.Lock 串行化后由 `asyncio.to_thread` 执行,不阻塞事件循环。

**这里不做 postgres.py 那样的建表前探测,是实测后的有意不对称**(issue #9):
SQLite 对已存在的表在**解析期**就把 `CREATE TABLE IF NOT EXISTS` 短路掉,
既不抢写锁也不检查可写性——实测同一时刻另一连接持 `BEGIN EXCLUSIVE`、或
文件 `chmod 444`,该语句均通过,而同条件下的 `INSERT` 与新表名建表分别报
database is locked / readonly database。故 PG 侧"权限检查早于存在性判断"
的坑在此不存在,加探测零收益。**别为了代码对称把它加回来**;需要对称的是
保证(表存在就不该因建表失败而失能),这一条两侧都已满足。
"""

from __future__ import annotations

import asyncio
import sqlite3
import threading
from pathlib import Path

from loguru import logger

_DDL = """
CREATE TABLE IF NOT EXISTS llm_calls (
    call_id            TEXT PRIMARY KEY,
    parent_call_id     TEXT,
    session_id         TEXT,
    model              TEXT    NOT NULL,
    provider           TEXT    NOT NULL,
    source_name        TEXT    NOT NULL,
    messages           TEXT    NOT NULL,
    response           TEXT    NOT NULL,
    thinking           TEXT    NOT NULL DEFAULT '',
    prompt_tokens      INTEGER NOT NULL,
    completion_tokens  INTEGER NOT NULL,
    usage_source       TEXT    NOT NULL,
    latency_ms         INTEGER NOT NULL,
    ttft_ms            REAL,
    max_inter_token_ms REAL,
    cache_hit          INTEGER NOT NULL DEFAULT 0,
    error              TEXT,
    cost               REAL,
    created_at         TEXT    NOT NULL DEFAULT (datetime('now')),
    cached_prompt_tokens INTEGER,
    model_reported       TEXT,
    sampling             TEXT,
    reasoning_tokens     INTEGER
);
"""

# 新列必须排在 created_at 之后: 旧表只能经 ALTER 追加到末尾,新建库若把它们
# 插在前面,两条路径的物理列序会分叉(列序断言测试无合规修法)。
_BACKFILL_COLUMNS = (
    ("cached_prompt_tokens", "INTEGER"),
    ("model_reported", "TEXT"),
    ("sampling", "TEXT"),
    ("reasoning_tokens", "INTEGER"),
)

_COLUMNS = (
    "call_id",
    "parent_call_id",
    "session_id",
    "model",
    "provider",
    "source_name",
    "messages",
    "response",
    "thinking",
    "prompt_tokens",
    "completion_tokens",
    "usage_source",
    "latency_ms",
    "ttft_ms",
    "max_inter_token_ms",
    "cache_hit",
    "error",
    "cost",
    "cached_prompt_tokens",
    "model_reported",
    "sampling",
    "reasoning_tokens",
)

_INSERT = (
    f"INSERT OR IGNORE INTO llm_calls ({', '.join(_COLUMNS)}) "
    f"VALUES ({', '.join('?' for _ in _COLUMNS)})"
)


class SQLiteRecorder:
    """TelemetryRecorder 端口的 SQLite 实现;初始化/写入失败全降级 warning。"""

    def __init__(self, db_path: Path | str) -> None:
        self._lock = threading.Lock()
        self._conn: sqlite3.Connection | None = None
        try:
            path = Path(db_path)
            path.parent.mkdir(parents=True, exist_ok=True)
            conn = sqlite3.connect(path, check_same_thread=False, timeout=10.0)
            conn.execute("PRAGMA journal_mode=WAL")
            conn.execute("PRAGMA busy_timeout=5000")
            conn.execute(_DDL)
            conn.commit()
            self._conn = conn
        except (OSError, sqlite3.Error) as exc:
            logger.warning("SQLite 遥测初始化失败,后续记录降级为 no-op: {}", exc)
        self._backfill_columns()

    def _backfill_columns(self) -> None:
        """给已存在的旧表补新列(issue #3);独立 try,失败只降级为逐行丢弃。

        必须放在 `self._conn` 赋值**之后**并先判空: 初始化失败时连接为 None,
        无守卫的补列会抛 AttributeError 逃出 `__init__`,把"静默降级"变成崩溃。
        补列失败也绝不清空 `self._conn`——那会让整个 recorder 永久 no-op,
        比逐行丢弃严重得多。
        """
        if self._conn is None:
            return
        try:
            existing = {row[1] for row in self._conn.execute("PRAGMA table_info(llm_calls)")}
        except sqlite3.Error as exc:
            logger.warning("SQLite 遥测列探测失败(写入将逐行降级): {}", exc)
            return
        for column, decl in _BACKFILL_COLUMNS:
            if column in existing:
                continue
            # 逐列独立 try: 一列撞上 duplicate 不得让后面的列漏补
            try:
                self._conn.execute(f"ALTER TABLE llm_calls ADD COLUMN {column} {decl}")
                self._conn.commit()
            except sqlite3.Error as exc:
                # duplicate column: 多进程共库时后到者必然撞上,属预期竞态,视为成功
                if "duplicate column" not in str(exc).lower():
                    logger.warning("SQLite 遥测补列失败(写入将逐行降级): {}", exc)

    async def record_llm_call(self, **fields: object) -> None:
        """写一行遥测;字段集合即 21 字段冻结签名(ports.TelemetryRecorder)。"""
        if self._conn is None:
            return
        row = tuple(fields[col] for col in _COLUMNS)
        try:
            await asyncio.to_thread(self._write, row)
        except (OSError, sqlite3.Error) as exc:
            logger.warning("SQLite 遥测写入失败(降级不冒泡): {}", exc)

    def _write(self, row: tuple) -> None:
        assert self._conn is not None  # 内部不变量: 调用方已判空
        with self._lock:
            self._conn.execute(_INSERT, row)
            self._conn.commit()

    def close(self) -> None:
        """幂等关闭持久连接。"""
        conn, self._conn = self._conn, None
        if conn is not None:
            with self._lock:
                conn.close()