chuk-ai-session-manager 0.7.1__py3-none-any.whl → 0.8__py3-none-any.whl

chuk_ai_session_manager/models/session_event.py

@@ -2,6 +2,7 @@
 """
 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
@@ -12,32 +13,34 @@ 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
 
-MessageT = TypeVar('MessageT')
+MessageT = TypeVar("MessageT")
+
 
 class SessionEvent(BaseModel, Generic[MessageT]):
     """
     A single event within a session.
-    
+
     Events track all interactions in a session including messages,
     tool calls, summaries, and other activities.
     """
+
     model_config = ConfigDict(arbitrary_types_allowed=True)
-    
+
     id: str = Field(default_factory=lambda: str(uuid4()))
     message: MessageT
     timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
-    
+
     # Make source and type have defaults for backward compatibility with tests
     source: EventSource = Field(default=EventSource.SYSTEM)
     type: EventType = Field(default=EventType.MESSAGE)
-    
+
     metadata: Dict[str, Any] = Field(default_factory=dict)
     parent_event_id: Optional[str] = None
     task_id: Optional[str] = None
-    
+
     # Token tracking
     token_usage: Optional[TokenUsage] = None
-    
+
     @classmethod
     async def create_with_tokens(
         cls,
@@ -47,11 +50,11 @@ class SessionEvent(BaseModel, Generic[MessageT]):
         model: str = "gpt-3.5-turbo",
         source: EventSource = EventSource.SYSTEM,
         type: EventType = EventType.MESSAGE,
-        **kwargs
+        **kwargs,
     ) -> SessionEvent[MessageT]:
         """
         Create a new SessionEvent with automatic token counting.
-        
+
         Args:
             message: The message content
             prompt: The prompt text for token counting
@@ -60,43 +63,37 @@ class SessionEvent(BaseModel, Generic[MessageT]):
             source: The event source
             type: The event type
             **kwargs: Additional fields for the event
-            
+
         Returns:
             A new SessionEvent instance with token usage calculated
         """
         # Create token usage
         token_usage = await TokenUsage.from_text(
-            prompt=prompt,
-            completion=completion,
-            model=model
+            prompt=prompt, completion=completion, model=model
         )
-        
+
         # Create the event
         event = cls(
-            message=message,
-            source=source,
-            type=type,
-            token_usage=token_usage,
-            **kwargs
+            message=message, source=source, type=type, token_usage=token_usage, **kwargs
         )
-        
+
         return event
-    
+
     async def update_token_usage(
-        self, 
+        self,
         prompt: Optional[str] = None,
         completion: Optional[str] = None,
         prompt_tokens: Optional[int] = None,
         completion_tokens: Optional[int] = None,
-        model: str = "gpt-3.5-turbo"
+        model: str = "gpt-3.5-turbo",
     ) -> None:
         """
         Update the token usage for this event.
-        
+
         This method supports two modes:
         1. Pass prompt/completion strings to calculate tokens
         2. Pass prompt_tokens/completion_tokens directly
-        
+
         Args:
             prompt: Optional prompt text to calculate tokens from
             completion: Optional completion text to calculate tokens from
@@ -107,139 +104,143 @@ class SessionEvent(BaseModel, Generic[MessageT]):
         if prompt is not None or completion is not None:
             # Calculate tokens from text
             self.token_usage = await TokenUsage.from_text(
-                prompt=prompt or "",
-                completion=completion,
-                model=model
+                prompt=prompt or "", completion=completion, model=model
             )
         elif prompt_tokens is not None or completion_tokens is not None:
             # Use provided token counts
             if not self.token_usage:
                 self.token_usage = TokenUsage(model=model)
-            
+
             if prompt_tokens is not None:
                 self.token_usage.prompt_tokens = prompt_tokens
             if completion_tokens is not None:
                 self.token_usage.completion_tokens = completion_tokens
-            
+
             # Update total
-            self.token_usage.total_tokens = self.token_usage.prompt_tokens + self.token_usage.completion_tokens
-            
+            self.token_usage.total_tokens = (
+                self.token_usage.prompt_tokens + self.token_usage.completion_tokens
+            )
+
             # Recalculate cost
-            self.token_usage.estimated_cost_usd = self.token_usage._calculate_cost_sync()
-    
+            self.token_usage.estimated_cost_usd = (
+                self.token_usage._calculate_cost_sync()
+            )
+
     async def set_metadata(self, key: str, value: Any) -> None:
         """
         Set a metadata value asynchronously.
-        
+
         Args:
             key: The metadata key
             value: The value to set
         """
         self.metadata[key] = value
-    
+
     async def get_metadata(self, key: str, default: Any = None) -> Any:
         """
         Get a metadata value asynchronously.
-        
+
         Args:
             key: The metadata key
             default: Default value if key not found
-            
+
         Returns:
             The metadata value or default
         """
         return self.metadata.get(key, default)
-    
+
     async def has_metadata(self, key: str) -> bool:
         """
         Check if a metadata key exists asynchronously.
-        
+
         Args:
             key: The metadata key to check
-            
+
         Returns:
             True if the key exists
         """
         return key in self.metadata
-    
+
     async def remove_metadata(self, key: str) -> None:
         """
         Remove a metadata key-value pair asynchronously.
-        
+
         Args:
             key: The metadata key to remove
         """
         if key in self.metadata:
             del self.metadata[key]
-    
+
     async def update_metadata(self, key: str, value: Any) -> None:
         """
         Update or add a metadata value asynchronously.
-        
+
         Args:
             key: The metadata key
             value: The new value
         """
         self.metadata[key] = value
-    
+
     async def merge_metadata(self, new_metadata: Dict[str, Any]) -> None:
         """
         Merge new metadata with existing metadata asynchronously.
-        
+
         Args:
             new_metadata: Dictionary of metadata to merge
         """
         self.metadata.update(new_metadata)
-    
+
     async def clear_metadata(self) -> None:
         """Clear all metadata asynchronously."""
         self.metadata.clear()
-    
+
     async def calculate_tokens(self, model: str = "gpt-3.5-turbo") -> int:
         """
         Calculate tokens for this event's message asynchronously.
-        
+
         Args:
             model: The model to use for token counting
-            
+
         Returns:
             The number of tokens in the message
         """
         if self.token_usage:
             return self.token_usage.total_tokens
-        
+
         # Calculate tokens from message
-        message_str = str(self.message) if not isinstance(self.message, str) else self.message
+        message_str = (
+            str(self.message) if not isinstance(self.message, str) else self.message
+        )
         return await TokenUsage.count_tokens(message_str, model)
-    
+
     def is_child_of(self, parent_event_id: str) -> bool:
         """
         Check if this event is a child of another event.
-        
+
         Args:
             parent_event_id: The ID of the potential parent event
-            
+
         Returns:
             True if this event is a child of the specified event
         """
         return self.parent_event_id == parent_event_id
-    
+
     def is_part_of_task(self, task_id: str) -> bool:
         """
         Check if this event is part of a specific task.
-        
+
         Args:
             task_id: The task ID to check
-            
+
         Returns:
             True if this event is part of the specified task
         """
         return self.task_id == task_id
-    
+
     async def to_dict(self) -> Dict[str, Any]:
         """
         Convert the event to a dictionary asynchronously.
-        
+
         Returns:
             Dictionary representation of the event
         """
@@ -249,22 +250,22 @@ class SessionEvent(BaseModel, Generic[MessageT]):
             "timestamp": self.timestamp.isoformat(),
             "source": self.source.value,
             "type": self.type.value,
-            "metadata": self.metadata
+            "metadata": self.metadata,
         }
-        
+
         if self.parent_event_id:
             result["parent_event_id"] = self.parent_event_id
-        
+
         if self.task_id:
             result["task_id"] = self.task_id
-        
+
         if self.token_usage:
             result["token_usage"] = {
                 "prompt_tokens": self.token_usage.prompt_tokens,
                 "completion_tokens": self.token_usage.completion_tokens,
                 "total_tokens": self.token_usage.total_tokens,
                 "model": self.token_usage.model,
-                "estimated_cost_usd": self.token_usage.estimated_cost_usd
+                "estimated_cost_usd": self.token_usage.estimated_cost_usd,
             }
-        
-        return result
+
+        return result

chuk_ai_session_manager/models/session_metadata.py

@@ -7,12 +7,13 @@ 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
@@ -21,17 +22,15 @@ class SessionMetadata(BaseModel):
     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:
+    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 {}
-        )
+        return cls(created_at=now, updated_at=now, properties=properties or {})

chuk_ai_session_manager/models/session_run.py

@@ -2,6 +2,7 @@
 """
 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
@@ -12,6 +13,7 @@ from pydantic import BaseModel, Field, ConfigDict
 
 class RunStatus(str, Enum):
     """Status of a session run."""
+
     PENDING = "pending"
     RUNNING = "running"
     COMPLETED = "completed"
@@ -21,6 +23,7 @@ class RunStatus(str, Enum):
 
 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()))
@@ -28,15 +31,14 @@ class SessionRun(BaseModel):
     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
+    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 {}
-        )
+        return cls(status=RunStatus.PENDING, metadata=metadata or {})
 
     async def mark_running(self) -> None:
         """Mark the run as started/running asynchronously."""
@@ -61,43 +63,40 @@ class SessionRun(BaseModel):
         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]:  
+
+    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
-        ]
-        
+        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 = {
@@ -105,11 +104,11 @@ class SessionRun(BaseModel):
             "status": self.status.value,
             "started_at": self.started_at.isoformat(),
             "metadata": self.metadata,
-            "tool_calls": self.tool_calls
+            "tool_calls": self.tool_calls,
         }
-        
+
         if self.ended_at:
             result["ended_at"] = self.ended_at.isoformat()
             result["duration"] = await self.get_duration()
-            
-        return result
+
+        return result