mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
3c77e52ce4
## Description Adds first-class GCP Vertex AI proxy routing for publisher REST endpoints so Vertex requests are forwarded to a configurable regional Vertex host instead of falling through to the generic OpenAI/Anthropic/Gemini passthrough selection. Fixes #792 ## Type of Change - [x] New feature (non-breaking change that adds functionality) - [x] Documentation update ## Changes Made - Added a `vertex` provider target with `VERTEX_TARGET_API_URL` and `--vertex-api-url` support. - Registered explicit Vertex publisher routes for Google `generateContent`, `streamGenerateContent`, `countTokens` and Anthropic publisher `rawPredict`, `streamRawPredict` passthrough. - Added startup banner/routing output for Vertex AI. - Added focused tests for provider target resolution, CLI/env config, banner output, and route delegation. - Added `wiki/vertex.md` with usage examples and Google Cloud source links. ## Sources - Vertex AI Gemini inference reference: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference - Google Cloud REST authentication: https://docs.cloud.google.com/docs/authentication/rest - Google Application Default Credentials: https://docs.cloud.google.com/docs/authentication/application-default-credentials ## Testing - [x] Linting passes (`python -m ruff check .`) - [x] New tests added for new functionality - [x] Focused unit tests pass - [ ] Full unit suite completed locally - [ ] Rust tests completed locally - [ ] Type checking passes locally ## Test Output ```text $ python -m ruff check . All checks passed! $ python -m pytest tests/test_provider_registry.py tests/test_provider_proxy_routes.py tests/test_cli_proxy_env.py tests/test_banner_upstream_targets.py -q 57 passed, 1 warning in 13.75s ``` Local limitations: - `python -m pytest tests scripts/tests -q` timed out after 1 hour on this Windows machine before completing. - `cargo test -p headroom-proxy --test integration_vertex_raw_predict` could not run because `cargo` is not installed on PATH in this environment. - The commit hook's `mypy` step fails locally on an existing Windows `fcntl` typing issue in `headroom/subscription/tracker.py`; `ruff`, `ruff-format`, and plugin-version hooks passed, and the commit was made with only `mypy` skipped. ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have made corresponding changes to the documentation - [x] I have added tests that prove the feature works - [ ] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable
517 lines
18 KiB
Python
517 lines
18 KiB
Python
"""Tests for CLI proxy env variable handling and backend validation.
|
|
|
|
Verifies that:
|
|
1. Provider target URL env vars are read by `headroom proxy`
|
|
2. litellm-* backends are accepted by both CLI and argparse paths
|
|
"""
|
|
|
|
import os
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
click = pytest.importorskip("click")
|
|
pytest.importorskip("fastapi")
|
|
|
|
from click.testing import CliRunner # noqa: E402
|
|
|
|
from headroom.cli.main import main # noqa: E402
|
|
|
|
|
|
@pytest.fixture
|
|
def runner():
|
|
return CliRunner()
|
|
|
|
|
|
class TestCLIProxyEnvVars:
|
|
"""Test that the CLI proxy command reads API URL env vars."""
|
|
|
|
def test_headroom_host_from_env(self, runner):
|
|
"""HEADROOM_HOST env var should be passed to ProxyConfig."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy"],
|
|
env={"HEADROOM_HOST": "0.0.0.0"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].host == "0.0.0.0"
|
|
|
|
def test_headroom_port_from_env(self, runner):
|
|
"""HEADROOM_PORT env var should be passed to ProxyConfig."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy"],
|
|
env={"HEADROOM_PORT": "9797"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].port == 9797
|
|
|
|
def test_headroom_budget_from_env(self, runner):
|
|
"""HEADROOM_BUDGET env var should be passed to ProxyConfig."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy"],
|
|
env={"HEADROOM_BUDGET": "100.5"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].budget_limit_usd == 100.5
|
|
|
|
def test_code_aware_enabled_from_env(self, runner):
|
|
"""HEADROOM_CODE_AWARE_ENABLED env var should be passed to ProxyConfig."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy"],
|
|
env={"HEADROOM_CODE_AWARE_ENABLED": "true"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].code_aware_enabled is True
|
|
|
|
def test_code_aware_enabled_defaults_false(self, runner):
|
|
"""Without HEADROOM_CODE_AWARE_ENABLED, code-aware stays disabled in the wrapper."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
env = {k: v for k, v in os.environ.items() if k != "HEADROOM_CODE_AWARE_ENABLED"}
|
|
|
|
with (
|
|
patch("headroom.proxy.server.run_server", mock_run_server),
|
|
patch.dict(os.environ, env, clear=True),
|
|
):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].code_aware_enabled is False
|
|
|
|
def test_code_aware_enabled_from_cli_flag(self, runner):
|
|
"""--code-aware should enable code-aware compression in the wrapper."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(main, ["proxy", "--code-aware"], catch_exceptions=False)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].code_aware_enabled is True
|
|
|
|
def test_code_aware_flag_overrides_env_var(self, runner):
|
|
"""--code-aware should win over HEADROOM_CODE_AWARE_ENABLED=false."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy", "--code-aware"],
|
|
env={"HEADROOM_CODE_AWARE_ENABLED": "false"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].code_aware_enabled is True
|
|
|
|
def test_openai_target_api_url_from_env(self, runner):
|
|
"""OPENAI_TARGET_API_URL env var should be passed to ProxyConfig."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy"],
|
|
env={"OPENAI_TARGET_API_URL": "http://my-vllm:4000"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].openai_api_url == "http://my-vllm:4000"
|
|
|
|
def test_gemini_target_api_url_from_env(self, runner):
|
|
"""GEMINI_TARGET_API_URL env var should be passed to ProxyConfig."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy"],
|
|
env={"GEMINI_TARGET_API_URL": "http://my-gemini:5000"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].gemini_api_url == "http://my-gemini:5000"
|
|
|
|
def test_vertex_target_api_url_from_env(self, runner):
|
|
"""VERTEX_TARGET_API_URL env var should be passed to ProxyConfig."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy"],
|
|
env={"VERTEX_TARGET_API_URL": "https://europe-west4-aiplatform.googleapis.com"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert (
|
|
captured_config["config"].vertex_api_url
|
|
== "https://europe-west4-aiplatform.googleapis.com"
|
|
)
|
|
|
|
def test_openai_api_url_cli_flag(self, runner):
|
|
"""--openai-api-url CLI flag should take precedence."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy", "--openai-api-url", "http://from-cli:4000"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].openai_api_url == "http://from-cli:4000"
|
|
|
|
def test_vertex_api_url_cli_flag(self, runner):
|
|
"""--vertex-api-url CLI flag should take precedence."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy", "--vertex-api-url", "https://us-east5-aiplatform.googleapis.com"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert (
|
|
captured_config["config"].vertex_api_url == "https://us-east5-aiplatform.googleapis.com"
|
|
)
|
|
|
|
def test_cli_flag_overrides_env_var(self, runner):
|
|
"""CLI flag should take precedence over env var."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy", "--openai-api-url", "http://from-cli:4000"],
|
|
env={"OPENAI_TARGET_API_URL": "http://from-env:4000"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].openai_api_url == "http://from-cli:4000"
|
|
|
|
def test_no_env_var_defaults_to_none(self, runner):
|
|
"""Without env var or flag, openai_api_url should be None."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
# Ensure the env var is not set
|
|
env = {k: v for k, v in os.environ.items() if k != "OPENAI_TARGET_API_URL"}
|
|
|
|
with (
|
|
patch("headroom.proxy.server.run_server", mock_run_server),
|
|
patch.dict(os.environ, env, clear=True),
|
|
):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].openai_api_url is None
|
|
|
|
def test_both_api_urls_from_env(self, runner):
|
|
"""Both OPENAI and GEMINI target URLs can be set via env."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy"],
|
|
env={
|
|
"OPENAI_TARGET_API_URL": "http://my-vllm:4000",
|
|
"GEMINI_TARGET_API_URL": "http://my-gemini:5000",
|
|
},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].openai_api_url == "http://my-vllm:4000"
|
|
assert captured_config["config"].gemini_api_url == "http://my-gemini:5000"
|
|
|
|
def test_retry_and_connect_timeout_cli_flags(self, runner):
|
|
"""Fast-fail CLI flags should map into ProxyConfig."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
[
|
|
"proxy",
|
|
"--retry-max-attempts",
|
|
"1",
|
|
"--connect-timeout-seconds",
|
|
"3",
|
|
],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].retry_max_attempts == 1
|
|
assert captured_config["config"].connect_timeout_seconds == 3
|
|
|
|
def test_production_scaling_env_vars(self, runner):
|
|
captured = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured["config"] = config
|
|
captured["kwargs"] = kwargs
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy"],
|
|
env={
|
|
"HEADROOM_WORKERS": "4",
|
|
"HEADROOM_LIMIT_CONCURRENCY": "250",
|
|
"HEADROOM_MAX_CONNECTIONS": "200",
|
|
"HEADROOM_MAX_KEEPALIVE": "50",
|
|
},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured["config"].max_connections == 200
|
|
assert captured["config"].max_keepalive_connections == 50
|
|
# Click CLI also passes `print_banner=False` to suppress the legacy
|
|
# run_server banner (cli/proxy.py prints its own). Assert the
|
|
# production-scaling keys we care about, not the full kwargs dict.
|
|
assert captured["kwargs"]["workers"] == 4
|
|
assert captured["kwargs"]["limit_concurrency"] == 250
|
|
assert captured["kwargs"].get("print_banner") is False
|
|
|
|
def test_production_scaling_cli_flags_override_env_vars(self, runner):
|
|
captured = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured["config"] = config
|
|
captured["kwargs"] = kwargs
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
[
|
|
"proxy",
|
|
"--workers",
|
|
"3",
|
|
"--limit-concurrency",
|
|
"125",
|
|
"--max-connections",
|
|
"150",
|
|
"--max-keepalive",
|
|
"25",
|
|
],
|
|
env={
|
|
"HEADROOM_WORKERS": "4",
|
|
"HEADROOM_LIMIT_CONCURRENCY": "250",
|
|
"HEADROOM_MAX_CONNECTIONS": "200",
|
|
"HEADROOM_MAX_KEEPALIVE": "50",
|
|
},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured["config"].max_connections == 150
|
|
assert captured["config"].max_keepalive_connections == 25
|
|
# Click CLI also passes `print_banner=False`. Assert production
|
|
# scaling keys explicitly rather than the full kwargs dict.
|
|
assert captured["kwargs"]["workers"] == 3
|
|
assert captured["kwargs"]["limit_concurrency"] == 125
|
|
assert captured["kwargs"].get("print_banner") is False
|
|
|
|
|
|
class TestCLIProxyBackend:
|
|
"""Test that litellm-* backends are accepted by the CLI."""
|
|
|
|
def test_litellm_hosted_vllm_backend_accepted(self, runner):
|
|
"""--backend litellm-hosted_vllm should be accepted (not rejected)."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy", "--backend", "litellm-hosted_vllm"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].backend == "litellm-hosted_vllm"
|
|
|
|
def test_litellm_vertex_backend_accepted(self, runner):
|
|
"""--backend litellm-vertex should be accepted."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy", "--backend", "litellm-vertex"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].backend == "litellm-vertex"
|
|
|
|
def test_litellm_backend_with_openai_url(self, runner):
|
|
"""Full vLLM setup: litellm backend + OPENAI_TARGET_API_URL."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
[
|
|
"proxy",
|
|
"--backend",
|
|
"litellm-hosted_vllm",
|
|
"--openai-api-url",
|
|
"http://my-vllm:4000",
|
|
],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].backend == "litellm-hosted_vllm"
|
|
assert captured_config["config"].openai_api_url == "http://my-vllm:4000"
|
|
|
|
|
|
class TestCLIAnyllmProviderEnv:
|
|
"""Test that HEADROOM_ANYLLM_PROVIDER env var is read by the CLI."""
|
|
|
|
def test_anyllm_provider_from_env(self, runner):
|
|
"""HEADROOM_ANYLLM_PROVIDER env var should override the default."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy", "--backend", "anyllm"],
|
|
env={"HEADROOM_ANYLLM_PROVIDER": "llamacpp"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].anyllm_provider == "llamacpp"
|
|
|
|
def test_anyllm_provider_cli_flag_works(self, runner):
|
|
"""--anyllm-provider flag should still work."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy", "--backend", "anyllm", "--anyllm-provider", "groq"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].anyllm_provider == "groq"
|
|
|
|
|
|
class TestArgparseBackendValidation:
|
|
"""Test that the argparse path (python -m headroom.proxy.server) accepts litellm-* backends."""
|
|
|
|
def test_argparse_accepts_litellm_backend(self):
|
|
"""The argparse --backend should accept litellm-hosted_vllm (no choices restriction)."""
|
|
import argparse
|
|
|
|
# Recreate the parser matching server.py's main() argparse setup
|
|
# We just need to verify argparse doesn't reject litellm-* values
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--backend", default="anthropic")
|
|
args = parser.parse_args(["--backend", "litellm-hosted_vllm"])
|
|
assert args.backend == "litellm-hosted_vllm"
|