generated from kgod/ai-review-template
90 lines
2.2 KiB
Python
90 lines
2.2 KiB
Python
"""图定义"""
|
|
from langgraph.graph import StateGraph, END
|
|
from .state import DetailAnalysisState
|
|
from .nodes import (
|
|
open_list_page,
|
|
click_first_job,
|
|
find_detail_area,
|
|
extract_field_selectors,
|
|
validate_data
|
|
)
|
|
|
|
|
|
def check_error(state: DetailAnalysisState) -> str:
|
|
"""检查是否有错误"""
|
|
if state.get("error"):
|
|
return "error"
|
|
return "continue"
|
|
|
|
|
|
def check_need_find_area(state: DetailAnalysisState) -> str:
|
|
"""检查是否需要找详情区域"""
|
|
# in_page 需要找弹窗区域，redirect/new_tab 直接用整个页面
|
|
if state.get("change_type") == "in_page":
|
|
return "find_area"
|
|
return "skip"
|
|
|
|
|
|
def check_validation(state: DetailAnalysisState) -> str:
|
|
"""检查验证结果"""
|
|
if state.get("is_valid"):
|
|
return "success"
|
|
|
|
retry_count = state.get("retry_count", 0)
|
|
if retry_count >= 3:
|
|
return "max_retry"
|
|
|
|
return "retry"
|
|
|
|
|
|
def create_graph():
|
|
"""创建流程图"""
|
|
graph = StateGraph(DetailAnalysisState)
|
|
|
|
# 添加节点
|
|
graph.add_node("open_list_page", open_list_page)
|
|
graph.add_node("click_first_job", click_first_job)
|
|
graph.add_node("find_detail_area", find_detail_area)
|
|
graph.add_node("extract_field_selectors", extract_field_selectors)
|
|
graph.add_node("validate_data", validate_data)
|
|
|
|
# 入口
|
|
graph.set_entry_point("open_list_page")
|
|
|
|
# 流程
|
|
graph.add_conditional_edges(
|
|
"open_list_page",
|
|
check_error,
|
|
{"error": END, "continue": "click_first_job"}
|
|
)
|
|
|
|
# 根据 change_type 决定是否找详情区域
|
|
graph.add_conditional_edges(
|
|
"click_first_job",
|
|
check_need_find_area,
|
|
{
|
|
"find_area": "find_detail_area",
|
|
"skip": "extract_field_selectors"
|
|
}
|
|
)
|
|
|
|
graph.add_conditional_edges(
|
|
"find_detail_area",
|
|
check_error,
|
|
{"error": END, "continue": "extract_field_selectors"}
|
|
)
|
|
|
|
graph.add_edge("extract_field_selectors", "validate_data")
|
|
|
|
graph.add_conditional_edges(
|
|
"validate_data",
|
|
check_validation,
|
|
{
|
|
"success": END,
|
|
"max_retry": END,
|
|
"retry": "extract_field_selectors"
|
|
}
|
|
)
|
|
|
|
return graph.compile()
|