clai should not require OPENAI_API_KEY (#1724)

samuelcolvin · web-flow · commit fa270d9467f8 · 2025-05-14T16:09:28.000-04:00
diff --git a/pydantic_ai_slim/pydantic_ai/_cli.py b/pydantic_ai_slim/pydantic_ai/_cli.py
@@ -102,7 +102,7 @@ def cli_exit(prog_name: str = 'pai'):  # pragma: no cover
     sys.exit(cli(prog_name=prog_name))
 
 
-def cli(args_list: Sequence[str] | None = None, *, prog_name: str = 'pai') -> int:
+def cli(args_list: Sequence[str] | None = None, *, prog_name: str = 'pai') -> int:  # noqa: C901
     """Run the CLI and return the exit code for the process."""
     parser = argparse.ArgumentParser(
         prog=prog_name,
@@ -122,7 +122,6 @@ def cli(args_list: Sequence[str] | None = None, *, prog_name: str = 'pai') -> in
         '--model',
         nargs='?',
         help='Model to use, in format "<provider>:<model>" e.g. "openai:gpt-4o" or "anthropic:claude-3-7-sonnet-latest". Defaults to "openai:gpt-4o".',
-        default='openai:gpt-4o',
     )
     # we don't want to autocomplete or list models that don't include the provider,
     # e.g. we want to show `openai:gpt-4o` but not `gpt-4o`
@@ -153,40 +152,49 @@ def cli(args_list: Sequence[str] | None = None, *, prog_name: str = 'pai') -> in
     args = parser.parse_args(args_list)
 
     console = Console()
-    console.print(
-        f'[green]{prog_name} - PydanticAI CLI v{__version__} using[/green] [magenta]{args.model}[/magenta]',
-        highlight=False,
-    )
+    name_version = f'[green]{prog_name} - PydanticAI CLI v{__version__}[/green]'
     if args.version:
+        console.print(name_version, highlight=False)
         return 0
     if args.list_models:
-        console.print('Available models:', style='green bold')
+        console.print(f'{name_version}\n\n[green]Available models:[/green]')
         for model in qualified_model_names:
             console.print(f'  {model}', highlight=False)
         return 0
 
     agent: Agent[None, str] = cli_agent
     if args.agent:
+        sys.path.append(os.getcwd())
         try:
-            current_path = os.getcwd()
-            sys.path.append(current_path)
-
             module_path, variable_name = args.agent.split(':')
-            module = importlib.import_module(module_path)
-            agent = getattr(module, variable_name)
-            if not isinstance(agent, Agent):
-                console.print(f'[red]Error: {args.agent} is not an Agent instance[/red]')
-                return 1
-            console.print(f'[green]Using custom agent:[/green] [magenta]{args.agent}[/magenta]', highlight=False)
         except ValueError:
             console.print('[red]Error: Agent must be specified in "module:variable" format[/red]')
             return 1
 
-    try:
-        agent.model = infer_model(args.model)
-    except UserError as e:
-        console.print(f'Error initializing [magenta]{args.model}[/magenta]:\n[red]{e}[/red]')
-        return 1
+        module = importlib.import_module(module_path)
+        agent = getattr(module, variable_name)
+        if not isinstance(agent, Agent):
+            console.print(f'[red]Error: {args.agent} is not an Agent instance[/red]')
+            return 1
+
+    model_arg_set = args.model is not None
+    if agent.model is None or model_arg_set:
+        try:
+            agent.model = infer_model(args.model or 'openai:gpt-4o')
+        except UserError as e:
+            console.print(f'Error initializing [magenta]{args.model}[/magenta]:\n[red]{e}[/red]')
+            return 1
+
+    model_name = agent.model if isinstance(agent.model, str) else f'{agent.model.system}:{agent.model.model_name}'
+    if args.agent and model_arg_set:
+        console.print(
+            f'{name_version} using custom agent [magenta]{args.agent}[/magenta] with [magenta]{model_name}[/magenta]',
+            highlight=False,
+        )
+    elif args.agent:
+        console.print(f'{name_version} using custom agent [magenta]{args.agent}[/magenta]', highlight=False)
+    else:
+        console.print(f'{name_version} with [magenta]{model_name}[/magenta]', highlight=False)
 
     stream = not args.no_stream
     if args.code_theme == 'light':
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -1,6 +1,7 @@
 import sys
+import types
 from io import StringIO
-from typing import Any
+from typing import Any, Callable
 
 import pytest
 from dirty_equals import IsInstance, IsStr
@@ -16,11 +17,13 @@
 from .conftest import TestEnv, try_import
 
 with try_import() as imports_successful:
+    from openai import OpenAIError
     from prompt_toolkit.input import create_pipe_input
     from prompt_toolkit.output import DummyOutput
     from prompt_toolkit.shortcuts import PromptSession
 
     from pydantic_ai._cli import cli, cli_agent, handle_slash_command
+    from pydantic_ai.models.openai import OpenAIModel
 
 pytestmark = pytest.mark.skipif(not imports_successful(), reason='install cli extras to run cli tests')
 
@@ -32,70 +35,90 @@ def test_cli_version(capfd: CaptureFixture[str]):
 
 def test_invalid_model(capfd: CaptureFixture[str]):
     assert cli(['--model', 'potato']) == 1
-    assert capfd.readouterr().out.splitlines() == snapshot(
-        [IsStr(), 'Error initializing potato:', 'Unknown model: potato']
-    )
+    assert capfd.readouterr().out.splitlines() == snapshot(['Error initializing potato:', 'Unknown model: potato'])
 
 
-def test_agent_flag(capfd: CaptureFixture[str], mocker: MockerFixture, env: TestEnv):
-    env.set('OPENAI_API_KEY', 'test')
+@pytest.fixture
+def create_test_module():
+    def _create_test_module(**namespace: Any) -> None:
+        assert 'test_module' not in sys.modules
 
-    # Create a dynamic module using types.ModuleType
-    import types
+        test_module = types.ModuleType('test_module')
+        for key, value in namespace.items():
+            setattr(test_module, key, value)
 
-    test_module = types.ModuleType('test_module')
+        sys.modules['test_module'] = 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)
+    try:
+        yield _create_test_module
+    finally:
+        if 'test_module' in sys.modules:
+            del sys.modules['test_module']
 
-    # 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')
+def test_agent_flag(
+    capfd: CaptureFixture[str],
+    mocker: MockerFixture,
+    env: TestEnv,
+    create_test_module: Callable[..., None],
+):
+    env.remove('OPENAI_API_KEY')
 
-        # Test CLI with custom agent
-        assert cli(['--agent', 'test_module:custom_agent', 'hello']) == 0
+    test_agent = Agent(TestModel(custom_output_text='Hello from custom agent'))
+    create_test_module(custom_agent=test_agent)
 
-        # Verify the output contains the custom agent message
-        assert 'Using custom agent: test_module:custom_agent' in capfd.readouterr().out
+    # Mock ask_agent to avoid actual execution but capture the agent
+    mock_ask = mocker.patch('pydantic_ai._cli.ask_agent')
 
-        # Verify ask_agent was called with our custom agent
-        mock_ask.assert_called_once()
-        assert mock_ask.call_args[0][0] is test_agent
+    # Test CLI with custom agent
+    assert cli(['--agent', 'test_module:custom_agent', 'hello']) == 0
 
-    finally:
-        # Clean up by removing the module from sys.modules
-        if 'test_module' in sys.modules:
-            del sys.modules['test_module']
+    # 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
 
-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
+def test_agent_flag_no_model(env: TestEnv, create_test_module: Callable[..., None]):
+    env.remove('OPENAI_API_KEY')
+    test_agent = Agent()
+    create_test_module(custom_agent=test_agent)
 
-    test_module = types.ModuleType('test_module')
+    msg = 'The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable'
+    with pytest.raises(OpenAIError, match=msg):
+        cli(['--agent', 'test_module:custom_agent', 'hello'])
 
-    # 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
+def test_agent_flag_set_model(
+    capfd: CaptureFixture[str],
+    mocker: MockerFixture,
+    env: TestEnv,
+    create_test_module: Callable[..., None],
+):
+    env.set('OPENAI_API_KEY', 'xxx')
 
-    try:
-        assert cli(['--agent', 'test_module:custom_agent', 'hello']) == 1
-        assert 'is not an Agent' in capfd.readouterr().out
+    custom_agent = Agent(TestModel(custom_output_text='Hello from custom agent'))
+    create_test_module(custom_agent=custom_agent)
 
-    finally:
-        # Clean up by removing the module from sys.modules
-        if 'test_module' in sys.modules:
-            del sys.modules['test_module']
+    mocker.patch('pydantic_ai._cli.ask_agent')
+
+    assert cli(['--agent', 'test_module:custom_agent', '--model', 'gpt-4o', 'hello']) == 0
+
+    assert 'using custom agent test_module:custom_agent with openai:gpt-4o' in capfd.readouterr().out
+
+    assert isinstance(custom_agent.model, OpenAIModel)
+
+
+def test_agent_flag_non_agent(
+    capfd: CaptureFixture[str], mocker: MockerFixture, create_test_module: Callable[..., None]
+):
+    test_agent = 'Not an Agent object'
+    create_test_module(custom_agent=test_agent)
+
+    assert cli(['--agent', 'test_module:custom_agent', 'hello']) == 1
+    assert 'is not an Agent' in capfd.readouterr().out
 
 
 def test_agent_flag_bad_module_variable_path(capfd: CaptureFixture[str], mocker: MockerFixture, env: TestEnv):
@@ -106,7 +129,7 @@ def test_agent_flag_bad_module_variable_path(capfd: CaptureFixture[str], mocker:
 def test_list_models(capfd: CaptureFixture[str]):
     assert cli(['--list-models']) == 0
     output = capfd.readouterr().out.splitlines()
-    assert output[:2] == snapshot([IsStr(regex='pai - PydanticAI CLI .* using openai:gpt-4o'), 'Available models:'])
+    assert output[:3] == snapshot([IsStr(regex='pai - PydanticAI CLI .*'), '', 'Available models:'])
 
     providers = (
         'openai',
@@ -119,7 +142,7 @@ def test_list_models(capfd: CaptureFixture[str]):
         'cohere',
         'deepseek',
     )
-    models = {line.strip().split(' ')[0] for line in output[2:]}
+    models = {line.strip().split(' ')[0] for line in output[3:]}
     for provider in providers:
         models = models - {model for model in models if model.startswith(provider)}
     assert models == set(), models
