Support different content inputs in TestModel (#1015)

Author: Kludex

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -282,7 +282,16 @@ def _estimate_string_tokens(content: str | Sequence[UserContent]) -> int:
         return 0
     if isinstance(content, str):
         return len(re.split(r'[\s",.:]+', content.strip()))
-    # TODO(Marcelo): We need to study how we can estimate the tokens for these types of content.
     else:  # pragma: no cover
-        assert isinstance(content, (AudioUrl, ImageUrl, BinaryContent))
-        return 0
+        tokens = 0
+        for part in content:
+            if isinstance(part, str):
+                tokens += len(re.split(r'[\s",.:]+', part.strip()))
+            # TODO(Marcelo): We need to study how we can estimate the tokens for these types of content.
+            if isinstance(part, (AudioUrl, ImageUrl)):
+                tokens += 0
+            elif isinstance(part, BinaryContent):
+                tokens += len(part.data)
+            else:
+                tokens += 0
+        return tokens

tests/models/test_model_test.py

@@ -13,6 +13,9 @@
 from pydantic_ai import Agent, ModelRetry, RunContext
 from pydantic_ai.exceptions import UnexpectedModelBehavior
 from pydantic_ai.messages import (
+    AudioUrl,
+    BinaryContent,
+    ImageUrl,
     ModelRequest,
     ModelResponse,
     RetryPromptPart,
@@ -22,6 +25,7 @@
     UserPromptPart,
 )
 from pydantic_ai.models.test import TestModel, _chars, _JsonSchemaTestData  # pyright: ignore[reportPrivateUsage]
+from pydantic_ai.usage import Usage
 
 from ..conftest import IsNow
 
@@ -271,3 +275,18 @@ def test_max_items():
     }
     data = _JsonSchemaTestData(json_schema).generate()
     assert data == snapshot([])
+
+
+@pytest.mark.parametrize(
+    'content',
+    [
+        AudioUrl(url='https://example.com'),
+        ImageUrl(url='https://example.com'),
+        BinaryContent(data=b'', media_type='image/png'),
+    ],
+)
+def test_different_content_input(content: AudioUrl | ImageUrl | BinaryContent):
+    agent = Agent()
+    result = agent.run_sync('x', model=TestModel(custom_result_text='custom'))
+    assert result.data == snapshot('custom')
+    assert result.usage() == snapshot(Usage(requests=1, request_tokens=51, response_tokens=1, total_tokens=52))