Add to_cli() method to Agent (#1642)

Co-authored-by: Marcelo Trylesinski <[email protected]>

Author: AndrewHannigan

clai/README.md

@@ -53,7 +53,7 @@ Either way, running `clai` will start an interactive session where you can chat
 ## Help
 
 ```
-usage: clai [-h] [-m [MODEL]] [-l] [-t [CODE_THEME]] [--no-stream] [--version] [prompt]
+usage: clai [-h] [-m [MODEL]] [-a AGENT] [-l] [-t [CODE_THEME]] [--no-stream] [--version] [prompt]
 
 PydanticAI CLI v...
 
@@ -69,6 +69,8 @@ options:
   -h, --help            show this help message and exit
   -m [MODEL], --model [MODEL]
                         Model to use, in format "<provider>:<model>" e.g. "openai:gpt-4o" or "anthropic:claude-3-7-sonnet-latest". Defaults to "openai:gpt-4o".
+  -a AGENT, --agent AGENT
+                        Custom Agent to use, in format "module:variable", e.g. "mymodule.submodule:my_agent"
   -l, --list-models     List all available models and exit
   -t [CODE_THEME], --code-theme [CODE_THEME]
                         Which colors to use for code, can be "dark", "light" or any theme from pygments.org/styles/. Defaults to "dark" which works well on dark terminals.

docs/cli.md

@@ -64,3 +64,47 @@ uvx clai --model anthropic:claude-3-7-sonnet-latest
 ```
 
 (a full list of models available can be printed with `uvx clai --list-models`)
+
+### Custom Agents
+
+You can specify a custom agent using the `--agent` flag with a module path and variable name:
+
+```python {title="custom_agent.py" test="skip"}
+from pydantic_ai import Agent
+
+agent = Agent('openai:gpt-4o', instructions='You always respond in Italian.')
+```
+
+Then run:
+
+```bash
+uvx clai --agent custom_agent:agent "What's the weather today?"
+```
+
+The format must be `module:variable` where:
+
+- `module` is the importable Python module path
+- `variable` is the name of the Agent instance in that module
+
+
+Additionally, you can directly launch CLI mode from an `Agent` instance using `Agent.to_cli_sync()`:
+
+```python {title="agent_to_cli_sync.py" test="skip" hl_lines=4}
+from pydantic_ai import Agent
+
+agent = Agent('openai:gpt-4o', instructions='You always respond in Italian.')
+agent.to_cli_sync()
+```
+
+You can also use the async interface with `Agent.to_cli()`:
+
+```python {title="agent_to_cli.py" test="skip" hl_lines=6}
+from pydantic_ai import Agent
+
+agent = Agent('openai:gpt-4o', instructions='You always respond in Italian.')
+
+async def main():
+    await agent.to_cli()
+```
+
+_(You'll need to add `asyncio.run(main())` to run `main`)_

pydantic_ai_slim/pydantic_ai/_cli.py

@@ -2,6 +2,7 @@
 
 import argparse
 import asyncio
+import importlib
 import sys
 from asyncio import CancelledError
 from collections.abc import Sequence
@@ -12,6 +13,9 @@
 
 from typing_inspection.introspection import get_literal_values
 
+from pydantic_ai.result import OutputDataT
+from pydantic_ai.tools import AgentDepsT
+
 from . import __version__
 from .agent import Agent
 from .exceptions import UserError
@@ -123,6 +127,11 @@ def cli(args_list: Sequence[str] | None = None, *, prog_name: str = 'pai') -> in
     # e.g. we want to show `openai:gpt-4o` but not `gpt-4o`
     qualified_model_names = [n for n in get_literal_values(KnownModelName.__value__) if ':' in n]
     arg.completer = argcomplete.ChoicesCompleter(qualified_model_names)  # type: ignore[reportPrivateUsage]
+    parser.add_argument(
+        '-a',
+        '--agent',
+        help='Custom Agent to use, in format "module:variable", e.g. "mymodule.submodule:my_agent"',
+    )
     parser.add_argument(
         '-l',
         '--list-models',
@@ -155,8 +164,22 @@ def cli(args_list: Sequence[str] | None = None, *, prog_name: str = 'pai') -> in
             console.print(f'  {model}', highlight=False)
         return 0
 
+    agent: Agent[None, str] = cli_agent
+    if args.agent:
+        try:
+            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:
-        cli_agent.model = infer_model(args.model)
+        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
@@ -171,21 +194,27 @@ def cli(args_list: Sequence[str] | None = None, *, prog_name: str = 'pai') -> in
 
     if prompt := cast(str, args.prompt):
         try:
-            asyncio.run(ask_agent(cli_agent, prompt, stream, console, code_theme))
+            asyncio.run(ask_agent(agent, prompt, stream, console, code_theme))
         except KeyboardInterrupt:
             pass
         return 0
 
     # doing this instead of `PromptSession[Any](history=` allows mocking of PromptSession in tests
     session: PromptSession[Any] = PromptSession(history=FileHistory(str(PROMPT_HISTORY_PATH)))
     try:
-        return asyncio.run(run_chat(session, stream, cli_agent, console, code_theme, prog_name))
+        return asyncio.run(run_chat(session, stream, agent, console, code_theme, prog_name))
     except KeyboardInterrupt:  # pragma: no cover
         return 0
 
 
 async def run_chat(
-    session: PromptSession[Any], stream: bool, agent: Agent, console: Console, code_theme: str, prog_name: str
+    session: PromptSession[Any],
+    stream: bool,
+    agent: Agent[AgentDepsT, OutputDataT],
+    console: Console,
+    code_theme: str,
+    prog_name: str,
+    deps: AgentDepsT = None,
 ) -> int:
     multiline = False
     messages: list[ModelMessage] = []
@@ -207,30 +236,31 @@ async def run_chat(
                 return exit_value
         else:
             try:
-                messages = await ask_agent(agent, text, stream, console, code_theme, messages)
+                messages = await ask_agent(agent, text, stream, console, code_theme, deps, messages)
             except CancelledError:  # pragma: no cover
                 console.print('[dim]Interrupted[/dim]')
 
 
 async def ask_agent(
-    agent: Agent,
+    agent: Agent[AgentDepsT, OutputDataT],
     prompt: str,
     stream: bool,
     console: Console,
     code_theme: str,
+    deps: AgentDepsT = None,
     messages: list[ModelMessage] | None = None,
 ) -> list[ModelMessage]:
     status = Status('[dim]Working on it…[/dim]', console=console)
 
     if not stream:
         with status:
-            result = await agent.run(prompt, message_history=messages)
-        content = result.output
+            result = await agent.run(prompt, message_history=messages, deps=deps)
+        content = str(result.output)
         console.print(Markdown(content, code_theme=code_theme))
         return result.all_messages()
 
     with status, ExitStack() as stack:
-        async with agent.iter(prompt, message_history=messages) as agent_run:
+        async with agent.iter(prompt, message_history=messages, deps=deps) as agent_run:
             live = Live('', refresh_per_second=15, console=console, vertical_overflow='ellipsis')
             async for node in agent_run:
                 if Agent.is_model_request_node(node):

pydantic_ai_slim/pydantic_ai/agent.py

@@ -12,7 +12,7 @@
 
 from opentelemetry.trace import NoOpTracer, use_span
 from pydantic.json_schema import GenerateJsonSchema
-from typing_extensions import Literal, Never, TypeIs, TypeVar, deprecated
+from typing_extensions import Literal, Never, Self, TypeIs, TypeVar, deprecated
 
 from pydantic_graph import End, Graph, GraphRun, GraphRunContext
 from pydantic_graph._utils import get_event_loop
@@ -1688,6 +1688,51 @@ async def run_mcp_servers(self) -> AsyncIterator[None]:
         finally:
             await exit_stack.aclose()
 
+    async def to_cli(self: Self, deps: AgentDepsT = None) -> None:
+        """Run the agent in a CLI chat interface.
+
+        Example:
+        ```python {title="agent_to_cli.py" test="skip"}
+        from pydantic_ai import Agent
+
+        agent = Agent('openai:gpt-4o', instructions='You always respond in Italian.')
+
+        async def main():
+            await agent.to_cli()
+        ```
+        """
+        from prompt_toolkit import PromptSession
+        from prompt_toolkit.history import FileHistory
+        from rich.console import Console
+
+        from pydantic_ai._cli import PROMPT_HISTORY_PATH, run_chat
+
+        # TODO(Marcelo): We need to refactor the CLI code to be able to be able to just pass `agent`, `deps` and
+        # `prog_name` from here.
+
+        session: PromptSession[Any] = PromptSession(history=FileHistory(str(PROMPT_HISTORY_PATH)))
+        await run_chat(
+            session=session,
+            stream=True,
+            agent=self,
+            deps=deps,
+            console=Console(),
+            code_theme='monokai',
+            prog_name='pydantic-ai',
+        )
+
+    def to_cli_sync(self: Self, deps: AgentDepsT = None) -> None:
+        """Run the agent in a CLI chat interface with the non-async interface.
+
+        ```python {title="agent_to_cli_sync.py" test="skip"}
+        from pydantic_ai import Agent
+
+        agent = Agent('openai:gpt-4o', instructions='You always respond in Italian.')
+        agent.to_cli_sync()
+        ```
+        """
+        return get_event_loop().run_until_complete(self.to_cli(deps=deps))
+
 
 @dataclasses.dataclass(repr=False)
 class AgentRun(Generic[AgentDepsT, OutputDataT]):

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,
+    )