from langgraph.graph import StateGraph

def build_graph():
    graph = StateGraph(dict)

    def llm_step(state):
        question = state.get("question")
        if not question:
            raise ValueError("State is missing 'question'")
        # This is where you integrate your LLM call.
        llm_output = f"Dummy answer to: {question}"
        state["llm_output"] = llm_output
        return state

    def formatter_step(state):
        llm_output = state.get("llm_output")
        if not llm_output:
            raise ValueError("State is missing 'llm_output'")
        # For now, we just pass it through.
        state["final_answer"] = llm_output.strip()
        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