Rohit Malhotra 25d9cf2890
[Refactor]: Add LLMRegistry for llm services (#9589)
Co-authored-by: openhands <openhands@all-hands.dev>
Co-authored-by: Graham Neubig <neubig@gmail.com>
Co-authored-by: Engel Nyst <enyst@users.noreply.github.com>
2025-08-18 02:11:20 -04:00

747 lines
26 KiB
Python

import asyncio
import logging
import os
import sys
from prompt_toolkit import print_formatted_text
from prompt_toolkit.formatted_text import HTML
from prompt_toolkit.shortcuts import clear
import openhands.agenthub # noqa F401 (we import this to get the agents registered)
import openhands.cli.suppress_warnings # noqa: F401
from openhands.cli.commands import (
check_folder_security_agreement,
handle_commands,
)
from openhands.cli.settings import modify_llm_settings_basic
from openhands.cli.shell_config import (
ShellConfigManager,
add_aliases_to_shell_config,
alias_setup_declined,
aliases_exist_in_shell_config,
mark_alias_setup_declined,
)
from openhands.cli.tui import (
UsageMetrics,
cli_confirm,
display_agent_running_message,
display_banner,
display_event,
display_initial_user_prompt,
display_initialization_animation,
display_runtime_initialization_message,
display_welcome_message,
read_confirmation_input,
read_prompt_input,
start_pause_listener,
stop_pause_listener,
update_streaming_output,
)
from openhands.cli.utils import (
update_usage_metrics,
)
from openhands.cli.vscode_extension import attempt_vscode_extension_install
from openhands.controller import AgentController
from openhands.controller.agent import Agent
from openhands.core.config import (
OpenHandsConfig,
setup_config_from_args,
)
from openhands.core.config.condenser_config import NoOpCondenserConfig
from openhands.core.config.mcp_config import (
OpenHandsMCPConfigImpl,
)
from openhands.core.config.utils import finalize_config
from openhands.core.logger import openhands_logger as logger
from openhands.core.loop import run_agent_until_done
from openhands.core.schema import AgentState
from openhands.core.schema.exit_reason import ExitReason
from openhands.core.setup import (
create_agent,
create_controller,
create_memory,
create_runtime,
generate_sid,
initialize_repository_for_runtime,
)
from openhands.events import EventSource, EventStreamSubscriber
from openhands.events.action import (
ChangeAgentStateAction,
MessageAction,
)
from openhands.events.event import Event
from openhands.events.observation import (
AgentStateChangedObservation,
)
from openhands.io import read_task
from openhands.mcp import add_mcp_tools_to_agent
from openhands.mcp.error_collector import mcp_error_collector
from openhands.memory.condenser.impl.llm_summarizing_condenser import (
LLMSummarizingCondenserConfig,
)
from openhands.microagent.microagent import BaseMicroagent
from openhands.runtime import get_runtime_cls
from openhands.runtime.base import Runtime
from openhands.storage.settings.file_settings_store import FileSettingsStore
from openhands.utils.utils import create_registry_and_convo_stats
async def cleanup_session(
loop: asyncio.AbstractEventLoop,
agent: Agent,
runtime: Runtime,
controller: AgentController,
) -> None:
"""Clean up all resources from the current session."""
event_stream = runtime.event_stream
end_state = controller.get_state()
end_state.save_to_session(
event_stream.sid,
event_stream.file_store,
event_stream.user_id,
)
try:
current_task = asyncio.current_task(loop)
pending = [task for task in asyncio.all_tasks(loop) if task is not current_task]
if pending:
done, pending_set = await asyncio.wait(set(pending), timeout=2.0)
pending = list(pending_set)
for task in pending:
task.cancel()
agent.reset()
runtime.close()
await controller.close()
except Exception as e:
logger.error(f'Error during session cleanup: {e}')
async def run_session(
loop: asyncio.AbstractEventLoop,
config: OpenHandsConfig,
settings_store: FileSettingsStore,
current_dir: str,
task_content: str | None = None,
conversation_instructions: str | None = None,
session_name: str | None = None,
skip_banner: bool = False,
conversation_id: str | None = None,
) -> bool:
reload_microagents = False
new_session_requested = False
exit_reason = ExitReason.INTENTIONAL
sid = conversation_id or generate_sid(config, session_name)
is_loaded = asyncio.Event()
is_paused = asyncio.Event() # Event to track agent pause requests
always_confirm_mode = False # Flag to enable always confirm mode
# Show runtime initialization message
display_runtime_initialization_message(config.runtime)
# Show Initialization loader
loop.run_in_executor(
None, display_initialization_animation, 'Initializing...', is_loaded
)
llm_registry, convo_stats, config = create_registry_and_convo_stats(
config,
sid,
None,
)
agent = create_agent(config, llm_registry)
runtime = create_runtime(
config,
llm_registry,
sid=sid,
headless_mode=True,
agent=agent,
)
def stream_to_console(output: str) -> None:
# Instead of printing to stdout, pass the string to the TUI module
update_streaming_output(output)
runtime.subscribe_to_shell_stream(stream_to_console)
controller, initial_state = create_controller(agent, runtime, config, convo_stats)
event_stream = runtime.event_stream
usage_metrics = UsageMetrics()
async def prompt_for_next_task(agent_state: str) -> None:
nonlocal reload_microagents, new_session_requested, exit_reason
while True:
next_message = await read_prompt_input(
config, agent_state, multiline=config.cli_multiline_input
)
if not next_message.strip():
continue
(
close_repl,
reload_microagents,
new_session_requested,
exit_reason,
) = await handle_commands(
next_message,
event_stream,
usage_metrics,
sid,
config,
current_dir,
settings_store,
agent_state,
)
if close_repl:
return
async def on_event_async(event: Event) -> None:
nonlocal reload_microagents, is_paused, always_confirm_mode
display_event(event, config)
update_usage_metrics(event, usage_metrics)
if isinstance(event, AgentStateChangedObservation):
if event.agent_state not in [AgentState.RUNNING, AgentState.PAUSED]:
await stop_pause_listener()
if isinstance(event, AgentStateChangedObservation):
if event.agent_state in [
AgentState.AWAITING_USER_INPUT,
AgentState.FINISHED,
]:
# If the agent is paused, do not prompt for input as it's already handled by PAUSED state change
if is_paused.is_set():
return
# Reload microagents after initialization of repo.md
if reload_microagents:
microagents: list[BaseMicroagent] = (
runtime.get_microagents_from_selected_repo(None)
)
memory.load_user_workspace_microagents(microagents)
reload_microagents = False
await prompt_for_next_task(event.agent_state)
if event.agent_state == AgentState.AWAITING_USER_CONFIRMATION:
# If the agent is paused, do not prompt for confirmation
# The confirmation step will re-run after the agent has been resumed
if is_paused.is_set():
return
if always_confirm_mode:
event_stream.add_event(
ChangeAgentStateAction(AgentState.USER_CONFIRMED),
EventSource.USER,
)
return
confirmation_status = await read_confirmation_input(config)
if confirmation_status in ('yes', 'always'):
event_stream.add_event(
ChangeAgentStateAction(AgentState.USER_CONFIRMED),
EventSource.USER,
)
else: # 'no' or alternative instructions
# Tell the agent the proposed action was rejected
event_stream.add_event(
ChangeAgentStateAction(AgentState.USER_REJECTED),
EventSource.USER,
)
# Notify the user
print_formatted_text(
HTML(
'<skyblue>Okay, please tell me what I should do next/instead.</skyblue>'
)
)
# Set the always_confirm_mode flag if the user wants to always confirm
if confirmation_status == 'always':
always_confirm_mode = True
if event.agent_state == AgentState.PAUSED:
is_paused.clear() # Revert the event state before prompting for user input
await prompt_for_next_task(event.agent_state)
if event.agent_state == AgentState.RUNNING:
display_agent_running_message()
start_pause_listener(loop, is_paused, event_stream)
def on_event(event: Event) -> None:
loop.create_task(on_event_async(event))
event_stream.subscribe(EventStreamSubscriber.MAIN, on_event, sid)
await runtime.connect()
# Initialize repository if needed
repo_directory = None
if config.sandbox.selected_repo:
repo_directory = initialize_repository_for_runtime(
runtime,
selected_repository=config.sandbox.selected_repo,
)
# when memory is created, it will load the microagents from the selected repository
memory = create_memory(
runtime=runtime,
event_stream=event_stream,
sid=sid,
selected_repository=config.sandbox.selected_repo,
repo_directory=repo_directory,
conversation_instructions=conversation_instructions,
working_dir=os.getcwd(),
)
# Add MCP tools to the agent
if agent.config.enable_mcp:
# Clear any previous errors and enable collection
mcp_error_collector.clear_errors()
mcp_error_collector.enable_collection()
# Add OpenHands' MCP server by default
_, openhands_mcp_stdio_servers = (
OpenHandsMCPConfigImpl.create_default_mcp_server_config(
config.mcp_host, config, None
)
)
runtime.config.mcp.stdio_servers.extend(openhands_mcp_stdio_servers)
await add_mcp_tools_to_agent(agent, runtime, memory)
# Disable collection after startup
mcp_error_collector.disable_collection()
# Clear loading animation
is_loaded.set()
# Clear the terminal
clear()
# Show OpenHands banner and session ID if not skipped
if not skip_banner:
display_banner(session_id=sid)
welcome_message = ''
# Display number of MCP servers configured
if agent.config.enable_mcp:
total_mcp_servers = (
len(runtime.config.mcp.stdio_servers)
+ len(runtime.config.mcp.sse_servers)
+ len(runtime.config.mcp.shttp_servers)
)
if total_mcp_servers > 0:
mcp_line = f'Using {len(runtime.config.mcp.stdio_servers)} stdio MCP servers, {len(runtime.config.mcp.sse_servers)} SSE MCP servers and {len(runtime.config.mcp.shttp_servers)} SHTTP MCP servers.'
# Check for MCP errors and add indicator to the same line
if agent.config.enable_mcp and mcp_error_collector.has_errors():
mcp_line += (
' ✗ MCP errors detected (type /mcp → select View errors to view)'
)
welcome_message += mcp_line + '\n\n'
welcome_message += 'What do you want to build?' # from the application
initial_message = '' # from the user
if task_content:
initial_message = task_content
# If we loaded a state, we are resuming a previous session
if initial_state is not None:
logger.info(f'Resuming session: {sid}')
if initial_state.last_error:
# If the last session ended in an error, provide a message.
error_message = initial_state.last_error
# Check if it's an authentication error
if 'ERROR_LLM_AUTHENTICATION' in error_message:
# Start with base authentication error message
welcome_message = 'Authentication error with the LLM provider. Please check your API key.'
# Add OpenHands-specific guidance if using an OpenHands model
llm_config = config.get_llm_config()
if llm_config.model.startswith('openhands/'):
welcome_message += "\nIf you're using OpenHands models, get a new API key from https://app.all-hands.dev/settings/api-keys"
else:
# For other errors, use the standard message
initial_message = (
'NOTE: the last session ended with an error.'
"Let's get back on track. Do NOT resume your task. Ask me about it."
)
else:
# If we are resuming, we already have a task
initial_message = ''
welcome_message += '\nLoading previous conversation.'
# Show OpenHands welcome
display_welcome_message(welcome_message)
# The prompt_for_next_task will be triggered if the agent enters AWAITING_USER_INPUT.
# If the restored state is already AWAITING_USER_INPUT, on_event_async will handle it.
if initial_message:
display_initial_user_prompt(initial_message)
event_stream.add_event(MessageAction(content=initial_message), EventSource.USER)
else:
# No session restored, no initial action: prompt for the user's first message
asyncio.create_task(prompt_for_next_task(''))
await run_agent_until_done(
controller, runtime, memory, [AgentState.STOPPED, AgentState.ERROR]
)
await cleanup_session(loop, agent, runtime, controller)
if exit_reason == ExitReason.INTENTIONAL:
print_formatted_text('✅ Session terminated successfully.\n')
else:
print_formatted_text(f'⚠️ Session was interrupted: {exit_reason.value}\n')
return new_session_requested
async def run_setup_flow(config: OpenHandsConfig, settings_store: FileSettingsStore):
"""Run the setup flow to configure initial settings.
Returns:
bool: True if settings were successfully configured, False otherwise.
"""
# Display the banner with ASCII art first
display_banner(session_id='setup')
print_formatted_text(
HTML('<grey>No settings found. Starting initial setup...</grey>\n')
)
# Use the existing settings modification function for basic setup
await modify_llm_settings_basic(config, settings_store)
# Ask if user wants to configure search API settings
print_formatted_text('')
setup_search = cli_confirm(
config,
'Would you like to configure Search API settings (optional)?',
['Yes', 'No'],
)
if setup_search == 0: # Yes
from openhands.cli.settings import modify_search_api_settings
await modify_search_api_settings(config, settings_store)
def run_alias_setup_flow(config: OpenHandsConfig) -> None:
"""Run the alias setup flow to configure shell aliases.
Prompts the user to set up aliases for 'openhands' and 'oh' commands.
Handles existing aliases by offering to keep or remove them.
Args:
config: OpenHands configuration
"""
print_formatted_text('')
print_formatted_text(HTML('<gold>🚀 Welcome to OpenHands CLI!</gold>'))
print_formatted_text('')
# Show the normal setup flow
print_formatted_text(
HTML('<grey>Would you like to set up convenient shell aliases?</grey>')
)
print_formatted_text('')
print_formatted_text(
HTML('<grey>This will add the following aliases to your shell profile:</grey>')
)
print_formatted_text(
HTML(
'<grey> • <b>openhands</b> → uvx --python 3.12 --from openhands-ai openhands</grey>'
)
)
print_formatted_text(
HTML(
'<grey> • <b>oh</b> → uvx --python 3.12 --from openhands-ai openhands</grey>'
)
)
print_formatted_text('')
print_formatted_text(
HTML(
'<ansiyellow>⚠️ Note: This requires uv to be installed first.</ansiyellow>'
)
)
print_formatted_text(
HTML(
'<ansiyellow> Installation guide: https://docs.astral.sh/uv/getting-started/installation</ansiyellow>'
)
)
print_formatted_text('')
# Use cli_confirm to get user choice
choice = cli_confirm(
config,
'Set up shell aliases?',
['Yes, set up aliases', 'No, skip this step'],
)
if choice == 0: # User chose "Yes"
success = add_aliases_to_shell_config()
if success:
print_formatted_text('')
print_formatted_text(
HTML('<ansigreen>✅ Aliases added successfully!</ansigreen>')
)
# Get the appropriate reload command using the shell config manager
shell_manager = ShellConfigManager()
reload_cmd = shell_manager.get_reload_command()
print_formatted_text(
HTML(
f'<grey>Run <b>{reload_cmd}</b> (or restart your terminal) to use the new aliases.</grey>'
)
)
else:
print_formatted_text('')
print_formatted_text(
HTML(
'<ansired>❌ Failed to add aliases. You can set them up manually later.</ansired>'
)
)
else: # User chose "No"
# Mark that the user has declined alias setup
mark_alias_setup_declined()
print_formatted_text('')
print_formatted_text(
HTML(
'<grey>Skipped alias setup. You can run this setup again anytime.</grey>'
)
)
print_formatted_text('')
async def main_with_loop(loop: asyncio.AbstractEventLoop, args) -> None:
"""Runs the agent in CLI mode."""
# Set log level from command line argument if provided
if args.log_level and isinstance(args.log_level, str):
log_level = getattr(logging, str(args.log_level).upper())
logger.setLevel(log_level)
else:
# Set default log level to WARNING if no LOG_LEVEL environment variable is set
# (command line argument takes precedence over environment variable)
env_log_level = os.getenv('LOG_LEVEL')
if not env_log_level:
logger.setLevel(logging.WARNING)
# If `config.toml` does not exist in current directory, use the file under home directory
if not os.path.exists(args.config_file):
home_config_file = os.path.join(
os.path.expanduser('~'), '.openhands', 'config.toml'
)
logger.info(
f'Config file {args.config_file} does not exist, using default config file in home directory: {home_config_file}.'
)
args.config_file = home_config_file
# Load config from toml and override with command line arguments
config: OpenHandsConfig = setup_config_from_args(args)
# Attempt to install VS Code extension if applicable (one-time attempt)
attempt_vscode_extension_install()
# Load settings from Settings Store
# TODO: Make this generic?
settings_store = await FileSettingsStore.get_instance(config=config, user_id=None)
settings = await settings_store.load()
# Track if we've shown the banner during setup
banner_shown = False
# If settings don't exist, automatically enter the setup flow
if not settings:
# Clear the terminal before showing the banner
clear()
await run_setup_flow(config, settings_store)
banner_shown = True
settings = await settings_store.load()
# Use settings from settings store if available and override with command line arguments
if settings:
# settings.agent is not None because we check for it in setup_config_from_args
assert settings.agent is not None
config.default_agent = settings.agent
# Handle LLM configuration with proper precedence:
# 1. CLI parameters (-l) have highest precedence (already handled in setup_config_from_args)
# 2. config.toml in current directory has next highest precedence (already loaded)
# 3. ~/.openhands/settings.json has lowest precedence (handled here)
# Only apply settings from settings.json if:
# - No LLM config was specified via CLI arguments (-l)
# - The current LLM config doesn't have model or API key set (indicating it wasn't loaded from config.toml)
llm_config = config.get_llm_config()
if (
not args.llm_config
and (not llm_config.model or not llm_config.api_key)
and settings.llm_model
and settings.llm_api_key
):
logger.debug('Using LLM configuration from settings.json')
llm_config.model = settings.llm_model
llm_config.api_key = settings.llm_api_key
llm_config.base_url = settings.llm_base_url
config.set_llm_config(llm_config)
config.security.confirmation_mode = (
settings.confirmation_mode if settings.confirmation_mode else False
)
# Load search API key from settings if available and not already set from config.toml
if settings.search_api_key and not config.search_api_key:
config.search_api_key = settings.search_api_key
logger.debug('Using search API key from settings.json')
if settings.enable_default_condenser:
# TODO: Make this generic?
llm_config = config.get_llm_config()
agent_config = config.get_agent_config(config.default_agent)
agent_config.condenser = LLMSummarizingCondenserConfig(
llm_config=llm_config,
type='llm',
)
config.set_agent_config(agent_config)
config.enable_default_condenser = True
else:
agent_config = config.get_agent_config(config.default_agent)
agent_config.condenser = NoOpCondenserConfig(type='noop')
config.set_agent_config(agent_config)
config.enable_default_condenser = False
# Determine if CLI defaults should be overridden
val_override = args.override_cli_mode
should_override_cli_defaults = (
val_override is True
or (isinstance(val_override, str) and val_override.lower() in ('true', '1'))
or (isinstance(val_override, int) and val_override == 1)
)
if not should_override_cli_defaults:
config.runtime = 'cli'
if not config.workspace_base:
config.workspace_base = os.getcwd()
config.security.confirmation_mode = True
# Need to finalize config again after setting runtime to 'cli'
# This ensures Jupyter plugin is disabled for CLI runtime
finalize_config(config)
# Check if we should show the alias setup flow
# Only show it if:
# 1. Aliases don't exist in the shell configuration
# 2. User hasn't previously declined alias setup
# 3. We're in an interactive environment (not during tests or CI)
should_show_alias_setup = (
not aliases_exist_in_shell_config()
and not alias_setup_declined()
and sys.stdin.isatty()
)
if should_show_alias_setup:
# Clear the terminal if we haven't shown a banner yet (i.e., setup flow didn't run)
if not banner_shown:
clear()
run_alias_setup_flow(config)
# Don't set banner_shown = True here, so the ASCII art banner will still be shown
# TODO: Set working directory from config or use current working directory?
current_dir = config.workspace_base
if not current_dir:
raise ValueError('Workspace base directory not specified')
if not check_folder_security_agreement(config, current_dir):
# User rejected, exit application
return
# Read task from file, CLI args, or stdin
if args.file:
# For CLI usage, we want to enhance the file content with a prompt
# that instructs the agent to read and understand the file first
with open(args.file, 'r', encoding='utf-8') as file:
file_content = file.read()
# Create a prompt that instructs the agent to read and understand the file first
task_str = f"""The user has tagged a file '{args.file}'.
Please read and understand the following file content first:
```
{file_content}
```
After reviewing the file, please ask the user what they would like to do with it."""
else:
task_str = read_task(args, config.cli_multiline_input)
# Setup the runtime
get_runtime_cls(config.runtime).setup(config)
# Run the first session
new_session_requested = await run_session(
loop,
config,
settings_store,
current_dir,
task_str,
session_name=args.name,
skip_banner=banner_shown,
conversation_id=args.conversation,
)
# If a new session was requested, run it
while new_session_requested:
new_session_requested = await run_session(
loop, config, settings_store, current_dir, None
)
# Teardown the runtime
get_runtime_cls(config.runtime).teardown(config)
def run_cli_command(args):
"""Run the CLI command with proper error handling and cleanup."""
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
loop.run_until_complete(main_with_loop(loop, args))
except KeyboardInterrupt:
print_formatted_text('⚠️ Session was interrupted: interrupted\n')
except ConnectionRefusedError as e:
print_formatted_text(f'Connection refused: {e}')
sys.exit(1)
finally:
try:
# Cancel all running tasks
pending = asyncio.all_tasks(loop)
for task in pending:
task.cancel()
# Wait for all tasks to complete with a timeout
loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
loop.close()
except Exception as e:
print_formatted_text(f'Error during cleanup: {e}')
sys.exit(1)