lucidicai 1.3.2__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',
 ]
 
 
@@ -122,6 +311,7 @@ def init(
     tags: Optional[list] = None,
     masking_function = None,
     auto_end: Optional[bool] = True,
+    capture_uncaught: Optional[bool] = True,
 ) -> str:
     """
     Initialize the Lucidic client.
@@ -189,6 +379,17 @@ def init(
     
     # 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
@@ -232,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
 
 
@@ -252,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(
@@ -274,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:
@@ -330,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

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()
 
@@ -134,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]
@@ -149,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/model_pricing.py

@@ -3,6 +3,12 @@ import logging
 logger = logging.getLogger("Lucidic")
 
 MODEL_PRICING = {
+
+    # OpenAI GPT-5 Series (Verified 2025)
+    "gpt-5": {"input": 10.0, "output": 10.0},
+    "gpt-5-mini": {"input": 0.250, "output": 2.0},
+    "gpt-5-nano": {"input": 0.05, "output": 0.4},
+
     # OpenAI GPT-4o Series (Verified 2025)
     "gpt-4o": {"input": 2.5, "output": 10.0},
     "gpt-4o-mini": {"input": 0.15, "output": 0.6},
@@ -190,6 +196,7 @@ PROVIDER_AVERAGES = {
     "together": {"input": 0.15, "output": 0.15},    # Together AI average
     "perplexity": {"input": 0.4, "output": 1.5},    # Perplexity average
     "grok": {"input": 2.4, "output": 12},           # Grok average
+    "groq": {"input": 0.3, "output": 0.6},          # Groq average (placeholder)
 }
 
 def get_provider_from_model(model: str) -> str:
@@ -218,6 +225,8 @@ def get_provider_from_model(model: str) -> str:
         return "perplexity"
     elif any(grok in model_lower for grok in ["grok", "xAI"]):
         return "grok"
+    elif "groq" in model_lower:
+        return "groq"
     else:
         return "unknown"
 
@@ -228,6 +237,8 @@ def normalize_model_name(model: str) -> str:
     model_lower = model.lower()
     # Remove provider prefixes (generalizable pattern: any_provider/)
     model_lower = re.sub(r'^[^/]+/', '', model_lower)
+    # Strip Google/Vertex prefixes
+    model_lower = model_lower.replace('publishers/google/models/', '').replace('models/', '')
     
     # Strip date suffixes (20240229, 20241022, etc.) but preserve model versions like o1-mini, o3-mini
     # Pattern: remove -YYYYMMDD or -YYYY-MM-DD at the end

lucidicai/telemetry/lucidic_exporter.py

@@ -8,6 +8,7 @@ from opentelemetry.trace import StatusCode
 from opentelemetry.semconv_ai import SpanAttributes
 
 from lucidicai.client import Client
+from lucidicai.context import current_session_id
 from lucidicai.model_pricing import calculate_cost
 from lucidicai.image_upload import extract_base64_images
 
@@ -28,9 +29,6 @@ class LucidicSpanExporter(SpanExporter):
         """Export spans by converting them to Lucidic events"""
         try:
             client = Client()
-            if not client.session:
-                logger.debug("No active session, skipping span export")
-                return SpanExportResult.SUCCESS
                 
             for span in spans:
                 self._process_span(span, client)
@@ -100,6 +98,19 @@ class LucidicSpanExporter(SpanExporter):
                    attributes.get(SpanAttributes.LLM_REQUEST_MODEL) or \
                    attributes.get('gen_ai.request.model') or 'unknown'
             
+            # Resolve target session id for this span
+            target_session_id = attributes.get('lucidic.session_id')
+            if not target_session_id:
+                try:
+                    target_session_id = current_session_id.get(None)
+                except Exception:
+                    target_session_id = None
+            if not target_session_id:
+                if getattr(client, 'session', None) and getattr(client.session, 'session_id', None):
+                    target_session_id = client.session.session_id
+            if not target_session_id:
+                return None
+
             # Create event
             event_kwargs = {
                 'description': description,
@@ -115,7 +126,7 @@ class LucidicSpanExporter(SpanExporter):
             if step_id:
                 event_kwargs['step_id'] = step_id
                 
-            return client.session.create_event(**event_kwargs)
+            return client.create_event_for_session(target_session_id, **event_kwargs)
             
         except Exception as e:
             logger.error(f"Failed to create event from span: {e}")
@@ -143,6 +154,7 @@ class LucidicSpanExporter(SpanExporter):
             if cost is not None:
                 update_kwargs['cost_added'] = cost
                 
+            # Route update to the same session; event_id is globally unique so server resolves it
             client.session.update_event(**update_kwargs)
             
         except Exception as e: