fix(google-genai): migrate Google Generative AI instrumentation to googleapis/python-genai

packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/__init__.py

@@ -4,7 +4,13 @@
 import types
 from typing import Collection
 
-from google.generativeai.types.generation_types import GenerateContentResponse
+try:
+    from google.generativeai.types.generation_types import GenerateContentResponse
+except ImportError:
+    try:
+        from google.genai.types import GenerateContentResponse
+    except ImportError:
+        GenerateContentResponse = None
 from opentelemetry import context as context_api
 from opentelemetry._events import get_event_logger
 from opentelemetry.instrumentation.google_generativeai.config import Config
@@ -66,21 +72,6 @@
     },
 ]
 
-WRAPPED_METHODS = [
-    {
-        "package": "google.genai.models",
-        "object": "Models",
-        "method": "generate_content",
-        "span_name": "gemini.generate_content",
-    },
-    {
-        "package": "google.genai.models",
-        "object": "AsyncModels",
-        "method": "generate_content",
-        "span_name": "gemini.generate_content",
-    },
-]
-
 
 def is_streaming_response(response):
     return isinstance(response, types.GeneratorType)
@@ -312,7 +303,7 @@ def _instrument(self, **kwargs):
                 __name__, __version__, event_logger_provider=event_logger_provider
             )
 
-        for wrapped_method in WRAPPED_METHODS:
+        for wrapped_method in self._wrapped_methods():
             wrap_package = wrapped_method.get("package")
             wrap_object = wrapped_method.get("object")
             wrap_method = wrapped_method.get("method")

packages/opentelemetry-instrumentation-google-generativeai/tests/conftest.py

@@ -2,7 +2,15 @@
 
 import os
 
-import google.generativeai as genai
+try:
+    import google.generativeai as genai_legacy
+except ImportError:
+    genai_legacy = None
+
+try:
+    from google import genai as genai_new
+except ImportError:
+    genai_new = None
 import pytest
 from opentelemetry.instrumentation.google_generativeai import (
     GoogleGenerativeAiInstrumentor,
@@ -53,9 +61,15 @@ def fixture_event_logger_provider(log_exporter):
 
 @pytest.fixture
 def genai_client():
-    genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
-
-    return genai.GenerativeModel("gemini-1.5-flash")
+    # Prefer new client over legacy
+    if genai_new:
+        client = genai_new.Client(api_key=os.environ["GOOGLE_API_KEY"])
+        return client.models
+    elif genai_legacy:
+        genai_legacy.configure(api_key=os.environ["GOOGLE_API_KEY"])
+        return genai_legacy.GenerativeModel("gemini-1.5-flash")
+    else:
+        pytest.skip("Neither google-genai nor google-generativeai is installed")
 
 
 @pytest.fixture(scope="function")

packages/opentelemetry-instrumentation-google-generativeai/tests/test_import_compatibility.py

@@ -0,0 +1,49 @@
+"""Test that both legacy and new Google GenAI libraries can be imported and instrumented."""
+
+from opentelemetry.instrumentation.google_generativeai import (
+    GoogleGenerativeAiInstrumentor,
+)
+from opentelemetry.instrumentation.google_generativeai.utils import is_package_installed
+
+
+def test_legacy_library_detection():
+    """Test that the legacy library can be detected."""
+    has_legacy = is_package_installed("google.generativeai")
+    assert has_legacy, "google.generativeai should be installed"
+
+
+def test_new_library_detection():
+    """Test that the new library can be detected."""
+    has_new = is_package_installed("google.genai")
+    assert has_new, "google.genai should be installed"
+
+
+def test_instrumentor_initialization():
+    """Test that the instrumentor can be initialized correctly."""
+    instrumentor = GoogleGenerativeAiInstrumentor()
+
+    # Should have appropriate dependencies
+    deps = instrumentor.instrumentation_dependencies()
+    assert len(deps) > 0, "Should have instrumentation dependencies"
+
+    # Should prefer new library over legacy
+    if is_package_installed("google.genai"):
+        assert "google-genai" in deps[0]
+    else:
+        assert "google-generativeai" in deps[0]
+
+
+def test_wrapped_methods_selection():
+    """Test that the correct wrapped methods are selected based on available library."""
+    instrumentor = GoogleGenerativeAiInstrumentor()
+    methods = instrumentor._wrapped_methods()
+
+    assert len(methods) > 0, "Should have wrapped methods"
+
+    # Check that the right package is being wrapped
+    for method in methods:
+        package = method.get("package", "")
+        if is_package_installed("google.genai"):
+            assert "google.genai" in package
+        else:
+            assert "google.generativeai" in package