Actually use GAPIC client since thats what we use under the hood

Also this is what LangChain uses

Author: aabmass

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

@@ -78,10 +78,15 @@ def _instrument(self, **kwargs: Any):
         )
 
         wrap_function_wrapper(
-            module="vertexai.generative_models._generative_models",
-            # Patching this base class also instruments the vertexai.preview.generative_models
-            # package
-            name="_GenerativeModel.generate_content",
+            module="google.cloud.aiplatform_v1beta1.services.prediction_service.client",
+            name="PredictionServiceClient.generate_content",
+            wrapper=generate_content_create(
+                tracer, event_logger, is_content_enabled()
+            ),
+        )
+        wrap_function_wrapper(
+            module="google.cloud.aiplatform_v1.services.prediction_service.client",
+            name="PredictionServiceClient.generate_content",
             wrapper=generate_content_create(
                 tracer, event_logger, is_content_enabled()
             ),

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

@@ -14,7 +14,15 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Iterable,
+    MutableSequence,
+    Optional,
+    Union,
+)
 
 from opentelemetry._events import EventLogger
 from opentelemetry.instrumentation.vertexai.utils import (
@@ -25,41 +33,49 @@
 from opentelemetry.trace import SpanKind, Tracer
 
 if TYPE_CHECKING:
+    from google.cloud.aiplatform_v1.types import (
+        content,
+        prediction_service,
+    )
     from vertexai.generative_models import (
         GenerationResponse,
-        Tool,
-        ToolConfig,
     )
     from vertexai.generative_models._generative_models import (
-        ContentsType,
-        GenerationConfigType,
-        SafetySettingsType,
         _GenerativeModel,
     )
 
 
 # Use parameter signature from
-# https://github.com/googleapis/python-aiplatform/blob/v1.76.0/vertexai/generative_models/_generative_models.py#L595
+# https://github.com/googleapis/python-aiplatform/blob/v1.76.0/google/cloud/aiplatform_v1/services/prediction_service/client.py#L2088
 # to handle named vs positional args robustly
 def _extract_params(
-    contents: ContentsType,
+    request: Optional[
+        Union[prediction_service.GenerateContentRequest, dict[Any, Any]]
+    ] = None,
     *,
-    generation_config: Optional[GenerationConfigType] = None,
-    safety_settings: Optional[SafetySettingsType] = None,
-    tools: Optional[list[Tool]] = None,
-    tool_config: Optional[ToolConfig] = None,
-    labels: Optional[dict[str, str]] = None,
-    stream: bool = False,
+    model: Optional[str] = None,
+    contents: Optional[MutableSequence[content.Content]] = None,
     **_kwargs: Any,
 ) -> GenerateContentParams:
+    # Request vs the named parameters are mututally exclusive or the RPC will fail
+    if not request:
+        return GenerateContentParams(
+            model=model or "",
+            contents=contents,
+        )
+
+    if isinstance(request, dict):
+        return GenerateContentParams(**request)
+
     return GenerateContentParams(
-        contents=contents,
-        generation_config=generation_config,
-        safety_settings=safety_settings,
-        tools=tools,
-        tool_config=tool_config,
-        labels=labels,
-        stream=stream,
+        model=request.model,
+        contents=request.contents,
+        system_instruction=request.system_instruction,
+        tools=request.tools,
+        tool_config=request.tool_config,
+        labels=request.labels,
+        safety_settings=request.safety_settings,
+        generation_config=request.generation_config,
     )
 
 
@@ -77,7 +93,7 @@ def traced_method(
         kwargs: Any,
     ):
         params = _extract_params(*args, **kwargs)
-        span_attributes = get_genai_request_attributes(instance, params)
+        span_attributes = get_genai_request_attributes(params)
 
         span_name = get_span_name(span_attributes)
         with tracer.start_as_current_span(

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

@@ -14,16 +14,14 @@
 
 from __future__ import annotations
 
+import re
 from dataclasses import dataclass
 from os import environ
 from typing import (
     TYPE_CHECKING,
-    Dict,
-    List,
     Mapping,
     Optional,
-    TypedDict,
-    cast,
+    Sequence,
 )
 
 from opentelemetry.semconv._incubating.attributes import (
@@ -32,96 +30,77 @@
 from opentelemetry.util.types import AttributeValue
 
 if TYPE_CHECKING:
-    from vertexai.generative_models import Tool, ToolConfig
-    from vertexai.generative_models._generative_models import (
-        ContentsType,
-        GenerationConfigType,
-        SafetySettingsType,
-        _GenerativeModel,
-    )
+    from google.cloud.aiplatform_v1.types import content, tool
 
 
 @dataclass(frozen=True)
 class GenerateContentParams:
-    contents: ContentsType
-    generation_config: Optional[GenerationConfigType]
-    safety_settings: Optional[SafetySettingsType]
-    tools: Optional[List["Tool"]]
-    tool_config: Optional["ToolConfig"]
-    labels: Optional[Dict[str, str]]
-    stream: bool
-
-
-class GenerationConfigDict(TypedDict, total=False):
-    temperature: Optional[float]
-    top_p: Optional[float]
-    top_k: Optional[int]
-    max_output_tokens: Optional[int]
-    stop_sequences: Optional[List[str]]
-    presence_penalty: Optional[float]
-    frequency_penalty: Optional[float]
-    seed: Optional[int]
-    # And more fields which aren't needed yet
+    model: str
+    contents: Optional[Sequence[content.Content]] = None
+    system_instruction: Optional[content.Content | None] = None
+    tools: Optional[Sequence[tool.Tool]] = None
+    tool_config: Optional[tool.ToolConfig] = None
+    labels: Optional[Mapping[str, str]] = None
+    safety_settings: Optional[Sequence[content.SafetySetting]] = None
+    generation_config: Optional[content.GenerationConfig] = None
 
 
 def get_genai_request_attributes(
-    instance: _GenerativeModel,
     params: GenerateContentParams,
     operation_name: GenAIAttributes.GenAiOperationNameValues = GenAIAttributes.GenAiOperationNameValues.CHAT,
 ):
-    model = _get_model_name(instance)
-    generation_config = _get_generation_config(instance, params)
-    attributes = {
+    model = _get_model_name(params.model)
+    generation_config = params.generation_config
+    attributes: dict[str, AttributeValue] = {
         GenAIAttributes.GEN_AI_OPERATION_NAME: operation_name.value,
         GenAIAttributes.GEN_AI_SYSTEM: GenAIAttributes.GenAiSystemValues.VERTEX_AI.value,
         GenAIAttributes.GEN_AI_REQUEST_MODEL: model,
-        GenAIAttributes.GEN_AI_REQUEST_TEMPERATURE: generation_config.get(
-            "temperature"
-        ),
-        GenAIAttributes.GEN_AI_REQUEST_TOP_P: generation_config.get("top_p"),
-        GenAIAttributes.GEN_AI_REQUEST_MAX_TOKENS: generation_config.get(
-            "max_output_tokens"
-        ),
-        GenAIAttributes.GEN_AI_REQUEST_PRESENCE_PENALTY: generation_config.get(
-            "presence_penalty"
-        ),
-        GenAIAttributes.GEN_AI_REQUEST_FREQUENCY_PENALTY: generation_config.get(
-            "frequency_penalty"
-        ),
-        GenAIAttributes.GEN_AI_OPENAI_REQUEST_SEED: generation_config.get(
-            "seed"
-        ),
-        GenAIAttributes.GEN_AI_REQUEST_STOP_SEQUENCES: generation_config.get(
-            "stop_sequences"
-        ),
     }
 
-    # filter out None values
-    return {k: v for k, v in attributes.items() if v is not None}
-
-
-def _get_generation_config(
-    instance: _GenerativeModel,
-    params: GenerateContentParams,
-) -> GenerationConfigDict:
-    generation_config = params.generation_config or instance._generation_config
-    if generation_config is None:
-        return {}
-    if isinstance(generation_config, dict):
-        return cast(GenerationConfigDict, generation_config)
-    return cast(GenerationConfigDict, generation_config.to_dict())
-
-
-_RESOURCE_PREFIX = "publishers/google/models/"
-
+    if not generation_config:
+        return attributes
+
+    # Check for optional fields
+    # https://proto-plus-python.readthedocs.io/en/stable/fields.html#optional-fields
+    if "temperature" in generation_config:
+        attributes[GenAIAttributes.GEN_AI_REQUEST_TEMPERATURE] = (
+            generation_config.temperature
+        )
+    if "top_p" in generation_config:
+        attributes[GenAIAttributes.GEN_AI_REQUEST_TOP_P] = (
+            generation_config.top_p
+        )
+    if "max_output_tokens" in generation_config:
+        attributes[GenAIAttributes.GEN_AI_REQUEST_MAX_TOKENS] = (
+            generation_config.max_output_tokens
+        )
+    if "presence_penalty" in generation_config:
+        attributes[GenAIAttributes.GEN_AI_REQUEST_PRESENCE_PENALTY] = (
+            generation_config.presence_penalty
+        )
+    if "frequency_penalty" in generation_config:
+        attributes[GenAIAttributes.GEN_AI_REQUEST_FREQUENCY_PENALTY] = (
+            generation_config.frequency_penalty
+        )
+    if "seed" in generation_config:
+        attributes[GenAIAttributes.GEN_AI_OPENAI_REQUEST_SEED] = (
+            generation_config.seed
+        )
+    if "stop_sequences" in generation_config:
+        attributes[GenAIAttributes.GEN_AI_REQUEST_STOP_SEQUENCES] = (
+            generation_config.stop_sequences
+        )
+
+    return attributes
+
+
+_MODEL_STRIP_RE = re.compile(
+    r"^projects/(.*)/locations/(.*)/publishers/google/models/"
+)
 
-def _get_model_name(instance: _GenerativeModel) -> str:
-    model_name = instance._model_name
 
-    # Can use str.removeprefix() once 3.8 is dropped
-    if model_name.startswith(_RESOURCE_PREFIX):
-        model_name = model_name[len(_RESOURCE_PREFIX) :]
-    return model_name
+def _get_model_name(model: str) -> str:
+    return _MODEL_STRIP_RE.sub("", model)
 
 
 OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT = (

instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_chat_completion_extra_client_level_params.yaml

@@ -55,7 +55,7 @@ interactions:
                 ]
               },
               "finishReason": 2,
-              "avgLogprobs": -0.006723951548337936
+              "avgLogprobs": -0.006721805781126022
             }
           ],
           "usageMetadata": {