VertexAI support for async calling (#3386)

aabmass · emdneto · web-flow · commit 8b2558f22edb · 2025-03-28T17:44:28.000Z
* VertexAI support for async calling

* Clarify things a bit

* Fix python 3.8 and 3.9 tests

* Remove unnecessary record check

* Update instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/__init__.py

Co-authored-by: Emídio Neto &lt;9735060+emdneto@users.noreply.github.com&gt;

* ruff

---------

Co-authored-by: Emídio Neto &lt;9735060+emdneto@users.noreply.github.com&gt;
diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/CHANGELOG.md b/instrumentation-genai/opentelemetry-instrumentation-vertexai/CHANGELOG.md
@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Implement uninstrument for `opentelemetry-instrumentation-vertexai`
   ([#3328](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3328))
+- VertexAI support for async calling
+  ([#3386](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3386))
 
 ## Version 2.0b0 (2025-02-24)
 
diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/__init__.py b/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/__init__.py
@@ -39,6 +39,8 @@
 ---
 """
 
+from __future__ import annotations
+
 from typing import Any, Collection
 
 from wrapt import (
@@ -49,32 +51,54 @@
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import unwrap
 from opentelemetry.instrumentation.vertexai.package import _instruments
-from opentelemetry.instrumentation.vertexai.patch import (
-    generate_content_create,
-)
+from opentelemetry.instrumentation.vertexai.patch import MethodWrappers
 from opentelemetry.instrumentation.vertexai.utils import is_content_enabled
 from opentelemetry.semconv.schemas import Schemas
 from opentelemetry.trace import get_tracer
 
 
-def _client_classes():
+def _methods_to_wrap(
+    method_wrappers: MethodWrappers,
+):
     # This import is very slow, do it lazily in case instrument() is not called
-
     # pylint: disable=import-outside-toplevel
     from google.cloud.aiplatform_v1.services.prediction_service import (
+        async_client,
         client,
     )
+    from google.cloud.aiplatform_v1beta1.services.prediction_service import (
+        async_client as async_client_v1beta1,
+    )
     from google.cloud.aiplatform_v1beta1.services.prediction_service import (
         client as client_v1beta1,
     )
 
-    return (
+    for client_class in (
         client.PredictionServiceClient,
         client_v1beta1.PredictionServiceClient,
-    )
+    ):
+        yield (
+            client_class,
+            client_class.generate_content.__name__,  # pyright: ignore[reportUnknownMemberType]
+            method_wrappers.generate_content,
+        )
+
+    for client_class in (
+        async_client.PredictionServiceAsyncClient,
+        async_client_v1beta1.PredictionServiceAsyncClient,
+    ):
+        yield (
+            client_class,
+            client_class.generate_content.__name__,  # pyright: ignore[reportUnknownMemberType]
+            method_wrappers.agenerate_content,
+        )
 
 
 class VertexAIInstrumentor(BaseInstrumentor):
+    def __init__(self) -> None:
+        super().__init__()
+        self._methods_to_unwrap: list[tuple[Any, str]] = []
+
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 
@@ -95,15 +119,19 @@ def _instrument(self, **kwargs: Any):
             event_logger_provider=event_logger_provider,
         )
 
-        for client_class in _client_classes():
+        method_wrappers = MethodWrappers(
+            tracer, event_logger, is_content_enabled()
+        )
+        for client_class, method_name, wrapper in _methods_to_wrap(
+            method_wrappers
+        ):
             wrap_function_wrapper(
                 client_class,
-                name="generate_content",
-                wrapper=generate_content_create(
-                    tracer, event_logger, is_content_enabled()
-                ),
+                name=method_name,
+                wrapper=wrapper,
             )
+            self._methods_to_unwrap.append((client_class, method_name))
 
     def _uninstrument(self, **kwargs: Any) -> None:
-        for client_class in _client_classes():
-            unwrap(client_class, "generate_content")
+        for client_class, method_name in self._methods_to_unwrap:
+            unwrap(client_class, method_name)
diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/patch.py b/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/patch.py
@@ -14,9 +14,11 @@
 
 from __future__ import annotations
 
+from contextlib import contextmanager
 from typing import (
     TYPE_CHECKING,
     Any,
+    Awaitable,
     Callable,
     MutableSequence,
 )
@@ -87,17 +89,17 @@ def _extract_params(
     )
 
 
-def generate_content_create(
-    tracer: Tracer, event_logger: EventLogger, capture_content: bool
-):
-    """Wrap the `generate_content` method of the `GenerativeModel` class to trace it."""
+class MethodWrappers:
+    def __init__(
+        self, tracer: Tracer, event_logger: EventLogger, capture_content: bool
+    ) -> None:
+        self.tracer = tracer
+        self.event_logger = event_logger
+        self.capture_content = capture_content
 
-    def traced_method(
-        wrapped: Callable[
-            ...,
-            prediction_service.GenerateContentResponse
-            | prediction_service_v1beta1.GenerateContentResponse,
-        ],
+    @contextmanager
+    def _with_instrumentation(
+        self,
         instance: client.PredictionServiceClient
         | client_v1beta1.PredictionServiceClient,
         args: Any,
@@ -111,32 +113,82 @@ def traced_method(
         }
 
         span_name = get_span_name(span_attributes)
-        with tracer.start_as_current_span(
+
+        with self.tracer.start_as_current_span(
             name=span_name,
             kind=SpanKind.CLIENT,
             attributes=span_attributes,
         ) as span:
             for event in request_to_events(
-                params=params, capture_content=capture_content
+                params=params, capture_content=self.capture_content
             ):
-                event_logger.emit(event)
+                self.event_logger.emit(event)
 
             # TODO: set error.type attribute
             # https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-spans.md
-            response = wrapped(*args, **kwargs)
-            # TODO: handle streaming
-            # if is_streaming(kwargs):
-            #     return StreamWrapper(
-            #         result, span, event_logger, capture_content
-            #     )
-
-            if span.is_recording():
-                span.set_attributes(get_genai_response_attributes(response))
-            for event in response_to_events(
-                response=response, capture_content=capture_content
-            ):
-                event_logger.emit(event)
 
+            def handle_response(
+                response: prediction_service.GenerateContentResponse
+                | prediction_service_v1beta1.GenerateContentResponse,
+            ) -> None:
+                if span.is_recording():
+                    # When streaming, this is called multiple times so attributes would be
+                    # overwritten. In practice, it looks the API only returns the interesting
+                    # attributes on the last streamed response. However, I couldn't find
+                    # documentation for this and setting attributes shouldn't be too expensive.
+                    span.set_attributes(
+                        get_genai_response_attributes(response)
+                    )
+
+                for event in response_to_events(
+                    response=response, capture_content=self.capture_content
+                ):
+                    self.event_logger.emit(event)
+
+            yield handle_response
+
+    def generate_content(
+        self,
+        wrapped: Callable[
+            ...,
+            prediction_service.GenerateContentResponse
+            | prediction_service_v1beta1.GenerateContentResponse,
+        ],
+        instance: client.PredictionServiceClient
+        | client_v1beta1.PredictionServiceClient,
+        args: Any,
+        kwargs: Any,
+    ) -> (
+        prediction_service.GenerateContentResponse
+        | prediction_service_v1beta1.GenerateContentResponse
+    ):
+        with self._with_instrumentation(
+            instance, args, kwargs
+        ) as handle_response:
+            response = wrapped(*args, **kwargs)
+            handle_response(response)
             return response
 
-    return traced_method
+    async def agenerate_content(
+        self,
+        wrapped: Callable[
+            ...,
+            Awaitable[
+                prediction_service.GenerateContentResponse
+                | prediction_service_v1beta1.GenerateContentResponse
+            ],
+        ],
+        instance: client.PredictionServiceClient
+        | client_v1beta1.PredictionServiceClient,
+        args: Any,
+        kwargs: Any,
+    ) -> (
+        prediction_service.GenerateContentResponse
+        | prediction_service_v1beta1.GenerateContentResponse
+    ):
+        with self._with_instrumentation(
+            instance, args, kwargs
+        ) as handle_response:
+            response = await wrapped(*args, **kwargs)
+            handle_response(response)
+            return response
diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/requirements.latest.txt b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/requirements.latest.txt
@@ -44,8 +44,8 @@ charset-normalizer==3.4.0
 Deprecated==1.2.15
 docstring_parser==0.16
 exceptiongroup==1.2.2
-google-api-core==2.23.0
-google-auth==2.36.0
+google-api-core[grpc, async_rest]==2.23.0
+google-auth[aiohttp]==2.36.0
 google-cloud-aiplatform==1.79.0
 google-cloud-bigquery==3.27.0
 google-cloud-core==2.4.1
diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/requirements.oldest.txt b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/requirements.oldest.txt
@@ -22,8 +22,8 @@ charset-normalizer==3.4.0
 Deprecated==1.2.14
 docstring_parser==0.16
 exceptiongroup==1.2.2
-google-api-core==2.23.0
-google-auth==2.36.0
+google-api-core[grpc, async_rest]==2.23.0
+google-auth[aiohttp]==2.36.0
 google-cloud-aiplatform==1.79.0
 google-cloud-bigquery==3.27.0
 google-cloud-core==2.4.1
diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_chat_completions.py b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_chat_completions.py
