pydantic-ai 0.1.9__tar.gz → 0.1.11__tar.gz

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai
-Version: 0.1.9
+Version: 0.1.11
 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
@@ -28,9 +28,9 @@ Classifier: Topic :: Internet
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.9
-Requires-Dist: pydantic-ai-slim[anthropic,bedrock,cli,cohere,evals,groq,mcp,mistral,openai,vertexai]==0.1.9
+Requires-Dist: pydantic-ai-slim[anthropic,bedrock,cli,cohere,evals,groq,mcp,mistral,openai,vertexai]==0.1.11
 Provides-Extra: examples
-Requires-Dist: pydantic-ai-examples==0.1.9; extra == 'examples'
+Requires-Dist: pydantic-ai-examples==0.1.11; extra == 'examples'
 Provides-Extra: logfire
 Requires-Dist: logfire>=3.11.0; extra == 'logfire'
 Description-Content-Type: text/markdown

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/pyproject.toml

@@ -60,14 +60,13 @@ Documentation = "https://ai.pydantic.dev"
 Changelog = "https://github.com/pydantic/pydantic-ai/releases"
 
 [project.scripts]
-pai = "pydantic_ai._cli:app"
+pai = "pydantic_ai._cli:cli_exit"  # TODO remove this when clai has been out for a while
 
 [tool.uv.sources]
 pydantic-ai-slim = { workspace = true }
 pydantic-evals = { workspace = true }
 pydantic-graph = { workspace = true }
 pydantic-ai-examples = { workspace = true }
-mcp-run-python = { workspace = true }
 
 [tool.uv.workspace]
 members = [
@@ -75,6 +74,7 @@ members = [
     "pydantic_evals",
     "pydantic_graph",
     "mcp-run-python",
+    "clai",
     "examples",
 ]
 
@@ -106,6 +106,7 @@ include = [
     "pydantic_graph/**/*.py",
     "mcp-run-python/**/*.py",
     "examples/**/*.py",
+    "clai/**/*.py",
     "tests/**/*.py",
     "docs/**/*.py",
 ]
@@ -114,7 +115,7 @@ include = [
 extend-select = [
     "Q",
     "RUF100",
-    "RUF018",  # https://docs.astral.sh/ruff/rules/assignment-in-assert/
+    "RUF018", # https://docs.astral.sh/ruff/rules/assignment-in-assert/
     "C90",
     "UP",
     "I",
@@ -163,6 +164,7 @@ include = [
     "mcp-run-python",
     "tests",
     "examples",
+    "clai",
 ]
 venvPath = ".venv"
 # see https://github.com/microsoft/pyright/issues/7771 - we don't want to error on decorated functions in tests
@@ -193,7 +195,7 @@ filterwarnings = [
     "ignore:websockets.server.WebSocketServerProtocol is deprecated:DeprecationWarning",
     # random resource warnings; I suspect these are coming from vendor SDKs when running examples..
     "ignore:unclosed <socket:ResourceWarning",
-    "ignore:unclosed event loop:ResourceWarning"
+    "ignore:unclosed event loop:ResourceWarning",
 ]
 
 # https://coverage.readthedocs.io/en/latest/config.html#run
@@ -215,7 +217,10 @@ show_missing = true
 ignore_errors = true
 precision = 2
 exclude_lines = [
+    # `# pragma: no cover` is standard marker for code that's not covered, this will error if code is covered
     'pragma: no cover',
+    # use `# pragma: lax no cover` if you want to ignore cases where (some of) the code is covered
+    'pragma: lax no cover',
     'raise NotImplementedError',
     'if TYPE_CHECKING:',
     'if typing.TYPE_CHECKING:',

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/conftest.py

@@ -59,7 +59,7 @@ class TestEnv:
     def remove(self, name: str) -> None:
         self.envars[name] = os.environ.pop(name, None)
 
-    def reset(self) -> None:  # pragma: no cover
+    def reset(self) -> None:
         for name, value in self.envars.items():
             if value is None:
                 os.environ.pop(name, None)
@@ -100,7 +100,7 @@ async def client_with_handler() -> AsyncIterator[ClientWithHandler]:
     try:
         yield create_client
     finally:
-        if client:  # pragma: no cover
+        if client:
             await client.aclose()
 
 
@@ -155,7 +155,7 @@ def create_module(tmp_path: Path, request: pytest.FixtureRequest) -> Callable[[s
 
 
 @contextmanager
-def try_import() -> Iterator[Callable[[], bool]]:  # pragma: no cover
+def try_import() -> Iterator[Callable[[], bool]]:
     import_success = False
 
     def check_import() -> bool:
@@ -276,7 +276,7 @@ def mistral_api_key() -> str:
 
 
 @pytest.fixture(scope='session')
-def bedrock_provider():  # pragma: no cover
+def bedrock_provider():
     try:
         import boto3
 
@@ -304,7 +304,7 @@ def model(
     co_api_key: str,
     gemini_api_key: str,
     bedrock_provider: BedrockProvider,
-) -> Model:  # pragma: no cover
+) -> Model:  # pragma: lax no cover
     try:
         if request.param == 'openai':
             from pydantic_ai.models.openai import OpenAIModel

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/evals/test_evaluator_common.py

@@ -40,7 +40,7 @@ if TYPE_CHECKING or imports_successful():
             self.expected_output = expected_output
             self.inputs = inputs
             self.duration = duration
-else:  # pragma: no cover
+else:
     MockContext = object
 
 

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/evals/test_otel.py

@@ -879,7 +879,7 @@ async def test_context_subtree_invalid_tracer_provider(mocker: MockerFixture):
 
     mocker.patch('pydantic_evals.otel._context_in_memory_span_exporter.get_tracer_provider', return_value=None)
     with pytest.raises(TypeError) as exc_info:
-        with context_subtree():  # pragma: no cover
+        with context_subtree():
             pass
     assert str(exc_info.value) == snapshot(
         "Expected `tracer_provider` to have an `add_span_processor` method; got an instance of <class 'NoneType'>. For help resolving this, please create an issue at https://github.com/pydantic/pydantic-ai/issues."

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/graph/test_mermaid.py

@@ -259,7 +259,7 @@ def httpx_with_handler() -> Iterator[HttpxWithHandler]:
     try:
         yield create_client
     finally:
-        if client:  # pragma: no cover
+        if client:
             client.close()
 
 

pydantic_ai-0.1.11/tests/models/cassettes/test_anthropic/test_extra_headers.yaml

@@ -0,0 +1,57 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '202'
+      content-type:
+      - application/json
+      host:
+      - api.anthropic.com
+      xProxy-Limit-IDs:
+      - monthly_budget
+    method: POST
+    parsed_body:
+      max_tokens: 1024
+      messages:
+      - content:
+        - text: hello
+          type: text
+        role: user
+      model: claude-3-opus-latest
+      stream: false
+    uri: https://api.anthropic.com/v1/messages
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '328'
+      content-type:
+      - application/json
+      transfer-encoding:
+      - chunked
+    parsed_body:
+      content:
+      - text: Hello! How can I assist you today?
+        type: text
+      id: msg_01U58nruzfn9BrXrrF2hhb4m
+      model: claude-3-5-haiku-latest
+      role: assistant
+      stop_reason: end_turn
+      stop_sequence: null
+      type: message
+      usage:
+        cache_creation_input_tokens: 0
+        cache_read_input_tokens: 0
+        input_tokens: 20
+        output_tokens: 10
+    status:
+      code: 200
+      message: OK
+version: 1

pydantic_ai-0.1.11/tests/models/cassettes/test_groq/test_extra_headers.yaml

@@ -0,0 +1,71 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '187'
+      content-type:
+      - application/json
+      host:
+      - api.groq.com
+      xProxy-Limit-IDs:
+      - monthly_budget
+    method: POST
+    parsed_body:
+      messages:
+      - content: hello
+        role: user
+      model: llama-3.3-70b-versatile
+      n: 1
+      stream: false
+    uri: https://api.groq.com/openai/v1/chat/completions
+  response:
+    headers:
+      alt-svc:
+      - h3=":443"; ma=86400
+      cache-control:
+      - private, max-age=0, no-store, no-cache, must-revalidate
+      connection:
+      - keep-alive
+      content-length:
+      - '570'
+      content-type:
+      - application/json
+      transfer-encoding:
+      - chunked
+      vary:
+      - Origin, Accept-Encoding
+    parsed_body:
+      choices:
+      - finish_reason: stop
+        index: 0
+        logprobs: null
+        message:
+          content: Hello! How can I assist you today?
+          role: assistant
+      created: 1744043573
+      id: chatcmpl-7586b6a9-fb4b-4ec7-86a0-59f0a77844cf
+      model: llama-3.3-70b-versatile
+      object: chat.completion
+      system_fingerprint: fp_72a5dc99ee
+      usage:
+        completion_time: 0.029090909
+        completion_tokens: 8
+        prompt_time: 0.002665957
+        prompt_tokens: 48
+        queue_time: 0.100731848
+        total_time: 0.031756866
+        total_tokens: 56
+      usage_breakdown:
+        models: null
+      x_groq:
+        id: req_01jr8hj0hzeq9b86xqb5dn7wqs
+    status:
+      code: 200
+      message: OK
+version: 1

pydantic_ai-0.1.11/tests/models/cassettes/test_openai/test_extra_headers.yaml

@@ -0,0 +1,80 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '103'
+      content-type:
+      - application/json
+      host:
+      - api.openai.com
+      xProxy-Limit-IDs:
+      - monthly_budget
+    method: POST
+    parsed_body:
+      messages:
+      - content: hello
+        role: user
+      model: gpt-4o
+      n: 1
+      stream: false
+    uri: https://api.openai.com/v1/chat/completions
+  response:
+    headers:
+      access-control-expose-headers:
+      - X-Request-ID
+      alt-svc:
+      - h3=":443"; ma=86400
+      connection:
+      - keep-alive
+      content-length:
+      - '835'
+      content-type:
+      - application/json
+      openai-organization:
+      - pydantic-28gund
+      openai-processing-ms:
+      - '584'
+      openai-version:
+      - '2020-10-01'
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      transfer-encoding:
+      - chunked
+    parsed_body:
+      choices:
+      - finish_reason: stop
+        index: 0
+        logprobs: null
+        message:
+          annotations: []
+          content: Hello! How can I assist you today?
+          refusal: null
+          role: assistant
+      created: 1743073438
+      id: chatcmpl-BFfJeRdAVFPUVWxV3OYH1tSR5KvrI
+      model: gpt-4o-2024-08-06
+      object: chat.completion
+      service_tier: default
+      system_fingerprint: fp_898ac29719
+      usage:
+        completion_tokens: 10
+        completion_tokens_details:
+          accepted_prediction_tokens: 0
+          audio_tokens: 0
+          reasoning_tokens: 0
+          rejected_prediction_tokens: 0
+        prompt_tokens: 8
+        prompt_tokens_details:
+          audio_tokens: 0
+          cached_tokens: 0
+        total_tokens: 18
+    status:
+      code: 200
+      message: OK
+version: 1

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/models/test_anthropic.py

@@ -634,6 +634,19 @@ async def test_image_url_input(allow_model_requests: None, anthropic_api_key: st
     )
 
 
+@pytest.mark.vcr()
+async def test_extra_headers(allow_model_requests: None, anthropic_api_key: str):
+    # This test doesn't do anything, it's just here to ensure that calls with `extra_headers` don't cause errors, including type.
+    m = AnthropicModel('claude-3-5-haiku-latest', provider=AnthropicProvider(api_key=anthropic_api_key))
+    agent = Agent(
+        m,
+        model_settings=AnthropicModelSettings(
+            anthropic_metadata={'user_id': '123'}, extra_headers={'Extra-Header-Key': 'Extra-Header-Value'}
+        ),
+    )
+    await agent.run('hello')
+
+
 @pytest.mark.vcr()
 async def test_image_url_input_invalid_mime_type(allow_model_requests: None, anthropic_api_key: str):
     m = AnthropicModel('claude-3-5-haiku-latest', provider=AnthropicProvider(api_key=anthropic_api_key))

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/models/test_bedrock.py

@@ -358,7 +358,7 @@ async def test_bedrock_model_iter_stream(allow_model_requests: None, bedrock_pro
         Args:
             city: The city name.
         """
-        return '30°C'  # pragma: no cover
+        return '30°C'
 
     event_parts: list[Any] = []
     async with agent.iter(user_prompt='What is the temperature of the capital of France?') as agent_run:

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/models/test_cohere.py

@@ -66,9 +66,7 @@ class MockAsyncClientV2:
     def create_mock(cls, completions: MockChatResponse | Sequence[MockChatResponse]) -> AsyncClientV2:
         return cast(AsyncClientV2, cls(completions=completions))
 
-    async def chat(  # pragma: no cover
-        self, *_args: Any, **_kwargs: Any
-    ) -> ChatResponse:
+    async def chat(self, *_args: Any, **_kwargs: Any) -> ChatResponse:
         assert self.completions is not None
         if isinstance(self.completions, Sequence):
             raise_if_exception(self.completions[self.index])

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/models/test_gemini.py

@@ -420,7 +420,7 @@ async def get_gemini_client(
 
 def gemini_response(content: _GeminiContent, finish_reason: Literal['STOP'] | None = 'STOP') -> _GeminiResponse:
     candidate = _GeminiCandidates(content=content, index=0, safety_ratings=[])
-    if finish_reason:  # pragma: no cover
+    if finish_reason:
         candidate['finish_reason'] = finish_reason
     return _GeminiResponse(candidates=[candidate], usage_metadata=example_usage(), model_version='gemini-1.5-flash-123')
 

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/models/test_groq.py

@@ -46,7 +46,7 @@ with try_import() as imports_successful:
     from groq.types.chat.chat_completion_message_tool_call import Function
     from groq.types.completion_usage import CompletionUsage
 
-    from pydantic_ai.models.groq import GroqModel
+    from pydantic_ai.models.groq import GroqModel, GroqModelSettings
     from pydantic_ai.providers.groq import GroqProvider
 
     # note: we use Union here so that casting works with Python 3.9
@@ -489,7 +489,7 @@ async def test_no_content(allow_model_requests: None):
 
     with pytest.raises(UnexpectedModelBehavior, match='Received empty model response'):
         async with agent.run_stream(''):
-            pass  # pragma: no cover
+            pass
 
 
 async def test_no_delta(allow_model_requests: None):
@@ -504,6 +504,14 @@ async def test_no_delta(allow_model_requests: None):
         assert result.is_complete
 
 
+@pytest.mark.vcr()
+async def test_extra_headers(allow_model_requests: None, groq_api_key: str):
+    # This test doesn't do anything, it's just here to ensure that calls with `extra_headers` don't cause errors, including type.
+    m = GroqModel('llama-3.3-70b-versatile', provider=GroqProvider(api_key=groq_api_key))
+    agent = Agent(m, model_settings=GroqModelSettings(extra_headers={'Extra-Header-Key': 'Extra-Header-Value'}))
+    await agent.run('hello')
+
+
 @pytest.mark.vcr()
 async def test_image_url_input(allow_model_requests: None, groq_api_key: str):
     m = GroqModel('meta-llama/llama-4-scout-17b-16e-instruct', provider=GroqProvider(api_key=groq_api_key))

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/models/test_mistral.py

@@ -92,7 +92,7 @@ class MockMistralAI:
     ) -> Mistral:
         return cast(Mistral, cls(stream=completions_streams))
 
-    async def chat_completions_create(  # pragma: no cover
+    async def chat_completions_create(  # pragma: lax no cover
         self, *_args: Any, stream: bool = False, **_kwargs: Any
     ) -> MistralChatCompletionResponse | MockAsyncStream[MockCompletionEvent]:
         if stream or self.stream:

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/models/test_model_function.py

@@ -92,7 +92,7 @@ def test_simple():
     )
 
 
-async def weather_model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:  # pragma: no cover
+async def weather_model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:  # pragma: lax no cover
     assert info.allow_text_output
     assert {t.name for t in info.function_tools} == {'get_location', 'get_weather'}
     last = messages[-1].parts[-1]
@@ -194,7 +194,7 @@ def test_weather():
     assert result.output == 'Sunny in Ipswich'
 
 
-async def call_function_model(messages: list[ModelMessage], _: AgentInfo) -> ModelResponse:  # pragma: no cover
+async def call_function_model(messages: list[ModelMessage], _: AgentInfo) -> ModelResponse:  # pragma: lax no cover
     last = messages[-1].parts[-1]
     if isinstance(last, UserPromptPart):
         if isinstance(last.content, str) and last.content.startswith('{'):

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/models/test_openai.py

@@ -99,7 +99,7 @@ class MockOpenAI:
     ) -> AsyncOpenAI:
         return cast(AsyncOpenAI, cls(stream=stream))
 
-    async def chat_completions_create(  # pragma: no cover
+    async def chat_completions_create(  # pragma: lax no cover
         self, *_args: Any, stream: bool = False, **kwargs: Any
     ) -> chat.ChatCompletion | MockAsyncStream[MockChatCompletionChunk]:
         self.chat_completion_kwargs.append({k: v for k, v in kwargs.items() if v is not NOT_GIVEN})
@@ -519,7 +519,7 @@ async def test_no_content(allow_model_requests: None):
 
     with pytest.raises(UnexpectedModelBehavior, match='Received empty model response'):
         async with agent.run_stream(''):
-            pass  # pragma: no cover
+            pass
 
 
 async def test_no_delta(allow_model_requests: None):
@@ -862,6 +862,14 @@ async def test_multiple_agent_tool_calls(allow_model_requests: None, gemini_api_
     assert result.output == snapshot('The capital of England is London.')
 
 
+@pytest.mark.vcr()
+async def test_extra_headers(allow_model_requests: None, openai_api_key: str):
+    # This test doesn't do anything, it's just here to ensure that calls with `extra_headers` don't cause errors, including type.
+    m = OpenAIModel('gpt-4o', provider=OpenAIProvider(api_key=openai_api_key))
+    agent = Agent(m, model_settings=OpenAIModelSettings(extra_headers={'Extra-Header-Key': 'Extra-Header-Value'}))
+    await agent.run('hello')
+
+
 @pytest.mark.vcr()
 async def test_user_id(allow_model_requests: None, openai_api_key: str):
     # This test doesn't do anything, it's just here to ensure that calls with `user` don't cause errors, including type.

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/models/test_openai_responses.py

@@ -68,7 +68,7 @@ async def test_openai_responses_output_type(allow_model_requests: None, openai_a
 
     agent = Agent(model=model, output_type=MyOutput)
     result = await agent.run('Give me the name and age of Brazil, Argentina, and Chile.')
-    assert result.output == snapshot({'name': 'Brazil', 'age': 2023})  # pragma: no cover
+    assert result.output == snapshot({'name': 'Brazil', 'age': 2023})
 
 
 async def test_openai_responses_reasoning_effort(allow_model_requests: None, openai_api_key: str):

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/providers/test_google_vertex.py

@@ -147,7 +147,7 @@ def save_service_account(service_account_path: Path, project_id: str) -> None:
 def vertex_provider_auth(mocker: MockerFixture) -> None:
     # Locally, we authenticate via `gcloud` CLI, so we don't need to patch anything.
     if not os.getenv('CI'):
-        return  # pragma: no cover
+        return  # pragma: lax no cover
 
     @dataclass
     class NoOpCredentials:

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/providers/test_provider_names.py

@@ -37,7 +37,7 @@ with try_import() as imports_successful:
     ]
 
 if not imports_successful():
-    test_infer_provider_params = []  # pragma: no cover
+    test_infer_provider_params = []
 
 pytestmark = pytest.mark.skipif(not imports_successful(), reason='need to install all extra packages')
 

{pydantic_ai-0.1.9 → pydantic_ai-0.1.11}/tests/test_cli.py

@@ -1,5 +1,3 @@
-import os
-import sys
 from io import StringIO
 from typing import Any
 
@@ -31,40 +29,6 @@ def test_cli_version(capfd: CaptureFixture[str]):
     assert capfd.readouterr().out.startswith('pai - PydanticAI CLI')
 
 
-@pytest.mark.skipif(not os.getenv('CI', False), reason="Marcelo can't make this test pass locally")
-@pytest.mark.skipif(sys.version_info >= (3, 13), reason='slightly different output with 3.13')
-def test_cli_help(capfd: CaptureFixture[str]):
-    with pytest.raises(SystemExit) as exc:
-        cli(['--help'])
-    assert exc.value.code == 0
-
-    assert capfd.readouterr().out.splitlines() == snapshot(
-        [
-            'usage: pai [-h] [-m [MODEL]] [-l] [-t [CODE_THEME]] [--no-stream] [--version] [prompt]',
-            '',
-            IsStr(),
-            '',
-            'Special prompt:',
-            '* `/exit` - exit the interactive mode',
-            '* `/markdown` - show the last markdown output of the last question',
-            '* `/multiline` - toggle multiline mode',
-            '',
-            'positional arguments:',
-            '  prompt                AI Prompt, if omitted fall into interactive mode',
-            '',
-            IsStr(),
-            '  -h, --help            show this help message and exit',
-            '  -m [MODEL], --model [MODEL]',
-            '                        Model to use, in format "<provider>:<model>" e.g. "openai:gpt-4o". Defaults to "openai:gpt-4o".',
-            '  -l, --list-models     List all available models and exit',
-            '  -t [CODE_THEME], --code-theme [CODE_THEME]',
-            '                        Which colors to use for code, can be "dark", "light" or any theme from pygments.org/styles/. Defaults to "monokai".',
-            '  --no-stream           Whether to stream responses from the model',
-            '  --version             Show version and exit',
-        ]
-    )
-
-
 def test_invalid_model(capfd: CaptureFixture[str]):
     assert cli(['--model', 'potato']) == 1
     assert capfd.readouterr().out.splitlines() == snapshot(
@@ -169,7 +133,7 @@ def test_code_theme_unset(mocker: MockerFixture, env: TestEnv):
     mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
     cli([])
     mock_run_chat.assert_awaited_once_with(
-        IsInstance(PromptSession), True, IsInstance(Agent), IsInstance(Console), 'monokai'
+        IsInstance(PromptSession), True, IsInstance(Agent), IsInstance(Console), 'monokai', 'pai'
     )
 
 
@@ -178,7 +142,7 @@ def test_code_theme_light(mocker: MockerFixture, env: TestEnv):
     mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
     cli(['--code-theme=light'])
     mock_run_chat.assert_awaited_once_with(
-        IsInstance(PromptSession), True, IsInstance(Agent), IsInstance(Console), 'default'
+        IsInstance(PromptSession), True, IsInstance(Agent), IsInstance(Console), 'default', 'pai'
     )
 
 
@@ -187,5 +151,5 @@ def test_code_theme_dark(mocker: MockerFixture, env: TestEnv):
     mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
     cli(['--code-theme=dark'])
     mock_run_chat.assert_awaited_once_with(
-        IsInstance(PromptSession), True, IsInstance(Agent), IsInstance(Console), 'monokai'
+        IsInstance(PromptSession), True, IsInstance(Agent), IsInstance(Console), 'monokai', 'pai'
     )