lucidicai 1.3.2__py3-none-any.whl → 2.0.1__py3-none-any.whl

lucidicai/telemetry/extract.py

@@ -0,0 +1,192 @@
+"""Extraction utilities matching TypeScript SDK for span attribute processing."""
+import json
+from typing import List, Dict, Any, Optional
+
+
+def detect_is_llm_span(span) -> bool:
+    """Check if span is LLM-related - matches TypeScript logic."""
+    name = (span.name or "").lower()
+    patterns = ['openai', 'anthropic', 'chat', 'completion', 'embedding', 'llm', 
+                'gemini', 'claude', 'bedrock', 'vertex', 'cohere', 'groq']
+    
+    if any(p in name for p in patterns):
+        return True
+    
+    if hasattr(span, 'attributes') and span.attributes:
+        for key in span.attributes:
+            if isinstance(key, str) and (key.startswith('gen_ai.') or key.startswith('llm.')):
+                return True
+    
+    return False
+
+
+def extract_images(attrs: Dict[str, Any]) -> List[str]:
+    """Extract images from span attributes - matches TypeScript logic.
+    
+    Looks for images in:
+    - gen_ai.prompt.{i}.content arrays with image_url items
+    - Direct image attributes
+    """
+    images = []
+    
+    # Check indexed prompt content for images (gen_ai.prompt.{i}.content)
+    for i in range(50):
+        key = f"gen_ai.prompt.{i}.content"
+        if key in attrs:
+            content = attrs[key]
+            
+            # Parse if JSON string
+            if isinstance(content, str):
+                try:
+                    content = json.loads(content)
+                except:
+                    continue
+            
+            # Extract images from content array
+            if isinstance(content, list):
+                for item in content:
+                    if isinstance(item, dict):
+                        if item.get("type") == "image_url":
+                            image_url = item.get("image_url", {})
+                            if isinstance(image_url, dict):
+                                url = image_url.get("url", "")
+                                if url.startswith("data:image"):
+                                    images.append(url)
+                        elif item.get("type") == "image":
+                            # Anthropic format
+                            source = item.get("source", {})
+                            if isinstance(source, dict):
+                                data = source.get("data", "")
+                                media_type = source.get("media_type", "image/jpeg")
+                                if data:
+                                    images.append(f"data:{media_type};base64,{data}")
+    
+    # Also check direct gen_ai.prompt list
+    prompt_list = attrs.get("gen_ai.prompt")
+    if isinstance(prompt_list, list):
+        for msg in prompt_list:
+            if isinstance(msg, dict):
+                content = msg.get("content")
+                if isinstance(content, list):
+                    for item in content:
+                        if isinstance(item, dict) and item.get("type") == "image_url":
+                            image_url = item.get("image_url", {})
+                            if isinstance(image_url, dict):
+                                url = image_url.get("url", "")
+                                if url.startswith("data:image"):
+                                    images.append(url)
+    
+    return images
+
+
+def extract_prompts(attrs: Dict[str, Any]) -> Optional[List[Dict]]:
+    """Extract prompts as message list from span attributes.
+    
+    Returns messages in format: [{"role": "user", "content": "..."}]
+    """
+    messages = []
+    
+    # Check indexed format (gen_ai.prompt.{i}.role/content)
+    for i in range(50):
+        role_key = f"gen_ai.prompt.{i}.role"
+        content_key = f"gen_ai.prompt.{i}.content"
+        
+        if role_key not in attrs and content_key not in attrs:
+            break
+            
+        role = attrs.get(role_key, "user")
+        content = attrs.get(content_key, "")
+        
+        # Parse content if it's JSON
+        if isinstance(content, str):
+            try:
+                parsed = json.loads(content)
+                content = parsed
+            except:
+                pass
+        
+        # Format content
+        if isinstance(content, list):
+            # Content array format (with text/image items)
+            text_parts = []
+            for item in content:
+                if isinstance(item, dict) and item.get("type") == "text":
+                    text_parts.append(item.get("text", ""))
+            if text_parts:
+                content = " ".join(text_parts)
+        
+        messages.append({"role": role, "content": content})
+    
+    if messages:
+        return messages
+    
+    # Check for direct message list
+    prompt_list = attrs.get("gen_ai.prompt") or attrs.get("gen_ai.messages")
+    if isinstance(prompt_list, list):
+        return prompt_list
+    
+    # Check AI SDK format
+    ai_prompt = attrs.get("ai.prompt.messages")
+    if isinstance(ai_prompt, str):
+        try:
+            parsed = json.loads(ai_prompt)
+            if isinstance(parsed, list):
+                return parsed
+        except:
+            pass
+    
+    return None
+
+
+def extract_completions(span, attrs: Dict[str, Any]) -> Optional[str]:
+    """Extract completion/response text from span attributes."""
+    completions = []
+    
+    # Check indexed format (gen_ai.completion.{i}.content)
+    for i in range(50):
+        key = f"gen_ai.completion.{i}.content"
+        if key not in attrs:
+            break
+        content = attrs[key]
+        if isinstance(content, str):
+            completions.append(content)
+        else:
+            try:
+                completions.append(json.dumps(content))
+            except:
+                completions.append(str(content))
+    
+    if completions:
+        return "\n".join(completions)
+    
+    # Check direct completion attribute
+    completion = attrs.get("gen_ai.completion") or attrs.get("llm.completions")
+    if isinstance(completion, str):
+        return completion
+    elif isinstance(completion, list) and completion:
+        return "\n".join(str(c) for c in completion)
+    
+    # Check AI SDK format
+    ai_completion = attrs.get("ai.response.text")
+    if isinstance(ai_completion, str):
+        return ai_completion
+    
+    # Check for error status
+    if hasattr(span, 'status'):
+        from opentelemetry.trace import StatusCode
+        if span.status.status_code == StatusCode.ERROR:
+            return f"Error: {span.status.description or 'Unknown error'}"
+    
+    return None
+
+
+def extract_model(attrs: Dict[str, Any]) -> Optional[str]:
+    """Extract model name from span attributes."""
+    return (
+        attrs.get("gen_ai.response.model") or
+        attrs.get("gen_ai.request.model") or
+        attrs.get("llm.response.model") or
+        attrs.get("llm.request.model") or
+        attrs.get("ai.model.id") or
+        attrs.get("ai.model.name")
+    )

lucidicai/telemetry/litellm_bridge.py

@@ -16,6 +16,7 @@ except ImportError:
 
 from lucidicai.client import Client
 from lucidicai.model_pricing import calculate_cost
+from lucidicai.context import current_parent_event_id
 
 logger = logging.getLogger("Lucidic")
 DEBUG = os.getenv("LUCIDIC_DEBUG", "False") == "True"
@@ -25,7 +26,7 @@ class LucidicLiteLLMCallback(CustomLogger):
     """
     Custom callback for LiteLLM that bridges to Lucidic's event system.
     
-    This callback integrates LiteLLM's logging with Lucidic's session/step/event hierarchy,
+    This callback integrates LiteLLM's logging with Lucidic's session/event hierarchy,
     enabling automatic tracking of all LiteLLM-supported providers.
     """
     
@@ -102,7 +103,7 @@ class LucidicLiteLLMCallback(CustomLogger):
                 traceback.print_exc()
     
     def log_success_event(self, kwargs, response_obj, start_time, end_time):
-        """Called on successful LLM completion"""
+        """Called on successful LLM completion -> create typed LLM_GENERATION event"""
         # Generate unique callback ID
         callback_id = f"success_{id(kwargs)}_{start_time}"
         self._register_callback(callback_id)
@@ -141,33 +142,30 @@ class LucidicLiteLLMCallback(CustomLogger):
             # Extract any images from multimodal requests
             images = self._extract_images_from_messages(messages)
             
-            # Create event
-            event_kwargs = {
-                "description": description,
-                "result": result,
-                "model": f"{provider}/{model}" if "/" not in model else model,
-                "is_finished": True
-            }
-            
-            if cost is not None:
-                event_kwargs["cost_added"] = cost
-                
-            if images:
-                event_kwargs["screenshots"] = images
-                
-            # Add metadata
-            metadata = {
-                "provider": provider,
-                "duration_ms": (end_time - start_time) * 1000,
-                "litellm": True
-            }
-            if usage:
-                metadata["usage"] = usage
-                
-            event_kwargs["metadata"] = metadata
-            
-            # Create the event
-            client.session.create_event(**event_kwargs)
+            # Create LLM_GENERATION typed event
+            parent_id = None
+            try:
+                parent_id = current_parent_event_id.get(None)
+            except Exception:
+                parent_id = None
+
+            # occurred_at/duration from datetimes
+            occ_dt = start_time if isinstance(start_time, datetime) else None
+            duration_secs = (end_time - start_time).total_seconds() if isinstance(start_time, datetime) and isinstance(end_time, datetime) else None
+
+            client.create_event(
+                type="llm_generation",
+                provider=provider,
+                model=model,
+                messages=messages,
+                output=result,
+                input_tokens=(usage or {}).get("prompt_tokens", 0),
+                output_tokens=(usage or {}).get("completion_tokens", 0),
+                cost=cost,
+                parent_event_id=parent_id,
+                occurred_at=occ_dt,
+                duration=duration_secs,
+            )
             
             if DEBUG:
                 logger.info(f"LiteLLM Bridge: Created event for {model} completion")
@@ -207,21 +205,24 @@ class LucidicLiteLLMCallback(CustomLogger):
             # Format error
             error_msg = str(response_obj) if response_obj else "Unknown error"
             
-            # Create failed event
-            event_kwargs = {
-                "description": description,
-                "result": f"Error: {error_msg}",
-                "model": f"{provider}/{model}" if "/" not in model else model,
-                "is_finished": True,
-                "metadata": {
-                    "provider": provider,
-                    "duration_ms": (end_time - start_time) * 1000,
-                    "litellm": True,
-                    "error": True
-                }
-            }
-            
-            client.session.create_event(**event_kwargs)
+            # Create error typed event under current parent if any
+            parent_id = None
+            try:
+                parent_id = current_parent_event_id.get(None)
+            except Exception:
+                parent_id = None
+            occ_dt = start_time if isinstance(start_time, datetime) else None
+            duration_secs = (end_time - start_time).total_seconds() if isinstance(start_time, datetime) and isinstance(end_time, datetime) else None
+
+            client.create_event(
+                type="error_traceback",
+                error=error_msg,
+                traceback="",
+                parent_event_id=parent_id,
+                occurred_at=occ_dt,
+                duration=duration_secs,
+                metadata={"provider": provider, "litellm": True}
+            )
             
             if DEBUG:
                 logger.info(f"LiteLLM Bridge: Created error event for {model}")
@@ -366,4 +367,38 @@ class LucidicLiteLLMCallback(CustomLogger):
                                 if url.startswith("data:image"):
                                     images.append(url)
                                     
-        return images
+        return images
+
+
+def setup_litellm_callback():
+    """Registers the LucidicLiteLLMCallback with LiteLLM if available.
+    
+    This function ensures only one instance of the callback is registered,
+    preventing duplicates across multiple SDK initializations.
+    """
+    try:
+        import litellm
+    except ImportError:
+        logger.info("[LiteLLM] litellm not installed, skipping callback setup")
+        return
+    
+    # Initialize callbacks list if needed
+    if not hasattr(litellm, 'callbacks'):
+        litellm.callbacks = []
+    elif litellm.callbacks is None:
+        litellm.callbacks = []
+    
+    # Check for existing registration to prevent duplicates
+    for existing in litellm.callbacks:
+        if isinstance(existing, LucidicLiteLLMCallback):
+            if DEBUG:
+                logger.debug("[LiteLLM] Callback already registered")
+            return
+    
+    # Register new callback
+    try:
+        cb = LucidicLiteLLMCallback()
+        litellm.callbacks.append(cb)
+        logger.info("[LiteLLM] Registered Lucidic callback for event tracking")
+    except Exception as e:
+        logger.error(f"[LiteLLM] Failed to register callback: {e}")