agentreplay 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

agentreplay/decorators.py

@@ -0,0 +1,545 @@
+# Copyright 2025 Sushanth (https://github.com/sushanthpy)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Decorator-based tracing for Agentreplay.
+
+Provides @traceable and @observe decorators for easy function instrumentation.
+
+Example:
+    >>> from agentreplay import init, traceable
+    >>> 
+    >>> init()
+    >>> 
+    >>> @traceable
+    >>> def my_function(query: str) -> str:
+    ...     return f"Result for {query}"
+    >>> 
+    >>> result = my_function("hello")  # Automatically traced!
+"""
+
+import functools
+import inspect
+import time
+import logging
+from typing import (
+    Optional, Callable, TypeVar, Any, Dict, Union, 
+    overload, Awaitable
+)
+try:
+    from typing import ParamSpec
+except ImportError:
+    from typing_extensions import ParamSpec
+from contextvars import ContextVar
+
+logger = logging.getLogger(__name__)
+
+# Type variables for generic decorators
+P = ParamSpec("P")
+R = TypeVar("R")
+
+# Context variable for current span
+_current_span: ContextVar[Optional[Any]] = ContextVar("current_span", default=None)
+
+
+# =============================================================================
+# Span Kind
+# =============================================================================
+
+class SpanKind:
+    """Span kind constants for categorizing operations."""
+    CHAIN = "chain"
+    LLM = "llm"
+    TOOL = "tool"
+    RETRIEVER = "retriever"
+    EMBEDDING = "embedding"
+    GUARDRAIL = "guardrail"
+    CACHE = "cache"
+    HTTP = "http"
+    DB = "db"
+
+
+# =============================================================================
+# Active Span
+# =============================================================================
+
+class ActiveSpan:
+    """Active span with methods to add data.
+    
+    This is yielded by the trace() context manager and passed to
+    decorated functions.
+    """
+    
+    def __init__(
+        self,
+        name: str,
+        kind: str = SpanKind.CHAIN,
+        span_id: Optional[str] = None,
+        parent_id: Optional[str] = None,
+        trace_id: Optional[str] = None,
+    ):
+        self.name = name
+        self.kind = kind
+        self.span_id = span_id or self._generate_id()
+        self.parent_id = parent_id
+        self.trace_id = trace_id or self._generate_id()
+        self.start_time = time.time()
+        self.end_time: Optional[float] = None
+        self.attributes: Dict[str, Any] = {}
+        self.events: list = []
+        self.input_data: Optional[Any] = None
+        self.output_data: Optional[Any] = None
+        self.error: Optional[Exception] = None
+        self.token_usage: Dict[str, int] = {}
+        self._ended = False
+    
+    @staticmethod
+    def _generate_id() -> str:
+        """Generate unique span ID."""
+        import uuid
+        return uuid.uuid4().hex[:16]
+    
+    def set_input(self, data: Any) -> "ActiveSpan":
+        """Set input data."""
+        self.input_data = data
+        return self
+    
+    def set_output(self, data: Any) -> "ActiveSpan":
+        """Set output data."""
+        self.output_data = data
+        return self
+    
+    def set_attribute(self, key: str, value: Any) -> "ActiveSpan":
+        """Set a span attribute."""
+        self.attributes[key] = value
+        return self
+    
+    def set_attributes(self, attributes: Dict[str, Any]) -> "ActiveSpan":
+        """Set multiple attributes."""
+        self.attributes.update(attributes)
+        return self
+    
+    def add_event(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> "ActiveSpan":
+        """Add an event to the span."""
+        self.events.append({
+            "name": name,
+            "timestamp": time.time(),
+            "attributes": attributes or {},
+        })
+        return self
+    
+    def set_error(self, error: Exception) -> "ActiveSpan":
+        """Set error on span."""
+        self.error = error
+        self.attributes["error.type"] = type(error).__name__
+        self.attributes["error.message"] = str(error)
+        import traceback
+        self.attributes["error.stack"] = traceback.format_exc()
+        return self
+    
+    def set_token_usage(
+        self,
+        prompt_tokens: Optional[int] = None,
+        completion_tokens: Optional[int] = None,
+        total_tokens: Optional[int] = None,
+    ) -> "ActiveSpan":
+        """Set token usage for LLM calls."""
+        if prompt_tokens is not None:
+            self.token_usage["prompt"] = prompt_tokens
+            self.attributes["gen_ai.usage.prompt_tokens"] = prompt_tokens
+        if completion_tokens is not None:
+            self.token_usage["completion"] = completion_tokens
+            self.attributes["gen_ai.usage.completion_tokens"] = completion_tokens
+        if total_tokens is not None:
+            self.token_usage["total"] = total_tokens
+            self.attributes["gen_ai.usage.total_tokens"] = total_tokens
+        return self
+    
+    def set_model(self, model: str, provider: Optional[str] = None) -> "ActiveSpan":
+        """Set model information."""
+        self.attributes["gen_ai.request.model"] = model
+        if provider:
+            self.attributes["gen_ai.system"] = provider
+        return self
+    
+    def end(self) -> None:
+        """End the span and send to backend."""
+        if self._ended:
+            return
+        
+        self._ended = True
+        self.end_time = time.time()
+        
+        # Send to backend
+        self._send()
+    
+    def _send(self) -> None:
+        """Send span to Agentreplay backend."""
+        try:
+            from agentreplay.sdk import get_batching_client, is_initialized, get_config
+            
+            if not is_initialized():
+                return
+            
+            config = get_config()
+            if not config.enabled:
+                return
+            
+            # Build edge
+            from agentreplay.models import AgentFlowEdge, SpanType
+            
+            # Map kind to SpanType
+            span_type_map = {
+                SpanKind.CHAIN: SpanType.ROOT,
+                SpanKind.LLM: SpanType.TOOL_CALL,
+                SpanKind.TOOL: SpanType.TOOL_CALL,
+                SpanKind.RETRIEVER: SpanType.TOOL_CALL,
+                SpanKind.EMBEDDING: SpanType.TOOL_CALL,
+            }
+            span_type = span_type_map.get(self.kind, SpanType.ROOT)
+            
+            # Calculate duration
+            duration_us = int((self.end_time - self.start_time) * 1_000_000) if self.end_time else 0
+            
+            # Build payload
+            payload = {}
+            if self.input_data is not None and config.capture_input:
+                payload["input"] = self._safe_serialize(self.input_data)
+            if self.output_data is not None and config.capture_output:
+                payload["output"] = self._safe_serialize(self.output_data)
+            if self.attributes:
+                payload["attributes"] = self.attributes
+            if self.events:
+                payload["events"] = self.events
+            if self.error:
+                payload["error"] = {
+                    "type": type(self.error).__name__,
+                    "message": str(self.error),
+                }
+            
+            edge = AgentFlowEdge(
+                tenant_id=config.tenant_id,
+                project_id=config.project_id,
+                agent_id=config.agent_id,
+                session_id=int(self.trace_id[:8], 16) if self.trace_id else 0,
+                span_type=span_type,
+                timestamp_us=int(self.start_time * 1_000_000),
+                duration_us=duration_us,
+                token_count=self.token_usage.get("total", 0),
+                payload=payload,
+            )
+            
+            # Send via batching client
+            client = get_batching_client()
+            client.insert(edge)
+            
+        except Exception as e:
+            logger.debug(f"Failed to send span: {e}")
+    
+    def _safe_serialize(self, data: Any, max_size: int = 10000) -> Any:
+        """Safely serialize data with size limits."""
+        import json
+        
+        try:
+            serialized = json.dumps(data, default=str)
+            if len(serialized) > max_size:
+                return {"__truncated": True, "__preview": serialized[:1000]}
+            return data
+        except Exception:
+            return str(data)[:max_size]
+    
+    def __enter__(self) -> "ActiveSpan":
+        """Context manager entry."""
+        # Set as current span
+        self._token = _current_span.set(self)
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Context manager exit."""
+        # Capture error if any
+        if exc_val is not None:
+            self.set_error(exc_val)
+        
+        # End span
+        self.end()
+        
+        # Reset current span
+        _current_span.reset(self._token)
+
+
+# =============================================================================
+# Get Current Span
+# =============================================================================
+
+def get_current_span() -> Optional[ActiveSpan]:
+    """Get the currently active span.
+    
+    Returns:
+        ActiveSpan if inside a traced context, None otherwise
+    """
+    return _current_span.get()
+
+
+# =============================================================================
+# Traceable Decorator
+# =============================================================================
+
+@overload
+def traceable(func: Callable[P, R]) -> Callable[P, R]: ...
+
+@overload
+def traceable(
+    *,
+    name: Optional[str] = None,
+    kind: str = SpanKind.CHAIN,
+    capture_input: bool = True,
+    capture_output: bool = True,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+
+
+def traceable(
+    func: Optional[Callable[P, R]] = None,
+    *,
+    name: Optional[str] = None,
+    kind: str = SpanKind.CHAIN,
+    capture_input: bool = True,
+    capture_output: bool = True,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> Union[Callable[P, R], Callable[[Callable[P, R]], Callable[P, R]]]:
+    """Decorator to trace a function.
+    
+    Works with both sync and async functions. Automatically captures
+    inputs, outputs, errors, and timing.
+    
+    Args:
+        func: Function to decorate (when used without parentheses)
+        name: Span name (default: function name)
+        kind: Span kind (chain, llm, tool, retriever, etc.)
+        capture_input: Whether to capture function inputs
+        capture_output: Whether to capture function output
+        metadata: Additional metadata to attach
+        
+    Returns:
+        Decorated function
+        
+    Example:
+        >>> @traceable
+        >>> def simple_function():
+        ...     return "hello"
+        
+        >>> @traceable(name="my_operation", kind="tool")
+        >>> def tool_function(query: str):
+        ...     return search(query)
+        
+        >>> @traceable(capture_input=False)  # Don't capture sensitive inputs
+        >>> def sensitive_function(password: str):
+        ...     return authenticate(password)
+    """
+    def decorator(fn: Callable[P, R]) -> Callable[P, R]:
+        span_name = name or fn.__name__
+        
+        # Check if async
+        if inspect.iscoroutinefunction(fn):
+            @functools.wraps(fn)
+            async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+                # Get parent span
+                parent = get_current_span()
+                
+                # Create span
+                span = ActiveSpan(
+                    name=span_name,
+                    kind=kind,
+                    parent_id=parent.span_id if parent else None,
+                    trace_id=parent.trace_id if parent else None,
+                )
+                
+                # Add metadata
+                if metadata:
+                    span.set_attributes(metadata)
+                
+                # Capture input
+                if capture_input:
+                    try:
+                        input_data = _capture_args(fn, args, kwargs)
+                        span.set_input(input_data)
+                    except Exception:
+                        pass
+                
+                # Execute with span context
+                with span:
+                    try:
+                        result = await fn(*args, **kwargs)
+                        
+                        # Capture output
+                        if capture_output:
+                            span.set_output(result)
+                        
+                        return result
+                    except Exception as e:
+                        span.set_error(e)
+                        raise
+            
+            return async_wrapper
+        else:
+            @functools.wraps(fn)
+            def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+                # Get parent span
+                parent = get_current_span()
+                
+                # Create span
+                span = ActiveSpan(
+                    name=span_name,
+                    kind=kind,
+                    parent_id=parent.span_id if parent else None,
+                    trace_id=parent.trace_id if parent else None,
+                )
+                
+                # Add metadata
+                if metadata:
+                    span.set_attributes(metadata)
+                
+                # Capture input
+                if capture_input:
+                    try:
+                        input_data = _capture_args(fn, args, kwargs)
+                        span.set_input(input_data)
+                    except Exception:
+                        pass
+                
+                # Execute with span context
+                with span:
+                    try:
+                        result = fn(*args, **kwargs)
+                        
+                        # Capture output
+                        if capture_output:
+                            span.set_output(result)
+                        
+                        return result
+                    except Exception as e:
+                        span.set_error(e)
+                        raise
+            
+            return sync_wrapper
+    
+    # Handle @traceable vs @traceable()
+    if func is not None:
+        return decorator(func)
+    return decorator
+
+
+# Alias for Langfuse-style API
+observe = traceable
+
+
+def _capture_args(fn: Callable, args: tuple, kwargs: dict) -> Dict[str, Any]:
+    """Capture function arguments as a dict."""
+    sig = inspect.signature(fn)
+    params = list(sig.parameters.keys())
+    
+    result = {}
+    for i, arg in enumerate(args):
+        if i < len(params):
+            result[params[i]] = arg
+        else:
+            result[f"arg_{i}"] = arg
+    
+    result.update(kwargs)
+    return result
+
+
+# =============================================================================
+# Trace Context Manager
+# =============================================================================
+
+def trace(
+    name: str,
+    *,
+    kind: str = SpanKind.CHAIN,
+    input: Optional[Any] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> ActiveSpan:
+    """Create a trace span as a context manager.
+    
+    Args:
+        name: Span name
+        kind: Span kind (chain, llm, tool, retriever, etc.)
+        input: Input data to record
+        metadata: Additional metadata
+        
+    Returns:
+        ActiveSpan context manager
+        
+    Example:
+        >>> with trace("retrieve_documents", kind="retriever") as span:
+        ...     docs = vector_db.search(query)
+        ...     span.set_output({"count": len(docs)})
+        ...     return docs
+    """
+    # Get parent span
+    parent = get_current_span()
+    
+    # Create span
+    span = ActiveSpan(
+        name=name,
+        kind=kind,
+        parent_id=parent.span_id if parent else None,
+        trace_id=parent.trace_id if parent else None,
+    )
+    
+    # Set input
+    if input is not None:
+        span.set_input(input)
+    
+    # Set metadata
+    if metadata:
+        span.set_attributes(metadata)
+    
+    return span
+
+
+def start_span(
+    name: str,
+    *,
+    kind: str = SpanKind.CHAIN,
+    input: Optional[Any] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> ActiveSpan:
+    """Start a manual span (must call span.end()).
+    
+    Use trace() context manager when possible. This is for cases
+    where you need manual control over span lifetime.
+    
+    Args:
+        name: Span name
+        kind: Span kind
+        input: Input data
+        metadata: Additional metadata
+        
+    Returns:
+        ActiveSpan (call .end() when done)
+        
+    Example:
+        >>> span = start_span("long_operation", kind="tool")
+        >>> try:
+        ...     result = do_something()
+        ...     span.set_output(result)
+        >>> except Exception as e:
+        ...     span.set_error(e)
+        ...     raise
+        >>> finally:
+        ...     span.end()
+    """
+    return trace(name, kind=kind, input=input, metadata=metadata)

agentreplay/install_pth.py

@@ -50,8 +50,12 @@ def get_site_packages():
         import agentreplay
         agentreplay_dir = os.path.dirname(agentreplay.__file__)
         site_packages = os.path.dirname(agentreplay_dir)
-        if site_packages not in paths:
-            paths.insert(0, site_packages)
+        
+        # Prioritize the directory where agentreplay is actually installed
+        if site_packages in paths:
+            paths.remove(site_packages)
+        paths.insert(0, site_packages)
+            
     except ImportError:
         pass