chuk-ai-session-manager 0.1.1__py3-none-any.whl

chuk_ai_session_manager/models/session.py

@@ -0,0 +1,316 @@
+# chuk_ai_session_manager/models/session.py
+"""
+Session model for the chuk session manager with improved async support.
+"""
+from __future__ import annotations
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional, Generic, TypeVar, Union
+from uuid import uuid4
+from pydantic import BaseModel, Field, model_validator
+import asyncio
+
+# Import models that Session depends on
+from chuk_ai_session_manager.models.session_metadata import SessionMetadata
+from chuk_ai_session_manager.models.session_event import SessionEvent
+from chuk_ai_session_manager.models.token_usage import TokenUsage, TokenSummary
+# Import SessionRun and RunStatus directly to avoid circular import
+from chuk_ai_session_manager.models.session_run import SessionRun, RunStatus
+
+MessageT = TypeVar('MessageT')
+
+class Session(BaseModel, Generic[MessageT]):
+    """A standalone conversation session with hierarchical support and async methods."""
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    metadata: SessionMetadata = Field(default_factory=SessionMetadata)
+
+    parent_id: Optional[str] = None
+    child_ids: List[str] = Field(default_factory=list)
+
+    task_ids: List[str] = Field(default_factory=list)
+    runs: List[SessionRun] = Field(default_factory=list)
+    events: List[SessionEvent[MessageT]] = Field(default_factory=list)
+    state: Dict[str, Any] = Field(default_factory=dict)
+    
+    # Token tracking
+    token_summary: TokenSummary = Field(default_factory=TokenSummary)
+
+    @model_validator(mode="after")
+    def _sync_hierarchy(cls, model: Session) -> Session:
+        """After creation, sync this session with its parent in the store.
+        
+        Note: This is synchronous for compatibility with Pydantic. 
+        For async parent syncing, use async_init() after creation.
+        """
+        # This validator will be called during model creation,
+        # but won't actually sync with storage - that requires async
+        return model
+    
+    async def async_init(self) -> None:
+        """
+        Initialize async components of the session.
+        
+        Call this after creating a new session to properly set up
+        parent-child relationships in the async storage.
+        """
+        if self.parent_id:
+            # Import here to avoid circular import
+            from chuk_ai_session_manager.storage import SessionStoreProvider
+            store = SessionStoreProvider.get_store()
+            parent = await store.get(self.parent_id)
+            if parent and self.id not in parent.child_ids:
+                parent.child_ids.append(self.id)
+                await store.save(parent)
+
+    @property
+    def last_update_time(self) -> datetime:
+        """Return timestamp of most recent event, or session creation."""
+        if not self.events:
+            return self.metadata.created_at
+        return max(evt.timestamp for evt in self.events)
+
+    @property
+    def active_run(self) -> Optional[SessionRun]:
+        """Return the currently running SessionRun, if any."""
+        for run in reversed(self.runs):
+            if run.status == RunStatus.RUNNING:
+                return run
+        return None
+    
+    @property
+    def total_tokens(self) -> int:
+        """Return the total number of tokens used in this session."""
+        return self.token_summary.total_tokens
+    
+    @property
+    def total_cost(self) -> float:
+        """Return the total estimated cost of this session."""
+        return self.token_summary.total_estimated_cost_usd
+
+    async def add_child(self, child_id: str) -> None:
+        """Add a child session ID and save the session."""
+        if child_id not in self.child_ids:
+            self.child_ids.append(child_id)
+            # Save the updated session
+            from chuk_ai_session_manager.storage import SessionStoreProvider
+            store = SessionStoreProvider.get_store()
+            await store.save(self)
+
+    async def remove_child(self, child_id: str) -> None:
+        """Remove a child session ID and save the session."""
+        if child_id in self.child_ids:
+            self.child_ids.remove(child_id)
+            # Save the updated session
+            from chuk_ai_session_manager.storage import SessionStoreProvider
+            store = SessionStoreProvider.get_store()
+            await store.save(self)
+
+    async def ancestors(self) -> List[Session]:
+        """Fetch ancestor sessions from store asynchronously."""
+        result: List[Session] = []
+        current = self.parent_id
+        
+        # Import here to avoid circular import
+        from chuk_ai_session_manager.storage import SessionStoreProvider
+        store = SessionStoreProvider.get_store()
+        
+        while current:
+            parent = await store.get(current)
+            if not parent:
+                break
+            result.append(parent)
+            current = parent.parent_id
+        return result
+
+    async def descendants(self) -> List[Session]:
+        """Fetch all descendant sessions from store in DFS order asynchronously."""
+        result: List[Session] = []
+        stack = list(self.child_ids)
+        
+        # Import here to avoid circular import
+        from chuk_ai_session_manager.storage import SessionStoreProvider
+        store = SessionStoreProvider.get_store()
+        
+        while stack:
+            cid = stack.pop()
+            child = await store.get(cid)
+            if child:
+                result.append(child)
+                stack.extend(child.child_ids)
+        return result
+    
+    async def add_event(self, event: SessionEvent[MessageT]) -> None:
+        """
+        Add an event to the session and update token tracking asynchronously.
+        
+        Args:
+            event: The event to add
+        """
+        # Add the event
+        self.events.append(event)
+        
+        # Update token summary if this event has token usage
+        if event.token_usage:
+            await self.token_summary.add_usage(event.token_usage)
+    
+    async def add_event_and_save(self, event: SessionEvent[MessageT]) -> None:
+        """
+        Add an event to the session, update token tracking, and save the session.
+        
+        Args:
+            event: The event to add
+        """
+        # Add the event asynchronously
+        await self.add_event(event)
+        
+        # Save the session
+        from chuk_ai_session_manager.storage import SessionStoreProvider
+        store = SessionStoreProvider.get_store()
+        await store.save(self)
+    
+    async def get_token_usage_by_source(self) -> Dict[str, TokenSummary]:
+        """
+        Get token usage statistics grouped by event source asynchronously.
+        
+        Returns:
+            A dictionary mapping event sources to token summaries
+        """
+        result: Dict[str, TokenSummary] = {}
+        
+        for event in self.events:
+            if not event.token_usage:
+                continue
+                
+            source = event.source.value
+            if source not in result:
+                result[source] = TokenSummary()
+                
+            await result[source].add_usage(event.token_usage)
+            
+        return result
+    
+    async def get_token_usage_by_run(self) -> Dict[str, TokenSummary]:
+        """
+        Get token usage statistics grouped by run asynchronously.
+        
+        Returns:
+            A dictionary mapping run IDs to token summaries
+        """
+        result: Dict[str, TokenSummary] = {}
+        
+        # Add an entry for events without a run
+        result["no_run"] = TokenSummary()
+        
+        for event in self.events:
+            if not event.token_usage:
+                continue
+                
+            run_id = event.task_id or "no_run"
+            if run_id not in result:
+                result[run_id] = TokenSummary()
+                
+            await result[run_id].add_usage(event.token_usage)
+            
+        return result
+    
+    async def count_message_tokens(
+        self, 
+        message: Union[str, Dict[str, Any]], 
+        model: str = "gpt-3.5-turbo"
+    ) -> int:
+        """
+        Count tokens in a message asynchronously.
+        
+        Args:
+            message: The message to count tokens for (string or dict)
+            model: The model to use for counting
+            
+        Returns:
+            The number of tokens in the message
+        """
+        # If message is already a string, count directly
+        if isinstance(message, str):
+            return await TokenUsage.count_tokens(message, model)
+        
+        # If it's a dict (like OpenAI messages), extract content
+        if isinstance(message, dict) and "content" in message:
+            return await TokenUsage.count_tokens(message["content"], model)
+            
+        # If it's some other object, convert to string and count
+        return await TokenUsage.count_tokens(str(message), model)
+    
+    async def set_state(self, key: str, value: Any) -> None:
+        """
+        Set a state value asynchronously.
+        
+        Args:
+            key: The state key to set
+            value: The value to set
+        """
+        self.state[key] = value
+        
+        # Auto-save if needed (could be added as an option)
+        # from chuk_ai_session_manager.storage import SessionStoreProvider
+        # store = SessionStoreProvider.get_store()
+        # await store.save(self)
+    
+    async def get_state(self, key: str, default: Any = None) -> Any:
+        """
+        Get a state value asynchronously.
+        
+        Args:
+            key: The state key to retrieve
+            default: Default value to return if key not found
+            
+        Returns:
+            The state value or default if not found
+        """
+        return self.state.get(key, default)
+    
+    async def has_state(self, key: str) -> bool:
+        """
+        Check if a state key exists asynchronously.
+        
+        Args:
+            key: The state key to check
+            
+        Returns:
+            True if the key exists in state
+        """
+        return key in self.state
+    
+    async def remove_state(self, key: str) -> None:
+        """
+        Remove a state key-value pair asynchronously.
+        
+        Args:
+            key: The state key to remove
+        """
+        if key in self.state:
+            del self.state[key]
+            
+            # Auto-save if needed (could be added as an option)
+            # from chuk_ai_session_manager.storage import SessionStoreProvider
+            # store = SessionStoreProvider.get_store()
+            # await store.save(self)
+
+    @classmethod
+    async def create(cls, parent_id: Optional[str] = None, **kwargs) -> Session:
+        """
+        Create a new session asynchronously, handling parent-child relationships.
+        
+        Args:
+            parent_id: Optional parent session ID
+            **kwargs: Additional arguments for Session initialization
+            
+        Returns:
+            A new Session instance with parent-child relationships set up
+        """
+        session = cls(parent_id=parent_id, **kwargs)
+        await session.async_init()
+        
+        # Save the new session
+        from chuk_ai_session_manager.storage import SessionStoreProvider
+        store = SessionStoreProvider.get_store()
+        await store.save(session)
+        
+        return session

chuk_ai_session_manager/models/session_event.py

@@ -0,0 +1,166 @@
+# chuk_ai_session_manager/models/session_event.py
+"""
+Session event model for the chuk session manager with improved async support.
+"""
+from __future__ import annotations
+from datetime import datetime, timezone
+from typing import Any, Dict, Generic, Optional, TypeVar, Union
+from uuid import uuid4
+from pydantic import BaseModel, Field, ConfigDict
+
+# session manager
+from chuk_ai_session_manager.models.event_source import EventSource
+from chuk_ai_session_manager.models.event_type import EventType
+from chuk_ai_session_manager.models.token_usage import TokenUsage
+
+# Generic type for event message content
+MessageT = TypeVar('MessageT')
+
+class SessionEvent(BaseModel, Generic[MessageT]):
+    """An event in a session."""
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    message: Optional[MessageT] = None
+    task_id: Optional[str] = None
+    type: EventType = EventType.MESSAGE
+    source: EventSource = EventSource.LLM
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    
+    # Field for token usage tracking
+    token_usage: Optional[TokenUsage] = None
+    
+    @classmethod
+    async def create_with_tokens(
+        cls,
+        message: MessageT,
+        prompt: str,
+        completion: Optional[str] = None,
+        model: str = "gpt-3.5-turbo",
+        source: EventSource = EventSource.LLM,
+        type: EventType = EventType.MESSAGE,
+        task_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> SessionEvent[MessageT]:
+        """
+        Create a session event with automatic token counting asynchronously.
+        
+        Args:
+            message: The message content
+            prompt: The prompt text used (for token counting)
+            completion: The completion text (for token counting)
+            model: The model used for this interaction
+            source: The source of this event
+            type: The type of this event
+            task_id: Optional task ID this event is associated with
+            metadata: Optional additional metadata
+            
+        Returns:
+            A new SessionEvent with token usage information
+        """
+        # Use the async method of TokenUsage
+        token_usage = await TokenUsage.from_text(prompt, completion, model)
+        
+        # Create the event
+        event = cls(
+            message=message,
+            task_id=task_id,
+            type=type,
+            source=source,
+            metadata=metadata or {},
+            token_usage=token_usage
+        )
+        
+        return event
+    
+    async def update_token_usage(
+        self, 
+        prompt: Optional[str] = None, 
+        completion: Optional[str] = None,
+        model: Optional[str] = None
+    ) -> None:
+        """
+        Update token usage information for this event.
+        
+        Args:
+            prompt: The prompt text to count tokens for
+            completion: The completion text to count tokens for
+            model: The model to use for counting tokens
+        """
+        # If we don't have token_usage yet, create it
+        if self.token_usage is None:
+            self.token_usage = TokenUsage(model=model or "")
+        
+        # If model is provided, update it
+        if model and not self.token_usage.model:
+            self.token_usage.model = model
+            
+        # Calculate tokens if text is provided
+        if prompt:
+            # Use async method for token counting
+            prompt_tokens = await TokenUsage.count_tokens(prompt, self.token_usage.model)
+            self.token_usage.prompt_tokens = prompt_tokens
+            
+        if completion:
+            # Use async method for token counting
+            completion_tokens = await TokenUsage.count_tokens(completion, self.token_usage.model)
+            self.token_usage.completion_tokens = completion_tokens
+            
+        # Recalculate totals
+        self.token_usage.total_tokens = self.token_usage.prompt_tokens + self.token_usage.completion_tokens
+        if self.token_usage.model:
+            # Use async method for cost calculation
+            self.token_usage.estimated_cost_usd = await self.token_usage.calculate_cost()
+            
+    # Metadata async methods with clean names
+    async def get_metadata(self, key: str, default: Any = None) -> Any:
+        """Get a metadata value.
+        
+        Args:
+            key: The metadata key to retrieve
+            default: Default value to return if key not found
+            
+        Returns:
+            The metadata value or default if not found
+        """
+        return self.metadata.get(key, default)
+
+    async def set_metadata(self, key: str, value: Any) -> None:
+        """Set a metadata value.
+        
+        Args:
+            key: The metadata key to set
+            value: The value to set
+        """
+        self.metadata[key] = value
+        
+    async def has_metadata(self, key: str) -> bool:
+        """Check if a metadata key exists.
+        
+        Args:
+            key: The metadata key to check
+            
+        Returns:
+            True if the key exists in metadata
+        """
+        return key in self.metadata
+        
+    async def remove_metadata(self, key: str) -> None:
+        """Remove a metadata key-value pair.
+        
+        Args:
+            key: The metadata key to remove
+        """
+        if key in self.metadata:
+            del self.metadata[key]
+    
+    # Alternative async method for updating metadata for backward compatibility
+    async def update_metadata(self, key: str, value: Any) -> None:
+        """Update a metadata value (alias for set_metadata).
+        
+        Args:
+            key: The metadata key to set
+            value: The value to set
+        """
+        await self.set_metadata(key, value)

chuk_ai_session_manager/models/session_metadata.py

@@ -0,0 +1,37 @@
+# chuk_ai_session_manager/models/session_metadata.py
+from __future__ import annotations
+from datetime import datetime, timezone
+from typing import Any, Dict, Optional
+from pydantic import BaseModel, Field
+
+
+class SessionMetadata(BaseModel):
+    """Core metadata associated with a session."""
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    
+    # Free-form properties for session-level identifiers and custom data
+    properties: Dict[str, Any] = Field(default_factory=dict)
+    
+    async def set_property(self, key: str, value: Any) -> None:
+        """Add or update a custom metadata property asynchronously."""
+        self.properties[key] = value
+        self.updated_at = datetime.now(timezone.utc)
+
+    async def get_property(self, key: str) -> Any:
+        """Retrieve a metadata property by key asynchronously."""
+        return self.properties.get(key)
+        
+    async def update_timestamp(self) -> None:
+        """Update the updated_at timestamp asynchronously."""
+        self.updated_at = datetime.now(timezone.utc)
+        
+    @classmethod
+    async def create(cls, properties: Optional[Dict[str, Any]] = None) -> SessionMetadata:
+        """Create a new SessionMetadata instance asynchronously."""
+        now = datetime.now(timezone.utc)
+        return cls(
+            created_at=now,
+            updated_at=now,
+            properties=properties or {}
+        )

chuk_ai_session_manager/models/session_run.py

@@ -0,0 +1,115 @@
+# chuk_ai_session_manager/session_run.py
+"""
+Session run model for the chuk session manager with improved async support.
+"""
+from __future__ import annotations
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any, Dict, Optional, List
+from uuid import uuid4
+from pydantic import BaseModel, Field, ConfigDict
+
+
+class RunStatus(str, Enum):
+    """Status of a session run."""
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+class SessionRun(BaseModel):
+    """A single execution or "run" within a session."""
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    started_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    ended_at: Optional[datetime] = None
+    status: RunStatus = RunStatus.PENDING
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    tool_calls: List[str] = Field(default_factory=list)  # IDs of associated tool call events
+
+    @classmethod
+    async def create(cls, metadata: Optional[Dict[str, Any]] = None) -> SessionRun:
+        """Create a new session run asynchronously."""
+        return cls(
+            status=RunStatus.PENDING,
+            metadata=metadata or {}
+        )
+
+    async def mark_running(self) -> None:
+        """Mark the run as started/running asynchronously."""
+        self.status = RunStatus.RUNNING
+        self.started_at = datetime.now(timezone.utc)
+
+    async def mark_completed(self) -> None:
+        """Mark the run as completed successfully asynchronously."""
+        self.status = RunStatus.COMPLETED
+        self.ended_at = datetime.now(timezone.utc)
+
+    async def mark_failed(self, reason: Optional[str] = None) -> None:
+        """Mark the run as failed asynchronously."""
+        self.status = RunStatus.FAILED
+        self.ended_at = datetime.now(timezone.utc)
+        if reason:
+            await self.set_metadata("failure_reason", reason)
+
+    async def mark_cancelled(self, reason: Optional[str] = None) -> None:
+        """Mark the run as cancelled asynchronously."""
+        self.status = RunStatus.CANCELLED
+        self.ended_at = datetime.now(timezone.utc)
+        if reason:
+            await self.set_metadata("cancel_reason", reason)
+        
+    async def set_metadata(self, key: str, value: Any) -> None:
+        """Set a metadata value asynchronously."""
+        self.metadata[key] = value
+        
+    async def get_metadata(self, key: str, default: Any = None) -> Any:
+        """Get a metadata value asynchronously."""
+        return self.metadata.get(key, default)
+        
+    async def has_metadata(self, key: str) -> bool:
+        """Check if a metadata key exists asynchronously."""
+        return key in self.metadata
+        
+    async def remove_metadata(self, key: str) -> None:
+        """Remove a metadata key-value pair asynchronously."""
+        if key in self.metadata:
+            del self.metadata[key]
+        
+    async def get_duration(self) -> Optional[float]:
+        """Get the duration of the run in seconds asynchronously."""
+        if self.ended_at is None:
+            return None
+        return (self.ended_at - self.started_at).total_seconds()
+    
+    async def add_tool_call(self, tool_call_id: str) -> None:
+        """Associate a tool call event with this run asynchronously."""
+        if tool_call_id not in self.tool_calls:
+            self.tool_calls.append(tool_call_id)
+            
+    async def get_tool_calls(self, session: Any) -> List[Any]:  
+        """Get all tool call events associated with this run asynchronously."""
+        # We use Any type to avoid circular imports
+        return [
+            event for event in session.events 
+            if event.id in self.tool_calls
+        ]
+        
+    async def to_dict(self) -> Dict[str, Any]:
+        """Convert the run to a dictionary asynchronously."""
+        result = {
+            "id": self.id,
+            "status": self.status.value,
+            "started_at": self.started_at.isoformat(),
+            "metadata": self.metadata,
+            "tool_calls": self.tool_calls
+        }
+        
+        if self.ended_at:
+            result["ended_at"] = self.ended_at.isoformat()
+            result["duration"] = await self.get_duration()
+            
+        return result