Improve test coverage

Author: DouweM

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import dataclasses
+import hashlib
 from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar
@@ -546,6 +547,13 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT
     )
 
 
+def multi_modal_content_identifier(identifier: str | bytes) -> str:
+    """Generate stable identifier for multi-modal content to help LLM in finding a specific file in tool call responses."""
+    if isinstance(identifier, str):
+        identifier = identifier.encode('utf-8')
+    return hashlib.sha1(identifier).hexdigest()[:6]
+
+
 async def process_function_tools(  # noqa C901
     tool_calls: list[_messages.ToolCallPart],
     output_tool_name: str | None,
@@ -671,7 +679,11 @@ async def process_function_tools(  # noqa C901
                     processed_contents: list[Any] = []
                     for content in contents:
                         if isinstance(content, _messages.MultiModalContentTypes):
-                            identifier = content.identifier
+                            if isinstance(content, _messages.BinaryContent):
+                                identifier = multi_modal_content_identifier(content.data)
+                            else:
+                                identifier = multi_modal_content_identifier(content.url)
+
                             user_parts.append(
                                 _messages.UserPromptPart(
                                     content=[f'This is file {identifier}:', content],

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -96,6 +96,9 @@ async def call_tool(
 
         Returns:
             The result of the tool call.
+
+        Raises:
+            ModelRetry: If the tool call fails.
         """
         result = await self._client.call_tool(tool_name, arguments)
 
@@ -131,27 +134,20 @@ async def call_tool(
             else:
                 assert_never(part)
 
-        if result.isError:
-            raise ModelRetry('\n'.join(text_parts) or 'Unknown error')
-
-        if text_parts and not binary_parts and not json_parts:
-            return '\n'.join(text_parts)
+        text = '\n'.join(text_parts)
 
-        if json_parts and not text_parts and not binary_parts:
-            if len(json_parts) == 1:
-                return json_parts[0]
-            return json_parts
+        if result.isError:
+            raise ModelRetry(text or 'Unknown error')
 
-        if binary_parts and not text_parts and not json_parts:
-            if len(binary_parts) == 1:
-                return binary_parts[0]
-            return binary_parts
+        parts: list[str | BinaryContent | dict[str, Any] | list[Any]] = []
+        if text:
+            parts.append(text)
+        parts.extend(json_parts)
+        parts.extend(binary_parts)
 
-        return [
-            *text_parts,
-            *json_parts,
-            *binary_parts,
-        ]
+        if len(parts) == 1:
+            return parts[0]
+        return parts
 
     async def __aenter__(self) -> Self:
         self._exit_stack = AsyncExitStack()

pydantic_ai_slim/pydantic_ai/messages.py

@@ -1,6 +1,5 @@
 from __future__ import annotations as _annotations
 
-import hashlib
 import uuid
 from collections.abc import Sequence
 from dataclasses import dataclass, field, replace
@@ -45,12 +44,6 @@
 VideoFormat: TypeAlias = Literal['mkv', 'mov', 'mp4', 'webm', 'flv', 'mpeg', 'mpg', 'wmv', 'three_gp']
 
 
-def _multi_modal_content_identifier(identifier: str | bytes) -> str:
-    if isinstance(identifier, str):
-        identifier = identifier.encode('utf-8')
-    return hashlib.sha1(identifier).hexdigest()[:6]
-
-
 @dataclass
 class SystemPromptPart:
     """A system prompt, generally written by the application developer.
@@ -87,10 +80,6 @@ class VideoUrl:
     kind: Literal['video-url'] = 'video-url'
     """Type identifier, this is available on all parts as a discriminator."""
 
-    @property
-    def identifier(self) -> str:
-        return _multi_modal_content_identifier(self.url)
-
     @property
     def media_type(self) -> VideoMediaType:  # pragma: no cover
         """Return the media type of the video, based on the url."""
@@ -132,10 +121,6 @@ class AudioUrl:
     kind: Literal['audio-url'] = 'audio-url'
     """Type identifier, this is available on all parts as a discriminator."""
 
-    @property
-    def identifier(self) -> str:
-        return _multi_modal_content_identifier(self.url)
-
     @property
     def media_type(self) -> AudioMediaType:
         """Return the media type of the audio file, based on the url."""
@@ -157,10 +142,6 @@ class ImageUrl:
     kind: Literal['image-url'] = 'image-url'
     """Type identifier, this is available on all parts as a discriminator."""
 
-    @property
-    def identifier(self) -> str:
-        return _multi_modal_content_identifier(self.url)
-
     @property
     def media_type(self) -> ImageMediaType:
         """Return the media type of the image, based on the url."""
@@ -194,10 +175,6 @@ class DocumentUrl:
     kind: Literal['document-url'] = 'document-url'
     """Type identifier, this is available on all parts as a discriminator."""
 
-    @property
-    def identifier(self) -> str:
-        return _multi_modal_content_identifier(self.url)
-
     @property
     def media_type(self) -> str:
         """Return the media type of the document, based on the url."""
@@ -228,10 +205,6 @@ class BinaryContent:
     kind: Literal['binary'] = 'binary'
     """Type identifier, this is available on all parts as a discriminator."""
 
-    @property
-    def identifier(self) -> str:
-        return _multi_modal_content_identifier(self.data)
-
     @property
     def is_audio(self) -> bool:
         """Return `True` if the media type is an audio type."""