pydantic-ai 0.0.33__tar.gz → 0.0.34__tar.gz

{pydantic_ai-0.0.33 → pydantic_ai-0.0.34}/PKG-INFO

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

{pydantic_ai-0.0.33 → pydantic_ai-0.0.34}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pydantic-ai"
-version = "0.0.33"
+version = "0.0.34"
 description = "Agent Framework / shim to use Pydantic with LLMs"
 authors = [
     { name = "Samuel Colvin", email = "samuel@pydantic.dev" },
@@ -35,9 +35,8 @@ classifiers = [
     "Framework :: Pydantic :: 2",
 ]
 requires-python = ">=3.9"
-
 dependencies = [
-    "pydantic-ai-slim[openai,vertexai,groq,anthropic,mistral,cohere]==0.0.33",
+    "pydantic-ai-slim[openai,vertexai,groq,anthropic,mistral,cohere,cli]==0.0.34",
 ]
 
 [project.urls]
@@ -47,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.33"]
+examples = ["pydantic-ai-examples==0.0.34"]
 logfire = ["logfire>=2.3"]
 
 [tool.uv.sources]
@@ -165,13 +164,14 @@ exclude_lines = [
     'typing.assert_never',
     '$\s*assert_never\(',
     'if __name__ == .__main__.:',
+    'except ImportError as _import_error:',
 ]
 
 [tool.logfire]
 ignore_no_config = true
 
 [tool.inline-snapshot]
-format-command="ruff format --stdin-filename {filename}"
+format-command = "ruff format --stdin-filename {filename}"
 
 [tool.inline-snapshot.shortcuts]
 snap-fix = ["create", "fix"]

{pydantic_ai-0.0.33 → pydantic_ai-0.0.34}/tests/models/test_instrumented.py

@@ -8,6 +8,8 @@ import pytest
 from dirty_equals import IsJson
 from inline_snapshot import snapshot
 from logfire_api import DEFAULT_LOGFIRE_INSTANCE
+from opentelemetry._events import NoOpEventLoggerProvider
+from opentelemetry.trace import NoOpTracerProvider
 
 from pydantic_ai.messages import (
     ModelMessage,
@@ -25,6 +27,7 @@ from pydantic_ai.messages import (
     UserPromptPart,
 )
 from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse
+from pydantic_ai.models.instrumented import InstrumentationSettings, InstrumentedModel
 from pydantic_ai.settings import ModelSettings
 from pydantic_ai.usage import Usage
 
@@ -32,11 +35,6 @@ from ..conftest import try_import
 
 with try_import() as imports_successful:
     from logfire.testing import CaptureLogfire
-    from opentelemetry._events import NoOpEventLoggerProvider
-    from opentelemetry.trace import NoOpTracerProvider
-
-    from pydantic_ai.models.instrumented import InstrumentedModel
-
 
 pytestmark = [
     pytest.mark.skipif(not imports_successful(), reason='logfire not installed'),
@@ -103,10 +101,9 @@ class MyResponseStream(StreamedResponse):
         return datetime(2022, 1, 1)
 
 
-@pytest.mark.anyio
 @requires_logfire_events
 async def test_instrumented_model(capfire: CaptureLogfire):
-    model = InstrumentedModel(MyModel(), event_mode='logs')
+    model = InstrumentedModel(MyModel(), InstrumentationSettings(event_mode='logs'))
     assert model.system == 'my_system'
     assert model.model_name == 'my_model'
 
@@ -193,7 +190,7 @@ async def test_instrumented_model(capfire: CaptureLogfire):
                 'trace_flags': 1,
             },
             {
-                'body': {'content': 'tool_return_content', 'role': 'tool', 'id': 'tool_call_3'},
+                'body': {'content': 'tool_return_content', 'role': 'tool', 'id': 'tool_call_3', 'name': 'tool3'},
                 'severity_number': 9,
                 'severity_text': None,
                 'attributes': {
@@ -216,6 +213,7 @@ Fix the errors and try again.\
 """,
                     'role': 'tool',
                     'id': 'tool_call_4',
+                    'name': 'tool4',
                 },
                 'severity_number': 9,
                 'severity_text': None,
@@ -311,9 +309,11 @@ Fix the errors and try again.\
     )
 
 
-@pytest.mark.anyio
 async def test_instrumented_model_not_recording():
-    model = InstrumentedModel(MyModel(), NoOpTracerProvider(), NoOpEventLoggerProvider())
+    model = InstrumentedModel(
+        MyModel(),
+        InstrumentationSettings(tracer_provider=NoOpTracerProvider(), event_logger_provider=NoOpEventLoggerProvider()),
+    )
 
     messages: list[ModelMessage] = [ModelRequest(parts=[SystemPromptPart('system_prompt')])]
     await model.request(
@@ -327,10 +327,9 @@ async def test_instrumented_model_not_recording():
     )
 
 
-@pytest.mark.anyio
 @requires_logfire_events
 async def test_instrumented_model_stream(capfire: CaptureLogfire):
-    model = InstrumentedModel(MyModel(), event_mode='logs')
+    model = InstrumentedModel(MyModel(), InstrumentationSettings(event_mode='logs'))
 
     messages: list[ModelMessage] = [
         ModelRequest(
@@ -410,10 +409,9 @@ 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(MyModel(), event_mode='logs')
+    model = InstrumentedModel(MyModel(), InstrumentationSettings(event_mode='logs'))
 
     messages: list[ModelMessage] = [
         ModelRequest(
@@ -505,9 +503,8 @@ 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')
+    model = InstrumentedModel(MyModel(), InstrumentationSettings(event_mode='attributes'))
     assert model.system == 'my_system'
     assert model.model_name == 'my_model'
 
@@ -577,6 +574,7 @@ async def test_instrumented_model_attributes_mode(capfire: CaptureLogfire):
                                     'event.name': 'gen_ai.tool.message',
                                     'content': 'tool_return_content',
                                     'role': 'tool',
+                                    'name': 'tool3',
                                     'id': 'tool_call_3',
                                     'gen_ai.message.index': 0,
                                     'gen_ai.system': 'my_system',
@@ -589,6 +587,7 @@ retry_prompt1
 Fix the errors and try again.\
 """,
                                     'role': 'tool',
+                                    'name': 'tool4',
                                     'id': 'tool_call_4',
                                     'gen_ai.message.index': 0,
                                     'gen_ai.system': 'my_system',

pydantic_ai-0.0.34/tests/test_cli.py

@@ -0,0 +1,52 @@
+import pytest
+from dirty_equals import IsStr
+from inline_snapshot import snapshot
+from pytest import CaptureFixture
+
+from .conftest import try_import
+
+with try_import() as imports_successful:
+    from pydantic_ai._cli import cli
+
+pytestmark = pytest.mark.skipif(not imports_successful(), reason='install cli extras to run cli tests')
+
+
+def test_cli_version(capfd: CaptureFixture[str]):
+    assert cli(['--version']) == 0
+    assert capfd.readouterr().out.startswith('pai - PydanticAI CLI')
+
+
+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] [--model [MODEL]] [--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',
+            '  --model [MODEL]  Model to use, it should be "<provider>:<model>" e.g. "openai:gpt-4o". If omitted it will default to "openai:gpt-4o"',
+            '  --no-stream      Whether to stream responses from OpenAI',
+            '  --version        Show version and exit',
+        ]
+    )
+
+
+def test_invalid_model(capfd: CaptureFixture[str]):
+    assert cli(['--model', 'invalid_model']) == 1
+    assert capfd.readouterr().out == snapshot("""\
+pai - PydanticAI CLI v0.0.34
+Invalid model "invalid_model"
+""")

{pydantic_ai-0.0.33 → pydantic_ai-0.0.34}/tests/test_logfire.py

@@ -9,6 +9,7 @@ from inline_snapshot import snapshot
 from typing_extensions import NotRequired, TypedDict
 
 from pydantic_ai import Agent
+from pydantic_ai.models.instrumented import InstrumentationSettings, InstrumentedModel
 from pydantic_ai.models.test import TestModel
 
 try:
@@ -58,8 +59,17 @@ def get_logfire_summary(capfire: CaptureLogfire) -> Callable[[], LogfireSummary]
 
 
 @pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
-def test_logfire(get_logfire_summary: Callable[[], LogfireSummary]) -> None:
-    my_agent = Agent(model=TestModel(), instrument=True)
+@pytest.mark.parametrize(
+    'instrument',
+    [
+        True,
+        False,
+        InstrumentationSettings(event_mode='attributes'),
+        InstrumentationSettings(event_mode='logs'),
+    ],
+)
+def test_logfire(get_logfire_summary: Callable[[], LogfireSummary], instrument: InstrumentationSettings | bool) -> None:
+    my_agent = Agent(model=TestModel(), instrument=instrument)
 
     @my_agent.tool_plain
     async def my_ret(x: int) -> str:
@@ -69,6 +79,10 @@ def test_logfire(get_logfire_summary: Callable[[], LogfireSummary]) -> None:
     assert result.data == snapshot('{"my_ret":"1"}')
 
     summary = get_logfire_summary()
+    if instrument is False:
+        assert summary.traces == []
+        return
+
     assert summary.traces == snapshot(
         [
             {
@@ -120,6 +134,7 @@ def test_logfire(get_logfire_summary: Callable[[], LogfireSummary]) -> None:
                             'content': '1',
                             'role': 'tool',
                             'id': None,
+                            'name': 'my_ret',
                             'gen_ai.message.index': 2,
                             'event.name': 'gen_ai.tool.message',
                         },
@@ -150,7 +165,44 @@ def test_logfire(get_logfire_summary: Callable[[], LogfireSummary]) -> None:
             'logfire.msg': 'preparing model request params',
         }
     )
-    assert summary.attributes[2] == snapshot(
+    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']}
+        assert attribute_mode_attributes == snapshot(
+            {
+                'events': IsJson(
+                    snapshot(
+                        [
+                            {
+                                'event.name': 'gen_ai.user.message',
+                                'content': 'Hello',
+                                'role': 'user',
+                                'gen_ai.message.index': 0,
+                                'gen_ai.system': 'test',
+                            },
+                            {
+                                'event.name': 'gen_ai.choice',
+                                'index': 0,
+                                'message': {
+                                    'role': 'assistant',
+                                    'tool_calls': [
+                                        {
+                                            'id': None,
+                                            'type': 'function',
+                                            'function': {'name': 'my_ret', 'arguments': {'x': 0}},
+                                        }
+                                    ],
+                                },
+                                'gen_ai.system': 'test',
+                            },
+                        ]
+                    )
+                ),
+                'logfire.json_schema': '{"type": "object", "properties": {"events": {"type": "array"}}}',
+            }
+        )
+
+    assert chat_span_attributes == snapshot(
         {
             'gen_ai.operation.name': 'chat',
             'gen_ai.system': 'test',
@@ -160,34 +212,32 @@ def test_logfire(get_logfire_summary: Callable[[], LogfireSummary]) -> None:
             'gen_ai.response.model': 'test',
             'gen_ai.usage.input_tokens': 51,
             'gen_ai.usage.output_tokens': 4,
-            'events': IsJson(
-                snapshot(
-                    [
-                        {
-                            'event.name': 'gen_ai.user.message',
-                            'content': 'Hello',
-                            'role': 'user',
-                            'gen_ai.message.index': 0,
-                            'gen_ai.system': 'test',
-                        },
-                        {
-                            'event.name': 'gen_ai.choice',
-                            'index': 0,
-                            'message': {
-                                'role': 'assistant',
-                                'tool_calls': [
-                                    {
-                                        'id': None,
-                                        'type': 'function',
-                                        'function': {'name': 'my_ret', 'arguments': {'x': 0}},
-                                    }
-                                ],
-                            },
-                            'gen_ai.system': 'test',
-                        },
-                    ]
-                )
-            ),
-            'logfire.json_schema': '{"type": "object", "properties": {"events": {"type": "array"}}}',
         }
     )
+
+
+def test_instrument_all():
+    model = TestModel()
+    agent = Agent()
+
+    def get_model():
+        return agent._get_model(model)  # type: ignore
+
+    Agent.instrument_all(False)
+    assert get_model() is model
+
+    Agent.instrument_all()
+    m = get_model()
+    assert isinstance(m, InstrumentedModel)
+    assert m.wrapped is model
+    assert m.options.event_mode == InstrumentationSettings().event_mode
+
+    options = InstrumentationSettings(event_mode='logs')
+    Agent.instrument_all(options)
+    m = get_model()
+    assert isinstance(m, InstrumentedModel)
+    assert m.wrapped is model
+    assert m.options is options
+
+    Agent.instrument_all(False)
+    assert get_model() is model