pydantic-ai-slim 0.0.46__tar.gz → 0.0.48__tar.gz

{pydantic_ai_slim-0.0.46 → pydantic_ai_slim-0.0.48}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.46
+Version: 0.0.48
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License-Expression: MIT
@@ -29,7 +29,7 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==0.0.46
+Requires-Dist: pydantic-graph==0.0.48
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: anthropic
@@ -41,13 +41,15 @@ Requires-Dist: argcomplete>=3.5.0; extra == 'cli'
 Requires-Dist: prompt-toolkit>=3; extra == 'cli'
 Requires-Dist: rich>=13; extra == 'cli'
 Provides-Extra: cohere
-Requires-Dist: cohere>=5.13.11; extra == 'cohere'
+Requires-Dist: cohere>=5.13.11; (platform_system != 'Emscripten') and extra == 'cohere'
 Provides-Extra: duckduckgo
 Requires-Dist: duckduckgo-search>=7.0.0; extra == 'duckduckgo'
+Provides-Extra: evals
+Requires-Dist: pydantic-evals==0.0.48; extra == 'evals'
 Provides-Extra: groq
 Requires-Dist: groq>=0.15.0; extra == 'groq'
 Provides-Extra: logfire
-Requires-Dist: logfire>=2.3; extra == 'logfire'
+Requires-Dist: logfire>=3.11.0; extra == 'logfire'
 Provides-Extra: mcp
 Requires-Dist: mcp>=1.4.1; (python_version >= '3.10') and extra == 'mcp'
 Provides-Extra: mistral

pydantic_ai_slim-0.0.48/pydantic_ai/__main__.py

@@ -0,0 +1,6 @@
+"""This means `python -m pydantic_ai` should run the CLI."""
+
+from ._cli import app
+
+if __name__ == '__main__':
+    app()

{pydantic_ai_slim-0.0.46 → pydantic_ai_slim-0.0.48}/pydantic_ai/_agent_graph.py

@@ -331,10 +331,7 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
         ctx.state.run_step += 1
 
         model_settings = merge_model_settings(ctx.deps.model_settings, None)
-        with ctx.deps.tracer.start_as_current_span(
-            'preparing model request params', attributes=dict(run_step=ctx.state.run_step)
-        ):
-            model_request_parameters = await _prepare_request_parameters(ctx)
+        model_request_parameters = await _prepare_request_parameters(ctx)
         return model_settings, model_request_parameters
 
     def _finish_handling(
@@ -374,9 +371,8 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
     ) -> Union[ModelRequestNode[DepsT, NodeRunEndT], End[result.FinalResult[NodeRunEndT]]]:  # noqa UP007
         async with self.stream(ctx):
             pass
-
-        assert (next_node := self._next_node) is not None, 'the stream should set `self._next_node` before it ends'
-        return next_node
+        assert self._next_node is not None, 'the stream should set `self._next_node` before it ends'
+        return self._next_node
 
     @asynccontextmanager
     async def stream(
@@ -635,19 +631,24 @@ async def process_function_tools(
                 )
                 output_parts.append(part)
         else:
-            output_parts.append(_unknown_tool(call.tool_name, ctx))
+            output_parts.append(_unknown_tool(call.tool_name, call.tool_call_id, ctx))
 
     if not calls_to_run:
         return
 
     # Run all tool tasks in parallel
     results_by_index: dict[int, _messages.ModelRequestPart] = {}
-    tool_names = [call.tool_name for _, call in calls_to_run]
     with ctx.deps.tracer.start_as_current_span(
-        'running tools', attributes={'tools': tool_names, 'logfire.msg': f'running tools: {", ".join(tool_names)}'}
+        'running tools',
+        attributes={
+            'tools': [call.tool_name for _, call in calls_to_run],
+            'logfire.msg': f'running {len(calls_to_run)} tool{"" if len(calls_to_run) == 1 else "s"}',
+        },
     ):
-        # TODO: Should we wrap each individual tool call in a dedicated span?
-        tasks = [asyncio.create_task(tool.run(call, run_context), name=call.tool_name) for tool, call in calls_to_run]
+        tasks = [
+            asyncio.create_task(tool.run(call, run_context, ctx.deps.tracer), name=call.tool_name)
+            for tool, call in calls_to_run
+        ]
         pending = tasks
         while pending:
             done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
@@ -697,6 +698,7 @@ async def _tool_from_mcp_server(
 
 def _unknown_tool(
     tool_name: str,
+    tool_call_id: str,
     ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
 ) -> _messages.RetryPromptPart:
     ctx.state.increment_retries(ctx.deps.max_result_retries)
@@ -709,7 +711,11 @@ def _unknown_tool(
     else:
         msg = 'No tools available.'
 
-    return _messages.RetryPromptPart(content=f'Unknown tool name: {tool_name!r}. {msg}')
+    return _messages.RetryPromptPart(
+        tool_name=tool_name,
+        tool_call_id=tool_call_id,
+        content=f'Unknown tool name: {tool_name!r}. {msg}',
+    )
 
 
 async def _validate_result(

{pydantic_ai_slim-0.0.46 → pydantic_ai_slim-0.0.48}/pydantic_ai/_cli.py

@@ -4,6 +4,7 @@ import argparse
 import asyncio
 import sys
 from collections.abc import Sequence
+from contextlib import ExitStack
 from datetime import datetime, timezone
 from importlib.metadata import version
 from pathlib import Path
@@ -11,9 +12,10 @@ from typing import cast
 
 from typing_inspection.introspection import get_literal_values
 
+from pydantic_ai.agent import Agent
 from pydantic_ai.exceptions import UserError
+from pydantic_ai.messages import ModelMessage, PartDeltaEvent, TextPartDelta
 from pydantic_ai.models import KnownModelName
-from pydantic_graph.nodes import End
 
 try:
     import argcomplete
@@ -24,8 +26,9 @@ try:
     from prompt_toolkit.history import FileHistory
     from rich.console import Console, ConsoleOptions, RenderResult
     from rich.live import Live
-    from rich.markdown import CodeBlock, Markdown
+    from rich.markdown import CodeBlock, Heading, Markdown
     from rich.status import Status
+    from rich.style import Style
     from rich.syntax import Syntax
     from rich.text import Text
 except ImportError as _import_error:
@@ -34,13 +37,16 @@ except ImportError as _import_error:
         'you can use the `cli` optional group — `pip install "pydantic-ai-slim[cli]"`'
     ) from _import_error
 
-from pydantic_ai.agent import Agent
-from pydantic_ai.messages import ModelMessage, PartDeltaEvent, TextPartDelta
 
 __version__ = version('pydantic-ai-slim')
 
 
 class SimpleCodeBlock(CodeBlock):
+    """Customised code blocks in markdown.
+
+    This avoids a background color which messes up copy-pasting and sets the language name as dim prefix and suffix.
+    """
+
     def __rich_console__(self, console: Console, options: ConsoleOptions) -> RenderResult:  # pragma: no cover
         code = str(self.text).rstrip()
         yield Text(self.lexer_name, style='dim')
@@ -48,7 +54,18 @@ class SimpleCodeBlock(CodeBlock):
         yield Text(f'/{self.lexer_name}', style='dim')
 
 
-Markdown.elements['fence'] = SimpleCodeBlock
+class LeftHeading(Heading):
+    """Customised headings in markdown to stop centering and prepend markdown style hashes."""
+
+    def __rich_console__(self, console: Console, options: ConsoleOptions) -> RenderResult:  # pragma: no cover
+        # note we use `Style(bold=True)` not `self.style_name` here to disable underlining which is ugly IMHO
+        yield Text(f'{"#" * int(self.tag[1:])} {self.text.plain}', style=Style(bold=True))
+
+
+Markdown.elements.update(
+    fence=SimpleCodeBlock,
+    heading_open=LeftHeading,
+)
 
 
 def cli(args_list: Sequence[str] | None = None) -> int:  # noqa: C901  # pragma: no cover
@@ -65,28 +82,53 @@ Special prompt:
         formatter_class=argparse.RawTextHelpFormatter,
     )
     parser.add_argument('prompt', nargs='?', help='AI Prompt, if omitted fall into interactive mode')
-    parser.add_argument(
+    arg = parser.add_argument(
+        '-m',
         '--model',
         nargs='?',
-        help='Model to use, it should be "<provider>:<model>" e.g. "openai:gpt-4o". If omitted it will default to "openai:gpt-4o"',
+        help='Model to use, in format "<provider>:<model>" e.g. "openai:gpt-4o". Defaults to "openai:gpt-4o".',
         default='openai:gpt-4o',
-    ).completer = argcomplete.ChoicesCompleter(list(get_literal_values(KnownModelName)))  # type: ignore[reportPrivateUsage]
-    parser.add_argument('--no-stream', action='store_true', help='Whether to stream responses from OpenAI')
+    )
+    # we don't want to autocomplete or list models that don't include the provider,
+    # e.g. we want to show `openai:gpt-4o` but not `gpt-4o`
+    qualified_model_names = [n for n in get_literal_values(KnownModelName.__value__) if ':' in n]
+    arg.completer = argcomplete.ChoicesCompleter(qualified_model_names)  # type: ignore[reportPrivateUsage]
+    parser.add_argument(
+        '-l',
+        '--list-models',
+        action='store_true',
+        help='List all available models and exit',
+    )
+    parser.add_argument(
+        '-t',
+        '--code-theme',
+        nargs='?',
+        help='Which colors to use for code, can be "dark", "light" or any theme from pygments.org/styles/. Defaults to "monokai".',
+        default='monokai',
+    )
+    parser.add_argument('--no-stream', action='store_true', help='Whether to stream responses from the model')
     parser.add_argument('--version', action='store_true', help='Show version and exit')
 
     argcomplete.autocomplete(parser)
     args = parser.parse_args(args_list)
 
     console = Console()
-    console.print(f'pai - PydanticAI CLI v{__version__}', style='green bold', highlight=False)
+    console.print(
+        f'[green]pai - PydanticAI CLI v{__version__} using[/green] [magenta]{args.model}[/magenta]', highlight=False
+    )
     if args.version:
         return 0
+    if args.list_models:
+        console.print('Available models:', style='green bold')
+        for model in qualified_model_names:
+            console.print(f'  {model}', highlight=False)
+        return 0
 
     now_utc = datetime.now(timezone.utc)
     tzname = now_utc.astimezone().tzinfo.tzname(now_utc)  # type: ignore
     try:
         agent = Agent(
-            model=args.model or 'openai:gpt-4o',
+            model=args.model,
             system_prompt=f"""\
     Help the user by responding to their request, the output should be concise and always written in markdown.
     The current date and time is {datetime.now()} {tzname}.
@@ -97,10 +139,16 @@ Special prompt:
         return 1
 
     stream = not args.no_stream
+    if args.code_theme == 'light':
+        code_theme = 'default'
+    elif args.code_theme == 'dark':
+        code_theme = 'monokai'
+    else:
+        code_theme = args.code_theme
 
     if prompt := cast(str, args.prompt):
         try:
-            asyncio.run(ask_agent(agent, prompt, stream, console))
+            asyncio.run(ask_agent(agent, prompt, stream, console, code_theme))
         except KeyboardInterrupt:
             pass
         return 0
@@ -121,37 +169,46 @@ Special prompt:
             continue
 
         ident_prompt = text.lower().strip(' ').replace(' ', '-').lstrip(' ')
-        if ident_prompt == '/markdown':
-            try:
-                parts = messages[-1].parts
-            except IndexError:
-                console.print('[dim]No markdown output available.[/dim]')
-                continue
-            for part in parts:
-                if part.part_kind == 'text':
-                    last_content = part.content
-                    console.print('[dim]Last markdown output of last question:[/dim]\n')
-                    console.print(Syntax(last_content, lexer='markdown', background_color='default'))
-
-            continue
-        if ident_prompt == '/multiline':
-            multiline = not multiline
-            if multiline:
-                console.print(
-                    'Enabling multiline mode. '
-                    '[dim]Press [Meta+Enter] or [Esc] followed by [Enter] to accept input.[/dim]'
-                )
+        if ident_prompt.startswith('/'):
+            if ident_prompt == '/markdown':
+                try:
+                    parts = messages[-1].parts
+                except IndexError:
+                    console.print('[dim]No markdown output available.[/dim]')
+                    continue
+                console.print('[dim]Markdown output of last question:[/dim]\n')
+                for part in parts:
+                    if part.part_kind == 'text':
+                        console.print(
+                            Syntax(
+                                part.content,
+                                lexer='markdown',
+                                theme=code_theme,
+                                word_wrap=True,
+                                background_color='default',
+                            )
+                        )
+
+            elif ident_prompt == '/multiline':
+                multiline = not multiline
+                if multiline:
+                    console.print(
+                        'Enabling multiline mode. '
+                        '[dim]Press [Meta+Enter] or [Esc] followed by [Enter] to accept input.[/dim]'
+                    )
+                else:
+                    console.print('Disabling multiline mode.')
+            elif ident_prompt == '/exit':
+                console.print('[dim]Exiting…[/dim]')
+                return 0
             else:
-                console.print('Disabling multiline mode.')
-            continue
-        if ident_prompt == '/exit':
-            console.print('[dim]Exiting…[/dim]')
-            return 0
-
-        try:
-            messages = asyncio.run(ask_agent(agent, text, stream, console, messages))
-        except KeyboardInterrupt:
-            return 0
+                console.print(f'[red]Unknown command[/red] [magenta]`{ident_prompt}`[/magenta]')
+        else:
+            try:
+                messages = asyncio.run(ask_agent(agent, text, stream, console, code_theme, messages))
+            except KeyboardInterrupt:
+                console.print('[dim]Interrupted[/dim]')
+                messages = []
 
 
 async def ask_agent(
@@ -159,48 +216,34 @@ async def ask_agent(
     prompt: str,
     stream: bool,
     console: Console,
+    code_theme: str,
     messages: list[ModelMessage] | None = None,
 ) -> list[ModelMessage]:  # pragma: no cover
-    status: None | Status = Status('[dim]Working on it…[/dim]', console=console)
-    live = Live('', refresh_per_second=15, console=console)
-    status.start()
-
-    async with agent.iter(prompt, message_history=messages) as agent_run:
-        console.print('\nResponse:', style='green')
-
-        content: str = ''
-        interrupted = False
-        try:
-            node = agent_run.next_node
-            while not isinstance(node, End):
-                node = await agent_run.next(node)
+    status = Status('[dim]Working on it…[/dim]', console=console)
+
+    if not stream:
+        with status:
+            result = await agent.run(prompt, message_history=messages)
+        content = result.data
+        console.print(Markdown(content, code_theme=code_theme))
+        return result.all_messages()
+
+    with status, ExitStack() as stack:
+        async with agent.iter(prompt, message_history=messages) as agent_run:
+            live = Live('', refresh_per_second=15, console=console, vertical_overflow='visible')
+            content: str = ''
+            async for node in agent_run:
                 if Agent.is_model_request_node(node):
                     async with node.stream(agent_run.ctx) as handle_stream:
-                        # NOTE(Marcelo): It took me a lot of time to figure out how to stop `status` and start `live`
-                        # in a context manager, so I had to do it manually with `stop` and `start` methods.
-                        # PR welcome to simplify this code.
-                        if status is not None:
-                            status.stop()
-                            status = None
-                        if not live.is_started:
-                            live.start()
+                        status.stop()  # stopping multiple times is idempotent
+                        stack.enter_context(live)  # entering multiple times is idempotent
+
                         async for event in handle_stream:
                             if isinstance(event, PartDeltaEvent) and isinstance(event.delta, TextPartDelta):
-                                if stream:
-                                    content += event.delta.content_delta
-                                    live.update(Markdown(content))
-        except KeyboardInterrupt:
-            interrupted = True
-        finally:
-            live.stop()
-
-        if interrupted:
-            console.print('[dim]Interrupted[/dim]')
+                                content += event.delta.content_delta
+                                live.update(Markdown(content, code_theme=code_theme))
 
-        assert agent_run.result
-        if not stream:
-            content = agent_run.result.data
-            console.print(Markdown(content))
+        assert agent_run.result is not None
         return agent_run.result.all_messages()
 
 

{pydantic_ai_slim-0.0.46 → pydantic_ai_slim-0.0.48}/pydantic_ai/_result.py

@@ -13,7 +13,7 @@ from typing_inspection.introspection import is_union_origin
 from . import _utils, messages as _messages
 from .exceptions import ModelRetry
 from .result import ResultDataT, ResultDataT_inv, ResultValidatorFunc
-from .tools import AgentDepsT, RunContext, ToolDefinition
+from .tools import AgentDepsT, GenerateToolJsonSchema, RunContext, ToolDefinition
 
 T = TypeVar('T')
 """An invariant TypeVar."""
@@ -159,13 +159,20 @@ class ResultTool(Generic[ResultDataT]):
             self.type_adapter = TypeAdapter(response_type)
             outer_typed_dict_key: str | None = None
             # noinspection PyArgumentList
-            parameters_json_schema = _utils.check_object_json_schema(self.type_adapter.json_schema())
+            parameters_json_schema = _utils.check_object_json_schema(
+                self.type_adapter.json_schema(schema_generator=GenerateToolJsonSchema)
+            )
         else:
-            response_data_typed_dict = TypedDict('response_data_typed_dict', {'response': response_type})  # noqa
+            response_data_typed_dict = TypedDict(  # noqa: UP013
+                'response_data_typed_dict',
+                {'response': response_type},  # pyright: ignore[reportInvalidTypeForm]
+            )
             self.type_adapter = TypeAdapter(response_data_typed_dict)
             outer_typed_dict_key = 'response'
             # noinspection PyArgumentList
-            parameters_json_schema = _utils.check_object_json_schema(self.type_adapter.json_schema())
+            parameters_json_schema = _utils.check_object_json_schema(
+                self.type_adapter.json_schema(schema_generator=GenerateToolJsonSchema)
+            )
             # including `response_data_typed_dict` as a title here doesn't add anything and could confuse the LLM
             parameters_json_schema.pop('title')
 

{pydantic_ai_slim-0.0.46 → pydantic_ai_slim-0.0.48}/pydantic_ai/_utils.py

@@ -40,7 +40,7 @@ def is_model_like(type_: Any) -> bool:
     return (
         isinstance(type_, type)
         and not isinstance(type_, GenericAlias)
-        and (issubclass(type_, BaseModel) or is_dataclass(type_) or is_typeddict(type_))
+        and (issubclass(type_, BaseModel) or is_dataclass(type_) or is_typeddict(type_))  # pyright: ignore[reportUnknownArgumentType]
     )
 
 

{pydantic_ai_slim-0.0.46 → pydantic_ai_slim-0.0.48}/pydantic_ai/agent.py

@@ -94,9 +94,11 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
     ```
     """
 
-    # we use dataclass fields in order to conveniently know what attributes are available
-    model: models.Model | models.KnownModelName | None
-    """The default model configured for this agent."""
+    model: models.Model | models.KnownModelName | str | None
+    """The default model configured for this agent.
+
+    We allow str here since the actual list of allowed models changes frequently.
+    """
 
     name: str | None
     """The name of the agent, used for logging.
@@ -142,7 +144,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
 
     def __init__(
         self,
-        model: models.Model | models.KnownModelName | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
         *,
         result_type: type[ResultDataT] = str,
         system_prompt: str | Sequence[str] = (),
@@ -163,7 +165,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
 
         Args:
             model: The default model to use for this agent, if not provide,
-                you must provide the model when calling it.
+                you must provide the model when calling it. We allow str here since the actual list of allowed models changes frequently.
             result_type: The type of the result data, used to validate the result data, defaults to `str`.
             system_prompt: Static system prompts to use for this agent, you can also register system
                 prompts via a function with [`system_prompt`][pydantic_ai.Agent.system_prompt].
@@ -212,16 +214,16 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
 
         self._result_tool_name = result_tool_name
         self._result_tool_description = result_tool_description
-        self._result_schema: _result.ResultSchema[ResultDataT] | None = _result.ResultSchema[result_type].build(
+        self._result_schema = _result.ResultSchema[result_type].build(
             result_type, result_tool_name, result_tool_description
         )
-        self._result_validators: list[_result.ResultValidator[AgentDepsT, ResultDataT]] = []
+        self._result_validators = []
 
         self._system_prompts = (system_prompt,) if isinstance(system_prompt, str) else tuple(system_prompt)
-        self._system_prompt_functions: list[_system_prompt.SystemPromptRunner[AgentDepsT]] = []
-        self._system_prompt_dynamic_functions: dict[str, _system_prompt.SystemPromptRunner[AgentDepsT]] = {}
+        self._system_prompt_functions = []
+        self._system_prompt_dynamic_functions = {}
 
-        self._function_tools: dict[str, Tool[AgentDepsT]] = {}
+        self._function_tools = {}
 
         self._default_retries = retries
         self._max_result_retries = result_retries if result_retries is not None else retries
@@ -244,7 +246,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         *,
         result_type: None = None,
         message_history: list[_messages.ModelMessage] | None = None,
-        model: models.Model | models.KnownModelName | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
@@ -259,7 +261,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         *,
         result_type: type[RunResultDataT],
         message_history: list[_messages.ModelMessage] | None = None,
-        model: models.Model | models.KnownModelName | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
@@ -273,7 +275,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         *,
         result_type: type[RunResultDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
-        model: models.Model | models.KnownModelName | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
@@ -327,8 +329,8 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             async for _ in agent_run:
                 pass
 
-        assert (final_result := agent_run.result) is not None, 'The graph run did not finish properly'
-        return final_result
+        assert agent_run.result is not None, 'The graph run did not finish properly'
+        return agent_run.result
 
     @asynccontextmanager
     async def iter(
@@ -337,7 +339,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         *,
         result_type: type[RunResultDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
-        model: models.Model | models.KnownModelName | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
@@ -498,7 +500,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         user_prompt: str | Sequence[_messages.UserContent],
         *,
         message_history: list[_messages.ModelMessage] | None = None,
-        model: models.Model | models.KnownModelName | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
@@ -513,7 +515,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         *,
         result_type: type[RunResultDataT] | None,
         message_history: list[_messages.ModelMessage] | None = None,
-        model: models.Model | models.KnownModelName | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
@@ -527,7 +529,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         *,
         result_type: type[RunResultDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
-        model: models.Model | models.KnownModelName | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
@@ -588,7 +590,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         *,
         result_type: None = None,
         message_history: list[_messages.ModelMessage] | None = None,
-        model: models.Model | models.KnownModelName | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
@@ -603,7 +605,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         *,
         result_type: type[RunResultDataT],
         message_history: list[_messages.ModelMessage] | None = None,
-        model: models.Model | models.KnownModelName | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
@@ -618,7 +620,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         *,
         result_type: type[RunResultDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
-        model: models.Model | models.KnownModelName | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
@@ -757,12 +759,12 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         self,
         *,
         deps: AgentDepsT | _utils.Unset = _utils.UNSET,
-        model: models.Model | models.KnownModelName | _utils.Unset = _utils.UNSET,
+        model: models.Model | models.KnownModelName | str | _utils.Unset = _utils.UNSET,
     ) -> Iterator[None]:
         """Context manager to temporarily override agent dependencies and model.
 
         This is particularly useful when testing.
-        You can find an example of this [here](../testing-evals.md#overriding-model-via-pytest-fixtures).
+        You can find an example of this [here](../testing.md#overriding-model-via-pytest-fixtures).
 
         Args:
             deps: The dependencies to use instead of the dependencies passed to the agent run.
@@ -774,11 +776,9 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         else:
             override_deps_before = _utils.UNSET
 
-        # noinspection PyTypeChecker
         if _utils.is_set(model):
             override_model_before = self._override_model
-            # noinspection PyTypeChecker
-            self._override_model = _utils.Some(models.infer_model(model))  # pyright: ignore[reportArgumentType]
+            self._override_model = _utils.Some(models.infer_model(model))
         else:
             override_model_before = _utils.UNSET
 
@@ -1154,7 +1154,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
 
         self._function_tools[tool.name] = tool
 
-    def _get_model(self, model: models.Model | models.KnownModelName | None) -> models.Model:
+    def _get_model(self, model: models.Model | models.KnownModelName | str | None) -> models.Model:
         """Create a model configured for this agent.
 
         Args:
@@ -1168,7 +1168,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             # we don't want `override()` to cover up errors from the model not being defined, hence this check
             if model is None and self.model is None:
                 raise exceptions.UserError(
-                    '`model` must be set either when creating the agent or when calling it. '
+                    '`model` must either be set on the agent or included when calling it. '
                     '(Even when `override(model=...)` is customizing the model that will actually be called)'
                 )
             model_ = some_model.value
@@ -1178,7 +1178,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             # noinspection PyTypeChecker
             model_ = self.model = models.infer_model(self.model)
         else:
-            raise exceptions.UserError('`model` must be set either when creating the agent or when calling it.')
+            raise exceptions.UserError('`model` must either be set on the agent or included when calling it.')
 
         instrument = self.instrument
         if instrument is None: