agentreplay 0.1.2__py3-none-any.whl

agentreplay/sampling.py

@@ -0,0 +1,298 @@
+# 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.
+
+"""Sampling strategies for trace collection.
+
+This module provides OTEL-compatible sampling strategies to control
+trace data volume while maintaining statistical validity.
+
+Supported samplers:
+    - AlwaysOnSampler: Sample every trace (100%)
+    - AlwaysOffSampler: Sample no traces (0%)
+    - TraceIdRatioBasedSampler: Sample based on trace ID hash
+    - ParentBasedSampler: Respect parent span's sampling decision
+
+Example:
+    >>> from agentreplay.sampling import TraceIdRatioBasedSampler
+    >>> 
+    >>> # Sample 10% of traces
+    >>> sampler = TraceIdRatioBasedSampler(0.1)
+    >>> 
+    >>> # Check if trace should be sampled
+    >>> if sampler.should_sample(trace_id=0x123abc):
+    ...     # Record trace
+    ...     pass
+"""
+
+import logging
+from typing import Optional
+from abc import ABC, abstractmethod
+
+logger = logging.getLogger(__name__)
+
+
+class Sampler(ABC):
+    """Abstract base class for sampling strategies."""
+    
+    @abstractmethod
+    def should_sample(
+        self,
+        trace_id: int,
+        parent_sampled: Optional[bool] = None,
+    ) -> bool:
+        """Determine if a trace should be sampled.
+        
+        Args:
+            trace_id: 128-bit trace identifier
+            parent_sampled: Whether parent span was sampled (if known)
+            
+        Returns:
+            True if trace should be sampled
+        """
+        pass
+    
+    @abstractmethod
+    def get_description(self) -> str:
+        """Get human-readable description of sampler.
+        
+        Returns:
+            Description string
+        """
+        pass
+
+
+class AlwaysOnSampler(Sampler):
+    """Sample all traces (100%).
+    
+    Use this for development or low-volume production workloads.
+    """
+    
+    def should_sample(
+        self,
+        trace_id: int,
+        parent_sampled: Optional[bool] = None,
+    ) -> bool:
+        """Always returns True."""
+        return True
+    
+    def get_description(self) -> str:
+        """Get description."""
+        return "AlwaysOnSampler"
+
+
+class AlwaysOffSampler(Sampler):
+    """Sample no traces (0%).
+    
+    Use this to completely disable tracing.
+    """
+    
+    def should_sample(
+        self,
+        trace_id: int,
+        parent_sampled: Optional[bool] = None,
+    ) -> bool:
+        """Always returns False."""
+        return False
+    
+    def get_description(self) -> str:
+        """Get description."""
+        return "AlwaysOffSampler"
+
+
+class TraceIdRatioBasedSampler(Sampler):
+    """Sample traces based on trace ID hash.
+    
+    Uses deterministic sampling: traces with the same ID always get
+    the same sampling decision. This ensures consistent sampling across
+    distributed services.
+    
+    Args:
+        rate: Sampling rate between 0.0 and 1.0
+        
+    Example:
+        >>> # Sample 10% of traces
+        >>> sampler = TraceIdRatioBasedSampler(0.1)
+        >>> 
+        >>> # Same trace ID always gets same decision
+        >>> assert sampler.should_sample(0x123) == sampler.should_sample(0x123)
+    """
+    
+    def __init__(self, rate: float):
+        """Initialize ratio-based sampler.
+        
+        Args:
+            rate: Sampling rate (0.0 = none, 1.0 = all)
+            
+        Raises:
+            ValueError: If rate is not in [0.0, 1.0]
+        """
+        if not 0.0 <= rate <= 1.0:
+            raise ValueError(f"Sampling rate must be in [0.0, 1.0], got {rate}")
+        
+        self.rate = rate
+        # Calculate threshold for comparison
+        # Use upper 64 bits of trace_id for sampling decision
+        self.threshold = int(rate * (2**64 - 1))
+        
+        logger.info(f"TraceIdRatioBasedSampler initialized: rate={rate:.2%}")
+    
+    def should_sample(
+        self,
+        trace_id: int,
+        parent_sampled: Optional[bool] = None,
+    ) -> bool:
+        """Determine if trace should be sampled based on trace ID.
+        
+        Uses the upper 64 bits of the 128-bit trace ID for sampling decision.
+        This ensures uniform distribution and deterministic decisions.
+        
+        Args:
+            trace_id: 128-bit trace identifier
+            parent_sampled: Ignored (not used for ratio-based sampling)
+            
+        Returns:
+            True if trace should be sampled
+        """
+        if self.rate == 1.0:
+            return True
+        if self.rate == 0.0:
+            return False
+        
+        # Extract upper 64 bits of trace_id
+        upper_64 = (trace_id >> 64) & ((1 << 64) - 1)
+        
+        # Compare with threshold
+        return upper_64 < self.threshold
+    
+    def get_description(self) -> str:
+        """Get description."""
+        return f"TraceIdRatioBasedSampler(rate={self.rate:.2%})"
+
+
+class ParentBasedSampler(Sampler):
+    """Sample based on parent span's sampling decision.
+    
+    If parent span was sampled, sample this span too. This ensures
+    complete traces are captured (no partial traces with missing spans).
+    
+    Falls back to root_sampler for root spans (no parent).
+    
+    Args:
+        root_sampler: Sampler to use for root spans (no parent)
+        
+    Example:
+        >>> # Use 10% sampling for root spans, but always sample if parent was sampled
+        >>> root_sampler = TraceIdRatioBasedSampler(0.1)
+        >>> sampler = ParentBasedSampler(root_sampler)
+    """
+    
+    def __init__(self, root_sampler: Sampler):
+        """Initialize parent-based sampler.
+        
+        Args:
+            root_sampler: Sampler for root spans
+        """
+        self.root_sampler = root_sampler
+        logger.info(f"ParentBasedSampler initialized with root: {root_sampler.get_description()}")
+    
+    def should_sample(
+        self,
+        trace_id: int,
+        parent_sampled: Optional[bool] = None,
+    ) -> bool:
+        """Determine if trace should be sampled.
+        
+        Args:
+            trace_id: 128-bit trace identifier
+            parent_sampled: Whether parent span was sampled
+            
+        Returns:
+            True if trace should be sampled
+        """
+        # If parent sampling decision is known, use it
+        if parent_sampled is not None:
+            return parent_sampled
+        
+        # No parent (root span), use root sampler
+        return self.root_sampler.should_sample(trace_id, parent_sampled=None)
+    
+    def get_description(self) -> str:
+        """Get description."""
+        return f"ParentBasedSampler(root={self.root_sampler.get_description()})"
+
+
+def create_sampler_from_config(sampler_name: str, sampler_arg: str = "1.0") -> Sampler:
+    """Create sampler from OTEL environment variable values.
+    
+    Args:
+        sampler_name: OTEL_TRACES_SAMPLER value
+        sampler_arg: OTEL_TRACES_SAMPLER_ARG value
+        
+    Returns:
+        Configured Sampler instance
+        
+    Example:
+        >>> import os
+        >>> sampler = create_sampler_from_config(
+        ...     os.getenv("OTEL_TRACES_SAMPLER", "always_on"),
+        ...     os.getenv("OTEL_TRACES_SAMPLER_ARG", "1.0")
+        ... )
+    """
+    sampler_name = sampler_name.lower().strip()
+    
+    if sampler_name == "always_on":
+        return AlwaysOnSampler()
+    
+    elif sampler_name == "always_off":
+        return AlwaysOffSampler()
+    
+    elif sampler_name == "traceidratio":
+        try:
+            rate = float(sampler_arg)
+            rate = max(0.0, min(1.0, rate))  # Clamp to [0, 1]
+            return TraceIdRatioBasedSampler(rate)
+        except ValueError:
+            logger.warning(f"Invalid sampler_arg: {sampler_arg}, using 1.0")
+            return TraceIdRatioBasedSampler(1.0)
+    
+    elif sampler_name == "parentbased_always_on":
+        return ParentBasedSampler(AlwaysOnSampler())
+    
+    elif sampler_name == "parentbased_always_off":
+        return ParentBasedSampler(AlwaysOffSampler())
+    
+    elif sampler_name == "parentbased_traceidratio":
+        try:
+            rate = float(sampler_arg)
+            rate = max(0.0, min(1.0, rate))
+            root_sampler = TraceIdRatioBasedSampler(rate)
+            return ParentBasedSampler(root_sampler)
+        except ValueError:
+            logger.warning(f"Invalid sampler_arg: {sampler_arg}, using 1.0")
+            root_sampler = TraceIdRatioBasedSampler(1.0)
+            return ParentBasedSampler(root_sampler)
+    
+    else:
+        logger.warning(f"Unknown sampler: {sampler_name}, using AlwaysOnSampler")
+        return AlwaysOnSampler()
+
+
+__all__ = [
+    "Sampler",
+    "AlwaysOnSampler",
+    "AlwaysOffSampler",
+    "TraceIdRatioBasedSampler",
+    "ParentBasedSampler",
+    "create_sampler_from_config",
+]

agentreplay/session.py

@@ -0,0 +1,164 @@
+# 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.
+
+"""Session management for conversational applications."""
+
+import time
+from typing import Optional
+from agentreplay.client import AgentreplayClient
+from agentreplay.models import SpanType
+from agentreplay.span import Span
+
+
+class Session:
+    """Session manager for chatbot and conversational applications.
+    
+    Automatically tracks session_id and message counts, simplifying
+    session lifecycle management for multi-turn conversations.
+    
+    Args:
+        client: Agentreplay client instance
+        session_id: Optional session identifier (auto-generated if not provided)
+        agent_id: Optional agent identifier (uses client default if not provided)
+        
+    Attributes:
+        session_id: Session identifier
+        message_count: Number of messages/traces in this session
+        start_time: Session start timestamp (microseconds)
+        
+    Example:
+        >>> client = AgentreplayClient(url="http://localhost:8080", tenant_id=1)
+        >>> session = Session(client)
+        >>> 
+        >>> # Track conversation turns
+        >>> with session.trace(SpanType.LLM) as turn1:
+        ...     turn1.set_attribute("prompt", "Hello")
+        ...     turn1.set_token_count(50)
+        ...
+        >>> with session.trace(SpanType.LLM) as turn2:
+        ...     turn2.set_attribute("prompt", "How are you?")
+        ...     turn2.set_token_count(60)
+        ...
+        >>> print(f"Session {session.session_id} had {session.message_count} turns")
+        >>> session.end()
+    """
+
+    def __init__(
+        self,
+        client: AgentreplayClient,
+        session_id: Optional[int] = None,
+        agent_id: Optional[int] = None,
+    ):
+        """Initialize session manager.
+        
+        Args:
+            client: Agentreplay client
+            session_id: Session identifier (auto-generated if None)
+            agent_id: Agent identifier (uses client default if None)
+        """
+        self.client = client
+        self.session_id = session_id if session_id is not None else int(time.time() * 1000)
+        self.agent_id = agent_id if agent_id is not None else client.agent_id
+        self.message_count = 0
+        self.start_time = int(time.time() * 1_000_000)
+        self._end_time: Optional[int] = None  # Track when session ended
+        self._ended = False
+
+    def __enter__(self) -> "Session":
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Context manager exit - end session."""
+        self.end()
+
+    def trace(
+        self,
+        span_type: SpanType = SpanType.ROOT,
+        **metadata,
+    ) -> Span:
+        """Create a trace within this session.
+        
+        Automatically sets session_id and tracks message count.
+        
+        Args:
+            span_type: Type of span (default: ROOT)
+            **metadata: Additional metadata to attach to span
+            
+        Returns:
+            Span context manager
+            
+        Example:
+            >>> session = Session(client)
+            >>> with session.trace(SpanType.LLM, model="gpt-4") as span:
+            ...     span.set_token_count(100)
+        """
+        if self._ended:
+            raise RuntimeError("Cannot create trace in ended session")
+
+        self.message_count += 1
+
+        # Add session metadata
+        full_metadata = {
+            "message_num": self.message_count,
+            "session_duration_us": int(time.time() * 1_000_000) - self.start_time,
+            **metadata,
+        }
+
+        return self.client.trace(
+            span_type=span_type,
+            agent_id=self.agent_id,
+            session_id=self.session_id,
+        )
+
+    def end(self) -> None:
+        """Mark session as ended.
+        
+        Optionally send session summary metrics to backend.
+        """
+        if self._ended:
+            return
+
+        self._end_time = int(time.time() * 1_000_000)
+        self._ended = True
+
+        # Compute session statistics
+        duration_us = self._end_time - self.start_time
+
+        # Could send session summary span
+        # with self.client.trace(
+        #     span_type=SpanType.ROOT,
+        #     agent_id=self.agent_id,
+        #     session_id=self.session_id
+        # ) as summary:
+        #     summary.set_attribute("session_ended", True)
+        #     summary.set_attribute("total_messages", self.message_count)
+        #     summary.set_attribute("duration_us", duration_us)
+
+    @property
+    def is_active(self) -> bool:
+        """Check if session is still active."""
+        return not self._ended
+
+    @property
+    def duration_seconds(self) -> float:
+        """Get session duration in seconds.
+        
+        Returns accurate duration even after session has ended.
+        """
+        if self._ended and self._end_time is not None:
+            # Return actual session duration
+            return (self._end_time - self.start_time) / 1_000_000
+        # Session still active, return current duration
+        return (int(time.time() * 1_000_000) - self.start_time) / 1_000_000

agentreplay/sitecustomize.py

@@ -0,0 +1,73 @@
+# 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.
+
+"""
+Agentreplay automatic instrumentation via sitecustomize.
+
+This file is automatically loaded by Python if it's in the site-packages directory.
+It enables zero-code instrumentation - just set env vars and run!
+
+Usage:
+    export AGENTREPLAY_ENABLED=true
+    export AGENTREPLAY_URL=http://localhost:47100
+    python my_app.py  # Automatically instrumented!
+
+Environment Variables:
+    AGENTREPLAY_ENABLED: Set to 'true' to enable auto-instrumentation
+    AGENTREPLAY_URL: Agentreplay server URL (default: http://localhost:47100)
+    AGENTREPLAY_TENANT_ID: Tenant ID (default: 1)
+    AGENTREPLAY_PROJECT_ID: Project ID (default: 0)
+    AGENTREPLAY_DEBUG: Set to 'true' for verbose logging
+    OTEL_SERVICE_NAME: Service name for traces
+    OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT: Capture message content
+
+Note:
+    This module is loaded very early in the Python startup process.
+    It must handle all errors gracefully to avoid breaking user applications.
+"""
+
+import os
+import sys
+
+# Only auto-instrument if explicitly enabled
+if os.getenv('AGENTREPLAY_ENABLED', '').lower() == 'true':
+    try:
+        # Import and initialize BEFORE any user code runs
+        from agentreplay.bootstrap import init_otel_instrumentation
+        
+        init_otel_instrumentation(
+            service_name=os.getenv('OTEL_SERVICE_NAME', os.path.basename(sys.argv[0])),
+            agentreplay_url=os.getenv('AGENTREPLAY_URL', 'http://localhost:47100'),
+            tenant_id=int(os.getenv('AGENTREPLAY_TENANT_ID', '1')),
+            project_id=int(os.getenv('AGENTREPLAY_PROJECT_ID', '0')),
+            capture_content=os.getenv('OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT', 'false').lower() == 'true'
+        )
+        
+        # Silent by default, verbose if DEBUG enabled
+        if os.getenv('AGENTREPLAY_DEBUG', '').lower() == 'true':
+            print("[Agentreplay] ✓ Auto-instrumentation enabled", file=sys.stderr)
+            print(f"[Agentreplay]   Service: {os.getenv('OTEL_SERVICE_NAME', os.path.basename(sys.argv[0]))}", file=sys.stderr)
+            print(f"[Agentreplay]   URL: {os.getenv('AGENTREPLAY_URL', 'http://localhost:47100')}", file=sys.stderr)
+            print(f"[Agentreplay]   Project: {os.getenv('AGENTREPLAY_PROJECT_ID', '0')}", file=sys.stderr)
+        
+    except ImportError as e:
+        if os.getenv('AGENTREPLAY_DEBUG', '').lower() == 'true':
+            print(f"[Agentreplay] ✗ Failed to auto-instrument: {e}", file=sys.stderr)
+            print("[Agentreplay]   Install: pip install opentelemetry-api opentelemetry-sdk", file=sys.stderr)
+    
+    except Exception as e:
+        if os.getenv('AGENTREPLAY_DEBUG', '').lower() == 'true':
+            print(f"[Agentreplay] ✗ Auto-instrumentation error: {e}", file=sys.stderr)
+            import traceback
+            traceback.print_exc(file=sys.stderr)