diff --git a/src/nsc/runtime/models.py b/src/nsc/runtime/models.py index 71438f5..5be01c4 100644 --- a/src/nsc/runtime/models.py +++ b/src/nsc/runtime/models.py @@ -2,18 +2,37 @@ tier → 模型 的映射在 config/models.yaml(配置是生成物,路由策略是资产)。 成本统计走 litellm.completion_cost;Langfuse trace 在配置缺失时静默降级。 +SW-01:每次调用的 prompt/response 落 SQLite transcripts(Lab ADR-0001 §接口), +best-effort——写库失败静默降级,绝不影响路由本身。 """ from __future__ import annotations +import json import os +import sqlite3 import time from dataclasses import dataclass +from datetime import UTC, datetime from pathlib import Path from typing import Any import yaml +_TRANSCRIPT_SCHEMA = """ +CREATE TABLE IF NOT EXISTS transcripts ( + ts TEXT NOT NULL, + caller TEXT NOT NULL, + model TEXT NOT NULL, + prompt TEXT NOT NULL, + response TEXT NOT NULL, + tokens_in INTEGER NOT NULL DEFAULT 0, + tokens_out INTEGER NOT NULL DEFAULT 0, + cost_usd REAL NOT NULL DEFAULT 0, + experiment_id TEXT NOT NULL DEFAULT '' +) +""" + @dataclass(slots=True) class LLMResult: @@ -36,7 +55,13 @@ def _has_json(text: str) -> bool: class ModelRouter: """按 tier 路由到具体模型,带重试与成本统计。""" - def __init__(self, config_path: str | Path = "config/models.yaml") -> None: + def __init__( + self, + config_path: str | Path = "config/models.yaml", + *, + transcript_db: str | Path | None = None, + experiment_id: str = "", + ) -> None: cfg = yaml.safe_load(Path(config_path).read_text("utf-8")) self.tiers: dict[str, dict[str, Any]] = cfg.get("tiers", {}) self.budgets: dict[str, float] = cfg.get("budgets", {}) @@ -44,6 +69,64 @@ def __init__(self, config_path: str | Path = "config/models.yaml") -> None: self.attempts: int = int(retry.get("attempts", 3)) # litellm 不认识的新模型(如 LongCat-2.0)用配置价兜底成本统计。 self.cost_per_mtok: dict[str, float] = cfg.get("cost_usd_per_mtok", {}) or {} + # SW-01 transcript 台账:库路径/实验号可用环境变量接线(Lab subprocess 场景)。 + self.transcript_db = Path( + transcript_db or os.environ.get("NSC_TRANSCRIPT_DB") or "out/transcripts.db" + ) + self.experiment_id = experiment_id or os.environ.get("NSC_EXPERIMENT_ID", "") + self._tconn: sqlite3.Connection | None = None + + def _transcript_conn(self) -> sqlite3.Connection | None: + """懒建连接;任何 IO 异常 → 返回 None(本功能 best-effort)。""" + if self._tconn is None: + try: + self.transcript_db.parent.mkdir(parents=True, exist_ok=True) + # timeout=0:transcripts 是 best-effort 记账,库被并发写锁住时立刻 + # 失败走静默路径,绝不为它阻塞路由(SW-01 review:busy timeout 会加路由延迟) + self._tconn = sqlite3.connect(str(self.transcript_db), timeout=0.0) + self._tconn.execute(_TRANSCRIPT_SCHEMA) + self._tconn.commit() + except Exception: + self._tconn = None + return self._tconn + + def _record_transcript( + self, + tier: str, + model_id: str, + messages: list[dict[str, str]], + text: str, + tokens_in: int, + tokens_out: int, + cost: float, + ) -> None: + conn = self._transcript_conn() + if conn is None: + return + try: + conn.execute( + "INSERT INTO transcripts (ts, caller, model, prompt, response," + " tokens_in, tokens_out, cost_usd, experiment_id)" + " VALUES (?,?,?,?,?,?,?,?,?)", + ( + datetime.now(UTC).isoformat(), + tier, + model_id, + json.dumps(messages, ensure_ascii=False), + text, + tokens_in, + tokens_out, + cost, + self.experiment_id, + ), + ) + conn.commit() + except Exception: + # best-effort:失败即回滚并弃置连接,防止半开事务长期持锁拖垮后续写入 + try: + conn.rollback() + finally: + self._tconn = None def resolve(self, tier: str) -> dict[str, Any]: if tier not in self.tiers: @@ -119,6 +202,7 @@ def complete( tokens_in * float(self.cost_per_mtok.get("input", 0.0)) + tokens_out * float(self.cost_per_mtok.get("output", 0.0)) ) / 1_000_000 + self._record_transcript(tier, cfg["model"], messages, text, tokens_in, tokens_out, cost) return LLMResult( text=text, model_id=cfg["model"], diff --git a/tests/test_transcripts.py b/tests/test_transcripts.py new file mode 100644 index 0000000..ef25a9d --- /dev/null +++ b/tests/test_transcripts.py @@ -0,0 +1,104 @@ +"""SW-01 transcript 持久化:ModelRouter 每次 LLM 调用的 prompt/response 落 SQLite。 + +表结构对齐 Lab 仓 ADR-0001 §接口:(ts, caller, model, prompt, response, +tokens_in, tokens_out, cost_usd, experiment_id)。库位置默认 out/transcripts.db, +可用 NSC_TRANSCRIPT_DB / 构造参数改道(Lab 经 subprocess 调 SW 时靠环境变量接线)。 +""" + +from __future__ import annotations + +import json +import sqlite3 +import types +from pathlib import Path + +import pytest + + +class _FakeResp: + """litellm.completion 的最小桩(无网络)。""" + + def __init__(self) -> None: + msg = types.SimpleNamespace(content='{"ok": 1}', reasoning_content="") + self.choices = [types.SimpleNamespace(message=msg)] + self.usage = types.SimpleNamespace(prompt_tokens=11, completion_tokens=7) + + +@pytest.fixture() +def fake_litellm(monkeypatch): + calls: list[dict] = [] + + def _completion(**kwargs): + calls.append(kwargs) + return _FakeResp() + + import litellm + + monkeypatch.setattr(litellm, "completion", _completion) + return calls + + +def _rows(db: Path) -> list[dict]: + conn = sqlite3.connect(str(db)) + try: + cur = conn.execute( + "SELECT ts, caller, model, prompt, response, tokens_in, tokens_out," + " cost_usd, experiment_id FROM transcripts" + ) + cols = [c[0] for c in cur.description or []] + return [dict(zip(cols, r, strict=True)) for r in cur.fetchall()] + finally: + conn.close() + + +def test_complete_writes_transcript_row(tmp_path, fake_litellm, monkeypatch): + from nsc.runtime.models import ModelRouter + + db = tmp_path / "t.db" + monkeypatch.setenv("NSC_TRANSCRIPT_DB", str(db)) + router = ModelRouter(experiment_id="exp-42") + res = router.complete( + "tier_bulk", + [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hello"}, + ], + json_mode=True, + seed=1, + ) + assert res.text == '{"ok": 1}' + + rows = _rows(db) + assert len(rows) == 1 + row = rows[0] + assert row["caller"] == "tier_bulk" + # 不硬编码具体模型 ID(models.yaml 会演化):断言落库值 == 路由实际选中的模型 + assert row["model"] == router.resolve("tier_bulk")["model"] + assert json.loads(row["prompt"]) == [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hello"}, + ] + assert row["response"] == '{"ok": 1}' + assert row["tokens_in"] == 11 + assert row["tokens_out"] == 7 + assert row["experiment_id"] == "exp-42" + assert row["cost_usd"] > 0 # cost_per_mtok 兜底价 + + +def test_transcript_db_via_constructor(tmp_path, fake_litellm, monkeypatch): + from nsc.runtime.models import ModelRouter + + monkeypatch.delenv("NSC_TRANSCRIPT_DB", raising=False) + db = tmp_path / "ctor.db" + router = ModelRouter(transcript_db=db) + router.complete("tier_bulk", [{"role": "user", "content": "hi"}]) + assert len(_rows(db)) == 1 + + +def test_transcript_failure_never_breaks_routing(tmp_path, fake_litellm, monkeypatch): + """transcript 是 best-effort 台账:写库失败不得影响 LLM 路由本身。""" + from nsc.runtime.models import ModelRouter + + router = ModelRouter(transcript_db=tmp_path) # 目录当 db 路径 → 打不开 + res = router.complete("tier_bulk", [{"role": "user", "content": "hi"}]) + assert res.text == '{"ok": 1}'