shekkari21/agent-from-scratch

1"""Core data models for the agent framework."""2 3from typing import Literal, Union, List, Dict, Optional, Any4from pydantic import BaseModel, Field5from dataclasses import dataclass, field6import uuid7from datetime import datetime8 9 10class Message(BaseModel):11    """A text message in the conversation."""12    type: Literal["message"] = "message"13    role: Literal["system", "user", "assistant"]14    content: str15 16 17class ToolCall(BaseModel):18    """LLM's request to execute a tool."""19    type: Literal["tool_call"] = "tool_call"20    tool_call_id: str21    name: str22    arguments: dict23 24 25class ToolResult(BaseModel):26    """Result from tool execution."""27    type: Literal["tool_result"] = "tool_result"28    tool_call_id: str29    name: str30    status: Literal["success", "error"]31    content: list32 33 34ContentItem = Union[Message, ToolCall, ToolResult]35 36class ToolConfirmation(BaseModel):37    """User's decision on a pending tool call."""38    39    tool_call_id: str40    approved: bool41    modified_arguments: dict | None = None42    reason: str | None = None  # Reason for rejection (if not approved)43 44class PendingToolCall(BaseModel):45    """A tool call awaiting user confirmation."""46    47    tool_call: ToolCall48    confirmation_message: str49 50class Event(BaseModel):51    """A recorded occurrence during agent execution."""52    id: str = Field(default_factory=lambda: str(uuid.uuid4()))53    execution_id: str54    timestamp: float = Field(default_factory=lambda: datetime.now().timestamp())55    author: str  # "user" or agent name56    content: List[ContentItem] = Field(default_factory=list)57 58 59@dataclass60class ExecutionContext:61    """Central storage for all execution state."""62    63    execution_id: str = field(default_factory=lambda: str(uuid.uuid4()))64    events: List[Event] = field(default_factory=list)65    current_step: int = 066    state: Dict[str, Any] = field(default_factory=dict)67    final_result: Optional[str | BaseModel] = None68    session_id: Optional[str] = None  # Link to session for persistence69    70    def add_event(self, event: Event):71        """Append an event to the execution history."""72        self.events.append(event)73    74    def increment_step(self):75        """Move to the next execution step."""76        self.current_step += 177 78class Session(BaseModel):79    """Container for persistent conversation state across multiple run() calls."""80    81    session_id: str82    user_id: str | None = None83    events: list[Event] = Field(default_factory=list)84    state: dict[str, Any] = Field(default_factory=dict)85    created_at: datetime = Field(default_factory=datetime.now)86    updated_at: datetime = Field(default_factory=datetime.now)87 88from abc import ABC, abstractmethod89 90class BaseSessionManager(ABC):91    """Abstract base class for session management."""92    93    @abstractmethod94    async def create(95        self, 96        session_id: str, 97        user_id: str | None = None98    ) -> Session:99        """Create a new session."""100        pass101    102    @abstractmethod103    async def get(self, session_id: str) -> Session | None:104        """Retrieve a session by ID. Returns None if not found."""105        pass106    107    @abstractmethod108    async def save(self, session: Session) -> None:109        """Persist session changes to storage."""110        pass111    112    async def get_or_create(113        self, 114        session_id: str, 115        user_id: str | None = None116    ) -> Session:117        """Get existing session or create new one."""118        session = await self.get(session_id)119        if session is None:120            session = await self.create(session_id, user_id)121        return session122 123class InMemorySessionManager(BaseSessionManager):124    """In-memory session storage for development and testing."""125    126    def __init__(self):127        self._sessions: dict[str, Session] = {}128    129    async def create(130        self, 131        session_id: str, 132        user_id: str | None = None133    ) -> Session:134        """Create a new session."""135        if session_id in self._sessions:136            raise ValueError(f"Session {session_id} already exists")137        138        session = Session(139            session_id=session_id,140            user_id=user_id141        )142        self._sessions[session_id] = session143        return session144    145    async def get(self, session_id: str) -> Session | None:146        """Retrieve a session by ID."""147        return self._sessions.get(session_id)148    149    async def save(self, session: Session) -> None:150        """Save session to storage."""151        self._sessions[session.session_id] = session