275 lines
10 KiB
Python
275 lines
10 KiB
Python
"""原始报文录制 - LangChain 回调处理器，捕获每次 LLM 调用的完整请求/响应"""
|
||
|
||
import json
|
||
import os
|
||
import time
|
||
import traceback
|
||
import uuid
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
from langchain_core.callbacks import AsyncCallbackHandler
|
||
from langchain_core.messages import BaseMessage
|
||
from langchain_core.outputs import LLMResult
|
||
|
||
|
||
class RawTraceCallback(AsyncCallbackHandler):
|
||
"""录制每次 LLM 调用的完整原始请求和响应到 JSONL 文件"""
|
||
|
||
def __init__(self, session_id: str, raw_logs_dir: Path):
|
||
super().__init__()
|
||
self.session_id = session_id
|
||
self.raw_logs_dir = raw_logs_dir
|
||
self.raw_logs_dir.mkdir(exist_ok=True, parents=True)
|
||
|
||
# 当前调用的跟踪信息 (key: run_id, value: trace_data)
|
||
self._traces: Dict[str, Dict[str, Any]] = {}
|
||
|
||
def _serialize_message(self, msg: BaseMessage) -> Dict[str, Any]:
|
||
"""将 LangChain 消息序列化为 JSON-safe dict"""
|
||
try:
|
||
result = {
|
||
"type": msg.type,
|
||
"content": msg.content if hasattr(msg, "content") else "",
|
||
}
|
||
# 提取额外字段
|
||
if hasattr(msg, "name"):
|
||
result["name"] = msg.name
|
||
if hasattr(msg, "tool_calls") and msg.tool_calls:
|
||
result["tool_calls"] = msg.tool_calls
|
||
if hasattr(msg, "additional_kwargs") and msg.additional_kwargs:
|
||
result["additional_kwargs"] = msg.additional_kwargs
|
||
return result
|
||
except Exception:
|
||
return {"type": "unknown", "content": str(msg)}
|
||
|
||
async def on_llm_start(
|
||
self,
|
||
serialized: Dict[str, Any],
|
||
prompts: List[str],
|
||
*,
|
||
run_id: uuid.UUID,
|
||
parent_run_id: Optional[uuid.UUID] = None,
|
||
tags: Optional[List[str]] = None,
|
||
metadata: Optional[Dict[str, Any]] = None,
|
||
**kwargs: Any,
|
||
) -> None:
|
||
"""LLM 调用开始 - 记录请求信息"""
|
||
try:
|
||
# 提取 messages (如果是 chat model)
|
||
messages_raw = kwargs.get("invocation_params", {}).get("messages", [])
|
||
messages_serialized = []
|
||
|
||
if messages_raw:
|
||
for msg in messages_raw:
|
||
if isinstance(msg, BaseMessage):
|
||
messages_serialized.append(self._serialize_message(msg))
|
||
elif isinstance(msg, dict):
|
||
messages_serialized.append(msg)
|
||
else:
|
||
messages_serialized.append({"content": str(msg)})
|
||
elif prompts:
|
||
# 如果没有 messages，使用 prompts
|
||
messages_serialized = [{"type": "human", "content": p} for p in prompts]
|
||
|
||
# 提取模型参数
|
||
invocation_params = kwargs.get("invocation_params", {})
|
||
model_name = invocation_params.get("model", "unknown")
|
||
temperature = invocation_params.get("temperature", None)
|
||
|
||
trace_data = {
|
||
"trace_id": str(run_id),
|
||
"session_id": self.session_id,
|
||
"timestamp_start": time.time(),
|
||
"request": {
|
||
"model": model_name,
|
||
"messages": messages_serialized,
|
||
"temperature": temperature,
|
||
"invocation_params": {
|
||
k: v for k, v in invocation_params.items()
|
||
if k not in ["messages", "api_key"] # 过滤敏感信息
|
||
},
|
||
},
|
||
}
|
||
|
||
self._traces[str(run_id)] = trace_data
|
||
|
||
except Exception:
|
||
traceback.print_exc()
|
||
|
||
async def on_llm_end(
|
||
self,
|
||
response: LLMResult,
|
||
*,
|
||
run_id: uuid.UUID,
|
||
parent_run_id: Optional[uuid.UUID] = None,
|
||
**kwargs: Any,
|
||
) -> None:
|
||
"""LLM 调用结束 - 记录响应信息并落盘"""
|
||
try:
|
||
trace_id = str(run_id)
|
||
if trace_id not in self._traces:
|
||
return
|
||
|
||
trace_data = self._traces[trace_id]
|
||
timestamp_end = time.time()
|
||
duration_ms = int((timestamp_end - trace_data["timestamp_start"]) * 1000)
|
||
|
||
# 提取响应内容
|
||
generations = response.generations[0] if response.generations else []
|
||
response_messages = []
|
||
|
||
# 提取 finish_reason 和 usage (优先从 message.response_metadata，支持流式)
|
||
finish_reason = "unknown"
|
||
usage_metadata = None
|
||
|
||
for gen in generations:
|
||
msg_dict = {
|
||
"text": gen.text if hasattr(gen, "text") else "",
|
||
}
|
||
|
||
# 提取 message (AIMessage)
|
||
if hasattr(gen, "message") and gen.message:
|
||
msg = gen.message
|
||
msg_dict["message"] = self._serialize_message(msg)
|
||
|
||
# 优先从 message.response_metadata 提取 finish_reason (流式模式)
|
||
if hasattr(msg, "response_metadata") and msg.response_metadata:
|
||
if "finish_reason" in msg.response_metadata and finish_reason == "unknown":
|
||
finish_reason = msg.response_metadata["finish_reason"]
|
||
# 提取流式聚合后的 usage_metadata
|
||
if "usage_metadata" in msg.response_metadata and not usage_metadata:
|
||
usage_metadata = msg.response_metadata["usage_metadata"]
|
||
|
||
# 也从 message.usage_metadata 提取（LangChain 新版本）
|
||
if hasattr(msg, "usage_metadata") and msg.usage_metadata and not usage_metadata:
|
||
usage_metadata = msg.usage_metadata
|
||
|
||
# 其次从 generation_info 提取 finish_reason
|
||
if finish_reason == "unknown" and hasattr(gen, "generation_info") and gen.generation_info:
|
||
if "finish_reason" in gen.generation_info:
|
||
finish_reason = gen.generation_info["finish_reason"]
|
||
|
||
response_messages.append(msg_dict)
|
||
|
||
# 提取 llm_output (包含 usage 等信息)
|
||
llm_output = response.llm_output or {}
|
||
|
||
trace_data["timestamp_end"] = timestamp_end
|
||
trace_data["duration_ms"] = duration_ms
|
||
trace_data["response"] = {
|
||
"generations": response_messages,
|
||
"llm_output": llm_output,
|
||
}
|
||
|
||
# 提取 token usage：优先使用 usage_metadata（流式），其次 llm_output
|
||
if usage_metadata:
|
||
# usage_metadata 通常是字典或对象，尝试兼容两种格式
|
||
if isinstance(usage_metadata, dict):
|
||
trace_data["usage"] = {
|
||
"input_tokens": usage_metadata.get("input_tokens", 0),
|
||
"output_tokens": usage_metadata.get("output_tokens", 0),
|
||
"total_tokens": usage_metadata.get("total_tokens", 0),
|
||
}
|
||
else:
|
||
# 对象格式（有属性访问）
|
||
trace_data["usage"] = {
|
||
"input_tokens": getattr(usage_metadata, "input_tokens", 0),
|
||
"output_tokens": getattr(usage_metadata, "output_tokens", 0),
|
||
"total_tokens": getattr(usage_metadata, "total_tokens", 0),
|
||
}
|
||
else:
|
||
# 回退到 llm_output.token_usage（非流式）
|
||
token_usage = llm_output.get("token_usage", {})
|
||
trace_data["usage"] = {
|
||
"input_tokens": token_usage.get("prompt_tokens", 0),
|
||
"output_tokens": token_usage.get("completion_tokens", 0),
|
||
"total_tokens": token_usage.get("total_tokens", 0),
|
||
}
|
||
|
||
# 最后从 llm_output 兜底 finish_reason（非流式）
|
||
if finish_reason == "unknown" and "finish_reason" in llm_output:
|
||
finish_reason = llm_output["finish_reason"]
|
||
|
||
# 统一字段：finish_reason（向后兼容，保留 stop_reason）
|
||
trace_data["finish_reason"] = finish_reason
|
||
trace_data["stop_reason"] = finish_reason # 向后兼容旧前端代码
|
||
|
||
# 追加写入 JSONL 文件
|
||
log_file = self.raw_logs_dir / f"{self.session_id}.jsonl"
|
||
with open(log_file, "a", encoding="utf-8") as f:
|
||
f.write(json.dumps(trace_data, ensure_ascii=False) + "\n")
|
||
|
||
# 清理已完成的 trace
|
||
del self._traces[trace_id]
|
||
|
||
except Exception:
|
||
traceback.print_exc()
|
||
|
||
async def on_llm_error(
|
||
self,
|
||
error: Exception,
|
||
*,
|
||
run_id: uuid.UUID,
|
||
parent_run_id: Optional[uuid.UUID] = None,
|
||
**kwargs: Any,
|
||
) -> None:
|
||
"""LLM 调用出错 - 记录错误"""
|
||
try:
|
||
trace_id = str(run_id)
|
||
if trace_id not in self._traces:
|
||
return
|
||
|
||
trace_data = self._traces[trace_id]
|
||
timestamp_end = time.time()
|
||
duration_ms = int((timestamp_end - trace_data["timestamp_start"]) * 1000)
|
||
|
||
trace_data["timestamp_end"] = timestamp_end
|
||
trace_data["duration_ms"] = duration_ms
|
||
trace_data["error"] = {
|
||
"type": type(error).__name__,
|
||
"message": str(error),
|
||
}
|
||
|
||
# 写入文件
|
||
log_file = self.raw_logs_dir / f"{self.session_id}.jsonl"
|
||
with open(log_file, "a", encoding="utf-8") as f:
|
||
f.write(json.dumps(trace_data, ensure_ascii=False) + "\n")
|
||
|
||
del self._traces[trace_id]
|
||
|
||
except Exception:
|
||
traceback.print_exc()
|
||
|
||
|
||
def load_raw_traces(session_id: str, raw_logs_dir: Path, limit: int = 10) -> List[Dict[str, Any]]:
|
||
"""读取指定 session 的原始报文记录
|
||
|
||
Args:
|
||
session_id: 会话 ID
|
||
raw_logs_dir: 原始日志目录
|
||
limit: 返回最近 N 条记录
|
||
|
||
Returns:
|
||
原始报文列表，按时间倒序
|
||
"""
|
||
log_file = raw_logs_dir / f"{session_id}.jsonl"
|
||
if not log_file.exists():
|
||
return []
|
||
|
||
traces = []
|
||
try:
|
||
with open(log_file, "r", encoding="utf-8") as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if line:
|
||
traces.append(json.loads(line))
|
||
|
||
# 按时间倒序，返回最近 limit 条
|
||
traces.sort(key=lambda x: x.get("timestamp_start", 0), reverse=True)
|
||
return traces[:limit]
|
||
|
||
except Exception:
|
||
traceback.print_exc()
|
||
return []
|