Files
PolyGateway/tests/unit/test_telemetry.py
T

636 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""遥测子系统测试: SQLiteRecorder(20 列)+ TelemetryEmitter 单一 helper + TelemetryMW。"""
import asyncio
import sqlite3
import subprocess
from pathlib import Path
import pytest
from polygateway.errors import CircuitOpenError, RequestRejectedError
from polygateway.middleware.telemetry import TelemetryEmitter, TelemetryMW
from polygateway.pricing import ModelPrice, PricingTable
from polygateway.telemetry.sqlite import SQLiteRecorder
from polygateway.types import ChatRequest, LLMResponse, SourceConfig
_REQ = ChatRequest(messages=[{"role": "user", "content": "hi"}], session_id="sess-1")
_EXPECTED_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",
"created_at",
"cached_prompt_tokens",
"model_reported",
]
def _resp(**overrides):
base = {
"content": "ok",
"thinking": "",
"model": "m",
"provider": "p",
"prompt_tokens": 1,
"completion_tokens": 2,
"latency_ms": 30,
"ttft_ms": None,
"max_inter_token_ms": None,
"cache_hit": False,
"call_id": "cid-1",
"source_name": "s1",
"usage_source": "measured",
}
base.update(overrides)
return LLMResponse(**base)
def _source():
return SourceConfig(
name="s1",
provider="p",
base_url="https://gw.example/v1",
api_key="sk",
model="m",
timeout_s=10.0,
)
# 输出单价 8 元/百万: 改前 `unavailable` 行按兜底的 0/4000 换算恰好是 0.032
_PRICING = PricingTable({"m": ModelPrice(input_per_1m=1.0, output_per_1m=8.0)})
async def _record_minimal(recorder, call_id="c1", **overrides):
fields = {
"call_id": call_id,
"parent_call_id": None,
"session_id": "sess-1",
"model": "m",
"provider": "p",
"source_name": "s1",
"messages": "[]",
"response": "ok",
"thinking": "",
"prompt_tokens": 1,
"completion_tokens": 2,
"usage_source": "measured",
"latency_ms": 10,
"ttft_ms": None,
"max_inter_token_ms": None,
"cache_hit": False,
"error": None,
"cost": None,
"cached_prompt_tokens": None,
"model_reported": None,
}
fields.update(overrides)
await recorder.record_llm_call(**fields)
class TestSQLiteRecorder:
async def test_schema_has_frozen_columns(self, tmp_path):
recorder = SQLiteRecorder(tmp_path / "t.db")
await _record_minimal(recorder)
recorder.close()
cols = [
r[1] for r in sqlite3.connect(tmp_path / "t.db").execute("PRAGMA table_info(llm_calls)")
]
assert cols == _EXPECTED_COLUMNS
async def test_call_id_idempotent(self, tmp_path):
recorder = SQLiteRecorder(tmp_path / "t.db")
await _record_minimal(recorder, call_id="dup")
await _record_minimal(recorder, call_id="dup", response="second")
recorder.close()
rows = (
sqlite3.connect(tmp_path / "t.db")
.execute("SELECT response FROM llm_calls WHERE call_id='dup'")
.fetchall()
)
assert rows == [("ok",)] # INSERT OR IGNORE: 第二次静默忽略
async def test_concurrent_writes_all_land(self, tmp_path):
recorder = SQLiteRecorder(tmp_path / "t.db")
await asyncio.gather(*(_record_minimal(recorder, call_id=f"c{i}") for i in range(50)))
recorder.close()
(count,) = (
sqlite3.connect(tmp_path / "t.db").execute("SELECT COUNT(*) FROM llm_calls").fetchone()
)
assert count == 50
async def test_unwritable_path_degrades_silently(self):
recorder = SQLiteRecorder(Path("/nonexistent-root/deep/t.db"))
await _record_minimal(recorder) # 不抛
recorder.close()
async def test_observability_columns_round_trip(self, tmp_path):
recorder = SQLiteRecorder(tmp_path / "t.db")
await _record_minimal(recorder, call_id="c-hit", cached_prompt_tokens=64)
await _record_minimal(recorder, call_id="c-zero", cached_prompt_tokens=0)
await _record_minimal(recorder, call_id="c-none", model_reported="MiniMax-Text-01")
recorder.close()
rows = dict(
sqlite3.connect(tmp_path / "t.db")
.execute("SELECT call_id, cached_prompt_tokens FROM llm_calls")
.fetchall()
)
assert rows["c-hit"] == 64
assert rows["c-zero"] == 0 # 真实零命中,读回仍是 0 而非 NULL
assert rows["c-none"] is None
class TestSQLiteColumnBackfill:
"""issue #3: 已存在的 18 列旧表必须自动补列,否则每行写入都被丢弃。"""
_LEGACY_DDL = """
CREATE TABLE 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'))
);
"""
async def test_legacy_table_is_upgraded_in_place(self, tmp_path):
db = tmp_path / "legacy.db"
legacy = sqlite3.connect(db)
legacy.execute(self._LEGACY_DDL)
legacy.commit()
legacy.close()
recorder = SQLiteRecorder(db)
await _record_minimal(recorder, cached_prompt_tokens=7, model_reported="m-real")
recorder.close()
conn = sqlite3.connect(db)
cols = [r[1] for r in conn.execute("PRAGMA table_info(llm_calls)")]
assert cols == _EXPECTED_COLUMNS # ALTER 追加到末尾,与新建库列序一致
assert conn.execute(
"SELECT cached_prompt_tokens, model_reported FROM llm_calls"
).fetchone() == (7, "m-real")
async def test_backfill_failure_keeps_the_recorder_usable(self, tmp_path):
"""补列失败只能逐行降级,绝不能把 recorder 整体变成 no-op(设计 D1 纪律)。
把 llm_calls 建成同名 view: `CREATE TABLE IF NOT EXISTS` 遇 view 静默
no-op(不抛),随后的 ALTER 才抛 "Cannot add a column to a view"——正是
补列失败这条分支。`_conn` 必须保持非 None,否则整个 recorder 永久失能。
"""
db = tmp_path / "view.db"
conn = sqlite3.connect(db)
conn.execute("CREATE TABLE real_rows (call_id TEXT)")
conn.execute("CREATE VIEW llm_calls AS SELECT call_id FROM real_rows")
conn.commit()
conn.close()
recorder = SQLiteRecorder(db) # 不得抛
assert recorder._conn is not None # 补列失败 ≠ recorder 失能(D1 纪律)
await _record_minimal(recorder) # 不得抛
recorder.close()
class _FakePgConn:
"""记录执行过的语句;可让 ALTER 抛错以模拟权限不足。"""
def __init__(self, existing: list[str], *, fail_alter: bool = False):
self.existing = existing
self.fail_alter = fail_alter
self.statements: list[str] = []
async def execute(self, sql, *args):
self.statements.append(sql)
if sql.startswith("ALTER TABLE") and self.fail_alter:
raise RuntimeError("must be owner of table llm_calls")
async def fetch(self, sql, *args):
self.statements.append(sql)
return [{"attname": name} for name in self.existing]
class _FakePgPool:
def __init__(self, conn):
self._conn = conn
def acquire(self):
conn = self._conn
class _Ctx:
async def __aenter__(self):
return conn
async def __aexit__(self, *exc):
return False
return _Ctx()
class TestPostgresBackfillDiscipline:
"""PG 补列必须与 SQLite 侧对称: 失败只逐行降级,且稳态不抢排他锁(issue #3)。"""
_LEGACY = ["call_id", "cost", "created_at"]
_CURRENT = ["call_id", "cost", "created_at", "cached_prompt_tokens", "model_reported"]
def _recorder(self, conn):
from polygateway.telemetry.postgres import PostgresRecorder
return PostgresRecorder("postgresql://u:p@h:5432/polygateway", pool=_FakePgPool(conn))
async def test_alter_failure_does_not_disable_the_recorder(self):
"""ALTER 失败(如账号只有 INSERT 权限)不得置 _failed —— 那会让遥测全灭。"""
conn = _FakePgConn(self._LEGACY, fail_alter=True)
recorder = self._recorder(conn)
await _record_minimal(recorder) # 不得抛
assert recorder._failed is False
assert any(s.startswith("INSERT INTO llm_calls") for s in conn.statements)
async def test_no_alter_when_columns_already_exist(self):
"""ADD COLUMN IF NOT EXISTS 即使列已存在也会先抢 ACCESS EXCLUSIVE 锁,
而遥测是内联 await——稳态下必须一条 ALTER 都不发,否则每个进程的首次
写入都会去锁共享审计表。
"""
conn = _FakePgConn(self._CURRENT)
await _record_minimal(self._recorder(conn))
assert not [s for s in conn.statements if s.startswith("ALTER TABLE")]
async def test_missing_columns_are_added_once(self):
conn = _FakePgConn(self._LEGACY)
await _record_minimal(self._recorder(conn))
altered = [s for s in conn.statements if s.startswith("ALTER TABLE")]
assert len(altered) == 2
assert all("IF NOT EXISTS" not in s for s in altered) # 探测已确认缺列,无需再判
class _MemoryRecorder:
def __init__(self):
self.rows = []
async def record_llm_call(self, **fields):
self.rows.append(fields)
class TestEmitterRecorderContract:
"""emitter 的实参键集合必须与两个后端的 _COLUMNS 完全一致(issue #3)。
两个后端的 `row = tuple(fields[col] for col in _COLUMNS)` 都在 try **之外**,
emitter 漏传一个键就抛 KeyError,被 `_record` 的 except Exception 吞成 warning
→ 遥测静默全丢。而 8 个 `**fields` 形态的 fake 一个都拦不住,故显式断言。
"""
async def test_emitter_supplies_exactly_the_backend_columns(self):
from polygateway.telemetry.postgres import _COLUMNS as PG_COLUMNS
from polygateway.telemetry.sqlite import _COLUMNS as SQLITE_COLUMNS
rec = _MemoryRecorder()
await TelemetryEmitter(rec).emit_attempt(
request=_REQ,
source=_source(),
call_id="cid-1",
latency_ms=42,
response=_resp(),
error=None,
)
assert set(rec.rows[0]) == set(SQLITE_COLUMNS) == set(PG_COLUMNS)
@pytest.mark.parametrize("emit", ["attempt", "cache_hit", "terminal_failure"])
async def test_every_entry_point_supplies_the_same_keys(self, emit):
from polygateway.telemetry.sqlite import _COLUMNS as SQLITE_COLUMNS
rec = _MemoryRecorder()
emitter = TelemetryEmitter(rec)
if emit == "attempt":
await emitter.emit_attempt(
request=_REQ,
source=_source(),
call_id="c",
latency_ms=1,
response=None,
error="boom",
)
elif emit == "cache_hit":
await emitter.emit_cache_hit(request=_REQ, response=_resp())
else:
await emitter.emit_terminal_failure(
request=_REQ, call_id="c", latency_ms=1, error="dead"
)
assert set(rec.rows[0]) == set(SQLITE_COLUMNS)
class TestEmitterObservabilityFields:
"""issue #3: 三个入口各自的取值口径(设计 §5 表)。"""
async def test_attempt_carries_the_response_values(self):
rec = _MemoryRecorder()
await TelemetryEmitter(rec).emit_attempt(
request=_REQ,
source=_source(),
call_id="cid-1",
latency_ms=42,
response=_resp(cached_prompt_tokens=64, model_reported="m-real"),
error=None,
)
assert rec.rows[0]["cached_prompt_tokens"] == 64
assert rec.rows[0]["model_reported"] == "m-real"
async def test_failed_attempt_has_no_provider_facts(self):
rec = _MemoryRecorder()
await TelemetryEmitter(rec).emit_attempt(
request=_REQ,
source=_source(),
call_id="cid-2",
latency_ms=7,
response=None,
error="boom",
)
assert rec.rows[0]["cached_prompt_tokens"] is None
assert rec.rows[0]["model_reported"] is None
async def test_cache_hit_replays_the_recorded_values(self):
"""决策 B1: 命中行原样回放,故命中率统计必须带 WHERE cache_hit = false。"""
rec = _MemoryRecorder()
await TelemetryEmitter(rec).emit_cache_hit(
request=_REQ, response=_resp(cached_prompt_tokens=64, model_reported="m-real")
)
row = rec.rows[0]
assert row["cache_hit"] is True
assert row["cached_prompt_tokens"] == 64 and row["model_reported"] == "m-real"
async def test_terminal_failure_records_none(self):
rec = _MemoryRecorder()
await TelemetryEmitter(rec).emit_terminal_failure(
request=_REQ, call_id="c", latency_ms=1, error="dead"
)
assert rec.rows[0]["cached_prompt_tokens"] is None
assert rec.rows[0]["model_reported"] is None
class TestCostWithCachedTier:
"""issue #3: 命中部分按缓存单价计费,避免 cost 系统性高估。"""
_TABLE = PricingTable(
{"m": ModelPrice(input_per_1m=10.0, output_per_1m=20.0, cached_input_per_1m=2.0)}
)
async def test_cached_hit_lowers_the_recorded_cost(self):
rec = _MemoryRecorder()
emitter = TelemetryEmitter(rec, pricing=self._TABLE)
full = _resp(prompt_tokens=1_000_000, completion_tokens=0)
await emitter.emit_attempt(
request=_REQ,
source=_source(),
call_id="c1",
latency_ms=1,
response=full,
error=None,
)
await emitter.emit_attempt(
request=_REQ,
source=_source(),
call_id="c2",
latency_ms=1,
response=_resp(
prompt_tokens=1_000_000, completion_tokens=0, cached_prompt_tokens=600_000
),
error=None,
)
assert rec.rows[0]["cost"] == pytest.approx(10.0)
assert rec.rows[1]["cost"] == pytest.approx(5.2) # 400k×10 + 600k×2
async def test_cache_hit_row_still_costs_zero(self):
"""缓存命中未产生新调用 → cost 恒 0.0,该短路必须排在任何换算之前。"""
rec = _MemoryRecorder()
await TelemetryEmitter(rec, pricing=self._TABLE).emit_cache_hit(
request=_REQ,
response=_resp(prompt_tokens=1_000_000, cached_prompt_tokens=600_000),
)
assert rec.rows[0]["cost"] == 0.0
async def test_unavailable_usage_still_costs_none(self):
rec = _MemoryRecorder()
await TelemetryEmitter(rec, pricing=self._TABLE).emit_attempt(
request=_REQ,
source=_source(),
call_id="c",
latency_ms=1,
response=_resp(usage_source="unavailable", cached_prompt_tokens=5),
error=None,
)
assert rec.rows[0]["cost"] is None
class TestEmitter:
async def test_attempt_success_row(self):
rec = _MemoryRecorder()
emitter = TelemetryEmitter(rec)
await emitter.emit_attempt(
request=_REQ,
source=_source(),
call_id="cid-1",
latency_ms=42,
response=_resp(),
error=None,
)
row = rec.rows[0]
assert row["call_id"] == "cid-1" and row["error"] is None
assert row["session_id"] == "sess-1" and row["source_name"] == "s1"
assert row["response"] == "ok" and row["cost"] is None
async def test_attempt_failure_row(self):
rec = _MemoryRecorder()
emitter = TelemetryEmitter(rec)
await emitter.emit_attempt(
request=_REQ,
source=_source(),
call_id="cid-2",
latency_ms=7,
response=None,
error="TransientError: boom",
)
row = rec.rows[0]
assert row["error"].startswith("TransientError")
# 失败尝试没有任何用量信息可言 → unavailable(设计 §3.2 #6)
assert row["response"] == "" and row["usage_source"] == "unavailable"
assert row["cost"] is None
async def test_terminal_failure_row_is_unavailable(self):
rec = _MemoryRecorder()
await TelemetryEmitter(rec, pricing=_PRICING).emit_terminal_failure(
request=_REQ, call_id="cid-t", latency_ms=5, error="cancelled"
)
row = rec.rows[0]
assert row["usage_source"] == "unavailable" and row["cost"] is None
@pytest.mark.parametrize(("prompt", "completion"), [(0, 0), (0, 4000)])
async def test_unavailable_success_row_has_null_cost(self, prompt, completion):
"""产生了真实调用但用量不可得 → cost 记 NULL(设计 §3.1 不变式)。
参数第二组是改前兜底写出的 `0/4000` 形态: 那时换算出 0.032 的假金额。
"""
rec = _MemoryRecorder()
await TelemetryEmitter(rec, pricing=_PRICING).emit_attempt(
request=_REQ,
source=_source(),
call_id="cid-u",
latency_ms=42,
response=_resp(
usage_source="unavailable", prompt_tokens=prompt, completion_tokens=completion
),
error=None,
)
assert rec.rows[0]["cost"] is None
async def test_measured_row_still_priced(self):
"""对照组: 同一价格表下 measured 行照常换算,证明 None 不是价格表没接上。"""
rec = _MemoryRecorder()
await TelemetryEmitter(rec, pricing=_PRICING).emit_attempt(
request=_REQ,
source=_source(),
call_id="cid-m",
latency_ms=42,
response=_resp(prompt_tokens=0, completion_tokens=4000),
error=None,
)
assert rec.rows[0]["cost"] == pytest.approx(0.032)
async def test_cache_hit_keeps_zero_cost_even_when_unavailable(self):
"""缓存命中未产生新调用,0.0 是事实而非未知 → 短路必须排在 cache_hit 之后。"""
rec = _MemoryRecorder()
await TelemetryEmitter(rec, pricing=_PRICING).emit_cache_hit(
request=_REQ,
response=_resp(cache_hit=True, usage_source="unavailable", completion_tokens=4000),
)
assert rec.rows[0]["cache_hit"] is True and rec.rows[0]["cost"] == 0.0
async def test_multimodal_messages_digested_before_storage(self):
rec = _MemoryRecorder()
emitter = TelemetryEmitter(rec)
big = "data:image/png;base64," + "A" * 100_000
req = ChatRequest(
messages=[
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": big}},
],
}
]
)
await emitter.emit_attempt(
request=req,
source=_source(),
call_id="c",
latency_ms=1,
response=None,
error="x",
)
assert len(rec.rows[0]["messages"]) < 500 # base64 不整段进库(VT R12)
async def test_recorder_failure_swallowed(self):
class Broken:
async def record_llm_call(self, **fields):
raise OSError("disk full")
emitter = TelemetryEmitter(Broken())
await emitter.emit_attempt(
request=_REQ,
source=_source(),
call_id="c",
latency_ms=1,
response=_resp(),
error=None,
) # 不抛(降级不冒泡)
class TestTelemetryMW:
async def test_cache_hit_recorded(self):
rec = _MemoryRecorder()
mw = TelemetryMW(TelemetryEmitter(rec))
async def terminal(request):
return _resp(cache_hit=True, latency_ms=0, call_id="cache-cid")
resp = await mw(_REQ, terminal)
assert resp.cache_hit
assert len(rec.rows) == 1
assert rec.rows[0]["cache_hit"] is True and rec.rows[0]["latency_ms"] == 0
async def test_normal_success_not_double_recorded(self):
"""成功尝试由 RetryMW 逐次记录;最外层不得重复记。"""
rec = _MemoryRecorder()
mw = TelemetryMW(TelemetryEmitter(rec))
async def terminal(request):
return _resp(cache_hit=False)
await mw(_REQ, terminal)
assert rec.rows == []
async def test_scope_level_failure_recorded(self):
rec = _MemoryRecorder()
mw = TelemetryMW(TelemetryEmitter(rec))
async def terminal(request):
raise CircuitOpenError(scope="llm", retry_after_s=30.0)
with pytest.raises(CircuitOpenError):
await mw(_REQ, terminal)
assert len(rec.rows) == 1 and "circuit_open" in rec.rows[0]["error"]
async def test_attempt_level_failure_not_double_recorded(self):
"""RequestRejected 已被 RetryMW 逐次记录 → 最外层跳过。"""
rec = _MemoryRecorder()
mw = TelemetryMW(TelemetryEmitter(rec))
async def terminal(request):
raise RequestRejectedError("400")
with pytest.raises(RequestRejectedError):
await mw(_REQ, terminal)
assert rec.rows == []
def test_single_emitter_discipline():
"""铁律执法: record_llm_call 在 src/ 的调用点只允许出现在 telemetry emitter。"""
out = subprocess.run(
["grep", "-rln", "record_llm_call(", "src/polygateway"],
capture_output=True,
text=True,
cwd=Path(__file__).resolve().parents[2],
).stdout.splitlines()
callers = [
p
for p in out
if not p.endswith(("ports.py", "telemetry/sqlite.py", "telemetry/postgres.py"))
]
assert callers == ["src/polygateway/middleware/telemetry.py"]