Added Vertex AI spans for request parameters

instrumentation-genai/opentelemetry-instrumentation-vertexai/CHANGELOG.md

@@ -7,5 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+- Added Vertex AI spans for request parameters
+  ([#3192](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3192))
 - Initial VertexAI instrumentation
   ([#3123](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3123))

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

@@ -41,9 +41,17 @@
 
 from typing import Any, Collection
 
+from wrapt import (
+    wrap_function_wrapper,  # type: ignore[reportUnknownVariableType]
+)
+
 from opentelemetry._events import get_event_logger
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.vertexai.package import _instruments
+from opentelemetry.instrumentation.vertexai.patch import (
+    generate_content_create,
+)
+from opentelemetry.instrumentation.vertexai.utils import is_content_enabled
 from opentelemetry.semconv.schemas import Schemas
 from opentelemetry.trace import get_tracer
 
@@ -55,20 +63,34 @@ def instrumentation_dependencies(self) -> Collection[str]:
     def _instrument(self, **kwargs: Any):
         """Enable VertexAI instrumentation."""
         tracer_provider = kwargs.get("tracer_provider")
-        _tracer = get_tracer(
+        tracer = get_tracer(
             __name__,
             "",
             tracer_provider,
             schema_url=Schemas.V1_28_0.value,
         )
         event_logger_provider = kwargs.get("event_logger_provider")
-        _event_logger = get_event_logger(
+        event_logger = get_event_logger(
             __name__,
             "",
             schema_url=Schemas.V1_28_0.value,
             event_logger_provider=event_logger_provider,
         )
-        # TODO: implemented in later PR
+
+        wrap_function_wrapper(
+            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()
+            ),
+        )
 
     def _uninstrument(self, **kwargs: Any) -> None:
         """TODO: implemented in later PR"""

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

@@ -11,3 +11,115 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Iterable,
+    MutableSequence,
+    Optional,
+    Union,
+)
+
+from opentelemetry._events import EventLogger
+from opentelemetry.instrumentation.vertexai.utils import (
+    GenerateContentParams,
+    get_genai_request_attributes,
+    get_span_name,
+)
+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,
+    )
+    from vertexai.generative_models._generative_models import (
+        _GenerativeModel,
+    )
+
+
+# Use parameter signature from
+# 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,
+    *,
+    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(
+        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,
+    )
+
+
+def generate_content_create(
+    tracer: Tracer, event_logger: EventLogger, capture_content: bool
+):
+    """Wrap the `generate_content` method of the `GenerativeModel` class to trace it."""
+
+    def traced_method(
+        wrapped: Callable[
+            ..., GenerationResponse | Iterable[GenerationResponse]
+        ],
+        instance: _GenerativeModel,
+        args: Any,
+        kwargs: Any,
+    ):
+        params = _extract_params(*args, **kwargs)
+        span_attributes = get_genai_request_attributes(params)
+
+        span_name = get_span_name(span_attributes)
+        with tracer.start_as_current_span(
+            name=span_name,
+            kind=SpanKind.CLIENT,
+            attributes=span_attributes,
+        ) as _span:
+            # TODO: emit request events
+            # if span.is_recording():
+            #     for message in kwargs.get("messages", []):
+            #         event_logger.emit(
+            #             message_to_event(message, capture_content)
+            #         )
+
+            result = wrapped(*args, **kwargs)
+            # TODO: handle streaming
+            # if is_streaming(kwargs):
+            #     return StreamWrapper(
+            #         result, span, event_logger, capture_content
+            #     )
+
+            # TODO: add response attributes and events
+            # if span.is_recording():
+            #     _set_response_attributes(
+            #         span, result, event_logger, capture_content
+            #     )
+            return result
+
+    return traced_method

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

@@ -0,0 +1,122 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from os import environ
+from typing import (
+    TYPE_CHECKING,
+    Mapping,
+    Optional,
+    Sequence,
+)
+
+from opentelemetry.semconv._incubating.attributes import (
+    gen_ai_attributes as GenAIAttributes,
+)
+from opentelemetry.util.types import AttributeValue
+
+if TYPE_CHECKING:
+    from google.cloud.aiplatform_v1.types import content, tool
+
+
+@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
+
+
+def get_genai_request_attributes(
+    params: GenerateContentParams,
+    operation_name: GenAIAttributes.GenAiOperationNameValues = GenAIAttributes.GenAiOperationNameValues.CHAT,
+):
+    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,
+    }
+
+    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(model: str) -> str:
+    return _MODEL_STRIP_RE.sub("", model)
+
+
+OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT = (
+    "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"
+)
+
+
+def is_content_enabled() -> bool:
+    capture_content = environ.get(
+        OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, "false"
+    )
+
+    return capture_content.lower() == "true"
+
+
+def get_span_name(span_attributes: Mapping[str, AttributeValue]):
+    name = span_attributes.get(GenAIAttributes.GEN_AI_OPERATION_NAME, "")
+    model = span_attributes.get(GenAIAttributes.GEN_AI_REQUEST_MODEL, "")
+    return f"{name} {model}"

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

@@ -0,0 +1,82 @@
+interactions:
+- request:
+    body: |-
+      {
+        "contents": [
+          {
+            "role": "user",
+            "parts": [
+              {
+                "text": "Say this is a test"
+              }
+            ]
+          }
+        ],
+        "generationConfig": {
+          "temperature": 0.2,
+          "topP": 0.95,
+          "topK": 2.0,
+          "maxOutputTokens": 5,
+          "stopSequences": [
+            "\n\n\n"
+          ],
+          "presencePenalty": -1.5,
+          "frequencyPenalty": 1.0,
+          "seed": 12345
+        }
+      }
+    headers:
+      Accept:
+      - '*/*'
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '376'
+      Content-Type:
+      - application/json
+      User-Agent:
+      - python-requests/2.32.3
+    method: POST
+    uri: https://us-central1-aiplatform.googleapis.com/v1/projects/fake-project/locations/us-central1/publishers/google/models/gemini-1.5-flash-002:generateContent?%24alt=json%3Benum-encoding%3Dint
+  response:
+    body:
+      string: |-
+        {
+          "candidates": [
+            {
+              "content": {
+                "role": "model",
+                "parts": [
+                  {
+                    "text": "Okay, I understand."
+                  }
+                ]
+              },
+              "finishReason": 2,
+              "avgLogprobs": -0.006721805781126022
+            }
+          ],
+          "usageMetadata": {
+            "promptTokenCount": 5,
+            "candidatesTokenCount": 5,
+            "totalTokenCount": 10
+          },
+          "modelVersion": "gemini-1.5-flash-002"
+        }
+    headers:
+      Content-Type:
+      - application/json; charset=UTF-8
+      Transfer-Encoding:
+      - chunked
+      Vary:
+      - Origin
+      - X-Origin
+      - Referer
+      content-length:
+      - '407'
+    status:
+      code: 200
+      message: OK
+version: 1