151 lines
4.9 KiB
Python
151 lines
4.9 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 backend.app.logger import info
|
||
from ._utils import dispatch_custom_event, make_react_event
|
||
|
||
|
||
def wrap_subgraph_for_error_handling(subgraph, name: str):
|
||
"""
|
||
包装子图，使其错误能传递给主图
|
||
|
||
Args:
|
||
subgraph: 编译好的子图
|
||
name: 子图名称（用于错误标识）
|
||
|
||
Returns: 包装后的节点函数
|
||
"""
|
||
async def wrapped_node(state: MainGraphState, config: Optional[RunnableConfig] = None) -> MainGraphState:
|
||
# 发送子图开始事件
|
||
try:
|
||
await dispatch_custom_event(
|
||
"react_reasoning",
|
||
make_react_event(
|
||
state.reasoning_step,
|
||
f"{name}_subgraph_start",
|
||
1.0,
|
||
f"开始执行 {name} 子图..."
|
||
),
|
||
config
|
||
)
|
||
except Exception as e:
|
||
info(f"[{name}_subgraph] 无法发送开始事件: {e}")
|
||
|
||
try:
|
||
# 调用子图（异步，传 config）
|
||
result = await subgraph.ainvoke(state, config=config)
|
||
|
||
# 更新主图状态
|
||
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()
|
||
})
|
||
|
||
# 发送子图完成事件
|
||
try:
|
||
await dispatch_custom_event(
|
||
"react_reasoning",
|
||
make_react_event(
|
||
state.reasoning_step,
|
||
f"{name}_subgraph_complete",
|
||
1.0,
|
||
f"{name} 子图执行完成"
|
||
),
|
||
config
|
||
)
|
||
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
|
||
|
||
# 发送子图错误事件
|
||
try:
|
||
await dispatch_custom_event(
|
||
"react_reasoning",
|
||
make_react_event(
|
||
state.reasoning_step,
|
||
f"{name}_subgraph_error",
|
||
1.0,
|
||
f"{name} 子图执行失败: {str(e)}"
|
||
),
|
||
config
|
||
)
|
||
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"
|
||
),
|
||
}
|