mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat: add Vertex AI proxy routing (#793)
## 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
This commit is contained in:
JD Davis
GitHub
parent
3db6cd430f
commit
3c77e52ce4
18 changed files with 1193 additions and 72 deletions
|
|
@ -441,6 +441,11 @@ def _selected_context_tool() -> str:
|
|||
default=None,
|
||||
help="Custom Cloud Code Assist API URL for compatibility endpoints (env: CLOUDCODE_TARGET_API_URL)",
|
||||
)
|
||||
@click.option(
|
||||
"--vertex-api-url",
|
||||
default=None,
|
||||
help=("Custom Vertex AI regional API URL for publisher endpoints (env: VERTEX_TARGET_API_URL)"),
|
||||
)
|
||||
@click.option(
|
||||
"--region",
|
||||
default="us-west-2",
|
||||
|
|
@ -518,6 +523,7 @@ def proxy(
|
|||
openai_api_url: str | None,
|
||||
gemini_api_url: str | None,
|
||||
cloudcode_api_url: str | None,
|
||||
vertex_api_url: str | None,
|
||||
region: str,
|
||||
bedrock_region: str | None,
|
||||
bedrock_profile: str | None,
|
||||
|
|
@ -578,6 +584,7 @@ def proxy(
|
|||
openai_api_url=openai_api_url,
|
||||
gemini_api_url=gemini_api_url,
|
||||
cloudcode_api_url=cloudcode_api_url,
|
||||
vertex_api_url=vertex_api_url,
|
||||
environ=os.environ,
|
||||
)
|
||||
|
||||
|
|
@ -634,6 +641,7 @@ def proxy(
|
|||
openai_api_url=provider_api_overrides.openai,
|
||||
gemini_api_url=provider_api_overrides.gemini,
|
||||
cloudcode_api_url=provider_api_overrides.cloudcode,
|
||||
vertex_api_url=provider_api_overrides.vertex,
|
||||
mode=effective_mode,
|
||||
optimize=not no_optimize,
|
||||
cache_enabled=not no_cache,
|
||||
|
|
@ -729,6 +737,7 @@ def proxy(
|
|||
anthropic_url = provider_api_targets.anthropic
|
||||
openai_url = provider_api_targets.openai
|
||||
cloudcode_url = provider_api_targets.cloudcode
|
||||
vertex_url = provider_api_targets.vertex
|
||||
backend_section = ""
|
||||
|
||||
if config.backend == "anyllm" or config.backend.startswith("anyllm-"):
|
||||
|
|
@ -851,6 +860,7 @@ Routing:
|
|||
/v1/chat/completions → {openai_url}
|
||||
/v1/responses → {openai_url} (HTTP + WebSocket)
|
||||
/v1internal:streamGenerateContent → {cloudcode_url}
|
||||
/v1/projects/.../publishers/... → {vertex_url}
|
||||
|
||||
Usage:
|
||||
Claude Code: ANTHROPIC_BASE_URL=http://{config.host}:{config.port} claude
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ def _api_target(proxy: Any, provider_name: str) -> str:
|
|||
"openai": "OPENAI_API_URL",
|
||||
"gemini": "GEMINI_API_URL",
|
||||
"cloudcode": "CLOUDCODE_API_URL",
|
||||
"vertex": "VERTEX_API_URL",
|
||||
}
|
||||
legacy_attr = legacy_attrs[provider_name]
|
||||
return cast(str, getattr(proxy, legacy_attr, proxy.provider_runtime.api_target(provider_name)))
|
||||
|
|
@ -315,6 +316,14 @@ async def _handle_chatgpt_model_metadata(
|
|||
def register_provider_routes(app: FastAPI, proxy: Any) -> None:
|
||||
"""Register provider-specific proxy endpoints."""
|
||||
|
||||
async def vertex_publisher_passthrough(request: Request, publisher: str, action: str):
|
||||
return await proxy.handle_passthrough(
|
||||
request,
|
||||
_api_target(proxy, "vertex"),
|
||||
action,
|
||||
f"vertex:{publisher}",
|
||||
)
|
||||
|
||||
@app.post("/v1/messages")
|
||||
async def anthropic_messages(request: Request):
|
||||
return await proxy.handle_anthropic_messages(request)
|
||||
|
|
@ -467,6 +476,112 @@ def register_provider_routes(app: FastAPI, proxy: Any) -> None:
|
|||
async def google_cloudcode_stream_generate_content_v1(request: Request):
|
||||
return await proxy.handle_google_cloudcode_stream(request)
|
||||
|
||||
@app.post(
|
||||
"/{api_version}/projects/{project}/locations/{location}/publishers/{publisher}/models/{model}:generateContent"
|
||||
)
|
||||
async def vertex_generate_content(
|
||||
request: Request,
|
||||
api_version: str,
|
||||
project: str,
|
||||
location: str,
|
||||
publisher: str,
|
||||
model: str,
|
||||
):
|
||||
del api_version, project, location
|
||||
if publisher == "google":
|
||||
return await proxy.handle_gemini_generate_content(
|
||||
request,
|
||||
model,
|
||||
_api_target(proxy, "vertex"),
|
||||
"vertex:google",
|
||||
)
|
||||
return await vertex_publisher_passthrough(request, publisher, "generateContent")
|
||||
|
||||
@app.post(
|
||||
"/{api_version}/projects/{project}/locations/{location}/publishers/{publisher}/models/{model}:streamGenerateContent"
|
||||
)
|
||||
async def vertex_stream_generate_content(
|
||||
request: Request,
|
||||
api_version: str,
|
||||
project: str,
|
||||
location: str,
|
||||
publisher: str,
|
||||
model: str,
|
||||
):
|
||||
del api_version, project, location
|
||||
if publisher == "google":
|
||||
return await proxy.handle_gemini_generate_content(
|
||||
request,
|
||||
model,
|
||||
_api_target(proxy, "vertex"),
|
||||
"vertex:google",
|
||||
)
|
||||
return await vertex_publisher_passthrough(request, publisher, "streamGenerateContent")
|
||||
|
||||
@app.post(
|
||||
"/{api_version}/projects/{project}/locations/{location}/publishers/{publisher}/models/{model}:countTokens"
|
||||
)
|
||||
async def vertex_count_tokens(
|
||||
request: Request,
|
||||
api_version: str,
|
||||
project: str,
|
||||
location: str,
|
||||
publisher: str,
|
||||
model: str,
|
||||
):
|
||||
del api_version, project, location
|
||||
if publisher == "google":
|
||||
return await proxy.handle_gemini_count_tokens(
|
||||
request,
|
||||
model,
|
||||
_api_target(proxy, "vertex"),
|
||||
"vertex:google",
|
||||
)
|
||||
return await vertex_publisher_passthrough(request, publisher, "countTokens")
|
||||
|
||||
@app.post(
|
||||
"/{api_version}/projects/{project}/locations/{location}/publishers/{publisher}/models/{model}:rawPredict"
|
||||
)
|
||||
async def vertex_raw_predict(
|
||||
request: Request,
|
||||
api_version: str,
|
||||
project: str,
|
||||
location: str,
|
||||
publisher: str,
|
||||
model: str,
|
||||
):
|
||||
del api_version, project, location
|
||||
if publisher == "anthropic":
|
||||
return await proxy.handle_anthropic_messages(
|
||||
request,
|
||||
_api_target(proxy, "vertex"),
|
||||
"vertex:anthropic",
|
||||
model,
|
||||
)
|
||||
return await vertex_publisher_passthrough(request, publisher, "rawPredict")
|
||||
|
||||
@app.post(
|
||||
"/{api_version}/projects/{project}/locations/{location}/publishers/{publisher}/models/{model}:streamRawPredict"
|
||||
)
|
||||
async def vertex_stream_raw_predict(
|
||||
request: Request,
|
||||
api_version: str,
|
||||
project: str,
|
||||
location: str,
|
||||
publisher: str,
|
||||
model: str,
|
||||
):
|
||||
del api_version, project, location
|
||||
if publisher == "anthropic":
|
||||
return await proxy.handle_anthropic_messages(
|
||||
request,
|
||||
_api_target(proxy, "vertex"),
|
||||
"vertex:anthropic",
|
||||
model,
|
||||
True,
|
||||
)
|
||||
return await vertex_publisher_passthrough(request, publisher, "streamRawPredict")
|
||||
|
||||
@app.get("/v1/models")
|
||||
async def list_models(request: Request):
|
||||
chatgpt_response = await _handle_chatgpt_model_metadata(
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from headroom.providers.codex import DEFAULT_API_URL as DEFAULT_OPENAI_API_URL
|
|||
from headroom.providers.gemini import DEFAULT_API_URL as DEFAULT_GEMINI_API_URL
|
||||
|
||||
DEFAULT_CLOUDCODE_API_URL = "https://cloudcode-pa.googleapis.com"
|
||||
DEFAULT_VERTEX_API_URL = "https://us-central1-aiplatform.googleapis.com"
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from headroom.backends.base import Backend
|
||||
|
|
@ -30,6 +31,7 @@ class ProviderApiOverrides:
|
|||
openai: str | None = None
|
||||
gemini: str | None = None
|
||||
cloudcode: str | None = None
|
||||
vertex: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
|
|
@ -40,6 +42,7 @@ class ProviderApiTargets:
|
|||
openai: str = DEFAULT_OPENAI_API_URL
|
||||
gemini: str = DEFAULT_GEMINI_API_URL
|
||||
cloudcode: str = DEFAULT_CLOUDCODE_API_URL
|
||||
vertex: str = DEFAULT_VERTEX_API_URL
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
|
|
@ -56,6 +59,7 @@ class ProxyProviderRuntime:
|
|||
"openai": self.api_targets.openai,
|
||||
"gemini": self.api_targets.gemini,
|
||||
"cloudcode": self.api_targets.cloudcode,
|
||||
"vertex": self.api_targets.vertex,
|
||||
}[provider_name]
|
||||
|
||||
def pipeline_provider(self, provider_name: str) -> Provider:
|
||||
|
|
@ -95,6 +99,7 @@ def resolve_api_overrides(
|
|||
openai_api_url: str | None,
|
||||
gemini_api_url: str | None,
|
||||
cloudcode_api_url: str | None,
|
||||
vertex_api_url: str | None = None,
|
||||
environ: Mapping[str, str] | None = None,
|
||||
) -> ProviderApiOverrides:
|
||||
"""Resolve provider API URL overrides from CLI/config inputs and environment."""
|
||||
|
|
@ -106,6 +111,7 @@ def resolve_api_overrides(
|
|||
openai=openai_api_url or env.get("OPENAI_TARGET_API_URL"),
|
||||
gemini=gemini_api_url or env.get("GEMINI_TARGET_API_URL"),
|
||||
cloudcode=cloudcode_api_url or env.get("CLOUDCODE_TARGET_API_URL"),
|
||||
vertex=vertex_api_url or env.get("VERTEX_TARGET_API_URL"),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -116,6 +122,7 @@ def resolve_api_targets(overrides: ProviderApiOverrides) -> ProviderApiTargets:
|
|||
openai=_normalize_api_url(overrides.openai, default=DEFAULT_OPENAI_API_URL),
|
||||
gemini=_normalize_api_url(overrides.gemini, default=DEFAULT_GEMINI_API_URL),
|
||||
cloudcode=_normalize_api_url(overrides.cloudcode, default=DEFAULT_CLOUDCODE_API_URL),
|
||||
vertex=_normalize_api_url(overrides.vertex, default=DEFAULT_VERTEX_API_URL),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ if TYPE_CHECKING:
|
|||
|
||||
import httpx
|
||||
|
||||
from headroom.copilot_auth import build_copilot_upstream_url
|
||||
from headroom.pipeline import PipelineStage, summarize_routing_markers
|
||||
from headroom.proxy.auth_mode import classify_auth_mode, classify_client
|
||||
from headroom.proxy.compression_decision import CompressionDecision
|
||||
|
|
@ -404,6 +405,10 @@ class AnthropicHandlerMixin:
|
|||
async def handle_anthropic_messages(
|
||||
self,
|
||||
request: Request,
|
||||
upstream_base_url: str | None = None,
|
||||
provider_name: str = "anthropic",
|
||||
model_override: str | None = None,
|
||||
force_stream: bool = False,
|
||||
) -> Response | StreamingResponse:
|
||||
"""Handle Anthropic /v1/messages endpoint."""
|
||||
if not hasattr(self, "pipeline_extensions"):
|
||||
|
|
@ -590,19 +595,22 @@ class AnthropicHandlerMixin:
|
|||
},
|
||||
},
|
||||
)
|
||||
model = body.get("model", "unknown")
|
||||
model = body.get("model") or model_override or "unknown"
|
||||
messages = body.get("messages", [])
|
||||
pipeline_provider = provider_name
|
||||
pipeline_path = request.url.path if upstream_base_url else "/v1/messages"
|
||||
pipeline_stream = bool(body.get("stream", False) or force_stream)
|
||||
with stage_timer.measure("deep_copy"):
|
||||
original_client_messages = copy.deepcopy(messages)
|
||||
input_event = self.pipeline_extensions.emit(
|
||||
PipelineStage.INPUT_RECEIVED,
|
||||
operation="proxy.request",
|
||||
request_id=request_id,
|
||||
provider="anthropic",
|
||||
provider=pipeline_provider,
|
||||
model=model,
|
||||
messages=messages,
|
||||
tools=body.get("tools"),
|
||||
metadata={"path": "/v1/messages", "stream": body.get("stream", False)},
|
||||
metadata={"path": pipeline_path, "stream": pipeline_stream},
|
||||
)
|
||||
if input_event.messages is not None:
|
||||
messages = input_event.messages
|
||||
|
|
@ -626,7 +634,7 @@ class AnthropicHandlerMixin:
|
|||
},
|
||||
)
|
||||
|
||||
stream = body.get("stream", False)
|
||||
stream = pipeline_stream
|
||||
|
||||
# Bypass: skip ALL compression, TOIN learning, and CCR injection
|
||||
# when the caller explicitly opts out via header.
|
||||
|
|
@ -701,7 +709,7 @@ class AnthropicHandlerMixin:
|
|||
rate_key = f"{api_key[:16]}:{client_ip}" if api_key else client_ip
|
||||
allowed, wait_seconds = await self.rate_limiter.check_request(rate_key)
|
||||
if not allowed:
|
||||
await self.metrics.record_rate_limited(provider="anthropic")
|
||||
await self.metrics.record_rate_limited(provider=provider_name)
|
||||
# Unit 4: release the pre-upstream semaphore before we
|
||||
# bail out of the handler via HTTPException — FastAPI's
|
||||
# exception handler will NOT run our ``finally``.
|
||||
|
|
@ -781,10 +789,10 @@ class AnthropicHandlerMixin:
|
|||
PipelineStage.INPUT_CACHED,
|
||||
operation="proxy.request",
|
||||
request_id=request_id,
|
||||
provider="anthropic",
|
||||
provider=pipeline_provider,
|
||||
model=model,
|
||||
messages=messages,
|
||||
metadata={"cache_hit": True, "path": "/v1/messages"},
|
||||
metadata={"cache_hit": True, "path": pipeline_path},
|
||||
)
|
||||
optimization_latency = (time.time() - start_time) * 1000
|
||||
|
||||
|
|
@ -798,7 +806,7 @@ class AnthropicHandlerMixin:
|
|||
await self._record_request_outcome(
|
||||
RequestOutcome(
|
||||
request_id=request_id,
|
||||
provider="anthropic",
|
||||
provider=provider_name,
|
||||
model=model,
|
||||
original_tokens=0,
|
||||
optimized_tokens=0,
|
||||
|
|
@ -839,7 +847,7 @@ class AnthropicHandlerMixin:
|
|||
messages, _security_ctx = self.security.scan_request(
|
||||
messages,
|
||||
{
|
||||
"provider": "anthropic",
|
||||
"provider": provider_name,
|
||||
"model": model,
|
||||
"request_id": str(request_id),
|
||||
"user_id": headers.get("x-api-key", "")[:16],
|
||||
|
|
@ -872,7 +880,7 @@ class AnthropicHandlerMixin:
|
|||
_hook_ctx = CompressContext(
|
||||
model=model,
|
||||
user_query=extract_user_query(messages),
|
||||
provider="anthropic",
|
||||
provider=provider_name,
|
||||
)
|
||||
try:
|
||||
messages = self.config.hooks.pre_compress(messages, _hook_ctx)
|
||||
|
|
@ -1176,7 +1184,7 @@ class AnthropicHandlerMixin:
|
|||
PipelineStage.INPUT_ROUTED,
|
||||
operation="proxy.request",
|
||||
request_id=request_id,
|
||||
provider="anthropic",
|
||||
provider=pipeline_provider,
|
||||
model=model,
|
||||
messages=optimized_messages,
|
||||
metadata={
|
||||
|
|
@ -1195,7 +1203,7 @@ class AnthropicHandlerMixin:
|
|||
PipelineStage.INPUT_COMPRESSED,
|
||||
operation="proxy.request",
|
||||
request_id=request_id,
|
||||
provider="anthropic",
|
||||
provider=pipeline_provider,
|
||||
model=model,
|
||||
messages=optimized_messages,
|
||||
metadata={
|
||||
|
|
@ -1227,7 +1235,7 @@ class AnthropicHandlerMixin:
|
|||
transforms_applied=transforms_applied,
|
||||
model=model,
|
||||
user_query=_hook_ctx.user_query if self.config.hooks else "",
|
||||
provider="anthropic",
|
||||
provider=provider_name,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -1627,7 +1635,7 @@ class AnthropicHandlerMixin:
|
|||
PipelineStage.INPUT_REMEMBERED,
|
||||
operation="proxy.request",
|
||||
request_id=request_id,
|
||||
provider="anthropic",
|
||||
provider=pipeline_provider,
|
||||
model=model,
|
||||
messages=optimized_messages,
|
||||
tools=tools,
|
||||
|
|
@ -1654,12 +1662,12 @@ class AnthropicHandlerMixin:
|
|||
PipelineStage.PRE_SEND,
|
||||
operation="proxy.request",
|
||||
request_id=request_id,
|
||||
provider="anthropic",
|
||||
provider=pipeline_provider,
|
||||
model=model,
|
||||
messages=optimized_messages,
|
||||
tools=tools,
|
||||
headers=headers,
|
||||
metadata={"path": "/v1/messages", "stream": stream},
|
||||
metadata={"path": pipeline_path, "stream": stream},
|
||||
)
|
||||
previous_presend_messages = optimized_messages
|
||||
if presend_event.messages is not None:
|
||||
|
|
@ -1705,11 +1713,11 @@ class AnthropicHandlerMixin:
|
|||
PipelineStage.POST_SEND,
|
||||
operation="proxy.request",
|
||||
request_id=request_id,
|
||||
provider="anthropic",
|
||||
provider=pipeline_provider,
|
||||
model=model,
|
||||
messages=body["messages"],
|
||||
tools=tools,
|
||||
metadata={"path": "/v1/messages", "stream": True},
|
||||
metadata={"path": pipeline_path, "stream": True},
|
||||
)
|
||||
await _finalize_pre_upstream()
|
||||
return await self._stream_response_bedrock(
|
||||
|
|
@ -1735,13 +1743,13 @@ class AnthropicHandlerMixin:
|
|||
PipelineStage.POST_SEND,
|
||||
operation="proxy.request",
|
||||
request_id=request_id,
|
||||
provider="anthropic",
|
||||
provider=pipeline_provider,
|
||||
model=model,
|
||||
messages=body["messages"],
|
||||
tools=tools,
|
||||
response=backend_response.body,
|
||||
metadata={
|
||||
"path": "/v1/messages",
|
||||
"path": pipeline_path,
|
||||
"stream": False,
|
||||
"status_code": backend_response.status_code,
|
||||
},
|
||||
|
|
@ -1750,11 +1758,11 @@ class AnthropicHandlerMixin:
|
|||
PipelineStage.RESPONSE_RECEIVED,
|
||||
operation="proxy.request",
|
||||
request_id=request_id,
|
||||
provider="anthropic",
|
||||
provider=pipeline_provider,
|
||||
model=model,
|
||||
response=backend_response.body,
|
||||
metadata={
|
||||
"path": "/v1/messages",
|
||||
"path": pipeline_path,
|
||||
"stream": False,
|
||||
"status_code": backend_response.status_code,
|
||||
},
|
||||
|
|
@ -1848,8 +1856,15 @@ class AnthropicHandlerMixin:
|
|||
},
|
||||
)
|
||||
|
||||
# Direct Anthropic API
|
||||
url = f"{self.ANTHROPIC_API_URL}/v1/messages"
|
||||
# Direct Anthropic API, or a provider-compatible Anthropic
|
||||
# Messages endpoint such as Vertex AI publisher rawPredict.
|
||||
url = (
|
||||
build_copilot_upstream_url(upstream_base_url, request.url.path)
|
||||
if upstream_base_url
|
||||
else f"{self.ANTHROPIC_API_URL}/v1/messages"
|
||||
)
|
||||
if upstream_base_url and request.url.query:
|
||||
url = f"{url}?{request.url.query}"
|
||||
|
||||
try:
|
||||
if stream:
|
||||
|
|
@ -1857,11 +1872,11 @@ class AnthropicHandlerMixin:
|
|||
PipelineStage.POST_SEND,
|
||||
operation="proxy.request",
|
||||
request_id=request_id,
|
||||
provider="anthropic",
|
||||
provider=pipeline_provider,
|
||||
model=model,
|
||||
messages=body["messages"],
|
||||
tools=tools,
|
||||
metadata={"path": "/v1/messages", "stream": True},
|
||||
metadata={"path": pipeline_path, "stream": True},
|
||||
)
|
||||
await _finalize_pre_upstream()
|
||||
return await self._stream_response(
|
||||
|
|
@ -1885,6 +1900,7 @@ class AnthropicHandlerMixin:
|
|||
body_mutated=body_mutation_tracker.mutated,
|
||||
mutation_reasons=body_mutation_tracker.reasons,
|
||||
memory_request_ctx=memory_request_ctx,
|
||||
outcome_provider=provider_name,
|
||||
)
|
||||
else:
|
||||
async with stage_timer.measure("upstream_connect"):
|
||||
|
|
@ -1904,13 +1920,13 @@ class AnthropicHandlerMixin:
|
|||
PipelineStage.POST_SEND,
|
||||
operation="proxy.request",
|
||||
request_id=request_id,
|
||||
provider="anthropic",
|
||||
provider=pipeline_provider,
|
||||
model=model,
|
||||
messages=body["messages"],
|
||||
tools=tools,
|
||||
response=response,
|
||||
metadata={
|
||||
"path": "/v1/messages",
|
||||
"path": pipeline_path,
|
||||
"stream": False,
|
||||
"status_code": response.status_code,
|
||||
},
|
||||
|
|
@ -1919,11 +1935,11 @@ class AnthropicHandlerMixin:
|
|||
PipelineStage.RESPONSE_RECEIVED,
|
||||
operation="proxy.request",
|
||||
request_id=request_id,
|
||||
provider="anthropic",
|
||||
provider=pipeline_provider,
|
||||
model=model,
|
||||
response=response,
|
||||
metadata={
|
||||
"path": "/v1/messages",
|
||||
"path": pipeline_path,
|
||||
"stream": False,
|
||||
"status_code": response.status_code,
|
||||
},
|
||||
|
|
@ -2316,7 +2332,7 @@ class AnthropicHandlerMixin:
|
|||
await self._record_request_outcome(
|
||||
RequestOutcome(
|
||||
request_id=request_id,
|
||||
provider="anthropic",
|
||||
provider=provider_name,
|
||||
model=model,
|
||||
original_tokens=original_tokens,
|
||||
optimized_tokens=optimized_tokens,
|
||||
|
|
@ -2395,7 +2411,7 @@ class AnthropicHandlerMixin:
|
|||
# Retry-After header. The outer finally still runs.
|
||||
raise
|
||||
except Exception as e:
|
||||
await self.metrics.record_failed(provider="anthropic")
|
||||
await self.metrics.record_failed(provider=provider_name)
|
||||
# Log full error details internally for debugging
|
||||
logger.error(f"[{request_id}] Request failed: {type(e).__name__}: {e}")
|
||||
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ if TYPE_CHECKING:
|
|||
from fastapi import Request
|
||||
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
||||
|
||||
from headroom.copilot_auth import build_copilot_upstream_url
|
||||
from headroom.proxy.auth_mode import classify_client
|
||||
from headroom.proxy.compression_decision import CompressionDecision
|
||||
from headroom.proxy.helpers import extract_tags
|
||||
|
|
@ -145,6 +146,8 @@ class GeminiHandlerMixin:
|
|||
self,
|
||||
request: Request,
|
||||
model: str,
|
||||
upstream_base_url: str | None = None,
|
||||
provider_name: str = "gemini",
|
||||
) -> Response | StreamingResponse:
|
||||
"""Handle Gemini native /v1beta/models/{model}:generateContent endpoint.
|
||||
|
||||
|
|
@ -275,7 +278,7 @@ class GeminiHandlerMixin:
|
|||
rate_key = headers.get("x-goog-api-key", "default")[:20]
|
||||
allowed, wait_seconds = await self.rate_limiter.check_request(rate_key)
|
||||
if not allowed:
|
||||
await self.metrics.record_rate_limited(provider="gemini")
|
||||
await self.metrics.record_rate_limited(provider=provider_name)
|
||||
raise HTTPException(
|
||||
status_code=429,
|
||||
detail=f"Rate limited. Retry after {wait_seconds:.1f}s",
|
||||
|
|
@ -294,17 +297,32 @@ class GeminiHandlerMixin:
|
|||
if len(preserved_indices) == len(contents):
|
||||
# All content has non-text parts, skip compression entirely
|
||||
# Just forward the request as-is
|
||||
url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:generateContent"
|
||||
query_params = dict(request.query_params)
|
||||
is_streaming = query_params.get("alt") == "sse"
|
||||
if "key" in query_params:
|
||||
is_streaming = query_params.get("alt") == "sse" or request.url.path.endswith(
|
||||
":streamGenerateContent"
|
||||
)
|
||||
if upstream_base_url:
|
||||
url = build_copilot_upstream_url(upstream_base_url, request.url.path)
|
||||
if is_streaming:
|
||||
url = url.replace(":generateContent", ":streamGenerateContent")
|
||||
if request.url.query:
|
||||
url = f"{url}?{request.url.query}"
|
||||
else:
|
||||
url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:generateContent"
|
||||
if "key" in query_params and not upstream_base_url:
|
||||
url += f"?key={query_params['key']}"
|
||||
|
||||
if is_streaming:
|
||||
stream_url = (
|
||||
f"{self.GEMINI_API_URL}/v1beta/models/{model}:streamGenerateContent?alt=sse"
|
||||
)
|
||||
if "key" in query_params:
|
||||
if upstream_base_url:
|
||||
stream_url = url
|
||||
separator = "&" if "?" in stream_url else "?"
|
||||
if "alt=" not in request.url.query:
|
||||
stream_url = f"{stream_url}{separator}alt=sse"
|
||||
else:
|
||||
stream_url = (
|
||||
f"{self.GEMINI_API_URL}/v1beta/models/{model}:streamGenerateContent?alt=sse"
|
||||
)
|
||||
if "key" in query_params and not upstream_base_url:
|
||||
stream_url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:streamGenerateContent?key={query_params['key']}&alt=sse"
|
||||
return await self._stream_response(
|
||||
stream_url,
|
||||
|
|
@ -319,12 +337,49 @@ class GeminiHandlerMixin:
|
|||
[],
|
||||
tags,
|
||||
0,
|
||||
outcome_provider=provider_name,
|
||||
)
|
||||
else:
|
||||
response = await self._retry_request("POST", url, headers, body)
|
||||
total_latency = (time.time() - start_time) * 1000
|
||||
total_input_tokens = 0
|
||||
output_tokens = 0
|
||||
cache_read_tokens = 0
|
||||
try:
|
||||
resp_json = response.json()
|
||||
usage = resp_json.get("usageMetadata", {})
|
||||
total_input_tokens = usage.get("promptTokenCount", 0)
|
||||
output_tokens = usage.get("candidatesTokenCount", 0)
|
||||
cache_read_tokens = usage.get("cachedContentTokenCount", 0)
|
||||
except (json.JSONDecodeError, ValueError, KeyError, TypeError, AttributeError):
|
||||
pass
|
||||
await self._record_request_outcome(
|
||||
RequestOutcome(
|
||||
request_id=request_id,
|
||||
provider=provider_name,
|
||||
model=model,
|
||||
original_tokens=total_input_tokens,
|
||||
optimized_tokens=total_input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
tokens_saved=0,
|
||||
attempted_input_tokens=total_input_tokens,
|
||||
cache_read_tokens=cache_read_tokens,
|
||||
uncached_input_tokens=max(0, total_input_tokens - cache_read_tokens),
|
||||
total_latency_ms=total_latency,
|
||||
num_messages=len(contents),
|
||||
tags=tags or {},
|
||||
client=client,
|
||||
)
|
||||
)
|
||||
response_headers = dict(response.headers)
|
||||
response_headers.pop("content-encoding", None)
|
||||
response_headers.pop("content-length", None)
|
||||
response_headers["x-headroom-tokens-before"] = str(total_input_tokens)
|
||||
response_headers["x-headroom-tokens-after"] = str(total_input_tokens)
|
||||
response_headers["x-headroom-tokens-saved"] = "0"
|
||||
response_headers["x-headroom-model"] = model
|
||||
if cache_read_tokens > 0:
|
||||
response_headers["x-headroom-cached"] = "true"
|
||||
return Response(
|
||||
content=response.content,
|
||||
status_code=response.status_code,
|
||||
|
|
@ -458,24 +513,41 @@ class GeminiHandlerMixin:
|
|||
elif "systemInstruction" in body:
|
||||
del body["systemInstruction"]
|
||||
|
||||
# Build URL - model is extracted from path
|
||||
url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:generateContent"
|
||||
|
||||
# Check if streaming requested via query param
|
||||
query_params = dict(request.query_params)
|
||||
is_streaming = query_params.get("alt") == "sse"
|
||||
is_streaming = query_params.get("alt") == "sse" or request.url.path.endswith(
|
||||
":streamGenerateContent"
|
||||
)
|
||||
|
||||
# Build URL - model is extracted from path. Vertex publisher
|
||||
# routes use the request's full path under the Vertex base URL;
|
||||
# native Gemini uses the public Gemini API shape.
|
||||
if upstream_base_url:
|
||||
url = build_copilot_upstream_url(upstream_base_url, request.url.path)
|
||||
if is_streaming:
|
||||
url = url.replace(":generateContent", ":streamGenerateContent")
|
||||
if request.url.query:
|
||||
url = f"{url}?{request.url.query}"
|
||||
else:
|
||||
url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:generateContent"
|
||||
|
||||
# Preserve API key in query params if present
|
||||
if "key" in query_params:
|
||||
if "key" in query_params and not upstream_base_url:
|
||||
url += f"?key={query_params['key']}"
|
||||
|
||||
try:
|
||||
if is_streaming:
|
||||
# For streaming, use streamGenerateContent endpoint
|
||||
stream_url = (
|
||||
f"{self.GEMINI_API_URL}/v1beta/models/{model}:streamGenerateContent?alt=sse"
|
||||
)
|
||||
if "key" in query_params:
|
||||
if upstream_base_url:
|
||||
stream_url = url
|
||||
separator = "&" if "?" in stream_url else "?"
|
||||
if "alt=" not in request.url.query:
|
||||
stream_url = f"{stream_url}{separator}alt=sse"
|
||||
else:
|
||||
stream_url = (
|
||||
f"{self.GEMINI_API_URL}/v1beta/models/{model}:streamGenerateContent?alt=sse"
|
||||
)
|
||||
if "key" in query_params and not upstream_base_url:
|
||||
stream_url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:streamGenerateContent?key={query_params['key']}&alt=sse"
|
||||
|
||||
return await self._stream_response(
|
||||
|
|
@ -491,6 +563,7 @@ class GeminiHandlerMixin:
|
|||
transforms_applied,
|
||||
tags,
|
||||
optimization_latency,
|
||||
outcome_provider=provider_name,
|
||||
)
|
||||
else:
|
||||
response = await self._retry_request("POST", url, headers, body)
|
||||
|
|
@ -530,7 +603,7 @@ class GeminiHandlerMixin:
|
|||
# without per-handler conditionals.
|
||||
outcome = RequestOutcome(
|
||||
request_id=request_id,
|
||||
provider="gemini",
|
||||
provider=provider_name,
|
||||
model=model,
|
||||
original_tokens=original_tokens,
|
||||
optimized_tokens=total_input_tokens,
|
||||
|
|
@ -580,7 +653,7 @@ class GeminiHandlerMixin:
|
|||
headers=response_headers,
|
||||
)
|
||||
except Exception as e:
|
||||
await self.metrics.record_failed(provider="gemini")
|
||||
await self.metrics.record_failed(provider=provider_name)
|
||||
logger.error(f"[{request_id}] Gemini request failed: {type(e).__name__}: {e}")
|
||||
return JSONResponse(
|
||||
status_code=502,
|
||||
|
|
@ -826,6 +899,8 @@ class GeminiHandlerMixin:
|
|||
self,
|
||||
request: Request,
|
||||
model: str,
|
||||
upstream_base_url: str | None = None,
|
||||
provider_name: str = "gemini",
|
||||
) -> Response:
|
||||
"""Handle Gemini /v1beta/models/{model}:countTokens endpoint with compression.
|
||||
|
||||
|
|
@ -889,9 +964,14 @@ class GeminiHandlerMixin:
|
|||
if len(preserved_indices) == len(contents):
|
||||
# All content has non-text parts, skip compression entirely
|
||||
# Just forward the countTokens request as-is
|
||||
url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:countTokens"
|
||||
if upstream_base_url:
|
||||
url = build_copilot_upstream_url(upstream_base_url, request.url.path)
|
||||
if request.url.query:
|
||||
url = f"{url}?{request.url.query}"
|
||||
else:
|
||||
url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:countTokens"
|
||||
query_params = dict(request.query_params)
|
||||
if "key" in query_params:
|
||||
if "key" in query_params and not upstream_base_url:
|
||||
url += f"?key={query_params['key']}"
|
||||
|
||||
response = await self._retry_request("POST", url, headers, body)
|
||||
|
|
@ -961,11 +1041,16 @@ class GeminiHandlerMixin:
|
|||
del body["systemInstruction"]
|
||||
|
||||
# Build URL
|
||||
url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:countTokens"
|
||||
if upstream_base_url:
|
||||
url = build_copilot_upstream_url(upstream_base_url, request.url.path)
|
||||
if request.url.query:
|
||||
url = f"{url}?{request.url.query}"
|
||||
else:
|
||||
url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:countTokens"
|
||||
|
||||
# Preserve API key in query params if present
|
||||
query_params = dict(request.query_params)
|
||||
if "key" in query_params:
|
||||
if "key" in query_params and not upstream_base_url:
|
||||
url += f"?key={query_params['key']}"
|
||||
|
||||
try:
|
||||
|
|
@ -993,7 +1078,7 @@ class GeminiHandlerMixin:
|
|||
await self._record_request_outcome(
|
||||
RequestOutcome(
|
||||
request_id=request_id,
|
||||
provider="gemini",
|
||||
provider=provider_name,
|
||||
model=model,
|
||||
original_tokens=original_tokens,
|
||||
optimized_tokens=compressed_tokens,
|
||||
|
|
@ -1028,7 +1113,7 @@ class GeminiHandlerMixin:
|
|||
headers=response_headers,
|
||||
)
|
||||
except Exception as e:
|
||||
await self.metrics.record_failed(provider="gemini")
|
||||
await self.metrics.record_failed(provider=provider_name)
|
||||
logger.error(f"[{request_id}] Gemini countTokens failed: {type(e).__name__}: {e}")
|
||||
return JSONResponse(
|
||||
status_code=502,
|
||||
|
|
|
|||
|
|
@ -60,6 +60,56 @@ _OPENAI_RESPONSES_UNIT_EXECUTOR_LOCK = threading.RLock()
|
|||
_OPENAI_RESPONSES_UNIT_EXECUTOR: ThreadPoolExecutor | None = None
|
||||
|
||||
|
||||
def _usage_int(value: Any) -> int:
|
||||
try:
|
||||
return max(int(value), 0)
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
def _passthrough_usage_from_json(payload: Any) -> dict[str, int]:
|
||||
"""Normalize usage from pass-through provider response shapes."""
|
||||
if not isinstance(payload, dict):
|
||||
return {}
|
||||
|
||||
usage_meta = payload.get("usageMetadata")
|
||||
if isinstance(usage_meta, dict):
|
||||
return {
|
||||
"input_tokens": _usage_int(usage_meta.get("promptTokenCount")),
|
||||
"output_tokens": _usage_int(usage_meta.get("candidatesTokenCount")),
|
||||
"cache_read_input_tokens": _usage_int(usage_meta.get("cachedContentTokenCount")),
|
||||
}
|
||||
|
||||
usage = payload.get("usage")
|
||||
if isinstance(usage, dict):
|
||||
input_tokens = usage.get("input_tokens")
|
||||
if input_tokens is None:
|
||||
input_tokens = usage.get("prompt_tokens")
|
||||
output_tokens = usage.get("output_tokens")
|
||||
if output_tokens is None:
|
||||
output_tokens = usage.get("completion_tokens")
|
||||
details = usage.get("prompt_tokens_details") or usage.get("input_tokens_details") or {}
|
||||
cache_read = details.get("cached_tokens") if isinstance(details, dict) else None
|
||||
return {
|
||||
"input_tokens": _usage_int(input_tokens),
|
||||
"output_tokens": _usage_int(output_tokens),
|
||||
"cache_read_input_tokens": _usage_int(usage.get("cache_read_input_tokens", cache_read)),
|
||||
"cache_creation_input_tokens": _usage_int(usage.get("cache_creation_input_tokens")),
|
||||
}
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
def _passthrough_model_from_path(path: str, endpoint_name: str) -> str:
|
||||
marker = "/models/"
|
||||
if marker in path:
|
||||
model_part = path.split(marker, 1)[1].split("/", 1)[0]
|
||||
model = model_part.split(":", 1)[0]
|
||||
if model:
|
||||
return model
|
||||
return f"passthrough:{endpoint_name}"
|
||||
|
||||
|
||||
def _openai_responses_unit_parallelism() -> int:
|
||||
raw = os.getenv(_OPENAI_RESPONSES_UNIT_PARALLELISM_ENV)
|
||||
if raw is None or raw.strip() == "":
|
||||
|
|
@ -5818,6 +5868,14 @@ class OpenAIHandlerMixin:
|
|||
"""
|
||||
from fastapi.responses import Response
|
||||
|
||||
if endpoint_name in {"streamGenerateContent", "streamRawPredict"} and provider:
|
||||
return await self._handle_streaming_passthrough(
|
||||
request=request,
|
||||
base_url=base_url,
|
||||
endpoint_name=endpoint_name,
|
||||
provider=provider,
|
||||
)
|
||||
|
||||
start_time = time.time()
|
||||
path = request.url.path
|
||||
url = build_copilot_upstream_url(base_url, path)
|
||||
|
|
@ -5880,21 +5938,36 @@ class OpenAIHandlerMixin:
|
|||
|
||||
# Passthrough request: forwarded upstream with no transforms.
|
||||
# Still recorded so dashboards see traffic on the passthrough
|
||||
# endpoints. Funnel handles the "no tokens, no cache" shape
|
||||
# via zero defaults.
|
||||
# endpoints. When the upstream exposes provider-native usage
|
||||
# fields, normalize them so dashboard totals do not collapse to
|
||||
# zero for Vertex/Gemini and other pass-through endpoints.
|
||||
if endpoint_name and provider:
|
||||
latency_ms = (time.time() - start_time) * 1000
|
||||
request_id = await self._next_request_id()
|
||||
usage: dict[str, int] = {}
|
||||
if response.headers.get("content-type", "").lower().startswith("application/json"):
|
||||
try:
|
||||
usage = _passthrough_usage_from_json(response.json())
|
||||
except (json.JSONDecodeError, ValueError, TypeError):
|
||||
usage = {}
|
||||
input_tokens = usage.get("input_tokens", 0)
|
||||
output_tokens = usage.get("output_tokens", 0)
|
||||
cache_read_tokens = usage.get("cache_read_input_tokens", 0)
|
||||
cache_write_tokens = usage.get("cache_creation_input_tokens", 0)
|
||||
uncached_input_tokens = max(0, input_tokens - cache_read_tokens - cache_write_tokens)
|
||||
await self._record_request_outcome(
|
||||
RequestOutcome(
|
||||
request_id=request_id,
|
||||
provider=provider,
|
||||
model=f"passthrough:{endpoint_name}",
|
||||
original_tokens=0,
|
||||
optimized_tokens=0,
|
||||
output_tokens=0,
|
||||
model=_passthrough_model_from_path(path, endpoint_name),
|
||||
original_tokens=input_tokens,
|
||||
optimized_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
tokens_saved=0,
|
||||
attempted_input_tokens=0,
|
||||
attempted_input_tokens=input_tokens,
|
||||
cache_read_tokens=cache_read_tokens,
|
||||
cache_write_tokens=cache_write_tokens,
|
||||
uncached_input_tokens=uncached_input_tokens,
|
||||
total_latency_ms=latency_ms,
|
||||
tags=tags,
|
||||
client=client,
|
||||
|
|
@ -5906,3 +5979,184 @@ class OpenAIHandlerMixin:
|
|||
status_code=response.status_code,
|
||||
headers=response_headers,
|
||||
)
|
||||
|
||||
async def _handle_streaming_passthrough(
|
||||
self,
|
||||
request: Request,
|
||||
base_url: str,
|
||||
endpoint_name: str,
|
||||
provider: str,
|
||||
) -> Response:
|
||||
"""Stream pass-through responses without buffering the upstream body."""
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
|
||||
from headroom.proxy.helpers import MAX_SSE_BUFFER_SIZE
|
||||
|
||||
start_time = time.time()
|
||||
path = request.url.path
|
||||
url = build_copilot_upstream_url(base_url, path)
|
||||
if request.url.query:
|
||||
url = f"{url}?{request.url.query}"
|
||||
|
||||
headers = dict(request.headers.items())
|
||||
headers.pop("host", None)
|
||||
headers.pop("accept-encoding", None)
|
||||
client = classify_client(headers)
|
||||
tags = extract_tags(headers)
|
||||
|
||||
from headroom.proxy.helpers import _strip_internal_headers, log_outbound_headers
|
||||
|
||||
_pre_strip_count_pt = sum(1 for k in headers if k.lower().startswith("x-headroom-"))
|
||||
headers = _strip_internal_headers(headers)
|
||||
log_outbound_headers(
|
||||
forwarder="streaming_passthrough",
|
||||
stripped_count=_pre_strip_count_pt,
|
||||
request_id=None,
|
||||
)
|
||||
|
||||
body = await request.body()
|
||||
headers = await apply_copilot_api_auth(headers, url=url)
|
||||
request_id = await self._next_request_id()
|
||||
stream_provider = "gemini" if provider == "vertex:google" else "anthropic"
|
||||
stream_state: dict[str, Any] = {
|
||||
"input_tokens": None,
|
||||
"output_tokens": None,
|
||||
"cache_read_input_tokens": 0,
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_creation_ephemeral_5m_input_tokens": 0,
|
||||
"cache_creation_ephemeral_1h_input_tokens": 0,
|
||||
"total_bytes": 0,
|
||||
"sse_buffer": bytearray(),
|
||||
"ttfb_ms": None,
|
||||
}
|
||||
|
||||
assert self.http_client is not None, "http_client must be initialized before streaming"
|
||||
try:
|
||||
upstream_request = self.http_client.build_request(
|
||||
request.method,
|
||||
url,
|
||||
headers=headers,
|
||||
content=body,
|
||||
)
|
||||
upstream_response = await self.http_client.send(upstream_request, stream=True)
|
||||
except (httpx.ConnectError, httpx.TimeoutException) as e:
|
||||
logger.warning(
|
||||
"Streaming passthrough failed before upstream response: %s %s -> %s: %s",
|
||||
request.method,
|
||||
path,
|
||||
url,
|
||||
e,
|
||||
)
|
||||
return Response(
|
||||
content=json.dumps(
|
||||
{
|
||||
"error": {
|
||||
"type": "connection_error",
|
||||
"message": f"Failed to connect to upstream API: {e}",
|
||||
}
|
||||
}
|
||||
),
|
||||
status_code=502,
|
||||
media_type="application/json",
|
||||
)
|
||||
|
||||
response_headers = dict(upstream_response.headers)
|
||||
response_headers.pop("content-length", None)
|
||||
response_headers.pop("transfer-encoding", None)
|
||||
response_headers.pop("connection", None)
|
||||
response_headers.pop("content-encoding", None)
|
||||
|
||||
if upstream_response.status_code >= 400:
|
||||
try:
|
||||
error_content = await upstream_response.aread()
|
||||
finally:
|
||||
await upstream_response.aclose()
|
||||
return Response(
|
||||
content=error_content,
|
||||
status_code=upstream_response.status_code,
|
||||
headers=response_headers,
|
||||
)
|
||||
|
||||
def _absorb_usage(usage: dict[str, int] | None) -> None:
|
||||
if not usage:
|
||||
return
|
||||
if "input_tokens" in usage:
|
||||
stream_state["input_tokens"] = usage["input_tokens"]
|
||||
if "output_tokens" in usage:
|
||||
stream_state["output_tokens"] = usage["output_tokens"]
|
||||
if "cache_read_input_tokens" in usage:
|
||||
stream_state["cache_read_input_tokens"] = usage["cache_read_input_tokens"]
|
||||
if "cache_creation_input_tokens" in usage:
|
||||
stream_state["cache_creation_input_tokens"] = usage["cache_creation_input_tokens"]
|
||||
if "cache_creation_ephemeral_5m_input_tokens" in usage:
|
||||
stream_state["cache_creation_ephemeral_5m_input_tokens"] = usage[
|
||||
"cache_creation_ephemeral_5m_input_tokens"
|
||||
]
|
||||
if "cache_creation_ephemeral_1h_input_tokens" in usage:
|
||||
stream_state["cache_creation_ephemeral_1h_input_tokens"] = usage[
|
||||
"cache_creation_ephemeral_1h_input_tokens"
|
||||
]
|
||||
|
||||
async def generate():
|
||||
try:
|
||||
async with contextlib.aclosing(upstream_response) as response:
|
||||
async for chunk in response.aiter_bytes():
|
||||
if stream_state["ttfb_ms"] is None:
|
||||
stream_state["ttfb_ms"] = (time.time() - start_time) * 1000
|
||||
stream_state["total_bytes"] += len(chunk)
|
||||
stream_state["sse_buffer"].extend(chunk)
|
||||
if len(stream_state["sse_buffer"]) > MAX_SSE_BUFFER_SIZE:
|
||||
tail = bytes(stream_state["sse_buffer"][-MAX_SSE_BUFFER_SIZE // 2 :])
|
||||
stream_state["sse_buffer"] = bytearray(tail)
|
||||
|
||||
_absorb_usage(
|
||||
self._parse_sse_usage_from_buffer(stream_state, stream_provider)
|
||||
)
|
||||
yield chunk
|
||||
finally:
|
||||
buf = stream_state["sse_buffer"]
|
||||
if len(buf) > 0:
|
||||
buf.extend(b"\n\n")
|
||||
_absorb_usage(self._parse_sse_usage_from_buffer(stream_state, stream_provider))
|
||||
|
||||
input_tokens = stream_state["input_tokens"] or 0
|
||||
output_tokens = stream_state["output_tokens"] or 0
|
||||
cache_read_tokens = stream_state["cache_read_input_tokens"] or 0
|
||||
cache_write_tokens = stream_state["cache_creation_input_tokens"] or 0
|
||||
uncached_input_tokens = max(
|
||||
0,
|
||||
input_tokens - cache_read_tokens - cache_write_tokens,
|
||||
)
|
||||
await self._record_request_outcome(
|
||||
RequestOutcome(
|
||||
request_id=request_id,
|
||||
provider=provider,
|
||||
model=_passthrough_model_from_path(path, endpoint_name),
|
||||
original_tokens=input_tokens,
|
||||
optimized_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
tokens_saved=0,
|
||||
attempted_input_tokens=input_tokens,
|
||||
cache_read_tokens=cache_read_tokens,
|
||||
cache_write_tokens=cache_write_tokens,
|
||||
cache_write_5m_tokens=stream_state[
|
||||
"cache_creation_ephemeral_5m_input_tokens"
|
||||
],
|
||||
cache_write_1h_tokens=stream_state[
|
||||
"cache_creation_ephemeral_1h_input_tokens"
|
||||
],
|
||||
uncached_input_tokens=uncached_input_tokens,
|
||||
total_latency_ms=(time.time() - start_time) * 1000,
|
||||
ttfb_ms=stream_state["ttfb_ms"] or 0,
|
||||
tags=tags,
|
||||
client=client,
|
||||
)
|
||||
)
|
||||
|
||||
media_type = upstream_response.headers.get("content-type") or "text/event-stream"
|
||||
return StreamingResponse(
|
||||
generate(),
|
||||
status_code=upstream_response.status_code,
|
||||
headers=response_headers,
|
||||
media_type=media_type,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -646,6 +646,7 @@ class StreamingMixin:
|
|||
*,
|
||||
body: dict,
|
||||
provider: str,
|
||||
outcome_provider: str | None = None,
|
||||
model: str,
|
||||
request_id: str,
|
||||
original_tokens: int,
|
||||
|
|
@ -665,6 +666,7 @@ class StreamingMixin:
|
|||
) -> None:
|
||||
from headroom.proxy.outcome import RequestOutcome
|
||||
|
||||
outcome_provider = outcome_provider or provider
|
||||
total_latency = (time.time() - start_time) * 1000
|
||||
|
||||
# Per-chunk SSE parsing only flushes events terminated by ``\n\n``.
|
||||
|
|
@ -710,7 +712,7 @@ class StreamingMixin:
|
|||
effective_optimized_tokens = optimized_tokens
|
||||
effective_original_tokens = original_tokens
|
||||
if (
|
||||
provider == "openai"
|
||||
provider in {"openai", "gemini"}
|
||||
and isinstance(provider_input_tokens, int)
|
||||
and provider_input_tokens > 0
|
||||
):
|
||||
|
|
@ -763,7 +765,7 @@ class StreamingMixin:
|
|||
# happening (issue #455).
|
||||
outcome = RequestOutcome.from_stream(
|
||||
body=body,
|
||||
provider=provider,
|
||||
provider=outcome_provider,
|
||||
model=model,
|
||||
request_id=request_id,
|
||||
original_tokens=effective_original_tokens,
|
||||
|
|
@ -809,6 +811,7 @@ class StreamingMixin:
|
|||
body_mutated: bool = True,
|
||||
mutation_reasons: list[str] | None = None,
|
||||
memory_request_ctx: Any | None = None,
|
||||
outcome_provider: str | None = None,
|
||||
) -> Response | StreamingResponse:
|
||||
"""Stream response with metrics tracking and memory tool handling.
|
||||
|
||||
|
|
@ -1045,6 +1048,7 @@ class StreamingMixin:
|
|||
await self._finalize_stream_response(
|
||||
body=body,
|
||||
provider=provider,
|
||||
outcome_provider=outcome_provider,
|
||||
model=model,
|
||||
request_id=request_id,
|
||||
original_tokens=original_tokens,
|
||||
|
|
@ -1306,6 +1310,7 @@ class StreamingMixin:
|
|||
await self._finalize_stream_response(
|
||||
body=body,
|
||||
provider=provider,
|
||||
outcome_provider=outcome_provider,
|
||||
model=model,
|
||||
request_id=request_id,
|
||||
original_tokens=original_tokens,
|
||||
|
|
|
|||
|
|
@ -96,6 +96,7 @@ class ProxyConfig:
|
|||
openai_api_url: str | None = None # Custom OpenAI API URL override
|
||||
gemini_api_url: str | None = None # Custom Gemini API URL override
|
||||
cloudcode_api_url: str | None = None # Custom Cloud Code Assist API URL override
|
||||
vertex_api_url: str | None = None # Custom Vertex AI regional API URL override
|
||||
|
||||
# Backend: "anthropic" (direct API), "litellm-*" (via LiteLLM), or "anyllm" (via any-llm)
|
||||
backend: str = "anthropic"
|
||||
|
|
@ -323,4 +324,5 @@ class ProxyConfig:
|
|||
openai=self.openai_api_url,
|
||||
gemini=self.gemini_api_url,
|
||||
cloudcode=self.cloudcode_api_url,
|
||||
vertex=self.vertex_api_url,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -222,6 +222,29 @@ class RequestOutcome:
|
|||
"""
|
||||
from headroom.proxy.helpers import compute_turn_id
|
||||
|
||||
request_items = body.get("messages")
|
||||
turn_messages = request_items
|
||||
if request_items is None:
|
||||
request_items = body.get("contents", [])
|
||||
if isinstance(request_items, list):
|
||||
turn_messages = []
|
||||
for item in request_items:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
parts = item.get("parts")
|
||||
text = ""
|
||||
if isinstance(parts, list):
|
||||
text = "\n".join(
|
||||
str(part.get("text"))
|
||||
for part in parts
|
||||
if isinstance(part, dict) and part.get("text")
|
||||
)
|
||||
role = "assistant" if item.get("role") == "model" else "user"
|
||||
turn_messages.append({"role": role, "content": text})
|
||||
system = body.get("system")
|
||||
if system is None:
|
||||
system = body.get("systemInstruction")
|
||||
|
||||
return cls(
|
||||
request_id=request_id,
|
||||
provider=provider,
|
||||
|
|
@ -243,11 +266,11 @@ class RequestOutcome:
|
|||
pipeline_timing=pipeline_timing,
|
||||
transforms_applied=tuple(transforms_applied),
|
||||
waste_signals=waste_signals,
|
||||
num_messages=len(body.get("messages", [])),
|
||||
turn_id=compute_turn_id(model, body.get("system"), body.get("messages")),
|
||||
num_messages=len(request_items) if isinstance(request_items, list) else 0,
|
||||
turn_id=compute_turn_id(model, system, turn_messages),
|
||||
tags=tags or {},
|
||||
client=client,
|
||||
request_messages=body.get("messages") if log_full_messages else None,
|
||||
request_messages=request_items if log_full_messages else None,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -96,6 +96,7 @@ from headroom.providers.registry import (
|
|||
DEFAULT_CLOUDCODE_API_URL,
|
||||
DEFAULT_GEMINI_API_URL,
|
||||
DEFAULT_OPENAI_API_URL,
|
||||
DEFAULT_VERTEX_API_URL,
|
||||
build_proxy_provider_runtime,
|
||||
create_proxy_backend,
|
||||
format_backend_status,
|
||||
|
|
@ -302,6 +303,7 @@ class HeadroomProxy(
|
|||
OPENAI_API_URL = DEFAULT_OPENAI_API_URL
|
||||
GEMINI_API_URL = DEFAULT_GEMINI_API_URL
|
||||
CLOUDCODE_API_URL = DEFAULT_CLOUDCODE_API_URL
|
||||
VERTEX_API_URL = DEFAULT_VERTEX_API_URL
|
||||
|
||||
def __init__(self, config: ProxyConfig):
|
||||
self.config = config
|
||||
|
|
@ -321,6 +323,7 @@ class HeadroomProxy(
|
|||
HeadroomProxy.OPENAI_API_URL = api_targets.openai
|
||||
HeadroomProxy.GEMINI_API_URL = api_targets.gemini
|
||||
HeadroomProxy.CLOUDCODE_API_URL = api_targets.cloudcode
|
||||
HeadroomProxy.VERTEX_API_URL = api_targets.vertex
|
||||
self.anthropic_provider = self.provider_runtime.pipeline_provider("anthropic")
|
||||
self.openai_provider = self.provider_runtime.pipeline_provider("openai")
|
||||
|
||||
|
|
@ -2973,6 +2976,7 @@ def _proxy_config_from_env() -> ProxyConfig:
|
|||
port=_get_env_int("HEADROOM_PORT", 8787),
|
||||
openai_api_url=os.environ.get("OPENAI_TARGET_API_URL"),
|
||||
anthropic_api_url=os.environ.get("ANTHROPIC_TARGET_API_URL"),
|
||||
vertex_api_url=os.environ.get("VERTEX_TARGET_API_URL"),
|
||||
backend=_get_env_str("HEADROOM_BACKEND", "anthropic"),
|
||||
bedrock_region=_get_env_str("HEADROOM_BEDROCK_REGION", "us-west-2"),
|
||||
bedrock_profile=os.environ.get("AWS_PROFILE"),
|
||||
|
|
@ -3054,6 +3058,7 @@ def run_server(
|
|||
║ OpenAI: {api_targets.openai:<57}║
|
||||
║ Gemini: {api_targets.gemini:<57}║
|
||||
║ Cloud Code: {api_targets.cloudcode:<57}║
|
||||
║ Vertex AI: {api_targets.vertex:<57}║
|
||||
╠══════════════════════════════════════════════════════════════════════╣
|
||||
║ FEATURES: ║
|
||||
║ Optimization: {"ENABLED " if config.optimize else "DISABLED"} ║
|
||||
|
|
@ -3240,6 +3245,10 @@ if __name__ == "__main__":
|
|||
"--anthropic-api-url",
|
||||
help=f"Custom Anthropic API URL (default: {DEFAULT_ANTHROPIC_API_URL})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vertex-api-url",
|
||||
help=f"Custom Vertex AI regional API URL (default: {DEFAULT_VERTEX_API_URL})",
|
||||
)
|
||||
|
||||
# Backend (anthropic direct, bedrock, openrouter, anyllm, or litellm-<provider>)
|
||||
parser.add_argument(
|
||||
|
|
@ -3391,6 +3400,7 @@ if __name__ == "__main__":
|
|||
port=_get_env_int("HEADROOM_PORT", args.port),
|
||||
openai_api_url=_get_env_str("OPENAI_TARGET_API_URL", args.openai_api_url),
|
||||
anthropic_api_url=_get_env_str("ANTHROPIC_TARGET_API_URL", args.anthropic_api_url),
|
||||
vertex_api_url=_get_env_str("VERTEX_TARGET_API_URL", args.vertex_api_url),
|
||||
# Backend settings
|
||||
backend=_get_env_str("HEADROOM_BACKEND", args.backend), # type: ignore[arg-type]
|
||||
bedrock_region=_get_env_str("HEADROOM_BEDROCK_REGION", args.bedrock_region),
|
||||
|
|
|
|||
|
|
@ -100,6 +100,7 @@ nav:
|
|||
- Failure Learning: learn.md
|
||||
- Integrations:
|
||||
- Overview: integration-guide.md
|
||||
- Vertex AI: vertex.md
|
||||
- LangChain: langchain.md
|
||||
- Agno: agno.md
|
||||
- Strands: strands.md
|
||||
|
|
|
|||
|
|
@ -18,7 +18,10 @@ pytest.importorskip("fastapi")
|
|||
from headroom.providers.claude import DEFAULT_API_URL as DEFAULT_ANTHROPIC_API_URL # noqa: E402
|
||||
from headroom.providers.codex import DEFAULT_API_URL as DEFAULT_OPENAI_API_URL # noqa: E402
|
||||
from headroom.providers.gemini import DEFAULT_API_URL as DEFAULT_GEMINI_API_URL # noqa: E402
|
||||
from headroom.providers.registry import DEFAULT_CLOUDCODE_API_URL # noqa: E402
|
||||
from headroom.providers.registry import ( # noqa: E402
|
||||
DEFAULT_CLOUDCODE_API_URL,
|
||||
DEFAULT_VERTEX_API_URL,
|
||||
)
|
||||
from headroom.proxy.models import ProxyConfig # noqa: E402
|
||||
from headroom.proxy.server import run_server # noqa: E402
|
||||
|
||||
|
|
@ -48,6 +51,7 @@ class TestBannerUpstreamTargets:
|
|||
assert DEFAULT_OPENAI_API_URL in output
|
||||
assert DEFAULT_GEMINI_API_URL in output
|
||||
assert DEFAULT_CLOUDCODE_API_URL in output
|
||||
assert DEFAULT_VERTEX_API_URL in output
|
||||
|
||||
def test_custom_anthropic_target_in_banner(self):
|
||||
"""A custom Anthropic API URL should be resolved and shown in the banner."""
|
||||
|
|
@ -81,6 +85,13 @@ class TestBannerUpstreamTargets:
|
|||
|
||||
assert "https://custom-cloudcode.example.com" in output
|
||||
|
||||
def test_custom_vertex_target_in_banner(self):
|
||||
"""A custom Vertex AI API URL should be resolved and shown in the banner."""
|
||||
config = ProxyConfig(vertex_api_url="https://europe-west4-aiplatform.googleapis.com")
|
||||
output = self._capture_banner(config)
|
||||
|
||||
assert "https://europe-west4-aiplatform.googleapis.com" in output
|
||||
|
||||
def test_multiple_custom_targets_in_banner(self):
|
||||
"""Multiple custom targets should all appear correctly in the banner."""
|
||||
config = ProxyConfig(
|
||||
|
|
@ -88,6 +99,7 @@ class TestBannerUpstreamTargets:
|
|||
openai_api_url="https://openai.internal",
|
||||
gemini_api_url="https://gemini.internal",
|
||||
cloudcode_api_url="https://cloudcode.internal",
|
||||
vertex_api_url="https://vertex.internal",
|
||||
)
|
||||
output = self._capture_banner(config)
|
||||
|
||||
|
|
@ -95,6 +107,7 @@ class TestBannerUpstreamTargets:
|
|||
assert "https://openai.internal" in output
|
||||
assert "https://gemini.internal" in output
|
||||
assert "https://cloudcode.internal" in output
|
||||
assert "https://vertex.internal" in output
|
||||
|
||||
def test_banner_suppressed_when_disabled(self):
|
||||
"""When print_banner=False, upstream targets should NOT be printed."""
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""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`
|
||||
1. Provider target URL env vars are read by `headroom proxy`
|
||||
2. litellm-* backends are accepted by both CLI and argparse paths
|
||||
"""
|
||||
|
||||
|
|
@ -187,6 +187,27 @@ class TestCLIProxyEnvVars:
|
|||
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 = {}
|
||||
|
|
@ -204,6 +225,25 @@ class TestCLIProxyEnvVars:
|
|||
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 = {}
|
||||
|
|
|
|||
|
|
@ -21,12 +21,16 @@ def _app() -> Any:
|
|||
openai_api_url="https://api.openai.test",
|
||||
gemini_api_url="https://api.gemini.test",
|
||||
cloudcode_api_url="https://cloudcode.test",
|
||||
vertex_api_url="https://vertex.test",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_provider_passthrough_routes_forward_expected_targets(monkeypatch) -> None:
|
||||
calls: list[tuple[str, str, str, str]] = []
|
||||
gemini_calls: list[tuple[str, str, str, str]] = []
|
||||
gemini_count_calls: list[tuple[str, str, str, str]] = []
|
||||
anthropic_calls: list[tuple[str, str, str, str, bool]] = []
|
||||
|
||||
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))
|
||||
|
|
@ -40,7 +44,68 @@ def test_provider_passthrough_routes_forward_expected_targets(monkeypatch) -> No
|
|||
}
|
||||
)
|
||||
|
||||
async def fake_gemini_generate(
|
||||
self,
|
||||
request,
|
||||
model,
|
||||
upstream_base_url=None,
|
||||
provider_name="gemini",
|
||||
): # type: ignore[no-untyped-def]
|
||||
gemini_calls.append((request.url.path, model, upstream_base_url, provider_name))
|
||||
return JSONResponse(
|
||||
{
|
||||
"handler": "handle_gemini_generate_content",
|
||||
"path": request.url.path,
|
||||
"model": model,
|
||||
"upstream_base_url": upstream_base_url,
|
||||
"provider": provider_name,
|
||||
}
|
||||
)
|
||||
|
||||
async def fake_anthropic_messages(
|
||||
self,
|
||||
request,
|
||||
upstream_base_url=None,
|
||||
provider_name="anthropic",
|
||||
model_override=None,
|
||||
force_stream=False,
|
||||
): # type: ignore[no-untyped-def]
|
||||
anthropic_calls.append(
|
||||
(request.url.path, upstream_base_url, provider_name, model_override, force_stream)
|
||||
)
|
||||
return JSONResponse(
|
||||
{
|
||||
"handler": "handle_anthropic_messages",
|
||||
"path": request.url.path,
|
||||
"upstream_base_url": upstream_base_url,
|
||||
"provider": provider_name,
|
||||
"model": model_override,
|
||||
"force_stream": force_stream,
|
||||
}
|
||||
)
|
||||
|
||||
async def fake_gemini_count(
|
||||
self,
|
||||
request,
|
||||
model,
|
||||
upstream_base_url=None,
|
||||
provider_name="gemini",
|
||||
): # type: ignore[no-untyped-def]
|
||||
gemini_count_calls.append((request.url.path, model, upstream_base_url, provider_name))
|
||||
return JSONResponse(
|
||||
{
|
||||
"handler": "handle_gemini_count_tokens",
|
||||
"path": request.url.path,
|
||||
"model": model,
|
||||
"upstream_base_url": upstream_base_url,
|
||||
"provider": provider_name,
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setattr(HeadroomProxy, "handle_passthrough", fake_passthrough)
|
||||
monkeypatch.setattr(HeadroomProxy, "handle_gemini_generate_content", fake_gemini_generate)
|
||||
monkeypatch.setattr(HeadroomProxy, "handle_gemini_count_tokens", fake_gemini_count)
|
||||
monkeypatch.setattr(HeadroomProxy, "handle_anthropic_messages", fake_anthropic_messages)
|
||||
|
||||
with TestClient(_app()) as client:
|
||||
assert client.post("/v1/messages/count_tokens").json()["base_url"] == (
|
||||
|
|
@ -68,6 +133,34 @@ def test_provider_passthrough_routes_forward_expected_targets(monkeypatch) -> No
|
|||
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(
|
||||
"/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent"
|
||||
).json() == {
|
||||
"handler": "handle_gemini_generate_content",
|
||||
"path": "/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent",
|
||||
"model": "gemini-2.0-flash",
|
||||
"upstream_base_url": "https://vertex.test",
|
||||
"provider": "vertex:google",
|
||||
}
|
||||
assert client.post(
|
||||
"/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:countTokens"
|
||||
).json() == {
|
||||
"handler": "handle_gemini_count_tokens",
|
||||
"path": "/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:countTokens",
|
||||
"model": "gemini-2.0-flash",
|
||||
"upstream_base_url": "https://vertex.test",
|
||||
"provider": "vertex:google",
|
||||
}
|
||||
assert client.post(
|
||||
"/v1beta1/projects/p/locations/us-central1/publishers/anthropic/models/claude-3-5-sonnet@20240620:rawPredict"
|
||||
).json() == {
|
||||
"handler": "handle_anthropic_messages",
|
||||
"path": "/v1beta1/projects/p/locations/us-central1/publishers/anthropic/models/claude-3-5-sonnet@20240620:rawPredict",
|
||||
"upstream_base_url": "https://vertex.test",
|
||||
"provider": "vertex:anthropic",
|
||||
"model": "claude-3-5-sonnet@20240620",
|
||||
"force_stream": False,
|
||||
}
|
||||
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"
|
||||
|
|
@ -86,6 +179,9 @@ def test_provider_passthrough_routes_forward_expected_targets(monkeypatch) -> No
|
|||
] == ("https://api.gemini.test")
|
||||
|
||||
assert len(calls) >= 16
|
||||
assert len(gemini_calls) >= 1
|
||||
assert len(gemini_count_calls) >= 1
|
||||
assert len(anthropic_calls) >= 1
|
||||
|
||||
|
||||
def test_proxy_route_helpers_prefer_legacy_targets_and_gemini_passthrough() -> None:
|
||||
|
|
@ -97,6 +193,7 @@ def test_proxy_route_helpers_prefer_legacy_targets_and_gemini_passthrough() -> N
|
|||
"ANTHROPIC_API_URL": "https://legacy.anthropic.test",
|
||||
"OPENAI_API_URL": "https://legacy.openai.test",
|
||||
"GEMINI_API_URL": "https://legacy.gemini.test",
|
||||
"VERTEX_API_URL": "https://legacy.vertex.test",
|
||||
"provider_runtime": type(
|
||||
"Runtime",
|
||||
(),
|
||||
|
|
@ -109,6 +206,7 @@ def test_proxy_route_helpers_prefer_legacy_targets_and_gemini_passthrough() -> N
|
|||
)()
|
||||
|
||||
assert proxy_routes._api_target(proxy, "anthropic") == "https://legacy.anthropic.test"
|
||||
assert proxy_routes._api_target(proxy, "vertex") == "https://legacy.vertex.test"
|
||||
assert proxy_routes._select_passthrough_base_url(proxy, {"x-goog-api-key": "test"}) == (
|
||||
"https://legacy.gemini.test"
|
||||
)
|
||||
|
|
@ -152,6 +250,7 @@ def test_provider_specific_routes_delegate_to_expected_proxy_handlers(monkeypatc
|
|||
"handle_google_batch_create",
|
||||
"handle_google_batch_results",
|
||||
"handle_google_batch_passthrough",
|
||||
"handle_passthrough",
|
||||
):
|
||||
install(handler_name)
|
||||
|
||||
|
|
@ -190,6 +289,36 @@ def test_provider_specific_routes_delegate_to_expected_proxy_handlers(monkeypatc
|
|||
assert client.post("/v1beta/models/demo:countTokens").json()["handler"] == (
|
||||
"handle_gemini_count_tokens"
|
||||
)
|
||||
assert client.post(
|
||||
"/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:streamGenerateContent"
|
||||
).json() == {
|
||||
"handler": "handle_gemini_generate_content",
|
||||
"path": "/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:streamGenerateContent",
|
||||
"args": [
|
||||
"gemini-2.0-flash",
|
||||
"https://vertex.test",
|
||||
"vertex:google",
|
||||
],
|
||||
}
|
||||
assert client.post(
|
||||
"/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:countTokens"
|
||||
).json() == {
|
||||
"handler": "handle_gemini_count_tokens",
|
||||
"path": "/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:countTokens",
|
||||
"args": [
|
||||
"gemini-2.0-flash",
|
||||
"https://vertex.test",
|
||||
"vertex:google",
|
||||
],
|
||||
}
|
||||
assert client.post(
|
||||
"/v1beta1/projects/p/locations/us-central1/publishers/anthropic/models/claude-3-5-sonnet@20240620:streamRawPredict"
|
||||
).json()["args"] == [
|
||||
"https://vertex.test",
|
||||
"vertex:anthropic",
|
||||
"claude-3-5-sonnet@20240620",
|
||||
True,
|
||||
]
|
||||
assert client.post("/v1internal:streamGenerateContent").json()["handler"] == (
|
||||
"handle_google_cloudcode_stream"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -16,12 +16,14 @@ from headroom.proxy.models import ProxyConfig
|
|||
def test_resolve_api_overrides_prefers_explicit_values_over_environment(monkeypatch) -> None:
|
||||
monkeypatch.setenv("ANTHROPIC_TARGET_API_URL", "https://env.anthropic.example/v1")
|
||||
monkeypatch.setenv("OPENAI_TARGET_API_URL", "https://env.openai.example/v1")
|
||||
monkeypatch.setenv("VERTEX_TARGET_API_URL", "https://env-vertex-aiplatform.example/v1")
|
||||
|
||||
overrides = resolve_api_overrides(
|
||||
anthropic_api_url="https://cli.anthropic.example/v1",
|
||||
openai_api_url=None,
|
||||
gemini_api_url=None,
|
||||
cloudcode_api_url=None,
|
||||
vertex_api_url="https://cli-vertex-aiplatform.example/v1",
|
||||
)
|
||||
|
||||
assert overrides == ProviderApiOverrides(
|
||||
|
|
@ -29,6 +31,7 @@ def test_resolve_api_overrides_prefers_explicit_values_over_environment(monkeypa
|
|||
openai="https://env.openai.example/v1",
|
||||
gemini=None,
|
||||
cloudcode=None,
|
||||
vertex="https://cli-vertex-aiplatform.example/v1",
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -39,6 +42,7 @@ def test_resolve_api_targets_normalizes_trailing_v1() -> None:
|
|||
openai="https://openai.example/v1",
|
||||
gemini="https://gemini.example/v1",
|
||||
cloudcode="https://cloudcode.example/v1/",
|
||||
vertex="https://vertex.example/v1/",
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -46,6 +50,7 @@ def test_resolve_api_targets_normalizes_trailing_v1() -> None:
|
|||
assert targets.openai == "https://openai.example"
|
||||
assert targets.gemini == "https://gemini.example"
|
||||
assert targets.cloudcode == "https://cloudcode.example"
|
||||
assert targets.vertex == "https://vertex.example"
|
||||
|
||||
|
||||
def test_proxy_config_exposes_provider_api_overrides() -> None:
|
||||
|
|
@ -54,6 +59,7 @@ def test_proxy_config_exposes_provider_api_overrides() -> None:
|
|||
openai_api_url="https://openai.example",
|
||||
gemini_api_url=None,
|
||||
cloudcode_api_url="https://cloudcode.example",
|
||||
vertex_api_url="https://vertex.example",
|
||||
)
|
||||
|
||||
assert config.provider_api_overrides == ProviderApiOverrides(
|
||||
|
|
@ -61,6 +67,7 @@ def test_proxy_config_exposes_provider_api_overrides() -> None:
|
|||
openai="https://openai.example",
|
||||
gemini=None,
|
||||
cloudcode="https://cloudcode.example",
|
||||
vertex="https://vertex.example",
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -8,9 +8,14 @@ from types import SimpleNamespace
|
|||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin
|
||||
from headroom.proxy.handlers.openai import OpenAIHandlerMixin, _decode_openai_bearer_payload
|
||||
from headroom.proxy.handlers.openai import (
|
||||
OpenAIHandlerMixin,
|
||||
_decode_openai_bearer_payload,
|
||||
_passthrough_usage_from_json,
|
||||
)
|
||||
from headroom.proxy.helpers import _headroom_bypass_enabled
|
||||
from headroom.proxy.server import HeadroomProxy
|
||||
|
||||
|
|
@ -55,6 +60,110 @@ class _PassthroughRequest:
|
|||
return b""
|
||||
|
||||
|
||||
class _VertexPassthroughRequest:
|
||||
method = "POST"
|
||||
headers = {}
|
||||
url = SimpleNamespace(
|
||||
path="/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent",
|
||||
query="",
|
||||
)
|
||||
|
||||
async def body(self) -> bytes:
|
||||
return b'{"contents":[]}'
|
||||
|
||||
|
||||
class _VertexStreamPassthroughRequest:
|
||||
method = "POST"
|
||||
headers = {}
|
||||
url = SimpleNamespace(
|
||||
path="/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:streamGenerateContent",
|
||||
query="alt=sse",
|
||||
)
|
||||
|
||||
async def body(self) -> bytes:
|
||||
return b'{"contents":[]}'
|
||||
|
||||
|
||||
class _VertexGeminiImageRequest:
|
||||
method = "POST"
|
||||
headers = {}
|
||||
query_params = {}
|
||||
url = SimpleNamespace(
|
||||
path="/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent",
|
||||
query="",
|
||||
)
|
||||
|
||||
async def body(self) -> bytes:
|
||||
return json.dumps(
|
||||
{
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{
|
||||
"inlineData": {
|
||||
"mimeType": "image/png",
|
||||
"data": "aW1hZ2U=",
|
||||
}
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
).encode("utf-8")
|
||||
|
||||
|
||||
class _VertexUsageClient:
|
||||
async def request(self, **kwargs): # noqa: ANN001, ANN201
|
||||
request = httpx.Request(kwargs["method"], kwargs["url"], content=kwargs["content"])
|
||||
return httpx.Response(
|
||||
200,
|
||||
request=request,
|
||||
headers={"content-type": "application/json"},
|
||||
json={
|
||||
"candidates": [{"content": {"parts": [{"text": "ok"}]}}],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 11,
|
||||
"candidatesTokenCount": 7,
|
||||
"cachedContentTokenCount": 3,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class _AsyncChunks(httpx.AsyncByteStream):
|
||||
def __init__(self, chunks: list[bytes]) -> None:
|
||||
self._chunks = chunks
|
||||
|
||||
async def __aiter__(self): # noqa: ANN204
|
||||
for chunk in self._chunks:
|
||||
yield chunk
|
||||
|
||||
|
||||
class _VertexStreamClient:
|
||||
def __init__(self) -> None:
|
||||
self.sent_url = ""
|
||||
|
||||
def build_request(self, method, url, headers, content): # noqa: ANN001, ANN201
|
||||
self.sent_url = str(url)
|
||||
return httpx.Request(method, url, headers=headers, content=content)
|
||||
|
||||
async def send(self, request, stream=False): # noqa: ANN001, ANN201
|
||||
assert stream is True
|
||||
return httpx.Response(
|
||||
200,
|
||||
request=request,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
stream=_AsyncChunks(
|
||||
[
|
||||
b'data: {"candidates":[{"content":{"parts":[{"text":"hello"}]}}]}\n\n',
|
||||
b'data: {"usageMetadata":{"promptTokenCount":13,'
|
||||
b'"candidatesTokenCount":5,"cachedContentTokenCount":2}}\n\n',
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class _RetryThenSuccessClient:
|
||||
def __init__(self) -> None:
|
||||
self.attempts = 0
|
||||
|
|
@ -140,6 +249,207 @@ def test_openai_passthrough_connect_timeout_returns_502() -> None:
|
|||
assert "Failed to connect to upstream API" in payload["error"]["message"]
|
||||
|
||||
|
||||
def test_passthrough_usage_normalizes_vertex_usage_metadata() -> None:
|
||||
usage = _passthrough_usage_from_json(
|
||||
{
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 11,
|
||||
"candidatesTokenCount": 7,
|
||||
"cachedContentTokenCount": 3,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
assert usage == {
|
||||
"input_tokens": 11,
|
||||
"output_tokens": 7,
|
||||
"cache_read_input_tokens": 3,
|
||||
}
|
||||
|
||||
|
||||
def test_vertex_passthrough_records_usage_metadata_for_dashboard() -> None:
|
||||
handler = object.__new__(HeadroomProxy)
|
||||
handler.http_client = _VertexUsageClient()
|
||||
outcomes = []
|
||||
|
||||
async def next_request_id(): # noqa: ANN202
|
||||
return "req_vertex"
|
||||
|
||||
async def record(outcome): # noqa: ANN001, ANN202
|
||||
outcomes.append(outcome)
|
||||
|
||||
handler._next_request_id = next_request_id
|
||||
handler._record_request_outcome = record
|
||||
|
||||
response = asyncio.run(
|
||||
handler.handle_passthrough(
|
||||
_VertexPassthroughRequest(),
|
||||
"https://vertex.test",
|
||||
"generateContent",
|
||||
"vertex:google",
|
||||
)
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert len(outcomes) == 1
|
||||
outcome = outcomes[0]
|
||||
assert outcome.provider == "vertex:google"
|
||||
assert outcome.model == "gemini-2.0-flash"
|
||||
assert outcome.optimized_tokens == 11
|
||||
assert outcome.output_tokens == 7
|
||||
assert outcome.cache_read_tokens == 3
|
||||
|
||||
|
||||
def test_vertex_stream_passthrough_preserves_chunks_and_records_usage() -> None:
|
||||
handler = object.__new__(HeadroomProxy)
|
||||
handler.http_client = _VertexStreamClient()
|
||||
outcomes = []
|
||||
|
||||
async def next_request_id(): # noqa: ANN202
|
||||
return "req_vertex_stream"
|
||||
|
||||
async def record(outcome): # noqa: ANN001, ANN202
|
||||
outcomes.append(outcome)
|
||||
|
||||
handler._next_request_id = next_request_id
|
||||
handler._record_request_outcome = record
|
||||
|
||||
response = asyncio.run(
|
||||
handler.handle_passthrough(
|
||||
_VertexStreamPassthroughRequest(),
|
||||
"https://vertex.test",
|
||||
"streamGenerateContent",
|
||||
"vertex:google",
|
||||
)
|
||||
)
|
||||
|
||||
assert isinstance(response, StreamingResponse)
|
||||
|
||||
async def collect(): # noqa: ANN202
|
||||
return [chunk async for chunk in response.body_iterator]
|
||||
|
||||
chunks = asyncio.run(collect())
|
||||
|
||||
assert len(chunks) == 2
|
||||
assert chunks[0].startswith(b'data: {"candidates"')
|
||||
assert b'"usageMetadata"' in chunks[1]
|
||||
assert len(outcomes) == 1
|
||||
outcome = outcomes[0]
|
||||
assert outcome.provider == "vertex:google"
|
||||
assert outcome.model == "gemini-2.0-flash"
|
||||
assert outcome.optimized_tokens == 13
|
||||
assert outcome.output_tokens == 5
|
||||
assert outcome.cache_read_tokens == 2
|
||||
|
||||
|
||||
def test_stream_finalizer_records_vertex_provider_for_dashboard() -> None:
|
||||
handler = object.__new__(HeadroomProxy)
|
||||
handler.config = SimpleNamespace(log_full_messages=False)
|
||||
outcomes = []
|
||||
|
||||
async def record(outcome): # noqa: ANN001, ANN202
|
||||
outcomes.append(outcome)
|
||||
|
||||
handler._record_request_outcome = record
|
||||
|
||||
asyncio.run(
|
||||
handler._finalize_stream_response(
|
||||
body={"contents": [{"role": "user", "parts": [{"text": "hello"}]}]},
|
||||
provider="gemini",
|
||||
outcome_provider="vertex:google",
|
||||
model="gemini-2.0-flash",
|
||||
request_id="req_vertex_stream_final",
|
||||
original_tokens=20,
|
||||
optimized_tokens=12,
|
||||
tokens_saved=8,
|
||||
transforms_applied=["test-transform"],
|
||||
optimization_latency=3.0,
|
||||
stream_state={
|
||||
"input_tokens": 12,
|
||||
"output_tokens": 5,
|
||||
"cache_read_input_tokens": 2,
|
||||
"cache_creation_input_tokens": 0,
|
||||
"cache_creation_ephemeral_5m_input_tokens": 0,
|
||||
"cache_creation_ephemeral_1h_input_tokens": 0,
|
||||
"total_bytes": 100,
|
||||
"sse_buffer": bytearray(),
|
||||
"ttfb_ms": 4.0,
|
||||
},
|
||||
start_time=0.0,
|
||||
tags={"route": "vertex"},
|
||||
)
|
||||
)
|
||||
|
||||
assert len(outcomes) == 1
|
||||
outcome = outcomes[0]
|
||||
assert outcome.provider == "vertex:google"
|
||||
assert outcome.model == "gemini-2.0-flash"
|
||||
assert outcome.optimized_tokens == 12
|
||||
assert outcome.output_tokens == 5
|
||||
assert outcome.tokens_saved == 8
|
||||
assert outcome.cache_read_tokens == 2
|
||||
|
||||
|
||||
def test_vertex_gemini_non_text_generate_records_dashboard_outcome() -> None:
|
||||
handler = object.__new__(HeadroomProxy)
|
||||
handler.memory_handler = None
|
||||
handler.rate_limiter = None
|
||||
outcomes = []
|
||||
upstream_urls = []
|
||||
|
||||
async def next_request_id(): # noqa: ANN202
|
||||
return "req_vertex_image"
|
||||
|
||||
async def record(outcome): # noqa: ANN001, ANN202
|
||||
outcomes.append(outcome)
|
||||
|
||||
async def retry_request(method, url, headers, body): # noqa: ANN001, ANN202
|
||||
upstream_urls.append(url)
|
||||
request = httpx.Request(method, url, headers=headers)
|
||||
return httpx.Response(
|
||||
200,
|
||||
request=request,
|
||||
headers={"content-type": "application/json"},
|
||||
json={
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 31,
|
||||
"candidatesTokenCount": 4,
|
||||
"cachedContentTokenCount": 6,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
handler._next_request_id = next_request_id
|
||||
handler._record_request_outcome = record
|
||||
handler._retry_request = retry_request
|
||||
|
||||
response = asyncio.run(
|
||||
handler.handle_gemini_generate_content(
|
||||
_VertexGeminiImageRequest(),
|
||||
"gemini-2.0-flash",
|
||||
"https://vertex.test",
|
||||
"vertex:google",
|
||||
)
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert upstream_urls == [
|
||||
"https://vertex.test/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent"
|
||||
]
|
||||
assert response.headers["x-headroom-tokens-before"] == "31"
|
||||
assert response.headers["x-headroom-tokens-after"] == "31"
|
||||
assert response.headers["x-headroom-tokens-saved"] == "0"
|
||||
assert len(outcomes) == 1
|
||||
outcome = outcomes[0]
|
||||
assert outcome.provider == "vertex:google"
|
||||
assert outcome.model == "gemini-2.0-flash"
|
||||
assert outcome.original_tokens == 31
|
||||
assert outcome.optimized_tokens == 31
|
||||
assert outcome.output_tokens == 4
|
||||
assert outcome.cache_read_tokens == 6
|
||||
assert outcome.num_messages == 1
|
||||
|
||||
|
||||
def test_retry_request_retries_connect_timeout() -> None:
|
||||
proxy = object.__new__(HeadroomProxy)
|
||||
proxy.http_client = _RetryThenSuccessClient()
|
||||
|
|
|
|||
|
|
@ -116,6 +116,33 @@ def test_client_field_round_trips() -> None:
|
|||
assert o.client == "codex"
|
||||
|
||||
|
||||
def test_stream_outcome_derives_gemini_contents_metadata() -> None:
|
||||
outcome = RequestOutcome.from_stream(
|
||||
body={
|
||||
"systemInstruction": {"parts": [{"text": "sys"}]},
|
||||
"contents": [{"role": "user", "parts": [{"text": "hello"}]}],
|
||||
},
|
||||
provider="vertex:google",
|
||||
model="gemini-2.0-flash",
|
||||
request_id="req-gemini-stream",
|
||||
original_tokens=12,
|
||||
optimized_tokens=10,
|
||||
output_tokens=3,
|
||||
tokens_saved=2,
|
||||
transforms_applied=["compress"],
|
||||
total_latency_ms=25.0,
|
||||
overhead_ms=4.0,
|
||||
tags={"route": "vertex"},
|
||||
client="codex",
|
||||
log_full_messages=True,
|
||||
)
|
||||
|
||||
assert outcome.provider == "vertex:google"
|
||||
assert outcome.num_messages == 1
|
||||
assert outcome.request_messages == [{"role": "user", "parts": [{"text": "hello"}]}]
|
||||
assert outcome.turn_id is not None
|
||||
|
||||
|
||||
# ── classify_client — the harness ID source ─────────────────────────
|
||||
|
||||
|
||||
|
|
|
|||
67
wiki/vertex.md
Normal file
67
wiki/vertex.md
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
# Vertex AI
|
||||
|
||||
Headroom supports Google Cloud Vertex AI publisher endpoints through the proxy
|
||||
passthrough surface. Configure the proxy with a regional Vertex base URL, then
|
||||
send normal Vertex REST requests through Headroom.
|
||||
|
||||
Google documents Gemini generation on Vertex with `generateContent` and
|
||||
`streamGenerateContent`, and the request body uses the Vertex/Gemini `contents`
|
||||
shape. See Google Cloud's model inference reference:
|
||||
https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference
|
||||
|
||||
Google Cloud REST calls authenticate with a bearer access token. For local
|
||||
development, Google documents both `gcloud auth print-access-token` and
|
||||
`gcloud auth application-default print-access-token`; Application Default
|
||||
Credentials search `GOOGLE_APPLICATION_CREDENTIALS`, local ADC files, and
|
||||
attached service accounts in that order. See:
|
||||
|
||||
- https://docs.cloud.google.com/docs/authentication/rest
|
||||
- https://docs.cloud.google.com/docs/authentication/application-default-credentials
|
||||
|
||||
## Configure
|
||||
|
||||
Set the Vertex regional host explicitly:
|
||||
|
||||
```bash
|
||||
headroom proxy --vertex-api-url https://us-central1-aiplatform.googleapis.com
|
||||
```
|
||||
|
||||
The same setting is available through `VERTEX_TARGET_API_URL`.
|
||||
|
||||
## Gemini On Vertex
|
||||
|
||||
Send Vertex publisher paths through the proxy unchanged:
|
||||
|
||||
```bash
|
||||
ACCESS_TOKEN="$(gcloud auth print-access-token)"
|
||||
|
||||
curl -sS \
|
||||
-H "Authorization: Bearer ${ACCESS_TOKEN}" \
|
||||
-H "Content-Type: application/json" \
|
||||
http://127.0.0.1:8787/v1/projects/PROJECT_ID/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent \
|
||||
-d '{
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [{"text": "Summarize this repository in one paragraph."}]
|
||||
}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
Supported passthrough actions:
|
||||
|
||||
- `generateContent`
|
||||
- `streamGenerateContent`
|
||||
- `countTokens`
|
||||
|
||||
## Anthropic Publisher On Vertex
|
||||
|
||||
Headroom also forwards Anthropic publisher calls on Vertex:
|
||||
|
||||
- `rawPredict`
|
||||
- `streamRawPredict`
|
||||
|
||||
The Python proxy preserves caller-supplied Google bearer auth. The native Rust
|
||||
proxy path additionally resolves GCP ADC and injects the bearer token for the
|
||||
Anthropic publisher route.
|
||||
Loading…
Add table
Add a link
Reference in a new issue