diff --git a/src/cai/repl/commands/model.py b/src/cai/repl/commands/model.py index a0a64994..a16bc5d4 100644 --- a/src/cai/repl/commands/model.py +++ b/src/cai/repl/commands/model.py @@ -22,6 +22,10 @@ LITELLM_URL = ( "model_prices_and_context_window.json" ) +# Global cache shared between /model and /model-show commands +_GLOBAL_MODEL_CACHE = [] +_GLOBAL_MODEL_NUMBERS = {} + def get_predefined_model_categories() -> Dict[str, List[Dict[str, str]]]: """Get the predefined model categories as the single source of truth. @@ -155,6 +159,43 @@ def get_predefined_model_names() -> List[str]: return [model["name"] for model in get_all_predefined_models()] +def load_all_available_models() -> tuple[List[str], List[Dict[str, Any]]]: + """Load all available models (predefined + LiteLLM + Ollama) in consistent order. + + This ensures /model and /model-show use the same numbering. + + Returns: + Tuple of (all_model_names, ollama_models_data) + """ + # Predefined models + predefined = [model["name"] for model in get_all_predefined_models()] + + # LiteLLM models + litellm_names = [] + try: + response = requests.get(LITELLM_URL, timeout=5) + if response.status_code == 200: + litellm_names = sorted(response.json().keys()) + except Exception: # pylint: disable=broad-except + pass + + # Ollama models + ollama_data = [] + ollama_names = [] + try: + api_base = get_ollama_api_base() + response = requests.get(f"{api_base.replace('/v1', '')}/api/tags", timeout=1) + if response.status_code == 200: + data = response.json() + ollama_data = data.get('models', data.get('items', [])) + ollama_names = [m.get('name', '') for m in ollama_data if m.get('name')] + except Exception: # pylint: disable=broad-except + pass + + all_models = predefined + litellm_names + ollama_names + return all_models, ollama_data + + class ModelCommand(Command): """Command for viewing and changing the current LLM model.""" @@ -196,29 +237,21 @@ class ModelCommand(Command): Returns: bool: True if the model was changed successfully """ - # Get all predefined models from the shared source of truth - # pylint: disable=invalid-name - ALL_MODELS = get_all_predefined_models() - - # Also fetch LiteLLM model names to make numbering consistent with /model-show - litellm_model_names = [] - try: - response = requests.get(LITELLM_URL, timeout=5) - if response.status_code == 200: - litellm_data = response.json() - # Add LiteLLM model names (sorted for consistency with /model-show) - litellm_model_names = sorted(litellm_data.keys()) - except Exception: # pylint: disable=broad-except - # Silently fail if LiteLLM is not available - pass - - # Update cached models to include all models for number selection (consistent with /model-show) - predefined_model_names = [model["name"] for model in ALL_MODELS] - self.cached_models = predefined_model_names + litellm_model_names - self.cached_model_numbers = { + # Load all models with consistent numbering and update global cache + global _GLOBAL_MODEL_CACHE, _GLOBAL_MODEL_NUMBERS + _GLOBAL_MODEL_CACHE, ollama_models_data = load_all_available_models() + _GLOBAL_MODEL_NUMBERS = { str(i): model_name - for i, model_name in enumerate(self.cached_models, 1) + for i, model_name in enumerate(_GLOBAL_MODEL_CACHE, 1) } + self.cached_models = _GLOBAL_MODEL_CACHE + self.cached_model_numbers = _GLOBAL_MODEL_NUMBERS + + # Get predefined and litellm counts for display + ALL_MODELS = get_all_predefined_models() + predefined_model_names = [model["name"] for model in ALL_MODELS] + litellm_model_names = [m for m in self.cached_models[len(predefined_model_names):] + if m not in [d.get('name') for d in ollama_models_data]] if not args: # pylint: disable=too-many-nested-blocks # Display current model @@ -272,61 +305,35 @@ class ModelCommand(Command): model["description"] ) - # Ollama models (if available) - # pylint: disable=too-many-nested-blocks - try: - # Get Ollama models with a short timeout to prevent hanging - api_base = get_ollama_api_base() - ollama_base = api_base.replace('/v1', '') - response = requests.get( - f"{ollama_base}/api/tags", - timeout=1 - ) - - if response.status_code == 200: - data = response.json() - ollama_models = [] - - if 'models' in data: - ollama_models = data['models'] - else: - # Fallback for older Ollama versions - ollama_models = data.get('items', []) - - # Add Ollama models to the table with continuing numbers - # (after predefined models + LiteLLM models in numbering) - start_index = len(predefined_model_names) + len(litellm_model_names) + 1 - for i, model in enumerate(ollama_models, start_index): - model_name = model.get('name', '') - model_size = model.get('size', 0) - # Convert size to human-readable format - size_str = "" - if model_size: - size_mb = model_size / (1024 * 1024) - if model_size < 1024 * 1024 * 1024: - size_str = f"{size_mb:.1f} MB" - else: - size_gb = size_mb / 1024 - size_str = f"{size_gb:.1f} GB" - - # Ollama models are free to use locally - model_description = "Local model" - if size_str: - model_description += f" ({size_str})" - - model_table.add_row( - str(i), - model_name, - "Ollama", - "Local", - "Free", - "Free", - model_description - ) - # Add to cached models for numeric selection - self.cached_models.append(model_name) - self.cached_model_numbers[str(i)] = model_name - except Exception: # pylint: disable=broad-except + # Ollama models (display from already loaded data) + if ollama_models_data: + start_index = len(predefined_model_names) + len(litellm_model_names) + 1 + for i, model in enumerate(ollama_models_data, start_index): + model_name = model.get('name', '') + model_size = model.get('size', 0) + size_str = "" + if model_size: + size_mb = model_size / (1024 * 1024) + if model_size < 1024 * 1024 * 1024: + size_str = f"{size_mb:.1f} MB" + else: + size_gb = size_mb / 1024 + size_str = f"{size_gb:.1f} GB" + + model_description = "Local model" + if size_str: + model_description += f" ({size_str})" + + model_table.add_row( + str(i), + model_name, + "Ollama", + "Local", + "Free", + "Free", + model_description + ) + else: # pylint: disable=broad-except # Add a note about Ollama if we couldn't fetch models start_index = len(predefined_model_names) + len(litellm_model_names) + 1 model_table.add_row( @@ -436,6 +443,15 @@ class ModelShowCommand(Command): if args: # If there are still args left, use as search term search_term = args[0].lower() + # Load all models and update global cache for consistent numbering with /model + global _GLOBAL_MODEL_CACHE, _GLOBAL_MODEL_NUMBERS + all_model_names, ollama_models_data = load_all_available_models() + _GLOBAL_MODEL_CACHE = all_model_names + _GLOBAL_MODEL_NUMBERS = { + str(i): model_name + for i, model_name in enumerate(_GLOBAL_MODEL_CACHE, 1) + } + # Fetch model pricing data from LiteLLM GitHub repository try: with console.status( @@ -482,10 +498,14 @@ class ModelShowCommand(Command): # Count models for summary total_models = 0 displayed_models = 0 - model_index = 1 - # Process and display models + # Process and display models (use global cache for numbering) for model_name, model_info in sorted(model_data.items()): + # Find the model index from global cache + try: + model_index = _GLOBAL_MODEL_CACHE.index(model_name) + 1 + except ValueError: + continue # Model not in cache, skip total_models += 1 # Skip if showing only supported models and no function calling @@ -569,71 +589,46 @@ class ModelShowCommand(Command): features_str ) - model_index += 1 - - # Now add Ollama models if available - try: - # Get Ollama models with a short timeout - api_base = get_ollama_api_base() - api_tags = f"{api_base.replace('/v1', '')}/api/tags" - ollama_response = requests.get(api_tags, timeout=1) - - if ollama_response.status_code == 200: - ollama_data = ollama_response.json() - ollama_models = [] - - if 'models' in ollama_data: - ollama_models = ollama_data['models'] + # Add Ollama models to the table (already loaded in global cache) + for model in ollama_models_data: + model_name = model.get('name', '') + + # Skip if search term provided and not in model name + if search_term and search_term not in model_name.lower(): + continue + + # Find index from global cache + try: + model_index = _GLOBAL_MODEL_CACHE.index(model_name) + 1 + except ValueError: + continue + + total_models += 1 + displayed_models += 1 + + model_size = model.get('size', 0) + size_str = "" + if model_size: + size_mb = model_size / (1024 * 1024) + if model_size < 1024 * 1024 * 1024: + size_str = f"{size_mb:.1f} MB" else: - # Fallback for older Ollama versions - ollama_models = ollama_data.get('items', []) - - # Add Ollama models to the table - for model in ollama_models: - model_name = model.get('name', '') - - # Skip if search term provided and not in model name - if (search_term and - search_term not in model_name.lower()): - continue - - total_models += 1 - displayed_models += 1 - - model_size = model.get('size', 0) - # Convert size to human-readable format - size_str = "" - if model_size: - size_mb = model_size / (1024 * 1024) - if model_size < 1024 * 1024 * 1024: - size_str = f"{size_mb:.1f} MB" - else: - size_gb = size_mb / 1024 - size_str = f"{size_gb:.1f} GB" - - # Add row to table - model_description = "Local model" - if size_str: - model_description += f" ({size_str})" - - model_table.add_row( - str(model_index), - model_name, - "Ollama", - "Varies", - "Free", - "Free", - model_description - ) - - model_index += 1 - except Exception: # pylint: disable=broad-except - # Silently fail if Ollama is not available - # This is acceptable as Ollama is optional and we don't want to - # disrupt the user experience if it's not running - console.print( - "[dim]Ollama models not available[/dim]", - style="dim") + size_gb = size_mb / 1024 + size_str = f"{size_gb:.1f} GB" + + model_description = "Local model" + if size_str: + model_description += f" ({size_str})" + + model_table.add_row( + str(model_index), + model_name, + "Ollama", + "Varies", + "Free", + "Free", + model_description + ) # Display the table console.print(model_table)