diff --git a/tests/test_auto_mode_comprehensive.py b/tests/test_auto_mode_comprehensive.py index c33e500..d4736f0 100644 --- a/tests/test_auto_mode_comprehensive.py +++ b/tests/test_auto_mode_comprehensive.py @@ -263,6 +263,7 @@ class TestAutoModeComprehensive: "OPENAI_API_KEY": None, "XAI_API_KEY": None, "OPENROUTER_API_KEY": None, + "CUSTOM_API_URL": None, "DEFAULT_MODEL": "auto", } diff --git a/tests/test_model_enumeration.py b/tests/test_model_enumeration.py index 0a78b17..ef30b56 100644 --- a/tests/test_model_enumeration.py +++ b/tests/test_model_enumeration.py @@ -6,6 +6,7 @@ all expected models based on which providers are configured via environment vari """ import importlib +import json import os import pytest @@ -121,6 +122,18 @@ class TestModelEnumeration: assert found_count == 0, "Custom models should not be included without CUSTOM_API_URL" + def test_custom_models_not_exposed_with_openrouter_only(self): + """Ensure OpenRouter access alone does not surface custom-only endpoints.""" + self._setup_environment({"OPENROUTER_API_KEY": "test-openrouter-key"}) + + tool = AnalyzeTool() + models = tool._get_available_models() + + for alias in ("local-llama", "llama3.2"): + assert ( + alias not in models + ), f"Custom model alias '{alias}' should remain hidden without CUSTOM_API_URL" + def test_no_duplicates_with_overlapping_providers(self): """Test that models aren't duplicated when multiple providers offer the same model.""" self._setup_environment( @@ -165,6 +178,54 @@ class TestModelEnumeration: else: assert model_name not in models, f"Native model {model_name} should not be present without API key" + def test_openrouter_free_model_aliases_available(self, tmp_path, monkeypatch): + """Free OpenRouter variants should expose both canonical names and aliases.""" + # Configure environment with OpenRouter access only + self._setup_environment({"OPENROUTER_API_KEY": "test-openrouter-key"}) + + # Create a temporary custom model config with a free variant + custom_config = { + "models": [ + { + "model_name": "deepseek/deepseek-r1:free", + "aliases": ["deepseek-free", "r1-free"], + "context_window": 163840, + "max_output_tokens": 8192, + "supports_extended_thinking": False, + "supports_json_mode": True, + "supports_function_calling": False, + "supports_images": False, + "max_image_size_mb": 0.0, + "description": "DeepSeek R1 free tier variant", + } + ] + } + + config_path = tmp_path / "custom_models.json" + config_path.write_text(json.dumps(custom_config), encoding="utf-8") + monkeypatch.setenv("CUSTOM_MODELS_CONFIG_PATH", str(config_path)) + + # Reset cached registries so the temporary config is loaded + from tools.shared.base_tool import BaseTool + + monkeypatch.setattr(BaseTool, "_openrouter_registry_cache", None, raising=False) + + from providers.openrouter import OpenRouterProvider + + monkeypatch.setattr(OpenRouterProvider, "_registry", None, raising=False) + + # Rebuild the provider registry with OpenRouter registered + ModelProviderRegistry._instance = None + from providers.base import ProviderType + + ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider) + + tool = AnalyzeTool() + models = tool._get_available_models() + + assert "deepseek/deepseek-r1:free" in models, "Canonical free model name should be available" + assert "deepseek-free" in models, "Free model alias should be included for MCP validation" + # DELETED: test_auto_mode_behavior_with_environment_variables # This test was fundamentally broken due to registry corruption. diff --git a/tools/shared/base_tool.py b/tools/shared/base_tool.py index a41d447..87bb1b4 100644 --- a/tools/shared/base_tool.py +++ b/tools/shared/base_tool.py @@ -1199,17 +1199,48 @@ When recommending searches, be specific about what information you need and why # Get models from enabled providers only (those with valid API keys) all_models = ModelProviderRegistry.get_available_model_names() - # Add OpenRouter models if OpenRouter is configured + # Add OpenRouter models and their aliases when OpenRouter is configured openrouter_key = os.getenv("OPENROUTER_API_KEY") if openrouter_key and openrouter_key != "your_openrouter_api_key_here": try: - from config import OPENROUTER_MODELS + registry = self._get_openrouter_registry() - all_models.extend(OPENROUTER_MODELS) - except ImportError: - pass + # Include every known alias so MCP enum matches registry capabilities + for alias in registry.list_aliases(): + config = registry.resolve(alias) + if config and config.is_custom: + # Custom-only models require CUSTOM_API_URL; defer to custom block + continue + if alias not in all_models: + all_models.append(alias) + except Exception as exc: # pragma: no cover - logged for observability + import logging - return sorted(set(all_models)) + logging.debug(f"Failed to add OpenRouter models to enum: {exc}") + + # Add custom models (and their aliases) when a custom endpoint is available + custom_url = os.getenv("CUSTOM_API_URL") + if custom_url: + try: + registry = self._get_openrouter_registry() + for alias in registry.list_aliases(): + config = registry.resolve(alias) + if config and config.is_custom and alias not in all_models: + all_models.append(alias) + except Exception as exc: # pragma: no cover - logged for observability + import logging + + logging.debug(f"Failed to add custom models to enum: {exc}") + + # Remove duplicates while preserving insertion order + seen: set[str] = set() + unique_models: list[str] = [] + for model in all_models: + if model not in seen: + seen.add(model) + unique_models.append(model) + + return unique_models def _resolve_model_context(self, arguments: dict, request) -> tuple[str, Any]: """