enforce TextOutput deps type checking via OutputSpec

Author: dsfaccini

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -133,7 +133,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
     be merged with this value, with the runtime argument taking priority.
     """
 
-    _output_type: OutputSpec[OutputDataT]
+    _output_type: OutputSpec[OutputDataT, AgentDepsT]
 
     instrument: InstrumentationSettings | bool | None
     """Options to automatically instrument with OpenTelemetry."""
@@ -170,7 +170,7 @@ def __init__(
         self,
         model: models.Model | models.KnownModelName | str | None = None,
         *,
-        output_type: OutputSpec[OutputDataT] = str,
+        output_type: OutputSpec[OutputDataT, AgentDepsT] = str,
         instructions: Instructions[AgentDepsT] = None,
         system_prompt: str | Sequence[str] = (),
         deps_type: type[AgentDepsT] = NoneType,
@@ -198,7 +198,7 @@ def __init__(
         self,
         model: models.Model | models.KnownModelName | str | None = None,
         *,
-        output_type: OutputSpec[OutputDataT] = str,
+        output_type: OutputSpec[OutputDataT, AgentDepsT] = str,
         instructions: Instructions[AgentDepsT] = None,
         system_prompt: str | Sequence[str] = (),
         deps_type: type[AgentDepsT] = NoneType,
@@ -224,7 +224,7 @@ def __init__(
         self,
         model: models.Model | models.KnownModelName | str | None = None,
         *,
-        output_type: OutputSpec[OutputDataT] = str,
+        output_type: OutputSpec[OutputDataT, AgentDepsT] = str,
         instructions: Instructions[AgentDepsT] = None,
         system_prompt: str | Sequence[str] = (),
         deps_type: type[AgentDepsT] = NoneType,
@@ -419,7 +419,7 @@ def deps_type(self) -> type:
         return self._deps_type
 
     @property
-    def output_type(self) -> OutputSpec[OutputDataT]:
+    def output_type(self) -> OutputSpec[OutputDataT, AgentDepsT]:
         """The type of data output by agent runs, used to validate the data returned by the model, defaults to `str`."""
         return self._output_type
 

pydantic_ai_slim/pydantic_ai/output.py

@@ -35,6 +35,8 @@
 T = TypeVar('T')
 T_co = TypeVar('T_co', covariant=True)
 TextOutputAgentDepsT = TypeVar('TextOutputAgentDepsT', default=None, contravariant=True)
+# default=Any for backward compat
+OutputSpecDepsT = TypeVar('OutputSpecDepsT', default=Any)
 
 OutputDataT = TypeVar('OutputDataT', default=str, covariant=True)
 """Covariant type variable for the output data type of a run."""
@@ -356,14 +358,18 @@ def __get_pydantic_json_schema__(
 
 _OutputSpecItem = TypeAliasType(
     '_OutputSpecItem',
-    OutputTypeOrFunction[T_co] | ToolOutput[T_co] | NativeOutput[T_co] | PromptedOutput[T_co] | TextOutput[T_co, Any],
-    type_params=(T_co,),
+    OutputTypeOrFunction[T_co]
+    | ToolOutput[T_co]
+    | NativeOutput[T_co]
+    | PromptedOutput[T_co]
+    | TextOutput[T_co, OutputSpecDepsT],
+    type_params=(T_co, OutputSpecDepsT),
 )
 
 OutputSpec = TypeAliasType(
     'OutputSpec',
-    _OutputSpecItem[T_co] | Sequence['OutputSpec[T_co]'],
-    type_params=(T_co,),
+    _OutputSpecItem[T_co, OutputSpecDepsT] | Sequence['OutputSpec[T_co, OutputSpecDepsT]'],
+    type_params=(T_co, OutputSpecDepsT),
 )
 """Specification of the agent's output data.
 

tests/typed_agent.py

@@ -292,10 +292,8 @@ def my_method(self) -> bool:
 assert_type(text_output_with_ctx, TextOutput[re.Pattern[str], int])
 Agent('test', output_type=text_output_with_ctx, deps_type=int)
 Agent('test', output_type=text_output_with_ctx, deps_type=bool)  # bool is subclass of int, works with contravariant
-# NOTE: The following don't produce type errors because _OutputSpecItem uses TextOutput[T_co, Any]
-# which erases the deps type constraint.
-Agent('test', output_type=text_output_with_ctx, deps_type=str)
-Agent('test', output_type=text_output_with_ctx)
+Agent('test', output_type=text_output_with_ctx, deps_type=str)  # pyright: ignore[reportArgumentType,reportCallIssue]
+Agent('test', output_type=text_output_with_ctx)  # pyright: ignore[reportArgumentType,reportCallIssue]
 
 # prepare example from docs: