packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/__init__.py (1)

306-319: Async wrapper is not selected for google.genai AsyncModels.generate_content

The condition only checks for a legacy method name (“generate_content_async”). For the new library, AsyncModels.generate_content is a coroutine and must use the async wrapper. As-is, it will wrap an async function with a sync wrapper and break instrumentation.

Select the wrapper based on actual coroutine-ness (or Async* class) by introspecting the target:

@@
-        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")
-
-            wrap_function_wrapper(
-                wrap_package,
-                f"{wrap_object}.{wrap_method}",
-                (
-                    _awrap(tracer, event_logger, wrapped_method)
-                    if wrap_method == "generate_content_async"
-                    else _wrap(tracer, event_logger, wrapped_method)
-                ),
-            )
+        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")
+
+            # Pick async wrapper if the target is a coroutine (e.g., AsyncModels.generate_content)
+            import importlib, inspect
+            mod = importlib.import_module(wrap_package)
+            obj = getattr(mod, wrap_object)
+            fn = getattr(obj, wrap_method)
+            is_async = inspect.iscoroutinefunction(fn) or wrap_object.lower().startswith("async") or wrap_method.endswith("_async")
+
+            wrap_function_wrapper(
+                wrap_package,
+                f"{wrap_object}.{wrap_method}",
+                _awrap(tracer, event_logger, wrapped_method) if is_async else _wrap(tracer, event_logger, wrapped_method),
+            )

Note: Ensure importlib and inspect are imported at module scope or inline as shown.

packages/opentelemetry-instrumentation-google-generativeai/tests/conftest.py (1)

64-72: Prevent flaky KeyError on GOOGLE_API_KEY and make fixture independent of autouse ordering

Relying on an autouse fixture to set the env var can still be brittle in pytest’s fixture execution order. Make the fixture resilient by reading the key via getenv once and reusing it in both branches.

Apply this diff within the fixture to avoid KeyError—no behavior change when the env var is present:

@@
-    # Prefer new client over legacy
-    if genai_new:
-        client = genai_new.Client(api_key=os.environ["GOOGLE_API_KEY"])
+    # Prefer new client over legacy
+    api_key = os.environ.get("GOOGLE_API_KEY", "test_api_key")
+    if genai_new:
+        client = genai_new.Client(api_key=api_key)
         return client.models
     elif genai_legacy:
-        genai_legacy.configure(api_key=os.environ["GOOGLE_API_KEY"])
+        genai_legacy.configure(api_key=api_key)
         return genai_legacy.GenerativeModel("gemini-1.5-flash")

Optionally, to guarantee initialization order, make the fixture depend explicitly on the autouse environment fixture:

@pytest.fixture
def genai_client(environment):
    ...

packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/__init__.py (2)

279-285: Return a consistent type from instrumentation_dependencies

Minor consistency: return the same sequence type across branches. Using tuples everywhere avoids surprises.

@@
-        if is_package_installed("google.genai"):
-            return ("google-genai >= 0.1.0",)
-        elif is_package_installed("google.generativeai"):
-            return ["google-generativeai >= 0.5.0"]
-        else:
-            return []
+        if is_package_installed("google.genai"):
+            return ("google-genai >= 0.1.0",)
+        elif is_package_installed("google.generativeai"):
+            return ("google-generativeai >= 0.5.0",)
+        else:
+            return ()

---

7-13: Avoid runtime imports purely for typing to reduce fragility

If GenerateContentResponse is only used for annotations, consider deferring its import with TYPE_CHECKING or enabling postponed evaluation of annotations. This avoids setting it to None and prevents import churn at runtime.

Option A (preferred): use postponed evaluation

from __future__ import annotations

Option B: TYPE_CHECKING guard

from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from google.generativeai.types.generation_types import GenerateContentResponse as LegacyGenerateContentResponse
    from google.genai.types import GenerateContentResponse as NewGenerateContentResponse
    GenerateContentResponse = LegacyGenerateContentResponse | NewGenerateContentResponse

packages/opentelemetry-instrumentation-google-generativeai/tests/test_import_compatibility.py (5)

1-7: Import pytest to enable conditional skipping

You’re asserting both libraries are present; in typical user environments, only one may be installed. Add pytest and leverage skip conditions in the tests below.

@@
-"""Test that both legacy and new Google GenAI libraries can be imported and instrumented."""
+"""Test that both legacy and new Google GenAI libraries can be imported and instrumented."""
+
+import pytest

---

9-13: Skip legacy detection test when the legacy package isn’t installed

Asserting presence is brittle; skip instead to keep the suite green across environments.

@@
-def test_legacy_library_detection():
+@pytest.mark.skipif(
+    not is_package_installed("google.generativeai"),
+    reason="google.generativeai not 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"
+    assert is_package_installed("google.generativeai")

---

15-19: Skip new library detection test when the new package isn’t installed

Mirror the legacy case to avoid false negatives when only one lib is present.

@@
-def test_new_library_detection():
+@pytest.mark.skipif(
+    not is_package_installed("google.genai"),
+    reason="google.genai not 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"
+    assert is_package_installed("google.genai")

---

21-34: Skip initialization test when neither library is installed

Keep the test suite environment-agnostic.

@@
-def test_instrumentor_initialization():
+@pytest.mark.skipif(
+    not (is_package_installed("google.genai") or is_package_installed("google.generativeai")),
+    reason="Neither google.genai nor google.generativeai is installed",
+)
+def test_instrumentor_initialization():

---

36-50: Skip wrapped-methods selection test when neither library is installed

Prevents failures on environments lacking both packages.

@@
-def test_wrapped_methods_selection():
+@pytest.mark.skipif(
+    not (is_package_installed("google.genai") or is_package_installed("google.generativeai")),
+    reason="Neither google.genai nor google.generativeai is installed",
+)
+def test_wrapped_methods_selection():

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/event_emitter.py (2)

5-5: Avoid hard dependency on google-genai in a shared module used by both clients.

Unconditional import of GenerateContentResponse from google.genai.types forces the new library to be present even if only the legacy client is used. If fallback to legacy is still supported, decouple the type import from runtime by using a TYPE_CHECKING guard. This keeps typing without breaking environments that only have the legacy lib.

Is legacy fallback still expected? If yes, consider this change:

-from google.genai.types import GenerateContentResponse
+from typing import TYPE_CHECKING, Any
+
+# Avoid a hard dependency on google-genai solely for typing.
+if TYPE_CHECKING:
+    from google.genai.types import GenerateContentResponse  # pragma: no cover
+else:
+    GenerateContentResponse = Any  # type: ignore[misc, assignment]

---

60-69: Guard against finish_reason being None or non-enum.

Depending on the library/version, candidate.finish_reason may be an enum, string, or None. Accessing .name can raise AttributeError. Consider a defensive conversion to a string.

Apply this tweak:

-                finish_reason=candidate.finish_reason.name,
+                finish_reason=getattr(candidate.finish_reason, "name", str(candidate.finish_reason or "UNKNOWN")),

packages/opentelemetry-instrumentation-google-generativeai/tests/conftest.py (1)

56-57: Fixture returns Models; consider exposing async models too for broader test coverage.

Returning client.models works for current tests. If you plan to add async coverage, you may also return or add a fixture for AsyncModels (if available via the client API), to validate async wrapping end-to-end.

packages/opentelemetry-instrumentation-google-generativeai/tests/test_new_library_instrumentation.py (1)

8-41: Also assert async method wrapping to strengthen coverage.

Great checks for Models.generate_content. Add AsyncModels.generate_content_async pre/post assertions to ensure the async path is wrapped and unwrapped too.

Apply this diff:

     # Import the library
     from google import genai
-    from google.genai.models import Models
+    from google.genai.models import Models, AsyncModels
@@
     # Verify methods are not wrapped initially
     assert not hasattr(Models.generate_content, '__wrapped__')
+    assert not hasattr(AsyncModels.generate_content_async, '__wrapped__')
@@
         # Verify methods are now wrapped
         assert hasattr(Models.generate_content, '__wrapped__')
+        assert hasattr(AsyncModels.generate_content_async, '__wrapped__')
@@
         # Verify methods are unwrapped
         assert not hasattr(Models.generate_content, '__wrapped__')
+        assert not hasattr(AsyncModels.generate_content_async, '__wrapped__')

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.