162 lines
5.8 KiB
Python
162 lines
5.8 KiB
Python
"""
|
|
子图包装器 - 为子图添加错误处理和事件追踪
|
|
"""
|
|
|
|
from typing import Dict, Any, Optional
|
|
from datetime import datetime
|
|
|
|
from langchain_core.runnables.config import RunnableConfig
|
|
|
|
from .state import MainGraphState, ErrorRecord, ErrorSeverity
|
|
from ..logger import info
|
|
|
|
|
|
def wrap_subgraph_for_error_handling(subgraph, name: str):
|
|
"""
|
|
包装子图，使其错误能传递给主图
|
|
|
|
Args:
|
|
subgraph: 编译好的子图
|
|
name: 子图名称（用于错误标识）
|
|
|
|
Returns: 包装后的节点函数
|
|
"""
|
|
async def wrapped_node(state: MainGraphState, config: Optional[RunnableConfig] = None) -> MainGraphState:
|
|
# 发送子图开始事件
|
|
if config:
|
|
try:
|
|
from langchain_core.callbacks.manager import adispatch_custom_event
|
|
callbacks = config.get("callbacks")
|
|
if callbacks:
|
|
await adispatch_custom_event(
|
|
"react_reasoning",
|
|
{
|
|
"step": state.reasoning_step,
|
|
"action": f"{name}_subgraph_start",
|
|
"confidence": 1.0,
|
|
"reasoning": f"开始执行 {name} 子图..."
|
|
},
|
|
callbacks=callbacks
|
|
)
|
|
except Exception as e:
|
|
info(f"[{name}_subgraph] 无法发送开始事件: {e}")
|
|
|
|
try:
|
|
# 调用子图
|
|
result = subgraph.invoke(state)
|
|
|
|
# 更新主图状态
|
|
subgraph_result = None
|
|
if name == "contact":
|
|
state.contact_result = result
|
|
subgraph_result = result.get("final_result", "")
|
|
elif name == "dictionary":
|
|
state.dictionary_result = result
|
|
subgraph_result = result.get("final_result", "")
|
|
elif name == "news_analysis":
|
|
state.news_result = result
|
|
subgraph_result = result.get("final_result", "")
|
|
|
|
# 设置最终结果
|
|
if subgraph_result:
|
|
state.final_result = subgraph_result
|
|
else:
|
|
state.final_result = "子图执行完成"
|
|
|
|
# 标记成功
|
|
state.success = True
|
|
state.current_phase = "done"
|
|
state.reasoning_history.append({
|
|
"step": state.reasoning_step,
|
|
"action": "subgraph_completed",
|
|
"confidence": 1.0,
|
|
"reasoning": f"{name}子图执行完成",
|
|
"timestamp": datetime.now().isoformat()
|
|
})
|
|
|
|
# 发送子图完成事件
|
|
if config:
|
|
try:
|
|
from langchain_core.callbacks.manager import adispatch_custom_event
|
|
callbacks = config.get("callbacks")
|
|
if callbacks:
|
|
await adispatch_custom_event(
|
|
"react_reasoning",
|
|
{
|
|
"step": state.reasoning_step,
|
|
"action": f"{name}_subgraph_complete",
|
|
"confidence": 1.0,
|
|
"reasoning": f"{name} 子图执行完成"
|
|
},
|
|
callbacks=callbacks
|
|
)
|
|
except Exception as e:
|
|
info(f"[{name}_subgraph] 无法发送完成事件: {e}")
|
|
|
|
return state
|
|
|
|
except Exception as e:
|
|
# 捕获子图错误，传递给主图
|
|
error_record = ErrorRecord(
|
|
error_type=f"{name}SubgraphError",
|
|
error_message=str(e),
|
|
severity=ErrorSeverity.WARNING,
|
|
source=f"{name}_subgraph",
|
|
timestamp=datetime.now().isoformat(),
|
|
retry_count=0,
|
|
max_retries=1,
|
|
context={"user_query": state.user_query}
|
|
)
|
|
state.errors.append(error_record)
|
|
state.current_error = error_record
|
|
state.current_phase = "error_handling"
|
|
state.success = False
|
|
|
|
# 发送子图错误事件
|
|
if config:
|
|
try:
|
|
from langchain_core.callbacks.manager import adispatch_custom_event
|
|
callbacks = config.get("callbacks")
|
|
if callbacks:
|
|
await adispatch_custom_event(
|
|
"react_reasoning",
|
|
{
|
|
"step": state.reasoning_step,
|
|
"action": f"{name}_subgraph_error",
|
|
"confidence": 1.0,
|
|
"reasoning": f"{name} 子图执行失败: {str(e)}"
|
|
},
|
|
callbacks=callbacks
|
|
)
|
|
except Exception as e:
|
|
info(f"[{name}_subgraph] 无法发送错误事件: {e}")
|
|
|
|
return state
|
|
|
|
return wrapped_node
|
|
|
|
|
|
def create_subgraph_nodes(contact_graph, dictionary_graph, news_analysis_graph) -> Dict[str, Any]:
|
|
"""
|
|
创建所有子图节点的字典
|
|
|
|
Args:
|
|
contact_graph: 联系人子图
|
|
dictionary_graph: 词典子图
|
|
news_analysis_graph: 新闻分析子图
|
|
|
|
Returns:
|
|
子图节点字典 {name: wrapped_node}
|
|
"""
|
|
return {
|
|
"contact_subgraph": wrap_subgraph_for_error_handling(
|
|
contact_graph.compile(), "contact"
|
|
),
|
|
"dictionary_subgraph": wrap_subgraph_for_error_handling(
|
|
dictionary_graph.compile(), "dictionary"
|
|
),
|
|
"news_analysis_subgraph": wrap_subgraph_for_error_handling(
|
|
news_analysis_graph.compile(), "news_analysis"
|
|
),
|
|
}
|