mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-26 06:03:19 +00:00
fix(prometheus): keep caller-sent provider credentials and token-exchange 401s out of deployment metrics
A caller can send its own provider credential in the request body, such as azure_ad_token, azure_scope, the Azure Entra app fields, watsonx token, gigachat_access_token, OCI keys, AWS keys or Vertex credentials, at the top level or inside extra_body. A provider 401 on that credential is the caller's failure, so it stays out of the shared deployment's failure, request, LLM-failed and state series, as on main. A token-exchange deployment trades the caller's identity token on every call, so its 401s stay out as well, matching the router cooldown.
This commit is contained in:
yucheng-berri
parent
5b7d65bee6
commit
84fc741b9c
3 changed files with 323 additions and 4 deletions
|
|
@ -338,15 +338,57 @@ def _logged_fields(value: object) -> Mapping[str, object]:
|
|||
return MappingProxyType({})
|
||||
|
||||
|
||||
_CALLER_CREDENTIAL_FIELDS: Final = frozenset(
|
||||
{
|
||||
"aws_access_key_id",
|
||||
"aws_secret_access_key",
|
||||
"aws_session_token",
|
||||
"aws_web_identity_token",
|
||||
"azure_ad_token",
|
||||
"azure_password",
|
||||
"azure_scope",
|
||||
"azure_username",
|
||||
"client_id",
|
||||
"client_secret",
|
||||
"gigachat_access_token",
|
||||
"oci_fingerprint",
|
||||
"oci_key",
|
||||
"oci_key_file",
|
||||
"oci_tenancy",
|
||||
"oci_user",
|
||||
"tenant_id",
|
||||
"token",
|
||||
"user_config",
|
||||
"vertex_ai_credentials",
|
||||
"vertex_credentials",
|
||||
"watsonx_credentials",
|
||||
"wx_credentials",
|
||||
"zen_api_key",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _sent_a_provider_credential(caller_body: Mapping[str, object]) -> bool:
|
||||
"""True when the caller's request body, or its ``extra_body``, carries a provider credential field."""
|
||||
return not (
|
||||
_CALLER_CREDENTIAL_FIELDS.isdisjoint(caller_body)
|
||||
and _CALLER_CREDENTIAL_FIELDS.isdisjoint(_logged_fields(caller_body.get("extra_body")))
|
||||
)
|
||||
|
||||
|
||||
def _used_a_caller_credential(logged_call: Mapping[str, object]) -> bool:
|
||||
"""True when the provider was called with a credential that belongs to the caller rather than the
|
||||
deployment: an api_key the caller sent with the request, the caller's own Anthropic OAuth token, or the
|
||||
caller's own GitHub Copilot connection."""
|
||||
deployment: an api_key or other provider credential field the caller sent with the request, the caller's own
|
||||
Anthropic OAuth token, the caller's own GitHub Copilot connection, or the caller's identity token that a
|
||||
token-exchange deployment trades in on every call."""
|
||||
litellm_params: Final = _logged_fields(logged_call.get("litellm_params"))
|
||||
credential_fields: Final = _logged_fields(litellm_params.get("proxy_server_request")).get("credential_fields")
|
||||
proxy_server_request: Final = _logged_fields(litellm_params.get("proxy_server_request"))
|
||||
credential_fields: Final = proxy_server_request.get("credential_fields")
|
||||
logged_metadata: Final = _logged_fields(_logged_fields(logged_call.get("standard_logging_object")).get("metadata"))
|
||||
return (
|
||||
(isinstance(credential_fields, tuple) and "api_key" in credential_fields)
|
||||
or _sent_a_provider_credential(_logged_fields(proxy_server_request.get("body")))
|
||||
or bool(litellm_params.get("token_exchange_endpoint"))
|
||||
or logged_metadata.get("used_client_oauth_token") is True
|
||||
or github_copilot_user_session_from(litellm_params) is not None
|
||||
)
|
||||
|
|
|
|||
|
|
@ -32,6 +32,24 @@ RECOVERS: Final = "provider-rejects-once"
|
|||
HEALTHY: Final = "provider-healthy"
|
||||
AGENT_365_GUARDRAIL: Final = f"agent-365-{RUN}"
|
||||
CALLER_OAUTH_TOKEN: Final = f"sk-ant-oat01-caller-{RUN}"
|
||||
AZURE_API_VERSION: Final = "2024-02-01"
|
||||
CALLER_AZURE_CELLS: Final = (
|
||||
("chat", False),
|
||||
("embeddings", False),
|
||||
("responses", False),
|
||||
("messages", False),
|
||||
("responses", True),
|
||||
)
|
||||
CALLER_CREDENTIAL_CELLS: Final = (
|
||||
("caller-watsonx-token", "token"),
|
||||
("caller-gigachat-access-token", "gigachat_access_token"),
|
||||
)
|
||||
TOKEN_EXCHANGE_CELLS: Final = ("token-exchange", "token-exchange-credential")
|
||||
TOKEN_EXCHANGE_SETTINGS: Final = {
|
||||
"token_exchange_endpoint": "https://login.microsoftonline.com/00000000-0000-4000-8000-0000000c0b01/oauth2/v2.0/token",
|
||||
"client_id": "synthetic-obo-client-id",
|
||||
"client_secret": "synthetic-obo-client-secret",
|
||||
}
|
||||
|
||||
FAILURES: Final = "litellm_deployment_failure_responses_total"
|
||||
TOTALS: Final = "litellm_deployment_total_requests_total"
|
||||
|
|
@ -106,7 +124,8 @@ class Provider:
|
|||
self._release: Final = release
|
||||
|
||||
def __call__(self, request: Request) -> Reply:
|
||||
model: Final = str(json.loads(request.body)["model"])
|
||||
sent: Final = json.loads(request.body)
|
||||
model: Final = str(sent.get("model") or sent["model_id"])
|
||||
with self._lock:
|
||||
self._seen[model] += 1
|
||||
attempt: Final = self._seen[model]
|
||||
|
|
@ -145,6 +164,49 @@ def anthropic_deployment(cell: str, api_base: str) -> dict[str, JsonValue]:
|
|||
}
|
||||
|
||||
|
||||
def azure_deployment(cell: str, api_base: str, credential: dict[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_name": name(cell),
|
||||
"litellm_params": {
|
||||
"model": f"azure/{REJECTS}",
|
||||
"api_base": api_base,
|
||||
"api_version": AZURE_API_VERSION,
|
||||
**credential,
|
||||
},
|
||||
"model_info": {"id": identity(cell)},
|
||||
}
|
||||
|
||||
|
||||
def caller_credential_deployment(cell: str, field: str, api_base: str) -> dict[str, JsonValue]:
|
||||
provider: Final = "watsonx" if field == "token" else "gigachat"
|
||||
project: Final[dict[str, JsonValue]] = {"project_id": f"project-{RUN}"} if provider == "watsonx" else {}
|
||||
return {
|
||||
"model_name": name(cell),
|
||||
"litellm_params": {
|
||||
"model": f"{provider}/{REJECTS}",
|
||||
"api_base": api_base,
|
||||
field: f"deployment-credential-{cell}",
|
||||
**project,
|
||||
},
|
||||
"model_info": {"id": identity(cell)},
|
||||
}
|
||||
|
||||
|
||||
def token_exchange_deployment(cell: str) -> dict[str, JsonValue]:
|
||||
settings: Final[dict[str, JsonValue]] = (
|
||||
{"litellm_credential_name": f"m365-obo-{RUN}"} if cell.endswith("-credential") else {**TOKEN_EXCHANGE_SETTINGS}
|
||||
)
|
||||
return {
|
||||
"model_name": name(cell),
|
||||
"litellm_params": {"model": "microsoft_365_copilot/chat", **settings},
|
||||
"model_info": {"id": identity(cell)},
|
||||
}
|
||||
|
||||
|
||||
def caller_azure_cell(opt_in: str, route: str, inside_extra_body: bool) -> str:
|
||||
return f"{opt_in}-caller-azure-token-{route}{'-extra-body' if inside_extra_body else ''}"
|
||||
|
||||
|
||||
def config(directory: Path, upstream: Wire) -> Path:
|
||||
api_base: Final = f"{upstream.url}/v1"
|
||||
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
|
|
@ -158,6 +220,17 @@ def config(directory: Path, upstream: Wire) -> Path:
|
|||
deployment("recovers", RECOVERS, api_base),
|
||||
anthropic_deployment("anthropic-own-key", upstream.url),
|
||||
anthropic_deployment("caller-oauth", upstream.url),
|
||||
azure_deployment("azure-own-token", upstream.url, {"azure_ad_token": "deployment-entra-token-azure-own-token"}),
|
||||
*(
|
||||
azure_deployment(
|
||||
caller_azure_cell("deployment", route, inside_extra_body),
|
||||
upstream.url,
|
||||
{"configurable_clientside_auth_params": ["azure_ad_token"]},
|
||||
)
|
||||
for route, inside_extra_body in CALLER_AZURE_CELLS
|
||||
),
|
||||
*(caller_credential_deployment(cell, field, upstream.url) for cell, field in CALLER_CREDENTIAL_CELLS),
|
||||
*(token_exchange_deployment(cell) for cell in TOKEN_EXCHANGE_CELLS),
|
||||
{
|
||||
"model_name": name("caller-credential"),
|
||||
"litellm_params": {
|
||||
|
|
@ -168,6 +241,10 @@ def config(directory: Path, upstream: Wire) -> Path:
|
|||
"model_info": {"id": identity("caller-credential")},
|
||||
},
|
||||
]
|
||||
configuration["credential_list"] = [
|
||||
*configuration.get("credential_list", []),
|
||||
{"credential_name": f"m365-obo-{RUN}", "credential_values": {**TOKEN_EXCHANGE_SETTINGS}, "credential_info": {}},
|
||||
]
|
||||
configuration["guardrails"] = [
|
||||
{
|
||||
"guardrail_name": AGENT_365_GUARDRAIL,
|
||||
|
|
@ -213,6 +290,36 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
|||
yield Rig(owned, upstream)
|
||||
|
||||
|
||||
def opted_in_config(directory: Path, upstream: Wire) -> Path:
|
||||
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
configuration["model_list"] = [
|
||||
*configuration["model_list"],
|
||||
*(
|
||||
azure_deployment(caller_azure_cell("proxy-wide", route, inside_extra_body), upstream.url, {})
|
||||
for route, inside_extra_body in CALLER_AZURE_CELLS
|
||||
),
|
||||
]
|
||||
configuration["general_settings"]["allow_client_side_credentials"] = True
|
||||
configuration["litellm_settings"]["callbacks"] = ["prometheus"]
|
||||
configuration.setdefault("router_settings", {}).update({"num_retries": 0, "disable_cooldowns": True})
|
||||
path: Final = directory / "prometheus-caller-credentials-opted-in.yaml"
|
||||
path.write_text(yaml.safe_dump(configuration))
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def opted_in(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("prometheus-caller-credentials-opted-in")
|
||||
released: Final = threading.Event()
|
||||
released.set()
|
||||
with (
|
||||
wire_server(Provider(released)) as upstream,
|
||||
gateway_from_environment() as shared,
|
||||
owned_proxy(shared, directory, {}, config=opted_in_config(directory, upstream)) as owned,
|
||||
):
|
||||
yield Rig(owned, upstream)
|
||||
|
||||
|
||||
def series(samples: Sequence[Sample], metric: str, **labels: str) -> tuple[Sample, ...]:
|
||||
return tuple(
|
||||
sample
|
||||
|
|
@ -464,6 +571,85 @@ def test_a_provider_401_on_an_anthropic_oauth_token_the_caller_sent_stays_out_of
|
|||
}, recorded(samples, requested_model=name("caller-oauth"))
|
||||
|
||||
|
||||
def test_a_provider_401_on_an_azure_deployments_own_ad_token_is_recorded(rig: Rig) -> None:
|
||||
rig.upstream.drain()
|
||||
response: Final = send(rig.gateway, "azure-own-token")
|
||||
assert response.status_code == 401, response.text
|
||||
assert [request.headers.get("authorization") for request in rig.upstream.drain()] == [
|
||||
"Bearer deployment-entra-token-azure-own-token"
|
||||
]
|
||||
assert_recorded_as_a_deployment_failure(rig, "azure-own-token", "401")
|
||||
|
||||
|
||||
def caller_azure_request(cell: str, route: str, credential: dict[str, JsonValue]) -> tuple[str, dict[str, JsonValue]]:
|
||||
if route == "embeddings":
|
||||
return "/v1/embeddings", {"model": name(cell), "input": cell, **credential}
|
||||
if route == "responses":
|
||||
return "/v1/responses", {"model": name(cell), "input": cell, **credential}
|
||||
if route == "messages":
|
||||
return "/v1/messages", {"model": name(cell), "max_tokens": 16, "messages": messages(cell), **credential}
|
||||
return "/v1/chat/completions", {"model": name(cell), "messages": messages(cell), **credential}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "inside_extra_body"),
|
||||
CALLER_AZURE_CELLS,
|
||||
ids=[f"{route}{'-extra-body' if inside_extra_body else ''}" for route, inside_extra_body in CALLER_AZURE_CELLS],
|
||||
)
|
||||
@pytest.mark.parametrize("opt_in", ("deployment", "proxy-wide"))
|
||||
def test_a_provider_401_on_an_azure_ad_token_the_caller_sent_stays_out_of_the_deployment_metrics(
|
||||
request: pytest.FixtureRequest, opt_in: str, route: str, inside_extra_body: bool
|
||||
) -> None:
|
||||
opted: Final[Rig] = request.getfixturevalue("rig" if opt_in == "deployment" else "opted_in")
|
||||
cell: Final = caller_azure_cell(opt_in, route, inside_extra_body)
|
||||
token: Final = f"caller-entra-token-{uuid.uuid4().hex}"
|
||||
credential: Final[dict[str, JsonValue]] = (
|
||||
{"extra_body": {"azure_ad_token": token}} if inside_extra_body else {"azure_ad_token": token}
|
||||
)
|
||||
opted.upstream.drain()
|
||||
response: Final = opted.gateway.request("POST", *caller_azure_request(cell, route, credential))
|
||||
assert response.status_code == 401, response.text
|
||||
assert [reached.headers.get("authorization") for reached in opted.upstream.drain()] == [
|
||||
None if inside_extra_body else f"Bearer {token}"
|
||||
]
|
||||
samples: Final = until(opted.gateway, PROXY_FAILED, requested_model=name(cell))
|
||||
assert {metric: values(samples, metric, model_id=identity(cell)) for metric in DEPLOYMENT_SERIES} == {
|
||||
metric: [] for metric in DEPLOYMENT_SERIES
|
||||
}, recorded(samples, requested_model=name(cell))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("cell", "field"), CALLER_CREDENTIAL_CELLS, ids=[cell for cell, _ in CALLER_CREDENTIAL_CELLS])
|
||||
def test_a_provider_401_on_a_provider_credential_the_caller_sent_stays_out_of_the_deployment_metrics(
|
||||
rig: Rig, cell: str, field: str
|
||||
) -> None:
|
||||
caller_credential: Final = f"caller-credential-{uuid.uuid4().hex}"
|
||||
rig.upstream.drain()
|
||||
response: Final = rig.gateway.request(
|
||||
"POST", "/v1/chat/completions", {"model": name(cell), "messages": messages(cell), field: caller_credential}
|
||||
)
|
||||
assert response.status_code == 401, response.text
|
||||
assert [reached.headers.get("authorization") for reached in rig.upstream.drain()] == [f"Bearer {caller_credential}"]
|
||||
samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell))
|
||||
assert {metric: values(samples, metric, model_id=identity(cell)) for metric in DEPLOYMENT_SERIES} == {
|
||||
metric: [] for metric in DEPLOYMENT_SERIES
|
||||
}, recorded(samples, requested_model=name(cell))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cell", TOKEN_EXCHANGE_CELLS)
|
||||
def test_a_token_exchange_deployment_rejecting_a_caller_without_an_identity_token_stays_out_of_the_deployment_metrics(
|
||||
rig: Rig, cell: str
|
||||
) -> None:
|
||||
rig.upstream.drain()
|
||||
response: Final = send(rig.gateway, cell)
|
||||
assert response.status_code == 401, response.text
|
||||
assert "IdP-issued access token" in response.text, response.text
|
||||
assert rig.upstream.drain() == ()
|
||||
samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell))
|
||||
assert {metric: values(samples, metric, model_id=identity(cell)) for metric in DEPLOYMENT_SERIES} == {
|
||||
metric: [] for metric in DEPLOYMENT_SERIES
|
||||
}, recorded(samples, requested_model=name(cell))
|
||||
|
||||
|
||||
def _gated_key(scenario: Scenario, cell: str) -> str | None:
|
||||
if cell == "invalid-key":
|
||||
return f"sk-not-a-virtual-key-{uuid.uuid4().hex}"
|
||||
|
|
|
|||
|
|
@ -19,6 +19,34 @@ from litellm.llms.github_copilot.per_user_auth import GithubCopilotUserSession
|
|||
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
|
||||
|
||||
|
||||
PROVIDER_CREDENTIAL_FIELDS = [
|
||||
"aws_access_key_id",
|
||||
"aws_secret_access_key",
|
||||
"aws_session_token",
|
||||
"aws_web_identity_token",
|
||||
"azure_ad_token",
|
||||
"azure_password",
|
||||
"azure_scope",
|
||||
"azure_username",
|
||||
"client_id",
|
||||
"client_secret",
|
||||
"gigachat_access_token",
|
||||
"oci_fingerprint",
|
||||
"oci_key",
|
||||
"oci_key_file",
|
||||
"oci_tenancy",
|
||||
"oci_user",
|
||||
"tenant_id",
|
||||
"token",
|
||||
"user_config",
|
||||
"vertex_ai_credentials",
|
||||
"vertex_credentials",
|
||||
"watsonx_credentials",
|
||||
"wx_credentials",
|
||||
"zen_api_key",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def prometheus_logger():
|
||||
"""Create a PrometheusLogger instance for testing."""
|
||||
|
|
@ -193,6 +221,25 @@ class TestSkipMetricsValidation:
|
|||
({"litellm_params": {"proxy_server_request": {"credential_fields": ("headers",)}}}, False),
|
||||
({"litellm_params": {"proxy_server_request": {"credential_fields": "api_key"}}}, False),
|
||||
({"standard_logging_object": {"metadata": {"used_client_oauth_token": False}}}, False),
|
||||
({"litellm_params": {"proxy_server_request": {"body": {"azure_ad_token": "caller-token"}}}}, True),
|
||||
(
|
||||
{
|
||||
"litellm_params": {
|
||||
"proxy_server_request": {"body": {"extra_body": {"azure_ad_token": "caller-token"}}}
|
||||
}
|
||||
},
|
||||
True,
|
||||
),
|
||||
(
|
||||
{
|
||||
"litellm_params": {
|
||||
"azure_ad_token": "deployment-token",
|
||||
"proxy_server_request": {"body": {"model": "azure-deployment"}},
|
||||
}
|
||||
},
|
||||
False,
|
||||
),
|
||||
({"litellm_params": {"proxy_server_request": {"body": {"extra_body": "azure_ad_token"}}}}, False),
|
||||
],
|
||||
ids=[
|
||||
"caller-supplied-api-key",
|
||||
|
|
@ -201,6 +248,10 @@ class TestSkipMetricsValidation:
|
|||
"caller-supplied-headers-only",
|
||||
"malformed-credential-fields",
|
||||
"deployment-anthropic-key",
|
||||
"caller-supplied-azure-ad-token",
|
||||
"caller-supplied-azure-ad-token-in-extra-body",
|
||||
"deployment-azure-ad-token",
|
||||
"malformed-extra-body",
|
||||
],
|
||||
)
|
||||
def test_skip_a_provider_401_on_a_credential_the_caller_owns(self, prometheus_logger, logged_call, expected):
|
||||
|
|
@ -212,6 +263,46 @@ class TestSkipMetricsValidation:
|
|||
kwargs = {"exception": exception, **logged_call}
|
||||
assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is expected
|
||||
|
||||
@pytest.mark.parametrize("field", PROVIDER_CREDENTIAL_FIELDS)
|
||||
@pytest.mark.parametrize("inside_extra_body", [False, True], ids=["body", "extra-body"])
|
||||
def test_skip_a_provider_401_on_a_provider_credential_field_the_caller_sent(
|
||||
self, prometheus_logger, field, inside_extra_body
|
||||
):
|
||||
sent = {field: "caller-credential"}
|
||||
body = {"model": "shared-deployment", **({"extra_body": sent} if inside_extra_body else sent)}
|
||||
kwargs = {
|
||||
"exception": litellm.AuthenticationError(message="Unauthorized", llm_provider="azure", model="gpt-4o"),
|
||||
"litellm_params": {"proxy_server_request": {"body": body}},
|
||||
}
|
||||
assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is True
|
||||
|
||||
@pytest.mark.parametrize("field", PROVIDER_CREDENTIAL_FIELDS)
|
||||
def test_count_a_provider_401_on_a_provider_credential_field_the_deployment_owns(self, prometheus_logger, field):
|
||||
kwargs = {
|
||||
"exception": litellm.AuthenticationError(message="Unauthorized", llm_provider="azure", model="gpt-4o"),
|
||||
"litellm_params": {
|
||||
field: "deployment-credential",
|
||||
"proxy_server_request": {"body": {"model": "shared-deployment"}},
|
||||
},
|
||||
}
|
||||
assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is False
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"token_exchange_endpoint,expected",
|
||||
[("https://login.example.com/oauth2/v2.0/token", True), ("", False), (None, False)],
|
||||
ids=["token-exchange-deployment", "blank-endpoint", "no-endpoint"],
|
||||
)
|
||||
def test_skip_a_provider_401_on_a_token_exchange_deployment(
|
||||
self, prometheus_logger, token_exchange_endpoint, expected
|
||||
):
|
||||
kwargs = {
|
||||
"exception": litellm.AuthenticationError(
|
||||
message="Unauthorized", llm_provider="microsoft_365_copilot", model="copilot"
|
||||
),
|
||||
"litellm_params": {"token_exchange_endpoint": token_exchange_endpoint, "client_secret": "deployment"},
|
||||
}
|
||||
assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is expected
|
||||
|
||||
def test_no_skip_for_valid_request(self, prometheus_logger):
|
||||
assert prometheus_logger._should_skip_metrics_for_invalid_key() is False
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue