From 3c77e52ce431210e6045671cf5f7c66c79f90a32 Mon Sep 17 00:00:00 2001 From: JD Davis Date: Wed, 10 Jun 2026 01:05:30 -0500 Subject: [PATCH] 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 --- headroom/cli/proxy.py | 10 + headroom/providers/proxy_routes.py | 115 ++++++++++ headroom/providers/registry.py | 7 + headroom/proxy/handlers/anthropic.py | 80 ++++--- headroom/proxy/handlers/gemini.py | 135 ++++++++--- headroom/proxy/handlers/openai.py | 268 +++++++++++++++++++++- headroom/proxy/handlers/streaming.py | 9 +- headroom/proxy/models.py | 2 + headroom/proxy/outcome.py | 29 ++- headroom/proxy/server.py | 10 + mkdocs.yml | 1 + tests/test_banner_upstream_targets.py | 15 +- tests/test_cli_proxy_env.py | 42 +++- tests/test_provider_proxy_routes.py | 129 +++++++++++ tests/test_provider_registry.py | 7 + tests/test_proxy_handler_helpers.py | 312 +++++++++++++++++++++++++- tests/test_request_outcome.py | 27 +++ wiki/vertex.md | 67 ++++++ 18 files changed, 1193 insertions(+), 72 deletions(-) create mode 100644 wiki/vertex.md diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index 361e8cc7b..34e395bac 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -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 diff --git a/headroom/providers/proxy_routes.py b/headroom/providers/proxy_routes.py index 9610633e8..afdb7a0a4 100644 --- a/headroom/providers/proxy_routes.py +++ b/headroom/providers/proxy_routes.py @@ -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( diff --git a/headroom/providers/registry.py b/headroom/providers/registry.py index b6de4a039..2300739db 100644 --- a/headroom/providers/registry.py +++ b/headroom/providers/registry.py @@ -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), ) diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index ffa0545c3..29547a9d4 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -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}") diff --git a/headroom/proxy/handlers/gemini.py b/headroom/proxy/handlers/gemini.py index 0867eafae..26a545f8d 100644 --- a/headroom/proxy/handlers/gemini.py +++ b/headroom/proxy/handlers/gemini.py @@ -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, diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index b401aed4b..80802a2b8 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -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, + ) diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index df8e7aa69..09b081052 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -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, diff --git a/headroom/proxy/models.py b/headroom/proxy/models.py index f5363fde6..246df1f25 100644 --- a/headroom/proxy/models.py +++ b/headroom/proxy/models.py @@ -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, ) diff --git a/headroom/proxy/outcome.py b/headroom/proxy/outcome.py index 74dbe1676..37f5d2643 100644 --- a/headroom/proxy/outcome.py +++ b/headroom/proxy/outcome.py @@ -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, ) diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index ca2c91d20..613fe6fb6 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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-) 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), diff --git a/mkdocs.yml b/mkdocs.yml index dec34fbf7..5d41f8b3c 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -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 diff --git a/tests/test_banner_upstream_targets.py b/tests/test_banner_upstream_targets.py index f7f5a0e7c..ad7fe44bc 100644 --- a/tests/test_banner_upstream_targets.py +++ b/tests/test_banner_upstream_targets.py @@ -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.""" diff --git a/tests/test_cli_proxy_env.py b/tests/test_cli_proxy_env.py index f3129a01c..7636c68f9 100644 --- a/tests/test_cli_proxy_env.py +++ b/tests/test_cli_proxy_env.py @@ -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 = {} diff --git a/tests/test_provider_proxy_routes.py b/tests/test_provider_proxy_routes.py index 50a9df4c6..d34a09827 100644 --- a/tests/test_provider_proxy_routes.py +++ b/tests/test_provider_proxy_routes.py @@ -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" ) diff --git a/tests/test_provider_registry.py b/tests/test_provider_registry.py index c9d782cd4..8c95499a2 100644 --- a/tests/test_provider_registry.py +++ b/tests/test_provider_registry.py @@ -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", ) diff --git a/tests/test_proxy_handler_helpers.py b/tests/test_proxy_handler_helpers.py index fba2fba1c..2028bee26 100644 --- a/tests/test_proxy_handler_helpers.py +++ b/tests/test_proxy_handler_helpers.py @@ -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() diff --git a/tests/test_request_outcome.py b/tests/test_request_outcome.py index 182d479d4..35ec7285c 100644 --- a/tests/test_request_outcome.py +++ b/tests/test_request_outcome.py @@ -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 ───────────────────────── diff --git a/wiki/vertex.md b/wiki/vertex.md new file mode 100644 index 000000000..5e5e2cbe2 --- /dev/null +++ b/wiki/vertex.md @@ -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.