|
|
| from langgraph.graph import StateGraph, START, END |
| from langgraph.checkpoint.memory import InMemorySaver |
|
|
| from .state import AgentState |
| from .extractor import extract_parameters |
| from .selector import select_tool |
| from .validator import validate_tool |
| from .clarification import ask_user |
| from .executor import execute_tool |
| from .response import generate_response |
|
|
|
|
| def route_after_validation(state: AgentState): |
| if state["selected_tool"] is None: |
| return "no_tool" |
|
|
| if state["missing_params"]: |
| return "ask_user" |
|
|
| return "execute_tool" |
|
|
|
|
| builder = StateGraph(AgentState) |
|
|
| |
| builder.add_node("extract", extract_parameters) |
| builder.add_node("select_tool", select_tool) |
| builder.add_node("validate", validate_tool) |
| builder.add_node("ask_user", ask_user) |
| builder.add_node("execute_tool", execute_tool) |
| builder.add_node("generate_response", generate_response) |
|
|
| |
| builder.add_edge(START, "extract") |
| builder.add_edge("extract", "select_tool") |
| builder.add_edge("select_tool", "validate") |
|
|
| |
| builder.add_conditional_edges( |
| "validate", |
| route_after_validation, |
| { |
| "no_tool": END, |
| "ask_user": "ask_user", |
| "execute_tool": "execute_tool" |
| } |
| ) |
|
|
| |
| builder.add_edge("ask_user", END) |
| builder.add_edge("execute_tool", "generate_response") |
| builder.add_edge("generate_response", END) |
|
|
| graph = builder.compile( |
| checkpointer=InMemorySaver() |
| ) |
|
|
