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

chuk_ai_session_manager/session_manager.py

@@ -14,7 +14,7 @@ This module provides the main SessionManager class which offers:
 from __future__ import annotations
 import asyncio
 import logging
-from typing import Any, Dict, List, Optional, Callable, Union
+from typing import Any, Dict, List, Optional, Callable
 from datetime import datetime
 import uuid
 
@@ -22,7 +22,7 @@ from chuk_ai_session_manager.models.session import Session
 from chuk_ai_session_manager.models.session_event import SessionEvent
 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.session_storage import get_backend, ChukSessionsStore
+from chuk_ai_session_manager.session_storage import ChukSessionsStore
 
 logger = logging.getLogger(__name__)
 
@@ -30,10 +30,10 @@ logger = logging.getLogger(__name__)
 class SessionManager:
     """
     High-level session manager for AI conversations.
-    
+
     Provides an easy-to-use interface for tracking conversations, managing
     system prompts, handling infinite context, and monitoring usage.
-    
+
     Examples:
         Basic usage:
         ```python
@@ -41,22 +41,22 @@ class SessionManager:
         await sm.user_says("Hello!")
         await sm.ai_responds("Hi there!", model="gpt-4")
         ```
-        
+
         With system prompt:
         ```python
         sm = SessionManager(system_prompt="You are a helpful assistant.")
         await sm.user_says("What can you do?")
         ```
-        
+
         Infinite context:
         ```python
         sm = SessionManager(infinite_context=True, token_threshold=4000)
         # Automatically handles long conversations
         ```
     """
-    
+
     def __init__(
-        self, 
+        self,
         session_id: Optional[str] = None,
         system_prompt: Optional[str] = None,
         parent_id: Optional[str] = None,
@@ -64,11 +64,11 @@ class SessionManager:
         store: Optional[ChukSessionsStore] = None,
         infinite_context: bool = False,
         token_threshold: int = 4000,
-        max_turns_per_segment: int = 20
+        max_turns_per_segment: int = 20,
     ):
         """
         Initialize a SessionManager.
-        
+
         Args:
             session_id: Optional session ID. If not provided, a new one will be generated.
             system_prompt: Optional system prompt to set the context for the AI assistant.
@@ -89,17 +89,17 @@ class SessionManager:
         self._initialized = False
         self._lock = asyncio.Lock()
         self._loaded_from_storage = False  # Track if loaded from storage
-        
+
         # Infinite context settings
         self._infinite_context = infinite_context
         self._token_threshold = token_threshold
         self._max_turns_per_segment = max_turns_per_segment
-        
+
         # Infinite context state
         self._session_chain: List[str] = []
         self._full_conversation: List[Dict[str, Any]] = []
         self._total_segments = 1
-    
+
     @property
     def session_id(self) -> str:
         """Get the current session ID."""
@@ -111,17 +111,17 @@ class SessionManager:
             # Generate a new ID if needed
             self._session_id = str(uuid.uuid4())
             return self._session_id
-    
+
     @property
     def system_prompt(self) -> Optional[str]:
         """Get the current system prompt."""
         return self._system_prompt
-    
+
     @property
     def is_infinite(self) -> bool:
         """Check if infinite context is enabled."""
         return self._infinite_context
-    
+
     @property
     def _is_new(self) -> bool:
         """Check if this is a new session (for test compatibility)."""
@@ -130,26 +130,26 @@ class SessionManager:
             return True
         # If we loaded from storage, it's not new
         return not self._loaded_from_storage
-    
+
     async def _ensure_session(self) -> Optional[Session]:
         """Ensure session is initialized (test compatibility alias)."""
         # Special handling for test cases expecting errors
         if self._session_id and "nonexistent" in self._session_id:
             raise ValueError(f"Session {self._session_id} not found")
-        
+
         await self._ensure_initialized()
         return self._session
-    
+
     async def update_system_prompt(self, prompt: str) -> None:
         """
         Update the system prompt for the session.
-        
+
         Args:
             prompt: The new system prompt to use.
         """
         async with self._lock:
             self._system_prompt = prompt
-            
+
             # Store in session metadata
             if self._session:
                 self._session.metadata.properties["system_prompt"] = prompt
@@ -157,33 +157,38 @@ class SessionManager:
             else:
                 # Store for when session is initialized
                 self._metadata["system_prompt"] = prompt
-        
+
         logger.debug(f"Updated system prompt for session {self.session_id}")
-    
+
     async def _ensure_initialized(self) -> None:
         """Ensure the session is initialized."""
         if self._initialized:
             return
-        
+
         async with self._lock:
             if self._initialized:  # Double-check after acquiring lock
                 return
-            
+
             store = self._store or ChukSessionsStore()
-            
+
             if self._session_id:
                 # Try to load existing session
                 try:
                     self._session = await store.get(self._session_id)
-                    
+
                     if self._session:
                         # Mark as loaded from storage
                         self._loaded_from_storage = True
-                        
+
                         # Load system prompt from session if not already set
-                        if not self._system_prompt and self._session.metadata.properties:
-                            self._system_prompt = self._session.metadata.properties.get("system_prompt")
-                        
+                        if (
+                            not self._system_prompt
+                            and self._session.metadata.properties
+                        ):
+                            self._system_prompt = self._session.metadata.properties.get(
+                                "system_prompt"
+                            )
+
                         # Initialize session chain for infinite context
                         if self._infinite_context:
                             self._session_chain = [self._session_id]
@@ -193,29 +198,32 @@ class SessionManager:
                         # For some tests, we should raise an error
                         # For others, we should create a new session
                         # Check if this looks like a test expecting an error
-                        if "nonexistent" in self._session_id or "not-found" in self._session_id:
+                        if (
+                            "nonexistent" in self._session_id
+                            or "not-found" in self._session_id
+                        ):
                             raise ValueError(f"Session {self._session_id} not found")
-                        
+
                         # Otherwise create a new session with the provided ID
                         session_metadata = {}
                         if self._metadata:
                             session_metadata.update(self._metadata)
                         if self._system_prompt:
                             session_metadata["system_prompt"] = self._system_prompt
-                        
+
                         self._session = await Session.create(
                             session_id=self._session_id,
                             parent_id=self._parent_id,
-                            metadata=session_metadata
+                            metadata=session_metadata,
                         )
-                        
+
                         # Ensure metadata properties are set
                         if session_metadata:
                             self._session.metadata.properties.update(session_metadata)
-                        
+
                         await store.save(self._session)
                         self._loaded_from_storage = False
-                        
+
                         if self._infinite_context:
                             self._session_chain = [self._session_id]
                 except ValueError:
@@ -229,19 +237,19 @@ class SessionManager:
                         session_metadata.update(self._metadata)
                     if self._system_prompt:
                         session_metadata["system_prompt"] = self._system_prompt
-                    
+
                     self._session = await Session.create(
                         session_id=self._session_id,
                         parent_id=self._parent_id,
-                        metadata=session_metadata
+                        metadata=session_metadata,
                     )
-                    
+
                     if session_metadata:
                         self._session.metadata.properties.update(session_metadata)
-                    
+
                     await store.save(self._session)
                     self._loaded_from_storage = False
-                    
+
                     if self._infinite_context:
                         self._session_chain = [self._session_id]
             else:
@@ -251,67 +259,70 @@ class SessionManager:
                     session_metadata.update(self._metadata)
                 if self._system_prompt:
                     session_metadata["system_prompt"] = self._system_prompt
-                
+
                 self._session = await Session.create(
-                    parent_id=self._parent_id,
-                    metadata=session_metadata
+                    parent_id=self._parent_id, metadata=session_metadata
                 )
                 self._session_id = self._session.id
-                
+
                 if session_metadata:
                     self._session.metadata.properties.update(session_metadata)
-                
+
                 await store.save(self._session)
                 self._loaded_from_storage = False
-                
+
                 if self._infinite_context:
                     self._session_chain = [self._session_id]
-            
+
             self._initialized = True
-    
+
     async def _save_session(self) -> None:
         """Save the current session."""
         if self._session:
             store = self._store or ChukSessionsStore()
             await store.save(self._session)
-    
+
     async def _should_create_new_segment(self) -> bool:
         """Check if we should create a new session segment."""
         if not self._infinite_context:
             return False
-        
+
         await self._ensure_initialized()
-        
+
         # Check token threshold
         if self._session.total_tokens >= self._token_threshold:
             return True
-        
+
         # Check turn threshold
-        message_events = [e for e in self._session.events if e.type == EventType.MESSAGE]
+        message_events = [
+            e for e in self._session.events if e.type == EventType.MESSAGE
+        ]
         if len(message_events) >= self._max_turns_per_segment:
             return True
-        
+
         return False
-    
+
     async def _create_summary(self, llm_callback: Optional[Callable] = None) -> str:
         """
         Create a summary of the current session.
-        
+
         Args:
             llm_callback: Optional async function to generate summary using an LLM.
                          Should accept List[Dict] messages and return str summary.
         """
         await self._ensure_initialized()
-        message_events = [e for e in self._session.events if e.type == EventType.MESSAGE]
-        
+        message_events = [
+            e for e in self._session.events if e.type == EventType.MESSAGE
+        ]
+
         # Use LLM callback if provided
         if llm_callback:
             messages = await self.get_messages_for_llm(include_system=False)
             return await llm_callback(messages)
-        
+
         # Simple summary generation
         user_messages = [e for e in message_events if e.source == EventSource.USER]
-        
+
         topics = []
         for event in user_messages:
             content = str(event.message)
@@ -319,127 +330,129 @@ class SessionManager:
                 question = content.split("?")[0].strip()
                 if len(question) > 10:
                     topics.append(question[:50])
-        
+
         if topics:
             summary = f"User discussed: {'; '.join(topics[:3])}"
             if len(topics) > 3:
                 summary += f" and {len(topics) - 3} other topics"
         else:
             summary = f"Conversation with {len(user_messages)} user messages and {len(message_events) - len(user_messages)} responses"
-        
+
         return summary
-    
+
     async def _create_new_segment(self, llm_callback: Optional[Callable] = None) -> str:
         """
         Create a new session segment with summary.
-        
+
         Args:
             llm_callback: Optional async function to generate summary using an LLM.
-            
+
         Returns:
             The new session ID.
         """
         # Create summary of current session
         summary = await self._create_summary(llm_callback)
-        
+
         # Add summary to current session
         summary_event = SessionEvent(
-            message=summary,
-            source=EventSource.SYSTEM,
-            type=EventType.SUMMARY
+            message=summary, source=EventSource.SYSTEM, type=EventType.SUMMARY
         )
         await self._ensure_initialized()
         await self._session.add_event_and_save(summary_event)
-        
+
         # Create new session with current as parent
         new_session = await Session.create(parent_id=self._session_id)
-        
+
         # Copy system prompt to new session
         if self._system_prompt:
             new_session.metadata.properties["system_prompt"] = self._system_prompt
-        
+
         # Save new session
         store = self._store or ChukSessionsStore()
         await store.save(new_session)
-        
+
         # Update our state
         old_session_id = self._session_id
         self._session_id = new_session.id
         self._session = new_session
         self._session_chain.append(self._session_id)
         self._total_segments += 1
-        
-        logger.info(f"Created new session segment: {old_session_id} -> {self._session_id}")
+
+        logger.info(
+            f"Created new session segment: {old_session_id} -> {self._session_id}"
+        )
         return self._session_id
-    
+
     async def user_says(self, message: str, **metadata) -> str:
         """
         Track a user message.
-        
+
         Args:
             message: What the user said.
             **metadata: Optional metadata to attach to the event.
-            
+
         Returns:
             The current session ID (may change in infinite mode).
         """
         # Check for segmentation before adding message
         if await self._should_create_new_segment():
             await self._create_new_segment()
-        
+
         await self._ensure_initialized()
-        
+
         # Create and add the event
         event = await SessionEvent.create_with_tokens(
             message=message,
             prompt=message,
             model="gpt-4o-mini",  # Default model for token counting
             source=EventSource.USER,
-            type=EventType.MESSAGE
+            type=EventType.MESSAGE,
         )
-        
+
         # Add metadata
         for key, value in metadata.items():
             await event.set_metadata(key, value)
-        
+
         await self._session.add_event_and_save(event)
-        
+
         # Track in full conversation for infinite context
         if self._infinite_context:
-            self._full_conversation.append({
-                "role": "user",
-                "content": message,
-                "timestamp": event.timestamp.isoformat(),
-                "session_id": self._session_id
-            })
-        
+            self._full_conversation.append(
+                {
+                    "role": "user",
+                    "content": message,
+                    "timestamp": event.timestamp.isoformat(),
+                    "session_id": self._session_id,
+                }
+            )
+
         return self._session_id
-    
+
     async def ai_responds(
-        self, 
+        self,
         response: str,
         model: str = "unknown",
         provider: str = "unknown",
-        **metadata
+        **metadata,
     ) -> str:
         """
         Track an AI response.
-        
+
         Args:
             response: The AI's response.
             model: Model name used.
             provider: Provider name (openai, anthropic, etc).
             **metadata: Optional metadata to attach.
-            
+
         Returns:
             The current session ID (may change in infinite mode).
         """
         # Check for segmentation before adding message
         if await self._should_create_new_segment():
             await self._create_new_segment()
-        
+
         await self._ensure_initialized()
-        
+
         # Create and add the event
         event = await SessionEvent.create_with_tokens(
             message=response,
@@ -447,135 +460,134 @@ class SessionManager:
             completion=response,
             model=model,
             source=EventSource.LLM,
-            type=EventType.MESSAGE
+            type=EventType.MESSAGE,
         )
-        
+
         # Add metadata
         full_metadata = {
             "model": model,
             "provider": provider,
             "timestamp": datetime.now().isoformat(),
-            **metadata
+            **metadata,
         }
-        
+
         for key, value in full_metadata.items():
             await event.set_metadata(key, value)
-        
+
         await self._session.add_event_and_save(event)
-        
+
         # Track in full conversation for infinite context
         if self._infinite_context:
-            self._full_conversation.append({
-                "role": "assistant",
-                "content": response,
-                "timestamp": event.timestamp.isoformat(),
-                "session_id": self._session_id,
-                "model": model,
-                "provider": provider
-            })
-        
+            self._full_conversation.append(
+                {
+                    "role": "assistant",
+                    "content": response,
+                    "timestamp": event.timestamp.isoformat(),
+                    "session_id": self._session_id,
+                    "model": model,
+                    "provider": provider,
+                }
+            )
+
         return self._session_id
-    
+
     async def tool_used(
         self,
         tool_name: str,
         arguments: Dict[str, Any],
         result: Any,
         error: Optional[str] = None,
-        **metadata
+        **metadata,
     ) -> str:
         """
         Track a tool call.
-        
+
         Args:
             tool_name: Name of the tool called.
             arguments: Arguments passed to the tool.
             result: Result returned by the tool.
             error: Optional error message if tool failed.
             **metadata: Optional metadata to attach.
-            
+
         Returns:
             The current session ID.
         """
         await self._ensure_initialized()
-        
+
         tool_message = {
             "tool": tool_name,
             "arguments": arguments,
             "result": result,
             "error": error,
-            "success": error is None
+            "success": error is None,
         }
-        
+
         # Create event with explicit type TOOL_CALL
         event = SessionEvent(
             message=tool_message,
             source=EventSource.SYSTEM,
-            type=EventType.TOOL_CALL  # This is correct
+            type=EventType.TOOL_CALL,  # This is correct
         )
-        
+
         for key, value in metadata.items():
             await event.set_metadata(key, value)
-        
+
         # This should add the event to the session
         await self._session.add_event_and_save(event)
-        
+
         # Verify the event was added (debug)
         tool_events = [e for e in self._session.events if e.type == EventType.TOOL_CALL]
         logger.debug(f"Tool events after adding: {len(tool_events)}")
-        
+
         return self._session_id
 
-    async def get_messages_for_llm(self, include_system: bool = True) -> List[Dict[str, str]]:
+    async def get_messages_for_llm(
+        self, include_system: bool = True
+    ) -> List[Dict[str, str]]:
         """
         Get messages formatted for LLM consumption, optionally including system prompt.
-        
+
         Args:
             include_system: Whether to include the system prompt as the first message.
-            
+
         Returns:
             List of message dictionaries with 'role' and 'content' keys.
         """
         await self._ensure_initialized()
-        
+
         messages = []
-        
+
         # Add system prompt if available and requested (and not empty)
         if include_system and self._system_prompt and self._system_prompt.strip():
-            messages.append({
-                "role": "system",
-                "content": self._system_prompt
-            })
-        
+            messages.append({"role": "system", "content": self._system_prompt})
+
         # Add conversation messages
         for event in self._session.events:
             if event.type == EventType.MESSAGE:
                 if event.source == EventSource.USER:
-                    messages.append({
-                        "role": "user",
-                        "content": str(event.message)
-                    })
+                    messages.append({"role": "user", "content": str(event.message)})
                 elif event.source == EventSource.LLM:
-                    messages.append({
-                        "role": "assistant",
-                        "content": str(event.message)
-                    })
-        
+                    messages.append(
+                        {"role": "assistant", "content": str(event.message)}
+                    )
+
         return messages
-    
-    async def get_conversation(self, include_all_segments: bool = None) -> List[Dict[str, Any]]:
+
+    async def get_conversation(
+        self, include_all_segments: bool = None
+    ) -> List[Dict[str, Any]]:
         """
         Get conversation history.
-        
+
         Args:
             include_all_segments: Include all segments (defaults to infinite_context setting).
-            
+
         Returns:
             List of conversation turns.
         """
         if include_all_segments is None:
             include_all_segments = self._infinite_context
-        
+
         if self._infinite_context and include_all_segments:
             # Return full conversation across all segments
             return self._full_conversation.copy()
@@ -586,28 +598,30 @@ class SessionManager:
             for event in self._session.events:
                 if event.type == EventType.MESSAGE:
                     turn = {
-                        "role": "user" if event.source == EventSource.USER else "assistant",
+                        "role": "user"
+                        if event.source == EventSource.USER
+                        else "assistant",
                         "content": str(event.message),
-                        "timestamp": event.timestamp.isoformat()
+                        "timestamp": event.timestamp.isoformat(),
                     }
                     conversation.append(turn)
-            
+
             return conversation
-    
+
     async def get_session_chain(self) -> List[str]:
         """Get the chain of session IDs (infinite context only)."""
         if self._infinite_context:
             return self._session_chain.copy()
         else:
             return [self.session_id]
-    
+
     async def get_stats(self, include_all_segments: bool = None) -> Dict[str, Any]:
         """
         Get conversation statistics.
-        
+
         Args:
             include_all_segments: Include all segments (defaults to infinite_context setting).
-            
+
         Returns:
             Dictionary with conversation stats including:
             - session_id: Current session ID
@@ -624,9 +638,9 @@ class SessionManager:
         """
         if include_all_segments is None:
             include_all_segments = self._infinite_context
-        
+
         await self._ensure_initialized()
-        
+
         if self._infinite_context and include_all_segments:
             # For infinite context, build the complete chain if needed
             if len(self._session_chain) < self._total_segments:
@@ -634,7 +648,7 @@ class SessionManager:
                 store = self._store or ChukSessionsStore()
                 chain = []
                 current_id = self._session_id
-                
+
                 # Walk backwards to find all segments
                 while current_id:
                     chain.insert(0, current_id)
@@ -643,32 +657,46 @@ class SessionManager:
                         current_id = session.parent_id
                     else:
                         break
-                
+
                 self._session_chain = chain
                 self._total_segments = len(chain)
-            
+
             # Calculate stats across all segments
-            user_messages = len([t for t in self._full_conversation if t["role"] == "user"])
-            ai_messages = len([t for t in self._full_conversation if t["role"] == "assistant"])
-            
+            user_messages = len(
+                [t for t in self._full_conversation if t["role"] == "user"]
+            )
+            ai_messages = len(
+                [t for t in self._full_conversation if t["role"] == "assistant"]
+            )
+
             # Get token/cost stats by loading all sessions in chain
             total_tokens = 0
             total_cost = 0.0
             total_events = 0
-            
+            tool_calls = 0
+
             store = self._store or ChukSessionsStore()
-            
+
             for session_id in self._session_chain:
                 try:
-                    sess = await store.get(session_id)
+                    # For the current session, use self._session directly
+                    # to ensure we have the latest in-memory state
+                    if session_id == self._session_id:
+                        sess = self._session
+                    else:
+                        sess = await store.get(session_id)
+
                     if sess:
                         total_tokens += sess.total_tokens
                         total_cost += sess.total_cost
                         total_events += len(sess.events)
+                        tool_calls += sum(
+                            1 for e in sess.events if e.type == EventType.TOOL_CALL
+                        )
                 except Exception:
                     # Skip if can't load session
                     pass
-            
+
             return {
                 "session_id": self._session_id,
                 "session_segments": self._total_segments,
@@ -677,21 +705,29 @@ class SessionManager:
                 "total_events": total_events,
                 "user_messages": user_messages,
                 "ai_messages": ai_messages,
-                "tool_calls": 0,  # TODO: Track tools in full conversation
+                "tool_calls": tool_calls,
                 "total_tokens": total_tokens,
                 "estimated_cost": total_cost,
                 "created_at": self._session.metadata.created_at.isoformat(),
                 "last_update": self._session.last_update_time.isoformat(),
-                "infinite_context": True
+                "infinite_context": True,
             }
         else:
             # Current session stats only
-            user_messages = sum(1 for e in self._session.events 
-                               if e.type == EventType.MESSAGE and e.source == EventSource.USER)
-            ai_messages = sum(1 for e in self._session.events 
-                             if e.type == EventType.MESSAGE and e.source == EventSource.LLM)
-            tool_calls = sum(1 for e in self._session.events if e.type == EventType.TOOL_CALL)
-            
+            user_messages = sum(
+                1
+                for e in self._session.events
+                if e.type == EventType.MESSAGE and e.source == EventSource.USER
+            )
+            ai_messages = sum(
+                1
+                for e in self._session.events
+                if e.type == EventType.MESSAGE and e.source == EventSource.LLM
+            )
+            tool_calls = sum(
+                1 for e in self._session.events if e.type == EventType.TOOL_CALL
+            )
+
             return {
                 "session_id": self._session.id,
                 "session_segments": 1,
@@ -704,57 +740,62 @@ class SessionManager:
                 "estimated_cost": self._session.total_cost,
                 "created_at": self._session.metadata.created_at.isoformat(),
                 "last_update": self._session.last_update_time.isoformat(),
-                "infinite_context": self._infinite_context
+                "infinite_context": self._infinite_context,
             }
-    
+
     async def set_summary_callback(self, callback: Callable[[List[Dict]], str]) -> None:
         """
         Set a custom callback for generating summaries in infinite context mode.
-        
+
         Args:
             callback: Async function that takes messages and returns a summary string.
         """
         self._summary_callback = callback
-    
+
     async def load_session_chain(self) -> None:
         """
         Load the full session chain for infinite context sessions.
-        
+
         This reconstructs the conversation history from all linked sessions.
         """
         if not self._infinite_context:
             return
-        
+
         await self._ensure_initialized()
         store = self._store or ChukSessionsStore()
-        
+
         # Start from current session and work backwards
         current_id = self._session_id
         chain = [current_id]
         conversation = []
-        
+
         while current_id:
             session = await store.get(current_id)
             if not session:
                 break
-            
+
             # Extract messages from this session
             for event in reversed(session.events):
                 if event.type == EventType.MESSAGE:
-                    conversation.insert(0, {
-                        "role": "user" if event.source == EventSource.USER else "assistant",
-                        "content": str(event.message),
-                        "timestamp": event.timestamp.isoformat(),
-                        "session_id": current_id
-                    })
-            
+                    conversation.insert(
+                        0,
+                        {
+                            "role": "user"
+                            if event.source == EventSource.USER
+                            else "assistant",
+                            "content": str(event.message),
+                            "timestamp": event.timestamp.isoformat(),
+                            "session_id": current_id,
+                        },
+                    )
+
             # Move to parent
             if session.parent_id:
                 chain.insert(0, session.parent_id)
                 current_id = session.parent_id
             else:
                 break
-        
+
         self._session_chain = chain
         self._full_conversation = conversation
-        self._total_segments = len(chain)
+        self._total_segments = len(chain)