From 105cf3c3e8c034459774e41c82df573bed50b8a4 Mon Sep 17 00:00:00 2001 From: Aaron Abbott Date: Mon, 13 Jan 2025 21:18:45 +0000 Subject: [PATCH] Record spans and add tests for Vertex AI instrumentation --- .../instrumentation/vertexai/__init__.py | 23 ++- .../instrumentation/vertexai/patch.py | 101 +++++++++++ .../instrumentation/vertexai/utils.py | 162 ++++++++++++++++++ ...at_completion_extra_call_level_params.yaml | 82 +++++++++ ..._completion_extra_client_level_params.yaml | 82 +++++++++ .../test_vertexai_generate_content.yaml | 70 ++++++++ .../tests/conftest.py | 136 +++++++++++++-- .../tests/test_chat_completions.py | 110 ++++++++++++ .../tests/test_placeholder.py | 20 --- 9 files changed, 746 insertions(+), 40 deletions(-) create mode 100644 instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/utils.py create mode 100644 instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_chat_completion_extra_call_level_params.yaml create mode 100644 instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_chat_completion_extra_client_level_params.yaml create mode 100644 instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_vertexai_generate_content.yaml create mode 100644 instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_chat_completions.py delete mode 100644 instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_placeholder.py 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..7f7b88ff42 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,29 @@ 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="vertexai.generative_models._generative_models", + # Patching this base class also instruments the vertexai.preview.generative_models + # package + name="_GenerativeModel.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..152d24acca 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,104 @@ # 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, Optional + +from opentelemetry._events import EventLogger +from opentelemetry.instrumentation.vertexai.utils import ( + GenerateContentParams, + get_genai_request_attributes, + get_span_name, + handle_span_exception, +) +from opentelemetry.trace import SpanKind, Tracer + +if TYPE_CHECKING: + from vertexai.generative_models import ( + GenerationResponse, + Tool, + ToolConfig, + ) + from vertexai.generative_models._generative_models import ( + ContentsType, + GenerationConfigType, + SafetySettingsType, + _GenerativeModel, + ) + + +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, + ): + # Use exact parameter signature to handle named vs positional args robustly + def extract_params( + contents: ContentsType, + *, + 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, + ) -> GenerateContentParams: + return GenerateContentParams( + contents=contents, + generation_config=generation_config, + safety_settings=safety_settings, + tools=tools, + tool_config=tool_config, + labels=labels, + stream=stream, + ) + + params = extract_params(*args, **kwargs) + + span_attributes = get_genai_request_attributes(instance, params) + + span_name = get_span_name(span_attributes) + with tracer.start_as_current_span( + name=span_name, + kind=SpanKind.CLIENT, + attributes=span_attributes, + end_on_exit=False, + ) 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) + # ) + + try: + 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 + # ) + span.end() + return result + + except Exception as error: + handle_span_exception(span, error) + raise + + 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..aab4c9055c --- /dev/null +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/utils.py @@ -0,0 +1,162 @@ +# 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 + +from dataclasses import dataclass +from os import environ +from typing import ( + TYPE_CHECKING, + Dict, + List, + Mapping, + Optional, + TypedDict, + cast, +) + +from opentelemetry.semconv._incubating.attributes import ( + gen_ai_attributes as GenAIAttributes, +) +from opentelemetry.semconv.attributes import ( + error_attributes as ErrorAttributes, +) +from opentelemetry.trace import Span +from opentelemetry.trace.status import Status, StatusCode +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, + ) + + +@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 + + +def get_genai_request_attributes( + # TODO: use types + instance: _GenerativeModel, + params: GenerateContentParams, + operation_name: GenAIAttributes.GenAiOperationNameValues = GenAIAttributes.GenAiOperationNameValues.CHAT, +): + model = _get_model_name(instance) + generation_config = _get_generation_config(instance, params) + attributes = { + 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/" + + +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 + + +# TODO: Everything below here should be replaced with +# opentelemetry.instrumentation.genai_utils instead once it is released. +# https://github.com/open-telemetry/opentelemetry-python-contrib/issues/3191 + +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}" + + +def handle_span_exception(span: Span, error: Exception): + span.set_status(Status(StatusCode.ERROR, str(error))) + if span.is_recording(): + span.set_attribute( + ErrorAttributes.ERROR_TYPE, type(error).__qualname__ + ) + span.end() diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_chat_completion_extra_call_level_params.yaml b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_chat_completion_extra_call_level_params.yaml new file mode 100644 index 0000000000..7a2326c15e --- /dev/null +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_chat_completion_extra_call_level_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.006723951548337936 + } + ], + "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_chat_completion_extra_client_level_params.yaml b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_chat_completion_extra_client_level_params.yaml new file mode 100644 index 0000000000..7a2326c15e --- /dev/null +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_chat_completion_extra_client_level_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.006723951548337936 + } + ], + "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_vertexai_generate_content.yaml b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_vertexai_generate_content.yaml new file mode 100644 index 0000000000..9a32b09288 --- /dev/null +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_vertexai_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.0054498989331094845 + } + ], + "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/conftest.py b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/conftest.py index 8337188ece..e1b12866e9 100644 --- a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/conftest.py +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/conftest.py @@ -1,16 +1,34 @@ """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 ( @@ -30,6 +48,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 +70,104 @@ 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: + # Unfortunately I couldn't find a nice way to globally reset the global_config for each + # test because different vertex submodules reference the global instance directly + # https://github.com/googleapis/python-aiplatform/blob/v1.74.0/google/cloud/aiplatform/initializer.py#L687 + # so this config will leak if we don't call init() for each test. + + # 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. + if vcr.record_mode is RecordMode.NONE: + vertexai.init(credentials=AnonymousCredentials()) + else: + vertexai.init() + + +@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 +236,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..931d09b756 --- /dev/null +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_chat_completions.py @@ -0,0 +1,110 @@ +import pytest +from vertexai.generative_models import ( + Content, + GenerationConfig, + GenerativeModel, + Part, +) + +from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + + +@pytest.mark.vcr +def test_vertexai_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_chat_completion_extra_client_level_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.openai.request.seed": 12345, + "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.2, + "gen_ai.request.top_p": 0.95, + "gen_ai.system": "vertex_ai", + } + + +@pytest.mark.vcr() +def test_chat_completion_extra_call_level_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.openai.request.seed": 12345, + "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.2, + "gen_ai.request.top_p": 0.95, + "gen_ai.system": "vertex_ai", + } 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