from langgraph.graph import Graph
from llm_node import llm_node
from formatter_node import formatter_node

def build_graph():
    graph = Graph()

    def llm_step(state):
        question = state["question"]
        result = llm_node().run(question)
        state["llm_output"] = result
        return state

    def formatter_step(state):
        cleaned = formatter_node(state["llm_output"])
        state["final_answer"] = cleaned
        return state

    graph.add_node("llm", llm_step)
    graph.add_node("formatter", formatter_step)
    graph.set_entry_point("llm")
    graph.add_edge("llm", "formatter")
    graph.set_finish_point("formatter")

    return graph