lucidicai 2.1.2__py3-none-any.whl → 3.0.0__py3-none-any.whl

lucidicai/sdk/shutdown_manager.py

@@ -8,18 +8,20 @@ import signal
 import sys
 import threading
 import time
-from typing import Dict, Optional, Set, Callable
+from typing import Dict, Optional, Set, Callable, TYPE_CHECKING
 from dataclasses import dataclass
 
 from ..utils.logger import debug, info, warning, error, truncate_id
 
+if TYPE_CHECKING:
+    from ..client import LucidicAI
+
 
 @dataclass
 class SessionState:
     """State information for an active session."""
     session_id: str
     http_client: Optional[object] = None
-    event_queue: Optional[object] = None
     is_shutting_down: bool = False
     auto_end: bool = True
 
@@ -46,14 +48,18 @@ class ShutdownManager:
         # only initialize once
         if self._initialized:
             return
-            
+
         self._initialized = True
         self.active_sessions: Dict[str, SessionState] = {}
         self.is_shutting_down = False
         self.shutdown_complete = threading.Event()
         self.listeners_registered = False
         self._session_lock = threading.Lock()
-        
+
+        # Client registry for multi-client support
+        self._clients: Dict[str, "LucidicAI"] = {}
+        self._client_lock = threading.Lock()
+
         debug("[ShutdownManager] Initialized")
     
     def register_session(self, session_id: str, state: SessionState) -> None:
@@ -87,16 +93,39 @@ class ShutdownManager:
     
     def is_session_active(self, session_id: str) -> bool:
         """Check if a session is active.
-        
+
         Args:
             session_id: Session identifier
-            
+
         Returns:
             True if session is active
         """
         with self._session_lock:
             return session_id in self.active_sessions
-    
+
+    def register_client(self, client: "LucidicAI") -> None:
+        """Register a client for shutdown tracking.
+
+        Args:
+            client: The LucidicAI client to register
+        """
+        with self._client_lock:
+            self._clients[client._client_id] = client
+            debug(f"[ShutdownManager] Registered client {client._client_id[:8]}...")
+
+            # Ensure listeners are registered
+            self._ensure_listeners_registered()
+
+    def unregister_client(self, client_id: str) -> None:
+        """Unregister a client.
+
+        Args:
+            client_id: The client ID to unregister
+        """
+        with self._client_lock:
+            self._clients.pop(client_id, None)
+            debug(f"[ShutdownManager] Unregistered client {client_id[:8]}...")
+
     def _ensure_listeners_registered(self) -> None:
         """Register process exit listeners once."""
         if self.listeners_registered:
@@ -164,46 +193,78 @@ class ShutdownManager:
         if self.is_shutting_down:
             debug(f"[ShutdownManager] Already shutting down, ignoring {trigger}")
             return
-        
+
         self.is_shutting_down = True
-        
+
+        # Check if there's anything to clean up
         with self._session_lock:
             session_count = len(self.active_sessions)
-            if session_count == 0:
-                debug("[ShutdownManager] No active sessions to clean up")
-                self.shutdown_complete.set()
-                return
-                
-            info(f"[ShutdownManager] Shutdown initiated by {trigger}, ending {session_count} active session(s)")
-            
-            # perform shutdown in separate thread to avoid deadlocks
-            import threading
-            shutdown_thread = threading.Thread(
-                target=self._perform_shutdown,
-                name="ShutdownThread"
-            )
-            shutdown_thread.daemon = True
-            shutdown_thread.start()
-            
-            # wait for shutdown with timeout
-            if not self.shutdown_complete.wait(timeout=30):
-                warning("[ShutdownManager] Shutdown timeout after 30s")
+        with self._client_lock:
+            client_count = len(self._clients)
+
+        if session_count == 0 and client_count == 0:
+            debug("[ShutdownManager] No active sessions or clients to clean up")
+            # Reset flag so future shutdowns can proceed (e.g., if exception triggered
+            # shutdown before any sessions were created, then user creates sessions)
+            self.is_shutting_down = False
+            self.shutdown_complete.set()
+            return
+
+        info(f"[ShutdownManager] Shutdown initiated by {trigger}, ending {session_count} active session(s), {client_count} client(s)")
+
+        # perform shutdown in separate thread to avoid deadlocks
+        import threading
+        shutdown_thread = threading.Thread(
+            target=self._perform_shutdown,
+            name="ShutdownThread"
+        )
+        shutdown_thread.daemon = True
+        shutdown_thread.start()
+
+        # wait for shutdown with timeout - MUST be outside lock to avoid deadlock
+        if not self.shutdown_complete.wait(timeout=30):
+            warning("[ShutdownManager] Shutdown timeout after 30s")
     
     def _perform_shutdown(self) -> None:
-        """Perform the actual shutdown of all sessions."""
+        """Perform the actual shutdown of all sessions and clients."""
         debug("[ShutdownManager] _perform_shutdown thread started")
         try:
+            # First, flush all pending background events before ending sessions
+            # This ensures telemetry from the exporter is sent
+            try:
+                debug("[ShutdownManager] Flushing pending events before session cleanup")
+                from ..sdk.event import flush
+                flush(timeout=5.0)
+                debug("[ShutdownManager] Event flush complete")
+            except Exception as e:
+                error(f"[ShutdownManager] Error flushing events: {e}")
+
+            # Close all registered clients (this will also end their sessions)
+            clients_to_close = []
+            with self._client_lock:
+                clients_to_close = list(self._clients.values())
+
+            if clients_to_close:
+                debug(f"[ShutdownManager] Closing {len(clients_to_close)} registered client(s)")
+                for client in clients_to_close:
+                    try:
+                        debug(f"[ShutdownManager] Closing client {client._client_id[:8]}...")
+                        client.close()
+                    except Exception as e:
+                        error(f"[ShutdownManager] Error closing client: {e}")
+
+            # Also handle sessions registered directly (legacy support)
             sessions_to_end = []
-            
+
             with self._session_lock:
                 # collect sessions that need ending
                 for session_id, state in self.active_sessions.items():
                     if state.auto_end and not state.is_shutting_down:
                         state.is_shutting_down = True
                         sessions_to_end.append((session_id, state))
-            
+
             debug(f"[ShutdownManager] Found {len(sessions_to_end)} sessions to end")
-            
+
             # end all sessions
             for session_id, state in sessions_to_end:
                 try:
@@ -214,13 +275,14 @@ class ShutdownManager:
             
             # Final telemetry shutdown after all sessions are ended
             try:
-                from ..sdk.init import _sdk_state
-                if hasattr(_sdk_state, 'tracer_provider') and _sdk_state.tracer_provider:
+                from ..sdk.init import get_tracer_provider
+                tracer_provider = get_tracer_provider()
+                if tracer_provider:
                     debug("[ShutdownManager] Final OpenTelemetry shutdown")
                     try:
                         # Final flush and shutdown with longer timeout
-                        _sdk_state.tracer_provider.force_flush(timeout_millis=5000)
-                        _sdk_state.tracer_provider.shutdown()
+                        tracer_provider.force_flush(timeout_millis=5000)
+                        tracer_provider.shutdown()
                         debug("[ShutdownManager] OpenTelemetry shutdown complete")
                     except Exception as e:
                         error(f"[ShutdownManager] Error in final telemetry shutdown: {e}")
@@ -244,25 +306,20 @@ class ShutdownManager:
             session_id: Session identifier
             state: Session state
         """
-        # Flush OpenTelemetry spans first (before event queue)
+        # Flush OpenTelemetry spans first
         try:
-            # Get the global tracer provider if it exists
-            from ..sdk.init import _sdk_state
-            if hasattr(_sdk_state, 'tracer_provider') and _sdk_state.tracer_provider:
+            from ..sdk.init import get_tracer_provider
+            tracer_provider = get_tracer_provider()
+            if tracer_provider:
                 debug(f"[ShutdownManager] Flushing OpenTelemetry spans for session {truncate_id(session_id)}")
                 try:
                     # Force flush with 3 second timeout
-                    _sdk_state.tracer_provider.force_flush(timeout_millis=3000)
+                    tracer_provider.force_flush(timeout_millis=3000)
                 except Exception as e:
                     error(f"[ShutdownManager] Error flushing spans: {e}")
         except ImportError:
             pass  # SDK not initialized
         
-        # Skip event queue flush during shutdown to avoid hanging
-        # The queue worker is a daemon thread and will flush on its own
-        if state.event_queue:
-            debug(f"[ShutdownManager] Skipping event queue flush during shutdown for session {truncate_id(session_id)}")
-        
         # end session via API if http client present
         if state.http_client and session_id:
             try:

lucidicai/session_obj.py

@@ -0,0 +1,321 @@
+"""Session object - Represents an active session bound to a LucidicAI client.
+
+A Session is returned by client.sessions.create() and provides:
+- Context manager support for automatic session binding and cleanup
+- Methods to update and end the session
+- Access to the parent client for operations within the session
+"""
+
+from typing import Optional, List, Any, TYPE_CHECKING
+import logging
+
+from .sdk.context import current_session_id, current_client
+
+if TYPE_CHECKING:
+    from .client import LucidicAI
+
+logger = logging.getLogger("Lucidic")
+
+
+class Session:
+    """Represents an active session bound to a LucidicAI client.
+
+    Sessions track a unit of work (conversation, workflow, etc.) and provide:
+    - Automatic context binding when used as a context manager
+    - Methods to update session metadata and end the session
+    - Reference to the parent client for additional operations
+
+    Example:
+        client = LucidicAI(api_key="...", providers=["openai"])
+
+        # Using as context manager (recommended)
+        with client.sessions.create(session_name="My Session") as session:
+            # Session is active, LLM calls are tracked
+            response = openai_client.chat.completions.create(...)
+
+        # Manual usage
+        session = client.sessions.create(session_name="Manual Session", auto_end=False)
+        try:
+            # Do work
+            pass
+        finally:
+            session.end()
+    """
+
+    def __init__(
+        self,
+        client: "LucidicAI",
+        session_id: str,
+        session_name: Optional[str] = None,
+        auto_end: bool = True,
+    ):
+        """Initialize a Session.
+
+        Args:
+            client: The LucidicAI client that owns this session
+            session_id: The unique session identifier
+            session_name: Optional human-readable name for the session
+            auto_end: Whether to automatically end the session on context exit
+        """
+        self._client = client
+        self._session_id = session_id
+        self._session_name = session_name
+        self._auto_end = auto_end
+        self._ended = False
+        self._context_token = None
+        self._client_token = None
+
+    @property
+    def session_id(self) -> str:
+        """Get the session ID."""
+        return self._session_id
+
+    @property
+    def session_name(self) -> Optional[str]:
+        """Get the session name."""
+        return self._session_name
+
+    @property
+    def client(self) -> "LucidicAI":
+        """Get the parent client."""
+        return self._client
+
+    @property
+    def is_finished(self) -> bool:
+        """Check if the session has been ended."""
+        return self._ended
+
+    def update(
+        self,
+        task: Optional[str] = None,
+        session_eval: Optional[float] = None,
+        session_eval_reason: Optional[str] = None,
+        is_successful: Optional[bool] = None,
+        is_successful_reason: Optional[str] = None,
+    ) -> None:
+        """Update the session metadata.
+
+        Args:
+            task: Update the task description
+            session_eval: Evaluation score (0.0 to 1.0)
+            session_eval_reason: Reason for the evaluation score
+            is_successful: Whether the session was successful
+            is_successful_reason: Reason for success/failure status
+        """
+        if self._ended:
+            logger.warning(f"[Session] Attempted to update ended session {self._session_id[:8]}...")
+            return
+
+        updates = {}
+        if task is not None:
+            updates["task"] = task
+        if session_eval is not None:
+            updates["session_eval"] = session_eval
+        if session_eval_reason is not None:
+            updates["session_eval_reason"] = session_eval_reason
+        if is_successful is not None:
+            updates["is_successful"] = is_successful
+        if is_successful_reason is not None:
+            updates["is_successful_reason"] = is_successful_reason
+
+        if updates:
+            logger.debug(
+                f"[Session] update() called - session_id={self._session_id[:8]}..., updates={updates}"
+            )
+            self._client.sessions.update(self._session_id, updates)
+            logger.debug(f"[Session] update() completed for session {self._session_id[:8]}...")
+        else:
+            logger.debug(f"[Session] update() called with no updates for session {self._session_id[:8]}...")
+
+    async def aupdate(
+        self,
+        task: Optional[str] = None,
+        session_eval: Optional[float] = None,
+        session_eval_reason: Optional[str] = None,
+        is_successful: Optional[bool] = None,
+        is_successful_reason: Optional[str] = None,
+    ) -> None:
+        """Update the session metadata (async version).
+
+        Args:
+            task: Update the task description
+            session_eval: Evaluation score (0.0 to 1.0)
+            session_eval_reason: Reason for the evaluation score
+            is_successful: Whether the session was successful
+            is_successful_reason: Reason for success/failure status
+        """
+        if self._ended:
+            logger.warning(f"[Session] Attempted to update ended session {self._session_id[:8]}...")
+            return
+
+        updates = {}
+        if task is not None:
+            updates["task"] = task
+        if session_eval is not None:
+            updates["session_eval"] = session_eval
+        if session_eval_reason is not None:
+            updates["session_eval_reason"] = session_eval_reason
+        if is_successful is not None:
+            updates["is_successful"] = is_successful
+        if is_successful_reason is not None:
+            updates["is_successful_reason"] = is_successful_reason
+
+        if updates:
+            logger.debug(
+                f"[Session] aupdate() called - session_id={self._session_id[:8]}..., updates={updates}"
+            )
+            await self._client.sessions.aupdate(self._session_id, updates)
+            logger.debug(f"[Session] aupdate() completed for session {self._session_id[:8]}...")
+        else:
+            logger.debug(f"[Session] aupdate() called with no updates for session {self._session_id[:8]}...")
+
+    def end(
+        self,
+        is_successful: Optional[bool] = None,
+        is_successful_reason: Optional[str] = None,
+        session_eval: Optional[float] = None,
+        session_eval_reason: Optional[str] = None,
+    ) -> None:
+        """End the session.
+
+        Args:
+            is_successful: Whether the session was successful
+            is_successful_reason: Reason for success/failure status
+            session_eval: Evaluation score (0.0 to 1.0)
+            session_eval_reason: Reason for the evaluation score
+        """
+        if self._ended:
+            return
+
+        # Unbind context before ending
+        self._unbind_context()
+
+        self._client.sessions.end(
+            session_id=self._session_id,
+            is_successful=is_successful,
+            is_successful_reason=is_successful_reason,
+            session_eval=session_eval,
+            session_eval_reason=session_eval_reason,
+        )
+        self._ended = True
+
+    async def aend(
+        self,
+        is_successful: Optional[bool] = None,
+        is_successful_reason: Optional[str] = None,
+        session_eval: Optional[float] = None,
+        session_eval_reason: Optional[str] = None,
+    ) -> None:
+        """End the session (async version).
+
+        Args:
+            is_successful: Whether the session was successful
+            is_successful_reason: Reason for success/failure status
+            session_eval: Evaluation score (0.0 to 1.0)
+            session_eval_reason: Reason for the evaluation score
+        """
+        if self._ended:
+            return
+
+        # Unbind context before ending
+        self._unbind_context()
+
+        await self._client.sessions.aend(
+            session_id=self._session_id,
+            is_successful=is_successful,
+            is_successful_reason=is_successful_reason,
+            session_eval=session_eval,
+            session_eval_reason=session_eval_reason,
+        )
+        self._ended = True
+
+    def _bind_context(self) -> None:
+        """Bind this session to the current context."""
+        self._context_token = current_session_id.set(self._session_id)
+        self._client_token = current_client.set(self._client)
+
+    def _unbind_context(self) -> None:
+        """Unbind this session from the current context."""
+        if self._context_token:
+            current_session_id.reset(self._context_token)
+            self._context_token = None
+        if self._client_token:
+            current_client.reset(self._client_token)
+            self._client_token = None
+
+    def __enter__(self) -> "Session":
+        """Enter the session context - binds session to current context."""
+        # Only bind if not already bound (create_session() already binds context)
+        if self._context_token is None:
+            self._bind_context()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the session context - unbinds and optionally ends the session."""
+        # Flush telemetry before ending
+        try:
+            from .telemetry.telemetry_manager import get_telemetry_manager
+
+            manager = get_telemetry_manager()
+            if manager.is_telemetry_initialized:
+                logger.debug(
+                    f"[Session] Flushing telemetry for session {self._session_id}"
+                )
+                flush_success = manager.force_flush(timeout_millis=5000)
+                if not flush_success:
+                    logger.warning("[Session] Telemetry flush may be incomplete")
+        except Exception as e:
+            logger.debug(f"[Session] Error flushing telemetry: {e}")
+
+        # Unbind context
+        self._unbind_context()
+
+        # End session if auto_end is enabled
+        if self._auto_end and not self._ended:
+            try:
+                self.end()
+            except Exception as e:
+                # Don't mask the original exception
+                logger.debug(f"[Session] Error ending session: {e}")
+
+    async def __aenter__(self) -> "Session":
+        """Enter the async session context - binds session to current context."""
+        # Only bind if not already bound (acreate_session() already binds context)
+        if self._context_token is None:
+            self._bind_context()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the async session context - unbinds and optionally ends the session."""
+        import asyncio
+
+        # Flush telemetry before ending
+        try:
+            from .telemetry.telemetry_manager import get_telemetry_manager
+
+            manager = get_telemetry_manager()
+            if manager.is_telemetry_initialized:
+                logger.debug(
+                    f"[Session] Flushing telemetry for async session {self._session_id}"
+                )
+                flush_success = manager.force_flush(timeout_millis=5000)
+                if not flush_success:
+                    logger.warning("[Session] Telemetry flush may be incomplete")
+        except Exception as e:
+            logger.debug(f"[Session] Error flushing telemetry: {e}")
+
+        # Unbind context
+        self._unbind_context()
+
+        # End session if auto_end is enabled
+        if self._auto_end and not self._ended:
+            try:
+                await self.aend()
+            except Exception as e:
+                # Don't mask the original exception
+                logger.debug(f"[Session] Error ending async session: {e}")
+
+    def __repr__(self) -> str:
+        status = "ended" if self._ended else "active"
+        name_part = f", name={self._session_name!r}" if self._session_name else ""
+        return f"<Session(id={self._session_id[:8]}...{name_part}, status={status})>"

lucidicai/telemetry/context_capture_processor.py

@@ -31,13 +31,13 @@ class ContextCaptureProcessor(SpanProcessor):
             
             try:
                 session_id = current_session_id.get(None)
-            except Exception:
-                pass
-            
+            except Exception as e:
+                debug(f"[ContextCapture] Failed to get session_id from contextvar: {e}")
+
             try:
                 parent_event_id = current_parent_event_id.get(None)
-            except Exception:
-                pass
+            except Exception as e:
+                debug(f"[ContextCapture] Failed to get parent_event_id from contextvar: {e}")
             
             # If not found in contextvars, try OpenTelemetry baggage
             # This handles cases where spans are created in different threads
@@ -57,12 +57,19 @@ class ContextCaptureProcessor(SpanProcessor):
             if session_id:
                 span.set_attribute("lucidic.session_id", session_id)
                 debug(f"[ContextCapture] Set session_id attribute for span {span.name}")
-            
+
             if parent_event_id:
                 span.set_attribute("lucidic.parent_event_id", parent_event_id)
                 debug(f"[ContextCapture] Captured parent_event_id {truncate_id(parent_event_id)} for span {span.name}")
             else:
                 debug(f"[ContextCapture] No parent_event_id available for span {span.name}")
+
+            # Capture client_id for multi-client routing
+            from lucidicai.sdk.context import get_active_client
+            client = get_active_client()
+            if client:
+                span.set_attribute("lucidic.client_id", client._client_id)
+                debug(f"[ContextCapture] Set client_id attribute for span {span.name}")
             
         except Exception as e:
             # Never fail span creation due to context capture