diff --git a/code_puppy/command_line/core_commands.py b/code_puppy/command_line/core_commands.py index 0608baf89..554149589 100644 --- a/code_puppy/command_line/core_commands.py +++ b/code_puppy/command_line/core_commands.py @@ -413,14 +413,57 @@ def handle_agent_command(command: str) -> bool: return True -async def interactive_model_picker() -> str | None: - """Show an interactive arrow-key selector to pick a model (async version). +from code_puppy.command_line.model_picker_constants import ( + CURRENT_MODEL_PREFIX, + CURRENT_MODEL_SUFFIX, + OTHER_MODEL_PREFIX, +) + + +def build_model_choices(model_names: list[str], current_model: str) -> list[str]: + """Build formatted choice strings for the model picker. + + Args: + model_names: List of available model names + current_model: The currently active model name + + Returns: + List of formatted choice strings with indicators for the current model + """ + choices = [] + for model_name in model_names: + if model_name == current_model: + choices.append(f"{CURRENT_MODEL_PREFIX}{model_name}{CURRENT_MODEL_SUFFIX}") + else: + choices.append(f"{OTHER_MODEL_PREFIX}{model_name}") + return choices + + +def parse_model_choice(choice: str) -> str: + """Extract the model name from a formatted choice string. + + Args: + choice: A formatted choice string like '✓ gpt-4 (current)' or ' claude-3' + + Returns: + The clean model name (e.g., 'gpt-4' or 'claude-3') + """ + # Remove prefix ("✓ " or " ") + model = choice.lstrip(CURRENT_MODEL_PREFIX.strip()).lstrip(OTHER_MODEL_PREFIX).strip() + # Remove suffix (" (current)") + if model.endswith(CURRENT_MODEL_SUFFIX): + model = model[: -len(CURRENT_MODEL_SUFFIX)].strip() + return model + + +def interactive_model_picker() -> str | None: + """Show an interactive arrow-key selector to pick a model. Returns: The selected model name, or None if cancelled """ - import asyncio import sys + import time from rich.console import Console from rich.panel import Panel @@ -431,19 +474,12 @@ async def interactive_model_picker() -> str | None: load_model_names, ) from code_puppy.tools.command_runner import set_awaiting_user_input - from code_puppy.tools.common import arrow_select_async + from code_puppy.tools.common import arrow_select # Load available models model_names = load_model_names() current_model = get_active_model() - - # Build choices with current model indicator - choices = [] - for model_name in model_names: - if model_name == current_model: - choices.append(f"✓ {model_name} (current)") - else: - choices.append(f" {model_name}") + choices = build_model_choices(model_names, current_model) # Create panel content panel_content = Text() @@ -451,7 +487,6 @@ async def interactive_model_picker() -> str | None: panel_content.append("Current model: ", style="dim") panel_content.append(current_model, style="bold green") - # Display panel panel = Panel( panel_content, title="[bold white]Model Selection[/bold white]", @@ -459,48 +494,31 @@ async def interactive_model_picker() -> str | None: padding=(1, 2), ) - # Pause spinners BEFORE showing panel + # Pause spinners and show panel set_awaiting_user_input(True) - await asyncio.sleep(0.3) # Let spinners fully stop - - local_console = Console() - emit_info("") - local_console.print(panel) - emit_info("") - - # Flush output before prompt_toolkit takes control - sys.stdout.flush() - sys.stderr.flush() - await asyncio.sleep(0.1) - - selected_model = None + time.sleep(0.05) # Give spinner thread time to notice the flag (runs every 50ms) try: - # Final flush + local_console = Console() + emit_info("") + local_console.print(panel) + emit_info("") sys.stdout.flush() - # Show arrow-key selector (async version) - choice = await arrow_select_async( + # Show arrow-key selector (synchronous - no async complications!) + choice = arrow_select( "💭 Which model would you like to use?", choices, ) - - # Extract model name from choice (remove prefix and suffix) - if choice: - # Remove the "✓ " or " " prefix and " (current)" suffix if present - selected_model = choice.strip().lstrip("✓").strip() - if selected_model.endswith(" (current)"): - selected_model = selected_model[:-10].strip() + return parse_model_choice(choice) if choice else None except (KeyboardInterrupt, EOFError): emit_error("Cancelled by user") - selected_model = None + return None finally: set_awaiting_user_input(False) - return selected_model - @register_command( name="model", @@ -511,8 +529,6 @@ async def interactive_model_picker() -> str | None: ) def handle_model_command(command: str) -> bool: """Set the active model.""" - import asyncio - from code_puppy.command_line.model_picker_completion import ( get_active_model, load_model_names, @@ -525,17 +541,7 @@ def handle_model_command(command: str) -> bool: # If just /model or /m with no args, show interactive picker if len(tokens) == 1: try: - # Run the async picker using asyncio utilities - # Since we're called from an async context but this function is sync, - # we need to carefully schedule and wait for the coroutine - import concurrent.futures - - # Create a new event loop in a thread and run the picker there - with concurrent.futures.ThreadPoolExecutor() as executor: - future = executor.submit( - lambda: asyncio.run(interactive_model_picker()) - ) - selected_model = future.result(timeout=300) # 5 min timeout + selected_model = interactive_model_picker() if selected_model: set_active_model(selected_model) @@ -543,12 +549,26 @@ def handle_model_command(command: str) -> bool: else: emit_warning("Model selection cancelled") return True + + except (KeyboardInterrupt, EOFError): + # User cancelled - already handled in picker, just return + return True + + except RuntimeError as e: + # Likely called from async context - show helpful error + emit_warning(f"Cannot show interactive picker: {e}") + model_names = load_model_names() + emit_warning("Usage: /model or /m ") + emit_warning(f"Available models: {', '.join(model_names)}") + return True + except Exception as e: - # Fallback to old behavior if picker fails + # Unexpected error - log details for debugging import traceback + from code_puppy.error_logging import log_error + log_error(e, context="interactive_model_picker") emit_warning(f"Interactive picker failed: {e}") - emit_warning(f"Traceback: {traceback.format_exc()}") model_names = load_model_names() emit_warning("Usage: /model or /m ") emit_warning(f"Available models: {', '.join(model_names)}") diff --git a/code_puppy/command_line/model_picker_completion.py b/code_puppy/command_line/model_picker_completion.py index 81bbdf097..930505bb2 100644 --- a/code_puppy/command_line/model_picker_completion.py +++ b/code_puppy/command_line/model_picker_completion.py @@ -116,65 +116,69 @@ def _find_matching_model(rest: str, model_names: list[str]) -> Optional[str]: return None +def _process_model_command( + text: str, content: str, cmd_prefix: str, model_names: list[str] +) -> Optional[str]: + """Process a model command and return the remaining text. + + Args: + text: Original input text (preserves whitespace) + content: Stripped content for parsing + cmd_prefix: The command prefix to check for (e.g., '/model ' or '/m ') + model_names: List of available model names + + Returns: + Remaining text after stripping command and model, or None if no match + """ + if not content.lower().startswith(cmd_prefix): + return None + + # Extract command and rest + cmd = content.split(" ", 1)[0] # Get the actual command (preserves case) + rest = content[len(cmd):].strip() + + # Find the best matching model + model = _find_matching_model(rest, model_names) + if not model: + return None + + # Set the model + set_active_model(model) + + # Build pattern to find and remove from original text + cmd_and_model_pattern = cmd + " " + rest[: len(model)] + idx = text.find(cmd_and_model_pattern) + if idx == -1: + return None + + return (text[:idx] + text[idx + len(cmd_and_model_pattern) :]).strip() + + def update_model_in_input(text: str) -> Optional[str]: - # If input starts with /model or /m and a model name, set model and strip it out + """Parse and handle /model or /m commands in input text. + + If input starts with /model or /m followed by a model name, + sets that model as active and returns the remaining text. + + Args: + text: The input text to parse + + Returns: + Remaining text after stripping command and model, or None if no match + """ content = text.strip() model_names = load_model_names() - # Check for /model command (require space after /model, case-insensitive) - if content.lower().startswith("/model "): - # Find the actual /model command (case-insensitive) - model_cmd = content.split(" ", 1)[0] # Get the command part - rest = content[len(model_cmd) :].strip() # Remove the actual command - - # Find the best matching model - model = _find_matching_model(rest, model_names) - if model: - # Found a matching model - now extract it properly - set_active_model(model) - - # Find the actual model name in the original text (preserving case) - # We need to find where the model ends in the original rest string - model_end_idx = len(model) - - # Build the full command+model part to remove - cmd_and_model_pattern = model_cmd + " " + rest[:model_end_idx] - idx = text.find(cmd_and_model_pattern) - if idx != -1: - new_text = ( - text[:idx] + text[idx + len(cmd_and_model_pattern) :] - ).strip() - return new_text - return None - - # Check for /m command (case-insensitive) - elif content.lower().startswith("/m ") and not content.lower().startswith( - "/model " - ): - # Find the actual /m command (case-insensitive) - m_cmd = content.split(" ", 1)[0] # Get the command part - rest = content[len(m_cmd) :].strip() # Remove the actual command - - # Find the best matching model - model = _find_matching_model(rest, model_names) - if model: - # Found a matching model - now extract it properly - set_active_model(model) - - # Find the actual model name in the original text (preserving case) - # We need to find where the model ends in the original rest string - model_end_idx = len(model) - - # Build the full command+model part to remove - # Handle space variations in the original text - cmd_and_model_pattern = m_cmd + " " + rest[:model_end_idx] - idx = text.find(cmd_and_model_pattern) - if idx != -1: - new_text = ( - text[:idx] + text[idx + len(cmd_and_model_pattern) :] - ).strip() - return new_text - return None + # Try /model first (must check before /m since /model starts with /m) + result = _process_model_command(text, content, "/model ", model_names) + if result is not None: + return result + + # Try /m (but not if it's actually /model) + if not content.lower().startswith("/model"): + result = _process_model_command(text, content, "/m ", model_names) + if result is not None: + return result return None diff --git a/code_puppy/command_line/model_picker_constants.py b/code_puppy/command_line/model_picker_constants.py new file mode 100644 index 000000000..c7836063e --- /dev/null +++ b/code_puppy/command_line/model_picker_constants.py @@ -0,0 +1,13 @@ +"""Shared constants for model picker formatting. + +These constants are used by both the interactive model picker (core_commands.py) +and the model completion system (model_picker_completion.py) to ensure +consistent formatting and parsing of model choices. +""" + +# Prefixes for model choices in the picker +CURRENT_MODEL_PREFIX = "✓ " +OTHER_MODEL_PREFIX = " " + +# Suffix for the currently active model +CURRENT_MODEL_SUFFIX = " (current)" diff --git a/tests/command_line/test_model_picker_core.py b/tests/command_line/test_model_picker_core.py new file mode 100644 index 000000000..3d63f6004 --- /dev/null +++ b/tests/command_line/test_model_picker_core.py @@ -0,0 +1,492 @@ +"""Tests for model picker functions in core_commands.py.""" + +import pytest +from unittest.mock import patch, MagicMock + + +class TestModelPickerConstants: + """Test the shared constants module.""" + + def test_constants_are_defined(self): + from code_puppy.command_line.model_picker_constants import ( + CURRENT_MODEL_PREFIX, + CURRENT_MODEL_SUFFIX, + OTHER_MODEL_PREFIX, + ) + + assert CURRENT_MODEL_PREFIX == "✓ " + assert OTHER_MODEL_PREFIX == " " + assert CURRENT_MODEL_SUFFIX == " (current)" + + def test_prefix_lengths_match(self): + """Both prefixes should be the same length for alignment.""" + from code_puppy.command_line.model_picker_constants import ( + CURRENT_MODEL_PREFIX, + OTHER_MODEL_PREFIX, + ) + + assert len(CURRENT_MODEL_PREFIX) == len(OTHER_MODEL_PREFIX) + + +class TestBuildModelChoices: + """Test the build_model_choices function.""" + + def test_marks_current_model(self): + from code_puppy.command_line.core_commands import build_model_choices + + choices = build_model_choices(["gpt-4", "claude-3"], "gpt-4") + + assert len(choices) == 2 + assert "✓" in choices[0] # gpt-4 is current + assert "(current)" in choices[0] + assert "✓" not in choices[1] # claude-3 is not current + assert "(current)" not in choices[1] + + def test_all_models_not_current(self): + from code_puppy.command_line.core_commands import build_model_choices + + choices = build_model_choices(["gpt-4", "claude-3"], "gemini") + + for choice in choices: + assert "✓" not in choice + assert "(current)" not in choice + + def test_empty_model_list(self): + from code_puppy.command_line.core_commands import build_model_choices + + choices = build_model_choices([], "gpt-4") + assert choices == [] + + def test_single_model_is_current(self): + from code_puppy.command_line.core_commands import build_model_choices + + choices = build_model_choices(["gpt-4"], "gpt-4") + + assert len(choices) == 1 + assert "✓ gpt-4 (current)" == choices[0] + + def test_preserves_model_order(self): + from code_puppy.command_line.core_commands import build_model_choices + + models = ["alpha", "beta", "gamma"] + choices = build_model_choices(models, "beta") + + assert "alpha" in choices[0] + assert "beta" in choices[1] + assert "gamma" in choices[2] + + +class TestParseModelChoice: + """Test the parse_model_choice function.""" + + def test_parse_current_model(self): + from code_puppy.command_line.core_commands import parse_model_choice + + result = parse_model_choice("✓ gpt-4 (current)") + assert result == "gpt-4" + + def test_parse_non_current_model(self): + from code_puppy.command_line.core_commands import parse_model_choice + + result = parse_model_choice(" claude-3") + assert result == "claude-3" + + def test_parse_model_with_hyphens(self): + from code_puppy.command_line.core_commands import parse_model_choice + + result = parse_model_choice("✓ gpt-4-turbo-preview (current)") + assert result == "gpt-4-turbo-preview" + + def test_parse_model_with_numbers(self): + from code_puppy.command_line.core_commands import parse_model_choice + + result = parse_model_choice(" claude-3.5-sonnet") + assert result == "claude-3.5-sonnet" + + def test_roundtrip_current_model(self): + """Build then parse should return original model name.""" + from code_puppy.command_line.core_commands import ( + build_model_choices, + parse_model_choice, + ) + + models = ["gpt-4", "claude-3"] + choices = build_model_choices(models, "gpt-4") + + # Parse the current model choice + parsed = parse_model_choice(choices[0]) + assert parsed == "gpt-4" + + def test_roundtrip_non_current_model(self): + """Build then parse should return original model name.""" + from code_puppy.command_line.core_commands import ( + build_model_choices, + parse_model_choice, + ) + + models = ["gpt-4", "claude-3"] + choices = build_model_choices(models, "gpt-4") + + # Parse the non-current model choice + parsed = parse_model_choice(choices[1]) + assert parsed == "claude-3" + + def test_roundtrip_all_models(self): + """All models should roundtrip correctly.""" + from code_puppy.command_line.core_commands import ( + build_model_choices, + parse_model_choice, + ) + + models = ["gpt-4", "claude-3", "gemini-pro", "llama-2"] + current = "claude-3" + choices = build_model_choices(models, current) + + for i, model in enumerate(models): + parsed = parse_model_choice(choices[i]) + assert parsed == model, f"Model {model} did not roundtrip correctly" + + +class TestInteractiveModelPicker: + """Test the interactive_model_picker function.""" + + def test_returns_selected_model(self): + from code_puppy.command_line.core_commands import interactive_model_picker + + with ( + patch( + "code_puppy.command_line.model_picker_completion.load_model_names", + return_value=["gpt-4", "claude-3"], + ), + patch( + "code_puppy.command_line.model_picker_completion.get_active_model", + return_value="gpt-4", + ), + patch( + "code_puppy.tools.command_runner.set_awaiting_user_input" + ), + patch( + "code_puppy.messaging.emit_info" + ), + patch( + "code_puppy.messaging.emit_error" + ), + patch( + "code_puppy.tools.common.arrow_select", + return_value=" claude-3", + ), + patch("rich.console.Console"), + ): + result = interactive_model_picker() + assert result == "claude-3" + + def test_returns_current_model_when_selected(self): + from code_puppy.command_line.core_commands import interactive_model_picker + + with ( + patch( + "code_puppy.command_line.model_picker_completion.load_model_names", + return_value=["gpt-4", "claude-3"], + ), + patch( + "code_puppy.command_line.model_picker_completion.get_active_model", + return_value="gpt-4", + ), + patch( + "code_puppy.tools.command_runner.set_awaiting_user_input" + ), + patch( + "code_puppy.messaging.emit_info" + ), + patch( + "code_puppy.messaging.emit_error" + ), + patch( + "code_puppy.tools.common.arrow_select", + return_value="✓ gpt-4 (current)", + ), + patch("rich.console.Console"), + ): + result = interactive_model_picker() + assert result == "gpt-4" + + def test_returns_none_on_keyboard_interrupt(self): + from code_puppy.command_line.core_commands import interactive_model_picker + + with ( + patch( + "code_puppy.command_line.model_picker_completion.load_model_names", + return_value=["gpt-4"], + ), + patch( + "code_puppy.command_line.model_picker_completion.get_active_model", + return_value="gpt-4", + ), + patch( + "code_puppy.tools.command_runner.set_awaiting_user_input" + ), + patch( + "code_puppy.messaging.emit_info" + ), + patch( + "code_puppy.messaging.emit_error" + ), + patch( + "code_puppy.tools.common.arrow_select", + side_effect=KeyboardInterrupt(), + ), + patch("rich.console.Console"), + ): + result = interactive_model_picker() + assert result is None + + def test_returns_none_on_eof_error(self): + from code_puppy.command_line.core_commands import interactive_model_picker + + with ( + patch( + "code_puppy.command_line.model_picker_completion.load_model_names", + return_value=["gpt-4"], + ), + patch( + "code_puppy.command_line.model_picker_completion.get_active_model", + return_value="gpt-4", + ), + patch( + "code_puppy.tools.command_runner.set_awaiting_user_input" + ), + patch( + "code_puppy.messaging.emit_info" + ), + patch( + "code_puppy.messaging.emit_error" + ), + patch( + "code_puppy.tools.common.arrow_select", + side_effect=EOFError(), + ), + patch("rich.console.Console"), + ): + result = interactive_model_picker() + assert result is None + + def test_returns_none_when_arrow_select_returns_none(self): + from code_puppy.command_line.core_commands import interactive_model_picker + + with ( + patch( + "code_puppy.command_line.model_picker_completion.load_model_names", + return_value=["gpt-4"], + ), + patch( + "code_puppy.command_line.model_picker_completion.get_active_model", + return_value="gpt-4", + ), + patch( + "code_puppy.tools.command_runner.set_awaiting_user_input" + ), + patch( + "code_puppy.messaging.emit_info" + ), + patch( + "code_puppy.messaging.emit_error" + ), + patch( + "code_puppy.tools.common.arrow_select", + return_value=None, + ), + patch("rich.console.Console"), + ): + result = interactive_model_picker() + assert result is None + + def test_always_resets_awaiting_user_input(self): + """Verify set_awaiting_user_input(False) is called even on error.""" + from code_puppy.command_line.core_commands import interactive_model_picker + + mock_set_awaiting = MagicMock() + + with ( + patch( + "code_puppy.command_line.model_picker_completion.load_model_names", + return_value=["gpt-4"], + ), + patch( + "code_puppy.command_line.model_picker_completion.get_active_model", + return_value="gpt-4", + ), + patch( + "code_puppy.tools.command_runner.set_awaiting_user_input", + mock_set_awaiting, + ), + patch( + "code_puppy.messaging.emit_info" + ), + patch( + "code_puppy.messaging.emit_error" + ), + patch( + "code_puppy.tools.common.arrow_select", + side_effect=KeyboardInterrupt(), + ), + patch("rich.console.Console"), + ): + interactive_model_picker() + + # Should be called with True first, then False + calls = mock_set_awaiting.call_args_list + assert len(calls) == 2 + assert calls[0][0][0] is True + assert calls[1][0][0] is False + + +class TestProcessModelCommand: + """Test the _process_model_command helper function.""" + + def test_returns_none_when_prefix_not_matched(self): + from code_puppy.command_line.model_picker_completion import ( + _process_model_command, + ) + + result = _process_model_command( + "hello world", "hello world", "/model ", ["gpt-4"] + ) + assert result is None + + def test_returns_none_when_model_not_found(self): + from code_puppy.command_line.model_picker_completion import ( + _process_model_command, + ) + + with patch( + "code_puppy.command_line.model_picker_completion.set_active_model" + ): + result = _process_model_command( + "/model xyz", "/model xyz", "/model ", ["gpt-4"] + ) + assert result is None + + def test_sets_model_and_returns_remaining_text(self): + from code_puppy.command_line.model_picker_completion import ( + _process_model_command, + ) + + with patch( + "code_puppy.command_line.model_picker_completion.set_active_model" + ) as mock_set: + result = _process_model_command( + "/model gpt-4 hello", "/model gpt-4 hello", "/model ", ["gpt-4"] + ) + mock_set.assert_called_once_with("gpt-4") + assert result == "hello" + + def test_handles_m_command(self): + from code_puppy.command_line.model_picker_completion import ( + _process_model_command, + ) + + with patch( + "code_puppy.command_line.model_picker_completion.set_active_model" + ) as mock_set: + result = _process_model_command( + "/m gpt-4", "/m gpt-4", "/m ", ["gpt-4"] + ) + mock_set.assert_called_once_with("gpt-4") + assert result == "" + + +class TestUpdateModelInInputRefactored: + """Test the refactored update_model_in_input function.""" + + def test_model_command_works(self): + from code_puppy.command_line.model_picker_completion import ( + update_model_in_input, + ) + + with ( + patch( + "code_puppy.model_factory.ModelFactory.load_config", + return_value={"gpt-4": {}}, + ), + patch( + "code_puppy.command_line.model_picker_completion.set_model_and_reload_agent" + ) as mock_set, + ): + result = update_model_in_input("/model gpt-4") + mock_set.assert_called_once_with("gpt-4") + assert result == "" + + def test_m_command_works(self): + from code_puppy.command_line.model_picker_completion import ( + update_model_in_input, + ) + + with ( + patch( + "code_puppy.model_factory.ModelFactory.load_config", + return_value={"gpt-4": {}}, + ), + patch( + "code_puppy.command_line.model_picker_completion.set_model_and_reload_agent" + ) as mock_set, + ): + result = update_model_in_input("/m gpt-4") + mock_set.assert_called_once_with("gpt-4") + assert result == "" + + def test_model_preferred_over_m(self): + """When input starts with /model, it should use /model not /m.""" + from code_puppy.command_line.model_picker_completion import ( + update_model_in_input, + ) + + with ( + patch( + "code_puppy.model_factory.ModelFactory.load_config", + return_value={"gpt-4": {}}, + ), + patch( + "code_puppy.command_line.model_picker_completion.set_model_and_reload_agent" + ) as mock_set, + ): + # /model should be matched, not /m + result = update_model_in_input("/model gpt-4 hello") + mock_set.assert_called_once_with("gpt-4") + assert "hello" in result + + def test_m_with_trailing_text(self): + from code_puppy.command_line.model_picker_completion import ( + update_model_in_input, + ) + + with ( + patch( + "code_puppy.model_factory.ModelFactory.load_config", + return_value={"gpt-4": {}}, + ), + patch( + "code_puppy.command_line.model_picker_completion.set_model_and_reload_agent" + ), + ): + result = update_model_in_input("/m gpt-4 tell me a joke") + assert result is not None + assert "tell me a joke" in result + + def test_no_command_returns_none(self): + from code_puppy.command_line.model_picker_completion import ( + update_model_in_input, + ) + + assert update_model_in_input("hello world") is None + + def test_invalid_model_returns_none(self): + from code_puppy.command_line.model_picker_completion import ( + update_model_in_input, + ) + + with patch( + "code_puppy.model_factory.ModelFactory.load_config", + return_value={"gpt-4": {}}, + ): + assert update_model_in_input("/model xyz") is None + assert update_model_in_input("/m xyz") is None