mseep-agentops 0.4.18__py3-none-any.whl

agentops/instrumentation/common/attributes.py

@@ -0,0 +1,278 @@
+"""Common attribute processing utilities shared across all instrumentors.
+
+This utility ensures consistent attribute extraction and transformation across different
+instrumentation use cases.
+
+This module provides core utilities for extracting and formatting
+OpenTelemetry-compatible attributes from span data. These functions
+are provider-agnostic and used by all instrumentors in the AgentOps
+package.
+
+The module includes:
+
+1. Helper functions for attribute extraction and mapping
+2. Common attribute getters used across all providers
+3. Base trace and span attribute functions
+
+All functions follow a consistent pattern:
+- Accept span/trace data as input
+- Process according to semantic conventions
+- Return a dictionary of formatted attributes
+
+These utilities ensure consistent attribute handling across different
+LLM service instrumentors while maintaining separation of concerns.
+"""
+
+from typing import runtime_checkable, Protocol, Any, Optional, Dict, TypedDict
+from agentops.logging import logger
+from agentops.helpers import safe_serialize, get_agentops_version
+from agentops.semconv import (
+    CoreAttributes,
+    InstrumentationAttributes,
+    WorkflowAttributes,
+)
+
+
+# `AttributeMap` is a dictionary that maps target attribute keys to source attribute keys.
+# It is used to extract and transform attributes from a span or trace data object
+# into a standardized format following OpenTelemetry semantic conventions.
+#
+# Key-Value Format:
+# - Key (str): The target attribute key in the standardized output format
+# - Value (str): The source attribute key in the input data object
+#
+# Example Usage:
+# --------------
+#
+# Create your mapping:
+#     attribute_mapping: AttributeMap = {
+#         CoreAttributes.TRACE_ID: "trace_id",
+#         CoreAttributes.SPAN_ID: "span_id"
+#     }
+#
+# Extract the attributes:
+#     span_data = {
+#         "trace_id": "12345",
+#         "span_id": "67890",
+#     }
+#
+#     attributes = _extract_attributes_from_mapping(span_data, attribute_mapping)
+#     # >> {"trace.id": "12345", "span.id": "67890"}
+AttributeMap = Dict[str, str]  # target_attribute_key: source_attribute
+
+
+# `IndexedAttributeMap` differs from `AttributeMap` in that it allows for dynamic formatting of
+# target attribute keys using indices `i` and optionally `j`. This is particularly useful
+# when dealing with collections of similar attributes that should be uniquely identified
+# in the output.
+#
+# Key-Value Format:
+# - Key (IndexedAttribute): An object implementing the IndexedAttribute protocol with a format method
+# - Value (str): The source attribute key in the input data object
+#
+# Example Usage:
+# --------------
+#
+# Create your mapping:
+#     attribute_mapping: IndexedAttributeMap = {
+#         MessageAttributes.TOOL_CALL_ID: "id",
+#         MessageAttributes.TOOL_CALL_TYPE: "type"
+#     }
+#
+# Process tool calls:
+#     span_data = {
+#         "id": "tool_1",
+#         "type": "search",
+#     }
+#
+#     attributes = _extract_attributes_from_mapping_with_index(
+#         span_data, attribute_mapping, i=0)
+#     # >> {"gen_ai.request.tools.0.id": "tool_1", "gen_ai.request.tools.0.type": "search"}
+
+
+@runtime_checkable
+class IndexedAttribute(Protocol):
+    """
+    Protocol for objects that define a method to format indexed attributes using
+    only the provided indices `i` and optionally `j`. This allows for dynamic
+    formatting of attribute keys based on the indices.
+    """
+
+    def format(self, *, i: int, j: Optional[int] = None) -> str:
+        ...
+
+
+IndexedAttributeMap = Dict[IndexedAttribute, str]  # target_attribute_key: source_attribute
+
+
+class IndexedAttributeData(TypedDict, total=False):
+    """
+    Represents a dictionary structure for indexed attribute data.
+
+    Attributes:
+        i (int): The primary index value. This field is required.
+        j (Optional[int]): An optional secondary index value.
+    """
+
+    i: int
+    j: Optional[int] = None
+
+
+def _extract_attributes_from_mapping(span_data: Any, attribute_mapping: AttributeMap) -> AttributeMap:
+    """Helper function to extract attributes based on a mapping.
+
+    Args:
+        span_data: The span data object or dict to extract attributes from
+        attribute_mapping: Dictionary mapping target attributes to source attributes
+
+    Returns:
+        Dictionary of extracted attributes
+    """
+    attributes = {}
+    for target_attr, source_attr in attribute_mapping.items():
+        if hasattr(span_data, source_attr):
+            # Use getattr to handle properties
+            value = getattr(span_data, source_attr)
+        elif isinstance(span_data, dict) and source_attr in span_data:
+            # Use direct key access for dicts
+            value = span_data[source_attr]
+        else:
+            continue
+
+        # Skip if value is None or empty
+        if value is None or (isinstance(value, (list, dict, str)) and not value):
+            continue
+
+        # Serialize complex objects
+        elif isinstance(value, (dict, list, object)) and not isinstance(value, (str, int, float, bool)):
+            value = safe_serialize(value)
+
+        attributes[target_attr] = value
+
+    return attributes
+
+
+def _extract_attributes_from_mapping_with_index(
+    span_data: Any, attribute_mapping: IndexedAttributeMap, i: int, j: Optional[int] = None
+) -> AttributeMap:
+    """Helper function to extract attributes based on a mapping with index.
+
+    This function extends `_extract_attributes_from_mapping` by allowing for indexed keys in the attribute mapping.
+
+    Span data is expected to have keys which contain format strings for i/j, e.g. `my_attr_{i}` or `my_attr_{i}_{j}`.
+
+    Args:
+        span_data: The span data object or dict to extract attributes from
+        attribute_mapping: Dictionary mapping target attributes to source attributes, with format strings for i/j
+        i: The primary index to use in formatting the attribute keys
+        j: An optional secondary index (default is None)
+    Returns:
+        Dictionary of extracted attributes with formatted indexed keys.
+    """
+
+    # `i` is required for formatting the attribute keys, `j` is optional
+    format_kwargs: IndexedAttributeData = {"i": i}
+    if j is not None:
+        format_kwargs["j"] = j
+
+    # Update the attribute mapping to include the index for the span
+    attribute_mapping_with_index: AttributeMap = {}
+    for target_attr, source_attr in attribute_mapping.items():
+        attribute_mapping_with_index[target_attr.format(**format_kwargs)] = source_attr
+
+    return _extract_attributes_from_mapping(span_data, attribute_mapping_with_index)
+
+
+def get_common_attributes() -> AttributeMap:
+    """Get common instrumentation attributes used across traces and spans.
+
+    Returns:
+        Dictionary of common instrumentation attributes
+    """
+    return {
+        InstrumentationAttributes.NAME: "agentops",
+        InstrumentationAttributes.VERSION: get_agentops_version(),
+    }
+
+
+def get_base_trace_attributes(trace: Any) -> AttributeMap:
+    """Create the base attributes dictionary for an OpenTelemetry trace.
+
+    Args:
+        trace: The trace object to extract attributes from
+
+    Returns:
+        Dictionary containing base trace attributes
+    """
+    if not hasattr(trace, "trace_id"):
+        logger.warning("Cannot create trace attributes: missing trace_id")
+        return {}
+
+    attributes = {
+        WorkflowAttributes.WORKFLOW_NAME: trace.name,
+        CoreAttributes.TRACE_ID: trace.trace_id,
+        WorkflowAttributes.WORKFLOW_STEP_TYPE: "trace",
+        **get_common_attributes(),
+    }
+
+    # Add tags from the config to the trace attributes (these should only be added to the trace)
+    from agentops import get_client
+
+    config = get_client().config
+    tags = []
+    if config.default_tags:
+        # `default_tags` can either be a `set` or a `list`
+        tags = list(config.default_tags)
+
+    attributes[CoreAttributes.TAGS] = tags
+
+    return attributes
+
+
+def get_base_span_attributes(span: Any) -> AttributeMap:
+    """Create the base attributes dictionary for an OpenTelemetry span.
+
+    Args:
+        span: The span object to extract attributes from
+
+    Returns:
+        Dictionary containing base span attributes
+    """
+    span_id = getattr(span, "span_id", "unknown")
+    trace_id = getattr(span, "trace_id", "unknown")
+    parent_id = getattr(span, "parent_id", None)
+
+    attributes = {
+        CoreAttributes.TRACE_ID: trace_id,
+        CoreAttributes.SPAN_ID: span_id,
+        **get_common_attributes(),
+    }
+
+    if parent_id:
+        attributes[CoreAttributes.PARENT_ID] = parent_id
+
+    return attributes
+
+
+def extract_token_usage(response: Any) -> Dict[str, int]:
+    """Extract token usage information from a response.
+
+    Args:
+        response: The response object to extract token usage from
+
+    Returns:
+        Dictionary containing token usage information
+    """
+    usage = {}
+
+    # Try to extract token counts from response
+    if hasattr(response, "usage"):
+        usage_data = response.usage
+        if hasattr(usage_data, "prompt_tokens"):
+            usage["prompt_tokens"] = usage_data.prompt_tokens
+        if hasattr(usage_data, "completion_tokens"):
+            usage["completion_tokens"] = usage_data.completion_tokens
+        if hasattr(usage_data, "total_tokens"):
+            usage["total_tokens"] = usage_data.total_tokens
+
+    return usage

agentops/instrumentation/common/instrumentor.py

@@ -0,0 +1,147 @@
+"""Base instrumentor utilities for AgentOps instrumentation.
+
+This module provides base classes and utilities for creating instrumentors,
+reducing boilerplate code across different provider instrumentations.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Collection, Dict, List, Optional, Any, Callable
+from dataclasses import dataclass, field
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import Tracer, get_tracer
+from opentelemetry.metrics import Meter, get_meter
+
+from agentops.instrumentation.common.wrappers import WrapConfig, wrap, unwrap
+from agentops.logging import logger
+
+
+@dataclass
+class InstrumentorConfig:
+    """Configuration for an instrumentor."""
+
+    library_name: str
+    library_version: str
+    wrapped_methods: List[WrapConfig] = field(default_factory=list)
+    metrics_enabled: bool = True
+    dependencies: Collection[str] = field(default_factory=list)
+
+
+class CommonInstrumentor(BaseInstrumentor, ABC):
+    """Base class for AgentOps instrumentors with common functionality."""
+
+    def __init__(self, config: InstrumentorConfig):
+        super().__init__()
+        self.config = config
+        self._tracer: Optional[Tracer] = None
+        self._meter: Optional[Meter] = None
+        self._metrics: Dict[str, Any] = {}
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        """Return required dependencies."""
+        return self.config.dependencies
+
+    def _instrument(self, **kwargs):
+        """Instrument the target library."""
+        # Initialize tracer
+        tracer_provider = kwargs.get("tracer_provider")
+        self._tracer = get_tracer(self.config.library_name, self.config.library_version, tracer_provider)
+
+        # Initialize meter if metrics enabled
+        if self.config.metrics_enabled:
+            meter_provider = kwargs.get("meter_provider")
+            self._meter = get_meter(self.config.library_name, self.config.library_version, meter_provider)
+            self._metrics = self._create_metrics(self._meter)
+
+        # Perform custom initialization
+        self._initialize(**kwargs)
+
+        # Wrap all configured methods
+        self._wrap_methods()
+
+        # Perform custom wrapping
+        self._custom_wrap(**kwargs)
+
+    def _uninstrument(self, **kwargs):
+        """Remove instrumentation."""
+        # Unwrap all configured methods
+        for wrap_config in self.config.wrapped_methods:
+            try:
+                unwrap(wrap_config)
+            except Exception as e:
+                logger.debug(
+                    f"Failed to unwrap {wrap_config.package}.{wrap_config.class_name}.{wrap_config.method_name}: {e}"
+                )
+
+        # Perform custom unwrapping
+        self._custom_unwrap(**kwargs)
+
+        # Clear references
+        self._tracer = None
+        self._meter = None
+        self._metrics.clear()
+
+    def _wrap_methods(self):
+        """Wrap all configured methods."""
+        for wrap_config in self.config.wrapped_methods:
+            try:
+                wrap(wrap_config, self._tracer)
+            except (AttributeError, ModuleNotFoundError) as e:
+                logger.debug(
+                    f"Could not wrap {wrap_config.package}.{wrap_config.class_name}.{wrap_config.method_name}: {e}"
+                )
+
+    @abstractmethod
+    def _create_metrics(self, meter: Meter) -> Dict[str, Any]:
+        """Create metrics for the instrumentor.
+
+        Returns a dictionary of metric name to metric instance.
+        """
+        pass
+
+    def _initialize(self, **kwargs):
+        """Perform custom initialization.
+
+        Override in subclasses for custom initialization logic.
+        """
+        pass
+
+    def _custom_wrap(self, **kwargs):
+        """Perform custom wrapping beyond configured methods.
+
+        Override in subclasses for special wrapping needs.
+        """
+        pass
+
+    def _custom_unwrap(self, **kwargs):
+        """Perform custom unwrapping beyond configured methods.
+
+        Override in subclasses for special unwrapping needs.
+        """
+        pass
+
+
+def create_wrapper_factory(wrapper_func: Callable, *wrapper_args, **wrapper_kwargs) -> Callable:
+    """Create a factory function for wrapt-style wrappers.
+
+    This is useful for creating wrappers that need additional arguments
+    beyond the standard (wrapped, instance, args, kwargs).
+
+    Args:
+        wrapper_func: The wrapper function to call
+        *wrapper_args: Arguments to pass to the wrapper
+        **wrapper_kwargs: Keyword arguments to pass to the wrapper
+
+    Returns:
+        A factory function that returns the configured wrapper
+    """
+
+    def factory(tracer: Tracer):
+        def wrapper(wrapped, instance, args, kwargs):
+            return wrapper_func(
+                tracer, *wrapper_args, wrapped=wrapped, instance=instance, args=args, kwargs=kwargs, **wrapper_kwargs
+            )
+
+        return wrapper
+
+    return factory

agentops/instrumentation/common/metrics.py

@@ -0,0 +1,100 @@
+"""Common metrics utilities for AgentOps instrumentation.
+
+This module provides utilities for creating and managing standard metrics
+across different instrumentations.
+"""
+
+from typing import Dict, Any, Optional
+from opentelemetry.metrics import Meter, Histogram, Counter
+from agentops.semconv import Meters
+
+
+class StandardMetrics:
+    """Factory for creating standard metrics used across instrumentations."""
+
+    @staticmethod
+    def create_token_histogram(meter: Meter) -> Histogram:
+        """Create a histogram for token usage."""
+        return meter.create_histogram(
+            name=Meters.LLM_TOKEN_USAGE, unit="token", description="Measures number of input and output tokens used"
+        )
+
+    @staticmethod
+    def create_duration_histogram(meter: Meter) -> Histogram:
+        """Create a histogram for operation duration."""
+        return meter.create_histogram(
+            name=Meters.LLM_OPERATION_DURATION, unit="s", description="GenAI operation duration"
+        )
+
+    @staticmethod
+    def create_exception_counter(meter: Meter, name: str = Meters.LLM_COMPLETIONS_EXCEPTIONS) -> Counter:
+        """Create a counter for exceptions."""
+        return meter.create_counter(
+            name=name, unit="time", description="Number of exceptions occurred during operations"
+        )
+
+    @staticmethod
+    def create_choice_counter(meter: Meter) -> Counter:
+        """Create a counter for generation choices."""
+        return meter.create_counter(
+            name=Meters.LLM_GENERATION_CHOICES,
+            unit="choice",
+            description="Number of choices returned by completions call",
+        )
+
+    @staticmethod
+    def create_standard_metrics(meter: Meter) -> Dict[str, Any]:
+        """Create a standard set of metrics for LLM operations.
+
+        Returns:
+            Dictionary with metric names as keys and metric instances as values
+        """
+        return {
+            "token_histogram": StandardMetrics.create_token_histogram(meter),
+            "duration_histogram": StandardMetrics.create_duration_histogram(meter),
+            "exception_counter": StandardMetrics.create_exception_counter(meter),
+        }
+
+
+class MetricsRecorder:
+    """Utility class for recording metrics in a consistent way."""
+
+    def __init__(self, metrics: Dict[str, Any]):
+        self.metrics = metrics
+
+    def record_token_usage(
+        self,
+        prompt_tokens: Optional[int] = None,
+        completion_tokens: Optional[int] = None,
+        attributes: Optional[Dict[str, Any]] = None,
+    ):
+        """Record token usage metrics."""
+        token_histogram = self.metrics.get("token_histogram")
+        if not token_histogram:
+            return
+
+        attrs = attributes or {}
+
+        if prompt_tokens is not None:
+            token_histogram.record(prompt_tokens, attributes={**attrs, "token.type": "input"})
+
+        if completion_tokens is not None:
+            token_histogram.record(completion_tokens, attributes={**attrs, "token.type": "output"})
+
+    def record_duration(self, duration: float, attributes: Optional[Dict[str, Any]] = None):
+        """Record operation duration."""
+        duration_histogram = self.metrics.get("duration_histogram")
+        if duration_histogram:
+            duration_histogram.record(duration, attributes=attributes or {})
+
+    def record_exception(self, attributes: Optional[Dict[str, Any]] = None):
+        """Record an exception occurrence."""
+        exception_counter = self.metrics.get("exception_counter")
+        if exception_counter:
+            exception_counter.add(1, attributes=attributes or {})
+
+    def record_choices(self, count: int, attributes: Optional[Dict[str, Any]] = None):
+        """Record number of choices returned."""
+        choice_counter = self.metrics.get("choice_counter")
+        if choice_counter:
+            choice_counter.add(count, attributes=attributes or {})

agentops/instrumentation/common/objects.py

@@ -0,0 +1,26 @@
+from agentops.client.api.types import UploadedObjectResponse
+from agentops.instrumentation.common import AttributeMap, _extract_attributes_from_mapping
+
+
+UPLOADED_OBJECT_ATTRIBUTES: AttributeMap = {
+    "object_url": "url",
+    "object_size": "size",
+}
+
+
+def get_uploaded_object_attributes(uploaded_object: UploadedObjectResponse, prefix: str) -> AttributeMap:
+    """Extract attributes from an uploaded object.
+
+    This is a common function so we can standardize the data format we serialize
+    stored objects to.
+
+    Args:
+        uploaded_object: The uploaded object to extract attributes from.
+        prefix: The prefix to use for the attribute keys. Keys will be concatenated
+            with the prefix and a dot (.) separator.
+
+    Returns:
+        A dictionary of extracted attributes.
+    """
+    attribute_map = {f"{prefix}.{key}": value for key, value in UPLOADED_OBJECT_ATTRIBUTES.items()}
+    return _extract_attributes_from_mapping(uploaded_object, attribute_map)