mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
chore: merge origin/main into litellm_jwt_auto_register_map_existing_key
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
commit
50575e32ed
92 changed files with 5638 additions and 262 deletions
14
.github/e2e-stack/secrets_to_env.py
vendored
14
.github/e2e-stack/secrets_to_env.py
vendored
|
|
@ -9,6 +9,7 @@ from pydantic import TypeAdapter, ValidationError
|
|||
secrets_adapter: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str])
|
||||
ENV_NAME: Final = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
|
||||
MIN_MASKED_LENGTH: Final = 8
|
||||
ACTIONS_RUNNER_FLAG: Final = "GITHUB_ACTIONS"
|
||||
|
||||
|
||||
def main() -> int:
|
||||
|
|
@ -30,10 +31,15 @@ def main() -> int:
|
|||
f"these names or values cannot be represented in both bash and dotenv: {' '.join(sorted(unusable))}\n"
|
||||
)
|
||||
return 1
|
||||
for value in secrets.values():
|
||||
if len(value) >= MIN_MASKED_LENGTH:
|
||||
_ = sys.stdout.write(f"::add-mask::{value.replace('%', '%25')}\n")
|
||||
sys.stdout.flush()
|
||||
if os.environ.get(ACTIONS_RUNNER_FLAG) == "true":
|
||||
_ = sys.stdout.write(
|
||||
"".join(
|
||||
f"::add-mask::{value.replace('%', '%25')}\n"
|
||||
for value in secrets.values()
|
||||
if len(value) >= MIN_MASKED_LENGTH
|
||||
)
|
||||
)
|
||||
sys.stdout.flush()
|
||||
lines: Final = tuple(f"{key}='{value}'" for key, value in secrets.items() if value)
|
||||
try:
|
||||
with os.fdopen(os.open(env_path, os.O_WRONLY | os.O_APPEND | os.O_CREAT | os.O_NOFOLLOW, 0o600), "w") as handle:
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@
|
|||
"bash_20250124": null,
|
||||
"code-execution-2025-08-25": "code-execution-2025-08-25",
|
||||
"compact-2026-01-12": "compact-2026-01-12",
|
||||
"compact-2026-09-04": "compact-2026-09-04",
|
||||
"computer-use-2025-01-24": "computer-use-2025-01-24",
|
||||
"computer-use-2025-11-24": "computer-use-2025-11-24",
|
||||
"context-1m-2025-08-07": "context-1m-2025-08-07",
|
||||
|
|
|
|||
|
|
@ -115,6 +115,7 @@ from litellm.types.utils import (
|
|||
LlmProviders,
|
||||
LlmProvidersSet,
|
||||
ModelInfo,
|
||||
ModelInfoBase,
|
||||
PromptTokensDetailsWrapper,
|
||||
ServiceTier,
|
||||
StandardBuiltInToolsParams,
|
||||
|
|
@ -322,6 +323,48 @@ class OCRPricing(TypedDict, total=False):
|
|||
annotation_cost_per_page: ReadOnly[float | None]
|
||||
|
||||
|
||||
_WALL_CLOCK_PRICED_MODES: Final = frozenset({"chat", "completion", "embedding", "responses"})
|
||||
|
||||
|
||||
def _has_token_or_tiered_pricing(model_info: ModelInfoBase) -> bool:
|
||||
return (
|
||||
(model_info.get("input_cost_per_token") or 0.0) > 0
|
||||
or (model_info.get("output_cost_per_token") or 0.0) > 0
|
||||
or model_info.get("tiered_pricing") is not None
|
||||
)
|
||||
|
||||
|
||||
def _bills_wall_clock_seconds(model_info: ModelInfoBase) -> bool:
|
||||
mode: Final = model_info.get("mode")
|
||||
return mode is None or mode in _WALL_CLOCK_PRICED_MODES
|
||||
|
||||
|
||||
def _per_second_pricing_cost(
|
||||
model: str,
|
||||
custom_llm_provider: str | None,
|
||||
response_time_ms: float | None,
|
||||
) -> tuple[float, float] | None:
|
||||
try:
|
||||
model_info: Final = _cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)
|
||||
except Exception: # noqa: BLE001 # the lookup raises plain Exception for an unmapped model
|
||||
return None
|
||||
if _has_token_or_tiered_pricing(model_info) or not _bills_wall_clock_seconds(model_info):
|
||||
return None
|
||||
input_cost_per_second: Final = model_info.get("input_cost_per_second")
|
||||
output_cost_per_second: Final = model_info.get("output_cost_per_second")
|
||||
if input_cost_per_second is None and output_cost_per_second is None:
|
||||
return None
|
||||
seconds: Final = (response_time_ms or 0.0) / 1000
|
||||
verbose_logger.debug(
|
||||
"For model=%s - input_cost_per_second: %s; output_cost_per_second: %s; response time: %s",
|
||||
model,
|
||||
input_cost_per_second,
|
||||
output_cost_per_second,
|
||||
response_time_ms,
|
||||
)
|
||||
return (input_cost_per_second or 0.0) * seconds, (output_cost_per_second or 0.0) * seconds
|
||||
|
||||
|
||||
def cost_per_token(
|
||||
model: str = "",
|
||||
prompt_tokens: int = 0,
|
||||
|
|
@ -448,9 +491,6 @@ def cost_per_token(
|
|||
if response_cost is not None:
|
||||
return response_cost[0], response_cost[1]
|
||||
|
||||
# given
|
||||
prompt_tokens_cost_usd_dollar: float = 0
|
||||
completion_tokens_cost_usd_dollar: float = 0
|
||||
model_cost_ref: Final = litellm.model_cost
|
||||
# Only callers that explicitly pass `custom_llm_provider` get the
|
||||
# dedup/prefix-join treatment. When provider is omitted, preserve legacy
|
||||
|
|
@ -611,6 +651,14 @@ def cost_per_token(
|
|||
number_of_queries=number_of_queries or 1,
|
||||
optional_params=(getattr(response, "_hidden_params", None) if response else None),
|
||||
)
|
||||
elif (
|
||||
per_second_cost := _per_second_pricing_cost(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
response_time_ms=response_time_ms,
|
||||
)
|
||||
) is not None:
|
||||
return per_second_cost
|
||||
elif custom_llm_provider == "vertex_ai":
|
||||
cost_router: Final = google_cost_router(
|
||||
model=model_without_prefix,
|
||||
|
|
@ -685,12 +733,7 @@ def cost_per_token(
|
|||
)
|
||||
else:
|
||||
model_info: Final = _cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)
|
||||
|
||||
if (
|
||||
(model_info.get("input_cost_per_token") or 0.0) > 0
|
||||
or (model_info.get("output_cost_per_token") or 0.0) > 0
|
||||
or model_info.get("tiered_pricing") is not None
|
||||
):
|
||||
if _has_token_or_tiered_pricing(model_info):
|
||||
return generic_cost_per_token(
|
||||
model=model,
|
||||
usage=usage_block,
|
||||
|
|
@ -698,36 +741,8 @@ def cost_per_token(
|
|||
service_tier=service_tier,
|
||||
data_residency=data_residency,
|
||||
)
|
||||
|
||||
input_cost_per_second: Final = model_info.get("input_cost_per_second")
|
||||
if input_cost_per_second is not None and response_time_ms is not None:
|
||||
verbose_logger.debug(
|
||||
"For model=%s - input_cost_per_second: %s; response time: %s",
|
||||
model,
|
||||
input_cost_per_second,
|
||||
response_time_ms,
|
||||
)
|
||||
## COST PER SECOND ##
|
||||
prompt_tokens_cost_usd_dollar = input_cost_per_second * response_time_ms / 1000
|
||||
|
||||
output_cost_per_second: Final = model_info.get("output_cost_per_second")
|
||||
if output_cost_per_second is not None and response_time_ms is not None:
|
||||
verbose_logger.debug(
|
||||
"For model=%s - output_cost_per_second: %s; response time: %s",
|
||||
model,
|
||||
output_cost_per_second,
|
||||
response_time_ms,
|
||||
)
|
||||
## COST PER SECOND ##
|
||||
completion_tokens_cost_usd_dollar = output_cost_per_second * response_time_ms / 1000
|
||||
|
||||
verbose_logger.debug(
|
||||
"Returned custom cost for model=%s - prompt_tokens_cost_usd_dollar: %s, completion_tokens_cost_usd_dollar: %s",
|
||||
model,
|
||||
prompt_tokens_cost_usd_dollar,
|
||||
completion_tokens_cost_usd_dollar,
|
||||
)
|
||||
return prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar
|
||||
verbose_logger.debug("No per-token, tiered, or per-second pricing for model=%s; cost is 0", model)
|
||||
return 0.0, 0.0
|
||||
|
||||
|
||||
def get_replicate_completion_pricing(completion_response: dict, total_time=0.0):
|
||||
|
|
@ -1222,6 +1237,21 @@ def _split_responses_ws_logging_object_by_service_tier(
|
|||
)
|
||||
|
||||
|
||||
def _response_time_ms_for_cost(
|
||||
completion_response: object,
|
||||
litellm_logging_obj: LitellmLoggingObject | None,
|
||||
total_time: float | None,
|
||||
) -> float:
|
||||
stamped: Final = getattr(completion_response, "_response_ms", None)
|
||||
if isinstance(stamped, (int, float)):
|
||||
return float(stamped)
|
||||
if total_time:
|
||||
return total_time
|
||||
if litellm_logging_obj is not None:
|
||||
return litellm_logging_obj.get_response_ms()
|
||||
return 0.0
|
||||
|
||||
|
||||
def completion_cost(
|
||||
completion_response: object | None = None,
|
||||
model: str | None = None,
|
||||
|
|
@ -1443,8 +1473,6 @@ def completion_cost(
|
|||
prompt_tokens_details = _usage.get("prompt_tokens_details") or {}
|
||||
cache_read_input_tokens = prompt_tokens_details.get("cached_tokens", 0)
|
||||
|
||||
total_time = getattr(completion_response, "_response_ms", 0)
|
||||
|
||||
hidden_params = getattr(completion_response, "_hidden_params", None)
|
||||
if hidden_params is not None:
|
||||
custom_llm_provider = hidden_params.get("custom_llm_provider", custom_llm_provider or None)
|
||||
|
|
@ -1676,6 +1704,11 @@ def completion_cost(
|
|||
)
|
||||
|
||||
return MCPCostCalculator.calculate_mcp_tool_call_cost(litellm_logging_obj=litellm_logging_obj)
|
||||
response_time_ms = _response_time_ms_for_cost(
|
||||
completion_response=completion_response,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
total_time=total_time,
|
||||
)
|
||||
# Calculate cost based on prompt_tokens, completion_tokens
|
||||
if (
|
||||
"togethercomputer" in model or "together_ai" in model or custom_llm_provider == "together_ai"
|
||||
|
|
@ -1686,7 +1719,7 @@ def completion_cost(
|
|||
# see https://replicate.com/pricing
|
||||
elif (model in litellm.replicate_models or "replicate" in model) and model not in litellm.model_cost:
|
||||
# for unmapped replicate model, default to replicate's time tracking logic
|
||||
return get_replicate_completion_pricing(completion_response, total_time)
|
||||
return get_replicate_completion_pricing(completion_response, response_time_ms)
|
||||
|
||||
if model is None:
|
||||
raise ValueError(
|
||||
|
|
@ -1718,7 +1751,7 @@ def completion_cost(
|
|||
prompt_tokens=prompt_tokens or 0,
|
||||
completion_tokens=completion_tokens or 0,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
response_time_ms=total_time,
|
||||
response_time_ms=response_time_ms,
|
||||
region_name=None if explicit_pricing else region_name,
|
||||
custom_cost_per_second=custom_cost_per_second,
|
||||
custom_cost_per_token=custom_cost_per_token,
|
||||
|
|
|
|||
180
litellm/litellm_core_utils/bug_report.py
Normal file
180
litellm/litellm_core_utils/bug_report.py
Normal file
|
|
@ -0,0 +1,180 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import platform
|
||||
import traceback
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final, Literal
|
||||
from urllib.parse import urlencode
|
||||
|
||||
import litellm
|
||||
from litellm._version import version as litellm_version
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
ISSUE_URL_BASE: Final = "https://github.com/BerriAI/litellm/issues/new"
|
||||
MAX_URL_LENGTH: Final = 6000
|
||||
MAX_FRAMES: Final = 12
|
||||
DISABLE_ENV_VAR: Final = "LITELLM_DISABLE_BUG_REPORT_LINK"
|
||||
NOTICE_PREFIX: Final = "This looks like a bug in LiteLLM rather than in your request."
|
||||
KNOWN_PROVIDERS: Final = frozenset(provider.value for provider in LlmProviders)
|
||||
|
||||
Surface = Literal["sdk", "proxy"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BugReport:
|
||||
surface: Surface
|
||||
exception_type: str
|
||||
litellm_frames: tuple[str, ...]
|
||||
litellm_version: str
|
||||
python_version: str
|
||||
call_type: str | None
|
||||
custom_llm_provider: str | None
|
||||
stream: bool | None
|
||||
config_lines: tuple[str, ...]
|
||||
|
||||
|
||||
def bug_report_enabled() -> bool:
|
||||
return os.getenv(DISABLE_ENV_VAR, "").lower() != "true"
|
||||
|
||||
|
||||
def should_report_bug(exc: object) -> bool:
|
||||
return bug_report_enabled() and isinstance(exc, BaseException) and getattr(exc, "status_code", None) is None
|
||||
|
||||
|
||||
def _format_frame(frame: traceback.FrameSummary, package_dir: Path, package_parent: Path) -> str | None:
|
||||
frame_path: Final = Path(frame.filename).resolve()
|
||||
try:
|
||||
frame_path.relative_to(package_dir)
|
||||
relative_path: Final = frame_path.relative_to(package_parent)
|
||||
except ValueError:
|
||||
return None
|
||||
return f"{relative_path.as_posix()}:{frame.lineno} in {frame.name}"
|
||||
|
||||
|
||||
def _get_litellm_frames(exc: BaseException) -> tuple[str, ...]:
|
||||
if exc.__traceback__ is None:
|
||||
return ()
|
||||
package_dir: Final = Path(litellm.__file__).resolve().parent
|
||||
package_parent: Final = package_dir.parent
|
||||
return tuple(
|
||||
frame_text
|
||||
for frame in traceback.extract_tb(exc.__traceback__)
|
||||
if (frame_text := _format_frame(frame, package_dir, package_parent)) is not None
|
||||
)[-MAX_FRAMES:]
|
||||
|
||||
|
||||
def allowlisted(value: object, allowed: frozenset[str]) -> str | None:
|
||||
return value if isinstance(value, str) and value in allowed else None
|
||||
|
||||
|
||||
def build_bug_report(
|
||||
exc: BaseException,
|
||||
*,
|
||||
surface: Surface,
|
||||
call_type: str | None = None,
|
||||
custom_llm_provider: object = None,
|
||||
stream: object = None,
|
||||
config_lines: tuple[str, ...] = (),
|
||||
) -> BugReport:
|
||||
return BugReport(
|
||||
surface=surface,
|
||||
exception_type=type(exc).__name__,
|
||||
litellm_frames=_get_litellm_frames(exc),
|
||||
litellm_version=litellm_version,
|
||||
python_version=platform.python_version(),
|
||||
call_type=call_type,
|
||||
custom_llm_provider=allowlisted(custom_llm_provider, KNOWN_PROVIDERS),
|
||||
stream=stream if isinstance(stream, bool) else None,
|
||||
config_lines=config_lines,
|
||||
)
|
||||
|
||||
|
||||
def _domain(report: BugReport) -> str:
|
||||
if report.surface == "sdk":
|
||||
return "Python SDK: the litellm package itself"
|
||||
if report.custom_llm_provider is not None:
|
||||
return "LLM translation: a specific provider's request or response"
|
||||
return "Proxy core: startup, config, health checks, endpoints"
|
||||
|
||||
|
||||
def _title(report: BugReport, frames: tuple[str, ...]) -> str:
|
||||
location: Final = frames[-1].split(":", 1)[0] if frames else "litellm"
|
||||
return f"[Bug]: {report.exception_type} in {location}"
|
||||
|
||||
|
||||
def _description(report: BugReport, frames: tuple[str, ...], config_lines: tuple[str, ...]) -> str:
|
||||
frame_block: Final = "LiteLLM frames:\n```\n" + "\n".join(frames) + "\n```\n\n" if frames else ""
|
||||
stream_line: Final = "" if report.stream is None else f"Stream: {str(report.stream).lower()}\n"
|
||||
config_block: Final = (
|
||||
"\nConfig (true/false flags and LiteLLM-defined values only):\n```\n" + "\n".join(config_lines) + "\n```\n"
|
||||
if config_lines
|
||||
else ""
|
||||
)
|
||||
return (
|
||||
"Auto-generated by LiteLLM's bug report link. It carries no request data or error text. "
|
||||
"Please describe what you were doing, and paste the error message from your log below "
|
||||
"if it contains nothing sensitive.\n\n"
|
||||
"```\n\n```\n\n"
|
||||
f"Exception: `{report.exception_type}`\n\n"
|
||||
f"{frame_block}"
|
||||
f"Surface: {report.surface}\n"
|
||||
f"Endpoint / call: {report.call_type or 'unknown'}\n"
|
||||
f"Provider: {report.custom_llm_provider or 'unknown'}\n"
|
||||
f"LiteLLM: {report.litellm_version}\n"
|
||||
f"Python: {report.python_version}\n"
|
||||
f"{stream_line}"
|
||||
f"{config_block}"
|
||||
)
|
||||
|
||||
|
||||
def _issue_url(report: BugReport, frames: tuple[str, ...], config_lines: tuple[str, ...]) -> str:
|
||||
deployment: Final[tuple[tuple[str, str], ...]] = (
|
||||
(("deployment", "pip / Python SDK"),)
|
||||
if report.surface == "sdk"
|
||||
else (("deployment", "Docker"),)
|
||||
if os.path.exists("/.dockerenv")
|
||||
else ()
|
||||
)
|
||||
fields: Final = (
|
||||
("template", "bug_report.yml"),
|
||||
("labels", "bug"),
|
||||
("title", _title(report, frames)),
|
||||
("version", report.litellm_version),
|
||||
("domain", _domain(report)),
|
||||
("description", _description(report, frames, config_lines)),
|
||||
) + deployment
|
||||
return f"{ISSUE_URL_BASE}?{urlencode(fields)}"
|
||||
|
||||
|
||||
def bug_report_issue_url(report: BugReport) -> str:
|
||||
frames: Final = report.litellm_frames
|
||||
config_lines: Final = report.config_lines
|
||||
candidates: Final = (
|
||||
*((frames, config_lines[:count]) for count in range(len(config_lines), -1, -1)),
|
||||
*((frames[index:], ()) for index in range(1, len(frames) + 1)),
|
||||
)
|
||||
return next(
|
||||
(
|
||||
url
|
||||
for candidate_frames, candidate_config in candidates
|
||||
if len(url := _issue_url(report, candidate_frames, candidate_config)) <= MAX_URL_LENGTH
|
||||
),
|
||||
_issue_url(report, (), ()),
|
||||
)
|
||||
|
||||
|
||||
def strip_bug_report_notice(message: str) -> str:
|
||||
index: Final = message.find(NOTICE_PREFIX)
|
||||
if index == -1:
|
||||
return message
|
||||
head: Final = message[:index]
|
||||
return head.removesuffix("\n")
|
||||
|
||||
|
||||
def bug_report_notice(report: BugReport) -> str:
|
||||
return (
|
||||
f"{NOTICE_PREFIX} File it with one click "
|
||||
f"(prefilled, no request data or error text, review before submitting): {bug_report_issue_url(report)}"
|
||||
)
|
||||
|
|
@ -10,6 +10,11 @@ import httpx
|
|||
|
||||
import litellm
|
||||
from litellm._logging import _ENABLE_SECRET_REDACTION, _redact_string, verbose_logger
|
||||
from litellm.litellm_core_utils.bug_report import (
|
||||
bug_report_notice,
|
||||
build_bug_report,
|
||||
should_report_bug,
|
||||
)
|
||||
from litellm.litellm_core_utils.secret_redaction import redact_string
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
|
|
@ -2673,7 +2678,21 @@ def exception_type(
|
|||
)
|
||||
else:
|
||||
raise APIConnectionError(
|
||||
message=f"{original_exception}\n{_redact_string(traceback.format_exc())}",
|
||||
message=(
|
||||
f"{original_exception}\n{_redact_string(traceback.format_exc())}"
|
||||
+ (
|
||||
"\n"
|
||||
+ bug_report_notice(
|
||||
build_bug_report(
|
||||
original_exception,
|
||||
surface="sdk",
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
)
|
||||
if should_report_bug(original_exception)
|
||||
else ""
|
||||
)
|
||||
),
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), # stub the request
|
||||
|
|
|
|||
|
|
@ -1,10 +1,24 @@
|
|||
import re
|
||||
from collections.abc import Iterator, Mapping
|
||||
from collections.abc import Generator, Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm.types.utils import OTEL_SPAN_SCOPES, TRUSTED_CALLBACK_VARS_FIELD, StandardCallbackDynamicParams
|
||||
|
||||
_CLIENT_CALLBACK_METADATA_SLOTS: Final[tuple[str, ...]] = ("litellm_metadata", "metadata")
|
||||
_inherited_message_logging_disabled: Final[ContextVar[bool]] = ContextVar(
|
||||
"inherited_message_logging_disabled", default=False
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def inherit_message_logging_privacy(disabled: bool) -> Generator[None]:
|
||||
token: Final = _inherited_message_logging_disabled.set(_inherited_message_logging_disabled.get() or disabled)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_inherited_message_logging_disabled.reset(token)
|
||||
|
||||
|
||||
def iter_client_callback_metadata_dicts(
|
||||
|
|
@ -143,7 +157,7 @@ def get_trusted_callback_params(kwargs: Mapping[str, Any] | None) -> tuple[tuple
|
|||
|
||||
|
||||
def initialize_standard_callback_dynamic_params(
|
||||
kwargs: dict | None = None,
|
||||
kwargs: dict[str, object] | None = None,
|
||||
) -> StandardCallbackDynamicParams:
|
||||
"""
|
||||
Initialize the standard callback dynamic params from the kwargs
|
||||
|
|
@ -179,4 +193,10 @@ def initialize_standard_callback_dynamic_params(
|
|||
if param in _trusted_overlay_callback_params:
|
||||
standard_callback_dynamic_params[param] = trusted_value
|
||||
|
||||
if _inherited_message_logging_disabled.get():
|
||||
private_params: Final[StandardCallbackDynamicParams] = {
|
||||
**standard_callback_dynamic_params,
|
||||
"turn_off_message_logging": True,
|
||||
}
|
||||
return private_params
|
||||
return standard_callback_dynamic_params
|
||||
|
|
|
|||
|
|
@ -107,6 +107,10 @@ from litellm.litellm_core_utils.redact_messages import (
|
|||
redact_streaming_responses_for_custom_logger,
|
||||
should_redact_message_logging,
|
||||
)
|
||||
from litellm.litellm_core_utils.served_output_texts import (
|
||||
SERVED_OUTPUT_TEXTS_KEY,
|
||||
overlay_served_output_texts,
|
||||
)
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
from litellm.llms.base_llm.search.transformation import SearchResponse
|
||||
from litellm.responses.utils import ResponseAPILoggingUtils
|
||||
|
|
@ -496,6 +500,14 @@ def mask_api_base_credentials(api_base: str) -> str:
|
|||
return api_base[:key_end] + "*" * 5 + api_base[-4:]
|
||||
|
||||
|
||||
def _timestamp_seconds(moment: object) -> float | None:
|
||||
if isinstance(moment, datetime.datetime):
|
||||
return moment.timestamp()
|
||||
if isinstance(moment, (int, float)):
|
||||
return float(moment)
|
||||
return None
|
||||
|
||||
|
||||
class Logging(LiteLLMLoggingBaseClass):
|
||||
global \
|
||||
supabaseClient, \
|
||||
|
|
@ -1634,10 +1646,12 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
return response.mcp_tool_call_response
|
||||
|
||||
def get_response_ms(self) -> float:
|
||||
return (
|
||||
self.model_call_details.get("end_time", datetime.datetime.now())
|
||||
- self.model_call_details.get("start_time", datetime.datetime.now())
|
||||
).total_seconds() * 1000
|
||||
now: Final = datetime.datetime.now()
|
||||
start_seconds: Final = _timestamp_seconds(self.model_call_details.get("start_time", now))
|
||||
end_seconds: Final = _timestamp_seconds(self.model_call_details.get("end_time", now))
|
||||
if start_seconds is None or end_seconds is None:
|
||||
return 0.0
|
||||
return (end_seconds - start_seconds) * 1000
|
||||
|
||||
def set_cost_breakdown(
|
||||
self,
|
||||
|
|
@ -5833,15 +5847,12 @@ class StandardLoggingPayloadSetup:
|
|||
|
||||
modified_final_response_obj: Final = redact_message_input_output_from_logging(
|
||||
model_call_details=kwargs,
|
||||
result=final_response_obj,
|
||||
result=overlay_served_output_texts(final_response_obj, kwargs.get(SERVED_OUTPUT_TEXTS_KEY)),
|
||||
)
|
||||
|
||||
if modified_final_response_obj is not None and isinstance(modified_final_response_obj, BaseModel):
|
||||
final_response_obj = modified_final_response_obj.model_dump()
|
||||
else:
|
||||
final_response_obj = modified_final_response_obj
|
||||
|
||||
return final_response_obj
|
||||
return modified_final_response_obj.model_dump()
|
||||
return modified_final_response_obj
|
||||
|
||||
@staticmethod
|
||||
def get_additional_headers(
|
||||
|
|
|
|||
|
|
@ -251,6 +251,6 @@ def update_response_metadata(
|
|||
return
|
||||
|
||||
metadata: Final = ResponseMetadata(result)
|
||||
metadata.set_hidden_params(logging_obj, model, kwargs)
|
||||
metadata.set_timing_metrics(start_time, end_time, logging_obj, include_overhead)
|
||||
metadata.set_hidden_params(logging_obj, model, kwargs)
|
||||
metadata.apply()
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from litellm.litellm_core_utils.classifier_logging import without_classifier_aud
|
|||
from litellm.litellm_core_utils.core_helpers import (
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
)
|
||||
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
redact_vertex_ai_metadata_from_litellm_params,
|
||||
redact_vertex_ai_metadata_from_logged_object,
|
||||
|
|
@ -267,6 +268,7 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
|
|||
model_call_details["messages"] = [{"role": "user", "content": REDACTED_BY_LITELLM}]
|
||||
model_call_details["prompt"] = ""
|
||||
model_call_details["input"] = ""
|
||||
model_call_details.pop(SERVED_OUTPUT_TEXTS_KEY, None)
|
||||
standard_logging_object: Final = model_call_details.get("standard_logging_object")
|
||||
if isinstance(standard_logging_object, Mapping):
|
||||
model_call_details["standard_logging_object"] = _redact_standard_logging_object(standard_logging_object)
|
||||
|
|
|
|||
171
litellm/litellm_core_utils/served_output_texts.py
Normal file
171
litellm/litellm_core_utils/served_output_texts.py
Normal file
|
|
@ -0,0 +1,171 @@
|
|||
"""Assistant text the caller received, per choice, so the logging payload stores the response a
|
||||
post-call guardrail rewrote rather than the provider response the proxy assembled before it ran."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Final, Literal
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.types.utils import ModelResponse, ModelResponseStream
|
||||
|
||||
SERVED_OUTPUT_TEXTS_KEY: Final = "served_output_texts"
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object])
|
||||
_JSON_LIST: Final = TypeAdapter(list[object])
|
||||
_TEXTS: Final = TypeAdapter(tuple[str | None, ...])
|
||||
|
||||
ServedTexts = tuple[str | None, ...]
|
||||
|
||||
|
||||
class _TextBlock(BaseModel):
|
||||
type: str
|
||||
text: str | None = None
|
||||
|
||||
|
||||
class _AnthropicMessage(BaseModel):
|
||||
type: Literal["message"]
|
||||
content: list[_TextBlock]
|
||||
|
||||
|
||||
class _ResponsesOutputItem(BaseModel):
|
||||
type: str
|
||||
content: list[_TextBlock] = []
|
||||
|
||||
|
||||
class _ResponsesResponse(BaseModel):
|
||||
object: Literal["response"]
|
||||
output: list[_ResponsesOutputItem]
|
||||
|
||||
|
||||
class _ChatChoices(BaseModel):
|
||||
choices: list[object]
|
||||
|
||||
|
||||
def _as_json_object(response: object) -> dict[str, object] | None:
|
||||
candidate: Final = response.model_dump() if isinstance(response, BaseModel) else response
|
||||
try:
|
||||
return _JSON_OBJECT.validate_python(candidate)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _joined_block_texts(blocks: Sequence[_TextBlock], *, text_type: str) -> str | None:
|
||||
texts: Final = tuple(block.text for block in blocks if block.type == text_type and block.text is not None)
|
||||
return "".join(texts) if texts else None
|
||||
|
||||
|
||||
def _chat_texts(response: ModelResponse) -> ServedTexts | None:
|
||||
texts: Final = tuple(
|
||||
choice.message.content if isinstance(choice.message.content, str) else None for choice in response.choices
|
||||
)
|
||||
return texts if any(text is not None for text in texts) else None
|
||||
|
||||
|
||||
def _anthropic_message_text(response: dict[str, object]) -> str | None:
|
||||
try:
|
||||
message: Final = _AnthropicMessage.model_validate(response)
|
||||
except ValidationError:
|
||||
return None
|
||||
return _joined_block_texts(message.content, text_type="text")
|
||||
|
||||
|
||||
def _responses_api_text(response: dict[str, object]) -> str | None:
|
||||
try:
|
||||
parsed: Final = _ResponsesResponse.model_validate(response)
|
||||
except ValidationError:
|
||||
return None
|
||||
texts: Final = tuple(
|
||||
text
|
||||
for item in parsed.output
|
||||
if item.type == "message" and (text := _joined_block_texts(item.content, text_type="output_text")) is not None
|
||||
)
|
||||
return "".join(texts) if texts else None
|
||||
|
||||
|
||||
def _chat_dict_texts(response: dict[str, object]) -> ServedTexts | None:
|
||||
try:
|
||||
_ChatChoices.model_validate(response)
|
||||
return _chat_texts(ModelResponse(**response))
|
||||
except (ValidationError, TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def served_output_texts(response: object) -> ServedTexts | None:
|
||||
if isinstance(response, ModelResponse):
|
||||
return _chat_texts(response)
|
||||
mapping: Final = _as_json_object(response)
|
||||
if mapping is None:
|
||||
return None
|
||||
chat_texts: Final = _chat_dict_texts(mapping)
|
||||
if chat_texts is not None:
|
||||
return chat_texts
|
||||
anthropic_text: Final = _anthropic_message_text(mapping)
|
||||
text: Final = anthropic_text if anthropic_text is not None else _responses_api_text(mapping)
|
||||
return (text,) if text is not None else None
|
||||
|
||||
|
||||
def served_stream_output_texts(chunks: Sequence[object]) -> ServedTexts | None:
|
||||
if chunks and all(isinstance(chunk, ModelResponseStream) for chunk in chunks):
|
||||
return _chat_stream_texts(tuple(chunk for chunk in chunks if isinstance(chunk, ModelResponseStream)))
|
||||
from litellm.proxy.guardrails.anthropic_sse import assemble_anthropic_sse_stream, is_anthropic_sse_stream
|
||||
|
||||
if not is_anthropic_sse_stream(chunks):
|
||||
return None
|
||||
assembled: Final = assemble_anthropic_sse_stream(chunks)
|
||||
return _chat_texts(assembled) if assembled is not None else None
|
||||
|
||||
|
||||
def _chat_stream_choice_text(chunks: Sequence[ModelResponseStream], index: int) -> str | None:
|
||||
contents: Final = tuple(
|
||||
content
|
||||
for chunk in chunks
|
||||
for choice in chunk.choices
|
||||
if choice.index == index and isinstance(content := choice.delta.content, str)
|
||||
)
|
||||
return "".join(contents) if contents else None
|
||||
|
||||
|
||||
def _chat_stream_texts(chunks: Sequence[ModelResponseStream]) -> ServedTexts | None:
|
||||
choice_count: Final = max((choice.index + 1 for chunk in chunks for choice in chunk.choices), default=0)
|
||||
texts: Final = tuple(_chat_stream_choice_text(chunks, index) for index in range(choice_count))
|
||||
return texts if any(text is not None for text in texts) else None
|
||||
|
||||
|
||||
def record_served_output_texts(model_call_details: dict[str, object], texts: ServedTexts | None) -> None:
|
||||
if texts is None:
|
||||
return
|
||||
model_call_details[SERVED_OUTPUT_TEXTS_KEY] = texts # rebind-ok: model_call_details is the shared kwargs bag
|
||||
|
||||
|
||||
def overlay_served_output_texts(
|
||||
response_obj: dict[str, object] | str | list[object] | None, served_texts: object
|
||||
) -> dict[str, object] | str | list[object] | None:
|
||||
if not isinstance(response_obj, dict):
|
||||
return response_obj
|
||||
logged: Final = _as_json_object(response_obj)
|
||||
if logged is None:
|
||||
return response_obj
|
||||
try:
|
||||
texts: Final = _TEXTS.validate_python(served_texts)
|
||||
choices: Final = _JSON_LIST.validate_python(logged.get("choices"))
|
||||
except ValidationError:
|
||||
return response_obj
|
||||
return {
|
||||
**logged,
|
||||
"choices": [
|
||||
_choice_with_text(choice, texts[index]) if index < len(texts) else choice
|
||||
for index, choice in enumerate(choices)
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _choice_with_text(choice: object, text: str | None) -> object:
|
||||
choice_obj: Final = _as_json_object(choice)
|
||||
if choice_obj is None or text is None:
|
||||
return choice
|
||||
message: Final = _as_json_object(choice_obj.get("message"))
|
||||
if message is None or message.get("content") == text:
|
||||
return choice
|
||||
return {**choice_obj, "message": {**message, "content": text}}
|
||||
|
|
@ -96,6 +96,7 @@ from ..common_utils import (
|
|||
AnthropicModelInfo,
|
||||
eager_input_streaming_flag,
|
||||
process_anthropic_headers,
|
||||
requires_native_compaction_beta,
|
||||
strip_advisor_blocks_from_messages,
|
||||
)
|
||||
|
||||
|
|
@ -1770,7 +1771,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
)
|
||||
return tools
|
||||
|
||||
def _ensure_beta_header(self, headers: dict, beta_value: str) -> None:
|
||||
def _ensure_beta_header(self, headers: dict[str, str], beta_value: str) -> None:
|
||||
"""
|
||||
Ensure a beta header value is present in the anthropic-beta header.
|
||||
Merges with existing values instead of overriding them.
|
||||
|
|
@ -1779,13 +1780,17 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
headers: Dictionary of headers to update
|
||||
beta_value: The beta header value to add
|
||||
"""
|
||||
existing_beta: Final = headers.get("anthropic-beta")
|
||||
if existing_beta is None:
|
||||
headers["anthropic-beta"] = beta_value
|
||||
return
|
||||
existing_values: Final = [beta.strip() for beta in existing_beta.split(",")]
|
||||
if beta_value not in existing_values:
|
||||
headers["anthropic-beta"] = f"{existing_beta}, {beta_value}"
|
||||
existing_values: Final = tuple(
|
||||
beta.strip()
|
||||
for key, value in headers.items()
|
||||
if key.lower() == "anthropic-beta"
|
||||
for beta in value.split(",")
|
||||
if beta.strip()
|
||||
)
|
||||
for key in tuple(headers):
|
||||
if key.lower() == "anthropic-beta":
|
||||
headers.pop(key)
|
||||
headers["anthropic-beta"] = ", ".join(dict.fromkeys((*existing_values, beta_value)))
|
||||
|
||||
def _ensure_context_management_beta_header(self, headers: dict, context_management: object) -> None:
|
||||
"""
|
||||
|
|
@ -1823,7 +1828,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
|
||||
)
|
||||
|
||||
def update_headers_with_optional_anthropic_beta(self, headers: dict, optional_params: dict) -> dict:
|
||||
def update_headers_with_optional_anthropic_beta(
|
||||
self, headers: dict, optional_params: dict, messages: Sequence[object] = ()
|
||||
) -> dict:
|
||||
"""Update headers with optional anthropic beta."""
|
||||
|
||||
# Skip adding beta headers for Vertex requests
|
||||
|
|
@ -1832,6 +1839,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
if is_vertex_request:
|
||||
return headers
|
||||
|
||||
if requires_native_compaction_beta(self._resolved_provider, optional_params, messages):
|
||||
self._ensure_beta_header(headers, ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value)
|
||||
|
||||
_tools: Final = optional_params.get("tools", [])
|
||||
for tool in _tools:
|
||||
if tool.get("type", None) and tool.get("type").startswith(ANTHROPIC_HOSTED_TOOLS.WEB_FETCH.value):
|
||||
|
|
@ -1928,8 +1938,6 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
custom_llm_provider=self._resolved_provider,
|
||||
)
|
||||
|
||||
headers = self.update_headers_with_optional_anthropic_beta(headers=headers, optional_params=optional_params)
|
||||
|
||||
# === Tool-name sanitization (single chokepoint) ===
|
||||
# Anthropic enforces ^[a-zA-Z0-9_-]{1,128}$ on every tool name. We
|
||||
# sanitize *here* -- not in map_openai_params -- because:
|
||||
|
|
@ -1976,6 +1984,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
message=f"{e}\nReceived Messages={messages}",
|
||||
) # don't use verbose_logger.exception, if exception is raised
|
||||
|
||||
self.update_headers_with_optional_anthropic_beta(
|
||||
headers=headers, optional_params=optional_params, messages=anthropic_messages
|
||||
)
|
||||
|
||||
## Auto-strip advisor blocks from history if advisor tool is absent.
|
||||
## Prevents Anthropic 400: advisor_tool_result in history requires advisor tool.
|
||||
_all_tools: Final = optional_params.get("tools") or []
|
||||
|
|
|
|||
|
|
@ -79,6 +79,27 @@ _CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
|
|||
_CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/")
|
||||
|
||||
|
||||
def requires_native_compaction_beta(
|
||||
custom_llm_provider: str,
|
||||
optional_params: Mapping[str, object],
|
||||
messages: Sequence[object],
|
||||
) -> bool:
|
||||
return custom_llm_provider == "anthropic" and (
|
||||
optional_params.get("compaction") is not None
|
||||
or any(
|
||||
isinstance(block, Mapping)
|
||||
and block.get("type") == "compaction"
|
||||
and isinstance(block.get("signature"), str)
|
||||
and bool(block.get("signature"))
|
||||
for message in messages
|
||||
if isinstance(message, Mapping)
|
||||
for content in (message.get("content"),)
|
||||
if isinstance(content, (list, tuple))
|
||||
for block in content
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def supports_anthropic_cache_control(model: str, custom_llm_provider: str | None) -> bool:
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.utils import supports_prompt_caching
|
||||
|
|
|
|||
71
litellm/llms/anthropic/compaction.py
Normal file
71
litellm/llms/anthropic/compaction.py
Normal file
|
|
@ -0,0 +1,71 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.llms.compaction import CompactionProtocol
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES, AnthropicCompaction
|
||||
|
||||
_MAPPING: Final = TypeAdapter(Mapping[str, object])
|
||||
_OBJECTS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
|
||||
_HEADERS: Final = TypeAdapter(dict[str, str])
|
||||
_EMPTY: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_CONFLICTS: Final = ("context_management", "response_format", "stop", "stop_sequences", "tool_choice")
|
||||
|
||||
|
||||
def supports_native_compaction(params: Mapping[str, object]) -> bool:
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
if params.get("custom_llm_provider") not in (None, "anthropic", "openai"):
|
||||
return False
|
||||
model: Final = str(params.get("model", "")).removeprefix("openai/").removeprefix("anthropic/")
|
||||
try:
|
||||
return get_model_info(model=model, custom_llm_provider="anthropic").get("supports_anthropic_compaction") is True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def compatible_defaults(payload: Mapping[str, object]) -> bool:
|
||||
return all(payload.get(key) is None for key in _CONFLICTS)
|
||||
|
||||
|
||||
def request_kwargs() -> Mapping[str, object]:
|
||||
operation: Final[AnthropicCompaction] = {"type": "summarize"}
|
||||
return MappingProxyType(
|
||||
{
|
||||
"compaction": operation,
|
||||
"extra_headers": _HEADERS.validate_python(
|
||||
MappingProxyType({"anthropic-beta": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value})
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _native_blocks(protocol: CompactionProtocol, response: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
|
||||
if protocol == "messages":
|
||||
return (
|
||||
_OBJECTS.validate_python(response.get("content", ())) if response.get("stop_reason") == "compaction" else ()
|
||||
)
|
||||
choices: Final = _OBJECTS.validate_python(response.get("choices", ()))
|
||||
choice: Final = choices[0] if len(choices) == 1 else _EMPTY
|
||||
message: Final = _MAPPING.validate_python(choice.get("message", _EMPTY))
|
||||
fields: Final = _MAPPING.validate_python(message.get("provider_specific_fields") or _EMPTY)
|
||||
return _OBJECTS.validate_python(fields.get("compaction_blocks", ()))
|
||||
|
||||
|
||||
def extract_summary(protocol: CompactionProtocol, response: Mapping[str, object]) -> str | None:
|
||||
blocks: Final = _native_blocks(protocol, response)
|
||||
block: Final = blocks[0] if len(blocks) == 1 else _EMPTY
|
||||
content: Final = block.get("content")
|
||||
return (
|
||||
content
|
||||
if block.get("type") == "compaction"
|
||||
and isinstance(block.get("signature"), str)
|
||||
and block.get("signature")
|
||||
and isinstance(content, str)
|
||||
and content.strip()
|
||||
else None
|
||||
)
|
||||
|
|
@ -2,8 +2,11 @@ import copy
|
|||
import hashlib
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.llms.anthropic.experimental_pass_through.utils import (
|
||||
is_reasoning_auto_summary_enabled,
|
||||
|
|
@ -16,6 +19,7 @@ OPENAI_MAX_TOOL_NAME_LENGTH: Final = 64
|
|||
TOOL_NAME_HASH_LENGTH: Final = 8
|
||||
TOOL_NAME_PREFIX_LENGTH: Final = OPENAI_MAX_TOOL_NAME_LENGTH - TOOL_NAME_HASH_LENGTH - 1 # 55
|
||||
PROVIDERS_PROXYING_AN_UNKNOWN_BACKEND: Final = frozenset({"litellm_proxy"})
|
||||
_COMPACTION_BLOCK: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
def _optional_attr(source: object, name: str) -> object:
|
||||
|
|
@ -36,6 +40,20 @@ def _thought_signature(provider_specific_fields: object) -> str | None:
|
|||
return signature if isinstance(signature, str) else None
|
||||
|
||||
|
||||
def _compaction_blocks(provider_specific_fields: object) -> tuple[Mapping[str, object], ...]:
|
||||
fields: Final = _as_string_mapping(provider_specific_fields)
|
||||
raw_blocks: Final = fields.get("compaction_blocks") if fields is not None else None
|
||||
return (
|
||||
tuple(
|
||||
block
|
||||
for raw_block in raw_blocks
|
||||
if (block := _as_string_mapping(raw_block)) is not None and block.get("type") == "compaction"
|
||||
)
|
||||
if isinstance(raw_blocks, (list, tuple))
|
||||
else ()
|
||||
)
|
||||
|
||||
|
||||
_ANTHROPIC_TOOL_SCHEMA_KEYS: Final = frozenset(
|
||||
{"name", "type", "input_schema", "description", "cache_control", "strict"}
|
||||
)
|
||||
|
|
@ -1330,7 +1348,11 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
tool_name_mapping: dict[str, str] | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
new_content: Final[list[dict[str, Any]]] = []
|
||||
for choice in choices:
|
||||
for choice, compaction_blocks in (
|
||||
(choice, _compaction_blocks(_optional_attr(choice.message, "provider_specific_fields")))
|
||||
for choice in choices
|
||||
):
|
||||
new_content.extend(_COMPACTION_BLOCK.validate_python(block) for block in compaction_blocks)
|
||||
# Handle thinking blocks first
|
||||
if hasattr(choice.message, "thinking_blocks") and choice.message.thinking_blocks:
|
||||
for thinking_block in choice.message.thinking_blocks:
|
||||
|
|
@ -1365,7 +1387,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
)
|
||||
|
||||
# Handle text content
|
||||
if choice.message.content is not None:
|
||||
if choice.message.content is not None and (choice.message.content != "" or not compaction_blocks):
|
||||
new_content.append(
|
||||
AnthropicResponseContentBlockText(type="text", text=choice.message.content).model_dump()
|
||||
)
|
||||
|
|
@ -1545,21 +1567,35 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
openai_finish_reason=openai_finish_reason
|
||||
)
|
||||
anthropic_finish_reason: Final = (
|
||||
"refusal"
|
||||
"compaction"
|
||||
if len(anthropic_content) == 1 and anthropic_content[0].get("type") == "compaction"
|
||||
else "refusal"
|
||||
if refusal_text is not None and translated_finish_reason != "max_tokens"
|
||||
else translated_finish_reason
|
||||
)
|
||||
# extract usage
|
||||
usage: Final[Usage] = getattr(response, "usage")
|
||||
anthropic_usage: Final = self._translate_openai_usage_to_anthropic_usage(usage)
|
||||
|
||||
if polyfill_result is not None and polyfill_result.iterations_usage is not None:
|
||||
message_iteration: Final[UsageIteration] = {
|
||||
"type": "message",
|
||||
"input_tokens": anthropic_usage["input_tokens"],
|
||||
"output_tokens": usage.completion_tokens or 0,
|
||||
}
|
||||
anthropic_usage["iterations"] = list(polyfill_result.iterations_usage) + [message_iteration]
|
||||
message_usage: Final = self._translate_openai_usage_to_anthropic_usage(usage)
|
||||
polyfill_iterations: Final = polyfill_result.iterations_usage if polyfill_result is not None else None
|
||||
anthropic_usage: Final[AnthropicUsage] = (
|
||||
TypeAdapter(AnthropicUsage).validate_python(
|
||||
MappingProxyType(
|
||||
{
|
||||
**message_usage,
|
||||
"iterations": (
|
||||
*polyfill_iterations,
|
||||
UsageIteration(
|
||||
type="message",
|
||||
input_tokens=message_usage.get("input_tokens", 0),
|
||||
output_tokens=usage.completion_tokens or 0,
|
||||
),
|
||||
),
|
||||
}
|
||||
)
|
||||
)
|
||||
if polyfill_iterations is not None
|
||||
else message_usage
|
||||
)
|
||||
|
||||
translated_obj: Final = AnthropicMessagesResponse(
|
||||
id=response.id,
|
||||
|
|
|
|||
|
|
@ -580,7 +580,9 @@ def anthropic_messages_handler(
|
|||
)
|
||||
if anthropic_messages_provider_config is None:
|
||||
# Route to Responses API for OpenAI / Azure, chat/completions for everything else.
|
||||
if _should_route_to_responses_api(custom_llm_provider, original_model, model):
|
||||
if kwargs.get("compaction") is None and _should_route_to_responses_api(
|
||||
custom_llm_provider, original_model, model
|
||||
):
|
||||
return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler(
|
||||
max_tokens=max_tokens,
|
||||
messages=messages,
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ from ...common_utils import (
|
|||
AnthropicError,
|
||||
AnthropicModelInfo,
|
||||
optionally_handle_anthropic_oauth,
|
||||
requires_native_compaction_beta,
|
||||
strip_advisor_blocks_from_messages,
|
||||
strip_encrypted_reasoning_blocks_from_anthropic_messages,
|
||||
)
|
||||
|
|
@ -74,6 +75,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
"tool_choice",
|
||||
"thinking",
|
||||
"context_management",
|
||||
*(("compaction",) if self._resolved_provider == "anthropic" else ()),
|
||||
"output_format",
|
||||
"inference_geo",
|
||||
"speed",
|
||||
|
|
@ -637,6 +639,9 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
)
|
||||
beta_values.update(existing_beta)
|
||||
|
||||
if requires_native_compaction_beta(custom_llm_provider, optional_params, messages):
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value)
|
||||
|
||||
# Check for context management
|
||||
context_management_param: Final = optional_params.get("context_management")
|
||||
if context_management_param is not None:
|
||||
|
|
|
|||
48
litellm/llms/compaction.py
Normal file
48
litellm/llms/compaction.py
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
||||
CompactionProtocol: TypeAlias = Literal["chat", "messages"]
|
||||
_MAPPING: Final = TypeAdapter(Mapping[str, object])
|
||||
_MESSAGES: Final = TypeAdapter(list[AllMessageValues])
|
||||
|
||||
|
||||
class NativeCompactionProvider(Protocol):
|
||||
def supports_native_compaction(self, params: Mapping[str, object]) -> bool: ...
|
||||
|
||||
def compatible_defaults(self, payload: Mapping[str, object]) -> bool: ...
|
||||
|
||||
def request_kwargs(self) -> Mapping[str, object]: ...
|
||||
|
||||
def extract_summary(self, protocol: CompactionProtocol, response: Mapping[str, object]) -> str | None: ...
|
||||
|
||||
|
||||
def get_native_compaction_provider(params: Mapping[str, object]) -> NativeCompactionProvider | None:
|
||||
from litellm.llms.anthropic import compaction
|
||||
|
||||
return compaction if compaction.supports_native_compaction(params) else None
|
||||
|
||||
|
||||
async def dispatch(router: Router, protocol: CompactionProtocol, payload: Mapping[str, object]) -> Mapping[str, object]:
|
||||
if protocol == "messages":
|
||||
return _MAPPING.validate_python(
|
||||
await router.aanthropic_messages(custom_llm_provider=None, client=None, **payload)
|
||||
)
|
||||
response: Final = await router.acompletion(
|
||||
model=str(payload["model"]),
|
||||
messages=_MESSAGES.validate_python(payload["messages"]),
|
||||
stream=False,
|
||||
**MappingProxyType(
|
||||
{key: value for key, value in payload.items() if key not in ("model", "messages", "stream")}
|
||||
),
|
||||
)
|
||||
return _MAPPING.validate_python(response.model_dump())
|
||||
46
litellm/llms/custom_httpx/asgi_handler.py
Normal file
46
litellm/llms/custom_httpx/asgi_handler.py
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from starlette.types import ASGIApp, Receive, Scope, Send
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared cache retains its legacy parameter mapping
|
||||
)
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ASGITarget:
|
||||
app: ASGIApp
|
||||
root_path: str
|
||||
client: tuple[str, int] | None
|
||||
|
||||
|
||||
_target: Final[ContextVar[_ASGITarget]] = ContextVar("httpx_asgi_target")
|
||||
|
||||
|
||||
async def _dispatch(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
target: Final = _target.get()
|
||||
await target.app({**scope, "root_path": target.root_path, "client": target.client}, receive, send)
|
||||
|
||||
|
||||
_TRANSPORT: Final = httpx.ASGITransport(app=_dispatch, raise_app_exceptions=False)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def get_async_asgi_client(
|
||||
app: ASGIApp, root_path: str = "", client: tuple[str, int] | None = None
|
||||
) -> Generator[httpx.AsyncClient]:
|
||||
handler: Final = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.ASGI,
|
||||
params={"transport": _TRANSPORT, "timeout": httpx.Timeout(None), "follow_redirects": False},
|
||||
)
|
||||
token: Final = _target.set(_ASGITarget(app, root_path, client))
|
||||
try:
|
||||
yield handler.client
|
||||
finally:
|
||||
_target.reset(token)
|
||||
|
|
@ -614,11 +614,15 @@ class AsyncHTTPHandler:
|
|||
client_alias: str | None = None, # name for client in logs
|
||||
ssl_verify: VerifyTypes | None = None,
|
||||
shared_session: Optional["ClientSession"] = None,
|
||||
transport: httpx.AsyncBaseTransport | None = None,
|
||||
follow_redirects: bool = True,
|
||||
):
|
||||
self.timeout = timeout
|
||||
self.event_hooks = event_hooks
|
||||
self.ssl_verify = ssl_verify
|
||||
self.shared_session = shared_session
|
||||
self.transport = transport
|
||||
self.follow_redirects = follow_redirects
|
||||
self._owns_client = True
|
||||
self._client = self.create_client(
|
||||
timeout=timeout,
|
||||
|
|
@ -651,6 +655,16 @@ class AsyncHTTPHandler:
|
|||
ssl_verify: VerifyTypes | None = None,
|
||||
shared_session: Optional["ClientSession"] = None,
|
||||
) -> httpx.AsyncClient:
|
||||
if self.transport is not None:
|
||||
return httpx.AsyncClient(
|
||||
transport=self.transport,
|
||||
event_hooks=event_hooks,
|
||||
timeout=timeout if timeout is not None else _DEFAULT_TIMEOUT,
|
||||
headers=get_default_headers(),
|
||||
cookies=blocked_cookie_jar(),
|
||||
follow_redirects=self.follow_redirects,
|
||||
trust_env=False,
|
||||
)
|
||||
# Get unified SSL configuration
|
||||
ssl_config: Final = get_ssl_configuration(ssl_verify)
|
||||
|
||||
|
|
@ -680,7 +694,7 @@ class AsyncHTTPHandler:
|
|||
cert=cert,
|
||||
headers=default_headers,
|
||||
cookies=blocked_cookie_jar(),
|
||||
follow_redirects=True,
|
||||
follow_redirects=self.follow_redirects,
|
||||
http2=http2_enabled(),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -14516,6 +14516,7 @@
|
|||
"source": "https://docs.anthropic.com/en/docs/about-claude/pricing"
|
||||
},
|
||||
"claude-sonnet-5": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 4e-06,
|
||||
"cache_read_input_token_cost": 2e-07,
|
||||
|
|
@ -14555,6 +14556,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-sonnet-4-6": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 6e-06,
|
||||
"cache_read_input_token_cost": 3e-07,
|
||||
|
|
@ -14770,6 +14772,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-opus-4-6": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
|
|
@ -14808,6 +14811,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-opus-4-6-20260205": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
|
|
@ -14845,6 +14849,7 @@
|
|||
"prompt_cache_min_tokens": 4096
|
||||
},
|
||||
"claude-opus-4-7": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
|
|
@ -14884,6 +14889,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-opus-4-7-20260416": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
|
|
@ -14922,6 +14928,7 @@
|
|||
"prompt_cache_min_tokens": 2048
|
||||
},
|
||||
"claude-fable-5": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_1hr": 2e-05,
|
||||
"cache_read_input_token_cost": 1e-06,
|
||||
|
|
@ -14961,6 +14968,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-fable-5-1": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_1hr": 2e-05,
|
||||
"cache_read_input_token_cost": 2.5e-07,
|
||||
|
|
@ -15001,6 +15009,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-opus-5": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
|
|
@ -15043,6 +15052,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-opus-4-8": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
|
|
@ -62148,6 +62158,7 @@
|
|||
"supports_audio_output": true
|
||||
},
|
||||
"claude-mythos-5": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_1hr": 2e-05,
|
||||
"cache_read_input_token_cost": 1e-06,
|
||||
|
|
@ -62187,6 +62198,7 @@
|
|||
}
|
||||
},
|
||||
"claude-mythos-5-1": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_1hr": 2e-05,
|
||||
"cache_read_input_token_cost": 2.5e-07,
|
||||
|
|
@ -62227,6 +62239,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-mythos-preview": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_1hr": 2e-05,
|
||||
"cache_read_input_token_cost": 1e-06,
|
||||
|
|
@ -77544,7 +77557,7 @@
|
|||
"supports_pdf_input": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": false,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": false
|
||||
|
|
@ -77560,12 +77573,12 @@
|
|||
"output_cost_per_token": 2.5e-07,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_audio_input": false,
|
||||
"supports_function_calling": false,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": false,
|
||||
"supports_tool_choice": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": false
|
||||
}
|
||||
|
|
|
|||
|
|
@ -50,6 +50,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
|
|||
"""Represents a LiteLLM_MCPServerTable record"""
|
||||
|
||||
server_id: str
|
||||
is_config: bool = Field(default=False, description="Whether this server is defined in config and is read-only.")
|
||||
server_name: str | None = None
|
||||
alias: str | None = None
|
||||
description: str | None = None
|
||||
|
|
|
|||
|
|
@ -7081,6 +7081,7 @@ class MCPServerManager:
|
|||
def _build_mcp_server_table(self, server: MCPServer) -> LiteLLM_MCPServerTable:
|
||||
return LiteLLM_MCPServerTable(
|
||||
server_id=server.server_id,
|
||||
is_config=self.is_config_declared_server(server.server_id) and server.server_id not in self.registry,
|
||||
server_name=server.server_name,
|
||||
alias=server.alias,
|
||||
description=(server.mcp_info.get("description") if server.mcp_info else None),
|
||||
|
|
|
|||
|
|
@ -24767,6 +24767,12 @@
|
|||
"title": "Is Byok",
|
||||
"type": "boolean"
|
||||
},
|
||||
"is_config": {
|
||||
"default": false,
|
||||
"description": "Whether this server is defined in config and is read-only.",
|
||||
"title": "Is Config",
|
||||
"type": "boolean"
|
||||
},
|
||||
"issuer": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
@ -27817,6 +27823,12 @@
|
|||
"title": "Is Byok",
|
||||
"type": "boolean"
|
||||
},
|
||||
"is_config": {
|
||||
"default": false,
|
||||
"description": "Whether this server is defined in config and is read-only.",
|
||||
"title": "Is Config",
|
||||
"type": "boolean"
|
||||
},
|
||||
"issuer": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -4310,7 +4310,10 @@ def _can_object_call_model(
|
|||
)
|
||||
return True
|
||||
|
||||
potential_models: Final = [model]
|
||||
from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent
|
||||
|
||||
compaction_parent: Final = native_compaction_parent(model)
|
||||
potential_models: Final = [model, compaction_parent] if compaction_parent is not None else [model]
|
||||
if model in litellm.model_alias_map:
|
||||
potential_models.append(litellm.model_alias_map[model])
|
||||
elif llm_router and model in llm_router.model_group_alias:
|
||||
|
|
|
|||
|
|
@ -71,12 +71,7 @@ def _as_proxy_exception(e: Exception) -> ProxyException:
|
|||
if isinstance(e, ProxyException):
|
||||
return e
|
||||
if PrismaDBExceptionHandler.is_database_service_unavailable_error(e):
|
||||
return ProxyException(
|
||||
message=PrismaDBExceptionHandler.database_unavailable_message(e),
|
||||
type=ProxyErrorTypes.no_db_connection,
|
||||
param="None",
|
||||
code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
)
|
||||
return PrismaDBExceptionHandler.service_unavailable_proxy_exception(e)
|
||||
return ProxyException(
|
||||
message="Authentication Error, " + str(e),
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
|
|
|
|||
|
|
@ -70,6 +70,7 @@ from litellm.types.agents import AgentResponse
|
|||
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
|
||||
|
||||
from .auth_checks import (
|
||||
TeamNotFoundError,
|
||||
_allowed_routes_check,
|
||||
allowed_routes_check,
|
||||
get_actual_routes,
|
||||
|
|
@ -147,6 +148,12 @@ class _JWTProvisioning:
|
|||
team_id_upsert: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class HeaderTeam:
|
||||
header_value: str
|
||||
team_id: str
|
||||
|
||||
|
||||
class AgentLookup(Protocol):
|
||||
"""The registered-agent lookups a JWT agent claim is matched against."""
|
||||
|
||||
|
|
@ -1871,48 +1878,104 @@ class JWTAuthManager:
|
|||
return True
|
||||
|
||||
@staticmethod
|
||||
def get_team_id_from_header(
|
||||
request_headers: Mapping[str, str] | None,
|
||||
allowed_team_ids: set[str],
|
||||
fallback_to_db_teams: bool = False,
|
||||
) -> str | None:
|
||||
"""
|
||||
Extract team_id from x-litellm-team-id header if present.
|
||||
Validates that the team is in the user's allowed teams from JWT.
|
||||
|
||||
Args:
|
||||
request_headers: Dictionary of request headers
|
||||
allowed_team_ids: Set of team IDs the user is allowed to access (from JWT)
|
||||
fallback_to_db_teams: When True and the JWT carries no team claims
|
||||
(allowed_team_ids is empty), the header value is returned
|
||||
provisionally and validated against DB memberships later in
|
||||
auth_builder instead of being rejected here.
|
||||
|
||||
Returns:
|
||||
The team_id from header if valid, None otherwise
|
||||
|
||||
Raises:
|
||||
HTTPException: If team_id is provided but not in allowed_team_ids
|
||||
"""
|
||||
def _team_header_value(request_headers: Mapping[str, str] | None) -> str | None:
|
||||
if not request_headers:
|
||||
return None
|
||||
|
||||
# Normalize headers to lowercase for case-insensitive lookup
|
||||
normalized_headers: Final = {k.lower(): v for k, v in request_headers.items()}
|
||||
header_team_id: Final = normalized_headers.get("x-litellm-team-id")
|
||||
return normalized_headers.get("x-litellm-team-id")
|
||||
|
||||
if not header_team_id:
|
||||
@staticmethod
|
||||
def _raise_header_team_not_allowed(header_value: str, allowed_team_ids: set[str]) -> NoReturn:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"Team '{header_value}' from x-litellm-team-id header is not in your JWT's allowed teams. Allowed teams: {list(allowed_team_ids)}",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _team_id_by_alias(
|
||||
team_alias: str,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> str | None:
|
||||
if prisma_client is None:
|
||||
return None
|
||||
try:
|
||||
team: Final = await get_team_object_by_alias(
|
||||
team_alias=team_alias,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except HTTPException as exc:
|
||||
if exc.status_code >= 500:
|
||||
raise
|
||||
return None
|
||||
return team.team_id
|
||||
|
||||
@staticmethod
|
||||
async def resolve_team_from_header(
|
||||
request_headers: Mapping[str, str] | None,
|
||||
allowed_team_ids: set[str],
|
||||
fallback_to_db_teams: bool,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> HeaderTeam | None:
|
||||
"""
|
||||
The team named by x-litellm-team-id, which may carry a team id or a team
|
||||
alias. A value that is already an allowed team id (or, under the DB
|
||||
fallback, an existing team id) never costs an alias lookup; an alias is
|
||||
accepted only when the team it names would have been accepted by id.
|
||||
Under the DB fallback only a team row that is provably absent falls
|
||||
through to the alias lookup; a read that failed for any other reason
|
||||
keeps the membership denial the id path already gives.
|
||||
|
||||
Raises:
|
||||
HTTPException: 403 when neither the value nor the team it aliases is
|
||||
an allowed team, or the DB fallback's membership denial when the
|
||||
value names no team at all; a 5xx from the alias lookup itself
|
||||
is re-raised rather than reported as a denial
|
||||
"""
|
||||
header_value: Final = JWTAuthManager._team_header_value(request_headers)
|
||||
if not header_value:
|
||||
return None
|
||||
|
||||
defer_to_db_membership: Final = fallback_to_db_teams and not allowed_team_ids
|
||||
if not defer_to_db_membership and header_team_id not in allowed_team_ids:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"Team '{header_team_id}' from x-litellm-team-id header is not in your JWT's allowed teams. Allowed teams: {list(allowed_team_ids)}",
|
||||
)
|
||||
if fallback_to_db_teams and not allowed_team_ids:
|
||||
try:
|
||||
await get_team_object(
|
||||
team_id=header_value,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id_upsert=False,
|
||||
)
|
||||
except TeamNotFoundError:
|
||||
aliased_team_id: Final = await JWTAuthManager._team_id_by_alias(
|
||||
header_value, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj
|
||||
)
|
||||
if aliased_team_id is None:
|
||||
JWTAuthManager._raise_header_team_membership_denial(header_value)
|
||||
return HeaderTeam(header_value=header_value, team_id=aliased_team_id)
|
||||
except HTTPException:
|
||||
JWTAuthManager._raise_header_team_membership_denial(header_value)
|
||||
return HeaderTeam(header_value=header_value, team_id=header_value)
|
||||
|
||||
verbose_proxy_logger.debug("Using team_id from x-litellm-team-id header: %s", header_team_id)
|
||||
return header_team_id
|
||||
if header_value in allowed_team_ids:
|
||||
verbose_proxy_logger.debug("Using team_id from x-litellm-team-id header: %s", header_value)
|
||||
return HeaderTeam(header_value=header_value, team_id=header_value)
|
||||
|
||||
team_id_by_alias: Final = await JWTAuthManager._team_id_by_alias(
|
||||
header_value, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj
|
||||
)
|
||||
if team_id_by_alias is None or team_id_by_alias not in allowed_team_ids:
|
||||
JWTAuthManager._raise_header_team_not_allowed(header_value, allowed_team_ids)
|
||||
verbose_proxy_logger.debug("Using team_id %s for x-litellm-team-id alias: %s", team_id_by_alias, header_value)
|
||||
return HeaderTeam(header_value=header_value, team_id=team_id_by_alias)
|
||||
|
||||
@staticmethod
|
||||
async def map_user_to_teams(
|
||||
|
|
@ -2264,31 +2327,34 @@ class JWTAuthManager:
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _raise_header_team_membership_denial(team_id: str) -> NoReturn:
|
||||
def _raise_header_team_membership_denial(header_value: str) -> NoReturn:
|
||||
"""
|
||||
The single denial shape for a provisional x-litellm-team-id header,
|
||||
raised identically for nonexistent teams and for teams the user is not
|
||||
a member of, so the response does not reveal whether a team id exists.
|
||||
a member of, and naming only the value the caller sent, so the response
|
||||
reveals neither whether a team exists nor which id an alias maps to.
|
||||
"""
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=(f"Team '{team_id}' (from x-litellm-team-id header) is not in your team memberships."),
|
||||
detail=(f"Team '{header_value}' (from x-litellm-team-id header) is not in your team memberships."),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _validate_header_team_in_db_membership(
|
||||
team_id: str,
|
||||
user_object: LiteLLM_UserTable | None,
|
||||
header_value: str,
|
||||
) -> None:
|
||||
"""
|
||||
A provisional team_id from the x-litellm-team-id header (accepted without
|
||||
JWT-team validation when the JWT carries no team claims) must exist in the
|
||||
user's DB team memberships before it becomes request context.
|
||||
user's DB team memberships before it becomes request context. The denial
|
||||
names `header_value`, the id or alias the caller sent, not `team_id`.
|
||||
"""
|
||||
user_team_ids: Final = user_object.teams if user_object else []
|
||||
if team_id in user_team_ids:
|
||||
return
|
||||
JWTAuthManager._raise_header_team_membership_denial(team_id)
|
||||
JWTAuthManager._raise_header_team_membership_denial(header_value)
|
||||
|
||||
@staticmethod
|
||||
async def auth_builder(
|
||||
|
|
@ -2514,13 +2580,17 @@ class JWTAuthManager:
|
|||
if specific_team_id and not db_team_fallback:
|
||||
all_team_ids.add(specific_team_id)
|
||||
|
||||
header_team_id: Final = JWTAuthManager.get_team_id_from_header(
|
||||
header_team: Final = await JWTAuthManager.resolve_team_from_header(
|
||||
request_headers=request_headers,
|
||||
allowed_team_ids=all_team_ids,
|
||||
fallback_to_db_teams=db_team_fallback,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if header_team_id:
|
||||
team_id = header_team_id
|
||||
if header_team:
|
||||
team_id = header_team.team_id
|
||||
# A provisional header team (accepted only because the JWT carries no
|
||||
# team claims) is validated against DB membership further down; never
|
||||
# upsert it here or an attacker-supplied x-litellm-team-id would create
|
||||
|
|
@ -2538,7 +2608,7 @@ class JWTAuthManager:
|
|||
except HTTPException:
|
||||
if not db_team_fallback:
|
||||
raise
|
||||
JWTAuthManager._raise_header_team_membership_denial(team_id)
|
||||
JWTAuthManager._raise_header_team_membership_denial(header_team.header_value)
|
||||
elif not team_id and not db_team_fallback:
|
||||
## SPECIFIC TEAM ID
|
||||
(
|
||||
|
|
@ -2679,10 +2749,11 @@ class JWTAuthManager:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id_upsert=team_id_upsert,
|
||||
)
|
||||
elif db_team_fallback and team_id == header_team_id:
|
||||
elif db_team_fallback and header_team is not None and team_id == header_team.team_id:
|
||||
JWTAuthManager._validate_header_team_in_db_membership(
|
||||
team_id=team_id,
|
||||
user_object=user_object,
|
||||
header_value=header_team.header_value,
|
||||
)
|
||||
if not JWTAuthManager._is_team_route_allowed(
|
||||
route=route,
|
||||
|
|
@ -2692,7 +2763,8 @@ class JWTAuthManager:
|
|||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=(
|
||||
f"Team '{team_id}' (from x-litellm-team-id header) is not allowed to access route '{route}'."
|
||||
f"Team '{header_team.header_value}' (from x-litellm-team-id header) "
|
||||
f"is not allowed to access route '{route}'."
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
185
litellm/proxy/bug_report_config.py
Normal file
185
litellm/proxy/bug_report_config.py
Normal file
|
|
@ -0,0 +1,185 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import functools
|
||||
import inspect
|
||||
from collections.abc import Mapping, Sequence
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.bug_report import KNOWN_PROVIDERS, BugReport, allowlisted, build_bug_report
|
||||
from litellm.proxy._types import ConfigGeneralSettings
|
||||
from litellm.router_utils.routing_groups import VALID_ROUTING_STRATEGIES
|
||||
from litellm.types.caching import LiteLLMCacheType
|
||||
from litellm.types.guardrails import GuardrailEventHooks, LitellmParams, SupportedGuardrailIntegrations
|
||||
from litellm.types.secret_managers.main import KeyManagementSystem
|
||||
|
||||
_OBJECT_MAP: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
_OBJECT_LIST: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...])
|
||||
_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
|
||||
|
||||
def _object_map(value: object) -> Mapping[str, object]:
|
||||
try:
|
||||
return _OBJECT_MAP.validate_python(value)
|
||||
except ValidationError:
|
||||
return MappingProxyType({})
|
||||
|
||||
|
||||
def _object_list(value: object) -> Sequence[object]:
|
||||
try:
|
||||
return _OBJECT_LIST.validate_python(value)
|
||||
except ValidationError:
|
||||
return ()
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _known_values() -> frozenset[str]:
|
||||
from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry
|
||||
|
||||
return frozenset(
|
||||
(
|
||||
*VALID_ROUTING_STRATEGIES,
|
||||
*KNOWN_PROVIDERS,
|
||||
*litellm._known_custom_logger_compatible_callbacks, # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownArgumentType] # untyped List of the callback Literal's args, no public alias
|
||||
*CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE,
|
||||
*(member.value for member in LiteLLMCacheType),
|
||||
*(member.value for member in KeyManagementSystem),
|
||||
*(member.value for member in SupportedGuardrailIntegrations),
|
||||
*(member.value for member in GuardrailEventHooks),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _module_level_names(node: ast.stmt) -> tuple[str, ...]:
|
||||
match node:
|
||||
case ast.Assign(targets=targets):
|
||||
return tuple(target.id for target in targets if isinstance(target, ast.Name))
|
||||
case ast.AnnAssign(target=ast.Name(id=name)):
|
||||
return (name,)
|
||||
case ast.ImportFrom(names=aliases):
|
||||
return tuple(alias.asname or alias.name for alias in aliases)
|
||||
case _:
|
||||
return ()
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _litellm_settings_keys() -> frozenset[str]:
|
||||
tree: Final = ast.parse(Path(litellm.__file__).read_text())
|
||||
return frozenset(name for node in tree.body for name in _module_level_names(node))
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _router_settings_keys() -> frozenset[str]:
|
||||
from litellm.router import Router
|
||||
|
||||
return frozenset(name for name in inspect.signature(Router.__init__).parameters if name != "self") # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped params, only names are read
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _cache_params_keys() -> frozenset[str]:
|
||||
from litellm.caching.caching import Cache
|
||||
|
||||
return frozenset(name for name in inspect.signature(Cache.__init__).parameters if name != "self") # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped params, only names are read
|
||||
|
||||
|
||||
def _render_json(value: JsonValue) -> str | None:
|
||||
match value:
|
||||
case bool():
|
||||
return str(value).lower()
|
||||
case str():
|
||||
return value if value in _known_values() else None
|
||||
case list():
|
||||
known_items: Final = tuple(rendered for item in value if (rendered := _render_json(item)) is not None)
|
||||
return f"[{', '.join(known_items)}]" if known_items else None
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _render(value: object) -> str | None:
|
||||
try:
|
||||
return _render_json(_JSON.validate_python(value))
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _section_lines(section: str, values: Mapping[str, object], known_keys: frozenset[str]) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
f"{section}.{key} = {rendered}"
|
||||
for key, value in values.items()
|
||||
if key in known_keys and (rendered := _render(value)) is not None
|
||||
)
|
||||
|
||||
|
||||
def _guardrail_lines(guardrails: object) -> tuple[str, ...]:
|
||||
known_keys: Final = frozenset(LitellmParams.model_fields)
|
||||
return tuple(
|
||||
line
|
||||
for index, guardrail in enumerate(_object_list(guardrails))
|
||||
for line in _section_lines(
|
||||
f"guardrails[{index}].litellm_params", _object_map(_object_map(guardrail).get("litellm_params")), known_keys
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _deployment_provider(model: object) -> str | None:
|
||||
prefix: Final = model.split("/", 1)[0] if isinstance(model, str) and "/" in model else None
|
||||
return allowlisted(prefix, KNOWN_PROVIDERS)
|
||||
|
||||
|
||||
def _model_list_lines(model_list: object) -> tuple[str, ...]:
|
||||
providers: Final = tuple(
|
||||
sorted(
|
||||
frozenset(
|
||||
provider
|
||||
for deployment in _object_list(model_list)
|
||||
if (
|
||||
provider := _deployment_provider(
|
||||
_object_map(_object_map(deployment).get("litellm_params")).get("model")
|
||||
)
|
||||
)
|
||||
is not None
|
||||
)
|
||||
)
|
||||
)
|
||||
return (f"model_list[*].provider = [{', '.join(providers)}]",) if providers else ()
|
||||
|
||||
|
||||
def safe_config_lines(config: Mapping[str, object], general_settings: Mapping[str, object]) -> tuple[str, ...]:
|
||||
litellm_settings: Final = _object_map(config.get("litellm_settings"))
|
||||
return (
|
||||
*_section_lines("general_settings", general_settings, frozenset(ConfigGeneralSettings.model_fields)),
|
||||
*_section_lines("litellm_settings", litellm_settings, _litellm_settings_keys()),
|
||||
*_section_lines(
|
||||
"litellm_settings.cache_params", _object_map(litellm_settings.get("cache_params")), _cache_params_keys()
|
||||
),
|
||||
*_section_lines("router_settings", _object_map(config.get("router_settings")), _router_settings_keys()),
|
||||
*_guardrail_lines(config.get("guardrails")),
|
||||
*_model_list_lines(config.get("model_list")),
|
||||
)
|
||||
|
||||
|
||||
def build_proxy_bug_report(
|
||||
exc: BaseException,
|
||||
*,
|
||||
call_type: str | None = None,
|
||||
custom_llm_provider: object = None,
|
||||
stream: object = None,
|
||||
) -> BugReport:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
return build_bug_report(
|
||||
exc,
|
||||
surface="proxy",
|
||||
call_type=call_type,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
stream=stream,
|
||||
config_lines=safe_config_lines(
|
||||
proxy_server.proxy_config.config,
|
||||
_object_map(proxy_server.general_settings), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # bare dict global, validated by _object_map
|
||||
),
|
||||
)
|
||||
|
|
@ -19,6 +19,7 @@ from typing import (
|
|||
overload,
|
||||
runtime_checkable,
|
||||
)
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
|
|
@ -45,6 +46,12 @@ from litellm.constants import (
|
|||
UNSAFE_PROXY_RESPONSE_HEADERS,
|
||||
)
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.bug_report import (
|
||||
allowlisted,
|
||||
bug_report_notice,
|
||||
should_report_bug,
|
||||
strip_bug_report_notice,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
get_or_create_metadata_bucket,
|
||||
independent_snapshot,
|
||||
|
|
@ -61,16 +68,21 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import (
|
|||
get_response_headers,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.served_output_texts import (
|
||||
record_served_output_texts,
|
||||
served_output_texts,
|
||||
)
|
||||
from litellm.litellm_core_utils.streaming_handler import (
|
||||
backfill_missing_cache_usage_fields,
|
||||
)
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
can_key_call_resolved_model,
|
||||
request_skips_budget_checks,
|
||||
tag_max_budget_check_for_tags,
|
||||
)
|
||||
from litellm.proxy.auth.auth_utils import check_response_size_is_safe, get_request_route
|
||||
from litellm.proxy.bug_report_config import build_proxy_bug_report
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
get_logging_caching_headers,
|
||||
get_remaining_tokens_and_requests_from_request_data,
|
||||
|
|
@ -95,6 +107,7 @@ from litellm.proxy.common_utils.sse_keepalive import (
|
|||
)
|
||||
from litellm.proxy.dd_span_tagger import DDSpanTagger
|
||||
from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression
|
||||
from litellm.proxy.native_compaction import with_proxy_compaction_executor
|
||||
from litellm.proxy.route_llm_request import route_request
|
||||
from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails
|
||||
from litellm.router import Router
|
||||
|
|
@ -107,6 +120,10 @@ from litellm.types.router_weights import validate_router_weights
|
|||
_LateResponseT = TypeVar("_LateResponseT", bound=Response)
|
||||
_LlmCallT = TypeVar("_LlmCallT")
|
||||
|
||||
KNOWN_PROXY_ROUTES: Final = frozenset(
|
||||
route for member in LiteLLMRoutes for route in member.value if route.startswith("/")
|
||||
)
|
||||
|
||||
ProxyRouteType: TypeAlias = Literal[
|
||||
"acompletion",
|
||||
"aembedding",
|
||||
|
|
@ -2543,7 +2560,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
user_model=user_model,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
llm_call_task: Final = asyncio.create_task(llm_call)
|
||||
llm_call_task: Final = asyncio.create_task(with_proxy_compaction_executor(llm_call, request))
|
||||
tasks.append(llm_call_task)
|
||||
|
||||
llm_responses: Final = asyncio.gather(*tasks) # run the moderation check in parallel to the actual llm api call
|
||||
|
|
@ -2790,6 +2807,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
)
|
||||
record_served_output_texts(logging_obj.model_call_details, served_output_texts(response))
|
||||
except Exception:
|
||||
_exception_raised = True
|
||||
raise
|
||||
|
|
@ -3669,8 +3687,27 @@ class ProxyBaseLLMRequestProcessing:
|
|||
_code = _exc_status_code
|
||||
else:
|
||||
_code = status.HTTP_500_INTERNAL_SERVER_ERROR
|
||||
if should_report_bug(e):
|
||||
proxy_server_request: Final = self.data.get("proxy_server_request")
|
||||
request_url: Final = (
|
||||
proxy_server_request.get("url") if isinstance(proxy_server_request, Mapping) else None
|
||||
)
|
||||
request_path: Final = urlparse(str(request_url)).path if request_url is not None else None
|
||||
verbose_proxy_logger.error(
|
||||
bug_report_notice(
|
||||
build_proxy_bug_report(
|
||||
e,
|
||||
call_type=allowlisted(request_path, KNOWN_PROXY_ROUTES),
|
||||
custom_llm_provider=self.data.get("custom_llm_provider"),
|
||||
stream=self.data.get("stream"),
|
||||
)
|
||||
)
|
||||
)
|
||||
client_message: Final = getattr(e, "message", error_msg)
|
||||
raise ProxyException(
|
||||
message=redact_internal_details_from_client_message(getattr(e, "message", error_msg)),
|
||||
message=redact_internal_details_from_client_message(
|
||||
strip_bug_report_notice(client_message) if isinstance(client_message, str) else error_msg
|
||||
),
|
||||
type=openai_error_type(e, _code),
|
||||
param=openai_error_param(e),
|
||||
openai_code=getattr(e, "code", None),
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import re
|
||||
from collections.abc import Awaitable, Callable, Iterator
|
||||
from http import HTTPStatus
|
||||
from typing import Final, Protocol, TypeVar
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
|
@ -378,6 +379,15 @@ class PrismaDBExceptionHandler:
|
|||
"The proxy deployment needs attention."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def service_unavailable_proxy_exception(e: Exception) -> ProxyException:
|
||||
return ProxyException(
|
||||
message=PrismaDBExceptionHandler.database_unavailable_message(e),
|
||||
type=ProxyErrorTypes.no_db_connection,
|
||||
param="None",
|
||||
code=HTTPStatus.SERVICE_UNAVAILABLE.value,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def find_database_service_unavailable_error_in_chain(e: BaseException) -> Exception | None:
|
||||
"""The exception in the ``__cause__`` / ``__context__`` chain that
|
||||
|
|
|
|||
|
|
@ -4475,9 +4475,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
# 'metadata' and 'litellm_metadata' fields from litellm_params
|
||||
standard_logging_object: Final = kwargs.get("standard_logging_object") or {}
|
||||
request_metadata: Final = get_litellm_metadata_from_kwargs(kwargs)
|
||||
if request_metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY):
|
||||
# Internal sub-calls bill spend to the caller but are not the caller's
|
||||
# traffic; charging them here would let background evals eat TPM headroom.
|
||||
origin: Final = request_metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)
|
||||
if origin and origin != "autorouter_compaction":
|
||||
# Background evaluations keep their exemption; foreground compaction
|
||||
# is necessary caller traffic and consumes the caller's token limits.
|
||||
return []
|
||||
standard_logging_metadata: Final = standard_logging_object.get("metadata") or {}
|
||||
|
||||
|
|
|
|||
|
|
@ -47,6 +47,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
object_permission_cache_key,
|
||||
user_object_permission_id_cache_key,
|
||||
)
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage
|
||||
from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import (
|
||||
|
|
@ -2810,6 +2811,9 @@ async def ui_view_users(
|
|||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e):
|
||||
verbose_proxy_logger.warning("Database unavailable during user search: %s", type(e).__name__)
|
||||
raise PrismaDBExceptionHandler.service_unavailable_proxy_exception(e) from e
|
||||
verbose_proxy_logger.exception("Error searching users: %s", e)
|
||||
raise HTTPException(status_code=500, detail=f"Error searching users: {e}")
|
||||
|
||||
|
|
|
|||
90
litellm/proxy/native_compaction.py
Normal file
90
litellm/proxy/native_compaction.py
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
import asyncio
|
||||
from collections.abc import Awaitable, Mapping
|
||||
from contextvars import Context
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeVar
|
||||
|
||||
from fastapi import Request
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from starlette.types import ASGIApp
|
||||
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
|
||||
inherit_message_logging_privacy,
|
||||
initialize_standard_callback_dynamic_params,
|
||||
)
|
||||
from litellm.llms.custom_httpx.asgi_handler import get_async_asgi_client
|
||||
from litellm.proxy.litellm_pre_call_utils import UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS
|
||||
from litellm.router_strategy.complexity_router.context_compaction import (
|
||||
compaction_executor,
|
||||
native_compaction_call,
|
||||
)
|
||||
|
||||
_ResultT: Final = TypeVar("_ResultT")
|
||||
_ASGI_APP: Final = TypeAdapter[ASGIApp](ASGIApp)
|
||||
_ROOT_PATH: Final = TypeAdapter(str)
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
|
||||
_REMOVED_HEADERS: Final = frozenset(
|
||||
(
|
||||
b"content-length",
|
||||
b"x-litellm-call-id",
|
||||
b"x-litellm-num-retries",
|
||||
b"x-litellm-timeout",
|
||||
b"x-litellm-stream-timeout",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def with_proxy_compaction_executor(call: Awaitable[_ResultT], request: Request) -> _ResultT:
|
||||
async def execute(
|
||||
protocol: Literal["chat", "messages"], payload: Mapping[str, object], parent_model: str | None = None
|
||||
) -> Mapping[str, object]:
|
||||
logging_disabled: Final = initialize_standard_callback_dynamic_params().get("turn_off_message_logging") is True
|
||||
|
||||
async def dispatch() -> Mapping[str, object]:
|
||||
scope: Final = _JSON_OBJECT.validate_python(request.scope)
|
||||
root_path: Final = _ROOT_PATH.validate_python(scope.get("root_path", ""))
|
||||
path: Final = "/v1/chat/completions" if protocol == "chat" else "/v1/messages"
|
||||
url: Final = str(request.url.replace(path=root_path.rstrip("/") + path, query="", fragment=""))
|
||||
headers: Final = tuple(
|
||||
(name, value)
|
||||
for name, value in request.headers.raw
|
||||
if name.lower() not in _REMOVED_HEADERS
|
||||
and not (logging_disabled and name.decode("latin-1").lower() in UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS)
|
||||
)
|
||||
with (
|
||||
native_compaction_call(parent_model, str(payload["model"])),
|
||||
inherit_message_logging_privacy(logging_disabled),
|
||||
):
|
||||
with get_async_asgi_client(
|
||||
app=_ASGI_APP.validate_python(scope["app"]),
|
||||
root_path=root_path,
|
||||
client=request.client,
|
||||
) as client:
|
||||
async with client.stream(
|
||||
"POST", url, headers=headers, json=_JSON_OBJECT.validate_python(payload)
|
||||
) as response:
|
||||
if not response.is_success:
|
||||
raise BadRequestError(
|
||||
message=f"Native compaction child request failed (HTTP {response.status_code})",
|
||||
model="context_compaction",
|
||||
llm_provider="",
|
||||
)
|
||||
body: Final = await response.aread()
|
||||
try:
|
||||
return MappingProxyType(_JSON_OBJECT.validate_json(body))
|
||||
except ValidationError:
|
||||
raise BadRequestError(
|
||||
message="Native compaction child returned an invalid JSON object",
|
||||
model="context_compaction",
|
||||
llm_provider="",
|
||||
) from None
|
||||
|
||||
task: Final = Context().run(asyncio.create_task, dispatch())
|
||||
return await task
|
||||
|
||||
token: Final = compaction_executor.set(execute)
|
||||
try:
|
||||
return await call
|
||||
finally:
|
||||
compaction_executor.reset(token)
|
||||
|
|
@ -75,6 +75,11 @@ from litellm.constants import (
|
|||
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
|
||||
)
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.bug_report import (
|
||||
allowlisted,
|
||||
bug_report_notice,
|
||||
should_report_bug,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
_init_custom_logger_compatible_class,
|
||||
)
|
||||
|
|
@ -367,10 +372,12 @@ from litellm.proxy.auth.user_api_key_auth import (
|
|||
user_api_key_auth_websocket,
|
||||
)
|
||||
from litellm.proxy.batches_endpoints.endpoints import router as batches_router
|
||||
from litellm.proxy.bug_report_config import build_proxy_bug_report
|
||||
|
||||
## Import All Misc routes here ##
|
||||
from litellm.proxy.caching_routes import router as caching_router
|
||||
from litellm.proxy.common_request_processing import (
|
||||
KNOWN_PROXY_ROUTES,
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
_is_azure_model_router_request,
|
||||
_should_return_raw_model_name,
|
||||
|
|
@ -1487,23 +1494,29 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
|
|||
## Initialize shared aiohttp session for connection reuse
|
||||
shared_aiohttp_session = await _initialize_shared_aiohttp_session()
|
||||
|
||||
model_info_scheduler: Final = scheduler if scheduler is not None else AsyncIOScheduler()
|
||||
model_info_scheduler.add_job(
|
||||
ProxyStartupEvent.refresh_model_info,
|
||||
"interval",
|
||||
seconds=MODEL_INFO_REFRESH_SECONDS,
|
||||
id="refresh_model_info",
|
||||
next_run_time=datetime.now(timezone.utc),
|
||||
max_instances=1,
|
||||
replace_existing=True,
|
||||
model_info_refresh_disabled: Final = (
|
||||
"disable_model_info_refresh" in general_settings and general_settings["disable_model_info_refresh"] is True
|
||||
)
|
||||
if not model_info_scheduler.running:
|
||||
model_info_scheduler.start()
|
||||
model_info_scheduler: Final = (
|
||||
None if model_info_refresh_disabled else scheduler if scheduler is not None else AsyncIOScheduler()
|
||||
)
|
||||
if model_info_scheduler is not None:
|
||||
model_info_scheduler.add_job(
|
||||
ProxyStartupEvent.refresh_model_info,
|
||||
"interval",
|
||||
seconds=MODEL_INFO_REFRESH_SECONDS,
|
||||
id="refresh_model_info",
|
||||
next_run_time=datetime.now(timezone.utc),
|
||||
max_instances=1,
|
||||
replace_existing=True,
|
||||
)
|
||||
if not model_info_scheduler.running:
|
||||
model_info_scheduler.start()
|
||||
|
||||
# End of startup event
|
||||
yield
|
||||
|
||||
if model_info_scheduler.running:
|
||||
if model_info_scheduler is not None and model_info_scheduler.running:
|
||||
model_info_scheduler.remove_job("refresh_model_info")
|
||||
if model_info_scheduler is not scheduler:
|
||||
model_info_scheduler.shutdown(wait=False)
|
||||
|
|
@ -1964,7 +1977,21 @@ async def otel_request_validation_exception_handler(request: Request, exc: Reque
|
|||
async def otel_unhandled_exception_handler(request: Request, exc: Exception):
|
||||
if isinstance(exc, (ProxyException, HTTPException, RequestValidationError)):
|
||||
raise exc
|
||||
if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(exc):
|
||||
verbose_proxy_logger.warning("Database unavailable during request: %s", type(exc).__name__)
|
||||
return await openai_exception_handler(
|
||||
request=request, exc=PrismaDBExceptionHandler.service_unavailable_proxy_exception(exc)
|
||||
)
|
||||
verbose_proxy_logger.exception("Unhandled exception in request: %s", type(exc).__name__)
|
||||
if should_report_bug(exc):
|
||||
verbose_proxy_logger.error(
|
||||
bug_report_notice(
|
||||
build_proxy_bug_report(
|
||||
exc,
|
||||
call_type=allowlisted(request.url.path, KNOWN_PROXY_ROUTES),
|
||||
)
|
||||
)
|
||||
)
|
||||
_close_dangling_otel_server_span(request, 500, exc=exc)
|
||||
return JSONResponse(
|
||||
status_code=500,
|
||||
|
|
|
|||
|
|
@ -57,6 +57,11 @@ from litellm.constants import (
|
|||
SPEND_LOG_WRITE_BATCH_MAX_BYTES,
|
||||
SPEND_LOG_WRITE_BATCH_MAX_ROWS,
|
||||
)
|
||||
from litellm.litellm_core_utils.bug_report import (
|
||||
bug_report_notice,
|
||||
should_report_bug,
|
||||
strip_bug_report_notice,
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
CommonProxyErrors,
|
||||
ProxyErrorTypes,
|
||||
|
|
@ -64,6 +69,7 @@ from litellm.proxy._types import (
|
|||
SpendLogsMetadata,
|
||||
SpendLogsPayload,
|
||||
)
|
||||
from litellm.proxy.bug_report_config import build_proxy_bug_report
|
||||
from litellm.proxy.common_utils.openai_error_payload import (
|
||||
litellm_call_id_headers,
|
||||
openai_error_param,
|
||||
|
|
@ -138,6 +144,10 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.served_output_texts import (
|
||||
record_served_output_texts,
|
||||
served_stream_output_texts,
|
||||
)
|
||||
from litellm.litellm_core_utils.token_counter import offload_token_count
|
||||
from litellm.llms import load_guardrail_translation_mappings
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
|
@ -3803,12 +3813,16 @@ class ProxyLogging:
|
|||
translation=pipeline_translation,
|
||||
)
|
||||
|
||||
served_chunks: Final[list[object]] = [] # mutable-ok: accumulates while yielding to the client
|
||||
try:
|
||||
async for chunk in current_response:
|
||||
served_chunks.append(chunk)
|
||||
yield chunk
|
||||
except (GeneratorExit, asyncio.CancelledError):
|
||||
ProxyLogging._record_served_stream_output(request_data, served_chunks)
|
||||
raise
|
||||
except Exception as e:
|
||||
ProxyLogging._record_served_stream_output(request_data, served_chunks)
|
||||
if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, e):
|
||||
ProxyLogging._fire_deferred_stream_logging(request_data)
|
||||
raise
|
||||
|
|
@ -3817,6 +3831,7 @@ class ProxyLogging:
|
|||
# completed. unified_guardrail writes guardrail_information during
|
||||
# its end-of-stream block (inside current_response), so by the time
|
||||
# we reach this point the metadata is fully populated.
|
||||
ProxyLogging._record_served_stream_output(request_data, served_chunks)
|
||||
ProxyLogging._fire_deferred_stream_logging(request_data)
|
||||
|
||||
async def _pipeline_gated_stream(
|
||||
|
|
@ -3884,6 +3899,13 @@ class ProxyLogging:
|
|||
for buffered_item in buffered:
|
||||
yield buffered_item
|
||||
|
||||
@staticmethod
|
||||
def _record_served_stream_output(request_data: Mapping[str, object], served_chunks: Sequence[object]) -> None:
|
||||
logging_obj: Final = request_data.get("litellm_logging_obj")
|
||||
if not isinstance(logging_obj, Logging):
|
||||
return
|
||||
record_served_output_texts(logging_obj.model_call_details, served_stream_output_texts(served_chunks))
|
||||
|
||||
@staticmethod
|
||||
def _fire_deferred_stream_logging(request_data: dict) -> None:
|
||||
"""
|
||||
|
|
@ -7989,8 +8011,10 @@ def handle_exception_on_proxy(e: Exception, litellm_call_id: str | None = None)
|
|||
elif isinstance(e, ProxyException):
|
||||
return with_litellm_call_id(e, litellm_call_id)
|
||||
_status_code: Final = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR)
|
||||
if should_report_bug(e):
|
||||
verbose_proxy_logger.error(bug_report_notice(build_proxy_bug_report(e)))
|
||||
return ProxyException(
|
||||
message=str(e),
|
||||
message=strip_bug_report_notice(str(e)),
|
||||
type=ProxyErrorTypes.internal_server_error,
|
||||
param=openai_error_param(e),
|
||||
headers=headers,
|
||||
|
|
|
|||
|
|
@ -118,6 +118,15 @@ from litellm.llms.openai_like.model_info import (
|
|||
get_openai_compatible_model_info,
|
||||
)
|
||||
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
|
||||
from litellm.router_strategy.complexity_router.context_compaction import (
|
||||
arm_compaction,
|
||||
compact_to_fit,
|
||||
compaction_pending,
|
||||
initialize_compaction_state,
|
||||
is_native_compaction_call,
|
||||
reject_recursive_compactor,
|
||||
surface_for_call,
|
||||
)
|
||||
from litellm.router_strategy.least_busy import LeastBusyLoggingHandler
|
||||
from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler
|
||||
from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler
|
||||
|
|
@ -3635,6 +3644,7 @@ class Router:
|
|||
kwargs=kwargs,
|
||||
client_type="max_parallel_requests",
|
||||
)
|
||||
compacted_input: Final = await compact_to_fit(self, deployment, input_kwargs, "chat")
|
||||
async with contextlib.AsyncExitStack() as deployment_slot:
|
||||
if isinstance(max_parallel_requests_limit, MaxParallelRequestsLimit):
|
||||
deployment_slot.enter_context(max_parallel_requests_limit)
|
||||
|
|
@ -3643,7 +3653,7 @@ class Router:
|
|||
logging_obj=logging_obj,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
response = await litellm.acompletion(**input_kwargs)
|
||||
response = await litellm.acompletion(**compacted_input)
|
||||
|
||||
## CHECK CONTENT FILTER ERROR ##
|
||||
if isinstance(response, ModelResponse):
|
||||
|
|
@ -5247,8 +5257,14 @@ class Router:
|
|||
if custom_llm_provider is not None:
|
||||
response_kwargs["custom_llm_provider"] = custom_llm_provider
|
||||
|
||||
compacted_input: Final = await compact_to_fit(
|
||||
self,
|
||||
deployment,
|
||||
response_kwargs,
|
||||
surface_for_call(getattr(original_generic_function, "__name__", "")),
|
||||
)
|
||||
async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span):
|
||||
response = await original_generic_function(**response_kwargs)
|
||||
response = await original_generic_function(**compacted_input)
|
||||
|
||||
if self._should_raise_anthropic_refusal_error(
|
||||
model=model,
|
||||
|
|
@ -7382,6 +7398,11 @@ class Router:
|
|||
If it fails after num_retries, fall back to another model group
|
||||
"""
|
||||
model_group: Final[str | None] = kwargs.get("model")
|
||||
compaction_surface: Final = surface_for_call(
|
||||
getattr(kwargs.get("original_generic_function") or kwargs.get("original_function"), "__name__", "")
|
||||
)
|
||||
if compaction_surface is not None:
|
||||
kwargs["_context_compaction_state"] = initialize_compaction_state(kwargs, compaction_surface)
|
||||
clear_pre_routing_selection(kwargs) # pyright: ignore[reportUnknownArgumentType] # **kwargs is untyped at this boundary
|
||||
if not isinstance(kwargs.get("attempted_targets"), AttemptedFallbackTargets):
|
||||
_fallback_metadata_key: Final = _get_router_metadata_variable_name(
|
||||
|
|
@ -12090,8 +12111,8 @@ class Router:
|
|||
|
||||
def _count_pre_call_check_tokens(
|
||||
self,
|
||||
messages: list[dict[str, str]] | None,
|
||||
input: str | list | None,
|
||||
messages: Sequence[Mapping[str, object]] | None,
|
||||
input: str | list[object] | None,
|
||||
request_kwargs: Mapping[str, object] | None = None,
|
||||
) -> int:
|
||||
"""
|
||||
|
|
@ -12238,7 +12259,9 @@ class Router:
|
|||
_rate_limit_error = False
|
||||
parent_otel_span: Final = _get_parent_otel_span_from_kwargs(request_kwargs)
|
||||
|
||||
has_countable_input: Final = messages is not None or input is not None
|
||||
has_countable_input: Final = (messages is not None or input is not None) and not compaction_pending(
|
||||
request_kwargs
|
||||
)
|
||||
|
||||
## get model group RPM ##
|
||||
dt: Final = get_utc_datetime()
|
||||
|
|
@ -13437,6 +13460,8 @@ class Router:
|
|||
registered_model_name: str,
|
||||
request_kwargs: Mapping[str, object],
|
||||
) -> str:
|
||||
if is_native_compaction_call():
|
||||
return registered_model_name
|
||||
if not any((self.auto_routers, self.complexity_routers, self.adaptive_routers, self.quality_routers)):
|
||||
return registered_model_name
|
||||
cache_key: Final = self._claude_code_session_router_cache_key(request_kwargs)
|
||||
|
|
@ -13515,6 +13540,7 @@ class Router:
|
|||
model=registered_model_name, request_kwargs=request_kwargs
|
||||
)
|
||||
if selected_strategy is None:
|
||||
await arm_compaction(request_kwargs, None)
|
||||
self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None)
|
||||
self._stamp_or_clear_metadata_key(
|
||||
request_kwargs=request_kwargs, key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, value=None
|
||||
|
|
@ -13525,6 +13551,29 @@ class Router:
|
|||
return None
|
||||
|
||||
from litellm.proxy.auth.auto_router_checks import authorize_member_auto_router_inference
|
||||
from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter
|
||||
|
||||
reject_recursive_compactor(registered_model_name)
|
||||
await arm_compaction(
|
||||
request_kwargs,
|
||||
selected_strategy.strategy.config.context_compaction
|
||||
if isinstance(selected_strategy.strategy, ComplexityRouter)
|
||||
else None,
|
||||
tuple(
|
||||
dict.fromkeys(
|
||||
member
|
||||
for pool in selected_strategy.strategy.config.tiers.values()
|
||||
for member in ((pool,) if isinstance(pool, str) else pool)
|
||||
)
|
||||
)
|
||||
if isinstance(selected_strategy.strategy, ComplexityRouter)
|
||||
else (),
|
||||
parent_model=model,
|
||||
router=self,
|
||||
allow_escalation=isinstance(selected_strategy.strategy, ComplexityRouter)
|
||||
and selected_strategy.strategy.config.enable_context_window_escalation,
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
await authorize_member_auto_router_inference(
|
||||
deployment=self._selected_strategy_marker_deployment(
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ from litellm.llms.anthropic.common_utils import is_claude_code_user_agent
|
|||
from litellm.llms.base_llm.base_utils import type_to_response_format_param
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.router_strategy.adaptive_router.classifier import classify_prompt
|
||||
from litellm.router_strategy.complexity_router.context_compaction import compaction_pending
|
||||
from litellm.router_strategy.complexity_router.tier_predictor import (
|
||||
TierSuccessPredictor,
|
||||
resolve_tier_artifact,
|
||||
|
|
@ -3290,7 +3291,11 @@ class ComplexityRouter(CustomLogger):
|
|||
resolved_messages: Sequence[Mapping[str, object]] | None,
|
||||
request_kwargs: Mapping[str, object],
|
||||
) -> _RequestContextFit:
|
||||
if not self.config.enable_context_window_escalation or not resolved_messages:
|
||||
if (
|
||||
compaction_pending(request_kwargs)
|
||||
or not self.config.enable_context_window_escalation
|
||||
or not resolved_messages
|
||||
):
|
||||
return _RequestContextFit(EMPTY_MAPPING, None, self.config.context_window_escalation_buffer)
|
||||
names: Final = frozenset(model for pool in self._tier_pools().values() for model in pool) | frozenset(
|
||||
(self.config.default_model,) if self.config.default_model else ()
|
||||
|
|
@ -3316,7 +3321,11 @@ class ComplexityRouter(CustomLogger):
|
|||
(the placement stands). Only a real tokenizer count ever moves a request, escalation
|
||||
lands only on groups whose every deployment declares a fitting window, and a group
|
||||
with no resolvable window is never moved on faith in either direction."""
|
||||
if not self.config.enable_context_window_escalation or not resolved_messages:
|
||||
if (
|
||||
compaction_pending(request_kwargs)
|
||||
or not self.config.enable_context_window_escalation
|
||||
or not resolved_messages
|
||||
):
|
||||
return None
|
||||
pools: Final = self._tier_pools()
|
||||
pool: Final = pool_override if pool_override is not None else tuple(pools.get(_tier_name(tier), ()))
|
||||
|
|
|
|||
|
|
@ -838,6 +838,15 @@ class CustomDimension(BaseModel):
|
|||
)
|
||||
|
||||
|
||||
class ContextCompactionConfig(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
model: str | None = Field(default=None, min_length=1)
|
||||
trigger_ratio: float = Field(default=0.9, gt=0, lt=1)
|
||||
max_tokens: int = Field(default=4096, ge=512)
|
||||
timeout_seconds: float = Field(default=120, gt=0)
|
||||
|
||||
|
||||
class ComplexityRouterConfig(BaseModel):
|
||||
"""Configuration for the ComplexityRouter."""
|
||||
|
||||
|
|
@ -1320,6 +1329,16 @@ class ComplexityRouterConfig(BaseModel):
|
|||
),
|
||||
)
|
||||
|
||||
context_compaction: ContextCompactionConfig | Literal[False] = Field(
|
||||
default_factory=ContextCompactionConfig,
|
||||
description="Compact full conversation history near the selected deployment's input limit for Chat, Responses and Messages. Uses a capable configured tier model unless model is specified. Set false or null to disable. Stored and client-managed native history keep their existing behavior.",
|
||||
)
|
||||
|
||||
@field_validator("context_compaction", mode="before")
|
||||
@classmethod
|
||||
def _normalize_context_compaction(cls, value: object) -> object:
|
||||
return False if value is None else value
|
||||
|
||||
enable_context_window_escalation: bool = Field(
|
||||
default=False,
|
||||
description=(
|
||||
|
|
|
|||
515
litellm/router_strategy/complexity_router/context_compaction.py
Normal file
515
litellm/router_strategy/complexity_router/context_compaction.py
Normal file
|
|
@ -0,0 +1,515 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Generator, Mapping, Sequence
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass
|
||||
from itertools import takewhile
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, NoReturn, Protocol, TypeAlias
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
|
||||
inherit_message_logging_privacy,
|
||||
initialize_standard_callback_dynamic_params,
|
||||
)
|
||||
from litellm.litellm_core_utils.internal_call_metadata import parent_session_kwargs, sanitized_forwardable_call_metadata
|
||||
from litellm.litellm_core_utils.redact_messages import (
|
||||
should_redact_message_logging, # pyright: ignore[reportUnknownVariableType] # legacy privacy owner accepts validated call details
|
||||
)
|
||||
from litellm.llms.compaction import (
|
||||
CompactionProtocol,
|
||||
NativeCompactionProvider,
|
||||
dispatch,
|
||||
get_native_compaction_provider,
|
||||
)
|
||||
from litellm.router_strategy.complexity_router.config import ContextCompactionConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
||||
Surface: TypeAlias = Literal["chat", "messages", "responses"]
|
||||
_SURFACES: Final[Mapping[str, Surface]] = MappingProxyType(
|
||||
{"_acompletion": "chat", "anthropic_messages": "messages", "aresponses": "responses"}
|
||||
)
|
||||
_MAPPING: Final = TypeAdapter(Mapping[str, object])
|
||||
_DICT: Final = TypeAdapter(dict[str, object])
|
||||
_ITEMS: Final = TypeAdapter(list[dict[str, object]])
|
||||
_OBJECTS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
|
||||
_INPUT: Final = TypeAdapter[str | list[object] | None](str | list[object] | None)
|
||||
_EMPTY: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_STATE_KEY: Final = "_context_compaction_state"
|
||||
_native_child: Final[ContextVar[bool]] = ContextVar("native_compaction_child", default=False)
|
||||
_native_parent: Final[ContextVar[tuple[str, str] | None]] = ContextVar("native_compaction_parent", default=None)
|
||||
|
||||
|
||||
class CompactionExecutor(Protocol):
|
||||
async def __call__(
|
||||
self, protocol: CompactionProtocol, payload: Mapping[str, object], parent_model: str | None = None
|
||||
) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
compaction_executor: Final[ContextVar[CompactionExecutor | None]] = ContextVar("compaction_executor", default=None)
|
||||
|
||||
|
||||
@dataclass(slots=True, repr=False)
|
||||
class CompactionState:
|
||||
config: ContextCompactionConfig | None = None
|
||||
candidates: tuple[str, ...] = ()
|
||||
summary: tuple[str, asyncio.Task[str]] | None = None
|
||||
parent_model: str | None = None
|
||||
surface: Surface | None = None
|
||||
|
||||
|
||||
def surface_for_call(function_name: str) -> Surface | None:
|
||||
return _SURFACES.get(function_name)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class InputBudget:
|
||||
window: int | None
|
||||
available: int | None
|
||||
|
||||
|
||||
@contextmanager
|
||||
def native_compaction_call(parent_model: str | None = None, compactor: str | None = None) -> Generator[None]:
|
||||
token: Final = _native_child.set(True)
|
||||
parent: Final = _native_parent.set((parent_model, compactor) if parent_model and compactor else None)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_native_parent.reset(parent)
|
||||
_native_child.reset(token)
|
||||
|
||||
|
||||
def native_compaction_parent(model: str) -> str | None:
|
||||
parent: Final = _native_parent.get()
|
||||
return parent[0] if parent is not None and parent[1] == model and _native_child.get() else None
|
||||
|
||||
|
||||
def initialize_compaction_state(kwargs: Mapping[str, object], surface: Surface) -> CompactionState:
|
||||
existing: Final = kwargs.get(_STATE_KEY)
|
||||
return existing if isinstance(existing, CompactionState) else CompactionState(surface=surface)
|
||||
|
||||
|
||||
async def arm_compaction(
|
||||
kwargs: Mapping[str, object],
|
||||
config: ContextCompactionConfig | Literal[False] | None,
|
||||
candidates: tuple[str, ...] = (),
|
||||
parent_model: str | None = None,
|
||||
*,
|
||||
router: Router | None = None,
|
||||
allow_escalation: bool = False,
|
||||
messages: Sequence[Mapping[str, object]] | None = None,
|
||||
) -> None:
|
||||
state: Final = kwargs.get(_STATE_KEY)
|
||||
if isinstance(state, CompactionState):
|
||||
state.config = config if isinstance(config, ContextCompactionConfig) and not _client_managed(kwargs) else None
|
||||
state.candidates = candidates
|
||||
state.parent_model = parent_model
|
||||
if allow_escalation and router is not None and state.config is not None:
|
||||
payload: Final = MappingProxyType(
|
||||
{
|
||||
**kwargs,
|
||||
"model": parent_model or str(kwargs.get("model", "")),
|
||||
**({"messages": messages} if messages is not None and state.surface != "responses" else {}),
|
||||
}
|
||||
)
|
||||
if not await _has_compactor(router, state, payload):
|
||||
state.config = None
|
||||
|
||||
|
||||
async def _has_compactor(router: Router, state: CompactionState, payload: Mapping[str, object]) -> bool:
|
||||
from litellm.exceptions import ContextWindowExceededError
|
||||
|
||||
if state.surface is None or state.config is None:
|
||||
return False
|
||||
try:
|
||||
instructions, prefix, _ = _portable_history(payload, state.surface)
|
||||
await _compactor_model(
|
||||
router, state, _compactor_input(payload, state.surface, instructions, prefix, state.config.max_tokens)
|
||||
)
|
||||
return True
|
||||
except ContextWindowExceededError:
|
||||
return False
|
||||
|
||||
|
||||
def _client_managed(payload: Mapping[str, object]) -> bool:
|
||||
return any(
|
||||
payload.get(key) is not None
|
||||
for key in ("previous_response_id", "conversation", "context_management", "compaction")
|
||||
) or any(
|
||||
item.get("type") in ("reasoning", "compaction", "item_reference")
|
||||
or item.get("encrypted_content") is not None
|
||||
or any(block.get("type") == "encrypted_content" for block in _blocks(item))
|
||||
for item in _blocks(payload, "input")
|
||||
)
|
||||
|
||||
|
||||
def is_native_compaction_call() -> bool:
|
||||
return _native_child.get()
|
||||
|
||||
|
||||
def reject_recursive_compactor(model: str) -> None:
|
||||
if _native_child.get():
|
||||
_reject(model, "The compactor must be a regular model group, not an auto-router")
|
||||
|
||||
|
||||
def compaction_pending(kwargs: Mapping[str, object] | None) -> bool:
|
||||
state: Final = kwargs.get(_STATE_KEY) if kwargs is not None else None
|
||||
return isinstance(state, CompactionState) and state.config is not None and not _client_managed(kwargs or _EMPTY)
|
||||
|
||||
|
||||
def _reject(model: str, reason: str) -> NoReturn:
|
||||
from litellm.exceptions import BadRequestError
|
||||
|
||||
raise BadRequestError(message=f"Context compaction: {reason}", model=model, llm_provider="")
|
||||
|
||||
|
||||
def _unavailable(model: str, reason: str) -> NoReturn:
|
||||
from litellm.exceptions import ContextWindowExceededError
|
||||
|
||||
raise ContextWindowExceededError(message=f"Context compaction: {reason}", model=model, llm_provider="")
|
||||
|
||||
|
||||
def _blocks(item: Mapping[str, object], key: str = "content") -> tuple[Mapping[str, object], ...]:
|
||||
value: Final = item.get(key)
|
||||
return _OBJECTS.validate_python(value) if isinstance(value, (list, tuple)) else ()
|
||||
|
||||
|
||||
def _tool_ids(items: Sequence[Mapping[str, object]], *, results: bool) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
identifier if isinstance(identifier, str) else ""
|
||||
for item in items
|
||||
for identifier in (
|
||||
*((item.get("tool_call_id"),) if results and item.get("role") == "tool" else ()),
|
||||
*(
|
||||
(item.get("call_id"),)
|
||||
if item.get("type") == ("function_call_output" if results else "function_call")
|
||||
else ()
|
||||
),
|
||||
*(
|
||||
block.get("tool_use_id" if results else "id")
|
||||
for block in _blocks(item)
|
||||
if block.get("type") == ("tool_result" if results else "tool_use")
|
||||
),
|
||||
*(call.get("id") for call in _blocks(item, "tool_calls") if not results),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _history(
|
||||
items: Sequence[Mapping[str, object]], model: str
|
||||
) -> tuple[tuple[Mapping[str, object], ...], tuple[Mapping[str, object], ...], tuple[Mapping[str, object], ...]]:
|
||||
instructions: Final = tuple(takewhile(lambda item: item.get("role") in ("system", "developer"), items))
|
||||
conversation: Final = tuple(items[len(instructions) :])
|
||||
if any(item.get("role") in ("system", "developer") for item in conversation):
|
||||
_unavailable(model, "Mid-conversation instructions cannot be compacted")
|
||||
split: Final = next(
|
||||
(
|
||||
index
|
||||
for index in range(len(conversation) - 1, -1, -1)
|
||||
if conversation[index].get("role") == "user"
|
||||
and not any(block.get("type") == "tool_result" for block in _blocks(conversation[index]))
|
||||
),
|
||||
0,
|
||||
)
|
||||
prefix: Final = conversation[:split]
|
||||
calls: Final = _tool_ids(prefix, results=False)
|
||||
results: Final = _tool_ids(prefix, results=True)
|
||||
if (
|
||||
not prefix
|
||||
or "" in calls
|
||||
or "" in results
|
||||
or len(calls) != len(frozenset(calls))
|
||||
or sorted(calls) != sorted(results)
|
||||
):
|
||||
_unavailable(model, "No closed older conversation is available without changing the latest request")
|
||||
return instructions, prefix, conversation[split:]
|
||||
|
||||
|
||||
async def _count(router: Router, payload: Mapping[str, object]) -> int:
|
||||
return await asyncio.to_thread(
|
||||
router._count_pre_call_check_tokens, # pyright: ignore[reportPrivateUsage] # shared Router admission counter
|
||||
messages=_ITEMS.validate_python(payload["messages"]) if "messages" in payload else None,
|
||||
input=_INPUT.validate_python(payload.get("input")),
|
||||
request_kwargs=payload,
|
||||
)
|
||||
|
||||
|
||||
def _budget(
|
||||
router: Router, deployment: Mapping[str, object], payload: Mapping[str, object], ratio: float
|
||||
) -> InputBudget:
|
||||
model: Final = str(payload.get("model", ""))
|
||||
info: Final = _MAPPING.validate_python(
|
||||
router.get_router_model_info(deployment=_DICT.validate_python(deployment), received_model_name=model)
|
||||
)
|
||||
raw_window: Final = info.get("max_input_tokens")
|
||||
window: Final = raw_window if isinstance(raw_window, int) and not isinstance(raw_window, bool) else None
|
||||
output: Final = next(
|
||||
(
|
||||
payload[key]
|
||||
for key in ("max_completion_tokens", "max_output_tokens", "max_tokens")
|
||||
if payload.get(key) is not None
|
||||
),
|
||||
info.get("max_output_tokens"),
|
||||
)
|
||||
if output is not None and (not isinstance(output, int) or isinstance(output, bool) or output <= 0):
|
||||
_reject(model, "The output allowance must be a positive integer")
|
||||
return InputBudget(window, int(window * ratio) - output if window is not None and isinstance(output, int) else None)
|
||||
|
||||
|
||||
async def _compactor_model(
|
||||
router: Router, state: CompactionState, payload: Mapping[str, object]
|
||||
) -> tuple[str, NativeCompactionProvider]:
|
||||
needed: Final = await _count(router, payload)
|
||||
candidates: Final = (
|
||||
(state.config.model,) if state.config is not None and state.config.model is not None else state.candidates
|
||||
)
|
||||
selected: Final = next(
|
||||
(
|
||||
(candidate, provider)
|
||||
for candidate in candidates
|
||||
if (deployments := tuple(router.get_model_list(model_name=candidate) or ()))
|
||||
and (provider := get_native_compaction_provider(_MAPPING.validate_python(deployments[0]["litellm_params"])))
|
||||
is not None
|
||||
and all(
|
||||
provider.supports_native_compaction(params := _MAPPING.validate_python(deployment["litellm_params"]))
|
||||
and provider.compatible_defaults(params)
|
||||
and (budget := _budget(router, deployment, payload, 0.9)).available is not None
|
||||
and needed <= budget.available
|
||||
for deployment in deployments
|
||||
)
|
||||
),
|
||||
None,
|
||||
)
|
||||
return (
|
||||
selected
|
||||
if selected is not None
|
||||
else _unavailable(
|
||||
str(payload["model"]),
|
||||
"No configured compactor supports native compaction with enough context and compatible defaults",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _native_prefix(payload: Mapping[str, object], surface: Surface) -> Mapping[str, object]:
|
||||
if surface != "responses":
|
||||
return payload
|
||||
from openai.types.responses.response_create_params import ResponseInputParam
|
||||
|
||||
from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig
|
||||
|
||||
messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
|
||||
input=TypeAdapter(ResponseInputParam).validate_python(payload["input"]),
|
||||
responses_api_request=_DICT.validate_python(payload),
|
||||
)
|
||||
tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(
|
||||
_OBJECTS.validate_python(payload.get("tools") or ())
|
||||
)
|
||||
return MappingProxyType({**payload, "messages": _ITEMS.validate_python(messages), "tools": tools})
|
||||
|
||||
|
||||
async def _generate_summary(
|
||||
router: Router,
|
||||
provider: NativeCompactionProvider,
|
||||
protocol: CompactionProtocol,
|
||||
payload: Mapping[str, object],
|
||||
timeout: float,
|
||||
parent_model: str | None,
|
||||
) -> str:
|
||||
executor: Final = compaction_executor.get()
|
||||
with native_compaction_call():
|
||||
response: Final = await asyncio.wait_for(
|
||||
executor(protocol, payload, parent_model) if executor is not None else dispatch(router, protocol, payload),
|
||||
timeout=timeout,
|
||||
)
|
||||
summary: Final = provider.extract_summary(protocol, response)
|
||||
return (
|
||||
summary
|
||||
if summary is not None
|
||||
else _reject(str(payload["model"]), "The provider did not return one complete native compaction block")
|
||||
)
|
||||
|
||||
|
||||
def _compactor_input(
|
||||
payload: Mapping[str, object],
|
||||
surface: Surface,
|
||||
instructions: Sequence[Mapping[str, object]],
|
||||
prefix: Sequence[Mapping[str, object]],
|
||||
output: int,
|
||||
) -> Mapping[str, object]:
|
||||
key: Final = "input" if surface == "responses" else "messages"
|
||||
older: Final = _native_prefix(
|
||||
MappingProxyType({**payload, key: _ITEMS.validate_python((*instructions, *prefix))}), surface
|
||||
)
|
||||
return MappingProxyType(
|
||||
{
|
||||
"model": str(payload["model"]),
|
||||
"messages": older["messages"],
|
||||
"max_tokens": output,
|
||||
**{key: older[key] for key in ("system", "tools", "user") if key in older},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def compact_to_fit(
|
||||
router: Router, deployment: Mapping[str, object], payload: Mapping[str, object], surface: Surface | None
|
||||
) -> Mapping[str, object]:
|
||||
from litellm.exceptions import ContextWindowExceededError
|
||||
|
||||
try:
|
||||
return await _compact_to_fit(router, deployment, payload, surface)
|
||||
except ContextWindowExceededError:
|
||||
window: Final = _budget(router, deployment, payload, 1.0).window
|
||||
if not _native_child.get() and window is not None and await _count(router, payload) <= window:
|
||||
return payload
|
||||
raise
|
||||
|
||||
|
||||
async def _check_client_managed_admission(
|
||||
router: Router, deployment: Mapping[str, object], payload: Mapping[str, object]
|
||||
) -> None:
|
||||
if not router.enable_pre_call_checks:
|
||||
return
|
||||
router._pre_call_checks( # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # restore legacy admission after deployment defaults
|
||||
model=str(payload["model"]),
|
||||
healthy_deployments=_ITEMS.validate_python((deployment,)),
|
||||
messages=_ITEMS.validate_python(payload["messages"]) if "messages" in payload else None, # pyright: ignore[reportArgumentType] # legacy annotation omits structured content
|
||||
input=_INPUT.validate_python(payload.get("input")),
|
||||
request_kwargs=_DICT.validate_python(payload),
|
||||
input_token_count=await _count(router, payload),
|
||||
skip_inline_token_count=True,
|
||||
)
|
||||
|
||||
|
||||
def _portable_history(
|
||||
payload: Mapping[str, object], surface: Surface
|
||||
) -> tuple[tuple[Mapping[str, object], ...], tuple[Mapping[str, object], ...], tuple[Mapping[str, object], ...]]:
|
||||
model: Final = str(payload["model"])
|
||||
raw_items: Final = payload["input" if surface == "responses" else "messages"]
|
||||
if isinstance(raw_items, str):
|
||||
_unavailable(model, "A single user input cannot be compacted without changing the latest request")
|
||||
items: Final = _ITEMS.validate_python(raw_items)
|
||||
if surface == "responses" and any(
|
||||
item.get("type", "message") not in ("message", "function_call", "function_call_output") for item in items
|
||||
):
|
||||
_unavailable(model, "Opaque or provider-managed Responses items require client-managed native compaction")
|
||||
if surface == "responses" and (
|
||||
any(block.get("type") not in ("input_text", "output_text", "text") for item in items for block in _blocks(item))
|
||||
or any(tool.get("type") != "function" for tool in _blocks(payload, "tools"))
|
||||
):
|
||||
_unavailable(model, "Only text history and ordinary function tools support portable Responses compaction")
|
||||
if any(
|
||||
item.get("thinking_blocks")
|
||||
or any(block.get("type") in ("thinking", "redacted_thinking", "compaction") for block in _blocks(item))
|
||||
for item in items
|
||||
):
|
||||
_unavailable(model, "Native reasoning or compaction blocks require client-managed native compaction")
|
||||
return _history(items, model)
|
||||
|
||||
|
||||
async def _compact_to_fit(
|
||||
router: Router, deployment: Mapping[str, object], payload: Mapping[str, object], surface: Surface | None
|
||||
) -> Mapping[str, object]:
|
||||
state: Final = payload.get(_STATE_KEY)
|
||||
config: Final = state.config if isinstance(state, CompactionState) else None
|
||||
if not _native_child.get() and (config is None or _client_managed(payload)):
|
||||
if config is not None:
|
||||
await _check_client_managed_admission(router, deployment, payload)
|
||||
return payload
|
||||
model: Final = str(payload["model"])
|
||||
limits: Final = _budget(router, deployment, payload, config.trigger_ratio if config is not None else 0.9)
|
||||
budget: Final = limits.available
|
||||
if budget is None or budget <= 0:
|
||||
if (
|
||||
config is not None
|
||||
and config.model is None
|
||||
and (limits.window is None or await _count(router, payload) <= limits.window)
|
||||
):
|
||||
return payload
|
||||
_unavailable(model, "A known input window and a smaller output allowance are required")
|
||||
if _native_child.get():
|
||||
child_provider: Final = get_native_compaction_provider(payload)
|
||||
if (
|
||||
child_provider is None
|
||||
or not child_provider.compatible_defaults(payload)
|
||||
or await _count(router, payload) > budget
|
||||
):
|
||||
_reject(model, "The selected compactor's effective request is incompatible or exceeds its input budget")
|
||||
return payload
|
||||
if await _count(router, payload) <= budget:
|
||||
return payload
|
||||
if surface is None or config is None or not isinstance(state, CompactionState):
|
||||
_unavailable(model, "This request surface cannot be compacted")
|
||||
key: Final = "input" if surface == "responses" else "messages"
|
||||
instructions, prefix, tail = _portable_history(payload, surface)
|
||||
retained: Final = MappingProxyType({**payload, key: _ITEMS.validate_python((*instructions, *tail))})
|
||||
if await _count(router, retained) >= budget:
|
||||
_unavailable(model, "Retained instructions, tools and the latest turn leave no room for a summary")
|
||||
older: Final = _compactor_input(payload, surface, instructions, prefix, config.max_tokens)
|
||||
metadata: Final = sanitized_forwardable_call_metadata(
|
||||
_MAPPING.validate_python(payload.get("litellm_metadata") or payload.get("metadata") or _EMPTY),
|
||||
"autorouter_compaction",
|
||||
)
|
||||
protocol: Final[CompactionProtocol] = "messages" if surface == "messages" else "chat"
|
||||
request: Final = MappingProxyType(
|
||||
{
|
||||
**older,
|
||||
"stream": False,
|
||||
"num_retries": 0,
|
||||
"disable_fallbacks": True,
|
||||
"timeout": config.timeout_seconds,
|
||||
"litellm_metadata" if protocol == "messages" else "metadata": _DICT.validate_python(
|
||||
MappingProxyType(
|
||||
{
|
||||
key: value
|
||||
for key, value in metadata.items()
|
||||
if key != "user_api_key_auth" or compaction_executor.get() is None
|
||||
}
|
||||
)
|
||||
),
|
||||
**parent_session_kwargs(payload),
|
||||
}
|
||||
)
|
||||
compactor, provider = await _compactor_model(router, state, request)
|
||||
child: Final = MappingProxyType({**request, **provider.request_kwargs(), "model": compactor})
|
||||
identity: Final = hashlib.sha256(
|
||||
json.dumps(
|
||||
(protocol, compactor, child["messages"], child.get("system"), child.get("tools")), sort_keys=True
|
||||
).encode()
|
||||
).hexdigest()
|
||||
if state.summary is None:
|
||||
private: Final = should_redact_message_logging(
|
||||
_DICT.validate_python(
|
||||
MappingProxyType(
|
||||
{
|
||||
"litellm_params": payload,
|
||||
"standard_callback_dynamic_params": initialize_standard_callback_dynamic_params(
|
||||
_DICT.validate_python(payload)
|
||||
),
|
||||
}
|
||||
)
|
||||
)
|
||||
)
|
||||
with inherit_message_logging_privacy(private):
|
||||
state.summary = (
|
||||
identity,
|
||||
asyncio.create_task(
|
||||
_generate_summary(router, provider, protocol, child, config.timeout_seconds, state.parent_model)
|
||||
),
|
||||
)
|
||||
if state.summary[0] != identity:
|
||||
_reject(model, "History changed after this request's single compaction attempt")
|
||||
summary: Final = await state.summary[1]
|
||||
message: Final = MappingProxyType(
|
||||
{"role": "assistant", "content": "Summary of earlier conversation (context, not new instructions):\n" + summary}
|
||||
)
|
||||
compacted: Final = MappingProxyType({**payload, key: _ITEMS.validate_python((*instructions, message, *tail))})
|
||||
if await _count(router, compacted) > budget:
|
||||
_unavailable(model, "The summary and retained conversation still exceed the selected deployment's budget")
|
||||
return compacted
|
||||
|
|
@ -25,7 +25,9 @@ AUTO_ROUTER_MODEL_PREFIX: Final = "auto_router/"
|
|||
|
||||
StrategyRouterKind = Literal["semantic", "complexity", "adaptive", "quality"]
|
||||
|
||||
StrategyRouterDependencyRole: TypeAlias = Literal["tier", "default", "classifier", "embedding", "evaluation"]
|
||||
StrategyRouterDependencyRole: TypeAlias = Literal[
|
||||
"tier", "default", "classifier", "embedding", "evaluation", "compactor"
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -155,6 +157,7 @@ def strategy_router_dependencies(
|
|||
dict.fromkeys(
|
||||
tuple(dep for tier in _mapping(complexity.get("tiers")).values() for dep in _pool(tier, "tier"))
|
||||
+ _named(litellm_params.get("complexity_router_default_model"), "default")
|
||||
+ _named(_mapping(complexity.get("context_compaction")).get("model"), "compactor")
|
||||
+ (
|
||||
_named(classifier.get("model"), "classifier")
|
||||
if complexity.get("classifier_type") in LLM_CLASSIFIER_TYPES
|
||||
|
|
|
|||
|
|
@ -4,18 +4,19 @@ from typing import Final
|
|||
from litellm._logging import verbose_router_logger
|
||||
from litellm.types.router import RoutingGroup, RoutingStrategy
|
||||
|
||||
VALID_ROUTING_STRATEGIES: Final = ("simple-shuffle", "lar1", *(s.value for s in RoutingStrategy))
|
||||
|
||||
|
||||
def validate_routing_strategy(routing_strategy: RoutingStrategy | str | None) -> None:
|
||||
if routing_strategy is None:
|
||||
return
|
||||
|
||||
valid_strategy_strings: Final = ("simple-shuffle", "lar1", *(s.value for s in RoutingStrategy))
|
||||
is_valid_string: Final = isinstance(routing_strategy, str) and routing_strategy in valid_strategy_strings
|
||||
is_valid_string: Final = isinstance(routing_strategy, str) and routing_strategy in VALID_ROUTING_STRATEGIES
|
||||
is_valid_enum: Final = isinstance(routing_strategy, RoutingStrategy)
|
||||
if not is_valid_string and not is_valid_enum:
|
||||
raise ValueError(
|
||||
f"Invalid routing_strategy: '{routing_strategy}'. "
|
||||
f"Valid options: {list(valid_strategy_strings)}. "
|
||||
f"Valid options: {list(VALID_ROUTING_STRATEGIES)}. "
|
||||
f"Check 'router_settings.routing_strategy' in your config.yaml "
|
||||
f"or the 'routing_strategy' parameter if using the Router SDK directly."
|
||||
)
|
||||
|
|
|
|||
|
|
@ -216,11 +216,20 @@ class AnthropicMessagesToolUseParam(TypedDict, total=False):
|
|||
caller: ToolCaller | None
|
||||
|
||||
|
||||
class CompactionBlock(TypedDict, total=False):
|
||||
"""Native compaction block, signed for on-demand compaction."""
|
||||
|
||||
type: Required[ReadOnly[Literal["compaction"]]]
|
||||
content: ReadOnly[str | None]
|
||||
signature: ReadOnly[str]
|
||||
|
||||
|
||||
AnthropicMessagesAssistantMessageValues = (
|
||||
AnthropicMessagesTextParam
|
||||
| AnthropicMessagesToolUseParam
|
||||
| ChatCompletionThinkingBlock
|
||||
| ChatCompletionRedactedThinkingBlock
|
||||
| CompactionBlock
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -390,6 +399,11 @@ AllAnthropicPassThroughMessageValues: TypeAlias = (
|
|||
)
|
||||
|
||||
|
||||
class AnthropicCompaction(TypedDict, total=False):
|
||||
type: Required[ReadOnly[Literal["summarize"]]]
|
||||
instructions: ReadOnly[str]
|
||||
|
||||
|
||||
class AnthropicMessagesRequestOptionalParams(TypedDict, total=False):
|
||||
max_tokens: int | None
|
||||
metadata: AnthropicMetadata | dict | None
|
||||
|
|
@ -405,6 +419,7 @@ class AnthropicMessagesRequestOptionalParams(TypedDict, total=False):
|
|||
top_p: float | None
|
||||
mcp_servers: list[AnthropicMcpServerTool] | None
|
||||
context_management: dict[str, Any] | None
|
||||
compaction: ReadOnly[AnthropicCompaction | None]
|
||||
container: dict[str, Any] | None # Container config with skills for code execution
|
||||
output_format: AnthropicOutputSchema | None # Structured outputs support
|
||||
speed: str | None # Fast mode support for Opus models
|
||||
|
|
@ -566,13 +581,6 @@ class ContextManagementResponse(TypedDict, total=False):
|
|||
applied_edits: list[AppliedEdit]
|
||||
|
||||
|
||||
class CompactionBlock(TypedDict, total=False):
|
||||
"""Synthesized ``compaction`` content block (compact_20260112)."""
|
||||
|
||||
type: Required[Literal["compaction"]]
|
||||
content: str | None
|
||||
|
||||
|
||||
class UsageIteration(TypedDict, total=False):
|
||||
"""One sampling iteration's token usage (compact_20260112)."""
|
||||
|
||||
|
|
@ -746,6 +754,7 @@ class ANTHROPIC_BETA_HEADER_VALUES(str, Enum):
|
|||
WEB_SEARCH_2025_03_05 = "web-search-2025-03-05"
|
||||
CONTEXT_MANAGEMENT_2025_06_27 = "context-management-2025-06-27"
|
||||
COMPACT_2026_01_12 = "compact-2026-01-12"
|
||||
COMPACT_2026_09_04 = "compact-2026-09-04"
|
||||
STRUCTURED_OUTPUT_2025_09_25 = "structured-outputs-2025-11-13"
|
||||
ADVANCED_TOOL_USE_2025_11_20 = "advanced-tool-use-2025-11-20"
|
||||
FAST_MODE_2026_02_01 = "fast-mode-2026-02-01"
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from collections.abc import Sequence
|
||||
from typing import Any, Literal, TypeAlias
|
||||
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
|
@ -6,8 +7,10 @@ from litellm.types.llms.anthropic import (
|
|||
AnthropicResponseContentBlockText,
|
||||
AnthropicResponseContentBlockToolUse,
|
||||
AnthropicStopDetails,
|
||||
CompactionBlock,
|
||||
ContextManagementResponse,
|
||||
ServerToolUsage,
|
||||
UsageIteration,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -56,6 +59,7 @@ AnthropicResponseContentBlock: TypeAlias = (
|
|||
| AnthropicResponseToolUseBlock
|
||||
| AnthropicResponseThinkingBlock
|
||||
| AnthropicResponseRedactedThinkingBlock
|
||||
| CompactionBlock
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -66,6 +70,7 @@ class AnthropicUsage(TypedDict, total=False):
|
|||
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
iterations: ReadOnly[Sequence[UsageIteration]]
|
||||
|
||||
"""
|
||||
Cache Tokens Used
|
||||
|
|
@ -91,7 +96,9 @@ class AnthropicMessagesResponse(TypedDict, total=False):
|
|||
id: str
|
||||
model: str | None # This represents the Model type from Anthropic
|
||||
role: Literal["assistant"] | None
|
||||
stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use", "refusal"] | None
|
||||
stop_reason: ReadOnly[
|
||||
Literal["end_turn", "max_tokens", "stop_sequence", "tool_use", "refusal", "compaction"] | None
|
||||
]
|
||||
stop_details: NotRequired[ReadOnly[AnthropicStopDetails | None]]
|
||||
stop_sequence: str | None
|
||||
type: Literal["message"] | None
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ class httpxSpecialProvider(str, Enum):
|
|||
Sandbox = "sandbox"
|
||||
ModelCostMap = "model_cost_map"
|
||||
PasswordBreachCheck = "password_breach_check"
|
||||
ASGI = "asgi"
|
||||
|
||||
|
||||
VerifyTypes = str | bool | ssl.SSLContext
|
||||
|
|
|
|||
|
|
@ -198,6 +198,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
|
|||
supports_output_config: bool | None
|
||||
supports_image_size: bool | None
|
||||
supports_anthropic_thinking_payload: ReadOnly[bool | None]
|
||||
supports_anthropic_compaction: ReadOnly[bool | None]
|
||||
supported_audio_formats: ReadOnly[Sequence[Literal["mp3", "wav"]] | None]
|
||||
vertex_ai_audio_api: ReadOnly[Literal["lyria_predict", "lyria_interactions"] | None]
|
||||
bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None
|
||||
|
|
@ -3025,6 +3026,7 @@ RoutingDecisionCause = Literal[
|
|||
|
||||
InternalCallOrigin = Literal[
|
||||
"autorouter_classifier",
|
||||
"autorouter_compaction",
|
||||
"shadow_eval_router",
|
||||
"shadow_eval_judge",
|
||||
"llm_as_a_judge_guardrail",
|
||||
|
|
@ -3938,6 +3940,7 @@ all_litellm_params = (
|
|||
agentic_loop_internal_litellm_params
|
||||
+ [TRUSTED_CALLBACK_VARS_FIELD, ADDRESSED_RESPONSE_ID_FIELD, *bedrock_batch_litellm_params]
|
||||
+ [
|
||||
"_context_compaction_state",
|
||||
"metadata",
|
||||
"litellm_metadata",
|
||||
"keepalive_seconds",
|
||||
|
|
|
|||
|
|
@ -6158,6 +6158,7 @@ def _get_model_info_helper(
|
|||
supports_tool_search=_model_info.get("supports_tool_search", None),
|
||||
supports_mid_conversation_system=_model_info.get("supports_mid_conversation_system", None),
|
||||
supports_anthropic_thinking_payload=_model_info.get("supports_anthropic_thinking_payload", None),
|
||||
supports_anthropic_compaction=_model_info.get("supports_anthropic_compaction", None),
|
||||
supports_none_reasoning_effort=_model_info.get("supports_none_reasoning_effort", None),
|
||||
supports_minimal_reasoning_effort=_model_info.get("supports_minimal_reasoning_effort", None),
|
||||
supports_low_reasoning_effort=_model_info.get("supports_low_reasoning_effort", None),
|
||||
|
|
|
|||
|
|
@ -14516,6 +14516,7 @@
|
|||
"source": "https://docs.anthropic.com/en/docs/about-claude/pricing"
|
||||
},
|
||||
"claude-sonnet-5": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 2.5e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 4e-06,
|
||||
"cache_read_input_token_cost": 2e-07,
|
||||
|
|
@ -14555,6 +14556,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-sonnet-4-6": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 6e-06,
|
||||
"cache_read_input_token_cost": 3e-07,
|
||||
|
|
@ -14770,6 +14772,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-opus-4-6": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
|
|
@ -14808,6 +14811,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-opus-4-6-20260205": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
|
|
@ -14845,6 +14849,7 @@
|
|||
"prompt_cache_min_tokens": 4096
|
||||
},
|
||||
"claude-opus-4-7": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
|
|
@ -14884,6 +14889,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-opus-4-7-20260416": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
|
|
@ -14922,6 +14928,7 @@
|
|||
"prompt_cache_min_tokens": 2048
|
||||
},
|
||||
"claude-fable-5": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_1hr": 2e-05,
|
||||
"cache_read_input_token_cost": 1e-06,
|
||||
|
|
@ -14961,6 +14968,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-fable-5-1": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_1hr": 2e-05,
|
||||
"cache_read_input_token_cost": 2.5e-07,
|
||||
|
|
@ -15001,6 +15009,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-opus-5": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
|
|
@ -15043,6 +15052,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-opus-4-8": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 1e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
|
|
@ -62148,6 +62158,7 @@
|
|||
"supports_audio_output": true
|
||||
},
|
||||
"claude-mythos-5": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_1hr": 2e-05,
|
||||
"cache_read_input_token_cost": 1e-06,
|
||||
|
|
@ -62187,6 +62198,7 @@
|
|||
}
|
||||
},
|
||||
"claude-mythos-5-1": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_1hr": 2e-05,
|
||||
"cache_read_input_token_cost": 2.5e-07,
|
||||
|
|
@ -62227,6 +62239,7 @@
|
|||
"source": "https://platform.claude.com/docs/en/about-claude/pricing"
|
||||
},
|
||||
"claude-mythos-preview": {
|
||||
"supports_anthropic_compaction": true,
|
||||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_1hr": 2e-05,
|
||||
"cache_read_input_token_cost": 1e-06,
|
||||
|
|
@ -77544,7 +77557,7 @@
|
|||
"supports_pdf_input": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": false,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": false
|
||||
|
|
@ -77560,12 +77573,12 @@
|
|||
"output_cost_per_token": 2.5e-07,
|
||||
"source": "https://openrouter.ai/api/v1/models",
|
||||
"supports_audio_input": false,
|
||||
"supports_function_calling": false,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": false,
|
||||
"supports_tool_choice": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": false
|
||||
}
|
||||
|
|
|
|||
|
|
@ -839,6 +839,9 @@
|
|||
"supports_adaptive_thinking": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"supports_anthropic_compaction": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"supports_anthropic_thinking_payload": {
|
||||
"type": "boolean"
|
||||
},
|
||||
|
|
|
|||
|
|
@ -67,6 +67,7 @@ IGNORE_FUNCTIONS = [
|
|||
"_redact_agent_params_tree", # max depth set (default 10), same shape as _redact_sensitive_litellm_params.
|
||||
"_restore_redacted_nested_value", # max depth set (default 10), mirrors _redact_agent_params_tree on the write side.
|
||||
"_unqualified", # bounded by the qualifier depth of a static TypedDict annotation (Annotated, Required/NotRequired, ReadOnly around one type, no cycles possible).
|
||||
"_render_json", # bounded by the nesting depth of a pydantic-validated JsonValue from the operator's config (a finite JSON tree, no cycles possible).
|
||||
"completion_cost", # max depth 1: recursion only fires for mixed-tier Responses WS logging objects, and each split part carries a single service_tier so _split_responses_ws_logging_object_by_service_tier returns None.
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -106,10 +106,11 @@ def test_short_values_are_written_without_masking_every_digit_in_the_log(tmp_pat
|
|||
env_path: Final = tmp_path / ".env"
|
||||
|
||||
result: Final = subprocess.run(
|
||||
[sys.executable, str(SECRETS_TO_ENV), str(env_path)],
|
||||
[sys.executable, "-I", str(SECRETS_TO_ENV), str(env_path)],
|
||||
input='{"FLAG": "1", "API_KEY": "sk-0123456789abcdef"}',
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env={**os.environ, "GITHUB_ACTIONS": "true"},
|
||||
)
|
||||
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
|
@ -117,6 +118,24 @@ def test_short_values_are_written_without_masking_every_digit_in_the_log(tmp_pat
|
|||
assert env_path.read_text() == "FLAG='1'\nAPI_KEY='sk-0123456789abcdef'\n"
|
||||
|
||||
|
||||
def test_outside_actions_no_value_is_printed(tmp_path: Path) -> None:
|
||||
env_path: Final = tmp_path / ".env"
|
||||
local_env: Final = {key: value for key, value in os.environ.items() if key != "GITHUB_ACTIONS"}
|
||||
|
||||
result: Final = subprocess.run(
|
||||
[sys.executable, "-I", str(SECRETS_TO_ENV), str(env_path)],
|
||||
input='{"FLAG": "1", "API_KEY": "sk-0123456789abcdef"}',
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=local_env,
|
||||
)
|
||||
|
||||
assert result.returncode == 0, result.stderr
|
||||
assert result.stdout == ""
|
||||
assert "sk-0123456789abcdef" not in result.stderr
|
||||
assert env_path.read_text() == "FLAG='1'\nAPI_KEY='sk-0123456789abcdef'\n"
|
||||
|
||||
|
||||
def redact_output(tmp_path: Path, values: tuple[str, ...], text: str) -> tuple[subprocess.CompletedProcess[str], Path]:
|
||||
env_path: Final = tmp_path / ".env"
|
||||
_ = env_path.write_text("".join(f"{name}='{value}'\n" for name, value in zip(("A", "B", "C"), values)))
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
- {id: guardrail.presidio.pre_call.masks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "PII masking pre-call; data-leak blast radius"}
|
||||
- {id: guardrail.presidio.post_call.masks, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Mask PII in model output"}
|
||||
- {id: guardrail.presidio.post_call.masks_generated_output, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, chat_completions_stream, anthropic_messages_stream], source: "guardrail_hooks/presidio.py", rationale: "Mask model-generated credit-card output with the UI default scope"}
|
||||
- {id: guardrail.presidio.post_call.spend_log_stores_masked_output, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, chat_completions_stream, messages, anthropic_messages_stream, responses], source: "guardrail_hooks/presidio.py", fail_before_fix: proven, rationale: "When an output guardrail masks the response, the spend log stores the masked text the caller received rather than the raw model output, on every endpoint and both stream modes (LIT-8325)"}
|
||||
- {id: guardrail.presidio.logging_only.masks, module: guardrail, tier: P0, hook_point: logging_only, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Redact in logs without blocking"}
|
||||
- {id: guardrail.presidio.pre_call.logs_masked_entities, module: guardrail, tier: P0, hook_point: pre_call, assertions: [logs_masked_entities], exercised_on: [chat_completions], source: "guardrail_hooks/presidio.py", rationale: "A masking run must record itself on the spend log: the dashboard's guardrail panel renders the masked-entity counts and per-entity scores straight off metadata.guardrail_information, so a run that masks but records nothing leaves an operator unable to audit it"}
|
||||
- {id: guardrail.bedrock.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "AWS content guardrail blocks harmful input"}
|
||||
|
|
|
|||
|
|
@ -94,3 +94,15 @@
|
|||
- {id: llm.messages.together_ai.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: together_ai, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together over /v1/messages streaming"}
|
||||
- {id: llm.messages.together_ai.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: together_ai, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together tool calls over /v1/messages"}
|
||||
- {id: llm.messages.together_ai.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: together_ai, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together tool result round trip over /v1/messages"}
|
||||
- {id: llm.chat_completions.anthropic.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: chat_completions, route: anthropic, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic over /chat/completions: cost header and spend row agree"}
|
||||
- {id: llm.chat_completions.anthropic.multi_turn.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic tool result round trip over /chat/completions"}
|
||||
- {id: llm.messages.anthropic.multi_turn.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic tool result round trip over /v1/messages"}
|
||||
- {id: llm.messages.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "OpenAI models served on the Anthropic Messages contract"}
|
||||
- {id: llm.messages.openai.basic.stream.works, module: llm, tier: P0, subject_endpoint: messages, route: openai, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "OpenAI over /v1/messages streams the Anthropic event grammar"}
|
||||
- {id: llm.messages.openai.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: messages, route: openai, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "OpenAI over /v1/messages: cost header and spend row agree"}
|
||||
- {id: llm.messages.openai.tool_use.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: openai, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "OpenAI tool calls translated to Anthropic tool_use blocks"}
|
||||
- {id: llm.messages.openai.multi_turn.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: openai, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "OpenAI tool result round trip over /v1/messages"}
|
||||
- {id: llm.responses.openai.multi_turn.nonstream.works, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "OpenAI function_call_output round trip over /v1/responses"}
|
||||
- {id: llm.responses.anthropic.basic.stream.works, module: llm, tier: P1, subject_endpoint: responses, route: anthropic, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic over /v1/responses streams Responses events"}
|
||||
- {id: llm.responses.anthropic.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: responses, route: anthropic, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic over /v1/responses: cost header and spend row agree"}
|
||||
- {id: llm.responses.anthropic.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic function_call_output round trip over /v1/responses"}
|
||||
|
|
|
|||
|
|
@ -15,6 +15,8 @@
|
|||
- {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:1158-1166 _decode_jwt_with_public_key", rationale: "A genuine token whose signature bytes were altered fails verification with 401"}
|
||||
- {id: other.auth.jwt.unknown_team_denied, module: other, tier: P0, area: auth, assertions: [unknown_team_denied], source: "handle_jwt.py:1549-1632 find_team_with_model_access", rationale: "A verified JWT whose groups claim resolves to no existing team is denied with 403 naming the unresolved team, never silently admitted without a team. The proxy words it as a model-access denial, the same body an existing team without model access gets"}
|
||||
- {id: other.auth.jwt.virtual_key_unaffected, module: other, tier: P0, area: auth, assertions: [virtual_key_unaffected], source: "handle_jwt.py:213 is_jwt / user_api_key_auth.py:1332-1333", rationale: "enable_jwt_auth only routes three-segment bearer tokens into the JWT branch, so sk- virtual keys keep working on the same proxy"}
|
||||
- {id: other.auth.jwt.team_header_alias_binds_team, module: other, tier: P0, area: auth, assertions: [team_header_alias_binds_team], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", fail_before_fix: proven, rationale: "x-litellm-team-id carrying the team alias binds and attributes the same team as the team id, so a managed client can pin a stable alias instead of a uuid"}
|
||||
- {id: other.auth.jwt.team_header_non_member_alias_denied, module: other, tier: P0, area: auth, assertions: [team_header_non_member_alias_denied], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", rationale: "x-litellm-team-id naming the alias of a team the JWT does not grant is denied 403 with the same body as an unknown value, so the response does not reveal whether that team exists"}
|
||||
- {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"}
|
||||
- {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"}
|
||||
- {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"}
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
general_settings:
|
||||
master_key: os.environ/LITELLM_MASTER_KEY
|
||||
store_model_in_db: true
|
||||
disable_model_info_refresh: true
|
||||
|
|
|
|||
|
|
@ -381,6 +381,26 @@ class GuardrailsClient:
|
|||
),
|
||||
)
|
||||
|
||||
def messages_raw(
|
||||
self,
|
||||
key: str,
|
||||
model: str,
|
||||
text: str,
|
||||
*,
|
||||
guardrails: list[str] | None = None,
|
||||
max_tokens: int = 64,
|
||||
) -> StreamingResponse:
|
||||
return self.proxy.transport.send(
|
||||
"/v1/messages",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=AnthropicMessagesBody(
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=text)],
|
||||
max_tokens=max_tokens,
|
||||
guardrails=guardrails,
|
||||
),
|
||||
)
|
||||
|
||||
def messages_stream_raw(
|
||||
self,
|
||||
key: str,
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ this suite deliberately requires the detected-entity details to remain visible.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
|
|
@ -481,6 +482,149 @@ class TestPresidioCreditCardOutputMasking:
|
|||
_assert_eventually_masks_generated_card(fetch)
|
||||
|
||||
|
||||
def _wire_text(outcome: StreamingResponse) -> str:
|
||||
return "\n".join(outcome.stream_events) if outcome.is_streaming else outcome.body
|
||||
|
||||
|
||||
def _poll_until_generated_card_masked(fetch: Callable[[], StreamingResponse]) -> StreamingResponse:
|
||||
"""The raw HTTP outcome once the output masker is in effect on the serving
|
||||
worker: whichever wire shape the endpoint speaks, a masked body carries the
|
||||
CREDIT_CARD placeholder and no Luhn-valid card run. A raw card is a worker
|
||||
that has not loaded the guardrail yet, so it is polled through like any
|
||||
other unmasked answer."""
|
||||
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
|
||||
last: str = "<no successful response yet>"
|
||||
while True:
|
||||
outcome = fetch()
|
||||
if outcome.ok and not outcome.stream_error:
|
||||
wire = _wire_text(outcome)
|
||||
last = wire
|
||||
if MASKED_CREDIT_CARD_TOKEN in wire and not _contains_card_number(wire):
|
||||
return outcome
|
||||
if time.monotonic() >= deadline:
|
||||
pytest.fail(
|
||||
"presidio post_call output masking never masked the generated card within "
|
||||
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
|
||||
)
|
||||
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
|
||||
|
||||
|
||||
def _spend_log_response_text(client: GuardrailsClient, key: str, call_id: str) -> str:
|
||||
rows = client.proxy.poll_logs_for_key(
|
||||
key,
|
||||
predicate=lambda logged: any(row.litellm_call_id == call_id for row in logged),
|
||||
)
|
||||
row = next((row for row in rows if row.litellm_call_id == call_id), None)
|
||||
assert row is not None, f"no spend log row ever appeared for x-litellm-call-id {call_id}"
|
||||
return json.dumps(row.response)
|
||||
|
||||
|
||||
class TestPresidioSpendLogStoresMaskedOutput:
|
||||
"""The spend log stores the response the caller received, on every endpoint
|
||||
and both stream modes, when an output-only post_call Presidio guardrail masks
|
||||
a card number the model generated."""
|
||||
|
||||
_CELL: Final = "guardrail.presidio.post_call.spend_log_stores_masked_output"
|
||||
|
||||
def _assert_spend_log_is_masked(
|
||||
self,
|
||||
client: GuardrailsClient,
|
||||
resources: ResourceManager,
|
||||
key: str,
|
||||
*,
|
||||
name: str,
|
||||
fetch: Callable[[str, str], StreamingResponse],
|
||||
) -> None:
|
||||
_register_presidio(
|
||||
client,
|
||||
resources,
|
||||
name=name,
|
||||
mode="post_call",
|
||||
filter_scope="output",
|
||||
entities={"CREDIT_CARD": "MASK"},
|
||||
)
|
||||
prompt: Final = _credit_card_prompt(unique_marker())
|
||||
outcome = _poll_until_generated_card_masked(lambda: fetch(prompt, name))
|
||||
assert outcome.call_id, f"the served response must carry x-litellm-call-id: {dict(outcome.headers)}"
|
||||
|
||||
logged = _spend_log_response_text(client, key, outcome.call_id)
|
||||
assert not _contains_card_number(logged), (
|
||||
"the caller got the masked response but the spend log stored the raw model output: "
|
||||
f"{logged[:400]!r}"
|
||||
)
|
||||
assert MASKED_CREDIT_CARD_TOKEN in logged, (
|
||||
f"the spend log response carries neither the card nor the placeholder: {logged[:400]!r}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers(_CELL, exercised_on=["chat_completions"])
|
||||
def test_spend_log_stores_masked_output_on_chat_completions(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
self._assert_spend_log_is_masked(
|
||||
client,
|
||||
resources,
|
||||
scoped_key,
|
||||
name=f"e2e-presidio-log-card-chat-{unique_marker()}",
|
||||
fetch=lambda prompt, guardrail: client.chat_raw(
|
||||
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
|
||||
),
|
||||
)
|
||||
|
||||
@pytest.mark.covers(_CELL, exercised_on=["chat_completions_stream"])
|
||||
def test_spend_log_stores_masked_output_on_streaming_chat_completions(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
self._assert_spend_log_is_masked(
|
||||
client,
|
||||
resources,
|
||||
scoped_key,
|
||||
name=f"e2e-presidio-log-card-chat-stream-{unique_marker()}",
|
||||
fetch=lambda prompt, guardrail: client.chat_stream_raw(
|
||||
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
|
||||
),
|
||||
)
|
||||
|
||||
@pytest.mark.covers(_CELL, exercised_on=["messages"])
|
||||
def test_spend_log_stores_masked_output_on_anthropic_messages(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
self._assert_spend_log_is_masked(
|
||||
client,
|
||||
resources,
|
||||
scoped_key,
|
||||
name=f"e2e-presidio-log-card-messages-{unique_marker()}",
|
||||
fetch=lambda prompt, guardrail: client.messages_raw(
|
||||
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
|
||||
),
|
||||
)
|
||||
|
||||
@pytest.mark.covers(_CELL, exercised_on=["anthropic_messages_stream"])
|
||||
def test_spend_log_stores_masked_output_on_streaming_anthropic_messages(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
self._assert_spend_log_is_masked(
|
||||
client,
|
||||
resources,
|
||||
scoped_key,
|
||||
name=f"e2e-presidio-log-card-messages-stream-{unique_marker()}",
|
||||
fetch=lambda prompt, guardrail: client.messages_stream_raw(
|
||||
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
|
||||
),
|
||||
)
|
||||
|
||||
@pytest.mark.covers(_CELL, exercised_on=["responses"])
|
||||
def test_spend_log_stores_masked_output_on_responses(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
self._assert_spend_log_is_masked(
|
||||
client,
|
||||
resources,
|
||||
scoped_key,
|
||||
name=f"e2e-presidio-log-card-responses-{unique_marker()}",
|
||||
fetch=lambda prompt, guardrail: client.responses(scoped_key, MODEL, prompt, guardrails=[guardrail]),
|
||||
)
|
||||
|
||||
|
||||
_LOGGED_ENTITIES: dict[PiiEntity, PiiAction] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"}
|
||||
|
||||
_ENTITY_LIST_ADAPTER: Final = TypeAdapter(list[GuardrailEntityMatch])
|
||||
|
|
|
|||
544
tests/e2e/llm_translation/conversational_matrix.py
Normal file
544
tests/e2e/llm_translation/conversational_matrix.py
Normal file
|
|
@ -0,0 +1,544 @@
|
|||
"""The endpoint x deployment x auth matrix behind test_conversational_matrix_e2e.py.
|
||||
|
||||
One conversation, three wire formats. Each `Surface` speaks its own API through
|
||||
the customer SDK (chat completions and Responses through the OpenAI SDK, Messages
|
||||
through the Anthropic SDK) and folds what came back into the surface-neutral
|
||||
`Reply` / `StreamedReply`, so a single behavior test asserts the same contract on
|
||||
every cell. A new model, from an existing or a new provider, is one `Deployment` row
|
||||
in DEPLOYMENTS; a new way of handing the proxy a provider credential is one
|
||||
`AuthMethod`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, Protocol
|
||||
|
||||
import anthropic
|
||||
import openai
|
||||
import pytest
|
||||
from _pytest.mark.structures import ParameterSet
|
||||
from anthropic.types import (
|
||||
MessageParam,
|
||||
RawMessageStreamEvent,
|
||||
TextBlock,
|
||||
ToolChoiceToolParam,
|
||||
ToolParam,
|
||||
ToolResultBlockParam,
|
||||
ToolUseBlock,
|
||||
ToolUseBlockParam,
|
||||
)
|
||||
from e2e_config import provider_edge_base, unique_marker
|
||||
from lifecycle import ResourceManager
|
||||
from llm_translation.sdk_clients import NO_PROXY_CACHE, SdkClients, response_header
|
||||
from models import CredentialCreateBody, LiteLLMParamsBody
|
||||
from openai.types.chat import (
|
||||
ChatCompletionAssistantMessageParam,
|
||||
ChatCompletionChunk,
|
||||
ChatCompletionMessageFunctionToolCallParam,
|
||||
ChatCompletionMessageParam,
|
||||
ChatCompletionNamedToolChoiceParam,
|
||||
ChatCompletionToolMessageParam,
|
||||
ChatCompletionToolParam,
|
||||
)
|
||||
from openai.types.chat.chat_completion_message_function_tool_call import ChatCompletionMessageFunctionToolCall
|
||||
from openai.types.responses import (
|
||||
FunctionToolParam,
|
||||
ResponseFunctionToolCall,
|
||||
ResponseFunctionToolCallParam,
|
||||
ResponseInputParam,
|
||||
ResponseStreamEvent,
|
||||
ToolChoiceFunctionParam,
|
||||
)
|
||||
from openai.types.responses.response_input_param import FunctionCallOutput
|
||||
from proxy_client import ProxyClient
|
||||
from pydantic import BaseModel
|
||||
|
||||
SurfaceName = Literal["chat_completions", "messages", "responses"]
|
||||
AuthMethod = Literal["env_ref", "stored_credential"]
|
||||
Capability = Literal["basic", "tool_use", "multi_turn"]
|
||||
Streaming = Literal["stream", "nonstream"]
|
||||
Assertion = Literal["works", "cost_logged"]
|
||||
ToolMode = Literal["none", "forced", "offered"]
|
||||
|
||||
SURFACES: Final[tuple[SurfaceName, ...]] = ("chat_completions", "messages", "responses")
|
||||
AUTH_METHODS: Final[tuple[AuthMethod, ...]] = ("env_ref", "stored_credential")
|
||||
|
||||
MAX_OUTPUT_TOKENS: Final = 512
|
||||
INSTRUCTIONS: Final = "You are a terse assistant. Answer in one short sentence."
|
||||
GREETING_PROMPT: Final = "Say hello."
|
||||
WEATHER_PROMPT: Final = "What is the weather in Paris right now? Use the get_weather tool."
|
||||
WEATHER_REPORT: Final = "Paris: 22 degrees Celsius, clear skies"
|
||||
WEATHER_TOOL_NAME: Final = "get_weather"
|
||||
WEATHER_TOOL_DESCRIPTION: Final = "Current weather for a city"
|
||||
WEATHER_TOOL_SCHEMA: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {"location": {"type": "string", "description": "City name"}},
|
||||
"required": ["location"],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Deployment:
|
||||
"""One deployment target: the litellm backend string plus how to wire it."""
|
||||
|
||||
route: Literal["openai", "anthropic"]
|
||||
label: str
|
||||
backend: str
|
||||
api_key_env: str
|
||||
edge_mount: str
|
||||
edge_suffix: str
|
||||
|
||||
def api_base(self) -> str | None:
|
||||
base: Final = provider_edge_base(self.edge_mount)
|
||||
return None if base is None else f"{base}{self.edge_suffix}"
|
||||
|
||||
def api_key(self) -> str:
|
||||
key: Final = os.environ.get(self.api_key_env, "")
|
||||
assert key, f"{self.api_key_env} is not set in the test process environment"
|
||||
return key
|
||||
|
||||
|
||||
DEPLOYMENTS: Final[tuple[Deployment, ...]] = (
|
||||
Deployment(
|
||||
route="openai",
|
||||
label="gpt-4o-mini",
|
||||
backend="openai/gpt-4o-mini",
|
||||
api_key_env="OPENAI_API_KEY",
|
||||
edge_mount="openai",
|
||||
edge_suffix="/v1",
|
||||
),
|
||||
Deployment(
|
||||
route="openai",
|
||||
label="gpt-5.4-mini",
|
||||
backend="openai/gpt-5.4-mini",
|
||||
api_key_env="OPENAI_API_KEY",
|
||||
edge_mount="openai",
|
||||
edge_suffix="/v1",
|
||||
),
|
||||
Deployment(
|
||||
route="anthropic",
|
||||
label="claude-haiku-4-5",
|
||||
backend="anthropic/claude-haiku-4-5",
|
||||
api_key_env="ANTHROPIC_API_KEY",
|
||||
edge_mount="anthropic",
|
||||
edge_suffix="",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Cell:
|
||||
surface: SurfaceName
|
||||
deployment: Deployment
|
||||
auth: AuthMethod
|
||||
|
||||
@property
|
||||
def id(self) -> str:
|
||||
return f"{self.surface}-{self.deployment.label}-{self.auth}"
|
||||
|
||||
def registry_id(self, capability: Capability, streaming: Streaming, assertion: Assertion) -> str:
|
||||
return f"llm.{self.surface}.{self.deployment.route}.{capability}.{streaming}.{assertion}"
|
||||
|
||||
|
||||
CELLS: Final[tuple[Cell, ...]] = tuple(
|
||||
Cell(surface=surface, deployment=deployment, auth=auth)
|
||||
for surface in SURFACES
|
||||
for deployment in DEPLOYMENTS
|
||||
for auth in AUTH_METHODS
|
||||
)
|
||||
|
||||
|
||||
def cells_covering(capability: Capability, streaming: Streaming, assertion: Assertion) -> tuple[ParameterSet, ...]:
|
||||
"""Every cell as a pytest param carrying the registry id its test proves."""
|
||||
return tuple(
|
||||
pytest.param(cell, id=cell.id, marks=pytest.mark.covers(cell.registry_id(capability, streaming, assertion)))
|
||||
for cell in CELLS
|
||||
)
|
||||
|
||||
|
||||
DeploymentKey = tuple[str, AuthMethod]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Deployments:
|
||||
"""Model aliases registered on the proxy, one per (deployment, auth)."""
|
||||
|
||||
aliases: Mapping[DeploymentKey, str]
|
||||
|
||||
def alias(self, cell: Cell) -> str:
|
||||
return self.aliases[(cell.deployment.label, cell.auth)]
|
||||
|
||||
|
||||
def _litellm_params(deployment: Deployment, auth: AuthMethod, credential_name: str) -> LiteLLMParamsBody:
|
||||
match auth:
|
||||
case "env_ref":
|
||||
return LiteLLMParamsBody(
|
||||
model=deployment.backend, api_key=f"os.environ/{deployment.api_key_env}", api_base=deployment.api_base()
|
||||
)
|
||||
case "stored_credential":
|
||||
return LiteLLMParamsBody(
|
||||
model=deployment.backend, litellm_credential_name=credential_name, api_base=deployment.api_base()
|
||||
)
|
||||
|
||||
|
||||
def _register(proxy: ProxyClient, resources: ResourceManager, deployment: Deployment, auth: AuthMethod) -> str:
|
||||
marker: Final = unique_marker()
|
||||
credential_name: Final = f"e2e-matrix-{deployment.label}-{marker}"
|
||||
if auth == "stored_credential":
|
||||
proxy.create_credential(
|
||||
CredentialCreateBody(credential_name=credential_name, credential_values={"api_key": deployment.api_key()})
|
||||
)
|
||||
resources.defer(lambda: proxy.delete_credential(credential_name))
|
||||
alias: Final = f"e2e-matrix-{deployment.label}-{auth}-{marker}"
|
||||
model_id: Final = proxy.create_model(alias, _litellm_params(deployment, auth, credential_name))
|
||||
resources.defer(lambda: proxy.delete_model(model_id))
|
||||
return alias
|
||||
|
||||
|
||||
def register_deployments(proxy: ProxyClient) -> Iterator[Deployments]:
|
||||
resources: Final = ResourceManager(client=proxy)
|
||||
try:
|
||||
yield Deployments(
|
||||
aliases=MappingProxyType(
|
||||
{
|
||||
(deployment.label, auth): _register(proxy, resources, deployment, auth)
|
||||
for deployment in DEPLOYMENTS
|
||||
for auth in AUTH_METHODS
|
||||
}
|
||||
)
|
||||
)
|
||||
finally:
|
||||
resources.teardown()
|
||||
|
||||
|
||||
class WeatherArgs(BaseModel):
|
||||
location: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ToolCall:
|
||||
call_id: str
|
||||
name: str
|
||||
arguments: str
|
||||
|
||||
def parsed(self) -> WeatherArgs:
|
||||
return WeatherArgs.model_validate_json(self.arguments)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Usage:
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Reply:
|
||||
"""What every surface owes the caller for one non-streamed turn."""
|
||||
|
||||
response_id: str
|
||||
text: str
|
||||
tool_calls: tuple[ToolCall, ...]
|
||||
usage: Usage | None
|
||||
call_id_header: str | None
|
||||
cost_header: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class StreamedReply:
|
||||
"""The reassembled stream: its text, whether the surface's own terminal event
|
||||
arrived, and whether usage was reported anywhere in the stream."""
|
||||
|
||||
text: str
|
||||
finished: bool
|
||||
usage_reported: bool
|
||||
event_count: int
|
||||
|
||||
|
||||
class Surface(Protocol):
|
||||
@property
|
||||
def name(self) -> SurfaceName: ...
|
||||
|
||||
def reply(self, key: str, model: str, prompt: str, *, with_tool: bool = False) -> Reply: ...
|
||||
|
||||
def stream(self, key: str, model: str, prompt: str) -> StreamedReply: ...
|
||||
|
||||
def reply_to_tool_result(self, key: str, model: str, prompt: str, call: ToolCall, result: str) -> Reply: ...
|
||||
|
||||
|
||||
def _chat_tool() -> ChatCompletionToolParam:
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": WEATHER_TOOL_NAME,
|
||||
"description": WEATHER_TOOL_DESCRIPTION,
|
||||
"parameters": dict(WEATHER_TOOL_SCHEMA),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _messages_tool() -> ToolParam:
|
||||
return {
|
||||
"name": WEATHER_TOOL_NAME,
|
||||
"description": WEATHER_TOOL_DESCRIPTION,
|
||||
"input_schema": dict(WEATHER_TOOL_SCHEMA),
|
||||
}
|
||||
|
||||
|
||||
def _responses_tool() -> FunctionToolParam:
|
||||
return {
|
||||
"type": "function",
|
||||
"name": WEATHER_TOOL_NAME,
|
||||
"description": WEATHER_TOOL_DESCRIPTION,
|
||||
"parameters": dict(WEATHER_TOOL_SCHEMA),
|
||||
"strict": False,
|
||||
}
|
||||
|
||||
|
||||
def _chat_tool_choice() -> ChatCompletionNamedToolChoiceParam:
|
||||
return {"type": "function", "function": {"name": WEATHER_TOOL_NAME}}
|
||||
|
||||
|
||||
def _messages_tool_choice() -> ToolChoiceToolParam:
|
||||
return {"type": "tool", "name": WEATHER_TOOL_NAME, "disable_parallel_tool_use": True}
|
||||
|
||||
|
||||
def _responses_tool_choice() -> ToolChoiceFunctionParam:
|
||||
return {"type": "function", "name": WEATHER_TOOL_NAME}
|
||||
|
||||
|
||||
def _usage(input_tokens: int | None, output_tokens: int | None) -> Usage | None:
|
||||
if input_tokens is None or output_tokens is None:
|
||||
return None
|
||||
return Usage(input_tokens=input_tokens, output_tokens=output_tokens)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ChatCompletionsSurface:
|
||||
sdk: SdkClients
|
||||
name: SurfaceName = "chat_completions"
|
||||
|
||||
def _turn(self, key: str, model: str, messages: Sequence[ChatCompletionMessageParam], tool: ToolMode) -> Reply:
|
||||
raw: Final = self.sdk.openai(key).chat.completions.with_raw_response.create(
|
||||
model=model,
|
||||
messages=list(messages),
|
||||
max_completion_tokens=MAX_OUTPUT_TOKENS,
|
||||
tools=openai.omit if tool == "none" else [_chat_tool()],
|
||||
tool_choice=_chat_tool_choice() if tool == "forced" else openai.omit,
|
||||
parallel_tool_calls=False if tool == "forced" else openai.omit,
|
||||
extra_body=NO_PROXY_CACHE,
|
||||
)
|
||||
completion: Final = raw.parse()
|
||||
message: Final = completion.choices[0].message
|
||||
calls: Final = tuple(
|
||||
ToolCall(call_id=call.id, name=call.function.name, arguments=call.function.arguments)
|
||||
for call in message.tool_calls or ()
|
||||
if isinstance(call, ChatCompletionMessageFunctionToolCall)
|
||||
)
|
||||
return Reply(
|
||||
response_id=completion.id,
|
||||
text=message.content or "",
|
||||
tool_calls=calls,
|
||||
usage=None
|
||||
if completion.usage is None
|
||||
else _usage(completion.usage.prompt_tokens, completion.usage.completion_tokens),
|
||||
call_id_header=response_header(raw.headers, "x-litellm-call-id"),
|
||||
cost_header=response_header(raw.headers, "x-litellm-response-cost"),
|
||||
)
|
||||
|
||||
def reply(self, key: str, model: str, prompt: str, *, with_tool: bool = False) -> Reply:
|
||||
return self._turn(key, model, _chat_history(prompt), "forced" if with_tool else "none")
|
||||
|
||||
def stream(self, key: str, model: str, prompt: str) -> StreamedReply:
|
||||
chunks: Final[tuple[ChatCompletionChunk, ...]] = tuple(
|
||||
self.sdk.openai(key).chat.completions.create(
|
||||
model=model,
|
||||
messages=_chat_history(prompt),
|
||||
max_completion_tokens=MAX_OUTPUT_TOKENS,
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
extra_body=NO_PROXY_CACHE,
|
||||
)
|
||||
)
|
||||
return StreamedReply(
|
||||
text="".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices),
|
||||
finished=any(chunk.choices[0].finish_reason is not None for chunk in chunks if chunk.choices),
|
||||
usage_reported=any(chunk.usage is not None for chunk in chunks),
|
||||
event_count=len(chunks),
|
||||
)
|
||||
|
||||
def reply_to_tool_result(self, key: str, model: str, prompt: str, call: ToolCall, result: str) -> Reply:
|
||||
tool_call: Final[ChatCompletionMessageFunctionToolCallParam] = {
|
||||
"id": call.call_id,
|
||||
"type": "function",
|
||||
"function": {"name": call.name, "arguments": call.arguments},
|
||||
}
|
||||
assistant: Final[ChatCompletionAssistantMessageParam] = {"role": "assistant", "tool_calls": [tool_call]}
|
||||
tool_result: Final[ChatCompletionToolMessageParam] = {
|
||||
"role": "tool",
|
||||
"tool_call_id": call.call_id,
|
||||
"content": result,
|
||||
}
|
||||
return self._turn(key, model, (*_chat_history(prompt), assistant, tool_result), "offered")
|
||||
|
||||
|
||||
def _chat_history(prompt: str) -> tuple[ChatCompletionMessageParam, ...]:
|
||||
return ({"role": "system", "content": INSTRUCTIONS}, {"role": "user", "content": prompt})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MessagesSurface:
|
||||
sdk: SdkClients
|
||||
name: SurfaceName = "messages"
|
||||
|
||||
def _turn(self, key: str, model: str, messages: Sequence[MessageParam], tool: ToolMode) -> Reply:
|
||||
raw: Final = self.sdk.anthropic(key).messages.with_raw_response.create(
|
||||
model=model,
|
||||
max_tokens=MAX_OUTPUT_TOKENS,
|
||||
system=INSTRUCTIONS,
|
||||
messages=list(messages),
|
||||
tools=anthropic.omit if tool == "none" else [_messages_tool()],
|
||||
tool_choice=_messages_tool_choice() if tool == "forced" else anthropic.omit,
|
||||
extra_body=NO_PROXY_CACHE,
|
||||
)
|
||||
message: Final = raw.parse()
|
||||
return Reply(
|
||||
response_id=message.id,
|
||||
text="".join(block.text for block in message.content if isinstance(block, TextBlock)),
|
||||
tool_calls=tuple(
|
||||
ToolCall(call_id=block.id, name=block.name, arguments=json.dumps(block.input))
|
||||
for block in message.content
|
||||
if isinstance(block, ToolUseBlock)
|
||||
),
|
||||
usage=_usage(message.usage.input_tokens, message.usage.output_tokens),
|
||||
call_id_header=response_header(raw.headers, "x-litellm-call-id"),
|
||||
cost_header=response_header(raw.headers, "x-litellm-response-cost"),
|
||||
)
|
||||
|
||||
def reply(self, key: str, model: str, prompt: str, *, with_tool: bool = False) -> Reply:
|
||||
return self._turn(key, model, ({"role": "user", "content": prompt},), "forced" if with_tool else "none")
|
||||
|
||||
def stream(self, key: str, model: str, prompt: str) -> StreamedReply:
|
||||
events: Final[tuple[RawMessageStreamEvent, ...]] = tuple(
|
||||
self.sdk.anthropic(key).messages.create(
|
||||
model=model,
|
||||
max_tokens=MAX_OUTPUT_TOKENS,
|
||||
system=INSTRUCTIONS,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
stream=True,
|
||||
extra_body=NO_PROXY_CACHE,
|
||||
)
|
||||
)
|
||||
return StreamedReply(
|
||||
text="".join(
|
||||
event.delta.text
|
||||
for event in events
|
||||
if event.type == "content_block_delta" and event.delta.type == "text_delta"
|
||||
),
|
||||
finished=any(event.type == "message_stop" for event in events),
|
||||
usage_reported=any(event.type == "message_delta" and event.usage.output_tokens > 0 for event in events),
|
||||
event_count=len(events),
|
||||
)
|
||||
|
||||
def reply_to_tool_result(self, key: str, model: str, prompt: str, call: ToolCall, result: str) -> Reply:
|
||||
tool_use: Final[ToolUseBlockParam] = {
|
||||
"type": "tool_use",
|
||||
"id": call.call_id,
|
||||
"name": call.name,
|
||||
"input": call.parsed().model_dump(),
|
||||
}
|
||||
tool_result: Final[ToolResultBlockParam] = {
|
||||
"type": "tool_result",
|
||||
"tool_use_id": call.call_id,
|
||||
"content": result,
|
||||
}
|
||||
history: Final[tuple[MessageParam, ...]] = (
|
||||
{"role": "user", "content": prompt},
|
||||
{"role": "assistant", "content": [tool_use]},
|
||||
{"role": "user", "content": [tool_result]},
|
||||
)
|
||||
return self._turn(key, model, history, "offered")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ResponsesSurface:
|
||||
sdk: SdkClients
|
||||
name: SurfaceName = "responses"
|
||||
|
||||
def _turn(self, key: str, model: str, history: ResponseInputParam, tool: ToolMode) -> Reply:
|
||||
raw: Final = self.sdk.openai(key).responses.with_raw_response.create(
|
||||
model=model,
|
||||
input=history,
|
||||
instructions=INSTRUCTIONS,
|
||||
max_output_tokens=MAX_OUTPUT_TOKENS,
|
||||
tools=openai.omit if tool == "none" else [_responses_tool()],
|
||||
tool_choice=_responses_tool_choice() if tool == "forced" else openai.omit,
|
||||
parallel_tool_calls=False if tool == "forced" else openai.omit,
|
||||
extra_body=NO_PROXY_CACHE,
|
||||
)
|
||||
response: Final = raw.parse()
|
||||
return Reply(
|
||||
response_id=response.id,
|
||||
text=response.output_text,
|
||||
tool_calls=tuple(
|
||||
ToolCall(call_id=item.call_id, name=item.name, arguments=item.arguments)
|
||||
for item in response.output
|
||||
if isinstance(item, ResponseFunctionToolCall)
|
||||
),
|
||||
usage=None if response.usage is None else _usage(response.usage.input_tokens, response.usage.output_tokens),
|
||||
call_id_header=response_header(raw.headers, "x-litellm-call-id"),
|
||||
cost_header=response_header(raw.headers, "x-litellm-response-cost"),
|
||||
)
|
||||
|
||||
def reply(self, key: str, model: str, prompt: str, *, with_tool: bool = False) -> Reply:
|
||||
return self._turn(key, model, [{"role": "user", "content": prompt}], "forced" if with_tool else "none")
|
||||
|
||||
def stream(self, key: str, model: str, prompt: str) -> StreamedReply:
|
||||
events: Final[tuple[ResponseStreamEvent, ...]] = tuple(
|
||||
self.sdk.openai(key).responses.create(
|
||||
model=model,
|
||||
input=prompt,
|
||||
instructions=INSTRUCTIONS,
|
||||
max_output_tokens=MAX_OUTPUT_TOKENS,
|
||||
stream=True,
|
||||
extra_body=NO_PROXY_CACHE,
|
||||
)
|
||||
)
|
||||
return StreamedReply(
|
||||
text="".join(event.delta for event in events if event.type == "response.output_text.delta"),
|
||||
finished=bool(events) and events[-1].type == "response.completed",
|
||||
usage_reported=any(
|
||||
event.type == "response.completed" and event.response.usage is not None for event in events
|
||||
),
|
||||
event_count=len(events),
|
||||
)
|
||||
|
||||
def reply_to_tool_result(self, key: str, model: str, prompt: str, call: ToolCall, result: str) -> Reply:
|
||||
function_call: Final[ResponseFunctionToolCallParam] = {
|
||||
"type": "function_call",
|
||||
"call_id": call.call_id,
|
||||
"name": call.name,
|
||||
"arguments": call.arguments,
|
||||
}
|
||||
output: Final[FunctionCallOutput] = {
|
||||
"type": "function_call_output",
|
||||
"call_id": call.call_id,
|
||||
"output": result,
|
||||
}
|
||||
return self._turn(key, model, [{"role": "user", "content": prompt}, function_call, output], "offered")
|
||||
|
||||
|
||||
def build_surfaces(sdk: SdkClients) -> Mapping[SurfaceName, Surface]:
|
||||
return MappingProxyType[SurfaceName, Surface](
|
||||
{
|
||||
"chat_completions": ChatCompletionsSurface(sdk),
|
||||
"messages": MessagesSurface(sdk),
|
||||
"responses": ResponsesSurface(sdk),
|
||||
}
|
||||
)
|
||||
163
tests/e2e/llm_translation/test_conversational_matrix_e2e.py
Normal file
163
tests/e2e/llm_translation/test_conversational_matrix_e2e.py
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
"""The same conversation contract on every (endpoint, deployment, auth) cell.
|
||||
A deployment is one provider model (openai/gpt-4o-mini, anthropic/claude-haiku-4-5, ...).
|
||||
|
||||
/chat/completions, /v1/messages and /v1/responses each have their own
|
||||
translation code in the proxy, so a bug fixed on one surface tends to survive
|
||||
on the others. Every test here runs once per cell in `CELLS`
|
||||
(conversational_matrix.py), so a change to a shared helper is proven against all
|
||||
surfaces and providers at once, and a new model or provider is one row in `DEPLOYMENTS`.
|
||||
|
||||
Edge-wired: OpenAI and Anthropic traffic goes through the provider edge in
|
||||
record and replay, so the whole matrix replays with zero provider calls.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator, Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from e2e_config import unique_marker
|
||||
from lifecycle import ResourceManager
|
||||
from llm_translation.conversational_matrix import (
|
||||
GREETING_PROMPT,
|
||||
WEATHER_PROMPT,
|
||||
WEATHER_REPORT,
|
||||
WEATHER_TOOL_NAME,
|
||||
Cell,
|
||||
Deployments,
|
||||
Surface,
|
||||
SurfaceName,
|
||||
ToolCall,
|
||||
build_surfaces,
|
||||
cells_covering,
|
||||
register_deployments,
|
||||
)
|
||||
from llm_translation.sdk_clients import SdkClients
|
||||
from models import SpendLogRow
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
pytestmark = [pytest.mark.e2e, pytest.mark.replayable]
|
||||
|
||||
|
||||
def _approx_equal(actual: float, expected: float) -> bool:
|
||||
return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def deployments(proxy: ProxyClient) -> Iterator[Deployments]:
|
||||
yield from register_deployments(proxy)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def surfaces(sdk: SdkClients) -> Mapping[SurfaceName, Surface]:
|
||||
return build_surfaces(sdk)
|
||||
|
||||
|
||||
def _weather_call(surface: Surface, key: str, model: str) -> ToolCall:
|
||||
first: Final = surface.reply(key, model, WEATHER_PROMPT, with_tool=True)
|
||||
assert len(first.tool_calls) == 1, (
|
||||
f"{surface.name} forced tool_choice={WEATHER_TOOL_NAME} with parallel calls off, "
|
||||
f"got {len(first.tool_calls)} tool call(s): {first.tool_calls} text={first.text!r}"
|
||||
)
|
||||
call: Final = first.tool_calls[0]
|
||||
assert call.name == WEATHER_TOOL_NAME, f"{surface.name} called {call.name!r}, not the forced {WEATHER_TOOL_NAME!r}"
|
||||
assert call.call_id, f"{surface.name} tool call has no id, so the caller cannot answer it: {call}"
|
||||
assert "paris" in call.parsed().location.lower(), f"{surface.name} tool arguments lost the location: {call}"
|
||||
return call
|
||||
|
||||
|
||||
class TestConversationalMatrix:
|
||||
@pytest.mark.parametrize("cell", cells_covering("basic", "nonstream", "works"))
|
||||
def test_reply_carries_assistant_text_and_usage(
|
||||
self,
|
||||
cell: Cell,
|
||||
deployments: Deployments,
|
||||
surfaces: Mapping[SurfaceName, Surface],
|
||||
resources: ResourceManager,
|
||||
) -> None:
|
||||
surface: Final = surfaces[cell.surface]
|
||||
reply: Final = surface.reply(resources.key(), deployments.alias(cell), GREETING_PROMPT)
|
||||
|
||||
assert reply.response_id, f"{cell.id}: response has no id"
|
||||
assert reply.text.strip(), f"{cell.id}: response carried no assistant text"
|
||||
assert reply.usage is not None and reply.usage.input_tokens > 0 and reply.usage.output_tokens > 0, (
|
||||
f"{cell.id}: usage missing or zero, so the caller cannot account for this call: {reply.usage}"
|
||||
)
|
||||
assert reply.call_id_header, f"{cell.id}: x-litellm-call-id header missing"
|
||||
|
||||
@pytest.mark.parametrize("cell", cells_covering("basic", "stream", "works"))
|
||||
def test_stream_delivers_text_usage_and_a_terminal_event(
|
||||
self,
|
||||
cell: Cell,
|
||||
deployments: Deployments,
|
||||
surfaces: Mapping[SurfaceName, Surface],
|
||||
resources: ResourceManager,
|
||||
) -> None:
|
||||
surface: Final = surfaces[cell.surface]
|
||||
streamed: Final = surface.stream(resources.key(), deployments.alias(cell), GREETING_PROMPT)
|
||||
|
||||
assert streamed.event_count > 1, f"{cell.id}: stream arrived as {streamed.event_count} event(s), not a stream"
|
||||
assert streamed.text.strip(), f"{cell.id}: stream carried no text deltas"
|
||||
assert streamed.finished, f"{cell.id}: stream never sent its terminal event"
|
||||
assert streamed.usage_reported, f"{cell.id}: stream never reported usage"
|
||||
|
||||
@pytest.mark.parametrize("cell", cells_covering("basic", "nonstream", "cost_logged"))
|
||||
def test_cost_header_matches_the_spend_log(
|
||||
self,
|
||||
cell: Cell,
|
||||
deployments: Deployments,
|
||||
surfaces: Mapping[SurfaceName, Surface],
|
||||
resources: ResourceManager,
|
||||
proxy: ProxyClient,
|
||||
) -> None:
|
||||
key: Final = resources.key()
|
||||
surface: Final = surfaces[cell.surface]
|
||||
reply: Final = surface.reply(key, deployments.alias(cell), f"{GREETING_PROMPT} {unique_marker()}")
|
||||
|
||||
assert reply.cost_header is not None, f"{cell.id}: x-litellm-response-cost header missing"
|
||||
header_cost: Final = float(reply.cost_header)
|
||||
assert header_cost > 0, f"{cell.id}: x-litellm-response-cost is not positive: {header_cost}"
|
||||
|
||||
def _priced(rows: list[SpendLogRow]) -> bool:
|
||||
return any(row.spend is not None and row.spend > 0 for row in rows)
|
||||
|
||||
rows: Final = proxy.poll_logs_for_key(key, predicate=_priced)
|
||||
priced: Final = tuple(row for row in rows if row.spend is not None and row.spend > 0)
|
||||
assert len(priced) == 1, f"{cell.id}: expected exactly one priced spend row for a fresh key, got {rows}"
|
||||
row: Final = priced[0]
|
||||
assert (row.prompt_tokens or 0) > 0 and (row.completion_tokens or 0) > 0, (
|
||||
f"{cell.id}: spend row has no token counts, so the cost is not real usage: {row}"
|
||||
)
|
||||
assert row.spend is not None and _approx_equal(row.spend, header_cost), (
|
||||
f"{cell.id}: logged spend {row.spend} disagrees with x-litellm-response-cost {header_cost}"
|
||||
)
|
||||
assert row.model and cell.deployment.backend.endswith(row.model), (
|
||||
f"{cell.id}: spend row logged model {row.model!r}, not the deployment's {cell.deployment.backend!r}"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("cell", cells_covering("tool_use", "nonstream", "works"))
|
||||
def test_tool_call_is_returned_named_and_addressable(
|
||||
self,
|
||||
cell: Cell,
|
||||
deployments: Deployments,
|
||||
surfaces: Mapping[SurfaceName, Surface],
|
||||
resources: ResourceManager,
|
||||
) -> None:
|
||||
_weather_call(surfaces[cell.surface], resources.key(), deployments.alias(cell))
|
||||
|
||||
@pytest.mark.parametrize("cell", cells_covering("multi_turn", "nonstream", "works"))
|
||||
def test_tool_result_round_trip_reaches_the_model(
|
||||
self,
|
||||
cell: Cell,
|
||||
deployments: Deployments,
|
||||
surfaces: Mapping[SurfaceName, Surface],
|
||||
resources: ResourceManager,
|
||||
) -> None:
|
||||
key: Final = resources.key()
|
||||
model: Final = deployments.alias(cell)
|
||||
surface: Final = surfaces[cell.surface]
|
||||
call: Final = _weather_call(surface, key, model)
|
||||
|
||||
answer: Final = surface.reply_to_tool_result(key, model, WEATHER_PROMPT, call, WEATHER_REPORT)
|
||||
assert "22" in answer.text, f"{cell.id}: the model never saw the tool result: {answer.text!r}"
|
||||
|
|
@ -874,6 +874,8 @@ class SpendLogRow(BaseModel):
|
|||
request_tags: list[str] | None = None
|
||||
metadata: SpendLogMetadata | None = None
|
||||
proxy_server_request: JsonValue = None
|
||||
response: JsonValue = None
|
||||
litellm_call_id: str | None = None
|
||||
|
||||
|
||||
class SpendLogs(RootModel[list[SpendLogRow]]):
|
||||
|
|
|
|||
|
|
@ -15,9 +15,11 @@ from __future__ import annotations
|
|||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from e2e_http import NoBody, ProbeResult, Result
|
||||
from e2e_http import AuthHeaders, NoBody, ProbeResult, Result
|
||||
from idp import Keycloak, keycloak_from_env
|
||||
from models import (
|
||||
ChatBody,
|
||||
ChatResponse,
|
||||
JwtKeyMappingDeleteBody,
|
||||
JwtKeyMappingDeleteResponse,
|
||||
JwtKeyMappingListParams,
|
||||
|
|
@ -32,6 +34,14 @@ from models import (
|
|||
UserNewResponse,
|
||||
)
|
||||
from proxy_client import ProxyClient
|
||||
from pydantic import Field
|
||||
|
||||
|
||||
class TeamHeaders(AuthHeaders):
|
||||
"""Bearer auth plus ``x-litellm-team-id``, the header a JWT caller sends to
|
||||
pick one of the teams it belongs to."""
|
||||
|
||||
x_litellm_team_id: str = Field(serialization_alias="x-litellm-team-id")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -112,6 +122,18 @@ class OtherClient:
|
|||
response_type=JwtKeyMappingDeleteResponse,
|
||||
)
|
||||
|
||||
def chat_as_team(self, token: str, team: str, body: ChatBody) -> Result[ChatResponse]:
|
||||
"""POST /chat/completions under `token` with `x-litellm-team-id: team`."""
|
||||
return self.proxy.transport.post(
|
||||
"/chat/completions",
|
||||
headers=TeamHeaders(
|
||||
authorization=self.proxy.transport.bearer(token).authorization,
|
||||
x_litellm_team_id=team,
|
||||
),
|
||||
json=body,
|
||||
response_type=ChatResponse,
|
||||
)
|
||||
|
||||
def list_users_as(self, key: str) -> Result[UserListResponse]:
|
||||
"""GET /user/list under `key`. Admin-only, so it doubles as the master
|
||||
key's authorization proof: the master key (proxy admin) reads it, a
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from __future__ import annotations
|
|||
|
||||
import base64
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
|
@ -52,6 +53,31 @@ def identity(client: OtherClient, resources: ResourceManager) -> Identity:
|
|||
return provisioned
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BoundTeam:
|
||||
identity: Identity
|
||||
team_id: str
|
||||
team_alias: str
|
||||
|
||||
|
||||
def _team(client: OtherClient, resources: ResourceManager, *, marker: str, team_id: str) -> str:
|
||||
"""A litellm team whose alias differs from its id, so a header naming one
|
||||
cannot accidentally match the other."""
|
||||
team_alias: Final = f"e2e-jwt-alias-{marker}"
|
||||
created: Final = client.proxy.create_team(TeamNewBody(team_alias=team_alias, team_id=team_id))
|
||||
resources.defer(lambda: client.proxy.delete_team(created))
|
||||
return team_alias
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bound_team(client: OtherClient, resources: ResourceManager) -> BoundTeam:
|
||||
"""An identity whose single group is a real team, plus that team's alias."""
|
||||
marker: Final = unique_marker()
|
||||
provisioned: Final = _provision(client, resources, marker=marker)
|
||||
team_alias: Final = _team(client, resources, marker=marker, team_id=provisioned.group)
|
||||
return BoundTeam(identity=provisioned, team_id=provisioned.group, team_alias=team_alias)
|
||||
|
||||
|
||||
def _ping() -> ChatBody:
|
||||
return ChatBody(
|
||||
model=CHEAP_OPENAI_MODEL,
|
||||
|
|
@ -155,3 +181,58 @@ class TestJwtAuth:
|
|||
def test_plain_virtual_key_still_works_with_jwt_auth_enabled(self, client: OtherClient, scoped_key: str) -> None:
|
||||
response: Final = unwrap(client.proxy.chat(scoped_key, _ping()))
|
||||
assert response.choices, f"an sk- key must keep working on a proxy with enable_jwt_auth, got {response}"
|
||||
|
||||
|
||||
def _team_of_request(client: OtherClient, token: str, team: str) -> str | None:
|
||||
response: Final = unwrap(client.chat_as_team(token, team, _ping()))
|
||||
assert response.id is not None and response.choices, (
|
||||
f"chat with x-litellm-team-id={team!r} returned no completion: {response}"
|
||||
)
|
||||
rows: Final = client.proxy.poll_logs_for_request_id(response.id)
|
||||
assert rows, f"no spend log row for request {response.id} within the poll deadline"
|
||||
return rows[0].team_id
|
||||
|
||||
|
||||
def _denial(client: OtherClient, token: str, team: str) -> str:
|
||||
result: Final = client.chat_as_team(token, team, _ping())
|
||||
assert isinstance(result, UnknownApiError) and result.status_code == 403, (
|
||||
f"x-litellm-team-id={team!r} names no team the caller is in, so it must be rejected with 403, got {result}"
|
||||
)
|
||||
assert team in result.body, f"the 403 must name the header value it rejected ({team!r}), got {result.body[:300]}"
|
||||
return result.body
|
||||
|
||||
|
||||
class TestJwtTeamHeader:
|
||||
@pytest.mark.covers("other.auth.jwt.team_header_alias_binds_team")
|
||||
def test_team_header_with_the_team_alias_binds_the_same_team_as_the_team_id(
|
||||
self, client: OtherClient, bound_team: BoundTeam
|
||||
) -> None:
|
||||
token: Final = client.idp.access_token(bound_team.identity)
|
||||
assert bound_team.team_alias != bound_team.team_id
|
||||
|
||||
by_id: Final = _team_of_request(client, token, bound_team.team_id)
|
||||
assert by_id == bound_team.team_id, (
|
||||
f"precondition: x-litellm-team-id with the team id must bind {bound_team.team_id!r}, got {by_id!r}"
|
||||
)
|
||||
|
||||
by_alias: Final = _team_of_request(client, token, bound_team.team_alias)
|
||||
assert by_alias == bound_team.team_id, (
|
||||
f"x-litellm-team-id={bound_team.team_alias!r} must bind the same team as its id "
|
||||
f"{bound_team.team_id!r}, got {by_alias!r}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("other.auth.jwt.team_header_non_member_alias_denied")
|
||||
def test_team_header_with_the_alias_of_a_team_the_caller_is_not_in_is_rejected_like_an_unknown_value(
|
||||
self, client: OtherClient, resources: ResourceManager, bound_team: BoundTeam
|
||||
) -> None:
|
||||
token: Final = client.idp.access_token(bound_team.identity)
|
||||
other_marker: Final = unique_marker()
|
||||
other_alias: Final = _team(client, resources, marker=other_marker, team_id=f"e2e-jwt-other-{other_marker}")
|
||||
unknown: Final = f"e2e-jwt-unknown-{unique_marker()}"
|
||||
|
||||
for_other_alias: Final = _denial(client, token, other_alias)
|
||||
for_unknown: Final = _denial(client, token, unknown)
|
||||
assert for_other_alias.replace(other_alias, "<value>") == for_unknown.replace(unknown, "<value>"), (
|
||||
"a non-member alias and an unknown value must get the same denial body, so the response does not "
|
||||
f"reveal whether the team exists; got {for_other_alias[:300]!r} vs {for_unknown[:300]!r}"
|
||||
)
|
||||
|
|
|
|||
204
tests/test_litellm/litellm_core_utils/test_bug_report.py
Normal file
204
tests/test_litellm/litellm_core_utils/test_bug_report.py
Normal file
|
|
@ -0,0 +1,204 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import cast
|
||||
from urllib.parse import parse_qs, unquote_plus, urlparse
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm._version import version
|
||||
from litellm.exceptions import APIConnectionError, BadRequestError, InternalServerError
|
||||
from litellm.litellm_core_utils.bug_report import (
|
||||
DISABLE_ENV_VAR,
|
||||
ISSUE_URL_BASE,
|
||||
MAX_FRAMES,
|
||||
MAX_URL_LENGTH,
|
||||
allowlisted,
|
||||
bug_report_enabled,
|
||||
bug_report_issue_url,
|
||||
bug_report_notice,
|
||||
build_bug_report,
|
||||
should_report_bug,
|
||||
strip_bug_report_notice,
|
||||
)
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
|
||||
def test_build_bug_report_keeps_only_litellm_frames():
|
||||
with pytest.raises(BadRequestError) as raised:
|
||||
get_llm_provider(cast(str, None))
|
||||
report = build_bug_report(raised.value, surface="sdk")
|
||||
|
||||
assert report.litellm_frames
|
||||
assert all(frame.startswith("litellm/") for frame in report.litellm_frames)
|
||||
assert all("test_bug_report.py" not in frame for frame in report.litellm_frames)
|
||||
|
||||
|
||||
def test_issue_url_never_contains_the_exception_message():
|
||||
secret = "sk-abcdefghijklmnopqrstuvwxyz1234567890"
|
||||
prompt = "my social security number is 123-45-6789"
|
||||
report = build_bug_report(RuntimeError(f"{secret} {prompt}"), surface="sdk")
|
||||
url = bug_report_issue_url(report)
|
||||
query = parse_qs(urlparse(url).query)
|
||||
|
||||
assert url.startswith(ISSUE_URL_BASE)
|
||||
assert secret not in url and "123-45-6789" not in url and "social" not in url
|
||||
assert query["title"] == ["[Bug]: RuntimeError in litellm"]
|
||||
assert query["version"] == [version]
|
||||
assert query["template"] == ["bug_report.yml"]
|
||||
assert query["domain"] == ["Python SDK: the litellm package itself"]
|
||||
assert query["deployment"] == ["pip / Python SDK"]
|
||||
assert "Exception: `RuntimeError`" in query["description"][0]
|
||||
assert "Python:" in query["description"][0]
|
||||
|
||||
|
||||
def test_issue_url_drops_unknown_provider_and_call_type():
|
||||
report = build_bug_report(
|
||||
ValueError("boom"),
|
||||
surface="proxy",
|
||||
custom_llm_provider="acme-internal-gateway",
|
||||
)
|
||||
query = parse_qs(urlparse(bug_report_issue_url(report)).query)
|
||||
|
||||
assert report.custom_llm_provider is None
|
||||
assert "acme" not in bug_report_issue_url(report)
|
||||
assert "Provider: unknown" in query["description"][0]
|
||||
assert "Endpoint / call: unknown" in query["description"][0]
|
||||
|
||||
|
||||
def test_allowlisted_only_passes_exact_members():
|
||||
allowed = frozenset({"/v1/chat/completions"})
|
||||
|
||||
assert allowlisted("/v1/chat/completions", allowed) == "/v1/chat/completions"
|
||||
assert allowlisted("/v1/chat/completions/../../admin", allowed) is None
|
||||
assert allowlisted(None, allowed) is None
|
||||
assert allowlisted(["/v1/chat/completions"], allowed) is None
|
||||
assert allowlisted({"provider": "openai"}, allowed) is None
|
||||
|
||||
|
||||
def test_build_bug_report_survives_unhashable_provider_from_request_data():
|
||||
report = build_bug_report(
|
||||
KeyError("missing"),
|
||||
surface="proxy",
|
||||
custom_llm_provider={"name": "openai"},
|
||||
)
|
||||
|
||||
assert report.custom_llm_provider is None
|
||||
|
||||
|
||||
def test_frames_are_capped_and_url_is_bounded():
|
||||
namespace: dict[str, object] = {}
|
||||
exec(
|
||||
compile(
|
||||
"def recurse(depth):\n if depth == 0:\n raise RuntimeError('deep')\n recurse(depth - 1)\n",
|
||||
str(Path(litellm.__file__).with_name("fake_deep_module.py")),
|
||||
"exec",
|
||||
),
|
||||
namespace,
|
||||
)
|
||||
recurse = cast(Callable[[int], None], namespace["recurse"])
|
||||
|
||||
with pytest.raises(RuntimeError) as raised:
|
||||
recurse(200)
|
||||
report = build_bug_report(raised.value, surface="proxy")
|
||||
|
||||
assert len(report.litellm_frames) == MAX_FRAMES
|
||||
assert all(frame.startswith("litellm/fake_deep_module.py:") for frame in report.litellm_frames)
|
||||
assert len(bug_report_issue_url(report)) <= MAX_URL_LENGTH
|
||||
|
||||
|
||||
def test_issue_url_builds_without_a_traceback():
|
||||
exc = RuntimeError("no traceback")
|
||||
assert exc.__traceback__ is None
|
||||
report = build_bug_report(exc, surface="proxy")
|
||||
|
||||
assert report.litellm_frames == ()
|
||||
assert bug_report_issue_url(report).startswith(ISSUE_URL_BASE)
|
||||
|
||||
|
||||
def test_bug_report_can_be_disabled(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv(DISABLE_ENV_VAR, "true")
|
||||
|
||||
assert bug_report_enabled() is False
|
||||
assert should_report_bug(RuntimeError("boom")) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"exc",
|
||||
[
|
||||
InternalServerError(message="upstream 500", llm_provider="openai", model="gpt-4"),
|
||||
APIConnectionError(
|
||||
message="connection reset",
|
||||
llm_provider="openai",
|
||||
model="gpt-4",
|
||||
request=httpx.Request(method="POST", url="https://api.openai.com/v1/"),
|
||||
),
|
||||
BadRequestError(message="bad input", llm_provider="openai", model="gpt-4"),
|
||||
"not an exception",
|
||||
],
|
||||
)
|
||||
def test_should_report_bug_skips_provider_and_network_errors(exc: object):
|
||||
assert should_report_bug(exc) is False
|
||||
|
||||
|
||||
def test_should_report_bug_accepts_plain_python_errors():
|
||||
assert should_report_bug(KeyError("missing")) is True
|
||||
|
||||
|
||||
def test_proxy_known_provider_uses_translation_domain():
|
||||
report = build_bug_report(
|
||||
RuntimeError("proxy failure"),
|
||||
surface="proxy",
|
||||
call_type="/v1/chat/completions",
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
query = parse_qs(urlparse(bug_report_issue_url(report)).query)
|
||||
|
||||
assert report.custom_llm_provider == "openai"
|
||||
assert query["domain"] == ["LLM translation: a specific provider's request or response"]
|
||||
assert "Endpoint / call: /v1/chat/completions" in query["description"][0]
|
||||
|
||||
|
||||
def test_strip_bug_report_notice():
|
||||
report = build_bug_report(RuntimeError("boom"), surface="sdk")
|
||||
notice = bug_report_notice(report)
|
||||
|
||||
assert strip_bug_report_notice(f"boom\n\n{notice}") == "boom\n"
|
||||
assert strip_bug_report_notice("boom") == "boom"
|
||||
|
||||
|
||||
def test_issue_description_renders_stream_and_config_block():
|
||||
report = build_bug_report(
|
||||
RuntimeError("boom"),
|
||||
surface="proxy",
|
||||
stream=True,
|
||||
config_lines=("router_settings.routing_strategy = least-busy", "litellm_settings.drop_params = true"),
|
||||
)
|
||||
description = parse_qs(urlparse(bug_report_issue_url(report)).query)["description"][0]
|
||||
|
||||
assert "Stream: true\n" in description
|
||||
assert "```\nrouter_settings.routing_strategy = least-busy\nlitellm_settings.drop_params = true\n```" in description
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [None, "true", 1])
|
||||
def test_issue_description_omits_stream_unless_it_is_a_bool(stream: object):
|
||||
report = build_bug_report(RuntimeError("boom"), surface="proxy", stream=stream)
|
||||
|
||||
assert report.stream is None
|
||||
assert "Stream:" not in unquote_plus(bug_report_issue_url(report))
|
||||
|
||||
|
||||
def test_oversized_config_is_trimmed_from_the_end_before_any_frame():
|
||||
with pytest.raises(BadRequestError) as raised:
|
||||
get_llm_provider(cast(str, None))
|
||||
config_lines = tuple(f"general_settings.flag_{index:04d} = true" for index in range(400))
|
||||
report = build_bug_report(raised.value, surface="proxy", config_lines=config_lines)
|
||||
description = parse_qs(urlparse(url := bug_report_issue_url(report)).query)["description"][0]
|
||||
|
||||
assert len(url) <= MAX_URL_LENGTH
|
||||
assert all(frame in description for frame in report.litellm_frames)
|
||||
assert "general_settings.flag_0000 = true" in description
|
||||
assert "general_settings.flag_0399 = true" not in description
|
||||
|
|
@ -3,8 +3,6 @@ import openai
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
|
||||
|
||||
from litellm.litellm_core_utils.exception_mapping_utils import (
|
||||
ExceptionCheckers,
|
||||
_get_body_error_code,
|
||||
|
|
@ -974,6 +972,31 @@ def test_an_unmapped_exception_with_no_model_or_provider_is_a_connection_error(q
|
|||
assert "boom" in raised.value.message
|
||||
|
||||
|
||||
def test_unmapped_sdk_exception_includes_bug_report_link(quiet_exception_mapping):
|
||||
with pytest.raises(litellm.APIConnectionError) as raised:
|
||||
exception_type(
|
||||
model="my-model",
|
||||
custom_llm_provider="minimax",
|
||||
original_exception=ValueError("boom"),
|
||||
)
|
||||
|
||||
assert "https://github.com/BerriAI/litellm/issues/new?" in str(raised.value)
|
||||
assert "ValueError" in str(raised.value)
|
||||
|
||||
|
||||
def test_unmapped_sdk_exception_bug_report_link_can_be_disabled(quiet_exception_mapping, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_DISABLE_BUG_REPORT_LINK", "true")
|
||||
|
||||
with pytest.raises(litellm.APIConnectionError) as raised:
|
||||
exception_type(
|
||||
model="my-model",
|
||||
custom_llm_provider="minimax",
|
||||
original_exception=ValueError("boom"),
|
||||
)
|
||||
|
||||
assert "https://github.com/BerriAI/litellm/issues/new?" not in str(raised.value)
|
||||
|
||||
|
||||
def _raise_and_map(model: str | None, original_exception: Exception, custom_llm_provider: str | None) -> None:
|
||||
"""Calls exception_type() from inside the except block, as litellm/main.py does,
|
||||
so traceback.format_exc() has a real stack."""
|
||||
|
|
|
|||
|
|
@ -1,8 +1,12 @@
|
|||
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
|
||||
inherit_message_logging_privacy,
|
||||
initialize_standard_callback_dynamic_params,
|
||||
iter_client_callback_metadata_dicts,
|
||||
)
|
||||
|
|
@ -189,6 +193,20 @@ def test_empty_kwargs_returns_empty_params():
|
|||
assert dict(params) == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("child_privacy", (False, True))
|
||||
def test_inherited_privacy_only_strengthens_child_and_resets(child_privacy: bool) -> None:
|
||||
kwargs: Final = TypeAdapter(dict[str, object]).validate_python(
|
||||
MappingProxyType({"turn_off_message_logging": child_privacy})
|
||||
)
|
||||
with inherit_message_logging_privacy(False):
|
||||
assert initialize_standard_callback_dynamic_params(kwargs)["turn_off_message_logging"] is child_privacy
|
||||
with inherit_message_logging_privacy(True), inherit_message_logging_privacy(False):
|
||||
params: Final = initialize_standard_callback_dynamic_params(kwargs)
|
||||
assert initialize_standard_callback_dynamic_params(kwargs)["turn_off_message_logging"] is child_privacy
|
||||
assert params["turn_off_message_logging"] is True
|
||||
assert initialize_standard_callback_dynamic_params().get("turn_off_message_logging") is None
|
||||
|
||||
|
||||
def test_newrelic_callback_params_are_not_extracted_from_request_kwargs():
|
||||
kwargs = {
|
||||
"newrelic_api_key": "caller-key",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import contextlib
|
||||
import datetime
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
|
@ -305,14 +306,15 @@ def test_response_cost_calculator_uses_router_model_id_from_litellm_metadata():
|
|||
|
||||
|
||||
class TestZeroCostDiagnostic:
|
||||
DEPLOYMENT_ID: Final = "lit7898-per-second-priced-deployment"
|
||||
MODEL_GROUP: Final = "per-second-priced-chat"
|
||||
DEPLOYMENT_ID: Final = "lit7898-query-only-priced-deployment"
|
||||
MODEL_GROUP: Final = "query-only-priced-chat"
|
||||
QUERY_ONLY_PRICING: Final = {"input_cost_per_query": 0.00042}
|
||||
PER_SECOND_PRICING: Final = {"input_cost_per_second": 0.00042, "output_cost_per_second": 0.00042}
|
||||
FREE_PRICING: Final = {"input_cost_per_token": 0, "output_cost_per_token": 0}
|
||||
|
||||
@pytest.fixture(params=["per_second", "free"])
|
||||
@pytest.fixture(params=["query_only", "free"])
|
||||
def deployment_pricing(self, request: pytest.FixtureRequest) -> Iterator[Mapping[str, float]]:
|
||||
pricing: Final = self.PER_SECOND_PRICING if request.param == "per_second" else self.FREE_PRICING
|
||||
pricing: Final = self.QUERY_ONLY_PRICING if request.param == "query_only" else self.FREE_PRICING
|
||||
litellm.register_model(model_cost={self.DEPLOYMENT_ID: pricing}, persist_across_reloads=False)
|
||||
try:
|
||||
yield pricing
|
||||
|
|
@ -557,11 +559,11 @@ class TestZeroCostDiagnostic:
|
|||
priced_pricing: Final = {"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}
|
||||
usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30)
|
||||
litellm.register_model(
|
||||
model_cost={self.DEPLOYMENT_ID: self.PER_SECOND_PRICING, priced_id: priced_pricing},
|
||||
model_cost={self.DEPLOYMENT_ID: self.QUERY_ONLY_PRICING, priced_id: priced_pricing},
|
||||
persist_across_reloads=False,
|
||||
)
|
||||
try:
|
||||
logging_obj: Final = self._logging_obj(self.PER_SECOND_PRICING)
|
||||
logging_obj: Final = self._logging_obj(self.QUERY_ONLY_PRICING)
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
assert logging_obj._response_cost_calculator(result=self._response(usage)) == 0.0
|
||||
self._assert_flagged(logging_obj, caplog)
|
||||
|
|
@ -570,7 +572,7 @@ class TestZeroCostDiagnostic:
|
|||
assert logging_obj._response_cost_calculator(result=self._response(usage)) == pytest.approx(5e-05)
|
||||
assert logging_obj.model_call_details["zero_cost_diagnostic"] is None
|
||||
|
||||
self._route_to_deployment(logging_obj, self.PER_SECOND_PRICING)
|
||||
self._route_to_deployment(logging_obj, self.QUERY_ONLY_PRICING)
|
||||
assert logging_obj._response_cost_calculator(result=self._response(usage)) == 0.0
|
||||
|
||||
assert logging_obj.model_call_details["zero_cost_diagnostic"]["reason"] == "missing_pricing_key"
|
||||
|
|
@ -585,7 +587,7 @@ class TestZeroCostDiagnostic:
|
|||
dated_model: Final = "lit7898-nano-2026-03-17"
|
||||
requested_model: Final = "lit7898-nano"
|
||||
usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30)
|
||||
cost_map_entry: Final = {"litellm_provider": "openai", "mode": "chat", **self.PER_SECOND_PRICING}
|
||||
cost_map_entry: Final = {"litellm_provider": "openai", "mode": "chat", **self.QUERY_ONLY_PRICING}
|
||||
litellm.register_model(
|
||||
model_cost={dated_model: cost_map_entry, requested_model: cost_map_entry}, persist_across_reloads=False
|
||||
)
|
||||
|
|
@ -646,6 +648,24 @@ class TestZeroCostDiagnostic:
|
|||
assert logging_obj.model_call_details["zero_cost_diagnostic"] is None
|
||||
assert self._zero_cost_warnings(caplog) == []
|
||||
|
||||
def test_per_second_priced_deployment_bills_the_call_duration_and_stays_silent(
|
||||
self, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
per_second_id: Final = "lit8315-per-second-priced-deployment"
|
||||
usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30)
|
||||
litellm.register_model(model_cost={per_second_id: self.PER_SECOND_PRICING}, persist_across_reloads=False)
|
||||
try:
|
||||
logging_obj: Final = self._logging_obj(self.PER_SECOND_PRICING, deployment_id=per_second_id)
|
||||
response: Final = self._response(usage)
|
||||
response._response_ms = 1000.0
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.00084)
|
||||
|
||||
assert logging_obj.model_call_details["zero_cost_diagnostic"] is None
|
||||
assert self._zero_cost_warnings(caplog) == []
|
||||
finally:
|
||||
litellm.model_cost.pop(per_second_id, None)
|
||||
|
||||
@pytest.mark.parametrize("spilled_over", [True, False])
|
||||
def test_ptu_deployment_is_judged_by_the_entry_the_calculator_priced_with(
|
||||
self, spilled_over: bool, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
|
|
@ -663,7 +683,7 @@ class TestZeroCostDiagnostic:
|
|||
litellm.register_model(
|
||||
model_cost={
|
||||
router_model_id: {**self.FREE_PRICING, "litellm_provider": "azure", "mode": "chat"},
|
||||
served_model: {**self.PER_SECOND_PRICING, "litellm_provider": "azure", "mode": "chat"},
|
||||
served_model: {**self.QUERY_ONLY_PRICING, "litellm_provider": "azure", "mode": "chat"},
|
||||
},
|
||||
persist_across_reloads=False,
|
||||
)
|
||||
|
|
@ -2994,6 +3014,54 @@ def test_get_final_response_obj_with_empty_response_obj_and_list_init():
|
|||
assert result[1].name == "Object2"
|
||||
|
||||
|
||||
def test_get_final_response_obj_stores_the_text_a_post_call_guardrail_served():
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
||||
|
||||
raw = {
|
||||
"id": "x",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"message": {"role": "assistant", "content": "Card: 4111 1111 1111 1111"},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
logged = StandardLoggingPayloadSetup.get_final_response_obj(
|
||||
response_obj=raw, init_response_obj=raw, kwargs={SERVED_OUTPUT_TEXTS_KEY: ("Card: <CREDIT_CARD>",)}
|
||||
)
|
||||
untouched = StandardLoggingPayloadSetup.get_final_response_obj(response_obj=raw, init_response_obj=raw, kwargs={})
|
||||
|
||||
assert isinstance(logged, dict)
|
||||
assert logged["choices"][0]["message"]["content"] == "Card: <CREDIT_CARD>"
|
||||
assert logged["choices"][0]["finish_reason"] == "stop"
|
||||
assert untouched == raw
|
||||
|
||||
|
||||
def test_get_final_response_obj_redacts_the_served_text_when_message_logging_is_off(monkeypatch: pytest.MonkeyPatch):
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
||||
|
||||
monkeypatch.setattr(litellm, "turn_off_message_logging", True)
|
||||
raw = {
|
||||
"id": "x",
|
||||
"choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "Card: 4111"}}],
|
||||
}
|
||||
|
||||
logged = StandardLoggingPayloadSetup.get_final_response_obj(
|
||||
response_obj=raw,
|
||||
init_response_obj=raw,
|
||||
kwargs={SERVED_OUTPUT_TEXTS_KEY: ("Card: <CREDIT_CARD>",), "litellm_params": {}},
|
||||
)
|
||||
|
||||
assert isinstance(logged, dict)
|
||||
assert "<CREDIT_CARD>" not in json.dumps(logged), logged
|
||||
assert "4111" not in json.dumps(logged), logged
|
||||
|
||||
|
||||
def test_get_usage_as_dict():
|
||||
"""
|
||||
Test get_usage_as_dict returns usage as plain dict from response_obj or combined_usage_object.
|
||||
|
|
@ -3504,6 +3572,20 @@ def _make_logging_obj(stream: bool) -> LitellmLogging:
|
|||
)
|
||||
|
||||
|
||||
def test_get_response_ms_measures_a_float_start_time_against_a_datetime_end_time():
|
||||
"""The files paths construct the logging object with ``time.time()`` while the success
|
||||
handler stamps a datetime end, and the per-second cost path reads this window."""
|
||||
logging_obj = _make_logging_obj(stream=False)
|
||||
logging_obj.update_environment_variables(
|
||||
model="openai/codex-mini-latest", user="", optional_params={}, litellm_params={}
|
||||
)
|
||||
start_seconds = logging_obj.model_call_details["start_time"]
|
||||
assert isinstance(start_seconds, float)
|
||||
logging_obj.model_call_details["end_time"] = datetime.datetime.fromtimestamp(start_seconds + 1.5)
|
||||
|
||||
assert logging_obj.get_response_ms() == pytest.approx(1500)
|
||||
|
||||
|
||||
def test_get_assembled_streaming_response_returns_none_for_non_streaming():
|
||||
"""Non-streaming requests should return None so the streaming block is skipped."""
|
||||
import datetime
|
||||
|
|
|
|||
|
|
@ -1032,3 +1032,11 @@ def test_a_callback_that_redacts_itself_keeps_its_messages_but_not_the_classifie
|
|||
assert "classifier_input" not in stored
|
||||
assert stored["messages"] == payload["messages"]
|
||||
assert stored["response"] == payload["response"]
|
||||
|
||||
|
||||
def test_perform_redaction_drops_the_served_output_texts_from_the_callback_kwargs() -> None:
|
||||
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
||||
|
||||
details: Final = {"litellm_params": {}, SERVED_OUTPUT_TEXTS_KEY: ("Card: <CREDIT_CARD>",)}
|
||||
perform_redaction(details, None)
|
||||
assert SERVED_OUTPUT_TEXTS_KEY not in details
|
||||
|
|
|
|||
|
|
@ -0,0 +1,114 @@
|
|||
from litellm.litellm_core_utils.served_output_texts import (
|
||||
SERVED_OUTPUT_TEXTS_KEY,
|
||||
overlay_served_output_texts,
|
||||
record_served_output_texts,
|
||||
served_output_texts,
|
||||
served_stream_output_texts,
|
||||
)
|
||||
from litellm.types.utils import Choices, Delta, Message, ModelResponse, ModelResponseStream, StreamingChoices
|
||||
|
||||
RAW = "Card: 4111 1111 1111 1111"
|
||||
MASKED = "Card: <CREDIT_CARD>"
|
||||
|
||||
|
||||
def _chat_response(*texts: str) -> ModelResponse:
|
||||
return ModelResponse(
|
||||
choices=[Choices(index=i, message=Message(content=text, role="assistant")) for i, text in enumerate(texts)]
|
||||
)
|
||||
|
||||
|
||||
def _chat_dict(*texts: str) -> dict[str, object]:
|
||||
return {
|
||||
"id": "x",
|
||||
"choices": [
|
||||
{"index": i, "finish_reason": "stop", "message": {"role": "assistant", "content": text}}
|
||||
for i, text in enumerate(texts)
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _choice_texts(response: object) -> tuple[str | None, ...]:
|
||||
texts = served_output_texts(response)
|
||||
assert texts is not None, response
|
||||
return texts
|
||||
|
||||
|
||||
def _stream_chunk(text: str, index: int = 0) -> ModelResponseStream:
|
||||
return ModelResponseStream(choices=[StreamingChoices(index=index, delta=Delta(content=text))])
|
||||
|
||||
|
||||
def test_served_output_texts_reads_each_response_shape():
|
||||
assert served_output_texts(_chat_response(MASKED, "second")) == (MASKED, "second")
|
||||
assert served_output_texts(_chat_dict(MASKED)) == (MASKED,)
|
||||
assert served_output_texts(
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}],
|
||||
}
|
||||
) == ("ab",)
|
||||
responses_api: dict[str, object] = {
|
||||
"object": "response",
|
||||
"output": [
|
||||
{"type": "reasoning", "content": []},
|
||||
{"type": "message", "content": [{"type": "output_text", "text": MASKED}]},
|
||||
],
|
||||
}
|
||||
assert served_output_texts(responses_api) == (MASKED,)
|
||||
assert served_output_texts({"data": [{"embedding": [0.1]}]}) is None
|
||||
assert served_output_texts("plain") is None
|
||||
|
||||
|
||||
def test_served_stream_output_texts_joins_chat_chunks_and_reads_anthropic_sse():
|
||||
assert served_stream_output_texts([_stream_chunk("Card: "), _stream_chunk("<CREDIT_CARD>")]) == (MASKED,)
|
||||
sse = (
|
||||
'event: message_start\ndata: {"type":"message_start","message":{"id":"m","type":"message","role":"assistant",'
|
||||
'"content":[],"model":"x","usage":{"input_tokens":1,"output_tokens":1}}}\n\n',
|
||||
'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}\n\n',
|
||||
f'event: content_block_delta\ndata: {{"type":"content_block_delta","index":0,"delta":{{"type":"text_delta","text":"{MASKED}"}}}}\n\n',
|
||||
'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n',
|
||||
'event: message_stop\ndata: {"type":"message_stop"}\n\n',
|
||||
)
|
||||
assert served_stream_output_texts(tuple(chunk.encode() for chunk in sse)) == (MASKED,)
|
||||
assert served_stream_output_texts([]) is None
|
||||
assert served_stream_output_texts([b"not sse"]) is None
|
||||
|
||||
|
||||
def test_served_stream_output_texts_keeps_every_choice_index_when_chunks_carry_one_choice_each():
|
||||
chunks = [_stream_chunk("first ", 0), _stream_chunk(MASKED, 1), _stream_chunk("choice", 0)]
|
||||
assert served_stream_output_texts(chunks) == ("first choice", MASKED)
|
||||
|
||||
|
||||
def test_blanked_output_is_served_as_empty_text_and_overlaid():
|
||||
assert served_output_texts(_chat_response("")) == ("",)
|
||||
assert served_output_texts({"type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}]}) == (
|
||||
"",
|
||||
)
|
||||
assert served_stream_output_texts([_stream_chunk("")]) == ("",)
|
||||
assert _choice_texts(overlay_served_output_texts(_chat_dict(RAW), ("",))) == ("",)
|
||||
|
||||
|
||||
def test_overlay_replaces_logged_choice_text_with_served_text():
|
||||
logged = _chat_dict(RAW, RAW)
|
||||
overlaid = overlay_served_output_texts(logged, (MASKED,))
|
||||
assert _choice_texts(overlaid) == (MASKED, RAW)
|
||||
assert isinstance(overlaid, dict)
|
||||
assert overlaid["id"] == "x"
|
||||
assert _choice_texts(logged) == (RAW, RAW)
|
||||
|
||||
|
||||
def test_overlay_leaves_unreadable_inputs_untouched():
|
||||
logged = _chat_dict(RAW)
|
||||
assert overlay_served_output_texts(logged, None) is logged
|
||||
assert overlay_served_output_texts(logged, "not a tuple") is logged
|
||||
assert overlay_served_output_texts(logged, (None,)) == logged
|
||||
assert overlay_served_output_texts("text", (MASKED,)) == "text"
|
||||
assert overlay_served_output_texts({"data": []}, (MASKED,)) == {"data": []}
|
||||
|
||||
|
||||
def test_record_served_output_texts_only_writes_readable_texts():
|
||||
details: dict[str, object] = {}
|
||||
record_served_output_texts(details, None)
|
||||
assert SERVED_OUTPUT_TEXTS_KEY not in details
|
||||
record_served_output_texts(details, (MASKED,))
|
||||
assert details[SERVED_OUTPUT_TEXTS_KEY] == (MASKED,)
|
||||
|
|
@ -1,8 +1,12 @@
|
|||
|
||||
import pytest
|
||||
|
||||
import json
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
ANTHROPIC_MIN_THINKING_BUDGET_TOKENS,
|
||||
|
|
@ -3771,6 +3775,70 @@ def test_multiple_compaction_blocks():
|
|||
assert compaction_blocks[1]["content"] == "Second summary..."
|
||||
|
||||
|
||||
@pytest.mark.parametrize("messages_api,gateway,native_endpoint", [
|
||||
(False, False, False), (True, False, False), (False, True, False), (True, True, False), (True, True, True),
|
||||
])
|
||||
async def test_native_compaction_wire_roundtrip(
|
||||
messages_api: bool, gateway: bool, native_endpoint: bool,
|
||||
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
monkeypatch.setenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "True")
|
||||
monkeypatch.setattr(litellm.anthropic_beta_headers_manager, "_BETA_HEADERS_CONFIG", None)
|
||||
monkeypatch.setattr(litellm, "use_chat_completions_url_for_anthropic_messages", False)
|
||||
block: Final = {"type": "compaction", "content": "Exact summary", "signature": "opaque-signature"}
|
||||
operation: Final = {"type": "summarize", "instructions": "Keep identifiers"}
|
||||
usage: Final = {"input_tokens": 0, "output_tokens": 0,
|
||||
"iterations": [{"type": "compaction", "input_tokens": 103, "output_tokens": 165}]}
|
||||
chat_wire: Final = gateway and not native_endpoint
|
||||
base: Final = "https://gateway.test/v1" if gateway else "https://api.anthropic.com/v1"
|
||||
route: Final = respx_mock.post(f"{base}/{'chat/completions' if chat_wire else 'messages'}")
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
payload: Final = json.loads(request.content)
|
||||
assert len(request.headers.get_list("anthropic-beta")) == 1
|
||||
assert {value.strip() for value in request.headers["anthropic-beta"].split(",")} == {
|
||||
"compact-2026-09-04", "interleaved-thinking-2025-05-14",
|
||||
}
|
||||
if "compaction" in payload:
|
||||
assert payload["compaction"] == operation
|
||||
else:
|
||||
assert payload["messages"][0] == {"role": "assistant", "content": [block]}
|
||||
body: Final = (
|
||||
{"id": "chatcmpl_compact", "object": "chat.completion", "created": 1, "model": "claude-sonnet-5",
|
||||
"choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "",
|
||||
"provider_specific_fields": {"compaction_blocks": [block]}}}],
|
||||
"usage": {"prompt_tokens": 103, "completion_tokens": 165, "total_tokens": 268}}
|
||||
if chat_wire else
|
||||
{"id": "msg_compact", "type": "message", "role": "assistant", "model": "claude-sonnet-5",
|
||||
"content": [block], "stop_reason": "compaction", "usage": usage}
|
||||
)
|
||||
return httpx.Response(200, json=body)
|
||||
|
||||
route.mock(side_effect=respond)
|
||||
call: Final = litellm.anthropic.messages.acreate if messages_api else litellm.acompletion
|
||||
params: Final = dict(
|
||||
model=f"{'openai/' if gateway else ''}anthropic/claude-sonnet-5", api_key="test", max_tokens=512,
|
||||
api_base=base if gateway else "https://api.anthropic.com",
|
||||
extra_headers={"Anthropic-Beta": f"interleaved-thinking-2025-05-14{',compact-2026-09-04' if gateway else ''}"},
|
||||
model_info={"supported_endpoints": ["/v1/messages"]} if native_endpoint else {},
|
||||
)
|
||||
response: Final = await call(
|
||||
messages=[{"role": "user", "content": "Remember identifiers"}], compaction=operation, **params
|
||||
)
|
||||
message: Final = response if messages_api else response.choices[0].message.model_dump()
|
||||
blocks: Final = message["content"] if messages_api else message["provider_specific_fields"]["compaction_blocks"]
|
||||
assert blocks == [block]
|
||||
if messages_api:
|
||||
assert response["stop_reason"] == "compaction"
|
||||
if not chat_wire:
|
||||
assert response["usage"] == usage
|
||||
if not gateway:
|
||||
replay: Final = {"role": "assistant", "content": blocks} if messages_api else message
|
||||
await call(messages=[replay, {"role": "user", "content": "Continue"}], **params)
|
||||
assert route.call_count == (1 if gateway else 2)
|
||||
|
||||
|
||||
def test_compaction_block_request_transformation():
|
||||
"""
|
||||
Test that compaction blocks from provider_specific_fields are correctly
|
||||
|
|
|
|||
|
|
@ -60,6 +60,20 @@ def test_translate_openai_response_to_anthropic_empty_choices() -> None:
|
|||
assert result["usage"]["input_tokens"] == 10
|
||||
|
||||
|
||||
@pytest.mark.parametrize("text,count,expected_stop", [
|
||||
("", 1, "compaction"), (None, 1, "compaction"), ("Answer", 1, "max_tokens"),
|
||||
(" ", 1, "max_tokens"), ("", 2, "max_tokens"), ("", 0, "max_tokens"),
|
||||
])
|
||||
def test_native_compaction_response_roundtrip(text: str | None, count: int, expected_stop: str) -> None:
|
||||
block: Final = {"type": "compaction", "content": "Exact summary", "signature": "opaque-signature"}
|
||||
message: Final = Message(content=text, provider_specific_fields={"compaction_blocks": [block] * count})
|
||||
response: Final = ModelResponse(choices=[Choices(message=message, finish_reason="length")], usage=Usage())
|
||||
result: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response)
|
||||
expected_text: Final = [{"type": "text", "text": text}] if text is not None and (text != "" or not count) else []
|
||||
assert result["content"] == [*([block] * count), *expected_text]
|
||||
assert result["stop_reason"] == expected_stop
|
||||
|
||||
|
||||
def test_translate_chat_refusal_to_anthropic_response():
|
||||
response = ModelResponse(
|
||||
id="chatcmpl-refusal",
|
||||
|
|
|
|||
53
tests/test_litellm/llms/custom_httpx/test_asgi_handler.py
Normal file
53
tests/test_litellm/llms/custom_httpx/test_asgi_handler.py
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
import asyncio
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from starlette.applications import Starlette
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import JSONResponse, RedirectResponse
|
||||
from starlette.routing import Route
|
||||
|
||||
from litellm.llms.custom_httpx.asgi_handler import get_async_asgi_client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cached_client_isolates_concurrent_apps_and_request_credentials() -> None:
|
||||
ready: Final = (asyncio.Event(), asyncio.Event())
|
||||
|
||||
async def call(index: int) -> Mapping[str, object]:
|
||||
async def endpoint(request: Request) -> JSONResponse:
|
||||
ready[index].set()
|
||||
await ready[1 - index].wait()
|
||||
assert request.scope["root_path"] == f"/gateway-{index}"
|
||||
assert request.client == (f"192.0.2.{index + 1}", 4321)
|
||||
assert request.headers["authorization"] == f"Bearer key-{index}"
|
||||
return JSONResponse({"app": index}, headers={"set-cookie": f"session=app-{index}; Path=/"})
|
||||
|
||||
app: Final = Starlette(routes=[Route("/child", endpoint, methods=["POST"])])
|
||||
with get_async_asgi_client(app, f"/gateway-{index}", (f"192.0.2.{index + 1}", 4321)) as client:
|
||||
response: Final = await client.post(
|
||||
f"https://proxy.test/gateway-{index}/child", headers={"authorization": f"Bearer key-{index}"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert not client.cookies
|
||||
with get_async_asgi_client(app) as reused:
|
||||
assert reused is client
|
||||
return response.json()
|
||||
|
||||
results: Final = await asyncio.wait_for(asyncio.gather(call(0), call(1)), timeout=5)
|
||||
assert results == [{"app": 0}, {"app": 1}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_client_does_not_follow_redirects_or_environment_proxies(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HTTPS_PROXY", "http://unreachable.invalid:8080")
|
||||
|
||||
async def endpoint(request: Request) -> RedirectResponse:
|
||||
return RedirectResponse("https://external.invalid/credentials")
|
||||
|
||||
app: Final = Starlette(routes=[Route("/redirect", endpoint, methods=["POST"])])
|
||||
with get_async_asgi_client(app) as client:
|
||||
response: Final = await client.post("https://proxy.test/redirect", headers={"authorization": "Bearer fixture"})
|
||||
assert response.status_code == 307
|
||||
assert not response.history
|
||||
|
|
@ -14470,6 +14470,21 @@ async def test_request_selected_during_guardrail_runs_concurrently_with_tool(mon
|
|||
assert result.content[0].text == "executed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("in_config,in_db,expected", [(True, False, True), (False, True, False), (True, True, False)])
|
||||
async def test_server_response_identifies_read_only_config(in_config, in_db, expected):
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(server_id="source-server", name="source_server", transport=MCPTransport.http)
|
||||
manager.config_mcp_servers = {server.server_id: server} if in_config else {}
|
||||
manager.registry = {server.server_id: server} if in_db else {}
|
||||
|
||||
listed = await manager.get_all_mcp_servers_unfiltered()
|
||||
|
||||
assert len(listed) == 1
|
||||
assert listed[0].model_dump().get("is_config") is expected
|
||||
assert manager._build_mcp_server_table(server).model_dump().get("is_config") is expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("with_caller,legacy_factory", [(True, False), (False, False), (True, True)])
|
||||
async def test_client_sampling_does_not_fill_explicit_context_from_another_ambient_caller(with_caller, legacy_factory):
|
||||
|
|
|
|||
|
|
@ -28,10 +28,12 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
|
||||
from litellm.proxy.auth.auth_checks import TeamNotFoundError
|
||||
from litellm.proxy.auth.handle_jwt import (
|
||||
JWKS_FETCH_ATTEMPTS,
|
||||
STALE_CACHE_KEY_PREFIX,
|
||||
STALE_WRITTEN_AT_CACHE_KEY_PREFIX,
|
||||
HeaderTeam,
|
||||
JWKSUnreachableError,
|
||||
JWTAuthManager,
|
||||
JWTHandler,
|
||||
|
|
@ -1993,29 +1995,41 @@ async def test_auth_builder_oidc_enabled_falls_back_to_jwt_auth_for_jwt_tokens()
|
|||
assert result["user_object"] == user_object
|
||||
|
||||
|
||||
def test_get_team_id_from_header():
|
||||
"""Test get_team_id_from_header returns team when valid, None when missing, raises on invalid."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
# Valid team in allowed list
|
||||
result = JWTAuthManager.get_team_id_from_header(
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_from_header_returns_allowed_id_none_without_header_and_403_on_invalid():
|
||||
"""Without a DB, x-litellm-team-id resolves to the team when it names an allowed
|
||||
team id, to None when the header is absent, and to a 403 for any other value."""
|
||||
allowed = await JWTAuthManager.resolve_team_from_header(
|
||||
request_headers={"x-litellm-team-id": "team-1"},
|
||||
allowed_team_ids={"team-1", "team-2"},
|
||||
fallback_to_db_teams=False,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=MagicMock(),
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
assert result == "team-1"
|
||||
assert allowed == HeaderTeam(header_value="team-1", team_id="team-1")
|
||||
|
||||
# No header returns None
|
||||
result = JWTAuthManager.get_team_id_from_header(
|
||||
absent = await JWTAuthManager.resolve_team_from_header(
|
||||
request_headers={"authorization": "Bearer token"},
|
||||
allowed_team_ids={"team-1"},
|
||||
fallback_to_db_teams=False,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=MagicMock(),
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
assert result is None
|
||||
assert absent is None
|
||||
|
||||
# Invalid team raises 403
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
JWTAuthManager.get_team_id_from_header(
|
||||
await JWTAuthManager.resolve_team_from_header(
|
||||
request_headers={"x-litellm-team-id": "invalid-team"},
|
||||
allowed_team_ids={"team-1", "team-2"},
|
||||
fallback_to_db_teams=False,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=MagicMock(),
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
|
||||
|
|
@ -5309,32 +5323,23 @@ def test_build_decode_kwargs_warns_for_unscoped_global_fallback_in_mixed_deploym
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_get_team_id_from_header_defers_to_db_membership_only_without_jwt_claims():
|
||||
"""With fallback_to_db_teams=True, an x-litellm-team-id header is accepted
|
||||
provisionally only when the JWT carries no team claims (allowed set empty).
|
||||
When the JWT does carry team claims, the header must still be validated
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_from_header_defers_to_db_membership_only_without_jwt_claims():
|
||||
"""With fallback_to_db_teams=True, an x-litellm-team-id header naming an existing
|
||||
team is accepted provisionally only when the JWT carries no team claims (allowed
|
||||
set empty). When the JWT does carry team claims, the header must still be validated
|
||||
against them, and the flag-off behavior must keep rejecting unknown teams."""
|
||||
deferred = JWTAuthManager.get_team_id_from_header(
|
||||
request_headers={"x-litellm-team-id": "team-from-db"},
|
||||
allowed_team_ids=set(),
|
||||
fallback_to_db_teams=True,
|
||||
)
|
||||
assert deferred == "team-from-db"
|
||||
known_ids = frozenset({"team-from-db"})
|
||||
|
||||
deferred, _, _ = await _resolve_header("team-from-db", set(), True, _teams_by_id(known_ids), _team_alias_lookup_404)
|
||||
assert deferred == HeaderTeam(header_value="team-from-db", team_id="team-from-db")
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
JWTAuthManager.get_team_id_from_header(
|
||||
request_headers={"x-litellm-team-id": "team-x"},
|
||||
allowed_team_ids={"team-1", "team-2"},
|
||||
fallback_to_db_teams=True,
|
||||
)
|
||||
await _resolve_header("team-x", {"team-1", "team-2"}, True, _teams_by_id(known_ids), _team_alias_lookup_404)
|
||||
assert exc_info.value.status_code == 403
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
JWTAuthManager.get_team_id_from_header(
|
||||
request_headers={"x-litellm-team-id": "team-from-db"},
|
||||
allowed_team_ids=set(),
|
||||
fallback_to_db_teams=False,
|
||||
)
|
||||
await _resolve_header("team-from-db", set(), False, _teams_by_id(known_ids), _team_alias_lookup_404)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -5788,6 +5793,7 @@ def test_validate_header_team_in_db_membership_does_not_leak_team_ids():
|
|||
JWTAuthManager._validate_header_team_in_db_membership(
|
||||
team_id="outsider_team",
|
||||
user_object=user_object,
|
||||
header_value="outsider_team",
|
||||
)
|
||||
|
||||
detail = exc_info.value.detail
|
||||
|
|
@ -5797,6 +5803,43 @@ def test_validate_header_team_in_db_membership_does_not_leak_team_ids():
|
|||
assert "outsider_team" in detail
|
||||
|
||||
|
||||
async def _team_lookup_404(team_id, **kwargs):
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail=f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call.",
|
||||
)
|
||||
|
||||
|
||||
async def _team_alias_lookup_404(team_alias, **kwargs):
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"Team with alias '{team_alias}' doesn't exist in db. Create team via `/team/new` call."},
|
||||
)
|
||||
|
||||
|
||||
def _teams_by_alias(aliases: Mapping[str, str]):
|
||||
"""An alias lookup over `aliases` (alias -> team_id) that 404s like the real one otherwise."""
|
||||
|
||||
async def lookup(team_alias, **kwargs):
|
||||
if team_alias not in aliases:
|
||||
return await _team_alias_lookup_404(team_alias)
|
||||
return LiteLLM_TeamTable(team_id=aliases[team_alias], team_alias=team_alias)
|
||||
|
||||
return lookup
|
||||
|
||||
|
||||
def _teams_by_id(team_ids: frozenset[str]):
|
||||
"""A team lookup that knows exactly `team_ids` and, like the real one, reports
|
||||
any other id as provably absent."""
|
||||
|
||||
async def lookup(team_id, **kwargs):
|
||||
if team_id not in team_ids:
|
||||
raise TeamNotFoundError(team_id=team_id)
|
||||
return LiteLLM_TeamTable(team_id=team_id)
|
||||
|
||||
return lookup
|
||||
|
||||
|
||||
async def _run_auth_builder_with_header_team(
|
||||
jwt_auth_config: LiteLLM_JWTAuth,
|
||||
token: dict,
|
||||
|
|
@ -5804,6 +5847,8 @@ async def _run_auth_builder_with_header_team(
|
|||
user_object: LiteLLM_UserTable,
|
||||
fake_get_team,
|
||||
allowed_team_ids: set,
|
||||
fake_get_team_by_alias=_team_alias_lookup_404,
|
||||
route: str = "/chat/completions",
|
||||
):
|
||||
jwt_handler = JWTHandler()
|
||||
jwt_handler.litellm_jwtauth = jwt_auth_config
|
||||
|
|
@ -5848,14 +5893,19 @@ async def _run_auth_builder_with_header_team(
|
|||
new_callable=AsyncMock,
|
||||
side_effect=fake_get_team,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.auth.handle_jwt.get_team_object_by_alias",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=fake_get_team_by_alias,
|
||||
),
|
||||
):
|
||||
return await JWTAuthManager.auth_builder(
|
||||
api_key="test_jwt_token",
|
||||
jwt_handler=jwt_handler,
|
||||
request_data={"model": "gpt-4"},
|
||||
general_settings={"enforce_rbac": False},
|
||||
route="/chat/completions",
|
||||
prisma_client=None,
|
||||
route=route,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=None,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
|
|
@ -5863,13 +5913,6 @@ async def _run_auth_builder_with_header_team(
|
|||
)
|
||||
|
||||
|
||||
async def _team_lookup_404(team_id, **kwargs):
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail=f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call.",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_builder_header_team_not_found_matches_non_membership_denial() -> (
|
||||
None
|
||||
|
|
@ -6942,6 +6985,273 @@ async def test_auth_builder_header_team_enforces_team_allowed_routes_under_db_fa
|
|||
assert result["team_id"] == header_team
|
||||
|
||||
|
||||
async def _resolve_header(
|
||||
header_value: str,
|
||||
allowed_team_ids: set[str],
|
||||
fallback_to_db_teams: bool,
|
||||
fake_get_team,
|
||||
fake_get_team_by_alias,
|
||||
) -> tuple[HeaderTeam | None, AsyncMock, AsyncMock]:
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.auth.handle_jwt.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=fake_get_team,
|
||||
) as by_id,
|
||||
patch(
|
||||
"litellm.proxy.auth.handle_jwt.get_team_object_by_alias",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=fake_get_team_by_alias,
|
||||
) as by_alias,
|
||||
):
|
||||
resolved = await JWTAuthManager.resolve_team_from_header(
|
||||
request_headers={"X-LiteLLM-Team-Id": header_value},
|
||||
allowed_team_ids=allowed_team_ids,
|
||||
fallback_to_db_teams=fallback_to_db_teams,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
return resolved, by_id, by_alias
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_from_header_accepts_the_alias_of_an_allowed_team():
|
||||
"""x-litellm-team-id may carry the team alias instead of the team id (LIT-7181).
|
||||
The alias resolves to its team id before the allowed-teams check, so a
|
||||
caller whose JWT grants team_a gets team_a whether it sends the id or the
|
||||
alias, and the id path never pays for an alias lookup."""
|
||||
aliases = {"alias_a": "team_a", "alias_b": "team_b"}
|
||||
|
||||
by_alias_value, lookups_by_id, lookups_by_alias = await _resolve_header(
|
||||
"alias_a", {"team_a"}, False, _team_lookup_404, _teams_by_alias(aliases)
|
||||
)
|
||||
assert by_alias_value == HeaderTeam(header_value="alias_a", team_id="team_a")
|
||||
lookups_by_alias.assert_awaited_once()
|
||||
assert lookups_by_alias.await_args.kwargs["team_alias"] == "alias_a"
|
||||
lookups_by_id.assert_not_awaited()
|
||||
|
||||
by_id_value, lookups_by_id, lookups_by_alias = await _resolve_header(
|
||||
"team_a", {"team_a"}, False, _team_lookup_404, _teams_by_alias(aliases)
|
||||
)
|
||||
assert by_id_value == HeaderTeam(header_value="team_a", team_id="team_a")
|
||||
lookups_by_alias.assert_not_awaited()
|
||||
lookups_by_id.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_from_header_denies_aliases_of_teams_the_jwt_does_not_grant():
|
||||
"""An alias that exists but names a team outside the JWT's allowed teams is
|
||||
refused with the same 403 as an unknown value, and the detail names only
|
||||
what the caller sent, so the response reveals neither that the alias exists
|
||||
nor which team id it maps to."""
|
||||
aliases = {"alias_a": "team_a", "alias_b": "team_b"}
|
||||
|
||||
with pytest.raises(HTTPException) as other_team:
|
||||
await _resolve_header("alias_b", {"team_a"}, False, _team_lookup_404, _teams_by_alias(aliases))
|
||||
with pytest.raises(HTTPException) as unknown:
|
||||
await _resolve_header("no_such", {"team_a"}, False, _team_lookup_404, _teams_by_alias(aliases))
|
||||
|
||||
assert other_team.value.status_code == 403
|
||||
assert unknown.value.status_code == 403
|
||||
assert "team_b" not in other_team.value.detail
|
||||
assert other_team.value.detail.replace("alias_b", "<value>") == unknown.value.detail.replace("no_such", "<value>")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_from_header_under_db_fallback_tries_the_id_before_the_alias():
|
||||
"""Under fallback_to_db_teams a claimless JWT's header is provisional: a value
|
||||
that is an existing team id resolves to itself without an alias lookup, a
|
||||
value that is only an alias resolves to that team's id, and a value that is
|
||||
neither gets the membership denial the id path already uses."""
|
||||
known_ids = frozenset({"team_a"})
|
||||
aliases = {"alias_a": "team_a"}
|
||||
|
||||
as_id, _, lookups_by_alias = await _resolve_header(
|
||||
"team_a", set(), True, _teams_by_id(known_ids), _teams_by_alias(aliases)
|
||||
)
|
||||
assert as_id == HeaderTeam(header_value="team_a", team_id="team_a")
|
||||
lookups_by_alias.assert_not_awaited()
|
||||
|
||||
as_alias, _, _ = await _resolve_header("alias_a", set(), True, _teams_by_id(known_ids), _teams_by_alias(aliases))
|
||||
assert as_alias == HeaderTeam(header_value="alias_a", team_id="team_a")
|
||||
|
||||
with pytest.raises(HTTPException) as neither:
|
||||
await _resolve_header("ghost", set(), True, _teams_by_id(known_ids), _teams_by_alias(aliases))
|
||||
assert neither.value.status_code == 403
|
||||
assert neither.value.detail == ("Team 'ghost' (from x-litellm-team-id header) is not in your team memberships.")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_from_header_under_db_fallback_never_aliases_a_team_id_it_could_not_read():
|
||||
"""Only a team row the database provably lacks falls through to the alias
|
||||
lookup. When the id read fails for any other reason (the generic 404 the
|
||||
team lookup uses for an unreadable database) the value keeps the id path's
|
||||
membership denial, so an outage can never turn a team id into the team
|
||||
that happens to carry it as an alias."""
|
||||
aliases = {"team_a": "team_b"}
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, side_effect=_team_lookup_404),
|
||||
patch(
|
||||
"litellm.proxy.auth.handle_jwt.get_team_object_by_alias",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=_teams_by_alias(aliases),
|
||||
) as lookups_by_alias,
|
||||
pytest.raises(HTTPException) as unreadable,
|
||||
):
|
||||
await JWTAuthManager.resolve_team_from_header(
|
||||
request_headers={"x-litellm-team-id": "team_a"},
|
||||
allowed_team_ids=set(),
|
||||
fallback_to_db_teams=True,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert unreadable.value.status_code == 403
|
||||
assert unreadable.value.detail == ("Team 'team_a' (from x-litellm-team-id header) is not in your team memberships.")
|
||||
lookups_by_alias.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_from_header_treats_a_duplicate_alias_as_no_match_but_surfaces_lookup_errors():
|
||||
"""An alias two teams share cannot name one team, so it is refused like an
|
||||
unknown value (a 4xx from the lookup is a miss), while a lookup failure
|
||||
(5xx) is not disguised as a denial and propagates as is."""
|
||||
|
||||
async def duplicate_alias(team_alias, **kwargs):
|
||||
raise HTTPException(status_code=400, detail={"error": f"Multiple teams found with alias '{team_alias}'."})
|
||||
|
||||
async def db_down(team_alias, **kwargs):
|
||||
raise HTTPException(status_code=500, detail={"error": f"Error looking up team by alias '{team_alias}'"})
|
||||
|
||||
with pytest.raises(HTTPException) as duplicate:
|
||||
await _resolve_header("shared_alias", {"team_a"}, False, _team_lookup_404, duplicate_alias)
|
||||
assert duplicate.value.status_code == 403
|
||||
assert "Multiple teams" not in str(duplicate.value.detail)
|
||||
|
||||
with pytest.raises(HTTPException) as failure:
|
||||
await _resolve_header("alias_a", {"team_a"}, False, _team_lookup_404, db_down)
|
||||
assert failure.value.status_code == 500
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_builder_header_alias_binds_the_aliased_team_under_claims_and_db_fallback():
|
||||
"""End to end through auth_builder, x-litellm-team-id carrying a team alias
|
||||
binds the request to the aliased team (result team_id is the canonical id)
|
||||
both when the JWT grants that team by claim and when a claimless JWT relies
|
||||
on fallback_to_db_teams and DB membership."""
|
||||
aliases = {"alias_member": "team_member"}
|
||||
user_object = LiteLLM_UserTable(
|
||||
user_id="u_alias",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
teams=["team_member"],
|
||||
)
|
||||
|
||||
claims_config = LiteLLM_JWTAuth(team_ids_jwt_field="team_ids")
|
||||
by_claim = await _run_auth_builder_with_header_team(
|
||||
claims_config,
|
||||
{"sub": "u_alias", "scope": "", "team_ids": ["team_member"]},
|
||||
"alias_member",
|
||||
user_object,
|
||||
_teams_by_id(frozenset({"team_member"})),
|
||||
{"team_member"},
|
||||
_teams_by_alias(aliases),
|
||||
)
|
||||
assert by_claim["team_id"] == "team_member"
|
||||
assert by_claim["team_object"].team_id == "team_member"
|
||||
|
||||
fallback_config = LiteLLM_JWTAuth(fallback_to_db_teams=True)
|
||||
by_membership = await _run_auth_builder_with_header_team(
|
||||
fallback_config,
|
||||
{"sub": "u_alias", "scope": ""},
|
||||
"alias_member",
|
||||
user_object,
|
||||
_teams_by_id(frozenset({"team_member"})),
|
||||
set(),
|
||||
_teams_by_alias(aliases),
|
||||
)
|
||||
assert by_membership["team_id"] == "team_member"
|
||||
assert by_membership["team_object"].team_id == "team_member"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_builder_header_alias_of_a_non_member_team_is_denied_like_an_unknown_value_under_db_fallback():
|
||||
"""Under fallback_to_db_teams, an alias naming a team the user is not a member
|
||||
of is denied with the exact same 403 as an unknown value, naming the alias
|
||||
the caller sent rather than the team id it resolved to."""
|
||||
aliases = {"alias_other": "team_other"}
|
||||
known_ids = frozenset({"team_member", "team_other"})
|
||||
user_object = LiteLLM_UserTable(
|
||||
user_id="u_alias_outsider",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
teams=["team_member"],
|
||||
)
|
||||
config = LiteLLM_JWTAuth(fallback_to_db_teams=True)
|
||||
token = {"sub": "u_alias_outsider", "scope": ""}
|
||||
|
||||
with pytest.raises(HTTPException) as outsider_alias:
|
||||
await _run_auth_builder_with_header_team(
|
||||
config, token, "alias_other", user_object, _teams_by_id(known_ids), set(), _teams_by_alias(aliases)
|
||||
)
|
||||
with pytest.raises(HTTPException) as unknown:
|
||||
await _run_auth_builder_with_header_team(
|
||||
config, token, "alias_ghost", user_object, _teams_by_id(known_ids), set(), _teams_by_alias(aliases)
|
||||
)
|
||||
|
||||
assert outsider_alias.value.status_code == 403
|
||||
assert unknown.value.status_code == 403
|
||||
assert "team_other" not in outsider_alias.value.detail
|
||||
assert outsider_alias.value.detail.replace("alias_other", "<value>") == unknown.value.detail.replace(
|
||||
"alias_ghost", "<value>"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_builder_header_alias_under_db_fallback_keeps_the_team_allowed_routes_gate():
|
||||
"""Under fallback_to_db_teams, a member team selected by alias is still held
|
||||
to team_allowed_routes, and the denial names the alias the caller sent."""
|
||||
aliases = {"alias_member": "team_member"}
|
||||
user_object = LiteLLM_UserTable(
|
||||
user_id="u_alias_routes",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
teams=["team_member"],
|
||||
)
|
||||
config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_allowed_routes=["openai_routes"])
|
||||
token = {"sub": "u_alias_routes", "scope": ""}
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _run_auth_builder_with_header_team(
|
||||
config,
|
||||
token,
|
||||
"alias_member",
|
||||
user_object,
|
||||
_teams_by_id(frozenset({"team_member"})),
|
||||
set(),
|
||||
_teams_by_alias(aliases),
|
||||
route="/key/info",
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == (
|
||||
"Team 'alias_member' (from x-litellm-team-id header) is not allowed to access route '/key/info'."
|
||||
)
|
||||
|
||||
allowed = await _run_auth_builder_with_header_team(
|
||||
config,
|
||||
token,
|
||||
"alias_member",
|
||||
user_object,
|
||||
_teams_by_id(frozenset({"team_member"})),
|
||||
set(),
|
||||
_teams_by_alias(aliases),
|
||||
route="/chat/completions",
|
||||
)
|
||||
assert allowed["team_id"] == "team_member"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_flag():
|
||||
"""Reading the singular team claim during sync is scoped to fallback_to_db_teams.
|
||||
|
|
|
|||
|
|
@ -6183,10 +6183,10 @@ async def test_configured_estimate_blocks_the_overrun_the_static_floor_admits(mo
|
|||
assert await admitted({"default_estimated_output_tokens": 3000}) == 2
|
||||
|
||||
|
||||
def test_internal_call_origin_success_ops_are_skipped():
|
||||
"""Internal sub-calls (auto-router classifier, shadow eval shadow/judge) bill spend
|
||||
to the caller's key but must not consume its TPM counters: the same kwargs charge
|
||||
ops without the origin stamp and none with it."""
|
||||
@pytest.mark.parametrize("origin", ["shadow_eval_judge", "autorouter_compaction"])
|
||||
@pytest.mark.parametrize("rate_limit_type", ["input", "output", "total"])
|
||||
def test_internal_call_origin_success_ops_are_skipped(origin, rate_limit_type):
|
||||
"""Foreground compaction charges the same scopes as ordinary caller traffic."""
|
||||
handler = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(DualCache())
|
||||
)
|
||||
|
|
@ -6202,23 +6202,27 @@ def test_internal_call_origin_success_ops_are_skipped():
|
|||
def _kwargs(metadata: Dict[str, Any]) -> Dict[str, Any]:
|
||||
return {
|
||||
"standard_logging_object": {
|
||||
"metadata": {"user_api_key_hash": hash_token("sk-internal-origin")}
|
||||
"metadata": {
|
||||
"user_api_key_hash": hash_token("sk-internal-origin"),
|
||||
"user_api_key_team_id": "compaction-team",
|
||||
"user_api_key_project_id": "compaction-project",
|
||||
}
|
||||
},
|
||||
"litellm_params": {"metadata": metadata},
|
||||
"model": "gpt-4o-mini",
|
||||
}
|
||||
|
||||
charged = handler._build_success_event_pipeline_operations(
|
||||
kwargs=_kwargs({}), response_obj=response, rate_limit_type="output"
|
||||
kwargs=_kwargs({}), response_obj=response, rate_limit_type=rate_limit_type
|
||||
)
|
||||
skipped = handler._build_success_event_pipeline_operations(
|
||||
kwargs=_kwargs({INTERNAL_CALL_ORIGIN_METADATA_KEY: "shadow_eval_judge"}),
|
||||
kwargs=_kwargs({INTERNAL_CALL_ORIGIN_METADATA_KEY: origin}),
|
||||
response_obj=response,
|
||||
rate_limit_type="output",
|
||||
rate_limit_type=rate_limit_type,
|
||||
)
|
||||
|
||||
assert charged
|
||||
assert skipped == []
|
||||
assert skipped == (charged if origin == "autorouter_compaction" else [])
|
||||
|
||||
|
||||
def _conflicting_budget_bodies() -> Dict[str, Dict[str, object]]:
|
||||
|
|
|
|||
|
|
@ -1,8 +1,10 @@
|
|||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -17,12 +19,15 @@ from litellm.proxy._types import (
|
|||
LiteLLM_UserTableFiltered,
|
||||
LitellmUserRoles,
|
||||
NewUserRequest,
|
||||
ProxyErrorTypes,
|
||||
ProxyException,
|
||||
UpdateUserRequest,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import (
|
||||
LiteLLM_UserTableWithKeyCount,
|
||||
_authorize_user_list_request,
|
||||
_resolve_org_filter_for_user_search,
|
||||
_resolve_user_email_metadata,
|
||||
_update_internal_user_params,
|
||||
get_user_key_counts,
|
||||
|
|
@ -4653,3 +4658,124 @@ async def test_delete_user_evicts_cached_user_rows(mocker: MockerFixture) -> Non
|
|||
assert await cache.async_get_cache(key=deleted.user_id, model_type=LiteLLM_UserTable) is None
|
||||
assert await cache.async_get_cache(key=survivor.user_id, model_type=LiteLLM_UserTable) == survivor
|
||||
broadcast.assert_awaited_once_with(cache_key=deleted.user_id)
|
||||
|
||||
|
||||
_DB_OUTAGE_503_BODY: Final = {
|
||||
"error": {
|
||||
"message": "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.",
|
||||
"type": "no_db_connection",
|
||||
"param": "None",
|
||||
"code": "503",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _user_read_raising(mocker: MockerFixture, error: Exception) -> tuple[MagicMock, MagicMock]:
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=error)
|
||||
cache = MagicMock()
|
||||
cache.async_get_cache = AsyncMock(return_value=None)
|
||||
cache.async_set_cache = AsyncMock()
|
||||
mocker.patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True)
|
||||
return prisma_client, cache
|
||||
|
||||
|
||||
def _db_unavailable_fallback_identity(route: str) -> UserAPIKeyAuth:
|
||||
from litellm.proxy.auth.auth_exception_handler import DB_UNAVAILABLE_FALLBACK_USER_ID
|
||||
|
||||
return UserAPIKeyAuth(
|
||||
key_name="failed-to-connect-to-db",
|
||||
token="failed-to-connect-to-db",
|
||||
user_id=DB_UNAVAILABLE_FALLBACK_USER_ID,
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
request_route=route,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authorize_user_list_request_propagates_a_db_outage_instead_of_answering_403(mocker):
|
||||
prisma_client, cache = _user_read_raising(mocker, httpx.ConnectError("All connection attempts failed"))
|
||||
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
await _authorize_user_list_request(
|
||||
user_api_key_dict=_db_unavailable_fallback_identity("/user/list"),
|
||||
organization_ids=None,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_org_filter_for_user_search_propagates_a_db_outage_instead_of_answering_403(mocker):
|
||||
prisma_client, cache = _user_read_raising(mocker, httpx.ConnectError("All connection attempts failed"))
|
||||
mocker.patch(
|
||||
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached",
|
||||
return_value={"scope_user_search_to_org": True},
|
||||
)
|
||||
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
await _resolve_org_filter_for_user_search(
|
||||
user_api_key_dict=_db_unavailable_fallback_identity("/user/filter/ui"),
|
||||
team_id=None,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_users_answers_a_db_outage_as_503_no_db_connection_not_as_its_own_500(mocker, caplog):
|
||||
prisma_client, cache = _user_read_raising(mocker, httpx.ConnectError("All connection attempts failed"))
|
||||
mocker.patch(
|
||||
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached",
|
||||
return_value={"scope_user_search_to_org": True},
|
||||
)
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache)
|
||||
mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised:
|
||||
await ui_view_users(
|
||||
user_api_key_dict=_db_unavailable_fallback_identity("/user/filter/ui"),
|
||||
user_id=None,
|
||||
user_email="lit",
|
||||
team_id=None,
|
||||
page=1,
|
||||
page_size=50,
|
||||
)
|
||||
|
||||
assert raised.value.code == "503"
|
||||
assert raised.value.type == ProxyErrorTypes.no_db_connection
|
||||
assert isinstance(raised.value.__cause__, httpx.ConnectError)
|
||||
outage_logs: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and "ConnectError" in r.getMessage()]
|
||||
assert outage_logs == ["Database unavailable during user search: ConnectError"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "params"),
|
||||
[("/user/list", {}), ("/user/filter/ui", {"user_email": "lit"})],
|
||||
ids=["user_list", "user_filter_ui"],
|
||||
)
|
||||
def test_user_routes_answer_503_no_db_connection_when_the_callers_user_read_hits_a_db_outage(
|
||||
mocker, route: str, params: dict[str, str]
|
||||
):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
prisma_client, cache = _user_read_raising(mocker, httpx.ConnectError("All connection attempts failed"))
|
||||
mocker.patch(
|
||||
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached",
|
||||
return_value={"scope_user_search_to_org": True},
|
||||
)
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: _db_unavailable_fallback_identity(route)
|
||||
try:
|
||||
response = TestClient(app, raise_server_exceptions=False).get(route, params=params)
|
||||
finally:
|
||||
app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 503, response.text
|
||||
assert response.json() == _DB_OUTAGE_503_BODY
|
||||
|
|
|
|||
|
|
@ -7900,3 +7900,43 @@ class TestGetMcpToolsWireShape:
|
|||
assert tool["outputSchema"] == {"type": "integer"}
|
||||
assert "_meta" in tool
|
||||
assert not {"input_schema", "output_schema", "meta"} & tool.keys()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("role,expected_status", [
|
||||
(LitellmUserRoles.PROXY_ADMIN, 404),
|
||||
(LitellmUserRoles.INTERNAL_USER, 403),
|
||||
])
|
||||
async def test_config_server_edit_preserves_api_contract_without_creating_rows(role, expected_status):
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
|
||||
|
||||
manager = MCPServerManager()
|
||||
server = generate_mock_mcp_server_config_record(server_id="read-only-config")
|
||||
manager.config_mcp_servers = {server.server_id: server}
|
||||
original = server.model_dump()
|
||||
prisma = MagicMock()
|
||||
prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None)
|
||||
prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=None)
|
||||
with (
|
||||
patch.object(mgmt_endpoints, "global_mcp_server_manager", manager),
|
||||
patch.object(mgmt_endpoints, "get_prisma_client_or_throw", return_value=prisma),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await mgmt_endpoints.edit_mcp_server(
|
||||
payload=UpdateMCPServerRequest(server_id=server.server_id, description="UI edit"),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="actor", user_role=role),
|
||||
)
|
||||
|
||||
assert exc.value.status_code == expected_status
|
||||
if role == LitellmUserRoles.PROXY_ADMIN:
|
||||
assert exc.value.detail == {
|
||||
"error": f"MCP Server not found, passed server_id={server.server_id}"
|
||||
}
|
||||
prisma.db.litellm_mcpservertable.update.assert_awaited_once()
|
||||
else:
|
||||
prisma.db.litellm_mcpservertable.update.assert_not_awaited()
|
||||
prisma.db.litellm_mcpservertable.create.assert_not_called()
|
||||
prisma.db.litellm_mcpservertable.create_many.assert_not_called()
|
||||
prisma.tx.assert_not_called()
|
||||
assert server.model_dump() == original
|
||||
assert manager.registry == {}
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from collections.abc import Sequence
|
|||
from typing import Final, Optional, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
|
|
@ -42,6 +43,8 @@ from litellm.proxy.management_endpoints.team_endpoints import (
|
|||
_STRIP_DELETED_TEAM_FROM_USERS_SQL,
|
||||
GetTeamMemberPermissionsResponse,
|
||||
UpdateTeamMemberPermissionsRequest,
|
||||
_build_team_list_where_conditions,
|
||||
_get_org_admin_org_ids,
|
||||
_persist_deleted_team_records,
|
||||
_save_deleted_team_records,
|
||||
_transform_teams_to_deleted_records,
|
||||
|
|
@ -16569,3 +16572,81 @@ def test_team_member_update_request_rejects_unusable_temp_budget_increase(increa
|
|||
TeamMemberUpdateRequest(
|
||||
team_id="team-1", user_id="user-1", temp_budget_increase=increase, temp_budget_expiry="2030-01-01T00:00:00Z"
|
||||
)
|
||||
|
||||
|
||||
_DB_OUTAGE_503_BODY: Final = {
|
||||
"error": {
|
||||
"message": "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.",
|
||||
"type": "no_db_connection",
|
||||
"param": "None",
|
||||
"code": "503",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _user_read_raising(error: Exception) -> tuple[MagicMock, MagicMock]:
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=error)
|
||||
cache = MagicMock()
|
||||
cache.async_get_cache = AsyncMock(return_value=None)
|
||||
cache.async_set_cache = AsyncMock()
|
||||
return prisma_client, cache
|
||||
|
||||
|
||||
def _db_unavailable_fallback_identity(route: str) -> UserAPIKeyAuth:
|
||||
from litellm.proxy.auth.auth_exception_handler import DB_UNAVAILABLE_FALLBACK_USER_ID
|
||||
|
||||
return UserAPIKeyAuth(
|
||||
key_name="failed-to-connect-to-db",
|
||||
token="failed-to-connect-to-db",
|
||||
user_id=DB_UNAVAILABLE_FALLBACK_USER_ID,
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
request_route=route,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_org_admin_org_ids_propagates_a_db_outage_instead_of_answering_not_an_org_admin():
|
||||
prisma_client, cache = _user_read_raising(httpx.ConnectError("All connection attempts failed"))
|
||||
|
||||
with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True):
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
await _get_org_admin_org_ids(
|
||||
user_id="outage-probe-user",
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_team_list_where_conditions_propagates_a_db_outage_instead_of_answering_user_not_found():
|
||||
prisma_client, cache = _user_read_raising(httpx.ConnectError("All connection attempts failed"))
|
||||
|
||||
with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True):
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
await _build_team_list_where_conditions(
|
||||
prisma_client=prisma_client,
|
||||
team_id=None,
|
||||
team_alias=None,
|
||||
organization_id=None,
|
||||
user_id="outage-probe-user",
|
||||
use_deleted_table=False,
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
|
||||
|
||||
def test_list_team_v2_answers_503_no_db_connection_when_the_callers_user_read_hits_a_db_outage(monkeypatch):
|
||||
prisma_client, cache = _user_read_raising(httpx.ConnectError("All connection attempts failed"))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: _db_unavailable_fallback_identity("/v2/team/list")
|
||||
try:
|
||||
with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True):
|
||||
response = TestClient(app, raise_server_exceptions=False).get("/v2/team/list")
|
||||
finally:
|
||||
app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 503, response.text
|
||||
assert response.json() == _DB_OUTAGE_503_BODY
|
||||
|
|
|
|||
|
|
@ -11,8 +11,10 @@ from __future__ import annotations
|
|||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
|
|
@ -364,6 +366,39 @@ async def test_otel_unhandled_exception_handler_returns_500_generic_payload():
|
|||
}
|
||||
|
||||
|
||||
_DB_OUTAGE_503_BODY: Final = {
|
||||
"error": {
|
||||
"message": "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.",
|
||||
"type": "no_db_connection",
|
||||
"param": "None",
|
||||
"code": "503",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _raised_from(outer: Exception, cause: Exception) -> Exception:
|
||||
try:
|
||||
raise outer from cause
|
||||
except Exception as chained:
|
||||
return chained
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"exc",
|
||||
[
|
||||
httpx.ConnectError("All connection attempts failed"),
|
||||
_raised_from(RuntimeError("user read failed"), httpx.ConnectError("All connection attempts failed")),
|
||||
],
|
||||
ids=["raw_connect_error", "connect_error_as_cause"],
|
||||
)
|
||||
async def test_otel_unhandled_exception_handler_answers_a_db_outage_with_503_no_db_connection(exc):
|
||||
response = await otel_unhandled_exception_handler(request=_make_request(path="/v2/team/list"), exc=exc)
|
||||
|
||||
assert response.status_code == 503
|
||||
assert json.loads(response.body) == _DB_OUTAGE_503_BODY
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_otel_unhandled_exception_handler_reraises_proxy_exception_error():
|
||||
"""ProxyException / HTTPException / RequestValidationError are re-raised
|
||||
|
|
|
|||
|
|
@ -818,6 +818,35 @@ async def test_proxy_startup_event_prunes_dead_workers_live_gauges(tmp_path):
|
|||
assert counter.exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("disable_model_info_refresh", "job_scheduled"), [(True, False), (False, True)])
|
||||
async def test_proxy_startup_event_honors_disable_model_info_refresh(
|
||||
disable_model_info_refresh: bool, job_scheduled: bool
|
||||
) -> None:
|
||||
"""``general_settings.disable_model_info_refresh: true`` keeps the proxy from polling every
|
||||
OpenAI-compatible deployment's ``/v1/models`` in the background, so a proxy fronting a replay
|
||||
fixture (or a metered upstream) makes only the calls its clients asked for."""
|
||||
scheduler = AsyncIOScheduler()
|
||||
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} | {
|
||||
"LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY": "true"
|
||||
}
|
||||
with (
|
||||
patch.dict(os.environ, clean_env, clear=True),
|
||||
patch.object(ps, "scheduler", scheduler),
|
||||
patch.dict(ps.general_settings, {"disable_model_info_refresh": disable_model_info_refresh}),
|
||||
):
|
||||
try:
|
||||
async with proxy_startup_event(app=None):
|
||||
job = scheduler.get_job("refresh_model_info")
|
||||
finally:
|
||||
if scheduler.running:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
assert (job is not None) is job_scheduled, (
|
||||
f"disable_model_info_refresh={disable_model_info_refresh} but refresh_model_info job is {job}"
|
||||
)
|
||||
|
||||
|
||||
def test_otel_global_provider_published_after_callback_init():
|
||||
"""The OTel V2 global-provider publish must run after callback
|
||||
initialization in ``proxy_startup_event``.
|
||||
|
|
|
|||
154
tests/test_litellm/proxy/test_bug_report_config.py
Normal file
154
tests/test_litellm/proxy/test_bug_report_config.py
Normal file
|
|
@ -0,0 +1,154 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator, Mapping
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.bug_report_config import build_proxy_bug_report, safe_config_lines
|
||||
|
||||
CUSTOMER_STRINGS = (
|
||||
"acme",
|
||||
"sk-live-secret",
|
||||
"hunter2",
|
||||
"postgres://",
|
||||
"10.0.0.7",
|
||||
)
|
||||
|
||||
CUSTOMER_CONFIG: Mapping[str, object] = {
|
||||
"model_list": [
|
||||
{
|
||||
"model_name": "acme-prod-gpt4",
|
||||
"litellm_params": {
|
||||
"model": "azure/acme-gpt4o-deployment",
|
||||
"api_base": "https://acme-eastus.openai.azure.com",
|
||||
"api_key": "sk-live-secret-1",
|
||||
"rpm": 600,
|
||||
"acme_extra_param": "acme",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "acme-mini",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-live-secret-2"},
|
||||
},
|
||||
{
|
||||
"model_name": "acme-backup",
|
||||
"litellm_params": {"model": "azure/acme-backup-deployment", "api_key": "sk-live-secret-3"},
|
||||
},
|
||||
{"model_name": "acme-bare", "litellm_params": {"model": "acme-custom-model"}},
|
||||
],
|
||||
"litellm_settings": {
|
||||
"callbacks": ["langfuse", "acme_hooks.audit_logger"],
|
||||
"drop_params": True,
|
||||
"num_retries": 3,
|
||||
"acme_internal_flag": True,
|
||||
"cache": True,
|
||||
"cache_params": {
|
||||
"type": "redis",
|
||||
"host": "10.0.0.7",
|
||||
"port": 6379,
|
||||
"password": "hunter2",
|
||||
"acme_cache_option": "acme",
|
||||
},
|
||||
},
|
||||
"router_settings": {
|
||||
"routing_strategy": "latency-based-routing",
|
||||
"redis_host": "10.0.0.7",
|
||||
"acme_router_option": True,
|
||||
},
|
||||
"guardrails": [
|
||||
{"guardrail_name": "acme-pii-mask", "litellm_params": {"guardrail": "presidio", "mode": "pre_call"}},
|
||||
{
|
||||
"guardrail_name": "acme-policy",
|
||||
"litellm_params": {"guardrail": "acme_guardrails.PolicyCheck", "api_key": "sk-live-secret-4"},
|
||||
},
|
||||
],
|
||||
"environment_variables": {"ACME_PROD_OPENAI_KEY": "sk-live-secret-5", "ACME_TENANT": "acme"},
|
||||
}
|
||||
|
||||
CUSTOMER_GENERAL_SETTINGS: Mapping[str, object] = {
|
||||
"master_key": "sk-live-secret-master",
|
||||
"database_url": "postgres://user:hunter2@10.0.0.7/litellm",
|
||||
"key_management_system": "aws_secret_manager",
|
||||
"store_model_in_db": True,
|
||||
"health_check_interval": 300,
|
||||
"acme_sso_tenant": "acme-prod",
|
||||
}
|
||||
|
||||
|
||||
def test_safe_config_lines_keep_only_flags_and_litellm_defined_values():
|
||||
lines = safe_config_lines(CUSTOMER_CONFIG, CUSTOMER_GENERAL_SETTINGS)
|
||||
|
||||
assert lines == (
|
||||
"general_settings.key_management_system = aws_secret_manager",
|
||||
"general_settings.store_model_in_db = true",
|
||||
"litellm_settings.callbacks = [langfuse]",
|
||||
"litellm_settings.drop_params = true",
|
||||
"litellm_settings.cache = true",
|
||||
"litellm_settings.cache_params.type = redis",
|
||||
"router_settings.routing_strategy = latency-based-routing",
|
||||
"guardrails[0].litellm_params.guardrail = presidio",
|
||||
"guardrails[0].litellm_params.mode = pre_call",
|
||||
"model_list[*].provider = [azure, openai]",
|
||||
)
|
||||
assert not any(customer_string in "\n".join(lines) for customer_string in CUSTOMER_STRINGS)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("config", "expected_lines"),
|
||||
[
|
||||
({"router_settings": {"routing_strategy": "acme-strategy"}}, ()),
|
||||
({"litellm_settings": {"cache_params": {"type": "acme-cache"}}}, ()),
|
||||
(
|
||||
{"litellm_settings": {"success_callback": ["acme_logger", "langsmith"]}},
|
||||
("litellm_settings.success_callback = [langsmith]",),
|
||||
),
|
||||
({"litellm_settings": {"callbacks": ["acme_hooks.audit_logger"]}}, ()),
|
||||
],
|
||||
)
|
||||
def test_string_values_show_only_when_litellm_defines_them(
|
||||
config: Mapping[str, object], expected_lines: tuple[str, ...]
|
||||
):
|
||||
assert safe_config_lines(config, {}) == expected_lines
|
||||
|
||||
|
||||
def test_secrets_numbers_and_unknown_values_leave_no_line():
|
||||
general_settings: Mapping[str, object] = {
|
||||
"master_key": "sk-live-secret-master",
|
||||
"health_check_interval": 300,
|
||||
"store_model_in_db": object(),
|
||||
"alerting": {"acme": "webhook"},
|
||||
"background_health_checks": False,
|
||||
}
|
||||
|
||||
assert safe_config_lines({}, general_settings) == ("general_settings.background_health_checks = false",)
|
||||
|
||||
|
||||
def test_malformed_sections_produce_no_lines():
|
||||
config: Mapping[str, object] = {
|
||||
"litellm_settings": "acme",
|
||||
"router_settings": ["acme"],
|
||||
"guardrails": {"acme": {"litellm_params": {"guardrail": "presidio"}}},
|
||||
"model_list": "acme",
|
||||
"environment_variables": None,
|
||||
}
|
||||
|
||||
assert safe_config_lines(config, {}) == ()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def loaded_proxy_config(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
|
||||
previous_config = proxy_server.proxy_config.get_config_state()
|
||||
proxy_server.proxy_config.update_config_state(config=CUSTOMER_CONFIG)
|
||||
monkeypatch.setattr(proxy_server, "general_settings", dict(CUSTOMER_GENERAL_SETTINGS))
|
||||
yield
|
||||
proxy_server.proxy_config.update_config_state(config=previous_config)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("loaded_proxy_config")
|
||||
def test_build_proxy_bug_report_reads_the_loaded_proxy_config():
|
||||
report = build_proxy_bug_report(RuntimeError("boom"), stream=False)
|
||||
|
||||
assert report.surface == "proxy"
|
||||
assert report.stream is False
|
||||
assert report.config_lines == safe_config_lines(CUSTOMER_CONFIG, CUSTOMER_GENERAL_SETTINGS)
|
||||
|
|
@ -4,6 +4,7 @@ import datetime
|
|||
import json
|
||||
from types import MappingProxyType, SimpleNamespace
|
||||
from typing import AsyncGenerator, Callable, Final, Iterator, Optional, Sequence
|
||||
from urllib.parse import unquote_plus
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -13,6 +14,12 @@ from fastapi.responses import JSONResponse, StreamingResponse
|
|||
|
||||
import litellm
|
||||
from litellm._uuid import uuid
|
||||
from litellm.litellm_core_utils.bug_report import (
|
||||
DISABLE_ENV_VAR,
|
||||
ISSUE_URL_BASE,
|
||||
bug_report_notice,
|
||||
build_bug_report,
|
||||
)
|
||||
from litellm.constants import (
|
||||
CLIENT_REQUESTED_MODEL_SCOPE_KEY,
|
||||
MAX_LITELLM_CALL_ID_LENGTH,
|
||||
|
|
@ -4290,6 +4297,21 @@ class TestHandleLLMApiExceptionRetryAfter:
|
|||
proxy_exc = await self._invoke(ValueError("some other failure"))
|
||||
assert "retry-after" not in proxy_exc.headers
|
||||
|
||||
async def test_handle_llm_api_exception_strips_bug_report_notice_from_client_message(self, caplog):
|
||||
report = build_bug_report(RuntimeError("boom"), surface="sdk")
|
||||
notice = bug_report_notice(report)
|
||||
exc = litellm.APIConnectionError(
|
||||
message=f"boom\n{notice}",
|
||||
model="gpt-4o",
|
||||
llm_provider="openai",
|
||||
)
|
||||
|
||||
with caplog.at_level("ERROR"):
|
||||
proxy_exc = await self._invoke(exc)
|
||||
|
||||
assert ISSUE_URL_BASE not in proxy_exc.message
|
||||
assert ISSUE_URL_BASE in caplog.text
|
||||
|
||||
async def test_handle_llm_api_exception_retry_after_survives_callback_headers(self):
|
||||
from litellm.types.router import RouterRateLimitError
|
||||
|
||||
|
|
@ -9012,6 +9034,60 @@ class TestDetachedStreamFailureHook:
|
|||
assert [call["original_exception"] for call in recorder.calls] == [failure]
|
||||
|
||||
|
||||
class TestPostCallMaskedOutputReachesDeferredLogging:
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_streaming_records_the_masked_response_before_deferred_logging_fires(self, monkeypatch):
|
||||
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
||||
from litellm.types.utils import Choices, Message, ModelResponse
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_call_id = "lit-8325-call"
|
||||
logging_obj._defer_async_logging = False
|
||||
logging_obj._on_deferred_stream_complete = None
|
||||
logging_obj.cost_breakdown = None
|
||||
logging_obj.model_call_details = {}
|
||||
recorded_at_enqueue: dict[str, object] = {}
|
||||
logging_obj._enqueue_deferred_logging = lambda: recorded_at_enqueue.update(logging_obj.model_call_details)
|
||||
|
||||
processor = ProxyBaseLLMRequestProcessing(data={"model": "oa", "litellm_logging_obj": logging_obj})
|
||||
|
||||
def mask(data, user_api_key_dict, response):
|
||||
response.choices[0].message.content = "Card: <CREDIT_CARD>"
|
||||
return response
|
||||
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
|
||||
proxy_logging_obj.update_request_status = AsyncMock(return_value=None)
|
||||
proxy_logging_obj.post_call_success_hook = AsyncMock(side_effect=mask)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
|
||||
async def fake_route_request(**kwargs):
|
||||
async def call():
|
||||
return ModelResponse(
|
||||
choices=[Choices(index=0, message=Message(content="Card: 4111 1111 1111 1111", role="assistant"))]
|
||||
)
|
||||
|
||||
return call()
|
||||
|
||||
monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request)
|
||||
|
||||
result = await processor.base_process_llm_request(
|
||||
request=Request(scope={"type": "http", "headers": []}),
|
||||
fastapi_response=Response(),
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"),
|
||||
route_type="acompletion",
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
general_settings={},
|
||||
proxy_config=MagicMock(spec=ProxyConfig),
|
||||
select_data_generator=MagicMock(),
|
||||
is_streaming_request=False,
|
||||
skip_pre_call_logic=True,
|
||||
)
|
||||
|
||||
assert result.choices[0].message.content == "Card: <CREDIT_CARD>"
|
||||
assert recorded_at_enqueue[SERVED_OUTPUT_TEXTS_KEY] == ("Card: <CREDIT_CARD>",)
|
||||
|
||||
|
||||
class TestStreamingResponseHeadersFollowFallback:
|
||||
"""LIT-6767: the streaming branch has to publish the deployment that served the stream."""
|
||||
|
||||
|
|
@ -9374,6 +9450,94 @@ async def test_handle_llm_api_exception_forwards_litellm_response_headers_when_r
|
|||
assert exc_info.value.headers["llm_provider-x-request-id"] == "req_openai_400"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_llm_api_exception_logs_bug_report_for_unmapped_error(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
):
|
||||
monkeypatch.delenv(DISABLE_ENV_VAR, raising=False)
|
||||
processor = ProxyBaseLLMRequestProcessing(
|
||||
data={
|
||||
"proxy_server_request": {"url": "https://example.test/v1/chat/completions?debug=true"},
|
||||
"model": "acme-prod-gpt4",
|
||||
"custom_llm_provider": "openai",
|
||||
"stream": True,
|
||||
}
|
||||
)
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={})
|
||||
|
||||
with caplog.at_level("ERROR", logger="LiteLLM Proxy"):
|
||||
with pytest.raises(ProxyException):
|
||||
await processor._handle_llm_api_exception(
|
||||
e=RuntimeError("unmapped for user@example.com"),
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"),
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
issue_url = next(word for word in caplog.text.split() if word.startswith(ISSUE_URL_BASE))
|
||||
assert "Endpoint / call: /v1/chat/completions" in unquote_plus(issue_url)
|
||||
assert "Provider: openai" in unquote_plus(issue_url)
|
||||
assert "Stream: true" in unquote_plus(issue_url)
|
||||
assert "acme-prod-gpt4" not in unquote_plus(issue_url)
|
||||
assert "user@example.com" not in unquote_plus(issue_url)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_llm_api_exception_bug_report_drops_unknown_route(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
):
|
||||
monkeypatch.delenv(DISABLE_ENV_VAR, raising=False)
|
||||
processor = ProxyBaseLLMRequestProcessing(
|
||||
data={"proxy_server_request": {"url": "https://example.test/v1/files/file-customer-123/content"}}
|
||||
)
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={})
|
||||
|
||||
with caplog.at_level("ERROR", logger="LiteLLM Proxy"):
|
||||
with pytest.raises(ProxyException):
|
||||
await processor._handle_llm_api_exception(
|
||||
e=RuntimeError("unmapped"),
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"),
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
issue_url = next(word for word in caplog.text.split() if word.startswith(ISSUE_URL_BASE))
|
||||
assert "Endpoint / call: unknown" in unquote_plus(issue_url)
|
||||
assert "file-customer-123" not in unquote_plus(issue_url)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_llm_api_exception_skips_bug_report_for_provider_status(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
):
|
||||
monkeypatch.delenv(DISABLE_ENV_VAR, raising=False)
|
||||
|
||||
class ProviderRateLimitError(Exception):
|
||||
def __init__(self, message: str):
|
||||
super().__init__(message)
|
||||
self.status_code = 429
|
||||
|
||||
processor = ProxyBaseLLMRequestProcessing(data={})
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={})
|
||||
|
||||
with caplog.at_level("ERROR", logger="LiteLLM Proxy"):
|
||||
with pytest.raises(ProxyException):
|
||||
await processor._handle_llm_api_exception(
|
||||
e=ProviderRateLimitError("rate limited"),
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"),
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
assert ISSUE_URL_BASE not in caplog.text
|
||||
|
||||
|
||||
class TestBackgroundResponseRetrievalGovernance:
|
||||
"""LIT-7175: retrieving a background Response attaches the model's post_call policy pipelines."""
|
||||
|
||||
|
|
|
|||
163
tests/test_litellm/proxy/test_native_compaction.py
Normal file
163
tests/test_litellm/proxy/test_native_compaction.py
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
import asyncio
|
||||
from collections.abc import Awaitable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI, Request
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import inherit_message_logging_privacy
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
|
||||
from litellm.proxy import common_request_processing, proxy_server
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_model
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash, get_request_stash
|
||||
from litellm.proxy.native_compaction import with_proxy_compaction_executor
|
||||
from litellm.router import Router
|
||||
from litellm.router_strategy.complexity_router.context_compaction import compaction_executor, reject_recursive_compactor
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
_HEADERS: Final = (
|
||||
(b"authorization", b"Bearer sk-compaction-fixture"), (b"cookie", b"session=fixture"),
|
||||
(b"content-length", b"99999"), (b"x-litellm-call-id", b"parent"),
|
||||
(b"litellm-disable-message-redaction", b"true"), (b"x-litellm-num-retries", b"8"),
|
||||
(b"X-LiteLLM-Timeout", b"600"), (b"x-litellm-stream-timeout", b"500"),
|
||||
)
|
||||
|
||||
|
||||
async def _child(
|
||||
protocol: Literal["chat", "messages"] = "chat", forged: bool = False, parent_model: str | None = None
|
||||
) -> Mapping[str, object]:
|
||||
executor: Final = compaction_executor.get()
|
||||
assert executor is not None
|
||||
payload: Final = TypeAdapter(Mapping[str, object]).validate_json(
|
||||
b'{"model":"compactor","messages":[{"role":"user","content":"history"}],'
|
||||
b'"num_retries":0,"timeout":7,"stream_timeout":7,"disable_fallbacks":true,"stream":false,'
|
||||
b'"metadata":{"turn_off_message_logging":true}}'
|
||||
)
|
||||
return await executor(protocol, MappingProxyType({
|
||||
"litellm_metadata" if protocol == "messages" and key == "metadata" else key: value
|
||||
for key, value in payload.items() if forged or key != "metadata"
|
||||
}), parent_model)
|
||||
|
||||
|
||||
def _request(app: FastAPI) -> Request:
|
||||
return Request(TypeAdapter(dict[str, object]).validate_python(MappingProxyType({
|
||||
"type": "http", "app": app, "scheme": "https", "server": ("proxy.test", 443),
|
||||
"path": "/gateway/parent", "root_path": "/gateway", "query_string": b"parent=1",
|
||||
"client": ("192.0.2.1", 4321), "headers": _HEADERS,
|
||||
})))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("protocol", ("chat", "messages"))
|
||||
async def test_child_preserves_credentials_and_isolates_context(protocol: Literal["chat", "messages"]) -> None:
|
||||
app: Final = FastAPI()
|
||||
stash: Final = get_or_create_request_stash()
|
||||
|
||||
@app.post("/v1/chat/completions" if protocol == "chat" else "/v1/messages")
|
||||
async def endpoint(request: Request) -> Mapping[str, object]:
|
||||
assert get_request_stash() is None and compaction_executor.get() is None
|
||||
assert request.client == ("192.0.2.1", 4321) and request.url.scheme == "https"
|
||||
assert request.scope["root_path"] == "/gateway" and request.cookies["session"] == "fixture"
|
||||
assert request.headers["authorization"] == "Bearer sk-compaction-fixture" and not request.query_params
|
||||
assert "x-litellm-call-id" not in request.headers
|
||||
assert "litellm-disable-message-redaction" not in request.headers
|
||||
assert int(request.headers["content-length"]) == len(await request.body())
|
||||
with pytest.raises(BadRequestError, match="regular model group"):
|
||||
reject_recursive_compactor("auto-router")
|
||||
return MappingProxyType({"summary": "compacted"})
|
||||
|
||||
with inherit_message_logging_privacy(True):
|
||||
assert (await with_proxy_compaction_executor(_child(protocol), _request(app)))["summary"] == "compacted"
|
||||
assert get_request_stash() is stash and compaction_executor.get() is None
|
||||
reject_recursive_compactor("auto-router")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("protocol", ("chat", "messages"))
|
||||
@pytest.mark.parametrize("policy", ("allowed", "denied", "forged", "router_alias", "unrelated_alias"))
|
||||
async def test_real_proxy_child_auth_privacy_and_body_policy(
|
||||
monkeypatch: pytest.MonkeyPatch, protocol: Literal["chat", "messages"], policy: str,
|
||||
) -> None:
|
||||
cache: Final = DualCache()
|
||||
token: Final = proxy_server.hash_token("sk-compaction-fixture")
|
||||
models: Final = {"denied": ("answer",), "router_alias": ("auto",), "unrelated_alias": ("other-auto",)}.get(policy, ("compactor",))
|
||||
auth: Final = UserAPIKeyAuth.model_validate(MappingProxyType({"token": token, "models": models}))
|
||||
await cache.async_set_cache(key=token, value=auth)
|
||||
dispatched: Final = asyncio.Event()
|
||||
allowed: Final = policy in ("allowed", "router_alias")
|
||||
|
||||
async def route(
|
||||
data: Mapping[str, object], llm_router: Router | None, user_model: str | None,
|
||||
route_type: str, user_api_key_dict: UserAPIKeyAuth | None,
|
||||
) -> Awaitable[ModelResponse]:
|
||||
dispatched.set()
|
||||
assert allowed
|
||||
if policy == "router_alias":
|
||||
with pytest.raises(ProxyException):
|
||||
await can_key_call_model("unrelated-compactor", None, auth, None)
|
||||
assert (data["num_retries"], data["timeout"], data["stream_timeout"]) == (0, 7, 7)
|
||||
assert data["disable_fallbacks"] is True and data["stream"] is False
|
||||
logging: Final = data["litellm_logging_obj"]
|
||||
assert isinstance(logging, Logging)
|
||||
assert logging.standard_callback_dynamic_params.get("turn_off_message_logging") is True
|
||||
assert should_redact_message_logging(TypeAdapter(dict[str, object]).validate_python(MappingProxyType({
|
||||
"litellm_params": data, "standard_callback_dynamic_params": logging.standard_callback_dynamic_params,
|
||||
})))
|
||||
return asyncio.sleep(0, result=ModelResponse(id="private-summary", model="compactor"))
|
||||
|
||||
monkeypatch.setattr(proxy_server.app, "dependency_overrides", {})
|
||||
monkeypatch.setattr(proxy_server, "master_key", "sk-master-fixture")
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", object())
|
||||
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", None)
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {})
|
||||
monkeypatch.setattr(common_request_processing, "route_request", route)
|
||||
with inherit_message_logging_privacy(True):
|
||||
call: Final = with_proxy_compaction_executor(
|
||||
_child(protocol, policy == "forged", "auto" if policy.endswith("alias") else None), _request(proxy_server.app)
|
||||
)
|
||||
if allowed:
|
||||
assert (await call)["id"] == "private-summary"
|
||||
else:
|
||||
status: Final = 401 if policy == "forged" else 403
|
||||
with pytest.raises(BadRequestError, match=rf"child request failed \(HTTP {status}\)"):
|
||||
await call
|
||||
assert dispatched.is_set() is allowed
|
||||
assert compaction_executor.get() is None
|
||||
if policy.endswith("alias"):
|
||||
with pytest.raises(ProxyException):
|
||||
await can_key_call_model("compactor", None, auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("timeout", [False, True])
|
||||
async def test_cancelling_parent_cancels_and_drains_child(timeout: bool) -> None:
|
||||
app: Final = FastAPI()
|
||||
started: Final = asyncio.Event()
|
||||
stopped: Final = asyncio.Event()
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
async def endpoint() -> None:
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
stopped.set()
|
||||
|
||||
parent: Final = asyncio.create_task(with_proxy_compaction_executor(_child(), _request(app)))
|
||||
await asyncio.wait_for(started.wait(), timeout=5)
|
||||
if timeout:
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
await asyncio.wait_for(parent, timeout=0)
|
||||
else:
|
||||
parent.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await parent
|
||||
assert stopped.is_set() and compaction_executor.get() is None
|
||||
|
|
@ -347,6 +347,60 @@ async def test_post_call_stream_guardrail_keeps_own_iterator_on_chat_completions
|
|||
assert delivered_text != ""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_stream_records_masked_text_for_deferred_logging(monkeypatch):
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
monkeypatch.setattr(litellm, "callbacks", [_content_filter_guardrail("MASK")])
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
||||
logging_obj = _streaming_logging_obj()
|
||||
|
||||
async def fake_stream():
|
||||
yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="the zebra runs"))])
|
||||
yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")])
|
||||
|
||||
delivered_text = ""
|
||||
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
response=fake_stream(),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/chat/completions"),
|
||||
request_data={"model": "gpt-4o-mini", "metadata": {}, "litellm_logging_obj": logging_obj},
|
||||
):
|
||||
for choice in chunk.choices:
|
||||
delivered_text += choice.delta.content or ""
|
||||
|
||||
assert "zebra" not in delivered_text
|
||||
assert logging_obj.model_call_details[SERVED_OUTPUT_TEXTS_KEY] == (delivered_text,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_stream_records_the_served_text_when_the_client_disconnects(monkeypatch):
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
monkeypatch.setattr(litellm, "callbacks", [_content_filter_guardrail("MASK")])
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
||||
logging_obj = _streaming_logging_obj()
|
||||
|
||||
async def fake_stream():
|
||||
yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="the zebra runs"))])
|
||||
yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=" far"))])
|
||||
|
||||
stream = proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
response=fake_stream(),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/chat/completions"),
|
||||
request_data={"model": "gpt-4o-mini", "metadata": {}, "litellm_logging_obj": logging_obj},
|
||||
)
|
||||
first = await stream.__anext__()
|
||||
await stream.aclose()
|
||||
|
||||
delivered_text = "".join(choice.delta.content or "" for choice in first.choices)
|
||||
assert "zebra" not in delivered_text
|
||||
assert logging_obj.model_call_details[SERVED_OUTPUT_TEXTS_KEY] == (delivered_text,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unified_guardrail_iterator_accepts_explicit_guardrail():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -6,10 +6,12 @@ import pytest
|
|||
from fastapi import HTTPException
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.exceptions import InternalServerError
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.bug_report import ISSUE_URL_BASE
|
||||
from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
|
||||
|
|
@ -2418,3 +2420,16 @@ def test_mcp_auth_policy_uses_original_request_model(monkeypatch, model, expecte
|
|||
synthetic = proxy_logging._convert_mcp_to_llm_format(proxy_logging._create_mcp_request_object_from_kwargs(kwargs), kwargs)
|
||||
assert ("model-rule" in synthetic["metadata"]["guardrails"]) is expected
|
||||
assert "request-rule" in synthetic["metadata"]["guardrails"]
|
||||
|
||||
|
||||
def test_handle_exception_on_proxy_logs_bug_report_only_for_unmapped_500(caplog):
|
||||
with caplog.at_level("ERROR", logger="LiteLLM Proxy"):
|
||||
provider_result = handle_exception_on_proxy(
|
||||
InternalServerError(message="upstream 500", llm_provider="openai", model="gpt-4")
|
||||
)
|
||||
assert ISSUE_URL_BASE not in caplog.text
|
||||
internal_result = handle_exception_on_proxy(KeyError("missing"))
|
||||
|
||||
assert provider_result.code == internal_result.code == "500"
|
||||
assert ISSUE_URL_BASE in caplog.text
|
||||
assert ISSUE_URL_BASE not in internal_result.message
|
||||
|
|
|
|||
|
|
@ -0,0 +1,503 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import Iterator, Mapping
|
||||
from copy import deepcopy
|
||||
from functools import partial
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm.llms import compaction as native
|
||||
from litellm.router_strategy.complexity_router.config import ContextCompactionConfig
|
||||
from litellm.router_strategy.complexity_router.context_compaction import (
|
||||
CompactionState,
|
||||
Surface,
|
||||
arm_compaction,
|
||||
compact_to_fit,
|
||||
compaction_executor,
|
||||
)
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
|
||||
from litellm.types.router import Deployment
|
||||
|
||||
pytestmark: Final = [pytest.mark.asyncio, pytest.mark.usefixtures("local_model_cost_map")]
|
||||
SCHEMA: Final = {"type": "object", "properties": {"code": {"type": "string"}}}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def native_catalog(monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None) -> None:
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
monkeypatch.setenv("LITELLM_LICENSE", "")
|
||||
monkeypatch.setattr(litellm, "use_chat_completions_url_for_anthropic_messages", False)
|
||||
monkeypatch.setitem(litellm.model_cost, "summary-fixture", {
|
||||
"litellm_provider": "anthropic", "mode": "chat", "max_input_tokens": 32000,
|
||||
"max_output_tokens": 4096, "supports_anthropic_compaction": True,
|
||||
})
|
||||
|
||||
|
||||
def make_router(
|
||||
window: int | None = 512, settings: Mapping[str, object] | None = None, *, compactor_window: int = 32000,
|
||||
conflict: bool = False, output: int | None = 64,
|
||||
answer_defaults: Mapping[str, object] | None = None,
|
||||
context_fallback: bool = False,
|
||||
) -> litellm.Router:
|
||||
config: Final = {
|
||||
"tiers": {"SIMPLE": "small", "MEDIUM": "large", "COMPLEX": "large", "REASONING": "large"},
|
||||
"keyword_tier_rules": [{"keywords": ["answer", "tail result"], "tier": "SIMPLE"}],
|
||||
"enable_context_window_escalation": False, "max_tokens_from_tier_model": False,
|
||||
**(settings or {}),
|
||||
}
|
||||
return litellm.Router(model_list=[
|
||||
{"model_name": "auto", "litellm_params": {
|
||||
"model": "auto_router/complexity_router", "complexity_router_config": config,
|
||||
}},
|
||||
{"model_name": "small", "litellm_params": {
|
||||
"model": "openai/arbitrary-answer", "api_base": "https://answer.test/v1", "api_key": "answer-test", "max_retries": 0,
|
||||
**(answer_defaults or {}),
|
||||
}, "model_info": {"id": "pinned-answer", "max_input_tokens": window, "max_output_tokens": output}},
|
||||
{"model_name": "large", "litellm_params": {
|
||||
"model": "anthropic/summary-fixture", "api_base": "https://compact.test", "api_key": "compact-test",
|
||||
**({"stop": ["deployment policy"]} if conflict else {}),
|
||||
}, "model_info": {"id": "native-compactor", "max_input_tokens": compactor_window, "max_output_tokens": 4096}},
|
||||
{"model_name": "backup", "litellm_params": {
|
||||
"model": "anthropic/summary-fixture", "api_base": "https://compact.test", "api_key": "backup-test",
|
||||
}, "model_info": {"id": "backup-compactor", "max_input_tokens": 32000, "max_output_tokens": 4096}},
|
||||
], enable_pre_call_checks=True, num_retries=0, disable_cooldowns=True,
|
||||
retry_policy={"InternalServerErrorRetries": 1},
|
||||
context_window_fallbacks=[{"auto": ["large"]}] if context_fallback else [])
|
||||
|
||||
|
||||
def exchange(surface: Surface, phase: str) -> list[dict[str, object]]:
|
||||
identifier: Final = f"{phase}-call"
|
||||
result: Final = f"{phase} result"
|
||||
if surface == "responses":
|
||||
return [
|
||||
{"type": "function_call", "call_id": identifier, "name": "lookup", "arguments": "{}"},
|
||||
{"type": "function_call_output", "call_id": identifier, "output": result},
|
||||
]
|
||||
if surface == "messages":
|
||||
return [
|
||||
{"role": "assistant", "content": [{"type": "tool_use", "id": identifier, "name": "lookup", "input": {}}]},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": identifier, "content": result}]},
|
||||
]
|
||||
return [
|
||||
{"role": "assistant", "tool_calls": [
|
||||
{"id": identifier, "type": "function", "function": {"name": "lookup", "arguments": "{}"}},
|
||||
]},
|
||||
{"role": "tool", "tool_call_id": identifier, "content": result},
|
||||
]
|
||||
|
||||
|
||||
def history(surface: Surface) -> dict[str, object]:
|
||||
conversation: Final = [
|
||||
{"role": "user", "content": "Project code MAPLE-47. Background detail. " * 150},
|
||||
{"role": "assistant", "content": "Recorded"},
|
||||
*exchange(surface, "prefix"),
|
||||
{"role": "user", "content": "Answer with the project code"},
|
||||
*exchange(surface, "tail"),
|
||||
]
|
||||
function: Final = {"name": "lookup", "parameters": SCHEMA}
|
||||
if surface == "responses":
|
||||
return {"instructions": "Keep the code exact", "tools": [{"type": "function", **function}], "input": [
|
||||
{"role": "developer", "content": "Retain the original spelling"}, *conversation,
|
||||
]}
|
||||
if surface == "messages":
|
||||
return {"system": "Keep the code exact", "tools": [{"name": "lookup", "input_schema": SCHEMA}], "messages": [
|
||||
*conversation,
|
||||
]}
|
||||
return {"tools": [{"type": "function", "function": function}], "messages": [
|
||||
{"role": "system", "content": "Keep the code exact"}, *conversation,
|
||||
]}
|
||||
|
||||
|
||||
def native_reply(summary: str = "Project code MAPLE-47", signed: bool = True, truncated: bool = False) -> httpx.Response:
|
||||
return httpx.Response(200, json={
|
||||
"id": "msg_compact", "type": "message", "role": "assistant", "model": "summary-fixture",
|
||||
"content": [{"type": "compaction", "content": summary, **({"signature": "native-signature"} if signed else {})}],
|
||||
"stop_reason": "max_tokens" if truncated else "compaction", "usage": {"input_tokens": 0, "output_tokens": 0, "iterations": [
|
||||
{"type": "compaction", "input_tokens": 1200, "output_tokens": 20},
|
||||
]},
|
||||
})
|
||||
|
||||
|
||||
def answer_reply(request: httpx.Request, expected_model: str = "arbitrary-answer") -> httpx.Response:
|
||||
payload: Final = json.loads(request.content)
|
||||
assert payload["model"] == expected_model
|
||||
assert request.headers["authorization"] == "Bearer answer-test"
|
||||
if request.url.path.endswith("responses"):
|
||||
return httpx.Response(200, json={
|
||||
"id": "resp_answer", "object": "response", "created_at": 0, "status": "completed",
|
||||
"model": payload["model"], "output": [{"id": "msg_answer", "type": "message", "role": "assistant",
|
||||
"status": "completed", "content": [{"type": "output_text", "text": "MAPLE-47", "annotations": []}]}],
|
||||
"usage": {"input_tokens": 60, "output_tokens": 8, "total_tokens": 68},
|
||||
})
|
||||
return httpx.Response(200, json={
|
||||
"id": "answer", "object": "chat.completion", "created": 0, "model": payload["model"],
|
||||
"choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "MAPLE-47"}}],
|
||||
"usage": {"prompt_tokens": 60, "completion_tokens": 8, "total_tokens": 68},
|
||||
})
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def wire() -> Iterator[tuple[respx.Route, respx.Route]]:
|
||||
with respx.mock(assert_all_called=False) as transport:
|
||||
compactor: Final = transport.post("https://compact.test/v1/messages").mock(return_value=native_reply())
|
||||
answer: Final = transport.route(method="POST", host="answer.test").mock(side_effect=answer_reply)
|
||||
yield compactor, answer
|
||||
|
||||
|
||||
async def invoke(router: litellm.Router, surface: Surface, payload: Mapping[str, object], retries: int = 0) -> object:
|
||||
if surface == "responses":
|
||||
return await router.aresponses(model="auto", max_output_tokens=64, num_retries=retries, **payload)
|
||||
if surface == "messages":
|
||||
return await router.aanthropic_messages(model="auto", max_tokens=64, num_retries=retries, **payload)
|
||||
return await router.acompletion(model="auto", max_tokens=64, num_retries=retries, **payload)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", ["chat", "messages", "responses"])
|
||||
@pytest.mark.parametrize("near", [False, True])
|
||||
@pytest.mark.parametrize("configured", [False, True])
|
||||
async def test_all_surfaces_compact_and_keep_selected_answerer(
|
||||
wire: tuple[respx.Route, respx.Route], surface: Surface, near: bool, configured: bool,
|
||||
) -> None:
|
||||
payload: Final = history(surface)
|
||||
original: Final = deepcopy(payload)
|
||||
counted: Final = make_router()._count_pre_call_check_tokens(payload.get("messages"), payload.get("input"), payload)
|
||||
window: Final = int((counted + 32) / ContextCompactionConfig().trigger_ratio) + 1 if near else 512
|
||||
assert (counted < window) is near
|
||||
settings: Final = {"enable_context_window_escalation": True,
|
||||
**({"context_compaction": {"model": "large", "max_tokens": 512}} if configured else {})}
|
||||
router: Final = make_router(window, settings)
|
||||
captured: Final = asyncio.Queue[Mapping[str, object]]()
|
||||
compactor, answer = wire
|
||||
retry: Final = near and configured
|
||||
|
||||
def answer_after_retry(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(500, json={"error": {"message": "retry answer"}}) if answer.call_count == 0 else answer_reply(request)
|
||||
|
||||
if retry:
|
||||
answer.mock(side_effect=answer_after_retry)
|
||||
|
||||
async def execute(
|
||||
protocol: native.CompactionProtocol, request: Mapping[str, object], parent_model: str | None = None
|
||||
) -> Mapping[str, object]:
|
||||
assert parent_model == "auto"
|
||||
result: Final = await native.dispatch(router, protocol, request)
|
||||
captured.put_nowait(result)
|
||||
return result
|
||||
|
||||
token: Final = compaction_executor.set(execute)
|
||||
try:
|
||||
response: Final = await invoke(router, surface, payload, retries=int(retry))
|
||||
finally:
|
||||
compaction_executor.reset(token)
|
||||
assert compactor.call_count == captured.qsize() == 1
|
||||
assert answer.call_count == router.total_calls["openai/arbitrary-answer"] == 1 + int(retry)
|
||||
compact_request: Final = compactor.calls[0].request
|
||||
compact_body: Final = json.loads(compact_request.content)
|
||||
answer_body: Final = json.loads(answer.calls[0].request.content)
|
||||
assert answer_body == json.loads(answer.calls[-1].request.content)
|
||||
assert compact_body["model"] == "summary-fixture" and compact_body["compaction"] == {"type": "summarize"}
|
||||
assert ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value in compact_request.headers["anthropic-beta"].split(",")
|
||||
assert compact_body["max_tokens"] == (512 if configured else ContextCompactionConfig().max_tokens)
|
||||
assert "Background detail" in str(compact_body) and "tail-call" not in str(compact_body)
|
||||
assert "prefix-call" in str(compact_body) and "prefix result" in str(compact_body)
|
||||
assert "Keep the code exact" in str(compact_body["system"])
|
||||
assert compact_body["tools"][0]["input_schema"] == SCHEMA
|
||||
assert "MAPLE-47" in str(answer_body) and "Background detail" not in str(answer_body)
|
||||
assert "prefix-call" not in str(answer_body)
|
||||
assert "native-signature" not in str(answer_body) and "compaction" not in answer_body
|
||||
assert "tail-call" in str(answer_body) and "tail result" in str(answer_body)
|
||||
assert "MAPLE-47" in str(response)
|
||||
usage: Final = captured.get_nowait()["usage"]
|
||||
if surface == "messages":
|
||||
assert usage["iterations"][0]["input_tokens"] == 1200 and usage["iterations"][0]["output_tokens"] == 20
|
||||
else:
|
||||
assert usage["prompt_tokens"] == 1200 and usage["completion_tokens"] == 20
|
||||
if surface == "responses":
|
||||
assert answer_body["input"][-3:] == original["input"][-3:]
|
||||
assert answer_body["input"][0] == original["input"][0]
|
||||
assert answer_body["instructions"] == original["instructions"] and answer_body["tools"] == original["tools"]
|
||||
assert payload == original
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", ["chat", "messages", "responses"])
|
||||
@pytest.mark.parametrize("mode", ["fitting", "false", "null"])
|
||||
async def test_fitting_and_disabled_requests_do_not_compact(
|
||||
wire: tuple[respx.Route, respx.Route], surface: Surface, mode: str,
|
||||
) -> None:
|
||||
settings: Final = {} if mode == "fitting" else {"context_compaction": False if mode == "false" else None}
|
||||
payload: Final = history(surface)
|
||||
original: Final = deepcopy(payload)
|
||||
counted: Final = make_router()._count_pre_call_check_tokens(payload.get("messages"), payload.get("input"), payload)
|
||||
await invoke(make_router(20000 if mode == "fitting" else counted + 64, settings), surface, payload)
|
||||
compactor, answer = wire
|
||||
assert compactor.call_count == 0 and answer.call_count == 1
|
||||
assert "Background detail" in answer.calls[0].request.content.decode()
|
||||
assert payload == original
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", ["chat", "messages", "responses"])
|
||||
@pytest.mark.parametrize("reason", ["single", "unclosed", "no_compactor"])
|
||||
async def test_fitting_request_survives_unavailable_compaction(
|
||||
wire: tuple[respx.Route, respx.Route], surface: Surface, reason: str,
|
||||
) -> None:
|
||||
key: Final = "input" if surface == "responses" else "messages"
|
||||
items: Final = [{"role": "user", "content": "Answer with MAPLE-47. Detail. " * 80}]
|
||||
payload: Final = history(surface) if reason == "no_compactor" else {key: (
|
||||
items if reason == "single" else [*items, *exchange(surface, "open")[:1], {"role": "user", "content": "Answer"}]
|
||||
)}
|
||||
counted: Final = make_router()._count_pre_call_check_tokens(payload.get("messages"), payload.get("input"), payload)
|
||||
settings: Final = {"tiers": {"SIMPLE": "small", "MEDIUM": "small", "COMPLEX": "small", "REASONING": "small"}} if reason == "no_compactor" else {}
|
||||
await invoke(make_router(counted + 1, settings), surface, payload)
|
||||
compactor, answer = wire
|
||||
assert compactor.call_count == 0 and answer.call_count == 1
|
||||
assert "Detail" in str(answer.calls[0].request.content) or "Background detail" in str(answer.calls[0].request.content)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", ["chat", "messages", "responses"])
|
||||
async def test_uncompactable_overflow_uses_explicit_context_fallback(
|
||||
wire: tuple[respx.Route, respx.Route], surface: Surface,
|
||||
) -> None:
|
||||
compactor, answer = wire
|
||||
compactor.mock(return_value=httpx.Response(200, json={**native_reply().json(),
|
||||
"content": [{"type": "text", "text": "MAPLE-47"}], "stop_reason": "end_turn"}))
|
||||
payload: Final = {"input" if surface == "responses" else "messages": [
|
||||
{"role": "user", "content": "Answer with MAPLE-47. Detail. " * 150},
|
||||
]}
|
||||
await invoke(make_router(context_fallback=True), surface, payload)
|
||||
assert answer.call_count == 0 and compactor.call_count == 1
|
||||
assert "compaction" not in json.loads(compactor.calls[0].request.content)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("escalate", [False, True])
|
||||
async def test_no_native_compactor_respects_explicit_escalation(
|
||||
wire: tuple[respx.Route, respx.Route], monkeypatch: pytest.MonkeyPatch, escalate: bool,
|
||||
) -> None:
|
||||
monkeypatch.setitem(litellm.model_cost["summary-fixture"], "supports_anthropic_compaction", False)
|
||||
router: Final = make_router(settings={"enable_context_window_escalation": escalate})
|
||||
compactor, answer = wire
|
||||
compactor.mock(return_value=httpx.Response(200, json={**native_reply().json(),
|
||||
"content": [{"type": "text", "text": "MAPLE-47"}], "stop_reason": "end_turn"}))
|
||||
if not escalate:
|
||||
with pytest.raises(litellm.ContextWindowExceededError, match="No configured compactor"):
|
||||
await invoke(router, "chat", history("chat"))
|
||||
assert compactor.call_count == 0
|
||||
else:
|
||||
await invoke(router, "chat", history("chat"))
|
||||
assert compactor.call_count == 1
|
||||
assert "compaction" not in json.loads(compactor.calls[0].request.content)
|
||||
assert answer.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", ["chat", "messages", "responses"])
|
||||
async def test_undersized_native_compactor_does_not_block_explicit_escalation(
|
||||
wire: tuple[respx.Route, respx.Route], surface: Surface,
|
||||
) -> None:
|
||||
router: Final = make_router(compactor_window=512, settings={
|
||||
"enable_context_window_escalation": True,
|
||||
"tiers": {"SIMPLE": "small", "MEDIUM": "large", "COMPLEX": "wide", "REASONING": "wide"},
|
||||
})
|
||||
router.add_deployment(Deployment(
|
||||
model_name="wide", litellm_params={
|
||||
"model": "openai/wide-answer", "api_base": "https://answer.test/v1", "api_key": "answer-test",
|
||||
}, model_info={"id": "wide-answer", "max_input_tokens": 32000, "max_output_tokens": 64},
|
||||
))
|
||||
compactor, answer = wire
|
||||
answer.mock(side_effect=partial(answer_reply, expected_model="wide-answer"))
|
||||
await invoke(router, surface, history(surface))
|
||||
assert compactor.call_count == 0 and answer.call_count == 1
|
||||
assert json.loads(answer.calls[0].request.content)["model"] == "wide-answer"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("surface", "failure"),
|
||||
[(surface, failure) for surface in ("chat", "messages", "responses") for failure in ("unsigned", "oversized", "provider")]
|
||||
+ [("messages", "truncated")],
|
||||
)
|
||||
async def test_bad_native_result_never_reaches_answerer(
|
||||
wire: tuple[respx.Route, respx.Route], surface: Surface, failure: str,
|
||||
) -> None:
|
||||
compactor, answer = wire
|
||||
reply: Final = httpx.Response(500, json={"error": {"type": "api_error", "message": "failed"}}) if failure == "provider" else native_reply(
|
||||
"too large " * 2000 if failure == "oversized" else "MAPLE-47", signed=failure != "unsigned",
|
||||
truncated=failure == "truncated",
|
||||
)
|
||||
compactor.mock(return_value=reply)
|
||||
with pytest.raises((litellm.BadRequestError, litellm.InternalServerError)):
|
||||
await invoke(make_router(), surface, history(surface))
|
||||
assert compactor.call_count == 1 and answer.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["item", "content", "tools", "unclosed", "missing", "duplicate", "instructions", "retained"])
|
||||
@pytest.mark.parametrize("needed", [False, True])
|
||||
async def test_unsafe_responses_reject_only_when_compaction_needed(
|
||||
wire: tuple[respx.Route, respx.Route], failure: str, needed: bool,
|
||||
) -> None:
|
||||
payload: Final = history("responses")
|
||||
extra: Final = {
|
||||
"item": [{"type": "computer_call", "call_id": "opaque-tool"}],
|
||||
"content": [{"role": "assistant", "content": [{"type": "refusal", "refusal": "cannot"}]}],
|
||||
"unclosed": [{"type": "function_call", "call_id": "unclosed", "name": "lookup", "arguments": "{}"}],
|
||||
"missing": [{"type": "function_call", "name": "lookup", "arguments": "{}"}],
|
||||
"duplicate": exchange("responses", "prefix"),
|
||||
"instructions": [{"role": "developer", "content": "Changed instructions"}],
|
||||
}
|
||||
request: Final = {
|
||||
**payload, "input": [*payload["input"][:3], *extra.get(failure, []), *payload["input"][3:]],
|
||||
**({"tools": [{"type": "computer_use_preview", "display_width": 800, "display_height": 600}]} if failure == "tools" else {}),
|
||||
**({"instructions": "Keep every instruction " * 600} if failure == "retained" else {}),
|
||||
}
|
||||
if needed:
|
||||
with pytest.raises(litellm.BadRequestError, match="Context compaction"):
|
||||
await invoke(make_router(), "responses", request)
|
||||
assert all(route.call_count == 0 for route in wire)
|
||||
else:
|
||||
await invoke(make_router(20000), "responses", request)
|
||||
compactor, answer = wire
|
||||
assert compactor.call_count == 0 and answer.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("owned", [
|
||||
{"previous_response_id": "resp_parent"}, {"conversation": "conv_parent"},
|
||||
{"context_management": [{"type": "compaction", "compact_threshold": 1000}]}, {"compaction": {"type": "summarize"}},
|
||||
{"input": [{"type": "reasoning", "encrypted_content": "opaque"}]},
|
||||
{"input": [{"type": "reasoning", "summary": [{"type": "summary_text", "text": "prior reasoning"}]}]},
|
||||
{"input": [{"type": "compaction", "encrypted_content": "opaque"}]},
|
||||
{"input": [{"type": "item_reference", "id": "item_parent"}]},
|
||||
{"input": [{"role": "assistant", "content": "visible", "encrypted_content": "opaque"}]},
|
||||
{"input": [{"role": "user", "content": [{"type": "encrypted_content", "encrypted_content": "opaque"}]}]},
|
||||
])
|
||||
@pytest.mark.parametrize("arm_first", [False, True])
|
||||
async def test_client_owned_history_bypasses_compaction(
|
||||
wire: tuple[respx.Route, respx.Route], owned: Mapping[str, object], arm_first: bool,
|
||||
) -> None:
|
||||
router: Final = make_router()
|
||||
deployment: Final = router.get_deployment(model_id="pinned-answer")
|
||||
assert deployment is not None
|
||||
state: Final = CompactionState()
|
||||
request: Final = {**history("responses"), **owned, "model": "small", "max_tokens": 64, "_context_compaction_state": state}
|
||||
original: Final = deepcopy({key: value for key, value in request.items() if key != "_context_compaction_state"})
|
||||
before_defaults: Final = {"_context_compaction_state": state} if arm_first else request
|
||||
await arm_compaction(before_defaults, ContextCompactionConfig(), ("large",))
|
||||
counted: Final = router._count_pre_call_check_tokens(None, request["input"], request)
|
||||
if arm_first and counted > 512:
|
||||
with pytest.raises(litellm.ContextWindowExceededError):
|
||||
await compact_to_fit(router, deployment.model_dump(), request, "responses")
|
||||
else:
|
||||
result: Final = await compact_to_fit(router, deployment.model_dump(), request, "responses")
|
||||
assert result is request
|
||||
assert {key: value for key, value in request.items() if key != "_context_compaction_state"} == original
|
||||
assert all(route.call_count == 0 for route in wire)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", ["chat", "messages", "responses"])
|
||||
@pytest.mark.parametrize("source", ["request", "deployment"])
|
||||
async def test_client_managed_overflow_keeps_context_window_admission(
|
||||
wire: tuple[respx.Route, respx.Route], surface: Surface, source: str,
|
||||
) -> None:
|
||||
managed: Final = {"context_management": {"edits": []}}
|
||||
router: Final = make_router(answer_defaults=managed if source == "deployment" else None)
|
||||
payload: Final = {**history(surface), **(managed if source == "request" else {})}
|
||||
compactor, answer = wire
|
||||
if source == "deployment":
|
||||
with pytest.raises(litellm.ContextWindowExceededError):
|
||||
await invoke(router, surface, payload)
|
||||
assert compactor.call_count == 0
|
||||
else:
|
||||
await invoke(router, surface, payload)
|
||||
assert compactor.call_count == 1
|
||||
assert "compaction" not in json.loads(compactor.calls[0].request.content)
|
||||
assert answer.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", ["conflicting_defaults", "small_window", "capability_false", "capability_missing"])
|
||||
async def test_automatic_compactor_skips_conflicts_and_requires_capacity_and_capability(
|
||||
wire: tuple[respx.Route, respx.Route], monkeypatch: pytest.MonkeyPatch, case: str,
|
||||
) -> None:
|
||||
if case.startswith("capability"):
|
||||
metadata: Final = {key: value for key, value in litellm.model_cost["summary-fixture"].items()
|
||||
if key != "supports_anthropic_compaction"}
|
||||
monkeypatch.setitem(litellm.model_cost, "summary-fixture", {
|
||||
**metadata, **({"supports_anthropic_compaction": False} if case == "capability_false" else {}),
|
||||
})
|
||||
conflict: Final = case == "conflicting_defaults"
|
||||
settings: Final = {"tiers": {"SIMPLE": "small", "MEDIUM": "large", "COMPLEX": "backup", "REASONING": "backup"}} if conflict else {}
|
||||
router: Final = make_router(settings=settings, conflict=conflict, compactor_window=512 if case == "small_window" else 32000)
|
||||
if conflict:
|
||||
await invoke(router, "chat", history("chat"))
|
||||
compactor, answer = wire
|
||||
assert compactor.call_count == answer.call_count == 1
|
||||
assert compactor.calls[0].request.headers["x-api-key"] == "backup-test"
|
||||
assert "stop_sequences" not in json.loads(compactor.calls[0].request.content)
|
||||
else:
|
||||
with pytest.raises(litellm.BadRequestError, match="No configured compactor"):
|
||||
await invoke(router, "chat", history("chat"))
|
||||
assert all(route.call_count == 0 for route in wire)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("unknown_output", [False, True])
|
||||
@pytest.mark.parametrize("overflow", [False, True])
|
||||
async def test_unusable_output_budget_still_enforces_known_input_window(
|
||||
wire: tuple[respx.Route, respx.Route], unknown_output: bool, overflow: bool,
|
||||
) -> None:
|
||||
payload: Final = history("chat")
|
||||
counted: Final = make_router(output=None)._count_pre_call_check_tokens(payload["messages"], None, payload)
|
||||
window: Final = 512 if overflow else counted + 64
|
||||
output: Final = None if unknown_output else window
|
||||
router: Final = make_router(window, output=output)
|
||||
deployment: Final = router.get_deployment(model_id="pinned-answer")
|
||||
assert deployment is not None
|
||||
state: Final = CompactionState(config=ContextCompactionConfig(), candidates=("large",))
|
||||
request: Final = {**payload, "model": "small", "max_tokens": output, "_context_compaction_state": state}
|
||||
if overflow:
|
||||
with pytest.raises(litellm.BadRequestError, match="known input window"):
|
||||
await compact_to_fit(router, deployment.model_dump(), request, "chat")
|
||||
else:
|
||||
assert await compact_to_fit(router, deployment.model_dump(), request, "chat") is request
|
||||
assert all(route.call_count == 0 for route in wire)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("outcome", ["success", "timeout", "cancel"])
|
||||
async def test_retry_reuses_summary_or_terminal_cancellation(outcome: Literal["success", "timeout", "cancel"]) -> None:
|
||||
router: Final = make_router()
|
||||
deployment: Final = router.get_deployment(model_id="pinned-answer")
|
||||
assert deployment is not None
|
||||
state: Final = CompactionState(config=ContextCompactionConfig(model="large", max_tokens=512, timeout_seconds=0.02))
|
||||
request: Final = {**history("messages"), "model": "small", "max_tokens": 64, "_context_compaction_state": state}
|
||||
calls: Final = asyncio.Queue[None]()
|
||||
started: Final = asyncio.Event()
|
||||
stopped: Final = asyncio.Event()
|
||||
|
||||
async def execute(
|
||||
protocol: native.CompactionProtocol, payload: Mapping[str, object], parent_model: str | None = None
|
||||
) -> Mapping[str, object]:
|
||||
calls.put_nowait(None)
|
||||
started.set()
|
||||
try:
|
||||
return native_reply().json() if outcome == "success" else await asyncio.Future[Mapping[str, object]]()
|
||||
finally:
|
||||
stopped.set()
|
||||
|
||||
token: Final = compaction_executor.set(execute)
|
||||
try:
|
||||
first: Final = asyncio.create_task(compact_to_fit(router, deployment.model_dump(), request, "messages"))
|
||||
await asyncio.wait_for(started.wait(), timeout=2)
|
||||
if outcome == "cancel":
|
||||
first.cancel()
|
||||
if outcome == "success":
|
||||
assert await first == await compact_to_fit(router, deployment.model_dump(), request, "messages")
|
||||
changed: Final = {**request, "messages": [{"role": "user", "content": "new history"}, *request["messages"]]}
|
||||
with pytest.raises(litellm.BadRequestError, match="History changed"):
|
||||
await compact_to_fit(router, deployment.model_dump(), changed, "messages")
|
||||
else:
|
||||
error: Final = asyncio.CancelledError if outcome == "cancel" else asyncio.TimeoutError
|
||||
with pytest.raises(error):
|
||||
await first
|
||||
with pytest.raises(error):
|
||||
await compact_to_fit(router, deployment.model_dump(), request, "messages")
|
||||
assert calls.qsize() == 1 and stopped.is_set()
|
||||
finally:
|
||||
compaction_executor.reset(token)
|
||||
|
|
@ -15317,6 +15317,7 @@ class TestHealthFallbackDispatch:
|
|||
|
||||
router: Final = self._router(
|
||||
config={
|
||||
"context_compaction": False,
|
||||
"tiers": {"SIMPLE": "primary", "MEDIUM": "peer", "COMPLEX": "large"},
|
||||
"enable_context_window_escalation": True,
|
||||
}
|
||||
|
|
@ -15399,6 +15400,7 @@ class TestHealthFallbackDispatch:
|
|||
async def test_modality_default_must_also_fit_context(self, default_fits: bool) -> None:
|
||||
router: Final = self._router(
|
||||
config={
|
||||
"context_compaction": False,
|
||||
"modality_routing": True,
|
||||
"tiers": {"SIMPLE": "primary"},
|
||||
"enable_context_window_escalation": True,
|
||||
|
|
|
|||
|
|
@ -83,6 +83,7 @@ class TestAnthropicBetaHeadersFiltering:
|
|||
filtered = filter_and_transform_beta_headers(
|
||||
beta_headers=all_headers, provider=provider
|
||||
)
|
||||
assert ("compact-2026-09-04" in filtered) is (provider == "anthropic")
|
||||
|
||||
for header in unsupported_headers:
|
||||
assert (
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import datetime
|
||||
import time
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -29,6 +30,7 @@ from litellm.types.utils import (
|
|||
PromptTokensDetailsWrapper,
|
||||
Usage,
|
||||
)
|
||||
from litellm.types.videos.main import VideoObject
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -3038,21 +3040,23 @@ def test_completion_cost_logs_cache_and_reasoning_breakdown_for_custom_pricing()
|
|||
assert total == pytest.approx(100 * 1e-6 + 800 * 1e-7 + 100 * 1.25e-6 + 500 * 2e-6)
|
||||
|
||||
|
||||
def test_cost_per_token_per_second_pricing(monkeypatch):
|
||||
@pytest.mark.parametrize("custom_llm_provider", ["together_ai", "openai", "anthropic", "bedrock", "azure"])
|
||||
def test_cost_per_token_per_second_pricing(monkeypatch, custom_llm_provider: str):
|
||||
"""
|
||||
Models priced by duration (input/output_cost_per_second) with no per-token rates
|
||||
must be billed as cost_per_second * response_time_ms / 1000 in cost_per_token.
|
||||
must be billed as cost_per_second * response_time_ms / 1000 in cost_per_token,
|
||||
whether or not the provider has its own cost calculator.
|
||||
"""
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
||||
|
||||
model = "test-per-second-pricing-model"
|
||||
model = f"test-per-second-pricing-{custom_llm_provider}"
|
||||
litellm.register_model(
|
||||
model_cost={
|
||||
model: {
|
||||
"input_cost_per_second": 0.02,
|
||||
"output_cost_per_second": 0.04,
|
||||
"litellm_provider": "together_ai",
|
||||
"litellm_provider": custom_llm_provider,
|
||||
"mode": "chat",
|
||||
}
|
||||
}
|
||||
|
|
@ -3060,7 +3064,7 @@ def test_cost_per_token_per_second_pricing(monkeypatch):
|
|||
|
||||
prompt_cost, completion_cost_value = cost_per_token(
|
||||
model=model,
|
||||
custom_llm_provider="together_ai",
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
prompt_tokens=10,
|
||||
completion_tokens=20,
|
||||
response_time_ms=1500.0,
|
||||
|
|
@ -3070,6 +3074,143 @@ def test_cost_per_token_per_second_pricing(monkeypatch):
|
|||
assert completion_cost_value == pytest.approx(0.04 * 1.5)
|
||||
|
||||
|
||||
def test_cost_per_token_keeps_token_pricing_when_per_second_rates_are_also_set(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
||||
|
||||
model = "test-token-and-per-second-pricing-model"
|
||||
litellm.register_model(
|
||||
model_cost={
|
||||
model: {
|
||||
"input_cost_per_token": 1e-6,
|
||||
"output_cost_per_token": 2e-6,
|
||||
"input_cost_per_second": 0.02,
|
||||
"output_cost_per_second": 0.04,
|
||||
"litellm_provider": "openai",
|
||||
"mode": "chat",
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost_value = cost_per_token(
|
||||
model=model,
|
||||
custom_llm_provider="openai",
|
||||
prompt_tokens=10,
|
||||
completion_tokens=20,
|
||||
response_time_ms=1500.0,
|
||||
)
|
||||
|
||||
assert prompt_cost == pytest.approx(10 * 1e-6)
|
||||
assert completion_cost_value == pytest.approx(20 * 2e-6)
|
||||
|
||||
|
||||
def _logging_obj_with_call_window(duration_ms: float) -> Logging:
|
||||
start_time: Final = datetime.datetime(2026, 9, 21, 12, 0, 0)
|
||||
logging_obj: Final = Logging(
|
||||
model="gpt-5.4-nano",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
stream=False,
|
||||
call_type="completion",
|
||||
start_time=start_time,
|
||||
litellm_call_id="per-second-call-window",
|
||||
function_id="f",
|
||||
)
|
||||
logging_obj.model_call_details["start_time"] = start_time
|
||||
logging_obj.model_call_details["end_time"] = start_time + datetime.timedelta(milliseconds=duration_ms)
|
||||
return logging_obj
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("stamped_response_ms", "total_time", "logged_duration_ms", "expected_seconds"),
|
||||
[(None, 0.0, 1500.0, 1.5), (3000.0, 0.0, 1500.0, 3.0), (None, 2500.0, 1500.0, 2.5), (3000.0, 2500.0, 1500.0, 3.0)],
|
||||
)
|
||||
def test_completion_cost_per_second_deployment_bills_the_call_duration(
|
||||
monkeypatch,
|
||||
stamped_response_ms: float | None,
|
||||
total_time: float,
|
||||
logged_duration_ms: float,
|
||||
expected_seconds: float,
|
||||
):
|
||||
"""
|
||||
A deployment priced only per second bills the stamped ``_response_ms`` when there is one,
|
||||
then the caller's explicit ``total_time``, and the logging object's start/end window otherwise
|
||||
(a streamed response is never stamped).
|
||||
"""
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
||||
|
||||
deployment_id = "per-second-openai-deployment"
|
||||
litellm.register_model(
|
||||
model_cost={
|
||||
deployment_id: {
|
||||
"input_cost_per_second": 0.02,
|
||||
"output_cost_per_second": 0.04,
|
||||
"litellm_provider": "openai",
|
||||
"mode": "chat",
|
||||
}
|
||||
}
|
||||
)
|
||||
response = ModelResponse(
|
||||
model="gpt-5.4-nano",
|
||||
usage=Usage(prompt_tokens=11, completion_tokens=7, total_tokens=18),
|
||||
)
|
||||
response._response_ms = stamped_response_ms
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=response,
|
||||
model="openai/gpt-5.4-nano",
|
||||
custom_llm_provider="openai",
|
||||
custom_pricing=True,
|
||||
router_model_id=deployment_id,
|
||||
total_time=total_time,
|
||||
litellm_logging_obj=_logging_obj_with_call_window(logged_duration_ms),
|
||||
)
|
||||
|
||||
assert cost == pytest.approx((0.02 + 0.04) * expected_seconds)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["audio_transcription", "audio_speech", "video_generation", "realtime"])
|
||||
def test_cost_per_token_leaves_media_second_rates_to_their_dedicated_paths(monkeypatch, mode: str):
|
||||
"""
|
||||
A media-mode entry's per-second rates price audio or video seconds, which the dedicated
|
||||
transcription, speech, video, and realtime paths bill from the media itself, so a call that
|
||||
reaches the generic path with only a wall-clock duration must not bill them.
|
||||
"""
|
||||
model = f"test-media-per-second-{mode}"
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
model,
|
||||
{"input_cost_per_second": 0.02, "output_cost_per_second": 0.4, "litellm_provider": "openai", "mode": mode},
|
||||
)
|
||||
|
||||
assert cost_per_token(model=model, custom_llm_provider="openai", response_time_ms=2000.0) == (0.0, 0.0)
|
||||
|
||||
|
||||
def test_completion_cost_video_status_poll_bills_nothing_on_a_per_second_video_model(monkeypatch):
|
||||
"""
|
||||
Polling a video job returns a ``VideoObject`` with no stamped duration, so the cost path falls
|
||||
back to the logging object's call window; on a video model priced per output second that
|
||||
window must not be billed, or every status poll would charge for the seconds it took to answer.
|
||||
"""
|
||||
model = "test-veo-per-second-poll"
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
model,
|
||||
{"output_cost_per_second": 0.4, "litellm_provider": "vertex_ai", "mode": "video_generation"},
|
||||
)
|
||||
video = VideoObject(id="video_1", object="video", status="completed", model=model, progress=100)
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=video,
|
||||
model=model,
|
||||
custom_llm_provider="vertex_ai",
|
||||
call_type=CallTypes.video_retrieve.value,
|
||||
litellm_logging_obj=_logging_obj_with_call_window(2000.0),
|
||||
)
|
||||
|
||||
assert cost == 0.0
|
||||
|
||||
|
||||
def _batch_cache_usage() -> Usage:
|
||||
return Usage(
|
||||
prompt_tokens=11000,
|
||||
|
|
|
|||
|
|
@ -178,6 +178,14 @@ def test_potential_model_names_keeps_provider_prefixed_candidate():
|
|||
assert bare["provider_prefixed_model_name"] == bare["combined_model_name"] == "perplexity/glm-5.2"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("capability", [True, False, None])
|
||||
def test_get_model_info_anthropic_compaction(
|
||||
local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, capability: bool | None
|
||||
) -> None:
|
||||
monkeypatch.setitem(litellm.model_cost["claude-sonnet-5"], "supports_anthropic_compaction", capability)
|
||||
assert litellm.get_model_info("claude-sonnet-5")["supports_anthropic_compaction"] is capability
|
||||
|
||||
|
||||
def test_get_model_info_strips_openai_finetune_ids_without_a_custom_suffix(local_model_cost_map):
|
||||
info = litellm.get_model_info(model="ft:gpt-4o-2024-08-06:my-org::abc123", custom_llm_provider="openai")
|
||||
assert info["key"] == "ft:gpt-4o-2024-08-06"
|
||||
|
|
@ -861,6 +869,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
|
|||
"source": {"type": "string"},
|
||||
"comment": {"type": "string"},
|
||||
"supports_assistant_prefill": {"type": "boolean"},
|
||||
"supports_anthropic_compaction": {"type": "boolean"},
|
||||
"supports_audio_input": {"type": "boolean"},
|
||||
"supports_audio_output": {"type": "boolean"},
|
||||
"gemini_native_audio": {"type": "boolean"},
|
||||
|
|
|
|||
|
|
@ -7,10 +7,12 @@ through _hidden_params to the x-litellm-callback-duration-ms response header.
|
|||
|
||||
import asyncio
|
||||
import datetime
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
import litellm.litellm_core_utils.llm_response_utils.response_metadata as response_metadata_mod
|
||||
import litellm.proxy.common_request_processing as common_request_processing_mod
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
|
@ -22,7 +24,7 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
|
|||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.types.utils import ModelResponse
|
||||
from litellm.types.utils import ModelResponse, Usage
|
||||
|
||||
|
||||
class TestCallbackDurationMs:
|
||||
|
|
@ -583,3 +585,57 @@ class TestLoggingInitCallbackDuration:
|
|||
# Should still be set (deep copy of None is essentially a no-op)
|
||||
assert hasattr(obj, "callback_duration_ms")
|
||||
assert obj.callback_duration_ms >= 0
|
||||
|
||||
|
||||
def test_update_response_metadata_prices_per_second_deployment_from_its_stamped_duration(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
||||
|
||||
deployment_id: Final = "per-second-deployment-response-metadata"
|
||||
litellm.register_model(
|
||||
model_cost={
|
||||
deployment_id: {
|
||||
"input_cost_per_second": 0.02,
|
||||
"output_cost_per_second": 0.04,
|
||||
"litellm_provider": "openai",
|
||||
"mode": "chat",
|
||||
}
|
||||
}
|
||||
)
|
||||
start_time: Final = datetime.datetime(2026, 9, 21, 12, 0, 0)
|
||||
logging_obj: Final = Logging(
|
||||
model="gpt-5.4-nano",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
stream=False,
|
||||
call_type="completion",
|
||||
start_time=start_time,
|
||||
litellm_call_id="per-second-response-metadata",
|
||||
function_id="f",
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model="gpt-5.4-nano",
|
||||
litellm_params={
|
||||
"input_cost_per_second": 0.02,
|
||||
"output_cost_per_second": 0.04,
|
||||
"metadata": {"model_info": {"id": deployment_id}},
|
||||
},
|
||||
optional_params={},
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
logging_obj.model_call_details["end_time"] = start_time + datetime.timedelta(seconds=10)
|
||||
result: Final = ModelResponse(
|
||||
model="gpt-5.4-nano",
|
||||
usage=Usage(prompt_tokens=11, completion_tokens=7, total_tokens=18),
|
||||
)
|
||||
|
||||
update_response_metadata(
|
||||
result=result,
|
||||
logging_obj=logging_obj,
|
||||
model="gpt-5.4-nano",
|
||||
kwargs={"model_info": {"id": deployment_id}},
|
||||
start_time=start_time,
|
||||
end_time=start_time + datetime.timedelta(seconds=2),
|
||||
)
|
||||
|
||||
assert result._response_ms == pytest.approx(2000)
|
||||
assert result._hidden_params["response_cost"] == pytest.approx((0.02 + 0.04) * 2)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ import { describe, it, expect, vi, beforeEach } from "vitest";
|
|||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import { MCPServerView } from "./mcp_server_view";
|
||||
import * as networking from "@/components/networking";
|
||||
import { setSecureItem } from "@/utils/secureStorage";
|
||||
import { EDIT_OAUTH_UI_STATE_KEY } from "./mcp_server_edit";
|
||||
import type { MCPServer } from "@/components/mcp_tools/types";
|
||||
|
||||
vi.mock(".", () => ({
|
||||
|
|
@ -68,6 +70,7 @@ const openUserCredentials = async (props: Record<string, unknown>) => {
|
|||
describe("MCPServerView", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
sessionStorage.clear();
|
||||
});
|
||||
|
||||
// Name, alias and description each label the header and a Settings row, so
|
||||
|
|
@ -146,6 +149,37 @@ describe("MCPServerView", () => {
|
|||
expect(screen.queryByRole("button", { name: "Edit Settings" })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it.each([false, true])("keeps config settings read-only with isEditing=%s", async (isEditing) => {
|
||||
renderView({ is_config: true }, { isEditing });
|
||||
|
||||
await userEvent.click(screen.getByRole("tab", { name: "Settings" }));
|
||||
|
||||
expect(screen.getByRole("button", { name: "Edit Settings" })).toBeDisabled();
|
||||
expect(screen.getByText("Defined in config. Edit your YAML configuration to make changes")).toBeVisible();
|
||||
expect(screen.queryByText("edit form")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it.each([true, false])("honors config read-only state on OAuth return: %s", async (isConfig) => {
|
||||
setSecureItem(EDIT_OAUTH_UI_STATE_KEY, JSON.stringify({ serverId: "srv-1" }));
|
||||
renderView({ is_config: isConfig });
|
||||
|
||||
await userEvent.click(screen.getByRole("tab", { name: "Settings" }));
|
||||
|
||||
if (isConfig) {
|
||||
expect(screen.queryByText("edit form")).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Edit Settings" })).toBeDisabled();
|
||||
} else {
|
||||
expect(screen.getByText("edit form")).toBeVisible();
|
||||
}
|
||||
});
|
||||
|
||||
it("does not open the editor for a view-only admin", async () => {
|
||||
renderView({}, { isViewOnly: true, isEditing: true });
|
||||
await userEvent.click(screen.getByRole("tab", { name: "Settings" }));
|
||||
expect(screen.getByRole("button", { name: "Edit Settings" })).toBeDisabled();
|
||||
expect(screen.queryByText("edit form")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("opens on the tab named by initialTabIndex", async () => {
|
||||
renderView({}, { initialTabIndex: 1 });
|
||||
|
||||
|
|
|
|||
|
|
@ -62,7 +62,8 @@ export const MCPServerView: React.FC<MCPServerViewProps> = ({
|
|||
}) => {
|
||||
// Open the editing Settings tab on first render when returning from the edit OAuth
|
||||
// redirect, so the "token fetched" feedback shows where the user left off (Settings=2).
|
||||
const returningFromEditOAuth = isReturningFromEditOAuth(isProxyAdmin, mcpServer.server_id);
|
||||
const canEdit = isProxyAdmin && !isViewOnly && !mcpServer.is_config;
|
||||
const returningFromEditOAuth = isReturningFromEditOAuth(canEdit, mcpServer.server_id);
|
||||
const [editing, setEditing] = useState(isEditing || returningFromEditOAuth);
|
||||
const [showFullUrl, setShowFullUrl] = useState(false);
|
||||
const [copiedStates, setCopiedStates] = useState<Record<string, boolean>>({});
|
||||
|
|
@ -224,13 +225,18 @@ export const MCPServerView: React.FC<MCPServerViewProps> = ({
|
|||
<Card className="p-6">
|
||||
<div className="mb-4 flex items-center justify-between">
|
||||
<h2 className="text-lg font-medium">MCP Server Settings</h2>
|
||||
{editing ? null : (
|
||||
<Button variant="outline" onClick={() => setEditing(true)}>
|
||||
{editing && canEdit ? null : (
|
||||
<Button variant="outline" disabled={!canEdit} onClick={() => setEditing(true)}>
|
||||
Edit Settings
|
||||
</Button>
|
||||
)}
|
||||
</div>
|
||||
{editing ? (
|
||||
{mcpServer.is_config && (
|
||||
<p className="mb-4 text-sm text-muted-foreground">
|
||||
Defined in config. Edit your YAML configuration to make changes
|
||||
</p>
|
||||
)}
|
||||
{editing && canEdit ? (
|
||||
<MCPServerEdit
|
||||
mcpServer={mcpServer}
|
||||
accessToken={accessToken}
|
||||
|
|
|
|||
|
|
@ -407,6 +407,7 @@ export interface MCPToolsViewerProps {
|
|||
|
||||
export interface MCPServer {
|
||||
server_id: string;
|
||||
is_config?: boolean;
|
||||
server_name?: string | null;
|
||||
alias?: string | null;
|
||||
description?: string | null;
|
||||
|
|
|
|||
31
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
31
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -27678,6 +27678,26 @@ export interface components {
|
|||
*/
|
||||
pattern_type: "prebuilt" | "regex";
|
||||
};
|
||||
/** ContextCompactionConfig */
|
||||
ContextCompactionConfig: {
|
||||
/**
|
||||
* Max Tokens
|
||||
* @default 4096
|
||||
*/
|
||||
max_tokens: number;
|
||||
/** Model */
|
||||
model?: string | null;
|
||||
/**
|
||||
* Timeout Seconds
|
||||
* @default 120
|
||||
*/
|
||||
timeout_seconds: number;
|
||||
/**
|
||||
* Trigger Ratio
|
||||
* @default 0.9
|
||||
*/
|
||||
trigger_ratio: number;
|
||||
};
|
||||
/**
|
||||
* CoordinationRedisNode
|
||||
* @description A single startup node of a cluster-mode Redis used for proxy coordination.
|
||||
|
|
@ -30660,6 +30680,12 @@ export interface components {
|
|||
* @default false
|
||||
*/
|
||||
is_byok: boolean;
|
||||
/**
|
||||
* Is Config
|
||||
* @description Whether this server is defined in config and is read-only.
|
||||
* @default false
|
||||
*/
|
||||
is_config: boolean;
|
||||
/** Issuer */
|
||||
issuer?: string | null;
|
||||
/** Last Health Check */
|
||||
|
|
@ -36962,6 +36988,11 @@ export interface components {
|
|||
* @description Keywords indicating code-related content
|
||||
*/
|
||||
code_keywords?: string[] | null;
|
||||
/**
|
||||
* Context Compaction
|
||||
* @description Compact full conversation history near the selected deployment's input limit for Chat, Responses and Messages. Uses a capable configured tier model unless model is specified. Set false or null to disable. Stored and client-managed native history keep their existing behavior.
|
||||
*/
|
||||
context_compaction?: components["schemas"]["ContextCompactionConfig"] | false;
|
||||
/**
|
||||
* Context Window Escalation Buffer
|
||||
* @description Fraction of a model's declared context window the estimated prompt must fit within. The token count is an estimate, so fitting against the full window would dispatch prompts that the provider's own tokenizer then rejects; 0.95 leaves room for that drift plus the response tokens.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue