diff --git a/headroom/__init__.py b/headroom/__init__.py index 6af2a70e4..372142d61 100644 --- a/headroom/__init__.py +++ b/headroom/__init__.py @@ -142,7 +142,7 @@ from .transforms import ( TransformPipeline, ) -__version__ = "0.4.0" +__version__ = "0.4.1" __all__ = [ # Main client diff --git a/headroom/backends/anyllm.py b/headroom/backends/anyllm.py index 474d4c78e..592ed48bf 100644 --- a/headroom/backends/anyllm.py +++ b/headroom/backends/anyllm.py @@ -6,6 +6,7 @@ through a single interface. Auth and format translation handled automatically. from __future__ import annotations +import json import logging import uuid from collections.abc import AsyncIterator @@ -139,12 +140,18 @@ class AnyLLMBackend(Backend): if hasattr(message, "tool_calls") and message.tool_calls: for tc in message.tool_calls: + args = tc.function.arguments + if isinstance(args, str): + try: + args = json.loads(args) + except (json.JSONDecodeError, TypeError): + pass content.append( { "type": "tool_use", "id": tc.id, "name": tc.function.name, - "input": tc.function.arguments, + "input": args, } ) @@ -207,6 +214,10 @@ class AnyLLMBackend(Backend): kwargs["top_p"] = body["top_p"] if "stop_sequences" in body: kwargs["stop"] = body["stop_sequences"] + if "tools" in body: + kwargs["tools"] = body["tools"] + if "tool_choice" in body: + kwargs["tool_choice"] = body["tool_choice"] logger.debug(f"any-llm request: provider={self.provider}, model={original_model}") @@ -249,6 +260,14 @@ class AnyLLMBackend(Backend): kwargs["max_tokens"] = body["max_tokens"] if "temperature" in body: kwargs["temperature"] = body["temperature"] + if "top_p" in body: + kwargs["top_p"] = body["top_p"] + if "stop_sequences" in body: + kwargs["stop"] = body["stop_sequences"] + if "tools" in body: + kwargs["tools"] = body["tools"] + if "tool_choice" in body: + kwargs["tool_choice"] = body["tool_choice"] msg_id = f"msg_{uuid.uuid4().hex[:24]}" diff --git a/headroom/backends/litellm.py b/headroom/backends/litellm.py index 880712dca..8659bd04a 100644 --- a/headroom/backends/litellm.py +++ b/headroom/backends/litellm.py @@ -12,6 +12,7 @@ LiteLLM handles all the auth and format translation internally. from __future__ import annotations +import json import logging import uuid from collections.abc import AsyncIterator @@ -158,11 +159,29 @@ def _normalize_bedrock_profile_id(profile_id: str) -> str | None: _BEDROCK_MODEL_MAP: dict[str, str] = {} _VERTEX_MODEL_MAP = { + # Claude 4.6 (latest, no date suffix) + "claude-opus-4-6": "vertex_ai/claude-opus-4-6", + "claude-sonnet-4-6": "vertex_ai/claude-sonnet-4-6", + # Claude 4.5 + "claude-sonnet-4-5-20250929": "vertex_ai/claude-sonnet-4-5@20250929", + "claude-opus-4-5-20251101": "vertex_ai/claude-opus-4-5@20251101", + # Claude 4.1 + "claude-opus-4-1-20250805": "vertex_ai/claude-opus-4-1@20250805", + # Claude 4 + "claude-sonnet-4-20250514": "vertex_ai/claude-sonnet-4@20250514", + "claude-opus-4-20250514": "vertex_ai/claude-opus-4@20250514", + # Claude 3.7 + "claude-3-7-sonnet-20250219": "vertex_ai/claude-3-7-sonnet@20250219", + # Claude 3.5 "claude-3-5-sonnet-20241022": "vertex_ai/claude-3-5-sonnet-v2@20241022", "claude-3-5-sonnet-20240620": "vertex_ai/claude-3-5-sonnet@20240620", + "claude-3-5-haiku-20241022": "vertex_ai/claude-3-5-haiku@20241022", + # Claude 3 (haiku 3 deprecated, others retired) "claude-3-opus-20240229": "vertex_ai/claude-3-opus@20240229", "claude-3-sonnet-20240229": "vertex_ai/claude-3-sonnet@20240229", "claude-3-haiku-20240307": "vertex_ai/claude-3-haiku@20240307", + # Haiku 4.5 + "claude-haiku-4-5-20251001": "vertex_ai/claude-haiku-4-5@20251001", } @@ -223,6 +242,53 @@ def get_provider_config(provider: str) -> ProviderConfig: ) +def _convert_anthropic_tool(tool: dict[str, Any]) -> dict[str, Any]: + """Convert Anthropic tool format to OpenAI function format. + + Anthropic: {"name": "...", "description": "...", "input_schema": {...}} + OpenAI: {"type": "function", "function": {"name": "...", "description": "...", "parameters": {...}}} + """ + func: dict[str, Any] = {"name": tool.get("name", "")} + if "description" in tool: + func["description"] = tool["description"] + if "input_schema" in tool: + func["parameters"] = tool["input_schema"] + return {"type": "function", "function": func} + + +def _convert_tool_choice(choice: Any) -> Any: + """Convert Anthropic tool_choice to OpenAI format. + + Anthropic: {"type": "auto"}, {"type": "any"}, {"type": "tool", "name": "..."} + OpenAI: "auto", "required", {"type": "function", "function": {"name": "..."}} + """ + if isinstance(choice, str): + return choice + if isinstance(choice, dict): + choice_type = choice.get("type", "auto") + if choice_type == "auto": + return "auto" + if choice_type == "any": + return "required" + if choice_type == "tool": + return {"type": "function", "function": {"name": choice.get("name", "")}} + return "auto" + + +def _parse_tool_arguments(arguments: Any) -> Any: + """Parse tool call arguments from string to dict. + + LiteLLM/OpenAI returns arguments as a JSON string, + but Anthropic expects input as a parsed dict. + """ + if isinstance(arguments, str): + try: + return json.loads(arguments) + except (json.JSONDecodeError, TypeError): + return arguments + return arguments + + class LiteLLMBackend(Backend): """Backend using LiteLLM for multi-provider support. @@ -379,7 +445,7 @@ class LiteLLMBackend(Backend): "type": "tool_use", "id": tc.id, "name": tc.function.name, - "input": tc.function.arguments, + "input": _parse_tool_arguments(tc.function.arguments), } ) @@ -438,6 +504,12 @@ class LiteLLMBackend(Backend): if "stop_sequences" in body: kwargs["stop"] = body["stop_sequences"] + # Tools (convert Anthropic format to OpenAI format) + if "tools" in body: + kwargs["tools"] = [_convert_anthropic_tool(t) for t in body["tools"]] + if "tool_choice" in body: + kwargs["tool_choice"] = _convert_tool_choice(body["tool_choice"]) + # System prompt (Anthropic puts it in body, OpenAI in messages) if "system" in body: system = body["system"] @@ -517,6 +589,14 @@ class LiteLLMBackend(Backend): kwargs["max_tokens"] = body["max_tokens"] if "temperature" in body: kwargs["temperature"] = body["temperature"] + if "top_p" in body: + kwargs["top_p"] = body["top_p"] + if "stop_sequences" in body: + kwargs["stop"] = body["stop_sequences"] + if "tools" in body: + kwargs["tools"] = [_convert_anthropic_tool(t) for t in body["tools"]] + if "tool_choice" in body: + kwargs["tool_choice"] = _convert_tool_choice(body["tool_choice"]) if "system" in body: system = body["system"] if isinstance(system, str): diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index 101b65135..aae0c2d47 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -1,5 +1,7 @@ """Proxy server CLI commands.""" +import os + import click from .main import main @@ -87,6 +89,16 @@ from .main import main default="openai", help="Provider for any-llm backend: openai, mistral, groq, ollama, etc. (default: openai)", ) +@click.option( + "--openai-api-url", + default=None, + help="Custom OpenAI API URL for passthrough endpoints (env: OPENAI_TARGET_API_URL)", +) +@click.option( + "--gemini-api-url", + default=None, + help="Custom Gemini API URL for passthrough endpoints (env: GEMINI_TARGET_API_URL)", +) @click.option( "--region", default="us-west-2", @@ -127,6 +139,8 @@ def proxy( memory_top_k: int, backend: str, anyllm_provider: str, + openai_api_url: str | None, + gemini_api_url: str | None, region: str, bedrock_region: str | None, bedrock_profile: str | None, @@ -155,9 +169,18 @@ def proxy( click.echo(f"Details: {e}") raise SystemExit(1) from None + # Resolve API URL overrides: CLI flag > env var > None + effective_openai_api_url = openai_api_url or os.environ.get("OPENAI_TARGET_API_URL") + effective_gemini_api_url = gemini_api_url or os.environ.get("GEMINI_TARGET_API_URL") + + # Resolve anyllm provider: env var takes precedence over CLI default (matches argparse path) + effective_anyllm_provider = os.environ.get("HEADROOM_ANYLLM_PROVIDER") or anyllm_provider + config = ProxyConfig( host=host, port=port, + openai_api_url=effective_openai_api_url, + gemini_api_url=effective_gemini_api_url, optimize=not no_optimize, cache_enabled=not no_cache, rate_limit_enabled=not no_rate_limit, @@ -185,7 +208,7 @@ def proxy( backend=backend, bedrock_region=bedrock_region or region, bedrock_profile=bedrock_profile, - anyllm_provider=anyllm_provider, + anyllm_provider=effective_anyllm_provider, ) memory_status = "DISABLED" @@ -198,7 +221,7 @@ def proxy( if config.backend == "anyllm" or config.backend.startswith("anyllm-"): # any-llm backend - backend_status = f"{anyllm_provider.title()} via any-llm" + backend_status = f"{effective_anyllm_provider.title()} via any-llm" backend_section = """ Set credentials for your provider (e.g., OPENAI_API_KEY, MISTRAL_API_KEY) Providers: https://mozilla-ai.github.io/any-llm/providers/ diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 701a4e701..684e76fb8 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -7215,12 +7215,14 @@ if __name__ == "__main__": "--openai-api-url", help=f"Custom OpenAI API URL (default: {HeadroomProxy.OPENAI_API_URL})" ) - # Backend (anthropic direct, bedrock, openrouter, or anyllm) + # Backend (anthropic direct, bedrock, openrouter, anyllm, or litellm-) parser.add_argument( "--backend", - choices=["anthropic", "bedrock", "openrouter", "anyllm"], default="anthropic", - help="Backend for Anthropic API: 'anthropic' (direct), 'bedrock' (AWS), 'openrouter', or 'anyllm' (any-llm)", + help=( + "Backend: 'anthropic' (direct), 'bedrock' (AWS), 'openrouter', " + "'anyllm' (any-llm), or 'litellm-' (e.g., litellm-hosted_vllm, litellm-vertex)" + ), ) parser.add_argument( "--bedrock-region", diff --git a/pyproject.toml b/pyproject.toml index a9028f045..c7a1342fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "headroom-ai" -version = "0.4.0" +version = "0.4.1" description = "The Context Optimization Layer for LLM Applications - Cut costs by 50-90%" readme = "README.md" license = "Apache-2.0" diff --git a/tests/test_backend_bugs.py b/tests/test_backend_bugs.py new file mode 100644 index 000000000..c7cd4b969 --- /dev/null +++ b/tests/test_backend_bugs.py @@ -0,0 +1,292 @@ +"""Tests for backend bug fixes in LiteLLM and any-llm integrations. + +Tests tool forwarding, tool argument parsing, streaming param forwarding, +and Vertex AI model mapping. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +pytest.importorskip("litellm") + +from headroom.backends.litellm import ( + _VERTEX_MODEL_MAP, + LiteLLMBackend, + _convert_anthropic_tool, + _convert_tool_choice, + _parse_tool_arguments, +) + +# ============================================================================= +# Tool Format Conversion (Bug 1) +# ============================================================================= + + +class TestConvertAnthropicTool: + """Test Anthropic → OpenAI tool format conversion.""" + + def test_basic_tool_conversion(self): + anthropic_tool = { + "name": "get_weather", + "description": "Get the weather for a location", + "input_schema": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + } + result = _convert_anthropic_tool(anthropic_tool) + assert result == { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a location", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } + + def test_tool_without_description(self): + tool = {"name": "do_thing", "input_schema": {"type": "object"}} + result = _convert_anthropic_tool(tool) + assert result["function"]["name"] == "do_thing" + assert "description" not in result["function"] + assert result["function"]["parameters"] == {"type": "object"} + + def test_tool_without_input_schema(self): + tool = {"name": "simple_tool", "description": "No params"} + result = _convert_anthropic_tool(tool) + assert result["function"]["name"] == "simple_tool" + assert "parameters" not in result["function"] + + +class TestConvertToolChoice: + """Test Anthropic → OpenAI tool_choice conversion.""" + + def test_auto(self): + assert _convert_tool_choice({"type": "auto"}) == "auto" + + def test_any_to_required(self): + assert _convert_tool_choice({"type": "any"}) == "required" + + def test_specific_tool(self): + result = _convert_tool_choice({"type": "tool", "name": "get_weather"}) + assert result == {"type": "function", "function": {"name": "get_weather"}} + + def test_string_passthrough(self): + assert _convert_tool_choice("auto") == "auto" + assert _convert_tool_choice("none") == "none" + + +# ============================================================================= +# Tool Argument Parsing (Bug 2) +# ============================================================================= + + +class TestParseToolArguments: + """Test that tool arguments are parsed from JSON string to dict.""" + + def test_json_string_parsed(self): + result = _parse_tool_arguments('{"location": "Paris"}') + assert result == {"location": "Paris"} + + def test_dict_passthrough(self): + d = {"location": "Paris"} + result = _parse_tool_arguments(d) + assert result == d + + def test_invalid_json_returns_original(self): + result = _parse_tool_arguments("not json") + assert result == "not json" + + def test_empty_string(self): + result = _parse_tool_arguments("") + assert result == "" + + def test_none_passthrough(self): + result = _parse_tool_arguments(None) + assert result is None + + +# ============================================================================= +# LiteLLM send_message Tools Forwarding (Bug 1) +# ============================================================================= + + +class TestLiteLLMToolsForwarding: + """Test that tools are forwarded through LiteLLM send_message.""" + + @pytest.mark.asyncio + async def test_tools_forwarded_in_send_message(self): + """Tools should be converted and passed to litellm.acompletion.""" + mock_response = MagicMock() + mock_response.choices = [ + MagicMock( + message=MagicMock(content="Hello", tool_calls=None), + finish_reason="stop", + ) + ] + mock_response.usage = MagicMock(prompt_tokens=10, completion_tokens=5) + + with ( + patch("headroom.backends.litellm.acompletion", new_callable=AsyncMock) as mock_acomp, + patch("headroom.backends.litellm._fetch_bedrock_inference_profiles", return_value={}), + ): + mock_acomp.return_value = mock_response + + backend = LiteLLMBackend(provider="openrouter") + body = { + "model": "claude-3-5-sonnet-20241022", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 100, + "tools": [ + { + "name": "get_weather", + "description": "Get weather", + "input_schema": {"type": "object", "properties": {}}, + } + ], + "tool_choice": {"type": "auto"}, + } + + await backend.send_message(body, {}) + + call_kwargs = mock_acomp.call_args[1] + assert "tools" in call_kwargs + assert call_kwargs["tools"][0]["type"] == "function" + assert call_kwargs["tools"][0]["function"]["name"] == "get_weather" + assert call_kwargs["tool_choice"] == "auto" + + @pytest.mark.asyncio + async def test_tool_arguments_parsed_in_response(self): + """Tool call arguments should be parsed from JSON string to dict.""" + mock_tc = MagicMock() + mock_tc.id = "call_123" + mock_tc.function.name = "get_weather" + mock_tc.function.arguments = '{"location": "Paris"}' + + mock_response = MagicMock() + mock_response.choices = [ + MagicMock( + message=MagicMock(content=None, tool_calls=[mock_tc]), + finish_reason="tool_calls", + ) + ] + mock_response.usage = MagicMock(prompt_tokens=10, completion_tokens=5) + + with ( + patch("headroom.backends.litellm.acompletion", new_callable=AsyncMock) as mock_acomp, + patch("headroom.backends.litellm._fetch_bedrock_inference_profiles", return_value={}), + ): + mock_acomp.return_value = mock_response + + backend = LiteLLMBackend(provider="openrouter") + result = await backend.send_message( + {"model": "test", "messages": [{"role": "user", "content": "hi"}]}, + {}, + ) + + tool_block = result.body["content"][0] + assert tool_block["type"] == "tool_use" + assert tool_block["input"] == {"location": "Paris"} + assert isinstance(tool_block["input"], dict) + + +# ============================================================================= +# Streaming Params (Bugs 3-4) +# ============================================================================= + + +class TestLiteLLMStreamingParams: + """Test that streaming forwards all params.""" + + @pytest.mark.asyncio + async def test_streaming_forwards_all_params(self): + """stream_message should forward top_p, stop, and tools.""" + + # Create an async iterator for the mock streaming response + async def mock_stream(): + chunk = MagicMock() + chunk.choices = [MagicMock(delta=MagicMock(content="Hi"))] + yield chunk + + with ( + patch("headroom.backends.litellm.acompletion", new_callable=AsyncMock) as mock_acomp, + patch("headroom.backends.litellm._fetch_bedrock_inference_profiles", return_value={}), + ): + mock_acomp.return_value = mock_stream() + + backend = LiteLLMBackend(provider="openrouter") + body = { + "model": "test", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 100, + "temperature": 0.7, + "top_p": 0.9, + "stop_sequences": ["\n"], + "tools": [ + { + "name": "test_tool", + "description": "A test", + "input_schema": {"type": "object"}, + } + ], + } + + events = [] + async for event in backend.stream_message(body, {}): + events.append(event) + + call_kwargs = mock_acomp.call_args[1] + assert call_kwargs["top_p"] == 0.9 + assert call_kwargs["stop"] == ["\n"] + assert "tools" in call_kwargs + assert call_kwargs["tools"][0]["function"]["name"] == "test_tool" + + +# ============================================================================= +# Vertex AI Model Map (Bug 6) +# ============================================================================= + + +class TestVertexModelMap: + """Test that Vertex AI model map includes all current models. + + Model IDs sourced from: https://platform.claude.com/docs/en/build-with-claude/claude-on-vertex-ai + """ + + def test_claude_46_models(self): + assert _VERTEX_MODEL_MAP["claude-opus-4-6"] == "vertex_ai/claude-opus-4-6" + assert _VERTEX_MODEL_MAP["claude-sonnet-4-6"] == "vertex_ai/claude-sonnet-4-6" + + def test_claude_45_models(self): + assert ( + _VERTEX_MODEL_MAP["claude-sonnet-4-5-20250929"] + == "vertex_ai/claude-sonnet-4-5@20250929" + ) + assert _VERTEX_MODEL_MAP["claude-opus-4-5-20251101"] == "vertex_ai/claude-opus-4-5@20251101" + + def test_claude_4_models(self): + assert _VERTEX_MODEL_MAP["claude-sonnet-4-20250514"] == "vertex_ai/claude-sonnet-4@20250514" + assert _VERTEX_MODEL_MAP["claude-opus-4-20250514"] == "vertex_ai/claude-opus-4@20250514" + + def test_claude_35_models(self): + assert ( + _VERTEX_MODEL_MAP["claude-3-5-sonnet-20241022"] + == "vertex_ai/claude-3-5-sonnet-v2@20241022" + ) + assert ( + _VERTEX_MODEL_MAP["claude-3-5-haiku-20241022"] == "vertex_ai/claude-3-5-haiku@20241022" + ) + + def test_claude_haiku_45(self): + assert ( + _VERTEX_MODEL_MAP["claude-haiku-4-5-20251001"] == "vertex_ai/claude-haiku-4-5@20251001" + ) + + def test_claude_3_legacy(self): + assert "claude-3-haiku-20240307" in _VERTEX_MODEL_MAP diff --git a/tests/test_cli_proxy_env.py b/tests/test_cli_proxy_env.py new file mode 100644 index 000000000..fc92d977d --- /dev/null +++ b/tests/test_cli_proxy_env.py @@ -0,0 +1,259 @@ +"""Tests for CLI proxy env variable handling and backend validation. + +Verifies that: +1. OPENAI_TARGET_API_URL and GEMINI_TARGET_API_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_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): + 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): + 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_openai_api_url_cli_flag(self, runner): + """--openai-api-url CLI flag should take precedence.""" + captured_config = {} + + def mock_run_server(config): + 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_cli_flag_overrides_env_var(self, runner): + """CLI flag should take precedence over env var.""" + captured_config = {} + + def mock_run_server(config): + 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): + 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): + 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" + + +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): + 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): + 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): + 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): + 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): + 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"