Use isinstance checks for message kinds (#252)

Author: sydney-runkle

pydantic_ai_slim/pydantic_ai/agent.py

@@ -20,7 +20,6 @@
     models,
     result,
 )
-from .messages import TextPart, ToolCallPart
 from .result import ResultData
 from .settings import ModelSettings, merge_model_settings
 from .tools import (
@@ -795,7 +794,7 @@ async def _prepare_messages(
         self, deps: AgentDeps, user_prompt: str, message_history: list[_messages.Message] | None
     ) -> tuple[int, list[_messages.Message]]:
         # if message history includes system prompts, we don't want to regenerate them
-        if message_history and any(m.message_kind == 'system-prompt' for m in message_history):
+        if message_history and any(isinstance(m, _messages.SystemPrompt) for m in message_history):
             # shallow copy messages
             messages = message_history.copy()
         else:
@@ -816,9 +815,9 @@ async def _handle_model_response(
             A tuple of `(final_result, messages)`. If `final_result` is not `None`, the conversation should end.
         """
         texts: list[str] = []
-        tool_calls: list[ToolCallPart] = []
+        tool_calls: list[_messages.ToolCallPart] = []
         for item in model_response.parts:
-            if isinstance(item, TextPart):
+            if isinstance(item, _messages.TextPart):
                 texts.append(item.content)
             else:
                 tool_calls.append(item)
@@ -852,7 +851,7 @@ async def _handle_text_response(
             return None, [response]
 
     async def _handle_structured_response(
-        self, tool_calls: list[ToolCallPart], deps: AgentDeps
+        self, tool_calls: list[_messages.ToolCallPart], deps: AgentDeps
     ) -> tuple[_MarkFinalResult[ResultData] | None, list[_messages.Message]]:
         """Handle a structured response containing tool calls from the model for non-streaming responses."""
         assert tool_calls, 'Expected at least one tool call'
@@ -870,7 +869,7 @@ async def _handle_structured_response(
 
     async def _process_final_tool_calls(
         self,
-        tool_calls: list[ToolCallPart],
+        tool_calls: list[_messages.ToolCallPart],
         deps: AgentDeps,
     ) -> tuple[_MarkFinalResult[ResultData] | None, list[_messages.Message]]:
         """Process any final result tool calls and return the first valid result."""
@@ -916,7 +915,7 @@ async def _process_final_tool_calls(
 
     async def _process_function_tools(
         self,
-        tool_calls: list[ToolCallPart],
+        tool_calls: list[_messages.ToolCallPart],
         deps: AgentDeps,
     ) -> list[_messages.Message]:
         """Process function (non-final) tool calls in parallel."""
@@ -1013,7 +1012,7 @@ async def _handle_streamed_model_response(
             # we now run all tool functions in parallel
             tasks: list[asyncio.Task[_messages.Message]] = []
             for item in structured_msg.parts:
-                if isinstance(item, ToolCallPart):
+                if isinstance(item, _messages.ToolCallPart):
                     call = item
                     if tool := self._function_tools.get(call.tool_name):
                         tasks.append(asyncio.create_task(tool.run(deps, call), name=call.tool_name))

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -196,7 +196,7 @@ async def _messages_create(
         anthropic_messages: list[MessageParam] = []
 
         for m in messages:
-            if m.message_kind == 'system-prompt':
+            if isinstance(m, SystemPrompt):
                 system_prompt += m.content
             else:
                 anthropic_messages.append(self._map_message(m))

tests/models/test_model_function.py

@@ -82,7 +82,7 @@ async def weather_model(messages: list[Message], info: AgentInfo) -> ModelRespon
     assert info.allow_text_result
     assert {t.name for t in info.function_tools} == {'get_location', 'get_weather'}
     last = messages[-1]
-    if last.message_kind == 'user-prompt':
+    if isinstance(last, UserPrompt):
         return ModelResponse(
             parts=[
                 ToolCallPart.from_json(
@@ -91,11 +91,11 @@ async def weather_model(messages: list[Message], info: AgentInfo) -> ModelRespon
                 )
             ]
         )
-    elif last.message_kind == 'tool-return':
+    elif isinstance(last, ToolReturn):
         if last.tool_name == 'get_location':
             return ModelResponse(parts=[ToolCallPart.from_json('get_weather', last.model_response_str())])
         elif last.tool_name == 'get_weather':
-            location_name = next(m.content for m in messages if m.message_kind == 'user-prompt')
+            location_name = next(m.content for m in messages if isinstance(m, UserPrompt))
             return ModelResponse.from_text(f'{last.content} in {location_name}')
 
     raise ValueError(f'Unexpected message: {last}')
@@ -177,7 +177,7 @@ def test_weather(set_event_loop: None):
 
 async def call_function_model(messages: list[Message], _: AgentInfo) -> ModelResponse:  # pragma: no cover
     last = messages[-1]
-    if last.message_kind == 'user-prompt':
+    if isinstance(last, UserPrompt):
         if last.content.startswith('{'):
             details = json.loads(last.content)
             return ModelResponse(
@@ -188,7 +188,7 @@ async def call_function_model(messages: list[Message], _: AgentInfo) -> ModelRes
                     )
                 ]
             )
-    elif last.message_kind == 'tool-return':
+    elif isinstance(last, ToolReturn):
         return ModelResponse.from_text(pydantic_core.to_json(last).decode())
 
     raise ValueError(f'Unexpected message: {last}')

tests/test_examples.py

@@ -22,7 +22,10 @@
     ArgsDict,
     Message,
     ModelResponse,
+    RetryPrompt,
     ToolCallPart,
+    ToolReturn,
+    UserPrompt,
 )
 from pydantic_ai.models import KnownModelName, Model
 from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel
@@ -215,7 +218,7 @@ async def async_http_request(url: str, **kwargs: Any) -> httpx.Response:
 
 async def model_logic(messages: list[Message], info: AgentInfo) -> ModelResponse:  # pragma: no cover
     m = messages[-1]
-    if m.message_kind == 'user-prompt':
+    if isinstance(m, UserPrompt):
         if response := text_responses.get(m.content):
             if isinstance(response, str):
                 return ModelResponse.from_text(content=response)
@@ -225,28 +228,28 @@ async def model_logic(messages: list[Message], info: AgentInfo) -> ModelResponse
         if re.fullmatch(r'sql prompt \d+', m.content):
             return ModelResponse.from_text(content='SELECT 1')
 
-    elif m.message_kind == 'tool-return' and m.tool_name == 'roulette_wheel':
+    elif isinstance(m, ToolReturn) and m.tool_name == 'roulette_wheel':
         win = m.content == 'winner'
         return ModelResponse(parts=[ToolCallPart(tool_name='final_result', args=ArgsDict({'response': win}))])
-    elif m.message_kind == 'tool-return' and m.tool_name == 'roll_die':
+    elif isinstance(m, ToolReturn) and m.tool_name == 'roll_die':
         return ModelResponse(parts=[ToolCallPart(tool_name='get_player_name', args=ArgsDict({}))])
-    elif m.message_kind == 'tool-return' and m.tool_name == 'get_player_name':
+    elif isinstance(m, ToolReturn) and m.tool_name == 'get_player_name':
         return ModelResponse.from_text(content="Congratulations Anne, you guessed correctly! You're a winner!")
     if (
-        m.message_kind == 'retry-prompt'
+        isinstance(m, RetryPrompt)
         and isinstance(m.content, str)
         and m.content.startswith("No user found with name 'Joh")
     ):
         return ModelResponse(parts=[ToolCallPart(tool_name='get_user_by_name', args=ArgsDict({'name': 'John Doe'}))])
-    elif m.message_kind == 'tool-return' and m.tool_name == 'get_user_by_name':
+    elif isinstance(m, ToolReturn) and m.tool_name == 'get_user_by_name':
         args = {
             'message': 'Hello John, would you be free for coffee sometime next week? Let me know what works for you!',
             'user_id': 123,
         }
         return ModelResponse(parts=[ToolCallPart(tool_name='final_result', args=ArgsDict(args))])
-    elif m.message_kind == 'retry-prompt' and m.tool_name == 'calc_volume':
+    elif isinstance(m, RetryPrompt) and m.tool_name == 'calc_volume':
         return ModelResponse(parts=[ToolCallPart(tool_name='calc_volume', args=ArgsDict({'size': 6}))])
-    elif m.message_kind == 'tool-return' and m.tool_name == 'customer_balance':
+    elif isinstance(m, ToolReturn) and m.tool_name == 'customer_balance':
         args = {
             'support_advice': 'Hello John, your current account balance, including pending transactions, is $123.45.',
             'block_card': False,
@@ -262,7 +265,7 @@ async def stream_model_logic(
     messages: list[Message], info: AgentInfo
 ) -> AsyncIterator[str | DeltaToolCalls]:  # pragma: no cover
     m = messages[-1]
-    if m.message_kind == 'user-prompt':
+    if isinstance(m, UserPrompt):
         if response := text_responses.get(m.content):
             if isinstance(response, str):
                 words = response.split(' ')