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 index d2c7b98adc..8f293ecb00 100644 --- a/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/patch.py +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/patch.py @@ -18,10 +18,7 @@ TYPE_CHECKING, Any, Callable, - Iterable, MutableSequence, - Optional, - Union, ) from opentelemetry._events import EventLogger @@ -33,15 +30,19 @@ from opentelemetry.trace import SpanKind, Tracer if TYPE_CHECKING: + from google.cloud.aiplatform_v1.services.prediction_service import client from google.cloud.aiplatform_v1.types import ( content, prediction_service, ) - from vertexai.generative_models import ( - GenerationResponse, + from google.cloud.aiplatform_v1beta1.services.prediction_service import ( + client as client_v1beta1, ) - from vertexai.generative_models._generative_models import ( - _GenerativeModel, + from google.cloud.aiplatform_v1beta1.types import ( + content as content_v1beta1, + ) + from google.cloud.aiplatform_v1beta1.types import ( + prediction_service as prediction_service_v1beta1, ) @@ -49,12 +50,15 @@ # 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( - request: Optional[ - Union[prediction_service.GenerateContentRequest, dict[Any, Any]] - ] = None, + request: prediction_service.GenerateContentRequest + | prediction_service_v1beta1.GenerateContentRequest + | dict[Any, Any] + | None = None, *, - model: Optional[str] = None, - contents: Optional[MutableSequence[content.Content]] = None, + model: str | None = None, + contents: MutableSequence[content.Content] + | MutableSequence[content_v1beta1.Content] + | None = None, **_kwargs: Any, ) -> GenerateContentParams: # Request vs the named parameters are mututally exclusive or the RPC will fail @@ -86,9 +90,12 @@ def generate_content_create( def traced_method( wrapped: Callable[ - ..., GenerationResponse | Iterable[GenerationResponse] + ..., + prediction_service.GenerateContentResponse + | prediction_service_v1beta1.GenerateContentResponse, ], - instance: _GenerativeModel, + instance: client.PredictionServiceClient + | client_v1beta1.PredictionServiceClient, args: Any, kwargs: Any, ): diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/utils.py b/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/utils.py index c0bff12824..183a445f3d 100644 --- a/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/utils.py +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/utils.py @@ -20,7 +20,6 @@ from typing import ( TYPE_CHECKING, Mapping, - Optional, Sequence, ) @@ -31,18 +30,32 @@ if TYPE_CHECKING: from google.cloud.aiplatform_v1.types import content, tool + from google.cloud.aiplatform_v1beta1.types import ( + content as content_v1beta1, + ) + from google.cloud.aiplatform_v1beta1.types import ( + tool as tool_v1beta1, + ) @dataclass(frozen=True) class GenerateContentParams: 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 + contents: ( + Sequence[content.Content] | Sequence[content_v1beta1.Content] | None + ) = None + system_instruction: content.Content | content_v1beta1.Content | None = None + tools: Sequence[tool.Tool] | Sequence[tool_v1beta1.Tool] | None = None + tool_config: tool.ToolConfig | tool_v1beta1.ToolConfig | None = None + labels: Mapping[str, str] | None = None + safety_settings: ( + Sequence[content.SafetySetting] + | Sequence[content_v1beta1.SafetySetting] + | None + ) = None + generation_config: ( + content.GenerationConfig | content_v1beta1.GenerationConfig | None + ) = None def get_genai_request_attributes(