graph/builder.py

from __future__ import annotations

from functools import partial

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import END, START, StateGraph

from api.client import GaiaApiClient
from config import Settings
from graph.routes import route_after_execute, route_after_verify
from graph.state import AgentState
from models.openrouter_model import OpenRouterModels
from nodes.answer_generator import generate_answer
from nodes.formatter import format_answer
from nodes.inspector import inspect_task
from nodes.planner import create_plan, revise_plan
from nodes.tool_executor import execute_next_tool
from nodes.verifier import verify_answer


def build_graph(settings: Settings, api_client: GaiaApiClient):
    models = OpenRouterModels(settings)
    builder = StateGraph(AgentState)

    builder.add_node("inspect", partial(inspect_task, settings=settings, api_client=api_client))
    builder.add_node("plan", partial(create_plan, model=models.reasoning))
    builder.add_node("execute", partial(execute_next_tool, settings=settings, models=models))
    builder.add_node("generate", partial(generate_answer, model=models.reasoning))
    builder.add_node("verify", partial(verify_answer, model=models.reasoning))
    builder.add_node("replan", partial(revise_plan, model=models.reasoning))
    builder.add_node("format", format_answer)

    builder.add_edge(START, "inspect")
    builder.add_edge("inspect", "plan")
    builder.add_edge("plan", "execute")
    builder.add_conditional_edges(
        "execute",
        partial(route_after_execute, settings=settings),
        {"execute": "execute", "generate": "generate"},
    )
    builder.add_edge("generate", "verify")
    builder.add_conditional_edges(
        "verify",
        partial(route_after_verify, settings=settings),
        {"format": "format", "replan": "replan"},
    )
    builder.add_edge("replan", "execute")
    builder.add_edge("format", END)

    return builder.compile(checkpointer=InMemorySaver())