import logfire
from deep_research_agent.nodes import (
    BeginResearch,
    FinalReport,
    Researcher,
    Supervisor,
    WriteResearchBrief,
)
from deep_research_agent.state import ResearchState
from pydantic_graph import Graph

logfire.configure()
logfire.instrument_pydantic_ai()


async def query_agent(research_query: str):
    state = ResearchState()
    graph = Graph(
        nodes=(BeginResearch, WriteResearchBrief, Supervisor, Researcher, FinalReport)
    )

    result = await graph.run(
        start_node=BeginResearch(query=research_query),
        state=state,
    )

    return result.output