Video-Tree-TRM5/adapters/telemetry.py

"""SQLite 遥测记录器 — TelemetryRecorder Protocol 的生产实现。

通过 asyncio.to_thread 将 SQLite 同步写入桥接到异步接口，
确保事件循环不被阻塞。表在首次写入时懒初始化。
"""
from __future__ import annotations

import asyncio
import sqlite3
from pathlib import Path


class SQLiteTelemetryRecorder:
    """基于 SQLite 的 LLM 调用遥测记录器。

    Parameters
    ----------
    db_path : Path
        SQLite 数据库文件路径，父目录必须存在。
    """

    _CREATE_TABLE_SQL = """
    CREATE TABLE IF NOT EXISTS llm_calls (
        call_id            TEXT PRIMARY KEY,
        parent_call_id     TEXT,
        session_id         TEXT,
        model_name         TEXT    NOT NULL,
        provider           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,
        latency_ms         INTEGER NOT NULL,
        ttft_ms            REAL,
        max_inter_token_ms REAL,
        cache_hit          INTEGER NOT NULL DEFAULT 0,
        error              TEXT,
        created_at         TEXT    NOT NULL DEFAULT (datetime('now'))
    );
    """

    _INSERT_SQL = """
    INSERT INTO llm_calls (
        call_id, parent_call_id, session_id,
        model_name, provider, messages, response, thinking,
        prompt_tokens, completion_tokens, latency_ms,
        ttft_ms, max_inter_token_ms,
        cache_hit, error
    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
    """

    def __init__(self, db_path: Path) -> None:
        self._db_path = db_path
        self._table_ready = False

    def _ensure_table(self, conn: sqlite3.Connection) -> None:
        """懒初始化：首次写入时创建 llm_calls 表。"""
        if self._table_ready:
            return
        conn.execute(self._CREATE_TABLE_SQL)
        conn.commit()
        self._table_ready = True

    def _write(
        self,
        *,
        call_id: str,
        parent_call_id: str | None,
        session_id: str | None,
        model_name: str,
        provider: str,
        messages: str,
        response: str,
        thinking: str,
        prompt_tokens: int,
        completion_tokens: int,
        latency_ms: int,
        ttft_ms: float | None,
        max_inter_token_ms: float | None,
        cache_hit: bool,
        error: str | None,
    ) -> None:
        """同步写入一条 LLM 调用记录到 SQLite。"""
        conn = sqlite3.connect(str(self._db_path))
        try:
            self._ensure_table(conn)
            conn.execute(
                self._INSERT_SQL,
                (
                    call_id,
                    parent_call_id,
                    session_id,
                    model_name,
                    provider,
                    messages,
                    response,
                    thinking,
                    prompt_tokens,
                    completion_tokens,
                    latency_ms,
                    ttft_ms,
                    max_inter_token_ms,
                    int(cache_hit),
                    error,
                ),
            )
            conn.commit()
        finally:
            conn.close()

    async def record_llm_call(
        self,
        *,
        call_id: str,
        parent_call_id: str | None,
        session_id: str | None,
        model_name: str,
        provider: str,
        messages: str,
        response: str,
        thinking: str,
        prompt_tokens: int,
        completion_tokens: int,
        latency_ms: int,
        ttft_ms: float | None,
        max_inter_token_ms: float | None,
        cache_hit: bool,
        error: str | None,
    ) -> None:
        """异步记录一次 LLM 调用的遥测数据。

        通过 asyncio.to_thread 将阻塞的 SQLite 写入卸载到线程池，
        避免阻塞事件循环。

        Parameters
        ----------
        call_id : str
            本次调用唯一标识（UUID）。
        parent_call_id : str | None
            父调用 ID（agent step → LLM call 链路）。
        session_id : str | None
            epoch/step/question 关联 ID。
        model_name : str
            使用的模型名。
        provider : str
            API 端点标识。
        messages : str
            原始输入（JSON 字符串）。
        response : str
            原始输出。
        thinking : str
            模型思考过程。
        prompt_tokens : int
            输入 token 用量。
        completion_tokens : int
            输出 token 用量。
        latency_ms : int
            总延迟毫秒。
        ttft_ms : float | None
            首 token 延迟毫秒。
        max_inter_token_ms : float | None
            最大 token 间隔毫秒。
        cache_hit : bool
            是否命中 Redis 缓存。
        error : str | None
            异常信息（正常为 None）。
        """
        await asyncio.to_thread(
            self._write,
            call_id=call_id,
            parent_call_id=parent_call_id,
            session_id=session_id,
            model_name=model_name,
            provider=provider,
            messages=messages,
            response=response,
            thinking=thinking,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            latency_ms=latency_ms,
            ttft_ms=ttft_ms,
            max_inter_token_ms=max_inter_token_ms,
            cache_hit=cache_hit,
            error=error,
        )