mirror of
https://github.com/OpenHands/OpenHands.git
synced 2025-12-26 05:48:36 +08:00
Add search API key settings to CLI (#9976)
Co-authored-by: openhands <openhands@all-hands.dev>
This commit is contained in:
parent
16106e6262
commit
6f44b7352e
@ -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.
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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()
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user