mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
test: expand provider pipeline coverage
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
JerrettDavis
parent
364a07228f
commit
ce8a2f3cf8
4 changed files with 701 additions and 0 deletions
|
|
@ -33,6 +33,22 @@ class MutatingExtension:
|
|||
return event
|
||||
|
||||
|
||||
class ReplacingExtension:
|
||||
def on_pipeline_event(self, event):
|
||||
return type(event)(
|
||||
stage=event.stage,
|
||||
operation=event.operation,
|
||||
model=event.model,
|
||||
messages=[{"role": "user", "content": "replaced"}],
|
||||
metadata={"replaced": True},
|
||||
)
|
||||
|
||||
|
||||
class RaisingExtension:
|
||||
def on_pipeline_event(self, event):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
|
||||
class RecordingHooks(CompressionHooks):
|
||||
def __init__(self) -> None:
|
||||
self.stages: list[PipelineStage] = []
|
||||
|
|
@ -130,6 +146,77 @@ def test_pipeline_extension_manager_uses_canonical_stage_contract():
|
|||
assert event.messages == [{"role": "user", "content": "mutated"}]
|
||||
|
||||
|
||||
def test_pipeline_extension_manager_replaces_events_and_ignores_failures(caplog):
|
||||
recorder = RecordingExtension()
|
||||
manager = PipelineExtensionManager(
|
||||
extensions=[recorder, RaisingExtension(), ReplacingExtension(), object()],
|
||||
discover=False,
|
||||
)
|
||||
|
||||
with caplog.at_level("WARNING", logger="headroom.pipeline"):
|
||||
event = manager.emit(
|
||||
PipelineStage.PRE_SEND,
|
||||
operation="test",
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
)
|
||||
|
||||
assert manager.enabled is True
|
||||
assert recorder.stages == [PipelineStage.PRE_SEND]
|
||||
assert event.messages == [{"role": "user", "content": "replaced"}]
|
||||
assert event.metadata == {"replaced": True}
|
||||
|
||||
|
||||
def test_discover_pipeline_extensions_handles_load_and_init_failures(monkeypatch):
|
||||
pipeline_module = importlib.import_module("headroom.pipeline")
|
||||
|
||||
class Entry:
|
||||
def __init__(self, name, loader):
|
||||
self.name = name
|
||||
self._loader = loader
|
||||
|
||||
def load(self):
|
||||
return self._loader()
|
||||
|
||||
class ExtensionClass:
|
||||
def on_pipeline_event(self, event):
|
||||
return event
|
||||
|
||||
class FailingInit:
|
||||
def __init__(self):
|
||||
raise RuntimeError("init failed")
|
||||
|
||||
entries = [
|
||||
Entry("instance", lambda: RecordingExtension()),
|
||||
Entry("class", lambda: ExtensionClass),
|
||||
Entry("load-fail", lambda: (_ for _ in ()).throw(RuntimeError("load failed"))),
|
||||
Entry("init-fail", lambda: FailingInit),
|
||||
]
|
||||
|
||||
monkeypatch.setattr(
|
||||
pipeline_module.importlib.metadata,
|
||||
"entry_points",
|
||||
lambda group: entries if group == pipeline_module.ENTRY_POINT_GROUP else [],
|
||||
)
|
||||
|
||||
discovered = pipeline_module.discover_pipeline_extensions()
|
||||
|
||||
assert len(discovered) == 2
|
||||
assert hasattr(discovered[0], "on_pipeline_event")
|
||||
assert hasattr(discovered[1], "on_pipeline_event")
|
||||
|
||||
|
||||
def test_discover_pipeline_extensions_returns_empty_when_entrypoint_lookup_fails(monkeypatch):
|
||||
pipeline_module = importlib.import_module("headroom.pipeline")
|
||||
monkeypatch.setattr(
|
||||
pipeline_module.importlib.metadata,
|
||||
"entry_points",
|
||||
lambda group: (_ for _ in ()).throw(RuntimeError("lookup failed")),
|
||||
)
|
||||
|
||||
assert pipeline_module.discover_pipeline_extensions() == []
|
||||
|
||||
|
||||
def test_compress_emits_canonical_pipeline_events(monkeypatch):
|
||||
hooks = RecordingHooks()
|
||||
compress_module = importlib.import_module("headroom.compress")
|
||||
|
|
|
|||
92
tests/test_provider_copilot_wrap.py
Normal file
92
tests/test_provider_copilot_wrap.py
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import json
|
||||
import urllib.error
|
||||
from unittest.mock import patch
|
||||
|
||||
from headroom.providers.copilot.wrap import (
|
||||
build_launch_env,
|
||||
detect_running_proxy_backend,
|
||||
model_configured,
|
||||
provider_key_source,
|
||||
query_proxy_config,
|
||||
resolve_provider_type,
|
||||
validate_configuration,
|
||||
)
|
||||
|
||||
|
||||
def test_query_proxy_config_handles_success_and_invalid_payload() -> None:
|
||||
payload = io.BytesIO(json.dumps({"config": {"backend": "anyllm"}}).encode("utf-8"))
|
||||
payload_missing = io.BytesIO(json.dumps({"status": "ok"}).encode("utf-8"))
|
||||
|
||||
with patch("urllib.request.urlopen", return_value=payload):
|
||||
assert query_proxy_config(8787) == {"backend": "anyllm"}
|
||||
with patch("urllib.request.urlopen", return_value=payload_missing):
|
||||
assert query_proxy_config(8787) is None
|
||||
with patch("urllib.request.urlopen", side_effect=urllib.error.URLError("down")):
|
||||
assert query_proxy_config(8787) is None
|
||||
|
||||
|
||||
def test_detect_running_proxy_backend_requires_string_backend(monkeypatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
"headroom.providers.copilot.wrap.query_proxy_config",
|
||||
lambda port: {"backend": 123} if port == 8787 else None,
|
||||
)
|
||||
|
||||
assert detect_running_proxy_backend(8787) is None
|
||||
assert detect_running_proxy_backend(9999) is None
|
||||
|
||||
|
||||
def test_resolve_provider_type_prefers_explicit_and_env() -> None:
|
||||
assert resolve_provider_type("anthropic", "openai") == "openai"
|
||||
assert resolve_provider_type(None, "auto", {"HEADROOM_BACKEND": "anthropic"}) == "anthropic"
|
||||
assert resolve_provider_type(None, "auto", {"HEADROOM_BACKEND": "anyllm"}) == "openai"
|
||||
|
||||
|
||||
def test_validate_configuration_accepts_supported_combinations() -> None:
|
||||
validate_configuration(provider_type="openai", wire_api="responses", backend=None)
|
||||
validate_configuration(provider_type="openai", wire_api="completions", backend="anyllm")
|
||||
|
||||
|
||||
def test_provider_key_source_and_build_launch_env_cover_anthropic_and_openai() -> None:
|
||||
assert provider_key_source("anthropic") == "ANTHROPIC_API_KEY"
|
||||
assert provider_key_source("openai") == "OPENAI_API_KEY"
|
||||
|
||||
anthropic_env, anthropic_lines = build_launch_env(
|
||||
port=8787,
|
||||
provider_type="anthropic",
|
||||
wire_api="responses",
|
||||
environ={
|
||||
"ANTHROPIC_API_KEY": "sk-ant-test",
|
||||
"COPILOT_PROVIDER_WIRE_API": "stale",
|
||||
},
|
||||
)
|
||||
openai_env, openai_lines = build_launch_env(
|
||||
port=8787,
|
||||
provider_type="openai",
|
||||
wire_api=None,
|
||||
environ={"OPENAI_API_KEY": "sk-proj-test"},
|
||||
)
|
||||
|
||||
assert anthropic_env["COPILOT_PROVIDER_TYPE"] == "anthropic"
|
||||
assert anthropic_env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787"
|
||||
assert anthropic_env["COPILOT_PROVIDER_API_KEY"] == "sk-ant-test"
|
||||
assert "COPILOT_PROVIDER_WIRE_API" not in anthropic_env
|
||||
assert anthropic_lines == [
|
||||
"COPILOT_PROVIDER_TYPE=anthropic",
|
||||
"COPILOT_PROVIDER_BASE_URL=http://127.0.0.1:8787",
|
||||
]
|
||||
|
||||
assert openai_env["COPILOT_PROVIDER_TYPE"] == "openai"
|
||||
assert openai_env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/v1"
|
||||
assert openai_env["COPILOT_PROVIDER_WIRE_API"] == "completions"
|
||||
assert openai_env["COPILOT_PROVIDER_API_KEY"] == "sk-proj-test"
|
||||
assert openai_lines[-1] == "COPILOT_PROVIDER_WIRE_API=completions"
|
||||
|
||||
|
||||
def test_model_configured_detects_env_and_cli_variants() -> None:
|
||||
assert model_configured((), {"COPILOT_MODEL": "gpt-4o"}) is True
|
||||
assert model_configured(("--model", "gpt-4o"), {}) is True
|
||||
assert model_configured(("--model=gpt-4o",), {}) is True
|
||||
assert model_configured(("--other", "value"), {}) is False
|
||||
283
tests/test_provider_proxy_routes.py
Normal file
283
tests/test_provider_proxy_routes.py
Normal file
|
|
@ -0,0 +1,283 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from fastapi.responses import JSONResponse
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from headroom.proxy.server import HeadroomProxy, ProxyConfig, create_app
|
||||
|
||||
|
||||
def _app() -> Any:
|
||||
return create_app(
|
||||
ProxyConfig(
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
anthropic_api_url="https://api.anthropic.test",
|
||||
openai_api_url="https://api.openai.test",
|
||||
gemini_api_url="https://api.gemini.test",
|
||||
cloudcode_api_url="https://cloudcode.test",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_provider_passthrough_routes_forward_expected_targets(monkeypatch) -> None:
|
||||
calls: list[tuple[str, str, str, str]] = []
|
||||
|
||||
async def fake_passthrough(self, request, base_url, sub_path="", provider_name=""): # type: ignore[no-untyped-def]
|
||||
calls.append((request.method, request.url.path, base_url, provider_name))
|
||||
return JSONResponse(
|
||||
{
|
||||
"method": request.method,
|
||||
"path": request.url.path,
|
||||
"base_url": base_url,
|
||||
"sub_path": sub_path,
|
||||
"provider": provider_name,
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setattr(HeadroomProxy, "handle_passthrough", fake_passthrough)
|
||||
|
||||
with TestClient(_app()) as client:
|
||||
assert client.post("/v1/messages/count_tokens").json()["base_url"] == (
|
||||
"https://api.anthropic.test"
|
||||
)
|
||||
assert client.get("/v1/models", headers={"x-goog-api-key": "test"}).json()["base_url"] == (
|
||||
"https://api.openai.test"
|
||||
)
|
||||
assert client.get("/v1/models/demo").json()["sub_path"] == "models"
|
||||
assert (
|
||||
client.get(
|
||||
"/azure/models",
|
||||
headers={
|
||||
"api-key": "azure-key",
|
||||
"x-headroom-base-url": "https://azure.example/openai/",
|
||||
},
|
||||
).json()["base_url"]
|
||||
== "https://azure.example/openai"
|
||||
)
|
||||
assert client.post("/v1/embeddings").json()["provider"] == "openai"
|
||||
assert client.post("/v1/moderations").json()["sub_path"] == "moderations"
|
||||
assert client.post("/v1/images/generations").json()["sub_path"] == "images/generations"
|
||||
assert client.post("/v1/audio/transcriptions").json()["sub_path"] == "audio/transcriptions"
|
||||
assert client.post("/v1/audio/speech").json()["sub_path"] == "audio/speech"
|
||||
assert client.get("/v1beta/models").json()["provider"] == "gemini"
|
||||
assert client.get("/v1beta/models/demo").json()["sub_path"] == "models"
|
||||
assert client.post("/v1beta/models/demo:embedContent").json()["sub_path"] == "embedContent"
|
||||
assert client.post("/v1beta/cachedContents").json()["sub_path"] == "cachedContents"
|
||||
assert client.get("/v1beta/cachedContents").json()["sub_path"] == "cachedContents"
|
||||
assert client.get("/v1beta/cachedContents/cache-1").json()["sub_path"] == "cachedContents"
|
||||
assert client.delete("/v1beta/cachedContents/cache-1").json()["sub_path"] == (
|
||||
"cachedContents"
|
||||
)
|
||||
assert (
|
||||
client.get(
|
||||
"/unhandled/path",
|
||||
headers={"x-headroom-base-url": "https://custom.example/base/"},
|
||||
).json()["base_url"]
|
||||
== "https://custom.example/base"
|
||||
)
|
||||
assert client.get("/another/path", headers={"x-goog-api-key": "test"}).json()[
|
||||
"base_url"
|
||||
] == ("https://api.gemini.test")
|
||||
|
||||
assert len(calls) >= 16
|
||||
|
||||
|
||||
def test_proxy_route_helpers_prefer_legacy_targets_and_gemini_passthrough() -> None:
|
||||
proxy_routes = importlib.import_module("headroom.providers.proxy_routes")
|
||||
proxy = type(
|
||||
"Proxy",
|
||||
(),
|
||||
{
|
||||
"ANTHROPIC_API_URL": "https://legacy.anthropic.test",
|
||||
"OPENAI_API_URL": "https://legacy.openai.test",
|
||||
"GEMINI_API_URL": "https://legacy.gemini.test",
|
||||
"provider_runtime": type(
|
||||
"Runtime",
|
||||
(),
|
||||
{
|
||||
"api_target": staticmethod(lambda provider: f"https://runtime.{provider}.test"),
|
||||
"model_metadata_provider": staticmethod(lambda headers: "anthropic"),
|
||||
},
|
||||
)(),
|
||||
},
|
||||
)()
|
||||
|
||||
assert proxy_routes._api_target(proxy, "anthropic") == "https://legacy.anthropic.test"
|
||||
assert proxy_routes._select_passthrough_base_url(proxy, {"x-goog-api-key": "test"}) == (
|
||||
"https://legacy.gemini.test"
|
||||
)
|
||||
assert (
|
||||
proxy_routes._select_passthrough_base_url(
|
||||
proxy, {"api-key": "azure", "x-headroom-base-url": "https://azure.example/base/"}
|
||||
)
|
||||
== "https://azure.example/base"
|
||||
)
|
||||
assert proxy_routes._select_passthrough_base_url(proxy, {}) == "https://legacy.anthropic.test"
|
||||
|
||||
|
||||
def test_provider_specific_routes_delegate_to_expected_proxy_handlers(monkeypatch) -> None:
|
||||
delegated: list[tuple[str, str, tuple[str, ...]]] = []
|
||||
|
||||
def install(name: str) -> None:
|
||||
async def fake(self, request, *args): # type: ignore[no-untyped-def]
|
||||
delegated.append((name, request.url.path, tuple(str(arg) for arg in args)))
|
||||
return JSONResponse({"handler": name, "path": request.url.path, "args": list(args)})
|
||||
|
||||
monkeypatch.setattr(HeadroomProxy, name, fake)
|
||||
|
||||
for handler_name in (
|
||||
"handle_anthropic_messages",
|
||||
"handle_anthropic_batch_create",
|
||||
"handle_anthropic_batch_passthrough",
|
||||
"handle_anthropic_batch_results",
|
||||
"handle_openai_chat",
|
||||
"handle_openai_responses",
|
||||
"handle_batch_create",
|
||||
"handle_batch_list",
|
||||
"handle_batch_get",
|
||||
"handle_batch_cancel",
|
||||
"handle_gemini_generate_content",
|
||||
"handle_gemini_stream_generate_content",
|
||||
"handle_gemini_count_tokens",
|
||||
"handle_google_cloudcode_stream",
|
||||
"handle_databricks_invocations",
|
||||
"handle_google_batch_create",
|
||||
"handle_google_batch_results",
|
||||
"handle_google_batch_passthrough",
|
||||
):
|
||||
install(handler_name)
|
||||
|
||||
with TestClient(_app()) as client:
|
||||
assert client.post("/v1/messages").json()["handler"] == "handle_anthropic_messages"
|
||||
assert (
|
||||
client.post("/v1/messages/batches").json()["handler"] == "handle_anthropic_batch_create"
|
||||
)
|
||||
assert client.get("/v1/messages/batches").json()["handler"] == (
|
||||
"handle_anthropic_batch_passthrough"
|
||||
)
|
||||
assert client.get("/v1/messages/batches/b1").json()["args"] == ["b1"]
|
||||
assert client.get("/v1/messages/batches/b1/results").json()["handler"] == (
|
||||
"handle_anthropic_batch_results"
|
||||
)
|
||||
assert client.post("/v1/messages/batches/b1/cancel").json()["handler"] == (
|
||||
"handle_anthropic_batch_passthrough"
|
||||
)
|
||||
assert client.post("/v1/chat/completions").json()["handler"] == "handle_openai_chat"
|
||||
assert client.post("/v1/responses").json()["handler"] == "handle_openai_responses"
|
||||
assert client.post("/v1/codex/responses").json()["handler"] == "handle_openai_responses"
|
||||
assert client.post("/backend-api/responses").json()["handler"] == "handle_openai_responses"
|
||||
assert client.post("/backend-api/codex/responses").json()["handler"] == (
|
||||
"handle_openai_responses"
|
||||
)
|
||||
assert client.post("/v1/batches").json()["handler"] == "handle_batch_create"
|
||||
assert client.get("/v1/batches").json()["handler"] == "handle_batch_list"
|
||||
assert client.get("/v1/batches/b1").json()["handler"] == "handle_batch_get"
|
||||
assert client.post("/v1/batches/b1/cancel").json()["handler"] == "handle_batch_cancel"
|
||||
assert client.post("/v1beta/models/demo:generateContent").json()["handler"] == (
|
||||
"handle_gemini_generate_content"
|
||||
)
|
||||
assert client.post("/v1beta/models/demo:streamGenerateContent").json()["handler"] == (
|
||||
"handle_gemini_stream_generate_content"
|
||||
)
|
||||
assert client.post("/v1beta/models/demo:countTokens").json()["handler"] == (
|
||||
"handle_gemini_count_tokens"
|
||||
)
|
||||
assert client.post("/v1internal:streamGenerateContent").json()["handler"] == (
|
||||
"handle_google_cloudcode_stream"
|
||||
)
|
||||
assert client.post("/v1/v1internal:streamGenerateContent").json()["handler"] == (
|
||||
"handle_google_cloudcode_stream"
|
||||
)
|
||||
assert client.post("/serving-endpoints/demo/invocations").json()["handler"] == (
|
||||
"handle_databricks_invocations"
|
||||
)
|
||||
assert client.post("/v1beta/models/demo:batchGenerateContent").json()["handler"] == (
|
||||
"handle_google_batch_create"
|
||||
)
|
||||
assert client.get("/v1beta/batches/b1").json()["handler"] == "handle_google_batch_results"
|
||||
assert client.post("/v1beta/batches/b1:cancel").json()["handler"] == (
|
||||
"handle_google_batch_passthrough"
|
||||
)
|
||||
assert client.delete("/v1beta/batches/b1").json()["handler"] == (
|
||||
"handle_google_batch_passthrough"
|
||||
)
|
||||
|
||||
assert len(delegated) >= 24
|
||||
|
||||
|
||||
def test_openai_response_websocket_aliases_delegate_to_openai_ws_handler(monkeypatch) -> None:
|
||||
seen_paths: list[str] = []
|
||||
|
||||
async def fake_ws(self, websocket): # type: ignore[no-untyped-def]
|
||||
seen_paths.append(websocket.url.path)
|
||||
await websocket.accept()
|
||||
await websocket.send_json({"path": websocket.url.path})
|
||||
await websocket.close()
|
||||
|
||||
monkeypatch.setattr(HeadroomProxy, "handle_openai_responses_ws", fake_ws)
|
||||
|
||||
with TestClient(_app()) as client:
|
||||
for path in (
|
||||
"/v1/responses",
|
||||
"/v1/codex/responses",
|
||||
"/backend-api/responses",
|
||||
"/backend-api/codex/responses",
|
||||
):
|
||||
with client.websocket_connect(path) as websocket:
|
||||
assert websocket.receive_json() == {"path": path}
|
||||
|
||||
assert seen_paths == [
|
||||
"/v1/responses",
|
||||
"/v1/codex/responses",
|
||||
"/backend-api/responses",
|
||||
"/backend-api/codex/responses",
|
||||
]
|
||||
|
||||
|
||||
def test_openai_response_subpath_passthrough_returns_502_on_http_failure() -> None:
|
||||
class FailingAsyncClient:
|
||||
async def request(self, method, url, **kwargs): # type: ignore[no-untyped-def]
|
||||
raise RuntimeError(f"boom: {method} {url}")
|
||||
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
with TestClient(_app()) as client:
|
||||
client.app.state.proxy.http_client = FailingAsyncClient()
|
||||
response = client.post("/v1/responses/compact?trace=1", json={"model": "gpt-4o"})
|
||||
|
||||
assert response.status_code == 502
|
||||
assert "boom: POST https://api.openai.test/v1/responses/compact?trace=1" in response.text
|
||||
|
||||
|
||||
def test_openai_response_subpath_passthrough_uses_openai_target() -> None:
|
||||
class FakeAsyncClient:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, str, dict[str, str]]] = []
|
||||
|
||||
async def request(self, method, url, **kwargs): # type: ignore[no-untyped-def]
|
||||
self.calls.append((method, url, dict(kwargs.get("headers", {}))))
|
||||
return httpx.Response(200, json={"url": url})
|
||||
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
with TestClient(_app()) as client:
|
||||
fake = FakeAsyncClient()
|
||||
client.app.state.proxy.http_client = fake
|
||||
response = client.delete(
|
||||
"/v1/responses/items/resp_123?trace=7",
|
||||
headers={"Authorization": "Bearer sk-proj-test"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert len(fake.calls) == 1
|
||||
method, url, headers = fake.calls[0]
|
||||
assert method == "DELETE"
|
||||
assert url == "https://api.openai.test/v1/responses/items/resp_123?trace=7"
|
||||
assert headers["authorization"] == "Bearer sk-proj-test"
|
||||
239
tests/test_provider_registry_extended.py
Normal file
239
tests/test_provider_registry_extended.py
Normal file
|
|
@ -0,0 +1,239 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.providers.registry import (
|
||||
ProviderApiTargets,
|
||||
ProxyProviderRuntime,
|
||||
call_client_transport,
|
||||
create_proxy_backend,
|
||||
format_backend_status,
|
||||
)
|
||||
|
||||
|
||||
class DummyStorage:
|
||||
def __init__(self) -> None:
|
||||
self.saved: list[Any] = []
|
||||
|
||||
def save(self, metrics: Any) -> None:
|
||||
self.saved.append(metrics)
|
||||
|
||||
|
||||
class DummyClient:
|
||||
def __init__(self) -> None:
|
||||
self._storage = DummyStorage()
|
||||
self._wrapped_stream: tuple[Any, Any] | None = None
|
||||
self._original = SimpleNamespace(
|
||||
chat=SimpleNamespace(completions=SimpleNamespace(create=self._openai_create)),
|
||||
messages=SimpleNamespace(create=self._anthropic_create, stream=self._anthropic_stream),
|
||||
)
|
||||
self.openai_calls: list[dict[str, Any]] = []
|
||||
self.anthropic_calls: list[dict[str, Any]] = []
|
||||
|
||||
def _openai_create(self, **kwargs: Any) -> Any:
|
||||
self.openai_calls.append(kwargs)
|
||||
if kwargs["stream"]:
|
||||
return iter(["chunk-1", "chunk-2"])
|
||||
return SimpleNamespace(
|
||||
usage=SimpleNamespace(
|
||||
completion_tokens=7,
|
||||
prompt_tokens_details=SimpleNamespace(cached_tokens=3),
|
||||
)
|
||||
)
|
||||
|
||||
def _anthropic_create(self, **kwargs: Any) -> Any:
|
||||
self.anthropic_calls.append(kwargs)
|
||||
return SimpleNamespace(
|
||||
usage=SimpleNamespace(
|
||||
output_tokens=5,
|
||||
cache_read_input_tokens=2,
|
||||
)
|
||||
)
|
||||
|
||||
def _anthropic_stream(self, **kwargs: Any) -> Any:
|
||||
self.anthropic_calls.append(kwargs)
|
||||
return "anthropic-stream"
|
||||
|
||||
def _wrap_stream(self, stream: Any, metrics: Any) -> Any:
|
||||
self._wrapped_stream = (stream, metrics)
|
||||
return ("wrapped", stream)
|
||||
|
||||
|
||||
def test_proxy_provider_runtime_selects_targets_and_providers() -> None:
|
||||
runtime = ProxyProviderRuntime(
|
||||
api_targets=ProviderApiTargets(
|
||||
anthropic="https://anthropic.example",
|
||||
openai="https://openai.example",
|
||||
gemini="https://gemini.example",
|
||||
cloudcode="https://cloudcode.example",
|
||||
),
|
||||
pipeline_providers={
|
||||
"anthropic": SimpleNamespace(name="anthropic"),
|
||||
"openai": SimpleNamespace(name="openai"),
|
||||
},
|
||||
)
|
||||
|
||||
assert runtime.api_target("anthropic") == "https://anthropic.example"
|
||||
assert runtime.pipeline_provider("openai").name == "openai"
|
||||
assert runtime.model_metadata_provider({"Authorization": "Bearer sk-ant-api03-test"}) == (
|
||||
"anthropic"
|
||||
)
|
||||
assert runtime.select_passthrough_base_url({"x-goog-api-key": "test"}) == (
|
||||
"https://gemini.example"
|
||||
)
|
||||
assert (
|
||||
runtime.select_passthrough_base_url(
|
||||
{"api-key": "azure-key", "x-headroom-base-url": "https://azure.example/openai/"}
|
||||
)
|
||||
== "https://azure.example/openai"
|
||||
)
|
||||
assert runtime.select_passthrough_base_url({}) == "https://openai.example"
|
||||
|
||||
|
||||
def test_create_proxy_backend_uses_injected_backend_types() -> None:
|
||||
logger = logging.getLogger("test")
|
||||
|
||||
anyllm = create_proxy_backend(
|
||||
backend="anyllm",
|
||||
anyllm_provider="groq",
|
||||
bedrock_region=None,
|
||||
logger=logger,
|
||||
anyllm_backend_cls=lambda provider: {"kind": "anyllm", "provider": provider},
|
||||
)
|
||||
litellm = create_proxy_backend(
|
||||
backend="bedrock",
|
||||
anyllm_provider="ignored",
|
||||
bedrock_region="us-east-1",
|
||||
logger=logger,
|
||||
litellm_backend_cls=lambda provider, region: {
|
||||
"kind": "litellm",
|
||||
"provider": provider,
|
||||
"region": region,
|
||||
},
|
||||
)
|
||||
|
||||
assert anyllm == {"kind": "anyllm", "provider": "groq"}
|
||||
assert litellm == {"kind": "litellm", "provider": "bedrock", "region": "us-east-1"}
|
||||
|
||||
|
||||
def test_create_proxy_backend_handles_missing_or_direct_backends(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
logger = logging.getLogger("test")
|
||||
|
||||
direct = create_proxy_backend(
|
||||
backend="anthropic",
|
||||
anyllm_provider="ignored",
|
||||
bedrock_region=None,
|
||||
logger=logger,
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
missing = create_proxy_backend(
|
||||
backend="anyllm",
|
||||
anyllm_provider="groq",
|
||||
bedrock_region=None,
|
||||
logger=logger,
|
||||
anyllm_backend_cls=lambda provider: (_ for _ in ()).throw(ImportError("missing")),
|
||||
)
|
||||
|
||||
assert direct is None
|
||||
assert missing is None
|
||||
assert "any-llm backend not available" in caplog.text
|
||||
|
||||
|
||||
def test_format_backend_status_uses_litellm_provider_metadata(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
"headroom.backends.litellm.get_provider_config",
|
||||
lambda provider: SimpleNamespace(
|
||||
display_name=provider.upper(),
|
||||
uses_region=(provider == "bedrock"),
|
||||
),
|
||||
)
|
||||
|
||||
assert (
|
||||
format_backend_status(
|
||||
backend="litellm-bedrock",
|
||||
anyllm_provider="ignored",
|
||||
bedrock_region="us-west-2",
|
||||
)
|
||||
== "BEDROCK via LiteLLM (region=us-west-2)"
|
||||
)
|
||||
assert (
|
||||
format_backend_status(
|
||||
backend="litellm-openai",
|
||||
anyllm_provider="ignored",
|
||||
bedrock_region=None,
|
||||
)
|
||||
== "OPENAI via LiteLLM"
|
||||
)
|
||||
|
||||
|
||||
def test_call_client_transport_covers_openai_and_anthropic_paths() -> None:
|
||||
client = DummyClient()
|
||||
openai_metrics = SimpleNamespace(tokens_output=0, cached_tokens=0)
|
||||
anthropic_metrics = SimpleNamespace(tokens_output=0, cached_tokens=0)
|
||||
|
||||
openai_response = call_client_transport(
|
||||
"openai",
|
||||
client,
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
stream=False,
|
||||
metrics=openai_metrics,
|
||||
temperature=0,
|
||||
)
|
||||
openai_stream = call_client_transport(
|
||||
"openai",
|
||||
client,
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
stream=True,
|
||||
metrics=openai_metrics,
|
||||
)
|
||||
anthropic_response = call_client_transport(
|
||||
"anthropic",
|
||||
client,
|
||||
model="claude-sonnet",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
stream=False,
|
||||
metrics=anthropic_metrics,
|
||||
max_tokens=32,
|
||||
)
|
||||
anthropic_stream = call_client_transport(
|
||||
"anthropic",
|
||||
client,
|
||||
model="claude-sonnet",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
stream=True,
|
||||
metrics=anthropic_metrics,
|
||||
max_tokens=32,
|
||||
)
|
||||
|
||||
assert openai_response.usage.completion_tokens == 7
|
||||
assert openai_metrics.tokens_output == 7
|
||||
assert openai_metrics.cached_tokens == 3
|
||||
assert openai_stream == ("wrapped", client._wrapped_stream[0])
|
||||
assert anthropic_response.usage.output_tokens == 5
|
||||
assert anthropic_metrics.tokens_output == 5
|
||||
assert anthropic_metrics.cached_tokens == 2
|
||||
assert anthropic_stream == "anthropic-stream"
|
||||
assert len(client._storage.saved) == 3
|
||||
|
||||
|
||||
def test_call_client_transport_rejects_unknown_api_style() -> None:
|
||||
with pytest.raises(ValueError, match="Unsupported api_style"):
|
||||
call_client_transport(
|
||||
"unknown",
|
||||
DummyClient(),
|
||||
model="gpt-4o",
|
||||
messages=[],
|
||||
stream=False,
|
||||
metrics=SimpleNamespace(),
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue