from langgraph.graph import StateGraph
from langchain_community.tools import TavilySearchResults
from llm_node import llm_node
from formatter_node import formatter_node

def build_graph():
    graph = StateGraph(dict)
    search_tool = TavilySearchResults()

    def search_step(state):
        question = state.get("question")
        if not question:
            raise ValueError("Missing 'question'")
        result = search_tool.run(question)
        state["search_result"] = result
        return state

    def llm_step(state):
        question = state.get("question")
        search_result = state.get("search_result", "")
        answer = llm_node(question, search_result)
        state["llm_output"] = answer
        return state

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

    graph.add_node("search", search_step)
    graph.add_node("llm", llm_step)
    graph.add_node("formatter", formatter_step)

    graph.set_entry_point("search")
    graph.add_edge("search", "llm")
    graph.add_edge("llm", "formatter")
    graph.set_finish_point("formatter")

    return graph