Implement Google ADK Instrumentation for AgentOps (#1009)

Dwij1704 · dot-agi · web-flow · commit bd397dd32363 · 2025-05-28T03:16:01.000+05:30
* Implement Google ADK Instrumentation for AgentOps

* Fix import compatibility for NotRequired in instrumentation module by adding a fallback to typing_extensions. This ensures broader compatibility across Python versions.

* Update instrumentation_dependencies to require 'google-adk' instead of 'google-genai' for improved package management.

* Refactor LLM attribute extraction to handle system instructions as role messages. Updated message processing to include proper indexing and role normalization for contents, enhancing clarity and consistency in attribute assignment.

* Refactor attribute assignment in Google ADK instrumentation to use constants for span attributes. This change enhances code clarity and maintainability by replacing string literals with defined constants for agent, LLM, and tool attributes.

* Enhance Google ADK instrumentation by adding 'type' parameter to span creation for agent, LLM, and tool contexts. This improves clarity in span categorization and aligns with the existing attribute assignment structure.

* Refactor Google ADK instrumentation to remove 'type' parameter from span creation for agent, LLM, and tool contexts. This change simplifies span initialization while maintaining clarity in span categorization through attribute assignment.

* refactor version code into `__init__.py`

---------

Co-authored-by: Pratyush Shukla &lt;ps4534@nyu.edu&gt;
diff --git a/agentops/instrumentation/__init__.py b/agentops/instrumentation/__init__.py
@@ -13,6 +13,11 @@
 """
 
 from typing import Optional, Set, TypedDict
+
+try:
+    from typing import NotRequired
+except ImportError:
+    from typing_extensions import NotRequired
 from types import ModuleType
 from dataclasses import dataclass
 import importlib
@@ -36,8 +41,11 @@
 
 def _is_package_instrumented(package_name: str) -> bool:
     """Check if a package is already instrumented by looking at active instrumentors."""
+    # Handle package.module names by converting dots to underscores for comparison
+    normalized_name = package_name.replace(".", "_").lower()
     return any(
-        instrumentor.__class__.__name__.lower().startswith(package_name.lower())
+        instrumentor.__class__.__name__.lower().startswith(normalized_name)
+        or instrumentor.__class__.__name__.lower().startswith(package_name.split(".")[-1].lower())
         for instrumentor in _active_instrumentors
     )
 
@@ -65,17 +73,14 @@ def _should_instrument_package(package_name: str) -> bool:
     if package_name in AGENTIC_LIBRARIES:
         _uninstrument_providers()
         _has_agentic_library = True
-        logger.debug(f"Uninstrumented all providers due to agentic library {package_name} detection")
         return True
 
     # Skip providers if an agentic library is already instrumented
     if package_name in PROVIDERS and _has_agentic_library:
-        logger.debug(f"Skipping provider {package_name} instrumentation as an agentic library is already instrumented")
         return False
 
     # Skip if already instrumented
     if _is_package_instrumented(package_name):
-        logger.debug(f"Package {package_name} is already instrumented")
         return False
 
     return True
@@ -102,36 +107,62 @@ def _import_monitor(name: str, globals_dict=None, locals_dict=None, fromlist=(),
     Monitor imports and instrument packages as they are imported.
     This replaces the built-in import function to intercept package imports.
     """
-    global _instrumenting_packages
-    root = name.split(".", 1)[0]
+    global _instrumenting_packages, _has_agentic_library
 
-    # Skip providers if an agentic library is already instrumented
-    if _has_agentic_library and root in PROVIDERS:
+    # If an agentic library is already instrumented, skip all further instrumentation
+    if _has_agentic_library:
         return _original_builtins_import(name, globals_dict, locals_dict, fromlist, level)
 
-    # Check if this is a package we should instrument
-    if (
-        root in TARGET_PACKAGES
-        and root not in _instrumenting_packages
-        and not _is_package_instrumented(root)  # Check if already instrumented before adding
-    ):
-        logger.debug(f"Detected import of {root}")
-        _instrumenting_packages.add(root)
-        try:
-            _perform_instrumentation(root)
-        except Exception as e:
-            logger.error(f"Error instrumenting {root}: {str(e)}")
-        finally:
-            _instrumenting_packages.discard(root)
-
-    return _original_builtins_import(name, globals_dict, locals_dict, fromlist, level)
+    # First, do the actual import
+    module = _original_builtins_import(name, globals_dict, locals_dict, fromlist, level)
+
+    # Check for exact matches first (handles package.module like google.adk)
+    packages_to_check = set()
+
+    # Check the imported module itself
+    if name in TARGET_PACKAGES:
+        packages_to_check.add(name)
+    else:
+        # Check if any target package is a prefix of the import name
+        for target in TARGET_PACKAGES:
+            if name.startswith(target + ".") or name == target:
+                packages_to_check.add(target)
+
+    # For "from X import Y" style imports, also check submodules
+    if fromlist:
+        for item in fromlist:
+            full_name = f"{name}.{item}"
+            if full_name in TARGET_PACKAGES:
+                packages_to_check.add(full_name)
+            else:
+                # Check if any target package matches this submodule
+                for target in TARGET_PACKAGES:
+                    if full_name == target or full_name.startswith(target + "."):
+                        packages_to_check.add(target)
+
+    # Instrument all matching packages
+    for package_to_check in packages_to_check:
+        if package_to_check not in _instrumenting_packages and not _is_package_instrumented(package_to_check):
+            _instrumenting_packages.add(package_to_check)
+            try:
+                _perform_instrumentation(package_to_check)
+                # If we just instrumented an agentic library, stop
+                if _has_agentic_library:
+                    break
+            except Exception as e:
+                logger.error(f"Error instrumenting {package_to_check}: {str(e)}")
+            finally:
+                _instrumenting_packages.discard(package_to_check)
+
+    return module
 
 
 # Define the structure for instrumentor configurations
 class InstrumentorConfig(TypedDict):
     module_name: str
     class_name: str
     min_version: str
+    package_name: NotRequired[str]  # Optional: actual pip package name if different from module
 
 
 # Configuration for supported LLM providers
@@ -146,16 +177,17 @@ class InstrumentorConfig(TypedDict):
         "class_name": "AnthropicInstrumentor",
         "min_version": "0.32.0",
     },
-    "google.genai": {
-        "module_name": "agentops.instrumentation.google_generativeai",
-        "class_name": "GoogleGenerativeAIInstrumentor",
-        "min_version": "0.1.0",
-    },
     "ibm_watsonx_ai": {
         "module_name": "agentops.instrumentation.ibm_watsonx_ai",
         "class_name": "IBMWatsonXInstrumentor",
         "min_version": "0.1.0",
     },
+    "google.genai": {
+        "module_name": "agentops.instrumentation.google_generativeai",
+        "class_name": "GoogleGenerativeAIInstrumentor",
+        "min_version": "0.1.0",
+        "package_name": "google-genai",  # Actual pip package name
+    },
 }
 
 # Configuration for supported agentic libraries
@@ -171,6 +203,11 @@ class InstrumentorConfig(TypedDict):
         "class_name": "OpenAIAgentsInstrumentor",
         "min_version": "0.0.1",
     },
+    "google.adk": {
+        "module_name": "agentops.instrumentation.google_adk",
+        "class_name": "GoogleADKInstrumentor",
+        "min_version": "0.1.0",
+    },
 }
 
 # Combine all target packages for monitoring
@@ -190,6 +227,7 @@ class InstrumentorLoader:
     module_name: str
     class_name: str
     min_version: str
+    package_name: Optional[str] = None  # Optional: actual pip package name
 
     @property
     def module(self) -> ModuleType:
@@ -200,7 +238,11 @@ def module(self) -> ModuleType:
     def should_activate(self) -> bool:
         """Check if the package is available and meets version requirements."""
         try:
-            provider_name = self.module_name.split(".")[-1]
+            # Use explicit package_name if provided, otherwise derive from module_name
+            if self.package_name:
+                provider_name = self.package_name
+            else:
+                provider_name = self.module_name.split(".")[-1]
             module_version = version(provider_name)
             return module_version is not None and Version(module_version) >= parse(self.min_version)
         except ImportError:
@@ -233,24 +275,44 @@ def instrument_all():
     # Check if active_instrumentors is empty, as a proxy for not started.
     if not _active_instrumentors:
         builtins.__import__ = _import_monitor
-        global _instrumenting_packages
+        global _instrumenting_packages, _has_agentic_library
+
+        # If an agentic library is already instrumented, don't instrument anything else
+        if _has_agentic_library:
+            return
+
         for name in list(sys.modules.keys()):
+            # Stop if an agentic library gets instrumented during the loop
+            if _has_agentic_library:
+                break
+
             module = sys.modules.get(name)
             if not isinstance(module, ModuleType):
                 continue
 
-            root = name.split(".", 1)[0]
-            if _has_agentic_library and root in PROVIDERS:
-                continue
-
-            if root in TARGET_PACKAGES and root not in _instrumenting_packages and not _is_package_instrumented(root):
-                _instrumenting_packages.add(root)
+            # Check for exact matches first (handles package.module like google.adk)
+            package_to_check = None
+            if name in TARGET_PACKAGES:
+                package_to_check = name
+            else:
+                # Check if any target package is a prefix of the module name
+                for target in TARGET_PACKAGES:
+                    if name.startswith(target + ".") or name == target:
+                        package_to_check = target
+                        break
+
+            if (
+                package_to_check
+                and package_to_check not in _instrumenting_packages
+                and not _is_package_instrumented(package_to_check)
+            ):
+                _instrumenting_packages.add(package_to_check)
                 try:
-                    _perform_instrumentation(root)
+                    _perform_instrumentation(package_to_check)
                 except Exception as e:
-                    logger.error(f"Error instrumenting {root}: {str(e)}")
+                    logger.error(f"Error instrumenting {package_to_check}: {str(e)}")
                 finally:
-                    _instrumenting_packages.discard(root)
+                    _instrumenting_packages.discard(package_to_check)
 
 
 def uninstrument_all():
@@ -269,8 +331,19 @@ def get_active_libraries() -> set[str]:
     Get all actively used libraries in the current execution context.
     Returns a set of package names that are currently imported and being monitored.
     """
-    return {
-        name.split(".")[0]
-        for name, module in sys.modules.items()
-        if isinstance(module, ModuleType) and name.split(".")[0] in TARGET_PACKAGES
-    }
+    active_libs = set()
+    for name, module in sys.modules.items():
+        if not isinstance(module, ModuleType):
+            continue
+
+        # Check for exact matches first
+        if name in TARGET_PACKAGES:
+            active_libs.add(name)
+        else:
+            # Check if any target package is a prefix of the module name
+            for target in TARGET_PACKAGES:
+                if name.startswith(target + ".") or name == target:
+                    active_libs.add(target)
+                    break
+
+    return active_libs
diff --git a/agentops/instrumentation/google_adk/__init__.py b/agentops/instrumentation/google_adk/__init__.py
@@ -0,0 +1,20 @@
+"""Google ADK Instrumentation for AgentOps
+
+This module provides instrumentation for Google's Agent Development Kit (ADK),
+capturing agent execution, LLM calls, tool calls, and other ADK-specific events.
+"""
+
+from importlib.metadata import version, PackageNotFoundError
+
+try:
+    __version__ = version("google-adk")
+except PackageNotFoundError:
+    __version__ = "0.0.0"
+
+LIBRARY_NAME = "agentops.instrumentation.google_adk"
+LIBRARY_VERSION = __version__
+
+from agentops.instrumentation.google_adk.instrumentor import GoogleADKInstrumentor  # noqa: E402
+from agentops.instrumentation.google_adk import patch  # noqa: E402
+
+__all__ = ["LIBRARY_NAME", "LIBRARY_VERSION", "GoogleADKInstrumentor", "patch"]
diff --git a/agentops/instrumentation/google_adk/instrumentor.py b/agentops/instrumentation/google_adk/instrumentor.py
@@ -0,0 +1,78 @@
+"""Google ADK Instrumentation for AgentOps
+
+This module provides instrumentation for Google's Agent Development Kit (ADK).
+It uses a patching approach to:
+1. Disable ADK's built-in telemetry to prevent duplicate spans
+2. Create AgentOps spans that mirror ADK's telemetry structure
+3. Extract and properly index LLM messages and tool calls
+"""
+
+from typing import Collection
+from opentelemetry.trace import get_tracer
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.metrics import get_meter
+
+from agentops.logging import logger
+from agentops.instrumentation.google_adk import LIBRARY_NAME, LIBRARY_VERSION
+from agentops.instrumentation.google_adk.patch import patch_adk, unpatch_adk
+from agentops.semconv import Meters
+
+
+class GoogleADKInstrumentor(BaseInstrumentor):
+    """An instrumentor for Google Agent Development Kit (ADK).
+
+    This instrumentor patches Google ADK to:
+    - Prevent ADK from creating its own telemetry spans
+    - Create AgentOps spans for agent runs, LLM calls, and tool calls
+    - Properly extract and index message content and tool interactions
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        """Return packages required for instrumentation."""
+        return ["google-adk >= 0.1.0"]
+
+    def _instrument(self, **kwargs):
+        """Instrument the Google ADK.
+
+        This method:
+        1. Disables ADK's built-in telemetry
+        2. Patches key ADK methods to create AgentOps spans
+        3. Sets up metrics for tracking token usage and operation duration
+        """
+        # Set up tracer and meter
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(LIBRARY_NAME, LIBRARY_VERSION, tracer_provider)
+
+        meter_provider = kwargs.get("meter_provider")
+        meter = get_meter(LIBRARY_NAME, LIBRARY_VERSION, meter_provider)
+
+        # Create metrics
+        meter.create_histogram(
+            name=Meters.LLM_TOKEN_USAGE,
+            unit="token",
+            description="Measures number of input and output tokens used with Google ADK",
+        )
+
+        meter.create_histogram(
+            name=Meters.LLM_OPERATION_DURATION,
+            unit="s",
+            description="Google ADK operation duration",
+        )
+
+        meter.create_counter(
+            name=Meters.LLM_COMPLETIONS_EXCEPTIONS,
+            unit="time",
+            description="Number of exceptions occurred during Google ADK operations",
+        )
+
+        # Apply patches
+        patch_adk(tracer)
+        logger.info("Google ADK instrumentation enabled")
+
+    def _uninstrument(self, **kwargs):
+        """Remove instrumentation from Google ADK.
+
+        This method removes all patches and restores ADK's original behavior.
+        """
+        unpatch_adk()
+        logger.info("Google ADK instrumentation disabled")
diff --git a/agentops/instrumentation/google_adk/patch.py b/agentops/instrumentation/google_adk/patch.py
