llmops-observability 8.0.0__py3-none-any.whl

llmops_observability/trace_manager.py

@@ -0,0 +1,641 @@
+"""
+Trace Manager for LLMOps Observability
+Handles tracing and tracking of LLM operations with direct Langfuse integration
+Inspired by veriskGO's trace_manager with instant span sending to Langfuse
+"""
+from __future__ import annotations
+import uuid
+import threading
+import time
+import traceback
+import functools
+import inspect
+import sys
+import json
+from datetime import datetime, timezone
+from typing import Optional, Dict, Any, List, Union
+from .models import SpanContext
+from .config import get_langfuse_client
+
+
+# ============================================================
+# Safe Serialization Helpers (from veriskGO)
+# ============================================================
+
+# Maximum size for serialized output (200 KB)
+MAX_OUTPUT_SIZE = 200 * 1024  # 200 KB
+
+def serialize_value(value: Any, max_size: int = MAX_OUTPUT_SIZE) -> Any:
+    """Serialize value with size limit to prevent large data issues.
+    
+    Args:
+        value: Value to serialize
+        max_size: Maximum size in bytes (default 200 KB)
+    
+    Returns:
+        Serialized value, truncated if necessary
+    """
+    try:
+        # First, serialize to JSON
+        serialized_str = json.dumps(value, default=str)
+        serialized_bytes = serialized_str.encode('utf-8')
+        
+        # Check size
+        if len(serialized_bytes) <= max_size:
+            return json.loads(serialized_str)
+        
+        # Too large - return truncation info instead
+        preview_size = min(1000, max_size // 2)  # Show first 1KB as preview
+        preview = serialized_str[:preview_size]
+        
+        print(f"[LLMOps-Observability] Output truncated: {len(serialized_bytes)} bytes → {max_size} bytes limit")
+        
+        return {
+            "_truncated": True,
+            "_original_size_bytes": len(serialized_bytes),
+            "_original_size_mb": round(len(serialized_bytes) / (1024 * 1024), 2),
+            "_preview": preview + "...",
+            "_message": f"Output truncated (original: {round(len(serialized_bytes) / (1024 * 1024), 2)} MB, limit: {round(max_size / 1024, 0)} KB)"
+        }
+    except Exception as e:
+        return str(value)
+
+
+def safe_locals(d: Dict[str, Any]) -> Dict[str, Any]:
+    """Safely serialize local variables"""
+    return {k: serialize_value(v) for k, v in d.items() if not k.startswith("_")}
+
+
+# ============================================================
+# Core Trace Manager
+# ============================================================
+
+class TraceManager:
+    _lock = threading.Lock()
+    _active: Dict[str, Any] = {
+        "trace_id": None,
+        "trace_obj": None,  # Langfuse trace object reference
+        "spans": [],
+        "stack": [],  # Stack of active spans for nesting
+        "metadata": {},
+        "observation_stack": [],  # Stack of Langfuse observation contexts
+    }
+
+    @classmethod
+    def has_active_trace(cls) -> bool:
+        """Check if there's an active trace"""
+        return cls._active["trace_id"] is not None
+
+    @classmethod
+    def _id(cls) -> str:
+        """Generate a unique ID compatible with Langfuse (32 lowercase hex chars)"""
+        return uuid.uuid4().hex
+
+    @classmethod
+    def _now(cls) -> str:
+        """Get current timestamp in ISO format"""
+        return datetime.now(timezone.utc).isoformat()
+
+    @classmethod
+    def start_trace(
+        cls,
+        name: str,
+        metadata: Optional[Dict[str, Any]] = None,
+        user_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+    ) -> str:
+        """
+        Start a new trace (similar to veriskGO API).
+        
+        Args:
+            name: Name of the trace
+            metadata: Optional metadata dictionary
+            user_id: Optional user ID
+            session_id: Optional session ID
+            tags: Optional list of tags
+            
+        Returns:
+            str: Trace ID
+        """
+        with cls._lock:
+            if cls._active["trace_id"]:
+                print("[LLMOps-Observability] WARNING: Trace already active. Ending previous trace.")
+                cls._end_trace_internal()
+
+            trace_id = cls._id()
+            
+            # Store trace information
+            cls._active["trace_id"] = trace_id
+            cls._active["trace_obj"] = {
+                "name": name,
+                "user_id": user_id,
+                "session_id": session_id,
+                "metadata": metadata or {},
+                "tags": tags or [],
+                "trace_input": None,
+                "trace_output": None,
+            }
+            cls._active["spans"] = []
+            cls._active["stack"] = []  # Initialize stack for nested spans
+            cls._active["metadata"] = metadata or {}
+            
+            print(f"[LLMOps-Observability] Trace started: {name} (ID: {trace_id})")
+            return trace_id
+
+    @classmethod
+    def _end_trace_internal(cls):
+        """Internal method to end trace without lock (already within lock context)"""
+        trace_id = cls._active["trace_id"]
+        
+        if trace_id:
+            # Flush to Langfuse immediately
+            langfuse = get_langfuse_client()
+            langfuse.flush()
+            
+            # Clear the active trace
+            cls._active["trace_id"] = None
+            cls._active["trace_obj"] = None
+            cls._active["spans"] = []
+            cls._active["stack"] = []
+            cls._active["metadata"] = {}
+            cls._active["observation_stack"] = []
+
+    @classmethod
+    def end_trace(cls, final_output: Optional[Any] = None) -> Optional[str]:
+        """
+        End the current trace and flush to Langfuse.
+        
+        Args:
+            final_output: Optional final output for the trace
+            
+        Returns:
+            Optional[str]: Trace ID if successful
+        """
+        with cls._lock:
+            trace_id = cls._active["trace_id"]
+            
+            if not trace_id:
+                return None
+            
+            if final_output and cls._active["trace_obj"]:
+                cls._active["trace_obj"]["trace_output"] = final_output
+            
+            cls._end_trace_internal()
+            
+            print(f"[LLMOps-Observability] Trace ended and flushed: {trace_id}")
+            return trace_id
+
+    @classmethod
+    def finalize_and_send(
+        cls,
+        *,
+        user_id: str,
+        session_id: str,
+        trace_name: str,
+        trace_input: dict,
+        trace_output: dict,
+        extra_spans: list = [],
+    ):
+        """
+        Finalize and send the trace with input/output (compatible with veriskGO API).
+        
+        This method provides compatibility with veriskGO's API by accepting
+        trace_input and trace_output, then ending the trace.
+        
+        Args:
+            user_id: User ID for the trace
+            session_id: Session ID for the trace
+            trace_name: Name of the trace
+            trace_input: Input data for the trace
+            trace_output: Output data for the trace
+            extra_spans: Additional spans (not used in this implementation)
+        
+        Returns:
+            bool: True if successful, False otherwise
+        """
+        with cls._lock:
+            trace_id = cls._active.get("trace_id")
+            trace_obj = cls._active.get("trace_obj")
+            
+            if not trace_id or not trace_obj:
+                print("[LLMOps-Observability] ERROR: No active trace to finalize.")
+                return False
+            
+            # Update trace object with input/output
+            trace_obj["trace_input"] = serialize_value(trace_input)
+            trace_obj["trace_output"] = serialize_value(trace_output)
+            trace_obj["user_id"] = user_id
+            trace_obj["session_id"] = session_id
+            trace_obj["name"] = trace_name
+            
+            # Store updated trace_obj
+            cls._active["trace_obj"] = trace_obj
+        
+        # Send trace-level data to Langfuse using the correct API
+        try:
+            langfuse = get_langfuse_client()
+            
+            # Create trace-level observation using start_as_current_observation
+            with langfuse.start_as_current_observation(
+                as_type="span",
+                name=trace_name,
+                trace_context={"trace_id": trace_id},
+                input=trace_obj["trace_input"],
+                output=trace_obj["trace_output"],
+            ) as root:
+                # Update trace with user_id, session_id, and metadata
+                root.update_trace(
+                    name=trace_name,
+                    user_id=user_id,
+                    session_id=session_id,
+                    metadata=trace_obj.get("metadata", {}),
+                    input=trace_obj["trace_input"],
+                    output=trace_obj["trace_output"],
+                )
+            
+            # Flush immediately (no batching)
+            langfuse.flush()
+            
+        except Exception as e:
+            print(f"[LLMOps-Observability] Error sending trace to Langfuse: {e}")
+            traceback.print_exc()
+        
+        # End the trace
+        with cls._lock:
+            cls._end_trace_internal()
+        
+        print(f"[LLMOps-Observability] Trace finalized and sent: {trace_name} (ID: {trace_id})")
+        return True
+
+
+    @classmethod
+    def start_observation_context(cls, span_name: str, span_type: str, input_data: Any):
+        """
+        Start a Langfuse observation context that will be the parent for nested calls.
+        
+        Args:
+            span_name: Name of the span
+            span_type: Type ("span" or "generation")
+            input_data: Input data for the span
+            
+        Returns:
+            observation context manager
+        """
+        with cls._lock:
+            trace_id = cls._active.get("trace_id")
+            trace_obj = cls._active.get("trace_obj")
+            current_obs_stack = cls._active.get("observation_stack", [])
+            
+            if not trace_id or not trace_obj:
+                return None
+        
+        langfuse = get_langfuse_client()
+        
+        # Serialize input
+        serialized_input = serialize_value(input_data)
+        
+        # Check if we have a parent observation
+        parent_obs = current_obs_stack[-1] if current_obs_stack else None
+        
+        # Start observation context from parent or from langfuse root
+        if parent_obs:
+            # Create child observation from parent
+            if span_type == "generation":
+                obs_ctx = parent_obs.start_as_current_observation(
+                    as_type="generation",
+                    name=span_name,
+                    input=serialized_input,
+                )
+            else:
+                obs_ctx = parent_obs.start_as_current_observation(
+                    as_type="span",
+                    name=span_name,
+                    input=serialized_input,
+                )
+        else:
+            # Create root observation from langfuse client
+            # Build trace context
+            trace_context = {"trace_id": trace_id}
+            if trace_obj.get("user_id"):
+                trace_context["user_id"] = trace_obj.get("user_id")
+            if trace_obj.get("session_id"):
+                trace_context["session_id"] = trace_obj.get("session_id")
+            
+            if span_type == "generation":
+                obs_ctx = langfuse.start_as_current_observation(
+                    as_type="generation",
+                    name=span_name,
+                    trace_context=trace_context,
+                    input=serialized_input,
+                )
+            else:
+                obs_ctx = langfuse.start_as_current_observation(
+                    as_type="span",
+                    name=span_name,
+                    trace_context=trace_context,
+                    input=serialized_input,
+                )
+        
+        return obs_ctx
+    
+    @classmethod
+    def push_observation(cls, obs):
+        """Push an observation onto the stack when entering its context"""
+        with cls._lock:
+            cls._active["observation_stack"].append(obs)
+    
+    @classmethod
+    def pop_observation(cls):
+        """Pop an observation from the stack when exiting its context"""
+        with cls._lock:
+            if cls._active["observation_stack"]:
+                cls._active["observation_stack"].pop()
+
+
+# ============================================================
+# Local Capture Utilities (from veriskGO)
+# ============================================================
+
+def capture_function_locals(func, capture_locals=True, capture_self=True):
+    """
+    Capture local variables before and after function execution.
+    Uses sys.settrace for frame inspection (similar to veriskGO).
+    
+    Args:
+        func: Function to capture locals from
+        capture_locals: Whether to capture locals (True, False, or list of var names)
+        capture_self: Whether to capture 'self' variable
+    
+    Returns:
+        Tuple of (tracer, locals_before, locals_after)
+    """
+    locals_before = {}
+    locals_after = {}
+
+    if not capture_locals:
+        return None, locals_before, locals_after
+
+    target_code = func.__code__
+    target_name = func.__name__
+    target_module = func.__module__
+
+    entered = False
+
+    def tracer(frame, event, arg):
+        nonlocal entered
+
+        if frame.f_code is target_code and frame.f_globals.get("__name__") == target_module:
+            try:
+                if not entered:
+                    entered = True
+                    f_locals = frame.f_locals
+                    
+                    # Filter specific variables if capture_locals is a list
+                    if isinstance(capture_locals, list):
+                        f_locals = {k: v for k, v in f_locals.items() if k in capture_locals}
+
+                    if not capture_self and "self" in f_locals:
+                        f_locals = {k: v for k, v in f_locals.items() if k != "self"}
+                    locals_before.update(safe_locals(f_locals))
+
+                if event == "return":
+                    f_locals = frame.f_locals
+                    
+                    # Filter specific variables if capture_locals is a list
+                    if isinstance(capture_locals, list):
+                        f_locals = {k: v for k, v in f_locals.items() if k in capture_locals}
+                    
+                    if not capture_self and "self" in f_locals:
+                        f_locals = {k: v for k, v in f_locals.items() if k != "self"}
+                    locals_after.update(safe_locals(f_locals))
+                    locals_after["_return"] = serialize_value(arg)
+            except Exception as e:
+                print(f"[LLMOps-Observability] TRACER ERROR: {e}")
+                traceback.print_exc()
+
+        return tracer
+
+    return tracer, locals_before, locals_after
+
+
+# ============================================================
+# Decorator: track_function (Enhanced with veriskGO features)
+# ============================================================
+
+def track_function(
+    name: Optional[str] = None,
+    *,
+    tags: Optional[Dict[str, Any]] = None,
+    capture_locals: Union[bool, List[str]] = False,
+    capture_self: bool = False,
+):
+    """
+    Decorator to track function execution with Langfuse (enhanced version from veriskGO).
+    
+    Features:
+    - Captures input arguments and output
+    - Optionally captures local variables before/after execution
+    - Manages span stack for proper nesting
+    - Sends directly to Langfuse (no batching)
+    
+    Usage:
+        @track_function()
+        def process_data(input_data):
+            # Your code here
+            return result
+            
+        @track_function(name="custom_name", tags={"version": "1.0"}, capture_locals=True)
+        def detailed_function(x, y):
+            result = x + y
+            return result
+            
+        @track_function(capture_locals=["important_var"])
+        async def async_function():
+            important_var = "tracked"
+            return result
+    
+    Args:
+        name: Optional custom name for the span
+        tags: Optional metadata tags
+        capture_locals: Capture local variables (True/False or list of var names)
+        capture_self: Whether to capture 'self' variable (default False)
+    """
+    def decorator(func):
+        span_name = name or func.__name__
+        is_async = inspect.iscoroutinefunction(func)
+        
+        if is_async:
+            @functools.wraps(func)
+            async def async_wrapper(*args, **kwargs):
+                if not TraceManager.has_active_trace():
+                    return await func(*args, **kwargs)
+
+                # Build input data
+                input_data = {
+                    "args": serialize_value(args),
+                    "kwargs": serialize_value(kwargs),
+                }
+
+                # Start observation context (this will be parent for nested calls)
+                obs_ctx = TraceManager.start_observation_context(span_name, "span", input_data)
+                
+                if not obs_ctx:
+                    return await func(*args, **kwargs)
+
+                # Setup local variable capture
+                tracer, locals_before, locals_after = capture_function_locals(
+                    func, capture_locals=capture_locals, capture_self=capture_self
+                )
+                
+                error = None
+                result = None
+                start_time = time.time()
+                
+                # Use the observation context properly with 'with' statement
+                # This keeps the context active during function execution
+                with obs_ctx as obs:
+                    # Push observation onto stack so nested calls become children
+                    TraceManager.push_observation(obs)
+                    
+                    if tracer:
+                        sys.settrace(tracer)
+                    
+                    try:
+                        result = await func(*args, **kwargs)
+                    except Exception as e:
+                        error = e
+                        raise
+                    finally:
+                        if tracer:
+                            sys.settrace(None)
+                        
+                        duration_ms = int((time.time() - start_time) * 1000)
+                        
+                        # Build output
+                        if error:
+                            output_data = {
+                                "status": "error",
+                                "error": str(error),
+                                "stacktrace": traceback.format_exc(),
+                                "locals_before": locals_before,
+                                "locals_after": locals_after,
+                            }
+                        else:
+                            output_data = {
+                                "status": "success",
+                                "latency_ms": duration_ms,
+                                "locals_before": locals_before,
+                                "locals_after": locals_after,
+                                "output": serialize_value(result),
+                            }
+                        
+                        # Update observation with output
+                        obs.update(
+                            output=serialize_value(output_data),
+                            metadata=tags or {},
+                            level="ERROR" if error else "DEFAULT",
+                            status_message=str(error) if error else None,
+                        )
+                        
+                        # Note: flush happens after context exit
+                
+                # Flush after exiting context
+                langfuse = get_langfuse_client()
+                langfuse.flush()
+                
+                status_str = " (error)" if error else ""
+                print(f"[LLMOps-Observability] Span sent{status_str}: {span_name} ({duration_ms}ms)")
+                
+                return result
+            
+            return async_wrapper
+        else:
+            @functools.wraps(func)
+            def sync_wrapper(*args, **kwargs):
+                if not TraceManager.has_active_trace():
+                    return func(*args, **kwargs)
+
+                # Build input data
+                input_data = {
+                    "args": serialize_value(args),
+                    "kwargs": serialize_value(kwargs),
+                }
+
+                # Start observation context (this will be parent for nested calls)
+                obs_ctx = TraceManager.start_observation_context(span_name, "span", input_data)
+                
+                if not obs_ctx:
+                    return func(*args, **kwargs)
+
+                # Setup local variable capture
+                tracer, locals_before, locals_after = capture_function_locals(
+                    func, capture_locals=capture_locals, capture_self=capture_self
+                )
+                
+                error = None
+                result = None
+                start_time = time.time()
+                
+                # Use the observation context properly with 'with' statement
+                # This keeps the context active during function execution
+                with obs_ctx as obs:
+                    # Push observation onto stack so nested calls become children
+                    TraceManager.push_observation(obs)
+                    
+                    if tracer:
+                        sys.settrace(tracer)
+                    
+                    try:
+                        result = func(*args, **kwargs)
+                    except Exception as e:
+                        error = e
+                        raise
+                    finally:
+                        # Pop observation from stack
+                        TraceManager.pop_observation()
+                        
+                        if tracer:
+                            sys.settrace(None)
+                        
+                        duration_ms = int((time.time() - start_time) * 1000)
+                        
+                        # Build output
+                        if error:
+                            output_data = {
+                                "status": "error",
+                                "error": str(error),
+                                "stacktrace": traceback.format_exc(),
+                                "locals_before": locals_before,
+                                "locals_after": locals_after,
+                            }
+                        else:
+                            output_data = {
+                                "status": "success",
+                                "latency_ms": duration_ms,
+                                "locals_before": locals_before,
+                                "locals_after": locals_after,
+                                "output": serialize_value(result),
+                            }
+                        
+                        # Update observation with output
+                        obs.update(
+                            output=serialize_value(output_data),
+                            metadata=tags or {},
+                            level="ERROR" if error else "DEFAULT",
+                            status_message=str(error) if error else None,
+                        )
+                
+                # Flush after exiting context
+                langfuse = get_langfuse_client()
+                langfuse.flush()
+                
+                status_str = " (error)" if error else ""
+                print(f"[LLMOps-Observability] Span sent{status_str}: {span_name} ({duration_ms}ms)")
+                
+                return result
+            
+            return sync_wrapper
+    
+    return decorator
+