""" Explore 工具 - 并行探索多个方案 启动多个 Sub-Trace 并行执行不同的探索方向，汇总结果返回。 """ import asyncio from typing import List, Optional, Dict, Any from datetime import datetime from .models import Trace, Message from .trace_id import generate_sub_trace_id from .goal_models import Goal async def explore_tool( current_trace_id: str, current_goal_id: str, branches: List[str], background: Optional[str] = None, store=None, run_agent=None ) -> str: """ 并行探索多个方向，汇总结果 Args: current_trace_id: 当前主 Trace ID current_goal_id: 当前 Goal ID branches: 探索方向列表（每个元素是一个探索任务描述） background: 可选，背景信息（如果提供则用作各 Sub-Trace 的初始 context） store: TraceStore 实例 run_agent: 运行 Agent 的函数 Returns: 汇总结果字符串 Example: >>> result = await explore_tool( ... current_trace_id="abc123", ... current_goal_id="2", ... branches=["JWT 方案", "Session 方案"], ... store=store, ... run_agent=run_agent_func ... ) """ if not store: raise ValueError("store parameter is required") if not run_agent: raise ValueError("run_agent parameter is required") # 1. 创建 agent_call Goal goal = Goal( id=current_goal_id, type="agent_call", description=f"并行探索 {len(branches)} 个方案", reason="探索多个可行方案", agent_call_mode="explore", sub_trace_ids=[], status="in_progress" ) # 更新 Goal（标记为 agent_call） await store.update_goal(current_trace_id, current_goal_id, type="agent_call", agent_call_mode="explore", status="in_progress") # 2. 为每个分支创建 Sub-Trace sub_traces = [] sub_trace_ids = [] for i, desc in enumerate(branches): # 生成 Sub-Trace ID sub_trace_id = generate_sub_trace_id(current_trace_id, "explore") # 创建 Sub-Trace sub_trace = Trace( trace_id=sub_trace_id, mode="agent", task=desc, parent_trace_id=current_trace_id, parent_goal_id=current_goal_id, agent_type="explore", context={ "allowed_tools": ["read", "grep", "glob"], # 探索模式：只读权限 "max_turns": 20, "background": background }, status="running", created_at=datetime.now() ) # 保存 Sub-Trace await store.create_trace(sub_trace) sub_traces.append(sub_trace) sub_trace_ids.append(sub_trace_id) # 推送 sub_trace_started 事件 await store.append_event(current_trace_id, "sub_trace_started", { "trace_id": sub_trace_id, "parent_trace_id": current_trace_id, "parent_goal_id": current_goal_id, "agent_type": "explore", "task": desc }) # 更新主 Goal 的 sub_trace_ids await store.update_goal(current_trace_id, current_goal_id, sub_trace_ids=sub_trace_ids) # 3. 并行执行所有 Sub-Traces results = await asyncio.gather( *[run_agent(st, background=background) for st in sub_traces], return_exceptions=True ) # 4. 收集元数据并汇总结果 sub_trace_metadata = {} summary_parts = ["## 探索结果\n"] for i, (sub_trace, result) in enumerate(zip(sub_traces, results), 1): branch_name = chr(ord('A') + i - 1) # A, B, C... if isinstance(result, Exception): # 处理异常情况 summary_parts.append(f"### 方案 {branch_name}: {sub_trace.task}") summary_parts.append(f"⚠️ 执行出错: {str(result)}\n") sub_trace_metadata[sub_trace.trace_id] = { "task": sub_trace.task, "status": "failed", "summary": f"执行出错: {str(result)}", "last_message": None, "stats": { "message_count": 0, "total_tokens": 0, "total_cost": 0.0 } } else: # 获取 Sub-Trace 的最终状态 updated_trace = await store.get_trace(sub_trace.trace_id) # 获取最后一条 assistant 消息 messages = await store.get_trace_messages(sub_trace.trace_id) last_message = None for msg in reversed(messages): if msg.role == "assistant": last_message = msg break # 构建元数据 # 优先使用 result 中的 summary，否则使用最后一条消息的内容 summary_text = None if isinstance(result, dict) and result.get("summary"): summary_text = result.get("summary") elif last_message and last_message.content: # 使用最后一条消息的内容作为 summary（截断至 200 字符） content_text = last_message.content if isinstance(content_text, dict) and "text" in content_text: content_text = content_text["text"] elif not isinstance(content_text, str): content_text = str(content_text) summary_text = content_text[:200] if content_text else "执行完成" else: summary_text = "执行完成" sub_trace_metadata[sub_trace.trace_id] = { "task": sub_trace.task, "status": updated_trace.status if updated_trace else "unknown", "summary": summary_text, "last_message": { "role": last_message.role, "description": last_message.description, "content": last_message.content[:500] if last_message.content else None, "created_at": last_message.created_at.isoformat() } if last_message else None, "stats": { "message_count": updated_trace.total_messages if updated_trace else 0, "total_tokens": updated_trace.total_tokens if updated_trace else 0, "total_cost": updated_trace.total_cost if updated_trace else 0.0 } } # 组装摘要文本 summary_parts.append(f"### 方案 {branch_name}: {sub_trace.task}") if updated_trace and updated_trace.status == "completed": summary_parts.append(f"{summary_text}\n") summary_parts.append(f"📊 统计: {updated_trace.total_messages} 条消息, " f"{updated_trace.total_tokens} tokens, " f"成本 ${updated_trace.total_cost:.4f}\n") else: summary_parts.append(f"未完成\n") # 推送 sub_trace_completed 事件 await store.append_event(current_trace_id, "sub_trace_completed", { "trace_id": sub_trace.trace_id, "status": "completed" if not isinstance(result, Exception) else "failed", "summary": result.get("summary", "") if isinstance(result, dict) else "" }) summary_parts.append("\n---") summary_parts.append(f"已完成 {len(branches)} 个方案的探索，请根据结果选择继续的方向。") summary = "\n".join(summary_parts) # 5. 完成主 Goal，保存元数据 await store.update_goal(current_trace_id, current_goal_id, status="completed", summary=f"探索了 {len(branches)} 个方案", sub_trace_metadata=sub_trace_metadata) return summary def create_explore_tool_schema() -> Dict[str, Any]: """ 创建 explore 工具的 JSON Schema Returns: 工具的 JSON Schema """ return { "type": "function", "function": { "name": "explore", "description": "并行探索多个方向，汇总结果。用于需要对比多个方案或尝试不同实现方式的场景。", "parameters": { "type": "object", "properties": { "branches": { "type": "array", "items": {"type": "string"}, "description": "探索方向列表，每个元素是一个探索任务的描述", "minItems": 2, "maxItems": 5 }, "background": { "type": "string", "description": "可选的背景信息，用于初始化各 Sub-Trace 的上下文" } }, "required": ["branches"] } } }