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

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

@@ -0,0 +1,88 @@
+"""Common attribute extraction for Google Generative AI instrumentation."""
+
+from typing import Dict, Any
+
+from agentops.logging import logger
+from agentops.semconv import InstrumentationAttributes, SpanAttributes
+from agentops.instrumentation.common.attributes import (
+    AttributeMap,
+    get_common_attributes,
+    _extract_attributes_from_mapping,
+)
+from agentops.instrumentation.providers.google_genai import LIBRARY_NAME, LIBRARY_VERSION
+
+# Common mapping for config parameters
+REQUEST_CONFIG_ATTRIBUTES: AttributeMap = {
+    SpanAttributes.LLM_REQUEST_TEMPERATURE: "temperature",
+    SpanAttributes.LLM_REQUEST_MAX_TOKENS: "max_output_tokens",
+    SpanAttributes.LLM_REQUEST_TOP_P: "top_p",
+    SpanAttributes.LLM_REQUEST_TOP_K: "top_k",
+    SpanAttributes.LLM_REQUEST_SEED: "seed",
+    SpanAttributes.LLM_REQUEST_SYSTEM_INSTRUCTION: "system_instruction",
+    SpanAttributes.LLM_REQUEST_PRESENCE_PENALTY: "presence_penalty",
+    SpanAttributes.LLM_REQUEST_FREQUENCY_PENALTY: "frequency_penalty",
+    SpanAttributes.LLM_REQUEST_STOP_SEQUENCES: "stop_sequences",
+    SpanAttributes.LLM_REQUEST_CANDIDATE_COUNT: "candidate_count",
+}
+
+
+def get_common_instrumentation_attributes() -> AttributeMap:
+    """Get common instrumentation attributes for the Google Generative AI instrumentation.
+
+    This combines the generic AgentOps attributes with Google Generative AI specific library attributes.
+
+    Returns:
+        Dictionary of common instrumentation attributes
+    """
+    attributes = get_common_attributes()
+    attributes.update(
+        {
+            InstrumentationAttributes.LIBRARY_NAME: LIBRARY_NAME,
+            InstrumentationAttributes.LIBRARY_VERSION: LIBRARY_VERSION,
+        }
+    )
+    return attributes
+
+
+def extract_request_attributes(kwargs: Dict[str, Any]) -> AttributeMap:
+    """Extract request attributes from the function arguments.
+
+    Extracts common request parameters that apply to both content generation
+    and chat completions, focusing on model parameters and generation settings.
+
+    Args:
+        kwargs: Request keyword arguments
+
+    Returns:
+        Dictionary of extracted request attributes
+    """
+    attributes = {}
+
+    if "model" in kwargs:
+        model = kwargs["model"]
+
+        # Handle string model names
+        if isinstance(model, str):
+            attributes[SpanAttributes.LLM_REQUEST_MODEL] = model
+        # Handle model objects with _model_name or name attribute
+        elif hasattr(model, "_model_name"):
+            attributes[SpanAttributes.LLM_REQUEST_MODEL] = model._model_name
+        elif hasattr(model, "name"):
+            attributes[SpanAttributes.LLM_REQUEST_MODEL] = model.name
+
+    config = kwargs.get("config")
+
+    if config:
+        try:
+            attributes.update(
+                _extract_attributes_from_mapping(
+                    config.__dict__ if hasattr(config, "__dict__") else config, REQUEST_CONFIG_ATTRIBUTES
+                )
+            )
+        except Exception as e:
+            logger.debug(f"Error extracting config parameters: {e}")
+
+    if "stream" in kwargs:
+        attributes[SpanAttributes.LLM_REQUEST_STREAMING] = kwargs["stream"]
+
+    return attributes

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

@@ -0,0 +1,284 @@
+"""Model 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,
+)
+
+
+def _extract_content_from_prompt(content: Any) -> str:
+    """Extract prompt text from content.
+
+    Handles the various content formats that Google's Generative AI SDK accepts,
+    including strings, ContentDict, lists of parts, etc.
+
+    Args:
+        content: The content object to extract text from
+
+    Returns:
+        Extracted text as a string
+    """
+    # Direct string case
+    if isinstance(content, str):
+        return content
+
+    # Lists of parts/content
+    if isinstance(content, list):
+        text = ""
+        for item in content:
+            if isinstance(item, str):
+                text += item + "\n"
+            elif isinstance(item, dict) and "text" in item:
+                if item.get("text") is not None:
+                    text += str(item["text"]) + "\n"
+            elif hasattr(item, "text"):
+                part_text = getattr(item, "text", None)
+                if part_text:
+                    text += part_text + "\n"
+            # Handle content as a list with mixed types
+            elif hasattr(item, "parts"):
+                parts = item.parts
+                for part in parts:
+                    if isinstance(part, str):
+                        text += part + "\n"
+                    elif hasattr(part, "text"):
+                        part_text = getattr(part, "text", None)
+                        if part_text:
+                            text += part_text + "\n"
+        return text
+
+    # Dict with text key
+    if isinstance(content, dict) and "text" in content:
+        return content["text"]
+
+    # Content object with text attribute
+    if hasattr(content, "text"):
+        return content.text
+
+    # Content object with parts attribute
+    if hasattr(content, "parts"):
+        text = ""
+        for part in content.parts:
+            if isinstance(part, str):
+                text += part + "\n"
+            elif hasattr(part, "text"):
+                part_text = getattr(part, "text", None)
+                if part_text:
+                    text += part_text + "\n"
+        return text
+
+    # Other object types - try to convert to string
+    try:
+        return str(content)
+    except Exception:
+        return ""
+
+
+def _set_prompt_attributes(attributes: AttributeMap, args: Tuple, kwargs: Dict[str, Any]) -> None:
+    """Extract and set prompt attributes from the request.
+
+    Respects privacy controls and handles the various ways prompts can be specified
+    in the Google Generative AI API.
+
+    Args:
+        attributes: The attribute dictionary to update
+        args: Positional arguments to the method
+        kwargs: Keyword arguments to the method
+    """
+
+    content = None
+    if args and len(args) > 0:
+        content = args[0]
+    elif "contents" in kwargs:
+        content = kwargs["contents"]
+    elif "content" in kwargs:
+        content = kwargs["content"]
+
+    if content is None:
+        return
+
+    if isinstance(content, list):
+        for i, item in enumerate(content):
+            try:
+                extracted_text = _extract_content_from_prompt(item)
+                if extracted_text:
+                    attributes[MessageAttributes.PROMPT_CONTENT.format(i=i)] = extracted_text
+                    role = "user"
+                    if isinstance(item, dict) and "role" in item:
+                        role = item["role"]
+                    elif hasattr(item, "role"):
+                        role = item.role
+                    attributes[MessageAttributes.PROMPT_ROLE.format(i=i)] = role
+            except Exception as e:
+                logger.debug(f"Error extracting prompt content at index {i}: {e}")
+    else:
+        try:
+            extracted_text = _extract_content_from_prompt(content)
+            if extracted_text:
+                attributes[MessageAttributes.PROMPT_CONTENT.format(i=0)] = extracted_text
+                attributes[MessageAttributes.PROMPT_ROLE.format(i=0)] = "user"
+        except Exception as e:
+            logger.debug(f"Error extracting prompt content: {e}")
+
+
+def _set_response_attributes(attributes: AttributeMap, response: Any) -> None:
+    """Extract and set response attributes from the completion response.
+
+    Args:
+        attributes: The attribute dictionary to update
+        response: The response from the API
+    """
+    if response is None:
+        return
+
+    if hasattr(response, "model"):
+        attributes[SpanAttributes.LLM_RESPONSE_MODEL] = response.model
+
+    if hasattr(response, "usage_metadata"):
+        usage = response.usage_metadata
+        if hasattr(usage, "prompt_token_count"):
+            attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] = usage.prompt_token_count
+        if hasattr(usage, "candidates_token_count"):
+            attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] = usage.candidates_token_count
+        if hasattr(usage, "total_token_count"):
+            attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = usage.total_token_count
+
+    try:
+        if hasattr(response, "text"):
+            attributes[MessageAttributes.COMPLETION_CONTENT.format(i=0)] = response.text
+            attributes[MessageAttributes.COMPLETION_ROLE.format(i=0)] = "assistant"
+        elif hasattr(response, "candidates"):
+            # List of candidates
+            for i, candidate in enumerate(response.candidates):
+                if hasattr(candidate, "content") and hasattr(candidate.content, "parts"):
+                    parts = candidate.content.parts
+                    text = ""
+                    for part in parts:
+                        if isinstance(part, str):
+                            text += part
+                        elif hasattr(part, "text"):
+                            part_text = getattr(part, "text", None)
+                            if part_text:
+                                text += part_text
+
+                    attributes[MessageAttributes.COMPLETION_CONTENT.format(i=i)] = text
+                    attributes[MessageAttributes.COMPLETION_ROLE.format(i=i)] = "assistant"
+
+                if hasattr(candidate, "finish_reason"):
+                    attributes[MessageAttributes.COMPLETION_FINISH_REASON.format(i=i)] = candidate.finish_reason
+    except Exception as e:
+        logger.debug(f"Error extracting completion content: {e}")
+
+
+def get_model_attributes(
+    args: Optional[Tuple] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
+    return_value: Optional[Any] = None,
+) -> AttributeMap:
+    """Extract attributes for GenerativeModel methods.
+
+    This function handles attribute extraction for the general model operations,
+    focusing on the common parameters and pattern shared by multiple methods.
+
+    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)
+
+    if args or kwargs:
+        _set_prompt_attributes(attributes, args or (), kwargs or {})
+
+    if return_value is not None:
+        _set_response_attributes(attributes, return_value)
+
+    return attributes
+
+
+def get_generate_content_attributes(
+    args: Optional[Tuple] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
+    return_value: Optional[Any] = None,
+) -> AttributeMap:
+    """Extract attributes for the generate_content method.
+
+    This specialized extractor handles the generate_content method,
+    which is the primary way to interact with Gemini models.
+
+    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
+    """
+    return get_model_attributes(args, kwargs, return_value)
+
+
+def get_token_counting_attributes(
+    args: Optional[Tuple] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
+    return_value: Optional[Any] = None,
+) -> AttributeMap:
+    """Extract attributes for token counting operations.
+
+    This specialized extractor handles token counting operations.
+
+    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] = "token_count"
+
+    # Process kwargs if available
+    if kwargs:
+        kwargs_attributes = extract_request_attributes(kwargs)
+        attributes.update(kwargs_attributes)
+
+    # Set token count from response
+    if return_value is not None:
+        if hasattr(return_value, "total_tokens"):
+            attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = return_value.total_tokens
+        elif hasattr(return_value, "total_token_count"):
+            attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = return_value.total_token_count
+
+    return attributes
+
+
+def get_stream_attributes(stream: Any) -> AttributeMap:
+    """Extract attributes from a stream object.
+
+    Args:
+        stream: The stream object to extract attributes from
+
+    Returns:
+        Dictionary of extracted attributes
+    """
+    attributes = {}
+
+    if hasattr(stream, "model"):
+        attributes[SpanAttributes.LLM_RESPONSE_MODEL] = stream.model
+
+    return attributes

agentops/instrumentation/providers/google_genai/instrumentor.py

@@ -0,0 +1,170 @@
+"""Google Generative AI Instrumentation for AgentOps
+
+This module provides instrumentation for the Google Generative AI API, implementing OpenTelemetry
+instrumentation for Gemini model requests and responses.
+
+We focus on instrumenting the following key endpoints:
+- ChatSession.send_message - Chat message API
+- Streaming responses - Special handling for streaming responses
+"""
+
+from typing import List, Dict, Any
+from wrapt import wrap_function_wrapper
+from opentelemetry.metrics import Meter
+
+from agentops.logging import logger
+from agentops.instrumentation.common import CommonInstrumentor, StandardMetrics, InstrumentorConfig
+from agentops.instrumentation.common.wrappers import WrapConfig
+from agentops.instrumentation.providers.google_genai.attributes.model import (
+    get_generate_content_attributes,
+    get_token_counting_attributes,
+)
+from agentops.instrumentation.providers.google_genai.stream_wrapper import (
+    generate_content_stream_wrapper,
+    generate_content_stream_async_wrapper,
+)
+
+# Library info for tracer/meter
+LIBRARY_NAME = "agentops.instrumentation.google_genai"
+LIBRARY_VERSION = "0.1.0"
+
+# Methods to wrap for instrumentation
+WRAPPED_METHODS: List[WrapConfig] = [
+    # Client-based API methods
+    WrapConfig(
+        trace_name="gemini.generate_content",
+        package="google.genai.models",
+        class_name="Models",
+        method_name="generate_content",
+        handler=get_generate_content_attributes,
+    ),
+    WrapConfig(
+        trace_name="gemini.count_tokens",
+        package="google.genai.models",
+        class_name="Models",
+        method_name="count_tokens",
+        handler=get_token_counting_attributes,
+    ),
+    WrapConfig(
+        trace_name="gemini.compute_tokens",
+        package="google.genai.models",
+        class_name="Models",
+        method_name="compute_tokens",
+        handler=get_token_counting_attributes,
+    ),
+    # Async client-based API methods
+    WrapConfig(
+        trace_name="gemini.generate_content",
+        package="google.genai.models",
+        class_name="AsyncModels",
+        method_name="generate_content",
+        handler=get_generate_content_attributes,
+        is_async=True,
+    ),
+    WrapConfig(
+        trace_name="gemini.count_tokens",
+        package="google.genai.models",
+        class_name="AsyncModels",
+        method_name="count_tokens",
+        handler=get_token_counting_attributes,
+        is_async=True,
+    ),
+    WrapConfig(
+        trace_name="gemini.compute_tokens",
+        package="google.genai.models",
+        class_name="AsyncModels",
+        method_name="compute_tokens",
+        handler=get_token_counting_attributes,
+        is_async=True,
+    ),
+]
+
+# Streaming methods that need special handling
+STREAMING_METHODS = [
+    # Client API
+    {
+        "module": "google.genai.models",
+        "class_method": "Models.generate_content_stream",
+        "wrapper": generate_content_stream_wrapper,
+        "is_async": False,
+    },
+    {
+        "module": "google.genai.models",
+        "class_method": "AsyncModels.generate_content_stream",
+        "wrapper": generate_content_stream_async_wrapper,
+        "is_async": True,
+    },
+]
+
+
+class GoogleGenaiInstrumentor(CommonInstrumentor):
+    """An instrumentor for Google Generative AI (Gemini) API.
+
+    This class provides instrumentation for Google's Generative AI API by wrapping key methods
+    in the client library and capturing telemetry data. It supports both synchronous and
+    asynchronous API calls, including streaming responses.
+
+    It captures metrics including token usage, operation duration, and exceptions.
+    """
+
+    def __init__(self):
+        """Initialize the Google GenAI instrumentor."""
+        config = InstrumentorConfig(
+            library_name=LIBRARY_NAME,
+            library_version=LIBRARY_VERSION,
+            wrapped_methods=WRAPPED_METHODS,
+            metrics_enabled=True,
+            dependencies=["google-genai >= 0.1.0"],
+        )
+        super().__init__(config)
+
+    def _create_metrics(self, meter: Meter) -> Dict[str, Any]:
+        """Create standard LLM metrics for Google GenAI operations.
+
+        Args:
+            meter: The OpenTelemetry meter to use for creating metrics.
+
+        Returns:
+            Dictionary containing the created metrics.
+        """
+        return StandardMetrics.create_standard_metrics(meter)
+
+    def _custom_wrap(self, **kwargs):
+        """Perform custom wrapping for streaming methods.
+
+        Args:
+            **kwargs: Configuration options for instrumentation.
+        """
+        # Special handling for streaming responses
+        for stream_method in STREAMING_METHODS:
+            try:
+                wrap_function_wrapper(
+                    stream_method["module"],
+                    stream_method["class_method"],
+                    stream_method["wrapper"](self._tracer),
+                )
+                logger.debug(
+                    f"Successfully wrapped streaming method {stream_method['module']}.{stream_method['class_method']}"
+                )
+            except (AttributeError, ModuleNotFoundError) as e:
+                logger.debug(f"Failed to wrap {stream_method['module']}.{stream_method['class_method']}: {e}")
+
+        logger.info("Google Generative AI instrumentation enabled")
+
+    def _custom_unwrap(self, **kwargs):
+        """Remove custom wrapping for streaming methods.
+
+        Args:
+            **kwargs: Configuration options for uninstrumentation.
+        """
+        # Unwrap streaming methods
+        from opentelemetry.instrumentation.utils import unwrap as otel_unwrap
+
+        for stream_method in STREAMING_METHODS:
+            try:
+                otel_unwrap(stream_method["module"], stream_method["class_method"])
+                logger.debug(f"Unwrapped streaming method {stream_method['module']}.{stream_method['class_method']}")
+            except (AttributeError, ModuleNotFoundError) as e:
+                logger.debug(f"Failed to unwrap {stream_method['module']}.{stream_method['class_method']}: {e}")
+
+        logger.info("Google Generative AI instrumentation disabled")