diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/CHANGELOG.md b/instrumentation-genai/opentelemetry-instrumentation-vertexai/CHANGELOG.md index 4d786c7840..4e43fbff19 100644 --- a/instrumentation-genai/opentelemetry-instrumentation-vertexai/CHANGELOG.md +++ b/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)) 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 index 9437184ff0..40d1cb48ac 100644 --- a/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/__init__.py +++ b/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""" 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 b0a6f42841..36a31045b5 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 @@ -11,3 +11,124 @@ # 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, + MutableSequence, +) + +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.services.prediction_service import client + from google.cloud.aiplatform_v1.types import ( + content, + prediction_service, + ) + from google.cloud.aiplatform_v1beta1.services.prediction_service import ( + client as client_v1beta1, + ) + from google.cloud.aiplatform_v1beta1.types import ( + content as content_v1beta1, + ) + from google.cloud.aiplatform_v1beta1.types import ( + prediction_service as prediction_service_v1beta1, + ) + + +# 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: prediction_service.GenerateContentRequest + | prediction_service_v1beta1.GenerateContentRequest + | dict[Any, Any] + | None = 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 + 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[ + ..., + prediction_service.GenerateContentResponse + | prediction_service_v1beta1.GenerateContentResponse, + ], + instance: client.PredictionServiceClient + | client_v1beta1.PredictionServiceClient, + 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) + # ) + + # TODO: set error.type attribute + # https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-spans.md + 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 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 new file mode 100644 index 0000000000..96d7125028 --- /dev/null +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/utils.py @@ -0,0 +1,139 @@ +# 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, + 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 + 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: ( + 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( + 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 + ) + # Uncomment once GEN_AI_REQUEST_SEED is released in 1.30 + # https://github.com/open-telemetry/semantic-conventions/pull/1710 + # if "seed" in generation_config: + # attributes[GenAIAttributes.GEN_AI_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]) -> str: + name = span_attributes[GenAIAttributes.GEN_AI_OPERATION_NAME] + model = span_attributes[GenAIAttributes.GEN_AI_REQUEST_MODEL] + if not model: + return f"{name}" + return f"{name} {model}" diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_generate_content.yaml b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_generate_content.yaml new file mode 100644 index 0000000000..69856f9308 --- /dev/null +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_generate_content.yaml @@ -0,0 +1,70 @@ +interactions: +- request: + body: |- + { + "contents": [ + { + "role": "user", + "parts": [ + { + "text": "Say this is a test" + } + ] + } + ] + } + headers: + Accept: + - '*/*' + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + Content-Length: + - '141' + 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. I'm ready for your test. Please proceed.\n" + } + ] + }, + "finishReason": 1, + "avgLogprobs": -0.005692833348324424 + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 19, + "totalTokenCount": 24 + }, + "modelVersion": "gemini-1.5-flash-002" + } + headers: + Content-Type: + - application/json; charset=UTF-8 + Transfer-Encoding: + - chunked + Vary: + - Origin + - X-Origin + - Referer + content-length: + - '453' + status: + code: 200 + message: OK +version: 1 diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_generate_content_extra_params.yaml b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_generate_content_extra_params.yaml new file mode 100644 index 0000000000..e6547166bc --- /dev/null +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_generate_content_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 diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_generate_content_invalid_temperature.yaml b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_generate_content_invalid_temperature.yaml new file mode 100644 index 0000000000..600635f9b2 --- /dev/null +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_generate_content_invalid_temperature.yaml @@ -0,0 +1,59 @@ +interactions: +- request: + body: |- + { + "contents": [ + { + "role": "user", + "parts": [ + { + "text": "Say this is a test" + } + ] + } + ], + "generationConfig": { + "temperature": 1000.0 + } + } + headers: + Accept: + - '*/*' + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + Content-Length: + - '196' + 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: |- + { + "error": { + "code": 400, + "message": "Unable to submit request because it has a temperature value of 1000 but the supported range is from 0 (inclusive) to 2.0001 (exclusive). Update the value and try again.", + "status": "INVALID_ARGUMENT", + "details": [] + } + } + headers: + Content-Type: + - application/json; charset=UTF-8 + Transfer-Encoding: + - chunked + Vary: + - Origin + - X-Origin + - Referer + content-length: + - '809' + status: + code: 400 + message: Bad Request +version: 1 diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_generate_content_missing_model.yaml b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_generate_content_missing_model.yaml new file mode 100644 index 0000000000..efe3e152ce --- /dev/null +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_generate_content_missing_model.yaml @@ -0,0 +1,56 @@ +interactions: +- request: + body: |- + { + "contents": [ + { + "role": "user", + "parts": [ + { + "text": "Say this is a test" + } + ] + } + ] + } + headers: + Accept: + - '*/*' + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + Content-Length: + - '141' + 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-does-not-exist:generateContent?%24alt=json%3Benum-encoding%3Dint + response: + body: + string: |- + { + "error": { + "code": 404, + "message": "Publisher Model `projects/otel-starter-project/locations/us-central1/publishers/google/models/gemini-does-not-exist` not found.", + "status": "NOT_FOUND", + "details": [] + } + } + headers: + Content-Type: + - application/json; charset=UTF-8 + Transfer-Encoding: + - chunked + Vary: + - Origin + - X-Origin + - Referer + content-length: + - '672' + status: + code: 404 + message: Not Found +version: 1 diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/conftest.py b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/conftest.py index 8337188ece..b76a108805 100644 --- a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/conftest.py +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/conftest.py @@ -1,22 +1,42 @@ """Unit tests configuration module.""" import json +import os +import re +from typing import Any, Mapping, MutableMapping import pytest +import vertexai import yaml - +from google.auth.credentials import AnonymousCredentials +from vcr import VCR +from vcr.record_mode import RecordMode +from vcr.request import Request + +from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor +from opentelemetry.instrumentation.vertexai.utils import ( + OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, +) from opentelemetry.sdk._events import EventLoggerProvider from opentelemetry.sdk._logs import LoggerProvider from opentelemetry.sdk._logs.export import ( InMemoryLogExporter, SimpleLogRecordProcessor, ) +from opentelemetry.sdk.metrics import ( + MeterProvider, +) +from opentelemetry.sdk.metrics.export import ( + InMemoryMetricReader, +) from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( InMemorySpanExporter, ) +FAKE_PROJECT = "fake-project" + @pytest.fixture(scope="function", name="span_exporter") def fixture_span_exporter(): @@ -30,6 +50,12 @@ def fixture_log_exporter(): yield exporter +@pytest.fixture(scope="function", name="metric_reader") +def fixture_metric_reader(): + exporter = InMemoryMetricReader() + yield exporter + + @pytest.fixture(scope="function", name="tracer_provider") def fixture_tracer_provider(span_exporter): provider = TracerProvider() @@ -46,17 +72,103 @@ def fixture_event_logger_provider(log_exporter): return event_logger_provider +@pytest.fixture(scope="function", name="meter_provider") +def fixture_meter_provider(metric_reader): + return MeterProvider( + metric_readers=[metric_reader], + ) + + +@pytest.fixture(autouse=True) +def vertexai_init(vcr: VCR) -> None: + # When not recording (in CI), don't do any auth. That prevents trying to read application + # default credentials from the filesystem or metadata server and oauth token exchange. This + # is not the interesting part of our instrumentation to test. + credentials = None + project = None + if vcr.record_mode == RecordMode.NONE: + credentials = AnonymousCredentials() + project = FAKE_PROJECT + vertexai.init( + api_transport="rest", credentials=credentials, project=project + ) + + +@pytest.fixture +def instrument_no_content( + tracer_provider, event_logger_provider, meter_provider +): + os.environ.update( + {OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT: "False"} + ) + + instrumentor = VertexAIInstrumentor() + instrumentor.instrument( + tracer_provider=tracer_provider, + event_logger_provider=event_logger_provider, + meter_provider=meter_provider, + ) + + yield instrumentor + os.environ.pop(OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, None) + instrumentor.uninstrument() + + +@pytest.fixture +def instrument_with_content( + tracer_provider, event_logger_provider, meter_provider +): + os.environ.update( + {OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT: "True"} + ) + instrumentor = VertexAIInstrumentor() + instrumentor.instrument( + tracer_provider=tracer_provider, + event_logger_provider=event_logger_provider, + meter_provider=meter_provider, + ) + + yield instrumentor + os.environ.pop(OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, None) + instrumentor.uninstrument() + + @pytest.fixture(scope="module") def vcr_config(): + filter_header_regexes = [ + r"X-.*", + "Server", + "Date", + "Expires", + "Authorization", + ] + + def filter_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + return { + key: val + for key, val in headers.items() + if not any( + re.match(filter_re, key, re.IGNORECASE) + for filter_re in filter_header_regexes + ) + } + + def before_record_cb(request: Request): + request.headers = filter_headers(request.headers) + request.uri = re.sub( + r"/projects/[^/]+/", "/projects/fake-project/", request.uri + ) + return request + + def before_response_cb(response: MutableMapping[str, Any]): + response["headers"] = filter_headers(response["headers"]) + return response + return { - "filter_headers": [ - ("cookie", "test_cookie"), - ("authorization", "Bearer test_vertexai_api_key"), - ("vertexai-organization", "test_vertexai_org_id"), - ("vertexai-project", "test_vertexai_project_id"), - ], "decode_compressed_response": True, - "before_record_response": scrub_response_headers, + "before_record_request": before_record_cb, + "before_record_response": before_response_cb, + "ignore_hosts": ["oauth2.googleapis.com"], } @@ -125,12 +237,3 @@ def deserialize(cassette_string): def fixture_vcr(vcr): vcr.register_serializer("yaml", PrettyPrintJSONBody) return vcr - - -def scrub_response_headers(response): - """ - This scrubs sensitive response headers. Note they are case-sensitive! - """ - response["headers"]["vertexai-organization"] = "test_vertexai_org_id" - response["headers"]["Set-Cookie"] = "test_set_cookie" - return response diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_chat_completions.py b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_chat_completions.py new file mode 100644 index 0000000000..63a2e2c2d1 --- /dev/null +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_chat_completions.py @@ -0,0 +1,173 @@ +import pytest +from google.api_core.exceptions import BadRequest, NotFound +from vertexai.generative_models import ( + Content, + GenerationConfig, + GenerativeModel, + Part, +) + +from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.trace import StatusCode + + +@pytest.mark.vcr +def test_generate_content( + span_exporter: InMemorySpanExporter, + instrument_with_content: VertexAIInstrumentor, +): + model = GenerativeModel("gemini-1.5-flash-002") + model.generate_content( + [ + Content(role="user", parts=[Part.from_text("Say this is a test")]), + ] + ) + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == "chat gemini-1.5-flash-002" + assert dict(spans[0].attributes) == { + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "gemini-1.5-flash-002", + "gen_ai.system": "vertex_ai", + } + + +@pytest.mark.vcr +def test_generate_content_empty_model( + span_exporter: InMemorySpanExporter, + instrument_with_content: VertexAIInstrumentor, +): + model = GenerativeModel("") + try: + model.generate_content( + [ + Content( + role="user", parts=[Part.from_text("Say this is a test")] + ) + ], + ) + except ValueError: + pass + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == "chat" + # Captures invalid params + assert dict(spans[0].attributes) == { + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "", + "gen_ai.system": "vertex_ai", + } + assert_span_error(spans[0]) + + +@pytest.mark.vcr +def test_generate_content_missing_model( + span_exporter: InMemorySpanExporter, + instrument_with_content: VertexAIInstrumentor, +): + model = GenerativeModel("gemini-does-not-exist") + try: + model.generate_content( + [ + Content( + role="user", parts=[Part.from_text("Say this is a test")] + ) + ], + ) + except NotFound: + pass + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == "chat gemini-does-not-exist" + # Captures invalid params + assert dict(spans[0].attributes) == { + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "gemini-does-not-exist", + "gen_ai.system": "vertex_ai", + } + assert_span_error(spans[0]) + + +@pytest.mark.vcr +def test_generate_content_invalid_temperature( + span_exporter: InMemorySpanExporter, + instrument_with_content: VertexAIInstrumentor, +): + model = GenerativeModel("gemini-1.5-flash-002") + try: + # Temperature out of range causes error + model.generate_content( + [ + Content( + role="user", parts=[Part.from_text("Say this is a test")] + ) + ], + generation_config=GenerationConfig(temperature=1000), + ) + except BadRequest: + pass + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == "chat gemini-1.5-flash-002" + assert dict(spans[0].attributes) == { + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "gemini-1.5-flash-002", + "gen_ai.request.temperature": 1000.0, + "gen_ai.system": "vertex_ai", + } + assert_span_error(spans[0]) + + +@pytest.mark.vcr() +def test_generate_content_extra_params(span_exporter, instrument_no_content): + generation_config = GenerationConfig( + top_k=2, + top_p=0.95, + temperature=0.2, + stop_sequences=["\n\n\n"], + max_output_tokens=5, + presence_penalty=-1.5, + frequency_penalty=1.0, + seed=12345, + ) + model = GenerativeModel("gemini-1.5-flash-002") + model.generate_content( + [ + Content(role="user", parts=[Part.from_text("Say this is a test")]), + ], + generation_config=generation_config, + ) + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert dict(spans[0].attributes) == { + "gen_ai.operation.name": "chat", + "gen_ai.request.frequency_penalty": 1.0, + "gen_ai.request.max_tokens": 5, + "gen_ai.request.model": "gemini-1.5-flash-002", + "gen_ai.request.presence_penalty": -1.5, + "gen_ai.request.stop_sequences": ("\n\n\n",), + "gen_ai.request.temperature": 0.20000000298023224, + "gen_ai.request.top_p": 0.949999988079071, + "gen_ai.system": "vertex_ai", + } + + +def assert_span_error(span: ReadableSpan) -> None: + # Sets error status + assert span.status.status_code == StatusCode.ERROR + + # TODO: check thate error.type is set + # https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-spans.md + + # Records exception event + error_events = [e for e in span.events if e.name == "exception"] + assert error_events != [] diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_placeholder.py b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_placeholder.py deleted file mode 100644 index c910bfa0bf..0000000000 --- a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_placeholder.py +++ /dev/null @@ -1,20 +0,0 @@ -# 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. - -# TODO: adapt tests from OpenLLMetry here along with tests from -# instrumentation-genai/opentelemetry-instrumentation-openai-v2/tests/test_chat_completions.py - - -def test_placeholder(): - assert True diff --git a/tox.ini b/tox.ini index 806c74b426..22e56a835f 100644 --- a/tox.ini +++ b/tox.ini @@ -802,7 +802,7 @@ commands = test-instrumentation-openai-v2: pytest {toxinidir}/instrumentation-genai/opentelemetry-instrumentation-openai-v2/tests {posargs} lint-instrumentation-openai-v2: sh -c "cd instrumentation-genai && pylint --rcfile ../.pylintrc opentelemetry-instrumentation-openai-v2" - test-instrumentation-vertexai: pytest {toxinidir}/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests {posargs} + test-instrumentation-vertexai: pytest {toxinidir}/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests --vcr-record=none {posargs} lint-instrumentation-vertexai: sh -c "cd instrumentation-genai && pylint --rcfile ../.pylintrc opentelemetry-instrumentation-vertexai" test-instrumentation-sio-pika: pytest {toxinidir}/instrumentation/opentelemetry-instrumentation-pika/tests {posargs}