194 lines
7.5 KiB
Python
194 lines
7.5 KiB
Python
#!/usr/bin/env python3
|
||
"""把一次 LLM 调用的输入、输出和调用明细原子落库。
|
||
|
||
调用入口只负责采集供应商响应；本模块负责把同一尝试的 raw lease、prompt、response
|
||
和 ``example_llm_call`` 放进一个事务。数据库失败时调用方不能拿到一份看似成功、却
|
||
没有审计证据的结果。
|
||
"""
|
||
import json
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
|
||
HERE = Path(__file__).resolve().parent
|
||
if str(HERE) not in sys.path:
|
||
sys.path.insert(0, str(HERE))
|
||
from persist_raw import _bare_sha256, _check_no_secrets # noqa: E402
|
||
|
||
|
||
CREATOR = "llm"
|
||
|
||
|
||
def _json_text(value) -> str:
|
||
"""生成稳定、可重放且不带请求头的 JSON 原文。"""
|
||
return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), default=str)
|
||
|
||
|
||
def _usage_int(usage, key):
|
||
"""读取供应商 usage，缺失或脏值按 0 记账。"""
|
||
value = (usage or {}).get(key, 0)
|
||
try:
|
||
return max(0, int(value or 0))
|
||
except (TypeError, ValueError):
|
||
return 0
|
||
|
||
|
||
def _insert_raw(conn, *, lease_id, kind, content, run_id, role, creator):
|
||
"""在调用事务内幂等插入 raw，避免 append-only 表上的冲突 UPDATE。"""
|
||
sha = _bare_sha256(content)
|
||
row = conn.execute(
|
||
"INSERT INTO example_raw_content(lease_id, kind, run_id, role, content_sha256, content, creator) "
|
||
"VALUES (%s,%s,%s,%s,%s,%s,%s) ON CONFLICT (lease_id, content_sha256) DO NOTHING "
|
||
"RETURNING id",
|
||
(lease_id, kind, run_id, role, sha, content, creator),
|
||
).fetchone()
|
||
if row:
|
||
return row[0], sha
|
||
row = conn.execute(
|
||
"SELECT id FROM example_raw_content WHERE lease_id=%s AND content_sha256=%s",
|
||
(lease_id, sha),
|
||
).fetchone()
|
||
if not row:
|
||
raise RuntimeError(f"raw {kind} 插入后无法回读 id")
|
||
return row[0], sha
|
||
|
||
|
||
def _reuse_raw(conn, *, lease_id, prompt_sha, response_sha):
|
||
"""迁移已有 raw 时只回读并校验，不复制全文或改变 append-only 数据。"""
|
||
prompt_row = conn.execute(
|
||
"SELECT id FROM example_raw_content WHERE lease_id=%s AND kind='prompt' AND content_sha256=%s",
|
||
(lease_id, prompt_sha),
|
||
).fetchone()
|
||
if not prompt_row:
|
||
raise ValueError(f"raw lease={lease_id} 不含匹配的 prompt")
|
||
response_id = None
|
||
if response_sha:
|
||
response_row = conn.execute(
|
||
"SELECT id FROM example_raw_content WHERE lease_id=%s AND kind='response' AND content_sha256=%s",
|
||
(lease_id, response_sha),
|
||
).fetchone()
|
||
if not response_row:
|
||
raise ValueError(f"raw lease={lease_id} 不含匹配的 response")
|
||
response_id = response_row[0]
|
||
return prompt_row[0], response_id
|
||
|
||
|
||
def persist_call(event, *, creator=CREATOR, dry_run=False):
|
||
"""原子写入一条成功调用。
|
||
|
||
``event`` 由 llm.chat 生成，至少包含 prompt/response、模型、usage 和调用方字段。
|
||
返回各证据行 id；dry-run 只验证事务并回滚。
|
||
"""
|
||
from db import connect
|
||
|
||
prompt = event.get("prompt")
|
||
response = event.get("response")
|
||
if not isinstance(prompt, str) or not prompt:
|
||
raise ValueError("LLM raw prompt 不能为空")
|
||
if response is not None and not isinstance(response, str):
|
||
raise ValueError("LLM raw response 必须是字符串或 NULL")
|
||
_check_no_secrets(prompt)
|
||
if response:
|
||
_check_no_secrets(response)
|
||
|
||
requested = str(event.get("requested_model_id") or "")
|
||
actual = str(event.get("actual_model_id") or "")
|
||
caller = str(event.get("caller") or "")
|
||
if not requested or not actual:
|
||
raise ValueError("LLM 调用明细缺 requested_model_id/actual_model_id")
|
||
if len(requested) > 64 or len(actual) > 64 or len(caller) > 64:
|
||
raise ValueError("LLM 调用明细字段超长")
|
||
|
||
usage = event.get("usage") or {}
|
||
cached = _usage_int((usage.get("prompt_tokens_details") or {}), "cached_tokens")
|
||
prompt_sha = _bare_sha256(prompt)
|
||
response_sha = _bare_sha256(response) if response else None
|
||
content_hashes = {"prompt": prompt_sha}
|
||
if response_sha:
|
||
content_hashes["response"] = response_sha
|
||
|
||
with connect() as conn:
|
||
try:
|
||
if event.get("raw_lease_id") is not None:
|
||
lease_id = int(event["raw_lease_id"])
|
||
prompt_id, response_id = _reuse_raw(
|
||
conn, lease_id=lease_id, prompt_sha=prompt_sha, response_sha=response_sha
|
||
)
|
||
else:
|
||
lease_id = conn.execute(
|
||
"INSERT INTO example_raw_lease(run_id, source_version, content_hashes, purpose, status, creator) "
|
||
"VALUES (%s,%s,%s::jsonb,%s,%s,%s) RETURNING id",
|
||
(
|
||
event.get("run_id"),
|
||
event.get("source_version"),
|
||
json.dumps(content_hashes, ensure_ascii=False),
|
||
event.get("purpose") or "llm_call",
|
||
"closed",
|
||
creator,
|
||
),
|
||
).fetchone()[0]
|
||
prompt_id, _ = _insert_raw(
|
||
conn,
|
||
lease_id=lease_id,
|
||
kind="prompt",
|
||
content=prompt,
|
||
run_id=event.get("run_id"),
|
||
role=event.get("role") or caller,
|
||
creator=creator,
|
||
)
|
||
response_id = None
|
||
if response:
|
||
response_id, _ = _insert_raw(
|
||
conn,
|
||
lease_id=lease_id,
|
||
kind="response",
|
||
content=response,
|
||
run_id=event.get("run_id"),
|
||
role=event.get("role") or caller,
|
||
creator=creator,
|
||
)
|
||
call_id = conn.execute(
|
||
"INSERT INTO example_llm_call(window_key, run_id, caller, requested_model_id, actual_model_id, "
|
||
"model_match, in_tokens, cached_tokens, out_tokens, cost_usd, stop_reason, duration_ms, "
|
||
"prompt_sha256, raw_content_id, creator) "
|
||
"VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) RETURNING id",
|
||
(
|
||
event.get("window_key"),
|
||
event.get("run_id"),
|
||
caller,
|
||
requested,
|
||
actual,
|
||
requested == actual,
|
||
_usage_int(usage, "prompt_tokens"),
|
||
cached,
|
||
_usage_int(usage, "completion_tokens"),
|
||
event.get("cost_usd") or 0,
|
||
str(event.get("stop_reason") or "")[:32] or None,
|
||
event.get("duration_ms"),
|
||
prompt_sha,
|
||
response_id,
|
||
creator,
|
||
),
|
||
).fetchone()[0]
|
||
if dry_run:
|
||
conn.rollback()
|
||
return {
|
||
"status": "dry_run_ok",
|
||
"lease_id": lease_id,
|
||
"prompt_id": prompt_id,
|
||
"response_id": response_id,
|
||
"call_id": call_id,
|
||
"note": "试跑已回滚，未落库",
|
||
}
|
||
conn.commit()
|
||
return {
|
||
"status": "written",
|
||
"lease_id": lease_id,
|
||
"prompt_id": prompt_id,
|
||
"response_id": response_id,
|
||
"call_id": call_id,
|
||
}
|
||
except Exception:
|
||
conn.rollback()
|
||
raise
|