lucidicai 1.3.1__py3-none-any.whl → 1.3.5__py3-none-any.whl

lucidicai/__init__.py

@@ -2,6 +2,9 @@ import atexit
 import logging
 import os
 import signal
+import sys
+import traceback
+import threading
 from typing import List, Literal, Optional
 
 from .client import Client
@@ -25,8 +28,35 @@ from .telemetry.otel_init import LucidicTelemetry
 
 # Import decorators
 from .decorators import step, event
+from .context import (
+    set_active_session,
+    bind_session,
+    bind_session_async,
+    clear_active_session,
+    current_session_id,
+    session,
+    session_async,
+    run_session,
+    run_in_session,
+)
 
-ProviderType = Literal["openai", "anthropic", "langchain", "pydantic_ai", "openai_agents", "litellm"]
+ProviderType = Literal[
+    "openai",
+    "anthropic",
+    "langchain",
+    "pydantic_ai",
+    "openai_agents",
+    "litellm",
+    "bedrock",
+    "aws_bedrock",
+    "amazon_bedrock",
+    "google",
+    "google_generativeai",
+    "vertexai",
+    "vertex_ai",
+    "cohere",
+    "groq",
+]
 
 # Configure logging
 logger = logging.getLogger("Lucidic")
@@ -38,6 +68,137 @@ if not logger.handlers:
     logger.setLevel(logging.INFO)
 
 
+# Crash/exit capture configuration
+MAX_ERROR_DESCRIPTION_LENGTH = 16384
+_crash_handlers_installed = False
+_original_sys_excepthook = None
+_original_threading_excepthook = None
+_shutdown_lock = threading.Lock()
+_is_shutting_down = False
+
+
+def _mask_and_truncate(text: Optional[str]) -> Optional[str]:
+    """Apply masking and truncate to a safe length. Best effort; never raises."""
+    if text is None:
+        return text
+    try:
+        masked = Client().mask(text)
+    except Exception:
+        masked = text
+    if masked is None:
+        return masked
+    return masked[:MAX_ERROR_DESCRIPTION_LENGTH]
+
+
+def _post_fatal_event(exit_code: int, description: str, extra: Optional[dict] = None) -> None:
+    """Best-effort creation of a final Lucidic event on fatal paths.
+
+    - Idempotent using a process-wide shutdown flag to avoid duplicates when
+      multiple hooks fire (signal + excepthook).
+    - Swallows all exceptions to avoid interfering with shutdown.
+    """
+    global _is_shutting_down
+    with _shutdown_lock:
+        if _is_shutting_down:
+            return
+        _is_shutting_down = True
+    try:
+        client = Client()
+        session = getattr(client, 'session', None)
+        if not session or getattr(session, 'is_finished', False):
+            return
+        arguments = {"exit_code": exit_code}
+        if extra:
+            try:
+                arguments.update(extra)
+            except Exception:
+                pass
+
+        event_id = session.create_event(
+            description=_mask_and_truncate(description),
+            result=f"process exited with code {exit_code}",
+            function_name="__process_exit__",
+            arguments=arguments,
+        )
+        session.update_event(event_id=event_id, is_finished=True)
+    except Exception:
+        # Never raise during shutdown
+        pass
+
+
+def _install_crash_handlers() -> None:
+    """Install global uncaught exception handlers (idempotent)."""
+    global _crash_handlers_installed, _original_sys_excepthook, _original_threading_excepthook
+    if _crash_handlers_installed:
+        return
+
+    _original_sys_excepthook = sys.excepthook
+
+    def _sys_hook(exc_type, exc, tb):
+        try:
+            trace_str = ''.join(traceback.format_exception(exc_type, exc, tb))
+        except Exception:
+            trace_str = f"Uncaught exception: {getattr(exc_type, '__name__', str(exc_type))}: {exc}"
+
+        # Emit final event and end the session as unsuccessful
+        _post_fatal_event(1, trace_str, {
+            "exception_type": getattr(exc_type, "__name__", str(exc_type)),
+            "exception_message": str(exc),
+            "thread_name": threading.current_thread().name,
+        })
+        try:
+            # Prevent auto_end double work
+            client = Client()
+            try:
+                client.auto_end = False
+            except Exception:
+                pass
+            # End session explicitly as unsuccessful
+            end_session()
+        except Exception:
+            pass
+        # Best-effort force flush and shutdown telemetry
+        try:
+            telemetry = LucidicTelemetry()
+            if telemetry.is_initialized():
+                try:
+                    telemetry.force_flush()
+                except Exception:
+                    pass
+                try:
+                    telemetry.uninstrument_all()
+                except Exception:
+                    pass
+        except Exception:
+            pass
+        # Chain to original to preserve default printing/behavior
+        try:
+            _original_sys_excepthook(exc_type, exc, tb)
+        except Exception:
+            # Avoid recursion/errors in fatal path
+            pass
+
+    sys.excepthook = _sys_hook
+
+    # For Python 3.8+, only treat main-thread exceptions as fatal (process-exiting)
+    if hasattr(threading, 'excepthook'):
+        _original_threading_excepthook = threading.excepthook
+
+        def _thread_hook(args):
+            try:
+                if args.thread is threading.main_thread():
+                    _sys_hook(args.exc_type, args.exc_value, args.exc_traceback)
+            except Exception:
+                pass
+            try:
+                _original_threading_excepthook(args)
+            except Exception:
+                pass
+
+        threading.excepthook = _thread_hook
+
+    _crash_handlers_installed = True
+
 def _setup_providers(client: Client, providers: List[ProviderType]) -> None:
     """Set up providers for the client, avoiding duplication
     
@@ -81,6 +242,26 @@ def _setup_providers(client: Client, providers: List[ProviderType]) -> None:
         elif provider == "litellm":
             client.set_provider(OTelLiteLLMHandler())
             setup_providers.add("litellm")
+        elif provider in ("bedrock", "aws_bedrock", "amazon_bedrock"):
+            from .telemetry.otel_handlers import OTelBedrockHandler
+            client.set_provider(OTelBedrockHandler())
+            setup_providers.add("bedrock")
+        elif provider in ("google", "google_generativeai"):
+            from .telemetry.otel_handlers import OTelGoogleGenerativeAIHandler
+            client.set_provider(OTelGoogleGenerativeAIHandler())
+            setup_providers.add("google")
+        elif provider in ("vertexai", "vertex_ai"):
+            from .telemetry.otel_handlers import OTelVertexAIHandler
+            client.set_provider(OTelVertexAIHandler())
+            setup_providers.add("vertexai")
+        elif provider == "cohere":
+            from .telemetry.otel_handlers import OTelCohereHandler
+            client.set_provider(OTelCohereHandler())
+            setup_providers.add("cohere")
+        elif provider == "groq":
+            from .telemetry.otel_handlers import OTelGroqHandler
+            client.set_provider(OTelGroqHandler())
+            setup_providers.add("groq")
 
 __all__ = [
     'Client',
@@ -105,6 +286,14 @@ __all__ = [
     'InvalidOperationError',
     'step',
     'event',
+    'set_active_session',
+    'bind_session',
+    'bind_session_async',
+    'clear_active_session',
+    'session',
+    'session_async',
+    'run_session',
+    'run_in_session',
 ]
 
 
@@ -117,10 +306,12 @@ def init(
     providers: Optional[List[ProviderType]] = [],
     production_monitoring: Optional[bool] = False,
     mass_sim_id: Optional[str] = None,
+    experiment_id: Optional[str] = None,
     rubrics: Optional[list] = None,
     tags: Optional[list] = None,
     masking_function = None,
     auto_end: Optional[bool] = True,
+    capture_uncaught: Optional[bool] = True,
 ) -> str:
     """
     Initialize the Lucidic client.
@@ -133,6 +324,7 @@ def init(
         task: Task description.
         providers: List of provider types ("openai", "anthropic", "langchain", "pydantic_ai").
         mass_sim_id: Optional mass simulation ID, if session is to be part of a mass simulation.
+        experiment_id: Optional experiment ID, if session is to be part of an experiment.
         rubrics: Optional rubrics for evaluation, list of strings.
         tags: Optional tags for the session, list of strings.
         masking_function: Optional function to mask sensitive data.
@@ -180,12 +372,24 @@ def init(
         tags=tags,
         production_monitoring=production_monitoring,
         session_id=session_id,
+        experiment_id=experiment_id,
     )
     if masking_function:
         client.masking_function = masking_function
     
     # Set the auto_end flag on the client
     client.auto_end = auto_end
+    # Bind this session id to the current execution context for async-safety
+    try:
+        set_active_session(real_session_id)
+    except Exception:
+        pass
+    # Install crash handlers unless explicitly disabled
+    try:
+        if capture_uncaught:
+            _install_crash_handlers()
+    except Exception:
+        pass
     
     logger.info("Session initialized successfully")
     return real_session_id
@@ -229,6 +433,11 @@ def continue_session(
     client.auto_end = auto_end
     
     logger.info(f"Session {session_id} continuing...")
+    # Bind this session id to the current execution context for async-safety
+    try:
+        set_active_session(session_id)
+    except Exception:
+        pass
     return session_id  # For consistency
 
 
@@ -249,10 +458,20 @@ def update_session(
         is_successful: Whether the session was successful.
         is_successful_reason: Session success reason.
     """
+    # Prefer context-bound session over global active session
     client = Client()
-    if not client.session:
+    target_sid = None
+    try:
+        target_sid = current_session_id.get(None)
+    except Exception:
+        target_sid = None
+    if not target_sid and client.session:
+        target_sid = client.session.session_id
+    if not target_sid:
         return
-    client.session.update_session(**locals())
+    # Use ephemeral session facade to avoid mutating global state
+    session = client.session if (client.session and client.session.session_id == target_sid) else Session(agent_id=client.agent_id, session_id=target_sid)
+    session.update_session(**locals())
 
 
 def end_session(
@@ -271,17 +490,31 @@ def end_session(
         is_successful_reason: Session success reason.
     """
     client = Client()
-    if not client.session:
+    # Prefer context-bound session id
+    target_sid = None
+    try:
+        target_sid = current_session_id.get(None)
+    except Exception:
+        target_sid = None
+    if not target_sid and client.session:
+        target_sid = client.session.session_id
+    if not target_sid:
         return
-    
-    # Wait for any pending LiteLLM callbacks before ending session
-    for provider in client.providers:
-        if hasattr(provider, '_callback') and hasattr(provider._callback, 'wait_for_pending_callbacks'):
-            logger.info("Waiting for LiteLLM callbacks to complete before ending session...")
-            provider._callback.wait_for_pending_callbacks(timeout=5.0)
-    
-    client.session.update_session(is_finished=True, **locals())
-    client.clear()
+
+    # If ending the globally active session, keep existing cleanup behavior
+    if client.session and client.session.session_id == target_sid:
+        # Wait for any pending LiteLLM callbacks before ending session
+        for provider in client.providers:
+            if hasattr(provider, '_callback') and hasattr(provider._callback, 'wait_for_pending_callbacks'):
+                logger.info("Waiting for LiteLLM callbacks to complete before ending session...")
+                provider._callback.wait_for_pending_callbacks(timeout=5.0)
+        client.session.update_session(is_finished=True, **locals())
+        client.clear()
+        return
+
+    # Otherwise, end the specified session id without clearing global state
+    temp = Session(agent_id=client.agent_id, session_id=target_sid)
+    temp.update_session(is_finished=True, **locals())
 
 
 def reset_sdk() -> None:
@@ -327,6 +560,20 @@ def _auto_end_session():
 
 def _signal_handler(signum, frame):
     """Handle interruption signals"""
+    # Best-effort final event for signal exits
+    try:
+        try:
+            name = signal.Signals(signum).name
+        except Exception:
+            name = str(signum)
+        try:
+            stack_str = ''.join(traceback.format_stack(frame)) if frame else ''
+        except Exception:
+            stack_str = ''
+        desc = _mask_and_truncate(f"Received signal {name}\n{stack_str}")
+        _post_fatal_event(128 + signum, desc, {"signal": name, "signum": signum})
+    except Exception:
+        pass
     _auto_end_session()
     _cleanup_telemetry()
     # Re-raise the signal for default handling
@@ -487,7 +734,9 @@ def create_event(
     result: Optional[str] = None,
     cost_added: Optional[float] = None, 
     model: Optional[str] = None,
-    screenshots: Optional[List[str]] = None
+    screenshots: Optional[List[str]] = None,
+    function_name: Optional[str] = None,
+    arguments: Optional[dict] = None,
 ) -> str:
     """
     Create a new event in the current step. Current step must not be finished.
@@ -498,6 +747,8 @@ def create_event(
         cost_added: Cost added by the event.
         model: Model used for the event.
         screenshots: List of screenshots encoded in base64.
+        function_name: Name of the function that created the event.
+        arguments: Arguments of the function that created the event.
     """
 
     client = Client()
@@ -512,7 +763,9 @@ def update_event(
     result: Optional[str] = None,
     cost_added: Optional[float] = None, 
     model: Optional[str] = None,
-    screenshots: Optional[List[str]] = None
+    screenshots: Optional[List[str]] = None,
+    function_name: Optional[str] = None,
+    arguments: Optional[dict] = None,
 ) -> None:
     """
     Update the event with the given ID in the current step.
@@ -524,6 +777,8 @@ def update_event(
         cost_added: Cost added by the event.
         model: Model used for the event.
         screenshots: List of screenshots encoded in base64.
+        function_name: Name of the function that created the event.
+        arguments: Arguments of the function that created the event.
     """
     client = Client()
     if not client.session:
@@ -537,7 +792,9 @@ def end_event(
     result: Optional[str] = None,
     cost_added: Optional[float] = None, 
     model: Optional[str] = None,
-    screenshots: Optional[List[str]] = None
+    screenshots: Optional[List[str]] = None,
+    function_name: Optional[str] = None,
+    arguments: Optional[dict] = None,
 ) -> None:
     """
     End the latest event in the current step.
@@ -548,6 +805,9 @@ def end_event(
         result: Result of the event.
         cost_added: Cost added by the event.
         model: Model used for the event.
+        screenshots: List of screenshots encoded in base64.
+        function_name: Name of the function that created the event.
+        arguments: Arguments of the function that created the event.
     """
     client = Client()
     if not client.session:

lucidicai/client.py

@@ -69,6 +69,10 @@ class Client:
 
     def set_provider(self, provider: BaseProvider) -> None:
         """Set the LLM provider to track"""
+        # Avoid duplicate provider registration of the same class
+        for existing in self.providers:
+            if type(existing) is type(provider):
+                return
         self.providers.append(provider)
         provider.override()
 
@@ -85,6 +89,7 @@ class Client:
         tags: Optional[list] = None,
         production_monitoring: Optional[bool] = False,
         session_id: Optional[str] = None,
+        experiment_id: Optional[str] = None,
     ) -> None:
         if session_id:
             # Check if it's a known session ID, maybe custom and maybe real
@@ -107,6 +112,7 @@ class Client:
             "session_name": session_name,
             "task": task,
             "mass_sim_id": mass_sim_id,
+            "experiment_id": experiment_id,
             "rubrics": rubrics,
             "tags": tags,
             "session_id": session_id
@@ -124,6 +130,7 @@ class Client:
             session_id=real_session_id,
             session_name=session_name,
             mass_sim_id=mass_sim_id,
+            experiment_id=experiment_id,
             task=task,
             rubrics=rubrics,
             tags=tags,
@@ -131,6 +138,16 @@ class Client:
         self.initialized = True
         return self.session.session_id
 
+    def create_event_for_session(self, session_id: str, **kwargs) -> str:
+        """Create an event for a specific session id without mutating global session.
+
+        This avoids cross-thread races by not switching the active session on
+        the singleton client. It constructs an ephemeral Session facade to send
+        requests under the provided session id.
+        """
+        temp_session = Session(agent_id=self.agent_id, session_id=session_id)
+        return temp_session.create_event(**kwargs)
+
     def continue_session(self, session_id: str):
         if session_id in self.custom_session_id_translations:
             session_id = self.custom_session_id_translations[session_id]
@@ -146,7 +163,8 @@ class Client:
             agent_id=self.agent_id,
             session_id=real_session_id
         )
-        logger.info(f"Session {data.get('session_name', '')} continuing...")
+        import logging as _logging
+        _logging.getLogger('Lucidic').info(f"Session {data.get('session_name', '')} continuing...")
         return self.session.session_id
 
     def init_mass_sim(self, **kwargs) -> str:

lucidicai/context.py

@@ -0,0 +1,119 @@
+"""Async-safe context helpers for session (and step, extensible).
+
+This module exposes context variables and helpers to bind a Lucidic
+session to the current execution context (threads/async tasks), so
+OpenTelemetry spans can be deterministically attributed to the correct
+session under concurrency.
+"""
+
+from contextlib import contextmanager, asynccontextmanager
+import contextvars
+from typing import Optional, Iterator, AsyncIterator, Callable, Any, Dict
+import logging
+import os
+
+
+# Context variable for the active Lucidic session id
+current_session_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
+    "lucidic.session_id", default=None
+)
+
+
+def set_active_session(session_id: Optional[str]) -> None:
+    """Bind the given session id to the current execution context."""
+    current_session_id.set(session_id)
+
+
+def clear_active_session() -> None:
+    """Clear any active session binding in the current execution context."""
+    current_session_id.set(None)
+
+
+@contextmanager
+def bind_session(session_id: str) -> Iterator[None]:
+    """Context manager to temporarily bind an active session id."""
+    token = current_session_id.set(session_id)
+    try:
+        yield
+    finally:
+        current_session_id.reset(token)
+
+
+@asynccontextmanager
+async def bind_session_async(session_id: str) -> AsyncIterator[None]:
+    """Async context manager to temporarily bind an active session id."""
+    token = current_session_id.set(session_id)
+    try:
+        yield
+    finally:
+        current_session_id.reset(token)
+
+
+@contextmanager
+def session(**init_params) -> Iterator[None]:
+    """All-in-one context manager: init → bind → yield → clear → end.
+
+    Notes:
+    - Ignores any provided auto_end parameter and ends the session on context exit.
+    - If LUCIDIC_DEBUG is true, logs a warning about ignoring auto_end.
+    """
+    # Lazy import to avoid circular imports
+    import lucidicai as lai  # type: ignore
+
+    # Force auto_end to False inside a context manager to control explicit end
+    user_auto_end = init_params.get('auto_end', None)
+    init_params = dict(init_params)
+    init_params['auto_end'] = False
+
+    if os.getenv('LUCIDIC_DEBUG', 'False') == 'True' and user_auto_end is not None:
+        logging.getLogger('Lucidic').warning('session(...) ignores auto_end and will end the session at context exit')
+
+    session_id = lai.init(**init_params)
+    token = current_session_id.set(session_id)
+    try:
+        yield
+    finally:
+        current_session_id.reset(token)
+        try:
+            lai.end_session()
+        except Exception:
+            # Avoid masking the original exception from the with-block
+            pass
+
+
+@asynccontextmanager
+async def session_async(**init_params) -> AsyncIterator[None]:
+    """Async counterpart of session(...)."""
+    import lucidicai as lai  # type: ignore
+
+    user_auto_end = init_params.get('auto_end', None)
+    init_params = dict(init_params)
+    init_params['auto_end'] = False
+
+    if os.getenv('LUCIDIC_DEBUG', 'False') == 'True' and user_auto_end is not None:
+        logging.getLogger('Lucidic').warning('session_async(...) ignores auto_end and will end the session at context exit')
+
+    session_id = lai.init(**init_params)
+    token = current_session_id.set(session_id)
+    try:
+        yield
+    finally:
+        current_session_id.reset(token)
+        try:
+            lai.end_session()
+        except Exception:
+            pass
+
+
+def run_session(fn: Callable[..., Any], *fn_args: Any, init_params: Optional[Dict[str, Any]] = None, **fn_kwargs: Any) -> Any:
+    """Run a callable within a full Lucidic session lifecycle context."""
+    with session(**(init_params or {})):
+        return fn(*fn_args, **fn_kwargs)
+
+
+def run_in_session(session_id: str, fn: Callable[..., Any], *fn_args: Any, **fn_kwargs: Any) -> Any:
+    """Run a callable with a bound session id. Does not end the session."""
+    with bind_session(session_id):
+        return fn(*fn_args, **fn_kwargs)
+
+

lucidicai/decorators.py

@@ -5,6 +5,7 @@ import inspect
 import json
 import logging
 from typing import Any, Callable, Optional, TypeVar, Union
+from collections.abc import Iterable
 
 from .client import Client
 from .errors import LucidicNotInitializedError
@@ -219,27 +220,44 @@ def event(
             
             # Build event description from inputs if not provided
             event_desc = description
+            function_name = func.__name__
+            
+            # Get function signature
+            sig = inspect.signature(func)
+            bound_args = sig.bind(*args, **kwargs)
+            bound_args.apply_defaults()
+
+            def serialize(value):
+                if isinstance(value, str):
+                    return value
+                if isinstance(value, int):
+                    return value
+                if isinstance(value, float):
+                    return value
+                if isinstance(value, bool):
+                    return value
+                if isinstance(value, dict):
+                    return {k: serialize(v) for k, v in value.items()}
+                if isinstance(value, Iterable):
+                    return [serialize(v) for v in value]
+                return str(value)
+
+            # Construct JSONable object of args
+            args_dict = {
+                param_name: serialize(param_value)  # Recursive - maybe change later
+                for param_name, param_value in bound_args.arguments.items()
+            }
+            
             if not event_desc:
-                # Get function signature
-                sig = inspect.signature(func)
-                bound_args = sig.bind(*args, **kwargs)
-                bound_args.apply_defaults()
-                
-                # Create string representation of inputs
-                input_parts = []
-                for param_name, param_value in bound_args.arguments.items():
-                    try:
-                        input_parts.append(f"{param_name}={repr(param_value)}")
-                    except Exception:
-                        input_parts.append(f"{param_name}=<{type(param_value).__name__}>")
-                
-                event_desc = f"{func.__name__}({', '.join(input_parts)})"
+                event_desc = f"Function {function_name}({json.dumps(args_dict)})"
             
             # Create the event
             event_id = create_event(
                 description=event_desc,
                 model=model,
-                cost_added=cost_added
+                cost_added=cost_added,
+                function_name=function_name,
+                arguments=args_dict, 
             )
             tok = _current_event.set(event_id)
             try:

lucidicai/event.py

@@ -48,7 +48,9 @@ class Event:
             "cost_added": kwargs.get("cost_added", None),
             "model": kwargs.get("model", None), 
             "nscreenshots": len(self.screenshots) + num_new_screenshots,
-            "duration": kwargs.get("duration", None)
+            "duration": kwargs.get("duration", None),
+            "function_name": kwargs.get("function_name", None),
+            "arguments": kwargs.get("arguments", None),
         }
     
     def _upload_screenshots(self, **kwargs) -> None: