Pretty format yaml cassettes

Author: aabmass

instrumentation-genai/opentelemetry-instrumentation-vertexai-v2/tests/cassettes/test_vertexai_generate_content.yaml

@@ -1,10 +1,24 @@
 interactions:
 - request:
-    body: "{\n  \"contents\": [\n    {\n      \"role\": \"user\",\n      \"parts\":
-      [\n        {\n          \"fileData\": {\n            \"mimeType\": \"image/jpeg\",\n
-      \           \"fileUri\": \"gs://generativeai-downloads/images/scones.jpg\"\n
-      \         }\n        },\n        {\n          \"text\": \"what is shown in this
-      image?\"\n        }\n      ]\n    }\n  ]\n}"
+    body: |-
+      {
+        "contents": [
+          {
+            "role": "user",
+            "parts": [
+              {
+                "fileData": {
+                  "mimeType": "image/jpeg",
+                  "fileUri": "gs://generativeai-downloads/images/scones.jpg"
+                }
+              },
+              {
+                "text": "what is shown in this image?"
+              }
+            ]
+          }
+        ]
+      }
     headers:
       Accept:
       - '*/*'
@@ -22,26 +36,60 @@ interactions:
     uri: https://us-central1-aiplatform.googleapis.com/v1beta1/projects/fake-project/locations/us-central1/publishers/google/models/gemini-pro-vision:generateContent?%24alt=json%3Benum-encoding%3Dint
   response:
     body:
-      string: "{\n  \"candidates\": [\n    {\n      \"content\": {\n        \"role\":
-        \"model\",\n        \"parts\": [\n          {\n            \"text\": \" The
-        image shows a table with a cup of coffee, a bowl of blueberries, and a plate
-        of scones. There are also some flowers on the table.\"\n          }\n        ]\n
-        \     },\n      \"finishReason\": 1,\n      \"safetyRatings\": [\n        {\n
-        \         \"category\": 1,\n          \"probability\": 1,\n          \"probabilityScore\":
-        0.02331543,\n          \"severity\": 1,\n          \"severityScore\": 0.05493164\n
-        \       },\n        {\n          \"category\": 2,\n          \"probability\":
-        1,\n          \"probabilityScore\": 0.026367188,\n          \"severity\":
-        1,\n          \"severityScore\": 0.05493164\n        },\n        {\n          \"category\":
-        3,\n          \"probability\": 1,\n          \"probabilityScore\": 0.046142578,\n
-        \         \"severity\": 1,\n          \"severityScore\": 0.030639648\n        },\n
-        \       {\n          \"category\": 4,\n          \"probability\": 1,\n          \"probabilityScore\":
-        0.080566406,\n          \"severity\": 1,\n          \"severityScore\": 0.095214844\n
-        \       }\n      ],\n      \"avgLogprobs\": -0.11595650642148909\n    }\n
-        \ ],\n  \"usageMetadata\": {\n    \"promptTokenCount\": 265,\n    \"candidatesTokenCount\":
-        31,\n    \"totalTokenCount\": 296\n  },\n  \"modelVersion\": \"gemini-pro-vision\"\n}\n"
+      string: |-
+        {
+          "candidates": [
+            {
+              "content": {
+                "role": "model",
+                "parts": [
+                  {
+                    "text": " The image shows a table with a cup of coffee, a bowl of blueberries, and a plate of scones with blueberries on top. There are also pink flowers on the table."
+                  }
+                ]
+              },
+              "finishReason": 1,
+              "safetyRatings": [
+                {
+                  "category": 1,
+                  "probability": 1,
+                  "probabilityScore": 0.025512695,
+                  "severity": 1,
+                  "severityScore": 0.06933594
+                },
+                {
+                  "category": 2,
+                  "probability": 1,
+                  "probabilityScore": 0.026367188,
+                  "severity": 1,
+                  "severityScore": 0.07080078
+                },
+                {
+                  "category": 3,
+                  "probability": 1,
+                  "probabilityScore": 0.041503906,
+                  "severity": 1,
+                  "severityScore": 0.03466797
+                },
+                {
+                  "category": 4,
+                  "probability": 1,
+                  "probabilityScore": 0.091308594,
+                  "severity": 1,
+                  "severityScore": 0.09033203
+                }
+              ],
+              "avgLogprobs": -0.09557106835501535
+            }
+          ],
+          "usageMetadata": {
+            "promptTokenCount": 265,
+            "candidatesTokenCount": 35,
+            "totalTokenCount": 300
+          },
+          "modelVersion": "gemini-pro-vision"
+        }
     headers:
-      Cache-Control:
-      - private
       Content-Type:
       - application/json; charset=UTF-8
       Transfer-Encoding:
@@ -51,7 +99,7 @@ interactions:
       - X-Origin
       - Referer
       content-length:
-      - '1275'
+      - '1299'
     status:
       code: 200
       message: OK

instrumentation-genai/opentelemetry-instrumentation-vertexai-v2/tests/conftest.py

@@ -1,10 +1,12 @@
 """Unit tests configuration module."""
 
+import json
 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
@@ -97,3 +99,70 @@ def before_response_cb(response: MutableMapping[str, Any]):
         "before_record_response": before_response_cb,
         "ignore_hosts": ["oauth2.googleapis.com"],
     }
+
+
+class LiteralBlockScalar(str):
+    """Formats the string as a literal block scalar, preserving whitespace and
+    without interpreting escape characters"""
+
+
+def literal_block_scalar_presenter(dumper, data):
+    """Represents a scalar string as a literal block, via '|' syntax"""
+    return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
+
+
+yaml.add_representer(LiteralBlockScalar, literal_block_scalar_presenter)
+
+
+def process_string_value(string_value):
+    """Pretty-prints JSON or returns long strings as a LiteralBlockScalar"""
+    try:
+        json_data = json.loads(string_value)
+        return LiteralBlockScalar(json.dumps(json_data, indent=2))
+    except (ValueError, TypeError):
+        if len(string_value) > 80:
+            return LiteralBlockScalar(string_value)
+    return string_value
+
+
+def convert_body_to_literal(data):
+    """Searches the data for body strings, attempting to pretty-print JSON"""
+    if isinstance(data, dict):
+        for key, value in data.items():
+            # Handle response body case (e.g., response.body.string)
+            if key == "body" and isinstance(value, dict) and "string" in value:
+                value["string"] = process_string_value(value["string"])
+
+            # Handle request body case (e.g., request.body)
+            elif key == "body" and isinstance(value, str):
+                data[key] = process_string_value(value)
+
+            else:
+                convert_body_to_literal(value)
+
+    elif isinstance(data, list):
+        for idx, choice in enumerate(data):
+            data[idx] = convert_body_to_literal(choice)
+
+    return data
+
+
+class PrettyPrintJSONBody:
+    """This makes request and response body recordings more readable."""
+
+    @staticmethod
+    def serialize(cassette_dict):
+        cassette_dict = convert_body_to_literal(cassette_dict)
+        return yaml.dump(
+            cassette_dict, default_flow_style=False, allow_unicode=True
+        )
+
+    @staticmethod
+    def deserialize(cassette_string):
+        return yaml.load(cassette_string, Loader=yaml.Loader)
+
+
+@pytest.fixture(scope="module", autouse=True)
+def fixture_vcr(vcr):
+    vcr.register_serializer("yaml", PrettyPrintJSONBody)
+    return vcr

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

@@ -30,7 +30,7 @@ def test_vertexai_generate_content(exporter):
         "gen_ai.operation.name": "text_completion",
         "gen_ai.request.model": "gemini-pro-vision",
         "gen_ai.response.model": "gemini-pro-vision",
-        "gen_ai.usage.output_tokens": 31,
+        "gen_ai.usage.output_tokens": 35,
         "gen_ai.usage.input_tokens": 265,
     }