Improve model names in tests

Author: dmontagu

Diff for: docs/agents.md

@@ -136,7 +136,7 @@ async def main():
         HandleResponseNode(
             model_response=ModelResponse(
                 parts=[TextPart(content='Paris', part_kind='text')],
-                model_name='function:model_logic',
+                model_name='gpt-4o',
                 timestamp=datetime.datetime(...),
                 kind='response',
             )
@@ -197,7 +197,7 @@ async def main():
             HandleResponseNode(
                 model_response=ModelResponse(
                     parts=[TextPart(content='Paris', part_kind='text')],
-                    model_name='function:model_logic',
+                    model_name='gpt-4o',
                     timestamp=datetime.datetime(...),
                     kind='response',
                 )
@@ -612,7 +612,7 @@ with capture_run_messages() as messages:  # (2)!
                         part_kind='tool-call',
                     )
                 ],
-                model_name='function:model_logic',
+                model_name='gpt-4o',
                 timestamp=datetime.datetime(...),
                 kind='response',
             ),
@@ -637,7 +637,7 @@ with capture_run_messages() as messages:  # (2)!
                         part_kind='tool-call',
                     )
                 ],
-                model_name='function:model_logic',
+                model_name='gpt-4o',
                 timestamp=datetime.datetime(...),
                 kind='response',
             ),

Diff for: docs/message-history.md

@@ -62,7 +62,7 @@ print(result.all_messages())
                 part_kind='text',
             )
         ],
-        model_name='function:model_logic',
+        model_name='gpt-4o',
         timestamp=datetime.datetime(...),
         kind='response',
     ),
@@ -193,7 +193,7 @@ print(result2.all_messages())
                 part_kind='text',
             )
         ],
-        model_name='function:model_logic',
+        model_name='gpt-4o',
         timestamp=datetime.datetime(...),
         kind='response',
     ),
@@ -214,7 +214,7 @@ print(result2.all_messages())
                 part_kind='text',
             )
         ],
-        model_name='function:model_logic',
+        model_name='gpt-4o',
         timestamp=datetime.datetime(...),
         kind='response',
     ),
@@ -273,7 +273,7 @@ print(result2.all_messages())
                 part_kind='text',
             )
         ],
-        model_name='function:model_logic',
+        model_name='gpt-4o',
         timestamp=datetime.datetime(...),
         kind='response',
     ),
@@ -294,7 +294,7 @@ print(result2.all_messages())
                 part_kind='text',
             )
         ],
-        model_name='function:model_logic',
+        model_name='gemini-1.5-pro',
         timestamp=datetime.datetime(...),
         kind='response',
     ),

Diff for: docs/models.md

@@ -694,15 +694,15 @@ print(response.all_messages())
     ),
     ModelResponse(
         parts=[TextPart(content='Paris', part_kind='text')],
-        model_name='function:model_logic',
+        model_name='claude-3-5-sonnet-latest',
         timestamp=datetime.datetime(...),
         kind='response',
     ),
 ]
 """
 ```
 
-The `ModelResponse` message above indicates that the result was returned by the Anthropic model, which is the second model specified in the `FallbackModel`.
+The `ModelResponse` message above indicates in the `model_name` field that the result was returned by the Anthropic model, which is the second model specified in the `FallbackModel`.
 
 !!! note
     Each model's options should be configured individually. For example, `base_url`, `api_key`, and custom clients should be set on each model itself, not on the `FallbackModel`.

Diff for: docs/tools.md

@@ -89,7 +89,7 @@ print(dice_result.all_messages())
                 tool_name='roll_die', args={}, tool_call_id=None, part_kind='tool-call'
             )
         ],
-        model_name='function:model_logic',
+        model_name='gemini-1.5-flash',
         timestamp=datetime.datetime(...),
         kind='response',
     ),
@@ -114,7 +114,7 @@ print(dice_result.all_messages())
                 part_kind='tool-call',
             )
         ],
-        model_name='function:model_logic',
+        model_name='gemini-1.5-flash',
         timestamp=datetime.datetime(...),
         kind='response',
     ),
@@ -137,7 +137,7 @@ print(dice_result.all_messages())
                 part_kind='text',
             )
         ],
-        model_name='function:model_logic',
+        model_name='gemini-1.5-flash',
         timestamp=datetime.datetime(...),
         kind='response',
     ),

Diff for: pydantic_ai_slim/pydantic_ai/agent.py

@@ -365,7 +365,7 @@ async def main():
                 HandleResponseNode(
                     model_response=ModelResponse(
                         parts=[TextPart(content='Paris', part_kind='text')],
-                        model_name='function:model_logic',
+                        model_name='gpt-4o',
                         timestamp=datetime.datetime(...),
                         kind='response',
                     )
@@ -1214,7 +1214,7 @@ async def main():
             HandleResponseNode(
                 model_response=ModelResponse(
                     parts=[TextPart(content='Paris', part_kind='text')],
-                    model_name='function:model_logic',
+                    model_name='gpt-4o',
                     timestamp=datetime.datetime(...),
                     kind='response',
                 )
@@ -1357,7 +1357,7 @@ async def main():
                     HandleResponseNode(
                         model_response=ModelResponse(
                             parts=[TextPart(content='Paris', part_kind='text')],
-                            model_name='function:model_logic',
+                            model_name='gpt-4o',
                             timestamp=datetime.datetime(...),
                             kind='response',
                         )

Diff for: pydantic_ai_slim/pydantic_ai/models/function.py

@@ -48,22 +48,31 @@ class FunctionModel(Model):
     _system: str | None = field(default=None, repr=False)
 
     @overload
-    def __init__(self, function: FunctionDef) -> None: ...
+    def __init__(self, function: FunctionDef, *, model_name: str | None = None) -> None: ...
 
     @overload
-    def __init__(self, *, stream_function: StreamFunctionDef) -> None: ...
+    def __init__(self, *, stream_function: StreamFunctionDef, model_name: str | None = None) -> None: ...
 
     @overload
-    def __init__(self, function: FunctionDef, *, stream_function: StreamFunctionDef) -> None: ...
+    def __init__(
+        self, function: FunctionDef, *, stream_function: StreamFunctionDef, model_name: str | None = None
+    ) -> None: ...
 
-    def __init__(self, function: FunctionDef | None = None, *, stream_function: StreamFunctionDef | None = None):
+    def __init__(
+        self,
+        function: FunctionDef | None = None,
+        *,
+        stream_function: StreamFunctionDef | None = None,
+        model_name: str | None = None,
+    ):
         """Initialize a `FunctionModel`.
 
         Either `function` or `stream_function` must be provided, providing both is allowed.
 
         Args:
             function: The function to call for non-streamed requests.
             stream_function: The function to call for streamed requests.
+            model_name: The name of the model. If not provided, a name is generated from the function names.
         """
         if function is None and stream_function is None:
             raise TypeError('Either `function` or `stream_function` must be provided')
@@ -72,7 +81,7 @@ def __init__(self, function: FunctionDef | None = None, *, stream_function: Stre
 
         function_name = self.function.__name__ if self.function is not None else ''
         stream_function_name = self.stream_function.__name__ if self.stream_function is not None else ''
-        self._model_name = f'function:{function_name}:{stream_function_name}'
+        self._model_name = model_name or f'function:{function_name}:{stream_function_name}'
 
     async def request(
         self,
@@ -95,7 +104,7 @@ async def request(
             response_ = await _utils.run_in_executor(self.function, messages, agent_info)
             assert isinstance(response_, ModelResponse), response_
             response = response_
-        response.model_name = f'function:{self.function.__name__}'
+        response.model_name = self._model_name
         # TODO is `messages` right here? Should it just be new messages?
         return response, _estimate_usage(chain(messages, [response]))
 

Diff for: tests/test_examples.py

@@ -16,6 +16,7 @@
 from pytest_examples import CodeExample, EvalExample, find_examples
 from pytest_mock import MockerFixture
 
+from pydantic_ai import ModelHTTPError
 from pydantic_ai._utils import group_by_temporal
 from pydantic_ai.exceptions import UnexpectedModelBehavior
 from pydantic_ai.messages import (
@@ -27,8 +28,10 @@
     ToolReturnPart,
     UserPromptPart,
 )
-from pydantic_ai.models import KnownModelName, Model
+from pydantic_ai.models import KnownModelName, Model, infer_model
+from pydantic_ai.models.fallback import FallbackModel
 from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel
+from pydantic_ai.models.openai import OpenAIModel
 from pydantic_ai.models.test import TestModel
 
 from .conftest import ClientWithHandler, TestEnv
@@ -77,6 +80,7 @@ def test_docs_examples(
     env.set('OPENAI_API_KEY', 'testing')
     env.set('GEMINI_API_KEY', 'testing')
     env.set('GROQ_API_KEY', 'testing')
+    env.set('CO_API_KEY', 'testing')
 
     sys.path.append('tests/example_modules')
 
@@ -377,12 +381,32 @@ async def stream_model_logic(
 
 
 def mock_infer_model(model: Model | KnownModelName) -> Model:
+    if model == 'test':
+        return TestModel()
+
+    if isinstance(model, str):
+        # Use the non-mocked model inference to ensure we get the same model name the user would
+        model = infer_model(model)
+
+    if isinstance(model, FallbackModel):
+        # When a fallback model is encountered, replace any OpenAIModel with a model that will raise a ModelHTTPError.
+        # Otherwise, do the usual inference.
+        def raise_http_error(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:  # pragma: no cover
+            raise ModelHTTPError(401, 'Invalid API Key')
+
+        mock_fallback_models: list[Model] = []
+        for m in model.models:
+            if isinstance(m, OpenAIModel):
+                # Raise an HTTP error for OpenAIModel
+                mock_fallback_models.append(FunctionModel(raise_http_error, model_name=m.model_name))
+            else:
+                mock_fallback_models.append(mock_infer_model(m))
+        return FallbackModel(*mock_fallback_models)
     if isinstance(model, (FunctionModel, TestModel)):
         return model
-    elif model == 'test':
-        return TestModel()
     else:
-        return FunctionModel(model_logic, stream_function=stream_model_logic)
+        model_name = model if isinstance(model, str) else model.model_name
+        return FunctionModel(model_logic, stream_function=stream_model_logic, model_name=model_name)
 
 
 def mock_group_by_temporal(aiter: Any, soft_max_interval: float | None) -> Any: