mirror of
https://github.com/usestrix/strix.git
synced 2026-11-01 07:02:32 +00:00
feat(llm): send app attribution headers to Vercel AI Gateway
This commit is contained in:
parent
ffa35e22d3
commit
98e3df1425
3 changed files with 28 additions and 17 deletions
|
|
@ -713,7 +713,7 @@ def configure_sdk_model_defaults(settings: Settings) -> None:
|
|||
if codex.subscription_model(llm.model):
|
||||
return
|
||||
_configure_litellm_compatibility()
|
||||
_configure_openrouter_attribution(llm.model)
|
||||
_configure_attribution(llm.model)
|
||||
if llm.api_key:
|
||||
set_default_openai_key(llm.api_key, use_for_tracing=False)
|
||||
_configure_litellm_default("api_key", llm.api_key)
|
||||
|
|
@ -884,9 +884,13 @@ def _openrouter_error_type(error: object) -> object:
|
|||
return None
|
||||
|
||||
|
||||
OPENROUTER_ATTRIBUTION_HEADERS = {
|
||||
APP_ATTRIBUTION_HEADERS = {
|
||||
"HTTP-Referer": "https://strix.ai",
|
||||
"X-Title": "Strix",
|
||||
}
|
||||
|
||||
OPENROUTER_ATTRIBUTION_HEADERS = {
|
||||
**APP_ATTRIBUTION_HEADERS,
|
||||
"X-OpenRouter-Categories": "cli-agent",
|
||||
}
|
||||
|
||||
|
|
@ -895,20 +899,25 @@ def is_openrouter_model(model_name: str | None) -> bool:
|
|||
return bool(model_name) and "openrouter/" in (model_name or "").strip().lower()
|
||||
|
||||
|
||||
def _configure_openrouter_attribution(model_name: str | None) -> None:
|
||||
def is_vercel_ai_gateway_model(model_name: str | None) -> bool:
|
||||
return (model_name or "").strip().lower().startswith("vercel_ai_gateway/")
|
||||
|
||||
|
||||
def attribution_headers(model_name: str | None) -> dict[str, str]:
|
||||
if is_openrouter_model(model_name):
|
||||
return OPENROUTER_ATTRIBUTION_HEADERS
|
||||
if is_vercel_ai_gateway_model(model_name):
|
||||
return APP_ATTRIBUTION_HEADERS
|
||||
return {}
|
||||
|
||||
|
||||
def _configure_attribution(model_name: str | None) -> None:
|
||||
import litellm
|
||||
|
||||
current: object = litellm.headers
|
||||
existing: dict[str, str] = current if isinstance(current, dict) else {}
|
||||
if not is_openrouter_model(model_name):
|
||||
if any(key in existing for key in OPENROUTER_ATTRIBUTION_HEADERS):
|
||||
remaining = {
|
||||
k: v for k, v in existing.items() if k not in OPENROUTER_ATTRIBUTION_HEADERS
|
||||
}
|
||||
litellm.headers = remaining or None # type: ignore[assignment]
|
||||
return
|
||||
|
||||
litellm.headers = {**existing, **OPENROUTER_ATTRIBUTION_HEADERS} # type: ignore[assignment]
|
||||
remaining = {k: v for k, v in existing.items() if k not in OPENROUTER_ATTRIBUTION_HEADERS}
|
||||
litellm.headers = {**remaining, **attribution_headers(model_name)} or None # type: ignore[assignment]
|
||||
|
||||
|
||||
def _configure_extra_headers(llm: LlmSettings) -> None:
|
||||
|
|
|
|||
|
|
@ -10,12 +10,11 @@ from openai.types.shared import Reasoning
|
|||
|
||||
from strix.config.models import (
|
||||
DEFAULT_MODEL_RETRY,
|
||||
OPENROUTER_ATTRIBUTION_HEADERS,
|
||||
attribution_headers,
|
||||
bedrock_route_supports_prompt_caching,
|
||||
is_bedrock_route,
|
||||
is_claude_model,
|
||||
is_known_openai_bare_model,
|
||||
is_openrouter_model,
|
||||
model_supports_reasoning,
|
||||
request_timeout_extra_args,
|
||||
routes_through_litellm,
|
||||
|
|
@ -287,9 +286,7 @@ def make_model_settings(
|
|||
def _request_headers(
|
||||
model_name: str, extra_headers: dict[str, str] | None
|
||||
) -> dict[str, str] | None:
|
||||
headers: dict[str, str] = {}
|
||||
if is_openrouter_model(model_name):
|
||||
headers.update(OPENROUTER_ATTRIBUTION_HEADERS)
|
||||
headers: dict[str, str] = dict(attribution_headers(model_name))
|
||||
if extra_headers:
|
||||
headers.update(extra_headers)
|
||||
return headers or None
|
||||
|
|
|
|||
|
|
@ -474,6 +474,11 @@ def test_openrouter_attribution_rides_on_the_request_headers() -> None:
|
|||
}
|
||||
|
||||
|
||||
def test_vercel_ai_gateway_attribution_rides_on_the_request_headers() -> None:
|
||||
headers = make_model_settings(None, model_name="vercel_ai_gateway/zai/glm-5.3").extra_headers
|
||||
assert headers == {"HTTP-Referer": "https://strix.ai", "X-Title": "Strix"}
|
||||
|
||||
|
||||
def test_openrouter_attribution_absent_for_other_providers() -> None:
|
||||
assert make_model_settings(None, model_name="anthropic/claude-sonnet-4-5").extra_headers is None
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue