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

lucidicai/providers/lucidic_span_processor.py

@@ -0,0 +1,648 @@
+"""Custom span processor for real-time Lucidic event handling"""
+import os
+import logging
+import json
+from typing import Optional, Dict, Any
+from opentelemetry import context as otel_context
+from opentelemetry.sdk.trace import Span, SpanProcessor
+from opentelemetry.trace import StatusCode
+from opentelemetry.semconv_ai import SpanAttributes
+
+from lucidicai.client import Client
+from lucidicai.model_pricing import calculate_cost
+from .image_storage import get_stored_images, clear_stored_images
+from .text_storage import get_stored_text, clear_stored_texts
+
+logger = logging.getLogger("Lucidic")
+DEBUG = os.getenv("LUCIDIC_DEBUG", "False") == "True"
+
+
+class LucidicSpanProcessor(SpanProcessor):
+    """
+    Real-time span processor that creates Lucidic events as spans start
+    and updates them as spans end, maintaining the current SDK behavior
+    """
+    
+    def __init__(self):
+        self.span_to_event = {}  # Map span_id to event_id
+        self.span_contexts = {}  # Store span start data
+        
+    def on_start(self, span: Span, parent_context: Optional[otel_context.Context] = None) -> None:
+        """Called when a span is started - create Lucidic event immediately"""
+        try:
+            if DEBUG:
+                logger.info(f"[SpanProcessor] on_start called for span: {span.name}")
+                logger.info(f"[SpanProcessor ASDRFTYHJUKLTGHJ] Span attributes at start: {dict(span.attributes or {})}")
+                
+            client = Client()
+            if not client.session:
+                logger.debug("No active session, skipping span tracking")
+                return
+                
+            # Only process LLM spans
+            if not self._is_llm_span(span):
+                if DEBUG:
+                    logger.info(f"[SpanProcessor] Skipping non-LLM span: {span.name}")
+                return
+            
+            # Store span info for processing - we'll create the event on_end when all attributes are available
+            span_id = span.get_span_context().span_id
+            self.span_contexts[span_id] = {
+                'start_time': span.start_time,
+                'name': span.name,
+                'attributes': dict(span.attributes or {}),
+                'span': span
+            }
+            
+            if DEBUG:
+                logger.info(f"[SpanProcessor] Stored span {span_id} for later processing")
+                
+        except Exception as e:
+            logger.error(f"Error in on_start: {e}")
+            if DEBUG:
+                import traceback
+                traceback.print_exc()
+    
+    def on_end(self, span: Span) -> None:
+        """Called when a span ends - create and complete the Lucidic event"""
+        try:
+            span_id = span.get_span_context().span_id
+            
+            if DEBUG:
+                logger.info(f"[SpanProcessor] on_end called for span: {span.name}")
+                logger.info(f"[SpanProcessor] Span attributes at end: {dict(span.attributes or {})}")
+                logger.info(f"[SpanProcessor] Tracked span contexts: {list(self.span_contexts.keys())}")
+                
+                # Log any attributes that might contain message data
+                attrs = dict(span.attributes or {})
+                for key, value in attrs.items():
+                    if 'message' in key.lower() or 'prompt' in key.lower() or 'content' in key.lower():
+                        logger.info(f"[SpanProcessor] Found potential message attr: {key} = {value[:200] if isinstance(value, str) else value}")
+            
+            # Check if we have context for this span
+            if span_id not in self.span_contexts:
+                if DEBUG:
+                    logger.warning(f"[SpanProcessor] No context found for span {span_id}")
+                return
+                
+            client = Client()
+            if not client.session:
+                return
+            
+            span_context = self.span_contexts.pop(span_id, {})
+            
+            # Create event with all the attributes now available
+            event_id = self._create_event_from_span_end(span, client)
+            
+            if DEBUG:
+                logger.info(f"[SpanProcessor] Created and completed event {event_id} for span {span_id}")
+            
+            # Clear thread-local images and texts after processing
+            clear_stored_images()
+            clear_stored_texts()
+            
+        except Exception as e:
+            logger.error(f"Error in on_end: {e}")
+            if DEBUG:
+                import traceback
+                traceback.print_exc()
+    
+    def _is_llm_span(self, span: Span) -> bool:
+        """Check if this is an LLM-related span with actual LLM content"""
+        # Check if it's an agent span without LLM content
+        if span.attributes:
+            attrs = dict(span.attributes)
+            
+            # Skip agent spans that don't have prompt/completion attributes
+            if attrs.get('gen_ai.operation.name') == 'agent':
+                # Check if it has actual LLM content
+                has_prompts = any(k for k in attrs.keys() if 'prompt' in k.lower())
+                has_completions = any(k for k in attrs.keys() if 'completion' in k.lower())
+                if not has_prompts and not has_completions:
+                    if DEBUG:
+                        logger.info(f"[SpanProcessor] Skipping agent span without LLM content: {span.name}")
+                    return False
+        
+        # Check span name
+        span_name_lower = span.name.lower()
+        llm_patterns = ['openai', 'anthropic', 'chat', 'completion', 'embedding', 'gemini', 'claude']
+        
+        if any(pattern in span_name_lower for pattern in llm_patterns):
+            return True
+            
+        # Check attributes
+        if 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 _create_event_from_span_start(self, span: Span, client: Client) -> Optional[str]:
+        """Create event when span starts"""
+        try:
+            attributes = dict(span.attributes or {})
+            
+            # Extract description
+            if DEBUG:
+                logger.info(f"[SpanProcessor -- DEBUG] Extracting Description from span start: {span}")
+            description = self._extract_description(span, attributes)
+            
+            # Extract images
+            images = self._extract_images(attributes)
+            
+            # Get model
+            model = (
+                attributes.get(SpanAttributes.LLM_REQUEST_MODEL) or
+                attributes.get('gen_ai.request.model') or
+                attributes.get('llm.model') or
+                'unknown'
+            )
+            
+            # Initial result based on whether it's streaming
+            is_streaming = attributes.get(SpanAttributes.LLM_IS_STREAMING, False) or \
+                         attributes.get('llm.is_streaming', False)
+            initial_result = None if is_streaming else "Waiting for response..."
+            
+            # Apply masking to description if configured
+            if client.masking_function:
+                description = client.mask(description)
+            
+            # Create event - session.create_event will handle temporary step creation if needed
+            event_kwargs = {
+                'description': description,
+                'result': initial_result,
+                'model': model
+            }
+
+            if DEBUG:
+                logger.info(f"[SpanProcessor -- DEBUG -- _create_event_from_span_start -- ORIGINAL] event_kwargs: {event_kwargs}")
+            
+            if images:
+                event_kwargs['screenshots'] = images
+                
+            # Check for step context
+            step_id = attributes.get('lucidic.step_id')
+            if step_id:
+                event_kwargs['step_id'] = step_id
+                
+            return client.session.create_event(**event_kwargs)
+            
+        except Exception as e:
+            logger.error(f"Failed to create event: {e}")
+            return None
+    
+    def _create_event_from_span_end(self, span: Span, client: Client) -> Optional[str]:
+        """Create and complete event when span ends with all attributes available"""
+        try:
+            attributes = dict(span.attributes or {})
+            
+            if DEBUG:
+                logger.info(f"[SpanProcessor] Creating event from span end with {len(attributes)} attributes")
+            
+            # Extract all information
+            if False and DEBUG:
+                logger.info(f"[SpanProcessor -- DEBUG] Extracting Description attributes: {attributes}")
+            description = self._extract_description(span, attributes)
+
+            if False and DEBUG:
+                logger.info(f"[SpanProcessor -- DEBUG] Extracting Result attributes: {attributes}")
+            raw_result = self._extract_result(span, attributes)
+
+            if False and DEBUG:
+                logger.info(f"[SpanProcessor -- DEBUG] Extracting Images: span: attributes: {attributes}")
+            images = self._extract_images(attributes)
+
+            if False and DEBUG:
+                logger.info(f"[SpanProcessor -- DEBUG] Extracting Model: span: {span} attributes: {attributes}")
+            model = (
+                attributes.get(SpanAttributes.LLM_RESPONSE_MODEL) or
+                attributes.get(SpanAttributes.LLM_REQUEST_MODEL) or
+                attributes.get('gen_ai.response.model') or
+                attributes.get('gen_ai.request.model') or
+                'unknown'
+            )
+
+            # Format result as Input/Output
+            # The description contains the input (prompts), raw_result contains the output (completions)
+            formatted_result = f"{raw_result}"
+
+            if False and DEBUG:
+                logger.info(f"[SpanProcessor -- DEBUG] description: {description}, result: {formatted_result}, model: {model}, images: {images}")
+            
+            # Apply masking
+            if client.masking_function:
+                formatted_result = client.mask(formatted_result)
+            
+            # Calculate cost
+            cost = self._calculate_cost(attributes)
+            
+            # Check success
+            is_successful = span.status.status_code != StatusCode.ERROR
+            
+            # Create event with all data
+            event_kwargs = {
+                'description': description,
+                'result': formatted_result,
+                'model': model,
+                'is_finished': True
+            }
+            
+            if images:
+                event_kwargs['screenshots'] = images
+                
+            if cost is not None:
+                event_kwargs['cost_added'] = cost
+                
+            # Check for step context
+            step_id = attributes.get('lucidic.step_id')
+            if step_id:
+                event_kwargs['step_id'] = step_id
+                
+            # Create the event (already completed)
+            event_id = client.session.create_event(**event_kwargs)
+                
+            return event_id
+            
+        except Exception as e:
+            logger.error(f"Failed to create event from span end: {e}")
+            if DEBUG:
+                import traceback
+                traceback.print_exc()
+            return None
+    
+    def _update_event_from_span_end(self, span: Span, event_id: str, client: Client) -> None:
+        """Update event when span ends"""
+        try:
+            attributes = dict(span.attributes or {})
+            
+            # Extract response
+            result = self._extract_result(span, attributes)
+            
+            # Apply masking to result if configured
+            if client.masking_function:
+                result = client.mask(result)
+            
+            # Calculate cost
+            cost = self._calculate_cost(attributes)
+            
+            # Check success
+            is_successful = span.status.status_code != StatusCode.ERROR
+            
+            # Update event
+            update_kwargs = {
+                'event_id': event_id,
+                'result': result,
+                'is_finished': True
+            }
+            
+            if cost is not None:
+                update_kwargs['cost_added'] = cost
+                
+            # Update model if we got a response model
+            response_model = attributes.get(SpanAttributes.LLM_RESPONSE_MODEL) or \
+                           attributes.get('gen_ai.response.model')
+            if response_model:
+                update_kwargs['model'] = response_model
+
+            if DEBUG:
+                logger.info(f"[SpanProcessor -- DEBUG] update_kwargs: {update_kwargs}")
+                
+            client.session.update_event(**update_kwargs)
+            
+        except Exception as e:
+            logger.error(f"Failed to update event: {e}")
+    
+    def _extract_description(self, span: Span, attributes: Dict[str, Any]) -> str:
+        """Extract description from span"""
+        if DEBUG:
+            logger.info(f"[SpanProcessor] Extracting description from attributes: {list(attributes.keys())}")
+        
+        # Try to reconstruct messages from indexed attributes (OpenLLMetry format)
+        messages = self._extract_indexed_messages(attributes)
+        if messages:
+            if DEBUG:
+                logger.info(f"[SpanProcessor] Reconstructed {len(messages)} messages from indexed attributes")
+            return self._format_messages(messages)
+        
+        # Try prompts first (other formats)
+        prompts = attributes.get(SpanAttributes.LLM_PROMPTS) or \
+                 attributes.get('gen_ai.prompt') or \
+                 attributes.get('llm.prompts')
+                 
+        if prompts:
+            if DEBUG:
+                logger.info(f"[SpanProcessor] Found prompts: {prompts}")
+            return self._format_prompts(prompts)
+            
+        # Try messages
+        messages = attributes.get('gen_ai.messages') or \
+                  attributes.get('llm.messages')
+                  
+        if messages:
+            if DEBUG:
+                logger.info(f"[SpanProcessor] Found messages: {messages}")
+            return self._format_messages(messages)
+            
+        # Fallback
+        if DEBUG:
+            logger.warning(f"[SpanProcessor] No prompts/messages found, using fallback")
+        return f"LLM Request: {span.name}"
+    
+    def _extract_indexed_messages(self, attributes: Dict[str, Any]) -> list:
+        """Extract messages from indexed attributes (gen_ai.prompt.0.role, gen_ai.prompt.0.content, etc.)"""
+        messages = []
+        i = 0
+        
+        # Keep extracting messages until we don't find any more
+        while True:
+            prefix = f"gen_ai.prompt.{i}"
+            role = attributes.get(f"{prefix}.role")
+            
+            if not role:
+                break
+                
+            message = {"role": role}
+            
+            # Get content
+            content = attributes.get(f"{prefix}.content")
+            if content:
+                # Try to parse JSON content (for multimodal)
+                try:
+                    import json
+                    parsed_content = json.loads(content)
+                    message["content"] = parsed_content
+                except:
+                    message["content"] = content
+            else:
+                # Content might be missing for multimodal messages due to size limits
+                # Check if we have stored text and/or images in thread-local storage
+                stored_text = get_stored_text(i)
+                stored_images = get_stored_images()
+                
+                if stored_text or stored_images:
+                    if DEBUG:
+                        logger.info(f"[SpanProcessor] No content for message {i}, but found stored text/images")
+                    
+                    # Create synthetic content with both text and images
+                    synthetic_content = []
+                    
+                    # Add text if available
+                    if stored_text:
+                        synthetic_content.append({
+                            "type": "text",
+                            "text": stored_text
+                        })
+                    
+                    # Add images if available
+                    if stored_images and i == 0:  # Assume first message might have images
+                        for idx, img in enumerate(stored_images):
+                            synthetic_content.append({
+                                "type": "image_url",
+                                "image_url": {"url": img}
+                            })
+                    
+                    if synthetic_content:
+                        message["content"] = synthetic_content
+            
+            messages.append(message)
+            i += 1
+            
+        return messages
+    
+    def _extract_indexed_completions(self, attributes: Dict[str, Any]) -> list:
+        """Extract completions from indexed attributes"""
+        completions = []
+        i = 0
+        
+        while True:
+            prefix = f"gen_ai.completion.{i}"
+            role = attributes.get(f"{prefix}.role")
+            content = attributes.get(f"{prefix}.content")
+            
+            if not role and not content:
+                break
+                
+            completion = {}
+            if role:
+                completion["role"] = role
+            if content:
+                completion["content"] = content
+                
+            if completion:
+                completions.append(completion)
+            
+            i += 1
+            
+        return completions
+    
+    def _extract_result(self, span: Span, attributes: Dict[str, Any]) -> str:
+        """Extract result from span"""
+        if DEBUG:
+            logger.info(f"[SpanProcessor -- _extract_result -- DEBUG] Extracting result from attributes: {attributes}")
+
+        # Try indexed completions first (OpenLLMetry format)
+        completions = self._extract_indexed_completions(attributes)
+        if completions:
+            if DEBUG:
+                logger.info(f"[SpanProcessor] Found {len(completions)} indexed completions")
+            # Format completions
+            results = []
+            for comp in completions:
+                if "content" in comp:
+                    results.append(str(comp["content"]))
+            if results:
+                return "\n".join(results)
+        
+        # Try completions
+        completions = attributes.get(SpanAttributes.LLM_COMPLETIONS) or \
+                     attributes.get('gen_ai.completion') or \
+                     attributes.get('llm.completions')
+                     
+        if completions:
+            if isinstance(completions, list):
+                return "\n".join(str(c) for c in completions)
+            else:
+                return str(completions)
+                
+        # Check for error
+        if span.status.status_code == StatusCode.ERROR:
+            return f"Error: {span.status.description or 'Unknown error'}"
+            
+        # Check streaming
+        if attributes.get(SpanAttributes.LLM_IS_STREAMING):
+            content = attributes.get('llm.response.content') or \
+                     attributes.get('gen_ai.response.content')
+            if content:
+                return content
+                
+        return "Response received"
+    
+    def _extract_images(self, attributes: Dict[str, Any]) -> list:
+        """Extract images from multimodal prompts"""
+        images = []
+
+        if DEBUG:
+            logger.info(f"[SpanProcessor -- _extract_images -- DEBUG] Extracting images from attributes: {attributes}")
+        
+        # First check indexed messages (OpenLLMetry format)
+        messages = self._extract_indexed_messages(attributes)
+        for msg in messages:
+            if isinstance(msg, dict):
+                images.extend(self._extract_images_from_message(msg))
+        
+        # Check for multimodal content in prompts
+        prompts = attributes.get(SpanAttributes.LLM_PROMPTS) or \
+                 attributes.get('gen_ai.prompt')
+                 
+        if isinstance(prompts, list):
+            for prompt in prompts:
+                if isinstance(prompt, dict):
+                    images.extend(self._extract_images_from_message(prompt))
+                    
+        # Check messages too
+        messages = attributes.get('gen_ai.messages') or \
+                  attributes.get('llm.messages')
+                  
+        if isinstance(messages, list):
+            for msg in messages:
+                if isinstance(msg, dict):
+                    images.extend(self._extract_images_from_message(msg))
+        
+        # If no images found but we have stored images in thread-local, retrieve them
+        stored_images = get_stored_images()
+        if not images and stored_images:
+            if DEBUG:
+                logger.info(f"[SpanProcessor] No images found in attributes, checking thread-local storage: {len(stored_images)} images")
+            for img in stored_images:
+                if img and not img.startswith('data:'):
+                    images.append(f"data:image/jpeg;base64,{img}")
+                else:
+                    images.append(img)
+        
+        if DEBUG and images:
+            logger.info(f"[SpanProcessor] Extracted {len(images)} images")
+                    
+        return images
+    
+    def _extract_images_from_message(self, message: dict) -> list:
+        """Extract images from a single message"""
+        images = []
+        content = message.get('content', '')
+
+        if DEBUG:
+            logger.info(f"[SpanProcessor -- _extract_images_from_message -- DEBUG] Extracting images from message: {message}, content: {content}")
+
+        # Handle case where content might be a JSON string
+        if isinstance(content, str) and content.strip().startswith('['):
+            try:
+                parsed_content = json.loads(content)
+                if isinstance(parsed_content, list):
+                    content = parsed_content
+            except json.JSONDecodeError:
+                # If parsing fails, keep content as string
+                pass
+
+        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)
+                        elif url.startswith('lucidic_image_'):
+                            # This is a placeholder - retrieve from thread-local storage
+                            image = self._retrieve_image_from_placeholder(url)
+                            if image:
+                                images.append(image)
+                            
+        return images
+    
+    def _retrieve_image_from_placeholder(self, placeholder: str) -> Optional[str]:
+        """Retrieve image from thread-local storage using placeholder"""
+        try:
+            base64_data = get_image_by_placeholder(placeholder)
+            if base64_data:
+                # Ensure it has proper data URI format
+                if not base64_data.startswith('data:'):
+                    # Add data URI prefix if missing
+                    base64_data = f"data:image/jpeg;base64,{base64_data}"
+                return base64_data
+        except Exception as e:
+            if DEBUG:
+                logger.error(f"[SpanProcessor] Failed to retrieve image from placeholder: {e}")
+        return None
+    
+    def _format_prompts(self, prompts: Any) -> str:
+        """Format prompts into description"""
+        if isinstance(prompts, str):
+            return prompts
+        elif isinstance(prompts, list):
+            return self._format_messages(prompts)
+        else:
+            return "Model request"
+    
+    def _format_messages(self, messages: list) -> str:
+        """Format message list"""
+        formatted = []
+        
+        for msg in messages:
+            if isinstance(msg, dict):
+                role = msg.get('role', 'unknown')
+                content = msg.get('content', '')
+                
+                if isinstance(content, str):
+                    formatted.append(f"{role}: {content}")
+                elif isinstance(content, list):
+                    # Extract text from multimodal
+                    texts = []
+                    for item in content:
+                        if isinstance(item, dict) and item.get('type') == 'text':
+                            texts.append(item.get('text', ''))
+                    if texts:
+                        formatted.append(f"{role}: {' '.join(texts)}")
+            elif isinstance(msg, str):
+                formatted.append(msg)
+                
+        return '\n'.join(formatted) if formatted else "Model request"
+    
+    def _calculate_cost(self, attributes: Dict[str, Any]) -> Optional[float]:
+        """Calculate cost from token usage"""
+        prompt_tokens = (
+            attributes.get(SpanAttributes.LLM_USAGE_PROMPT_TOKENS) or
+            attributes.get('gen_ai.usage.prompt_tokens') or
+            attributes.get('gen_ai.usage.input_tokens') or
+            0
+        )
+        
+        completion_tokens = (
+            attributes.get(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS) or
+            attributes.get('gen_ai.usage.completion_tokens') or
+            attributes.get('gen_ai.usage.output_tokens') or
+            0
+        )
+        
+        total_tokens = prompt_tokens + completion_tokens
+        
+        if total_tokens > 0:
+            model = (
+                attributes.get(SpanAttributes.LLM_RESPONSE_MODEL) or
+                attributes.get(SpanAttributes.LLM_REQUEST_MODEL) or
+                attributes.get('gen_ai.response.model') or
+                attributes.get('gen_ai.request.model')
+            )
+            
+            if model:
+                return calculate_cost(model, {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": total_tokens})
+                
+        return None
+    
+    def shutdown(self, timeout_millis: int = 30000) -> None:
+        """Shutdown processor"""
+        if self.span_to_event:
+            logger.warning(f"Shutting down with {len(self.span_to_event)} incomplete spans")
+            
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Force flush - no-op for this processor"""
+        return True