pydantic-ai 0.2.1__tar.gz → 0.2.2__tar.gz

{pydantic_ai-0.2.1 → pydantic_ai-0.2.2}/PKG-INFO

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

{pydantic_ai-0.2.1 → pydantic_ai-0.2.2}/tests/test_cli.py

@@ -1,3 +1,4 @@
+import sys
 from io import StringIO
 from typing import Any
 
@@ -36,6 +37,72 @@ def test_invalid_model(capfd: CaptureFixture[str]):
     )
 
 
+def test_agent_flag(capfd: CaptureFixture[str], mocker: MockerFixture, env: TestEnv):
+    env.set('OPENAI_API_KEY', 'test')
+
+    # Create a dynamic module using types.ModuleType
+    import types
+
+    test_module = types.ModuleType('test_module')
+
+    # Create and add agent to the module
+    test_agent = Agent()
+    test_agent.model = TestModel(custom_output_text='Hello from custom agent')
+    setattr(test_module, 'custom_agent', test_agent)
+
+    # Register the module in sys.modules
+    sys.modules['test_module'] = test_module
+
+    try:
+        # Mock ask_agent to avoid actual execution but capture the agent
+        mock_ask = mocker.patch('pydantic_ai._cli.ask_agent')
+
+        # Test CLI with custom agent
+        assert cli(['--agent', 'test_module:custom_agent', 'hello']) == 0
+
+        # Verify the output contains the custom agent message
+        assert 'Using custom agent: test_module:custom_agent' in capfd.readouterr().out
+
+        # Verify ask_agent was called with our custom agent
+        mock_ask.assert_called_once()
+        assert mock_ask.call_args[0][0] is test_agent
+
+    finally:
+        # Clean up by removing the module from sys.modules
+        if 'test_module' in sys.modules:
+            del sys.modules['test_module']
+
+
+def test_agent_flag_non_agent(capfd: CaptureFixture[str], mocker: MockerFixture, env: TestEnv):
+    env.set('OPENAI_API_KEY', 'test')
+
+    # Create a dynamic module using types.ModuleType
+    import types
+
+    test_module = types.ModuleType('test_module')
+
+    # Create and add agent to the module
+    test_agent = 'Not an Agent object'
+    setattr(test_module, 'custom_agent', test_agent)
+
+    # Register the module in sys.modules
+    sys.modules['test_module'] = test_module
+
+    try:
+        assert cli(['--agent', 'test_module:custom_agent', 'hello']) == 1
+        assert 'is not an Agent' in capfd.readouterr().out
+
+    finally:
+        # Clean up by removing the module from sys.modules
+        if 'test_module' in sys.modules:
+            del sys.modules['test_module']
+
+
+def test_agent_flag_bad_module_variable_path(capfd: CaptureFixture[str], mocker: MockerFixture, env: TestEnv):
+    assert cli(['--agent', 'bad_path', 'hello']) == 1
+    assert 'Agent must be specified in "module:variable" format' in capfd.readouterr().out
+
+
 def test_list_models(capfd: CaptureFixture[str]):
     assert cli(['--list-models']) == 0
     output = capfd.readouterr().out.splitlines()
@@ -153,3 +220,34 @@ def test_code_theme_dark(mocker: MockerFixture, env: TestEnv):
     mock_run_chat.assert_awaited_once_with(
         IsInstance(PromptSession), True, IsInstance(Agent), IsInstance(Console), 'monokai', 'pai'
     )
+
+
+def test_agent_to_cli_sync(mocker: MockerFixture, env: TestEnv):
+    env.set('OPENAI_API_KEY', 'test')
+    mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
+    cli_agent.to_cli_sync()
+    mock_run_chat.assert_awaited_once_with(
+        session=IsInstance(PromptSession),
+        stream=True,
+        agent=IsInstance(Agent),
+        console=IsInstance(Console),
+        code_theme='monokai',
+        prog_name='pydantic-ai',
+        deps=None,
+    )
+
+
+@pytest.mark.anyio
+async def test_agent_to_cli_async(mocker: MockerFixture, env: TestEnv):
+    env.set('OPENAI_API_KEY', 'test')
+    mock_run_chat = mocker.patch('pydantic_ai._cli.run_chat')
+    await cli_agent.to_cli()
+    mock_run_chat.assert_awaited_once_with(
+        session=IsInstance(PromptSession),
+        stream=True,
+        agent=IsInstance(Agent),
+        console=IsInstance(Console),
+        code_theme='monokai',
+        prog_name='pydantic-ai',
+        deps=None,
+    )