Add boilerplate for async wrappers

Author: aabmass

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

@@ -62,7 +62,13 @@ def _methods_to_wrap(
 ):
     # 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 client
+    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,
     )
@@ -77,6 +83,16 @@ def _methods_to_wrap(
             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:

instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/patch.py

@@ -18,6 +18,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    AsyncIterable,
     Callable,
     Iterable,
     MutableSequence,
@@ -168,3 +169,25 @@ def generate_content(
             response = wrapped(*args, **kwargs)
             handle_response(response)
             return response
+
+    async def agenerate_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