from langgraph.graph import StateGraph

class AgentState(dict):
    """
    Holds the agent state passed between graph nodes.
    Always includes:
    - question: str
    - llm_output: str (added after LLM node)
    - final_answer: str (added after formatter)
    """
    def __init__(self, initial_data):
        super().__init__()
        if "question" not in initial_data:
            raise ValueError("AgentState requires a 'question' key at init")
        self.update(initial_data)

from llm_node import llm_node
from formatter_node import formatter_node

def build_graph():
    graph = StateGraph(AgentState)

    def llm_step(state):
        question = state.get("question")
        if not question:
            raise ValueError("Missing 'question' in state")
        llm_output = llm_node(question)
        state["llm_output"] = llm_output
        return state

    def formatter_step(state):
        llm_output = state.get("llm_output")
        if not llm_output:
            raise ValueError("Missing 'llm_output' in state")
        final_answer = formatter_node(llm_output)
        state["final_answer"] = final_answer
        return state

    graph.add_node("llm", llm_step)
    graph.add_node("formatter", formatter_step)
    graph.set_entry_point("llm")  # ✅ no input_key
    graph.add_edge("llm", "formatter")
    graph.set_finish_point("formatter")

    return graph