"""
Pydantic models for TAgent state management and structured outputs.
"""

from typing import Dict, Any
from pydantic import BaseModel, Field


class AgentState(BaseModel):
    """Represents agent state as a typed dictionary."""

    data: Dict[str, Any] = {}


class StructuredResponse(BaseModel):
    """Schema for structured outputs generated by LLMs."""

    model_config = {
        "populate_by_name": True
    }  # Allow both 'params' and 'parameters' field names

    action: str = Field(
        description=(
            "Action to be taken, must be one of: "
            "'plan', 'execute', 'summarize', 'evaluate'"
        )
    )
    params: Dict[str, Any] = Field(
        default_factory=dict,
        description="Parameters for the action",
        alias="parameters",
    )
    reasoning: str = Field(default="", description="Reasoning for the action")
