From 005b48a57bba3c676c66310142f50c9611cb252c Mon Sep 17 00:00:00 2001 From: Max Muoto Date: Fri, 2 May 2025 17:20:34 -0500 Subject: [PATCH 1/2] Refactor type guards in Agent class to use TypeIs instead of TypeGuard for improved type narrowing --- pydantic_ai_slim/pydantic_ai/agent.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/agent.py b/pydantic_ai_slim/pydantic_ai/agent.py index 2871fa4e3..c41982234 100644 --- a/pydantic_ai_slim/pydantic_ai/agent.py +++ b/pydantic_ai_slim/pydantic_ai/agent.py @@ -12,7 +12,7 @@ from opentelemetry.trace import NoOpTracer, use_span from pydantic.json_schema import GenerateJsonSchema -from typing_extensions import Literal, Never, TypeGuard, TypeVar, deprecated +from typing_extensions import Literal, Never, TypeIs, TypeVar, deprecated from pydantic_graph import End, Graph, GraphRun, GraphRunContext from pydantic_graph._utils import get_event_loop @@ -1049,6 +1049,7 @@ def override( override_model_before = self._override_model self._override_model = _utils.Some(models.infer_model(model)) else: + model override_model_before = _utils.UNSET try: @@ -1614,7 +1615,7 @@ def _prepare_output_schema( @staticmethod def is_model_request_node( node: _agent_graph.AgentNode[T, S] | End[result.FinalResult[S]], - ) -> TypeGuard[_agent_graph.ModelRequestNode[T, S]]: + ) -> TypeIs[_agent_graph.ModelRequestNode[T, S]]: """Check if the node is a `ModelRequestNode`, narrowing the type if it is. This method preserves the generic parameters while narrowing the type, unlike a direct call to `isinstance`. @@ -1624,7 +1625,7 @@ def is_model_request_node( @staticmethod def is_call_tools_node( node: _agent_graph.AgentNode[T, S] | End[result.FinalResult[S]], - ) -> TypeGuard[_agent_graph.CallToolsNode[T, S]]: + ) -> TypeIs[_agent_graph.CallToolsNode[T, S]]: """Check if the node is a `CallToolsNode`, narrowing the type if it is. This method preserves the generic parameters while narrowing the type, unlike a direct call to `isinstance`. @@ -1634,7 +1635,7 @@ def is_call_tools_node( @staticmethod def is_user_prompt_node( node: _agent_graph.AgentNode[T, S] | End[result.FinalResult[S]], - ) -> TypeGuard[_agent_graph.UserPromptNode[T, S]]: + ) -> TypeIs[_agent_graph.UserPromptNode[T, S]]: """Check if the node is a `UserPromptNode`, narrowing the type if it is. This method preserves the generic parameters while narrowing the type, unlike a direct call to `isinstance`. @@ -1644,7 +1645,7 @@ def is_user_prompt_node( @staticmethod def is_end_node( node: _agent_graph.AgentNode[T, S] | End[result.FinalResult[S]], - ) -> TypeGuard[End[result.FinalResult[S]]]: + ) -> TypeIs[End[result.FinalResult[S]]]: """Check if the node is a `End`, narrowing the type if it is. This method preserves the generic parameters while narrowing the type, unlike a direct call to `isinstance`. From a30c188a31ce9367494405d15b3f96846fb3152a Mon Sep 17 00:00:00 2001 From: Max Muoto Date: Fri, 2 May 2025 17:21:00 -0500 Subject: [PATCH 2/2] Remove --- pydantic_ai_slim/pydantic_ai/agent.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pydantic_ai_slim/pydantic_ai/agent.py b/pydantic_ai_slim/pydantic_ai/agent.py index c41982234..ce0fdea70 100644 --- a/pydantic_ai_slim/pydantic_ai/agent.py +++ b/pydantic_ai_slim/pydantic_ai/agent.py @@ -1049,7 +1049,6 @@ def override( override_model_before = self._override_model self._override_model = _utils.Some(models.infer_model(model)) else: - model override_model_before = _utils.UNSET try: