Add pytest-recording/vcr and use it for a test (#969)

Author: dmontagu

Makefile

@@ -66,6 +66,10 @@ testcov: test ## Run tests and generate a coverage report
 update-examples: ## Update documentation examples
 	uv run -m pytest --update-examples tests/test_examples.py
 
+.PHONY: update-vcr-tests
+update-vcr-tests: ## Update tests using VCR that hit LLM APIs; note you'll need to set API keys as appropriate
+	uv run -m pytest --record-mode=rewrite tests
+
 # `--no-strict` so you can build the docs without insiders packages
 .PHONY: docs
 docs: ## Build the documentation

pydantic_ai_slim/pyproject.toml

@@ -63,6 +63,7 @@ dev = [
     "pytest-examples>=0.0.14",
     "pytest-mock>=3.14.0",
     "pytest-pretty>=1.2.0",
+    "pytest-recording>=0.13.2",
     "diff-cover>=9.2.0",
 ]
 

tests/conftest.py

@@ -17,20 +17,23 @@
 import pytest
 from _pytest.assertion.rewrite import AssertionRewritingHook
 from typing_extensions import TypeAlias
+from vcr import VCR
 
 import pydantic_ai.models
 
-__all__ = 'IsNow', 'IsFloat', 'TestEnv', 'ClientWithHandler', 'try_import'
+__all__ = 'IsDatetime', 'IsFloat', 'IsNow', 'IsStr', 'TestEnv', 'ClientWithHandler', 'try_import'
 
 
 pydantic_ai.models.ALLOW_MODEL_REQUESTS = False
 
 if TYPE_CHECKING:
 
-    def IsNow(*args: Any, **kwargs: Any) -> datetime: ...
+    def IsDatetime(*args: Any, **kwargs: Any) -> datetime: ...
     def IsFloat(*args: Any, **kwargs: Any) -> float: ...
+    def IsNow(*args: Any, **kwargs: Any) -> datetime: ...
+    def IsStr(*args: Any, **kwargs: Any) -> str: ...
 else:
-    from dirty_equals import IsFloat, IsNow as _IsNow
+    from dirty_equals import IsDatetime, IsFloat, IsNow as _IsNow, IsStr
 
     def IsNow(*args: Any, **kwargs: Any):
         # Increase the default value of `delta` to 10 to reduce test flakiness on overburdened machines
@@ -179,3 +182,18 @@ def set_event_loop() -> Iterator[None]:
     asyncio.set_event_loop(new_loop)
     yield
     new_loop.close()
+
+
+def pytest_recording_configure(config: Any, vcr: VCR):
+    from . import json_body_serializer
+
+    vcr.register_serializer('yaml', json_body_serializer)
+
+
+@pytest.fixture(scope='module')
+def vcr_config():
+    return {
+        # Note: additional header filtering is done inside the serializer
+        'filter_headers': ['authorization', 'x-api-key'],
+        'decode_compressed_response': True,
+    }

tests/json_body_serializer.py

@@ -0,0 +1,77 @@
+# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false
+import json
+from typing import TYPE_CHECKING, Any
+
+import yaml
+
+if TYPE_CHECKING:
+    from yaml import Dumper, Loader
+else:
+    try:
+        from yaml import CDumper as Dumper, CLoader as Loader
+    except ImportError:
+        from yaml import Dumper, Loader
+
+FILTERED_HEADER_PREFIXES = ['anthropic-', 'cf-', 'x-']
+FILTERED_HEADERS = {'authorization', 'date', 'request-id', 'server', 'user-agent', 'via'}
+
+
+class LiteralDumper(Dumper):
+    """
+    A custom dumper that will represent multi-line strings using literal style.
+    """
+
+
+def str_presenter(dumper: Dumper, data: str):
+    """If the string contains newlines, represent it as a literal block."""
+    if '\n' in data:
+        return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')
+    return dumper.represent_scalar('tag:yaml.org,2002:str', data)
+
+
+# Register the custom presenter on our dumper
+LiteralDumper.add_representer(str, str_presenter)
+
+
+def deserialize(cassette_string: str):
+    cassette_dict = yaml.load(cassette_string, Loader=Loader)
+    for interaction in cassette_dict['interactions']:
+        for kind, data in interaction.items():
+            parsed_body = data.pop('parsed_body', None)
+            if parsed_body is not None:
+                dumped_body = json.dumps(parsed_body)
+                data['body'] = {'string': dumped_body} if kind == 'response' else dumped_body
+    return cassette_dict
+
+
+def serialize(cassette_dict: Any):
+    for interaction in cassette_dict['interactions']:
+        for _kind, data in interaction.items():
+            headers: dict[str, list[str]] = data.get('headers', {})
+            # make headers lowercase
+            headers = {k.lower(): v for k, v in headers.items()}
+            # filter headers by name
+            headers = {k: v for k, v in headers.items() if k not in FILTERED_HEADERS}
+            # filter headers by prefix
+            headers = {
+                k: v for k, v in headers.items() if not any(k.startswith(prefix) for prefix in FILTERED_HEADER_PREFIXES)
+            }
+            # update headers on source object
+            data['headers'] = headers
+
+            content_type = headers.get('content-type', None)
+            if content_type != ['application/json']:
+                continue
+
+            # Parse the body as JSON
+            body: Any = data.get('body', None)
+            assert body is not None, data
+            if isinstance(body, dict):
+                # Responses will have the body under a field called 'string'
+                body = body.get('string')
+            if body is not None:
+                data['parsed_body'] = json.loads(body)
+                del data['body']
+
+    # Use our custom dumper
+    return yaml.dump(cassette_dict, Dumper=LiteralDumper, allow_unicode=True, width=120)

tests/models/cassettes/test_anthropic/test_multiple_parallel_tool_calls.yaml

@@ -0,0 +1,212 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '793'
+      content-type:
+      - application/json
+      host:
+      - api.anthropic.com
+    method: POST
+    parsed_body:
+      max_tokens: 1024
+      messages:
+      - content:
+        - text: Alice, Bob, Charlie and Daisy are a family. Who is the youngest?
+          type: text
+        role: user
+      model: claude-3-5-haiku-latest
+      stream: false
+      system: "\n    Use the `retrieve_entity_info` tool to get information about a specific person.\n    If you need to use
+        `retrieve_entity_info` to get information about multiple people, try\n    to call them in parallel as much as possible.\n
+        \   Think step by step and then provide a single most probable concise answer.\n    "
+      tool_choice:
+        type: auto
+      tools:
+      - description: Get the knowledge about the given entity.
+        input_schema:
+          additionalProperties: false
+          properties:
+            name:
+              title: Name
+              type: string
+          required:
+          - name
+          type: object
+        name: retrieve_entity_info
+    uri: https://api.anthropic.com/v1/messages
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '835'
+      content-type:
+      - application/json
+      transfer-encoding:
+      - chunked
+    parsed_body:
+      content:
+      - text: Let me retrieve the information about each family member to determine their ages.
+        type: text
+      - id: toolu_01B6ALG3PUWnoTbSG77WfHsW
+        input:
+          name: Alice
+        name: retrieve_entity_info
+        type: tool_use
+      - id: toolu_018ar6FXUarSmz9v81ajKN5q
+        input:
+          name: Bob
+        name: retrieve_entity_info
+        type: tool_use
+      - id: toolu_01DQDAMdPsj6Seitxpc29HZF
+        input:
+          name: Charlie
+        name: retrieve_entity_info
+        type: tool_use
+      - id: toolu_01Vzma2bAahRtnZi69djzStJ
+        input:
+          name: Daisy
+        name: retrieve_entity_info
+        type: tool_use
+      id: msg_01B6EyfNCSVqtFxbXvfmQXaX
+      model: claude-3-5-haiku-20241022
+      role: assistant
+      stop_reason: tool_use
+      stop_sequence: null
+      type: message
+      usage:
+        cache_creation_input_tokens: 0
+        cache_read_input_tokens: 0
+        input_tokens: 429
+        output_tokens: 186
+    status:
+      code: 200
+      message: OK
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '1928'
+      content-type:
+      - application/json
+      host:
+      - api.anthropic.com
+    method: POST
+    parsed_body:
+      max_tokens: 1024
+      messages:
+      - content:
+        - text: Alice, Bob, Charlie and Daisy are a family. Who is the youngest?
+          type: text
+        role: user
+      - content:
+        - text: Let me retrieve the information about each family member to determine their ages.
+          type: text
+        - id: toolu_01B6ALG3PUWnoTbSG77WfHsW
+          input:
+            name: Alice
+          name: retrieve_entity_info
+          type: tool_use
+        - id: toolu_018ar6FXUarSmz9v81ajKN5q
+          input:
+            name: Bob
+          name: retrieve_entity_info
+          type: tool_use
+        - id: toolu_01DQDAMdPsj6Seitxpc29HZF
+          input:
+            name: Charlie
+          name: retrieve_entity_info
+          type: tool_use
+        - id: toolu_01Vzma2bAahRtnZi69djzStJ
+          input:
+            name: Daisy
+          name: retrieve_entity_info
+          type: tool_use
+        role: assistant
+      - content:
+        - content: alice is bob's wife
+          is_error: false
+          tool_use_id: toolu_01B6ALG3PUWnoTbSG77WfHsW
+          type: tool_result
+        - content: bob is alice's husband
+          is_error: false
+          tool_use_id: toolu_018ar6FXUarSmz9v81ajKN5q
+          type: tool_result
+        - content: charlie is alice's son
+          is_error: false
+          tool_use_id: toolu_01DQDAMdPsj6Seitxpc29HZF
+          type: tool_result
+        - content: daisy is bob's daughter and charlie's younger sister
+          is_error: false
+          tool_use_id: toolu_01Vzma2bAahRtnZi69djzStJ
+          type: tool_result
+        role: user
+      model: claude-3-5-haiku-latest
+      stream: false
+      system: "\n    Use the `retrieve_entity_info` tool to get information about a specific person.\n    If you need to use
+        `retrieve_entity_info` to get information about multiple people, try\n    to call them in parallel as much as possible.\n
+        \   Think step by step and then provide a single most probable concise answer.\n    "
+      tool_choice:
+        type: auto
+      tools:
+      - description: Get the knowledge about the given entity.
+        input_schema:
+          additionalProperties: false
+          properties:
+            name:
+              title: Name
+              type: string
+          required:
+          - name
+          type: object
+        name: retrieve_entity_info
+    uri: https://api.anthropic.com/v1/messages
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '618'
+      content-type:
+      - application/json
+      transfer-encoding:
+      - chunked
+    parsed_body:
+      content:
+      - text: |-
+          Based on the retrieved information, we can see the family relationships:
+          - Alice and Bob are married
+          - Charlie is their son
+          - Daisy is their daughter and Charlie's younger sister
+
+          Since Daisy is described as Charlie's younger sister, Daisy is the youngest in this family.
+
+          The answer is: Daisy is the youngest.
+        type: text
+      id: msg_018ApbFsve4yLKjivLPAzqvA
+      model: claude-3-5-haiku-20241022
+      role: assistant
+      stop_reason: end_turn
+      stop_sequence: null
+      type: message
+      usage:
+        cache_creation_input_tokens: 0
+        cache_read_input_tokens: 0
+        input_tokens: 761
+        output_tokens: 78
+    status:
+      code: 200
+      message: OK
+version: 1