graph/state.py

from __future__ import annotations

from typing import Any, TypedDict

from api.schemas import GaiaQuestion


class PlanStep(TypedDict, total=False):
    action: str
    instruction: str
    query: str
    expression: str
    status: str
    result: str


class AgentState(TypedDict, total=False):
    task_id: str
    question: str
    file_name: str | None
    file_path: str | None
    file_type: str
    plan: list[PlanStep]
    step_index: int
    evidence: list[dict[str, Any]]
    tool_results: list[dict[str, Any]]
    draft_answer: str
    final_answer: str
    verification_passed: bool
    feedback: str
    iteration: int
    retry_count: int
    error: str


def initial_state(question: GaiaQuestion) -> AgentState:
    return {
        "task_id": question.task_id,
        "question": question.question,
        "file_name": question.file_name,
        "file_path": None,
        "file_type": "none",
        "plan": [],
        "step_index": 0,
        "evidence": [],
        "tool_results": [],
        "draft_answer": "",
        "final_answer": "",
        "verification_passed": False,
        "feedback": "",
        "iteration": 0,
        "retry_count": 0,
        "error": "",
    }