state.py

from dataclasses import dataclass, field
from typing import List, Dict, Any, Optional
import json
from dataclasses import dataclass, field, asdict
from langchain.schema import SystemMessage, HumanMessage, AIMessage, BaseMessage


@dataclass
class AgentState:
    """Single source‑of‑truth context for one user query run."""

    user_question: str
    task_id: Optional[str] = None
    messages: List[BaseMessage] = field(default_factory=list)

    next_action: Optional[str] = None      # wiki | ocr | audio | final
    query: Optional[str] = None            # wiki search term
    tool_calls: int = 0

    final_answer: Optional[str] = None

    def add(self, *msgs: BaseMessage):
        self.messages.extend(msgs)