from gaia_agent.graph import build_graph
from gaia_agent.observability import trace_agent_run


class GaiaAgent:
    def __init__(self, llm=None):
        self.llm = llm
        print("GaiaAgent initialized.")

    def __call__(
        self,
        question: str,
        *,
        session_id: str | None = None,
        user_id: str | None = None,
        task_id: str | None = None,
        file_name: str | None = None,
        file_path: str | None = None,
        level: str | None = None,
    ) -> str:
        print(f"Agent received question (first 80 chars): {question[:80]}...")
        with trace_agent_run(
            question,
            session_id=session_id,
            user_id=user_id,
            task_id=task_id,
        ) as trace:
            graph = build_graph(trace=trace, llm=self.llm)
            initial_state = {
                "question": question,
                "task_id": task_id or "",
                "file_name": file_name or "",
                "level": level or "",
            }
            if file_path:
                initial_state["file_path"] = file_path
            result = graph.invoke(initial_state)
        final_answer = result["final_answer"]
        print(f"Agent returning answer: {final_answer}")
        return final_answer