From a8b22804559419091e1651a70651535f0c5aee5c Mon Sep 17 00:00:00 2001 From: "Gabriel Braga L." Date: Tue, 9 Jun 2026 16:11:11 -0300 Subject: [PATCH] Add custom proxy support for OpenAI-compatible APIs in general This removes the model name check for custom proxys, if for example the naming there is different then the ones used in the HF Router. Thus giving more custom-ability for the user. --- README.md | 17 ++++++++++++++--- agent/core/llm_params.py | 19 +++++++++++++------ agent/core/local_models.py | 17 ++++++++++++++--- agent/core/model_switcher.py | 11 +++++++---- tests/unit/test_cli_local_models.py | 3 +++ tests/unit/test_llm_params.py | 22 ++++++++++++++++++++++ 6 files changed, 73 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 75c2cea20..a0ab26a15 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 d2f821c2b..ab8b99643 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 9f8a9491d..5b68bb502 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 b47b5d496..2b0323e62 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 7898e984f..70ef71e88 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 b3134cc27..67fa9d44d 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)