mseep-agentops 0.4.18__py3-none-any.whl → 0.4.22__py3-none-any.whl

agentops/instrumentation/providers/anthropic/stream_wrapper.py

@@ -0,0 +1,436 @@
+"""Anthropic stream wrapper implementation.
+
+This module provides wrappers for Anthropic's streaming functionality,
+focusing on the MessageStreamManager for both sync and async operations.
+It instruments streams to collect telemetry data for monitoring and analysis.
+"""
+
+import logging
+from typing import TypeVar
+
+from opentelemetry import context as context_api
+from opentelemetry.trace import SpanKind
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+
+from agentops.semconv import SpanAttributes, LLMRequestTypeValues, CoreAttributes, MessageAttributes
+from agentops.instrumentation.common.wrappers import _with_tracer_wrapper
+from agentops.instrumentation.providers.anthropic.attributes.message import (
+    get_message_request_attributes,
+    get_stream_attributes,
+)
+from agentops.instrumentation.providers.anthropic.event_handler_wrapper import EventHandleWrapper
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+@_with_tracer_wrapper
+def messages_stream_wrapper(tracer, wrapped, instance, args, kwargs):
+    """Wrapper for the Messages.stream method.
+
+    This wrapper creates spans for tracking stream performance and injects
+    an event handler wrapper to capture streaming events.
+
+    Args:
+        tracer: The OpenTelemetry tracer to use
+        wrapped: The original stream method
+        instance: The instance the method is bound to
+        args: Positional arguments to the method
+        kwargs: Keyword arguments to the method
+
+    Returns:
+        A wrapped stream manager that captures telemetry data
+    """
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    span = tracer.start_span(
+        "anthropic.messages.stream",
+        kind=SpanKind.CLIENT,
+        attributes={SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.CHAT.value},
+    )
+
+    request_attributes = get_message_request_attributes(kwargs)
+    for key, value in request_attributes.items():
+        span.set_attribute(key, value)
+
+    span.set_attribute(SpanAttributes.LLM_REQUEST_STREAMING, True)
+
+    original_event_handler = kwargs.get("event_handler")
+
+    if original_event_handler is not None:
+        wrapped_handler = EventHandleWrapper(original_handler=original_event_handler, span=span)
+        kwargs["event_handler"] = wrapped_handler
+
+    try:
+
+        class TracedStreamManager:
+            """A wrapper for Anthropic's MessageStreamManager that adds telemetry.
+
+            This class wraps the original stream manager to capture metrics about
+            the streaming process, including token counts, content, and errors.
+            """
+
+            def __init__(self, original_manager):
+                """Initialize with the original manager.
+
+                Args:
+                    original_manager: The Anthropic MessageStreamManager to wrap
+                """
+                self.original_manager = original_manager
+                self.stream = None
+
+            def __enter__(self):
+                """Context manager entry that initializes stream monitoring.
+
+                Returns:
+                    The original stream with instrumentation added
+                """
+                self.stream = self.original_manager.__enter__()
+
+                try:
+                    stream_attributes = get_stream_attributes(self.stream)
+                    for key, value in stream_attributes.items():
+                        span.set_attribute(key, value)
+                except Exception as e:
+                    logger.debug(f"Error getting stream attributes: {e}")
+
+                # Set the event handler on the stream if provided
+                if original_event_handler is not None:
+                    self.stream.event_handler = kwargs["event_handler"]
+                else:
+                    try:
+                        original_text_stream = self.stream.text_stream
+                        token_count = 0
+
+                        class InstrumentedTextStream:
+                            """A wrapper for Anthropic's text stream that counts tokens."""
+
+                            def __iter__(self):
+                                """Iterate through text chunks, counting tokens.
+
+                                Yields:
+                                    Text chunks from the original stream
+                                """
+                                nonlocal token_count
+                                for text in original_text_stream:
+                                    token_count += len(text.split())
+                                    span.set_attribute(SpanAttributes.LLM_USAGE_STREAMING_TOKENS, token_count)
+                                    yield text
+
+                        self.stream.text_stream = InstrumentedTextStream()
+                    except Exception as e:
+                        logger.debug(f"Error patching text_stream: {e}")
+
+                return self.stream
+
+            def __exit__(self, exc_type, exc_val, exc_tb):
+                """Context manager exit that records final metrics.
+
+                Args:
+                    exc_type: Exception type, if an exception occurred
+                    exc_val: Exception value, if an exception occurred
+                    exc_tb: Exception traceback, if an exception occurred
+
+                Returns:
+                    Result of the original context manager's __exit__
+                """
+                try:
+                    if exc_type is not None:
+                        span.record_exception(exc_val)
+                        span.set_attribute(CoreAttributes.ERROR_MESSAGE, str(exc_val))
+                        span.set_attribute(CoreAttributes.ERROR_TYPE, exc_type.__name__)
+
+                    try:
+                        final_message = None
+
+                        if hasattr(self.original_manager, "_MessageStreamManager__stream") and hasattr(
+                            self.original_manager._MessageStreamManager__stream,
+                            "_MessageStream__final_message_snapshot",
+                        ):
+                            final_message = self.original_manager._MessageStreamManager__stream._MessageStream__final_message_snapshot
+
+                        if final_message:
+                            if hasattr(final_message, "content"):
+                                content_text = ""
+                                if isinstance(final_message.content, list):
+                                    for content_block in final_message.content:
+                                        if hasattr(content_block, "text"):
+                                            content_text += content_block.text
+
+                                if content_text:
+                                    span.set_attribute(MessageAttributes.COMPLETION_TYPE.format(i=0), "text")
+                                    span.set_attribute(MessageAttributes.COMPLETION_ROLE.format(i=0), "assistant")
+                                    span.set_attribute(MessageAttributes.COMPLETION_CONTENT.format(i=0), content_text)
+
+                            if hasattr(final_message, "usage"):
+                                usage = final_message.usage
+                                if hasattr(usage, "input_tokens"):
+                                    span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, usage.input_tokens)
+
+                                if hasattr(usage, "output_tokens"):
+                                    span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, usage.output_tokens)
+
+                                if hasattr(usage, "input_tokens") and hasattr(usage, "output_tokens"):
+                                    total_tokens = usage.input_tokens + usage.output_tokens
+                                    span.set_attribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, total_tokens)
+                    except Exception as e:
+                        logger.debug(f"Failed to extract final message data: {e}")
+                finally:
+                    if span.is_recording():
+                        span.end()
+                    return self.original_manager.__exit__(exc_type, exc_val, exc_tb)
+
+        stream_manager = wrapped(*args, **kwargs)
+
+        return TracedStreamManager(stream_manager)
+
+    except Exception as e:
+        span.record_exception(e)
+        span.set_attribute(CoreAttributes.ERROR_MESSAGE, str(e))
+        span.set_attribute(CoreAttributes.ERROR_TYPE, e.__class__.__name__)
+        span.end()
+        raise
+
+
+class AsyncStreamContextManagerWrapper:
+    """A wrapper that implements both async context manager and awaitable protocols.
+
+    This wrapper allows the instrumented async stream to be used either with
+    'async with' or by awaiting it first, preserving compatibility with
+    different usage patterns.
+    """
+
+    def __init__(self, coro):
+        """Initialize with a coroutine.
+
+        Args:
+            coro: The coroutine that will return a stream manager
+        """
+        self._coro = coro
+        self._stream_manager = None
+
+    def __await__(self):
+        """Make this wrapper awaitable.
+
+        This allows users to do:
+        stream_manager = await client.messages.stream(...)
+
+        Returns:
+            An awaitable that yields the traced stream manager
+        """
+
+        async def get_stream_manager():
+            self._stream_manager = await self._coro
+            return self._stream_manager
+
+        return get_stream_manager().__await__()
+
+    async def __aenter__(self):
+        """Async context manager enter.
+
+        This allows users to do:
+        async with client.messages.stream(...) as stream:
+
+        Returns:
+            The result of the stream manager's __aenter__
+        """
+        if self._stream_manager is None:
+            self._stream_manager = await self._coro
+
+        return await self._stream_manager.__aenter__()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit.
+
+        Args:
+            exc_type: Exception type
+            exc_val: Exception value
+            exc_tb: Exception traceback
+
+        Returns:
+            The result of the stream manager's __aexit__
+        """
+        if self._stream_manager is not None:
+            return await self._stream_manager.__aexit__(exc_type, exc_val, exc_tb)
+        return False
+
+
+@_with_tracer_wrapper
+def messages_stream_async_wrapper(tracer, wrapped, instance, args, kwargs):
+    """Wrapper for the async Messages.stream method.
+
+    This wrapper creates spans for tracking stream performance and injects
+    an event handler wrapper to capture streaming events in async contexts.
+
+    Args:
+        tracer: The OpenTelemetry tracer to use
+        wrapped: The original async stream method
+        instance: The instance the method is bound to
+        args: Positional arguments to the method
+        kwargs: Keyword arguments to the method
+
+    Returns:
+        An object that can be used with async with or awaited
+    """
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    span = tracer.start_span(
+        "anthropic.messages.stream",
+        kind=SpanKind.CLIENT,
+        attributes={SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.CHAT.value},
+    )
+
+    request_attributes = get_message_request_attributes(kwargs)
+    for key, value in request_attributes.items():
+        span.set_attribute(key, value)
+
+    span.set_attribute(SpanAttributes.LLM_REQUEST_STREAMING, True)
+
+    original_event_handler = kwargs.get("event_handler")
+
+    if original_event_handler is not None:
+        wrapped_handler = EventHandleWrapper(original_handler=original_event_handler, span=span)
+        kwargs["event_handler"] = wrapped_handler
+
+    async def _wrapped_stream():
+        """Async wrapper function for the stream method.
+
+        Returns:
+            A traced async stream manager
+        """
+        try:
+            # Don't await wrapped(*args, **kwargs) - it returns an async context manager, not a coroutine
+            stream_manager = wrapped(*args, **kwargs)
+
+            class TracedAsyncStreamManager:
+                """A wrapper for Anthropic's AsyncMessageStreamManager that adds telemetry.
+
+                This class wraps the original async stream manager to capture metrics
+                about the streaming process, including token counts, content, and errors.
+                """
+
+                def __init__(self, original_manager):
+                    """Initialize with the original manager.
+
+                    Args:
+                        original_manager: The Anthropic AsyncMessageStreamManager to wrap
+                    """
+                    self.original_manager = original_manager
+                    self.stream = None
+
+                async def __aenter__(self):
+                    """Async context manager entry that initializes stream monitoring.
+
+                    Returns:
+                        The original stream with instrumentation added
+                    """
+                    self.stream = await self.original_manager.__aenter__()
+
+                    try:
+                        stream_attributes = get_stream_attributes(self.stream)
+                        for key, value in stream_attributes.items():
+                            span.set_attribute(key, value)
+                    except Exception as e:
+                        logger.debug(f"Error getting async stream attributes: {e}")
+
+                    if original_event_handler is None:
+                        try:
+                            original_text_stream = self.stream.text_stream
+                            token_count = 0
+
+                            class InstrumentedAsyncTextStream:
+                                """A wrapper for Anthropic's async text stream that counts tokens."""
+
+                                async def __aiter__(self):
+                                    """Async iterate through text chunks, counting tokens.
+
+                                    Yields:
+                                        Text chunks from the original async stream
+                                    """
+                                    nonlocal token_count
+                                    async for text in original_text_stream:
+                                        token_count += len(text.split())
+                                        span.set_attribute(SpanAttributes.LLM_USAGE_STREAMING_TOKENS, token_count)
+                                        yield text
+
+                            self.stream.text_stream = InstrumentedAsyncTextStream()
+                        except Exception as e:
+                            logger.debug(f"Error patching async text_stream: {e}")
+
+                    return self.stream
+
+                async def __aexit__(self, exc_type, exc_val, exc_tb):
+                    """Async context manager exit that records final metrics.
+
+                    Args:
+                        exc_type: Exception type, if an exception occurred
+                        exc_val: Exception value, if an exception occurred
+                        exc_tb: Exception traceback, if an exception occurred
+
+                    Returns:
+                        Result of the original async context manager's __aexit__
+                    """
+                    try:
+                        if exc_type is not None:
+                            span.record_exception(exc_val)
+                            span.set_attribute(CoreAttributes.ERROR_MESSAGE, str(exc_val))
+                            span.set_attribute(CoreAttributes.ERROR_TYPE, exc_type.__name__)
+
+                        try:
+                            final_message = None
+
+                            if hasattr(self.original_manager, "_AsyncMessageStreamManager__stream") and hasattr(
+                                self.original_manager._AsyncMessageStreamManager__stream,
+                                "_AsyncMessageStream__final_message_snapshot",
+                            ):
+                                final_message = self.original_manager._AsyncMessageStreamManager__stream._AsyncMessageStream__final_message_snapshot
+
+                            if final_message:
+                                if hasattr(final_message, "content"):
+                                    content_text = ""
+                                    if isinstance(final_message.content, list):
+                                        for content_block in final_message.content:
+                                            if hasattr(content_block, "text"):
+                                                content_text += content_block.text
+
+                                    if content_text:
+                                        span.set_attribute(MessageAttributes.COMPLETION_TYPE.format(i=0), "text")
+                                        span.set_attribute(MessageAttributes.COMPLETION_ROLE.format(i=0), "assistant")
+                                        span.set_attribute(
+                                            MessageAttributes.COMPLETION_CONTENT.format(i=0), content_text
+                                        )
+
+                                if hasattr(final_message, "usage"):
+                                    usage = final_message.usage
+                                    if hasattr(usage, "input_tokens"):
+                                        span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, usage.input_tokens)
+
+                                    if hasattr(usage, "output_tokens"):
+                                        span.set_attribute(
+                                            SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, usage.output_tokens
+                                        )
+
+                                    if hasattr(usage, "input_tokens") and hasattr(usage, "output_tokens"):
+                                        total_tokens = usage.input_tokens + usage.output_tokens
+                                        span.set_attribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, total_tokens)
+                        except Exception as e:
+                            logger.debug(f"Failed to extract final async message data: {e}")
+                    finally:
+                        if span.is_recording():
+                            span.end()
+                        return await self.original_manager.__aexit__(exc_type, exc_val, exc_tb)
+
+            return TracedAsyncStreamManager(stream_manager)
+
+        except Exception as e:
+            span.record_exception(e)
+            span.set_attribute(CoreAttributes.ERROR_MESSAGE, str(e))
+            span.set_attribute(CoreAttributes.ERROR_TYPE, e.__class__.__name__)
+            span.end()
+            raise
+
+    # Return a wrapper that implements both async context manager and awaitable protocols
+    return AsyncStreamContextManagerWrapper(_wrapped_stream())

agentops/instrumentation/providers/google_genai/README.md

@@ -0,0 +1,33 @@
+# Google Generative AI (Gemini) Instrumentation
+
+This module provides OpenTelemetry instrumentation for Google's Generative AI (Gemini) API. The instrumentation allows you to trace all API calls made using the `google-genai` Python SDK, capturing:
+
+- Model parameters (temperature, max_tokens, etc.)
+- Prompt content (with privacy controls)
+- Response text and token usage
+- Streaming metrics
+- Token counting
+- Performance and error data
+
+## Supported Features
+
+The instrumentation covers all major API methods including:
+
+### Client-Based API
+- `client.models.generate_content`
+- `client.models.generate_content_stream`
+- `client.models.count_tokens`
+- `client.models.compute_tokens`
+- And their corresponding async variants
+
+## Metrics
+
+The instrumentation captures the following metrics:
+
+- Input tokens used
+- Output tokens generated
+- Total tokens consumed
+- Operation duration
+- Exception counts
+
+These metrics are available as OpenTelemetry span attributes and can be viewed in your observability platform of choice when properly configured.

agentops/instrumentation/providers/google_genai/__init__.py

@@ -0,0 +1,24 @@
+"""Google Generative AI (Gemini) API instrumentation.
+
+This module provides instrumentation for the Google Generative AI (Gemini) API,
+including content generation, streaming, and chat functionality.
+"""
+
+import logging
+from agentops.instrumentation.common import LibraryInfo
+
+logger = logging.getLogger(__name__)
+
+# Library information
+_library_info = LibraryInfo(name="google-genai")
+LIBRARY_NAME = _library_info.name
+LIBRARY_VERSION = _library_info.version
+
+# Import after defining constants to avoid circular imports
+from agentops.instrumentation.providers.google_genai.instrumentor import GoogleGenaiInstrumentor  # noqa: E402
+
+__all__ = [
+    "LIBRARY_NAME",
+    "LIBRARY_VERSION",
+    "GoogleGenaiInstrumentor",
+]

agentops/instrumentation/providers/google_genai/attributes/__init__.py

@@ -0,0 +1,25 @@
+"""Attribute extractors for Google Generative AI instrumentation."""
+
+from agentops.instrumentation.providers.google_genai.attributes.common import (
+    get_common_instrumentation_attributes,
+    extract_request_attributes,
+)
+from agentops.instrumentation.providers.google_genai.attributes.model import (
+    get_model_attributes,
+    get_generate_content_attributes,
+    get_stream_attributes,
+    get_token_counting_attributes,
+)
+from agentops.instrumentation.providers.google_genai.attributes.chat import (
+    get_chat_attributes,
+)
+
+__all__ = [
+    "get_common_instrumentation_attributes",
+    "extract_request_attributes",
+    "get_model_attributes",
+    "get_generate_content_attributes",
+    "get_stream_attributes",
+    "get_chat_attributes",
+    "get_token_counting_attributes",
+]

agentops/instrumentation/providers/google_genai/attributes/chat.py

@@ -0,0 +1,125 @@
+"""Chat attribute extraction for Google Generative AI instrumentation."""
+
+from typing import Dict, Any, Optional, Tuple
+
+from agentops.logging import logger
+from agentops.semconv import SpanAttributes, LLMRequestTypeValues, MessageAttributes
+from agentops.instrumentation.common.attributes import AttributeMap
+from agentops.instrumentation.providers.google_genai.attributes.common import (
+    extract_request_attributes,
+    get_common_instrumentation_attributes,
+)
+from agentops.instrumentation.providers.google_genai.attributes.model import (
+    _extract_content_from_prompt,
+    _set_response_attributes,
+)
+
+
+def _extract_message_content(message: Any) -> str:
+    """Extract text content from a chat message.
+
+    Handles the various message formats in the Gemini chat API.
+
+    Args:
+        message: The message to extract content from
+
+    Returns:
+        Extracted text as a string
+    """
+    if isinstance(message, str):
+        return message
+
+    if isinstance(message, dict):
+        if "content" in message:
+            return _extract_content_from_prompt(message["content"])
+        if "text" in message:
+            return message["text"]
+
+    if hasattr(message, "content"):
+        return _extract_content_from_prompt(message.content)
+
+    if hasattr(message, "text"):
+        return message.text
+
+    return ""
+
+
+def _set_chat_history_attributes(attributes: AttributeMap, args: Tuple, kwargs: Dict[str, Any]) -> None:
+    """Extract and set chat history attributes from the request.
+
+    Args:
+        attributes: The attribute dictionary to update
+        args: Positional arguments to the method
+        kwargs: Keyword arguments to the method
+    """
+    messages = []
+    if "message" in kwargs:
+        messages = [kwargs["message"]]
+    elif args and len(args) > 0:
+        messages = [args[0]]
+    elif "messages" in kwargs:
+        messages = kwargs["messages"]
+
+    if not messages:
+        return
+
+    for i, message in enumerate(messages):
+        try:
+            content = _extract_message_content(message)
+            if content:
+                role = "user"
+
+                if isinstance(message, dict) and "role" in message:
+                    role = message["role"]
+                elif hasattr(message, "role"):
+                    role = message.role
+
+                attributes[MessageAttributes.PROMPT_CONTENT.format(i=i)] = content
+                attributes[MessageAttributes.PROMPT_ROLE.format(i=i)] = role
+        except Exception as e:
+            logger.debug(f"Error extracting chat message at index {i}: {e}")
+
+
+def get_chat_attributes(
+    args: Optional[Tuple] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
+    return_value: Optional[Any] = None,
+) -> AttributeMap:
+    """Extract attributes for chat session methods.
+
+    This function handles attribute extraction for chat session operations,
+    particularly the send_message method.
+
+    Args:
+        args: Positional arguments to the method
+        kwargs: Keyword arguments to the method
+        return_value: Return value from the method
+
+    Returns:
+        Dictionary of extracted attributes
+    """
+    attributes = get_common_instrumentation_attributes()
+    attributes[SpanAttributes.LLM_SYSTEM] = "Gemini"
+    attributes[SpanAttributes.LLM_REQUEST_TYPE] = LLMRequestTypeValues.CHAT.value
+
+    if kwargs:
+        kwargs_attributes = extract_request_attributes(kwargs)
+        attributes.update(kwargs_attributes)
+
+    chat_session = None
+    if args and len(args) >= 1:
+        chat_session = args[0]
+
+    if chat_session and hasattr(chat_session, "model"):
+        if isinstance(chat_session.model, str):
+            attributes[SpanAttributes.LLM_REQUEST_MODEL] = chat_session.model
+        elif hasattr(chat_session.model, "name"):
+            attributes[SpanAttributes.LLM_REQUEST_MODEL] = chat_session.model.name
+
+    if args or kwargs:
+        _set_chat_history_attributes(attributes, args or (), kwargs or {})
+
+    if return_value is not None:
+        _set_response_attributes(attributes, return_value)
+
+    return attributes