File size: 2,040 Bytes
8a1fc38 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """
MathGraph Definition
Builds the StateGraph and connects nodes with conditional routing.
"""
import logging
from langgraph.graph import StateGraph, START, END
from backend.src.graph.state import MathAgentState
from backend.src.graph.nodes.nodes import (
classify_node,
retrieve_rag_node,
retrieve_memory_node,
solve_node,
verify_node,
format_node
)
logger = logging.getLogger("math_assistant.graph.math_graph")
def should_retry(state: MathAgentState):
"""Conditional edge logic after verification."""
if state["is_correct"]:
return "format"
elif state["retries"] >= 2:
logger.warning(f"Max retries reached. Forcing format. Feedback: {state.get('verification_feedback')}")
return "format"
else:
logger.info(f"Verification failed. Retrying solve... (Retry {state['retries']})")
return "solve"
# 1. Initialize StateGraph
graph_builder = StateGraph(MathAgentState)
# 2. Add Nodes
graph_builder.add_node("classify", classify_node)
graph_builder.add_node("retrieve_rag", retrieve_rag_node)
graph_builder.add_node("retrieve_memory", retrieve_memory_node)
graph_builder.add_node("solve", solve_node)
graph_builder.add_node("verify", verify_node)
graph_builder.add_node("format", format_node)
# 3. Add Edges
# Start -> Classify
graph_builder.add_edge(START, "classify")
# Classify -> Parallel Retrievals
graph_builder.add_edge("classify", "retrieve_rag")
graph_builder.add_edge("classify", "retrieve_memory")
# Parallel Retrievals -> Solve
# LangGraph waits for all incoming edges to complete before firing 'solve'
graph_builder.add_edge("retrieve_rag", "solve")
graph_builder.add_edge("retrieve_memory", "solve")
# Solve -> Verify
graph_builder.add_edge("solve", "verify")
# Verify -> Conditional Routing
graph_builder.add_conditional_edges(
"verify",
should_retry,
{
"format": "format",
"solve": "solve"
}
)
# Format -> End
graph_builder.add_edge("format", END)
# 4. Compile Graph
math_graph = graph_builder.compile()
|