mseep-agentops 0.4.18__py3-none-any.whl

agentops/instrumentation/common/span_management.py

@@ -0,0 +1,176 @@
+"""Common span management utilities for AgentOps instrumentation.
+
+This module provides utilities for creating and managing spans with
+consistent attributes and error handling.
+"""
+
+import time
+from contextlib import contextmanager
+from typing import Optional, Dict, Any, Callable, Tuple
+from functools import wraps
+
+from opentelemetry.trace import Tracer, Span, SpanKind, Status, StatusCode, get_current_span
+from opentelemetry.sdk.resources import SERVICE_NAME, TELEMETRY_SDK_NAME, DEPLOYMENT_ENVIRONMENT
+from opentelemetry import context as context_api
+
+from agentops.logging import logger
+from agentops.semconv import CoreAttributes
+
+
+class SpanAttributeManager:
+    """Manages common span attributes across instrumentations."""
+
+    def __init__(self, service_name: str = "agentops", deployment_environment: str = "production"):
+        self.service_name = service_name
+        self.deployment_environment = deployment_environment
+
+    def set_common_attributes(self, span: Span):
+        """Set common attributes on a span."""
+        span.set_attribute(TELEMETRY_SDK_NAME, "agentops")
+        span.set_attribute(SERVICE_NAME, self.service_name)
+        span.set_attribute(DEPLOYMENT_ENVIRONMENT, self.deployment_environment)
+
+    def set_config_tags(self, span: Span):
+        """Set tags from AgentOps config on a span."""
+        # Import locally to avoid circular dependency
+        from agentops import get_client
+
+        client = get_client()
+        if client and client.config.default_tags and len(client.config.default_tags) > 0:
+            tag_list = list(client.config.default_tags)
+            span.set_attribute(CoreAttributes.TAGS, tag_list)
+
+
+@contextmanager
+def create_span(
+    tracer: Tracer,
+    name: str,
+    kind: SpanKind = SpanKind.CLIENT,
+    attributes: Optional[Dict[str, Any]] = None,
+    set_common_attributes: bool = True,
+    attribute_manager: Optional[SpanAttributeManager] = None,
+):
+    """Context manager for creating spans with consistent error handling.
+
+    Args:
+        tracer: The tracer to use for creating the span
+        name: The name of the span
+        kind: The kind of span to create
+        attributes: Initial attributes to set on the span
+        set_common_attributes: Whether to set common attributes
+        attribute_manager: Optional attribute manager for setting common attributes
+
+    Yields:
+        The created span
+    """
+    with tracer.start_as_current_span(name, kind=kind, attributes=attributes or {}) as span:
+        try:
+            if set_common_attributes and attribute_manager:
+                attribute_manager.set_common_attributes(span)
+            yield span
+            span.set_status(Status(StatusCode.OK))
+        except Exception as e:
+            span.set_status(Status(StatusCode.ERROR, str(e)))
+            span.record_exception(e)
+            logger.error(f"Error in span {name}: {e}")
+            raise
+
+
+def timed_span(tracer: Tracer, name: str, record_duration: Optional[Callable[[float], None]] = None, **span_kwargs):
+    """Decorator for creating timed spans around functions.
+
+    Args:
+        tracer: The tracer to use
+        name: The name of the span
+        record_duration: Optional callback to record duration
+        **span_kwargs: Additional arguments for span creation
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            start_time = time.time()
+            with create_span(tracer, name, **span_kwargs):
+                result = func(*args, **kwargs)
+                if record_duration:
+                    duration = time.time() - start_time
+                    record_duration(duration)
+                return result
+
+        return wrapper
+
+    return decorator
+
+
+class StreamingSpanManager:
+    """Manages spans for streaming operations."""
+
+    def __init__(self, tracer: Tracer):
+        self.tracer = tracer
+        self._active_spans: Dict[Any, Span] = {}
+
+    def start_streaming_span(self, stream_id: Any, name: str, **span_kwargs) -> Span:
+        """Start a span for a streaming operation."""
+        span = self.tracer.start_span(name, **span_kwargs)
+        self._active_spans[stream_id] = span
+        return span
+
+    def get_streaming_span(self, stream_id: Any) -> Optional[Span]:
+        """Get an active streaming span."""
+        return self._active_spans.get(stream_id)
+
+    def end_streaming_span(self, stream_id: Any, status: Optional[Status] = None):
+        """End a streaming span."""
+        span = self._active_spans.pop(stream_id, None)
+        if span:
+            if status:
+                span.set_status(status)
+            else:
+                span.set_status(Status(StatusCode.OK))
+            span.end()
+
+
+def extract_parent_context(parent_span: Optional[Span] = None) -> Any:
+    """Extract parent context for span creation.
+
+    Args:
+        parent_span: Optional parent span to use
+
+    Returns:
+        Context to use as parent
+    """
+    if parent_span:
+        from opentelemetry.trace import set_span_in_context
+
+        return set_span_in_context(parent_span)
+    return context_api.get_current()
+
+
+def safe_set_attribute(span: Span, key: str, value: Any, max_length: int = 1000):
+    """Safely set an attribute on a span, handling None values and truncating long strings."""
+    if value is None:
+        return
+
+    if isinstance(value, str) and len(value) > max_length:
+        value = value[: max_length - 3] + "..."
+
+    try:
+        span.set_attribute(key, value)
+    except Exception as e:
+        logger.debug(f"Failed to set span attribute {key}: {e}")
+
+
+def get_span_context_info(span: Optional[Span] = None) -> Tuple[str, str]:
+    """Get trace and span IDs from a span for debugging.
+
+    Returns:
+        Tuple of (trace_id, span_id) as strings
+    """
+    if not span:
+        span = get_current_span()
+
+    span_context = span.get_span_context()
+    trace_id = format(span_context.trace_id, "032x") if span_context.trace_id else "unknown"
+    span_id = format(span_context.span_id, "016x") if span_context.span_id else "unknown"
+
+    return trace_id, span_id

agentops/instrumentation/common/streaming.py

@@ -0,0 +1,218 @@
+"""Common streaming utilities for handling streaming responses.
+
+This module provides utilities for instrumenting streaming API responses
+in a consistent way across different providers.
+"""
+
+from typing import Optional, Any, Dict, Callable
+from abc import ABC
+import time
+
+from opentelemetry.trace import Tracer, Span, Status, StatusCode
+
+from agentops.logging import logger
+from agentops.instrumentation.common.span_management import safe_set_attribute
+from agentops.instrumentation.common.token_counting import TokenUsage, TokenUsageExtractor
+
+
+class BaseStreamWrapper(ABC):
+    """Base class for wrapping streaming responses."""
+
+    def __init__(
+        self,
+        stream: Any,
+        span: Span,
+        extract_chunk_content: Callable[[Any], Optional[str]],
+        extract_chunk_attributes: Optional[Callable[[Any], Dict[str, Any]]] = None,
+    ):
+        self.stream = stream
+        self.span = span
+        self.extract_chunk_content = extract_chunk_content
+        self.extract_chunk_attributes = extract_chunk_attributes or (lambda x: {})
+
+        self.start_time = time.time()
+        self.first_token_time: Optional[float] = None
+        self.chunks_received = 0
+        self.accumulated_content = []
+        self.token_usage = TokenUsage()
+
+    def _process_chunk(self, chunk: Any):
+        """Process a single chunk from the stream."""
+        # Record time to first token
+        if self.first_token_time is None:
+            self.first_token_time = time.time()
+            time_to_first_token = self.first_token_time - self.start_time
+            safe_set_attribute(self.span, "streaming.time_to_first_token", time_to_first_token)
+
+        self.chunks_received += 1
+
+        # Extract content from chunk
+        content = self.extract_chunk_content(chunk)
+        if content:
+            self.accumulated_content.append(content)
+
+        # Extract and set additional attributes
+        attributes = self.extract_chunk_attributes(chunk)
+        for key, value in attributes.items():
+            safe_set_attribute(self.span, key, value)
+
+        # Try to extract token usage if available
+        if hasattr(chunk, "usage") or hasattr(chunk, "usage_metadata"):
+            chunk_usage = TokenUsageExtractor.extract_from_response(chunk)
+            # Accumulate token counts
+            if chunk_usage.prompt_tokens:
+                self.token_usage.prompt_tokens = chunk_usage.prompt_tokens
+            if chunk_usage.completion_tokens:
+                self.token_usage.completion_tokens = (
+                    self.token_usage.completion_tokens or 0
+                ) + chunk_usage.completion_tokens
+
+    def _finalize(self):
+        """Finalize the stream processing."""
+        try:
+            # Set final content
+            final_content = "".join(self.accumulated_content)
+            safe_set_attribute(self.span, "streaming.final_content", final_content)
+            safe_set_attribute(self.span, "streaming.chunk_count", self.chunks_received)
+
+            # Set timing metrics
+            total_time = time.time() - self.start_time
+            safe_set_attribute(self.span, "streaming.total_duration", total_time)
+
+            if self.first_token_time:
+                generation_time = time.time() - self.first_token_time
+                safe_set_attribute(self.span, "streaming.generation_duration", generation_time)
+
+            # Set token usage
+            for attr_name, value in self.token_usage.to_attributes().items():
+                safe_set_attribute(self.span, attr_name, value)
+
+            self.span.set_status(Status(StatusCode.OK))
+        except Exception as e:
+            logger.error(f"Error finalizing stream: {e}")
+            self.span.set_status(Status(StatusCode.ERROR, str(e)))
+            self.span.record_exception(e)
+        finally:
+            self.span.end()
+
+
+class SyncStreamWrapper(BaseStreamWrapper):
+    """Wrapper for synchronous streaming responses."""
+
+    def __iter__(self):
+        try:
+            for chunk in self.stream:
+                self._process_chunk(chunk)
+                yield chunk
+        except Exception as e:
+            self.span.set_status(Status(StatusCode.ERROR, str(e)))
+            self.span.record_exception(e)
+            raise
+        finally:
+            self._finalize()
+
+
+class AsyncStreamWrapper(BaseStreamWrapper):
+    """Wrapper for asynchronous streaming responses."""
+
+    async def __aiter__(self):
+        try:
+            async for chunk in self.stream:
+                self._process_chunk(chunk)
+                yield chunk
+        except Exception as e:
+            self.span.set_status(Status(StatusCode.ERROR, str(e)))
+            self.span.record_exception(e)
+            raise
+        finally:
+            self._finalize()
+
+
+def create_stream_wrapper_factory(
+    tracer: Tracer,
+    span_name: str,
+    extract_chunk_content: Callable[[Any], Optional[str]],
+    extract_chunk_attributes: Optional[Callable[[Any], Dict[str, Any]]] = None,
+    initial_attributes: Optional[Dict[str, Any]] = None,
+) -> Callable:
+    """Create a factory function for wrapping streaming methods.
+
+    Args:
+        tracer: The tracer to use for creating spans
+        span_name: Name for the streaming span
+        extract_chunk_content: Function to extract content from chunks
+        extract_chunk_attributes: Optional function to extract attributes from chunks
+        initial_attributes: Initial attributes to set on the span
+
+    Returns:
+        A wrapper function suitable for use with wrapt
+    """
+
+    def wrapper(wrapped, instance, args, kwargs):
+        # Start the span
+        span = tracer.start_span(span_name)
+
+        # Set initial attributes
+        if initial_attributes:
+            for key, value in initial_attributes.items():
+                safe_set_attribute(span, key, value)
+
+        try:
+            # Call the wrapped method
+            stream = wrapped(*args, **kwargs)
+
+            # Determine if it's async or sync
+            if hasattr(stream, "__aiter__"):
+                return AsyncStreamWrapper(stream, span, extract_chunk_content, extract_chunk_attributes)
+            else:
+                return SyncStreamWrapper(stream, span, extract_chunk_content, extract_chunk_attributes)
+        except Exception as e:
+            span.set_status(Status(StatusCode.ERROR, str(e)))
+            span.record_exception(e)
+            span.end()
+            raise
+
+    return wrapper
+
+
+class StreamingResponseHandler:
+    """Handles common patterns for streaming responses."""
+
+    @staticmethod
+    def extract_openai_chunk_content(chunk: Any) -> Optional[str]:
+        """Extract content from OpenAI-style streaming chunks."""
+        if hasattr(chunk, "choices") and chunk.choices:
+            delta = getattr(chunk.choices[0], "delta", None)
+            if delta and hasattr(delta, "content"):
+                return delta.content
+        return None
+
+    @staticmethod
+    def extract_anthropic_chunk_content(chunk: Any) -> Optional[str]:
+        """Extract content from Anthropic-style streaming chunks."""
+        if hasattr(chunk, "type"):
+            if chunk.type == "content_block_delta":
+                if hasattr(chunk, "delta") and hasattr(chunk.delta, "text"):
+                    return chunk.delta.text
+            elif chunk.type == "message_delta":
+                if hasattr(chunk, "delta") and hasattr(chunk.delta, "content"):
+                    return chunk.delta.content
+        return None
+
+    @staticmethod
+    def extract_generic_chunk_content(chunk: Any) -> Optional[str]:
+        """Extract content from generic streaming chunks."""
+        # Try common patterns
+        if hasattr(chunk, "content"):
+            return str(chunk.content)
+        elif hasattr(chunk, "text"):
+            return str(chunk.text)
+        elif hasattr(chunk, "delta"):
+            delta = chunk.delta
+            if hasattr(delta, "content"):
+                return str(delta.content)
+            elif hasattr(delta, "text"):
+                return str(delta.text)
+        elif isinstance(chunk, str):
+            return chunk
+        return None

agentops/instrumentation/common/token_counting.py

@@ -0,0 +1,177 @@
+"""Common token counting and usage extraction utilities.
+
+This module provides utilities for extracting and recording token usage
+information from various response formats.
+"""
+
+from typing import Dict, Any, Optional
+from dataclasses import dataclass
+
+from agentops.logging import logger
+from agentops.semconv import SpanAttributes
+
+
+@dataclass
+class TokenUsage:
+    """Represents token usage information."""
+
+    prompt_tokens: Optional[int] = None
+    completion_tokens: Optional[int] = None
+    total_tokens: Optional[int] = None
+    cached_prompt_tokens: Optional[int] = None
+    cached_read_tokens: Optional[int] = None
+    reasoning_tokens: Optional[int] = None
+
+    def to_attributes(self) -> Dict[str, int]:
+        """Convert to span attributes dictionary.
+
+        Only metrics greater than zero are included so that non‑LLM spans do
+        not contain empty token usage fields.
+        """
+        attributes = {}
+
+        if self.prompt_tokens:
+            attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] = self.prompt_tokens
+
+        if self.completion_tokens:
+            attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] = self.completion_tokens
+
+        if self.total_tokens:
+            attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = self.total_tokens
+
+        if self.cached_prompt_tokens:
+            attributes[SpanAttributes.LLM_USAGE_CACHE_CREATION_INPUT_TOKENS] = self.cached_prompt_tokens
+
+        if self.cached_read_tokens:
+            attributes[SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS] = self.cached_read_tokens
+
+        if self.reasoning_tokens:
+            attributes[SpanAttributes.LLM_USAGE_REASONING_TOKENS] = self.reasoning_tokens
+
+        return attributes
+
+
+class TokenUsageExtractor:
+    """Extracts token usage from various response formats."""
+
+    @staticmethod
+    def extract_from_response(response: Any) -> TokenUsage:
+        """Extract token usage from a generic response object.
+
+        Handles various response formats from different providers.
+        """
+        usage = TokenUsage()
+
+        # Try direct usage attribute
+        if hasattr(response, "usage"):
+            usage_data = response.usage
+            usage = TokenUsageExtractor._extract_from_usage_object(usage_data)
+
+        # Try usage_metadata (Anthropic style)
+        elif hasattr(response, "usage_metadata"):
+            usage_data = response.usage_metadata
+            usage = TokenUsageExtractor._extract_from_usage_object(usage_data)
+
+        # Try token_usage attribute (CrewAI style)
+        elif hasattr(response, "token_usage"):
+            usage = TokenUsageExtractor._extract_from_crewai_format(response.token_usage)
+
+        # Try direct attributes on response
+        elif hasattr(response, "prompt_tokens") or hasattr(response, "completion_tokens"):
+            usage = TokenUsageExtractor._extract_from_attributes(response)
+
+        return usage
+
+    @staticmethod
+    def _extract_from_usage_object(usage_data: Any) -> TokenUsage:
+        """Extract from a usage object with standard attributes."""
+        if not usage_data:
+            return TokenUsage()
+
+        return TokenUsage(
+            prompt_tokens=getattr(usage_data, "prompt_tokens", None),
+            completion_tokens=getattr(usage_data, "completion_tokens", None),
+            total_tokens=getattr(usage_data, "total_tokens", None),
+            cached_prompt_tokens=getattr(usage_data, "cached_prompt_tokens", None),
+            cached_read_tokens=getattr(usage_data, "cache_read_input_tokens", None),
+            reasoning_tokens=getattr(usage_data, "reasoning_tokens", None),
+        )
+
+    @staticmethod
+    def _extract_from_crewai_format(token_usage_str: str) -> TokenUsage:
+        """Extract from CrewAI's string format (e.g., 'prompt_tokens=100 completion_tokens=50')."""
+        usage = TokenUsage()
+
+        try:
+            metrics = {}
+            for item in str(token_usage_str).split():
+                if "=" in item:
+                    key, value = item.split("=")
+                    try:
+                        metrics[key] = int(value)
+                    except ValueError:
+                        pass
+
+            usage.prompt_tokens = metrics.get("prompt_tokens")
+            usage.completion_tokens = metrics.get("completion_tokens")
+            usage.total_tokens = metrics.get("total_tokens")
+            usage.cached_prompt_tokens = metrics.get("cached_prompt_tokens")
+
+        except Exception as e:
+            logger.debug(f"Failed to parse CrewAI token usage: {e}")
+
+        return usage
+
+    @staticmethod
+    def _extract_from_attributes(response: Any) -> TokenUsage:
+        """Extract from direct attributes on the response."""
+        return TokenUsage(
+            prompt_tokens=getattr(response, "prompt_tokens", None),
+            completion_tokens=getattr(response, "completion_tokens", None),
+            total_tokens=getattr(response, "total_tokens", None),
+        )
+
+
+def calculate_token_efficiency(usage: TokenUsage) -> Optional[float]:
+    """Calculate token efficiency ratio (completion/prompt).
+
+    Returns:
+        Efficiency ratio or None if cannot be calculated
+    """
+    if usage.prompt_tokens and usage.completion_tokens and usage.prompt_tokens > 0:
+        return usage.completion_tokens / usage.prompt_tokens
+    return None
+
+
+def calculate_cache_efficiency(usage: TokenUsage) -> Optional[float]:
+    """Calculate cache efficiency ratio (cached/total prompt).
+
+    Returns:
+        Cache ratio or None if cannot be calculated
+    """
+    if usage.prompt_tokens and usage.cached_prompt_tokens and usage.prompt_tokens > 0:
+        return usage.cached_prompt_tokens / usage.prompt_tokens
+    return None
+
+
+def set_token_usage_attributes(span: Any, response: Any):
+    """Extract and set token usage attributes on a span.
+
+    Args:
+        span: The span to set attributes on
+        response: The response object to extract usage from
+    """
+    usage = TokenUsageExtractor.extract_from_response(response)
+
+    # Set basic token attributes
+    for attr_name, value in usage.to_attributes().items():
+        span.set_attribute(attr_name, value)
+
+    # Calculate and set efficiency metrics
+    efficiency = calculate_token_efficiency(usage)
+    if efficiency is not None:
+        span.set_attribute("llm.token_efficiency", f"{efficiency:.4f}")
+
+    cache_efficiency = calculate_cache_efficiency(usage)
+    if cache_efficiency is not None:
+        span.set_attribute("llm.cache_efficiency", f"{cache_efficiency:.4f}")

agentops/instrumentation/common/version.py

@@ -0,0 +1,71 @@
+"""Version utilities for AgentOps instrumentation.
+
+This module provides common functionality for retrieving and managing
+library versions across all instrumentation modules.
+"""
+
+import logging
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+def get_library_version(package_name: str, default_version: str = "unknown") -> str:
+    """Get the version of a library package.
+
+    Attempts to retrieve the installed version of a package using importlib.metadata.
+    Falls back to the default version if the version cannot be determined.
+
+    Args:
+        package_name: The name of the package to get the version for (as used in pip/importlib.metadata)
+        default_version: The default version to return if the package version cannot be found
+
+    Returns:
+        The version string of the package or the default version
+
+    Examples:
+        >>> get_library_version("openai")
+        "1.0.0"
+
+        >>> get_library_version("nonexistent-package")
+        "unknown"
+
+        >>> get_library_version("ibm-watsonx-ai", "1.3.11")
+        "1.3.11"  # If not found
+    """
+    try:
+        from importlib.metadata import version
+
+        return version(package_name)
+    except (ImportError, Exception) as e:
+        logger.debug(f"Could not find {package_name} version: {e}")
+        return default_version
+
+
+class LibraryInfo:
+    """Container for library information used in instrumentation.
+
+    This class provides a standardized way to store and access library
+    information (name and version) across all instrumentors.
+
+    Attributes:
+        name: The library name used for identification
+        version: The library version string
+        package_name: The package name used in pip/importlib.metadata (optional)
+    """
+
+    def __init__(self, name: str, package_name: Optional[str] = None, default_version: str = "unknown"):
+        """Initialize library information.
+
+        Args:
+            name: The library name used for identification
+            package_name: The package name used in pip/importlib.metadata.
+                         If not provided, uses the library name.
+            default_version: Default version if package version cannot be determined
+        """
+        self.name = name
+        self.package_name = package_name or name
+        self.version = get_library_version(self.package_name, default_version)
+
+    def __repr__(self) -> str:
+        return f"LibraryInfo(name={self.name!r}, version={self.version!r})"