lucidicai 1.2.16__py3-none-any.whl → 1.2.18__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,16 +99,14 @@ __all__ = [
     'LucidicNotInitializedError',
     'PromptError',
     'InvalidOperationError',
-    'LucidicLangchainHandler',
-    'AnthropicHandler',
-    'OpenAIHandler',
-    'OpenAIAgentsHandler',
-    'PydanticAIHandler'
+    'step',
+    'event',
 ]
 
 
 def init(
-    session_name: str,
+    session_name: Optional[str] = None,
+    session_id: Optional[str] = None,
     lucidic_api_key: Optional[str] = None,
     agent_id: Optional[str] = None,
     task: Optional[str] = None,
@@ -106,12 +116,14 @@ def init(
     rubrics: Optional[list] = None,
     tags: Optional[list] = None,
     masking_function = None,
+    auto_end: Optional[bool] = True,
 ) -> str:
     """
     Initialize the Lucidic client.
     
     Args:
-        session_name: The name of the session.
+        session_name: The display name of the session.
+        session_id: Custom ID of the session. If not provided, a random ID will be generated.
         lucidic_api_key: API key for authentication. If not provided, will use the LUCIDIC_API_KEY environment variable.
         agent_id: Agent ID. If not provided, will use the LUCIDIC_AGENT_ID environment variable.
         task: Task description.
@@ -120,6 +132,7 @@ def init(
         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.
@@ -147,6 +160,10 @@ def init(
         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(
@@ -156,9 +173,14 @@ def init(
         rubrics=rubrics,
         tags=tags,
         production_monitoring=production_monitoring,
+        custom_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("Session initialized successfully")
     return session_id
 
@@ -169,6 +191,7 @@ def continue_session(
     agent_id: Optional[str] = None,
     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)
@@ -178,21 +201,27 @@ def continue_session(
         agent_id = os.getenv("LUCIDIC_AGENT_ID", None)
         if agent_id is None:
             raise APIKeyVerificationError("Lucidic agent ID not specified. Make sure to either pass your agent ID into lai.init() or set the LUCIDIC_AGENT_ID environment variable.")
-    try:
-        client = Client()
-        if client.session:
-            raise InvalidOperationError("[Lucidic] Session already in progress. Please call lai.end_session() first.")
-    except LucidicNotInitializedError:
-        client = Client(
-            lucidic_api_key=lucidic_api_key,
-            agent_id=agent_id,
-        )
+    
+    client = Client()
+    if client.session:
+        raise InvalidOperationError("[Lucidic] Session already in progress. Please call lai.end_session() or lai.reset_sdk() first.")
+    # if not yet initialized or still the NullClient -> create a real client when init is called
+    if not getattr(client, 'initialized', False):
+        client = Client(lucidic_api_key=lucidic_api_key, 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
 
@@ -249,9 +278,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

@@ -10,7 +10,7 @@ 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
 
@@ -32,6 +32,7 @@ class Client:
         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
@@ -76,7 +77,8 @@ class Client:
         task: Optional[str] = None,
         rubrics: Optional[list] = None,
         tags: Optional[list] = None,
-        production_monitoring: Optional[bool] = False
+        production_monitoring: Optional[bool] = False,
+        custom_session_id: Optional[str] = None,
     ) -> None:
         self.session = Session(
             agent_id=self.agent_id,
@@ -85,16 +87,20 @@ class Client:
             task=task,
             rubrics=rubrics,
             tags=tags,
-            production_monitoring=production_monitoring
+            production_monitoring=production_monitoring,
+            custom_session_id=custom_session_id
         )
         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
         )
+        if self.session.session_id != session_id:
+            # Custom session ID provided
+            self.session.custom_session_id = session_id
         self.initialized = True
         return self.session.session_id
 

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/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