File size: 3,569 Bytes
75bea1c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
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