diff --git a/README.md b/README.md index 75c2cea2..a0ab26a1 100644 --- a/README.md +++ b/README.md @@ -85,10 +85,11 @@ Inside interactive mode, switch with `/model`: /model ollama/llama3.1:8b /model lm_studio/google/gemma-3-4b /model llamacpp/llama-3.1-8b-instruct +/model custom-proxy/Azure AI/gpt-5.5 ``` -Supported local prefixes are `ollama/`, `vllm/`, `lm_studio/`, and -`llamacpp/`. +Supported local/custom prefixes are `ollama/`, `vllm/`, `lm_studio/`, +`llamacpp/`, and `custom-proxy/`. ```bash LOCAL_LLM_BASE_URL=http://localhost:8000 @@ -99,7 +100,17 @@ Set `LOCAL_LLM_BASE_URL` and optional `LOCAL_LLM_API_KEY` to use one shared local endpoint, or override a specific provider with its matching `*_BASE_URL` / `*_API_KEY` variable, such as `OLLAMA_BASE_URL` or `VLLM_API_KEY`. Provider-specific variables take precedence over the shared local variables. -Base URLs may include or omit `/v1`. +Base URLs may include or omit `/v1` for the built-in local providers. + +For a custom OpenAI-compatible proxy, set `CUSTOM_PROXY_BASE_URL` to the exact +OpenAI-compatible base URL and optional `CUSTOM_PROXY_API_KEY`, then use +`custom-proxy/`: + +```bash +CUSTOM_PROXY_BASE_URL="https://proxy.test/api/v1/" +CUSTOM_PROXY_API_KEY= +ml-intern --model "custom-proxy/Azure AI/gpt-5.5" "your prompt" +``` **CLI tool runtime:** diff --git a/agent/core/llm_params.py b/agent/core/llm_params.py index d2f821c2..ab8b9964 100644 --- a/agent/core/llm_params.py +++ b/agent/core/llm_params.py @@ -75,6 +75,11 @@ def _resolve_local_model_params( or os.environ.get(LOCAL_MODEL_BASE_URL_ENV) or provider["base_url_default"] ) + if not raw_base: + raise ValueError( + f"Set {provider['base_url_env']} or {LOCAL_MODEL_BASE_URL_ENV} " + f"to use {model_name}" + ) api_key = ( os.environ.get(provider["api_key_env"]) or os.environ.get(LOCAL_MODEL_API_KEY_ENV) @@ -82,7 +87,9 @@ def _resolve_local_model_params( ) return { "model": f"openai/{local_name}", - "api_base": _local_api_base(raw_base), + "api_base": _local_api_base(raw_base) + if provider.get("base_url_mode") != "exact" + else raw_base.strip().rstrip("/"), "api_key": api_key, } @@ -96,11 +103,11 @@ def _resolve_llm_params( """ Build LiteLLM kwargs for a given model id. - • ``ollama/``, ``vllm/``, ``lm_studio/``, and - ``llamacpp/`` — local OpenAI-compatible endpoints. The id prefix - selects a configurable localhost base URL, and the model suffix is sent - to LiteLLM as ``openai/``. These endpoints don't receive - ``reasoning_effort``. + • ``ollama/``, ``vllm/``, ``lm_studio/``, + ``llamacpp/``, and ``custom-proxy/`` — local or custom + OpenAI-compatible endpoints. The id prefix selects a configurable base + URL, and the model suffix is sent to LiteLLM as ``openai/``. + These endpoints don't receive ``reasoning_effort``. • Anything else is treated as an HF Router id. We hit the auto-routing OpenAI-compatible endpoint at ``https://router.huggingface.co/v1``. diff --git a/agent/core/local_models.py b/agent/core/local_models.py index 9f8a9491..5b68bb50 100644 --- a/agent/core/local_models.py +++ b/agent/core/local_models.py @@ -21,6 +21,12 @@ "base_url_default": "http://localhost:8080", "api_key_env": "LLAMACPP_API_KEY", }, + "custom-proxy/": { + "base_url_env": "CUSTOM_PROXY_BASE_URL", + "base_url_default": "", + "api_key_env": "CUSTOM_PROXY_API_KEY", + "base_url_mode": "exact", + }, } LOCAL_MODEL_PREFIXES = tuple(LOCAL_MODEL_PROVIDERS) @@ -48,10 +54,15 @@ def local_model_name(model_id: str) -> str | None: def is_local_model_id(model_id: str) -> bool: - """Return True for non-empty, whitespace-free local model ids.""" - if not model_id or any(char.isspace() for char in model_id): + """Return True for valid local/custom model ids.""" + if not model_id: + return False + name = local_model_name(model_id) + if name is None: return False - return local_model_name(model_id) is not None + if model_id.startswith("custom-proxy/"): + return bool(name.strip()) + return not any(char.isspace() for char in model_id) def is_reserved_local_model_id(model_id: str) -> bool: diff --git a/agent/core/model_switcher.py b/agent/core/model_switcher.py index b47b5d49..2b0323e6 100644 --- a/agent/core/model_switcher.py +++ b/agent/core/model_switcher.py @@ -59,7 +59,8 @@ def is_valid_model_id(model_id: str) -> bool: """Loose format check — lets users pick any model id. Accepts: - • ollama/, vllm/, lm_studio/, llamacpp/ + • ollama/, vllm/, lm_studio/, llamacpp/, + custom-proxy/ • /[:] (HF router; tag = provider or policy) • huggingface//[:] (same, optional LiteLLM prefix) @@ -163,8 +164,9 @@ def print_model_listing(config, console) -> None: console.print( "\n[dim]Paste any HF model id (e.g. 'MiniMaxAI/MiniMax-M2.7').\n" "Add ':fastest', ':cheapest', ':preferred', or ':' to override routing.\n" - "Use 'ollama/', 'vllm/', 'lm_studio/', or " - "'llamacpp/' for local OpenAI-compatible endpoints.[/dim]" + "Use 'ollama/', 'vllm/', 'lm_studio/', " + "'llamacpp/', or 'custom-proxy/' for local/custom " + "OpenAI-compatible endpoints.[/dim]" ) @@ -173,7 +175,8 @@ def print_invalid_id(arg: str, console) -> None: console.print( "[dim]Expected:\n" " • /[:tag] (HF router — paste from huggingface.co)\n" - " • ollama/ | vllm/ | lm_studio/ | llamacpp/[/dim]" + " • ollama/ | vllm/ | lm_studio/ | " + "llamacpp/ | custom-proxy/[/dim]" ) diff --git a/tests/unit/test_cli_local_models.py b/tests/unit/test_cli_local_models.py index 7898e984..70ef71e8 100644 --- a/tests/unit/test_cli_local_models.py +++ b/tests/unit/test_cli_local_models.py @@ -11,6 +11,7 @@ def test_local_model_helper_accepts_supported_prefixes(): assert is_local_model_id("vllm/meta-llama/Llama-3.1-8B-Instruct") assert is_local_model_id("lm_studio/google/gemma-3-4b") assert is_local_model_id("llamacpp/unsloth/Qwen3.5-2B") + assert is_local_model_id("custom-proxy/Azure AI/gpt-5.5") def test_model_switcher_accepts_supported_local_prefixes(): @@ -18,6 +19,7 @@ def test_model_switcher_accepts_supported_local_prefixes(): assert model_switcher.is_valid_model_id("vllm/meta-llama/Llama-3.1-8B") assert model_switcher.is_valid_model_id("lm_studio/google/gemma-3-4b") assert model_switcher.is_valid_model_id("llamacpp/llama-3.1-8b") + assert model_switcher.is_valid_model_id("custom-proxy/Azure AI/gpt-5.5") def test_model_switcher_rejects_empty_or_whitespace_local_ids(): @@ -25,6 +27,7 @@ def test_model_switcher_rejects_empty_or_whitespace_local_ids(): assert not model_switcher.is_valid_model_id("vllm/") assert not model_switcher.is_valid_model_id("lm_studio/") assert not model_switcher.is_valid_model_id("llamacpp/") + assert not model_switcher.is_valid_model_id("custom-proxy/") assert not model_switcher.is_valid_model_id("ollama/llama 3.1") diff --git a/tests/unit/test_llm_params.py b/tests/unit/test_llm_params.py index b3134cc2..67fa9d44 100644 --- a/tests/unit/test_llm_params.py +++ b/tests/unit/test_llm_params.py @@ -137,6 +137,28 @@ def test_resolve_llamacpp_params_strips_provider_prefix(monkeypatch): assert params["api_base"] == "http://localhost:8080/v1" +def test_resolve_custom_proxy_params_uses_explicit_endpoint(monkeypatch): + monkeypatch.setenv( + "CUSTOM_PROXY_BASE_URL", + "https://proxy.test/api/v1/", + ) + monkeypatch.setenv("CUSTOM_PROXY_API_KEY", "proxy-secret") + + params = _resolve_llm_params("custom-proxy/Azure AI/gpt-5.5") + + assert params["model"] == "openai/Azure AI/gpt-5.5" + assert params["api_base"] == "https://proxy.test/api/v1" + assert params["api_key"] == "proxy-secret" + + +def test_resolve_custom_proxy_requires_base_url(monkeypatch): + monkeypatch.delenv("CUSTOM_PROXY_BASE_URL", raising=False) + monkeypatch.delenv("LOCAL_LLM_BASE_URL", raising=False) + + with pytest.raises(ValueError, match="CUSTOM_PROXY_BASE_URL"): + _resolve_llm_params("custom-proxy/Azure AI/gpt-5.5") + + def test_local_params_reject_reasoning_effort_in_strict_mode(): with pytest.raises(UnsupportedEffortError, match="reasoning_effort"): _resolve_llm_params("ollama/llama3.1", reasoning_effort="high", strict=True)