mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-26 06:03:19 +00:00
fix(prometheus): count provider-side 401s in LLM failure and deployment metrics
The invalid virtual key skip no longer applies to an error the LLM provider returned while the deployment was called with its own credential, so a 401 for a bad deployment credential increments litellm_llm_api_failed_requests_metric, litellm_deployment_failure_responses and litellm_deployment_total_requests and marks litellm_deployment_state as a partial outage again. Proxy-gate rejections and a failing credential that belongs to the caller (a rejected per-user credential, an api_key sent with the request, a per-user GitHub Copilot session) keep the skip
This commit is contained in:
yucheng-berri
parent
35970c3eac
commit
16aa0ef1c6
6 changed files with 1071 additions and 60 deletions
|
|
@ -14,6 +14,7 @@ from functools import partial
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast
|
||||
|
||||
import openai
|
||||
from pydantic import BaseModel, Field, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict, assert_never
|
||||
|
||||
|
|
@ -22,11 +23,13 @@ from litellm._internal_context import with_service_target
|
|||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.constants import (
|
||||
AUTH_OBJECTS_TARGET,
|
||||
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
|
||||
PROMETHEUS_OVERFLOW_SERIES_LABEL_VALUE,
|
||||
PROXY_LLM_PROVIDER_FALLBACK,
|
||||
PROXY_REJECTED_BEFORE_ROUTING_KEY,
|
||||
)
|
||||
from litellm.exceptions import (
|
||||
CallerCredentialAuthenticationError,
|
||||
validate_rate_limit_category,
|
||||
validate_rate_limit_type,
|
||||
)
|
||||
|
|
@ -49,6 +52,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
from litellm.litellm_core_utils.service_tier_utils import (
|
||||
get_service_tier_from_standard_logging_payload,
|
||||
)
|
||||
from litellm.llms.github_copilot.per_user_auth import github_copilot_user_session_from
|
||||
from litellm.models.end_user import LiteLLM_EndUserTable
|
||||
from litellm.models.team import LiteLLM_TeamTable
|
||||
from litellm.models.user import LiteLLM_UserTable
|
||||
|
|
@ -324,6 +328,36 @@ def _get_proxy_llm_router() -> Router | None:
|
|||
return llm_router
|
||||
|
||||
|
||||
_LOGGED_FIELDS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _logged_fields(value: object) -> Mapping[str, object]:
|
||||
try:
|
||||
return _LOGGED_FIELDS.validate_python(value)
|
||||
except ValidationError:
|
||||
return MappingProxyType({})
|
||||
|
||||
|
||||
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, or the caller's own GitHub Copilot connection."""
|
||||
litellm_params: Final = _logged_fields(logged_call.get("litellm_params"))
|
||||
credential_fields: Final = _logged_fields(litellm_params.get("proxy_server_request")).get("credential_fields")
|
||||
return (
|
||||
isinstance(credential_fields, tuple) and "api_key" in credential_fields
|
||||
) or github_copilot_user_session_from(litellm_params) is not None
|
||||
|
||||
|
||||
def _is_deployment_failure(exception: openai.APIError, logged_call: Mapping[str, object]) -> bool:
|
||||
"""True for an error raised while calling the deployment with its own credential, on a request the proxy
|
||||
did not reject itself."""
|
||||
return not (
|
||||
isinstance(exception, CallerCredentialAuthenticationError)
|
||||
or bool(logged_call.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL))
|
||||
or _used_a_caller_credential(logged_call)
|
||||
)
|
||||
|
||||
|
||||
def _bounded_requested_model_label(requested_model: object, router_originated: bool = False) -> str | None:
|
||||
"""
|
||||
Bound ``requested_model`` label cardinality: names the router recognizes
|
||||
|
|
@ -2769,6 +2803,8 @@ class PrometheusLogger(CustomLogger):
|
|||
information from various callback function signatures and determines if the request
|
||||
represents an invalid API key attempt that should be filtered from metrics.
|
||||
|
||||
An error raised while calling the deployment with its own credential is never skipped.
|
||||
|
||||
Args:
|
||||
kwargs: Dictionary potentially containing exception and other data
|
||||
user_api_key_dict: User API key authentication object (currently unused)
|
||||
|
|
@ -2788,6 +2824,11 @@ class PrometheusLogger(CustomLogger):
|
|||
if exception is None and kwargs:
|
||||
exception = kwargs.get("exception")
|
||||
|
||||
if isinstance(exception, openai.APIError) and _is_deployment_failure(
|
||||
exception, _LOGGED_FIELDS.validate_python(kwargs or {})
|
||||
):
|
||||
return False
|
||||
|
||||
if self._is_invalid_api_key_request(status_code, exception=exception):
|
||||
verbose_logger.debug(
|
||||
"Skipping Prometheus metrics for invalid API key request: status_code=%s, exception=%s",
|
||||
|
|
|
|||
122
tests/integration/_support/copilot_session_prometheus_worker.py
Normal file
122
tests/integration/_support/copilot_session_prometheus_worker.py
Normal file
|
|
@ -0,0 +1,122 @@
|
|||
"""One SDK process for the per-user GitHub Copilot metric cells: a provider 401 on a caller's own Copilot session
|
||||
and one on the deployment's own key, then a report of the deployment series on that process's default registry.
|
||||
Run as a script by ``run_copilot_session_worker`` so each process registers the Prometheus metrics exactly once."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import litellm
|
||||
import openai
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.llms.github_copilot.per_user_auth import GithubCopilotUserSession
|
||||
from prometheus_client import REGISTRY
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
COPILOT_DEPLOYMENT: Final = "copilot-per-user-deployment"
|
||||
OPERATOR_DEPLOYMENT: Final = "operator-key-deployment"
|
||||
CALLER_SESSION_TOKEN: Final = "caller-copilot-session-token"
|
||||
OPERATOR_KEY: Final = "rotated-operator-key"
|
||||
FAILURES: Final = "litellm_deployment_failure_responses_total"
|
||||
TOTALS: Final = "litellm_deployment_total_requests_total"
|
||||
LLM_FAILED: Final = "litellm_llm_api_failed_requests_metric_total"
|
||||
STATE: Final = "litellm_deployment_state"
|
||||
DEPLOYMENT_SERIES: Final = frozenset({FAILURES, TOTALS, LLM_FAILED, STATE})
|
||||
WORKER_TIMEOUT_SECONDS: Final = 120
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Outcome:
|
||||
deployment: str
|
||||
status: int | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Series:
|
||||
metric: str
|
||||
model_id: str
|
||||
value: float
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class WorkerReport:
|
||||
outcomes: tuple[Outcome, ...]
|
||||
series: tuple[Series, ...]
|
||||
|
||||
def values(self, metric: str, model_id: str) -> tuple[float, ...]:
|
||||
return tuple(s.value for s in self.series if s.metric == metric and s.model_id == model_id)
|
||||
|
||||
|
||||
_REPORT: Final = TypeAdapter(WorkerReport)
|
||||
|
||||
|
||||
def run_copilot_session_worker(api_base: str) -> WorkerReport:
|
||||
completed: Final = subprocess.run(
|
||||
(sys.executable, "-P", __file__, api_base),
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=WORKER_TIMEOUT_SECONDS,
|
||||
check=False,
|
||||
)
|
||||
assert completed.returncode == 0, completed.stderr[-3000:]
|
||||
return _REPORT.validate_json(completed.stdout.splitlines()[-1])
|
||||
|
||||
|
||||
def _deployment_series() -> tuple[Series, ...]:
|
||||
return tuple(
|
||||
Series(metric=sample.name, model_id=str(sample.labels.get("model_id", "")), value=float(sample.value))
|
||||
for family in REGISTRY.collect()
|
||||
for sample in family.samples
|
||||
if sample.name in DEPLOYMENT_SERIES
|
||||
)
|
||||
|
||||
|
||||
def _router(api_base: str) -> litellm.Router:
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": COPILOT_DEPLOYMENT,
|
||||
"litellm_params": {
|
||||
"model": "github_copilot/gpt-4o",
|
||||
"litellm_credential_name": "copilot-per-user",
|
||||
"github_copilot_auth_type": "per_user_oauth",
|
||||
},
|
||||
"model_info": {"id": COPILOT_DEPLOYMENT},
|
||||
},
|
||||
{
|
||||
"model_name": OPERATOR_DEPLOYMENT,
|
||||
"litellm_params": {"model": "openai/gpt-5.6", "api_base": api_base, "api_key": OPERATOR_KEY},
|
||||
"model_info": {"id": OPERATOR_DEPLOYMENT},
|
||||
},
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
|
||||
async def _attempt(router: litellm.Router, model: str, **extra: object) -> Outcome:
|
||||
try:
|
||||
await router.acompletion(model=model, messages=[{"role": "user", "content": f"{model} worker"}], **extra)
|
||||
except openai.APIStatusError as error:
|
||||
return Outcome(deployment=model, status=error.status_code)
|
||||
return Outcome(deployment=model, status=None)
|
||||
|
||||
|
||||
async def main(api_base: str) -> None:
|
||||
litellm.callbacks = [PrometheusLogger()] # test-quality-ok: this worker process exits right after the report
|
||||
session: Final = GithubCopilotUserSession(token=CALLER_SESSION_TOKEN, api_base=api_base)
|
||||
router: Final = _router(api_base)
|
||||
outcomes: Final = (
|
||||
await _attempt(router, COPILOT_DEPLOYMENT, github_copilot_user_session=session),
|
||||
await _attempt(router, OPERATOR_DEPLOYMENT),
|
||||
)
|
||||
report: Final = WorkerReport(outcomes=outcomes, series=_deployment_series())
|
||||
sys.stdout.write(json.dumps(_REPORT.dump_python(report)) + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main(sys.argv[1]))
|
||||
|
|
@ -0,0 +1,488 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
import uuid
|
||||
from collections import Counter
|
||||
from collections.abc import Iterator, Sequence
|
||||
from dataclasses import dataclass
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment
|
||||
from integration._support.database import read_rows, write_rows
|
||||
from integration._support.mcp import ENTRY_POINTS, McpCaller, McpPeer, echo_tool, register_mcp, scripted_peer
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.prometheus_series import Sample, scrape
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
RUN: Final = uuid.uuid4().hex[:8]
|
||||
REJECTS: Final = "provider-rejects-credential"
|
||||
HELD_REJECTS: Final = "provider-rejects-credential-when-released"
|
||||
KEY_NOT_VALID: Final = "provider-key-not-valid"
|
||||
RATE_LIMITS: Final = "provider-rate-limits"
|
||||
RECOVERS: Final = "provider-rejects-once"
|
||||
HEALTHY: Final = "provider-healthy"
|
||||
AGENT_365_GUARDRAIL: Final = f"agent-365-{RUN}"
|
||||
|
||||
FAILURES: Final = "litellm_deployment_failure_responses_total"
|
||||
TOTALS: Final = "litellm_deployment_total_requests_total"
|
||||
LLM_FAILED: Final = "litellm_llm_api_failed_requests_metric_total"
|
||||
STATE: Final = "litellm_deployment_state"
|
||||
PROXY_FAILED: Final = "litellm_proxy_failed_requests_metric_total"
|
||||
SUCCESSES: Final = "litellm_deployment_success_responses_total"
|
||||
DEPLOYMENT_SERIES: Final = (FAILURES, TOTALS, LLM_FAILED, STATE)
|
||||
|
||||
ENDPOINT_CELLS: Final = (
|
||||
"chat",
|
||||
"chat-stream",
|
||||
"messages",
|
||||
"messages-stream",
|
||||
"responses",
|
||||
"responses-stream",
|
||||
"embeddings",
|
||||
)
|
||||
REJECTED_CELLS: Final = (*ENDPOINT_CELLS, "sdk-sync", "sdk-async", "repeated", "fallback", "burst", "caller-api-key")
|
||||
GATED_CELLS: Final = ("invalid-key", "missing-key", "expired-key", "blocked-key", "no-model-access")
|
||||
HEALTHY_CELLS: Final = (*GATED_CELLS, "healthy", "key-rpm-limit")
|
||||
GUARDRAIL_ROWS: Final = (
|
||||
"SELECT COALESCE(jsonb_path_query_first(metadata, "
|
||||
"'$.guardrail_information[*] ? (@.guardrail_name == $name).guardrail_status', "
|
||||
"jsonb_build_object('name', %s::text)) #>> '{}', 'none') AS status "
|
||||
'FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND call_type = %s'
|
||||
)
|
||||
|
||||
|
||||
def name(cell: str) -> str:
|
||||
return f"prometheus-{cell}-{RUN}"
|
||||
|
||||
|
||||
def identity(cell: str) -> str:
|
||||
return f"{name(cell)}-deployment"
|
||||
|
||||
|
||||
def _error(status: int, message: str, kind: str, code: str) -> Reply:
|
||||
return Reply(status=status, body=json.dumps({"error": {"message": message, "type": kind, "code": code}}).encode())
|
||||
|
||||
|
||||
def _completion() -> Reply:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
class Provider:
|
||||
"""Answers by upstream model name. `HELD_REJECTS` waits for `release` before rejecting the credential."""
|
||||
|
||||
def __init__(self, release: threading.Event) -> None:
|
||||
self._lock: Final = threading.Lock()
|
||||
self._seen: Final[Counter[str]] = Counter()
|
||||
self._release: Final = release
|
||||
|
||||
def __call__(self, request: Request) -> Reply:
|
||||
model: Final = str(json.loads(request.body)["model"])
|
||||
with self._lock:
|
||||
self._seen[model] += 1
|
||||
attempt: Final = self._seen[model]
|
||||
if model == HELD_REJECTS:
|
||||
assert self._release.wait(timeout=30), "held provider reply was never released"
|
||||
if model in (REJECTS, HELD_REJECTS) or (model == RECOVERS and attempt == 1):
|
||||
return _error(401, "Incorrect API key provided: sk-rotat*****", "invalid_request_error", "invalid_api_key")
|
||||
if model == KEY_NOT_VALID:
|
||||
return _error(400, "API key not valid. Please pass a valid API key.", "invalid_request_error", "400")
|
||||
if model == RATE_LIMITS:
|
||||
return _error(429, "Rate limit reached for requests", "rate_limit_error", "rate_limit_exceeded")
|
||||
return _completion()
|
||||
|
||||
|
||||
def deployment(cell: str, upstream_model: str, api_base: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_name": name(cell),
|
||||
"litellm_params": {
|
||||
"model": f"openai/{upstream_model}",
|
||||
"api_key": f"sk-deployment-credential-{cell}",
|
||||
"api_base": api_base,
|
||||
},
|
||||
"model_info": {"id": identity(cell)},
|
||||
}
|
||||
|
||||
|
||||
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())
|
||||
configuration["model_list"] = [
|
||||
*configuration["model_list"],
|
||||
*(deployment(cell, REJECTS, api_base) for cell in REJECTED_CELLS),
|
||||
*(deployment(cell, HEALTHY, api_base) for cell in HEALTHY_CELLS),
|
||||
deployment("held", HELD_REJECTS, api_base),
|
||||
deployment("key-not-valid", KEY_NOT_VALID, api_base),
|
||||
deployment("rate-limited", RATE_LIMITS, api_base),
|
||||
deployment("recovers", RECOVERS, api_base),
|
||||
{
|
||||
"model_name": name("caller-credential"),
|
||||
"litellm_params": {
|
||||
"model": "github_copilot/gpt-4o",
|
||||
"litellm_credential_name": f"copilot-per-user-{RUN}",
|
||||
"github_copilot_auth_type": "per_user_oauth",
|
||||
},
|
||||
"model_info": {"id": identity("caller-credential")},
|
||||
},
|
||||
]
|
||||
configuration["guardrails"] = [
|
||||
{
|
||||
"guardrail_name": AGENT_365_GUARDRAIL,
|
||||
"litellm_params": {
|
||||
"guardrail": "agent_365",
|
||||
"mode": "pre_mcp_call",
|
||||
"default_on": True,
|
||||
"tenant_id": "00000000-0000-4000-8000-0000000a3650",
|
||||
"client_id": "synthetic-client-id",
|
||||
"client_secret": "synthetic-client-secret",
|
||||
},
|
||||
}
|
||||
]
|
||||
configuration["litellm_settings"]["callbacks"] = ["prometheus"]
|
||||
configuration.setdefault("router_settings", {}).update({"num_retries": 0, "disable_cooldowns": True})
|
||||
path: Final = directory / "prometheus-provider-auth-failure.yaml"
|
||||
path.write_text(yaml.safe_dump(configuration))
|
||||
return path
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
gateway: Gateway
|
||||
upstream: Wire
|
||||
|
||||
@property
|
||||
def api_base(self) -> str:
|
||||
return f"{self.upstream.url}/v1"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("prometheus-provider-auth-failure")
|
||||
released: Final = threading.Event()
|
||||
released.set()
|
||||
with (
|
||||
wire_server(Provider(released)) as upstream,
|
||||
gateway_from_environment() as shared,
|
||||
owned_proxy(
|
||||
shared, directory, {"PROXY_CONFIG_RELOAD_INTERVAL_SECONDS": "2"}, config=config(directory, upstream)
|
||||
) as owned,
|
||||
):
|
||||
yield Rig(owned, upstream)
|
||||
|
||||
|
||||
def series(samples: Sequence[Sample], metric: str, **labels: str) -> tuple[Sample, ...]:
|
||||
return tuple(
|
||||
sample
|
||||
for sample in samples
|
||||
if sample.name == metric and all(sample.labels.get(label) == value for label, value in labels.items())
|
||||
)
|
||||
|
||||
|
||||
def values(samples: Sequence[Sample], metric: str, **labels: str) -> list[float]:
|
||||
return [sample.value for sample in series(samples, metric, **labels)]
|
||||
|
||||
|
||||
def until(gateway: Gateway, metric: str, **labels: str) -> tuple[Sample, ...]:
|
||||
return eventually(lambda: scrape(gateway), lambda samples: bool(series(samples, metric, **labels)), 30)
|
||||
|
||||
|
||||
def messages(cell: str) -> list[JsonValue]:
|
||||
return [{"role": "user", "content": f"{cell} {uuid.uuid4().hex}"}]
|
||||
|
||||
|
||||
def send(gateway: Gateway, cell: str, model: str | None = None, key: str | None = None) -> httpx.Response:
|
||||
asked: Final = model or name(cell)
|
||||
stream: Final = cell.endswith("-stream")
|
||||
if cell.startswith("messages"):
|
||||
anthropic: Final = {"model": asked, "max_tokens": 16, "messages": messages(cell), "stream": stream}
|
||||
return gateway.request("POST", "/v1/messages", anthropic, key=key)
|
||||
if cell.startswith("responses"):
|
||||
return gateway.request("POST", "/v1/responses", {"model": asked, "input": cell, "stream": stream}, key=key)
|
||||
if cell == "embeddings":
|
||||
return gateway.request("POST", "/v1/embeddings", {"model": asked, "input": cell}, key=key)
|
||||
return gateway.request(
|
||||
"POST", "/v1/chat/completions", {"model": asked, "messages": messages(cell), "stream": stream}, key=key
|
||||
)
|
||||
|
||||
|
||||
def assert_reached_with_the_deployment_credential(rig: Rig, cell: str, attempts: int = 1) -> None:
|
||||
reached: Final = rig.upstream.drain()
|
||||
assert [request.headers["authorization"] for request in reached] == [
|
||||
f"Bearer sk-deployment-credential-{cell}"
|
||||
] * attempts, [request.target for request in reached]
|
||||
|
||||
|
||||
def recorded(samples: Sequence[Sample], **labels: str) -> list[tuple[str, dict[str, str], float]]:
|
||||
return [(sample.name, dict(sample.labels), sample.value) for sample in series_of(samples, **labels)]
|
||||
|
||||
|
||||
def series_of(samples: Sequence[Sample], **labels: str) -> tuple[Sample, ...]:
|
||||
return tuple(
|
||||
sample
|
||||
for sample in samples
|
||||
if sample.name.startswith("litellm_")
|
||||
and all(sample.labels.get(label) == value for label, value in labels.items())
|
||||
)
|
||||
|
||||
|
||||
def assert_recorded_as_a_deployment_failure(rig: Rig, cell: str, status: str, count: float = 1.0) -> None:
|
||||
samples: Final = until(rig.gateway, FAILURES, model_id=identity(cell))
|
||||
observed: Final = recorded(samples, model_id=identity(cell))
|
||||
assert values(samples, FAILURES, model_id=identity(cell), exception_status=status) == [count], observed
|
||||
assert values(samples, TOTALS, model_id=identity(cell)) == [count], observed
|
||||
assert values(samples, LLM_FAILED, model_id=identity(cell)) == [count], observed
|
||||
assert values(samples, STATE, model_id=identity(cell)) == [1.0], observed
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cell", ENDPOINT_CELLS)
|
||||
def test_a_provider_401_on_every_endpoint_counts_as_a_deployment_failure_and_a_partial_outage(
|
||||
rig: Rig, cell: str
|
||||
) -> None:
|
||||
rig.upstream.drain()
|
||||
response: Final = send(rig.gateway, cell)
|
||||
assert response.status_code == 401, response.text
|
||||
assert_reached_with_the_deployment_credential(rig, cell)
|
||||
assert_recorded_as_a_deployment_failure(rig, cell, "401")
|
||||
|
||||
|
||||
def test_a_provider_401_seen_by_the_sync_openai_client_is_recorded(rig: Rig) -> None:
|
||||
rig.upstream.drain()
|
||||
with (
|
||||
openai.OpenAI(
|
||||
base_url=f"{rig.gateway.client.base_url}/v1",
|
||||
api_key=rig.gateway.key,
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(trust_env=False),
|
||||
) as client,
|
||||
pytest.raises(openai.AuthenticationError, match="Incorrect API key provided"),
|
||||
):
|
||||
client.chat.completions.create(model=name("sdk-sync"), messages=[{"role": "user", "content": "sync"}])
|
||||
assert_reached_with_the_deployment_credential(rig, "sdk-sync")
|
||||
assert_recorded_as_a_deployment_failure(rig, "sdk-sync", "401")
|
||||
|
||||
|
||||
def test_a_provider_401_seen_by_the_async_openai_client_is_recorded(rig: Rig) -> None:
|
||||
async def ask() -> None:
|
||||
async with openai.AsyncOpenAI(
|
||||
base_url=f"{rig.gateway.client.base_url}/v1",
|
||||
api_key=rig.gateway.key,
|
||||
max_retries=0,
|
||||
http_client=httpx.AsyncClient(trust_env=False),
|
||||
) as client:
|
||||
await client.chat.completions.create(model=name("sdk-async"), messages=[{"role": "user", "content": "a"}])
|
||||
|
||||
rig.upstream.drain()
|
||||
with pytest.raises(openai.AuthenticationError, match="Incorrect API key provided"):
|
||||
asyncio.run(ask())
|
||||
assert_reached_with_the_deployment_credential(rig, "sdk-async")
|
||||
assert_recorded_as_a_deployment_failure(rig, "sdk-async", "401")
|
||||
|
||||
|
||||
def test_each_repeated_provider_401_is_counted_once(rig: Rig) -> None:
|
||||
rig.upstream.drain()
|
||||
assert [send(rig.gateway, "repeated").status_code for _ in range(3)] == [401, 401, 401]
|
||||
assert_reached_with_the_deployment_credential(rig, "repeated", attempts=3)
|
||||
assert_recorded_as_a_deployment_failure(rig, "repeated", "401", count=3.0)
|
||||
|
||||
|
||||
def test_a_provider_401_rescued_by_a_fallback_still_counts_against_the_rejected_deployment(rig: Rig) -> None:
|
||||
rig.upstream.drain()
|
||||
response: Final = rig.gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": name("fallback"), "messages": messages("fallback"), "fallbacks": [name("healthy")]},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert [request.headers["authorization"] for request in rig.upstream.drain()] == [
|
||||
"Bearer sk-deployment-credential-fallback",
|
||||
"Bearer sk-deployment-credential-healthy",
|
||||
]
|
||||
assert_recorded_as_a_deployment_failure(rig, "fallback", "401")
|
||||
rescued: Final = eventually(
|
||||
lambda: scrape(rig.gateway), lambda samples: values(samples, SUCCESSES, model_id=identity("healthy")) == [1.0]
|
||||
)
|
||||
assert values(rescued, FAILURES, model_id=identity("healthy")) == []
|
||||
|
||||
|
||||
def test_a_provider_auth_message_on_a_non_401_status_is_recorded(rig: Rig) -> None:
|
||||
rig.upstream.drain()
|
||||
response: Final = send(rig.gateway, "key-not-valid")
|
||||
assert response.status_code == 400, response.text
|
||||
assert_reached_with_the_deployment_credential(rig, "key-not-valid")
|
||||
assert_recorded_as_a_deployment_failure(rig, "key-not-valid", "400")
|
||||
|
||||
|
||||
def test_a_provider_429_is_recorded_as_before(rig: Rig) -> None:
|
||||
rig.upstream.drain()
|
||||
response: Final = send(rig.gateway, "rate-limited")
|
||||
assert response.status_code == 429, response.text
|
||||
assert_reached_with_the_deployment_credential(rig, "rate-limited")
|
||||
assert_recorded_as_a_deployment_failure(rig, "rate-limited", "429")
|
||||
|
||||
|
||||
def test_a_deployment_marked_down_by_a_provider_401_is_healthy_again_after_a_success(rig: Rig) -> None:
|
||||
rig.upstream.drain()
|
||||
assert send(rig.gateway, "recovers").status_code == 401
|
||||
assert_recorded_as_a_deployment_failure(rig, "recovers", "401")
|
||||
recovered: Final = send(rig.gateway, "recovers")
|
||||
assert recovered.status_code == 200, recovered.text
|
||||
assert len(rig.upstream.drain()) == 2
|
||||
eventually(
|
||||
lambda: scrape(rig.gateway),
|
||||
lambda samples: values(samples, STATE, model_id=identity("recovers"), api_base=rig.api_base) == [0.0],
|
||||
)
|
||||
|
||||
|
||||
def test_a_per_caller_credential_rejection_stays_out_of_the_deployment_metrics(rig: Rig) -> None:
|
||||
rig.upstream.drain()
|
||||
response: Final = send(rig.gateway, "caller-credential")
|
||||
assert response.status_code == 401, response.text
|
||||
assert "Connect GitHub Copilot" in response.text, response.text
|
||||
assert rig.upstream.drain() == ()
|
||||
samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name("caller-credential"))
|
||||
assert values(samples, FAILURES, model_id=identity("caller-credential")) == []
|
||||
assert values(samples, TOTALS, model_id=identity("caller-credential")) == []
|
||||
assert values(samples, LLM_FAILED, model_id=identity("caller-credential")) == []
|
||||
assert values(samples, STATE, model_id=identity("caller-credential")) == []
|
||||
|
||||
|
||||
def deployment_series(samples: Sequence[Sample]) -> frozenset[tuple[str, tuple[tuple[str, str], ...]]]:
|
||||
return frozenset((sample.name, sample.identity()) for sample in samples if sample.name in DEPLOYMENT_SERIES)
|
||||
|
||||
|
||||
def test_a_provider_401_on_an_api_key_the_caller_sent_stays_out_of_the_deployment_metrics(rig: Rig) -> None:
|
||||
caller_key: Final = f"sk-caller-supplied-{uuid.uuid4().hex}"
|
||||
before: Final = deployment_series(scrape(rig.gateway))
|
||||
rig.upstream.drain()
|
||||
response: Final = rig.gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": name("caller-api-key"), "messages": messages("caller-api-key"), "api_key": caller_key},
|
||||
)
|
||||
assert response.status_code == 401, response.text
|
||||
assert [request.headers["authorization"] for request in rig.upstream.drain()] == [f"Bearer {caller_key}"]
|
||||
samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name("caller-api-key"))
|
||||
assert deployment_series(samples) - before == frozenset(), recorded(samples, requested_model=name("caller-api-key"))
|
||||
|
||||
|
||||
def _gated_key(scenario: Scenario, cell: str) -> str | None:
|
||||
if cell == "invalid-key":
|
||||
return f"sk-not-a-virtual-key-{uuid.uuid4().hex}"
|
||||
if cell == "missing-key":
|
||||
return None
|
||||
if cell == "no-model-access":
|
||||
return scenario.key(models=[name("healthy")])
|
||||
key: Final = scenario.key(models=[name(cell)])
|
||||
if cell == "expired-key":
|
||||
write_rows(
|
||||
"UPDATE \"LiteLLM_VerificationToken\" SET expires = now() - interval '1 hour' WHERE token = %s",
|
||||
(sha256(key.encode()).hexdigest(),),
|
||||
)
|
||||
return key
|
||||
scenario.gateway.post("/key/block", {"key": key})
|
||||
return key
|
||||
|
||||
|
||||
def _gated(gateway: Gateway, cell: str, key: str | None) -> httpx.Response:
|
||||
if key is None:
|
||||
return gateway.client.post("/v1/chat/completions", json={"model": name(cell), "messages": messages(cell)})
|
||||
return send(gateway, cell, key=key)
|
||||
|
||||
|
||||
def recorded_before(cell: str) -> list[tuple[str, str, str, float]]:
|
||||
if cell in ("invalid-key", "expired-key"):
|
||||
return []
|
||||
return sorted([(FAILURES, "", "", 1.0), (TOTALS, "", "", 1.0), (LLM_FAILED, identity(cell), "", 1.0)])
|
||||
|
||||
|
||||
def deployment_rows(samples: Sequence[Sample], cell: str) -> list[tuple[str, str, str, float]]:
|
||||
return sorted(
|
||||
(sample.name, sample.labels.get("model_id", ""), sample.labels.get("api_base", ""), sample.value)
|
||||
for sample in samples
|
||||
if sample.name in DEPLOYMENT_SERIES
|
||||
and name(cell) in (sample.labels.get("requested_model"), sample.labels.get("model"))
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cell", GATED_CELLS)
|
||||
def test_a_request_the_proxy_rejects_itself_never_marks_the_deployment_down_and_is_recorded_as_before(
|
||||
rig: Rig, cell: str
|
||||
) -> None:
|
||||
with rig.gateway.scenario() as scenario:
|
||||
key: Final = _gated_key(scenario, cell)
|
||||
rig.upstream.drain()
|
||||
response: Final = _gated(rig.gateway, cell, key)
|
||||
assert response.status_code == (403 if cell == "no-model-access" else 401), response.text
|
||||
assert rig.upstream.drain() == ()
|
||||
samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell))
|
||||
assert deployment_rows(samples, cell) == recorded_before(cell), recorded(samples, requested_model=name(cell))
|
||||
|
||||
|
||||
def test_a_key_rate_limit_rejected_at_the_proxy_is_still_counted_as_before(rig: Rig) -> None:
|
||||
with rig.gateway.scenario() as scenario:
|
||||
key: Final = scenario.key(models=[name("key-rpm-limit")], rpm_limit=1)
|
||||
rig.upstream.drain()
|
||||
admitted: Final = send(rig.gateway, "key-rpm-limit", key=key)
|
||||
assert admitted.status_code == 200, admitted.text
|
||||
limited: Final = send(rig.gateway, "key-rpm-limit", key=key)
|
||||
assert limited.status_code == 429, limited.text
|
||||
assert len(rig.upstream.drain()) == 1
|
||||
samples: Final = until(rig.gateway, FAILURES, requested_model=name("key-rpm-limit"), exception_status="429")
|
||||
assert values(samples, FAILURES, requested_model=name("key-rpm-limit"), exception_status="429") == [1.0]
|
||||
eventually(
|
||||
lambda: scrape(rig.gateway),
|
||||
lambda observed: (
|
||||
values(observed, STATE, model_id=identity("key-rpm-limit"), api_base=rig.api_base) == [0.0]
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def mcp_peer() -> Iterator[McpPeer]:
|
||||
with scripted_peer(echo_tool("add")) as peer:
|
||||
yield peer
|
||||
|
||||
|
||||
def test_an_mcp_tool_call_the_proxy_rejects_for_the_caller_stays_out_of_the_llm_failure_metrics(
|
||||
rig: Rig, mcp_peer: McpPeer
|
||||
) -> None:
|
||||
alias: Final = "a365" + RUN
|
||||
with rig.gateway.scenario() as scenario:
|
||||
server: Final = register_mcp(scenario, mcp_peer, alias)
|
||||
key: Final = scenario.key(object_permission={"mcp_servers": [server]})
|
||||
hashed: Final = sha256(key.encode()).hexdigest()
|
||||
outcomes: Final = {
|
||||
entry: McpCaller(rig.gateway, key, entry, alias).call(f"{alias}-add", {"entry": entry}, server)
|
||||
for entry in ENTRY_POINTS
|
||||
}
|
||||
assert all(
|
||||
outcome.error is not None and "Agent 365 guardrail rejected" in outcome.raw for outcome in outcomes.values()
|
||||
), {entry: outcome.raw for entry, outcome in outcomes.items()}
|
||||
eventually(
|
||||
lambda: read_rows(GUARDRAIL_ROWS, (AGENT_365_GUARDRAIL, hashed, "call_mcp_tool")),
|
||||
lambda rows: [row["status"] for row in rows] == ["guardrail_intervened"] * len(ENTRY_POINTS),
|
||||
seconds=70,
|
||||
)
|
||||
samples: Final = scrape(rig.gateway)
|
||||
assert values(samples, LLM_FAILED, hashed_api_key=hashed) == []
|
||||
assert values(samples, FAILURES, hashed_api_key=hashed) == []
|
||||
assert values(samples, TOTALS, hashed_api_key=hashed) == []
|
||||
|
|
@ -0,0 +1,147 @@
|
|||
import threading
|
||||
import uuid
|
||||
from collections.abc import Iterator, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import psutil
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.process import OwnedProxy, owned_proxy_process
|
||||
from integration._support.prometheus_series import Sample, scrape
|
||||
from integration._support.wire import Wire, wire_server
|
||||
from test_prometheus_provider_auth_failure import (
|
||||
FAILURES,
|
||||
LLM_FAILED,
|
||||
STATE,
|
||||
TOTALS,
|
||||
Provider,
|
||||
config,
|
||||
identity,
|
||||
name,
|
||||
send,
|
||||
values,
|
||||
)
|
||||
|
||||
BURST: Final = 12
|
||||
ENDPOINTS: Final = ("chat", "chat-stream", "messages", "responses")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ChaosRig:
|
||||
owned: OwnedProxy
|
||||
upstream: Wire
|
||||
release: threading.Event
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _rig(gateway: Gateway, tmp_path: Path) -> Iterator[ChaosRig]:
|
||||
release: Final = threading.Event()
|
||||
prom_dir: Final = tmp_path / "prom"
|
||||
prom_dir.mkdir()
|
||||
with (
|
||||
wire_server(Provider(release)) as upstream,
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"PROMETHEUS_MULTIPROC_DIR": str(prom_dir)},
|
||||
config=config(tmp_path, upstream),
|
||||
workers=2,
|
||||
) as owned,
|
||||
):
|
||||
try:
|
||||
yield ChaosRig(owned, upstream, release)
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
|
||||
def _status(gateway: Gateway, cell: str, model: str, key: str | None = None) -> int:
|
||||
try:
|
||||
return send(gateway, cell, model=model, key=key).status_code
|
||||
except httpx.HTTPError:
|
||||
return -1
|
||||
|
||||
|
||||
def _holds_a_connection_to(pid: int, port: int) -> bool:
|
||||
return any(connection.raddr[1:] == (port,) for connection in psutil.Process(pid).net_connections(kind="tcp"))
|
||||
|
||||
|
||||
def _workers_calling_upstream(rig: ChaosRig) -> frozenset[int]:
|
||||
port: Final = int(rig.upstream.url.rsplit(":", 1)[1])
|
||||
return frozenset(pid for pid in rig.owned.worker_pids() if _holds_a_connection_to(pid, port))
|
||||
|
||||
|
||||
def _counted(samples: Sequence[Sample], cell: str) -> tuple[float, float, float]:
|
||||
return (
|
||||
sum(values(samples, FAILURES, model_id=identity(cell))),
|
||||
sum(values(samples, TOTALS, model_id=identity(cell))),
|
||||
sum(values(samples, LLM_FAILED, model_id=identity(cell))),
|
||||
)
|
||||
|
||||
|
||||
def test_a_burst_of_provider_401s_and_proxy_rejects_across_two_workers_counts_only_the_provider_401s(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with _rig(gateway, tmp_path) as rig:
|
||||
rig.release.set()
|
||||
workers: Final = eventually(rig.owned.worker_pids, lambda pids: len(pids) == 2)
|
||||
invalid: Final = f"sk-not-a-virtual-key-{uuid.uuid4().hex}"
|
||||
calls: Final[tuple[tuple[str, str, str | None], ...]] = tuple(
|
||||
(ENDPOINTS[index % len(ENDPOINTS)], name("burst"), None) for index in range(BURST)
|
||||
) + tuple(("chat", name("invalid-key"), invalid) for _ in range(BURST))
|
||||
|
||||
def call(cell: str, model: str, key: str | None) -> int:
|
||||
return _status(rig.owned.gateway, cell, model, key)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
statuses: Final = tuple(pool.map(call, *zip(*calls, strict=True)))
|
||||
assert statuses == (401,) * (2 * BURST), statuses
|
||||
assert len(rig.upstream.drain()) == BURST
|
||||
samples: Final = eventually(
|
||||
lambda: scrape(rig.owned.gateway),
|
||||
lambda observed: _counted(observed, "burst") == (BURST, BURST, BURST),
|
||||
seconds=70,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
assert _counted(samples, "burst") == (BURST, BURST, BURST), (workers, _counted(samples, "burst"))
|
||||
assert set(values(samples, STATE, model_id=identity("burst"))) == {1.0}
|
||||
assert values(samples, FAILURES, requested_model=name("invalid-key")) == []
|
||||
assert values(samples, LLM_FAILED, model=name("invalid-key")) == []
|
||||
|
||||
|
||||
def test_a_worker_killed_while_provider_401s_are_in_flight_loses_no_answered_failure(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with _rig(gateway, tmp_path) as rig:
|
||||
eventually(rig.owned.worker_pids, lambda pids: len(pids) == 2)
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
futures: Final = tuple(pool.submit(_status, rig.owned.gateway, "chat", name("held")) for _ in range(BURST))
|
||||
eventually(lambda: rig.upstream.received.qsize() >= 8, bool, seconds=30)
|
||||
victim: Final = min(eventually(lambda: _workers_calling_upstream(rig), bool))
|
||||
with rig.owned.replacing_worker(victim) as kill:
|
||||
kill()
|
||||
rig.release.set()
|
||||
statuses: Final = tuple(future.result() for future in futures)
|
||||
reached: Final = len(rig.upstream.drain())
|
||||
answered: Final = statuses.count(401)
|
||||
lost: Final = statuses.count(-1)
|
||||
assert answered + lost == BURST and lost > 0, statuses
|
||||
samples: Final = eventually(
|
||||
lambda: scrape(rig.owned.gateway),
|
||||
lambda observed: answered <= _counted(observed, "held")[0] <= reached,
|
||||
seconds=70,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
counted: Final = _counted(samples, "held")
|
||||
assert answered <= counted[0] <= reached and counted[0] == counted[1] == counted[2], (statuses, counted)
|
||||
follow_up: Final = _status(rig.owned.gateway, "chat", name("held"))
|
||||
assert follow_up == 401
|
||||
recounted: Final = eventually(
|
||||
lambda: _counted(scrape(rig.owned.gateway), "held"),
|
||||
lambda observed: observed[0] == counted[0] + 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert recounted == (counted[0] + 1,) * 3
|
||||
|
|
@ -0,0 +1,71 @@
|
|||
"""A provider 401 on a caller's own per-user GitHub Copilot session is the caller's credential failing, so the
|
||||
shared Copilot deployment keeps no failure, request or outage series for it, while a provider 401 on the
|
||||
deployment's own key in the same process is counted against that deployment."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Iterator
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.copilot_session_prometheus_worker import (
|
||||
CALLER_SESSION_TOKEN,
|
||||
COPILOT_DEPLOYMENT,
|
||||
DEPLOYMENT_SERIES,
|
||||
FAILURES,
|
||||
LLM_FAILED,
|
||||
OPERATOR_DEPLOYMENT,
|
||||
OPERATOR_KEY,
|
||||
STATE,
|
||||
TOTALS,
|
||||
Outcome,
|
||||
WorkerReport,
|
||||
run_copilot_session_worker,
|
||||
)
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def ran() -> Iterator[tuple[WorkerReport, tuple[Request, ...]]]:
|
||||
def reject(request: Request) -> Reply:
|
||||
assert request.method == "POST", request.method
|
||||
return Reply(
|
||||
status=401,
|
||||
body=json.dumps({"error": {"message": "bad credentials", "type": "authentication_error"}}).encode(),
|
||||
)
|
||||
|
||||
with wire_server(reject) as upstream:
|
||||
report: Final = run_copilot_session_worker(upstream.url)
|
||||
yield report, upstream.drain()
|
||||
|
||||
|
||||
def test_both_calls_reach_the_provider_with_their_own_credential_and_fail_with_its_401(
|
||||
ran: tuple[WorkerReport, tuple[Request, ...]],
|
||||
) -> None:
|
||||
report, received = ran
|
||||
assert report.outcomes == (
|
||||
Outcome(deployment=COPILOT_DEPLOYMENT, status=401),
|
||||
Outcome(deployment=OPERATOR_DEPLOYMENT, status=401),
|
||||
)
|
||||
assert [r.headers.get("authorization") for r in received] == [
|
||||
f"Bearer {CALLER_SESSION_TOKEN}",
|
||||
f"Bearer {OPERATOR_KEY}",
|
||||
]
|
||||
|
||||
|
||||
def test_a_provider_401_on_a_callers_copilot_session_stays_out_of_the_shared_deployment_metrics(
|
||||
ran: tuple[WorkerReport, tuple[Request, ...]],
|
||||
) -> None:
|
||||
report, _ = ran
|
||||
assert {s.metric for s in report.series if s.model_id == COPILOT_DEPLOYMENT} & DEPLOYMENT_SERIES == set()
|
||||
|
||||
|
||||
def test_a_provider_401_on_the_deployments_own_key_counts_as_a_deployment_failure(
|
||||
ran: tuple[WorkerReport, tuple[Request, ...]],
|
||||
) -> None:
|
||||
report, _ = ran
|
||||
assert report.values(FAILURES, OPERATOR_DEPLOYMENT) == (1.0,)
|
||||
assert report.values(TOTALS, OPERATOR_DEPLOYMENT) == (1.0,)
|
||||
assert report.values(LLM_FAILED, OPERATOR_DEPLOYMENT) == (1.0,)
|
||||
assert report.values(STATE, OPERATOR_DEPLOYMENT) == (1.0,)
|
||||
|
|
@ -2,7 +2,8 @@
|
|||
Unit tests for Prometheus invalid API key request filtering.
|
||||
|
||||
Tests the 401 detection helpers, that LLM-level metrics skip invalid API key
|
||||
requests, and that the proxy-level failed request counter still records them.
|
||||
requests rejected at the proxy gate while recording a provider's 401, and that
|
||||
the proxy-level failed request counter still records them.
|
||||
"""
|
||||
|
||||
from unittest.mock import Mock, patch
|
||||
|
|
@ -11,7 +12,10 @@ import pytest
|
|||
from fastapi import HTTPException
|
||||
from prometheus_client import REGISTRY
|
||||
|
||||
import litellm
|
||||
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.llms.github_copilot.per_user_auth import GithubCopilotUserSession
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
|
||||
|
||||
|
||||
|
|
@ -48,18 +52,13 @@ class TestExtractStatusCode:
|
|||
(ExceptionWithStatusCode, 401, 401),
|
||||
],
|
||||
)
|
||||
def test_extract_from_exception(
|
||||
self, prometheus_logger, exception_class, code_value, expected
|
||||
):
|
||||
def test_extract_from_exception(self, prometheus_logger, exception_class, code_value, expected):
|
||||
exception = exception_class(code_value)
|
||||
assert prometheus_logger._extract_status_code(exception=exception) == expected
|
||||
|
||||
def test_extract_from_kwargs(self, prometheus_logger):
|
||||
exception = ExceptionWithCode("401")
|
||||
assert (
|
||||
prometheus_logger._extract_status_code(kwargs={"exception": exception})
|
||||
== 401
|
||||
)
|
||||
assert prometheus_logger._extract_status_code(kwargs={"exception": exception}) == 401
|
||||
|
||||
def test_extract_from_enum_values(self, prometheus_logger):
|
||||
enum_values = Mock(status_code="401")
|
||||
|
|
@ -79,50 +78,133 @@ class TestInvalidAPIKeyDetection:
|
|||
],
|
||||
)
|
||||
def test_status_code_detection(self, prometheus_logger, status_code, expected):
|
||||
assert (
|
||||
prometheus_logger._is_invalid_api_key_request(status_code=status_code)
|
||||
== expected
|
||||
)
|
||||
assert prometheus_logger._is_invalid_api_key_request(status_code=status_code) == expected
|
||||
|
||||
def test_auth_error_message_detection(self, prometheus_logger):
|
||||
exception = AssertionError(
|
||||
"LiteLLM Virtual Key expected. Received=invalid-key-12345, expected to start with 'sk-'."
|
||||
)
|
||||
assert (
|
||||
prometheus_logger._is_invalid_api_key_request(
|
||||
status_code=None, exception=exception
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert prometheus_logger._is_invalid_api_key_request(status_code=None, exception=exception) is True
|
||||
|
||||
def test_non_auth_exception_not_detected(self, prometheus_logger):
|
||||
exception = ValueError("Some other error")
|
||||
assert (
|
||||
prometheus_logger._is_invalid_api_key_request(
|
||||
status_code=None, exception=exception
|
||||
)
|
||||
is False
|
||||
)
|
||||
assert prometheus_logger._is_invalid_api_key_request(status_code=None, exception=exception) is False
|
||||
|
||||
|
||||
class TestSkipMetricsValidation:
|
||||
"""Test high-level validation method that orchestrates detection and extraction."""
|
||||
|
||||
def test_skip_for_401_exception(self, prometheus_logger):
|
||||
"""Test full flow: extraction -> detection -> skip decision."""
|
||||
exception = ExceptionWithCode("401")
|
||||
assert (
|
||||
prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception)
|
||||
is True
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"exception,proxy_gate,expected",
|
||||
[
|
||||
(ExceptionWithCode("401"), True, True),
|
||||
(AssertionError("expected to start with 'sk-'"), True, True),
|
||||
(ExceptionWithStatusCode(429), True, False),
|
||||
(ExceptionWithCode("401"), False, True),
|
||||
(AssertionError("expected to start with 'sk-'"), False, True),
|
||||
(
|
||||
litellm.AuthenticationError(
|
||||
message="Incorrect API key provided: sk-rotat*****",
|
||||
llm_provider="openai",
|
||||
model="gpt-4o-mini",
|
||||
),
|
||||
False,
|
||||
False,
|
||||
),
|
||||
(
|
||||
litellm.AuthenticationError(
|
||||
message="API key not valid. Please pass a valid API key.",
|
||||
llm_provider="gemini",
|
||||
model="gemini-2.5-flash",
|
||||
),
|
||||
False,
|
||||
False,
|
||||
),
|
||||
(
|
||||
litellm.AuthenticationError(
|
||||
message="Authentication Error, Invalid proxy server token passed.",
|
||||
llm_provider="openai",
|
||||
model="gpt-4o-mini",
|
||||
),
|
||||
True,
|
||||
True,
|
||||
),
|
||||
(
|
||||
litellm.CallerCredentialAuthenticationError(
|
||||
message="GitHub Copilot connection is missing for this user",
|
||||
llm_provider="github_copilot",
|
||||
model="gpt-4o",
|
||||
),
|
||||
False,
|
||||
True,
|
||||
),
|
||||
(
|
||||
HTTPException(status_code=401, detail="caller did not present a token"),
|
||||
False,
|
||||
True,
|
||||
),
|
||||
(
|
||||
litellm.BadRequestError(
|
||||
message="API key not valid. Please pass a valid API key.",
|
||||
llm_provider="gemini",
|
||||
model="gemini-2.5-flash",
|
||||
),
|
||||
False,
|
||||
False,
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"proxy-gate-401",
|
||||
"proxy-gate-auth-message",
|
||||
"proxy-gate-429",
|
||||
"unmarked-proxy-401",
|
||||
"unmarked-auth-message",
|
||||
"provider-401",
|
||||
"provider-auth-message",
|
||||
"proxy-gate-litellm-auth-error",
|
||||
"caller-credential-401",
|
||||
"unmarked-http-exception-401",
|
||||
"provider-400-auth-message",
|
||||
],
|
||||
)
|
||||
def test_skip_except_for_provider_errors_the_proxy_did_not_raise(
|
||||
self, prometheus_logger, exception, proxy_gate, expected
|
||||
):
|
||||
kwargs = {"exception": exception}
|
||||
if proxy_gate:
|
||||
kwargs[LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL] = True
|
||||
assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is expected
|
||||
|
||||
def test_skip_for_auth_error_message(self, prometheus_logger):
|
||||
"""Test full flow: exception message -> detection -> skip decision."""
|
||||
exception = AssertionError("expected to start with 'sk-'")
|
||||
assert (
|
||||
prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception)
|
||||
is True
|
||||
@pytest.mark.parametrize(
|
||||
"litellm_params,expected",
|
||||
[
|
||||
({"proxy_server_request": {"credential_fields": ("api_key",)}}, True),
|
||||
(
|
||||
{
|
||||
"github_copilot_user_session": GithubCopilotUserSession(
|
||||
token="caller-session-token", api_base="https://api.githubcopilot.com"
|
||||
)
|
||||
},
|
||||
True,
|
||||
),
|
||||
({"proxy_server_request": {"credential_fields": ("headers",)}}, False),
|
||||
({"proxy_server_request": {"credential_fields": "api_key"}}, False),
|
||||
],
|
||||
ids=[
|
||||
"caller-supplied-api-key",
|
||||
"caller-copilot-session",
|
||||
"caller-supplied-headers-only",
|
||||
"malformed-credential-fields",
|
||||
],
|
||||
)
|
||||
def test_skip_a_provider_401_on_a_credential_the_caller_owns(self, prometheus_logger, litellm_params, expected):
|
||||
exception = litellm.AuthenticationError(
|
||||
message="Incorrect API key provided: sk-calle*****",
|
||||
llm_provider="openai",
|
||||
model="gpt-4o-mini",
|
||||
)
|
||||
kwargs = {"exception": exception, "litellm_params": litellm_params}
|
||||
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
|
||||
|
|
@ -147,19 +229,13 @@ class TestAsyncHooks:
|
|||
),
|
||||
],
|
||||
)
|
||||
async def test_post_call_failure_hook_counts_401_without_key_hash(
|
||||
self, prometheus_logger, exception
|
||||
):
|
||||
async def test_post_call_failure_hook_counts_401_without_key_hash(self, prometheus_logger, exception):
|
||||
unauthenticated = UserAPIKeyAuth(request_route="/v1/chat/completions")
|
||||
unauthenticated.api_key = "notakeyatall"
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
prometheus_logger, "litellm_proxy_failed_requests_metric"
|
||||
) as mock_failed,
|
||||
patch.object(
|
||||
prometheus_logger, "litellm_proxy_total_requests_metric"
|
||||
) as mock_total,
|
||||
patch.object(prometheus_logger, "litellm_proxy_failed_requests_metric") as mock_failed,
|
||||
patch.object(prometheus_logger, "litellm_proxy_total_requests_metric") as mock_total,
|
||||
):
|
||||
await prometheus_logger.async_post_call_failure_hook(
|
||||
request_data={"model": "test-model"},
|
||||
|
|
@ -176,9 +252,7 @@ class TestAsyncHooks:
|
|||
mock_total.labels.return_value.inc.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_failure_hook_keeps_resolved_identity_labels_for_401(
|
||||
self, prometheus_logger
|
||||
):
|
||||
async def test_post_call_failure_hook_keeps_resolved_identity_labels_for_401(self, prometheus_logger):
|
||||
expired_key = UserAPIKeyAuth(
|
||||
api_key="sk-expired",
|
||||
key_alias="expired-alias",
|
||||
|
|
@ -191,9 +265,7 @@ class TestAsyncHooks:
|
|||
code=401,
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
prometheus_logger, "litellm_proxy_failed_requests_metric"
|
||||
) as mock_failed:
|
||||
with patch.object(prometheus_logger, "litellm_proxy_failed_requests_metric") as mock_failed:
|
||||
await prometheus_logger.async_post_call_failure_hook(
|
||||
request_data={"model": "test-model"},
|
||||
original_exception=exception,
|
||||
|
|
@ -207,7 +279,7 @@ class TestAsyncHooks:
|
|||
assert failed_labels["team"] == "team-1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_failure_event_skips_401(self, prometheus_logger):
|
||||
async def test_log_failure_event_skips_401_rejected_at_the_proxy_gate(self, prometheus_logger):
|
||||
exception = ExceptionWithCode("401")
|
||||
kwargs = {
|
||||
"model": "test-model",
|
||||
|
|
@ -220,20 +292,90 @@ class TestAsyncHooks:
|
|||
},
|
||||
"exception": exception,
|
||||
"litellm_params": {},
|
||||
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True,
|
||||
}
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
prometheus_logger, "litellm_llm_api_failed_requests_metric"
|
||||
) as mock_failed,
|
||||
patch.object(
|
||||
prometheus_logger, "set_llm_deployment_failure_metrics"
|
||||
) as mock_deployment,
|
||||
patch.object(prometheus_logger, "litellm_llm_api_failed_requests_metric") as mock_failed,
|
||||
patch.object(prometheus_logger, "set_llm_deployment_failure_metrics") as mock_deployment,
|
||||
):
|
||||
|
||||
await prometheus_logger.async_log_failure_event(
|
||||
kwargs=kwargs, response_obj=None, start_time=None, end_time=None
|
||||
)
|
||||
|
||||
mock_failed.labels.assert_not_called()
|
||||
mock_deployment.assert_not_called()
|
||||
|
||||
|
||||
def _provider_auth_failure_kwargs() -> dict:
|
||||
exception = litellm.AuthenticationError(
|
||||
message="Incorrect API key provided: sk-rotat*****",
|
||||
llm_provider="openai",
|
||||
model="gpt-4o-mini",
|
||||
)
|
||||
return {
|
||||
"model": "gpt-4o-mini",
|
||||
"standard_logging_object": {
|
||||
"metadata": {
|
||||
"user_api_key_hash": "test-key",
|
||||
"user_api_key_alias": "test-alias",
|
||||
"user_api_key_user_id": "test-user",
|
||||
"user_api_key_team_id": "test-team",
|
||||
"user_api_key_team_alias": "test-team-alias",
|
||||
},
|
||||
"model_group": "rotated-credential",
|
||||
"model_id": "deployment-1",
|
||||
"api_base": "https://api.openai.com/v1",
|
||||
"custom_llm_provider": "openai",
|
||||
},
|
||||
"exception": exception,
|
||||
"litellm_params": {"custom_llm_provider": "openai", "metadata": {}},
|
||||
}
|
||||
|
||||
|
||||
def _metric_values(name: str, model_id: str) -> list[float]:
|
||||
return [
|
||||
sample.value
|
||||
for sample in _registry_samples()
|
||||
if sample.name == name and sample.labels.get("model_id") == model_id
|
||||
]
|
||||
|
||||
|
||||
def _registry_samples():
|
||||
for metric in REGISTRY.collect():
|
||||
yield from metric.samples
|
||||
|
||||
|
||||
class TestProviderAuthFailures:
|
||||
"""A 401 raised by the provider is an upstream failure and is recorded."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_failure_event_records_a_provider_401(self, prometheus_logger):
|
||||
await prometheus_logger.async_log_failure_event(
|
||||
kwargs=_provider_auth_failure_kwargs(),
|
||||
response_obj=None,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
|
||||
assert _metric_values("litellm_llm_api_failed_requests_metric_total", "deployment-1") == [1.0]
|
||||
assert _metric_values("litellm_deployment_failure_responses_total", "deployment-1") == [1.0]
|
||||
assert _metric_values("litellm_deployment_state", "deployment-1") == [1.0]
|
||||
|
||||
def test_deployment_failure_metrics_record_a_provider_401(self, prometheus_logger):
|
||||
prometheus_logger.set_llm_deployment_failure_metrics(_provider_auth_failure_kwargs())
|
||||
|
||||
failures = [
|
||||
sample for sample in _registry_samples() if sample.name == "litellm_deployment_failure_responses_total"
|
||||
]
|
||||
assert [
|
||||
(
|
||||
sample.labels["model_id"],
|
||||
sample.labels["exception_status"],
|
||||
sample.labels["api_provider"],
|
||||
sample.value,
|
||||
)
|
||||
for sample in failures
|
||||
] == [("deployment-1", "401", "openai", 1.0)]
|
||||
assert _metric_values("litellm_deployment_total_requests_total", "deployment-1") == [1.0]
|
||||
assert _metric_values("litellm_deployment_state", "deployment-1") == [1.0]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue