Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion clai/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ Either way, running `clai` will start an interactive session where you can chat
## Help

```
usage: clai [-h] [-m [MODEL]] [-a AGENT] [-l] [-t [CODE_THEME]] [--no-stream] [--version] [prompt]
usage: clai [-h] [-m [MODEL]] [-a AGENT] [-l] [-t [CODE_THEME]] [--no-stream] [--version] {web} ... [prompt]

Pydantic AI CLI v...

Expand All @@ -65,6 +65,8 @@ Special prompts:
* `/cp` - copy the last response to clipboard

positional arguments:
{web} Available commands
web Launch web chat UI for discovered agents
prompt AI Prompt, if omitted fall into interactive mode

options:
Expand Down
6 changes: 6 additions & 0 deletions clai/clai/chat/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
"""Chat UI module for clai."""

from .agent_discovery import find_agents
from .cli import run_chat_command

__all__ = ['find_agents', 'run_chat_command']
186 changes: 186 additions & 0 deletions clai/clai/chat/agent_discovery.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,186 @@
"""Agent discovery using AST parsing to find pydantic_ai.Agent objects."""
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is cool but I think it's too much magic, at least for the first version of this feature (I wouldn't throw the code away; we could consider it as a separate PR later). I think the existing -a AGENT flag is sufficient for now.


from __future__ import annotations

import ast
from pathlib import Path
from typing import NamedTuple


class AgentInfo(NamedTuple):
"""Information about a discovered agent."""

module_path: str # e.g., "chat.golden_gate_bridge:agent"
file_path: Path
agent_name: str


# Directories to exclude from search
EXCLUDE_DIRS = {
'.venv',
'venv',
'env',
'node_modules',
'__pycache__',
'.pytest_cache',
'.git',
'.hg',
'.svn',
'dist',
'build',
'.eggs',
'.tox',
'.nox',
'.mypy_cache',
'.ruff_cache',
'.pydantic-work',
}


def _should_exclude_dir(dir_path: Path) -> bool:
return dir_path.name in EXCLUDE_DIRS or dir_path.name.startswith('.')


def _file_to_module_path(file_path: Path, root_dir: Path) -> str | None:
"""Convert a file path to a Python module path.

Args:
file_path: The file path to convert.
root_dir: The root directory of the project.

Returns:
The module path (e.g., "chat.my_agent") or None if conversion fails.
"""
try:
rel_path = file_path.relative_to(root_dir)

if rel_path.suffix != '.py':
return None

parts = list(rel_path.parts[:-1]) + [rel_path.stem]
module_path = '.'.join(parts)

return module_path
except (ValueError, AttributeError):
return None


def _parse_file_for_agents(file_path: Path, root_dir: Path) -> list[AgentInfo]: # noqa: C901
"""Parse a Python file for pydantic_ai.Agent instances using AST.

Args:
file_path: The Python file to parse.
root_dir: The root directory of the project.

Returns:
List of AgentInfo objects for agents found in the file.
"""
try:
content = file_path.read_text(encoding='utf-8')
except (OSError, UnicodeDecodeError):
# Skip files we can't read
return []

try:
tree = ast.parse(content, filename=str(file_path))
except SyntaxError:
# Skip files with syntax errors
return []

# Track imports of Agent class
agent_names: set[str] = set()

for node in ast.walk(tree):
# Handle: from pydantic_ai import Agent
if isinstance(node, ast.ImportFrom):
if node.module and 'pydantic_ai' in node.module:
for alias in node.names:
if alias.name == 'Agent':
agent_names.add(alias.asname or 'Agent')

# Handle: import pydantic_ai (then pydantic_ai.Agent)
elif isinstance(node, ast.Import):
for alias in node.names:
if 'pydantic_ai' in alias.name:
# Track as the alias or original name
agent_names.add(alias.asname or alias.name)

if not agent_names:
# No pydantic_ai imports, skip this file
return []

# Find agent instances
agents: list[AgentInfo] = []
module_path = _file_to_module_path(file_path, root_dir)

if not module_path:
return []

for node in ast.walk(tree):
# Look for assignments like: agent = Agent(...)
if isinstance(node, ast.Assign):
# Check if the value is a Call to Agent
if isinstance(node.value, ast.Call):
call_name = None

# Handle direct call: Agent(...)
if isinstance(node.value.func, ast.Name):
if node.value.func.id in agent_names:
call_name = node.value.func.id

# Handle attribute call: pydantic_ai.Agent(...)
elif isinstance(node.value.func, ast.Attribute):
if (
node.value.func.attr == 'Agent'
and isinstance(node.value.func.value, ast.Name)
and node.value.func.value.id in agent_names
):
call_name = 'Agent'

if call_name:
# Extract variable names being assigned
for target in node.targets:
if isinstance(target, ast.Name):
agent_name = target.id
agents.append(
AgentInfo(
module_path=f'{module_path}:{agent_name}',
file_path=file_path,
agent_name=agent_name,
)
)

return agents


def find_agents(root_dir: Path | None = None) -> list[AgentInfo]:
"""Find all pydantic_ai.Agent instances in the current directory.

Args:
root_dir: The root directory to search. Defaults to current working directory.

Returns:
List of AgentInfo objects for all discovered agents.
"""
if root_dir is None:
root_dir = Path.cwd()

agents: list[AgentInfo] = []
visited_files: set[Path] = set()

for path in root_dir.rglob('*.py'):
if path in visited_files:
continue
visited_files.add(path)

try:
# Check if any parent directory should be excluded
if any(_should_exclude_dir(parent) for parent in path.parents):
continue
except (OSError, RuntimeError):
continue

file_agents = _parse_file_for_agents(path, root_dir)
agents.extend(file_agents)

return agents
166 changes: 166 additions & 0 deletions clai/clai/chat/cli.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,166 @@
"""CLI command for launching a web chat UI for discovered agents."""

from __future__ import annotations

import importlib.util
import sys
from pathlib import Path

from pydantic_ai import Agent
from pydantic_ai.builtin_tools import AbstractBuiltinTool
from pydantic_ai.ui.web import AIModel, BuiltinTool, create_chat_app

from .agent_discovery import AgentInfo, find_agents


def load_agent_options(
config_path: Path,
) -> tuple[list[AIModel] | None, dict[str, AbstractBuiltinTool] | None, list[BuiltinTool] | None]:
"""Load agent options from a config file.

Args:
config_path: Path to the config file (e.g., agent_options.py)

Returns:
Tuple of (models, builtin_tools, builtin_tool_defs) or (None, None, None) if not found
"""
if not config_path.exists():
return None, None, None

try:
spec = importlib.util.spec_from_file_location('agent_options_config', config_path)
if spec is None or spec.loader is None:
print(f'Warning: Could not load config from {config_path}')
return None, None, None

module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)

models = getattr(module, 'AI_MODELS', None)
builtin_tools = getattr(module, 'BUILTIN_TOOLS', None)
builtin_tool_defs = getattr(module, 'BUILTIN_TOOL_DEFS', None)

return models, builtin_tools, builtin_tool_defs

except Exception as e:
print(f'Warning: Error loading config from {config_path}: {e}')
return None, None, None


def select_agent(agents: list[AgentInfo]) -> AgentInfo | None:
"""Prompt user to select an agent from the list."""
if not agents:
print('No agents found in the current directory.')
return None

if len(agents) == 1:
print(f'Found agent: {agents[0].agent_name} in {agents[0].file_path}')
return agents[0]

print('Multiple agents found:')
for i, agent_info in enumerate(agents, 1):
print(f' {i}. {agent_info.agent_name} ({agent_info.file_path})')

while True:
try:
choice = input('\nSelect an agent (enter number): ').strip()
index = int(choice) - 1
if 0 <= index < len(agents):
return agents[index]
else:
print(f'Please enter a number between 1 and {len(agents)}')
except (ValueError, KeyboardInterrupt):
print('\nSelection cancelled.')
return None


def load_agent(agent_info: AgentInfo) -> Agent | None:
"""Load an agent from the given agent info."""
sys.path.insert(0, str(Path.cwd()))

try:
spec = importlib.util.spec_from_file_location(agent_info.module_path, agent_info.file_path)
if spec is None or spec.loader is None:
print(f'Error: Could not load module from {agent_info.file_path}')
return None

module = importlib.util.module_from_spec(spec)
sys.modules[agent_info.module_path] = module
spec.loader.exec_module(module)

agent = getattr(module, agent_info.agent_name, None)
if agent is None:
print(f'Error: Agent {agent_info.agent_name} not found in module')
return None

if not isinstance(agent, Agent):
print(f'Error: {agent_info.agent_name} is not an Agent instance')
return None

return agent # pyright: ignore[reportUnknownVariableType]

except Exception as e:
print(f'Error loading agent: {e}')
return None


def run_chat_command(
root_dir: Path | None = None,
host: str = '127.0.0.1',
port: int = 8000,
config_path: Path | None = None,
auto_config: bool = True,
) -> int:
"""Run the chat command to discover and serve an agent.

Args:
root_dir: Directory to search for agents (defaults to current directory)
host: Host to bind the server to
port: Port to bind the server to
config_path: Path to agent_options.py config file
auto_config: Auto-discover agent_options.py in current directory
"""
search_dir = root_dir or Path.cwd()

print(f'Searching for agents in {search_dir}...')
agents = find_agents(search_dir)

selected = select_agent(agents)
if selected is None:
return 1

agent = load_agent(selected)
if agent is None:
return 1

models, builtin_tools, builtin_tool_defs = None, None, None
if config_path:
print(f'Loading config from {config_path}...')
models, builtin_tools, builtin_tool_defs = load_agent_options(config_path)
elif auto_config:
default_config = Path.cwd() / 'agent_options.py'
if default_config.exists():
print(f'Found config file: {default_config}')
models, builtin_tools, builtin_tool_defs = load_agent_options(default_config)

app = create_chat_app(agent, models=models, builtin_tools=builtin_tools, builtin_tool_defs=builtin_tool_defs)

print(f'\nStarting chat UI for {selected.agent_name}...')
print(f'Open your browser at: http://{host}:{port}')
print('Press Ctrl+C to stop the server\n')

try:
import uvicorn

uvicorn.run(app, host=host, port=port)
return 0
except KeyboardInterrupt:
print('\nServer stopped.')
return 0
except ImportError:
print('Error: uvicorn is required to run the chat UI')
print('Install it with: pip install uvicorn')
return 1
except Exception as e:
print(f'Error starting server: {e}')
return 1
Loading
Loading