src/agent/models.py

from __future__ import annotations
"""Data models for the agent."""

from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any


class IntentType(str, Enum):
    """Types of user intents."""

    FACTUAL_QUERY = "factual_query"
    CURRENT_EVENTS = "current_events"
    RESEARCH = "research"
    CALCULATION = "calculation"
    TEMPORAL = "temporal"
    HOW_TO = "how_to"
    COMPARISON = "comparison"
    OPINION_SEEKING = "opinion_seeking"
    CLARIFICATION = "clarification"
    CHITCHAT = "chitchat"


class WorkflowStrategy(str, Enum):
    """Workflow execution strategies."""

    DIRECT_ANSWER = "direct_answer"
    SINGLE_SEARCH = "single_search"
    MULTI_SEARCH = "multi_search"
    SEARCH_AND_SCRAPE = "search_and_scrape"
    ITERATIVE_RESEARCH = "iterative_research"
    CALCULATION_BASED = "calculation_based"
    COMPARISON_ANALYSIS = "comparison_analysis"


@dataclass
class Intent:
    """Parsed intent from user query."""

    intent_type: IntentType
    confidence: float
    secondary_intents: list[IntentType] = field(default_factory=list)
    entities: dict[str, Any] = field(default_factory=dict)
    requires_web_search: bool = False
    complexity: str = "simple"  # simple, moderate, complex


@dataclass
class PlanStep:
    """A single step in the execution plan."""

    step_number: int
    action: str
    tool: str | None
    parameters: dict[str, Any]
    purpose: str
    depends_on: list[int] = field(default_factory=list)


@dataclass
class ExecutionPlan:
    """Plan for executing a user query."""

    strategy: WorkflowStrategy
    reasoning: str
    steps: list[PlanStep]
    max_iterations: int = 5
    fallback_strategy: WorkflowStrategy | None = None
    success_criteria: str = ""


@dataclass
class ThoughtStep:
    """A thought step in the reasoning process."""

    iteration: int
    thought: str
    action: str
    action_input: dict[str, Any]
    observation: str | None = None
    timestamp: datetime = field(default_factory=datetime.now)


@dataclass
class Citation:
    """A source citation."""

    title: str
    url: str
    snippet: str = ""


@dataclass
class AgentResponse:
    """Final response from the agent."""

    answer: str
    confidence: float
    sources: list[Citation] = field(default_factory=list)
    reasoning_summary: str | None = None
    follow_up_questions: list[str] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def formatted_answer(self) -> str:
        """Get answer with source citations appended."""
        if not self.sources:
            return self.answer

        citations = "\n\n**Sources:**\n"
        for i, source in enumerate(self.sources, 1):
            citations += f"[{i}] [{source.title}]({source.url})\n"

        return self.answer + citations


@dataclass
class AgentState:
    """Current state of the agent execution."""

    query: str
    intent: Intent | None = None
    plan: ExecutionPlan | None = None
    current_step: int = 0
    iteration: int = 0
    thought_history: list[ThoughtStep] = field(default_factory=list)
    tool_results: list[dict[str, Any]] = field(default_factory=list)
    working_memory: dict[str, Any] = field(default_factory=dict)
    error: str | None = None

    @property
    def is_complete(self) -> bool:
        """Check if execution is complete."""
        if self.error:
            return True
        if self.plan and self.current_step >= len(self.plan.steps):
            return True
        return False