agentreplay 0.1.2__py3-none-any.whl

agentreplay/auto_instrument/openai.py

@@ -0,0 +1,431 @@
+# 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.
+
+"""OpenAI-specific instrumentation for streaming and tool calls.
+
+This module provides custom wrappers for OpenAI API calls to:
+1. Handle streaming responses (sync and async)
+2. Capture tool calls and their results
+3. Inject agent context into spans
+4. Respect content capture settings
+
+The wrappers are designed to work alongside the official OpenTelemetry
+OpenAI instrumentation, adding Agentreplay-specific enhancements.
+"""
+
+import logging
+from typing import Iterator, AsyncIterator, Optional, Any, Dict, List
+import json
+
+logger = logging.getLogger(__name__)
+
+# Configuration from environment
+import os
+CAPTURE_CONTENT = os.getenv("AGENTREPLAY_CAPTURE_CONTENT", "true").lower() in {
+    "1", "true", "yes"
+}
+MAX_CONTENT_LENGTH = int(os.getenv("AGENTREPLAY_MAX_CONTENT_LENGTH", "10000"))
+
+
+def is_streaming(response: Any) -> bool:
+    """Check if an OpenAI response is a stream.
+    
+    Args:
+        response: OpenAI API response
+    
+    Returns:
+        True if response is a stream, False otherwise
+    """
+    # Check for stream attribute or iterator protocol
+    if hasattr(response, "__iter__") and not isinstance(response, (str, bytes, dict)):
+        return True
+    if hasattr(response, "__aiter__"):
+        return True
+    return False
+
+
+class _StreamWrapper:
+    """Wrapper for synchronous OpenAI streaming responses.
+    
+    This wrapper:
+    - Yields chunks to the caller transparently
+    - Accumulates content for span attributes
+    - Handles tool calls in streaming mode
+    - Respects MAX_CONTENT_LENGTH
+    
+    Example:
+        >>> stream = client.chat.completions.create(..., stream=True)
+        >>> wrapped = _StreamWrapper(stream, span)
+        >>> for chunk in wrapped:
+        ...     print(chunk.choices[0].delta.content)
+    """
+    
+    def __init__(self, stream: Iterator, span: Optional[Any] = None):
+        """Initialize stream wrapper.
+        
+        Args:
+            stream: Original OpenAI stream
+            span: OpenTelemetry span to annotate (optional)
+        """
+        self.stream = stream
+        self.span = span
+        self.accumulated_content = []
+        # Tool calls accumulator: Dict[int, Dict] keyed by tool call index
+        # OpenAI streams tool calls as deltas that need to be merged by index
+        self.tool_calls_by_index: Dict[int, Dict[str, Any]] = {}
+        self.total_length = 0
+        self.chunk_count = 0
+        self._capture_content = CAPTURE_CONTENT
+    
+    def __iter__(self):
+        return self
+    
+    def __next__(self):
+        """Get next chunk from stream and capture metadata."""
+        try:
+            chunk = next(self.stream)
+            self.chunk_count += 1
+            
+            # Extract content if available
+            if hasattr(chunk, "choices") and len(chunk.choices) > 0:
+                choice = chunk.choices[0]
+                
+                # Capture text content
+                if hasattr(choice, "delta") and hasattr(choice.delta, "content"):
+                    content = choice.delta.content
+                    if content and self._capture_content:
+                        self.accumulated_content.append(content)
+                        self.total_length += len(content)
+                        
+                        # Stop accumulating if we exceed max length
+                        if self.total_length > MAX_CONTENT_LENGTH:
+                            self.accumulated_content.append(
+                                f"... (truncated, total {self.total_length} chars)"
+                            )
+                            self._capture_content = False  # Disable for rest of stream
+                
+                # Capture tool calls - merge deltas by index
+                if hasattr(choice.delta, "tool_calls") and choice.delta.tool_calls:
+                    for tool_call in choice.delta.tool_calls:
+                        # Get the index for this tool call (OpenAI sends this)
+                        idx = getattr(tool_call, "index", 0)
+                        
+                        # Initialize if first delta for this index
+                        if idx not in self.tool_calls_by_index:
+                            self.tool_calls_by_index[idx] = {
+                                "id": None,
+                                "name": None,
+                                "arguments": "",
+                            }
+                        
+                        acc = self.tool_calls_by_index[idx]
+                        
+                        # Merge id (usually only in first delta)
+                        if hasattr(tool_call, "id") and tool_call.id:
+                            acc["id"] = tool_call.id
+                        
+                        # Merge function name and arguments
+                        if hasattr(tool_call, "function"):
+                            func = tool_call.function
+                            if hasattr(func, "name") and func.name:
+                                acc["name"] = func.name
+                            if hasattr(func, "arguments") and func.arguments:
+                                # Arguments come as fragments, concatenate them
+                                acc["arguments"] += func.arguments
+            
+            return chunk
+            
+        except StopIteration:
+            # Stream ended, finalize span
+            self._finalize_span()
+            raise
+    
+    def _finalize_span(self):
+        """Add accumulated data to span when stream completes."""
+        if not self.span:
+            return
+        
+        try:
+            # Add accumulated content
+            if self.accumulated_content:
+                full_content = "".join(self.accumulated_content)
+                if self._capture_content:
+                    self.span.set_attribute("llm.response.content", full_content[:MAX_CONTENT_LENGTH])
+                self.span.set_attribute("llm.response.length", self.total_length)
+            
+            # Add tool calls (merged by index)
+            if self.tool_calls_by_index:
+                # Sort by index for consistent ordering
+                sorted_tool_calls = [
+                    self.tool_calls_by_index[idx] 
+                    for idx in sorted(self.tool_calls_by_index.keys())
+                ]
+                self.span.set_attribute("llm.tool_calls.count", len(sorted_tool_calls))
+                for i, tool_call in enumerate(sorted_tool_calls[:10]):  # Max 10
+                    if tool_call.get("name"):
+                        self.span.set_attribute(f"llm.tool_call.{i}.name", tool_call["name"])
+                    if tool_call.get("id"):
+                        self.span.set_attribute(f"llm.tool_call.{i}.id", tool_call["id"])
+                    if tool_call.get("arguments") and self._capture_content:
+                        args = tool_call["arguments"][:500]  # Truncate args
+                        self.span.set_attribute(f"llm.tool_call.{i}.arguments", args)
+            
+            # Add streaming metadata
+            self.span.set_attribute("llm.streaming", True)
+            self.span.set_attribute("llm.stream.chunks", self.chunk_count)
+            
+        except Exception as e:
+            logger.debug(f"Failed to finalize stream span: {e}")
+
+
+class _AsyncStreamWrapper:
+    """Wrapper for asynchronous OpenAI streaming responses.
+    
+    Similar to _StreamWrapper but for async/await code.
+    
+    Example:
+        >>> stream = await client.chat.completions.create(..., stream=True)
+        >>> wrapped = _AsyncStreamWrapper(stream, span)
+        >>> async for chunk in wrapped:
+        ...     print(chunk.choices[0].delta.content)
+    """
+    
+    def __init__(self, stream: AsyncIterator, span: Optional[Any] = None):
+        """Initialize async stream wrapper.
+        
+        Args:
+            stream: Original OpenAI async stream
+            span: OpenTelemetry span to annotate (optional)
+        """
+        self.stream = stream
+        self.span = span
+        self.accumulated_content = []
+        # Tool calls accumulator: Dict[int, Dict] keyed by tool call index
+        # OpenAI streams tool calls as deltas that need to be merged by index
+        self.tool_calls_by_index: Dict[int, Dict[str, Any]] = {}
+        self.total_length = 0
+        self.chunk_count = 0
+        self._capture_content = CAPTURE_CONTENT
+    
+    def __aiter__(self):
+        return self
+    
+    async def __anext__(self):
+        """Get next chunk from async stream and capture metadata."""
+        try:
+            chunk = await self.stream.__anext__()
+            self.chunk_count += 1
+            
+            # Extract content if available
+            if hasattr(chunk, "choices") and len(chunk.choices) > 0:
+                choice = chunk.choices[0]
+                
+                # Capture text content
+                if hasattr(choice, "delta") and hasattr(choice.delta, "content"):
+                    content = choice.delta.content
+                    if content and self._capture_content:
+                        self.accumulated_content.append(content)
+                        self.total_length += len(content)
+                        
+                        # Stop accumulating if we exceed max length
+                        if self.total_length > MAX_CONTENT_LENGTH:
+                            self.accumulated_content.append(
+                                f"... (truncated, total {self.total_length} chars)"
+                            )
+                            self._capture_content = False
+                
+                # Capture tool calls - merge deltas by index
+                if hasattr(choice.delta, "tool_calls") and choice.delta.tool_calls:
+                    for tool_call in choice.delta.tool_calls:
+                        # Get the index for this tool call (OpenAI sends this)
+                        idx = getattr(tool_call, "index", 0)
+                        
+                        # Initialize if first delta for this index
+                        if idx not in self.tool_calls_by_index:
+                            self.tool_calls_by_index[idx] = {
+                                "id": None,
+                                "name": None,
+                                "arguments": "",
+                            }
+                        
+                        acc = self.tool_calls_by_index[idx]
+                        
+                        # Merge id (usually only in first delta)
+                        if hasattr(tool_call, "id") and tool_call.id:
+                            acc["id"] = tool_call.id
+                        
+                        # Merge function name and arguments
+                        if hasattr(tool_call, "function"):
+                            func = tool_call.function
+                            if hasattr(func, "name") and func.name:
+                                acc["name"] = func.name
+                            if hasattr(func, "arguments") and func.arguments:
+                                # Arguments come as fragments, concatenate them
+                                acc["arguments"] += func.arguments
+            
+            return chunk
+            
+        except StopAsyncIteration:
+            # Stream ended, finalize span
+            self._finalize_span()
+            raise
+    
+    def _finalize_span(self):
+        """Add accumulated data to span when stream completes."""
+        if not self.span:
+            return
+        
+        try:
+            # Add accumulated content
+            if self.accumulated_content:
+                full_content = "".join(self.accumulated_content)
+                if self._capture_content:
+                    self.span.set_attribute("llm.response.content", full_content[:MAX_CONTENT_LENGTH])
+                self.span.set_attribute("llm.response.length", self.total_length)
+            
+            # Add tool calls (merged by index)
+            if self.tool_calls_by_index:
+                # Sort by index for consistent ordering
+                sorted_tool_calls = [
+                    self.tool_calls_by_index[idx] 
+                    for idx in sorted(self.tool_calls_by_index.keys())
+                ]
+                self.span.set_attribute("llm.tool_calls.count", len(sorted_tool_calls))
+                for i, tool_call in enumerate(sorted_tool_calls[:10]):  # Max 10
+                    if tool_call.get("name"):
+                        self.span.set_attribute(f"llm.tool_call.{i}.name", tool_call["name"])
+                    if tool_call.get("id"):
+                        self.span.set_attribute(f"llm.tool_call.{i}.id", tool_call["id"])
+                    if tool_call.get("arguments") and self._capture_content:
+                        args = tool_call["arguments"][:500]  # Truncate args
+                        self.span.set_attribute(f"llm.tool_call.{i}.arguments", args)
+            
+            # Add streaming metadata
+            self.span.set_attribute("llm.streaming", True)
+            self.span.set_attribute("llm.stream.chunks", self.chunk_count)
+            
+        except Exception as e:
+            logger.debug(f"Failed to finalize async stream span: {e}")
+
+
+def _inject_agent_context(span: Any):
+    """Inject current agent context into span attributes.
+    
+    Reads from contextvars set by AgentContext and adds them to the span.
+    
+    Args:
+        span: OpenTelemetry span
+    """
+    try:
+        from agentreplay.context import (
+            get_current_agent_id,
+            get_current_session_id,
+            get_current_workflow_id,
+            get_current_user_id,
+        )
+        
+        agent_id = get_current_agent_id()
+        if agent_id:
+            span.set_attribute("agentreplay.agent_id", agent_id)
+        
+        session_id = get_current_session_id()
+        if session_id:
+            span.set_attribute("agentreplay.session_id", session_id)
+        
+        workflow_id = get_current_workflow_id()
+        if workflow_id:
+            span.set_attribute("agentreplay.workflow_id", workflow_id)
+        
+        user_id = get_current_user_id()
+        if user_id:
+            span.set_attribute("agentreplay.user_id", user_id)
+    
+    except ImportError:
+        # Context module not available
+        pass
+    except Exception as e:
+        logger.debug(f"Failed to inject agent context: {e}")
+
+
+def extract_tool_calls(response: Any) -> List[Dict[str, Any]]:
+    """Extract tool calls from OpenAI response.
+    
+    Args:
+        response: OpenAI chat completion response
+    
+    Returns:
+        List of tool call dictionaries with id, name, arguments
+    """
+    tool_calls = []
+    
+    try:
+        if not hasattr(response, "choices") or not response.choices:
+            return tool_calls
+        
+        choice = response.choices[0]
+        if not hasattr(choice, "message") or not hasattr(choice.message, "tool_calls"):
+            return tool_calls
+        
+        if not choice.message.tool_calls:
+            return tool_calls
+        
+        for tool_call in choice.message.tool_calls:
+            if hasattr(tool_call, "function"):
+                tool_calls.append({
+                    "id": getattr(tool_call, "id", None),
+                    "type": getattr(tool_call, "type", "function"),
+                    "name": getattr(tool_call.function, "name", None),
+                    "arguments": getattr(tool_call.function, "arguments", None),
+                })
+    
+    except Exception as e:
+        logger.debug(f"Failed to extract tool calls: {e}")
+    
+    return tool_calls
+
+
+def annotate_span_with_tool_calls(span: Any, tool_calls: List[Dict[str, Any]]):
+    """Add tool call information to span attributes.
+    
+    Args:
+        span: OpenTelemetry span
+        tool_calls: List of tool call dictionaries
+    """
+    if not tool_calls:
+        return
+    
+    try:
+        span.set_attribute("llm.tool_calls.count", len(tool_calls))
+        
+        for i, tool_call in enumerate(tool_calls[:10]):  # Max 10 tool calls
+            prefix = f"llm.tool_call.{i}"
+            
+            if tool_call.get("id"):
+                span.set_attribute(f"{prefix}.id", tool_call["id"])
+            
+            if tool_call.get("name"):
+                span.set_attribute(f"{prefix}.name", tool_call["name"])
+            
+            if tool_call.get("type"):
+                span.set_attribute(f"{prefix}.type", tool_call["type"])
+            
+            if tool_call.get("arguments") and CAPTURE_CONTENT:
+                args = tool_call["arguments"]
+                if isinstance(args, str):
+                    # Truncate if too long
+                    args = args[:500]
+                    span.set_attribute(f"{prefix}.arguments", args)
+    
+    except Exception as e:
+        logger.debug(f"Failed to annotate span with tool calls: {e}")

agentreplay/batching.py

@@ -0,0 +1,270 @@
+# 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.
+
+"""Batching client for high-throughput trace ingestion."""
+
+import threading
+import time
+from collections import deque
+from typing import Deque, List, Optional
+from agentreplay.client import AgentreplayClient
+from agentreplay.models import AgentFlowEdge
+
+
+class BatchingAgentreplayClient:
+    """Client wrapper that batches spans for efficient ingestion.
+
+    Automatically buffers spans and flushes when batch size is reached
+    or flush interval expires, reducing HTTP overhead by 100x.
+
+    **CRITICAL FIX**: Now includes max_buffer_size to prevent OOM.
+    When buffer is full, oldest edges are dropped (sampling behavior).
+
+    Args:
+        client: Underlying AgentreplayClient
+        batch_size: Number of spans to buffer before flushing (default: 100)
+        flush_interval: Seconds between automatic flushes (default: 5.0)
+        max_buffer_size: Maximum buffer size before dropping edges (default: 10000)
+
+    Example:
+        >>> client = AgentreplayClient(url="http://localhost:8080", tenant_id=1)
+        >>> batching_client = BatchingAgentreplayClient(
+        ...     client,
+        ...     batch_size=100,
+        ...     max_buffer_size=10000  # Prevent OOM
+        ... )
+        >>>
+        >>> # Spans are buffered
+        >>> for i in range(1000):
+        ...     edge = AgentFlowEdge(
+        ...         tenant_id=1,
+        ...         agent_id=1,
+        ...         session_id=42,
+        ...         span_type=SpanType.ROOT
+        ...     )
+        ...     batching_client.insert(edge)  # Buffered, not sent immediately
+        ...
+        >>> # Flush remaining spans
+        >>> batching_client.flush()
+        >>> batching_client.close()
+    """
+
+    def __init__(
+        self,
+        client: AgentreplayClient,
+        batch_size: int = 100,
+        flush_interval: float = 5.0,
+        max_buffer_size: int = 10000,
+    ):
+        """Initialize batching client."""
+        self.client = client
+        self.batch_size = batch_size
+        self.flush_interval = flush_interval
+        self.max_buffer_size = max_buffer_size
+        self._buffer: List[AgentFlowEdge] = []
+        self._retry_queue: Deque[List[AgentFlowEdge]] = deque()  # Failed batches awaiting retry
+        self._max_retry_batches = 10  # Limit retry queue to prevent unbounded growth
+        self._lock = threading.Lock()
+        self._running = True
+        self._dropped_count = 0  # Track dropped edges for monitoring
+        self._flush_thread = threading.Thread(target=self._auto_flush, daemon=True)
+        self._flush_thread.start()
+
+    def __enter__(self) -> "BatchingAgentreplayClient":
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Context manager exit - flush and close."""
+        self.flush()
+        self.close()
+
+    def insert(self, edge: AgentFlowEdge) -> AgentFlowEdge:
+        """Buffer a single edge for batched insertion.
+
+        **CRITICAL FIX**: Now enforces max_buffer_size to prevent OOM.
+        If buffer is full, drops oldest edges (FIFO sampling).
+
+        Args:
+            edge: Edge to buffer
+
+        Returns:
+            The same edge (for consistency with AgentreplayClient API)
+        """
+        with self._lock:
+            # CRITICAL FIX: Enforce max buffer size to prevent OOM
+            if len(self._buffer) >= self.max_buffer_size:
+                # Drop oldest edge (FIFO sampling)
+                self._buffer.pop(0)
+                self._dropped_count += 1
+
+                # Log warning every 1000 drops
+                if self._dropped_count % 1000 == 0:
+                    print(
+                        f"WARNING: Dropped {self._dropped_count} edges due to full buffer. "
+                        f"Backend may be slow or down. Consider increasing max_buffer_size "
+                        f"or reducing ingestion rate."
+                    )
+
+            self._buffer.append(edge)
+
+            # Flush if batch size reached
+            if len(self._buffer) >= self.batch_size:
+                self._flush_unlocked()
+
+        return edge
+
+    def flush(self) -> int:
+        """Manually flush all buffered spans.
+        
+        Returns:
+            Number of spans flushed
+        """
+        with self._lock:
+            return self._flush_unlocked()
+
+    def _flush_unlocked(self) -> int:
+        """Flush buffer without acquiring lock (caller must hold lock).
+        
+        Returns:
+            Number of spans flushed
+        """
+        if not self._buffer:
+            return 0
+
+        # Send entire batch in one HTTP request
+        try:
+            self.client.insert_batch(self._buffer)
+            count = len(self._buffer)
+            self._buffer = []
+            return count
+        except Exception as e:
+            # Log error but don't lose spans
+            print(f"Error flushing batch: {e}")
+            return 0
+
+    def _auto_flush(self) -> None:
+        """Background thread that flushes buffer periodically.
+        
+        CRITICAL FIX: Transactional buffer management - data is only removed
+        from buffer after successful network I/O. Failed batches are queued
+        for retry to prevent data loss.
+        """
+        while self._running:
+            time.sleep(self.flush_interval)
+            if self._running:  # Check again after sleep
+                # 1. First, try to send any previously failed batches
+                self._process_retry_queue()
+                
+                # 2. Grab data to send (under lock) but DON'T clear buffer yet
+                batch_to_send = []
+                with self._lock:
+                    if self._buffer:
+                        batch_to_send = self._buffer[:]  # Copy for sending
+                
+                # 3. Send I/O outside the lock (won't block application threads)
+                if batch_to_send:
+                    success = False
+                    try:
+                        self.client.insert_batch(batch_to_send)
+                        success = True
+                    except Exception as e:
+                        # Log error but don't crash the thread
+                        print(f"Error flushing batch: {e}")
+                    
+                    # 4. ONLY clear buffer after confirmed success (transactional)
+                    with self._lock:
+                        if success:
+                            # Remove only the items we successfully sent
+                            # (new items may have been added during I/O)
+                            self._buffer = self._buffer[len(batch_to_send):]
+                        else:
+                            # Failed: queue batch for retry, clear from main buffer
+                            # to prevent duplicate sends
+                            self._buffer = self._buffer[len(batch_to_send):]
+                            self._queue_for_retry(batch_to_send)
+
+    def _process_retry_queue(self) -> None:
+        """Process failed batches from retry queue."""
+        while self._retry_queue:
+            # Pop from front (FIFO)
+            with self._lock:
+                if not self._retry_queue:
+                    break
+                batch = self._retry_queue.popleft()
+            
+            try:
+                self.client.insert_batch(batch)
+                # Success - batch is now sent, continue to next
+            except Exception as e:
+                # Still failing - re-queue at the back for later retry
+                print(f"Retry failed for batch of {len(batch)} edges: {e}")
+                with self._lock:
+                    if len(self._retry_queue) < self._max_retry_batches:
+                        self._retry_queue.append(batch)
+                    else:
+                        # Drop batch to prevent unbounded growth
+                        self._dropped_count += len(batch)
+                        print(f"WARNING: Dropped batch of {len(batch)} edges after max retries")
+                break  # Stop processing retry queue on failure
+
+    def _queue_for_retry(self, batch: List[AgentFlowEdge]) -> None:
+        """Add failed batch to retry queue (caller must NOT hold lock)."""
+        if len(self._retry_queue) < self._max_retry_batches:
+            self._retry_queue.append(batch)
+        else:
+            # Drop oldest retry batch to make room
+            dropped = self._retry_queue.popleft()
+            self._dropped_count += len(dropped)
+            self._retry_queue.append(batch)
+            print(f"WARNING: Dropped oldest retry batch ({len(dropped)} edges) to make room")
+
+    def close(self) -> None:
+        """Stop auto-flush thread and flush remaining spans including retry queue."""
+        self._running = False
+        if self._flush_thread.is_alive():
+            self._flush_thread.join(timeout=self.flush_interval + 1.0)
+        
+        # Flush main buffer
+        self.flush()
+        
+        # Attempt to flush retry queue (best effort)
+        retry_attempts = 0
+        max_close_retries = 3
+        while self._retry_queue and retry_attempts < max_close_retries:
+            retry_attempts += 1
+            with self._lock:
+                if not self._retry_queue:
+                    break
+                batch = self._retry_queue.popleft()
+            try:
+                self.client.insert_batch(batch)
+            except Exception as e:
+                print(f"Failed to flush retry queue on close (attempt {retry_attempts}): {e}")
+                # Re-queue for next attempt
+                with self._lock:
+                    self._retry_queue.appendleft(batch)
+                break
+        
+        # Report any remaining data that couldn't be sent
+        remaining = sum(len(b) for b in self._retry_queue)
+        if remaining > 0:
+            print(f"WARNING: {remaining} edges in retry queue could not be sent on close")
+
+    def __del__(self) -> None:
+        """Destructor - ensure spans are flushed."""
+        try:
+            self.close()
+        except:
+            pass  # Ignore errors during cleanup