213 lines
7.2 KiB
Python
213 lines
7.2 KiB
Python
"""SQLite 遥测记录器 — TelemetryRecorder Protocol 的生产实现。
|
||
|
||
通过 asyncio.to_thread 将 SQLite 同步写入桥接到异步接口，
|
||
确保事件循环不被阻塞。表在首次写入时懒初始化。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import sqlite3
|
||
import threading
|
||
from pathlib import Path
|
||
|
||
from loguru import logger
|
||
|
||
|
||
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 OR IGNORE 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:
|
||
"""建单持久连接 + 进程内 Lock（对齐 app/harness/log.py:HarnessLog 并发写模式）。
|
||
|
||
把并发控制拉到进程内（threading.Lock 串行化写），消除"每次新连接并发写同一
|
||
db、靠 SQLite busy_timeout 跨连接协调"在高频下撑爆 timeout → database is locked
|
||
的根因。check_same_thread=False：record_llm_call 经 asyncio.to_thread 在线程池
|
||
不同线程调用，共享连接跨线程访问需此 flag，串行性由 self._lock 保证。
|
||
|
||
遥测哲学（P5）：连接初始化失败降级不冒泡（self._conn=None，写入直接丢弃 warning），
|
||
绝不因遥测故障拖垮 LLM 调用 / 训练。
|
||
"""
|
||
self._db_path = db_path
|
||
self._lock = threading.Lock()
|
||
self._conn: sqlite3.Connection | None = None
|
||
try:
|
||
Path(db_path).parent.mkdir(parents=True, exist_ok=True)
|
||
conn = sqlite3.connect(str(db_path), check_same_thread=False)
|
||
conn.execute("PRAGMA journal_mode=WAL")
|
||
conn.execute("PRAGMA busy_timeout=5000")
|
||
conn.execute(self._CREATE_TABLE_SQL)
|
||
conn.commit()
|
||
self._conn = conn
|
||
except sqlite3.Error as exc:
|
||
logger.warning("遥测连接初始化失败（已降级，后续写入丢弃）: {}", exc)
|
||
self._conn = None
|
||
|
||
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 调用记录（单持久连接 + Lock 串行化，对齐 HarnessLog）。
|
||
|
||
三层防御：
|
||
1. INSERT OR IGNORE — call_id 主键冲突静默忽略（幂等）
|
||
2. 进程内 threading.Lock 串行化写 — 消除并发锁竞争（非依赖 SQLite busy_timeout）
|
||
3. try/except sqlite3.Error — DB 错误降级不冒泡，遥测失败绝不拖垮 LLM 调用
|
||
"""
|
||
if self._conn is None:
|
||
logger.warning("遥测连接不可用（已降级），丢弃 call_id={}", call_id)
|
||
return
|
||
try:
|
||
with self._lock:
|
||
self._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,
|
||
),
|
||
)
|
||
self._conn.commit()
|
||
except sqlite3.Error as exc:
|
||
logger.warning("遥测写入失败（已降级），call_id={}: {}", call_id, exc)
|
||
|
||
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,
|
||
)
|