Add GoogleGenerativeAIInstrumentor (#932)

* Add GoogleGenerativeAIInstrumentor and top_k span attribute

* Add LLM_REQUEST_SEED attribute to SpanAttributes class

* Refactor attribute handling in Google Generative AI instrumentation and added new request configuration attributes to SpanAttributes

* Refactor chat message attribute assignment to use MessageAttributes constants for improved clarity and consistency

* Update provider import name for GoogleGenerativeAIInstrumentor from "google.generativeai" to "google.genai" for consistency.

* Update library name from "google-generativeai" to "google-genai" for consistency with recent changes.

* Remove token count tracking from content stream wrappers to streamline processing and improve performance.

* Add token count tracking to content stream wrappers for improved text processing

* Remove token count tracking from content stream wrappers to streamline processing and improve performance.

Author: Dwij1704

agentops/instrumentation/__init__.py

@@ -72,6 +72,11 @@ def get_instance(self) -> BaseInstrumentor:
         class_name="OpenAIAgentsInstrumentor",
         provider_import_name="agents",
     ),
+    InstrumentorLoader(
+        module_name="agentops.instrumentation.google_generativeai",
+        class_name="GoogleGenerativeAIInstrumentor",
+        provider_import_name="google.genai",
+    ),
 ]
 
 

agentops/instrumentation/google_generativeai/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/google_generativeai/__init__.py

@@ -0,0 +1,38 @@
+"""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 typing import Collection
+
+def get_version() -> str:
+    """Get the version of the Google Generative AI SDK, or 'unknown' if not found
+    
+    Attempts to retrieve the installed version of the Google Generative AI SDK using importlib.metadata.
+    Falls back to 'unknown' if the version cannot be determined.
+    
+    Returns:
+        The version string of the Google Generative AI SDK or 'unknown'
+    """
+    try:
+        from importlib.metadata import version
+        return version("google-genai")
+    except ImportError:
+        logger.debug("Could not find Google Generative AI SDK version")
+        return "unknown"
+
+LIBRARY_NAME = "google-genai"
+LIBRARY_VERSION: str = get_version()
+
+logger = logging.getLogger(__name__)
+
+# Import after defining constants to avoid circular imports
+from agentops.instrumentation.google_generativeai.instrumentor import GoogleGenerativeAIInstrumentor  # noqa: E402
+
+__all__ = [
+    "LIBRARY_NAME",
+    "LIBRARY_VERSION", 
+    "GoogleGenerativeAIInstrumentor",
+] 

agentops/instrumentation/google_generativeai/attributes/__init__.py

@@ -0,0 +1,25 @@
+"""Attribute extractors for Google Generative AI instrumentation."""
+
+from agentops.instrumentation.google_generativeai.attributes.common import (
+    get_common_instrumentation_attributes,
+    extract_request_attributes,
+)
+from agentops.instrumentation.google_generativeai.attributes.model import (
+    get_model_attributes,
+    get_generate_content_attributes,
+    get_stream_attributes,
+    get_token_counting_attributes,
+)
+from agentops.instrumentation.google_generativeai.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/google_generativeai/attributes/chat.py

@@ -0,0 +1,125 @@
+"""Chat attribute extraction for Google Generative AI instrumentation."""
+
+from typing import Dict, Any, Optional, Tuple, List, Union
+
+from agentops.logging import logger
+from agentops.semconv import SpanAttributes, LLMRequestTypeValues, MessageAttributes
+from agentops.instrumentation.common.attributes import AttributeMap
+from agentops.instrumentation.google_generativeai.attributes.common import (
+    extract_request_attributes,
+    get_common_instrumentation_attributes,
+)
+from agentops.instrumentation.google_generativeai.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 

agentops/instrumentation/google_generativeai/attributes/common.py

@@ -0,0 +1,80 @@
+"""Common attribute extraction for Google Generative AI instrumentation."""
+
+from typing import Dict, Any, Optional
+
+from agentops.logging import logger
+from agentops.semconv import InstrumentationAttributes, SpanAttributes, LLMRequestTypeValues
+from agentops.instrumentation.common.attributes import AttributeMap, get_common_attributes, _extract_attributes_from_mapping
+from agentops.instrumentation.google_generativeai 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