pydantic-ai 0.0.43__tar.gz → 0.0.44__tar.gz

{pydantic_ai-0.0.43 → pydantic_ai-0.0.44}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai
-Version: 0.0.43
+Version: 0.0.44
 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,groq,mcp,mistral,openai,vertexai]==0.0.43
+Requires-Dist: pydantic-ai-slim[anthropic,bedrock,cli,cohere,groq,mcp,mistral,openai,vertexai]==0.0.44
 Provides-Extra: examples
-Requires-Dist: pydantic-ai-examples==0.0.43; extra == 'examples'
+Requires-Dist: pydantic-ai-examples==0.0.44; extra == 'examples'
 Provides-Extra: logfire
 Requires-Dist: logfire>=2.3; extra == 'logfire'
 Description-Content-Type: text/markdown

{pydantic_ai-0.0.43 → pydantic_ai-0.0.44}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pydantic-ai"
-version = "0.0.43"
+version = "0.0.44"
 description = "Agent Framework / shim to use Pydantic with LLMs"
 authors = [
     { name = "Samuel Colvin", email = "samuel@pydantic.dev" },
@@ -36,7 +36,7 @@ classifiers = [
 ]
 requires-python = ">=3.9"
 dependencies = [
-    "pydantic-ai-slim[openai,vertexai,groq,anthropic,mistral,cohere,bedrock,cli,mcp]==0.0.43",
+    "pydantic-ai-slim[openai,vertexai,groq,anthropic,mistral,cohere,bedrock,cli,mcp]==0.0.44",
 ]
 
 [project.urls]
@@ -46,7 +46,7 @@ Documentation = "https://ai.pydantic.dev"
 Changelog = "https://github.com/pydantic/pydantic-ai/releases"
 
 [project.optional-dependencies]
-examples = ["pydantic-ai-examples==0.0.43"]
+examples = ["pydantic-ai-examples==0.0.44"]
 logfire = ["logfire>=2.3"]
 
 [tool.uv.sources]

{pydantic_ai-0.0.43 → pydantic_ai-0.0.44}/tests/conftest.py

@@ -246,6 +246,11 @@ def anthropic_api_key() -> str:
     return os.getenv('ANTHROPIC_API_KEY', 'mock-api-key')
 
 
+@pytest.fixture(scope='session')
+def co_api_key() -> str:
+    return os.getenv('CO_API_KEY', 'mock-api-key')
+
+
 @pytest.fixture
 def mock_snapshot_id(mocker: MockerFixture):
     i = 0

pydantic_ai-0.0.44/tests/graph/test_utils.py

@@ -0,0 +1,20 @@
+from threading import Thread
+
+from pydantic_graph._utils import run_until_complete
+
+
+def test_run_until_complete_in_main_thread():
+    async def run(): ...
+
+    run_until_complete(run())
+
+
+def test_run_until_complete_in_thread():
+    async def run(): ...
+
+    def get_and_close_event_loop():
+        run_until_complete(run())
+
+    thread = Thread(target=get_and_close_event_loop)
+    thread.start()
+    thread.join()

pydantic_ai-0.0.44/tests/models/cassettes/test_cohere/test_request_simple_success_with_vcr.yaml

@@ -0,0 +1,64 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '93'
+      content-type:
+      - application/json
+      host:
+      - api.cohere.com
+    method: POST
+    parsed_body:
+      messages:
+      - content: hello
+        role: user
+      model: command-r7b-12-2024
+      stream: false
+    uri: https://api.cohere.com/v2/chat
+  response:
+    headers:
+      access-control-expose-headers:
+      - X-Debug-Trace-ID
+      alt-svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      cache-control:
+      - no-cache, no-store, no-transform, must-revalidate, private, max-age=0
+      content-length:
+      - '286'
+      content-type:
+      - application/json
+      expires:
+      - Thu, 01 Jan 1970 00:00:00 UTC
+      num_chars:
+      - '2583'
+      num_tokens:
+      - '10'
+      pragma:
+      - no-cache
+      vary:
+      - Origin
+    parsed_body:
+      finish_reason: COMPLETE
+      id: f17a5f6c-1734-4098-bd0d-733ef000ac7b
+      message:
+        content:
+        - text: Hello! How can I assist you today?
+          type: text
+        role: assistant
+      usage:
+        billed_units:
+          input_tokens: 1
+          output_tokens: 9
+        tokens:
+          input_tokens: 496
+          output_tokens: 11
+    status:
+      code: 200
+      message: OK
+version: 1

pydantic_ai-0.0.44/tests/models/cassettes/test_openai/test_max_completion_tokens[gpt-4.5-preview].yaml

@@ -0,0 +1,78 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '123'
+      content-type:
+      - application/json
+      host:
+      - api.openai.com
+    method: POST
+    parsed_body:
+      max_completion_tokens: 100
+      messages:
+      - content: hello
+        role: user
+      model: gpt-4.5-preview
+      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:
+      - '807'
+      content-type:
+      - application/json
+      openai-organization:
+      - pydantic-28gund
+      openai-processing-ms:
+      - '1408'
+      openai-version:
+      - '2020-10-01'
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      transfer-encoding:
+      - chunked
+    parsed_body:
+      choices:
+      - finish_reason: stop
+        index: 0
+        message:
+          annotations: []
+          content: Hello! How can I help you today?
+          refusal: null
+          role: assistant
+      created: 1742636225
+      id: chatcmpl-BDpZplWguNLn40wA5mIpCR3OIzvYP
+      model: gpt-4.5-preview-2025-02-27
+      object: chat.completion
+      service_tier: default
+      system_fingerprint: null
+      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.0.44/tests/models/cassettes/test_openai/test_max_completion_tokens[gpt-4o-mini].yaml

@@ -0,0 +1,79 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '119'
+      content-type:
+      - application/json
+      host:
+      - api.openai.com
+    method: POST
+    parsed_body:
+      max_completion_tokens: 100
+      messages:
+      - content: hello
+        role: user
+      model: gpt-4o-mini
+      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:
+      - '840'
+      content-type:
+      - application/json
+      openai-organization:
+      - pydantic-28gund
+      openai-processing-ms:
+      - '278'
+      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: 1742636224
+      id: chatcmpl-BDpZoxw4i90ZesN8iyrwLmGWRJ5lz
+      model: gpt-4o-mini-2024-07-18
+      object: chat.completion
+      service_tier: default
+      system_fingerprint: fp_b8bc95a0ac
+      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.0.44/tests/models/cassettes/test_openai/test_max_completion_tokens[o3-mini].yaml

@@ -0,0 +1,78 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '115'
+      content-type:
+      - application/json
+      host:
+      - api.openai.com
+    method: POST
+    parsed_body:
+      max_completion_tokens: 100
+      messages:
+      - content: hello
+        role: user
+      model: o3-mini
+      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:
+      - '817'
+      content-type:
+      - application/json
+      openai-organization:
+      - pydantic-28gund
+      openai-processing-ms:
+      - '1895'
+      openai-version:
+      - '2020-10-01'
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      transfer-encoding:
+      - chunked
+    parsed_body:
+      choices:
+      - finish_reason: stop
+        index: 0
+        message:
+          annotations: []
+          content: Hello there! How can I help you today?
+          refusal: null
+          role: assistant
+      created: 1742636222
+      id: chatcmpl-BDpZm1SiItIXIcDA0xRV9QGZhD97e
+      model: o3-mini-2025-01-31
+      object: chat.completion
+      service_tier: default
+      system_fingerprint: fp_617f206dd9
+      usage:
+        completion_tokens: 85
+        completion_tokens_details:
+          accepted_prediction_tokens: 0
+          audio_tokens: 0
+          reasoning_tokens: 64
+          rejected_prediction_tokens: 0
+        prompt_tokens: 7
+        prompt_tokens_details:
+          audio_tokens: 0
+          cached_tokens: 0
+        total_tokens: 92
+    status:
+      code: 200
+      message: OK
+version: 1

{pydantic_ai-0.0.43 → pydantic_ai-0.0.44}/tests/models/test_cohere.py

@@ -38,6 +38,7 @@ with try_import() as imports_successful:
     from cohere.core.api_error import ApiError
 
     from pydantic_ai.models.cohere import CohereModel
+    from pydantic_ai.providers.cohere import CohereProvider
 
     # note: we use Union here for compatibility with Python 3.9
     MockChatResponse = Union[ChatResponse, Exception]
@@ -49,7 +50,7 @@ pytestmark = [
 
 
 def test_init():
-    m = CohereModel('command-r7b-12-2024', api_key='foobar')
+    m = CohereModel('command-r7b-12-2024', provider=CohereProvider(api_key='foobar'))
     assert m.model_name == 'command-r7b-12-2024'
     assert m.system == 'cohere'
     assert m.base_url == 'https://api.cohere.com'
@@ -96,7 +97,7 @@ async def test_request_simple_success(allow_model_requests: None):
         )
     )
     mock_client = MockAsyncClientV2.create_mock(c)
-    m = CohereModel('command-r7b-12-2024', cohere_client=mock_client)
+    m = CohereModel('command-r7b-12-2024', provider=CohereProvider(cohere_client=mock_client))
     agent = Agent(m)
 
     result = await agent.run('hello')
@@ -135,7 +136,7 @@ async def test_request_simple_usage(allow_model_requests: None):
         ),
     )
     mock_client = MockAsyncClientV2.create_mock(c)
-    m = CohereModel('command-r7b-12-2024', cohere_client=mock_client)
+    m = CohereModel('command-r7b-12-2024', provider=CohereProvider(cohere_client=mock_client))
     agent = Agent(m)
 
     result = await agent.run('Hello')
@@ -169,7 +170,7 @@ async def test_request_structured_response(allow_model_requests: None):
         )
     )
     mock_client = MockAsyncClientV2.create_mock(c)
-    m = CohereModel('command-r7b-12-2024', cohere_client=mock_client)
+    m = CohereModel('command-r7b-12-2024', provider=CohereProvider(cohere_client=mock_client))
     agent = Agent(m, result_type=list[int])
 
     result = await agent.run('Hello')
@@ -243,7 +244,7 @@ async def test_request_tool_call(allow_model_requests: None):
         ),
     ]
     mock_client = MockAsyncClientV2.create_mock(responses)
-    m = CohereModel('command-r7b-12-2024', cohere_client=mock_client)
+    m = CohereModel('command-r7b-12-2024', provider=CohereProvider(cohere_client=mock_client))
     agent = Agent(m, system_prompt='this is the system prompt')
 
     @agent.tool_plain
@@ -326,7 +327,7 @@ async def test_request_tool_call(allow_model_requests: None):
 async def test_multimodal(allow_model_requests: None):
     c = completion_message(AssistantMessageResponse(content=[TextAssistantMessageResponseContentItem(text='world')]))
     mock_client = MockAsyncClientV2.create_mock(c)
-    m = CohereModel('command-r7b-12-2024', cohere_client=mock_client)
+    m = CohereModel('command-r7b-12-2024', provider=CohereProvider(cohere_client=mock_client))
     agent = Agent(m)
 
     with pytest.raises(RuntimeError, match='Cohere does not yet support multi-modal inputs.'):
@@ -347,8 +348,16 @@ def test_model_status_error(allow_model_requests: None) -> None:
             body={'error': 'test error'},
         )
     )
-    m = CohereModel('command-r', cohere_client=mock_client)
+    m = CohereModel('command-r', provider=CohereProvider(cohere_client=mock_client))
     agent = Agent(m)
     with pytest.raises(ModelHTTPError) as exc_info:
         agent.run_sync('hello')
     assert str(exc_info.value) == snapshot("status_code: 500, model_name: command-r, body: {'error': 'test error'}")
+
+
+@pytest.mark.vcr()
+async def test_request_simple_success_with_vcr(allow_model_requests: None, co_api_key: str):
+    m = CohereModel('command-r7b-12-2024', provider=CohereProvider(api_key=co_api_key))
+    agent = Agent(m)
+    result = await agent.run('hello')
+    assert result.data == snapshot('Hello! How can I assist you today?')

{pydantic_ai-0.0.43 → pydantic_ai-0.0.44}/tests/models/test_fallback.py

@@ -136,6 +136,7 @@ def test_first_failed_instrumented(capfire: CaptureLogfire) -> None:
                 'end_time': 5000000000,
                 'attributes': {
                     'gen_ai.operation.name': 'chat',
+                    'model_request_parameters': '{"function_tools": [], "allow_text_result": true, "result_tools": []}',
                     'logfire.span_type': 'span',
                     'logfire.msg': 'chat fallback:function:failure_response:,function:success_response:',
                     'gen_ai.usage.input_tokens': 51,
@@ -160,7 +161,7 @@ def test_first_failed_instrumented(capfire: CaptureLogfire) -> None:
                             },
                         ]
                     ),
-                    'logfire.json_schema': '{"type": "object", "properties": {"events": {"type": "array"}}}',
+                    'logfire.json_schema': '{"type": "object", "properties": {"events": {"type": "array"}, "model_request_parameters": {"type": "object"}}}',
                 },
             },
             {
@@ -233,6 +234,7 @@ async def test_first_failed_instrumented_stream(capfire: CaptureLogfire) -> None
                 'end_time': 5000000000,
                 'attributes': {
                     'gen_ai.operation.name': 'chat',
+                    'model_request_parameters': '{"function_tools": [], "allow_text_result": true, "result_tools": []}',
                     'logfire.span_type': 'span',
                     'logfire.msg': 'chat fallback:function::failure_response_stream,function::success_response_stream',
                     'gen_ai.system': 'function',
@@ -241,7 +243,7 @@ async def test_first_failed_instrumented_stream(capfire: CaptureLogfire) -> None
                     'gen_ai.usage.output_tokens': 2,
                     'gen_ai.response.model': 'function::success_response_stream',
                     'events': '[{"content": "input", "role": "user", "gen_ai.system": "function", "gen_ai.message.index": 0, "event.name": "gen_ai.user.message"}, {"index": 0, "message": {"role": "assistant", "content": "hello world"}, "gen_ai.system": "function", "event.name": "gen_ai.choice"}]',
-                    'logfire.json_schema': '{"type": "object", "properties": {"events": {"type": "array"}}}',
+                    'logfire.json_schema': '{"type": "object", "properties": {"events": {"type": "array"}, "model_request_parameters": {"type": "object"}}}',
                 },
             },
             {

{pydantic_ai-0.0.43 → pydantic_ai-0.0.44}/tests/models/test_instrumented.py

@@ -152,6 +152,8 @@ async def test_instrumented_model(capfire: CaptureLogfire):
                     'gen_ai.request.model': 'my_model',
                     'server.address': 'example.com',
                     'server.port': 8000,
+                    'model_request_parameters': '{"function_tools": [], "allow_text_result": true, "result_tools": []}',
+                    'logfire.json_schema': '{"type": "object", "properties": {"model_request_parameters": {"type": "object"}}}',
                     'gen_ai.request.temperature': 1,
                     'logfire.msg': 'chat my_model',
                     'logfire.span_type': 'span',
@@ -374,6 +376,8 @@ async def test_instrumented_model_stream(capfire: CaptureLogfire):
                     'gen_ai.request.model': 'my_model',
                     'server.address': 'example.com',
                     'server.port': 8000,
+                    'model_request_parameters': '{"function_tools": [], "allow_text_result": true, "result_tools": []}',
+                    'logfire.json_schema': '{"type": "object", "properties": {"model_request_parameters": {"type": "object"}}}',
                     'gen_ai.request.temperature': 1,
                     'logfire.msg': 'chat my_model',
                     'logfire.span_type': 'span',
@@ -457,6 +461,8 @@ async def test_instrumented_model_stream_break(capfire: CaptureLogfire):
                     'gen_ai.request.model': 'my_model',
                     'server.address': 'example.com',
                     'server.port': 8000,
+                    'model_request_parameters': '{"function_tools": [], "allow_text_result": true, "result_tools": []}',
+                    'logfire.json_schema': '{"type": "object", "properties": {"model_request_parameters": {"type": "object"}}}',
                     'gen_ai.request.temperature': 1,
                     'logfire.msg': 'chat my_model',
                     'logfire.span_type': 'span',
@@ -559,6 +565,7 @@ async def test_instrumented_model_attributes_mode(capfire: CaptureLogfire):
                     'gen_ai.request.model': 'my_model',
                     'server.address': 'example.com',
                     'server.port': 8000,
+                    'model_request_parameters': '{"function_tools": [], "allow_text_result": true, "result_tools": []}',
                     'gen_ai.request.temperature': 1,
                     'logfire.msg': 'chat my_model',
                     'logfire.span_type': 'span',
@@ -652,7 +659,7 @@ Fix the errors and try again.\
                             ]
                         )
                     ),
-                    'logfire.json_schema': '{"type": "object", "properties": {"events": {"type": "array"}}}',
+                    'logfire.json_schema': '{"type": "object", "properties": {"events": {"type": "array"}, "model_request_parameters": {"type": "object"}}}',
                 },
             },
         ]

{pydantic_ai-0.0.43 → pydantic_ai-0.0.44}/tests/models/test_openai.py

@@ -28,7 +28,7 @@ from pydantic_ai.messages import (
 from pydantic_ai.result import Usage
 from pydantic_ai.settings import ModelSettings
 
-from ..conftest import IsNow, TestEnv, raise_if_exception, try_import
+from ..conftest import IsNow, IsStr, TestEnv, raise_if_exception, try_import
 from .mock_async_stream import MockAsyncStream
 
 with try_import() as imports_successful:
@@ -684,3 +684,13 @@ def test_model_status_error(allow_model_requests: None) -> None:
     with pytest.raises(ModelHTTPError) as exc_info:
         agent.run_sync('hello')
     assert str(exc_info.value) == snapshot("status_code: 500, model_name: gpt-4o, body: {'error': 'test error'}")
+
+
+@pytest.mark.vcr()
+@pytest.mark.parametrize('model_name', ['o3-mini', 'gpt-4o-mini', 'gpt-4.5-preview'])
+async def test_max_completion_tokens(allow_model_requests: None, model_name: str, openai_api_key: str):
+    m = OpenAIModel(model_name, provider=OpenAIProvider(api_key=openai_api_key))
+    agent = Agent(m, model_settings=ModelSettings(max_tokens=100))
+
+    result = await agent.run('hello')
+    assert result.data == IsStr()

pydantic_ai-0.0.44/tests/providers/test_cohere.py

@@ -0,0 +1,54 @@
+from __future__ import annotations as _annotations
+
+import os
+from unittest.mock import patch
+
+import httpx
+import pytest
+
+from ..conftest import try_import
+
+with try_import() as imports_successful:
+    from cohere import AsyncClientV2
+    from cohere.core.http_client import AsyncHttpClient
+
+    from pydantic_ai.providers.cohere import CohereProvider
+
+
+pytestmark = pytest.mark.skipif(not imports_successful(), reason='cohere not installed')
+
+
+def test_cohere_provider() -> None:
+    provider = CohereProvider(api_key='api-key')
+    assert provider.name == 'cohere'
+    assert provider.base_url == 'https://api.cohere.com'
+    assert isinstance(provider.client, AsyncClientV2)
+    assert provider.client._client_wrapper._token == 'api-key'  # type: ignore[reportPrivateUsage]
+
+
+def test_cohere_provider_need_api_key() -> None:
+    with patch.dict(os.environ, {}, clear=True):
+        with pytest.raises(ValueError, match='CO_API_KEY'):
+            CohereProvider()
+
+
+def test_cohere_provider_pass_http_client() -> None:
+    http_client = httpx.AsyncClient()
+    provider = CohereProvider(http_client=http_client, api_key='api-key')
+    # The AsyncClientV2 wraps our httpx client in an AsyncHttpClient
+    # So we just check that the httpx_client is an instance of AsyncHttpClient
+    assert isinstance(provider.client._client_wrapper.httpx_client, AsyncHttpClient)  # type: ignore[reportPrivateUsage]
+
+
+def test_cohere_provider_pass_cohere_client() -> None:
+    cohere_client = AsyncClientV2(api_key='test-api-key')
+    provider = CohereProvider(cohere_client=cohere_client)
+    assert provider.client == cohere_client
+
+
+def test_cohere_provider_with_env_base_url(monkeypatch: pytest.MonkeyPatch) -> None:
+    custom_base_url = 'https://custom.cohere.com/'
+    # Test with environment variable for base_url
+    monkeypatch.setenv('CO_BASE_URL', custom_base_url)
+    provider = CohereProvider(api_key='api-key')
+    assert provider.base_url == custom_base_url

{pydantic_ai-0.0.43 → pydantic_ai-0.0.44}/tests/providers/test_provider_names.py

@@ -12,6 +12,7 @@ from ..conftest import try_import
 
 with try_import() as imports_successful:
     from pydantic_ai.providers.anthropic import AnthropicProvider
+    from pydantic_ai.providers.cohere import CohereProvider
     from pydantic_ai.providers.deepseek import DeepSeekProvider
     from pydantic_ai.providers.google_gla import GoogleGLAProvider
     from pydantic_ai.providers.google_vertex import GoogleVertexProvider
@@ -21,6 +22,7 @@ with try_import() as imports_successful:
 
     test_infer_provider_params = [
         ('anthropic', AnthropicProvider, 'ANTHROPIC_API_KEY'),
+        ('cohere', CohereProvider, 'CO_API_KEY'),
         ('deepseek', DeepSeekProvider, 'DEEPSEEK_API_KEY'),
         ('openai', OpenAIProvider, None),
         ('google-vertex', GoogleVertexProvider, None),

{pydantic_ai-0.0.43 → pydantic_ai-0.0.44}/tests/test_examples.py

@@ -86,6 +86,8 @@ def test_docs_examples(  # noqa: C901
     env.set('GEMINI_API_KEY', 'testing')
     env.set('GROQ_API_KEY', 'testing')
     env.set('CO_API_KEY', 'testing')
+    env.set('MISTRAL_API_KEY', 'testing')
+    env.set('ANTHROPIC_API_KEY', 'testing')
 
     sys.path.append('tests/example_modules')
 

{pydantic_ai-0.0.43 → pydantic_ai-0.0.44}/tests/test_logfire.py

@@ -167,7 +167,7 @@ def test_logfire(get_logfire_summary: Callable[[], LogfireSummary], instrument:
     )
     chat_span_attributes = summary.attributes[2]
     if instrument is True or instrument.event_mode == 'attributes':
-        attribute_mode_attributes = {k: chat_span_attributes.pop(k) for k in ['events', 'logfire.json_schema']}
+        attribute_mode_attributes = {k: chat_span_attributes.pop(k) for k in ['events']}
         assert attribute_mode_attributes == snapshot(
             {
                 'events': IsJson(
@@ -198,7 +198,6 @@ def test_logfire(get_logfire_summary: Callable[[], LogfireSummary], instrument:
                         ]
                     )
                 ),
-                'logfire.json_schema': '{"type": "object", "properties": {"events": {"type": "array"}}}',
             }
         )
 
@@ -207,6 +206,28 @@ def test_logfire(get_logfire_summary: Callable[[], LogfireSummary], instrument:
             'gen_ai.operation.name': 'chat',
             'gen_ai.system': 'test',
             'gen_ai.request.model': 'test',
+            'model_request_parameters': IsJson(
+                snapshot(
+                    {
+                        'function_tools': [
+                            {
+                                'name': 'my_ret',
+                                'description': '',
+                                'parameters_json_schema': {
+                                    'additionalProperties': False,
+                                    'properties': {'x': {'type': 'integer'}},
+                                    'required': ['x'],
+                                    'type': 'object',
+                                },
+                                'outer_typed_dict_key': None,
+                            }
+                        ],
+                        'allow_text_result': True,
+                        'result_tools': [],
+                    }
+                )
+            ),
+            'logfire.json_schema': IsJson(),
             'logfire.span_type': 'span',
             'logfire.msg': 'chat test',
             'gen_ai.response.model': 'test',

{pydantic_ai-0.0.43 → pydantic_ai-0.0.44}/tests/test_tools.py

@@ -122,7 +122,6 @@ def sphinx_style_docstring(foo: int, /) -> str:  # pragma: no cover
     """Sphinx style docstring.
 
     :param foo: The foo thing.
-    :return: The result.
     """
     return str(foo)
 
@@ -187,6 +186,152 @@ def test_docstring_numpy(docstring_format: Literal['numpy', 'auto']):
     )
 
 
+def test_google_style_with_returns():
+    agent = Agent(FunctionModel(get_json_schema))
+
+    def my_tool(x: int) -> str:  # pragma: no cover
+        """A function that does something.
+
+        Args:
+            x: The input value.
+
+        Returns:
+            str: The result as a string.
+        """
+        return str(x)
+
+    agent.tool_plain(my_tool)
+    result = agent.run_sync('Hello')
+    json_schema = json.loads(result.data)
+    assert json_schema == snapshot(
+        {
+            'name': 'my_tool',
+            'description': """\
+<summary>A function that does something.</summary>
+<returns>
+<type>str</type>
+<description>The result as a string.</description>
+</returns>\
+""",
+            'parameters_json_schema': {
+                'additionalProperties': False,
+                'properties': {'x': {'description': 'The input value.', 'type': 'integer'}},
+                'required': ['x'],
+                'type': 'object',
+            },
+            'outer_typed_dict_key': None,
+        }
+    )
+
+
+def test_sphinx_style_with_returns():
+    agent = Agent(FunctionModel(get_json_schema))
+
+    def my_tool(x: int) -> str:  # pragma: no cover
+        """A sphinx function with returns.
+
+        :param x: The input value.
+        :rtype: str
+        :return: The result as a string with type.
+        """
+        return str(x)
+
+    agent.tool_plain(docstring_format='sphinx')(my_tool)
+    result = agent.run_sync('Hello')
+    json_schema = json.loads(result.data)
+    assert json_schema == snapshot(
+        {
+            'name': 'my_tool',
+            'description': """\
+<summary>A sphinx function with returns.</summary>
+<returns>
+<type>str</type>
+<description>The result as a string with type.</description>
+</returns>\
+""",
+            'parameters_json_schema': {
+                'additionalProperties': False,
+                'properties': {'x': {'description': 'The input value.', 'type': 'integer'}},
+                'required': ['x'],
+                'type': 'object',
+            },
+            'outer_typed_dict_key': None,
+        }
+    )
+
+
+def test_numpy_style_with_returns():
+    agent = Agent(FunctionModel(get_json_schema))
+
+    def my_tool(x: int) -> str:  # pragma: no cover
+        """A numpy function with returns.
+
+        Parameters
+        ----------
+        x : int
+            The input value.
+
+        Returns
+        -------
+        str
+            The result as a string with type.
+        """
+        return str(x)
+
+    agent.tool_plain(docstring_format='numpy')(my_tool)
+    result = agent.run_sync('Hello')
+    json_schema = json.loads(result.data)
+    assert json_schema == snapshot(
+        {
+            'name': 'my_tool',
+            'description': """\
+<summary>A numpy function with returns.</summary>
+<returns>
+<type>str</type>
+<description>The result as a string with type.</description>
+</returns>\
+""",
+            'parameters_json_schema': {
+                'additionalProperties': False,
+                'properties': {'x': {'description': 'The input value.', 'type': 'integer'}},
+                'required': ['x'],
+                'type': 'object',
+            },
+            'outer_typed_dict_key': None,
+        }
+    )
+
+
+def only_returns_type() -> str:  # pragma: no cover
+    """
+
+    Returns:
+        str: The result as a string.
+    """
+    return 'foo'
+
+
+def test_only_returns_type():
+    agent = Agent(FunctionModel(get_json_schema))
+    agent.tool_plain(only_returns_type)
+
+    result = agent.run_sync('Hello')
+    json_schema = json.loads(result.data)
+    assert json_schema == snapshot(
+        {
+            'name': 'only_returns_type',
+            'description': """\
+<returns>
+<type>str</type>
+<description>The result as a string.</description>
+</returns>\
+""",
+            'parameters_json_schema': {'additionalProperties': False, 'properties': {}, 'type': 'object'},
+            'outer_typed_dict_key': None,
+        }
+    )
+
+
 def unknown_docstring(**kwargs: int) -> str:  # pragma: no cover
     """Unknown style docstring."""
     return str(kwargs)
@@ -572,11 +717,7 @@ agent = Agent('test', tools=[ctx_tool], deps_type=int)
 
 
 async def tool_without_return_annotation_in_docstring() -> str:  # pragma: no cover
-    """A tool that documents what it returns but doesn't have a return annotation in the docstring.
-
-    Returns:
-        A value.
-    """
+    """A tool that documents what it returns but doesn't have a return annotation in the docstring."""
 
     return ''
 
@@ -591,8 +732,7 @@ def test_suppress_griffe_logging(caplog: LogCaptureFixture):
     json_schema = json.loads(result.data)
     assert json_schema == snapshot(
         {
-            'description': "A tool that documents what it returns but doesn't have a "
-            'return annotation in the docstring.',
+            'description': "A tool that documents what it returns but doesn't have a return annotation in the docstring.",
             'name': 'tool_without_return_annotation_in_docstring',
             'outer_typed_dict_key': None,
             'parameters_json_schema': {'additionalProperties': False, 'properties': {}, 'type': 'object'},

pydantic_ai-0.0.43/tests/graph/test_utils.py

@@ -1,13 +0,0 @@
-from threading import Thread
-
-from pydantic_graph._utils import get_event_loop
-
-
-def test_get_event_loop_in_thread():
-    def get_and_close_event_loop():
-        event_loop = get_event_loop()
-        event_loop.close()
-
-    thread = Thread(target=get_and_close_event_loop)
-    thread.start()
-    thread.join()