Skip to content

Commit

Permalink
Pretty format yaml cassettes
Browse files Browse the repository at this point in the history
  • Loading branch information
aabmass committed Dec 6, 2024
1 parent b25a9e4 commit 5dfe192
Show file tree
Hide file tree
Showing 3 changed files with 143 additions and 26 deletions.
Original file line number Diff line number Diff line change
@@ -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:
- '*/*'
Expand All @@ -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:
Expand All @@ -51,7 +99,7 @@ interactions:
- X-Origin
- Referer
content-length:
- '1275'
- '1299'
status:
code: 200
message: OK
Expand Down
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}

Expand Down

0 comments on commit 5dfe192

Please sign in to comment.