pydantic-ai 0.0.29__tar.gz → 0.0.30__tar.gz

{pydantic_ai-0.0.29 → pydantic_ai-0.0.30}/PKG-INFO

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

{pydantic_ai-0.0.29 → pydantic_ai-0.0.30}/pyproject.toml

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

{pydantic_ai-0.0.29 → pydantic_ai-0.0.30}/tests/models/test_instrumented.py

@@ -5,6 +5,7 @@ from contextlib import asynccontextmanager
 from datetime import datetime
 
 import pytest
+from dirty_equals import IsJson
 from inline_snapshot import snapshot
 from logfire_api import DEFAULT_LOGFIRE_INSTANCE
 
@@ -105,7 +106,7 @@ class MyResponseStream(StreamedResponse):
 @pytest.mark.anyio
 @requires_logfire_events
 async def test_instrumented_model(capfire: CaptureLogfire):
-    model = InstrumentedModel.from_logfire(MyModel())
+    model = InstrumentedModel.from_logfire(MyModel(), event_mode='logs')
     assert model.system == 'my_system'
     assert model.model_name == 'my_model'
 
@@ -323,7 +324,7 @@ async def test_instrumented_model_not_recording():
 @pytest.mark.anyio
 @requires_logfire_events
 async def test_instrumented_model_stream(capfire: CaptureLogfire):
-    model = InstrumentedModel.from_logfire(MyModel())
+    model = InstrumentedModel.from_logfire(MyModel(), event_mode='logs')
 
     messages: list[ModelMessage] = [
         ModelRequest(
@@ -405,7 +406,7 @@ async def test_instrumented_model_stream(capfire: CaptureLogfire):
 @pytest.mark.anyio
 @requires_logfire_events
 async def test_instrumented_model_stream_break(capfire: CaptureLogfire):
-    model = InstrumentedModel.from_logfire(MyModel())
+    model = InstrumentedModel.from_logfire(MyModel(), event_mode='logs')
 
     messages: list[ModelMessage] = [
         ModelRequest(
@@ -494,3 +495,140 @@ async def test_instrumented_model_stream_break(capfire: CaptureLogfire):
             },
         ]
     )
+
+
+@pytest.mark.anyio
+async def test_instrumented_model_attributes_mode(capfire: CaptureLogfire):
+    model = InstrumentedModel(MyModel(), event_mode='attributes')
+    assert model.system == 'my_system'
+    assert model.model_name == 'my_model'
+
+    messages = [
+        ModelRequest(
+            parts=[
+                SystemPromptPart('system_prompt'),
+                UserPromptPart('user_prompt'),
+                ToolReturnPart('tool3', 'tool_return_content', 'tool_call_3'),
+                RetryPromptPart('retry_prompt1', tool_name='tool4', tool_call_id='tool_call_4'),
+                RetryPromptPart('retry_prompt2'),
+                {},  # test unexpected parts  # type: ignore
+            ]
+        ),
+        ModelResponse(
+            parts=[
+                TextPart('text3'),
+            ]
+        ),
+    ]
+    await model.request(
+        messages,
+        model_settings=ModelSettings(temperature=1),
+        model_request_parameters=ModelRequestParameters(
+            function_tools=[],
+            allow_text_result=True,
+            result_tools=[],
+        ),
+    )
+
+    assert capfire.exporter.exported_spans_as_dict() == snapshot(
+        [
+            {
+                'name': 'chat my_model',
+                'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
+                'parent': None,
+                'start_time': 1000000000,
+                'end_time': 2000000000,
+                'attributes': {
+                    'gen_ai.operation.name': 'chat',
+                    'gen_ai.system': 'my_system',
+                    'gen_ai.request.model': 'my_model',
+                    'gen_ai.request.temperature': 1,
+                    'logfire.msg': 'chat my_model',
+                    'logfire.span_type': 'span',
+                    'gen_ai.response.model': 'my_model_123',
+                    'gen_ai.usage.input_tokens': 100,
+                    'gen_ai.usage.output_tokens': 200,
+                    'events': IsJson(
+                        snapshot(
+                            [
+                                {
+                                    'event.name': 'gen_ai.system.message',
+                                    'content': 'system_prompt',
+                                    'role': 'system',
+                                    'gen_ai.system': 'my_system',
+                                },
+                                {
+                                    'event.name': 'gen_ai.user.message',
+                                    'content': 'user_prompt',
+                                    'role': 'user',
+                                    'gen_ai.system': 'my_system',
+                                },
+                                {
+                                    'event.name': 'gen_ai.tool.message',
+                                    'content': 'tool_return_content',
+                                    'role': 'tool',
+                                    'id': 'tool_call_3',
+                                    'gen_ai.system': 'my_system',
+                                },
+                                {
+                                    'event.name': 'gen_ai.tool.message',
+                                    'content': """\
+retry_prompt1
+
+Fix the errors and try again.\
+""",
+                                    'role': 'tool',
+                                    'id': 'tool_call_4',
+                                    'gen_ai.system': 'my_system',
+                                },
+                                {
+                                    'event.name': 'gen_ai.user.message',
+                                    'content': """\
+retry_prompt2
+
+Fix the errors and try again.\
+""",
+                                    'role': 'user',
+                                    'gen_ai.system': 'my_system',
+                                },
+                                {
+                                    'event.name': 'gen_ai.assistant.message',
+                                    'role': 'assistant',
+                                    'content': 'text3',
+                                    'gen_ai.system': 'my_system',
+                                },
+                                {
+                                    'event.name': 'gen_ai.choice',
+                                    'index': 0,
+                                    'message': {
+                                        'role': 'assistant',
+                                        'content': 'text1',
+                                        'tool_calls': [
+                                            {
+                                                'id': 'tool_call_1',
+                                                'type': 'function',
+                                                'function': {'name': 'tool1', 'arguments': 'args1'},
+                                            },
+                                            {
+                                                'id': 'tool_call_2',
+                                                'type': 'function',
+                                                'function': {'name': 'tool2', 'arguments': {'args2': 3}},
+                                            },
+                                        ],
+                                    },
+                                    'gen_ai.system': 'my_system',
+                                },
+                                {
+                                    'event.name': 'gen_ai.choice',
+                                    'index': 0,
+                                    'message': {'role': 'assistant', 'content': 'text2'},
+                                    'gen_ai.system': 'my_system',
+                                },
+                            ]
+                        )
+                    ),
+                    'logfire.json_schema': '{"type": "object", "properties": {"events": {"type": "array"}}}',
+                },
+            },
+        ]
+    )

{pydantic_ai-0.0.29 → pydantic_ai-0.0.30}/tests/models/test_model_test.py

@@ -13,6 +13,9 @@ from pydantic import BaseModel, Field
 from pydantic_ai import Agent, ModelRetry, RunContext
 from pydantic_ai.exceptions import UnexpectedModelBehavior
 from pydantic_ai.messages import (
+    AudioUrl,
+    BinaryContent,
+    ImageUrl,
     ModelRequest,
     ModelResponse,
     RetryPromptPart,
@@ -22,6 +25,7 @@ from pydantic_ai.messages import (
     UserPromptPart,
 )
 from pydantic_ai.models.test import TestModel, _chars, _JsonSchemaTestData  # pyright: ignore[reportPrivateUsage]
+from pydantic_ai.usage import Usage
 
 from ..conftest import IsNow
 
@@ -271,3 +275,18 @@ def test_max_items():
     }
     data = _JsonSchemaTestData(json_schema).generate()
     assert data == snapshot([])
+
+
+@pytest.mark.parametrize(
+    'content',
+    [
+        AudioUrl(url='https://example.com'),
+        ImageUrl(url='https://example.com'),
+        BinaryContent(data=b'', media_type='image/png'),
+    ],
+)
+def test_different_content_input(content: AudioUrl | ImageUrl | BinaryContent):
+    agent = Agent()
+    result = agent.run_sync('x', model=TestModel(custom_result_text='custom'))
+    assert result.data == snapshot('custom')
+    assert result.usage() == snapshot(Usage(requests=1, request_tokens=51, response_tokens=1, total_tokens=52))

{pydantic_ai-0.0.29 → pydantic_ai-0.0.30}/tests/test_examples.py

@@ -182,6 +182,9 @@ text_responses: dict[str, str | ToolCallPart] = {
     'What is the weather like in West London and in Wiltshire?': (
         'The weather in West London is raining, while in Wiltshire it is sunny.'
     ),
+    'What will the weather be like in Paris on Tuesday?': ToolCallPart(
+        tool_name='weather_forecast', args={'location': 'Paris', 'forecast_date': '2030-01-01'}, tool_call_id='0001'
+    ),
     'Tell me a joke.': 'Did you hear about the toothpaste scandal? They called it Colgate.',
     'Explain?': 'This is an excellent joke invented by Samuel Colvin, it needs no explanation.',
     'What is the capital of France?': 'Paris',
@@ -270,6 +273,13 @@ text_responses: dict[str, str | ToolCallPart] = {
     ),
 }
 
+tool_responses: dict[tuple[str, str], str] = {
+    (
+        'weather_forecast',
+        'The forecast in Paris on 2030-01-01 is 24°C and sunny.',
+    ): 'It will be warm and sunny in Paris on Tuesday.',
+}
+
 
 async def model_logic(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:  # pragma: no cover  # noqa: C901
     m = messages[-1].parts[-1]
@@ -348,35 +358,53 @@ async def model_logic(messages: list[ModelMessage], info: AgentInfo) -> ModelRes
         raise RuntimeError(f'Unexpected message: {m}')
 
 
-async def stream_model_logic(
+async def stream_model_logic(  # noqa C901
     messages: list[ModelMessage], info: AgentInfo
 ) -> AsyncIterator[str | DeltaToolCalls]:  # pragma: no cover
-    m = messages[-1].parts[-1]
-    if isinstance(m, UserPromptPart):
-        assert isinstance(m.content, str)
-        if response := text_responses.get(m.content):
-            if isinstance(response, str):
-                words = response.split(' ')
-                chunk: list[str] = []
-                for work in words:
-                    chunk.append(work)
-                    if len(chunk) == 3:
-                        yield ' '.join(chunk) + ' '
-                        chunk.clear()
-                if chunk:
-                    yield ' '.join(chunk)
-                return
-            else:
-                json_text = response.args_as_json_str()
-
-                yield {1: DeltaToolCall(name=response.tool_name)}
-                for chunk_index in range(0, len(json_text), 15):
-                    text_chunk = json_text[chunk_index : chunk_index + 15]
-                    yield {1: DeltaToolCall(json_args=text_chunk)}
-                return
+    async def stream_text_response(r: str) -> AsyncIterator[str]:
+        if isinstance(r, str):
+            words = r.split(' ')
+            chunk: list[str] = []
+            for word in words:
+                chunk.append(word)
+                if len(chunk) == 3:
+                    yield ' '.join(chunk) + ' '
+                    chunk.clear()
+            if chunk:
+                yield ' '.join(chunk)
+
+    async def stream_tool_call_response(r: ToolCallPart) -> AsyncIterator[DeltaToolCalls]:
+        json_text = r.args_as_json_str()
+
+        yield {1: DeltaToolCall(name=r.tool_name, tool_call_id=r.tool_call_id)}
+        for chunk_index in range(0, len(json_text), 15):
+            text_chunk = json_text[chunk_index : chunk_index + 15]
+            yield {1: DeltaToolCall(json_args=text_chunk)}
+
+    async def stream_part_response(r: str | ToolCallPart) -> AsyncIterator[str | DeltaToolCalls]:
+        if isinstance(r, str):
+            async for chunk in stream_text_response(r):
+                yield chunk
+        else:
+            async for chunk in stream_tool_call_response(r):
+                yield chunk
+
+    last_part = messages[-1].parts[-1]
+    if isinstance(last_part, UserPromptPart):
+        assert isinstance(last_part.content, str)
+        if response := text_responses.get(last_part.content):
+            async for chunk in stream_part_response(response):
+                yield chunk
+            return
+    elif isinstance(last_part, ToolReturnPart):
+        assert isinstance(last_part.content, str)
+        if response := tool_responses.get((last_part.tool_name, last_part.content)):
+            async for chunk in stream_part_response(response):
+                yield chunk
+            return
 
     sys.stdout.write(str(debug.format(messages, info)))
-    raise RuntimeError(f'Unexpected message: {m}')
+    raise RuntimeError(f'Unexpected message: {last_part}')
 
 
 def mock_infer_model(model: Model | KnownModelName) -> Model:

{pydantic_ai-0.0.29 → pydantic_ai-0.0.30}/tests/test_streaming.py

@@ -2,14 +2,18 @@ from __future__ import annotations as _annotations
 
 import datetime
 import json
+import re
 from collections.abc import AsyncIterator
+from copy import deepcopy
 from datetime import timezone
+from typing import Union
 
 import pytest
 from inline_snapshot import snapshot
 from pydantic import BaseModel
 
 from pydantic_ai import Agent, UnexpectedModelBehavior, UserError, capture_run_messages
+from pydantic_ai.agent import AgentRun
 from pydantic_ai.messages import (
     ModelMessage,
     ModelRequest,
@@ -22,7 +26,8 @@ from pydantic_ai.messages import (
 )
 from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel
 from pydantic_ai.models.test import TestModel
-from pydantic_ai.result import Usage
+from pydantic_ai.result import AgentStream, FinalResult, Usage
+from pydantic_graph import End
 
 from .conftest import IsNow
 
@@ -739,3 +744,109 @@ async def test_custom_result_type_default_structured() -> None:
     async with agent.run_stream('test', result_type=str) as result:
         response = await result.get_data()
         assert response == snapshot('success (no tool calls)')
+
+
+async def test_iter_stream_output():
+    m = TestModel(custom_result_text='The cat sat on the mat.')
+
+    agent = Agent(m)
+
+    @agent.result_validator
+    def result_validator_simple(data: str) -> str:
+        # Make a substitution in the validated results
+        return re.sub('cat sat', 'bat sat', data)
+
+    run: AgentRun
+    stream: AgentStream
+    messages: list[str] = []
+
+    stream_usage: Usage | None = None
+    with agent.iter('Hello') as run:
+        async for node in run:
+            if agent.is_model_request_node(node):
+                async with node.stream(run.ctx) as stream:
+                    async for chunk in stream.stream_output(debounce_by=None):
+                        messages.append(chunk)
+                stream_usage = deepcopy(stream.usage())
+    assert run.next_node == End(data=FinalResult(data='The bat sat on the mat.', tool_name=None))
+    assert (
+        run.usage()
+        == stream_usage
+        == Usage(requests=1, request_tokens=51, response_tokens=7, total_tokens=58, details=None)
+    )
+
+    assert messages == [
+        '',
+        'The ',
+        'The cat ',
+        'The bat sat ',
+        'The bat sat on ',
+        'The bat sat on the ',
+        'The bat sat on the mat.',
+        'The bat sat on the mat.',
+    ]
+
+
+async def test_iter_stream_responses():
+    m = TestModel(custom_result_text='The cat sat on the mat.')
+
+    agent = Agent(m)
+
+    @agent.result_validator
+    def result_validator_simple(data: str) -> str:
+        # Make a substitution in the validated results
+        return re.sub('cat sat', 'bat sat', data)
+
+    run: AgentRun
+    stream: AgentStream
+    messages: list[ModelResponse] = []
+    with agent.iter('Hello') as run:
+        async for node in run:
+            if agent.is_model_request_node(node):
+                async with node.stream(run.ctx) as stream:
+                    async for chunk in stream.stream_responses(debounce_by=None):
+                        messages.append(chunk)
+
+    assert messages == [
+        ModelResponse(
+            parts=[TextPart(content=text, part_kind='text')],
+            model_name='test',
+            timestamp=IsNow(tz=timezone.utc),
+            kind='response',
+        )
+        for text in [
+            '',
+            '',
+            'The ',
+            'The cat ',
+            'The cat sat ',
+            'The cat sat on ',
+            'The cat sat on the ',
+            'The cat sat on the mat.',
+        ]
+    ]
+
+    # Note: as you can see above, the result validator is not applied to the streamed responses, just the final result:
+    assert run.result is not None
+    assert run.result.data == 'The bat sat on the mat.'
+
+
+async def test_stream_iter_structured_validator() -> None:
+    class NotResultType(BaseModel):
+        not_value: str
+
+    agent = Agent[None, Union[ResultType, NotResultType]]('test', result_type=Union[ResultType, NotResultType])  # pyright: ignore[reportArgumentType]
+
+    @agent.result_validator
+    def result_validator(data: ResultType | NotResultType) -> ResultType | NotResultType:
+        assert isinstance(data, ResultType)
+        return ResultType(value=data.value + ' (validated)')
+
+    outputs: list[ResultType] = []
+    with agent.iter('test') as run:
+        async for node in run:
+            if agent.is_model_request_node(node):
+                async with node.stream(run.ctx) as stream:
+                    async for output in stream.stream_output(debounce_by=None):
+                        outputs.append(output)
+    assert outputs == [ResultType(value='a (validated)'), ResultType(value='a (validated)')]