Keerthisujana's picture
Initial deploy of AI ResearchPilot
6a5d1ee
Raw
History Blame Contribute Delete
2.95 kB
"""
graph.py
---------
Wires the five agents into a real LangGraph StateGraph implementing:
Planner -> Retriever -> Evidence -> Reasoning -> Verification
^ |
|__________ (insufficient, iterate) ___|
|
(sufficient / max iters)
v
Final Answer
This is the "PLAN -> RETRIEVE -> CHECK -> REASON -> VERIFY -> RETRIEVE AGAIN
IF NECESSARY -> ANSWER" loop from the spec, expressed as an actual
conditional edge in LangGraph rather than a fixed linear pipeline.
"""
from __future__ import annotations
from langgraph.graph import END, StateGraph
from src.agents.evidence import evaluate_evidence
from src.agents.final_answer import finalize
from src.agents.planner import plan
from src.agents.reasoning import reason
from src.agents.retrieval import retrieve
from src.agents.state import ResearchState
from src.agents.verification import verify
def _route_after_verification(state: ResearchState) -> str:
if state.get("error"):
return "end"
if state.get("verification_passed"):
return "finalize"
# Not passed but verify() only leaves verification_passed=False when a retry is queued
return "retrieve"
def _route_after_plan(state: ResearchState) -> str:
return "end" if state.get("error") else "retrieve"
def build_graph():
graph = StateGraph(ResearchState)
graph.add_node("plan", plan)
graph.add_node("retrieve", retrieve)
graph.add_node("evaluate_evidence", evaluate_evidence)
graph.add_node("reason", reason)
graph.add_node("verify", verify)
graph.add_node("finalize", finalize)
graph.set_entry_point("plan")
graph.add_conditional_edges("plan", _route_after_plan, {"retrieve": "retrieve", "end": END})
graph.add_edge("retrieve", "evaluate_evidence")
graph.add_edge("evaluate_evidence", "reason")
graph.add_edge("reason", "verify")
graph.add_conditional_edges(
"verify", _route_after_verification, {"retrieve": "retrieve", "finalize": "finalize", "end": END}
)
graph.add_edge("finalize", END)
return graph.compile()
_COMPILED_GRAPH = None
def get_graph():
global _COMPILED_GRAPH
if _COMPILED_GRAPH is None:
_COMPILED_GRAPH = build_graph()
return _COMPILED_GRAPH
def run_research(question: str, max_iterations: int = 3) -> ResearchState:
initial_state: ResearchState = {
"question": question,
"max_iterations": max_iterations,
"iteration": 0,
"trace": [],
}
graph = get_graph()
# recursion_limit accounts for the loop: each iteration touches ~4 nodes
result = graph.invoke(initial_state, config={"recursion_limit": 6 * max_iterations + 10})
return result