"""原始报文录制 - 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 []