Add search API key settings to CLI (#9976)

Co-authored-by: openhands <openhands@all-hands.dev>
This commit is contained in:
Xingyao Wang 2025-07-30 14:03:29 -04:00 committed by GitHub
parent 16106e6262
commit 6f44b7352e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 195 additions and 4 deletions

View File

@ -15,6 +15,7 @@ from openhands.cli.settings import (
display_settings,
modify_llm_settings_advanced,
modify_llm_settings_basic,
modify_search_api_settings,
)
from openhands.cli.tui import (
COLOR_GREY,
@ -271,8 +272,9 @@ async def handle_settings_command(
config,
'\nWhich settings would you like to modify?',
[
'Basic',
'Advanced',
'LLM (Basic)',
'LLM (Advanced)',
'Search API (Optional)',
'Go back',
],
)
@ -281,6 +283,8 @@ async def handle_settings_command(
await modify_llm_settings_basic(config, settings_store)
elif modify_settings == 1:
await modify_llm_settings_advanced(config, settings_store)
elif modify_settings == 2:
await modify_search_api_settings(config, settings_store)
# FIXME: Currently there's an issue with the actual 'resume' behavior.

View File

@ -417,6 +417,19 @@ async def run_setup_flow(config: OpenHandsConfig, settings_store: FileSettingsSt
# 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.
@ -590,6 +603,11 @@ async def main_with_loop(loop: asyncio.AbstractEventLoop) -> None:
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()

View File

@ -74,6 +74,10 @@ def display_settings(config: OpenHandsConfig) -> None:
' Memory Condensation',
'Enabled' if config.enable_default_condenser else 'Disabled',
),
(
' Search API Key',
'********' if config.search_api_key else 'Not Set',
),
(
' Configuration File',
str(Path(config.file_store_path) / 'settings.json'),
@ -488,3 +492,75 @@ async def modify_llm_settings_advanced(
settings.enable_default_condenser = enable_memory_condensation
await settings_store.store(settings)
async def modify_search_api_settings(
config: OpenHandsConfig, settings_store: FileSettingsStore
) -> None:
"""Modify search API settings."""
session = PromptSession(key_bindings=kb_cancel())
search_api_key = None
try:
print_formatted_text(
HTML(
'\n<grey>Configure Search API Key for enhanced search capabilities.</grey>'
)
)
print_formatted_text(
HTML('<grey>You can get a Tavily API key from: https://tavily.com/</grey>')
)
print_formatted_text('')
# Show current status
current_key_status = '********' if config.search_api_key else 'Not Set'
print_formatted_text(
HTML(
f'<grey>Current Search API Key: </grey><green>{current_key_status}</green>'
)
)
print_formatted_text('')
# Ask if user wants to modify
modify_key = cli_confirm(
config,
'Do you want to modify the Search API Key?',
['Set/Update API Key', 'Remove API Key', 'Keep current setting'],
)
if modify_key == 0: # Set/Update API Key
search_api_key = await get_validated_input(
session,
'Enter Tavily Search API Key. You can get it from https://www.tavily.com/ (starts with tvly-, CTRL-c to cancel): ',
validator=lambda x: x.startswith('tvly-') if x.strip() else False,
error_message='Search API Key must start with "tvly-"',
)
elif modify_key == 1: # Remove API Key
search_api_key = '' # Empty string to remove the key
else: # Keep current setting
return
except (
UserCancelledError,
KeyboardInterrupt,
EOFError,
):
return # Return on exception
save_settings = save_settings_confirmation(config)
if not save_settings:
return
# Update config
config.search_api_key = SecretStr(search_api_key) if search_api_key else None
# Update settings store
settings = await settings_store.load()
if not settings:
settings = Settings()
settings.search_api_key = SecretStr(search_api_key) if search_api_key else None
await settings_store.store(settings)

View File

@ -549,8 +549,8 @@ class TestHandleSettingsCommand:
config = MagicMock(spec=OpenHandsConfig)
settings_store = MagicMock(spec=FileSettingsStore)
# Mock user selecting "Go back"
mock_cli_confirm.return_value = 2
# Mock user selecting "Go back" (now option 4, index 3)
mock_cli_confirm.return_value = 3
# Call the function under test
await handle_settings_command(config, settings_store)

View File

@ -9,6 +9,7 @@ from openhands.cli.settings import (
display_settings,
modify_llm_settings_advanced,
modify_llm_settings_basic,
modify_search_api_settings,
)
from openhands.cli.tui import UserCancelledError
from openhands.core.config import OpenHandsConfig
@ -46,6 +47,7 @@ class TestDisplaySettings:
config.security = security_mock
config.enable_default_condenser = True
config.search_api_key = SecretStr('tvly-test-key')
return config
@pytest.fixture
@ -65,6 +67,7 @@ class TestDisplaySettings:
config.security = security_mock
config.enable_default_condenser = True
config.search_api_key = SecretStr('tvly-test-key')
return config
@patch('openhands.cli.settings.print_container')
@ -90,6 +93,8 @@ class TestDisplaySettings:
assert 'Enabled' in settings_text
assert 'Memory Condensation:' in settings_text
assert 'Enabled' in settings_text
assert 'Search API Key:' in settings_text
assert '********' in settings_text # Search API key should be masked
assert 'Configuration File' in settings_text
assert str(Path(app_config.file_store_path)) in settings_text
@ -625,3 +630,91 @@ class TestModifyLLMSettingsAdvanced:
# Verify settings were not changed
app_config.set_llm_config.assert_not_called()
settings_store.store.assert_not_called()
class TestModifySearchApiSettings:
@pytest.fixture
def app_config(self):
config = MagicMock(spec=OpenHandsConfig)
config.search_api_key = SecretStr('tvly-existing-key')
return config
@pytest.fixture
def settings_store(self):
store = MagicMock(spec=FileSettingsStore)
store.load = AsyncMock(return_value=Settings())
store.store = AsyncMock()
return store
@pytest.mark.asyncio
@patch('openhands.cli.settings.PromptSession')
@patch('openhands.cli.settings.cli_confirm')
@patch('openhands.cli.settings.print_formatted_text')
async def test_modify_search_api_settings_set_new_key(
self, mock_print, mock_confirm, mock_session, app_config, settings_store
):
# Setup mocks
session_instance = MagicMock()
session_instance.prompt_async = AsyncMock(return_value='tvly-new-key')
mock_session.return_value = session_instance
# Mock user confirmations: Set/Update API Key, then Save
mock_confirm.side_effect = [0, 0]
# Call the function
await modify_search_api_settings(app_config, settings_store)
# Verify config was updated
assert app_config.search_api_key.get_secret_value() == 'tvly-new-key'
# Verify settings were saved
settings_store.store.assert_called_once()
args, kwargs = settings_store.store.call_args
settings = args[0]
assert settings.search_api_key.get_secret_value() == 'tvly-new-key'
@pytest.mark.asyncio
@patch('openhands.cli.settings.PromptSession')
@patch('openhands.cli.settings.cli_confirm')
@patch('openhands.cli.settings.print_formatted_text')
async def test_modify_search_api_settings_remove_key(
self, mock_print, mock_confirm, mock_session, app_config, settings_store
):
# Setup mocks
session_instance = MagicMock()
mock_session.return_value = session_instance
# Mock user confirmations: Remove API Key, then Save
mock_confirm.side_effect = [1, 0]
# Call the function
await modify_search_api_settings(app_config, settings_store)
# Verify config was updated to None
assert app_config.search_api_key is None
# Verify settings were saved
settings_store.store.assert_called_once()
args, kwargs = settings_store.store.call_args
settings = args[0]
assert settings.search_api_key is None
@pytest.mark.asyncio
@patch('openhands.cli.settings.PromptSession')
@patch('openhands.cli.settings.cli_confirm')
@patch('openhands.cli.settings.print_formatted_text')
async def test_modify_search_api_settings_keep_current(
self, mock_print, mock_confirm, mock_session, app_config, settings_store
):
# Setup mocks
session_instance = MagicMock()
mock_session.return_value = session_instance
# Mock user confirmation: Keep current setting
mock_confirm.return_value = 2
# Call the function
await modify_search_api_settings(app_config, settings_store)
# Verify settings were not changed
settings_store.store.assert_not_called()