pydantic-ai 1.0.6__tar.gz → 1.0.7__tar.gz

{pydantic_ai-1.0.6 → pydantic_ai-1.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai
-Version: 1.0.6
+Version: 1.0.7
 Summary: Agent Framework / shim to use Pydantic with LLMs
 Project-URL: Homepage, https://ai.pydantic.dev
 Project-URL: Source, https://github.com/pydantic/pydantic-ai
@@ -27,13 +27,13 @@ Classifier: Topic :: Internet
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.10
-Requires-Dist: pydantic-ai-slim[ag-ui,anthropic,bedrock,cli,cohere,evals,google,groq,huggingface,logfire,mcp,mistral,openai,retries,temporal,vertexai]==1.0.6
+Requires-Dist: pydantic-ai-slim[ag-ui,anthropic,bedrock,cli,cohere,evals,google,groq,huggingface,logfire,mcp,mistral,openai,retries,temporal,vertexai]==1.0.7
 Provides-Extra: a2a
 Requires-Dist: fasta2a>=0.4.1; extra == 'a2a'
 Provides-Extra: dbos
-Requires-Dist: pydantic-ai-slim[dbos]==1.0.6; extra == 'dbos'
+Requires-Dist: pydantic-ai-slim[dbos]==1.0.7; extra == 'dbos'
 Provides-Extra: examples
-Requires-Dist: pydantic-ai-examples==1.0.6; extra == 'examples'
+Requires-Dist: pydantic-ai-examples==1.0.7; extra == 'examples'
 Description-Content-Type: text/markdown
 
 <div align="center">

{pydantic_ai-1.0.6 → pydantic_ai-1.0.7}/tests/mcp_server.py

@@ -11,6 +11,7 @@ from mcp.types import (
     SamplingMessage,
     TextContent,
     TextResourceContents,
+    ToolAnnotations,
 )
 from pydantic import AnyUrl, BaseModel
 
@@ -18,7 +19,7 @@ mcp = FastMCP('Pydantic AI MCP Server')
 log_level = 'unset'
 
 
-@mcp.tool()
+@mcp.tool(annotations=ToolAnnotations(title='Celsius to Fahrenheit'))
 async def celsius_to_fahrenheit(celsius: float) -> float:
     """Convert Celsius to Fahrenheit.
 

{pydantic_ai-1.0.6 → pydantic_ai-1.0.7}/tests/test_a2a.py

@@ -623,10 +623,10 @@ async def test_a2a_multiple_tasks_same_context():
                                 content='Final result processed.',
                                 tool_call_id=IsStr(),
                                 timestamp=IsDatetime(),
-                            )
-                        ]
+                            ),
+                            UserPromptPart(content='Second message', timestamp=IsDatetime()),
+                        ],
                     ),
-                    ModelRequest(parts=[UserPromptPart(content='Second message', timestamp=IsDatetime())]),
                 ]
             )
 

{pydantic_ai-1.0.6 → pydantic_ai-1.0.7}/tests/test_ag_ui.py

@@ -21,7 +21,16 @@ from pydantic import BaseModel
 from pydantic_ai._run_context import RunContext
 from pydantic_ai.agent import Agent, AgentRunResult
 from pydantic_ai.exceptions import UserError
-from pydantic_ai.messages import ModelMessage
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
 from pydantic_ai.models.function import (
     AgentInfo,
     DeltaThinkingCalls,
@@ -34,7 +43,7 @@ from pydantic_ai.models.test import TestModel
 from pydantic_ai.output import OutputDataT
 from pydantic_ai.tools import AgentDepsT, ToolDefinition
 
-from .conftest import IsSameStr
+from .conftest import IsDatetime, IsSameStr
 
 has_ag_ui: bool = False
 with contextlib.suppress(ImportError):
@@ -59,6 +68,7 @@ with contextlib.suppress(ImportError):
         SSE_CONTENT_TYPE,
         OnCompleteFunc,
         StateDeps,
+        _messages_from_ag_ui,  # type: ignore[reportPrivateUsage]
         run_ag_ui,
     )
 
@@ -1347,3 +1357,127 @@ async def test_callback_with_error() -> None:
     assert len(events) > 0
     assert events[0]['type'] == 'RUN_STARTED'
     assert any(event['type'] == 'RUN_ERROR' for event in events)
+
+
+async def test_messages_from_ag_ui() -> None:
+    messages = [
+        SystemMessage(
+            id='msg_1',
+            content='System message',
+        ),
+        DeveloperMessage(
+            id='msg_2',
+            content='Developer message',
+        ),
+        UserMessage(
+            id='msg_3',
+            content='User message',
+        ),
+        UserMessage(
+            id='msg_4',
+            content='User message',
+        ),
+        AssistantMessage(
+            id='msg_5',
+            content='Assistant message',
+        ),
+        AssistantMessage(
+            id='msg_6',
+            tool_calls=[
+                ToolCall(
+                    id='tool_call_1',
+                    function=FunctionCall(
+                        name='tool_call_1',
+                        arguments='{}',
+                    ),
+                ),
+            ],
+        ),
+        AssistantMessage(
+            id='msg_7',
+            tool_calls=[
+                ToolCall(
+                    id='tool_call_2',
+                    function=FunctionCall(
+                        name='tool_call_2',
+                        arguments='{}',
+                    ),
+                ),
+            ],
+        ),
+        ToolMessage(
+            id='msg_8',
+            content='Tool message',
+            tool_call_id='tool_call_1',
+        ),
+        ToolMessage(
+            id='msg_9',
+            content='Tool message',
+            tool_call_id='tool_call_2',
+        ),
+        UserMessage(
+            id='msg_10',
+            content='User message',
+        ),
+        AssistantMessage(
+            id='msg_11',
+            content='Assistant message',
+        ),
+    ]
+
+    assert _messages_from_ag_ui(messages) == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    SystemPromptPart(
+                        content='System message',
+                        timestamp=IsDatetime(),
+                    ),
+                    SystemPromptPart(
+                        content='Developer message',
+                        timestamp=IsDatetime(),
+                    ),
+                    UserPromptPart(
+                        content='User message',
+                        timestamp=IsDatetime(),
+                    ),
+                    UserPromptPart(
+                        content='User message',
+                        timestamp=IsDatetime(),
+                    ),
+                ]
+            ),
+            ModelResponse(
+                parts=[
+                    TextPart(content='Assistant message'),
+                    ToolCallPart(tool_name='tool_call_1', args='{}', tool_call_id='tool_call_1'),
+                    ToolCallPart(tool_name='tool_call_2', args='{}', tool_call_id='tool_call_2'),
+                ],
+                timestamp=IsDatetime(),
+            ),
+            ModelRequest(
+                parts=[
+                    ToolReturnPart(
+                        tool_name='tool_call_1',
+                        content='Tool message',
+                        tool_call_id='tool_call_1',
+                        timestamp=IsDatetime(),
+                    ),
+                    ToolReturnPart(
+                        tool_name='tool_call_2',
+                        content='Tool message',
+                        tool_call_id='tool_call_2',
+                        timestamp=IsDatetime(),
+                    ),
+                    UserPromptPart(
+                        content='User message',
+                        timestamp=IsDatetime(),
+                    ),
+                ]
+            ),
+            ModelResponse(
+                parts=[TextPart(content='Assistant message')],
+                timestamp=IsDatetime(),
+            ),
+        ]
+    )

{pydantic_ai-1.0.6 → pydantic_ai-1.0.7}/tests/test_agent.py

@@ -1797,7 +1797,7 @@ def test_run_with_history_new():
             ),
         ]
     )
-    assert result2._new_message_index == snapshot(4)  # pyright: ignore[reportPrivateUsage]
+    assert result2.new_messages() == result2.all_messages()[-2:]
     assert result2.output == snapshot('{"ret_a":"a-apple"}')
     assert result2._output_tool_name == snapshot(None)  # pyright: ignore[reportPrivateUsage]
     assert result2.usage() == snapshot(RunUsage(requests=1, input_tokens=55, output_tokens=13))
@@ -1854,7 +1854,7 @@ def test_run_with_history_new():
             ),
         ]
     )
-    assert result3._new_message_index == snapshot(4)  # pyright: ignore[reportPrivateUsage]
+    assert result3.new_messages() == result3.all_messages()[-2:]
     assert result3.output == snapshot('{"ret_a":"a-apple"}')
     assert result3._output_tool_name == snapshot(None)  # pyright: ignore[reportPrivateUsage]
     assert result3.usage() == snapshot(RunUsage(requests=1, input_tokens=55, output_tokens=13))
@@ -1983,7 +1983,7 @@ def test_run_with_history_new_structured():
         ]
     )
     assert result2.output == snapshot(Response(a=0))
-    assert result2._new_message_index == snapshot(5)  # pyright: ignore[reportPrivateUsage]
+    assert result2.new_messages() == result2.all_messages()[-3:]
     assert result2._output_tool_name == snapshot('final_result')  # pyright: ignore[reportPrivateUsage]
     assert result2.usage() == snapshot(RunUsage(requests=1, input_tokens=59, output_tokens=13))
     new_msg_part_kinds = [(m.kind, [p.part_kind for p in m.parts]) for m in result2.all_messages()]
@@ -2003,34 +2003,61 @@ def test_run_with_history_new_structured():
 
 
 def test_run_with_history_ending_on_model_request_and_no_user_prompt():
+    m = TestModel()
+    agent = Agent(m)
+
+    @agent.system_prompt(dynamic=True)
+    async def system_prompt(ctx: RunContext) -> str:
+        return f'System prompt: user prompt length = {len(ctx.prompt or [])}'
+
     messages: list[ModelMessage] = [
-        ModelRequest(parts=[UserPromptPart(content='Hello')], instructions='Original instructions'),
+        ModelRequest(
+            parts=[
+                SystemPromptPart(content='System prompt', dynamic_ref=system_prompt.__qualname__),
+                UserPromptPart(content=['Hello', ImageUrl('https://example.com/image.jpg')]),
+                UserPromptPart(content='How goes it?'),
+            ],
+            instructions='Original instructions',
+        ),
     ]
 
-    m = TestModel()
-    agent = Agent(m, instructions='New instructions')
+    @agent.instructions
+    async def instructions(ctx: RunContext) -> str:
+        assert ctx.prompt == ['Hello', ImageUrl('https://example.com/image.jpg'), 'How goes it?']
+        return 'New instructions'
 
     result = agent.run_sync(message_history=messages)
     assert result.all_messages() == snapshot(
         [
             ModelRequest(
                 parts=[
+                    SystemPromptPart(
+                        content='System prompt: user prompt length = 3',
+                        timestamp=IsDatetime(),
+                        dynamic_ref=IsStr(),
+                    ),
                     UserPromptPart(
-                        content='Hello',
+                        content=['Hello', ImageUrl(url='https://example.com/image.jpg', identifier='39cfc4')],
                         timestamp=IsDatetime(),
-                    )
+                    ),
+                    UserPromptPart(
+                        content='How goes it?',
+                        timestamp=IsDatetime(),
+                    ),
                 ],
                 instructions='New instructions',
             ),
             ModelResponse(
                 parts=[TextPart(content='success (no tool calls)')],
-                usage=RequestUsage(input_tokens=51, output_tokens=4),
+                usage=RequestUsage(input_tokens=61, output_tokens=4),
                 model_name='test',
                 timestamp=IsDatetime(),
             ),
         ]
     )
 
+    assert result.new_messages() == result.all_messages()[-1:]
+
 
 def test_run_with_history_ending_on_model_response_with_tool_calls_and_no_user_prompt():
     """Test that an agent run with message_history ending on ModelResponse starts with CallToolsNode."""
@@ -2084,6 +2111,8 @@ def test_run_with_history_ending_on_model_response_with_tool_calls_and_no_user_p
         ]
     )
 
+    assert result.new_messages() == result.all_messages()[-2:]
+
 
 def test_run_with_history_ending_on_model_response_with_tool_calls_and_user_prompt():
     """Test that an agent run raises error when message_history ends on ModelResponse with tool calls and there's a new prompt."""
@@ -2136,6 +2165,8 @@ def test_run_with_history_ending_on_model_response_without_tool_calls_or_user_pr
         ]
     )
 
+    assert result.new_messages() == []
+
 
 def test_empty_tool_calls():
     def empty(_: list[ModelMessage], _info: AgentInfo) -> ModelResponse:
@@ -2941,6 +2972,8 @@ def test_dynamic_false_no_reevaluate():
         ]
     )
 
+    assert res_two.new_messages() == res_two.all_messages()[-2:]
+
 
 def test_dynamic_true_reevaluate_system_prompt():
     """When dynamic is true, the system prompt is reevaluated
@@ -3023,6 +3056,8 @@ def test_dynamic_true_reevaluate_system_prompt():
         ]
     )
 
+    assert res_two.new_messages() == res_two.all_messages()[-2:]
+
 
 def test_dynamic_system_prompt_no_changes():
     """Test coverage for _reevaluate_dynamic_prompts branch where no parts are changed
@@ -3553,6 +3588,8 @@ def test_instructions_with_message_history():
         ]
     )
 
+    assert result.new_messages() == result.all_messages()[-2:]
+
 
 def test_instructions_parameter_with_sequence():
     def instructions() -> str:
@@ -4675,7 +4712,11 @@ async def test_hitl_tool_approval():
                         content='File \'new_file.py\' created with content: print("Hello, world!")',
                         tool_call_id='create_file',
                         timestamp=IsDatetime(),
-                    ),
+                    )
+                ]
+            ),
+            ModelRequest(
+                parts=[
                     ToolReturnPart(
                         tool_name='delete_file',
                         content="File 'ok_to_delete.py' deleted",
@@ -4700,6 +4741,33 @@ async def test_hitl_tool_approval():
     )
     assert result.output == snapshot('Done!')
 
+    assert result.new_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    ToolReturnPart(
+                        tool_name='delete_file',
+                        content="File 'ok_to_delete.py' deleted",
+                        tool_call_id='ok_to_delete',
+                        timestamp=IsDatetime(),
+                    ),
+                    ToolReturnPart(
+                        tool_name='delete_file',
+                        content='File cannot be deleted',
+                        tool_call_id='never_delete',
+                        timestamp=IsDatetime(),
+                    ),
+                ]
+            ),
+            ModelResponse(
+                parts=[TextPart(content='Done!')],
+                usage=RequestUsage(input_tokens=78, output_tokens=24),
+                model_name='function:model_function:',
+                timestamp=IsDatetime(),
+            ),
+        ]
+    )
+
 
 async def test_run_with_deferred_tool_results_errors():
     agent = Agent('test')
@@ -4708,7 +4776,7 @@ async def test_run_with_deferred_tool_results_errors():
 
     with pytest.raises(
         UserError,
-        match='Tool call results were provided, but the last message in the history was a `ModelRequest` with user parts not tied to preliminary tool results.',
+        match='Tool call results were provided, but the message history does not contain a `ModelResponse`.',
     ):
         await agent.run(
             'Hello again',
@@ -4753,6 +4821,15 @@ async def test_run_with_deferred_tool_results_errors():
             deferred_tool_results=DeferredToolResults(approvals={'create_file': True}),
         )
 
+    with pytest.raises(
+        UserError, match='Cannot provide a new user prompt when the message history contains unprocessed tool calls.'
+    ):
+        await agent.run(
+            'Hello again',
+            message_history=message_history,
+            deferred_tool_results=DeferredToolResults(approvals={'create_file': True}),
+        )
+
     message_history: list[ModelMessage] = [
         ModelRequest(parts=[UserPromptPart(content='Hello')]),
         ModelResponse(
@@ -4772,7 +4849,6 @@ async def test_run_with_deferred_tool_results_errors():
 
     with pytest.raises(UserError, match="Tool call 'run_me' was already executed and its result cannot be overridden."):
         await agent.run(
-            'Hello again',
             message_history=message_history,
             deferred_tool_results=DeferredToolResults(
                 calls={'run_me': 'Failure', 'defer_me': 'Failure'},
@@ -4783,7 +4859,6 @@ async def test_run_with_deferred_tool_results_errors():
         UserError, match="Tool call 'run_me_too' was already executed and its result cannot be overridden."
     ):
         await agent.run(
-            'Hello again',
             message_history=message_history,
             deferred_tool_results=DeferredToolResults(
                 calls={'run_me_too': 'Success', 'defer_me': 'Failure'},
@@ -4802,3 +4877,101 @@ def test_tool_requires_approval_error():
         @agent.tool_plain(requires_approval=True)
         def delete_file(path: str) -> None:
             pass
+
+
+async def test_consecutive_model_responses_in_history():
+    received_messages: list[ModelMessage] | None = None
+
+    def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+        nonlocal received_messages
+        received_messages = messages
+        return ModelResponse(
+            parts=[
+                TextPart('All right then, goodbye!'),
+            ]
+        )
+
+    history: list[ModelMessage] = [
+        ModelRequest(parts=[UserPromptPart(content='Hello...')]),
+        ModelResponse(parts=[TextPart(content='...world!')]),
+        ModelResponse(parts=[TextPart(content='Anything else I can help with?')]),
+    ]
+
+    m = FunctionModel(llm)
+    agent = Agent(m)
+    result = await agent.run('No thanks', message_history=history)
+
+    assert result.all_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='Hello...',
+                        timestamp=IsDatetime(),
+                    )
+                ]
+            ),
+            ModelResponse(
+                parts=[TextPart(content='...world!'), TextPart(content='Anything else I can help with?')],
+                timestamp=IsDatetime(),
+            ),
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='No thanks',
+                        timestamp=IsDatetime(),
+                    )
+                ]
+            ),
+            ModelResponse(
+                parts=[TextPart(content='All right then, goodbye!')],
+                usage=RequestUsage(input_tokens=54, output_tokens=12),
+                model_name='function:llm:',
+                timestamp=IsDatetime(),
+            ),
+        ]
+    )
+
+    assert result.new_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='No thanks',
+                        timestamp=IsDatetime(),
+                    )
+                ]
+            ),
+            ModelResponse(
+                parts=[TextPart(content='All right then, goodbye!')],
+                usage=RequestUsage(input_tokens=54, output_tokens=12),
+                model_name='function:llm:',
+                timestamp=IsDatetime(),
+            ),
+        ]
+    )
+
+    assert received_messages == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='Hello...',
+                        timestamp=IsDatetime(),
+                    )
+                ]
+            ),
+            ModelResponse(
+                parts=[TextPart(content='...world!'), TextPart(content='Anything else I can help with?')],
+                timestamp=IsDatetime(),
+            ),
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='No thanks',
+                        timestamp=IsDatetime(),
+                    )
+                ]
+            ),
+        ]
+    )

{pydantic_ai-1.0.6 → pydantic_ai-1.0.7}/tests/test_history_processor.py

@@ -1,5 +1,5 @@
 from collections.abc import AsyncIterator
-from typing import Any, cast
+from typing import Any
 
 import pytest
 from inline_snapshot import snapshot
@@ -174,8 +174,18 @@ async def test_history_processor_messages_sent_to_provider(
     )
     assert received_messages == snapshot(
         [
-            ModelRequest(parts=[UserPromptPart(content='Previous question', timestamp=IsDatetime())]),
-            ModelRequest(parts=[UserPromptPart(content='New question', timestamp=IsDatetime())]),
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='Previous question',
+                        timestamp=IsDatetime(),
+                    ),
+                    UserPromptPart(
+                        content='New question',
+                        timestamp=IsDatetime(),
+                    ),
+                ]
+            )
         ]
     )
 
@@ -244,8 +254,18 @@ async def test_async_history_processor(function_model: FunctionModel, received_m
     await agent.run('Question 2', message_history=message_history)
     assert received_messages == snapshot(
         [
-            ModelRequest(parts=[UserPromptPart(content='Question 1', timestamp=IsDatetime())]),
-            ModelRequest(parts=[UserPromptPart(content='Question 2', timestamp=IsDatetime())]),
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='Question 1',
+                        timestamp=IsDatetime(),
+                    ),
+                    UserPromptPart(
+                        content='Question 2',
+                        timestamp=IsDatetime(),
+                    ),
+                ]
+            )
         ]
     )
 
@@ -271,8 +291,18 @@ async def test_history_processor_on_streamed_run(function_model: FunctionModel,
 
     assert received_messages == snapshot(
         [
-            ModelRequest(parts=[UserPromptPart(content='Question 1', timestamp=IsDatetime())]),
-            ModelRequest(parts=[UserPromptPart(content='Question 2', timestamp=IsDatetime())]),
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='Question 1',
+                        timestamp=IsDatetime(),
+                    ),
+                    UserPromptPart(
+                        content='Question 2',
+                        timestamp=IsDatetime(),
+                    ),
+                ]
+            )
         ]
     )
 
@@ -367,9 +397,19 @@ async def test_history_processor_mixed_signatures(function_model: FunctionModel,
     await agent.run('Question 2', message_history=message_history, deps=Deps())
 
     # Should have filtered responses and added prefix
-    assert len(received_messages) == 2
-    for msg in received_messages:
-        assert isinstance(msg, ModelRequest)
-        user_part = msg.parts[0]
-        assert isinstance(user_part, UserPromptPart)
-        assert cast(str, user_part.content).startswith('TEST: ')
+    assert received_messages == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='TEST: Question 1',
+                        timestamp=IsDatetime(),
+                    ),
+                    UserPromptPart(
+                        content='TEST: Question 2',
+                        timestamp=IsDatetime(),
+                    ),
+                ]
+            )
+        ]
+    )

{pydantic_ai-1.0.6 → pydantic_ai-1.0.7}/tests/test_logfire.py

@@ -387,6 +387,7 @@ def test_logfire(
                                 'strict': None,
                                 'sequential': False,
                                 'kind': 'function',
+                                'metadata': None,
                             }
                         ],
                         'builtin_tools': [],
@@ -780,6 +781,7 @@ def test_instructions_with_structured_output_exclude_content_v2(
                                 'strict': None,
                                 'sequential': False,
                                 'kind': 'output',
+                                'metadata': None,
                             }
                         ],
                         'allow_text_output': False,

{pydantic_ai-1.0.6 → pydantic_ai-1.0.7}/tests/test_mcp.py

@@ -1252,6 +1252,22 @@ async def test_tool_returning_multiple_items(allow_model_requests: None, agent:
         )
 
 
+async def test_tool_metadata_extraction():
+    """Test that MCP tool metadata is properly extracted into ToolDefinition."""
+
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    async with server:
+        ctx = RunContext(deps=None, model=TestModel(), usage=RunUsage())
+        tools = [tool.tool_def for tool in (await server.get_tools(ctx)).values()]
+        # find `celsius_to_fahrenheit`
+        celsius_to_fahrenheit = next(tool for tool in tools if tool.name == 'celsius_to_fahrenheit')
+        assert celsius_to_fahrenheit.metadata is not None
+        assert celsius_to_fahrenheit.metadata.get('annotations') is not None
+        assert celsius_to_fahrenheit.metadata.get('annotations', {}).get('title', None) == 'Celsius to Fahrenheit'
+        assert celsius_to_fahrenheit.metadata.get('output_schema') is not None
+        assert celsius_to_fahrenheit.metadata.get('output_schema', {}).get('type', None) == 'object'
+
+
 async def test_client_sampling(run_context: RunContext[int]):
     server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
     server.sampling_model = TestModel(custom_output_text='sampling model response')