Add vcr tests for vertexai

Author: aabmass

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.00556742831280357
+            }
+          ],
+          "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:
+      - '451'
+    status:
+      code: 200
+      message: OK
+version: 1

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.utils import (
+    OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT,
+)
+from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor
 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,
 )
 
+# from opentelemetry.instrumentation.vertexai_v2 import VertexAIInstrumentor
+
 
 @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,110 @@ 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(scope="function", name="metric_reader")
+def fixture_metric_reader():
+    exporter = InMemoryMetricReader()
+    yield exporter
+
+
+@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 +244,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

instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_gemini.py

@@ -0,0 +1,37 @@
+from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
+    InMemorySpanExporter,
+)
+import pytest
+from vertexai.generative_models import Content, GenerativeModel, Part
+
+# from opentelemetry.semconv_ai import SpanAttributes
+
+
+@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 [span.name for span in spans] == ["chat gemini-1.5-flash-002"]
+
+    vertexai_span = spans[0]
+    assert len(spans) == 1
+
+    assert dict(vertexai_span.attributes) == {
+        "gen_ai.operation.name": "chat",
+        "gen_ai.request.model": "gemini-1.5-flash-002",
+        "gen_ai.system": "vertex_ai",
+    }
+    # missing some e.g.
+    # "gen_ai.response.model": "gemini-pro-vision",
+    # "gen_ai.usage.output_tokens": 35,
+    # "gen_ai.usage.input_tokens": 265,