pydantic-ai-slim 0.0.20__tar.gz → 0.0.21__tar.gz

{pydantic_ai_slim-0.0.20 → pydantic_ai_slim-0.0.21}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.20
+Version: 0.0.21
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License-Expression: MIT
@@ -34,7 +34,7 @@ Requires-Dist: anthropic>=0.40.0; extra == 'anthropic'
 Provides-Extra: cohere
 Requires-Dist: cohere>=5.13.11; extra == 'cohere'
 Provides-Extra: graph
-Requires-Dist: pydantic-graph==0.0.20; extra == 'graph'
+Requires-Dist: pydantic-graph==0.0.21; extra == 'graph'
 Provides-Extra: groq
 Requires-Dist: groq>=0.12.0; extra == 'groq'
 Provides-Extra: logfire

{pydantic_ai_slim-0.0.20 → pydantic_ai_slim-0.0.21}/pydantic_ai/_parts_manager.py

@@ -221,7 +221,7 @@ class ModelResponsePartsManager:
             ModelResponseStreamEvent: A `PartStartEvent` indicating that a new tool call part
             has been added to the manager, or replaced an existing part.
         """
-        new_part = ToolCallPart.from_raw_args(tool_name=tool_name, args=args, tool_call_id=tool_call_id)
+        new_part = ToolCallPart(tool_name=tool_name, args=args, tool_call_id=tool_call_id)
         if vendor_part_id is None:
             # vendor_part_id is None, so we unconditionally append a new ToolCallPart to the end of the list
             new_part_index = len(self._parts)

{pydantic_ai_slim-0.0.20 → pydantic_ai_slim-0.0.21}/pydantic_ai/_result.py

@@ -201,14 +201,10 @@ class ResultTool(Generic[ResultDataT]):
         """
         try:
             pyd_allow_partial: Literal['off', 'trailing-strings'] = 'trailing-strings' if allow_partial else 'off'
-            if isinstance(tool_call.args, _messages.ArgsJson):
-                result = self.type_adapter.validate_json(
-                    tool_call.args.args_json or '', experimental_allow_partial=pyd_allow_partial
-                )
+            if isinstance(tool_call.args, str):
+                result = self.type_adapter.validate_json(tool_call.args, experimental_allow_partial=pyd_allow_partial)
             else:
-                result = self.type_adapter.validate_python(
-                    tool_call.args.args_dict, experimental_allow_partial=pyd_allow_partial
-                )
+                result = self.type_adapter.validate_python(tool_call.args, experimental_allow_partial=pyd_allow_partial)
         except ValidationError as e:
             if wrap_validation_errors:
                 m = _messages.RetryPromptPart(

{pydantic_ai_slim-0.0.20 → pydantic_ai_slim-0.0.21}/pydantic_ai/_utils.py

@@ -8,7 +8,7 @@ from dataclasses import dataclass, is_dataclass
 from datetime import datetime, timezone
 from functools import partial
 from types import GenericAlias
-from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union
 
 from pydantic import BaseModel
 from pydantic.json_schema import JsonSchemaValue
@@ -79,61 +79,6 @@ def is_set(t_or_unset: T | Unset) -> TypeGuard[T]:
     return t_or_unset is not UNSET
 
 
-Left = TypeVar('Left')
-Right = TypeVar('Right')
-
-
-class Either(Generic[Left, Right]):
-    """Two member Union that records which member was set, this is analogous to Rust enums with two variants.
-
-    Usage:
-
-    ```python
-    if left_thing := either.left:
-        use_left(left_thing.value)
-    else:
-        use_right(either.right)
-    ```
-    """
-
-    __slots__ = '_left', '_right'
-
-    @overload
-    def __init__(self, *, left: Left) -> None: ...
-
-    @overload
-    def __init__(self, *, right: Right) -> None: ...
-
-    def __init__(self, left: Left | Unset = UNSET, right: Right | Unset = UNSET) -> None:
-        if left is not UNSET:
-            assert right is UNSET, '`Either` must receive exactly one argument - `left` or `right`'
-            self._left: Option[Left] = Some(cast(Left, left))
-        else:
-            assert right is not UNSET, '`Either` must receive exactly one argument - `left` or `right`'
-            self._left = None
-            self._right = cast(Right, right)
-
-    @property
-    def left(self) -> Option[Left]:
-        return self._left
-
-    @property
-    def right(self) -> Right:
-        return self._right
-
-    def is_left(self) -> bool:
-        return self._left is not None
-
-    def whichever(self) -> Left | Right:
-        return self._left.value if self._left is not None else self.right
-
-    def __repr__(self):
-        if left := self._left:
-            return f'Either(left={left.value!r})'
-        else:
-            return f'Either(right={self.right!r})'
-
-
 @asynccontextmanager
 async def group_by_temporal(
     aiterable: AsyncIterable[T], soft_max_interval: float | None

{pydantic_ai_slim-0.0.20 → pydantic_ai_slim-0.0.21}/pydantic_ai/agent.py

@@ -60,7 +60,7 @@ EndStrategy = Literal['early', 'exhaustive']
 - `'early'`: Stop processing other tool calls once a final result is found
 - `'exhaustive'`: Process all tool calls even after finding a final result
 """
-RunResultData = TypeVar('RunResultData')
+RunResultDataT = TypeVar('RunResultDataT')
 """Type variable for the result data of a run where `result_type` was customized on the run call."""
 
 
@@ -214,7 +214,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData],
+        result_type: type[RunResultDataT],
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDepsT = None,
@@ -222,7 +222,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> result.RunResult[RunResultData]: ...
+    ) -> result.RunResult[RunResultDataT]: ...
 
     async def run(
         self,
@@ -234,7 +234,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         model_settings: ModelSettings | None = None,
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
-        result_type: type[RunResultData] | None = None,
+        result_type: type[RunResultDataT] | None = None,
         infer_name: bool = True,
     ) -> result.RunResult[Any]:
         """Run the agent with a user prompt in async mode.
@@ -352,7 +352,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData] | None,
+        result_type: type[RunResultDataT] | None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDepsT = None,
@@ -360,13 +360,13 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> result.RunResult[RunResultData]: ...
+    ) -> result.RunResult[RunResultDataT]: ...
 
     def run_sync(
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData] | None = None,
+        result_type: type[RunResultDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDepsT = None,
@@ -442,7 +442,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData],
+        result_type: type[RunResultDataT],
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDepsT = None,
@@ -450,14 +450,14 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> AbstractAsyncContextManager[result.StreamedRunResult[AgentDepsT, RunResultData]]: ...
+    ) -> AbstractAsyncContextManager[result.StreamedRunResult[AgentDepsT, RunResultDataT]]: ...
 
     @asynccontextmanager
     async def run_stream(
         self,
         user_prompt: str,
         *,
-        result_type: type[RunResultData] | None = None,
+        result_type: type[RunResultDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDepsT = None,
@@ -572,7 +572,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
                                 # there are result validators that might convert the result data from an overridden
                                 # `result_type` to a type that is not valid as such.
                                 result_validators = cast(
-                                    list[_result.ResultValidator[AgentDepsT, RunResultData]], self._result_validators
+                                    list[_result.ResultValidator[AgentDepsT, RunResultDataT]], self._result_validators
                                 )
 
                                 yield result.StreamedRunResult(
@@ -999,7 +999,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         return model_
 
     async def _prepare_model(
-        self, run_context: RunContext[AgentDepsT], result_schema: _result.ResultSchema[RunResultData] | None
+        self, run_context: RunContext[AgentDepsT], result_schema: _result.ResultSchema[RunResultDataT] | None
     ) -> models.AgentModel:
         """Build tools and create an agent model."""
         function_tools: list[ToolDefinition] = []
@@ -1035,8 +1035,8 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
                                 )
 
     def _prepare_result_schema(
-        self, result_type: type[RunResultData] | None
-    ) -> _result.ResultSchema[RunResultData] | None:
+        self, result_type: type[RunResultDataT] | None
+    ) -> _result.ResultSchema[RunResultDataT] | None:
         if result_type is not None:
             if self._result_validators:
                 raise exceptions.UserError('Cannot set a custom run `result_type` when the agent has result validators')
@@ -1053,7 +1053,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         run_context: RunContext[AgentDepsT],
     ) -> list[_messages.ModelMessage]:
         try:
-            ctx_messages = _messages_ctx_var.get()
+            ctx_messages = get_captured_run_messages()
         except LookupError:
             messages: list[_messages.ModelMessage] = []
         else:
@@ -1080,8 +1080,8 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         self,
         model_response: _messages.ModelResponse,
         run_context: RunContext[AgentDepsT],
-        result_schema: _result.ResultSchema[RunResultData] | None,
-    ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
+    ) -> tuple[_MarkFinalResult[RunResultDataT] | None, list[_messages.ModelRequestPart]]:
         """Process a non-streamed response from the model.
 
         Returns:
@@ -1110,11 +1110,11 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             raise exceptions.UnexpectedModelBehavior('Received empty model response')
 
     async def _handle_text_response(
-        self, text: str, run_context: RunContext[AgentDepsT], result_schema: _result.ResultSchema[RunResultData] | None
-    ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
+        self, text: str, run_context: RunContext[AgentDepsT], result_schema: _result.ResultSchema[RunResultDataT] | None
+    ) -> tuple[_MarkFinalResult[RunResultDataT] | None, list[_messages.ModelRequestPart]]:
         """Handle a plain text response from the model for non-streaming responses."""
         if self._allow_text_result(result_schema):
-            result_data_input = cast(RunResultData, text)
+            result_data_input = cast(RunResultDataT, text)
             try:
                 result_data = await self._validate_result(result_data_input, run_context, None)
             except _result.ToolRetryError as e:
@@ -1133,13 +1133,13 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         self,
         tool_calls: list[_messages.ToolCallPart],
         run_context: RunContext[AgentDepsT],
-        result_schema: _result.ResultSchema[RunResultData] | None,
-    ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
+    ) -> tuple[_MarkFinalResult[RunResultDataT] | None, list[_messages.ModelRequestPart]]:
         """Handle a structured response containing tool calls from the model for non-streaming responses."""
         assert tool_calls, 'Expected at least one tool call'
 
         # first look for the result tool call
-        final_result: _MarkFinalResult[RunResultData] | None = None
+        final_result: _MarkFinalResult[RunResultDataT] | None = None
 
         parts: list[_messages.ModelRequestPart] = []
         if result_schema is not None:
@@ -1168,7 +1168,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         tool_calls: list[_messages.ToolCallPart],
         result_tool_name: str | None,
         run_context: RunContext[AgentDepsT],
-        result_schema: _result.ResultSchema[RunResultData] | None,
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
     ) -> list[_messages.ModelRequestPart]:
         """Process function (non-result) tool calls in parallel.
 
@@ -1227,7 +1227,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         self,
         streamed_response: models.StreamedResponse,
         run_context: RunContext[AgentDepsT],
-        result_schema: _result.ResultSchema[RunResultData] | None,
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
     ) -> _MarkFinalResult[models.StreamedResponse] | tuple[_messages.ModelResponse, list[_messages.ModelRequestPart]]:
         """Process a streamed response from the model.
 
@@ -1282,15 +1282,15 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
 
     async def _validate_result(
         self,
-        result_data: RunResultData,
+        result_data: RunResultDataT,
         run_context: RunContext[AgentDepsT],
         tool_call: _messages.ToolCallPart | None,
-    ) -> RunResultData:
+    ) -> RunResultDataT:
         if self._result_validators:
             agent_result_data = cast(ResultDataT, result_data)
             for validator in self._result_validators:
                 agent_result_data = await validator.validate(agent_result_data, tool_call, run_context)
-            return cast(RunResultData, agent_result_data)
+            return cast(RunResultDataT, agent_result_data)
         else:
             return result_data
 
@@ -1315,7 +1315,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
     def _unknown_tool(
         self,
         tool_name: str,
-        result_schema: _result.ResultSchema[RunResultData] | None,
+        result_schema: _result.ResultSchema[RunResultDataT] | None,
     ) -> _messages.RetryPromptPart:
         names = list(self._function_tools.keys())
         if result_schema:
@@ -1358,7 +1358,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
                             return
 
     @staticmethod
-    def _allow_text_result(result_schema: _result.ResultSchema[RunResultData] | None) -> bool:
+    def _allow_text_result(result_schema: _result.ResultSchema[RunResultDataT] | None) -> bool:
         return result_schema is None or result_schema.allow_text_result
 
     @property
@@ -1413,6 +1413,10 @@ def capture_run_messages() -> Iterator[list[_messages.ModelMessage]]:
             _messages_ctx_var.reset(token)
 
 
+def get_captured_run_messages() -> _RunMessages:
+    return _messages_ctx_var.get()
+
+
 @dataclasses.dataclass
 class _MarkFinalResult(Generic[ResultDataT]):
     """Marker class to indicate that the result is the final result.

{pydantic_ai_slim-0.0.20 → pydantic_ai_slim-0.0.21}/pydantic_ai/messages.py

@@ -6,7 +6,6 @@ from typing import Annotated, Any, Literal, Union, cast, overload
 
 import pydantic
 import pydantic_core
-from typing_extensions import Self, assert_never
 
 from ._utils import now_utc as _now_utc
 from .exceptions import UnexpectedModelBehavior
@@ -168,22 +167,6 @@ class TextPart:
         return bool(self.content)
 
 
-@dataclass
-class ArgsJson:
-    """Tool arguments as a JSON string."""
-
-    args_json: str
-    """A JSON string of arguments."""
-
-
-@dataclass
-class ArgsDict:
-    """Tool arguments as a Python dictionary."""
-
-    args_dict: dict[str, Any]
-    """A python dictionary of arguments."""
-
-
 @dataclass
 class ToolCallPart:
     """A tool call from a model."""
@@ -191,10 +174,10 @@ class ToolCallPart:
     tool_name: str
     """The name of the tool to call."""
 
-    args: ArgsJson | ArgsDict
+    args: str | dict[str, Any]
     """The arguments to pass to the tool.
 
-    Either as JSON or a Python dictionary depending on how data was returned.
+    This is stored either as a JSON string or a Python dictionary depending on how data was received.
     """
 
     tool_call_id: str | None = None
@@ -203,24 +186,14 @@ class ToolCallPart:
     part_kind: Literal['tool-call'] = 'tool-call'
     """Part type identifier, this is available on all parts as a discriminator."""
 
-    @classmethod
-    def from_raw_args(cls, tool_name: str, args: str | dict[str, Any], tool_call_id: str | None = None) -> Self:
-        """Create a `ToolCallPart` from raw arguments, converting them to `ArgsJson` or `ArgsDict`."""
-        if isinstance(args, str):
-            return cls(tool_name, ArgsJson(args), tool_call_id)
-        elif isinstance(args, dict):
-            return cls(tool_name, ArgsDict(args), tool_call_id)
-        else:
-            assert_never(args)
-
     def args_as_dict(self) -> dict[str, Any]:
         """Return the arguments as a Python dictionary.
 
         This is just for convenience with models that require dicts as input.
         """
-        if isinstance(self.args, ArgsDict):
-            return self.args.args_dict
-        args = pydantic_core.from_json(self.args.args_json)
+        if isinstance(self.args, dict):
+            return self.args
+        args = pydantic_core.from_json(self.args)
         assert isinstance(args, dict), 'args should be a dict'
         return cast(dict[str, Any], args)
 
@@ -229,16 +202,18 @@ class ToolCallPart:
 
         This is just for convenience with models that require JSON strings as input.
         """
-        if isinstance(self.args, ArgsJson):
-            return self.args.args_json
-        return pydantic_core.to_json(self.args.args_dict).decode()
+        if isinstance(self.args, str):
+            return self.args
+        return pydantic_core.to_json(self.args).decode()
 
     def has_content(self) -> bool:
         """Return `True` if the arguments contain any data."""
-        if isinstance(self.args, ArgsDict):
-            return any(self.args.args_dict.values())
+        if isinstance(self.args, dict):
+            # TODO: This should probably return True if you have the value False, or 0, etc.
+            #   It makes sense to me to ignore empty strings, but not sure about empty lists or dicts
+            return any(self.args.values())
         else:
-            return bool(self.args.args_json)
+            return bool(self.args)
 
 
 ModelResponsePart = Annotated[Union[TextPart, ToolCallPart], pydantic.Discriminator('part_kind')]
@@ -331,7 +306,7 @@ class ToolCallPartDelta:
         if self.tool_name_delta is None or self.args_delta is None:
             return None
 
-        return ToolCallPart.from_raw_args(
+        return ToolCallPart(
             self.tool_name_delta,
             self.args_delta,
             self.tool_call_id,
@@ -396,7 +371,7 @@ class ToolCallPartDelta:
 
         # If we now have enough data to create a full ToolCallPart, do so
         if delta.tool_name_delta is not None and delta.args_delta is not None:
-            return ToolCallPart.from_raw_args(
+            return ToolCallPart(
                 delta.tool_name_delta,
                 delta.args_delta,
                 delta.tool_call_id,
@@ -412,15 +387,15 @@ class ToolCallPartDelta:
             part = replace(part, tool_name=tool_name)
 
         if isinstance(self.args_delta, str):
-            if not isinstance(part.args, ArgsJson):
+            if not isinstance(part.args, str):
                 raise UnexpectedModelBehavior(f'Cannot apply JSON deltas to non-JSON tool arguments ({part=}, {self=})')
-            updated_json = part.args.args_json + self.args_delta
-            part = replace(part, args=ArgsJson(updated_json))
+            updated_json = part.args + self.args_delta
+            part = replace(part, args=updated_json)
         elif isinstance(self.args_delta, dict):
-            if not isinstance(part.args, ArgsDict):
+            if not isinstance(part.args, dict):
                 raise UnexpectedModelBehavior(f'Cannot apply dict deltas to non-dict tool arguments ({part=}, {self=})')
-            updated_dict = {**(part.args.args_dict or {}), **self.args_delta}
-            part = replace(part, args=ArgsDict(updated_dict))
+            updated_dict = {**(part.args or {}), **self.args_delta}
+            part = replace(part, args=updated_dict)
 
         if self.tool_call_id:
             # Replace the tool_call_id entirely if given

{pydantic_ai_slim-0.0.20 → pydantic_ai_slim-0.0.21}/pydantic_ai/models/__init__.py

@@ -12,9 +12,10 @@ from contextlib import asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
 from functools import cache
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING
 
 import httpx
+from typing_extensions import Literal
 
 from .._parts_manager import ModelResponsePartsManager
 from ..exceptions import UserError
@@ -27,58 +28,6 @@ if TYPE_CHECKING:
 
 
 KnownModelName = Literal[
-    'openai:gpt-4o',
-    'openai:gpt-4o-mini',
-    'openai:gpt-4-turbo',
-    'openai:gpt-4',
-    'openai:o1-preview',
-    'openai:o1-mini',
-    'openai:o1',
-    'openai:gpt-3.5-turbo',
-    'groq:llama-3.3-70b-versatile',
-    'groq:llama-3.1-70b-versatile',
-    'groq:llama3-groq-70b-8192-tool-use-preview',
-    'groq:llama3-groq-8b-8192-tool-use-preview',
-    'groq:llama-3.1-70b-specdec',
-    'groq:llama-3.1-8b-instant',
-    'groq:llama-3.2-1b-preview',
-    'groq:llama-3.2-3b-preview',
-    'groq:llama-3.2-11b-vision-preview',
-    'groq:llama-3.2-90b-vision-preview',
-    'groq:llama3-70b-8192',
-    'groq:llama3-8b-8192',
-    'groq:mixtral-8x7b-32768',
-    'groq:gemma2-9b-it',
-    'groq:gemma-7b-it',
-    'google-gla:gemini-1.5-flash',
-    'google-gla:gemini-1.5-pro',
-    'google-gla:gemini-2.0-flash-exp',
-    'google-vertex:gemini-1.5-flash',
-    'google-vertex:gemini-1.5-pro',
-    'google-vertex:gemini-2.0-flash-exp',
-    'mistral:mistral-small-latest',
-    'mistral:mistral-large-latest',
-    'mistral:codestral-latest',
-    'mistral:mistral-moderation-latest',
-    'ollama:codellama',
-    'ollama:deepseek-r1',
-    'ollama:gemma',
-    'ollama:gemma2',
-    'ollama:llama3',
-    'ollama:llama3.1',
-    'ollama:llama3.2',
-    'ollama:llama3.2-vision',
-    'ollama:llama3.3',
-    'ollama:mistral',
-    'ollama:mistral-nemo',
-    'ollama:mixtral',
-    'ollama:phi3',
-    'ollama:phi4',
-    'ollama:qwq',
-    'ollama:qwen',
-    'ollama:qwen2',
-    'ollama:qwen2.5',
-    'ollama:starcoder2',
     'anthropic:claude-3-5-haiku-latest',
     'anthropic:claude-3-5-sonnet-latest',
     'anthropic:claude-3-opus-latest',
@@ -98,6 +47,104 @@ KnownModelName = Literal[
     'cohere:command-r-plus-04-2024',
     'cohere:command-r-plus-08-2024',
     'cohere:command-r7b-12-2024',
+    'google-gla:gemini-1.0-pro',
+    'google-gla:gemini-1.5-flash',
+    'google-gla:gemini-1.5-flash-8b',
+    'google-gla:gemini-1.5-pro',
+    'google-gla:gemini-2.0-flash-exp',
+    'google-vertex:gemini-1.0-pro',
+    'google-vertex:gemini-1.5-flash',
+    'google-vertex:gemini-1.5-flash-8b',
+    'google-vertex:gemini-1.5-pro',
+    'google-vertex:gemini-2.0-flash-exp',
+    'gpt-3.5-turbo',
+    'gpt-3.5-turbo-0125',
+    'gpt-3.5-turbo-0301',
+    'gpt-3.5-turbo-0613',
+    'gpt-3.5-turbo-1106',
+    'gpt-3.5-turbo-16k',
+    'gpt-3.5-turbo-16k-0613',
+    'gpt-4',
+    'gpt-4-0125-preview',
+    'gpt-4-0314',
+    'gpt-4-0613',
+    'gpt-4-1106-preview',
+    'gpt-4-32k',
+    'gpt-4-32k-0314',
+    'gpt-4-32k-0613',
+    'gpt-4-turbo',
+    'gpt-4-turbo-2024-04-09',
+    'gpt-4-turbo-preview',
+    'gpt-4-vision-preview',
+    'gpt-4o',
+    'gpt-4o-2024-05-13',
+    'gpt-4o-2024-08-06',
+    'gpt-4o-2024-11-20',
+    'gpt-4o-audio-preview',
+    'gpt-4o-audio-preview-2024-10-01',
+    'gpt-4o-audio-preview-2024-12-17',
+    'gpt-4o-mini',
+    'gpt-4o-mini-2024-07-18',
+    'gpt-4o-mini-audio-preview',
+    'gpt-4o-mini-audio-preview-2024-12-17',
+    'groq:gemma2-9b-it',
+    'groq:llama-3.1-8b-instant',
+    'groq:llama-3.2-11b-vision-preview',
+    'groq:llama-3.2-1b-preview',
+    'groq:llama-3.2-3b-preview',
+    'groq:llama-3.2-90b-vision-preview',
+    'groq:llama-3.3-70b-specdec',
+    'groq:llama-3.3-70b-versatile',
+    'groq:llama3-70b-8192',
+    'groq:llama3-8b-8192',
+    'groq:mixtral-8x7b-32768',
+    'mistral:codestral-latest',
+    'mistral:mistral-large-latest',
+    'mistral:mistral-moderation-latest',
+    'mistral:mistral-small-latest',
+    'o1',
+    'o1-2024-12-17',
+    'o1-mini',
+    'o1-mini-2024-09-12',
+    'o1-preview',
+    'o1-preview-2024-09-12',
+    'openai:chatgpt-4o-latest',
+    'openai:gpt-3.5-turbo',
+    'openai:gpt-3.5-turbo-0125',
+    'openai:gpt-3.5-turbo-0301',
+    'openai:gpt-3.5-turbo-0613',
+    'openai:gpt-3.5-turbo-1106',
+    'openai:gpt-3.5-turbo-16k',
+    'openai:gpt-3.5-turbo-16k-0613',
+    'openai:gpt-4',
+    'openai:gpt-4-0125-preview',
+    'openai:gpt-4-0314',
+    'openai:gpt-4-0613',
+    'openai:gpt-4-1106-preview',
+    'openai:gpt-4-32k',
+    'openai:gpt-4-32k-0314',
+    'openai:gpt-4-32k-0613',
+    'openai:gpt-4-turbo',
+    'openai:gpt-4-turbo-2024-04-09',
+    'openai:gpt-4-turbo-preview',
+    'openai:gpt-4-vision-preview',
+    'openai:gpt-4o',
+    'openai:gpt-4o-2024-05-13',
+    'openai:gpt-4o-2024-08-06',
+    'openai:gpt-4o-2024-11-20',
+    'openai:gpt-4o-audio-preview',
+    'openai:gpt-4o-audio-preview-2024-10-01',
+    'openai:gpt-4o-audio-preview-2024-12-17',
+    'openai:gpt-4o-mini',
+    'openai:gpt-4o-mini-2024-07-18',
+    'openai:gpt-4o-mini-audio-preview',
+    'openai:gpt-4o-mini-audio-preview-2024-12-17',
+    'openai:o1',
+    'openai:o1-2024-12-17',
+    'openai:o1-mini',
+    'openai:o1-mini-2024-09-12',
+    'openai:o1-preview',
+    'openai:o1-preview-2024-09-12',
     'test',
 ]
 """Known model names that can be used with the `model` parameter of [`Agent`][pydantic_ai.Agent].
@@ -291,10 +338,6 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .mistral import MistralModel
 
         return MistralModel(model[8:])
-    elif model.startswith('ollama:'):
-        from .ollama import OllamaModel
-
-        return OllamaModel(model[7:])
     elif model.startswith('anthropic'):
         from .anthropic import AnthropicModel