feat(adapters): 实现 SQLiteTelemetryRecorder
asyncio.to_thread 桥接 SQLite，字段与 TelemetryRecorder Protocol 一一对应。 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,186 @@
|
||||
"""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,
|
||||
)
|
||||
@@ -0,0 +1,106 @@
|
||||
"""adapters/telemetry.py 单元测试 — SQLiteTelemetryRecorder。"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from adapters.telemetry import SQLiteTelemetryRecorder
|
||||
from core.protocols import TelemetryRecorder
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture()
|
||||
def db_path(tmp_path):
|
||||
"""返回临时数据库路径。"""
|
||||
return tmp_path / "telemetry_test.db"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def recorder(db_path):
|
||||
"""构造 SQLiteTelemetryRecorder 实例。"""
|
||||
return SQLiteTelemetryRecorder(db_path=db_path)
|
||||
|
||||
|
||||
def _make_call_kwargs(*, cache_hit: bool = False, error: str | None = None):
|
||||
"""构造 record_llm_call 的标准参数字典。"""
|
||||
return dict(
|
||||
call_id=str(uuid.uuid4()),
|
||||
parent_call_id=None,
|
||||
session_id="sess-001",
|
||||
model_name="gpt-4o",
|
||||
provider="openai",
|
||||
messages='[{"role":"user","content":"hi"}]',
|
||||
response="hello",
|
||||
thinking="",
|
||||
prompt_tokens=10,
|
||||
completion_tokens=5,
|
||||
latency_ms=120,
|
||||
ttft_ms=45.2,
|
||||
max_inter_token_ms=12.3,
|
||||
cache_hit=cache_hit,
|
||||
error=error,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_satisfies_protocol(recorder):
|
||||
"""SQLiteTelemetryRecorder 满足 TelemetryRecorder Protocol。"""
|
||||
assert isinstance(recorder, TelemetryRecorder)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_creates_table_and_inserts(recorder, db_path):
|
||||
"""首次写入应懒创建表并成功插入一条记录。"""
|
||||
kwargs = _make_call_kwargs()
|
||||
await recorder.record_llm_call(**kwargs)
|
||||
|
||||
conn = sqlite3.connect(str(db_path))
|
||||
conn.row_factory = sqlite3.Row
|
||||
rows = conn.execute("SELECT * FROM llm_calls").fetchall()
|
||||
conn.close()
|
||||
|
||||
assert len(rows) == 1
|
||||
row = rows[0]
|
||||
assert row["call_id"] == kwargs["call_id"]
|
||||
assert row["model_name"] == "gpt-4o"
|
||||
assert row["prompt_tokens"] == 10
|
||||
assert row["completion_tokens"] == 5
|
||||
assert row["cache_hit"] == 0 # False → INTEGER 0
|
||||
assert row["error"] is None
|
||||
assert row["created_at"] is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_with_error(recorder, db_path):
|
||||
"""error 字段非 None 时应正确存储。"""
|
||||
kwargs = _make_call_kwargs(error="RateLimitError: 429")
|
||||
await recorder.record_llm_call(**kwargs)
|
||||
|
||||
conn = sqlite3.connect(str(db_path))
|
||||
conn.row_factory = sqlite3.Row
|
||||
row = conn.execute("SELECT error FROM llm_calls WHERE call_id = ?", (kwargs["call_id"],)).fetchone()
|
||||
conn.close()
|
||||
|
||||
assert row["error"] == "RateLimitError: 429"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_cache_hit(recorder, db_path):
|
||||
"""cache_hit=True 时应存储为 INTEGER 1。"""
|
||||
kwargs = _make_call_kwargs(cache_hit=True)
|
||||
await recorder.record_llm_call(**kwargs)
|
||||
|
||||
conn = sqlite3.connect(str(db_path))
|
||||
conn.row_factory = sqlite3.Row
|
||||
row = conn.execute("SELECT cache_hit FROM llm_calls WHERE call_id = ?", (kwargs["call_id"],)).fetchone()
|
||||
conn.close()
|
||||
|
||||
assert row["cache_hit"] == 1
|
||||
Reference in New Issue
Block a user