mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Fix proxy backend bugs: env vars, tool forwarding, and provider support
- CLI now reads OPENAI_TARGET_API_URL, GEMINI_TARGET_API_URL, and HEADROOM_ANYLLM_PROVIDER environment variables - Add --openai-api-url and --gemini-api-url CLI flags - Remove --backend choices restriction so litellm-* backends work - Forward tools/tool_choice through LiteLLM and any-llm backends - Parse tool call arguments from JSON string to dict (Anthropic format) - Forward top_p, stop_sequences, tools in streaming paths - Update Vertex AI model map with Claude 3 through 4.6 (from official docs) - Bump version to 0.4.1
This commit is contained in:
chopratejas
parent
8f438b0674
commit
4bea17ba8a
8 changed files with 684 additions and 9 deletions
|
|
@ -142,7 +142,7 @@ from .transforms import (
|
|||
TransformPipeline,
|
||||
)
|
||||
|
||||
__version__ = "0.4.0"
|
||||
__version__ = "0.4.1"
|
||||
|
||||
__all__ = [
|
||||
# Main client
|
||||
|
|
|
|||
|
|
@ -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]}"
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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/
|
||||
|
|
|
|||
|
|
@ -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-<provider>)
|
||||
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-<provider>' (e.g., litellm-hosted_vllm, litellm-vertex)"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--bedrock-region",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
292
tests/test_backend_bugs.py
Normal file
292
tests/test_backend_bugs.py
Normal file
|
|
@ -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
|
||||
259
tests/test_cli_proxy_env.py
Normal file
259
tests/test_cli_proxy_env.py
Normal file
|
|
@ -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"
|
||||
Loading…
Add table
Add a link
Reference in a new issue