564 lines
20 KiB
Python
564 lines
20 KiB
Python
"""遥测子系统测试: 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 _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"]
|