-
Notifications
You must be signed in to change notification settings - Fork 1.4k
Add Agent.to_web() method and web chat UI #3456
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
dsfaccini
wants to merge
6
commits into
pydantic:main
Choose a base branch
from
dsfaccini:clai-chat
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
6 commits
Select commit
Hold shift + click to select a range
1cfa48e
add Agent.to_web() method and web chat UI module #3295
dsfaccini d83caa5
add "web" group
dsfaccini bfffd4b
try import create_chat_app
dsfaccini 055e120
fix tests to run on CI
dsfaccini 6bc8b16
fix example to use tool_plain
dsfaccini 32f7e1d
add clai web - tested with uv run clai web
dsfaccini File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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'] |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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.""" | ||
|
|
||
| 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 | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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 |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
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 AGENTflag is sufficient for now.