""" 流式事件处理器 处理 LangGraph 的各种流式事件并转换为前端可消费的格式 """ import json import logging from enum import Enum from typing import Any, Dict, Optional, AsyncGenerator, List from dataclasses import dataclass, field from app.core.timezone import get_now_iso logger = logging.getLogger(__name__) class StreamEventType(str, Enum): """流式事件类型""" # 🔥 LLM 思考相关 - 这些是最重要的！展示 LLM 的大脑活动 LLM_START = "llm_start" # LLM 开始思考 LLM_THOUGHT = "llm_thought" # LLM 思考内容 ⭐ 核心 LLM_DECISION = "llm_decision" # LLM 决策 ⭐ 核心 LLM_ACTION = "llm_action" # LLM 动作 LLM_OBSERVATION = "llm_observation" # LLM 观察结果 LLM_COMPLETE = "llm_complete" # LLM 完成 # LLM Token 流 (实时输出) THINKING_START = "thinking_start" # 开始思考 THINKING_TOKEN = "thinking_token" # 思考 Token (流式) THINKING_END = "thinking_end" # 思考结束 # 工具调用相关 - LLM 决定调用工具 TOOL_CALL_START = "tool_call_start" # 工具调用开始 TOOL_CALL_INPUT = "tool_call_input" # 工具输入参数 TOOL_CALL_OUTPUT = "tool_call_output" # 工具输出结果 TOOL_CALL_END = "tool_call_end" # 工具调用结束 TOOL_CALL_ERROR = "tool_call_error" # 工具调用错误 # 节点相关 NODE_START = "node_start" # 节点开始 NODE_END = "node_end" # 节点结束 # 阶段相关 PHASE_START = "phase_start" PHASE_END = "phase_end" # 发现相关 FINDING_NEW = "finding_new" # 新发现 FINDING_VERIFIED = "finding_verified" # 验证通过 # 状态相关 PROGRESS = "progress" INFO = "info" WARNING = "warning" ERROR = "error" # 任务相关 TASK_START = "task_start" TASK_COMPLETE = "task_complete" TASK_ERROR = "task_error" TASK_CANCEL = "task_cancel" # 心跳 HEARTBEAT = "heartbeat" @dataclass class StreamEvent: """流式事件""" event_type: StreamEventType data: Dict[str, Any] = field(default_factory=dict) timestamp: str = field(default_factory=get_now_iso) sequence: int = 0 # 可选字段 node_name: Optional[str] = None phase: Optional[str] = None tool_name: Optional[str] = None def to_sse(self) -> str: """转换为 SSE 格式""" event_data = { "type": self.event_type.value, "data": self.data, "timestamp": self.timestamp, "sequence": self.sequence, } if self.node_name: event_data["node"] = self.node_name if self.phase: event_data["phase"] = self.phase if self.tool_name: event_data["tool"] = self.tool_name return f"event: {self.event_type.value}\ndata: {json.dumps(event_data, ensure_ascii=False)}\n\n" def to_dict(self) -> Dict[str, Any]: """转换为字典""" return { "event_type": self.event_type.value, "data": self.data, "timestamp": self.timestamp, "sequence": self.sequence, "node_name": self.node_name, "phase": self.phase, "tool_name": self.tool_name, } class StreamHandler: """ 流式事件处理器 最佳实践: 1. 使用 astream_events 捕获所有 LangGraph 事件 2. 将内部事件转换为前端友好的格式 3. 支持多种事件类型的分发 """ def __init__(self, task_id: str): self.task_id = task_id self._sequence = 0 self._current_phase = None self._current_node = None self._thinking_buffer = [] self._tool_states: Dict[str, Dict] = {} def _next_sequence(self) -> int: """获取下一个序列号""" self._sequence += 1 return self._sequence async def process_langgraph_event(self, event: Dict[str, Any]) -> Optional[StreamEvent]: """ 处理 LangGraph 事件 支持的事件类型: - on_chain_start: 链/节点开始 - on_chain_end: 链/节点结束 - on_chain_stream: LLM Token 流 - on_chat_model_start: 模型开始 - on_chat_model_stream: 模型 Token 流 - on_chat_model_end: 模型结束 - on_tool_start: 工具开始 - on_tool_end: 工具结束 - on_custom_event: 自定义事件 """ event_kind = event.get("event", "") event_name = event.get("name", "") event_data = event.get("data", {}) # LLM Token 流 if event_kind == "on_chat_model_stream": return await self._handle_llm_stream(event_data, event_name) # LLM 开始 elif event_kind == "on_chat_model_start": return await self._handle_llm_start(event_data, event_name) # LLM 结束 elif event_kind == "on_chat_model_end": return await self._handle_llm_end(event_data, event_name) # 工具开始 elif event_kind == "on_tool_start": return await self._handle_tool_start(event_name, event_data) # 工具结束 elif event_kind == "on_tool_end": return await self._handle_tool_end(event_name, event_data) # 节点开始 elif event_kind == "on_chain_start" and self._is_node_event(event_name): return await self._handle_node_start(event_name, event_data) # 节点结束 elif event_kind == "on_chain_end" and self._is_node_event(event_name): return await self._handle_node_end(event_name, event_data) # 自定义事件 elif event_kind == "on_custom_event": return await self._handle_custom_event(event_name, event_data) return None def _is_node_event(self, name: str) -> bool: """判断是否是节点事件""" node_names = ["recon", "analysis", "verification", "report", "ReconNode", "AnalysisNode", "VerificationNode", "ReportNode"] return any(n.lower() in name.lower() for n in node_names) async def _handle_llm_start(self, data: Dict, name: str) -> StreamEvent: """处理 LLM 开始事件""" self._thinking_buffer = [] return StreamEvent( event_type=StreamEventType.THINKING_START, sequence=self._next_sequence(), node_name=self._current_node, phase=self._current_phase, data={ "model": name, "message": "🤔 正在思考...", }, ) async def _handle_llm_stream(self, data: Dict, name: str) -> Optional[StreamEvent]: """处理 LLM Token 流事件""" chunk = data.get("chunk") if not chunk: return None # 提取 Token 内容 content = "" if hasattr(chunk, "content"): content = chunk.content elif isinstance(chunk, dict): content = chunk.get("content", "") if not content: return None # 添加到缓冲区 self._thinking_buffer.append(content) return StreamEvent( event_type=StreamEventType.THINKING_TOKEN, sequence=self._next_sequence(), node_name=self._current_node, phase=self._current_phase, data={ "token": content, "accumulated": "".join(self._thinking_buffer), }, ) async def _handle_llm_end(self, data: Dict, name: str) -> StreamEvent: """处理 LLM 结束事件""" full_response = "".join(self._thinking_buffer) self._thinking_buffer = [] # 提取使用的 Token 数 usage = {} output = data.get("output") if output and hasattr(output, "usage_metadata"): usage = { "input_tokens": getattr(output.usage_metadata, "input_tokens", 0), "output_tokens": getattr(output.usage_metadata, "output_tokens", 0), } return StreamEvent( event_type=StreamEventType.THINKING_END, sequence=self._next_sequence(), node_name=self._current_node, phase=self._current_phase, data={ "response": full_response[:2000], # 截断长响应 "usage": usage, "message": "💡 思考完成", }, ) async def _handle_tool_start(self, tool_name: str, data: Dict) -> StreamEvent: """处理工具开始事件""" import time tool_input = data.get("input", {}) # 记录工具状态 self._tool_states[tool_name] = { "start_time": time.time(), "input": tool_input, } return StreamEvent( event_type=StreamEventType.TOOL_CALL_START, sequence=self._next_sequence(), node_name=self._current_node, phase=self._current_phase, tool_name=tool_name, data={ "tool_name": tool_name, "input": self._truncate_data(tool_input), "message": f"🔧 调用工具: {tool_name}", }, ) async def _handle_tool_end(self, tool_name: str, data: Dict) -> StreamEvent: """处理工具结束事件""" import time # 计算执行时间 duration_ms = 0 if tool_name in self._tool_states: start_time = self._tool_states[tool_name].get("start_time", time.time()) duration_ms = int((time.time() - start_time) * 1000) del self._tool_states[tool_name] # 提取输出 output = data.get("output", "") if hasattr(output, "content"): output = output.content return StreamEvent( event_type=StreamEventType.TOOL_CALL_END, sequence=self._next_sequence(), node_name=self._current_node, phase=self._current_phase, tool_name=tool_name, data={ "tool_name": tool_name, "output": self._truncate_data(output), "duration_ms": duration_ms, "message": f"✅ 工具 {tool_name} 完成 ({duration_ms}ms)", }, ) async def _handle_node_start(self, node_name: str, data: Dict) -> StreamEvent: """处理节点开始事件""" self._current_node = node_name # 映射节点到阶段 phase_map = { "recon": "reconnaissance", "analysis": "analysis", "verification": "verification", "report": "reporting", } for key, phase in phase_map.items(): if key in node_name.lower(): self._current_phase = phase break return StreamEvent( event_type=StreamEventType.NODE_START, sequence=self._next_sequence(), node_name=node_name, phase=self._current_phase, data={ "node": node_name, "phase": self._current_phase, "message": f"▶️ 开始节点: {node_name}", }, ) async def _handle_node_end(self, node_name: str, data: Dict) -> StreamEvent: """处理节点结束事件""" # 提取输出信息 output = data.get("output", {}) summary = {} if isinstance(output, dict): # 提取关键信息 if "findings" in output: summary["findings_count"] = len(output["findings"]) if "entry_points" in output: summary["entry_points_count"] = len(output["entry_points"]) if "high_risk_areas" in output: summary["high_risk_areas_count"] = len(output["high_risk_areas"]) if "verified_findings" in output: summary["verified_count"] = len(output["verified_findings"]) return StreamEvent( event_type=StreamEventType.NODE_END, sequence=self._next_sequence(), node_name=node_name, phase=self._current_phase, data={ "node": node_name, "phase": self._current_phase, "summary": summary, "message": f"⏹️ 节点完成: {node_name}", }, ) async def _handle_custom_event(self, event_name: str, data: Dict) -> StreamEvent: """处理自定义事件""" # 映射自定义事件名到事件类型 event_type_map = { "finding": StreamEventType.FINDING_NEW, "finding_verified": StreamEventType.FINDING_VERIFIED, "progress": StreamEventType.PROGRESS, "warning": StreamEventType.WARNING, "error": StreamEventType.ERROR, } event_type = event_type_map.get(event_name, StreamEventType.INFO) return StreamEvent( event_type=event_type, sequence=self._next_sequence(), node_name=self._current_node, phase=self._current_phase, data=data, ) def _truncate_data(self, data: Any, max_length: int = 1000) -> Any: """截断数据""" if isinstance(data, str): return data[:max_length] + "..." if len(data) > max_length else data elif isinstance(data, dict): return {k: self._truncate_data(v, max_length // 2) for k, v in list(data.items())[:10]} elif isinstance(data, list): return [self._truncate_data(item, max_length // len(data)) for item in data[:10]] else: return str(data)[:max_length] def create_progress_event( self, current: int, total: int, message: Optional[str] = None, ) -> StreamEvent: """创建进度事件""" percentage = (current / total * 100) if total > 0 else 0 return StreamEvent( event_type=StreamEventType.PROGRESS, sequence=self._next_sequence(), node_name=self._current_node, phase=self._current_phase, data={ "current": current, "total": total, "percentage": round(percentage, 1), "message": message or f"进度: {current}/{total}", }, ) def create_finding_event( self, finding: Dict[str, Any], is_verified: bool = False, ) -> StreamEvent: """创建发现事件""" event_type = StreamEventType.FINDING_VERIFIED if is_verified else StreamEventType.FINDING_NEW return StreamEvent( event_type=event_type, sequence=self._next_sequence(), node_name=self._current_node, phase=self._current_phase, data={ "title": finding.get("title", "Unknown"), "severity": finding.get("severity", "medium"), "vulnerability_type": finding.get("vulnerability_type", "other"), "file_path": finding.get("file_path"), "line_start": finding.get("line_start"), "is_verified": is_verified, "message": f"{'✅ 已验证' if is_verified else '🔍 新发现'}: [{finding.get('severity', 'medium').upper()}] {finding.get('title', 'Unknown')}", }, ) def create_heartbeat(self) -> StreamEvent: """创建心跳事件""" return StreamEvent( event_type=StreamEventType.HEARTBEAT, sequence=self._sequence, # 心跳不增加序列号 data={"message": "ping"}, )