lucidicai 1.2.15__py3-none-any.whl → 1.2.17__py3-none-any.whl

lucidicai/__init__.py

@@ -7,14 +7,24 @@ from typing import List, Literal, Optional
 from .client import Client
 from .errors import APIKeyVerificationError, InvalidOperationError, LucidicNotInitializedError, PromptError
 from .event import Event
-from .providers.anthropic_handler import AnthropicHandler
-from .providers.langchain import LucidicLangchainHandler
-from .providers.openai_handler import OpenAIHandler
-from .providers.openai_agents_handler import OpenAIAgentsHandler
-from .providers.pydantic_ai_handler import PydanticAIHandler
 from .session import Session
 from .step import Step
 
+# Import OpenTelemetry-based handlers
+from .telemetry.otel_handlers import (
+    OTelOpenAIHandler,
+    OTelAnthropicHandler,
+    OTelLangChainHandler,
+    OTelPydanticAIHandler,
+    OTelOpenAIAgentsHandler
+)
+
+# Import telemetry manager
+from .telemetry.otel_init import LucidicTelemetry
+
+# Import decorators
+from .decorators import step, event
+
 ProviderType = Literal["openai", "anthropic", "langchain", "pydantic_ai", "openai_agents"]
 
 # Configure logging
@@ -37,31 +47,33 @@ def _setup_providers(client: Client, providers: List[ProviderType]) -> None:
     # Track which providers have been set up to avoid duplication
     setup_providers = set()
     
+    # Initialize telemetry if using OpenTelemetry
+    if providers:
+        telemetry = LucidicTelemetry()
+        if not telemetry.is_initialized():
+            telemetry.initialize(agent_id=client.agent_id)
+    
     for provider in providers:
         if provider in setup_providers:
             continue
             
         if provider == "openai":
-            client.set_provider(OpenAIHandler())
+            client.set_provider(OTelOpenAIHandler())
             setup_providers.add("openai")
         elif provider == "anthropic":
-            client.set_provider(AnthropicHandler())
+            client.set_provider(OTelAnthropicHandler())
             setup_providers.add("anthropic")
         elif provider == "langchain":
+            client.set_provider(OTelLangChainHandler())
             logger.info("For LangChain, make sure to create a handler and attach it to your top-level Agent class.")
             setup_providers.add("langchain")
         elif provider == "pydantic_ai":
-            client.set_provider(PydanticAIHandler())
+            client.set_provider(OTelPydanticAIHandler())
             setup_providers.add("pydantic_ai")
         elif provider == "openai_agents":
             try:
-                # For OpenAI Agents SDK, we want both handlers
-                client.set_provider(OpenAIAgentsHandler())
+                client.set_provider(OTelOpenAIAgentsHandler())
                 setup_providers.add("openai_agents")
-                # Also enable OpenAI handler if not already set up
-                if "openai" not in setup_providers:
-                    client.set_provider(OpenAIHandler())
-                    setup_providers.add("openai")
             except Exception as e:
                 logger.error(f"Failed to set up OpenAI Agents provider: {e}")
                 raise
@@ -87,11 +99,8 @@ __all__ = [
     'LucidicNotInitializedError',
     'PromptError',
     'InvalidOperationError',
-    'LucidicLangchainHandler',
-    'AnthropicHandler',
-    'OpenAIHandler',
-    'OpenAIAgentsHandler',
-    'PydanticAIHandler'
+    'step',
+    'event',
 ]
 
 
@@ -101,9 +110,12 @@ def init(
     agent_id: Optional[str] = None,
     task: Optional[str] = None,
     providers: Optional[List[ProviderType]] = [],
+    production_monitoring: Optional[bool] = False,
     mass_sim_id: Optional[str] = None,
     rubrics: Optional[list] = None,
     tags: Optional[list] = None,
+    masking_function = None,
+    auto_end: Optional[bool] = True,
 ) -> str:
     """
     Initialize the Lucidic client.
@@ -117,6 +129,8 @@ def init(
         mass_sim_id: Optional mass simulation ID, if session is to be part of a mass simulation.
         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.
+        auto_end: If True, automatically end the session on process exit. Defaults to True.
     
     Raises:
         InvalidOperationError: If the client is already initialized.
@@ -137,6 +151,17 @@ def init(
     if not getattr(client, 'initialized', False):
         client = Client(lucidic_api_key=lucidic_api_key, agent_id=agent_id)
     
+    if not production_monitoring:
+        production_monitoring = os.getenv("LUCIDIC_PRODUCTION_MONITORING", False)
+        if production_monitoring == "True":
+            production_monitoring = True
+        else:
+            production_monitoring = False
+    
+    # Handle auto_end with environment variable support
+    if auto_end is None:
+        auto_end = os.getenv("LUCIDIC_AUTO_END", "True").lower() == "true"
+    
     # Set up providers
     _setup_providers(client, providers)
     session_id = client.init_session(
@@ -144,8 +169,15 @@ def init(
         mass_sim_id=mass_sim_id,
         task=task,
         rubrics=rubrics,
-        tags=tags
+        tags=tags,
+        production_monitoring=production_monitoring,
     )
+    if masking_function:
+        client.masking_function = masking_function
+    
+    # Set the auto_end flag on the client
+    client.auto_end = auto_end
+    
     logger.info("Session initialized successfully")
     return session_id
 
@@ -154,7 +186,9 @@ def continue_session(
     session_id: str,
     lucidic_api_key: Optional[str] = None,
     agent_id: Optional[str] = None,
-    providers: Optional[List[ProviderType]] = []
+    providers: Optional[List[ProviderType]] = [],
+    masking_function = None,
+    auto_end: Optional[bool] = True,
 ):
     if lucidic_api_key is None:
         lucidic_api_key = os.getenv("LUCIDIC_API_KEY", None)
@@ -174,9 +208,19 @@ def continue_session(
             agent_id=agent_id,
         )
     
+    # Handle auto_end with environment variable support
+    if auto_end is None:
+        auto_end = os.getenv("LUCIDIC_AUTO_END", "True").lower() == "true"
+    
     # Set up providers
     _setup_providers(client, providers)
     session_id = client.continue_session(session_id=session_id)
+    if masking_function:
+        client.masking_function = masking_function
+    
+    # Set the auto_end flag on the client
+    client.auto_end = auto_end
+    
     logger.info(f"Session {session_id} continuing...")
     return session_id  # For consistency
 
@@ -233,9 +277,55 @@ def reset_sdk() -> None:
     client = Client()
     if not client.initialized:
         return
+    
+    # Shutdown OpenTelemetry if it was initialized
+    telemetry = LucidicTelemetry()
+    if telemetry.is_initialized():
+        telemetry.uninstrument_all()
+    
     client.clear()
 
 
+def _cleanup_telemetry():
+    """Cleanup function for OpenTelemetry shutdown"""
+    try:
+        telemetry = LucidicTelemetry()
+        if telemetry.is_initialized():
+            telemetry.uninstrument_all()
+            logger.info("OpenTelemetry instrumentation cleaned up")
+    except Exception as e:
+        logger.error(f"Error during telemetry cleanup: {e}")
+
+
+def _auto_end_session():
+    """Automatically end session on exit if auto_end is enabled"""
+    try:
+        client = Client()
+        if hasattr(client, 'auto_end') and client.auto_end and client.session and not client.session.is_finished:
+            logger.info("Auto-ending active session on exit")
+            end_session()
+    except Exception as e:
+        logger.debug(f"Error during auto-end session: {e}")
+
+
+def _signal_handler(signum, frame):
+    """Handle interruption signals"""
+    _auto_end_session()
+    _cleanup_telemetry()
+    # Re-raise the signal for default handling
+    signal.signal(signum, signal.SIG_DFL)
+    os.kill(os.getpid(), signum)
+
+
+# Register cleanup functions (auto-end runs first due to LIFO order)
+atexit.register(_cleanup_telemetry)
+atexit.register(_auto_end_session)
+
+# Register signal handlers for graceful shutdown
+signal.signal(signal.SIGINT, _signal_handler)
+signal.signal(signal.SIGTERM, _signal_handler)
+
+
 def create_mass_sim(
     mass_sim_name: str,
     total_num_sessions: int,

lucidicai/client.py

@@ -4,12 +4,13 @@ from datetime import datetime, timezone
 from typing import Optional, Tuple
 
 import requests
+import logging
 from requests.adapters import HTTPAdapter, Retry
 from urllib3.util import Retry
 
 
 from .errors import APIKeyVerificationError, InvalidOperationError, LucidicNotInitializedError
-from .providers.base_providers import BaseProvider
+from .telemetry.base_provider import BaseProvider 
 from .session import Session
 from .singleton import singleton, clear_singletons
 
@@ -30,6 +31,8 @@ class Client:
         self.providers = []
         self.api_key = lucidic_api_key
         self.agent_id = agent_id
+        self.masking_function = None
+        self.auto_end = False  # Default to False until explicitly set during init
         self.request_session = requests.Session()
         retry_cfg = Retry(
             total=3,                     # 3 attempts in total
@@ -73,7 +76,8 @@ class Client:
         mass_sim_id: Optional[str] = None,
         task: Optional[str] = None,
         rubrics: Optional[list] = None,
-        tags: Optional[list] = None
+        tags: Optional[list] = None,
+        production_monitoring: Optional[bool] = False
     ) -> None:
         self.session = Session(
             agent_id=self.agent_id,
@@ -81,12 +85,13 @@ class Client:
             mass_sim_id=mass_sim_id,
             task=task,
             rubrics=rubrics,
-            tags=tags
+            tags=tags,
+            production_monitoring=production_monitoring
         )
         self.initialized = True
         return self.session.session_id
 
-    def continue_session(self, session_id: str) -> None:
+    def continue_session(self, session_id: str):
         self.session = Session(
             agent_id=self.agent_id,
             session_id=session_id
@@ -147,4 +152,16 @@ class Client:
             response.raise_for_status()
         except requests.exceptions.HTTPError as e:
             raise InvalidOperationError(f"Request to Lucidic AI Backend failed: {e.response.text}")
-        return response.json()
+        return response.json()
+
+    def mask(self, data):
+        if not self.masking_function:
+            return data
+        if not data:
+            return data
+        try:
+            return self.masking_function(data)
+        except Exception as e:
+            logger = logging.getLogger('Lucidic')
+            logger.error(f"Error in custom masking function: {repr(e)}")
+            return "<Error in custom masking function, this is a fully-masked placeholder>"

lucidicai/decorators.py

@@ -0,0 +1,357 @@
+"""Decorators for the Lucidic SDK to simplify step and event tracking."""
+import functools
+import contextvars
+import inspect
+import json
+import logging
+from typing import Any, Callable, Optional, TypeVar, Union
+
+from .client import Client
+from .errors import LucidicNotInitializedError
+
+logger = logging.getLogger("Lucidic")
+
+F = TypeVar('F', bound=Callable[..., Any])
+
+# Create context variables to store the current step and event
+_current_step = contextvars.ContextVar("current_step", default=None)
+_current_event = contextvars.ContextVar("current_event", default=None)
+
+def get_decorator_step():
+    return _current_step.get()
+
+def get_decorator_event():
+    return _current_event.get()
+
+
+def step(
+    state: Optional[str] = None,
+    action: Optional[str] = None,
+    goal: Optional[str] = None,
+    screenshot_path: Optional[str] = None,
+    eval_score: Optional[float] = None,
+    eval_description: Optional[str] = None
+) -> Callable[[F], F]:
+    """
+    Decorator that wraps a function with step tracking.
+    
+    The decorated function will be wrapped with create_step() at the start
+    and end_step() at the end, ensuring proper cleanup even on exceptions.
+    
+    Args:
+        state: State description for the step
+        action: Action description for the step  
+        goal: Goal description for the step
+        eval_score: Evaluation score for the step
+        eval_description: Evaluation description for the step
+        
+    Example:
+        @lai.step(
+            state="Processing user input",
+            action="Validate and parse request",
+            goal="Extract intent from user message"
+        )
+        def process_user_input(message: str) -> dict:
+            # Function logic here
+            return parsed_intent
+    """
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            # Check if SDK is initialized
+            try:
+                client = Client()
+                if not client.session:
+                    # No active session, run function normally
+                    logger.warning("No active session, running function normally")
+                    return func(*args, **kwargs)
+            except LucidicNotInitializedError:
+                # SDK not initialized, run function normally
+                logger.warning("Lucidic not initialized, running function normally")
+                return func(*args, **kwargs)
+            
+            # Create the step
+            step_params = {
+                'state': state,
+                'action': action,
+                'goal': goal,
+                'screenshot_path': screenshot_path,
+                'eval_score': eval_score,
+                'eval_description': eval_description
+            }
+            # Remove None values
+            step_params = {k: v for k, v in step_params.items() if v is not None}
+            
+            # Import here to avoid circular imports
+            from . import create_step, end_step
+            step_id = create_step(**step_params)
+            tok = _current_step.set(step_id)
+            
+            try:
+                # Execute the wrapped function
+                result = func(*args, **kwargs)
+                # End step successfully
+                end_step(step_id=step_id)
+                _current_step.reset(tok)
+                return result
+            except Exception as e:
+                # End step with error indication
+                try:
+                    end_step(
+                        step_id=step_id,
+                        eval_score=0.0,
+                        eval_description=f"Step failed with error: {str(e)}"
+                    )
+                    _current_step.reset(tok)
+                except Exception:
+                    # If end_step fails, just log it
+                    logger.error(f"Failed to end step {step_id} after error")
+                raise
+        
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            # Check if SDK is initialized
+            try:
+                client = Client()
+                if not client.session:
+                    # No active session, run function normally
+                    logger.warning("No active session, running function normally")
+                    return await func(*args, **kwargs)
+            except LucidicNotInitializedError:
+                # SDK not initialized, run function normally
+                logger.warning("Lucidic not initialized, running function normally")
+                return await func(*args, **kwargs)
+            
+            # Create the step
+            step_params = {
+                'state': state,
+                'action': action,
+                'goal': goal,
+                'screenshot_path': screenshot_path,
+                'eval_score': eval_score,
+                'eval_description': eval_description
+            }
+            # Remove None values
+            step_params = {k: v for k, v in step_params.items() if v is not None}
+            
+            # Import here to avoid circular imports
+            from . import create_step, end_step
+            
+            step_id = create_step(**step_params)
+            tok = _current_step.set(step_id)
+            try:
+                # Execute the wrapped function
+                result = await func(*args, **kwargs)
+                # End step successfully
+                end_step(step_id=step_id)
+                _current_step.reset(tok)
+                return result
+            except Exception as e:
+                # End step with error indication
+                try:
+                    end_step(
+                        step_id=step_id,
+                        eval_score=0.0,
+                        eval_description=f"Step failed with error: {str(e)}"
+                    )
+                    _current_step.reset(tok)
+                except Exception:
+                    # If end_step fails, just log it
+                    logger.error(f"Failed to end step {step_id} after error")
+                raise
+        
+        # Return appropriate wrapper based on function type
+        if inspect.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return sync_wrapper
+    
+    return decorator
+
+
+### -- TODO -- Updating even within function causes function result to not be recorded.
+def event(
+    description: Optional[str] = None,
+    result: Optional[str] = None,
+    model: Optional[str] = None,
+    cost_added: Optional[float] = 0
+) -> Callable[[F], F]:
+    """
+    Decorator that creates an event for a function call.
+    
+    The decorated function will create an event that captures:
+    - Function inputs (as string representation) if description not provided
+    - Function output (as string representation) if result not provided
+    
+    LLM calls within the function will create their own events as normal.
+    
+    Args:
+        description: Custom description for the event. If not provided,
+                    will use string representation of function inputs
+        result: Custom result for the event. If not provided,
+                will use string representation of function output
+        model: Model name if this function represents a model call (default: None)
+        cost_added: Cost to add for this event (default: 0)
+        
+    Example:
+        @lai.event(description="Parse user query", model="custom-parser")
+        def parse_query(query: str) -> dict:
+            # Function logic here
+            return {"intent": "search", "query": query}
+    """
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            # Check if SDK is initialized
+            try:
+                client = Client()
+                if not client.session:
+                    # No active session, run function normally
+                    logger.warning("No active session, running function normally")
+                    return func(*args, **kwargs)
+            except (LucidicNotInitializedError, AttributeError):
+                # SDK not initialized or no session, run function normally
+                logger.warning("Lucidic not initialized, running function normally")
+                return func(*args, **kwargs)
+            
+            # Import here to avoid circular imports
+            from . import create_event, end_event
+            
+            # Build event description from inputs if not provided
+            event_desc = description
+            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)})"
+            
+            # Create the event
+            event_id = create_event(
+                description=event_desc,
+                model=model,
+                cost_added=cost_added
+            )
+            tok = _current_event.set(event_id)
+            try:
+                # Execute the wrapped function
+                function_result = func(*args, **kwargs)
+                
+                # Build event result from output if not provided
+                event_result = result
+                if not event_result:
+                    try:
+                        event_result = repr(function_result)
+                    except Exception:
+                        event_result = str(function_result)
+                
+                # Update and end the event
+                end_event(
+                    event_id=event_id,
+                    result=event_result,
+                )
+                _current_event.reset(tok)
+                return function_result
+                
+            except Exception as e:
+                # Update event with error
+                try:
+                    end_event(
+                        event_id=event_id,
+                        result=f"Error: {str(e)}",
+                    )
+                    _current_event.reset(tok)
+                except Exception:
+                    logger.error(f"Failed to end event {event_id} after error")
+                raise
+        
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            # Check if SDK is initialized
+            try:
+                client = Client()
+                if not client.session:
+                    # No active session, run function normally
+                    logger.warning("No active session, running function normally")
+                    return await func(*args, **kwargs)
+            except (LucidicNotInitializedError, AttributeError):
+                # SDK not initialized or no session, run function normally
+                logger.warning("Lucidic not initialized, running function normally")
+                return await func(*args, **kwargs)
+            
+            # Import here to avoid circular imports
+            from . import create_event, end_event
+            
+            # Build event description from inputs if not provided
+            event_desc = description
+            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)})"
+            
+            # Create the event
+            event_id = create_event(
+                description=event_desc,
+                model=model,
+                cost_added=cost_added
+            )
+            tok = _current_event.set(event_id)
+            try:
+                # Execute the wrapped function
+                function_result = await func(*args, **kwargs)
+                
+                # Build event result from output if not provided
+                event_result = result
+                if not event_result:
+                    try:
+                        event_result = repr(function_result)
+                    except Exception:
+                        event_result = str(function_result)
+                
+                # Update and end the event
+                end_event(
+                    event_id=event_id,
+                    result=event_result,
+                )
+                _current_event.reset(tok)
+                return function_result
+                
+            except Exception as e:
+                # Update event with error
+                try:
+                    end_event(
+                        event_id=event_id,
+                        result=f"Error: {str(e)}",
+                    )
+                    _current_event.reset(tok)
+                except Exception:
+                    logger.error(f"Failed to end event {event_id} after error")
+                raise
+        
+        # Return appropriate wrapper based on function type
+        if inspect.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return sync_wrapper
+    
+    return decorator

lucidicai/event.py

@@ -40,8 +40,8 @@ class Event:
             self.is_finished = kwargs['is_finished']
         request_data = {
             "event_id": self.event_id,
-            "description": kwargs.get("description", None),
-            "result": kwargs.get("result", None),
+            "description": Client().mask(kwargs.get("description", None)),
+            "result": Client().mask(kwargs.get("result", None)),
             "is_finished": self.is_finished, 
             "cost_added": kwargs.get("cost_added", None),
             "model": kwargs.get("model", None), 

lucidicai/image_upload.py

@@ -85,4 +85,27 @@ def screenshot_path_to_jpeg(screenshot_path):
     buffered = io.BytesIO()
     img.save(buffered, format="JPEG")
     img_byte = buffered.getvalue()
-    return base64.b64encode(img_byte).decode('utf-8')
+    return base64.b64encode(img_byte).decode('utf-8')
+
+def extract_base64_images(data):
+    """Extract base64 image URLs from various data structures
+    
+    Args:
+        data: Can be a string, dict, list, or nested structure containing image data
+        
+    Returns:
+        List of base64 image data URLs (data:image/...)
+    """
+    images = []
+    
+    if isinstance(data, str):
+        if data.startswith('data:image'):
+            images.append(data)
+    elif isinstance(data, dict):
+        for value in data.values():
+            images.extend(extract_base64_images(value))
+    elif isinstance(data, list):
+        for item in data:
+            images.extend(extract_base64_images(item))
+            
+    return images