mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
Merge branch 'litellm_internal_staging' into litellm_auto_router_compression_split
This commit is contained in:
moe-berri
GitHub
commit
fc3da5e830
32 changed files with 1189 additions and 110 deletions
|
|
@ -433,9 +433,9 @@ _default_detect_secrets_config = {
|
|||
"name": "ZendeskSecretKeyDetector",
|
||||
"path": _custom_plugins_path + "/zendesk_secret_key.py",
|
||||
},
|
||||
{"name": "Base64HighEntropyString", "limit": 3.0},
|
||||
{"name": "Base64HighEntropyString", "limit": 4.5},
|
||||
{"name": "HexHighEntropyString", "limit": 3.0},
|
||||
]
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -466,16 +466,19 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
|
|||
|
||||
os.remove(temp_file.name)
|
||||
|
||||
detected_secrets = []
|
||||
for file in secrets.files:
|
||||
for found_secret in secrets[file]:
|
||||
if found_secret.secret_value is None:
|
||||
continue
|
||||
detected_secrets.append(
|
||||
{"type": found_secret.type, "value": found_secret.secret_value}
|
||||
)
|
||||
|
||||
return detected_secrets
|
||||
return [
|
||||
{"type": found_secret.type, "value": found_secret.secret_value}
|
||||
for file in sorted(secrets.files)
|
||||
for found_secret in sorted(
|
||||
secrets[file],
|
||||
key=lambda secret: (
|
||||
-len(secret.secret_value or ""),
|
||||
secret.type,
|
||||
secret.secret_value or "",
|
||||
),
|
||||
)
|
||||
if found_secret.secret_value is not None
|
||||
]
|
||||
|
||||
def redact_text(self, text: str, source: str = "message") -> str:
|
||||
"""Replace every detected secret in ``text`` with ``[REDACTED]`` and
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ This plugin searches for OpenAI API Keys.
|
|||
"""
|
||||
|
||||
import re
|
||||
from collections.abc import Generator
|
||||
|
||||
from detect_secrets.plugins.base import RegexBasedDetector
|
||||
|
||||
|
|
@ -16,4 +17,16 @@ class OpenAIApiKeyDetector(RegexBasedDetector):
|
|||
|
||||
@property
|
||||
def denylist(self) -> list[re.Pattern]:
|
||||
return [re.compile(r"""(sk-[a-zA-Z0-9]{5,})""")]
|
||||
return [
|
||||
re.compile(
|
||||
r"((?:(?<![a-zA-Z0-9])|(?<=%[0-9A-Fa-f]{2}))"
|
||||
r"sk[-_]"
|
||||
r"[a-zA-Z0-9_-]{5,}"
|
||||
r"(?![a-zA-Z0-9_-]))"
|
||||
)
|
||||
]
|
||||
|
||||
def analyze_string(self, string: str) -> Generator[str, None, None]:
|
||||
# the digit check lives outside the regex: a lookahead re-scans the token
|
||||
# from every `sk` inside it, which is quadratic on `-sk-sk-sk-...` input
|
||||
yield from (match for match in super().analyze_string(string) if re.search(r"[0-9]", match))
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.64"
|
||||
version = "0.1.65"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.64"
|
||||
version = "0.1.65"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.93"
|
||||
version = "0.4.94"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.93"
|
||||
version = "0.4.94"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -215,8 +215,8 @@ MAX_CALLBACKS: Final = get_env_int("LITELLM_MAX_CALLBACKS", 100)
|
|||
# so the deployment-level hook does not re-run them for the same request
|
||||
PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails"
|
||||
|
||||
# Metadata key listing compression guardrails an auto router's own compression
|
||||
# policy suppresses for this request. See litellm.proxy.guardrails.auto_router_compression.
|
||||
# Attribute stamped on log_guardrail_information wrappers so __init_subclass__ does not wrap them again
|
||||
LOGS_GUARDRAIL_INFORMATION_MARKER: Final = "_litellm_logs_guardrail_information"
|
||||
|
||||
# Generic fallback for unknown models
|
||||
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int(
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ dc: Final = DualCache()
|
|||
|
||||
from litellm.constants import (
|
||||
GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS,
|
||||
LOGS_GUARDRAIL_INFORMATION_MARKER,
|
||||
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
|
||||
)
|
||||
from litellm.exceptions import (
|
||||
|
|
@ -151,6 +152,13 @@ class CustomGuardrail(CustomLogger):
|
|||
|
||||
records_own_guardrail_information: ClassVar[bool] = False
|
||||
|
||||
def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: forwarded to cooperative __init_subclass__ hooks
|
||||
super().__init_subclass__(**kwargs)
|
||||
own_apply_guardrail: Final = cls.__dict__.get("apply_guardrail")
|
||||
if own_apply_guardrail is None or LOGS_GUARDRAIL_INFORMATION_MARKER in vars(own_apply_guardrail):
|
||||
return
|
||||
cls.apply_guardrail = log_guardrail_information(own_apply_guardrail)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
guardrail_name: str | None = None,
|
||||
|
|
@ -1579,4 +1587,5 @@ def log_guardrail_information(func):
|
|||
return async_wrapper(*args, **kwargs)
|
||||
return sync_wrapper(*args, **kwargs)
|
||||
|
||||
vars(wrapper)[LOGS_GUARDRAIL_INFORMATION_MARKER] = True # rebind-ok: stamps the wrapper this call just built
|
||||
return wrapper
|
||||
|
|
|
|||
|
|
@ -5,6 +5,20 @@ from typing import Final
|
|||
from fastapi import HTTPException
|
||||
|
||||
|
||||
class MCPServerURLCredentialsError(HTTPException):
|
||||
"""A fixed, sanitized URL-credential migration error safe for operator previews."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(
|
||||
status_code=500,
|
||||
detail=(
|
||||
"misconfigured: auth_type none cannot be used with credentials embedded in the upstream URL; "
|
||||
"remove them from the URL and configure Basic Auth with auth_type: basic and "
|
||||
"auth_value: username:password"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class MCPUpstreamAuthError(Exception):
|
||||
"""Raised when an upstream MCP server returns an authentication failure
|
||||
(typically HTTP 401) and the gateway should surface it transparently to
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from pydantic import SecretStr
|
|||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.experimental_mcp_client.client import strip_auth_scheme, to_basic_credentials
|
||||
from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError
|
||||
from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
|
||||
DEFAULT_CREDENTIAL_HEADER,
|
||||
|
|
@ -293,6 +294,8 @@ def raise_public(error: CredError) -> NoReturn:
|
|||
)
|
||||
case "misconfigured":
|
||||
raise HTTPException(status_code=500, detail=error.summary)
|
||||
case "url_credentials_not_allowed":
|
||||
raise MCPServerURLCredentialsError()
|
||||
case "upstream_unavailable":
|
||||
raise HTTPException(status_code=503, detail=error.summary)
|
||||
case "unsupported_mode":
|
||||
|
|
|
|||
|
|
@ -134,7 +134,7 @@ class UpstreamCredentialProvider:
|
|||
async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]:
|
||||
match server.config:
|
||||
case NoneConfig():
|
||||
return Ok(NoOpAuth())
|
||||
return self._none(server)
|
||||
case ApiKeyConfig() as config:
|
||||
return self._api_key(config)
|
||||
case PassthroughConfig():
|
||||
|
|
@ -151,6 +151,15 @@ class UpstreamCredentialProvider:
|
|||
return _not_implemented(AuthSpecKind.aws_sigv4)
|
||||
assert_never(server.config)
|
||||
|
||||
def _none(self, server: ServerSpec) -> Result[httpx.Auth, CredError]:
|
||||
try:
|
||||
resource: Final = httpx.URL(server.resource)
|
||||
except httpx.InvalidURL:
|
||||
return Ok(NoOpAuth())
|
||||
if resource.userinfo:
|
||||
return Error(CredError.of_url_credentials_not_allowed())
|
||||
return Ok(NoOpAuth())
|
||||
|
||||
async def has_user_token(self, subject: Subject, server: ServerSpec) -> bool:
|
||||
"""Whether a usable per-user token exists for this server (the preemptive 401's check).
|
||||
|
||||
|
|
|
|||
|
|
@ -95,6 +95,7 @@ class CredError:
|
|||
tag: Literal[
|
||||
"unauthorized",
|
||||
"misconfigured",
|
||||
"url_credentials_not_allowed",
|
||||
"upstream_unavailable",
|
||||
"unsupported_mode",
|
||||
"precondition_required",
|
||||
|
|
@ -103,6 +104,7 @@ class CredError:
|
|||
|
||||
unauthorized: Unauthorized = case() # no usable credential for this (subject, server) -> 401 challenge
|
||||
misconfigured: str = case() # the declared mode is missing required config -> 5xx (operator)
|
||||
url_credentials_not_allowed: None = case()
|
||||
upstream_unavailable: str = case() # the IdP / token endpoint could not be reached -> 503
|
||||
unsupported_mode: str = case() # a raw mode string did not parse into AuthSpecKind (boundary)
|
||||
precondition_required: str = case() # a required per-user value (e.g. an env var) has not been provided -> 412
|
||||
|
|
@ -129,6 +131,10 @@ class CredError:
|
|||
def of_misconfigured(detail: str) -> CredError:
|
||||
return CredError(misconfigured=detail)
|
||||
|
||||
@staticmethod
|
||||
def of_url_credentials_not_allowed() -> CredError:
|
||||
return CredError(url_credentials_not_allowed=None)
|
||||
|
||||
@staticmethod
|
||||
def of_upstream_unavailable(detail: str) -> CredError:
|
||||
return CredError(upstream_unavailable=detail)
|
||||
|
|
@ -154,6 +160,12 @@ class CredError:
|
|||
return f"unauthorized: {self.unauthorized.detail}"
|
||||
case "misconfigured":
|
||||
return f"misconfigured: {self.misconfigured}"
|
||||
case "url_credentials_not_allowed":
|
||||
return (
|
||||
"misconfigured: auth_type none cannot be used with credentials embedded in the upstream URL; "
|
||||
"remove them from the URL and configure Basic Auth with auth_type: basic and "
|
||||
"auth_value: username:password"
|
||||
)
|
||||
case "upstream_unavailable":
|
||||
return f"upstream unavailable: {self.upstream_unavailable}"
|
||||
case "unsupported_mode":
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from litellm.exceptions import (
|
|||
)
|
||||
from litellm.proxy._experimental.mcp_server.exceptions import (
|
||||
MCPServerListError,
|
||||
MCPServerURLCredentialsError,
|
||||
MCPUpstreamAuthError,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
|
||||
|
|
@ -75,6 +76,8 @@ _MCP_GUARDRAIL_REJECTIONS: Final = (
|
|||
|
||||
|
||||
def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str:
|
||||
if isinstance(exc, MCPServerURLCredentialsError):
|
||||
return str(exc.detail)
|
||||
if isinstance(exc, TimeoutError):
|
||||
return (
|
||||
f"Failed to connect to MCP server: no response from {url or 'the server'} "
|
||||
|
|
|
|||
|
|
@ -199,6 +199,12 @@ class PrismaDBExceptionHandler:
|
|||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def is_prisma_error(e: Exception) -> bool:
|
||||
import prisma
|
||||
|
||||
return isinstance(e, _exception_types(prisma.errors.PrismaError))
|
||||
|
||||
@staticmethod
|
||||
def is_deadlock_error(e: Exception) -> bool:
|
||||
"""True iff ``e`` is a Postgres deadlock (P2034 / 40P01) surfaced through prisma."""
|
||||
|
|
|
|||
|
|
@ -6386,10 +6386,17 @@ class ProxyUpdateSpend:
|
|||
)
|
||||
break
|
||||
except Exception as e:
|
||||
if not PrismaDBExceptionHandler.is_database_transport_error(e):
|
||||
if not _is_transient_spend_log_write_error(e):
|
||||
if PrismaDBExceptionHandler.is_prisma_error(e):
|
||||
await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True)
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - DB error writing spend logs, requeued %d rows for the next flush. error=%s",
|
||||
len(logs_to_process),
|
||||
str(e),
|
||||
)
|
||||
raise
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - DB connection error writing spend logs, retry %d/%d. logs_count=%d, error=%s",
|
||||
"Spend tracking - transient DB error writing spend logs, retry %d/%d. logs_count=%d, error=%s",
|
||||
i + 1,
|
||||
n_retry_times,
|
||||
len(logs_to_process),
|
||||
|
|
@ -6732,6 +6739,10 @@ async def _monitor_spend_logs_queue(
|
|||
MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256
|
||||
|
||||
|
||||
def _is_transient_spend_log_write_error(e: Exception) -> bool:
|
||||
return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e)
|
||||
|
||||
|
||||
async def _create_spend_logs_with_poison_isolation(
|
||||
repo: SpendLogsRepository,
|
||||
rows: Sequence[Mapping[str, object]],
|
||||
|
|
@ -6767,6 +6778,8 @@ async def _create_spend_logs_with_poison_isolation(
|
|||
raise
|
||||
if PrismaDBExceptionHandler.is_database_service_unavailable_error(e):
|
||||
raise
|
||||
if PrismaDBExceptionHandler.is_deadlock_error(e):
|
||||
raise
|
||||
budget_left: Final = max(failure_budget - 1, 0)
|
||||
if len(rows) == 1:
|
||||
request_id: Final = rows[0].get("request_id")
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import asyncio
|
|||
import contextvars
|
||||
from collections.abc import Coroutine, Generator, Iterable, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
|
||||
|
||||
|
|
@ -23,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
)
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig
|
||||
from litellm.responses.litellm_completion_transformation.handler import (
|
||||
LiteLLMCompletionTransformationHandler,
|
||||
)
|
||||
|
|
@ -403,8 +405,40 @@ def _bridges_to_chat_completions(
|
|||
return responses_api_provider_config is None or use_chat_completions_api is True
|
||||
|
||||
|
||||
def _deployment_passes_through_responses(model_info: object) -> bool:
|
||||
"""Whether ``model_info.supported_endpoints`` opts the deployment into native ``{api_base}/responses``."""
|
||||
if not isinstance(model_info, dict):
|
||||
return False
|
||||
supported_endpoints: Final = model_info.get("supported_endpoints")
|
||||
return isinstance(supported_endpoints, (list, tuple)) and "/v1/responses" in supported_endpoints
|
||||
|
||||
|
||||
def _deployment_model_info_after_prompt_swap(
|
||||
requested_provider: str | None, resolved_provider: str | None, model_info: object
|
||||
) -> object:
|
||||
"""Deployment metadata only describes the upstream while the prompt manager keeps its provider."""
|
||||
return model_info if resolved_provider == requested_provider else None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _AsyncPromptManagementOutcome:
|
||||
merged_optional_params: Mapping[str, object]
|
||||
deployment_model_info: object
|
||||
|
||||
|
||||
def _resolve_responses_api_provider_config(
|
||||
model: str, custom_llm_provider: str, model_info: object
|
||||
) -> BaseResponsesAPIConfig | None:
|
||||
provider_config: Final = ProviderConfigManager.get_provider_responses_api_config(
|
||||
model=model, provider=custom_llm_provider
|
||||
)
|
||||
if provider_config is not None or not _deployment_passes_through_responses(model_info):
|
||||
return provider_config
|
||||
return OpenAILikeResponsesConfig()
|
||||
|
||||
|
||||
def _will_bridge_to_chat_completions(
|
||||
model: str, custom_llm_provider: str | None, use_chat_completions_api: bool
|
||||
model: str, custom_llm_provider: str | None, use_chat_completions_api: bool, model_info: object
|
||||
) -> bool:
|
||||
"""``_bridges_to_chat_completions`` for callers running before the provider config is resolved.
|
||||
|
||||
|
|
@ -418,9 +452,7 @@ def _will_bridge_to_chat_completions(
|
|||
if custom_llm_provider is None:
|
||||
return True
|
||||
return _bridges_to_chat_completions(
|
||||
ProviderConfigManager.get_provider_responses_api_config(
|
||||
model=normalized_model[0], provider=custom_llm_provider
|
||||
),
|
||||
_resolve_responses_api_provider_config(normalized_model[0], custom_llm_provider, model_info),
|
||||
use_chat_completions_api or normalized_model[1],
|
||||
)
|
||||
|
||||
|
|
@ -527,7 +559,10 @@ async def aresponses(
|
|||
with _prompt_management_sees_a_provisional_message_list(
|
||||
kwargs,
|
||||
bridged=_will_bridge_to_chat_completions(
|
||||
model, custom_llm_provider, bool(kwargs.get("use_chat_completions_api"))
|
||||
model,
|
||||
custom_llm_provider,
|
||||
bool(kwargs.get("use_chat_completions_api")),
|
||||
kwargs.get("model_info"),
|
||||
),
|
||||
):
|
||||
(
|
||||
|
|
@ -552,6 +587,7 @@ async def aresponses(
|
|||
merged_input=merged_input,
|
||||
),
|
||||
)
|
||||
requested_provider: Final = custom_llm_provider
|
||||
if model != original_model:
|
||||
custom_llm_provider = _resolve_prompt_swapped_provider(
|
||||
original_model=original_model,
|
||||
|
|
@ -561,7 +597,12 @@ async def aresponses(
|
|||
prompt_id=prompt_id,
|
||||
)
|
||||
kwargs.pop("prompt_id", None)
|
||||
kwargs["_async_prompt_merged_params"] = merged_optional_params
|
||||
kwargs["_async_prompt_merged_params"] = _AsyncPromptManagementOutcome(
|
||||
merged_optional_params=merged_optional_params,
|
||||
deployment_model_info=_deployment_model_info_after_prompt_swap(
|
||||
requested_provider, custom_llm_provider, kwargs.get("model_info")
|
||||
),
|
||||
)
|
||||
|
||||
func: Final = partial(
|
||||
responses,
|
||||
|
|
@ -666,12 +707,14 @@ def _apply_prompt_management_to_responses_call(
|
|||
kwargs: dict[str, Any],
|
||||
local_vars: dict[str, object],
|
||||
use_chat_completions_api: bool,
|
||||
) -> tuple[str | ResponseInputParam, str, str | None]:
|
||||
async_merged: Final[Mapping[str, object] | None] = kwargs.pop("_async_prompt_merged_params", None)
|
||||
if async_merged is not None:
|
||||
for key, value in async_merged.items():
|
||||
) -> tuple[str | ResponseInputParam, str, str | None, object]:
|
||||
"""Returns the prompt-managed input, model and provider, plus the deployment metadata that still
|
||||
describes the upstream (``None`` once the prompt manager moved the request to another provider)."""
|
||||
async_outcome: Final[_AsyncPromptManagementOutcome | None] = kwargs.pop("_async_prompt_merged_params", None)
|
||||
if async_outcome is not None:
|
||||
for key, value in async_outcome.merged_optional_params.items():
|
||||
local_vars[key] = value
|
||||
return input, model, custom_llm_provider
|
||||
return input, model, custom_llm_provider, async_outcome.deployment_model_info
|
||||
|
||||
prompt_id: Final = cast(str | None, kwargs.get("prompt_id", None))
|
||||
prompt_variables: Final = cast(dict | None, kwargs.get("prompt_variables", None))
|
||||
|
|
@ -684,7 +727,9 @@ def _apply_prompt_management_to_responses_call(
|
|||
):
|
||||
with _prompt_management_sees_a_provisional_message_list(
|
||||
kwargs,
|
||||
bridged=_will_bridge_to_chat_completions(model, custom_llm_provider, use_chat_completions_api),
|
||||
bridged=_will_bridge_to_chat_completions(
|
||||
model, custom_llm_provider, use_chat_completions_api, kwargs.get("model_info")
|
||||
),
|
||||
):
|
||||
(
|
||||
model,
|
||||
|
|
@ -710,19 +755,28 @@ def _apply_prompt_management_to_responses_call(
|
|||
)
|
||||
local_vars["input"] = input
|
||||
local_vars["model"] = model
|
||||
if model != original_model:
|
||||
custom_llm_provider = _resolve_prompt_swapped_provider(
|
||||
resolved_provider: Final = (
|
||||
custom_llm_provider
|
||||
if model == original_model
|
||||
else _resolve_prompt_swapped_provider(
|
||||
original_model=original_model,
|
||||
swapped_model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs,
|
||||
prompt_id=prompt_id,
|
||||
)
|
||||
local_vars["custom_llm_provider"] = custom_llm_provider
|
||||
)
|
||||
local_vars["custom_llm_provider"] = resolved_provider
|
||||
for key, value in merged_optional_params.items():
|
||||
local_vars[key] = value
|
||||
return (
|
||||
input,
|
||||
model,
|
||||
resolved_provider,
|
||||
_deployment_model_info_after_prompt_swap(custom_llm_provider, resolved_provider, kwargs.get("model_info")),
|
||||
)
|
||||
|
||||
return input, model, custom_llm_provider
|
||||
return input, model, custom_llm_provider, kwargs.get("model_info")
|
||||
|
||||
|
||||
# Opt-in via model id (mirrors the `responses/` prefix pattern on chat completions).
|
||||
|
|
@ -1052,7 +1106,7 @@ def responses(
|
|||
)
|
||||
local_vars["custom_llm_provider"] = custom_llm_provider
|
||||
|
||||
input, model, custom_llm_provider = _apply_prompt_management_to_responses_call(
|
||||
input, model, custom_llm_provider, deployment_model_info = _apply_prompt_management_to_responses_call(
|
||||
input=input,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
|
|
@ -1123,9 +1177,8 @@ def responses(
|
|||
if custom_llm_provider is None:
|
||||
responses_api_provider_config = None
|
||||
else:
|
||||
responses_api_provider_config = ProviderConfigManager.get_provider_responses_api_config(
|
||||
model=model,
|
||||
provider=custom_llm_provider,
|
||||
responses_api_provider_config = _resolve_responses_api_provider_config(
|
||||
model, custom_llm_provider, deployment_model_info
|
||||
)
|
||||
|
||||
local_vars.update(kwargs)
|
||||
|
|
|
|||
|
|
@ -2597,14 +2597,20 @@ class Router:
|
|||
model_response: CustomStreamWrapper,
|
||||
messages: list[dict[str, str]],
|
||||
initial_kwargs: dict,
|
||||
deployment_slot: contextlib.AsyncExitStack | None = None,
|
||||
) -> CustomStreamWrapper:
|
||||
"""
|
||||
Helper to iterate over a streaming response.
|
||||
|
||||
Catches errors for fallbacks using the router's fallback system
|
||||
|
||||
`deployment_slot` holds the deployment's max_parallel_requests semaphore; it is
|
||||
released when the stream is exhausted, closed, or falls back to another deployment
|
||||
"""
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
|
||||
held_slot: Final = deployment_slot if deployment_slot is not None else contextlib.AsyncExitStack()
|
||||
|
||||
class FallbackStreamWrapper(CustomStreamWrapper):
|
||||
def __init__(self, async_generator: AsyncGenerator):
|
||||
# Copy attributes from the original model_response
|
||||
|
|
@ -2628,12 +2634,26 @@ class Router:
|
|||
async def __anext__(self):
|
||||
return await self._async_generator.__anext__()
|
||||
|
||||
async def close_model_response() -> None:
|
||||
if not hasattr(model_response, "aclose"):
|
||||
return
|
||||
try:
|
||||
await model_response.aclose()
|
||||
except BaseException as e:
|
||||
verbose_router_logger.debug(
|
||||
"stream_with_fallbacks: error closing model_response: %s",
|
||||
e,
|
||||
)
|
||||
|
||||
async def stream_with_fallbacks():
|
||||
fallback_response = None # Track for cleanup in finally
|
||||
try:
|
||||
async for item in model_response:
|
||||
yield item
|
||||
except MidStreamFallbackError as e:
|
||||
with anyio.CancelScope(shield=True):
|
||||
await close_model_response()
|
||||
await held_slot.aclose()
|
||||
if not e.is_pre_first_chunk and (
|
||||
e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)
|
||||
):
|
||||
|
|
@ -2707,14 +2727,8 @@ class Router:
|
|||
# (e.g. on client disconnect).
|
||||
# Shield from anyio cancellation so the awaits can complete.
|
||||
with anyio.CancelScope(shield=True):
|
||||
if hasattr(model_response, "aclose"):
|
||||
try:
|
||||
await model_response.aclose()
|
||||
except BaseException as e:
|
||||
verbose_router_logger.debug(
|
||||
"stream_with_fallbacks: error closing model_response: %s",
|
||||
e,
|
||||
)
|
||||
await close_model_response()
|
||||
await held_slot.aclose()
|
||||
if fallback_response is not None and hasattr(fallback_response, "aclose"):
|
||||
try:
|
||||
await fallback_response.aclose()
|
||||
|
|
@ -3379,61 +3393,53 @@ class Router:
|
|||
kwargs=kwargs,
|
||||
client_type="max_parallel_requests",
|
||||
)
|
||||
if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore):
|
||||
async with rpm_semaphore:
|
||||
"""
|
||||
- Check rpm limits before making the call
|
||||
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
|
||||
"""
|
||||
await self.async_routing_strategy_pre_call_checks(
|
||||
deployment=deployment,
|
||||
logging_obj=logging_obj,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
response = await _response
|
||||
else:
|
||||
async with contextlib.AsyncExitStack() as deployment_slot:
|
||||
if isinstance(rpm_semaphore, asyncio.Semaphore):
|
||||
await deployment_slot.enter_async_context(rpm_semaphore)
|
||||
await self.async_routing_strategy_pre_call_checks(
|
||||
deployment=deployment,
|
||||
logging_obj=logging_obj,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
|
||||
response = await _response
|
||||
|
||||
## CHECK CONTENT FILTER ERROR ##
|
||||
if isinstance(response, ModelResponse):
|
||||
_should_raise = self._should_raise_content_policy_error(model=model, response=response, kwargs=kwargs)
|
||||
if _should_raise:
|
||||
raise litellm.ContentPolicyViolationError(
|
||||
message="Response output was blocked.",
|
||||
model=model,
|
||||
llm_provider="",
|
||||
## CHECK CONTENT FILTER ERROR ##
|
||||
if isinstance(response, ModelResponse):
|
||||
_should_raise = self._should_raise_content_policy_error(
|
||||
model=model, response=response, kwargs=kwargs
|
||||
)
|
||||
if _should_raise:
|
||||
raise litellm.ContentPolicyViolationError(
|
||||
message="Response output was blocked.",
|
||||
model=model,
|
||||
llm_provider="",
|
||||
)
|
||||
|
||||
if (
|
||||
isinstance(response, CustomStreamWrapper)
|
||||
and response.completion_stream is None
|
||||
and response.make_call is not None
|
||||
):
|
||||
await response.fetch_stream()
|
||||
if (
|
||||
isinstance(response, CustomStreamWrapper)
|
||||
and response.completion_stream is None
|
||||
and response.make_call is not None
|
||||
):
|
||||
await response.fetch_stream()
|
||||
|
||||
self.success_calls[model_name] += 1
|
||||
verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
|
||||
# debug how often this deployment picked
|
||||
self._track_deployment_metrics(
|
||||
deployment=deployment,
|
||||
response=response,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
|
||||
if isinstance(response, CustomStreamWrapper):
|
||||
return await self._acompletion_streaming_iterator(
|
||||
model_response=response,
|
||||
messages=messages,
|
||||
initial_kwargs=input_kwargs_for_streaming_fallback,
|
||||
self.success_calls[model_name] += 1
|
||||
verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
|
||||
# debug how often this deployment picked
|
||||
self._track_deployment_metrics(
|
||||
deployment=deployment,
|
||||
response=response,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
|
||||
return response
|
||||
if isinstance(response, CustomStreamWrapper):
|
||||
return await self._acompletion_streaming_iterator(
|
||||
model_response=response,
|
||||
messages=messages,
|
||||
initial_kwargs=input_kwargs_for_streaming_fallback,
|
||||
deployment_slot=deployment_slot.pop_all(),
|
||||
)
|
||||
|
||||
return response
|
||||
except litellm.Timeout as e:
|
||||
deployment_request_timeout_param: Final = _timeout_debug_deployment_dict.get("litellm_params", {}).get(
|
||||
"request_timeout", None
|
||||
|
|
|
|||
|
|
@ -67,8 +67,8 @@ proxy = [
|
|||
"azure-identity>=1.25.2,<2.0",
|
||||
"azure-storage-blob>=12.28.0,<13.0",
|
||||
"mcp>=1.28.1,<2.0",
|
||||
"litellm-proxy-extras==0.4.93",
|
||||
"litellm-enterprise==0.1.64",
|
||||
"litellm-proxy-extras==0.4.94",
|
||||
"litellm-enterprise==0.1.65",
|
||||
"RestrictedPython>=8.5,<9.0",
|
||||
"rich>=13.9.4,<14.0",
|
||||
"InquirerPy>=0.3.4,<1.0",
|
||||
|
|
|
|||
|
|
@ -851,6 +851,15 @@ class ModelListEntry(BaseModel):
|
|||
id: str
|
||||
|
||||
|
||||
class ModelsListParams(BaseModel):
|
||||
"""Query for GET /v1/models. A wildcard route such as ``openai/gpt-5.4*`` is
|
||||
listed only under ``return_wildcard_routes``; without it the route is dropped
|
||||
and only its expansions remain, so a readiness poll for the pattern itself
|
||||
never resolves."""
|
||||
|
||||
return_wildcard_routes: bool = True
|
||||
|
||||
|
||||
class ModelsListResponse(BaseModel):
|
||||
"""GET /v1/models on the data plane: the deployments the gateway can actually
|
||||
serve right now. Used to confirm a freshly created model has propagated from
|
||||
|
|
|
|||
|
|
@ -55,6 +55,7 @@ from models import (
|
|||
ModelMode,
|
||||
ModelNewBody,
|
||||
ModelNewResponse,
|
||||
ModelsListParams,
|
||||
ModelsListResponse,
|
||||
ModelUpdateBody,
|
||||
OcrBody,
|
||||
|
|
@ -336,7 +337,7 @@ class ProxyClient:
|
|||
lambda poll_timeout: self.transport.get(
|
||||
"/v1/models",
|
||||
headers=headers,
|
||||
params=NoBody(),
|
||||
params=ModelsListParams(),
|
||||
response_type=ModelsListResponse,
|
||||
timeout=poll_timeout,
|
||||
),
|
||||
|
|
|
|||
|
|
@ -597,9 +597,20 @@ class TestSemanticAutoRouterResponses:
|
|||
)
|
||||
)
|
||||
assert answer.id, "/v1/responses through the semantic auto-router returned no response id"
|
||||
rows: Final = proxy.poll_logs_for_key(key, min_rows=1)
|
||||
rows: Final = proxy.poll_logs_for_key(
|
||||
key,
|
||||
min_rows=2,
|
||||
predicate=lambda logged: any(row.model == EMBEDDING_MODEL for row in logged),
|
||||
)
|
||||
embedding_rows: Final = tuple(row for row in rows if row.model == EMBEDDING_MODEL)
|
||||
assert embedding_rows, (
|
||||
"the routing embedding was not billed to the caller's key; "
|
||||
f"spend logs show {tuple(row.model for row in rows)}"
|
||||
)
|
||||
_assert_served_only_by(
|
||||
rows, CHEAP_SERVED | {semantic_auto_router.target}, "semantic auto-router /v1/responses string input"
|
||||
[row for row in rows if row.model != EMBEDDING_MODEL],
|
||||
CHEAP_SERVED | {semantic_auto_router.target},
|
||||
"semantic auto-router /v1/responses string input",
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -10,6 +10,8 @@ Covers the three defects from the ticket:
|
|||
handling live only on the native path).
|
||||
"""
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm_enterprise.enterprise_callbacks.secret_detection import (
|
||||
|
|
@ -19,12 +21,16 @@ from litellm.caching.caching import DualCache
|
|||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
AWS_KEY = "AKIAIOSFODNN7EXAMPLE"
|
||||
OPENAI_KEY = "sk-test-abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGH"
|
||||
SHORT_OPENAI_KEY = "sk-12345"
|
||||
UNICODE_DIGIT_SUFFIX = "sk-notification٣"
|
||||
STRIPE_LIVE_KEY = f"sk_live_{'1234567890' * 3}"
|
||||
URL_ENCODED_KEY = "Bearer%20sk-Ab3dEf6Gh7Ij8Kl9Mn0Pq2Rs3Tu4Vw5X"
|
||||
AWS_KEYS = [f"AKIAIOSFODNN7EXAMPL{suffix}" for suffix in "FEDCBA"]
|
||||
|
||||
|
||||
def _guardrail() -> _ENTERPRISE_SecretDetection:
|
||||
return _ENTERPRISE_SecretDetection(
|
||||
guardrail_name="hide-secrets", event_hook="pre_call", default_on=True
|
||||
)
|
||||
return _ENTERPRISE_SecretDetection(guardrail_name="hide-secrets", event_hook="pre_call", default_on=True)
|
||||
|
||||
|
||||
def _recorded(request_data: dict) -> dict:
|
||||
|
|
@ -33,6 +39,91 @@ def _recorded(request_data: dict) -> dict:
|
|||
return entries[0]
|
||||
|
||||
|
||||
def test_scan_message_preserves_benign_identifiers_and_xml_tags():
|
||||
guardrail = _guardrail()
|
||||
content = "<task-notification> model: claude-sonnet-4-5-20250929 </task-notification>"
|
||||
|
||||
assert guardrail.scan_message_for_secrets(content) == []
|
||||
assert guardrail.redact_text(content) == content
|
||||
assert guardrail.redact_text("result = compute(x) </task-notification>") == (
|
||||
"result = compute(x) </task-notification>"
|
||||
)
|
||||
|
||||
|
||||
def test_scan_message_preserves_quoted_benign_identifiers():
|
||||
guardrail = _guardrail()
|
||||
content = '{"content-type": "application/json", "model": "claude-sonnet-4-5-20250929"}'
|
||||
|
||||
assert guardrail.scan_message_for_secrets(content) == []
|
||||
assert guardrail.redact_text(content) == content
|
||||
|
||||
|
||||
def test_scan_message_redacts_every_openai_key_occurrence():
|
||||
guardrail = _guardrail()
|
||||
content = f"first {OPENAI_KEY}, second {OPENAI_KEY}"
|
||||
|
||||
assert guardrail.redact_text(content) == "first [REDACTED], second [REDACTED]"
|
||||
|
||||
|
||||
def test_scan_message_redacts_short_numeric_openai_like_values():
|
||||
guardrail = _guardrail()
|
||||
|
||||
assert guardrail.redact_text(f"value {SHORT_OPENAI_KEY}") == "value [REDACTED]"
|
||||
|
||||
|
||||
def test_scan_message_requires_ascii_digits_for_openai_like_values():
|
||||
guardrail = _guardrail()
|
||||
|
||||
assert guardrail.scan_message_for_secrets(UNICODE_DIGIT_SUFFIX) == []
|
||||
assert guardrail.redact_text(UNICODE_DIGIT_SUFFIX) == UNICODE_DIGIT_SUFFIX
|
||||
|
||||
|
||||
def test_scan_message_redacts_openai_key_after_separator():
|
||||
guardrail = _guardrail()
|
||||
|
||||
assert guardrail.redact_text(f"openai_{OPENAI_KEY} key-{OPENAI_KEY}") == (
|
||||
"openai_[REDACTED] key-[REDACTED]"
|
||||
)
|
||||
assert guardrail.redact_text(URL_ENCODED_KEY) == "Bearer%20[REDACTED]"
|
||||
|
||||
|
||||
def test_scan_message_does_not_stop_openai_key_at_token_characters():
|
||||
guardrail = _guardrail()
|
||||
|
||||
assert guardrail.redact_text("key sk-proj-abcde12345/extra") == "key [REDACTED]/extra"
|
||||
|
||||
|
||||
def test_scan_message_stays_linear_on_repeated_sk_separators():
|
||||
guardrail = _guardrail()
|
||||
content = "-sk-" * 25_000
|
||||
|
||||
started = time.perf_counter()
|
||||
assert guardrail.scan_message_for_secrets(content) == []
|
||||
assert time.perf_counter() - started < 2.0
|
||||
|
||||
|
||||
def test_scan_message_redacts_whole_stripe_live_key():
|
||||
guardrail = _guardrail()
|
||||
|
||||
assert guardrail.redact_text(f"stripe {STRIPE_LIVE_KEY} end") == "stripe [REDACTED] end"
|
||||
|
||||
|
||||
def test_scan_message_returns_matches_in_stable_order():
|
||||
guardrail = _guardrail()
|
||||
detected = guardrail.scan_message_for_secrets(" ".join(AWS_KEYS))
|
||||
|
||||
assert [secret["value"] for secret in detected] == sorted(AWS_KEYS)
|
||||
|
||||
|
||||
def test_scan_message_replaces_longest_overlapping_match_first():
|
||||
guardrail = _guardrail()
|
||||
content = f'token = "{OPENAI_KEY}/extra"'
|
||||
|
||||
detected = guardrail.scan_message_for_secrets(content)
|
||||
assert [secret["value"] for secret in detected] == [f"{OPENAI_KEY}/extra", OPENAI_KEY]
|
||||
assert guardrail.redact_text(content) == 'token = "[REDACTED]"'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_redacts_secrets():
|
||||
"""Playground path: the returned texts must carry [REDACTED], not the secret."""
|
||||
|
|
@ -199,9 +290,7 @@ async def test_apply_guardrail_without_texts_records_nothing():
|
|||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": "https://x/y.png"}}
|
||||
],
|
||||
"content": [{"type": "image_url", "image_url": {"url": "https://x/y.png"}}],
|
||||
}
|
||||
],
|
||||
"metadata": {},
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import asyncio
|
||||
from typing import TYPE_CHECKING, Literal, Optional
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -11,6 +12,9 @@ from litellm.integrations.custom_guardrail import (
|
|||
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
|
||||
class TestCustomGuardrailDeploymentHook:
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -2175,6 +2179,170 @@ class TestRecordsOwnGuardrailInformation:
|
|||
assert _guardrail_entries(request_data) == []
|
||||
|
||||
|
||||
class _UndecoratedGuardrail(CustomGuardrail):
|
||||
"""apply_guardrail written like the docs example: no @log_guardrail_information."""
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
|
||||
if any("forbidden" in text for text in inputs.get("texts") or []):
|
||||
raise GuardrailRaisedException(guardrail_name=self.guardrail_name, message="Content blocked")
|
||||
return inputs
|
||||
|
||||
|
||||
class _UndecoratedSelfRecordingGuardrail(CustomGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response={"custom": True},
|
||||
request_data=request_data,
|
||||
guardrail_status="success",
|
||||
start_time=0.0,
|
||||
end_time=0.0,
|
||||
duration=0.0,
|
||||
)
|
||||
return inputs
|
||||
|
||||
|
||||
class _InheritedApplyGuardrail(_UndecoratedGuardrail):
|
||||
pass
|
||||
|
||||
|
||||
class TestUndecoratedApplyGuardrailIsLogged:
|
||||
"""LIT-5983 regression: a custom guardrail that overrides apply_guardrail without the
|
||||
@log_guardrail_information decorator must still record guardrail information, and the
|
||||
auto-wrap must not double-record decorated or self-recording implementations."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_undecorated_success_is_recorded(self):
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
guardrail = _UndecoratedGuardrail(guardrail_name="docs-style", event_hook=GuardrailEventHooks.pre_call)
|
||||
request_data: dict = {"model": "gpt-4o"}
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=GenericGuardrailAPIInputs(texts=["hello"]),
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
entries = _guardrail_entries(request_data)
|
||||
assert len(entries) == 1
|
||||
assert entries[0]["guardrail_name"] == "docs-style"
|
||||
assert entries[0]["guardrail_mode"] == "pre_call"
|
||||
assert entries[0]["guardrail_status"] == "success"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_undecorated_block_is_recorded_and_reraised(self):
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
|
||||
guardrail = _UndecoratedGuardrail(guardrail_name="docs-style")
|
||||
request_data: dict = {"model": "gpt-4o"}
|
||||
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=GenericGuardrailAPIInputs(texts=["forbidden"]),
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
entries = _guardrail_entries(request_data)
|
||||
assert len(entries) == 1
|
||||
assert entries[0]["guardrail_name"] == "docs-style"
|
||||
assert entries[0]["guardrail_status"] == "guardrail_intervened"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_undecorated_bare_exception_is_recorded_as_failed_to_respond(self):
|
||||
class _BareExceptionGuardrail(CustomGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
raise Exception("Content blocked: policy violation")
|
||||
|
||||
guardrail = _BareExceptionGuardrail(guardrail_name="docs-style")
|
||||
request_data: dict = {"model": "gpt-4o"}
|
||||
|
||||
with pytest.raises(Exception, match="Content blocked"):
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=GenericGuardrailAPIInputs(texts=["x"]),
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
entries = _guardrail_entries(request_data)
|
||||
assert len(entries) == 1
|
||||
assert entries[0]["guardrail_status"] == "guardrail_failed_to_respond"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inherited_apply_guardrail_is_recorded_once(self):
|
||||
guardrail = _InheritedApplyGuardrail(guardrail_name="child")
|
||||
request_data: dict = {"model": "gpt-4o"}
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=GenericGuardrailAPIInputs(texts=["hello"]),
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert len(_guardrail_entries(request_data)) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_undecorated_self_recording_apply_guardrail_is_recorded_once(self):
|
||||
guardrail = _UndecoratedSelfRecordingGuardrail(guardrail_name="self-recording")
|
||||
request_data: dict = {"model": "gpt-4o"}
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=GenericGuardrailAPIInputs(texts=["hello"]),
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
entries = _guardrail_entries(request_data)
|
||||
assert len(entries) == 1
|
||||
assert entries[0]["guardrail_response"] == {"custom": True}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_apply_guardrail_is_not_recorded(self):
|
||||
guardrail = CustomGuardrail(guardrail_name="base")
|
||||
request_data: dict = {"model": "gpt-4o"}
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=GenericGuardrailAPIInputs(texts=["hello"]),
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert _guardrail_entries(request_data) == []
|
||||
|
||||
def test_subclass_keywords_reach_cooperative_init_subclass(self):
|
||||
class _LabelMixin:
|
||||
seen_label: str = ""
|
||||
|
||||
def __init_subclass__(cls, label: str = "", **kwargs: object) -> None:
|
||||
super().__init_subclass__(**kwargs)
|
||||
cls.seen_label = label
|
||||
|
||||
class _Labelled(CustomGuardrail, _LabelMixin, label="docs-style"):
|
||||
pass
|
||||
|
||||
assert _Labelled.seen_label == "docs-style"
|
||||
|
||||
|
||||
class _ApplyOnlyObserver(CustomGuardrail):
|
||||
"""Overrides only apply_guardrail, like panw_prisma_airs; inherits async_logging_hook."""
|
||||
|
||||
|
|
|
|||
|
|
@ -1075,6 +1075,37 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
|
|||
assert result == responses_so_far
|
||||
|
||||
|
||||
class TestUndecoratedGuardrailIsRecorded:
|
||||
"""LIT-5983 regression: the handler calls apply_guardrail bare, so a custom guardrail
|
||||
without @log_guardrail_information must still end up in the request's guardrail
|
||||
information on both the request and response paths."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_path_records_undecorated_guardrail(self):
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
guardrail = MockGuardrail(guardrail_name="docs-style")
|
||||
data = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}}
|
||||
|
||||
await handler.process_input_messages(data, guardrail)
|
||||
|
||||
entries = data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert [(e["guardrail_name"], e["guardrail_status"]) for e in entries] == [("docs-style", "success")]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_path_records_undecorated_guardrail(self):
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
guardrail = MockGuardrail(guardrail_name="docs-style")
|
||||
response = ModelResponse(
|
||||
choices=[Choices(finish_reason="stop", index=0, message=Message(content="hi", role="assistant"))]
|
||||
)
|
||||
request_data: dict = {"metadata": {}}
|
||||
|
||||
await handler.process_output_response(response, guardrail, request_data=request_data)
|
||||
|
||||
entries = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert [(e["guardrail_name"], e["guardrail_status"]) for e in entries] == [("docs-style", "success")]
|
||||
|
||||
|
||||
class TestGetStructuredMessages:
|
||||
"""Test the get_structured_messages method."""
|
||||
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from fastapi import HTTPException
|
|||
from pydantic import ValidationError
|
||||
|
||||
from litellm.experimental_mcp_client.client import MCPClient
|
||||
from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
|
||||
oauth_protected_resource_path,
|
||||
raise_public,
|
||||
|
|
@ -442,6 +443,17 @@ def test_raise_public_maps_each_error_to_its_status(error, status):
|
|||
assert exc_info.value.status_code == status
|
||||
|
||||
|
||||
def test_raise_public_marks_only_url_credentials_error_as_safe_for_preview():
|
||||
with pytest.raises(HTTPException) as generic_exc_info:
|
||||
raise_public(CredError.of_misconfigured("private operator detail"))
|
||||
assert not isinstance(generic_exc_info.value, MCPServerURLCredentialsError)
|
||||
|
||||
error = CredError.of_url_credentials_not_allowed()
|
||||
with pytest.raises(MCPServerURLCredentialsError) as url_exc_info:
|
||||
raise_public(error)
|
||||
assert url_exc_info.value.detail == error.summary
|
||||
|
||||
|
||||
def test_raise_public_emits_unauthorized_challenge():
|
||||
body = {"error": "byok_auth_required", "server_id": "s1"}
|
||||
error = CredError.of_unauthorized("needs key", www_authenticate='Bearer resource_metadata="/x"', body=body)
|
||||
|
|
|
|||
|
|
@ -116,6 +116,35 @@ async def test_none_mode_yields_a_no_op_auth():
|
|||
assert isinstance(result.ok, NoOpAuth)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_none_mode_rejects_url_userinfo():
|
||||
spec = ServerSpec(
|
||||
server_id="s",
|
||||
resource="https://lit-user:s3cr3t@upstream.example.com/mcp",
|
||||
config=NoneConfig(),
|
||||
)
|
||||
|
||||
result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, spec)
|
||||
|
||||
assert isinstance(result, Error)
|
||||
assert result.error.tag == "url_credentials_not_allowed"
|
||||
assert "Basic Auth" in result.error.summary
|
||||
assert "auth_type: basic" in result.error.summary
|
||||
assert "auth_value: username:password" in result.error.summary
|
||||
assert "lit-user" not in result.error.summary
|
||||
assert "s3cr3t" not in result.error.summary
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_none_mode_does_not_validate_non_credential_resource():
|
||||
spec = ServerSpec(server_id="s", resource="https://[::1", config=NoneConfig())
|
||||
|
||||
result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, spec)
|
||||
|
||||
assert isinstance(result, Ok)
|
||||
assert isinstance(result.ok, NoOpAuth)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_key_shared_emits_the_configured_header():
|
||||
config = ApiKeyConfig(
|
||||
|
|
|
|||
|
|
@ -75,6 +75,15 @@ def test_crederror_factory_sets_the_matching_tag(factory, expected_tag):
|
|||
assert "detail text" in err.summary
|
||||
|
||||
|
||||
def test_url_credentials_error_has_a_fixed_actionable_summary():
|
||||
err = CredError.of_url_credentials_not_allowed()
|
||||
|
||||
assert err.tag == "url_credentials_not_allowed"
|
||||
assert "Basic Auth" in err.summary
|
||||
assert "auth_type: basic" in err.summary
|
||||
assert "auth_value: username:password" in err.summary
|
||||
|
||||
|
||||
def test_apikeyconfig_requires_a_key_source():
|
||||
with pytest.raises(ValidationError):
|
||||
ApiKeyConfig() # type: ignore[call-arg]
|
||||
|
|
|
|||
|
|
@ -8966,6 +8966,24 @@ class TestCreateMcpClientV2Graft:
|
|||
assert isinstance(client._resolved_auth, NoOpAuth)
|
||||
assert client._mcp_auth_value is None
|
||||
|
||||
@pytest.mark.parametrize("auth_type", [None, MCPAuth.none])
|
||||
async def test_none_mode_rejects_url_userinfo(self, auth_type):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await MCPServerManager()._create_mcp_client(
|
||||
self._http_server(
|
||||
auth_type=auth_type,
|
||||
url="https://lit-user:s3cr3t@upstream.example.com/mcp",
|
||||
)
|
||||
)
|
||||
|
||||
detail = str(exc_info.value.detail)
|
||||
assert exc_info.value.status_code == 500
|
||||
assert "Basic Auth" in detail
|
||||
assert "auth_type: basic" in detail
|
||||
assert "auth_value: username:password" in detail
|
||||
assert "lit-user" not in detail
|
||||
assert "s3cr3t" not in detail
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"auth_type, token, expected_name, expected_value",
|
||||
[
|
||||
|
|
@ -11369,6 +11387,30 @@ class TestResolveOpenapiToolAuth:
|
|||
|
||||
assert "Authorization" not in (forwarded or {})
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_none_mode_without_url_keeps_spec_path_server_unauthenticated(self):
|
||||
server = MCPServer(
|
||||
server_id="openapi-only",
|
||||
name="report_api",
|
||||
server_name="report_api",
|
||||
url=None,
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.none,
|
||||
spec_path="https://api.example.com/openapi.json",
|
||||
)
|
||||
|
||||
resolved, forwarded = await MCPServerManager().resolve_openapi_upstream_auth(
|
||||
mcp_server=server,
|
||||
oauth2_headers=None,
|
||||
raw_headers=None,
|
||||
mcp_auth_header=None,
|
||||
user_api_key_auth=None,
|
||||
forwarded_headers={"X-Trace": "trace-id"},
|
||||
)
|
||||
|
||||
assert resolved is None
|
||||
assert forwarded == {"X-Trace": "trace-id"}
|
||||
|
||||
|
||||
class TestOpenApiHandlerRelaysUpstreamAuth:
|
||||
"""`_call_openapi_tool_handler` must not flatten a re-auth signal into a generic message.
|
||||
|
|
|
|||
|
|
@ -177,6 +177,36 @@ class TestExecuteWithMcpClient:
|
|||
assert "https://api.example.com/mcp/" in message
|
||||
assert "30s" in message
|
||||
|
||||
def test_connection_error_message_hides_arbitrary_http_exception_detail(self):
|
||||
message = rest_endpoints._connection_error_message(
|
||||
HTTPException(status_code=500, detail="secret upstream detail"),
|
||||
"https://api.example.com/mcp/",
|
||||
30.0,
|
||||
)
|
||||
|
||||
assert "secret upstream detail" not in message
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_none_mode_url_credentials_returns_actionable_redacted_error(self):
|
||||
async def unreached_operation(client):
|
||||
raise AssertionError("operation must not run for an invalid server configuration")
|
||||
|
||||
payload = NewMCPServerRequest(
|
||||
server_name="example",
|
||||
url="https://lit-user:s3cr3t@upstream.example.com/mcp",
|
||||
auth_type=MCPAuth.none,
|
||||
)
|
||||
|
||||
result = await rest_endpoints._execute_with_mcp_client(payload, unreached_operation)
|
||||
|
||||
message = str(result["message"])
|
||||
assert result["error"] is True
|
||||
assert "Basic Auth" in message
|
||||
assert "auth_type: basic" in message
|
||||
assert "auth_value: username:password" in message
|
||||
assert "lit-user" not in message
|
||||
assert "s3cr3t" not in message
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_forwards_static_headers(self, monkeypatch):
|
||||
"""Ensure static_headers are forwarded to the MCP client during test calls.
|
||||
|
|
|
|||
|
|
@ -1137,7 +1137,11 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source():
|
|||
mock_api.assert_called_once()
|
||||
kwargs = mock_api.call_args.kwargs
|
||||
assert kwargs["source"] == "OUTPUT"
|
||||
assert kwargs["request_data"] == {"model": "gpt-4o"}
|
||||
assert kwargs["request_data"]["model"] == "gpt-4o"
|
||||
recorded = kwargs["request_data"]["metadata"]["standard_logging_guardrail_information"]
|
||||
assert [(e["guardrail_name"], e["guardrail_status"]) for e in recorded] == [
|
||||
(guardrail.guardrail_name, "success")
|
||||
]
|
||||
synthetic = kwargs["response"]
|
||||
assert isinstance(synthetic, ModelResponse)
|
||||
assert len(synthetic.choices) == 2
|
||||
|
|
|
|||
|
|
@ -473,6 +473,110 @@ async def test_update_spend_logs_retries_and_requeues_batch_on_db_outage(
|
|||
assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"]
|
||||
|
||||
|
||||
def _deadlock_error() -> Exception:
|
||||
return _data_error(
|
||||
'Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: None, '
|
||||
'kind: QueryError(PostgresError { code: "40P01", message: "deadlock detected", severity: "ERROR" }) })'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_retries_deadlock_and_keeps_every_row(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""A 40P01 deadlock aborts the whole insert, so the same rows succeed on replay.
|
||||
Before the fix the deadlock surfaced as a plain ``DataError`` and went through
|
||||
poison-row isolation, which bisected the batch and dropped every row the
|
||||
deadlock happened to hit as if Postgres had rejected it.
|
||||
"""
|
||||
|
||||
async def _fake_sleep(_: float) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
|
||||
create_many = AsyncMock(side_effect=[_deadlock_error(), _deadlock_error(), None])
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = create_many
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
mock_prisma_client.spend_log_transactions = []
|
||||
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=2,
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")],
|
||||
)
|
||||
|
||||
attempts = tuple(
|
||||
tuple(row["request_id"] for row in call.kwargs["data"]) for call in create_many.await_args_list
|
||||
)
|
||||
assert attempts == (("a", "b"), ("a", "b"), ("a", "b"))
|
||||
assert mock_prisma_client.spend_log_transactions == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_requeues_batch_once_deadlock_retries_exhaust(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""If every retry deadlocks, the batch goes back to the head of the queue for
|
||||
the next flush instead of being dropped.
|
||||
"""
|
||||
|
||||
async def _fake_sleep(_: float) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_deadlock_error())
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")]
|
||||
|
||||
with pytest.raises(type(_deadlock_error())):
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=1,
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")],
|
||||
)
|
||||
|
||||
assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 2
|
||||
assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_requeues_batch_on_non_transport_db_error(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any
|
||||
) -> None:
|
||||
"""A DB error that is neither transport nor deadlock (here P2021, the table is
|
||||
gone mid-migration) is not retried in place, but the dequeued batch must not
|
||||
be lost either: it goes back to the head of the queue so it lands once the
|
||||
DB is healthy again.
|
||||
"""
|
||||
from prisma.errors import TableNotFoundError
|
||||
|
||||
err = TableNotFoundError(
|
||||
{"user_facing_error": {"error_code": "P2021", "message": "The table does not exist", "meta": {}}}
|
||||
)
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=err)
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")]
|
||||
|
||||
with pytest.raises(TableNotFoundError):
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=2,
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")],
|
||||
)
|
||||
|
||||
assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 1
|
||||
assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_requeue_after_outage_drops_oldest_logs_past_the_byte_budget(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any
|
||||
|
|
@ -549,8 +653,9 @@ async def test_flush_returns_the_bytes_it_took_off_the_queue(mock_prisma_client:
|
|||
async def test_update_spend_logs_does_not_requeue_non_transport_failures(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any
|
||||
) -> None:
|
||||
"""Only transport failures are worth replaying. A rejection the DB will keep
|
||||
rejecting must not be requeued, or it would wedge the queue forever.
|
||||
"""Only DB failures are worth replaying. A row the proxy itself cannot
|
||||
serialize would fail the same way on every flush, so requeueing it would
|
||||
wedge the head of the queue forever.
|
||||
"""
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=ValueError("bad payload"))
|
||||
proxy_logging = MagicMock()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,254 @@
|
|||
"""
|
||||
A deployment with `model_info.supported_endpoints` containing `/v1/responses` forwards
|
||||
`/v1/responses` natively to `{api_base}/responses`. Without it, generic OpenAI-compatible
|
||||
providers such as `custom_openai` keep bridging through `/v1/chat/completions`.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig
|
||||
from litellm.responses.main import _resolve_responses_api_provider_config
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
API_BASE = "https://backend.example/v1"
|
||||
RESPONSES_URL = f"{API_BASE}/responses"
|
||||
CHAT_URL = f"{API_BASE}/chat/completions"
|
||||
OPT_IN = {"supported_endpoints": ["/v1/chat/completions", "/v1/responses"]}
|
||||
|
||||
RESPONSES_BODY = {
|
||||
"id": "resp_native",
|
||||
"object": "response",
|
||||
"created_at": 1741476542,
|
||||
"status": "completed",
|
||||
"model": "my-model",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_1",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "native", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"parallel_tool_calls": True,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
|
||||
CHAT_BODY = {
|
||||
"id": "chatcmpl_bridged",
|
||||
"object": "chat.completion",
|
||||
"created": 1741476542,
|
||||
"model": "my-model",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "bridged"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
|
||||
SSE_BODY = (
|
||||
"event: response.created\n"
|
||||
f"data: {json.dumps({'type': 'response.created', 'response': RESPONSES_BODY})}\n\n"
|
||||
"event: response.completed\n"
|
||||
f"data: {json.dumps({'type': 'response.completed', 'response': RESPONSES_BODY})}\n\n"
|
||||
)
|
||||
|
||||
|
||||
def _mock_backend(router: respx.MockRouter) -> tuple[respx.Route, respx.Route]:
|
||||
responses_route = router.post(RESPONSES_URL).mock(return_value=httpx.Response(200, json=RESPONSES_BODY))
|
||||
chat_route = router.post(CHAT_URL).mock(return_value=httpx.Response(200, json=CHAT_BODY))
|
||||
return responses_route, chat_route
|
||||
|
||||
|
||||
SWAPPED_MODEL = "deepseek/deepseek-chat"
|
||||
SWAPPED_API_BASE = "https://api.deepseek.com/beta"
|
||||
|
||||
|
||||
def _prompt_manager_swapping_to(model: str) -> MagicMock:
|
||||
"""A logging object whose prompt hook rewrites the request's model, as a prompt manager does."""
|
||||
prompt_return = (model, [{"role": "user", "content": "hi"}], {})
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.__class__ = LiteLLMLoggingObj
|
||||
logging_obj.should_run_prompt_management_hooks.return_value = True
|
||||
logging_obj.get_chat_completion_prompt.return_value = prompt_return
|
||||
logging_obj.async_get_chat_completion_prompt = AsyncMock(return_value=prompt_return)
|
||||
logging_obj.model_call_details = {}
|
||||
return logging_obj
|
||||
|
||||
|
||||
def _mock_swap_targets(router: respx.MockRouter, monkeypatch) -> tuple[respx.Route, respx.Route]:
|
||||
"""The swapped provider's chat endpoint, plus the `/responses` it does not serve but a stale
|
||||
opt-in would send to."""
|
||||
monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-deepseek")
|
||||
swapped_chat_route = router.post(f"{SWAPPED_API_BASE}/chat/completions").mock(
|
||||
return_value=httpx.Response(200, json=CHAT_BODY)
|
||||
)
|
||||
stale_responses_route = router.post(f"{SWAPPED_API_BASE}/responses").mock(
|
||||
return_value=httpx.Response(200, json=RESPONSES_BODY)
|
||||
)
|
||||
return swapped_chat_route, stale_responses_route
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _respx_interceptable_httpx_client(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "num_retries", 0)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
yield
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_info, expected_type",
|
||||
[
|
||||
(OPT_IN, OpenAILikeResponsesConfig),
|
||||
({"supported_endpoints": ["/v1/chat/completions"]}, type(None)),
|
||||
({}, type(None)),
|
||||
(None, type(None)),
|
||||
("/v1/responses", type(None)),
|
||||
],
|
||||
)
|
||||
def test_resolver_opt_in_gates_openai_like_config(model_info, expected_type):
|
||||
config = _resolve_responses_api_provider_config("my-model", "custom_openai", model_info)
|
||||
assert type(config) is expected_type
|
||||
|
||||
|
||||
def test_resolver_keeps_native_provider_config():
|
||||
"""`openai/` already routes /v1/responses natively; the opt-in must not swap its config."""
|
||||
config = _resolve_responses_api_provider_config("gpt-4.1", "openai", OPT_IN)
|
||||
assert type(config) is OpenAIResponsesAPIConfig
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_opt_in_forwards_responses_natively():
|
||||
responses_route, chat_route = _mock_backend(respx.mock)
|
||||
|
||||
result = await litellm.aresponses(
|
||||
model="custom_openai/my-model",
|
||||
input="hi",
|
||||
api_base=API_BASE,
|
||||
api_key="sk-backend",
|
||||
model_info=OPT_IN,
|
||||
)
|
||||
|
||||
assert responses_route.call_count == 1
|
||||
assert chat_route.call_count == 0
|
||||
request = responses_route.calls.last.request
|
||||
assert request.headers["authorization"] == "Bearer sk-backend"
|
||||
assert json.loads(request.content)["input"] == "hi"
|
||||
assert isinstance(result, ResponsesAPIResponse)
|
||||
assert result.output[0].content[0].text == "native"
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_opt_in_forwards_streaming_responses_natively(monkeypatch):
|
||||
"""The router registers each deployment in `litellm.model_cost`; an unregistered model is
|
||||
treated as non-streaming and would be faked, so mirror that registration here."""
|
||||
monkeypatch.setitem(litellm.model_cost, "custom_openai/my-model", {"litellm_provider": "custom_openai"})
|
||||
responses_route = respx.post(RESPONSES_URL).mock(
|
||||
return_value=httpx.Response(200, text=SSE_BODY, headers={"content-type": "text/event-stream"})
|
||||
)
|
||||
chat_route = respx.post(CHAT_URL).mock(return_value=httpx.Response(200, json=CHAT_BODY))
|
||||
|
||||
stream = await litellm.aresponses(
|
||||
model="custom_openai/my-model",
|
||||
input="hi",
|
||||
stream=True,
|
||||
api_base=API_BASE,
|
||||
api_key="sk-backend",
|
||||
model_info=OPT_IN,
|
||||
)
|
||||
events = [event async for event in stream]
|
||||
|
||||
assert responses_route.call_count == 1
|
||||
assert chat_route.call_count == 0
|
||||
assert json.loads(responses_route.calls.last.request.content)["stream"] is True
|
||||
assert [event.type for event in events] == ["response.created", "response.completed"]
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_without_opt_in_still_bridges_through_chat_completions():
|
||||
responses_route, chat_route = _mock_backend(respx.mock)
|
||||
|
||||
result = await litellm.aresponses(
|
||||
model="custom_openai/my-model",
|
||||
input="hi",
|
||||
api_base=API_BASE,
|
||||
api_key="sk-backend",
|
||||
model_info={"supported_endpoints": ["/v1/chat/completions"]},
|
||||
)
|
||||
|
||||
assert chat_route.call_count == 1
|
||||
assert responses_route.call_count == 0
|
||||
assert isinstance(result, ResponsesAPIResponse)
|
||||
assert result.output[0].content[0].text == "bridged"
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_prompt_swap_to_other_provider_drops_deployment_opt_in(monkeypatch):
|
||||
"""When a prompt manager moves the request to another provider, the original deployment's
|
||||
`supported_endpoints` no longer describes the upstream, so the swapped provider bridges."""
|
||||
swapped_chat_route, stale_responses_route = _mock_swap_targets(respx.mock, monkeypatch)
|
||||
|
||||
result = await litellm.aresponses(
|
||||
model="custom_openai/my-model",
|
||||
input="hi",
|
||||
prompt_id="p1",
|
||||
litellm_logging_obj=_prompt_manager_swapping_to(SWAPPED_MODEL),
|
||||
model_info=OPT_IN,
|
||||
)
|
||||
|
||||
assert swapped_chat_route.call_count == 1
|
||||
assert stale_responses_route.call_count == 0
|
||||
assert isinstance(result, ResponsesAPIResponse)
|
||||
assert result.output[0].content[0].text == "bridged"
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_sync_prompt_swap_to_other_provider_drops_deployment_opt_in(monkeypatch):
|
||||
swapped_chat_route, stale_responses_route = _mock_swap_targets(respx.mock, monkeypatch)
|
||||
|
||||
result = litellm.responses(
|
||||
model="custom_openai/my-model",
|
||||
input="hi",
|
||||
prompt_id="p1",
|
||||
litellm_logging_obj=_prompt_manager_swapping_to(SWAPPED_MODEL),
|
||||
model_info=OPT_IN,
|
||||
)
|
||||
|
||||
assert swapped_chat_route.call_count == 1
|
||||
assert stale_responses_route.call_count == 0
|
||||
assert isinstance(result, ResponsesAPIResponse)
|
||||
assert result.output[0].content[0].text == "bridged"
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_mode_responses_chat_completion_reaches_native_responses(monkeypatch):
|
||||
"""A `mode: responses` deployment bridges chat completions into the Responses API; with
|
||||
the opt-in that inner call must reach `{api_base}/responses` instead of bouncing back
|
||||
to `/chat/completions`."""
|
||||
responses_route, chat_route = _mock_backend(respx.mock)
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"custom_openai/my-model",
|
||||
{"mode": "responses", "litellm_provider": "custom_openai"},
|
||||
)
|
||||
|
||||
result = await litellm.acompletion(
|
||||
model="custom_openai/my-model",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
api_base=API_BASE,
|
||||
api_key="sk-backend",
|
||||
model_info={"mode": "responses", **OPT_IN},
|
||||
)
|
||||
|
||||
assert responses_route.call_count == 1
|
||||
assert chat_route.call_count == 0
|
||||
assert isinstance(result, ModelResponse)
|
||||
assert result.choices[0].message.content == "native"
|
||||
|
|
@ -13148,3 +13148,144 @@ async def test_router_retry_policy_controls_upstream_attempt_count(
|
|||
await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}])
|
||||
|
||||
assert upstream.call_count == expected_upstream_calls
|
||||
|
||||
|
||||
class _InFlightTracker:
|
||||
def __init__(self) -> None:
|
||||
self.current = 0
|
||||
self.peak = 0
|
||||
|
||||
def enter(self) -> None:
|
||||
self.current += 1
|
||||
self.peak = max(self.peak, self.current)
|
||||
|
||||
def exit(self) -> None:
|
||||
self.current -= 1
|
||||
|
||||
|
||||
_SSE_CHUNKS: Final[tuple[bytes, ...]] = tuple(
|
||||
b'data: {"id":"c","object":"chat.completion.chunk","created":1,"model":"gpt-5.6",'
|
||||
b'"choices":[{"index":0,"delta":{"content":"x"},"finish_reason":null}]}\n\n'
|
||||
for _ in range(5)
|
||||
)
|
||||
|
||||
|
||||
class _CountingSSEStream(httpx.AsyncByteStream):
|
||||
def __init__(self, tracker: _InFlightTracker) -> None:
|
||||
self._tracker = tracker
|
||||
self._in_flight = False
|
||||
|
||||
def _finish(self) -> None:
|
||||
if self._in_flight:
|
||||
self._in_flight = False
|
||||
self._tracker.exit()
|
||||
|
||||
async def __aiter__(self):
|
||||
self._in_flight = True
|
||||
self._tracker.enter()
|
||||
try:
|
||||
for chunk in _SSE_CHUNKS:
|
||||
await asyncio.sleep(0.02)
|
||||
yield chunk
|
||||
finally:
|
||||
await self.aclose()
|
||||
yield b"data: [DONE]\n\n"
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await asyncio.sleep(0.02)
|
||||
self._finish()
|
||||
|
||||
|
||||
def _max_parallel_router(max_parallel_requests: int) -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-5.6",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-5.6",
|
||||
"api_key": "sk-fake",
|
||||
"api_base": "https://max-parallel.local/v1",
|
||||
"max_parallel_requests": max_parallel_requests,
|
||||
},
|
||||
}
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
async def test_router_max_parallel_requests_bounds_in_flight_upstream_calls(
|
||||
monkeypatch: pytest.MonkeyPatch, stream: bool
|
||||
):
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
tracker: Final = _InFlightTracker()
|
||||
router: Final = _max_parallel_router(max_parallel_requests=2)
|
||||
|
||||
async def upstream(request: httpx.Request) -> httpx.Response:
|
||||
if stream:
|
||||
return httpx.Response(
|
||||
200, headers={"content-type": "text/event-stream"}, stream=_CountingSSEStream(tracker)
|
||||
)
|
||||
tracker.enter()
|
||||
await asyncio.sleep(0.05)
|
||||
tracker.exit()
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "c",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-5.6",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "x"}, "finish_reason": "stop"}],
|
||||
},
|
||||
)
|
||||
|
||||
async def one_call() -> None:
|
||||
response = await router.acompletion(
|
||||
model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], stream=stream
|
||||
)
|
||||
if stream:
|
||||
async for _ in response:
|
||||
pass
|
||||
|
||||
with respx.mock(assert_all_called=True) as respx_mock:
|
||||
respx_mock.post("https://max-parallel.local/v1/chat/completions").mock(side_effect=upstream)
|
||||
await asyncio.wait_for(asyncio.gather(*(one_call() for _ in range(10))), timeout=10)
|
||||
|
||||
assert tracker.peak <= 2
|
||||
assert tracker.current == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_max_parallel_requests_slot_released_when_stream_closed_early(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
tracker: Final = _InFlightTracker()
|
||||
router: Final = _max_parallel_router(max_parallel_requests=1)
|
||||
|
||||
with respx.mock() as respx_mock:
|
||||
respx_mock.post("https://max-parallel.local/v1/chat/completions").mock(
|
||||
side_effect=lambda request: httpx.Response(
|
||||
200, headers={"content-type": "text/event-stream"}, stream=_CountingSSEStream(tracker)
|
||||
)
|
||||
)
|
||||
first: Final = await router.acompletion(
|
||||
model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], stream=True
|
||||
)
|
||||
await first.__anext__()
|
||||
|
||||
async def second_call() -> None:
|
||||
second = await router.acompletion(
|
||||
model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], stream=True
|
||||
)
|
||||
async for _ in second:
|
||||
pass
|
||||
|
||||
second_task: Final = asyncio.create_task(second_call())
|
||||
await asyncio.sleep(0.05)
|
||||
assert tracker.current == 1
|
||||
await first.aclose()
|
||||
await asyncio.wait_for(second_task, timeout=2)
|
||||
|
||||
assert tracker.peak == 1
|
||||
assert tracker.current == 0
|
||||
|
|
|
|||
6
uv.lock
generated
6
uv.lock
generated
|
|
@ -10,7 +10,7 @@ resolution-markers = [
|
|||
]
|
||||
|
||||
[options]
|
||||
exclude-newer = "2026-09-01T21:00:02.682921Z"
|
||||
exclude-newer = "2026-09-02T16:58:34.594994Z"
|
||||
exclude-newer-span = "P3D"
|
||||
|
||||
[manifest]
|
||||
|
|
@ -4771,12 +4771,12 @@ proxy-dev = [
|
|||
|
||||
[[package]]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.64"
|
||||
version = "0.1.65"
|
||||
source = { editable = "enterprise" }
|
||||
|
||||
[[package]]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.93"
|
||||
version = "0.4.94"
|
||||
source = { editable = "litellm-proxy-extras" }
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue