mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
Merge remote-tracking branch 'origin/main' into litellm_realtime_health_check_credential_name
This commit is contained in:
commit
1771255b32
143 changed files with 7177 additions and 1057 deletions
|
|
@ -2915,6 +2915,25 @@ jobs:
|
|||
exit 1
|
||||
fi
|
||||
|
||||
provider_replay_harness:
|
||||
docker:
|
||||
- *python312_image
|
||||
working_directory: ~/project
|
||||
resource_class: medium
|
||||
steps:
|
||||
- setup_litellm_test_deps
|
||||
- run:
|
||||
name: Test provider replay harness
|
||||
command: |
|
||||
mkdir -p test-results/provider-replay-harness
|
||||
uv run --no-sync pytest -q --noconftest -o addopts= -o pythonpath=tests/e2e -p no:rerunfailures \
|
||||
--junitxml=test-results/provider-replay-harness/junit.xml \
|
||||
tests/e2e/test_provider_edge.py tests/e2e/test_fixture_bundle.py \
|
||||
tests/e2e/test_fixture_canonical.py tests/e2e/test_fixture_mode.py \
|
||||
tests/code_coverage_tests/test_provider_replay_harness.py
|
||||
- store_test_results:
|
||||
path: test-results/provider-replay-harness
|
||||
|
||||
integration_contracts:
|
||||
parameters:
|
||||
suite:
|
||||
|
|
@ -2967,6 +2986,7 @@ workflows:
|
|||
only:
|
||||
- main
|
||||
- /litellm_.*/
|
||||
- provider_replay_harness
|
||||
- base_sdk_install:
|
||||
filters: *main_branches
|
||||
- local_testing_part1:
|
||||
|
|
|
|||
15
.github/e2e-stack/assert_tests_ran.py
vendored
15
.github/e2e-stack/assert_tests_ran.py
vendored
|
|
@ -3,6 +3,9 @@ import xml.etree.ElementTree as ET
|
|||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e"))
|
||||
from coverage_registry.management_cases import MANAGEMENT_CASES
|
||||
|
||||
|
||||
def main() -> int:
|
||||
selected: Final = tuple(sys.argv[2:])
|
||||
|
|
@ -16,6 +19,17 @@ def main() -> int:
|
|||
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
|
||||
)
|
||||
missing: Final = tuple(path for path in selected if path not in passed)
|
||||
required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected)
|
||||
passed_nodes: Final = frozenset(
|
||||
prop.get("value")
|
||||
for case in cases
|
||||
if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
|
||||
for prop in case.findall("./properties/property")
|
||||
if prop.get("name") == "management_node"
|
||||
)
|
||||
missing_nodes: Final = required_nodes - passed_nodes
|
||||
for node in sorted(missing_nodes):
|
||||
_ = sys.stdout.write(f"::error::required management case did not pass: {node}\n")
|
||||
for path in selected:
|
||||
collected: Final = sum(case.get("file") == path for case in cases)
|
||||
skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases)
|
||||
|
|
@ -27,6 +41,7 @@ def main() -> int:
|
|||
if (
|
||||
selected
|
||||
and not missing
|
||||
and not missing_nodes
|
||||
and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error"))
|
||||
):
|
||||
return 0
|
||||
|
|
|
|||
5
.github/e2e-stack/oidc-profile.sh
vendored
Executable file
5
.github/e2e-stack/oidc-profile.sh
vendored
Executable file
|
|
@ -0,0 +1,5 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
|
||||
cd "${REPO_ROOT}"
|
||||
exec uv run --no-sync python tests/e2e/idp.py "$@"
|
||||
2
.github/e2e-stack/select_tests.py
vendored
2
.github/e2e-stack/select_tests.py
vendored
|
|
@ -12,6 +12,8 @@ UNSUPPORTED: Final = re.compile(
|
|||
HARNESS: Final = re.compile(
|
||||
r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$"
|
||||
r"|^tests/e2e/idp_realm\.json$"
|
||||
r"|^tests/e2e/management/(management_client|jwt_actors|conftest)\.py$"
|
||||
r"|^tests/e2e/coverage_registry/management_cases\.py$"
|
||||
r"|^tests/e2e/gateway/"
|
||||
r"|^\.github/e2e-stack/"
|
||||
r"|^\.github/workflows/test-e2e-changed\.yml$"
|
||||
|
|
|
|||
2
.github/workflows/_test-unit-base.yml
vendored
2
.github/workflows/_test-unit-base.yml
vendored
|
|
@ -57,6 +57,7 @@ permissions:
|
|||
|
||||
env:
|
||||
UV_PYTHON: "3.12"
|
||||
LITELLM_LOCAL_MODEL_COST_MAP: "True"
|
||||
|
||||
jobs:
|
||||
run:
|
||||
|
|
@ -113,6 +114,7 @@ jobs:
|
|||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 8
|
||||
run: |
|
||||
diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
|
||||
uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'
|
||||
|
||||
|
|
|
|||
2
.github/workflows/test-e2e-changed.yml
vendored
2
.github/workflows/test-e2e-changed.yml
vendored
|
|
@ -183,7 +183,7 @@ jobs:
|
|||
log="${RUNNER_TEMP}/e2e-pass-${pass}.log"
|
||||
echo "::group::pass ${pass} of 3"
|
||||
set +e
|
||||
uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v -p no:cacheprovider \
|
||||
uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v --reruns 0 -p no:cacheprovider \
|
||||
-o junit_family=xunit1 --junitxml="${report}" > "${log}" 2>&1
|
||||
status=$?
|
||||
uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}"
|
||||
|
|
|
|||
|
|
@ -538,7 +538,7 @@ context_window_fallbacks: Optional[List] = None
|
|||
content_policy_fallbacks: Optional[List] = None
|
||||
allowed_fails: int = 3
|
||||
allow_dynamic_callback_disabling: bool = True
|
||||
num_retries_per_request: Optional[int] = None # cap on Router retries of one model group; resets per fallback hop
|
||||
num_retries_per_request: Optional[int] = None # for the request overall (incl. fallbacks + model retries)
|
||||
####### SECRET MANAGERS #####################
|
||||
secret_manager_client: Optional[Any] = (
|
||||
None # list of instantiated key management clients - e.g. azure kv, infisical, etc.
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import litellm
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
|
||||
from litellm.llms.vertex_ai.batches.transformation import vertex_prompt_tokens_details
|
||||
from litellm.types.llms.openai import Batch
|
||||
from litellm.types.utils import ModelInfo, Usage
|
||||
from litellm.utils import token_counter
|
||||
|
|
@ -356,6 +357,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
|
|||
prompt_tokens=_prompt,
|
||||
completion_tokens=_completion,
|
||||
total_tokens=_total,
|
||||
prompt_tokens_details=vertex_prompt_tokens_details(usage_metadata),
|
||||
)
|
||||
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -205,21 +205,35 @@ def _extract_anthropic_tool_exchange_spans(
|
|||
return spans, None
|
||||
|
||||
|
||||
def _message_has_cache_control(message: Mapping[str, object]) -> bool:
|
||||
if message.get("cache_control") is not None:
|
||||
return True
|
||||
content: Final = message.get("content")
|
||||
if isinstance(content, list):
|
||||
return any(isinstance(part, Mapping) and part.get("cache_control") is not None for part in content)
|
||||
return False
|
||||
|
||||
|
||||
def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
|
||||
"""
|
||||
Return indices of messages that must never be compressed:
|
||||
- All system messages
|
||||
- The last user message
|
||||
- The last assistant message
|
||||
- Any message carrying an Anthropic cache_control breakpoint
|
||||
|
||||
The last user message is what the model is being asked to act on right now,
|
||||
so compressing it replaces the live instruction with a marker. Compression
|
||||
guardrails share this policy; see the Headroom guardrail.
|
||||
guardrails share this policy; see the Headroom guardrail. A cache_control
|
||||
breakpoint pins the provider's prompt-cache prefix to that row's exact
|
||||
bytes, so rewriting a marked row anywhere in history turns the next
|
||||
request's cache read into a cache write.
|
||||
"""
|
||||
system_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "system")
|
||||
last_user: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "user")[-1:]
|
||||
last_assistant = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[-1:]
|
||||
return system_indices + last_user + last_assistant
|
||||
assistant_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")
|
||||
cache_control_indices: Final = tuple(index for index, msg in enumerate(messages) if _message_has_cache_control(msg))
|
||||
return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + cache_control_indices))
|
||||
|
||||
|
||||
def _combine_scores(
|
||||
|
|
@ -421,7 +435,7 @@ def compress(
|
|||
combined_scores = bm25_scores
|
||||
|
||||
# Protected messages are never compressed
|
||||
protected_indices: Final = get_protected_indices(normalized_messages)
|
||||
protected_indices: Final = get_protected_indices(original_messages)
|
||||
kept_indices: set[int] = set(protected_indices)
|
||||
|
||||
tool_exchange_spans: list[set[int]] = []
|
||||
|
|
|
|||
|
|
@ -1976,6 +1976,8 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset(
|
|||
|
||||
UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS
|
||||
|
||||
STRINGIFIED_NONE: Final[str] = "None"
|
||||
|
||||
# A retrieved response replays the usage of the call that created it, so pricing these
|
||||
# read/management routes like inference bills the same tokens twice.
|
||||
NON_INFERENCE_CALL_TYPES: Final[frozenset[str]] = frozenset(
|
||||
|
|
|
|||
|
|
@ -2278,6 +2278,19 @@ def default_video_cost_calculator(
|
|||
return 0.0
|
||||
|
||||
|
||||
def _batch_rate(
|
||||
model_info: ModelInfo,
|
||||
key: Literal[
|
||||
"input_cost_per_audio_token_batches",
|
||||
"input_cost_per_image_token_batches",
|
||||
"input_cost_per_video_token_batches",
|
||||
],
|
||||
fallback: float,
|
||||
) -> float:
|
||||
rate: Final = model_info.get(key)
|
||||
return fallback if rate is None else rate
|
||||
|
||||
|
||||
def batch_cost_calculator(
|
||||
usage: Usage,
|
||||
model: str,
|
||||
|
|
@ -2337,7 +2350,29 @@ def batch_cost_calculator(
|
|||
total_prompt_cost = 0.0
|
||||
total_completion_cost = 0.0
|
||||
if input_cost_per_token_batches is not None:
|
||||
total_prompt_cost = usage.prompt_tokens * input_cost_per_token_batches
|
||||
batch_details: Final = parse_prompt_tokens_details(usage)
|
||||
audio_tokens, image_tokens, video_tokens = (
|
||||
batch_details["audio_tokens"],
|
||||
batch_details["image_tokens"],
|
||||
batch_details["video_tokens"],
|
||||
)
|
||||
modality_rates: Final = (
|
||||
_batch_rate(model_info, "input_cost_per_audio_token_batches", input_cost_per_token_batches),
|
||||
_batch_rate(model_info, "input_cost_per_image_token_batches", input_cost_per_token_batches),
|
||||
_batch_rate(model_info, "input_cost_per_video_token_batches", input_cost_per_token_batches),
|
||||
)
|
||||
total_prompt_cost = sum(
|
||||
tokens * rate
|
||||
for tokens, rate in zip(
|
||||
(
|
||||
max((usage.prompt_tokens or 0) - audio_tokens - image_tokens - video_tokens, 0),
|
||||
audio_tokens,
|
||||
image_tokens,
|
||||
video_tokens,
|
||||
),
|
||||
(input_cost_per_token_batches, *modality_rates),
|
||||
)
|
||||
)
|
||||
elif input_cost_per_token:
|
||||
details: Final = parse_prompt_tokens_details(usage)
|
||||
cache_read_tokens: Final = details["cache_hit_tokens"]
|
||||
|
|
|
|||
|
|
@ -309,8 +309,8 @@ def max_retries_per_request_hit(kwargs: Mapping[str, object], num_retries_per_re
|
|||
metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs))
|
||||
if not isinstance(metadata, Mapping):
|
||||
return False
|
||||
attempted_retries: Final = metadata.get("attempted_retries")
|
||||
return type(attempted_retries) is int and 0 < attempted_retries and num_retries_per_request <= attempted_retries
|
||||
retry_count: Final = metadata.get("request_retry_count")
|
||||
return type(retry_count) is int and 0 < retry_count and num_retries_per_request <= retry_count
|
||||
|
||||
|
||||
def get_or_create_metadata_bucket(
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
import inspect
|
||||
import json
|
||||
import re
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Protocol, cast
|
||||
|
||||
import httpx
|
||||
|
|
@ -202,11 +205,17 @@ def _get_response_headers(original_exception: Exception) -> httpx.Headers | None
|
|||
return _response_headers
|
||||
|
||||
|
||||
def _accepted_init_kwargs(exception_class: type[Exception], candidates: Mapping[str, object]) -> Mapping[str, object]:
|
||||
accepted: Final = inspect.signature(exception_class).parameters
|
||||
return MappingProxyType({name: value for name, value in candidates.items() if name in accepted})
|
||||
|
||||
|
||||
def extract_and_raise_litellm_exception(
|
||||
response: Any | None,
|
||||
error_str: str,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
body: object | None = None,
|
||||
):
|
||||
"""
|
||||
Covers scenario where litellm sdk calling proxy.
|
||||
|
|
@ -216,32 +225,19 @@ def extract_and_raise_litellm_exception(
|
|||
Relevant Issue: https://github.com/BerriAI/litellm/issues/7259
|
||||
"""
|
||||
pattern: Final = r"litellm\.\w+Error"
|
||||
|
||||
# Search for the exception in the error string
|
||||
match: Final = re.search(pattern, error_str)
|
||||
|
||||
# Extract the exception if found
|
||||
if match:
|
||||
exception_name = match.group(0)
|
||||
exception_name = exception_name.strip().replace("litellm.", "")
|
||||
raised_exception_obj: Final = getattr(litellm, exception_name, None)
|
||||
if raised_exception_obj:
|
||||
# Try with response parameter first, fall back to without it
|
||||
# Some exceptions (e.g., APIConnectionError) don't accept response param
|
||||
try:
|
||||
raise raised_exception_obj(
|
||||
message=error_str,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=response,
|
||||
)
|
||||
except TypeError:
|
||||
# Exception doesn't accept response parameter
|
||||
raise raised_exception_obj(
|
||||
message=error_str,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
)
|
||||
if match is None:
|
||||
return
|
||||
exception_name: Final = match.group(0).removeprefix("litellm.")
|
||||
raised_exception_obj: Final = getattr(litellm, exception_name, None)
|
||||
if not raised_exception_obj:
|
||||
return
|
||||
raise raised_exception_obj(
|
||||
message=error_str,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
**_accepted_init_kwargs(raised_exception_obj, MappingProxyType({"response": response, "body": body})),
|
||||
)
|
||||
|
||||
|
||||
class _ProviderHTTPException(Protocol):
|
||||
|
|
@ -254,6 +250,23 @@ class _ProviderHTTPException(Protocol):
|
|||
llm_provider: str
|
||||
|
||||
|
||||
def _litellm_proxy_response(
|
||||
original_exception: _ProviderHTTPException, custom_llm_provider: str
|
||||
) -> httpx.Response | None:
|
||||
response: Final = getattr(original_exception, "response", None)
|
||||
if custom_llm_provider != "litellm_proxy" or not isinstance(response, httpx.Response) or response.headers:
|
||||
return response
|
||||
headers: Final = getattr(original_exception, "headers", None)
|
||||
if not isinstance(headers, Mapping) or not headers:
|
||||
return response
|
||||
pairs: Final = headers.multi_items() if isinstance(headers, httpx.Headers) else headers.items()
|
||||
return httpx.Response(
|
||||
status_code=response.status_code,
|
||||
headers=[(str(k), str(v)) for k, v in pairs],
|
||||
request=getattr(original_exception, "request", None),
|
||||
)
|
||||
|
||||
|
||||
def _map_openai_exception(
|
||||
*,
|
||||
model: str,
|
||||
|
|
@ -264,6 +277,7 @@ def _map_openai_exception(
|
|||
exception_provider: str,
|
||||
extra_information: str,
|
||||
) -> None:
|
||||
response: Final = _litellm_proxy_response(original_exception, custom_llm_provider)
|
||||
# custom_llm_provider is openai, make it OpenAI
|
||||
message = get_error_message(error_obj=original_exception)
|
||||
if message is None:
|
||||
|
|
@ -292,14 +306,14 @@ def _map_openai_exception(
|
|||
message=f"RateLimitError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
)
|
||||
elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str):
|
||||
raise ContextWindowExceededError(
|
||||
message=f"ContextWindowExceededError: {exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif "invalid_request_error" in error_str and "model_not_found" in error_str:
|
||||
|
|
@ -307,7 +321,7 @@ def _map_openai_exception(
|
|||
message=f"{exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif "A timeout occurred" in error_str:
|
||||
|
|
@ -326,8 +340,9 @@ def _map_openai_exception(
|
|||
message=f"ContentPolicyViolationError: {exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
elif "invalid_encrypted_content" in error_str or "could not be verified" in error_str:
|
||||
helpful_message: Final = (
|
||||
|
|
@ -345,7 +360,7 @@ def _map_openai_exception(
|
|||
message=helpful_message,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
|
|
@ -354,7 +369,7 @@ def _map_openai_exception(
|
|||
message=f"{exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
|
|
@ -372,7 +387,7 @@ def _map_openai_exception(
|
|||
message=f"RateLimitError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif (
|
||||
|
|
@ -383,7 +398,7 @@ def _map_openai_exception(
|
|||
message=f"AuthenticationError: {exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif "Mistral API raised a streaming error" in error_str:
|
||||
|
|
@ -402,15 +417,16 @@ def _map_openai_exception(
|
|||
message=f"{exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
elif original_exception.status_code == 401:
|
||||
raise AuthenticationError(
|
||||
message=f"AuthenticationError: {exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 404:
|
||||
|
|
@ -418,7 +434,7 @@ def _map_openai_exception(
|
|||
message=f"NotFoundError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 408:
|
||||
|
|
@ -433,7 +449,7 @@ def _map_openai_exception(
|
|||
message=f"{exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
|
|
@ -442,7 +458,7 @@ def _map_openai_exception(
|
|||
message=f"RateLimitError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 500:
|
||||
|
|
@ -450,7 +466,7 @@ def _map_openai_exception(
|
|||
message=f"InternalServerError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 502:
|
||||
|
|
@ -458,7 +474,7 @@ def _map_openai_exception(
|
|||
message=f"BadGatewayError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 503:
|
||||
|
|
@ -466,7 +482,7 @@ def _map_openai_exception(
|
|||
message=f"ServiceUnavailableError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 504: # gateway timeout error
|
||||
|
|
@ -2423,10 +2439,11 @@ def exception_type(
|
|||
custom_llm_provider == "litellm_proxy"
|
||||
): # handle special case where calling litellm proxy + exception str contains error message
|
||||
extract_and_raise_litellm_exception(
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=_litellm_proxy_response(mappable_exception, custom_llm_provider),
|
||||
error_str=error_str,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
if (
|
||||
custom_llm_provider == "openai"
|
||||
|
|
|
|||
|
|
@ -484,11 +484,12 @@ def apply_off_peak_pricing(model_info: ModelInfo, current_time: datetime | None,
|
|||
def _apply_off_peak_to_base_costs(
|
||||
model_info: ModelInfo,
|
||||
current_time: datetime | None,
|
||||
base_costs: tuple[float, float, float, float, float],
|
||||
base_costs: tuple[float, float, float, float | None, float],
|
||||
) -> tuple[float, float, float, float, float]:
|
||||
"""Apply off-peak rates to an already-resolved set of base costs, whichever pricing path
|
||||
produced them. The one-hour cache-creation rate passes through untouched, since
|
||||
off_peak_pricing has no field for it, and reasoning is left to _resolve_billed_reasoning_rate.
|
||||
produced them. off_peak_pricing has no field for the one-hour cache-creation rate, so a
|
||||
present one passes through untouched and an absent one resolves to the applied
|
||||
cache-creation rate. Reasoning is left to _resolve_billed_reasoning_rate.
|
||||
"""
|
||||
prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs
|
||||
rates: Final = apply_off_peak_pricing(
|
||||
|
|
@ -506,7 +507,7 @@ def _apply_off_peak_to_base_costs(
|
|||
rates.input_rate,
|
||||
rates.output_rate,
|
||||
rates.cache_creation_rate,
|
||||
cache_creation_above_1hr,
|
||||
rates.cache_creation_rate if cache_creation_above_1hr is None else cache_creation_above_1hr,
|
||||
rates.cache_read_rate,
|
||||
)
|
||||
|
||||
|
|
@ -532,6 +533,11 @@ def _get_token_base_cost(
|
|||
`missing_cache_read_uses_input` resolves an absent cache-read rate to the resolved
|
||||
input rate instead of 0.0; an explicit 0.0 rate stays a real price either way.
|
||||
|
||||
An absent cache-creation rate always resolves to the resolved input rate, the way the
|
||||
tiered table and custom deployment pricing already do, since a provider that publishes
|
||||
no write price bills cache writes as ordinary input. An absent 1h write rate resolves
|
||||
to the cache-creation rate, off-peak included. An explicit 0.0 stays a real price for both.
|
||||
|
||||
Returns:
|
||||
Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost)
|
||||
"""
|
||||
|
|
@ -554,10 +560,9 @@ def _get_token_base_cost(
|
|||
output_image_cost: Final = _get_cost_per_unit(model_info, "output_cost_per_image_token", None)
|
||||
if output_image_cost is not None:
|
||||
completion_base_cost = cast(float, output_image_cost)
|
||||
cache_creation_cost = cast(float, _get_cost_per_unit(model_info, cache_creation_cost_key))
|
||||
cache_creation_cost_above_1hr = cast(
|
||||
float,
|
||||
_get_cost_per_unit(model_info, "cache_creation_input_token_cost_above_1hr"),
|
||||
cache_creation_cost = _get_cost_per_unit(model_info, cache_creation_cost_key, default_value=None)
|
||||
cache_creation_cost_above_1hr = _get_cost_per_unit(
|
||||
model_info, "cache_creation_input_token_cost_above_1hr", default_value=None
|
||||
)
|
||||
cache_read_cost = _get_cost_per_unit(model_info, cache_read_cost_key, default_value=None)
|
||||
|
||||
|
|
@ -639,22 +644,10 @@ def _get_token_base_cost(
|
|||
else f"cache_read_input_token_cost_above_{threshold_str}_tokens"
|
||||
)
|
||||
|
||||
cache_creation_cost = cast(
|
||||
float,
|
||||
_get_cost_per_unit(
|
||||
model_info,
|
||||
cache_creation_tiered_key,
|
||||
cache_creation_cost,
|
||||
),
|
||||
)
|
||||
cache_creation_cost = _get_cost_per_unit(model_info, cache_creation_tiered_key, cache_creation_cost)
|
||||
|
||||
cache_creation_cost_above_1hr = cast(
|
||||
float,
|
||||
_get_cost_per_unit(
|
||||
model_info,
|
||||
cache_creation_1hr_tiered_key,
|
||||
cache_creation_cost_above_1hr,
|
||||
),
|
||||
cache_creation_cost_above_1hr = _get_cost_per_unit(
|
||||
model_info, cache_creation_1hr_tiered_key, cache_creation_cost_above_1hr
|
||||
)
|
||||
|
||||
cache_read_cost = _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost)
|
||||
|
|
@ -665,16 +658,16 @@ def _get_token_base_cost(
|
|||
except Exception:
|
||||
continue
|
||||
|
||||
input_rate_for_missing_cache_rates: Final = _off_peak_rate(
|
||||
_open_off_peak_block(model_info, current_time) or MappingProxyType({}),
|
||||
"input_cost_per_token",
|
||||
prompt_base_cost,
|
||||
)
|
||||
if cache_read_cost is None:
|
||||
cache_read_cost = (
|
||||
_off_peak_rate(
|
||||
_open_off_peak_block(model_info, current_time) or MappingProxyType({}),
|
||||
"input_cost_per_token",
|
||||
prompt_base_cost,
|
||||
)
|
||||
if missing_cache_read_uses_input
|
||||
else 0.0
|
||||
)
|
||||
cache_read_cost = input_rate_for_missing_cache_rates if missing_cache_read_uses_input else 0.0
|
||||
resolved_cache_creation_cost: Final = (
|
||||
input_rate_for_missing_cache_rates if cache_creation_cost is None else cache_creation_cost
|
||||
)
|
||||
|
||||
return _apply_off_peak_to_base_costs(
|
||||
model_info,
|
||||
|
|
@ -682,7 +675,7 @@ def _get_token_base_cost(
|
|||
(
|
||||
prompt_base_cost,
|
||||
completion_base_cost,
|
||||
cache_creation_cost,
|
||||
resolved_cache_creation_cost,
|
||||
cache_creation_cost_above_1hr,
|
||||
cache_read_cost,
|
||||
),
|
||||
|
|
@ -956,12 +949,16 @@ def _calculate_input_cost(
|
|||
)
|
||||
|
||||
### AUDIO COST
|
||||
if prompt_tokens_details["audio_tokens"]:
|
||||
if prompt_tokens_details["audio_tokens"] and not (
|
||||
prompt_tokens_details["audio_length_seconds"] and model_info.get("input_cost_per_audio_per_second") is not None
|
||||
):
|
||||
audio_cost_key: Final = _get_service_tier_cost_key("input_cost_per_audio_token", service_tier)
|
||||
prompt_cost += calculate_cost_component(model_info, audio_cost_key, prompt_tokens_details["audio_tokens"])
|
||||
|
||||
### IMAGE TOKEN COST
|
||||
if prompt_tokens_details["image_tokens"]:
|
||||
if prompt_tokens_details["image_tokens"] and not (
|
||||
prompt_tokens_details["image_count"] and model_info.get("input_cost_per_image") is not None
|
||||
):
|
||||
# For image token costs:
|
||||
# First check if input_cost_per_image_token is available. If not, default to generic input_cost_per_token.
|
||||
image_token_cost_key = "input_cost_per_image_token"
|
||||
|
|
@ -970,7 +967,9 @@ def _calculate_input_cost(
|
|||
prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"])
|
||||
|
||||
### VIDEO TOKEN COST
|
||||
if prompt_tokens_details["video_tokens"]:
|
||||
if prompt_tokens_details["video_tokens"] and not (
|
||||
prompt_tokens_details["video_length_seconds"] and model_info.get("input_cost_per_video_per_second") is not None
|
||||
):
|
||||
video_token_cost_key = "input_cost_per_video_token"
|
||||
if model_info.get(video_token_cost_key) is None:
|
||||
video_token_cost_key = "input_cost_per_token"
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from itertools import chain, repeat
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict, assert_never
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -44,6 +45,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
|
|||
scoped_structured_message_indices,
|
||||
stream_item_field,
|
||||
stream_item_fingerprint,
|
||||
unappliable_request_rewrite,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
|
||||
AnthropicPassthroughLoggingHandler,
|
||||
|
|
@ -103,9 +105,24 @@ class ToolResultBlockTextTarget:
|
|||
block_idx: int
|
||||
|
||||
|
||||
InputWriteBackTarget = (
|
||||
MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget
|
||||
)
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SystemStringTarget:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SystemBlockTextTarget:
|
||||
block_idx: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ToolUseInputTarget:
|
||||
msg_idx: int
|
||||
content_idx: int
|
||||
|
||||
|
||||
MessageTextTarget = MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget
|
||||
InputWriteBackTarget = SystemStringTarget | SystemBlockTextTarget | MessageTextTarget
|
||||
|
||||
|
||||
def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]:
|
||||
|
|
@ -146,10 +163,17 @@ class ScannedText:
|
|||
target: InputWriteBackTarget
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ScannedToolCall:
|
||||
tool_call: ChatCompletionToolCallChunk
|
||||
target: ToolUseInputTarget
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ExtractedInput:
|
||||
scanned: tuple[ScannedText, ...]
|
||||
images: tuple[str, ...]
|
||||
tool_calls: tuple[ScannedToolCall, ...] = ()
|
||||
|
||||
|
||||
EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=())
|
||||
|
|
@ -161,6 +185,74 @@ class _ToolCallShape:
|
|||
arguments: str
|
||||
|
||||
|
||||
def _is_client_tool_use(block: Mapping[str, object]) -> bool:
|
||||
return (
|
||||
block.get("type") == "tool_use"
|
||||
and isinstance(block.get("id"), str)
|
||||
and isinstance(block.get("name"), str)
|
||||
and isinstance(block.get("input"), dict)
|
||||
)
|
||||
|
||||
|
||||
def _write_back_system_block(system: object, block_idx: int, response: str) -> None:
|
||||
if not isinstance(system, list):
|
||||
return
|
||||
text_blocks: Final = tuple(block for block in system if isinstance(block, dict) and block.get("type") == "text")
|
||||
if block_idx < len(text_blocks):
|
||||
text_blocks[block_idx]["text"] = (
|
||||
response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
|
||||
|
||||
def _write_back_message_text(message: _WritableMessage, target: MessageTextTarget, response: str) -> None:
|
||||
content: Final = message.get("content", None)
|
||||
if content is None:
|
||||
return
|
||||
match target:
|
||||
case MessageContentTarget():
|
||||
if isinstance(content, str):
|
||||
message["content"] = response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
case ContentBlockTextTarget(content_idx=content_idx):
|
||||
if isinstance(content, list):
|
||||
content[content_idx]["text"] = (
|
||||
response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
case ToolResultStringTarget(content_idx=content_idx):
|
||||
if isinstance(content, list):
|
||||
content[content_idx]["content"] = (
|
||||
response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx):
|
||||
if isinstance(content, list):
|
||||
content[content_idx]["content"][block_idx]["text"] = (
|
||||
response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
case _:
|
||||
assert_never(target)
|
||||
|
||||
|
||||
_TOOL_USE_INPUT_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def _rewritten_tool_use_input(arguments: str) -> Mapping[str, object] | None:
|
||||
try:
|
||||
return _TOOL_USE_INPUT_ADAPTER.validate_json(arguments)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _write_back_tool_use(
|
||||
message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape, rewritten_input: Mapping[str, object]
|
||||
) -> None:
|
||||
content: Final = message.get("content", None)
|
||||
block: Final = content[target.content_idx] if isinstance(content, list) else None
|
||||
if not isinstance(block, dict):
|
||||
return
|
||||
block["input"] = rewritten_input # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
if shape.name is not None and shape.name != block.get("name"):
|
||||
block["name"] = shape.name # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _SSEFieldRewrite:
|
||||
"""One field of one nested section of a buffered SSE event, rewritten."""
|
||||
|
|
@ -452,9 +544,8 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
|
||||
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)
|
||||
|
||||
# Exclude only the trusted top-level prompt. In-sequence system entries are untrusted
|
||||
# and must stay aligned with texts_to_check for positional masking. When the top-level
|
||||
# prompt is included, the pre-existing count mismatch disables positional masking.
|
||||
# The top-level prompt is translated on its own below so it can be hoisted in front of
|
||||
# any mid-turn system entries and scanned first, aligned with that structured position.
|
||||
translation_source: Final = { # mutable-ok: API message payload
|
||||
key: value for key, value in data.items() if key != "system"
|
||||
}
|
||||
|
|
@ -490,7 +581,12 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
]
|
||||
)
|
||||
|
||||
# Step 1: Extract all text content and images
|
||||
# Step 1: Extract all text content, images, and tool calls
|
||||
top_level_system_scanned: Final = (
|
||||
()
|
||||
if hoisted_system_message is None or scan_only_tool_results
|
||||
else self._extract_top_level_system_text(hoisted_system_message)
|
||||
)
|
||||
extracted: Final = tuple(
|
||||
self._extract_input_text_and_images(
|
||||
message=message,
|
||||
|
|
@ -501,17 +597,27 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
)
|
||||
for msg_idx, message in enumerate(messages)
|
||||
)
|
||||
scanned: Final = tuple(item for one_message in extracted for item in one_message.scanned)
|
||||
scanned: Final = (
|
||||
*top_level_system_scanned,
|
||||
*(item for one_message in extracted for item in one_message.scanned),
|
||||
)
|
||||
texts_to_check: Final = [item.text for item in scanned] # mutable-ok: GenericGuardrailAPIInputs takes list[str]
|
||||
images_to_check: Final = [
|
||||
image for one_message in extracted for image in one_message.images
|
||||
] # mutable-ok: GenericGuardrailAPIInputs takes list[str]
|
||||
scanned_tool_calls: Final = tuple(item for one_message in extracted for item in one_message.tool_calls)
|
||||
tool_calls_to_check: Final = [
|
||||
item.tool_call for item in scanned_tool_calls
|
||||
] # mutable-ok: GenericGuardrailAPIInputs takes list[ChatCompletionToolCallChunk]
|
||||
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
|
||||
|
||||
# Step 2: Apply guardrail to all texts in batch
|
||||
if texts_to_check:
|
||||
# Step 2: Apply guardrail to all texts and tool calls in batch
|
||||
if texts_to_check or tool_calls_to_check:
|
||||
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
|
||||
if images_to_check:
|
||||
inputs["images"] = images_to_check
|
||||
if tool_calls_to_check:
|
||||
inputs["tool_calls"] = tool_calls_to_check
|
||||
if tools_to_check:
|
||||
inputs["tools"] = tools_to_check
|
||||
original_structured_messages: Final = structured_messages
|
||||
|
|
@ -570,9 +676,18 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
preserve_system_messages=has_midturn_system_message,
|
||||
)
|
||||
else:
|
||||
if guardrailed_texts and len(guardrailed_texts) != len(scanned):
|
||||
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
|
||||
self._apply_guardrail_tool_calls_to_input(
|
||||
messages=messages,
|
||||
scanned_tool_calls=scanned_tool_calls,
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
returned_tool_calls=guardrailed_inputs.get("tool_calls"),
|
||||
guardrail_name=guardrail_to_apply.guardrail_name,
|
||||
)
|
||||
# Step 3: Map guardrail responses back to original message structure
|
||||
await self._apply_guardrail_responses_to_input(
|
||||
messages=messages,
|
||||
data=data,
|
||||
responses=guardrailed_texts,
|
||||
scanned=scanned,
|
||||
)
|
||||
|
|
@ -598,6 +713,19 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
hoisted: Final = probe.get("messages") or [] # mutable-ok: API message payload
|
||||
return hoisted[0] if hoisted else None
|
||||
|
||||
@staticmethod
|
||||
def _extract_top_level_system_text(hoisted_system_message: AllMessageValues) -> tuple[ScannedText, ...]:
|
||||
content: Final = hoisted_system_message.get("content")
|
||||
if isinstance(content, str):
|
||||
return (ScannedText(content, SystemStringTarget()),)
|
||||
if not isinstance(content, list):
|
||||
return ()
|
||||
return tuple(
|
||||
ScannedText(text_str, SystemBlockTextTarget(block_idx))
|
||||
for block_idx, block in enumerate(content)
|
||||
if isinstance(block, dict) and isinstance(text_str := block.get("text"), str)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _openai_system_message_to_anthropic(
|
||||
message: Mapping[str, object],
|
||||
|
|
@ -852,9 +980,25 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
for content_idx, content_item in enumerate(content)
|
||||
if isinstance(content_item, dict)
|
||||
)
|
||||
tool_use_blocks: Final = (
|
||||
()
|
||||
if scan_only_tool_results
|
||||
else tuple(
|
||||
(content_idx, content_item)
|
||||
for content_idx, content_item in enumerate(content)
|
||||
if isinstance(content_item, dict) and _is_client_tool_use(content_item)
|
||||
)
|
||||
)
|
||||
return ExtractedInput(
|
||||
scanned=tuple(item for block in blocks for item in block.scanned),
|
||||
images=tuple(image for block in blocks for image in block.images),
|
||||
tool_calls=tuple(
|
||||
ScannedToolCall(
|
||||
tool_call=AnthropicConfig.convert_tool_use_to_openai_format(content_item, tool_call_idx),
|
||||
target=ToolUseInputTarget(msg_idx, content_idx),
|
||||
)
|
||||
for tool_call_idx, (content_idx, content_item) in enumerate(tool_use_blocks)
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
|
|
@ -940,43 +1084,59 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
async def _apply_guardrail_responses_to_input(
|
||||
self,
|
||||
messages: Sequence[_WritableMessage],
|
||||
responses: list[str],
|
||||
data: dict[str, object], # mutable-ok: API message payload
|
||||
responses: Sequence[str],
|
||||
scanned: tuple[ScannedText, ...],
|
||||
) -> None:
|
||||
"""
|
||||
Apply guardrail responses back to input messages.
|
||||
Apply guardrail responses back to the top-level system prompt and the input messages.
|
||||
"""
|
||||
raw_messages: Final = data.get("messages")
|
||||
messages: Final[Sequence[_WritableMessage]] = raw_messages if isinstance(raw_messages, list) else ()
|
||||
for item, guardrail_response in zip(scanned, responses):
|
||||
target = item.target
|
||||
message = messages[target.msg_idx]
|
||||
content = message.get("content", None)
|
||||
if content is None:
|
||||
continue
|
||||
|
||||
match target:
|
||||
case MessageContentTarget():
|
||||
if isinstance(content, str):
|
||||
message["content"] = (
|
||||
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
case ContentBlockTextTarget(content_idx=content_idx):
|
||||
if isinstance(content, list):
|
||||
content[content_idx]["text"] = (
|
||||
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
case ToolResultStringTarget(content_idx=content_idx):
|
||||
if isinstance(content, list):
|
||||
content[content_idx]["content"] = (
|
||||
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx):
|
||||
if isinstance(content, list):
|
||||
content[content_idx]["content"][block_idx]["text"] = (
|
||||
match item.target:
|
||||
case SystemStringTarget():
|
||||
if isinstance(data.get("system"), str):
|
||||
data["system"] = (
|
||||
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
|
||||
)
|
||||
case SystemBlockTextTarget(block_idx=block_idx):
|
||||
_write_back_system_block(data.get("system"), block_idx, guardrail_response)
|
||||
case (
|
||||
MessageContentTarget()
|
||||
| ContentBlockTextTarget()
|
||||
| ToolResultStringTarget()
|
||||
| ToolResultBlockTextTarget() as message_target
|
||||
):
|
||||
_write_back_message_text(messages[message_target.msg_idx], message_target, guardrail_response)
|
||||
case _:
|
||||
assert_never(target)
|
||||
assert_never(item.target)
|
||||
|
||||
@staticmethod
|
||||
def _apply_guardrail_tool_calls_to_input(
|
||||
messages: Sequence[_WritableMessage],
|
||||
scanned_tool_calls: tuple[ScannedToolCall, ...],
|
||||
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
|
||||
returned_tool_calls: Sequence[object] | None,
|
||||
guardrail_name: str | None,
|
||||
) -> None:
|
||||
post_guardrail_tool_calls: Final = _tool_call_shapes(
|
||||
returned_tool_calls
|
||||
if returned_tool_calls is not None and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
|
||||
else tuple(item.tool_call for item in scanned_tool_calls)
|
||||
)
|
||||
rewritten: Final = tuple(
|
||||
(item, after, _rewritten_tool_use_input(after.arguments))
|
||||
for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls)
|
||||
if before != after
|
||||
)
|
||||
applicable: Final = tuple(
|
||||
(item, after, rewritten_input) for item, after, rewritten_input in rewritten if rewritten_input is not None
|
||||
)
|
||||
if len(applicable) != len(rewritten):
|
||||
raise unappliable_request_rewrite(guardrail_name)
|
||||
for item, after, rewritten_input in applicable:
|
||||
_write_back_tool_use(messages[item.target.msg_idx], item.target, after, rewritten_input)
|
||||
|
||||
async def process_output_response(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Callable, Iterator, Sequence
|
||||
from typing import Final, TypeVar
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
|
@ -364,3 +364,67 @@ def merge_guardrailed_scoped_messages(
|
|||
yield from appended
|
||||
|
||||
return list(_merged())
|
||||
|
||||
|
||||
def _content_part_text(part: object) -> str | None:
|
||||
if not isinstance(part, Mapping):
|
||||
return None
|
||||
text: Final = part.get("text")
|
||||
return text if isinstance(text, str) else None
|
||||
|
||||
|
||||
def message_slot_texts(message: Mapping[str, object]) -> tuple[str, ...]:
|
||||
content: Final = message.get("content")
|
||||
if isinstance(content, str):
|
||||
return (content,)
|
||||
if isinstance(content, list):
|
||||
return tuple(text for part in content if (text := _content_part_text(part)) is not None)
|
||||
return ()
|
||||
|
||||
|
||||
def message_text_slot_count(message: AllMessageValues) -> int:
|
||||
return len(message_slot_texts(message))
|
||||
|
||||
|
||||
def _part_with_text(part: object, text: str) -> object:
|
||||
if not isinstance(part, Mapping):
|
||||
return part
|
||||
return {**part, "text": text} # mutable-ok: content parts stay JSON-plain dicts
|
||||
|
||||
|
||||
def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> Sequence[object]:
|
||||
remaining_texts: Final = iter(texts)
|
||||
return [ # mutable-ok: message content stays a JSON list
|
||||
_part_with_text(part, next(remaining_texts)) if _content_part_text(part) is not None else part
|
||||
for part in content
|
||||
]
|
||||
|
||||
|
||||
def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None:
|
||||
"""Swap one rewritten text into each text slot of a chat row, in order.
|
||||
|
||||
A slot is a string ``content`` or one list part carrying a string ``text``;
|
||||
images and other parts ride along untouched. Returns None unless the counts
|
||||
line up exactly, so a rewrite never lands on the wrong slot.
|
||||
"""
|
||||
if message_text_slot_count(message) != len(texts):
|
||||
return None
|
||||
content: Final = message.get("content")
|
||||
if not isinstance(content, (str, list)):
|
||||
return message
|
||||
rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts)
|
||||
rewritten: Final = {**message, "content": rewritten_content} # mutable-ok: chat rows stay JSON-plain dicts
|
||||
return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped
|
||||
|
||||
|
||||
class UnappliableRequestRewrite(Exception):
|
||||
def __init__(self, guardrail_name: str) -> None:
|
||||
super().__init__(
|
||||
f"Guardrail '{guardrail_name}' rewrote the request in a way this endpoint cannot apply, "
|
||||
"so the request was rejected rather than sent unrewritten"
|
||||
)
|
||||
self.guardrail_name: Final = guardrail_name
|
||||
|
||||
|
||||
def unappliable_request_rewrite(guardrail_name: str | None) -> UnappliableRequestRewrite:
|
||||
return UnappliableRequestRewrite(guardrail_name or "unknown")
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ BaseAWSLLM._sign_request after the request body is finalized.
|
|||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
from typing import Any, Final, cast # noqa: TID251 # map_openai_params returns the filtered params as a bare dict
|
||||
|
||||
import httpx
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
|
@ -32,6 +32,7 @@ from litellm.llms.bedrock_mantle.common_utils import (
|
|||
BedrockMantleAuthMixin,
|
||||
)
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.responses.additional_tools import HoistedAdditionalTools, hoist_additional_tools
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import (
|
||||
ResponseInputParam,
|
||||
|
|
@ -58,8 +59,6 @@ _BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES: Final = frozenset(
|
|||
_BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"})
|
||||
_BEDROCK_MANTLE_OPENAI_PATH_SUPPORTED_REASONING_SUMMARIES: Final = frozenset({"auto"})
|
||||
|
||||
_CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools"
|
||||
|
||||
_CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE: Final = "agent_message"
|
||||
_CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE: Final = "context_compaction"
|
||||
_CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: Final = "local_shell_call"
|
||||
|
|
@ -233,17 +232,14 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
|
|||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
remaining_input, hoisted_tools = self._hoist_codex_additional_tools(input)
|
||||
normalized_input: Final = self._normalize_codex_input_items(remaining_input)
|
||||
params: Final = cast( # cast-ok: the base signature leaves the params dict untyped
|
||||
"ResponsesAPIOptionalRequestParams", response_api_optional_request_params
|
||||
)
|
||||
hoisted: Final = hoist_additional_tools(input, params.get("tools"))
|
||||
normalized_input: Final = self._normalize_codex_input_items(hoisted.input)
|
||||
request_params: Final = (
|
||||
{
|
||||
**response_api_optional_request_params,
|
||||
"tools": [
|
||||
*(response_api_optional_request_params.get("tools") or []),
|
||||
*hoisted_tools,
|
||||
],
|
||||
}
|
||||
if hoisted_tools
|
||||
self._params_with_hoisted_tools(params, hoisted)
|
||||
if hoisted.hoisted
|
||||
else response_api_optional_request_params
|
||||
)
|
||||
return super().transform_responses_api_request(
|
||||
|
|
@ -254,41 +250,14 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
|
|||
headers=headers,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_codex_additional_tools_item(item: Any) -> bool:
|
||||
return isinstance(item, dict) and item.get("type") == _CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE
|
||||
|
||||
@staticmethod
|
||||
def _tools_of_additional_tools_item(item: "dict[str, Any]") -> "list[Any]":
|
||||
tools: Final = item.get("tools")
|
||||
return tools if isinstance(tools, list) else []
|
||||
|
||||
@classmethod
|
||||
def _hoist_codex_additional_tools(
|
||||
cls,
|
||||
input: "str | ResponseInputParam",
|
||||
) -> "tuple[str | ResponseInputParam, list[Any]]":
|
||||
"""Codex's "responses lite" wire mode ships tool definitions inside
|
||||
`input` as {"type": "additional_tools", "role": "developer",
|
||||
"tools": [...]} items. api.openai.com accepts that item type; Mantle
|
||||
rejects the whole request with 400 "Invalid 'input': value did not
|
||||
match any expected variant" but accepts the same tools at the top
|
||||
level, so move them there and strip the items from `input`.
|
||||
"""
|
||||
if not isinstance(input, list):
|
||||
return input, []
|
||||
additional_tools_items: Final = [item for item in input if cls._is_codex_additional_tools_item(item)]
|
||||
if not additional_tools_items:
|
||||
return input, []
|
||||
remaining_input: Final = [item for item in input if not cls._is_codex_additional_tools_item(item)]
|
||||
hoisted_tools = [tool for item in additional_tools_items for tool in cls._tools_of_additional_tools_item(item)]
|
||||
verbose_logger.debug(
|
||||
"Bedrock Mantle Responses API: hoisting %d tool(s) out of %d 'additional_tools' input item(s) "
|
||||
"into the top-level tools param (Mantle rejects that input item type).",
|
||||
len(hoisted_tools),
|
||||
len(additional_tools_items),
|
||||
)
|
||||
return remaining_input, cls._filter_unsupported_tools(hoisted_tools)
|
||||
def _params_with_hoisted_tools(
|
||||
cls, params: Mapping[str, object], hoisted: HoistedAdditionalTools
|
||||
) -> dict[str, object]:
|
||||
supported_tools: Final = cls._filter_unsupported_tools(list(hoisted.tools))
|
||||
if supported_tools:
|
||||
return {**params, "tools": supported_tools}
|
||||
return {key: value for key, value in params.items() if key != "tools"}
|
||||
|
||||
@staticmethod
|
||||
def _agent_message_text(item: "Mapping[str, object]") -> str:
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
|
|||
stream_item_field,
|
||||
stream_item_fingerprint,
|
||||
stream_item_items,
|
||||
unappliable_request_rewrite,
|
||||
)
|
||||
from litellm.main import stream_chunk_builder
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
|
|
@ -196,6 +197,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
else:
|
||||
# Step 3: Map guardrail responses back to original message structure
|
||||
if guardrailed_texts and texts_to_check:
|
||||
if len(guardrailed_texts) != len(text_task_mappings):
|
||||
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
|
||||
await self._apply_guardrail_responses_to_input_texts(
|
||||
messages=messages,
|
||||
responses=guardrailed_texts,
|
||||
|
|
@ -210,6 +213,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
task_mappings=tool_call_task_mappings,
|
||||
)
|
||||
|
||||
elif (
|
||||
not images_to_check
|
||||
and not guardrail_to_apply.records_own_guardrail_information
|
||||
and (not_run_reason := self._not_run_reason(messages)) is not None
|
||||
):
|
||||
guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response=not_run_reason,
|
||||
request_data=data,
|
||||
guardrail_status="not_run",
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"OpenAI Chat Completions: Processed input messages: %s",
|
||||
data.get("messages"),
|
||||
|
|
@ -217,6 +231,28 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
|
||||
return data
|
||||
|
||||
def _not_run_reason(
|
||||
self,
|
||||
messages: Sequence[dict[str, Any]], # mutable-ok: raw request messages consumed by _extract_inputs
|
||||
) -> str | None:
|
||||
"""Why nothing was scanned, or None when the only unscoped content is images, which this handler never scans."""
|
||||
texts: Final[list[str]] = [] # mutable-ok: filled by _extract_inputs
|
||||
images: Final[list[str]] = [] # mutable-ok: filled by _extract_inputs
|
||||
tool_calls: Final[list[ChatCompletionToolParam]] = [] # mutable-ok: filled by _extract_inputs
|
||||
for msg_idx, message in enumerate(messages):
|
||||
self._extract_inputs(
|
||||
message=message,
|
||||
msg_idx=msg_idx,
|
||||
texts_to_check=texts,
|
||||
images_to_check=images,
|
||||
tool_calls_to_check=tool_calls,
|
||||
text_task_mappings=[], # mutable-ok: required by _extract_inputs, unused here
|
||||
tool_call_task_mappings=[], # mutable-ok: required by _extract_inputs, unused here
|
||||
)
|
||||
if texts or tool_calls:
|
||||
return "no scannable content after message scoping"
|
||||
return None if images else "no scannable content"
|
||||
|
||||
def extract_request_tool_names(self, data: dict) -> list[str]:
|
||||
"""Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name)."""
|
||||
names: Final[list[str]] = []
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
|
|||
stream_item_field,
|
||||
stream_item_fingerprint,
|
||||
stream_item_items,
|
||||
unappliable_request_rewrite,
|
||||
)
|
||||
from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
|
|
@ -495,13 +496,13 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
data["instructions"] = written_back.instructions # rebind-ok: data is an out-param
|
||||
elif isinstance(input_data, str):
|
||||
guardrailed_texts: Final = guardrailed_inputs.get("texts") or ()
|
||||
if len(guardrailed_texts) > 1:
|
||||
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
|
||||
data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data # rebind-ok: data is an out-param
|
||||
else:
|
||||
rewritten_texts: Final = guardrailed_inputs.get("texts") or ()
|
||||
if len(rewritten_texts) != len(extracted.task_mappings):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite
|
||||
|
||||
raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown")
|
||||
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
|
||||
await self._apply_guardrail_responses_to_input(
|
||||
messages=input_data,
|
||||
responses=rewritten_texts,
|
||||
|
|
|
|||
|
|
@ -6,8 +6,10 @@ from typing import Final, TypeAlias
|
|||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.responses.litellm_completion_transformation.custom_tools import custom_tool_grammar_suffix
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
NAMESPACE_DESCRIPTION_SEPARATOR,
|
||||
NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS,
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
|
||||
|
|
@ -34,8 +36,8 @@ def _validated_tools(values: Iterable[object]) -> tuple[Tool, ...]:
|
|||
return tuple(tool for tool in validated if tool is not None)
|
||||
|
||||
|
||||
def _is_function(tool: Tool) -> bool:
|
||||
return tool.get("type") == "function"
|
||||
def _has_chat_tool(member: Tool) -> bool:
|
||||
return member.get("type") in NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS
|
||||
|
||||
|
||||
def _chat_tool_key(tool: Tool) -> str:
|
||||
|
|
@ -67,18 +69,19 @@ def _function_fields(tool: Tool) -> Tool:
|
|||
return function if function is not None else MappingProxyType({})
|
||||
|
||||
|
||||
def _without_namespace_prefix(key: str, value: object, prefix: str) -> object:
|
||||
if key != "description" or not isinstance(value, str) or not value.startswith(prefix):
|
||||
def _member_description(key: str, value: object, prefix: str, suffix: str) -> object:
|
||||
if key != "description" or not isinstance(value, str):
|
||||
return value
|
||||
return value[len(prefix) :]
|
||||
return value.replace(prefix, "", 1).replace(suffix, "", 1)
|
||||
|
||||
|
||||
def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_description: str) -> Tool:
|
||||
flattened_function: Final = _function_fields(flattened)
|
||||
prefix: Final = f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}" if namespace_description else ""
|
||||
suffix: Final = custom_tool_grammar_suffix(member.get("format")) if member.get("type") == "custom" else ""
|
||||
changed_function: Final = MappingProxyType(
|
||||
{
|
||||
key: _without_namespace_prefix(key, value, prefix)
|
||||
key: _member_description(key, value, prefix, suffix)
|
||||
for key, value in _function_fields(guardrailed).items()
|
||||
if flattened_function.get(key) != value
|
||||
}
|
||||
|
|
@ -93,8 +96,8 @@ def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_
|
|||
return {**member, **changed_extras, **changed_function} # mutable-ok: json.dumps rejects MappingProxyType
|
||||
|
||||
|
||||
def _rebuilt_function_members(
|
||||
function_members: Sequence[Tool],
|
||||
def _rebuilt_flattened_members(
|
||||
flattened_members: Sequence[Tool],
|
||||
flattened_group: Sequence[Tool],
|
||||
group_keys: Sequence[IndexedKey],
|
||||
guardrailed_by_key: Mapping[IndexedKey, Tool],
|
||||
|
|
@ -106,7 +109,7 @@ def _rebuilt_function_members(
|
|||
else member
|
||||
if guardrailed_by_key[key] == flattened
|
||||
else _rebuilt_member(member, flattened, guardrailed_by_key[key], namespace_description)
|
||||
for member, flattened, key in zip(function_members, flattened_group, group_keys)
|
||||
for member, flattened, key in zip(flattened_members, flattened_group, group_keys)
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -118,9 +121,9 @@ def _rebuilt_namespace(
|
|||
guardrailed_by_key: Mapping[IndexedKey, Tool],
|
||||
) -> tuple[Tool, ...]:
|
||||
namespace_description: Final = str(original.get("description") or "")
|
||||
rebuilt_functions: Final = iter(
|
||||
_rebuilt_function_members(
|
||||
tuple(member for member in members if _is_function(member)),
|
||||
rebuilt_flattened: Final = iter(
|
||||
_rebuilt_flattened_members(
|
||||
tuple(member for member in members if _has_chat_tool(member)),
|
||||
flattened_group,
|
||||
group_keys,
|
||||
guardrailed_by_key,
|
||||
|
|
@ -129,7 +132,7 @@ def _rebuilt_namespace(
|
|||
)
|
||||
rebuilt_members: Final = tuple(
|
||||
rebuilt
|
||||
for rebuilt in (next(rebuilt_functions) if _is_function(member) else member for member in members)
|
||||
for rebuilt in (next(rebuilt_flattened) if _has_chat_tool(member) else member for member in members)
|
||||
if rebuilt is not None
|
||||
)
|
||||
if not rebuilt_members:
|
||||
|
|
@ -149,7 +152,7 @@ def _merged_original(
|
|||
if guardrailed_group == tuple(flattened_group):
|
||||
return (original,)
|
||||
members: Final = _namespace_members(original) if original.get("type") == "namespace" else ()
|
||||
if members and sum(map(_is_function, members)) == len(flattened_group):
|
||||
if members and sum(map(_has_chat_tool, members)) == len(flattened_group):
|
||||
return _rebuilt_namespace(original, members, flattened_group, group_keys, guardrailed_by_key)
|
||||
if not guardrailed_group:
|
||||
return ()
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
from urllib.parse import unquote
|
||||
|
||||
|
|
@ -8,7 +9,36 @@ from litellm.llms.vertex_ai.common_utils import (
|
|||
)
|
||||
from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest
|
||||
from litellm.types.llms.vertex_ai import *
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
from litellm.types.utils import LiteLLMBatch, PromptTokensDetailsWrapper
|
||||
|
||||
|
||||
def vertex_prompt_tokens_details(
|
||||
usage_metadata: Mapping[str, object],
|
||||
) -> PromptTokensDetailsWrapper | None:
|
||||
raw_details: Final = usage_metadata.get("promptTokensDetails")
|
||||
if not isinstance(raw_details, list):
|
||||
return None
|
||||
|
||||
def _normalize(detail: object) -> tuple[str, int] | None:
|
||||
if not isinstance(detail, Mapping):
|
||||
return None
|
||||
modality: Final = detail.get("modality")
|
||||
token_count: Final = detail.get("tokenCount")
|
||||
if not isinstance(modality, str) or not isinstance(token_count, int):
|
||||
return None
|
||||
return modality.upper(), token_count
|
||||
|
||||
parsed_details: Final = tuple(_normalize(detail) for detail in raw_details)
|
||||
normalized: Final = tuple(detail for detail in parsed_details if detail is not None)
|
||||
if len(normalized) != len(parsed_details):
|
||||
return None
|
||||
|
||||
return PromptTokensDetailsWrapper(
|
||||
text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")),
|
||||
audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"),
|
||||
image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"),
|
||||
video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"),
|
||||
)
|
||||
|
||||
|
||||
class VertexAIBatchTransformation:
|
||||
|
|
|
|||
|
|
@ -298,8 +298,6 @@ def transform_openai_input_gemini_embed_content(
|
|||
|
||||
|
||||
_IMAGE_MIME_TYPES: Final = frozenset({"image/png", "image/jpeg"})
|
||||
_VIDEO_TOKENS_PER_SECOND: Final = 258.0
|
||||
_AUDIO_TOKENS_PER_SECOND: Final = 32.0
|
||||
_usage_metadata_adapter: Final = TypeAdapter(UsageMetadata)
|
||||
|
||||
|
||||
|
|
@ -339,11 +337,12 @@ def _is_image_element(
|
|||
return False
|
||||
|
||||
|
||||
def _count_input_images(
|
||||
def _is_image_only_input(
|
||||
input: GeminiEmbeddingInput,
|
||||
resolved_files: Mapping[str, Mapping[str, str]],
|
||||
) -> int:
|
||||
return sum(1 for element in _flatten_input(input) if _is_image_element(element, resolved_files))
|
||||
) -> bool:
|
||||
elements: Final = _flatten_input(input)
|
||||
return bool(elements) and all(_is_image_element(element, resolved_files) for element in elements)
|
||||
|
||||
|
||||
def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int:
|
||||
|
|
@ -372,30 +371,29 @@ def _usage_from_embed_content_response(
|
|||
total_tokens: Final = usage_metadata.get("totalTokenCount") or prompt_tokens
|
||||
|
||||
details: Final[Sequence[PromptTokensDetails]] = usage_metadata.get("promptTokensDetails") or ()
|
||||
if not details:
|
||||
return Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
text_tokens=0,
|
||||
image_tokens=prompt_tokens if _is_image_only_input(input, resolved_files) else 0,
|
||||
),
|
||||
)
|
||||
|
||||
text_tokens: Final = _tokens_for_modality(details, "TEXT")
|
||||
audio_tokens: Final = _tokens_for_modality(details, "AUDIO")
|
||||
image_tokens: Final = _tokens_for_modality(details, "IMAGE")
|
||||
video_tokens: Final = _tokens_for_modality(details, "VIDEO")
|
||||
image_count: Final = _count_input_images(input, resolved_files)
|
||||
|
||||
video_length_seconds: Final = video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0
|
||||
audio_length_seconds: Final = audio_tokens / _AUDIO_TOKENS_PER_SECOND if audio_tokens > 0 else 0.0
|
||||
|
||||
# generic_cost_per_token rewrites text_tokens to the full prompt minus
|
||||
# other modalities when both text_tokens and image_count are zero. For
|
||||
# video, that misallocates video tokens to text; a 1-token floor sidesteps
|
||||
# the rewrite and keeps billing on input_cost_per_video_per_second.
|
||||
needs_video_text_floor: Final = video_length_seconds > 0 and text_tokens == 0 and image_count == 0
|
||||
resolved_text_tokens: Final = 1 if needs_video_text_floor else text_tokens
|
||||
|
||||
return Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
text_tokens=resolved_text_tokens,
|
||||
text_tokens=text_tokens,
|
||||
audio_tokens=audio_tokens,
|
||||
image_count=image_count,
|
||||
video_length_seconds=video_length_seconds,
|
||||
audio_length_seconds=audio_length_seconds,
|
||||
image_tokens=image_tokens,
|
||||
video_tokens=video_tokens,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -415,8 +413,7 @@ def process_embed_content_response(
|
|||
model_response: EmbeddingResponse to populate
|
||||
model: Model name
|
||||
response_json: Raw JSON response from embedContent endpoint
|
||||
resolved_files: Mapping of file references (files/abc) to {mime_type, uri},
|
||||
used to bill resolved image references at the per-image rate
|
||||
resolved_files: Mapping of file references to resolved metadata
|
||||
|
||||
Returns:
|
||||
EmbeddingResponse with single embedding
|
||||
|
|
|
|||
|
|
@ -25601,10 +25601,14 @@
|
|||
"source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models"
|
||||
},
|
||||
"gemini-embedding-2-preview": {
|
||||
"input_cost_per_audio_per_second": 0.00016,
|
||||
"input_cost_per_image": 0.00012,
|
||||
"input_cost_per_audio_token": 6.5e-06,
|
||||
"input_cost_per_audio_token_batches": 3.25e-06,
|
||||
"input_cost_per_image_token": 4.5e-07,
|
||||
"input_cost_per_image_token_batches": 2.25e-07,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"input_cost_per_video_per_second": 0.00079,
|
||||
"input_cost_per_token_batches": 1e-07,
|
||||
"input_cost_per_video_token": 1.2e-05,
|
||||
"input_cost_per_video_token_batches": 6e-06,
|
||||
"litellm_provider": "vertex_ai-embedding-models",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
|
|
@ -25615,13 +25619,14 @@
|
|||
"uses_embed_content": true
|
||||
},
|
||||
"gemini-embedding-2": {
|
||||
"input_cost_per_audio_per_second": 0.00016,
|
||||
"input_cost_per_audio_token": 6.5e-06,
|
||||
"input_cost_per_image": 0.00012,
|
||||
"input_cost_per_audio_token_batches": 3.25e-06,
|
||||
"input_cost_per_image_token": 4.5e-07,
|
||||
"input_cost_per_image_token_batches": 2.25e-07,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"input_cost_per_token_batches": 1e-07,
|
||||
"input_cost_per_video_per_second": 0.00079,
|
||||
"input_cost_per_video_token": 1.2e-05,
|
||||
"input_cost_per_video_token_batches": 6e-06,
|
||||
"litellm_provider": "vertex_ai-embedding-models",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
|
|
@ -25633,10 +25638,14 @@
|
|||
"uses_embed_content": true
|
||||
},
|
||||
"vertex_ai/gemini-embedding-2-preview": {
|
||||
"input_cost_per_audio_per_second": 0.00016,
|
||||
"input_cost_per_image": 0.00012,
|
||||
"input_cost_per_audio_token": 6.5e-06,
|
||||
"input_cost_per_audio_token_batches": 3.25e-06,
|
||||
"input_cost_per_image_token": 4.5e-07,
|
||||
"input_cost_per_image_token_batches": 2.25e-07,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"input_cost_per_video_per_second": 0.00079,
|
||||
"input_cost_per_token_batches": 1e-07,
|
||||
"input_cost_per_video_token": 1.2e-05,
|
||||
"input_cost_per_video_token_batches": 6e-06,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
|
|
@ -25648,13 +25657,14 @@
|
|||
"uses_embed_content": true
|
||||
},
|
||||
"vertex_ai/gemini-embedding-2": {
|
||||
"input_cost_per_audio_per_second": 0.00016,
|
||||
"input_cost_per_audio_token": 6.5e-06,
|
||||
"input_cost_per_image": 0.00012,
|
||||
"input_cost_per_audio_token_batches": 3.25e-06,
|
||||
"input_cost_per_image_token": 4.5e-07,
|
||||
"input_cost_per_image_token_batches": 2.25e-07,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"input_cost_per_token_batches": 1e-07,
|
||||
"input_cost_per_video_per_second": 0.00079,
|
||||
"input_cost_per_video_token": 1.2e-05,
|
||||
"input_cost_per_video_token_batches": 6e-06,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
|
|
@ -25693,10 +25703,14 @@
|
|||
},
|
||||
"gemini/gemini-embedding-2-preview": {
|
||||
"deprecation_date": "2026-08-10",
|
||||
"input_cost_per_audio_per_second": 0.00016,
|
||||
"input_cost_per_image": 0.00012,
|
||||
"input_cost_per_audio_token": 6.5e-06,
|
||||
"input_cost_per_audio_token_batches": 3.25e-06,
|
||||
"input_cost_per_image_token": 4.5e-07,
|
||||
"input_cost_per_image_token_batches": 2.25e-07,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"input_cost_per_video_per_second": 0.00079,
|
||||
"input_cost_per_token_batches": 1e-07,
|
||||
"input_cost_per_video_token": 1.2e-05,
|
||||
"input_cost_per_video_token_batches": 6e-06,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
|
|
@ -25709,10 +25723,14 @@
|
|||
"tpm": 10000000
|
||||
},
|
||||
"gemini/gemini-embedding-2": {
|
||||
"input_cost_per_audio_per_second": 0.00016,
|
||||
"input_cost_per_image": 0.00012,
|
||||
"input_cost_per_audio_token": 6.5e-06,
|
||||
"input_cost_per_audio_token_batches": 3.25e-06,
|
||||
"input_cost_per_image_token": 4.5e-07,
|
||||
"input_cost_per_image_token_batches": 2.25e-07,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"input_cost_per_video_per_second": 0.00079,
|
||||
"input_cost_per_token_batches": 1e-07,
|
||||
"input_cost_per_video_token": 1.2e-05,
|
||||
"input_cost_per_video_token_batches": 6e-06,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
|
|
|
|||
|
|
@ -4,11 +4,12 @@ import os
|
|||
from collections.abc import Callable, Mapping
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple
|
||||
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple
|
||||
|
||||
import httpx
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
BeforeValidator,
|
||||
ConfigDict,
|
||||
Field,
|
||||
Json,
|
||||
|
|
@ -47,6 +48,7 @@ from litellm.types.proxy.carried_budget_state import (
|
|||
)
|
||||
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
|
||||
from litellm.types.router import RouterErrors, UpdateRouterConfig
|
||||
from litellm.types.router_weights import validate_router_settings_dict
|
||||
from litellm.types.secret_managers.main import KeyManagementSystem
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
|
|
@ -1981,8 +1983,14 @@ class OrgMember(MemberBase):
|
|||
|
||||
from litellm.models.team import TeamBase as TeamBase # noqa: E402
|
||||
|
||||
RouterSettingsDict = Annotated[
|
||||
dict[str, object],
|
||||
BeforeValidator(validate_router_settings_dict, json_schema_input_type=UpdateRouterConfig),
|
||||
]
|
||||
|
||||
|
||||
class NewTeamRequest(TeamBase):
|
||||
router_settings: RouterSettingsDict | None = None
|
||||
model_aliases: dict | None = None
|
||||
tags: list | None = None
|
||||
guardrails: list[str] | None = None
|
||||
|
|
@ -2080,7 +2088,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
|
|||
allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None
|
||||
enforced_batch_output_expires_after: dict | None = None
|
||||
enforced_file_expires_after: dict | None = None
|
||||
router_settings: dict | None = None
|
||||
router_settings: RouterSettingsDict | None = None
|
||||
access_group_ids: list[str] | None = None
|
||||
budget_limits: list[BudgetLimitEntry] | None = None # multiple concurrent budget windows
|
||||
default_team_member_models: list[str] | None = None # default allowed_models seeded onto new team members
|
||||
|
|
|
|||
|
|
@ -249,6 +249,7 @@ async def authenticate_user(
|
|||
|
||||
if os.getenv("DATABASE_URL") is not None:
|
||||
response = await generate_key_helper_fn(
|
||||
llm_router=None,
|
||||
request_type="key",
|
||||
**{
|
||||
"user_role": LitellmUserRoles.PROXY_ADMIN,
|
||||
|
|
@ -324,6 +325,7 @@ async def authenticate_user(
|
|||
await _rehash_password_if_needed(_user_row.user_id, password, _password)
|
||||
if os.getenv("DATABASE_URL") is not None:
|
||||
response = await generate_key_helper_fn(
|
||||
llm_router=None,
|
||||
request_type="key",
|
||||
**{
|
||||
"user_role": user_role,
|
||||
|
|
|
|||
|
|
@ -878,6 +878,7 @@ async def _auto_register_jwt_mapping(
|
|||
# the NOT NULL @id constraint. Every successful key-creation caller (e.g.
|
||||
# /key/generate) passes table_name="key" explicitly.
|
||||
key_data: Final = await generate_key_helper_fn(
|
||||
llm_router=None,
|
||||
request_type="key",
|
||||
table_name="key",
|
||||
team_id=team_id,
|
||||
|
|
|
|||
|
|
@ -580,8 +580,8 @@ What the command changed is recorded in `~/.litellm/claude_configure_state.json`
|
|||
`lite configure claude`, `lite login --config-claude`, `lite up` and `lite autoroute up` also install a status line (`~/.litellm/statusline.py`, registered as `statusLine` in `~/.claude/settings.json` unless you already run one) that shows which model the auto-router actually served the last turn and, once the proxy has recorded the session, what the session cost against the router's savings baseline:
|
||||
|
||||
```
|
||||
claude-auto · Routed to: claude-haiku-4-5 -63% vs Claude Opus 5
|
||||
LiteLLM ████████░░░░░░░░░░░░░░░░ $0.14
|
||||
Routed to: claude-haiku-4-5 -63% vs Claude Opus 5
|
||||
claude-auto ████████░░░░░░░░░░░░░░░░ $0.14
|
||||
Claude Opus 5 ████████████████████████ $0.38
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ import os
|
|||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import unicodedata
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from collections.abc import Callable, Mapping
|
||||
|
|
@ -42,7 +43,6 @@ FETCH_TIMEOUT_SECONDS: Final = 3
|
|||
BAR_WIDTH: Final = 24
|
||||
BAR_FULL: Final = "\u2588"
|
||||
BAR_EMPTY: Final = "\u2591"
|
||||
SEPARATOR: Final = " \u00b7 "
|
||||
TRANSCRIPT_SCAN_LIMIT_BYTES: Final = 4 * 1024 * 1024
|
||||
CLAUDE_BASE_URL_ENV_KEYS: Final = ("ANTHROPIC_BASE_URL",)
|
||||
CLAUDE_API_KEY_ENV_KEYS: Final = ("ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY")
|
||||
|
|
@ -50,7 +50,6 @@ CODEX_BASE_URL_ENV_KEYS: Final = ("OPENAI_BASE_URL",)
|
|||
CODEX_API_KEY_ENV_KEYS: Final = ("OPENAI_API_KEY",)
|
||||
CODEX_STOP_EVENT: Final = "Stop"
|
||||
SYNTHETIC_MODEL: Final = "<synthetic>"
|
||||
LITELLM_LABEL: Final = "LiteLLM"
|
||||
RESET: Final = "\033[0m"
|
||||
BOLD: Final = "\033[1m"
|
||||
DIM: Final = "\033[90m"
|
||||
|
|
@ -302,31 +301,37 @@ def _bar(fraction: float, color: str, width: int, use_color: bool) -> str:
|
|||
return f"{color}{BAR_FULL * filled}{DIM}{BAR_EMPTY * (width - filled)}{RESET}"
|
||||
|
||||
|
||||
def _display_width(label: str) -> int:
|
||||
return sum(
|
||||
2 if unicodedata.east_asian_width(character) in ("W", "F") else 1
|
||||
for character in label
|
||||
if unicodedata.category(character) not in ("Mn", "Me")
|
||||
)
|
||||
|
||||
|
||||
def render(model: str, session: Session | None, config_dir: Path, use_color: bool, bar_width: int = BAR_WIDTH) -> str:
|
||||
def paint(code: str, text: str) -> str:
|
||||
return f"{code}{text}{RESET}" if use_color else text
|
||||
|
||||
routed: Final = paint(BOLD, f"Routed to: {model}")
|
||||
if session is None:
|
||||
if session is None or session.baseline_model is None or session.baseline_spend <= 0:
|
||||
return routed
|
||||
header: Final = f"{session.router_name}{SEPARATOR}{routed}"
|
||||
if session.baseline_model is None or session.baseline_spend <= 0:
|
||||
return header
|
||||
reference: Final = baseline_label(session.baseline_model, config_dir)
|
||||
pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100
|
||||
delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}")
|
||||
peak: Final = max(session.spend, session.baseline_spend)
|
||||
label_width: Final = max(len(LITELLM_LABEL), len(reference))
|
||||
label_width: Final = max(_display_width(session.router_name), _display_width(reference))
|
||||
rows: Final = (
|
||||
(LITELLM_LABEL, session.spend, LITELLM_COLOR),
|
||||
(session.router_name, session.spend, LITELLM_COLOR),
|
||||
(reference, session.baseline_spend, BASELINE_COLOR),
|
||||
)
|
||||
lines: Final = (
|
||||
f"{paint(DIM, label.ljust(label_width))} {_bar(amount / peak, color, bar_width, use_color)} "
|
||||
f"{paint(DIM, label + ' ' * (label_width - _display_width(label)))} "
|
||||
f"{_bar(amount / peak, color, bar_width, use_color)} "
|
||||
f"{paint(DIM, f'${amount:.2f}')}"
|
||||
for label, amount, color in rows
|
||||
)
|
||||
return "\n".join((f"{header} {delta}", *lines))
|
||||
return "\n".join((f"{routed} {delta}", *lines))
|
||||
|
||||
|
||||
def color_enabled(env: Mapping[str, str]) -> bool:
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import httpx
|
|||
import orjson
|
||||
from fastapi import HTTPException, Request, status
|
||||
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
||||
from pydantic import ValidationError
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
import litellm
|
||||
|
|
@ -76,6 +77,7 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di
|
|||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.router import RouterRateLimitError
|
||||
from litellm.types.router_weights import validate_router_weights
|
||||
|
||||
_LateResponseT = TypeVar("_LateResponseT", bound=Response)
|
||||
_LlmCallT = TypeVar("_LlmCallT")
|
||||
|
|
@ -1939,6 +1941,13 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# This avoids expensive Router instantiation on each request
|
||||
if router_settings is not None:
|
||||
self.data["router_settings_override"] = router_settings
|
||||
try:
|
||||
self.data["_router_weights"] = validate_router_weights(router_settings.get("weights"))
|
||||
except ValidationError:
|
||||
self.data["_router_weights"] = None
|
||||
verbose_proxy_logger.warning(
|
||||
"Ignoring invalid saved router weights; update team/key router_settings"
|
||||
)
|
||||
alias_target: Final = await _resolve_per_request_model_group_alias(
|
||||
requested_model=self.data.get("model"),
|
||||
router_settings=router_settings,
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ from typing import Final
|
|||
|
||||
from fastapi import status
|
||||
|
||||
from litellm.constants import STRINGIFIED_NONE
|
||||
|
||||
_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType(
|
||||
{
|
||||
status.HTTP_401_UNAUTHORIZED: "authentication_error",
|
||||
|
|
@ -35,7 +37,7 @@ def openai_error_type(exc: object, status_code: int) -> str:
|
|||
"""OpenAI types ``error.type`` as a required string, so an exception carrying none
|
||||
falls back to the type its status code stands for."""
|
||||
carried: Final = attribute_of(exc, "type")
|
||||
if isinstance(carried, str):
|
||||
if isinstance(carried, str) and carried != STRINGIFIED_NONE:
|
||||
return carried
|
||||
mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code)
|
||||
if mapped is not None:
|
||||
|
|
@ -49,4 +51,4 @@ def openai_error_param(exc: object) -> str | None:
|
|||
"""OpenAI types ``error.param`` as nullable, so an exception carrying none
|
||||
serializes as JSON ``null``."""
|
||||
carried: Final = attribute_of(exc, "param")
|
||||
return carried if isinstance(carried, str) else None
|
||||
return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@ class ComplianceChecker:
|
|||
|
||||
def __init__(self, data: ComplianceCheckRequest):
|
||||
self.data = data
|
||||
self.guardrails = data.guardrail_information or []
|
||||
self.guardrails = tuple(g for g in data.guardrail_information or () if g.get("guardrail_status") != "not_run")
|
||||
|
||||
def _get_guardrails_by_mode(self, mode: str) -> list[dict]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@
|
|||
|
||||
import fnmatch
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
|
||||
|
||||
import httpx
|
||||
|
|
@ -24,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.llms.openai import ChatCompletionToolParam
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
GenericGuardrailAPIMetadata,
|
||||
GenericGuardrailAPIRequest,
|
||||
|
|
@ -150,6 +150,26 @@ def _extract_inbound_headers(
|
|||
return None
|
||||
|
||||
|
||||
def _structured_rows_to_write_back(
|
||||
original_rows: Sequence[AllMessageValues] | None,
|
||||
shown_rows: Sequence[AllMessageValues] | None,
|
||||
returned_rows: Sequence[AllMessageValues],
|
||||
) -> tuple[AllMessageValues, ...] | None:
|
||||
"""The request model drops row keys its message types do not declare, so a
|
||||
row the server echoes back verbatim is restored to the original row object.
|
||||
A server that echoes every row back unchanged has not rewritten anything
|
||||
per row, so its answer is read from texts, as it was before rows could be
|
||||
returned at all."""
|
||||
if original_rows is None or shown_rows is None or len(returned_rows) != len(original_rows):
|
||||
return tuple(returned_rows)
|
||||
if all(returned == shown for shown, returned in zip(shown_rows, returned_rows)):
|
||||
return None
|
||||
return tuple(
|
||||
original if returned == shown else returned
|
||||
for original, shown, returned in zip(original_rows, shown_rows, returned_rows)
|
||||
)
|
||||
|
||||
|
||||
class GenericGuardrailAPI(CustomGuardrail):
|
||||
"""
|
||||
Generic Guardrail API integration for LiteLLM.
|
||||
|
|
@ -322,6 +342,8 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
texts: list,
|
||||
images: list[str] | None,
|
||||
tools: list[ChatCompletionToolParam] | None,
|
||||
structured_messages: Sequence[AllMessageValues] | None,
|
||||
shown_messages: Sequence[AllMessageValues] | None,
|
||||
guardrail_response: GenericGuardrailAPIResponse,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
# Action is NONE or no modifications needed
|
||||
|
|
@ -336,6 +358,13 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
return_inputs["tools"] = guardrail_response.tools
|
||||
elif tools:
|
||||
return_inputs["tools"] = tools
|
||||
rows_to_write_back: Final = (
|
||||
_structured_rows_to_write_back(structured_messages, shown_messages, guardrail_response.structured_messages)
|
||||
if guardrail_response.structured_messages
|
||||
else None
|
||||
)
|
||||
if rows_to_write_back is not None:
|
||||
return_inputs["structured_messages"] = list(rows_to_write_back) # mutable-ok: guardrail inputs take a list
|
||||
if guardrail_response.stream_holdback_chars is not None:
|
||||
return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars
|
||||
return return_inputs
|
||||
|
|
@ -473,6 +502,8 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
texts=texts,
|
||||
images=images,
|
||||
tools=tools,
|
||||
structured_messages=structured_messages,
|
||||
shown_messages=guardrail_request.structured_messages,
|
||||
guardrail_response=guardrail_response,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1600,8 +1600,8 @@ class PanwPrismaAirsHandler(CustomGuardrail):
|
|||
|
||||
Args:
|
||||
texts: Flattened text entries from the framework.
|
||||
messages: Original request messages (request_data["messages"]),
|
||||
NOT structured_messages (which may have injected system content).
|
||||
messages: The structured messages the framework flattened into ``texts``,
|
||||
hoisted top-level system prompt included, so positions line up.
|
||||
|
||||
Returns a set of scannable indices, or None on count mismatch or no user/developer
|
||||
message (safety fallback to existing role-filter behavior).
|
||||
|
|
@ -1788,15 +1788,10 @@ class PanwPrismaAirsHandler(CustomGuardrail):
|
|||
structured_messages: Final = inputs.get("structured_messages")
|
||||
if structured_messages:
|
||||
# For Anthropic /v1/messages: default to latest-user-only scanning.
|
||||
# Uses request_data["messages"] (original format), NOT structured_messages
|
||||
# (which has injected system content from adapter translation).
|
||||
if self._use_latest_user_only(request_data, logging_obj):
|
||||
original_messages: Final = request_data.get("messages")
|
||||
if original_messages:
|
||||
scannable_indices = self._get_latest_user_text_indices(texts, original_messages)
|
||||
scannable_indices = self._get_latest_user_text_indices(texts, structured_messages)
|
||||
# Fall through to existing role filtering if:
|
||||
# - not Anthropic, OR flag explicitly False, OR
|
||||
# - no original messages, OR
|
||||
# - latest-user extraction returned None (no user / count mismatch)
|
||||
if scannable_indices is None:
|
||||
scannable_indices = self._get_scannable_text_indices(texts, structured_messages)
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import asyncio
|
|||
import base64
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional
|
||||
|
||||
import httpx
|
||||
|
|
@ -14,11 +15,13 @@ from litellm.integrations.custom_guardrail import (
|
|||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts, message_with_slot_texts
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -28,12 +31,36 @@ if TYPE_CHECKING:
|
|||
|
||||
_SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0
|
||||
_SANITIZE_FILE_QUEUED_STATUSES: Final = frozenset({"created", "in progress"})
|
||||
_PROTECT_ROLES: Final = frozenset({"system", "user", "assistant"})
|
||||
|
||||
|
||||
class PromptSecurityGuardrailMissingSecrets(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def _inputs_with_structured_messages(
|
||||
inputs: GenericGuardrailAPIInputs, rewritten_messages: Sequence[AllMessageValues] | None
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
if rewritten_messages is None:
|
||||
return inputs
|
||||
patched: Final[GenericGuardrailAPIInputs] = {
|
||||
**inputs,
|
||||
"structured_messages": list(rewritten_messages), # mutable-ok: the TypedDict field is declared as a list
|
||||
}
|
||||
return patched
|
||||
|
||||
|
||||
def _inputs_with_modifications(
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
modified_texts: list[str],
|
||||
rewritten_messages: Sequence[AllMessageValues] | None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
if not modified_texts:
|
||||
return _inputs_with_structured_messages(inputs, rewritten_messages)
|
||||
with_texts: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": modified_texts}
|
||||
return _inputs_with_structured_messages(with_texts, rewritten_messages)
|
||||
|
||||
|
||||
class _ProtectVerdict(TypedDict, total=False):
|
||||
"""One side (``prompt`` or ``response``) of an ``/api/protect`` verdict."""
|
||||
|
||||
|
|
@ -276,14 +303,39 @@ class PromptSecurityGuardrail(CustomGuardrail):
|
|||
detail="Blocked by Prompt Security, Violations: " + ", ".join(violations),
|
||||
)
|
||||
elif action == "modify":
|
||||
# Extract modified texts from modified_messages
|
||||
modified_messages: Final = result.get("modified_messages", [])
|
||||
modified_texts: Final = self._extract_texts_from_messages(modified_messages)
|
||||
if modified_texts:
|
||||
inputs["texts"] = modified_texts
|
||||
return _inputs_with_modifications(
|
||||
inputs,
|
||||
self._extract_texts_from_messages(modified_messages),
|
||||
self._structured_messages_with_modifications(structured_messages, modified_messages),
|
||||
)
|
||||
|
||||
return inputs
|
||||
|
||||
def _is_sent_to_protect(self, message: Mapping[str, object]) -> bool:
|
||||
return self.check_tool_results or message.get("role") in _PROTECT_ROLES
|
||||
|
||||
def _structured_messages_with_modifications(
|
||||
self,
|
||||
structured_messages: Sequence[AllMessageValues],
|
||||
modified_messages: Sequence[Mapping[str, object]],
|
||||
) -> tuple[AllMessageValues, ...] | None:
|
||||
sent_indices: Final = tuple(
|
||||
index for index, message in enumerate(structured_messages) if self._is_sent_to_protect(message)
|
||||
)
|
||||
if not sent_indices or len(sent_indices) != len(modified_messages):
|
||||
return None
|
||||
rewritten: Final = tuple(
|
||||
message_with_slot_texts(structured_messages[index], self._extract_texts_from_messages((modified,)))
|
||||
for index, modified in zip(sent_indices, modified_messages)
|
||||
)
|
||||
replacements: Final = MappingProxyType(
|
||||
{index: message for index, message in zip(sent_indices, rewritten) if message is not None}
|
||||
)
|
||||
if len(replacements) != len(sent_indices):
|
||||
return None
|
||||
return tuple(replacements.get(index, message) for index, message in enumerate(structured_messages))
|
||||
|
||||
async def _apply_guardrail_on_response(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
|
|
@ -347,19 +399,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
|
|||
return inputs
|
||||
|
||||
def _extract_texts_from_messages(self, messages: Sequence[Mapping[str, object]]) -> list[str]:
|
||||
"""Extract text content from messages."""
|
||||
texts: Final = []
|
||||
for message in messages:
|
||||
content = message.get("content")
|
||||
if isinstance(content, str):
|
||||
texts.append(content)
|
||||
elif isinstance(content, list):
|
||||
for item in content:
|
||||
if isinstance(item, dict) and item.get("type") == "text":
|
||||
text = item.get("text")
|
||||
if text:
|
||||
texts.append(text)
|
||||
return texts
|
||||
return [text for message in messages for text in message_slot_texts(message)]
|
||||
|
||||
async def _process_standalone_images(self, images: list[str], user_api_key_alias: str | None) -> None:
|
||||
"""Process standalone images from inputs (data URLs)."""
|
||||
|
|
@ -681,14 +721,13 @@ class PromptSecurityGuardrail(CustomGuardrail):
|
|||
|
||||
This allows checking tool results for indirect prompt injection when enabled.
|
||||
"""
|
||||
supported_roles: Final = ["system", "user", "assistant"]
|
||||
filtered_messages: Final = []
|
||||
transformed_count = 0
|
||||
filtered_count = 0
|
||||
|
||||
for message in messages:
|
||||
role = message.get("role", "")
|
||||
if role in supported_roles:
|
||||
if role in _PROTECT_ROLES:
|
||||
filtered_messages.append(message)
|
||||
else:
|
||||
if self.check_tool_results:
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ if TYPE_CHECKING:
|
|||
router: Final = APIRouter()
|
||||
|
||||
_EMPTY_UNITS: Final[Mapping[str, int]] = MappingProxyType({})
|
||||
_ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"passed": 0, "flagged": 1, "blocked": 2})
|
||||
_ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"not_run": 0, "passed": 1, "flagged": 2, "blocked": 3})
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
|
@ -325,7 +325,7 @@ class UsageDetailResponse(BaseModel):
|
|||
class UsageLogEntry(BaseModel):
|
||||
id: str
|
||||
timestamp: str
|
||||
action: str # blocked | passed | flagged
|
||||
action: str # blocked | passed | flagged | not_run
|
||||
score: float | None
|
||||
latency_ms: float | None
|
||||
model: str | None
|
||||
|
|
|
|||
|
|
@ -193,10 +193,12 @@ async def _upsert_rows_with_retry(
|
|||
|
||||
|
||||
def guardrail_status_to_action(status: str | None) -> str:
|
||||
"""Map StandardLogging guardrail_status to blocked/passed/flagged."""
|
||||
"""Map StandardLogging guardrail_status to blocked/passed/flagged/not_run."""
|
||||
if not status:
|
||||
return "passed"
|
||||
s: Final = (status or "").lower()
|
||||
if s == "not_run":
|
||||
return "not_run"
|
||||
if "intervened" in s or "block" in s:
|
||||
return "blocked"
|
||||
if "flagged" in s or "fail" in s or "error" in s:
|
||||
|
|
@ -354,37 +356,49 @@ async def process_spend_logs_guardrail_usage(
|
|||
"flagged_count": 0,
|
||||
}
|
||||
)
|
||||
index_rows: Final[list[dict[str, object]]] = []
|
||||
index_rows_by_key: Final[dict[tuple[str, str], dict[str, object]]] = {}
|
||||
|
||||
for payload in logs_to_process:
|
||||
request_id = payload.get("request_id")
|
||||
start_time = _parse_payload_start_time(payload)
|
||||
if not request_id or start_time is None:
|
||||
if not isinstance(request_id, str) or not request_id or start_time is None:
|
||||
continue
|
||||
date_key = _date_str(start_time)
|
||||
|
||||
for entry in _parse_guardrail_info_from_payload(payload):
|
||||
guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or ""
|
||||
if not guardrail_id:
|
||||
entries = _parse_guardrail_info_from_payload(payload)
|
||||
ids_by_name = MappingProxyType(
|
||||
{
|
||||
e["guardrail_name"]: e["guardrail_id"]
|
||||
for e in entries
|
||||
if e.get("guardrail_id") and isinstance(e.get("guardrail_name"), str) and e["guardrail_name"]
|
||||
}
|
||||
)
|
||||
for entry in entries:
|
||||
raw_name = entry.get("guardrail_name")
|
||||
guardrail_name = raw_name if isinstance(raw_name, str) else ""
|
||||
guardrail_id = entry.get("guardrail_id") or ids_by_name.get(guardrail_name) or guardrail_name
|
||||
if not isinstance(guardrail_id, str) or not guardrail_id:
|
||||
continue
|
||||
key = _MetricsKey(guardrail_id, date_key)
|
||||
daily_guardrail[key]["requests_evaluated"] += 1
|
||||
action = guardrail_status_to_action(entry.get("guardrail_status"))
|
||||
if action == "passed":
|
||||
daily_guardrail[key]["passed_count"] += 1
|
||||
elif action == "blocked":
|
||||
daily_guardrail[key]["blocked_count"] += 1
|
||||
else:
|
||||
daily_guardrail[key]["flagged_count"] += 1
|
||||
if action != "not_run":
|
||||
key = _MetricsKey(guardrail_id, date_key)
|
||||
daily_guardrail[key]["requests_evaluated"] += 1
|
||||
if action == "passed":
|
||||
daily_guardrail[key]["passed_count"] += 1
|
||||
elif action == "blocked":
|
||||
daily_guardrail[key]["blocked_count"] += 1
|
||||
else:
|
||||
daily_guardrail[key]["flagged_count"] += 1
|
||||
policy_id = entry.get("policy_id")
|
||||
index_rows.append(
|
||||
{
|
||||
prior = index_rows_by_key.get((request_id, guardrail_id))
|
||||
if prior is None or (prior["policy_id"] is None and policy_id is not None):
|
||||
index_rows_by_key[(request_id, guardrail_id)] = {
|
||||
"request_id": request_id,
|
||||
"guardrail_id": guardrail_id,
|
||||
"policy_id": policy_id,
|
||||
"start_time": start_time,
|
||||
}
|
||||
)
|
||||
index_rows: Final = tuple(index_rows_by_key.values())
|
||||
|
||||
async with pending.lock:
|
||||
pending_metrics: Final = pending.metrics
|
||||
|
|
|
|||
|
|
@ -221,6 +221,8 @@ LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset(
|
|||
)
|
||||
|
||||
_UNTRUSTED_ROOT_CONTROL_FIELDS: Final = (
|
||||
"weights",
|
||||
"_router_weights",
|
||||
"proxy_server_request",
|
||||
"standard_logging_object",
|
||||
"secret_fields",
|
||||
|
|
@ -334,7 +336,7 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg
|
|||
# and read by spend logs as fact; a client value has no legitimate meaning and no
|
||||
# key or team setting keeps it, so the strip is never gated.
|
||||
_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset(
|
||||
{"attempted_fallbacks", "original_model_group", CLIENT_OUTPUT_CEILING_METADATA_KEY}
|
||||
{"attempted_fallbacks", "original_model_group", "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY}
|
||||
)
|
||||
_ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override"
|
||||
|
||||
|
|
|
|||
|
|
@ -567,7 +567,7 @@ async def new_user(
|
|||
teams = check_if_default_team_set()
|
||||
organization_ids: Final = cast(list[str] | None, data_json.pop("organizations", None))
|
||||
|
||||
response: Final = await generate_key_helper_fn(request_type="user", **data_json)
|
||||
response: Final = await generate_key_helper_fn(request_type="user", **data_json, llm_router=None)
|
||||
# Admin UI Logic
|
||||
# Add User to Team and Organization
|
||||
# if team_id passed add this user to the team
|
||||
|
|
|
|||
|
|
@ -96,6 +96,7 @@ from litellm.proxy.management_endpoints.common_utils import (
|
|||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_add_model_to_db,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.router_weights import validate_router_settings_weights
|
||||
from litellm.proxy.management_helpers.access_group_key_sync import (
|
||||
sync_key_access_group_membership,
|
||||
sync_key_regeneration_access_group_membership,
|
||||
|
|
@ -201,6 +202,10 @@ class _KeyUpdateResult(TypedDict):
|
|||
data: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
class _StoredKeyRouterSettings(BaseModel):
|
||||
router_settings: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
class _KeyRowWhere(TypedDict):
|
||||
token: ReadOnly[str]
|
||||
|
||||
|
|
@ -1330,7 +1335,7 @@ async def _common_key_generation_helper(
|
|||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
response = await generate_key_helper_fn(request_type="key", **data_json, table_name="key")
|
||||
response = await generate_key_helper_fn(request_type="key", **data_json, table_name="key", llm_router=llm_router)
|
||||
|
||||
response["soft_budget"] = data.soft_budget # include the user-input soft budget in the response
|
||||
|
||||
|
|
@ -2234,7 +2239,26 @@ async def _update_key_row_with_soft_budget(
|
|||
async def prepare_key_update_data(
|
||||
data: UpdateKeyRequest | RegenerateKeyRequest,
|
||||
existing_key_row: LiteLLM_VerificationToken,
|
||||
*,
|
||||
prisma_client: PrismaClient | None = None,
|
||||
llm_router: Router | None = None,
|
||||
):
|
||||
if data.router_settings is not None or (
|
||||
"router_settings" not in data.model_fields_set
|
||||
and "team_id" in data.model_fields_set
|
||||
and data.team_id != existing_key_row.team_id
|
||||
):
|
||||
effective_settings: Final = (
|
||||
data.router_settings
|
||||
if data.router_settings is not None
|
||||
else _StoredKeyRouterSettings.model_validate(existing_key_row, from_attributes=True).router_settings
|
||||
)
|
||||
await validate_router_settings_weights(
|
||||
effective_settings,
|
||||
team_id=data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id,
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
data_json: Final[dict] = data.model_dump(exclude_unset=True)
|
||||
data_json.pop("key", None)
|
||||
data_json.pop("new_key", None)
|
||||
|
|
@ -2575,7 +2599,9 @@ async def _process_single_key_update(
|
|||
)
|
||||
|
||||
# Prepare update data
|
||||
non_default_values = await prepare_key_update_data(data=update_key_request, existing_key_row=existing_key_row)
|
||||
non_default_values = await prepare_key_update_data(
|
||||
data=update_key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router
|
||||
)
|
||||
|
||||
# Update key in database
|
||||
if prisma_client is None:
|
||||
|
|
@ -3093,7 +3119,9 @@ async def update_key_fn(
|
|||
|
||||
# Enforce upperbound key params on update (don't fill defaults)
|
||||
_enforce_upperbound_key_params(data, fill_defaults=False)
|
||||
non_default_values: Final = await prepare_key_update_data(data=data, existing_key_row=existing_key_row)
|
||||
non_default_values: Final = await prepare_key_update_data(
|
||||
data=data, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router
|
||||
)
|
||||
|
||||
# Only validate key_alias format if it's actually being changed
|
||||
new_key_alias: Final = non_default_values.get("key_alias", None)
|
||||
|
|
@ -4137,15 +4165,24 @@ async def generate_key_helper_fn(
|
|||
object_permission: LiteLLM_ObjectPermissionBase | None = None,
|
||||
auto_rotate: bool | None = None,
|
||||
rotation_interval: str | None = None,
|
||||
router_settings: dict | None = None,
|
||||
router_settings: dict[str, object] | None = None,
|
||||
access_group_ids: list[str] | None = None,
|
||||
budget_limits: list | None = None, # multiple concurrent budget windows
|
||||
*,
|
||||
llm_router: Router | None = None,
|
||||
):
|
||||
from litellm.proxy.proxy_server import premium_user, prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise Exception("Connect Proxy to database to generate keys - https://docs.litellm.ai/docs/proxy/virtual_keys ")
|
||||
|
||||
await validate_router_settings_weights(
|
||||
router_settings,
|
||||
team_id=team_id,
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
if token is None:
|
||||
if key is not None:
|
||||
token = key
|
||||
|
|
@ -5070,6 +5107,7 @@ async def _insert_deprecated_key(
|
|||
async def _execute_virtual_key_regeneration(
|
||||
*,
|
||||
prisma_client: PrismaClient,
|
||||
llm_router: Router | None = None,
|
||||
key_in_db: LiteLLM_VerificationToken,
|
||||
hashed_api_key: str,
|
||||
key: str,
|
||||
|
|
@ -5129,7 +5167,9 @@ async def _execute_virtual_key_regeneration(
|
|||
if data is not None:
|
||||
# Enforce upperbound key params on regenerate (don't fill defaults)
|
||||
_enforce_upperbound_key_params(data, fill_defaults=False)
|
||||
non_default_values = await prepare_key_update_data(data=data, existing_key_row=key_in_db)
|
||||
non_default_values = await prepare_key_update_data(
|
||||
data=data, existing_key_row=key_in_db, prisma_client=prisma_client, llm_router=llm_router
|
||||
)
|
||||
# Only validate key_alias format if it's actually being changed
|
||||
new_key_alias: Final = non_default_values.get("key_alias")
|
||||
if new_key_alias != key_in_db.key_alias:
|
||||
|
|
@ -5268,6 +5308,7 @@ async def regenerate_key_fn(
|
|||
try:
|
||||
from litellm.proxy.proxy_server import (
|
||||
hash_token,
|
||||
llm_router,
|
||||
master_key,
|
||||
premium_user,
|
||||
prisma_client,
|
||||
|
|
@ -5456,6 +5497,7 @@ async def regenerate_key_fn(
|
|||
|
||||
return await _execute_virtual_key_regeneration(
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
key_in_db=_key_in_db,
|
||||
hashed_api_key=hashed_api_key,
|
||||
key=key,
|
||||
|
|
|
|||
129
litellm/proxy/management_endpoints/router_weights.py
Normal file
129
litellm/proxy/management_endpoints/router_weights.py
Normal file
|
|
@ -0,0 +1,129 @@
|
|||
from abc import abstractmethod
|
||||
from collections.abc import Mapping
|
||||
from typing import Annotated, Final, Protocol
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import BaseModel, BeforeValidator, ValidationError
|
||||
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.types.router_weights import RouterWeights
|
||||
|
||||
|
||||
class _StoredModel(Protocol):
|
||||
@property
|
||||
@abstractmethod
|
||||
def model_id(self) -> str:
|
||||
pass
|
||||
|
||||
|
||||
class _ModelDb(Protocol):
|
||||
@property
|
||||
@abstractmethod
|
||||
def litellm_proxymodeltable(self) -> TableActions[_StoredModel]:
|
||||
pass
|
||||
|
||||
|
||||
class _PrismaClient(Protocol):
|
||||
@property
|
||||
@abstractmethod
|
||||
def db(self) -> _ModelDb:
|
||||
pass
|
||||
|
||||
|
||||
class _Router(Protocol):
|
||||
@abstractmethod
|
||||
def get_deployment(self, model_id: str) -> object | None:
|
||||
pass
|
||||
|
||||
|
||||
class _RouterWeightSettings(BaseModel):
|
||||
weights: RouterWeights | None = None
|
||||
|
||||
|
||||
class _RouterWeightModelInfo(BaseModel):
|
||||
team_id: str | None = None
|
||||
db_model: bool | None = None
|
||||
team_public_model_name: str | None = None
|
||||
|
||||
|
||||
def _router_weight_model_info(value: object) -> _RouterWeightModelInfo:
|
||||
if isinstance(value, str):
|
||||
return _RouterWeightModelInfo.model_validate_json(value)
|
||||
return _RouterWeightModelInfo.model_validate(value or {}, from_attributes=True)
|
||||
|
||||
|
||||
class _RouterWeightDeployment(BaseModel):
|
||||
model_name: str
|
||||
model_info: Annotated[_RouterWeightModelInfo, BeforeValidator(_router_weight_model_info)]
|
||||
|
||||
|
||||
def _validate_router_weight_reference(
|
||||
model_group: str,
|
||||
deployment_id: str,
|
||||
team_id: str | None,
|
||||
stored: _RouterWeightDeployment | None,
|
||||
configured: object | None,
|
||||
) -> None:
|
||||
reference: Final = (
|
||||
stored
|
||||
if stored is not None
|
||||
else (
|
||||
_RouterWeightDeployment.model_validate(configured, from_attributes=True) if configured is not None else None
|
||||
)
|
||||
)
|
||||
if (
|
||||
reference is None
|
||||
or (stored is None and reference.model_info.db_model)
|
||||
or (reference.model_info.team_id is not None and reference.model_info.team_id != team_id)
|
||||
):
|
||||
raise HTTPException(status_code=400, detail=f"Unknown deployment ID in router weights: {deployment_id}")
|
||||
canonical_group: Final = (
|
||||
reference.model_info.team_public_model_name if reference.model_info.team_id is not None else None
|
||||
) or reference.model_name
|
||||
if model_group != canonical_group:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Deployment {deployment_id} does not belong to model group {model_group}",
|
||||
)
|
||||
|
||||
|
||||
async def validate_router_settings_weights(
|
||||
router_settings: BaseModel | Mapping[str, object] | None,
|
||||
*,
|
||||
team_id: str | None,
|
||||
prisma_client: _PrismaClient | None,
|
||||
llm_router: _Router | None,
|
||||
) -> None:
|
||||
try:
|
||||
weights: Final = (
|
||||
_RouterWeightSettings.model_validate(router_settings, from_attributes=True).weights
|
||||
if router_settings is not None
|
||||
else None
|
||||
)
|
||||
except ValidationError:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Invalid router weights. Replace or clear router_settings.weights.",
|
||||
) from None
|
||||
if not weights:
|
||||
return
|
||||
deployment_ids: Final = frozenset(deployment_id for group in weights.values() for deployment_id in group)
|
||||
if not deployment_ids:
|
||||
return
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=503, detail="Database unavailable while validating router weights")
|
||||
stored_models: Final = await prisma_client.db.litellm_proxymodeltable.find_many(
|
||||
where={"model_id": {"in": list(deployment_ids)}}
|
||||
)
|
||||
stored_by_id: Final = {
|
||||
row.model_id: _RouterWeightDeployment.model_validate(row, from_attributes=True) for row in stored_models
|
||||
}
|
||||
for model_group, group_weights in weights.items():
|
||||
for deployment_id in group_weights:
|
||||
_validate_router_weight_reference(
|
||||
model_group,
|
||||
deployment_id,
|
||||
team_id,
|
||||
stored_by_id.get(deployment_id),
|
||||
llm_router.get_deployment(model_id=deployment_id) if llm_router is not None else None,
|
||||
)
|
||||
|
|
@ -112,6 +112,7 @@ from litellm.proxy.management_endpoints.common_utils import (
|
|||
from litellm.proxy.management_endpoints.organization_endpoints import (
|
||||
add_member_to_organization,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.router_weights import validate_router_settings_weights
|
||||
from litellm.proxy.management_endpoints.tag_management_endpoints import (
|
||||
get_daily_activity,
|
||||
)
|
||||
|
|
@ -1288,6 +1289,7 @@ async def new_team(
|
|||
create_audit_log_for_update,
|
||||
general_settings,
|
||||
litellm_proxy_admin_name,
|
||||
llm_router,
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
|
@ -1462,6 +1464,13 @@ async def new_team(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
await validate_router_settings_weights(
|
||||
data.router_settings,
|
||||
team_id=data.team_id,
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
## ADD TO MODEL TABLE
|
||||
_model_id = None
|
||||
if data.model_aliases is not None and isinstance(data.model_aliases, dict):
|
||||
|
|
@ -2075,6 +2084,13 @@ async def update_team(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
await validate_router_settings_weights(
|
||||
data.router_settings,
|
||||
team_id=data.team_id,
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
_existing_team_metadata: Final[object] = getattr(existing_team_row, "metadata", None)
|
||||
enforce_output_token_estimates_are_admin_only(
|
||||
data=data,
|
||||
|
|
|
|||
|
|
@ -3592,6 +3592,7 @@ class SSOAuthenticationHandler:
|
|||
verbose_proxy_logger.info("user_defined_values for creating ui key: %s", user_defined_values)
|
||||
|
||||
response: Final = await generate_key_helper_fn(
|
||||
llm_router=None,
|
||||
request_type="key",
|
||||
duration=LITELLM_UI_SESSION_DURATION,
|
||||
key_max_budget=litellm.max_ui_session_budget,
|
||||
|
|
|
|||
|
|
@ -58,15 +58,6 @@ class UndeliverableStreamRewrite(Exception):
|
|||
self.guardrail_name: Final = guardrail_name
|
||||
|
||||
|
||||
class UnappliableRequestRewrite(Exception):
|
||||
def __init__(self, guardrail_name: str) -> None:
|
||||
super().__init__(
|
||||
f"Guardrail '{guardrail_name}' rewrote the request in a way this endpoint cannot apply, "
|
||||
"so the request was rejected rather than sent unrewritten"
|
||||
)
|
||||
self.guardrail_name: Final = guardrail_name
|
||||
|
||||
|
||||
def _tool_call_shape(tool_call: object) -> tuple[object, object]:
|
||||
plain: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call
|
||||
function: Final = plain.get("function") if isinstance(plain, Mapping) else None
|
||||
|
|
|
|||
|
|
@ -9546,6 +9546,7 @@ class ProxyStartupEvent:
|
|||
gate the first duration window.
|
||||
"""
|
||||
await generate_key_helper_fn(
|
||||
llm_router=llm_router,
|
||||
request_type="user",
|
||||
table_name="user",
|
||||
user_id=LITELLM_PROXY_BUDGET_NAME,
|
||||
|
|
@ -16290,6 +16291,7 @@ async def _generate_onboarding_ui_session_token(user_obj: _UserTableRow) -> str:
|
|||
global master_key, general_settings
|
||||
|
||||
response: Final = await generate_key_helper_fn(
|
||||
llm_router=llm_router,
|
||||
request_type="key",
|
||||
**{
|
||||
"user_role": user_obj.user_role,
|
||||
|
|
|
|||
65
litellm/responses/additional_tools.py
Normal file
65
litellm/responses/additional_tools.py
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, cast # noqa: TID251 # validating the openai tool union strips vendor keys from raw tools
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS, ResponseInputParam
|
||||
|
||||
ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools"
|
||||
|
||||
|
||||
class _InputItemType(BaseModel):
|
||||
type: str = ""
|
||||
|
||||
|
||||
class _AdditionalToolsItem(BaseModel):
|
||||
tools: tuple[dict[str, object], ...] = ()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class HoistedAdditionalTools:
|
||||
input: str | ResponseInputParam
|
||||
tools: tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]
|
||||
hoisted: tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]
|
||||
|
||||
|
||||
def _is_additional_tools_item(item: object) -> bool:
|
||||
try:
|
||||
return _InputItemType.model_validate(item).type == ADDITIONAL_TOOLS_INPUT_ITEM_TYPE
|
||||
except ValidationError:
|
||||
return False
|
||||
|
||||
|
||||
def _tools_of_item(item: object) -> tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]:
|
||||
try:
|
||||
parsed: Final = _AdditionalToolsItem.model_validate(item)
|
||||
except ValidationError:
|
||||
return ()
|
||||
return tuple(
|
||||
cast(
|
||||
"ALL_RESPONSES_API_TOOL_PARAMS", tool
|
||||
) # cast-ok: nested tools carry the same raw tool JSON as top-level tools
|
||||
for tool in parsed.tools
|
||||
)
|
||||
|
||||
|
||||
def hoist_additional_tools(
|
||||
input: str | ResponseInputParam,
|
||||
tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None,
|
||||
) -> HoistedAdditionalTools:
|
||||
existing: Final = tuple(tools or ())
|
||||
if isinstance(input, str):
|
||||
return HoistedAdditionalTools(input=input, tools=existing, hoisted=())
|
||||
items: Final = tuple(item for item in input if _is_additional_tools_item(item))
|
||||
if not items:
|
||||
return HoistedAdditionalTools(input=input, tools=existing, hoisted=())
|
||||
hoisted: Final = tuple(tool for item in items for tool in _tools_of_item(item))
|
||||
verbose_logger.debug(
|
||||
"Responses API: hoisting %d tool(s) out of %d 'additional_tools' input item(s) into the top-level tools param.",
|
||||
len(hoisted),
|
||||
len(items),
|
||||
)
|
||||
remaining_input: Final = [item for item in input if not _is_additional_tools_item(item)]
|
||||
return HoistedAdditionalTools(input=remaining_input, tools=(*existing, *hoisted), hoisted=hoisted)
|
||||
|
|
@ -39,15 +39,38 @@ def openai_shaped_tool_call_item_id(item_type: str, tool_id: str) -> str:
|
|||
return f"{prefix}_{tool_id}"
|
||||
|
||||
|
||||
class _ToolNameFields(BaseModel):
|
||||
type: str = ""
|
||||
name: str = ""
|
||||
tools: tuple[object, ...] = ()
|
||||
|
||||
|
||||
def _tool_name_fields_of(tool: object) -> _ToolNameFields | None:
|
||||
try:
|
||||
return _ToolNameFields.model_validate(tool)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _custom_tool_name_of(tool: object) -> str | None:
|
||||
parsed: Final = _tool_name_fields_of(tool)
|
||||
if parsed is None or parsed.type != "custom" or not parsed.name:
|
||||
return None
|
||||
return parsed.name
|
||||
|
||||
|
||||
def _nested_tools_of(tool: object) -> tuple[object, ...]:
|
||||
parsed: Final = _tool_name_fields_of(tool)
|
||||
if parsed is None or parsed.type != "namespace":
|
||||
return ()
|
||||
return parsed.tools
|
||||
|
||||
|
||||
def extract_custom_tool_names(tools: Sequence[object] | None) -> set[str]:
|
||||
"""Extract names of tools originally defined as ``type: "custom"``."""
|
||||
if not tools:
|
||||
return set()
|
||||
names: Final[set[str]] = set()
|
||||
for tool in tools:
|
||||
if isinstance(tool, dict) and tool.get("type") == "custom" and "name" in tool:
|
||||
names.add(tool["name"])
|
||||
return names
|
||||
"""Extract names of ``type: "custom"`` tools, at the top level or one level inside a ``namespace`` tool."""
|
||||
top_level: Final = tuple(tools or ())
|
||||
nested: Final = tuple(nested_tool for tool in top_level for nested_tool in _nested_tools_of(tool))
|
||||
return {name for tool in (*top_level, *nested) if (name := _custom_tool_name_of(tool)) is not None}
|
||||
|
||||
|
||||
def is_custom_tool_call(tool_name: str, custom_tool_names: set[str]) -> bool:
|
||||
|
|
@ -143,7 +166,7 @@ def validated_allowed_callers(value: object) -> list[str] | None:
|
|||
raise ValueError("allowed_callers must be a list of strings") from exc
|
||||
|
||||
|
||||
def _grammar_suffix(fmt: object) -> str:
|
||||
def custom_tool_grammar_suffix(fmt: object) -> str:
|
||||
try:
|
||||
parsed: Final = _CustomToolFormat.model_validate(fmt)
|
||||
except ValidationError:
|
||||
|
|
@ -167,7 +190,9 @@ def convert_custom_tool_to_function_tool(tool: Mapping[str, object]) -> ChatComp
|
|||
raw_name: Final = tool.get("name")
|
||||
name: Final = raw_name if isinstance(raw_name, str) else ""
|
||||
raw_description: Final = tool.get("description")
|
||||
description = (raw_description if isinstance(raw_description, str) else "") + _grammar_suffix(tool.get("format"))
|
||||
description: Final = (raw_description if isinstance(raw_description, str) else "") + custom_tool_grammar_suffix(
|
||||
tool.get("format")
|
||||
)
|
||||
allowed_callers: Final = validated_allowed_callers(tool.get("allowed_callers"))
|
||||
function_chunk: Final = ChatCompletionToolParamFunctionChunk(
|
||||
name=name,
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from collections.abc import Coroutine, Mapping
|
|||
from typing import Final
|
||||
|
||||
import litellm
|
||||
from litellm.responses.additional_tools import hoist_additional_tools
|
||||
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
|
||||
LiteLLMCompletionStreamingIterator,
|
||||
)
|
||||
|
|
@ -37,11 +38,16 @@ class LiteLLMCompletionTransformationHandler:
|
|||
| BaseResponsesAPIStreamingIterator
|
||||
| Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]
|
||||
):
|
||||
hoisted: Final = hoist_additional_tools(input, responses_api_request.get("tools"))
|
||||
bridged_input: Final = hoisted.input
|
||||
bridged_request: Final[ResponsesAPIOptionalRequestParams] = (
|
||||
{**responses_api_request, "tools": list(hoisted.tools)} if hoisted.hoisted else responses_api_request
|
||||
)
|
||||
litellm_completion_request: Final[dict] = (
|
||||
LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
|
||||
model=model,
|
||||
input=input,
|
||||
responses_api_request=responses_api_request,
|
||||
input=bridged_input,
|
||||
responses_api_request=bridged_request,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
stream=stream,
|
||||
extra_headers=extra_headers,
|
||||
|
|
@ -52,8 +58,8 @@ class LiteLLMCompletionTransformationHandler:
|
|||
if _is_async:
|
||||
return self.async_response_api_handler(
|
||||
litellm_completion_request=litellm_completion_request,
|
||||
request_input=input,
|
||||
responses_api_request=responses_api_request,
|
||||
request_input=bridged_input,
|
||||
responses_api_request=bridged_request,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
|
@ -70,8 +76,8 @@ class LiteLLMCompletionTransformationHandler:
|
|||
responses_api_response: Final[ResponsesAPIResponse] = (
|
||||
LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(
|
||||
chat_completion_response=litellm_completion_response,
|
||||
request_input=input,
|
||||
responses_api_request=responses_api_request,
|
||||
request_input=bridged_input,
|
||||
responses_api_request=bridged_request,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -81,8 +87,8 @@ class LiteLLMCompletionTransformationHandler:
|
|||
return LiteLLMCompletionStreamingIterator(
|
||||
model=model,
|
||||
litellm_custom_stream_wrapper=litellm_completion_response,
|
||||
request_input=input,
|
||||
responses_api_request=responses_api_request,
|
||||
request_input=bridged_input,
|
||||
responses_api_request=bridged_request,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_metadata=kwargs.get("litellm_metadata", {}),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from litellm.main import stream_chunk_builder
|
|||
from litellm.responses.litellm_completion_transformation.custom_tools import (
|
||||
build_tool_call_item_kwargs,
|
||||
extract_custom_tool_names,
|
||||
is_custom_tool_call,
|
||||
serialize_tool_call_arguments,
|
||||
)
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
|
|
@ -166,6 +167,14 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
return tool_name, namespace
|
||||
return fn_name, None
|
||||
|
||||
def _tool_call_item_kwargs(self, call_id: str, fn_name: str, arguments: str, status: str) -> dict[str, str]:
|
||||
item_kwargs: Final = build_tool_call_item_kwargs(call_id, fn_name, arguments, status, self._custom_tool_names)
|
||||
if is_custom_tool_call(fn_name, self._custom_tool_names):
|
||||
return item_kwargs
|
||||
tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name)
|
||||
namespace_kwargs: Final = {"namespace": tool_namespace} if tool_namespace else {}
|
||||
return {**item_kwargs, "name": tool_name, **namespace_kwargs}
|
||||
|
||||
def _is_reasoning_end(self, chunk):
|
||||
delta: Final = chunk.choices[0].delta
|
||||
|
||||
|
|
@ -244,17 +253,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
else:
|
||||
fn_name = str(getattr(fn, "name", "") or "")
|
||||
fn_args_delta = serialize_tool_call_arguments(getattr(fn, "arguments", ""))
|
||||
tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name)
|
||||
output_index = self._get_or_assign_tool_output_index(call_id)
|
||||
|
||||
if call_id not in self._tool_args_by_call_id:
|
||||
self._tool_args_by_call_id[call_id] = ""
|
||||
self._sequence_number += 1
|
||||
names = self._custom_tool_names
|
||||
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names)
|
||||
item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, "", "in_progress")
|
||||
self._tool_item_id_by_call_id[call_id] = item_kwargs["id"]
|
||||
if tool_namespace:
|
||||
item_kwargs["namespace"] = tool_namespace
|
||||
event = OutputItemAddedEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
|
||||
output_index=output_index,
|
||||
|
|
@ -315,7 +320,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
else:
|
||||
fn_name = str(getattr(fn, "name", "") or "")
|
||||
fn_args = serialize_tool_call_arguments(getattr(fn, "arguments", ""))
|
||||
tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name)
|
||||
web_search_call = self._web_search_calls.get(call_id)
|
||||
if web_search_call is not None:
|
||||
if call_id not in self._queued_web_search_call_ids:
|
||||
|
|
@ -330,11 +334,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
if is_new_tool_call:
|
||||
self._tool_args_by_call_id[call_id] = ""
|
||||
self._sequence_number += 1
|
||||
names = self._custom_tool_names
|
||||
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names)
|
||||
item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, "", "in_progress")
|
||||
self._tool_item_id_by_call_id[call_id] = item_kwargs["id"]
|
||||
if tool_namespace:
|
||||
item_kwargs["namespace"] = tool_namespace
|
||||
event = OutputItemAddedEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
|
||||
output_index=output_index,
|
||||
|
|
@ -376,11 +377,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
self._pending_tool_events.append(done_event)
|
||||
|
||||
self._sequence_number += 1
|
||||
names = self._custom_tool_names
|
||||
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, final_args, "completed", names)
|
||||
item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, final_args, "completed")
|
||||
item_kwargs["id"] = self._tool_item_id_by_call_id.setdefault(call_id, item_kwargs["id"])
|
||||
if tool_namespace:
|
||||
item_kwargs["namespace"] = tool_namespace
|
||||
item_done_event = OutputItemDoneEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
|
||||
output_index=output_index,
|
||||
|
|
|
|||
|
|
@ -110,6 +110,7 @@ NamespaceTool: TypeAlias = Mapping[str, object]
|
|||
ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None
|
||||
ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool
|
||||
NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n"
|
||||
NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: Final = frozenset({"function", "custom"})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -1891,9 +1892,21 @@ class LiteLLMCompletionResponsesConfig:
|
|||
namespace_tool: NamespaceTool,
|
||||
nested: bool,
|
||||
) -> ChatCompletionToolParam | None:
|
||||
if nested and namespace_tool.get("type") != "function":
|
||||
tool_type: Final = namespace_tool.get("type")
|
||||
if nested and tool_type not in NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS:
|
||||
return None
|
||||
|
||||
raw_description: Final = str(namespace_tool.get("description") or "")
|
||||
description: Final = (
|
||||
f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}{raw_description}"
|
||||
if nested and namespace_description and raw_description
|
||||
else namespace_description
|
||||
if nested and namespace_description
|
||||
else raw_description
|
||||
)
|
||||
if nested and tool_type == "custom":
|
||||
return convert_custom_tool_to_function_tool({**namespace_tool, "description": description})
|
||||
|
||||
raw_parameters: Final = namespace_tool.get("parameters")
|
||||
parameters: Final = (
|
||||
MappingProxyType(raw_parameters) if isinstance(raw_parameters, Mapping) else MappingProxyType({})
|
||||
|
|
@ -1902,14 +1915,6 @@ class LiteLLMCompletionResponsesConfig:
|
|||
parameters if parameters and "type" in parameters else MappingProxyType({**parameters, "type": "object"})
|
||||
)
|
||||
tool_name: Final = str(namespace_tool.get("name") or "")
|
||||
raw_description: Final = str(namespace_tool.get("description") or "")
|
||||
description: Final = (
|
||||
f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}{raw_description}"
|
||||
if nested and namespace_description and raw_description
|
||||
else namespace_description
|
||||
if nested and namespace_description
|
||||
else raw_description
|
||||
)
|
||||
chat_tool_name: Final = f"{namespace}__{tool_name}" if nested else tool_name
|
||||
function: Final = ChatCompletionToolParamFunctionChunk(
|
||||
name=chat_tool_name,
|
||||
|
|
|
|||
|
|
@ -4849,6 +4849,7 @@ class Router:
|
|||
model=model,
|
||||
messages=messages,
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
request_kwargs=kwargs,
|
||||
)
|
||||
|
||||
data: Final = deployment["litellm_params"].copy()
|
||||
|
|
@ -5163,13 +5164,11 @@ class Router:
|
|||
return healthy_deployments[0]
|
||||
|
||||
# Use simple_shuffle for weighted selection
|
||||
return cast(
|
||||
GuardrailTypedDict,
|
||||
simple_shuffle(
|
||||
llm_router_instance=self,
|
||||
healthy_deployments=healthy_deployments,
|
||||
model=guardrail_name,
|
||||
),
|
||||
return simple_shuffle(
|
||||
resolve_model_alias=self._get_model_from_alias,
|
||||
healthy_deployments=healthy_deployments,
|
||||
model=guardrail_name,
|
||||
request_kwargs=None,
|
||||
)
|
||||
|
||||
async def _ageneric_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs):
|
||||
|
|
@ -8378,7 +8377,8 @@ class Router:
|
|||
|
||||
def log_retry(self, kwargs: dict, e: Exception) -> dict:
|
||||
"""
|
||||
When a retry or fallback happens, record which model group, deployment and attempt just failed and why
|
||||
When a retry or fallback happens, record which model group, deployment and attempt just failed and why,
|
||||
and count it toward the request-wide num_retries_per_request cap
|
||||
"""
|
||||
from litellm.types.router import RetryAttemptRecord
|
||||
|
||||
|
|
@ -8402,7 +8402,10 @@ class Router:
|
|||
else ()
|
||||
)
|
||||
breadcrumbs: Final = (*kept_breadcrumbs, attempt_record)
|
||||
earlier: Final = request_metadata.get("request_retry_count")
|
||||
request_retry_count: Final = (earlier if type(earlier) is int and 0 <= earlier else 0) + 1
|
||||
kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict
|
||||
kwargs[_metadata_var]["request_retry_count"] = request_retry_count # rebind-ok: same dict, read by the cap
|
||||
return kwargs
|
||||
|
||||
def _update_usage(self, deployment_id: str, parent_otel_span: Span | None) -> int:
|
||||
|
|
@ -13045,9 +13048,10 @@ class Router:
|
|||
start_time: Final = time.time()
|
||||
if strategy == "simple-shuffle":
|
||||
return simple_shuffle(
|
||||
llm_router_instance=self,
|
||||
resolve_model_alias=self._get_model_from_alias,
|
||||
healthy_deployments=healthy_deployments,
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
deployment: Final = await self._select_deployment_async(
|
||||
strategy=strategy,
|
||||
|
|
@ -13190,9 +13194,10 @@ class Router:
|
|||
start_time: Final = time.perf_counter()
|
||||
if strategy == "simple-shuffle":
|
||||
return simple_shuffle(
|
||||
llm_router_instance=self,
|
||||
resolve_model_alias=self._get_model_from_alias,
|
||||
healthy_deployments=pass_through_deployments,
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
deployment: Final = await self._select_deployment_async(
|
||||
strategy=strategy,
|
||||
|
|
@ -13888,9 +13893,10 @@ class Router:
|
|||
# if users pass rpm or tpm, we do a random weighted pick - based on rpm/tpm
|
||||
############## Check 'weight' param set for weighted pick #################
|
||||
return simple_shuffle(
|
||||
llm_router_instance=self,
|
||||
resolve_model_alias=self._get_model_from_alias,
|
||||
healthy_deployments=healthy_deployments,
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
deployment: Final = self._select_deployment_sync(
|
||||
strategy=strategy,
|
||||
|
|
@ -13958,6 +13964,7 @@ class Router:
|
|||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
|
||||
strategy, strategy_selector = self._get_routing_context(model, request_kwargs)
|
||||
|
|
@ -14040,9 +14047,10 @@ class Router:
|
|||
# 6. Apply load balancing strategy
|
||||
if strategy == "simple-shuffle":
|
||||
return simple_shuffle(
|
||||
llm_router_instance=self,
|
||||
resolve_model_alias=self._get_model_from_alias,
|
||||
healthy_deployments=pass_through_deployments,
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
deployment: Final = self._select_deployment_sync(
|
||||
strategy=strategy,
|
||||
|
|
|
|||
|
|
@ -1,71 +1,67 @@
|
|||
"""
|
||||
Returns a random deployment from the list of healthy deployments.
|
||||
"""Choose among eligible deployments using request weights, then global metrics."""
|
||||
|
||||
If weights are provided, it will return a deployment based on the weights.
|
||||
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import random
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from itertools import chain
|
||||
from typing import Final, TypeVar
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.types.router_weights import validate_router_weights
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router as _Router
|
||||
_DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object])
|
||||
_ROUTER_LOGGER: Final = logging.getLogger("LiteLLM Router")
|
||||
|
||||
LitellmRouter = _Router
|
||||
else:
|
||||
LitellmRouter = Any
|
||||
|
||||
def _metric_weight(deployment: Mapping[str, object], metric: str) -> float:
|
||||
params: Final = deployment.get("litellm_params")
|
||||
value: Final = params.get(metric) if isinstance(params, Mapping) else None
|
||||
if value is None:
|
||||
return 0.0
|
||||
if isinstance(value, (int, float)):
|
||||
return float(value)
|
||||
raise TypeError(f"Deployment {metric} must be numeric")
|
||||
|
||||
|
||||
def _scoped_weights(
|
||||
deployments: Sequence[Mapping[str, object]],
|
||||
model: str,
|
||||
request_kwargs: Mapping[str, object] | None,
|
||||
) -> tuple[float, ...]:
|
||||
settings: Final = validate_router_weights((request_kwargs or {}).get("_router_weights"))
|
||||
model_weights: Final = settings.get(model) if settings is not None else None
|
||||
if not model_weights:
|
||||
return ()
|
||||
return tuple(
|
||||
model_weights.get(str(info.get("id")), 0.0) if isinstance(info, Mapping) else 0.0
|
||||
for deployment in deployments
|
||||
for info in (deployment.get("model_info"),)
|
||||
)
|
||||
|
||||
|
||||
def simple_shuffle(
|
||||
llm_router_instance: LitellmRouter,
|
||||
healthy_deployments: list[Any] | dict[Any, Any],
|
||||
resolve_model_alias: Callable[[str], str | None],
|
||||
healthy_deployments: Sequence[_DeploymentT],
|
||||
model: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Returns a random deployment from the list of healthy deployments.
|
||||
|
||||
If weights are provided, it will return a deployment based on the weights.
|
||||
|
||||
If users pass `rpm` or `tpm`, we do a random weighted pick - based on `rpm`/`tpm`.
|
||||
|
||||
Args:
|
||||
llm_router_instance: LitellmRouter instance
|
||||
healthy_deployments: List of healthy deployments
|
||||
model: Model name
|
||||
|
||||
Returns:
|
||||
Dict: A single healthy deployment
|
||||
"""
|
||||
|
||||
############## Check if 'weight' or 'rpm' or 'tpm' param set for a weighted pick #################
|
||||
for weight_by in ["weight", "rpm", "tpm"]:
|
||||
if any(m["litellm_params"].get(weight_by) is not None for m in healthy_deployments):
|
||||
weights = [m["litellm_params"].get(weight_by, 0) for m in healthy_deployments]
|
||||
verbose_router_logger.debug("\nweight %s", weights)
|
||||
total_weight = sum(weights)
|
||||
if total_weight <= 0:
|
||||
# All remaining candidates have weight 0 for this metric (e.g.
|
||||
# after a weighted-failover exclusion left only zero-weight
|
||||
# backups). Skip to the next metric (rpm/tpm) which may still
|
||||
# provide a meaningful weighted pick; if none do, we fall
|
||||
# through to the uniform random pick at the end.
|
||||
continue
|
||||
weights = [weight / total_weight for weight in weights]
|
||||
verbose_router_logger.debug("\n weights %s by %s", weights, weight_by)
|
||||
# Perform weighted random pick
|
||||
selected_index = random.choices(range(len(weights)), weights=weights)[0]
|
||||
verbose_router_logger.debug("\n selected index, %s", selected_index)
|
||||
deployment = healthy_deployments[selected_index]
|
||||
verbose_router_logger.info(
|
||||
"get_available_deployment for model: %s, Selected deployment: %s for model: %s",
|
||||
model,
|
||||
llm_router_instance.print_deployment(deployment) or deployment[0],
|
||||
model,
|
||||
)
|
||||
return deployment or deployment[0]
|
||||
|
||||
############## No RPM/TPM passed, we do a random pick #################
|
||||
item: Final = random.choice(healthy_deployments)
|
||||
return item or item[0]
|
||||
request_kwargs: Mapping[str, object] | None,
|
||||
) -> _DeploymentT:
|
||||
resolved_model: Final = resolve_model_alias(model) or model
|
||||
weight_sets: Final = chain(
|
||||
(_scoped_weights(healthy_deployments, resolved_model, request_kwargs),),
|
||||
(
|
||||
tuple(_metric_weight(deployment, metric) for deployment in healthy_deployments)
|
||||
for metric in ("weight", "rpm", "tpm")
|
||||
),
|
||||
)
|
||||
for weights in weight_sets:
|
||||
largest = max(weights, default=0.0)
|
||||
if largest <= 0:
|
||||
continue
|
||||
normalized = tuple(weight / largest for weight in weights)
|
||||
if sum(normalized) <= 0:
|
||||
continue
|
||||
selected = random.choices(healthy_deployments, weights=normalized)[0]
|
||||
_ROUTER_LOGGER.info("Selected deployment for model %s: %s", model, selected.get("model_info"))
|
||||
return selected
|
||||
return random.choice(healthy_deployments)
|
||||
|
|
|
|||
|
|
@ -48,6 +48,7 @@ from litellm.exceptions import (
|
|||
ServiceUnavailableError,
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger, Span
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
encrypted_content_of_block,
|
||||
strip_encrypted_reasoning_from_messages,
|
||||
|
|
@ -215,11 +216,11 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
@staticmethod
|
||||
def _encryption_boundary_key(
|
||||
litellm_params: object,
|
||||
) -> tuple | None:
|
||||
) -> tuple[object, object] | None:
|
||||
"""
|
||||
``(api_base, api_key)`` pair identifying an Azure resource. Two
|
||||
deployments sharing both are interchangeable for ``encrypted_content``
|
||||
follow-ups; Azure rejects content produced by any other resource.
|
||||
``(api_base, api_key)`` identifies an upstream encryption boundary.
|
||||
The values are resolved from the deployment and its named credential
|
||||
without modifying the deployment.
|
||||
|
||||
Accepts any object exposing dict-style ``.get(key, default)``: plain
|
||||
dicts (the common case in ``healthy_deployments``) as well as
|
||||
|
|
@ -234,9 +235,25 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
return None
|
||||
api_base: Final = getter("api_base")
|
||||
api_key: Final = getter("api_key")
|
||||
if not api_base or not api_key:
|
||||
credential_name: Final = getter("litellm_credential_name")
|
||||
credential_values: Final[Mapping[str, object] | None] = (
|
||||
CredentialAccessor.get_credential_values(credential_name)
|
||||
if isinstance(credential_name, str) and credential_name
|
||||
else None
|
||||
)
|
||||
effective_api_base: Final = (
|
||||
credential_values.get("api_base")
|
||||
if credential_values is not None and "api_base" in credential_values
|
||||
else api_base
|
||||
)
|
||||
effective_api_key: Final = (
|
||||
credential_values.get("api_key")
|
||||
if credential_values is not None and "api_key" in credential_values
|
||||
else api_key
|
||||
)
|
||||
if not effective_api_base or not effective_api_key:
|
||||
return None
|
||||
return (api_base, api_key)
|
||||
return (effective_api_base, effective_api_key)
|
||||
|
||||
def _find_deployments_on_same_encryption_boundary(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
from typing import Any, Final, Literal
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Final, Literal, cast # noqa: TID251 # JSON chat rows have no typed constructor across roles
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from typing_extensions import TypedDict
|
||||
|
|
@ -158,12 +159,21 @@ def coerce_stream_holdback_value(value: Any) -> int:
|
|||
return 0
|
||||
|
||||
|
||||
def structured_messages_from_response(value: object) -> Sequence[AllMessageValues] | None:
|
||||
if not isinstance(value, list):
|
||||
return None
|
||||
if not all(isinstance(message, Mapping) and isinstance(message.get("role"), str) for message in value):
|
||||
return None
|
||||
return cast("Sequence[AllMessageValues]", value) # cast-ok: JSON rows checked for a role, the same trust texts get
|
||||
|
||||
|
||||
class GenericGuardrailAPIResponse:
|
||||
"""Response model for the Generic Guardrail API"""
|
||||
|
||||
texts: list[str] | None
|
||||
images: list[str] | None
|
||||
tools: list[GuardrailToolParam] | None
|
||||
structured_messages: Sequence[AllMessageValues] | None
|
||||
action: str
|
||||
blocked_reason: str | None
|
||||
stream_holdback_chars: list[int] | None
|
||||
|
|
@ -176,12 +186,14 @@ class GenericGuardrailAPIResponse:
|
|||
images: list[str] | None = None,
|
||||
tools: list[GuardrailToolParam] | None = None,
|
||||
stream_holdback_chars: list[int] | None = None,
|
||||
structured_messages: Sequence[AllMessageValues] | None = None,
|
||||
) -> None:
|
||||
self.action = action
|
||||
self.blocked_reason = blocked_reason
|
||||
self.texts = texts
|
||||
self.images = images
|
||||
self.tools = tools
|
||||
self.structured_messages = structured_messages
|
||||
# Number of trailing chars, indexed the same as ``texts``, that the
|
||||
# framework must withhold from streaming emission until the next
|
||||
# processing round (word-boundary safety for text transformations).
|
||||
|
|
@ -200,4 +212,5 @@ class GenericGuardrailAPIResponse:
|
|||
images=data.get("images"),
|
||||
tools=data.get("tools"),
|
||||
stream_holdback_chars=stream_holdback_chars,
|
||||
structured_messages=structured_messages_from_response(data.get("structured_messages")),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_c
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
|
||||
from litellm.types.router_weights import RouterWeights
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
|
@ -146,6 +147,7 @@ class UpdateRouterConfig(BaseModel):
|
|||
context_window_fallbacks: list[dict] | None = None
|
||||
model_group_alias: dict[str, str | dict] | None = {}
|
||||
enable_tag_filtering: bool | None = None
|
||||
weights: RouterWeights | None = None
|
||||
tag_routing_prefix: str | None = None
|
||||
optional_pre_call_checks: OptionalPreCallChecks | None = None
|
||||
|
||||
|
|
|
|||
30
litellm/types/router_weights.py
Normal file
30
litellm/types/router_weights.py
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Annotated, Final
|
||||
|
||||
from pydantic import AfterValidator, Field, TypeAdapter
|
||||
|
||||
|
||||
def _validate_positive_router_weights(weights: Mapping[str, Mapping[str, float]]) -> Mapping[str, Mapping[str, float]]:
|
||||
if any(group and not any(weight > 0 for weight in group.values()) for group in weights.values()):
|
||||
raise ValueError("Each nonempty weights group must contain at least one positive weight")
|
||||
return weights
|
||||
|
||||
|
||||
RouterWeightIdentifier = Annotated[str, Field(strict=True, min_length=1, pattern=r"\S")]
|
||||
RouterWeight = Annotated[float, Field(strict=True, ge=0, allow_inf_nan=False)]
|
||||
RouterWeights = Annotated[
|
||||
dict[RouterWeightIdentifier, dict[RouterWeightIdentifier, RouterWeight]],
|
||||
AfterValidator(_validate_positive_router_weights),
|
||||
]
|
||||
_ROUTER_WEIGHTS_ADAPTER: Final[TypeAdapter[RouterWeights | None]] = TypeAdapter(RouterWeights | None)
|
||||
_ROUTER_SETTINGS_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def validate_router_weights(value: object) -> RouterWeights | None:
|
||||
return _ROUTER_WEIGHTS_ADAPTER.validate_python(value)
|
||||
|
||||
|
||||
def validate_router_settings_dict(value: object) -> dict[str, object]:
|
||||
settings: Final = _ROUTER_SETTINGS_DICT_ADAPTER.validate_python(value)
|
||||
validate_router_weights(settings.get("weights"))
|
||||
return settings
|
||||
|
|
@ -283,8 +283,11 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
input_cost_per_video_token: float | None # for gemini omni models with video input
|
||||
input_cost_per_audio_per_second: float | None # only for vertex ai models
|
||||
input_cost_per_video_per_second: float | None # only for vertex ai models
|
||||
input_cost_per_audio_token_batches: ReadOnly[float | None]
|
||||
input_cost_per_image_token_batches: ReadOnly[float | None]
|
||||
input_cost_per_second: float | None # for OpenAI Speech models
|
||||
input_cost_per_token_batches: float | None
|
||||
input_cost_per_video_token_batches: ReadOnly[float | None]
|
||||
output_cost_per_token_batches: float | None
|
||||
output_cost_per_token: Required[float | None]
|
||||
output_cost_per_token_flex: float | None # OpenAI flex service tier pricing
|
||||
|
|
@ -3583,7 +3586,10 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
|
|||
input_cost_per_video_per_second_above_128k_tokens: float | None = None
|
||||
input_cost_per_video_per_second_above_15s_interval: float | None = None
|
||||
input_cost_per_video_per_second_above_8s_interval: float | None = None
|
||||
input_cost_per_audio_token_batches: float | None = None
|
||||
input_cost_per_image_token_batches: float | None = None
|
||||
input_cost_per_token_batches: float | None = None
|
||||
input_cost_per_video_token_batches: float | None = None
|
||||
output_cost_per_token_batches: float | None = None
|
||||
output_cost_per_token_flex: float | None = None
|
||||
output_cost_per_token_priority: float | None = None
|
||||
|
|
@ -3776,6 +3782,7 @@ all_litellm_params = (
|
|||
"id",
|
||||
"fallbacks",
|
||||
"routing_strategy",
|
||||
"_router_weights",
|
||||
"azure",
|
||||
"headers",
|
||||
"model_list",
|
||||
|
|
|
|||
|
|
@ -1208,30 +1208,13 @@ def _dispatch_success_logging(
|
|||
is_litellm_internal_call: bool,
|
||||
) -> None:
|
||||
if not is_litellm_internal_call:
|
||||
if getattr(logging_obj, "_defer_async_logging", False):
|
||||
|
||||
def _enqueue_deferred_logging() -> None:
|
||||
asyncio.create_task(
|
||||
_client_async_logging_helper(
|
||||
logging_obj=logging_obj,
|
||||
result=result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
is_completion_with_fallbacks=is_completion_with_fallbacks,
|
||||
)
|
||||
)
|
||||
|
||||
logging_obj._enqueue_deferred_logging = _enqueue_deferred_logging
|
||||
else:
|
||||
asyncio.create_task(
|
||||
_client_async_logging_helper(
|
||||
logging_obj=logging_obj,
|
||||
result=result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
is_completion_with_fallbacks=is_completion_with_fallbacks,
|
||||
)
|
||||
)
|
||||
_schedule_async_success_logging(
|
||||
logging_obj=logging_obj,
|
||||
result=result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
is_completion_with_fallbacks=is_completion_with_fallbacks,
|
||||
)
|
||||
|
||||
logging_obj.handle_sync_success_callbacks_for_async_calls(
|
||||
result=result,
|
||||
|
|
@ -1240,6 +1223,43 @@ def _dispatch_success_logging(
|
|||
)
|
||||
|
||||
|
||||
def _schedule_async_success_logging(
|
||||
logging_obj: LiteLLMLoggingObject,
|
||||
result: object,
|
||||
start_time: datetime.datetime,
|
||||
end_time: datetime.datetime,
|
||||
is_completion_with_fallbacks: bool,
|
||||
) -> None:
|
||||
"""Fire the async success log for ``result`` now, or park it on the logging object while
|
||||
the proxy defers logging past its post-call guardrails.
|
||||
|
||||
Nested @client wrappers (Anthropic Messages over the chat adapter, chat over the Responses
|
||||
bridge) each exit through here with the same logging object and their own shape of the same
|
||||
response. The immediate path already logs one request once, since the first task marks
|
||||
``has_logged_async_success`` and the later ones skip. The deferred slot keeps the same
|
||||
first-wins rule: the innermost wrapper's provider-shaped result is the one the spend log
|
||||
reads usage from, and a later wrapper never swaps in its client-shaped translation.
|
||||
"""
|
||||
|
||||
def _enqueue_async_logging() -> None:
|
||||
asyncio.create_task(
|
||||
_client_async_logging_helper(
|
||||
logging_obj=logging_obj,
|
||||
result=result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
is_completion_with_fallbacks=is_completion_with_fallbacks,
|
||||
)
|
||||
)
|
||||
|
||||
if not getattr(logging_obj, "_defer_async_logging", False):
|
||||
_enqueue_async_logging()
|
||||
return
|
||||
if getattr(logging_obj, "_enqueue_deferred_logging", None) is not None:
|
||||
return
|
||||
logging_obj._enqueue_deferred_logging = _enqueue_async_logging
|
||||
|
||||
|
||||
async def _client_async_logging_helper(
|
||||
logging_obj: LiteLLMLoggingObject,
|
||||
result,
|
||||
|
|
@ -5923,10 +5943,13 @@ def _get_model_info_helper(
|
|||
input_cost_per_audio_token=_model_info.get("input_cost_per_audio_token", None),
|
||||
input_cost_per_image_token=_model_info.get("input_cost_per_image_token", None),
|
||||
input_cost_per_video_token=_model_info.get("input_cost_per_video_token", None),
|
||||
input_cost_per_audio_token_batches=_model_info.get("input_cost_per_audio_token_batches", None),
|
||||
input_cost_per_image_token_batches=_model_info.get("input_cost_per_image_token_batches", None),
|
||||
input_cost_per_image=_model_info.get("input_cost_per_image", None),
|
||||
input_cost_per_audio_per_second=_model_info.get("input_cost_per_audio_per_second", None),
|
||||
input_cost_per_video_per_second=_model_info.get("input_cost_per_video_per_second", None),
|
||||
input_cost_per_token_batches=_model_info.get("input_cost_per_token_batches"),
|
||||
input_cost_per_video_token_batches=_model_info.get("input_cost_per_video_token_batches", None),
|
||||
output_cost_per_token_batches=_model_info.get("output_cost_per_token_batches"),
|
||||
output_cost_per_token=_output_cost_per_token,
|
||||
output_cost_per_token_flex=_model_info.get("output_cost_per_token_flex", None),
|
||||
|
|
|
|||
|
|
@ -25601,10 +25601,14 @@
|
|||
"source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models"
|
||||
},
|
||||
"gemini-embedding-2-preview": {
|
||||
"input_cost_per_audio_per_second": 0.00016,
|
||||
"input_cost_per_image": 0.00012,
|
||||
"input_cost_per_audio_token": 6.5e-06,
|
||||
"input_cost_per_audio_token_batches": 3.25e-06,
|
||||
"input_cost_per_image_token": 4.5e-07,
|
||||
"input_cost_per_image_token_batches": 2.25e-07,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"input_cost_per_video_per_second": 0.00079,
|
||||
"input_cost_per_token_batches": 1e-07,
|
||||
"input_cost_per_video_token": 1.2e-05,
|
||||
"input_cost_per_video_token_batches": 6e-06,
|
||||
"litellm_provider": "vertex_ai-embedding-models",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
|
|
@ -25615,13 +25619,14 @@
|
|||
"uses_embed_content": true
|
||||
},
|
||||
"gemini-embedding-2": {
|
||||
"input_cost_per_audio_per_second": 0.00016,
|
||||
"input_cost_per_audio_token": 6.5e-06,
|
||||
"input_cost_per_image": 0.00012,
|
||||
"input_cost_per_audio_token_batches": 3.25e-06,
|
||||
"input_cost_per_image_token": 4.5e-07,
|
||||
"input_cost_per_image_token_batches": 2.25e-07,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"input_cost_per_token_batches": 1e-07,
|
||||
"input_cost_per_video_per_second": 0.00079,
|
||||
"input_cost_per_video_token": 1.2e-05,
|
||||
"input_cost_per_video_token_batches": 6e-06,
|
||||
"litellm_provider": "vertex_ai-embedding-models",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
|
|
@ -25633,10 +25638,14 @@
|
|||
"uses_embed_content": true
|
||||
},
|
||||
"vertex_ai/gemini-embedding-2-preview": {
|
||||
"input_cost_per_audio_per_second": 0.00016,
|
||||
"input_cost_per_image": 0.00012,
|
||||
"input_cost_per_audio_token": 6.5e-06,
|
||||
"input_cost_per_audio_token_batches": 3.25e-06,
|
||||
"input_cost_per_image_token": 4.5e-07,
|
||||
"input_cost_per_image_token_batches": 2.25e-07,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"input_cost_per_video_per_second": 0.00079,
|
||||
"input_cost_per_token_batches": 1e-07,
|
||||
"input_cost_per_video_token": 1.2e-05,
|
||||
"input_cost_per_video_token_batches": 6e-06,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
|
|
@ -25648,13 +25657,14 @@
|
|||
"uses_embed_content": true
|
||||
},
|
||||
"vertex_ai/gemini-embedding-2": {
|
||||
"input_cost_per_audio_per_second": 0.00016,
|
||||
"input_cost_per_audio_token": 6.5e-06,
|
||||
"input_cost_per_image": 0.00012,
|
||||
"input_cost_per_audio_token_batches": 3.25e-06,
|
||||
"input_cost_per_image_token": 4.5e-07,
|
||||
"input_cost_per_image_token_batches": 2.25e-07,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"input_cost_per_token_batches": 1e-07,
|
||||
"input_cost_per_video_per_second": 0.00079,
|
||||
"input_cost_per_video_token": 1.2e-05,
|
||||
"input_cost_per_video_token_batches": 6e-06,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
|
|
@ -25693,10 +25703,14 @@
|
|||
},
|
||||
"gemini/gemini-embedding-2-preview": {
|
||||
"deprecation_date": "2026-08-10",
|
||||
"input_cost_per_audio_per_second": 0.00016,
|
||||
"input_cost_per_image": 0.00012,
|
||||
"input_cost_per_audio_token": 6.5e-06,
|
||||
"input_cost_per_audio_token_batches": 3.25e-06,
|
||||
"input_cost_per_image_token": 4.5e-07,
|
||||
"input_cost_per_image_token_batches": 2.25e-07,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"input_cost_per_video_per_second": 0.00079,
|
||||
"input_cost_per_token_batches": 1e-07,
|
||||
"input_cost_per_video_token": 1.2e-05,
|
||||
"input_cost_per_video_token_batches": 6e-06,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
|
|
@ -25709,10 +25723,14 @@
|
|||
"tpm": 10000000
|
||||
},
|
||||
"gemini/gemini-embedding-2": {
|
||||
"input_cost_per_audio_per_second": 0.00016,
|
||||
"input_cost_per_image": 0.00012,
|
||||
"input_cost_per_audio_token": 6.5e-06,
|
||||
"input_cost_per_audio_token_batches": 3.25e-06,
|
||||
"input_cost_per_image_token": 4.5e-07,
|
||||
"input_cost_per_image_token_batches": 2.25e-07,
|
||||
"input_cost_per_token": 2e-07,
|
||||
"input_cost_per_video_per_second": 0.00079,
|
||||
"input_cost_per_token_batches": 1e-07,
|
||||
"input_cost_per_video_token": 1.2e-05,
|
||||
"input_cost_per_video_token_batches": 6e-06,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
|
|
|
|||
|
|
@ -249,6 +249,10 @@
|
|||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_audio_token_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_audio_token_priority": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -276,6 +280,10 @@
|
|||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_image_token_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_pixel": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
|
|
@ -375,6 +383,14 @@
|
|||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"input_cost_per_video_token": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_video_token_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_dbu_cost_per_token": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
|
|
|
|||
|
|
@ -81,6 +81,25 @@ def test_missing_execution_evidence_fails(tmp_path: Path, contents: str) -> None
|
|||
assert result.returncode == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omitted_role", ("proxy_admin", "team_member", "internal_user_viewer"))
|
||||
def test_one_passing_management_case_cannot_hide_a_missing_actor(tmp_path: Path, omitted_role: str) -> None:
|
||||
suite: Final = ET.Element("testsuite")
|
||||
path: Final = "tests/e2e/management/test_jwt_management_e2e.py"
|
||||
case: Final = ET.SubElement(suite, "testcase", file=path)
|
||||
properties: Final = ET.SubElement(case, "properties")
|
||||
_ = ET.SubElement(
|
||||
properties,
|
||||
"property",
|
||||
name="management_node",
|
||||
value=f"{path}::TestJwtManagement::test_actor_subject_and_database_role[proxy_admin_viewer]",
|
||||
)
|
||||
report: Final = tmp_path / "report.xml"
|
||||
ET.ElementTree(suite).write(report)
|
||||
result: Final = subprocess.run([sys.executable, str(GATE), str(report), path], capture_output=True, text=True)
|
||||
assert result.returncode == 1
|
||||
assert f"test_actor_subject_and_database_role[{omitted_role}]" in result.stdout
|
||||
|
||||
|
||||
def test_short_values_are_written_without_masking_every_digit_in_the_log(tmp_path: Path) -> None:
|
||||
env_path: Final = tmp_path / ".env"
|
||||
|
||||
|
|
@ -141,6 +160,10 @@ def test_changed_suite_files_are_selected_unless_the_stack_cannot_run_them(
|
|||
(
|
||||
"tests/e2e/proxy_client.py",
|
||||
"tests/e2e/conftest.py",
|
||||
"tests/e2e/management/management_client.py",
|
||||
"tests/e2e/management/jwt_actors.py",
|
||||
"tests/e2e/management/conftest.py",
|
||||
"tests/e2e/coverage_registry/management_cases.py",
|
||||
"tests/e2e/pytest.ini",
|
||||
"tests/e2e/gateway/stage_mirror_ci_config.yml",
|
||||
".github/e2e-stack/up.sh",
|
||||
|
|
|
|||
329
tests/code_coverage_tests/test_provider_replay_harness.py
Normal file
329
tests/code_coverage_tests/test_provider_replay_harness.py
Normal file
|
|
@ -0,0 +1,329 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from fixture_bundle import BundleRecorder, LoadedBundle, load_bundle, prepare_bundle
|
||||
from fixture_mode import current_test_key
|
||||
from fixture_profile import MatchProfile
|
||||
from provider_edge import REPLAY_MISS_STATUS, RecordEdge, ReplayEdge, ReplaySource
|
||||
from test_provider_edge import (
|
||||
CHAT_PATH,
|
||||
SSE_CHUNKS,
|
||||
STREAM_BODY,
|
||||
UPLOAD_PATH,
|
||||
call_edge,
|
||||
chunked_provider,
|
||||
fake_provider,
|
||||
json_object,
|
||||
provider_url,
|
||||
raw_stream_post,
|
||||
running_edge,
|
||||
this_tests_files,
|
||||
)
|
||||
|
||||
|
||||
class TestStrictIdentity:
|
||||
@pytest.mark.parametrize("path", [CHAT_PATH, "/anthropic/v1/messages"])
|
||||
def test_roundtrip_rejects_semantic_changes(self, tmp_path: Path, path: str) -> None:
|
||||
recorder: Final = prepare_bundle(tmp_path / "strict", profile="stateless_v1")
|
||||
assert isinstance(recorder, BundleRecorder)
|
||||
original: Final = (
|
||||
b'{"model":"synthetic","messages":[{"role":"user",'
|
||||
b'"content":"2031-04-05 00000000-0000-0000-0000-000000000001"}],"options":[1,2]}'
|
||||
)
|
||||
headers: Final = {
|
||||
"content-type": "application/json",
|
||||
"accept": "application/json",
|
||||
"anthropic-version": "2023-06-01",
|
||||
"anthropic-beta": "feature-a",
|
||||
"openai-beta": "feature-b",
|
||||
"authorization": "Bearer synthetic-secret-one",
|
||||
}
|
||||
query: Final = "?part=one&part=two&blank="
|
||||
with fake_provider() as provider:
|
||||
mounts: Final = {"openai": provider_url(provider), "anthropic": provider_url(provider)}
|
||||
with running_edge(RecordEdge(recorder, threading.Lock()), mounts) as edge:
|
||||
captured: Final = call_edge(edge, "POST", path + query, body=original, headers=headers)
|
||||
assert captured.status_code == 200
|
||||
assert json_object(captured.body)["echo"] == original.decode()
|
||||
loaded: Final = load_bundle(recorder.root, profile="stateless_v1")
|
||||
assert isinstance(loaded, LoadedBundle)
|
||||
assert loaded.manifest.match_profile == "stateless_v1"
|
||||
source: Final = ReplaySource(loaded)
|
||||
with running_edge(ReplayEdge(source), mounts) as edge:
|
||||
cases: Final = (
|
||||
(original.replace(b"2031-04-05", b"2032-06-07"), headers, query, "body"),
|
||||
(original.replace(b"000000000001", b"000000000002"), headers, query, "body"),
|
||||
(original.replace(b"[1,2]", b"[2,1]"), headers, query, "body"),
|
||||
(original.replace(b"synthetic", b"other"), headers, query, "body"),
|
||||
(original, headers, "?part=three&part=two&blank=", "query"),
|
||||
(original, headers, "?part=two&part=one&blank=", "query"),
|
||||
*(
|
||||
(original, {k: v for k, v in headers.items() if k != name}, query, "headers")
|
||||
for name in ("accept", "anthropic-version", "anthropic-beta", "openai-beta")
|
||||
),
|
||||
*(
|
||||
(original, {**headers, name: value}, query, "headers")
|
||||
for name in ("accept", "anthropic-version", "anthropic-beta", "openai-beta")
|
||||
for value in ("different", "")
|
||||
),
|
||||
(original, {**headers, "authorization": "Basic synthetic-secret-two"}, query, "auth"),
|
||||
(original, {k: v for k, v in headers.items() if k != "authorization"}, query, "auth"),
|
||||
)
|
||||
for rejected, reason in (
|
||||
(call_edge(edge, "POST", path + changed_query, body=body, headers=changed_headers), reason)
|
||||
for body, changed_headers, changed_query, reason in cases
|
||||
):
|
||||
assert rejected.status_code == REPLAY_MISS_STATUS
|
||||
assert reason in rejected.body.decode()
|
||||
assert b"synthetic-secret" not in rejected.body
|
||||
reordered: Final = json.dumps(dict(reversed(list(json_object(original).items())))).encode()
|
||||
accepted: Final = call_edge(
|
||||
edge, "POST", path + query, body=reordered, headers={k.upper(): v for k, v in headers.items()}
|
||||
)
|
||||
assert accepted.status_code == 200
|
||||
assert accepted.body == captured.body
|
||||
assert source.leftover_error(current_test_key()) is None
|
||||
assert len(provider.hits) == 1
|
||||
assert "synthetic-secret" not in "".join(file.read_text() for file in recorder.root.rglob("*.json"))
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[
|
||||
b'{"value":null}',
|
||||
b'{"value":""}',
|
||||
b'{"value":false}',
|
||||
b'{"value":0}',
|
||||
b'{"value":[]}',
|
||||
b'{"value":{}}',
|
||||
b'{"value":0.123456789012345678901}',
|
||||
b'{"value":0.123456789012345678902}',
|
||||
b'{"value":1e400}',
|
||||
b'{"value":1}',
|
||||
b'{"value":1e0}',
|
||||
b'{"value":-0}',
|
||||
b'{"value":1e9999999999999999999}',
|
||||
],
|
||||
)
|
||||
def test_json_values_remain_distinct(self, tmp_path: Path, body: bytes) -> None:
|
||||
recorder: Final = prepare_bundle(tmp_path / "strict", profile="stateless_v1")
|
||||
assert isinstance(recorder, BundleRecorder)
|
||||
with fake_provider() as provider:
|
||||
mounts: Final = {"openai": provider_url(provider)}
|
||||
with running_edge(RecordEdge(recorder, threading.Lock()), mounts) as edge:
|
||||
assert (
|
||||
call_edge(
|
||||
edge, "POST", CHAT_PATH, body=body, headers={"content-type": "application/json"}
|
||||
).status_code
|
||||
== 200
|
||||
)
|
||||
loaded: Final = load_bundle(recorder.root, profile="stateless_v1")
|
||||
assert isinstance(loaded, LoadedBundle)
|
||||
with running_edge(ReplayEdge(ReplaySource(loaded)), mounts) as edge:
|
||||
values: Final = (
|
||||
b"{}",
|
||||
b'{"value":null}',
|
||||
b'{"value":""}',
|
||||
b'{"value":false}',
|
||||
b'{"value":0}',
|
||||
b'{"value":[]}',
|
||||
b'{"value":{}}',
|
||||
b'{"value":0.123456789012345678901}',
|
||||
b'{"value":0.123456789012345678902}',
|
||||
b'{"value":1e400}',
|
||||
b'{"value":1}',
|
||||
b'{"value":1e0}',
|
||||
b'{"value":-0}',
|
||||
b'{"value":1e9999999999999999999}',
|
||||
)
|
||||
for rejected in (
|
||||
call_edge(edge, "POST", CHAT_PATH, body=value, headers={"content-type": "application/json"})
|
||||
for value in values
|
||||
if value != body
|
||||
):
|
||||
assert rejected.status_code == REPLAY_MISS_STATUS
|
||||
assert b"body" in rejected.body
|
||||
assert (
|
||||
call_edge(
|
||||
edge, "POST", CHAT_PATH, body=body, headers={"content-type": "application/json"}
|
||||
).status_code
|
||||
== 200
|
||||
)
|
||||
assert len(provider.hits) == 1
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"path,body,headers",
|
||||
[
|
||||
(UPLOAD_PATH, b"{}", {"content-type": "application/json"}),
|
||||
(CHAT_PATH + "?part=%FF", b"{}", {"content-type": "application/json"}),
|
||||
(CHAT_PATH + "?part=%FE", b"{}", {"content-type": "application/json"}),
|
||||
(CHAT_PATH, b"opaque", {"content-type": "application/octet-stream"}),
|
||||
(CHAT_PATH, b"--boundary", {"content-type": "multipart/form-data; boundary=boundary"}),
|
||||
(CHAT_PATH, b'{"x":1,"x":2}', {"content-type": "application/json"}),
|
||||
(CHAT_PATH, b"{}", {"content-type": "application/json", "x-custom-behavior": "synthetic-private-value"}),
|
||||
],
|
||||
)
|
||||
def test_ineligible_capture_never_calls_provider(
|
||||
self, tmp_path: Path, path: str, body: bytes, headers: dict[str, str]
|
||||
) -> None:
|
||||
recorder: Final = prepare_bundle(tmp_path / "strict", profile="stateless_v1")
|
||||
assert isinstance(recorder, BundleRecorder)
|
||||
with fake_provider() as provider:
|
||||
with running_edge(RecordEdge(recorder, threading.Lock()), {"openai": provider_url(provider)}) as edge:
|
||||
result: Final = call_edge(edge, "POST", path, body=body, headers=headers)
|
||||
assert result.status_code == REPLAY_MISS_STATUS
|
||||
assert b"eligibility error" in result.body
|
||||
assert b"synthetic-private-value" not in result.body
|
||||
assert provider.hits == []
|
||||
assert this_tests_files(recorder.root) == []
|
||||
|
||||
def test_destination_is_part_of_actual_http_identity(self, tmp_path: Path) -> None:
|
||||
recorder: Final = prepare_bundle(tmp_path / "strict", profile="stateless_v1")
|
||||
assert isinstance(recorder, BundleRecorder)
|
||||
with fake_provider() as provider:
|
||||
with running_edge(RecordEdge(recorder, threading.Lock()), {"openai": provider_url(provider)}) as edge:
|
||||
assert (
|
||||
call_edge(
|
||||
edge, "POST", CHAT_PATH, body=b"{}", headers={"content-type": "application/json"}
|
||||
).status_code
|
||||
== 200
|
||||
)
|
||||
loaded: Final = load_bundle(recorder.root, profile="stateless_v1")
|
||||
assert isinstance(loaded, LoadedBundle)
|
||||
with running_edge(ReplayEdge(ReplaySource(loaded)), {"openai": provider_url(provider) + "/other"}) as edge:
|
||||
result: Final = call_edge(
|
||||
edge, "POST", CHAT_PATH, body=b"{}", headers={"content-type": "application/json"}
|
||||
)
|
||||
assert result.status_code == REPLAY_MISS_STATUS
|
||||
assert b"upstream" in result.body
|
||||
assert len(provider.hits) == 1
|
||||
|
||||
def test_credentials_are_not_identity_and_fresh_process_replays(self, tmp_path: Path) -> None:
|
||||
recorder: Final = prepare_bundle(tmp_path / "strict", profile="stateless_v1")
|
||||
assert isinstance(recorder, BundleRecorder)
|
||||
headers: Final = {
|
||||
"content-type": "application/json",
|
||||
"authorization": "bEaReR synthetic-token",
|
||||
"x-api-key": "synthetic-api-key",
|
||||
"cookie": "synthetic-cookie",
|
||||
}
|
||||
path: Final = CHAT_PATH + "?api_key=synthetic-query-secret&part=one&part=two"
|
||||
body: Final = b'{"model":"synthetic","messages":[]}'
|
||||
with fake_provider(echo_request=False) as provider:
|
||||
mounts: Final = {"openai": provider_url(provider)}
|
||||
with running_edge(RecordEdge(recorder, threading.Lock()), mounts) as edge:
|
||||
captured: Final = call_edge(edge, "POST", path, body=body, headers=headers)
|
||||
assert captured.status_code == 200
|
||||
seen_headers, seen_body = provider.requests[0]
|
||||
assert {k.lower(): v for k, v in seen_headers.items()}.items() >= headers.items()
|
||||
assert seen_body == body
|
||||
assert provider.hits == ["POST " + path.removeprefix("/openai")]
|
||||
artifacts: Final = "".join(file.read_text() for file in recorder.root.rglob("*.json"))
|
||||
for secret in ("synthetic-token", "synthetic-api-key", "synthetic-cookie", "synthetic-query-secret"):
|
||||
assert secret not in artifacts
|
||||
child: Final = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-c",
|
||||
"""
|
||||
import json, sys
|
||||
from pathlib import Path
|
||||
from fixture_bundle import LoadedBundle, load_bundle
|
||||
from provider_edge import ProviderRequestObservation, observed_provider_edge, replay_leftover_error
|
||||
from test_provider_edge import call_edge
|
||||
from fixture_profile import MatchProfile
|
||||
from fixture_mode import current_test_key
|
||||
loaded = load_bundle(Path(sys.argv[1]), profile="stateless_v1")
|
||||
assert isinstance(loaded, LoadedBundle)
|
||||
with observed_provider_edge(ProviderRequestObservation("synthetic"), mode_raw="replay", bundle_dir=Path(sys.argv[1]), bind_host="127.0.0.1", advertise_host="127.0.0.1", mounts={"openai": sys.argv[2]}) as edge:
|
||||
response = call_edge(edge, "POST", sys.argv[3], body=sys.argv[4].encode(), headers=json.loads(sys.argv[5]))
|
||||
assert response.status_code == 200
|
||||
print(response.body.decode())
|
||||
assert replay_leftover_error(mode_raw="replay", bundle_dir=Path(sys.argv[1]), test_key=current_test_key()) is None
|
||||
""",
|
||||
str(recorder.root),
|
||||
provider_url(provider),
|
||||
path.replace("synthetic-query-secret", "new-query-credential"),
|
||||
body.decode(),
|
||||
json.dumps({**headers, "authorization": "Bearer another-credential", "x-api-key": "another-key"}),
|
||||
],
|
||||
env={
|
||||
**os.environ,
|
||||
"PYTHONPATH": str(Path(__file__).resolve().parents[1] / "e2e"),
|
||||
"E2E_REPLAY_MATCH_PROFILE": "stateless_v1",
|
||||
},
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
assert child.returncode == 0, child.stderr
|
||||
assert child.stdout.strip().encode() == captured.body
|
||||
assert len(provider.hits) == 1
|
||||
|
||||
@pytest.mark.parametrize("profile,other", [("legacy", "stateless_v1"), ("stateless_v1", "legacy")])
|
||||
def test_profiles_cannot_load_each_others_bundles(
|
||||
self, tmp_path: Path, profile: MatchProfile, other: MatchProfile
|
||||
) -> None:
|
||||
from fixture_bundle import UnreadableBundle
|
||||
|
||||
recorder: Final = prepare_bundle(tmp_path / profile, profile=profile)
|
||||
assert isinstance(recorder, BundleRecorder)
|
||||
mismatch: Final = load_bundle(recorder.root, profile=other)
|
||||
assert isinstance(mismatch, UnreadableBundle)
|
||||
assert "profile mismatch" in mismatch.reason
|
||||
assert "re-record" in mismatch.reason
|
||||
|
||||
@pytest.mark.parametrize("abort_after", [None, 2])
|
||||
def test_strict_stream_preserves_chunks_and_truncation(self, tmp_path: Path, abort_after: int | None) -> None:
|
||||
recorder: Final = prepare_bundle(tmp_path / "strict", profile="stateless_v1")
|
||||
assert isinstance(recorder, BundleRecorder)
|
||||
with chunked_provider(abort_after=abort_after) as provider:
|
||||
mounts: Final = {"anthropic": provider_url(provider)}
|
||||
with running_edge(RecordEdge(recorder, threading.Lock()), mounts) as edge:
|
||||
_, captured, captured_ending = raw_stream_post(edge.port, "/anthropic/v1/messages", STREAM_BODY)
|
||||
loaded: Final = load_bundle(recorder.root, profile="stateless_v1")
|
||||
assert isinstance(loaded, LoadedBundle)
|
||||
source: Final = ReplaySource(loaded)
|
||||
with running_edge(ReplayEdge(source), mounts) as edge:
|
||||
_, replayed, ending = raw_stream_post(edge.port, "/anthropic/v1/messages", STREAM_BODY)
|
||||
assert captured == replayed == list(SSE_CHUNKS[:abort_after])
|
||||
assert ending == captured_ending
|
||||
assert (ending == "terminated") == (abort_after is None)
|
||||
assert source.leftover_error(current_test_key()) is None
|
||||
assert len(provider.hits) == 1
|
||||
|
||||
def test_auth_scheme_survives_missing_credentials(self, tmp_path: Path) -> None:
|
||||
recorder: Final = prepare_bundle(tmp_path / "strict", profile="stateless_v1")
|
||||
assert isinstance(recorder, BundleRecorder)
|
||||
headers: Final = {"content-type": "application/json", "authorization": "Bearer"}
|
||||
with fake_provider() as provider:
|
||||
mounts: Final = {"openai": provider_url(provider)}
|
||||
with running_edge(RecordEdge(recorder, threading.Lock()), mounts) as edge:
|
||||
assert call_edge(edge, "POST", CHAT_PATH, body=b"{}", headers=headers).status_code == 200
|
||||
loaded: Final = load_bundle(recorder.root, profile="stateless_v1")
|
||||
assert isinstance(loaded, LoadedBundle)
|
||||
with running_edge(ReplayEdge(ReplaySource(loaded)), mounts) as edge:
|
||||
for result in (
|
||||
call_edge(edge, "POST", CHAT_PATH, body=b"{}", headers={**headers, "authorization": scheme})
|
||||
for scheme in ("Basic", "Digest")
|
||||
):
|
||||
assert result.status_code == REPLAY_MISS_STATUS
|
||||
assert b"auth" in result.body
|
||||
assert (
|
||||
call_edge(
|
||||
edge,
|
||||
"POST",
|
||||
CHAT_PATH,
|
||||
body=b"{}",
|
||||
headers={**headers, "authorization": "bEaReR synthetic-token"},
|
||||
).status_code
|
||||
== 200
|
||||
)
|
||||
assert len(provider.hits) == 1
|
||||
|
|
@ -60,7 +60,11 @@ The suites run against a live proxy, so bring one up first by running the litell
|
|||
|
||||
Keycloak's password grant is a test-only provisioning shortcut, not a production login recommendation. The `litellm-e2e-admin` client adds the proxy's admin scope; the normal client does not. Never reuse this permissive realm outside an isolated test stack.
|
||||
|
||||
Management tests can use the shared `idp` and `jwt_identity` fixtures. Each test gets a unique Keycloak group/user and a matching proxy user/team. Setup and fallback cleanup use the master key; the operations and read-backs being tested must explicitly use `caller_key=idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID)` (or a member token). See `management/test_jwt_management_e2e.py` for create/read/update/clear/delete and tenant-denial examples. A group claim alone is not database team membership: permission tests explicitly add the member and prove an allowed read before asserting the denied write.
|
||||
Management tests can bind a credential once with `client.with_caller(Caller(...))`; direct calls, delegated helpers and replica read-backs then retain that caller. Explicit `caller_key` arguments override the binding. Keep the original master-backed client for bootstrap and cleanup. `actor_factory` lazily provisions database roles and tenant memberships, with `database_role` tokens carrying no groups and `group_scoped` actors retaining the existing team route gate. Token minting is explicit through `actor.mint_caller(idp)`. The factory runs requests without backend retries and reports cleanup failures. `coverage_registry/management_cases.py` records exact canary nodes and non-secret actor labels; the CI execution assertion rejects a missing or skipped actor row
|
||||
|
||||
For the opt-in browser profile, start the existing IdP first, then run `.github/e2e-stack/oidc-profile.sh "$PROXY_BASE_URL" <server-command>`. The wrapper creates a confidential client with an exact `/sso/callback` redirect and S256 PKCE, passes the client secret only through the child process environment, and removes the client on exit. It uses the existing generic OIDC handler with `GENERIC_USER_ID_ATTRIBUTE=sub`. Preserve the IdP's PostgreSQL data across restarts
|
||||
|
||||
`tests/e2e/ui/playwright.oidc.config.ts` uses an already running OIDC stack and separate storage/output files. Supply `E2E_OIDC_UI_URL`, `JWT_ISSUER`, `E2E_OIDC_USERNAME` and `E2E_OIDC_PASSWORD` for a seeded actor. Its setup follows the real login and callback path. The current Python canary qualifies browser-client configuration and token/userinfo identity mapping; browser journey specs under `ui/oidc/` are a separate coverage step
|
||||
|
||||
Every successful IdP create immediately registers cleanup, including partial setup failures. Cleanup failures emit warnings. Tokens are minted on demand, and the expiration test waits relative to the token's actual `exp` with a bounded clock-drift check. To check first-attempt behavior locally, run both files with `--reruns 0`:
|
||||
|
||||
|
|
@ -232,3 +236,15 @@ Before you push
|
|||
4. Capture screenshots of the test run and attach them to the PR as proof
|
||||
|
||||
5. If a test fails because it surfaced a real issue in the product, flag that explicitly in the PR rather than reworking the test until it passes
|
||||
|
||||
### Strict stateless replay matching
|
||||
|
||||
Set `E2E_REPLAY_MATCH_PROFILE=stateless_v1` for both recording and replay to bind OpenAI `/v1/chat/completions` and Anthropic `/v1/messages` requests to their upstream destination, ordered query pairs, semantic headers and literal JSON content. The default remains `legacy`. Strict bundles use format 5 and cannot load as legacy bundles; select the matching profile or re-record with `E2E_FIXTURE_MODE=record`. Missing profile metadata never enrolls a legacy bundle in strict matching
|
||||
|
||||
Strict matching preserves dates, UUIDs, hashes, model names, tool arguments, array order and omitted/null/empty/false/zero values. JSON object key order and header name casing may change. The strict body uses tagged JSON values so number precision and JSON types survive persistence, including exact numeric spelling and numbers larger than a floating-point value. Invalid UTF-8 query values fail eligibility. Duplicate JSON keys, unsupported endpoints, non-JSON bodies and unknown semantic headers fail eligibility before contacting a provider
|
||||
|
||||
The semantic header set is `content-type`, `accept`, `anthropic-version`, `anthropic-beta` and `openai-beta`, including missing versus present values. Authorization records presence and the case-insensitive scheme; `x-api-key` records presence only. Credential values and cookies are excluded. Credential query values are redacted while their position and field name remain in the identity. Never use real customer inputs in fixture qualification
|
||||
|
||||
Excluded transport and telemetry headers are `host`, `content-length`, `connection`, `accept-encoding`, `user-agent`, `traceparent`, `tracestate`, `x-request-id`, `x-client-request-id` and `x-stainless-*`. Inbound transfer-encoding is unsupported; send JSON with content-length framing. The destination represents host identity and the relay carries original body bytes. Replay does not verify credentials, SDK timeout/retry behavior, transport performance, model availability or stateful remote IDs. Live relay uses original request bytes and header values, never the stored identity
|
||||
|
||||
Strict replay harness regression tests live in `tests/code_coverage_tests/test_provider_replay_harness.py`. The CircleCI `provider_replay_harness` job runs them alongside the existing legacy harness files with `--noconftest -o pythonpath=tests/e2e`; they need only synthetic HTTP providers and temporary fixture storage
|
||||
|
|
|
|||
151
tests/e2e/coverage_registry/management_cases.py
Normal file
151
tests/e2e/coverage_registry/management_cases.py
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
from dataclasses import dataclass
|
||||
from typing import Final, Literal
|
||||
|
||||
CredentialKind = Literal["master", "idp_admin", "direct_jwt", "virtual_key", "dashboard_session"]
|
||||
DependencyProfile = Literal["management_only", "real_oidc_browser", "external_provider_required"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ManagementCase:
|
||||
node: str
|
||||
credential_kind: CredentialKind
|
||||
actor: str
|
||||
profile: str
|
||||
method: Literal["GET", "POST"]
|
||||
path: str
|
||||
operation_family: str
|
||||
dependency_profile: DependencyProfile = "management_only"
|
||||
|
||||
|
||||
JWT_FILE: Final = "tests/e2e/management/test_jwt_management_e2e.py"
|
||||
JWT_CLASS: Final = f"{JWT_FILE}::TestJwtManagement"
|
||||
ACTORS: Final = (
|
||||
"proxy_admin",
|
||||
"proxy_admin_viewer",
|
||||
"organization_admin",
|
||||
"team_admin",
|
||||
"team_member",
|
||||
"internal_user",
|
||||
"internal_user_viewer",
|
||||
"unrelated_user",
|
||||
)
|
||||
MANAGEMENT_CASES: Final = tuple(
|
||||
ManagementCase(
|
||||
node=f"{JWT_CLASS}::test_actor_subject_and_database_role[{role}]",
|
||||
credential_kind="direct_jwt",
|
||||
actor=role,
|
||||
profile="database_role",
|
||||
method="GET",
|
||||
path="/user/info",
|
||||
operation_family="identity",
|
||||
)
|
||||
for role in ACTORS
|
||||
) + (
|
||||
ManagementCase(
|
||||
node=f"{JWT_CLASS}::test_admin_viewer_reads_but_cannot_update",
|
||||
credential_kind="direct_jwt",
|
||||
actor="proxy_admin_viewer",
|
||||
profile="database_role",
|
||||
method="POST",
|
||||
path="/key/update",
|
||||
operation_family="denial",
|
||||
),
|
||||
ManagementCase(
|
||||
node=f"{JWT_CLASS}::test_admin_creates_reads_updates_clears_and_deletes_a_key[direct_jwt]",
|
||||
credential_kind="direct_jwt",
|
||||
actor="proxy_admin",
|
||||
profile="group_scoped",
|
||||
method="POST",
|
||||
path="/key/generate",
|
||||
operation_family="key_lifecycle",
|
||||
),
|
||||
ManagementCase(
|
||||
node=f"{JWT_CLASS}::test_admin_creates_reads_updates_clears_and_deletes_a_key[virtual_key]",
|
||||
credential_kind="virtual_key",
|
||||
actor="proxy_admin",
|
||||
profile="group_scoped",
|
||||
method="POST",
|
||||
path="/key/generate",
|
||||
operation_family="key_lifecycle",
|
||||
),
|
||||
ManagementCase(
|
||||
node=f"{JWT_CLASS}::test_two_actor_sets_keep_tenants_and_keys_isolated",
|
||||
credential_kind="direct_jwt",
|
||||
actor="team_member",
|
||||
profile="group_scoped",
|
||||
method="GET",
|
||||
path="/key/info",
|
||||
operation_family="tenant_isolation",
|
||||
),
|
||||
ManagementCase(
|
||||
node=f"{JWT_CLASS}::test_member_cannot_write_and_another_team_cannot_read_the_key",
|
||||
credential_kind="direct_jwt",
|
||||
actor="team_member",
|
||||
profile="group_scoped",
|
||||
method="POST",
|
||||
path="/key/update",
|
||||
operation_family="tenant_isolation",
|
||||
),
|
||||
ManagementCase(
|
||||
node=f"{JWT_CLASS}::test_multi_group_actor_keeps_exact_memberships",
|
||||
credential_kind="master",
|
||||
actor="bootstrap",
|
||||
profile="group_scoped",
|
||||
method="GET",
|
||||
path="/team/info",
|
||||
operation_family="memberships",
|
||||
),
|
||||
ManagementCase(
|
||||
node=f"{JWT_CLASS}::test_successful_actor_cleanup_removes_owned_state",
|
||||
credential_kind="master",
|
||||
actor="bootstrap",
|
||||
profile="failure_cleanup",
|
||||
method="GET",
|
||||
path="/team/info",
|
||||
operation_family="cleanup",
|
||||
),
|
||||
ManagementCase(
|
||||
node=f"{JWT_CLASS}::test_partial_setup_removes_previously_created_identities[group]",
|
||||
credential_kind="idp_admin",
|
||||
actor="idp_admin",
|
||||
profile="failure_cleanup",
|
||||
method="POST",
|
||||
path="/groups",
|
||||
operation_family="cleanup",
|
||||
),
|
||||
ManagementCase(
|
||||
node=f"{JWT_CLASS}::test_partial_setup_removes_previously_created_identities[user]",
|
||||
credential_kind="idp_admin",
|
||||
actor="idp_admin",
|
||||
profile="failure_cleanup",
|
||||
method="POST",
|
||||
path="/users",
|
||||
operation_family="cleanup",
|
||||
),
|
||||
ManagementCase(
|
||||
node=f"{JWT_CLASS}::test_oidc_browser_profile_identity_mapping",
|
||||
credential_kind="direct_jwt",
|
||||
actor="internal_user",
|
||||
profile="oidc_configuration",
|
||||
method="GET",
|
||||
path="/protocol/openid-connect/userinfo",
|
||||
operation_family="oidc_identity",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def canonical_node(node: str) -> str:
|
||||
return node if node.startswith("tests/e2e/") else f"tests/e2e/{node}"
|
||||
|
||||
|
||||
def case_properties(node: str) -> tuple[tuple[str, str], ...]:
|
||||
case: Final = next((case for case in MANAGEMENT_CASES if case.node == canonical_node(node)), None)
|
||||
if case is None:
|
||||
return ()
|
||||
return (
|
||||
("management_node", case.node),
|
||||
("credential_kind", case.credential_kind),
|
||||
("actor", case.actor),
|
||||
("auth_profile", case.profile),
|
||||
("dependency_profile", case.dependency_profile),
|
||||
)
|
||||
|
|
@ -90,3 +90,11 @@
|
|||
- {id: mgmt.mcp_toolset.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3098", rationale: "Narrowing the tools to one entry reads back exactly that entry"}
|
||||
- {id: mgmt.mcp_toolset.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:3098", fail_before_fix: proven, rationale: "An explicit null clears the stored description; the update used to drop null and keep the old value"}
|
||||
- {id: mgmt.mcp_toolset.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3149", rationale: "A deleted toolset is gone by id and from the list on every replica"}
|
||||
|
||||
- {id: mgmt.user.jwt.database_roles, module: mgmt, tier: P0, surface: api, assertions: [database_roles], source: "auth/handle_jwt.py", rationale: "User-only JWT subjects retain their seeded database roles and memberships"}
|
||||
- {id: mgmt.key.jwt.viewer_denied, module: mgmt, tier: P0, surface: api, assertions: [viewer_denied], source: "auth/route_checks.py", rationale: "An admin viewer can read a key but cannot update it or change stored state"}
|
||||
- {id: mgmt.user.oidc.identity_mapping, module: mgmt, tier: P0, surface: api, assertions: [identity_mapping], source: "tests/e2e/idp.py", rationale: "IdP configuration canary only: confidential-client token and userinfo subjects match the seeded user; application SSO is separate"}
|
||||
- {id: mgmt.team.jwt.tenant_isolation, module: mgmt, tier: P0, surface: api, assertions: [tenant_isolation], source: "auth/handle_jwt.py", rationale: "Isolated team actors read their own key and receive 403 for the other tenant key"}
|
||||
- {id: mgmt.team.jwt.multiple_memberships, module: mgmt, tier: P0, surface: api, assertions: [multiple_memberships], source: "auth/handle_jwt.py", rationale: "A multi-group actor has exactly the configured memberships without admin scope"}
|
||||
- {id: mgmt.user.jwt.cleanup, module: mgmt, tier: P0, surface: api, assertions: [cleanup], source: "management_endpoints/internal_user_endpoints.py", rationale: "Owned users teams organizations keys and IdP objects disappear after successful cleanup"}
|
||||
- {id: mgmt.user.jwt.partial_cleanup, module: mgmt, tier: P0, surface: api, assertions: [partial_cleanup], source: "auth/handle_jwt.py", rationale: "Partial identity setup removes the group and user created before failure"}
|
||||
|
|
|
|||
|
|
@ -16,9 +16,11 @@ requests itself imports.
|
|||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from collections.abc import Callable, Mapping
|
||||
from collections.abc import Callable, Generator, Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Generator, Generic, Iterator, Literal, NewType, Protocol, TypeVar, cast
|
||||
from typing import Final, Generic, Literal, NewType, Protocol, TypeVar, cast
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
|
@ -36,8 +38,8 @@ class Headers(BaseModel):
|
|||
|
||||
class AuthHeaders(Headers):
|
||||
# litellm accepts either; set whichever the call needs, leave the other None.
|
||||
authorization: str | None = None
|
||||
x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key")
|
||||
authorization: str | None = Field(default=None, repr=False)
|
||||
x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key", repr=False)
|
||||
|
||||
|
||||
class AnthropicHeaders(AuthHeaders):
|
||||
|
|
@ -168,6 +170,7 @@ class StreamingResponse(BaseModel):
|
|||
# the consumed body is elided, so this is the only place they surface.
|
||||
stream_error: str | None = None
|
||||
stream_done: bool = False
|
||||
stream_done_positions: tuple[int, ...] = ()
|
||||
|
||||
@property
|
||||
def ok(self) -> bool:
|
||||
|
|
@ -292,6 +295,22 @@ def _params(params: BaseModel | None) -> dict[str, str]:
|
|||
|
||||
TRANSIENT_STATUSES: frozenset[int] = frozenset({529})
|
||||
RETRY_ATTEMPTS: int = 3
|
||||
_QUALIFICATION: Final[ContextVar[bool]] = ContextVar("e2e_qualification", default=False)
|
||||
|
||||
|
||||
def retry_attempts(default: int) -> int:
|
||||
return 1 if _QUALIFICATION.get() else default
|
||||
|
||||
|
||||
@contextmanager
|
||||
def without_retries() -> Generator[None]:
|
||||
token: Final = _QUALIFICATION.set(True)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_QUALIFICATION.reset(token)
|
||||
|
||||
|
||||
RETRY_BACKOFF_SECONDS: float = 0.5
|
||||
|
||||
|
||||
|
|
@ -319,7 +338,7 @@ def request_with_retry[T: RetryableResponse](
|
|||
hang should surface as a hang instead of doubling the wall clock. Every
|
||||
retry prints, so flakiness stays visible in the run log instead of
|
||||
vanishing into green."""
|
||||
for attempt in range(1, RETRY_ATTEMPTS):
|
||||
for attempt in range(1, retry_attempts(RETRY_ATTEMPTS)):
|
||||
resp = issue()
|
||||
if resp.status_code not in TRANSIENT_STATUSES:
|
||||
return resp
|
||||
|
|
@ -414,6 +433,7 @@ def get_external[R: BaseModel](
|
|||
url: str,
|
||||
*,
|
||||
response_type: type[R],
|
||||
headers: BaseModel | None = None,
|
||||
timeout: float = 30.0,
|
||||
) -> Result[R]:
|
||||
"""GET an absolute URL outside the proxy (e.g. a public /.well-known document).
|
||||
|
|
@ -422,7 +442,7 @@ def get_external[R: BaseModel](
|
|||
try:
|
||||
resp = requests.get(
|
||||
url,
|
||||
headers={"Accept": "application/json"},
|
||||
headers={"Accept": "application/json", **(_headers(headers) if headers is not None else {})},
|
||||
timeout=timeout,
|
||||
)
|
||||
except requests.RequestException as exc:
|
||||
|
|
@ -628,6 +648,7 @@ def streaming_outcome(
|
|||
stream_events=[payload for payload, _ in events],
|
||||
stream_event_arrivals=[arrived for _, arrived in events],
|
||||
stream_done=any(payload == _SSE_DONE for payload, _ in payloads),
|
||||
stream_done_positions=tuple(index for index, (payload, _) in enumerate(payloads) if payload == _SSE_DONE),
|
||||
stream_error=next(
|
||||
(line.decode(errors="replace")[:300] for line, _ in stamped if _is_stream_error_line(line)),
|
||||
None,
|
||||
|
|
|
|||
|
|
@ -30,9 +30,11 @@ from datetime import datetime, timedelta, timezone
|
|||
from pathlib import Path
|
||||
from typing import Annotated, Final, Literal
|
||||
|
||||
from fixture_profile import MatchProfile, StrictIdentity
|
||||
from pydantic import BaseModel, Field, JsonValue
|
||||
|
||||
BUNDLE_FORMAT_VERSION: Final = 4
|
||||
STRICT_BUNDLE_FORMAT_VERSION: Final = 5
|
||||
MAX_BUNDLE_AGE: Final = timedelta(days=7)
|
||||
MANIFEST_FILENAME: Final = "manifest.json"
|
||||
|
||||
|
|
@ -41,6 +43,7 @@ class Manifest(BaseModel):
|
|||
format_version: int
|
||||
recorded_at: datetime
|
||||
harness_version: str
|
||||
match_profile: MatchProfile = "legacy"
|
||||
|
||||
|
||||
class RecordedRequest(BaseModel):
|
||||
|
|
@ -69,6 +72,7 @@ class RecordedRequest(BaseModel):
|
|||
file_name: str | None = None
|
||||
file_sha256: str | None = None
|
||||
file_bytes: int | None = None
|
||||
strict_identity: StrictIdentity | None = None
|
||||
|
||||
|
||||
class RecordedHttpResponse(BaseModel):
|
||||
|
|
@ -100,9 +104,7 @@ class RecordedStreamedResponse(BaseModel):
|
|||
truncated: str | None = None
|
||||
|
||||
|
||||
type RecordedResponse = Annotated[
|
||||
RecordedHttpResponse | RecordedStreamedResponse, Field(discriminator="kind")
|
||||
]
|
||||
type RecordedResponse = Annotated[RecordedHttpResponse | RecordedStreamedResponse, Field(discriminator="kind")]
|
||||
|
||||
|
||||
class Interaction(BaseModel):
|
||||
|
|
@ -152,6 +154,7 @@ class BundleRecorder:
|
|||
manifest, so record mode never reads (or merges into) an existing bundle."""
|
||||
|
||||
root: Path
|
||||
profile: MatchProfile = "legacy"
|
||||
_ordinals: dict[str, int] = field(default_factory=dict)
|
||||
|
||||
def record(self, *, test_key: str, request: RecordedRequest, response: RecordedResponse) -> None:
|
||||
|
|
@ -162,7 +165,12 @@ class BundleRecorder:
|
|||
directory.mkdir(parents=True, exist_ok=True)
|
||||
interaction = Interaction(request=request, response=response)
|
||||
target = directory / interaction_filename(ordinal, request)
|
||||
target.write_text(interaction.model_dump_json(indent=2), encoding="utf-8")
|
||||
target.write_text(
|
||||
interaction.model_dump_json(
|
||||
indent=2, exclude={"request": {"strict_identity"}} if self.profile == "legacy" else None
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -171,7 +179,7 @@ class UnsafeBundleDir:
|
|||
reason: str
|
||||
|
||||
|
||||
def prepare_bundle(root: Path) -> BundleRecorder | UnsafeBundleDir:
|
||||
def prepare_bundle(root: Path, *, profile: MatchProfile = "legacy") -> BundleRecorder | UnsafeBundleDir:
|
||||
"""Start a fresh bundle at ``root`` for record mode: wipe whatever bundle is
|
||||
there and write a new manifest. Refuses to wipe a directory that is neither
|
||||
empty nor a bundle (no manifest.json), so a mistyped E2E_FIXTURE_DIR can
|
||||
|
|
@ -188,12 +196,15 @@ def prepare_bundle(root: Path) -> BundleRecorder | UnsafeBundleDir:
|
|||
shutil.rmtree(root)
|
||||
root.mkdir(parents=True)
|
||||
manifest = Manifest(
|
||||
format_version=BUNDLE_FORMAT_VERSION,
|
||||
format_version=BUNDLE_FORMAT_VERSION if profile == "legacy" else STRICT_BUNDLE_FORMAT_VERSION,
|
||||
match_profile=profile,
|
||||
recorded_at=datetime.now(timezone.utc),
|
||||
harness_version=harness_version(),
|
||||
)
|
||||
(root / MANIFEST_FILENAME).write_text(manifest.model_dump_json(indent=2), encoding="utf-8")
|
||||
return BundleRecorder(root=root)
|
||||
(root / MANIFEST_FILENAME).write_text(
|
||||
manifest.model_dump_json(indent=2, exclude={"match_profile"} if profile == "legacy" else None), encoding="utf-8"
|
||||
)
|
||||
return BundleRecorder(root=root, profile=profile)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -226,25 +237,30 @@ def _read_manifest(root: Path) -> Manifest | UnreadableBundle:
|
|||
return UnreadableBundle(reason=f"{MANIFEST_FILENAME} is invalid: {exc}")
|
||||
|
||||
|
||||
def _supported_manifest(root: Path) -> Manifest | UnreadableBundle:
|
||||
def _supported_manifest(root: Path, profile: MatchProfile = "legacy") -> Manifest | UnreadableBundle:
|
||||
"""The manifest, refused when it was written under a different format version.
|
||||
A bundle is atomic (record wipes and rewrites the whole directory and never
|
||||
merges), so a foreign version is a hard reject rather than a partial read."""
|
||||
manifest = _read_manifest(root)
|
||||
if isinstance(manifest, UnreadableBundle):
|
||||
return manifest
|
||||
if manifest.format_version != BUNDLE_FORMAT_VERSION:
|
||||
expected_version: Final = BUNDLE_FORMAT_VERSION if profile == "legacy" else STRICT_BUNDLE_FORMAT_VERSION
|
||||
if manifest.match_profile != profile:
|
||||
return UnreadableBundle(
|
||||
reason="match profile mismatch; select the recorded E2E_REPLAY_MATCH_PROFILE or re-record"
|
||||
)
|
||||
if manifest.format_version != expected_version:
|
||||
return UnreadableBundle(
|
||||
reason=(
|
||||
f"format_version {manifest.format_version} != supported {BUNDLE_FORMAT_VERSION}; "
|
||||
f"format_version {manifest.format_version} != supported {expected_version}; "
|
||||
"re-record with E2E_FIXTURE_MODE=record"
|
||||
)
|
||||
)
|
||||
return manifest
|
||||
|
||||
|
||||
def check_freshness(root: Path, *, now: datetime) -> BundleFreshness:
|
||||
manifest = _supported_manifest(root)
|
||||
def check_freshness(root: Path, *, now: datetime, profile: MatchProfile = "legacy") -> BundleFreshness:
|
||||
manifest = _supported_manifest(root, profile)
|
||||
if isinstance(manifest, UnreadableBundle):
|
||||
return manifest
|
||||
recorded_at = (
|
||||
|
|
@ -269,16 +285,27 @@ class LoadedBundle:
|
|||
interactions: dict[str, tuple[Interaction, ...]]
|
||||
|
||||
|
||||
def load_bundle(root: Path) -> LoadedBundle | UnreadableBundle:
|
||||
manifest = _supported_manifest(root)
|
||||
def load_bundle(root: Path, *, profile: MatchProfile = "legacy") -> LoadedBundle | UnreadableBundle:
|
||||
manifest = _supported_manifest(root, profile)
|
||||
if isinstance(manifest, UnreadableBundle):
|
||||
return manifest
|
||||
interactions = {
|
||||
directory.name: tuple(
|
||||
Interaction.model_validate_json(file.read_text(encoding="utf-8"))
|
||||
for file in sorted(directory.glob("*.json"))
|
||||
)
|
||||
for directory in sorted(root.iterdir())
|
||||
if directory.is_dir()
|
||||
}
|
||||
try:
|
||||
interactions = {
|
||||
directory.name: tuple(
|
||||
Interaction.model_validate_json(file.read_text(encoding="utf-8"))
|
||||
for file in sorted(directory.glob("*.json"))
|
||||
)
|
||||
for directory in sorted(root.iterdir())
|
||||
if directory.is_dir()
|
||||
}
|
||||
except (ValueError, OSError):
|
||||
if profile == "legacy":
|
||||
raise
|
||||
return UnreadableBundle(reason="invalid stateless_v1 interaction; re-record with the selected profile")
|
||||
if any(
|
||||
(item.request.strict_identity is not None) != (profile == "stateless_v1")
|
||||
for items in interactions.values()
|
||||
for item in items
|
||||
):
|
||||
return UnreadableBundle(reason="request identity/profile mismatch; re-record with the selected profile")
|
||||
return LoadedBundle(manifest=manifest, interactions=interactions)
|
||||
|
|
|
|||
|
|
@ -23,9 +23,8 @@ from dataclasses import dataclass
|
|||
from functools import reduce
|
||||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue
|
||||
|
||||
from fixture_bundle import RecordedRequest
|
||||
from pydantic import JsonValue
|
||||
|
||||
VOLATILE_HEADER_NAMES: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
|
|
@ -123,6 +122,12 @@ class CanonicalRequest:
|
|||
|
||||
|
||||
def canonicalize(request: RecordedRequest) -> CanonicalRequest:
|
||||
if request.strict_identity is not None:
|
||||
return CanonicalRequest(
|
||||
method=request.method,
|
||||
path=request.path,
|
||||
content=json.dumps(request.strict_identity.model_dump(mode="json"), sort_keys=True, separators=(",", ":")),
|
||||
)
|
||||
file_identity: Final[JsonValue | None] = (
|
||||
None
|
||||
if request.file_name is None and request.file_sha256 is None
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ from fixture_bundle import (
|
|||
check_freshness,
|
||||
format_age,
|
||||
)
|
||||
from fixture_profile import match_profile
|
||||
|
||||
type FixtureMode = Literal["live", "record", "replay"]
|
||||
|
||||
|
|
@ -82,6 +83,7 @@ def fixture_mode_collection_error(mode_raw: str, bundle_dir: Path, *, now: datet
|
|||
Called at collection time (conftest pytest_sessionstart) so a stale or missing
|
||||
bundle fails the whole run up front, naming the bundle age, instead of failing
|
||||
every test individually."""
|
||||
match_profile()
|
||||
mode = parse_fixture_mode(mode_raw)
|
||||
match mode:
|
||||
case InvalidFixtureMode(value=value):
|
||||
|
|
@ -89,7 +91,7 @@ def fixture_mode_collection_error(mode_raw: str, bundle_dir: Path, *, now: datet
|
|||
case "live" | "record":
|
||||
return None
|
||||
case "replay":
|
||||
freshness = check_freshness(bundle_dir, now=now)
|
||||
freshness = check_freshness(bundle_dir, now=now, profile=match_profile())
|
||||
match freshness:
|
||||
case FreshBundle():
|
||||
return None
|
||||
|
|
@ -110,6 +112,7 @@ def fixture_mode_collection_error(mode_raw: str, bundle_dir: Path, *, now: datet
|
|||
def fixture_report_lines(mode_raw: str, bundle_dir: Path, *, now: datetime) -> list[str]:
|
||||
"""pytest report-header lines; empty in live mode so an unset
|
||||
E2E_FIXTURE_MODE keeps today's output byte-identical."""
|
||||
match_profile()
|
||||
mode = parse_fixture_mode(mode_raw)
|
||||
match mode:
|
||||
case InvalidFixtureMode() | "live":
|
||||
|
|
@ -117,7 +120,7 @@ def fixture_report_lines(mode_raw: str, bundle_dir: Path, *, now: datetime) -> l
|
|||
case "record":
|
||||
return [f"e2e fixture mode: record -> {bundle_dir}"]
|
||||
case "replay":
|
||||
freshness = check_freshness(bundle_dir, now=now)
|
||||
freshness = check_freshness(bundle_dir, now=now, profile=match_profile())
|
||||
match freshness:
|
||||
case FreshBundle(manifest=manifest):
|
||||
return [
|
||||
|
|
|
|||
179
tests/e2e/fixture_profile.py
Normal file
179
tests/e2e/fixture_profile.py
Normal file
|
|
@ -0,0 +1,179 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Literal
|
||||
from urllib.parse import parse_qsl, urlsplit
|
||||
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
|
||||
type MatchProfile = Literal["legacy", "stateless_v1"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NumberToken:
|
||||
literal: str
|
||||
|
||||
|
||||
type ExactJson = dict[str, ExactJson] | list[ExactJson] | str | bool | NumberToken | None
|
||||
|
||||
SEMANTIC_HEADERS: Final = frozenset({"content-type", "accept", "anthropic-version", "anthropic-beta", "openai-beta"})
|
||||
AUTH_HEADERS: Final = frozenset({"authorization", "x-api-key"})
|
||||
EXCLUDED_HEADERS: Final = frozenset(
|
||||
{
|
||||
"host",
|
||||
"content-length",
|
||||
"transfer-encoding",
|
||||
"connection",
|
||||
"accept-encoding",
|
||||
"user-agent",
|
||||
"traceparent",
|
||||
"tracestate",
|
||||
"x-request-id",
|
||||
"x-client-request-id",
|
||||
"cookie",
|
||||
}
|
||||
)
|
||||
CREDENTIAL_QUERY: Final = frozenset(
|
||||
{
|
||||
"api_key",
|
||||
"api-key",
|
||||
"apikey",
|
||||
"key",
|
||||
"token",
|
||||
"access_token",
|
||||
"signature",
|
||||
"password",
|
||||
"secret",
|
||||
"credentials",
|
||||
"authorization",
|
||||
"sig",
|
||||
"client_secret",
|
||||
"aws_access_key_id",
|
||||
"aws_secret_access_key",
|
||||
"aws_session_token",
|
||||
}
|
||||
)
|
||||
JSON_VALUE: Final[TypeAdapter[ExactJson]] = TypeAdapter(ExactJson)
|
||||
|
||||
|
||||
def match_profile() -> MatchProfile:
|
||||
raw: Final = os.environ.get("E2E_REPLAY_MATCH_PROFILE", "legacy")
|
||||
if raw in ("legacy", "stateless_v1"):
|
||||
return raw
|
||||
raise ValueError("E2E_REPLAY_MATCH_PROFILE must be legacy or stateless_v1")
|
||||
|
||||
|
||||
class StrictIdentity(BaseModel):
|
||||
upstream: str
|
||||
mount: str
|
||||
query: tuple[tuple[str, str], ...]
|
||||
headers: dict[str, str]
|
||||
auth: dict[str, str]
|
||||
body_present: bool
|
||||
body: JsonValue
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class IneligibleRequest:
|
||||
reason: str
|
||||
|
||||
|
||||
def _unique_object(pairs: list[tuple[str, ExactJson]]) -> dict[str, ExactJson]:
|
||||
if len({key for key, _ in pairs}) != len(pairs):
|
||||
raise ValueError("duplicate JSON object keys")
|
||||
return dict(pairs)
|
||||
|
||||
|
||||
def _invalid_constant(value: str) -> ExactJson:
|
||||
raise ValueError("nonfinite JSON number")
|
||||
|
||||
|
||||
def _exact_value(value: ExactJson) -> JsonValue:
|
||||
match value:
|
||||
case dict():
|
||||
return {"object": {key: _exact_value(item) for key, item in value.items()}}
|
||||
case list():
|
||||
return {"array": [_exact_value(item) for item in value]}
|
||||
case bool():
|
||||
return {"boolean": value}
|
||||
case NumberToken(literal=literal):
|
||||
return {"number": literal}
|
||||
case str():
|
||||
return {"string": value}
|
||||
case None:
|
||||
return None
|
||||
|
||||
|
||||
def strict_identity(
|
||||
*,
|
||||
method: str,
|
||||
path: str,
|
||||
query: str,
|
||||
headers: Mapping[str, str],
|
||||
body: bytes | None,
|
||||
mount: str,
|
||||
upstream_base: str,
|
||||
) -> StrictIdentity | IneligibleRequest:
|
||||
if (mount, path, method.upper()) not in {
|
||||
("openai", "/openai/v1/chat/completions", "POST"),
|
||||
("anthropic", "/anthropic/v1/messages", "POST"),
|
||||
}:
|
||||
return IneligibleRequest("unsupported endpoint or method")
|
||||
lowered: Final = {key.lower(): value for key, value in headers.items()}
|
||||
if len(lowered) != len(headers):
|
||||
return IneligibleRequest("duplicate header names")
|
||||
if any(
|
||||
key not in SEMANTIC_HEADERS | AUTH_HEADERS | EXCLUDED_HEADERS and not key.startswith("x-stainless-")
|
||||
for key in lowered
|
||||
):
|
||||
return IneligibleRequest("unsupported semantic header")
|
||||
if "transfer-encoding" in lowered:
|
||||
return IneligibleRequest("unsupported request transfer-encoding; send a content-length framed JSON body")
|
||||
authorization: Final = lowered.get("authorization")
|
||||
if authorization is not None and authorization.partition(" ")[0].lower() not in {"bearer", "basic", "digest"}:
|
||||
return IneligibleRequest("unsupported authorization scheme")
|
||||
destination: Final = urlsplit(upstream_base)
|
||||
if destination.username or destination.password or destination.query or destination.fragment:
|
||||
return IneligibleRequest("upstream destination contains credentials, query or fragment")
|
||||
if destination.scheme not in ("http", "https") or not destination.netloc:
|
||||
return IneligibleRequest("unsupported upstream destination")
|
||||
if body and lowered.get("content-type", "").split(";", 1)[0].strip().lower() != "application/json":
|
||||
return IneligibleRequest("unsupported body content-type; stateless_v1 requires JSON")
|
||||
try:
|
||||
parsed: Final = (
|
||||
JSON_VALUE.validate_python(
|
||||
json.loads(
|
||||
body,
|
||||
object_pairs_hook=_unique_object,
|
||||
parse_constant=_invalid_constant,
|
||||
parse_float=NumberToken,
|
||||
parse_int=NumberToken,
|
||||
)
|
||||
)
|
||||
if body
|
||||
else None
|
||||
)
|
||||
except (ValueError, UnicodeError):
|
||||
return IneligibleRequest("invalid JSON or duplicate JSON object keys")
|
||||
if body and not isinstance(parsed, dict):
|
||||
return IneligibleRequest("stateless inference requires a JSON object")
|
||||
try:
|
||||
query_pairs: Final = tuple(parse_qsl(query, keep_blank_values=True, errors="strict"))
|
||||
except UnicodeError:
|
||||
return IneligibleRequest("invalid UTF-8 query encoding")
|
||||
return StrictIdentity(
|
||||
upstream=upstream_base,
|
||||
mount=mount,
|
||||
query=tuple((key, "<credential>" if key.lower() in CREDENTIAL_QUERY else value) for key, value in query_pairs),
|
||||
headers={key: value for key, value in lowered.items() if key in SEMANTIC_HEADERS},
|
||||
auth={
|
||||
key: (value.partition(" ")[0].lower() if key == "authorization" else "present")
|
||||
for key, value in lowered.items()
|
||||
if key in AUTH_HEADERS
|
||||
},
|
||||
body_present=bool(body),
|
||||
body=_exact_value(parsed),
|
||||
)
|
||||
279
tests/e2e/idp.py
279
tests/e2e/idp.py
|
|
@ -2,11 +2,18 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import os
|
||||
import secrets
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass, field, replace
|
||||
from types import FrameType
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
|
|
@ -14,11 +21,15 @@ from e2e_http import (
|
|||
AuthHeaders,
|
||||
ExternalWrite,
|
||||
NetworkError,
|
||||
NoBody,
|
||||
Result,
|
||||
Success,
|
||||
UnknownApiError,
|
||||
delete_external,
|
||||
get_external,
|
||||
post_form_external,
|
||||
post_json_external,
|
||||
unwrap,
|
||||
)
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
|
@ -46,7 +57,9 @@ class TokenGrantForm(BaseModel):
|
|||
grant_type: Literal["password"] = "password"
|
||||
client_id: str
|
||||
username: str
|
||||
password: str
|
||||
password: str = Field(repr=False)
|
||||
client_secret: str | None = Field(default=None, repr=False)
|
||||
scope: str | None = None
|
||||
|
||||
|
||||
class TokenResponse(BaseModel):
|
||||
|
|
@ -63,7 +76,7 @@ class GroupCreateBody(BaseModel):
|
|||
|
||||
class PasswordCredential(BaseModel):
|
||||
type: Literal["password"] = "password"
|
||||
value: str
|
||||
value: str = Field(repr=False)
|
||||
temporary: bool = False
|
||||
|
||||
|
||||
|
|
@ -101,8 +114,20 @@ class Identity:
|
|||
user_id: str
|
||||
username: str
|
||||
password: str = field(repr=False)
|
||||
group: str
|
||||
group_id: str
|
||||
groups: tuple[str, ...]
|
||||
group_ids: tuple[str, ...]
|
||||
|
||||
@property
|
||||
def group(self) -> str:
|
||||
if len(self.groups) != 1:
|
||||
raise ValueError("A single-group identity is required")
|
||||
return self.groups[0]
|
||||
|
||||
@property
|
||||
def group_id(self) -> str:
|
||||
if len(self.group_ids) != 1:
|
||||
raise ValueError("A single-group identity is required")
|
||||
return self.group_ids[0]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -111,6 +136,10 @@ class Keycloak:
|
|||
realm: str
|
||||
admin_username: str
|
||||
admin_password: str = field(repr=False)
|
||||
strict_cleanup: bool = False
|
||||
|
||||
def with_strict_cleanup(self) -> Keycloak:
|
||||
return replace(self, strict_cleanup=True)
|
||||
|
||||
@property
|
||||
def issuer(self) -> str:
|
||||
|
|
@ -150,7 +179,9 @@ class Keycloak:
|
|||
f"group {name}",
|
||||
)
|
||||
|
||||
def create_user(self, *, username: str, email: str, password: str, group: str) -> str:
|
||||
def create_user(
|
||||
self, *, username: str, email: str, password: str, group: str | None = None, groups: tuple[str, ...] = ()
|
||||
) -> str:
|
||||
return created_id(
|
||||
post_json_external(
|
||||
self._admin_url("/users"),
|
||||
|
|
@ -158,7 +189,7 @@ class Keycloak:
|
|||
json=UserCreateBody(
|
||||
username=username,
|
||||
email=email,
|
||||
groups=(group,),
|
||||
groups=(group,) if group is not None else groups,
|
||||
credentials=(PasswordCredential(value=password),),
|
||||
),
|
||||
),
|
||||
|
|
@ -171,14 +202,28 @@ class Keycloak:
|
|||
def delete_group(self, group_id: str) -> None:
|
||||
self._delete(f"/groups/{group_id}")
|
||||
|
||||
def assert_absent(self, kind: Literal["users", "groups", "clients"], resource_id: str) -> None:
|
||||
result: Final = get_external(
|
||||
self._admin_url(f"/{kind}/{resource_id}"),
|
||||
headers=self._admin_headers(),
|
||||
response_type=NoBody,
|
||||
)
|
||||
assert isinstance(result, UnknownApiError) and result.status_code == 404, (
|
||||
f"Owned IdP {kind} still exists: {result}"
|
||||
)
|
||||
|
||||
def _delete(self, path: str) -> None:
|
||||
try:
|
||||
headers: Final = self._admin_headers()
|
||||
except pytest.fail.Exception as exc:
|
||||
if self.strict_cleanup:
|
||||
raise RuntimeError(f"Keycloak cleanup could not authenticate for {path}") from exc
|
||||
warnings.warn(f"Keycloak cleanup could not authenticate for {path}: {exc}", RuntimeWarning, stacklevel=2)
|
||||
return
|
||||
result: Final = delete_external(self._admin_url(path), headers=headers)
|
||||
if result.status_code not in (204, 404):
|
||||
if self.strict_cleanup:
|
||||
raise RuntimeError(f"Keycloak cleanup failed for {path}: HTTP {result.status_code}")
|
||||
warnings.warn(
|
||||
f"Keycloak cleanup failed for {path}: HTTP {result.status_code} {result.body[:300]}",
|
||||
RuntimeWarning,
|
||||
|
|
@ -188,15 +233,34 @@ class Keycloak:
|
|||
def provision(self, *, marker: str, group: str, defer: Callable[[Callable[[], object]], None]) -> Identity:
|
||||
"""Create `group` and a user in it, credentialed with a password generated
|
||||
for this test alone, and hand back the identity a token can be minted for."""
|
||||
group_id: Final = self.create_group(group)
|
||||
defer(lambda: self.delete_group(group_id))
|
||||
return self.provision_groups(marker=marker, groups=(group,), defer=defer)
|
||||
|
||||
def provision_groups(
|
||||
self, *, marker: str, groups: tuple[str, ...], defer: Callable[[Callable[[], object]], None]
|
||||
) -> Identity:
|
||||
def provision_group(name: str) -> str:
|
||||
created: Final = self.create_group(name)
|
||||
defer(lambda: self.delete_group(created))
|
||||
return created
|
||||
|
||||
group_ids: Final = tuple(provision_group(group) for group in groups)
|
||||
return self.provision_user(marker=marker, groups=groups, group_ids=group_ids, defer=defer)
|
||||
|
||||
def provision_user(
|
||||
self,
|
||||
*,
|
||||
marker: str,
|
||||
groups: tuple[str, ...],
|
||||
group_ids: tuple[str, ...],
|
||||
defer: Callable[[Callable[[], object]], None],
|
||||
) -> Identity:
|
||||
username: Final = f"e2e-jwt-user-{marker}"
|
||||
password: Final = secrets.token_urlsafe(24)
|
||||
user_id: Final = self.create_user(
|
||||
username=username, email=f"{username}@example.com", password=password, group=group
|
||||
username=username, email=f"{username}@example.com", password=password, groups=groups
|
||||
)
|
||||
defer(lambda: self.delete_user(user_id))
|
||||
return Identity(user_id=user_id, username=username, password=password, group=group, group_id=group_id)
|
||||
return Identity(user_id=user_id, username=username, password=password, groups=groups, group_ids=group_ids)
|
||||
|
||||
def access_token(
|
||||
self, identity: Identity, *, client_id: str = TESTS_CLIENT_ID, issuer_host: str | None = None
|
||||
|
|
@ -211,6 +275,65 @@ class Keycloak:
|
|||
)
|
||||
return self._token(result, f"a token for {identity.username}")
|
||||
|
||||
def discovery(self) -> Discovery:
|
||||
return unwrap(get_external(f"{self.issuer}/.well-known/openid-configuration", response_type=Discovery))
|
||||
|
||||
def browser_client(self, *, callback_url: str, defer: Callable[[Callable[[], object]], None]) -> BrowserClient:
|
||||
client: Final = BrowserClient(
|
||||
client_id=f"e2e-browser-{secrets.token_hex(8)}",
|
||||
secret=secrets.token_urlsafe(32),
|
||||
callback_url=callback_url,
|
||||
)
|
||||
resource_id: Final = created_id(
|
||||
post_json_external(
|
||||
self._admin_url("/clients"),
|
||||
headers=self._admin_headers(),
|
||||
json=BrowserClientBody(
|
||||
clientId=client.client_id,
|
||||
secret=client.secret,
|
||||
redirectUris=(callback_url,),
|
||||
),
|
||||
),
|
||||
"browser client",
|
||||
)
|
||||
defer(lambda: self._delete(f"/clients/{resource_id}"))
|
||||
configured: Final = unwrap(
|
||||
get_external(
|
||||
self._admin_url(f"/clients/{resource_id}"),
|
||||
headers=self._admin_headers(),
|
||||
response_type=BrowserClientBody,
|
||||
)
|
||||
)
|
||||
assert configured.redirect_uris == (callback_url,)
|
||||
assert configured.standard_flow_enabled and not configured.public_client
|
||||
assert configured.attributes.pkce == "S256"
|
||||
return client
|
||||
|
||||
def browser_token(self, identity: Identity, client: BrowserClient) -> str:
|
||||
return self._token(
|
||||
post_form_external(
|
||||
self.token_url(self.realm),
|
||||
form=TokenGrantForm(
|
||||
client_id=client.client_id,
|
||||
client_secret=client.secret,
|
||||
username=identity.username,
|
||||
password=identity.password,
|
||||
scope="openid email",
|
||||
),
|
||||
response_type=TokenResponse,
|
||||
),
|
||||
"browser-profile identity mapping",
|
||||
)
|
||||
|
||||
def userinfo(self, token: str) -> UserInfo:
|
||||
return unwrap(
|
||||
get_external(
|
||||
f"{self.issuer}/protocol/openid-connect/userinfo",
|
||||
headers=AuthHeaders(authorization=f"Bearer {token}"),
|
||||
response_type=UserInfo,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def keycloak_from_env() -> Keycloak:
|
||||
admin_username: Final = os.environ.get(KEYCLOAK_ADMIN_USER_ENV, "").strip()
|
||||
|
|
@ -226,3 +349,137 @@ def keycloak_from_env() -> Keycloak:
|
|||
admin_username=admin_username,
|
||||
admin_password=admin_password,
|
||||
)
|
||||
|
||||
|
||||
class TokenClaims(BaseModel):
|
||||
sub: str
|
||||
iss: str
|
||||
aud: str | tuple[str, ...]
|
||||
exp: int
|
||||
scope: str = ""
|
||||
groups: tuple[str, ...] = ()
|
||||
|
||||
|
||||
class Discovery(BaseModel):
|
||||
issuer: str
|
||||
authorization_endpoint: str
|
||||
token_endpoint: str
|
||||
userinfo_endpoint: str
|
||||
jwks_uri: str
|
||||
|
||||
|
||||
class UserInfo(BaseModel):
|
||||
sub: str
|
||||
email: str
|
||||
|
||||
|
||||
class BrowserAttributes(BaseModel):
|
||||
pkce: str = Field(default="S256", alias="pkce.code.challenge.method")
|
||||
|
||||
|
||||
class AudienceConfig(BaseModel):
|
||||
audience: str = Field(default="litellm-e2e", alias="included.custom.audience")
|
||||
access_token: str = Field(default="true", alias="access.token.claim")
|
||||
id_token: str = Field(default="false", alias="id.token.claim")
|
||||
|
||||
|
||||
class AudienceMapper(BaseModel):
|
||||
name: str = "litellm-audience"
|
||||
protocol: str = "openid-connect"
|
||||
mapper: str = Field(default="oidc-audience-mapper", alias="protocolMapper")
|
||||
config: AudienceConfig = Field(default_factory=AudienceConfig)
|
||||
|
||||
|
||||
class BrowserClientBody(BaseModel):
|
||||
client_id: str = Field(alias="clientId")
|
||||
secret: str = Field(repr=False)
|
||||
redirect_uris: tuple[str, ...] = Field(alias="redirectUris")
|
||||
enabled: bool = True
|
||||
public_client: bool = Field(default=False, alias="publicClient")
|
||||
standard_flow_enabled: bool = Field(default=True, alias="standardFlowEnabled")
|
||||
direct_access_grants_enabled: bool = Field(default=True, alias="directAccessGrantsEnabled")
|
||||
default_client_scopes: tuple[str, ...] = Field(default=("email", "basic"), alias="defaultClientScopes")
|
||||
attributes: BrowserAttributes = Field(default_factory=BrowserAttributes)
|
||||
protocol_mappers: tuple[AudienceMapper, ...] = Field(default=(AudienceMapper(),), alias="protocolMappers")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BrowserClient:
|
||||
client_id: str
|
||||
secret: str = field(repr=False)
|
||||
callback_url: str
|
||||
|
||||
def environment(self, discovery: Discovery) -> dict[str, str]:
|
||||
return {
|
||||
"GENERIC_CLIENT_ID": self.client_id,
|
||||
"GENERIC_CLIENT_SECRET": self.secret,
|
||||
"GENERIC_USER_ID_ATTRIBUTE": "sub",
|
||||
"GENERIC_AUTHORIZATION_ENDPOINT": discovery.authorization_endpoint,
|
||||
"GENERIC_TOKEN_ENDPOINT": discovery.token_endpoint,
|
||||
"GENERIC_USERINFO_ENDPOINT": discovery.userinfo_endpoint,
|
||||
"GENERIC_CLIENT_USE_PKCE": "true",
|
||||
"GENERIC_SCOPE": "openid email",
|
||||
}
|
||||
|
||||
|
||||
def token_claims(token: str) -> TokenClaims:
|
||||
payload: Final = token.split(".")[1]
|
||||
return TokenClaims.model_validate_json(base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4)))
|
||||
|
||||
|
||||
def _signal_process_group(process_id: int, signum: int) -> bool:
|
||||
try:
|
||||
os.killpg(process_id, signum)
|
||||
except ProcessLookupError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _stop_process_group(child: subprocess.Popen[bytes]) -> None:
|
||||
_signal_process_group(child.pid, signal.SIGTERM)
|
||||
deadline: Final = time.monotonic() + 5
|
||||
while _process_group_exists(child.pid):
|
||||
child.poll()
|
||||
if time.monotonic() >= deadline:
|
||||
_signal_process_group(child.pid, signal.SIGKILL)
|
||||
break
|
||||
time.sleep(0.05)
|
||||
child.wait()
|
||||
|
||||
|
||||
def _process_group_exists(process_id: int) -> bool:
|
||||
try:
|
||||
os.killpg(process_id, 0)
|
||||
except ProcessLookupError:
|
||||
return False
|
||||
except PermissionError:
|
||||
return True
|
||||
return True
|
||||
|
||||
|
||||
def run_oidc_profile(proxy_url: str, command: list[str]) -> int:
|
||||
idp: Final = keycloak_from_env().with_strict_cleanup()
|
||||
with ExitStack() as cleanup:
|
||||
|
||||
def terminate(signum: int, frame: FrameType | None) -> None:
|
||||
raise SystemExit(128 + signum)
|
||||
|
||||
previous: Final = signal.signal(signal.SIGTERM, terminate)
|
||||
cleanup.callback(signal.signal, signal.SIGTERM, previous)
|
||||
|
||||
def defer(callback: Callable[[], object]) -> None:
|
||||
cleanup.callback(callback)
|
||||
|
||||
client: Final = idp.browser_client(callback_url=f"{proxy_url.rstrip('/')}/sso/callback", defer=defer)
|
||||
environment: Final = {**os.environ, **client.environment(idp.discovery()), "PROXY_BASE_URL": proxy_url}
|
||||
with subprocess.Popen(command, env=environment, start_new_session=True) as child:
|
||||
try:
|
||||
return child.wait()
|
||||
finally:
|
||||
_stop_process_group(child)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) < 3:
|
||||
raise SystemExit("Usage: idp.py PROXY_URL COMMAND [ARG ...]; requires a running test IdP")
|
||||
raise SystemExit(run_oidc_profile(sys.argv[1], sys.argv[2:]))
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from __future__ import annotations
|
|||
from collections.abc import Iterable
|
||||
|
||||
import pytest
|
||||
from coverage_registry.management_cases import case_properties
|
||||
|
||||
# Hardcoded because the runner image copies tests/e2e/ to /app/e2e, so nothing
|
||||
# at runtime names this suite's place in the repo. test_junit_properties.py
|
||||
|
|
@ -94,7 +95,7 @@ def result_properties(item: pytest.Item) -> tuple[tuple[str, str], ...]:
|
|||
("package", package_from_nodeid(item.nodeid)),
|
||||
("covers", ",".join(covers_from_item(item))),
|
||||
("source", source_from_item(item)),
|
||||
)
|
||||
) + case_properties(item.nodeid)
|
||||
|
||||
|
||||
def attach_result_properties(item: pytest.Item) -> None:
|
||||
|
|
|
|||
|
|
@ -1,51 +1,90 @@
|
|||
"""Vendor §12.3: chat completions streaming SSE contract (LIT-4778).
|
||||
|
||||
Asserts a streamed /chat/completions response is SSE, carries content chunks,
|
||||
and terminates with the OpenAI [DONE] sentinel.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from e2e_config import unique_marker
|
||||
from e2e_config import provider_edge_base, unique_marker
|
||||
from e2e_http import require_successful_call
|
||||
from lifecycle import ResourceManager
|
||||
from models import ChatBody, ChatMessage, LiteLLMParamsBody
|
||||
from models import ChatBody, ChatMessage, ChatStreamOptions, LiteLLMParamsBody, Usage
|
||||
from proxy_client import ProxyClient
|
||||
from pydantic import BaseModel
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
pytestmark = [pytest.mark.e2e, pytest.mark.replayable]
|
||||
|
||||
|
||||
class _Delta(BaseModel):
|
||||
content: str | None = None
|
||||
|
||||
|
||||
class _Choice(BaseModel):
|
||||
index: int
|
||||
delta: _Delta
|
||||
finish_reason: str | None = None
|
||||
|
||||
|
||||
class _Chunk(BaseModel):
|
||||
choices: tuple[_Choice, ...]
|
||||
usage: Usage | None = None
|
||||
|
||||
|
||||
class TestChatStreamContract:
|
||||
@pytest.mark.covers("llm.chat_completions.openai.basic.stream.works")
|
||||
def test_chat_stream_is_sse_and_ends_with_done(self, proxy: ProxyClient, resources: ResourceManager) -> None:
|
||||
model = f"e2e-chat-stream-{unique_marker()}"
|
||||
model_id = proxy.create_model(
|
||||
model: Final = f"e2e-chat-stream-{unique_marker()}"
|
||||
base: Final = provider_edge_base("openai")
|
||||
model_id: Final = proxy.create_model(
|
||||
model,
|
||||
LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"),
|
||||
LiteLLMParamsBody(
|
||||
model="openai/gpt-5.6",
|
||||
api_key="os.environ/OPENAI_API_KEY",
|
||||
api_base=f"{base}/v1" if base else None,
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: proxy.delete_model(model_id))
|
||||
key = resources.key()
|
||||
|
||||
result = proxy.chat_stream(
|
||||
key: Final = resources.key()
|
||||
expected: Final = "The amber kite crosses the quiet lake."
|
||||
result: Final = proxy.chat_stream(
|
||||
key,
|
||||
ChatBody(
|
||||
model=model,
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="user",
|
||||
content=f"Reply with the single word ok. {unique_marker()}",
|
||||
role="user", content=f"Repeat exactly this sentence, with no additional text: {expected}"
|
||||
)
|
||||
],
|
||||
stream=True,
|
||||
max_completion_tokens=32,
|
||||
temperature=0.0,
|
||||
stream_options=ChatStreamOptions(include_usage=True),
|
||||
max_completion_tokens=256,
|
||||
reasoning_effort="none",
|
||||
),
|
||||
)
|
||||
require_successful_call(result)
|
||||
assert result.is_streaming, f"expected SSE content-type, got {result.content_type!r}"
|
||||
assert result.stream_events, "stream returned no data events"
|
||||
assert result.stream_done, (
|
||||
f"stream must terminate with [DONE]; "
|
||||
f"chunks={result.chunks} done={result.stream_done} events={len(result.stream_events)}"
|
||||
assert not result.stream_error, f"stream errored: {result.stream_error}"
|
||||
assert result.stream_done, "stream must terminate with [DONE]"
|
||||
assert result.stream_done_positions == (len(result.stream_events),), "[DONE] must occur once after all events"
|
||||
chunks: Final = tuple(_Chunk.model_validate_json(event) for event in result.stream_events)
|
||||
text_positions: Final = tuple(
|
||||
i for i, chunk in enumerate(chunks) if any(c.delta.content for c in chunk.choices)
|
||||
)
|
||||
terminal_positions: Final = tuple(
|
||||
i for i, chunk in enumerate(chunks) if any(c.finish_reason is not None for c in chunk.choices)
|
||||
)
|
||||
assert text_positions, "stream completed without meaningful text"
|
||||
assert len(terminal_positions) == 1, "expected exactly one terminal choice"
|
||||
assert text_positions[0] < terminal_positions[0], "meaningful text must arrive before termination"
|
||||
assert text_positions[-1] <= terminal_positions[0], "text arrived after termination"
|
||||
assert all(c.index == 0 for chunk in chunks for c in chunk.choices)
|
||||
assert tuple(c.finish_reason for c in chunks[terminal_positions[0]].choices) == ("stop",)
|
||||
text: Final = "".join(c.delta.content or "" for chunk in chunks for c in chunk.choices)
|
||||
assert text.strip() == expected, f"streamed answer was altered or incomplete: {text!r}"
|
||||
usage_positions: Final = tuple(i for i, chunk in enumerate(chunks) if chunk.usage is not None)
|
||||
assert usage_positions == (len(chunks) - 1,), "expected one final usage chunk"
|
||||
assert terminal_positions[0] < usage_positions[0], "usage must follow the terminal choice"
|
||||
usage: Final = chunks[-1].usage
|
||||
assert usage is not None
|
||||
assert usage.prompt_tokens is not None and usage.prompt_tokens > 0
|
||||
assert usage.completion_tokens is not None and usage.completion_tokens > 0
|
||||
assert usage.total_tokens == usage.prompt_tokens + usage.completion_tokens
|
||||
|
|
|
|||
|
|
@ -21,7 +21,12 @@ from e2e_http import assert_client_error, require_successful_call, unwrap
|
|||
from endpoints_client import EndpointsClient, MessagesResult
|
||||
from lifecycle import ResourceManager
|
||||
from models import (
|
||||
AnthropicAssistantTurn,
|
||||
AnthropicContentBlock,
|
||||
AnthropicCustomTool,
|
||||
AnthropicToolChoice,
|
||||
AnthropicToolResultBlock,
|
||||
AnthropicToolResultTurn,
|
||||
AnthropicMessagesBody,
|
||||
ChatMessage,
|
||||
JsonSchemaProperty,
|
||||
|
|
@ -29,7 +34,7 @@ from models import (
|
|||
SpendLogRow,
|
||||
ToolInputSchema,
|
||||
)
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
pytestmark = [pytest.mark.e2e, pytest.mark.replayable]
|
||||
|
||||
|
|
@ -284,8 +289,139 @@ class TestAnthropicMessages:
|
|||
result = endpoints_client.proxy.transport.send(
|
||||
"/v1/messages",
|
||||
headers=endpoints_client.proxy.transport.bearer(key),
|
||||
json=_OptionalMessagesBody(
|
||||
messages=[ChatMessage(role="user", content="hi")], max_tokens=50
|
||||
),
|
||||
json=_OptionalMessagesBody(messages=[ChatMessage(role="user", content="hi")], max_tokens=50),
|
||||
)
|
||||
assert_client_error(result, "messages missing model")
|
||||
|
||||
|
||||
class _BridgeDelta(BaseModel):
|
||||
type: str | None = None
|
||||
partial_json: str | None = None
|
||||
stop_reason: str | None = None
|
||||
|
||||
|
||||
class _BridgeEvent(BaseModel):
|
||||
type: str
|
||||
index: int | None = None
|
||||
content_block: AnthropicContentBlock | None = None
|
||||
delta: _BridgeDelta | None = None
|
||||
|
||||
|
||||
class _ParcelInput(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", strict=True)
|
||||
parcel: str
|
||||
shelf: int
|
||||
|
||||
|
||||
def _tool_from_stream(events: tuple[_BridgeEvent, ...]) -> AnthropicContentBlock:
|
||||
starts: Final = tuple(
|
||||
event
|
||||
for event in events
|
||||
if event.type == "content_block_start"
|
||||
and event.content_block is not None
|
||||
and event.content_block.type == "tool_use"
|
||||
)
|
||||
assert len(starts) == 1, "expected exactly one tool call"
|
||||
start: Final = starts[0]
|
||||
block: Final = start.content_block
|
||||
assert block is not None and block.id and start.index is not None
|
||||
fragments: Final = tuple(
|
||||
event
|
||||
for event in events
|
||||
if event.type == "content_block_delta" and event.delta is not None and event.delta.type == "input_json_delta"
|
||||
)
|
||||
assert fragments, "tool stream contained no argument fragments"
|
||||
assert all(event.index == start.index for event in fragments), "tool fragments changed index"
|
||||
positions: Final = tuple(i for i, event in enumerate(events) if event in fragments)
|
||||
stops: Final = tuple(
|
||||
i for i, event in enumerate(events) if event.type == "content_block_stop" and event.index == start.index
|
||||
)
|
||||
assert len(stops) == 1 and events.index(start) < positions[0] <= positions[-1] < stops[0]
|
||||
assert tuple(
|
||||
event.delta.stop_reason for event in events if event.type == "message_delta" and event.delta is not None
|
||||
) == ("tool_use",)
|
||||
terminal_positions: Final = tuple(i for i, event in enumerate(events) if event.type == "message_delta")
|
||||
assert len(terminal_positions) == 1 and stops[0] < terminal_positions[0] < len(events) - 1
|
||||
assert tuple(i for i, event in enumerate(events) if event.type == "message_stop") == (len(events) - 1,), (
|
||||
"tool stream did not terminate exactly once"
|
||||
)
|
||||
arguments: Final = _ParcelInput.model_validate_json(
|
||||
"".join(event.delta.partial_json or "" for event in fragments if event.delta is not None)
|
||||
)
|
||||
return AnthropicContentBlock(type="tool_use", id=block.id, name=block.name, input=arguments.model_dump())
|
||||
|
||||
|
||||
def _parcel_result(tool: AnthropicContentBlock, result: AnthropicToolResultBlock) -> AnthropicToolResultTurn:
|
||||
assert tool.id and result.tool_use_id == tool.id, "tool result ID does not match the emitted call"
|
||||
return AnthropicToolResultTurn(content=[result])
|
||||
|
||||
|
||||
def _request_tool(
|
||||
client: EndpointsClient, key: str, request: AnthropicMessagesBody, stream: bool
|
||||
) -> AnthropicContentBlock:
|
||||
if stream:
|
||||
response: Final = client.proxy.messages_stream(key, request)
|
||||
require_successful_call(response)
|
||||
assert response.is_streaming and not response.stream_error
|
||||
return _tool_from_stream(tuple(_BridgeEvent.model_validate_json(event) for event in response.stream_events))
|
||||
response_body: Final = unwrap(client.proxy.messages(key, request))
|
||||
blocks: Final = tuple(block for block in response_body.content or () if block.type == "tool_use")
|
||||
assert len(blocks) == 1
|
||||
return blocks[0]
|
||||
|
||||
|
||||
class TestOpenAIMessagesToolContinuation:
|
||||
@pytest.mark.parametrize("stream", [True, False], ids=["stream", "nonstream"])
|
||||
def test_required_tool_arguments_and_correlated_result(
|
||||
self, endpoints_client: EndpointsClient, resources: ResourceManager, stream: bool
|
||||
) -> None:
|
||||
model: Final = f"e2e-bridge-tool-{unique_marker()}"
|
||||
base: Final = provider_edge_base("openai")
|
||||
model_id: Final = endpoints_client.create_model(
|
||||
model,
|
||||
LiteLLMParamsBody(
|
||||
model="openai/gpt-5.6", api_key="os.environ/OPENAI_API_KEY", api_base=f"{base}/v1" if base else None
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: endpoints_client.delete_model(model_id))
|
||||
key: Final = resources.key(models=[model])
|
||||
tool: Final = AnthropicCustomTool(
|
||||
name="locate_parcel",
|
||||
description="Look up the receipt for a parcel on a shelf. Return the receipt verbatim.",
|
||||
input_schema=ToolInputSchema(
|
||||
properties={"parcel": JsonSchemaProperty(type="string"), "shelf": JsonSchemaProperty(type="integer")},
|
||||
required=["parcel", "shelf"],
|
||||
),
|
||||
)
|
||||
question: Final = ChatMessage(
|
||||
role="user",
|
||||
content="Call locate_parcel with parcel exactly amber-kite and shelf exactly 7. After the tool result, reply with only the receipt returned by the tool.",
|
||||
)
|
||||
request: Final = AnthropicMessagesBody(
|
||||
model=model,
|
||||
max_tokens=2048,
|
||||
messages=[question],
|
||||
tools=[tool],
|
||||
tool_choice=AnthropicToolChoice(type="tool", name=tool.name),
|
||||
stream=stream,
|
||||
)
|
||||
emitted: Final = _request_tool(endpoints_client, key, request, stream)
|
||||
assert emitted.id and emitted.name == "locate_parcel"
|
||||
assert emitted.input == {"parcel": "amber-kite", "shelf": 7}, "required tool arguments were lost or changed"
|
||||
receipt: Final = f"receipt-{unique_marker()}"
|
||||
result_turn: Final = _parcel_result(emitted, AnthropicToolResultBlock(tool_use_id=emitted.id, content=receipt))
|
||||
continuation: Final = unwrap(
|
||||
endpoints_client.proxy.messages(
|
||||
key,
|
||||
AnthropicMessagesBody(
|
||||
model=model,
|
||||
max_tokens=2048,
|
||||
tools=[tool],
|
||||
tool_choice=AnthropicToolChoice(type="none"),
|
||||
messages=[question, AnthropicAssistantTurn(content=[emitted]), result_turn],
|
||||
),
|
||||
)
|
||||
)
|
||||
answer: Final = "".join(block.text or "" for block in continuation.content or ())
|
||||
assert answer.strip() == receipt, "continuation did not consume the correlated tool result"
|
||||
assert all(block.type != "tool_use" for block in continuation.content or ())
|
||||
|
|
|
|||
|
|
@ -5,8 +5,14 @@ holds the shared ProxyClient so `resources` / `scoped_key` clean up keys, teams,
|
|||
users, and orgs this suite creates.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from collections.abc import Generator
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from e2e_http import without_retries
|
||||
from idp import Keycloak
|
||||
from lifecycle import ResourceManager
|
||||
from management.jwt_actors import ActorFactory
|
||||
from management_client import ManagementClient, build_client
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
|
@ -21,3 +27,14 @@ def pytest_configure(config: pytest.Config) -> None:
|
|||
@pytest.fixture(scope="session")
|
||||
def client(proxy: ProxyClient) -> ManagementClient:
|
||||
return build_client(proxy)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def actor_factory(proxy: ProxyClient, idp: Keycloak) -> Generator[ActorFactory]:
|
||||
bootstrap: Final = build_client(proxy)
|
||||
resources: Final = ResourceManager(client=proxy, strict_cleanup=True)
|
||||
with without_retries():
|
||||
try:
|
||||
yield ActorFactory(bootstrap=bootstrap, idp=idp, resources=resources)
|
||||
finally:
|
||||
resources.teardown()
|
||||
|
|
|
|||
175
tests/e2e/management/jwt_actors.py
Normal file
175
tests/e2e/management/jwt_actors.py
Normal file
|
|
@ -0,0 +1,175 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Literal
|
||||
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import NoBody, unwrap
|
||||
from idp import ADMIN_CLIENT_ID, TESTS_CLIENT_ID, Identity, Keycloak
|
||||
from lifecycle import ResourceManager
|
||||
from management.management_client import ManagementClient
|
||||
from models import (
|
||||
KeyGenerateBody,
|
||||
KeyGenerateResponse,
|
||||
OrgDeleteBody,
|
||||
OrgDeleteResponse,
|
||||
OrgMemberAddBody,
|
||||
OrgMemberEntry,
|
||||
OrgNewBody,
|
||||
TeamDeleteBody,
|
||||
TeamMemberAddBody,
|
||||
TeamMemberEntry,
|
||||
TeamNewBody,
|
||||
UserNewBody,
|
||||
UserRole,
|
||||
)
|
||||
from proxy_client import Caller
|
||||
|
||||
ActorRole = Literal[
|
||||
"proxy_admin",
|
||||
"proxy_admin_viewer",
|
||||
"organization_admin",
|
||||
"team_admin",
|
||||
"team_member",
|
||||
"internal_user",
|
||||
"internal_user_viewer",
|
||||
"unrelated_user",
|
||||
]
|
||||
ActorProfile = Literal["database_role", "group_scoped"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Tenant:
|
||||
organization_id: str
|
||||
team_id: str
|
||||
group_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Actor:
|
||||
identity: Identity
|
||||
role: ActorRole
|
||||
global_role: UserRole
|
||||
profile: ActorProfile
|
||||
tenants: tuple[Tenant, ...]
|
||||
|
||||
def mint_caller(self, idp: Keycloak) -> Caller:
|
||||
return Caller(
|
||||
credential=idp.access_token(
|
||||
self.identity, client_id=ADMIN_CLIENT_ID if self.role == "proxy_admin" else TESTS_CLIENT_ID
|
||||
),
|
||||
kind="direct_jwt",
|
||||
role=self.role,
|
||||
tenant=self.tenants[0].organization_id if self.tenants else None,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ActorFactory:
|
||||
bootstrap: ManagementClient
|
||||
idp: Keycloak
|
||||
resources: ResourceManager
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.bootstrap.proxy.caller is not None:
|
||||
raise ValueError("Actor bootstrap requires a separately held master client")
|
||||
|
||||
def key(self, tenant: Tenant | None = None, *, user_id: str | None = None) -> KeyGenerateResponse:
|
||||
created: Final = unwrap(
|
||||
self.bootstrap.generate_key(
|
||||
KeyGenerateBody(
|
||||
team_id=tenant.team_id if tenant is not None else None,
|
||||
user_id=user_id,
|
||||
key_alias=f"e2e-actor-key-{unique_marker()}",
|
||||
)
|
||||
)
|
||||
)
|
||||
self.resources.defer(lambda: self.bootstrap.delete_key_strict(created.key, missing_ok=True))
|
||||
return created
|
||||
|
||||
def tenant(self) -> Tenant:
|
||||
marker: Final = unique_marker()
|
||||
organization_id: Final = self.bootstrap.create_org(OrgNewBody(organization_alias=f"e2e-organization-{marker}"))
|
||||
self.resources.defer(
|
||||
lambda: unwrap(
|
||||
self.bootstrap.proxy.transport.delete(
|
||||
"/organization/delete",
|
||||
headers=self.bootstrap.proxy.management_headers(),
|
||||
json=OrgDeleteBody(organization_ids=[organization_id]),
|
||||
response_type=OrgDeleteResponse,
|
||||
)
|
||||
)
|
||||
)
|
||||
team_id: Final = self.bootstrap.proxy.create_team(
|
||||
TeamNewBody(team_alias=f"e2e-team-{marker}", organization_id=organization_id)
|
||||
)
|
||||
self.resources.defer(
|
||||
lambda: unwrap(
|
||||
self.bootstrap.proxy.transport.post(
|
||||
"/team/delete",
|
||||
headers=self.bootstrap.proxy.management_headers(),
|
||||
json=TeamDeleteBody(team_ids=[team_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
)
|
||||
)
|
||||
self.bootstrap.delete_team_member(team_id, self.bootstrap.user_info().user_id)
|
||||
group_id: Final = self.idp.create_group(team_id)
|
||||
self.resources.defer(lambda: self.idp.with_strict_cleanup().delete_group(group_id))
|
||||
return Tenant(organization_id=organization_id, team_id=team_id, group_id=group_id)
|
||||
|
||||
def create(
|
||||
self, role: ActorRole, *, tenants: tuple[Tenant, ...] = (), profile: ActorProfile = "database_role"
|
||||
) -> Actor:
|
||||
if role in ("team_admin", "team_member", "organization_admin") and not tenants:
|
||||
raise ValueError("A membership actor requires a tenant")
|
||||
identity: Final = self.idp.with_strict_cleanup().provision_user(
|
||||
marker=unique_marker(),
|
||||
groups=tuple(tenant.team_id for tenant in tenants) if profile == "group_scoped" else (),
|
||||
group_ids=tuple(tenant.group_id for tenant in tenants) if profile == "group_scoped" else (),
|
||||
defer=self.resources.defer,
|
||||
)
|
||||
global_role: Final[UserRole] = (
|
||||
role
|
||||
if role in ("proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer")
|
||||
else "internal_user"
|
||||
)
|
||||
self.bootstrap.create_user(
|
||||
UserNewBody(
|
||||
user_id=identity.user_id,
|
||||
user_email=f"{identity.username}@example.com",
|
||||
user_role=global_role,
|
||||
auto_create_key=False,
|
||||
)
|
||||
)
|
||||
self.resources.defer(lambda: self.bootstrap.delete_user_strict(identity.user_id))
|
||||
for tenant in tenants:
|
||||
unwrap(
|
||||
self.bootstrap.proxy.transport.post(
|
||||
"/organization/member_add",
|
||||
headers=self.bootstrap.proxy.management_headers(),
|
||||
json=OrgMemberAddBody(
|
||||
organization_id=tenant.organization_id,
|
||||
member=OrgMemberEntry(
|
||||
user_id=identity.user_id,
|
||||
role="org_admin" if role == "organization_admin" else "internal_user",
|
||||
),
|
||||
),
|
||||
response_type=NoBody,
|
||||
)
|
||||
)
|
||||
unwrap(
|
||||
self.bootstrap.proxy.transport.post(
|
||||
"/team/member_add",
|
||||
headers=self.bootstrap.proxy.management_headers(),
|
||||
json=TeamMemberAddBody(
|
||||
team_id=tenant.team_id,
|
||||
member=TeamMemberEntry(
|
||||
user_id=identity.user_id,
|
||||
role="admin" if role == "team_admin" else "user",
|
||||
),
|
||||
),
|
||||
response_type=NoBody,
|
||||
)
|
||||
)
|
||||
return Actor(identity=identity, role=role, global_role=global_role, profile=profile, tenants=tenants)
|
||||
|
|
@ -7,7 +7,8 @@ llm-only key hitting a management route).
|
|||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
import warnings
|
||||
from dataclasses import dataclass, field, replace
|
||||
|
||||
import jwt
|
||||
from e2e_config import MASTER_KEY
|
||||
|
|
@ -20,6 +21,7 @@ from e2e_http import (
|
|||
StreamingResponse,
|
||||
Success,
|
||||
UnknownApiError,
|
||||
retry_attempts,
|
||||
unwrap,
|
||||
)
|
||||
from models import (
|
||||
|
|
@ -81,7 +83,7 @@ from models import (
|
|||
UserNewResponse,
|
||||
UserUpdateBody,
|
||||
)
|
||||
from proxy_client import ProxyClient
|
||||
from proxy_client import Caller, ProxyClient
|
||||
|
||||
MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied"
|
||||
ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route"
|
||||
|
|
@ -98,7 +100,7 @@ class DashboardSession:
|
|||
its bearer on every subsequent call, the claims it renders the signed-in user
|
||||
from, and where it lands the browser."""
|
||||
|
||||
session_key: str
|
||||
session_key: str = field(repr=False)
|
||||
claims: UiSessionClaims
|
||||
redirect_url: str
|
||||
|
||||
|
|
@ -106,7 +108,10 @@ class DashboardSession:
|
|||
@dataclass(frozen=True, slots=True)
|
||||
class ManagementClient:
|
||||
proxy: ProxyClient
|
||||
master_key: str
|
||||
master_key: str = field(repr=False)
|
||||
|
||||
def with_caller(self, caller: Caller) -> ManagementClient:
|
||||
return replace(self, proxy=self.proxy.with_caller(caller))
|
||||
|
||||
def llm_only_key(self) -> str:
|
||||
return self.proxy.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"]))
|
||||
|
|
@ -117,7 +122,7 @@ class ManagementClient:
|
|||
dashboard creates it under the session key their sign-in minted). Returns
|
||||
the outcome rather than unwrapping it, so a caller can poll a route that is
|
||||
only transiently refusing."""
|
||||
headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key)
|
||||
headers = self.proxy.management_headers(caller_key)
|
||||
return self.proxy.transport.post(
|
||||
"/key/generate",
|
||||
headers=headers,
|
||||
|
|
@ -131,9 +136,9 @@ class ManagementClient:
|
|||
sign-in minted, never the master key). Returns the outcome rather than
|
||||
unwrapping it, so a caller can poll a route that is only transiently
|
||||
refusing; `update_key_models` is the unwrapping shorthand."""
|
||||
headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key)
|
||||
headers = self.proxy.management_headers(caller_key)
|
||||
last: Result[NoBody] = NetworkError(message="/key/update was never attempted")
|
||||
for attempt in range(_KEY_WRITE_ATTEMPTS):
|
||||
for attempt in range(retry_attempts(_KEY_WRITE_ATTEMPTS)):
|
||||
last = self.proxy.transport.post(
|
||||
"/key/update",
|
||||
headers=headers,
|
||||
|
|
@ -144,6 +149,7 @@ class ManagementClient:
|
|||
case UnknownApiError(body=error_body) if any(
|
||||
marker in error_body.lower() for marker in _TRANSIENT_BACKEND_MARKERS
|
||||
):
|
||||
warnings.warn(f"Transient backend response on attempt {attempt + 1}", RuntimeWarning, stacklevel=2)
|
||||
time.sleep(0.5 * (attempt + 1))
|
||||
continue
|
||||
case _:
|
||||
|
|
@ -153,25 +159,26 @@ class ManagementClient:
|
|||
def update_key_models(self, key: str, models: list[str]) -> None:
|
||||
_ = unwrap(self.update_key(KeyUpdateBody(key=key, models=models)))
|
||||
|
||||
def key_info_as(self, key: str, *, caller_key: str) -> Result[KeyInfoResponse]:
|
||||
def key_info_as(self, key: str, *, caller_key: str | None = None) -> Result[KeyInfoResponse]:
|
||||
return self.proxy.transport.get(
|
||||
"/key/info",
|
||||
headers=self.proxy.transport.bearer(caller_key),
|
||||
headers=self.proxy.management_headers(caller_key),
|
||||
params=KeyInfoParams(key=key),
|
||||
response_type=KeyInfoResponse,
|
||||
)
|
||||
|
||||
def delete_key_strict(self, key: str, *, caller_key: str | None = None) -> None:
|
||||
def delete_key_strict(self, key: str, *, caller_key: str | None = None, missing_ok: bool = False) -> None:
|
||||
"""Strict delete for the act phase of a test: a failed delete is a hard
|
||||
failure, unlike the warn-only ProxyClient.delete_key used at teardown."""
|
||||
_ = unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/key/delete",
|
||||
headers=self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key),
|
||||
json=KeyDeleteBody(keys=[key]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
result = self.proxy.transport.post(
|
||||
"/key/delete",
|
||||
headers=self.proxy.management_headers(caller_key),
|
||||
json=KeyDeleteBody(keys=[key]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
if missing_ok and isinstance(result, UnknownApiError) and result.status_code == 404:
|
||||
return
|
||||
_ = unwrap(result)
|
||||
|
||||
def delete_model_strict(self, model_id: str) -> None:
|
||||
"""Strict delete for the act phase of a test: a failed delete is a hard
|
||||
|
|
@ -179,7 +186,7 @@ class ManagementClient:
|
|||
_ = unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/model/delete",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=ModelDeleteBody(id=model_id),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -190,7 +197,7 @@ class ManagementClient:
|
|||
Connection button, probing the live provider with the supplied params."""
|
||||
return self.proxy.transport.post(
|
||||
"/health/test_connection",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=body,
|
||||
response_type=ConnectionTestResponse,
|
||||
timeout=120.0,
|
||||
|
|
@ -200,7 +207,7 @@ class ManagementClient:
|
|||
_ = unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/key/block",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=KeyBlockBody(key=key),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -209,7 +216,7 @@ class ManagementClient:
|
|||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/key/regenerate",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=KeyRegenerateBody(key=key, grace_period=grace_period),
|
||||
response_type=KeyGenerateResponse,
|
||||
)
|
||||
|
|
@ -219,7 +226,7 @@ class ManagementClient:
|
|||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
f"/key/{key}/reset_spend",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=KeyResetSpendBody(reset_to=reset_to),
|
||||
response_type=KeyResetSpendResponse,
|
||||
)
|
||||
|
|
@ -228,7 +235,7 @@ class ManagementClient:
|
|||
def key_list(self, key_alias: str, *, caller_key: str | None = None) -> Result[KeyListResponse]:
|
||||
"""GET /key/list, the Virtual Keys page's own inventory call. `caller_key` is
|
||||
who is asking: the master key by default, or a virtual key."""
|
||||
headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key)
|
||||
headers = self.proxy.management_headers(caller_key)
|
||||
return self.proxy.transport.get(
|
||||
"/key/list",
|
||||
headers=headers,
|
||||
|
|
@ -266,7 +273,7 @@ class ManagementClient:
|
|||
team_id = unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/team/new",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=body,
|
||||
response_type=TeamNewResponse,
|
||||
)
|
||||
|
|
@ -276,10 +283,10 @@ class ManagementClient:
|
|||
|
||||
def update_team(self, body: TeamUpdateBody) -> None:
|
||||
last: Result[NoBody] | None = None
|
||||
for attempt in range(5):
|
||||
for attempt in range(retry_attempts(5)):
|
||||
last = self.proxy.transport.post(
|
||||
"/team/update",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=body,
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -289,6 +296,7 @@ class ManagementClient:
|
|||
case UnknownApiError(body=body_text) if (
|
||||
"connecting to redis" in body_text.lower() or "name resolution" in body_text.lower()
|
||||
):
|
||||
warnings.warn(f"Transient backend response on attempt {attempt + 1}", RuntimeWarning, stacklevel=2)
|
||||
time.sleep(0.5 * (attempt + 1))
|
||||
continue
|
||||
case _:
|
||||
|
|
@ -299,7 +307,7 @@ class ManagementClient:
|
|||
def delete_team(self, team_id: str) -> None:
|
||||
_ = self.proxy.transport.post(
|
||||
"/team/delete",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=TeamDeleteBody(team_ids=[team_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -308,7 +316,7 @@ class ManagementClient:
|
|||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/team/info",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
params=TeamInfoParams(team_id=team_id),
|
||||
response_type=TeamInfoResponse,
|
||||
)
|
||||
|
|
@ -320,7 +328,7 @@ class ManagementClient:
|
|||
for entry in unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/team/list",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
params=NoBody(),
|
||||
response_type=TeamListResponse,
|
||||
)
|
||||
|
|
@ -328,14 +336,16 @@ class ManagementClient:
|
|||
)
|
||||
|
||||
def team_info_status(self, team_id: str) -> ProbeResult:
|
||||
return self.proxy.transport.probe("/team/info", params=TeamInfoParams(team_id=team_id))
|
||||
return self.proxy.transport.probe(
|
||||
"/team/info", params=TeamInfoParams(team_id=team_id), headers=self.proxy.management_headers()
|
||||
)
|
||||
|
||||
def _wait_for_team(self, team_id: str) -> None:
|
||||
last: Result[TeamInfoResponse] | None = None
|
||||
for _ in range(_TEAM_READY_ATTEMPTS):
|
||||
for _ in range(retry_attempts(_TEAM_READY_ATTEMPTS)):
|
||||
last = self.proxy.transport.get(
|
||||
"/team/info",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
params=TeamInfoParams(team_id=team_id),
|
||||
response_type=TeamInfoResponse,
|
||||
)
|
||||
|
|
@ -343,25 +353,29 @@ class ManagementClient:
|
|||
case Success():
|
||||
return
|
||||
case _:
|
||||
warnings.warn("Repeating team read while the team becomes available", RuntimeWarning, stacklevel=2)
|
||||
time.sleep(_TEAM_READY_SLEEP_SECONDS)
|
||||
assert last is not None
|
||||
raise AssertionError(last)
|
||||
|
||||
def add_team_member(self, team_id: str, user_id: str) -> None:
|
||||
last: Result[NoBody] | None = None
|
||||
for attempt in range(_TEAM_READY_ATTEMPTS):
|
||||
for attempt in range(retry_attempts(_TEAM_READY_ATTEMPTS)):
|
||||
last = self.proxy.transport.post(
|
||||
"/team/member_add",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=TeamMemberAddBody(team_id=team_id, member=TeamMemberEntry(role="user", user_id=user_id)),
|
||||
response_type=NoBody,
|
||||
)
|
||||
match last:
|
||||
case Success():
|
||||
return
|
||||
case UnknownApiError(body=body) if (
|
||||
"doesn't exist" in body and attempt + 1 < _TEAM_READY_ATTEMPTS
|
||||
case UnknownApiError(body=body) if "doesn't exist" in body and attempt + 1 < retry_attempts(
|
||||
_TEAM_READY_ATTEMPTS
|
||||
):
|
||||
warnings.warn(
|
||||
"Retrying team membership while the team becomes available", RuntimeWarning, stacklevel=2
|
||||
)
|
||||
time.sleep(_TEAM_READY_SLEEP_SECONDS)
|
||||
continue
|
||||
case _:
|
||||
|
|
@ -373,7 +387,7 @@ class ManagementClient:
|
|||
_ = unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/team/member_delete",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=TeamMemberDeleteBody(team_id=team_id, user_id=user_id),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -383,7 +397,7 @@ class ManagementClient:
|
|||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/user/new",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=body,
|
||||
response_type=UserNewResponse,
|
||||
)
|
||||
|
|
@ -393,7 +407,7 @@ class ManagementClient:
|
|||
_ = unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/customer/new",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=CustomerNewBody(user_id=user_id),
|
||||
response_type=CustomerResponse,
|
||||
)
|
||||
|
|
@ -404,7 +418,7 @@ class ManagementClient:
|
|||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/customer/info",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
params=CustomerInfoParams(end_user_id=end_user_id),
|
||||
response_type=CustomerResponse,
|
||||
)
|
||||
|
|
@ -413,7 +427,7 @@ class ManagementClient:
|
|||
def delete_customer(self, user_id: str) -> None:
|
||||
_ = self.proxy.transport.post(
|
||||
"/customer/delete",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=CustomerDeleteBody(user_ids=[user_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -422,7 +436,7 @@ class ManagementClient:
|
|||
_ = unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/user/update",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=body,
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -431,7 +445,7 @@ class ManagementClient:
|
|||
def delete_user(self, user_id: str) -> None:
|
||||
_ = self.proxy.transport.post(
|
||||
"/user/delete",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=UserDeleteBody(user_ids=[user_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -442,17 +456,17 @@ class ManagementClient:
|
|||
_ = unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/user/delete",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=UserDeleteBody(user_ids=[user_id]),
|
||||
response_type=UserDeleteResponse,
|
||||
)
|
||||
)
|
||||
|
||||
def user_info(self, user_id: str) -> UserInfoResponse:
|
||||
def user_info(self, user_id: str | None = None) -> UserInfoResponse:
|
||||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/user/info",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
params=UserInfoParams(user_id=user_id),
|
||||
response_type=UserInfoResponse,
|
||||
)
|
||||
|
|
@ -462,7 +476,7 @@ class ManagementClient:
|
|||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/user/list",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
params=UserListParams(user_ids=user_id),
|
||||
response_type=UserListResponse,
|
||||
)
|
||||
|
|
@ -472,7 +486,7 @@ class ManagementClient:
|
|||
listing = unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/user/list",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
params=UserListParams(user_ids=user_id),
|
||||
response_type=UserListResponse,
|
||||
)
|
||||
|
|
@ -483,7 +497,7 @@ class ManagementClient:
|
|||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/organization/new",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=body,
|
||||
response_type=OrgNewResponse,
|
||||
)
|
||||
|
|
@ -493,7 +507,7 @@ class ManagementClient:
|
|||
_ = unwrap(
|
||||
self.proxy.transport.patch(
|
||||
"/organization/update",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=body,
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -502,7 +516,7 @@ class ManagementClient:
|
|||
def delete_org(self, organization_id: str) -> None:
|
||||
_ = self.proxy.transport.delete(
|
||||
"/organization/delete",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=OrgDeleteBody(organization_ids=[organization_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -511,19 +525,24 @@ class ManagementClient:
|
|||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/organization/info",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
params=OrgInfoParams(organization_id=organization_id),
|
||||
response_type=OrgInfoResponse,
|
||||
)
|
||||
)
|
||||
|
||||
def org_info_status(self, organization_id: str) -> ProbeResult:
|
||||
return self.proxy.transport.probe("/organization/info", params=OrgInfoParams(organization_id=organization_id))
|
||||
return self.proxy.transport.probe(
|
||||
"/organization/info",
|
||||
params=OrgInfoParams(organization_id=organization_id),
|
||||
headers=self.proxy.management_headers(),
|
||||
)
|
||||
|
||||
def create_tag(self, body: TagNewBody) -> None:
|
||||
_ = unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/tag/new",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=body,
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -532,7 +551,7 @@ class ManagementClient:
|
|||
def delete_tag(self, name: str) -> None:
|
||||
_ = self.proxy.transport.post(
|
||||
"/tag/delete",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=TagDeleteBody(name=name),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -542,7 +561,7 @@ class ManagementClient:
|
|||
unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/tag/list",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
params=NoBody(),
|
||||
response_type=TagListResponse,
|
||||
)
|
||||
|
|
@ -553,7 +572,7 @@ class ManagementClient:
|
|||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/v1/mcp/server",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=body,
|
||||
response_type=McpServerRow,
|
||||
)
|
||||
|
|
@ -565,7 +584,7 @@ class ManagementClient:
|
|||
return unwrap(
|
||||
self.proxy.transport.put(
|
||||
"/v1/mcp/server",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=body,
|
||||
response_type=McpServerRow,
|
||||
)
|
||||
|
|
@ -576,7 +595,7 @@ class ManagementClient:
|
|||
unwrap it while a deferred teardown can ignore an already-deleted server."""
|
||||
return self.proxy.transport.delete(
|
||||
f"/v1/mcp/server/{server_id}",
|
||||
headers=self.proxy.transport.master,
|
||||
headers=self.proxy.management_headers(),
|
||||
json=NoBody(),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2,60 +2,247 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Final
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
from e2e_config import CHEAP_OPENAI_MODEL, unique_marker
|
||||
from e2e_config import CHEAP_OPENAI_MODEL, PROXY_BASE_URL, unique_marker
|
||||
from e2e_http import UnauthorizedError, UnknownApiError, unwrap
|
||||
from idp import ADMIN_CLIENT_ID, Identity, Keycloak
|
||||
from idp import ADMIN_CLIENT_ID, Identity, Keycloak, token_claims
|
||||
from lifecycle import ResourceManager
|
||||
from management.jwt_actors import ActorFactory, ActorRole
|
||||
from management_client import ManagementClient
|
||||
from models import KeyGenerateBody, KeyUpdateBody, TeamNewBody, UserNewBody
|
||||
from models import KeyGenerateBody, KeyUpdateBody, TeamNewBody, UserInfoParams, UserInfoResponse, UserNewBody
|
||||
from proxy_client import Caller
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
|
||||
class TestJwtManagement:
|
||||
@pytest.mark.covers("mgmt.key.jwt.lifecycle")
|
||||
def test_admin_creates_reads_updates_clears_and_deletes_a_key(
|
||||
self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager
|
||||
) -> None:
|
||||
admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID)
|
||||
alias: Final = f"e2e-jwt-key-{unique_marker()}"
|
||||
created: Final = unwrap(
|
||||
client.generate_key(
|
||||
KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group, models=[CHEAP_OPENAI_MODEL]),
|
||||
caller_key=admin,
|
||||
@pytest.mark.parametrize(
|
||||
"role",
|
||||
(
|
||||
"proxy_admin",
|
||||
"proxy_admin_viewer",
|
||||
"organization_admin",
|
||||
"team_admin",
|
||||
"team_member",
|
||||
"internal_user",
|
||||
"internal_user_viewer",
|
||||
"unrelated_user",
|
||||
),
|
||||
)
|
||||
@pytest.mark.covers("mgmt.user.jwt.database_roles")
|
||||
def test_actor_subject_and_database_role(self, actor_factory: ActorFactory, role: ActorRole) -> None:
|
||||
tenants: Final = (
|
||||
(actor_factory.tenant(),) if role in ("organization_admin", "team_admin", "team_member") else ()
|
||||
)
|
||||
actor: Final = actor_factory.create(role, tenants=tenants)
|
||||
caller: Final = actor.mint_caller(actor_factory.idp)
|
||||
claims: Final = token_claims(caller.credential)
|
||||
assert claims.sub == actor.identity.user_id
|
||||
assert claims.iss == actor_factory.idp.issuer
|
||||
assert claims.aud == "litellm-e2e" or "litellm-e2e" in claims.aud
|
||||
assert actor.identity.groups == ()
|
||||
assert ("litellm_proxy_admin" in claims.scope.split()) == (role == "proxy_admin")
|
||||
stored: Final = actor_factory.bootstrap.user_info(actor.identity.user_id)
|
||||
assert stored.user_id == actor.identity.user_id
|
||||
assert stored.user_info.user_role == actor.global_role
|
||||
bound: Final = actor_factory.bootstrap.with_caller(caller)
|
||||
own: Final = unwrap(
|
||||
bound.proxy.transport.get(
|
||||
"/user/info",
|
||||
headers=bound.proxy.management_headers(),
|
||||
params=UserInfoParams(),
|
||||
response_type=UserInfoResponse,
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: client.proxy.delete_key(created.key))
|
||||
assert own.user_id == actor.identity.user_id
|
||||
assert own.user_info.user_role == actor.global_role
|
||||
for tenant in tenants:
|
||||
info = actor_factory.bootstrap.team_info(tenant.team_id)
|
||||
assert info.organization_id == tenant.organization_id
|
||||
assert {(member.user_id, member.role) for member in info.members_with_roles} == {
|
||||
(actor.identity.user_id, "admin" if role == "team_admin" else "user")
|
||||
}
|
||||
assert {
|
||||
(member.user_id, member.user_role)
|
||||
for member in actor_factory.bootstrap.org_info(tenant.organization_id).members
|
||||
} == {(actor.identity.user_id, "org_admin" if role == "organization_admin" else "internal_user")}
|
||||
|
||||
original: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info
|
||||
assert original.key_alias == alias and original.team_id == jwt_identity.group
|
||||
@pytest.mark.covers("mgmt.key.jwt.viewer_denied")
|
||||
def test_admin_viewer_reads_but_cannot_update(self, actor_factory: ActorFactory) -> None:
|
||||
actor: Final = actor_factory.create("proxy_admin_viewer")
|
||||
viewer: Final = actor_factory.bootstrap.with_caller(actor.mint_caller(actor_factory.idp))
|
||||
alias: Final = f"e2e-viewer-{unique_marker()}"
|
||||
key: Final = actor_factory.key().key
|
||||
unwrap(actor_factory.bootstrap.update_key(KeyUpdateBody(key=key, key_alias=alias)))
|
||||
assert viewer.proxy.key_info(key).key_alias == alias
|
||||
denied: Final = viewer.update_key(KeyUpdateBody(key=key, key_alias="forbidden"))
|
||||
assert isinstance(denied, UnknownApiError) and denied.status_code == 403, f"viewer write was accepted: {denied}"
|
||||
assert "proxy_admin_viewer" in denied.body and "/key/update" in denied.body
|
||||
assert actor_factory.bootstrap.proxy.key_info(key).key_alias == alias
|
||||
|
||||
@pytest.mark.covers("mgmt.user.oidc.identity_mapping")
|
||||
def test_oidc_browser_profile_identity_mapping(self, actor_factory: ActorFactory) -> None:
|
||||
actor: Final = actor_factory.create("internal_user")
|
||||
idp: Final = actor_factory.idp.with_strict_cleanup()
|
||||
discovery: Final = idp.discovery()
|
||||
assert discovery.issuer == idp.issuer
|
||||
assert discovery.jwks_uri == idp.jwks_url
|
||||
callback: Final = f"{PROXY_BASE_URL}/sso/callback"
|
||||
browser: Final = idp.browser_client(callback_url=callback, defer=actor_factory.resources.defer)
|
||||
token: Final = idp.browser_token(actor.identity, browser)
|
||||
assert token_claims(token).sub == actor.identity.user_id
|
||||
userinfo: Final = idp.userinfo(token)
|
||||
assert userinfo.sub == actor.identity.user_id
|
||||
assert userinfo.email == f"{actor.identity.username}@example.com"
|
||||
assert browser.environment(discovery)["GENERIC_USER_ID_ATTRIBUTE"] == "sub"
|
||||
|
||||
@pytest.mark.covers("mgmt.key.jwt.lifecycle")
|
||||
@pytest.mark.parametrize("credential_kind", ("direct_jwt", "virtual_key"))
|
||||
def test_admin_creates_reads_updates_clears_and_deletes_a_key(
|
||||
self,
|
||||
actor_factory: ActorFactory,
|
||||
credential_kind: Literal["direct_jwt", "virtual_key"],
|
||||
) -> None:
|
||||
tenant: Final = actor_factory.tenant()
|
||||
actor: Final = actor_factory.create("proxy_admin", tenants=(tenant,), profile="group_scoped")
|
||||
virtual_key: Final = (
|
||||
actor_factory.key(user_id=actor.identity.user_id).key if credential_kind == "virtual_key" else None
|
||||
)
|
||||
admin: Final = virtual_key if virtual_key is not None else actor.mint_caller(actor_factory.idp).credential
|
||||
bound: Final = actor_factory.bootstrap.with_caller(
|
||||
Caller(credential=admin, kind=credential_kind, role="proxy_admin")
|
||||
)
|
||||
assert bound.user_info().user_id == actor.identity.user_id
|
||||
alias: Final = f"e2e-jwt-key-{unique_marker()}"
|
||||
created: Final = unwrap(
|
||||
bound.generate_key(
|
||||
KeyGenerateBody(key_alias=alias, team_id=tenant.team_id, models=[CHEAP_OPENAI_MODEL]),
|
||||
)
|
||||
)
|
||||
actor_factory.resources.defer(lambda: actor_factory.bootstrap.delete_key_strict(created.key, missing_ok=True))
|
||||
|
||||
original: Final = unwrap(bound.key_info_as(created.key)).info
|
||||
assert original.key_alias == alias and original.team_id == tenant.team_id
|
||||
assert original.models == [CHEAP_OPENAI_MODEL]
|
||||
|
||||
updated_alias: Final = f"{alias}-updated"
|
||||
unwrap(
|
||||
client.update_key(KeyUpdateBody(key=created.key, key_alias=updated_alias, rpm_limit=120), caller_key=admin)
|
||||
)
|
||||
updated: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info
|
||||
unwrap(bound.update_key(KeyUpdateBody(key=created.key, key_alias=updated_alias, rpm_limit=120)))
|
||||
updated: Final = unwrap(bound.key_info_as(created.key)).info
|
||||
assert updated.key_alias == updated_alias and updated.rpm_limit == 120
|
||||
assert updated.models == [CHEAP_OPENAI_MODEL], "omitted models must preserve the restriction"
|
||||
|
||||
unwrap(client.update_key(KeyUpdateBody(key=created.key, models=[]), caller_key=admin))
|
||||
cleared: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info
|
||||
unwrap(bound.update_key(KeyUpdateBody(key=created.key, models=[])))
|
||||
cleared: Final = unwrap(bound.key_info_as(created.key)).info
|
||||
assert cleared.models == [] and cleared.rpm_limit == 120
|
||||
|
||||
assert unwrap(client.key_list(updated_alias, caller_key=admin)).total_count == 1
|
||||
client.delete_key_strict(created.key, caller_key=admin)
|
||||
assert unwrap(client.key_list(updated_alias, caller_key=admin)).total_count == 0
|
||||
assert unwrap(bound.key_list(updated_alias)).total_count == 1
|
||||
bound.delete_key_strict(created.key)
|
||||
assert unwrap(bound.key_list(updated_alias)).total_count == 0
|
||||
|
||||
@pytest.mark.covers("mgmt.team.jwt.tenant_isolation")
|
||||
def test_two_actor_sets_keep_tenants_and_keys_isolated(self, actor_factory: ActorFactory) -> None:
|
||||
first: Final = actor_factory.tenant()
|
||||
second: Final = actor_factory.tenant()
|
||||
assert first.organization_id != second.organization_id and first.team_id != second.team_id
|
||||
actors: Final = tuple(
|
||||
actor_factory.create("team_member", tenants=(tenant,), profile="group_scoped") for tenant in (first, second)
|
||||
)
|
||||
assert actors[0].identity.user_id != actors[1].identity.user_id
|
||||
callers: Final = tuple(
|
||||
actor_factory.bootstrap.with_caller(actor.mint_caller(actor_factory.idp)) for actor in actors
|
||||
)
|
||||
keys: Final = tuple(actor_factory.key(tenant) for tenant in (first, second))
|
||||
assert keys[0].key != keys[1].key
|
||||
assert callers[0].proxy.key_info(keys[0].key).team_id == first.team_id
|
||||
assert callers[1].proxy.key_info(keys[1].key).team_id == second.team_id
|
||||
for caller, other_key in ((callers[0], keys[1].key), (callers[1], keys[0].key)):
|
||||
hidden = caller.key_info_as(other_key)
|
||||
assert isinstance(hidden, UnknownApiError) and hidden.status_code == 403
|
||||
assert tuple(actor.identity.groups for actor in actors) == ((first.team_id,), (second.team_id,))
|
||||
|
||||
@pytest.mark.covers("mgmt.team.jwt.multiple_memberships")
|
||||
def test_multi_group_actor_keeps_exact_memberships(self, actor_factory: ActorFactory) -> None:
|
||||
tenants: Final = (actor_factory.tenant(), actor_factory.tenant())
|
||||
actor: Final = actor_factory.create("team_member", tenants=tenants, profile="group_scoped")
|
||||
claims: Final = token_claims(actor.mint_caller(actor_factory.idp).credential)
|
||||
assert set(claims.groups) == {tenant.team_id for tenant in tenants}
|
||||
assert "litellm_proxy_admin" not in claims.scope.split()
|
||||
assert actor.identity.groups == tuple(tenant.team_id for tenant in tenants)
|
||||
for tenant in tenants:
|
||||
assert {
|
||||
(entry.user_id, entry.role)
|
||||
for entry in actor_factory.bootstrap.team_info(tenant.team_id).members_with_roles
|
||||
} == {(actor.identity.user_id, "user")}
|
||||
|
||||
@pytest.mark.covers("mgmt.user.jwt.cleanup")
|
||||
def test_successful_actor_cleanup_removes_owned_state(self, actor_factory: ActorFactory) -> None:
|
||||
resources: Final = ResourceManager(client=actor_factory.bootstrap.proxy, strict_cleanup=True)
|
||||
factory: Final = ActorFactory(bootstrap=actor_factory.bootstrap, idp=actor_factory.idp, resources=resources)
|
||||
try:
|
||||
tenant: Final = factory.tenant()
|
||||
actor: Final = factory.create("team_member", tenants=(tenant,), profile="group_scoped")
|
||||
key: Final = factory.key(tenant)
|
||||
alias: Final = factory.bootstrap.proxy.key_info(key.key).key_alias
|
||||
assert alias is not None
|
||||
finally:
|
||||
resources.teardown()
|
||||
assert factory.bootstrap.user_count(actor.identity.user_id) == 0
|
||||
assert factory.bootstrap.key_alias_count(alias) == 0
|
||||
assert factory.bootstrap.team_info_status(tenant.team_id).status_code == 404
|
||||
assert factory.bootstrap.org_info_status(tenant.organization_id).status_code == 404
|
||||
factory.idp.assert_absent("users", actor.identity.user_id)
|
||||
factory.idp.assert_absent("groups", tenant.group_id)
|
||||
|
||||
@pytest.mark.parametrize("stage", ("group", "user"))
|
||||
@pytest.mark.covers("mgmt.user.jwt.partial_cleanup")
|
||||
def test_partial_setup_removes_previously_created_identities(
|
||||
self,
|
||||
actor_factory: ActorFactory,
|
||||
stage: Literal["group", "user"],
|
||||
) -> None:
|
||||
idp: Final = actor_factory.idp.with_strict_cleanup()
|
||||
resources: Final = ResourceManager(client=actor_factory.bootstrap.proxy, strict_cleanup=True)
|
||||
marker: Final = unique_marker()
|
||||
group_id: Final = idp.create_group(f"e2e-partial-{marker}")
|
||||
resources.defer(lambda: idp.delete_group(group_id))
|
||||
try:
|
||||
identity: Final = (
|
||||
idp.provision_user(
|
||||
marker=marker,
|
||||
groups=(f"e2e-partial-{marker}",),
|
||||
group_ids=(group_id,),
|
||||
defer=resources.defer,
|
||||
)
|
||||
if stage == "user"
|
||||
else None
|
||||
)
|
||||
if identity is None:
|
||||
with pytest.raises(pytest.fail.Exception, match="HTTP 409"):
|
||||
idp.create_group(f"e2e-partial-{marker}")
|
||||
else:
|
||||
with pytest.raises(pytest.fail.Exception, match="HTTP 409"):
|
||||
idp.create_user(
|
||||
username=identity.username,
|
||||
email=f"{identity.username}@example.com",
|
||||
password=identity.password,
|
||||
groups=identity.groups,
|
||||
)
|
||||
finally:
|
||||
resources.teardown()
|
||||
idp.assert_absent("groups", group_id)
|
||||
if identity is not None:
|
||||
idp.assert_absent("users", identity.user_id)
|
||||
|
||||
@pytest.mark.covers("mgmt.key.jwt.member_denied", "mgmt.key.jwt.other_team_denied")
|
||||
def test_member_cannot_write_and_another_team_cannot_read_the_key(
|
||||
self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager
|
||||
) -> None:
|
||||
admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID)
|
||||
bound: Final = client.with_caller(Caller(credential=admin, kind="direct_jwt", role="proxy_admin"))
|
||||
member: Final = idp.access_token(jwt_identity)
|
||||
member_client: Final = client.with_caller(Caller(credential=member, kind="direct_jwt", role="team_member"))
|
||||
alias: Final = f"e2e-jwt-owned-{unique_marker()}"
|
||||
created: Final = unwrap(
|
||||
client.generate_key(KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group), caller_key=admin)
|
||||
|
|
@ -63,14 +250,14 @@ class TestJwtManagement:
|
|||
resources.defer(lambda: client.proxy.delete_key(created.key))
|
||||
|
||||
client.add_team_member(jwt_identity.group, jwt_identity.user_id)
|
||||
assert unwrap(client.key_info_as(created.key, caller_key=member)).info.key_alias == alias
|
||||
assert unwrap(member_client.key_info_as(created.key)).info.key_alias == alias
|
||||
|
||||
refused: Final = client.update_key(KeyUpdateBody(key=created.key, key_alias="forbidden"), caller_key=member)
|
||||
refused: Final = member_client.update_key(KeyUpdateBody(key=created.key, key_alias="forbidden"))
|
||||
assert isinstance(refused, UnauthorizedError), f"member write was accepted: {refused}"
|
||||
assert "does not have permissions for endpoint" in refused.body.lower(), (
|
||||
f"expected a permission denial: {refused}"
|
||||
)
|
||||
assert unwrap(client.key_info_as(created.key, caller_key=admin)).info.key_alias == alias
|
||||
assert unwrap(bound.key_info_as(created.key)).info.key_alias == alias
|
||||
|
||||
marker: Final = unique_marker()
|
||||
outsider: Final = idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}", defer=resources.defer)
|
||||
|
|
@ -88,4 +275,4 @@ class TestJwtManagement:
|
|||
assert isinstance(hidden, UnknownApiError) and hidden.status_code == 403, (
|
||||
f"another team must not read this key: {hidden}"
|
||||
)
|
||||
assert unwrap(client.key_info_as(created.key, caller_key=admin)).info.team_id == jwt_identity.group
|
||||
assert unwrap(bound.key_info_as(created.key)).info.team_id == jwt_identity.group
|
||||
|
|
|
|||
|
|
@ -283,10 +283,15 @@ class ChatToolResultTurn(BaseModel):
|
|||
type ChatTurn = ChatMessage | ChatAssistantTurn | ChatToolResultTurn
|
||||
|
||||
|
||||
class ChatStreamOptions(BaseModel):
|
||||
include_usage: bool
|
||||
|
||||
|
||||
class ChatBody(BaseModel):
|
||||
model: str
|
||||
messages: Sequence[ChatTurn]
|
||||
stream: bool = False
|
||||
stream_options: ChatStreamOptions | None = None
|
||||
max_tokens: int | None = None
|
||||
max_completion_tokens: int | None = None
|
||||
temperature: float | None = None
|
||||
|
|
@ -488,12 +493,18 @@ class AnthropicToolResultTurn(BaseModel):
|
|||
type AnthropicMessage = ChatMessage | AnthropicAssistantTurn | AnthropicToolResultTurn
|
||||
|
||||
|
||||
class AnthropicToolChoice(BaseModel):
|
||||
type: Literal["auto", "any", "tool", "none"]
|
||||
name: str | None = None
|
||||
|
||||
|
||||
class AnthropicMessagesBody(BaseModel):
|
||||
model: str
|
||||
messages: list[AnthropicMessage]
|
||||
max_tokens: int
|
||||
stream: bool | None = None
|
||||
tools: list[AnthropicTool] | None = None
|
||||
tool_choice: AnthropicToolChoice | None = None
|
||||
guardrails: list[str] | None = None
|
||||
cache: dict[str, bool] | None = {"no-cache": True}
|
||||
|
||||
|
|
@ -1091,13 +1102,13 @@ class UiLoginBody(BaseModel):
|
|||
|
||||
|
||||
class UiLoginResponse(BaseModel):
|
||||
token: str
|
||||
token: str = Field(repr=False)
|
||||
redirect_url: str
|
||||
|
||||
|
||||
class UiSessionClaims(BaseModel):
|
||||
user_id: str
|
||||
key: str
|
||||
key: str = Field(repr=False)
|
||||
user_role: str
|
||||
login_method: Literal["sso", "username_password"]
|
||||
exp: int
|
||||
|
|
@ -1135,6 +1146,7 @@ class TeamInfoParams(BaseModel):
|
|||
|
||||
|
||||
class TeamData(BaseModel):
|
||||
organization_id: str | None = None
|
||||
team_alias: str | None = None
|
||||
models: list[str] = []
|
||||
members_with_roles: list[TeamMemberEntry] = []
|
||||
|
|
@ -1175,6 +1187,7 @@ class UserNewBody(BaseModel):
|
|||
user_email: str
|
||||
user_role: UserRole
|
||||
user_id: str | None = None
|
||||
auto_create_key: bool | None = None
|
||||
|
||||
|
||||
class UserNewResponse(BaseModel):
|
||||
|
|
@ -1187,7 +1200,7 @@ class UserUpdateBody(BaseModel):
|
|||
|
||||
|
||||
class UserInfoParams(BaseModel):
|
||||
user_id: str
|
||||
user_id: str | None = None
|
||||
|
||||
|
||||
class UserData(BaseModel):
|
||||
|
|
@ -1240,16 +1253,36 @@ class OrgInfoParams(BaseModel):
|
|||
organization_id: str
|
||||
|
||||
|
||||
class OrgMembership(BaseModel):
|
||||
user_id: str
|
||||
user_role: str
|
||||
|
||||
|
||||
class OrgInfoResponse(BaseModel):
|
||||
organization_id: str
|
||||
organization_alias: str | None = None
|
||||
models: list[str] = []
|
||||
members: tuple[OrgMembership, ...] = ()
|
||||
|
||||
|
||||
class OrgMemberEntry(BaseModel):
|
||||
user_id: str
|
||||
role: Literal["org_admin", "internal_user"]
|
||||
|
||||
|
||||
class OrgMemberAddBody(BaseModel):
|
||||
organization_id: str
|
||||
member: OrgMemberEntry
|
||||
|
||||
|
||||
class OrgDeleteBody(BaseModel):
|
||||
organization_ids: list[str]
|
||||
|
||||
|
||||
class OrgDeleteResponse(RootModel[tuple[OrgInfoResponse, ...]]):
|
||||
pass
|
||||
|
||||
|
||||
# ---------- tags (management) ----------
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -92,6 +92,7 @@ from fixture_mode import (
|
|||
current_test_key,
|
||||
parse_fixture_mode,
|
||||
)
|
||||
from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType(
|
||||
|
|
@ -404,6 +405,12 @@ def _miss_message(test_key: str, slug: str, canonical: CanonicalRequest, bundle:
|
|||
f"under {slug}; re-record with E2E_FIXTURE_MODE=record"
|
||||
)
|
||||
closest, closest_file = _closest_recorded(canonical, recorded)
|
||||
if bundle.manifest.match_profile == "stateless_v1":
|
||||
expected: Final = _JSON.validate_json(closest.content)
|
||||
actual: Final = _JSON.validate_json(canonical.content)
|
||||
assert isinstance(expected, dict) and isinstance(actual, dict)
|
||||
changed: Final = ", ".join(key for key in expected if expected[key] != actual.get(key))
|
||||
return f"stateless_v1 replay mismatch: {changed or 'method/path'}; re-record with E2E_FIXTURE_MODE=record"
|
||||
diff: Final = "\n".join(
|
||||
islice(
|
||||
difflib.unified_diff(
|
||||
|
|
@ -785,11 +792,33 @@ def handle_edge_request(
|
|||
mount, _, upstream_path = split.path.lstrip("/").partition("/")
|
||||
upstream_base: Final = mounts.get(mount)
|
||||
if upstream_base is None:
|
||||
return _text_reply(
|
||||
404, f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(mounts))}"
|
||||
return _text_reply(404, f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(mounts))}")
|
||||
profile: Final = (
|
||||
backend.recorder.profile
|
||||
if isinstance(backend, RecordEdge)
|
||||
else backend.source.bundle.manifest.match_profile
|
||||
if isinstance(backend, ReplayEdge)
|
||||
else "legacy"
|
||||
)
|
||||
identity: Final = (
|
||||
strict_identity(
|
||||
method=method,
|
||||
path=split.path,
|
||||
query=split.query,
|
||||
headers=headers,
|
||||
body=body,
|
||||
mount=mount,
|
||||
upstream_base=upstream_base,
|
||||
)
|
||||
request: Final = edge_request(
|
||||
method, split.path, split.query, body, _header_value(headers, "content-type")
|
||||
if profile == "stateless_v1"
|
||||
else None
|
||||
)
|
||||
if isinstance(identity, IneligibleRequest):
|
||||
return _text_reply(REPLAY_MISS_STATUS, f"stateless_v1 eligibility error: {identity.reason}")
|
||||
request: Final = (
|
||||
RecordedRequest(method=method.lower(), path=split.path, headers={}, strict_identity=identity)
|
||||
if identity is not None
|
||||
else edge_request(method, split.path, split.query, body, _header_value(headers, "content-type"))
|
||||
)
|
||||
match backend:
|
||||
case LiveEdge():
|
||||
|
|
@ -837,6 +866,14 @@ class _EdgeHandler(BaseHTTPRequestHandler):
|
|||
body: Final = self.rfile.read(length) if length else None
|
||||
if edge_server.observation is not None:
|
||||
edge_server.observation.observe(body)
|
||||
strict: Final = (
|
||||
isinstance(edge_server.backend, RecordEdge) and edge_server.backend.recorder.profile == "stateless_v1"
|
||||
or isinstance(edge_server.backend, ReplayEdge)
|
||||
and edge_server.backend.source.bundle.manifest.match_profile == "stateless_v1"
|
||||
)
|
||||
if strict and len({name.lower() for name in self.headers.keys()}) != len(self.headers):
|
||||
self._write_reply(_text_reply(REPLAY_MISS_STATUS, "stateless_v1 eligibility error: duplicate headers"))
|
||||
return
|
||||
outcome: Final = handle_edge_request(
|
||||
edge_server.backend,
|
||||
edge_server.mounts,
|
||||
|
|
@ -955,16 +992,16 @@ def start_provider_edge(
|
|||
|
||||
|
||||
@functools.lru_cache(maxsize=8)
|
||||
def _shared_recorder(root: Path) -> BundleRecorder:
|
||||
prepared = prepare_bundle(root)
|
||||
def _shared_recorder(root: Path, profile: MatchProfile = "legacy") -> BundleRecorder:
|
||||
prepared = prepare_bundle(root, profile=profile)
|
||||
if isinstance(prepared, UnsafeBundleDir):
|
||||
raise ValueError(f"E2E_FIXTURE_DIR {prepared.path} {prepared.reason}")
|
||||
return prepared
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=8)
|
||||
def _shared_replay_source(root: Path) -> ReplaySource:
|
||||
loaded = load_bundle(root)
|
||||
def _shared_replay_source(root: Path, profile: MatchProfile = "legacy") -> ReplaySource:
|
||||
loaded = load_bundle(root, profile=profile)
|
||||
if isinstance(loaded, UnreadableBundle):
|
||||
raise ValueError(f"cannot replay from {root}: {loaded.reason}")
|
||||
return ReplaySource(bundle=loaded)
|
||||
|
|
@ -977,11 +1014,12 @@ def _shared_edge(
|
|||
bind_host: str,
|
||||
advertise_host: str,
|
||||
forward_timeout: float,
|
||||
profile: MatchProfile,
|
||||
) -> ProviderEdge:
|
||||
backend: Final[EdgeBackend] = (
|
||||
RecordEdge(recorder=_shared_recorder(bundle_dir), lock=threading.Lock())
|
||||
RecordEdge(recorder=_shared_recorder(bundle_dir, profile), lock=threading.Lock())
|
||||
if mode == "record"
|
||||
else ReplayEdge(source=_shared_replay_source(bundle_dir))
|
||||
else ReplayEdge(source=_shared_replay_source(bundle_dir, profile))
|
||||
)
|
||||
return start_provider_edge(
|
||||
backend,
|
||||
|
|
@ -998,7 +1036,7 @@ def replay_leftover_error(*, mode_raw: str, bundle_dir: Path, test_key: str) ->
|
|||
recording it no longer matches. Inert in every other mode."""
|
||||
if parse_fixture_mode(mode_raw) != "replay":
|
||||
return None
|
||||
return _shared_replay_source(bundle_dir).leftover_error(test_key)
|
||||
return _shared_replay_source(bundle_dir, match_profile()).leftover_error(test_key)
|
||||
|
||||
|
||||
def provider_edge_api_base(
|
||||
|
|
@ -1021,10 +1059,10 @@ def provider_edge_api_base(
|
|||
return None
|
||||
case "record" | "replay":
|
||||
if mount not in EDGE_MOUNTS:
|
||||
raise ValueError(
|
||||
f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(EDGE_MOUNTS))}"
|
||||
)
|
||||
return _shared_edge(mode, bundle_dir, bind_host, advertise_host, forward_timeout).api_base(mount)
|
||||
raise ValueError(f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(EDGE_MOUNTS))}")
|
||||
return _shared_edge(mode, bundle_dir, bind_host, advertise_host, forward_timeout, match_profile()).api_base(
|
||||
mount
|
||||
)
|
||||
case _:
|
||||
assert_never(mode)
|
||||
|
||||
|
|
@ -1037,9 +1075,9 @@ def _observed_backend(mode_raw: str, bundle_dir: Path) -> EdgeBackend:
|
|||
case "live":
|
||||
return LiveEdge()
|
||||
case "record":
|
||||
return RecordEdge(_shared_recorder(bundle_dir), threading.Lock())
|
||||
return RecordEdge(_shared_recorder(bundle_dir, match_profile()), threading.Lock())
|
||||
case "replay":
|
||||
return ReplayEdge(_shared_replay_source(bundle_dir))
|
||||
return ReplayEdge(_shared_replay_source(bundle_dir, match_profile()))
|
||||
case _:
|
||||
assert_never(mode)
|
||||
|
||||
|
|
|
|||
|
|
@ -11,14 +11,23 @@ from __future__ import annotations
|
|||
import time
|
||||
import warnings
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from functools import reduce
|
||||
from dataclasses import dataclass, field, replace
|
||||
from datetime import datetime
|
||||
from functools import reduce
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel
|
||||
from typing import Final, Literal
|
||||
|
||||
from e2e_config import (
|
||||
CONTROL_PLANE_BASE_URL,
|
||||
MASTER_KEY,
|
||||
POLL_INTERVAL,
|
||||
POLL_TIMEOUT,
|
||||
PROXY_BASE_URL,
|
||||
PROXY_REPLICA_URLS,
|
||||
REQUEST_TIMEOUT,
|
||||
SLOW_PROVIDER_TIMEOUT_SECONDS,
|
||||
settle_propagation,
|
||||
)
|
||||
from e2e_http import (
|
||||
AnthropicHeaders,
|
||||
AuthHeaders,
|
||||
|
|
@ -55,6 +64,7 @@ from models import (
|
|||
KeyInfoParams,
|
||||
KeyInfoResponse,
|
||||
LiteLLMParamsBody,
|
||||
MemorySummaryResponse,
|
||||
ModelDeleteBody,
|
||||
ModelInfoBody,
|
||||
ModelInfoEntry,
|
||||
|
|
@ -63,7 +73,6 @@ from models import (
|
|||
ModelNewBody,
|
||||
ModelNewResponse,
|
||||
ModelsListParams,
|
||||
MemorySummaryResponse,
|
||||
ModelsListResponse,
|
||||
ModelUpdateBody,
|
||||
OcrBody,
|
||||
|
|
@ -76,23 +85,13 @@ from models import (
|
|||
TeamDeleteBody,
|
||||
TeamNewBody,
|
||||
TeamNewResponse,
|
||||
UserDeleteBody,
|
||||
UserDeleteResponse,
|
||||
ToolsetCreateBody,
|
||||
ToolsetRow,
|
||||
ToolsetUpdateBody,
|
||||
UserDeleteBody,
|
||||
UserDeleteResponse,
|
||||
)
|
||||
from e2e_config import (
|
||||
CONTROL_PLANE_BASE_URL,
|
||||
MASTER_KEY,
|
||||
POLL_INTERVAL,
|
||||
POLL_TIMEOUT,
|
||||
PROXY_BASE_URL,
|
||||
PROXY_REPLICA_URLS,
|
||||
REQUEST_TIMEOUT,
|
||||
SLOW_PROVIDER_TIMEOUT_SECONDS,
|
||||
settle_propagation,
|
||||
)
|
||||
from pydantic import BaseModel
|
||||
from transport import HttpTransport, SplitTransport, Transport, is_control_plane_path
|
||||
|
||||
RowsPredicate = Callable[[list[SpendLogRow]], bool]
|
||||
|
|
@ -421,11 +420,23 @@ def converge_timeout_message(*, what: str, replica: str, timeout: float, last_re
|
|||
)
|
||||
|
||||
|
||||
CredentialKind = Literal["master", "direct_jwt", "virtual_key", "dashboard_session"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Caller:
|
||||
credential: str = field(repr=False)
|
||||
kind: CredentialKind
|
||||
role: str
|
||||
tenant: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProxyClient:
|
||||
transport: Transport
|
||||
replicas: Mapping[str, Transport]
|
||||
control_replicas: Mapping[str, Transport]
|
||||
caller: Caller | None = None
|
||||
poll_timeout: float = 120.0
|
||||
poll_interval: float = 5.0
|
||||
model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT
|
||||
|
|
@ -433,13 +444,24 @@ class ProxyClient:
|
|||
model_servable_interval: float = MODEL_SERVABLE_INTERVAL
|
||||
model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT
|
||||
|
||||
def with_caller(self, caller: Caller) -> ProxyClient:
|
||||
return replace(self, caller=caller)
|
||||
|
||||
def management_headers(self, caller_key: str | None = None, *, transport: Transport | None = None) -> AuthHeaders:
|
||||
selected: Final = self.transport if transport is None else transport
|
||||
if caller_key is not None:
|
||||
return selected.bearer(caller_key)
|
||||
if self.caller is not None:
|
||||
return selected.bearer(self.caller.credential)
|
||||
return selected.master
|
||||
|
||||
# ---- keys / customers (satisfies lifecycle.ResourceClient) ----------
|
||||
|
||||
def generate_key(self, body: KeyGenerateBody) -> str:
|
||||
return unwrap(
|
||||
self.transport.post(
|
||||
"/key/generate",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
json=body,
|
||||
response_type=KeyGenerateResponse,
|
||||
)
|
||||
|
|
@ -448,7 +470,7 @@ class ProxyClient:
|
|||
def delete_key(self, key: str) -> None:
|
||||
_ = self.transport.post(
|
||||
"/key/delete",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
json=KeyDeleteBody(keys=[key]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -458,7 +480,7 @@ class ProxyClient:
|
|||
return
|
||||
_ = self.transport.post(
|
||||
"/customer/delete",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
json=CustomerDeleteBody(user_ids=user_ids),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -467,7 +489,7 @@ class ProxyClient:
|
|||
return unwrap(
|
||||
self.transport.get(
|
||||
"/key/info",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
params=KeyInfoParams(key=key),
|
||||
response_type=KeyInfoResponse,
|
||||
)
|
||||
|
|
@ -477,7 +499,7 @@ class ProxyClient:
|
|||
return {
|
||||
url: transport.get(
|
||||
"/debug/memory/summary",
|
||||
headers=transport.master,
|
||||
headers=self.management_headers(transport=transport),
|
||||
params=NoBody(),
|
||||
response_type=MemorySummaryResponse,
|
||||
)
|
||||
|
|
@ -524,11 +546,12 @@ class ProxyClient:
|
|||
{replica: outcome.result for replica, outcome in outcomes.items() if isinstance(outcome, Converged)}
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _body_poller[R: BaseModel](
|
||||
transport: Transport, path: str, params: BaseModel, response_type: type[R]
|
||||
self, transport: Transport, path: str, params: BaseModel, response_type: type[R]
|
||||
) -> Poller[Result[R]]:
|
||||
return lambda: transport.get(path, headers=transport.master, params=params, response_type=response_type)
|
||||
return lambda: transport.get(
|
||||
path, headers=self.management_headers(transport=transport), params=params, response_type=response_type
|
||||
)
|
||||
|
||||
def model_info(self) -> list[ModelInfoEntry]:
|
||||
"""Every configured deployment with the price the proxy resolved for it
|
||||
|
|
@ -536,7 +559,7 @@ class ProxyClient:
|
|||
return unwrap(
|
||||
self.transport.get(
|
||||
"/model/info",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
params=NoBody(),
|
||||
response_type=ModelInfoResponse,
|
||||
)
|
||||
|
|
@ -546,7 +569,7 @@ class ProxyClient:
|
|||
return unwrap(
|
||||
self.transport.get(
|
||||
"/public/litellm_model_cost_map",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
params=NoBody(),
|
||||
response_type=CostMap,
|
||||
)
|
||||
|
|
@ -607,7 +630,7 @@ class ProxyClient:
|
|||
model_id = unwrap(
|
||||
self.transport.post(
|
||||
"/model/new",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
json=body,
|
||||
response_type=ModelNewResponse,
|
||||
)
|
||||
|
|
@ -623,7 +646,7 @@ class ProxyClient:
|
|||
|
||||
def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None:
|
||||
"""Block until every replica lists `model_name`, or fail at model_servable_timeout."""
|
||||
headers: Final = self.transport.master if listed_for is None else self.transport.bearer(listed_for)
|
||||
headers: Final = self.management_headers(listed_for)
|
||||
outcome: Final = await_servable_everywhere(
|
||||
{url: self._models_poller(transport, headers) for url, transport in self.replicas.items()},
|
||||
model_name=model_name,
|
||||
|
|
@ -666,7 +689,7 @@ class ProxyClient:
|
|||
unwrap(
|
||||
self.transport.post(
|
||||
"/model/update",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
json=ModelUpdateBody(
|
||||
litellm_params=litellm_params,
|
||||
model_info=ModelInfoBody(id=model_id),
|
||||
|
|
@ -678,7 +701,7 @@ class ProxyClient:
|
|||
def delete_model(self, model_id: str) -> None:
|
||||
result = self.transport.post(
|
||||
"/model/delete",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
json=ModelDeleteBody(id=model_id),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -747,11 +770,10 @@ class ProxyClient:
|
|||
f"GET {path} on {replica} still answers {self.poll_timeout}s after the delete; last read: {last}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]:
|
||||
def _reader[R: BaseModel](self, transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]:
|
||||
return lambda request_timeout: transport.get(
|
||||
path,
|
||||
headers=transport.master,
|
||||
headers=self.management_headers(transport=transport),
|
||||
params=NoBody(),
|
||||
response_type=response_type,
|
||||
timeout=request_timeout,
|
||||
|
|
@ -763,7 +785,7 @@ class ProxyClient:
|
|||
return unwrap(
|
||||
self.transport.post(
|
||||
"/v1/mcp/toolset",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
json=body,
|
||||
response_type=ToolsetRow,
|
||||
)
|
||||
|
|
@ -775,7 +797,7 @@ class ProxyClient:
|
|||
return unwrap(
|
||||
self.transport.put(
|
||||
"/v1/mcp/toolset",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
json=body,
|
||||
response_type=ToolsetRow,
|
||||
)
|
||||
|
|
@ -786,7 +808,7 @@ class ProxyClient:
|
|||
can unwrap it while a deferred teardown can ignore an already-deleted row."""
|
||||
return self.transport.delete(
|
||||
f"/v1/mcp/toolset/{toolset_id}",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
json=NoBody(),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -795,7 +817,7 @@ class ProxyClient:
|
|||
unwrap(
|
||||
self.transport.post(
|
||||
"/credentials",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
json=body,
|
||||
response_type=CredentialCreateResponse,
|
||||
)
|
||||
|
|
@ -804,7 +826,7 @@ class ProxyClient:
|
|||
def delete_credential(self, credential_name: str) -> None:
|
||||
result = self.transport.delete(
|
||||
f"/credentials/{credential_name}",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
json=NoBody(),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -815,7 +837,7 @@ class ProxyClient:
|
|||
return unwrap(
|
||||
self.transport.post(
|
||||
"/team/new",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
json=body,
|
||||
response_type=TeamNewResponse,
|
||||
)
|
||||
|
|
@ -824,7 +846,7 @@ class ProxyClient:
|
|||
def delete_team(self, team_id: str) -> None:
|
||||
result = self.transport.post(
|
||||
"/team/delete",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
json=TeamDeleteBody(team_ids=[team_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
|
@ -836,7 +858,7 @@ class ProxyClient:
|
|||
a user the proxy only upserts after a successful auth."""
|
||||
result = self.transport.post(
|
||||
"/user/delete",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
json=UserDeleteBody(user_ids=[user_id]),
|
||||
response_type=UserDeleteResponse,
|
||||
)
|
||||
|
|
@ -909,7 +931,7 @@ class ProxyClient:
|
|||
def spend_logs(self, params: SpendLogsParams) -> list[SpendLogRow]:
|
||||
result = self.transport.get(
|
||||
"/spend/logs",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
params=params,
|
||||
response_type=SpendLogs,
|
||||
)
|
||||
|
|
@ -924,7 +946,7 @@ class ProxyClient:
|
|||
return unwrap(
|
||||
self.transport.get(
|
||||
"/spend/logs/v2",
|
||||
headers=self.transport.master,
|
||||
headers=self.management_headers(),
|
||||
params=SpendLogsPageParams(
|
||||
start_date=start.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
end_date=end.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
|
|
@ -977,7 +999,7 @@ class ProxyClient:
|
|||
# ---- route probe ----------------------------------------------------
|
||||
|
||||
def probe(self, path: str, *, params: NoBody) -> ProbeResult:
|
||||
return self.transport.probe(path, params=params)
|
||||
return self.transport.probe(path, params=params, headers=self.management_headers())
|
||||
|
||||
|
||||
def build_proxy_client(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,113 @@
|
|||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from math import isclose
|
||||
from typing import Final
|
||||
|
||||
from e2e_config import provider_edge_base, unique_marker
|
||||
from e2e_http import unwrap
|
||||
from lifecycle import ResourceManager
|
||||
from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody, LiteLLMParamsBody, TeamNewBody
|
||||
from spend_e2e_client import SpendClient
|
||||
|
||||
INPUT_RATE: Final = 0.00004
|
||||
OUTPUT_RATE: Final = 0.00008
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TeamTraffic:
|
||||
team_id: str
|
||||
key: str
|
||||
responses: tuple[ChatResponse, ...]
|
||||
|
||||
@property
|
||||
def prompt_tokens(self) -> int:
|
||||
return sum(response.usage.prompt_tokens or 0 for response in self.responses if response.usage)
|
||||
|
||||
@property
|
||||
def completion_tokens(self) -> int:
|
||||
return sum(response.usage.completion_tokens or 0 for response in self.responses if response.usage)
|
||||
|
||||
@property
|
||||
def spend(self) -> float:
|
||||
return self.prompt_tokens * INPUT_RATE + self.completion_tokens * OUTPUT_RATE
|
||||
|
||||
|
||||
def create_traffic(client: SpendClient, resources: ResourceManager) -> tuple[TeamTraffic, ...]:
|
||||
base: Final = provider_edge_base("openai")
|
||||
model: Final = f"e2e-reconciliation-{unique_marker()}"
|
||||
model_id: Final = client.proxy.create_model(
|
||||
model,
|
||||
LiteLLMParamsBody(
|
||||
model="openai/gpt-5.6-luna",
|
||||
api_key="os.environ/OPENAI_API_KEY",
|
||||
api_base=None if base is None else f"{base}/v1",
|
||||
input_cost_per_token=INPUT_RATE,
|
||||
output_cost_per_token=OUTPUT_RATE,
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: client.proxy.delete_model(model_id))
|
||||
|
||||
def team_traffic() -> TeamTraffic:
|
||||
team: Final = client.proxy.create_team(TeamNewBody(team_alias=f"e2e-spend-{unique_marker()}"))
|
||||
resources.defer(lambda: client.proxy.delete_team(team))
|
||||
key: Final = client.proxy.generate_key(KeyGenerateBody(team_id=team, models=[model]))
|
||||
resources.defer(lambda: client.proxy.delete_key(key))
|
||||
|
||||
prompts: Final = tuple(f"Reply with one word. {index} {unique_marker()}" for index in range(7))
|
||||
|
||||
def call(index: int) -> ChatResponse:
|
||||
response: Final = unwrap(
|
||||
client.proxy.chat(
|
||||
key,
|
||||
ChatBody(
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=prompts[index])],
|
||||
max_completion_tokens=128,
|
||||
),
|
||||
)
|
||||
)
|
||||
assert response.id, "successful response must have an ID"
|
||||
assert response.usage is not None, "successful response must have usage"
|
||||
assert response.usage.prompt_tokens is not None and response.usage.prompt_tokens > 0
|
||||
assert response.usage.completion_tokens is not None and response.usage.completion_tokens > 0
|
||||
assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens
|
||||
assert not response.usage.cache_creation_input_tokens
|
||||
assert not response.usage.cache_read_input_tokens
|
||||
assert not response.usage.prompt_tokens_details or not response.usage.prompt_tokens_details.cached_tokens
|
||||
return response
|
||||
|
||||
sequential: Final = call(0)
|
||||
with ThreadPoolExecutor(max_workers=6) as pool:
|
||||
concurrent: Final = tuple(pool.map(call, range(1, 7)))
|
||||
return TeamTraffic(team, key, (sequential, *concurrent))
|
||||
|
||||
return tuple(team_traffic() for _ in range(2))
|
||||
|
||||
|
||||
def assert_logs_match(client: SpendClient, traffic: TeamTraffic) -> None:
|
||||
expected_ids: Final = frozenset(response.id for response in traffic.responses)
|
||||
assert len(expected_ids) == len(traffic.responses), "responses must have distinct IDs"
|
||||
rows: Final = client.poll_logs_for_key(
|
||||
traffic.key,
|
||||
min_rows=len(traffic.responses),
|
||||
predicate=lambda values: frozenset(row.request_id for row in values) == expected_ids,
|
||||
)
|
||||
assert frozenset(row.request_id for row in rows) == expected_ids, "stored IDs must equal returned response IDs"
|
||||
assert len(rows) == len(traffic.responses), "expected exactly one scoped spend row per response"
|
||||
by_id: Final = {row.request_id: row for row in rows}
|
||||
|
||||
def assert_response(response: ChatResponse) -> None:
|
||||
row: Final = by_id[response.id]
|
||||
usage: Final = response.usage
|
||||
assert usage is not None and usage.prompt_tokens is not None and usage.completion_tokens is not None
|
||||
assert row.team_id == traffic.team_id
|
||||
assert row.status == "success"
|
||||
assert row.cache_hit != "True"
|
||||
assert row.prompt_tokens == usage.prompt_tokens
|
||||
assert row.completion_tokens == usage.completion_tokens
|
||||
assert row.total_tokens == usage.total_tokens
|
||||
expected_cost: Final = usage.prompt_tokens * INPUT_RATE + usage.completion_tokens * OUTPUT_RATE
|
||||
assert row.spend is not None and isclose(row.spend, expected_cost, rel_tol=1e-6, abs_tol=1e-9)
|
||||
|
||||
for response in traffic.responses:
|
||||
assert_response(response)
|
||||
|
|
@ -17,13 +17,13 @@ fails the test; a pricing or token-count drift does not.
|
|||
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from math import isclose
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from e2e_http import Result, Success
|
||||
from e2e_http import Success
|
||||
from lifecycle import ResourceManager
|
||||
from models import ChatResponse, LiteLLMParamsBody, SpendLogs, SpendLogsParams
|
||||
from models import LiteLLMParamsBody, SpendLogs, SpendLogsParams
|
||||
from spend_e2e_client import SpendClient, SpendLogRow, is_ok, unique_marker, unwrap
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
|
@ -280,51 +280,22 @@ def test_key_spend_equals_sum_of_logs(client: SpendClient, scoped_key: str) -> N
|
|||
), f"key aggregate {key_spend} != sum of logs {logs_total}; rows: {_summarize(rows)}"
|
||||
|
||||
|
||||
@pytest.mark.replayable
|
||||
@pytest.mark.covers("quota_management.spend_tracking.concurrent_burst.loses_no_spend")
|
||||
def test_burst_of_concurrent_calls_loses_no_spend(
|
||||
client: SpendClient, scoped_key: str
|
||||
client: SpendClient, resources: ResourceManager
|
||||
) -> None:
|
||||
"""Six concurrent calls on one key: every call lands its own spend row under a
|
||||
distinct request_id and the key aggregate equals the sum of the rows.
|
||||
Sequential accuracy is covered by test_key_spend_equals_sum_of_logs; this pins
|
||||
the concurrent increment path (parallel writers racing on one key's counter),
|
||||
where a lost update can never be reproduced by sequential calls."""
|
||||
burst = 6
|
||||
from spend_reconciliation import TeamTraffic, assert_logs_match, create_traffic
|
||||
|
||||
def call(idx: int) -> Result[ChatResponse]:
|
||||
return client.chat(
|
||||
scoped_key,
|
||||
"gemini-2.5-flash",
|
||||
f"burst call {idx} {unique_marker()}",
|
||||
max_tokens=16,
|
||||
)
|
||||
traffic: Final = create_traffic(client, resources)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=burst) as pool:
|
||||
results = tuple(pool.map(call, range(burst)))
|
||||
failed = [r for r in results if not is_ok(r)]
|
||||
assert not failed, f"{len(failed)}/{burst} burst calls failed; first: {failed[0]}"
|
||||
def assert_team(team: TeamTraffic) -> None:
|
||||
assert_logs_match(client, team)
|
||||
key_spend: Final = client.poll_key_spend(team.key, minimum=team.spend * 0.999999)
|
||||
assert isclose(key_spend, team.spend, rel_tol=1e-6, abs_tol=1e-9)
|
||||
|
||||
rows = client.poll_logs_for_key(
|
||||
scoped_key,
|
||||
min_rows=burst,
|
||||
predicate=lambda rs: len([r for r in rs if (r.spend or 0) > 0]) >= burst,
|
||||
)
|
||||
costed = [r for r in rows if (r.spend or 0) > 0]
|
||||
assert len(costed) >= burst, (
|
||||
f"only {len(costed)}/{burst} burst calls produced a costed row - "
|
||||
f"rows lost under concurrency: {_summarize(rows)}"
|
||||
)
|
||||
request_ids = [r.request_id for r in costed]
|
||||
assert len(set(request_ids)) == len(request_ids), (
|
||||
f"concurrent rows collapsed onto shared request_ids: {_summarize(rows)}"
|
||||
)
|
||||
|
||||
logs_total = sum((r.spend or 0) for r in rows)
|
||||
key_spend = client.poll_key_spend(scoped_key, minimum=logs_total * 0.999)
|
||||
assert _approx_equal(key_spend, logs_total), (
|
||||
f"key aggregate {key_spend} != sum of {len(rows)} rows {logs_total} - "
|
||||
f"spend increments lost under concurrency: {_summarize(rows)}"
|
||||
)
|
||||
for team in traffic:
|
||||
assert_team(team)
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.spend_tracking.pagination.keeps_total")
|
||||
|
|
|
|||
|
|
@ -7,13 +7,18 @@ missing start/end dates are rejected.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from math import isclose
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from e2e_http import ProbeResult
|
||||
from models import DateRangeParams
|
||||
from lifecycle import ResourceManager
|
||||
from proxy_client import Converged, await_converged
|
||||
from pydantic import BaseModel
|
||||
from spend_e2e_client import SpendClient
|
||||
from spend_reconciliation import TeamTraffic, assert_logs_match, create_traffic
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
|
|
@ -24,22 +29,45 @@ class TeamDailyActivityParams(BaseModel):
|
|||
start_date: str | None = None
|
||||
end_date: str | None = None
|
||||
page: int = 1
|
||||
page_size: int = 1
|
||||
team_ids: str | None = None
|
||||
|
||||
|
||||
class TeamDailyActivityRow(BaseModel):
|
||||
date: str
|
||||
metrics: TeamDailyActivityMetrics
|
||||
breakdown: TeamDailyActivityBreakdown
|
||||
|
||||
|
||||
class TeamDailyActivityMetrics(BaseModel):
|
||||
spend: float
|
||||
total_tokens: int
|
||||
prompt_tokens: int
|
||||
completion_tokens: int
|
||||
api_requests: int
|
||||
successful_requests: int
|
||||
failed_requests: int
|
||||
|
||||
|
||||
class TeamDailyActivityEntity(BaseModel):
|
||||
metrics: TeamDailyActivityMetrics
|
||||
|
||||
|
||||
class TeamDailyActivityBreakdown(BaseModel):
|
||||
entities: dict[str, TeamDailyActivityEntity]
|
||||
|
||||
|
||||
class TeamDailyActivityMetadata(BaseModel):
|
||||
page: int
|
||||
total_pages: int
|
||||
has_more: bool
|
||||
total_spend: float
|
||||
total_prompt_tokens: int
|
||||
total_completion_tokens: int
|
||||
total_tokens: int
|
||||
total_api_requests: int
|
||||
total_successful_requests: int
|
||||
total_failed_requests: int
|
||||
|
||||
|
||||
class TeamDailyActivityResponse(BaseModel):
|
||||
|
|
@ -47,32 +75,128 @@ class TeamDailyActivityResponse(BaseModel):
|
|||
metadata: TeamDailyActivityMetadata
|
||||
|
||||
|
||||
def _range_days(days: int) -> DateRangeParams:
|
||||
end = datetime.now(timezone.utc).date()
|
||||
start = end - timedelta(days=days)
|
||||
return DateRangeParams(start_date=start.isoformat(), end_date=end.isoformat())
|
||||
|
||||
|
||||
def _probe(client: SpendClient, params: BaseModel) -> ProbeResult:
|
||||
return client.proxy.transport.probe(ROUTE, params=params)
|
||||
|
||||
|
||||
class TestTeamDailyActivity:
|
||||
@pytest.mark.replayable
|
||||
@pytest.mark.covers("mgmt.team.daily_activity.happy_path")
|
||||
@pytest.mark.parametrize("days", [1, 7, 30])
|
||||
def test_valid_date_range_returns_results_and_metadata(self, client: SpendClient, days: int) -> None:
|
||||
result = _probe(client, _range_days(days))
|
||||
assert result.status_code == 200, (
|
||||
f"{ROUTE} range={days}d must be 200, got {result.status_code}: {result.body[:600]}"
|
||||
def test_valid_date_range_returns_results_and_metadata(
|
||||
self, client: SpendClient, resources: ResourceManager
|
||||
) -> None:
|
||||
started: Final = datetime.now(timezone.utc).date()
|
||||
traffic: Final = create_traffic(client, resources)
|
||||
for team in traffic:
|
||||
assert_logs_match(client, team)
|
||||
ended: Final = datetime.now(timezone.utc).date()
|
||||
team_ids: Final = ",".join(team.team_id for team in traffic)
|
||||
|
||||
def fetch(
|
||||
page: int, start: str = (started - timedelta(days=1)).isoformat(), end: str = ended.isoformat()
|
||||
) -> TeamDailyActivityResponse:
|
||||
result: Final = _probe(
|
||||
client,
|
||||
TeamDailyActivityParams(
|
||||
start_date=start,
|
||||
end_date=end,
|
||||
page=page,
|
||||
page_size=1,
|
||||
team_ids=team_ids,
|
||||
),
|
||||
)
|
||||
assert result.status_code == 200, f"daily activity failed: {result.status_code} {result.body[:300]}"
|
||||
return TeamDailyActivityResponse.model_validate_json(result.body)
|
||||
|
||||
def pages() -> tuple[TeamDailyActivityResponse, ...]:
|
||||
first: Final = fetch(1)
|
||||
assert first.metadata.total_pages <= len(traffic) * 2, "unexpected extra scoped daily groups"
|
||||
return (first, *(fetch(page) for page in range(2, first.metadata.total_pages + 1)))
|
||||
|
||||
outcome: Final = await_converged(
|
||||
pages,
|
||||
converged=lambda values: (
|
||||
sum(page.metadata.total_api_requests for page in values) >= sum(len(team.responses) for team in traffic)
|
||||
),
|
||||
timeout=client.proxy.poll_timeout,
|
||||
interval=client.proxy.poll_interval,
|
||||
now=time.monotonic,
|
||||
sleep=time.sleep,
|
||||
)
|
||||
parsed = TeamDailyActivityResponse.model_validate_json(result.body)
|
||||
assert parsed.metadata.page == 1
|
||||
assert parsed.metadata.total_pages >= 1
|
||||
if parsed.results:
|
||||
first = parsed.results[0]
|
||||
assert first.date
|
||||
assert first.metrics.spend >= 0
|
||||
assert first.metrics.total_tokens >= 0
|
||||
observed: Final = outcome.result if isinstance(outcome, Converged) else outcome.last_result
|
||||
assert observed is not None, "daily aggregation must return a response before the deadline"
|
||||
|
||||
assert len(observed) >= 2, "two teams must exercise a page boundary"
|
||||
|
||||
def assert_page(index: int, page: TeamDailyActivityResponse) -> None:
|
||||
assert page.metadata.page == index
|
||||
assert page.metadata.total_pages == len(observed)
|
||||
assert page.metadata.has_more == (index < len(observed))
|
||||
assert len(page.results) == 1, "each fetched daily group must appear in results"
|
||||
row: Final = page.results[0]
|
||||
assert started <= datetime.fromisoformat(row.date).date() <= ended
|
||||
assert len(row.breakdown.entities) == 1
|
||||
assert row.metrics.total_tokens == page.metadata.total_tokens
|
||||
assert row.metrics.prompt_tokens == page.metadata.total_prompt_tokens
|
||||
assert row.metrics.completion_tokens == page.metadata.total_completion_tokens
|
||||
assert row.metrics.api_requests == page.metadata.total_api_requests
|
||||
assert row.metrics.successful_requests == page.metadata.total_successful_requests
|
||||
assert row.metrics.failed_requests == page.metadata.total_failed_requests
|
||||
assert isclose(row.metrics.spend, page.metadata.total_spend, rel_tol=1e-6, abs_tol=1e-9)
|
||||
|
||||
for index, page in enumerate(observed, 1):
|
||||
assert_page(index, page)
|
||||
|
||||
entities: Final = tuple(
|
||||
(team_id, entity.metrics)
|
||||
for page in observed
|
||||
for row in page.results
|
||||
for team_id, entity in row.breakdown.entities.items()
|
||||
)
|
||||
assert frozenset(team_id for team_id, _ in entities) == frozenset(team.team_id for team in traffic)
|
||||
|
||||
def assert_team(team: TeamTraffic) -> None:
|
||||
metrics: Final = tuple(metrics for team_id, metrics in entities if team_id == team.team_id)
|
||||
assert sum(m.api_requests for m in metrics) == len(team.responses)
|
||||
assert sum(m.successful_requests for m in metrics) == len(team.responses)
|
||||
assert sum(m.failed_requests for m in metrics) == 0
|
||||
assert sum(m.prompt_tokens for m in metrics) == team.prompt_tokens
|
||||
assert sum(m.completion_tokens for m in metrics) == team.completion_tokens
|
||||
assert sum(m.total_tokens for m in metrics) == team.prompt_tokens + team.completion_tokens
|
||||
assert isclose(sum(m.spend for m in metrics), team.spend, rel_tol=1e-6, abs_tol=1e-9)
|
||||
|
||||
for team in traffic:
|
||||
assert_team(team)
|
||||
|
||||
assert isclose(
|
||||
sum(page.metadata.total_spend for page in observed),
|
||||
sum(team.spend for team in traffic),
|
||||
rel_tol=1e-6,
|
||||
abs_tol=1e-9,
|
||||
)
|
||||
assert sum(page.metadata.total_tokens for page in observed) == sum(
|
||||
team.prompt_tokens + team.completion_tokens for team in traffic
|
||||
)
|
||||
|
||||
for days in (7, 30):
|
||||
assert (
|
||||
tuple(fetch(page, (started - timedelta(days=days)).isoformat()) for page in range(1, len(observed) + 1))
|
||||
== observed
|
||||
), f"{days}-day activity must preserve the same isolated groups and totals"
|
||||
|
||||
empty_date: Final = (started - timedelta(days=7)).isoformat()
|
||||
empty: Final = fetch(1, empty_date, empty_date)
|
||||
assert empty.results == []
|
||||
assert empty.metadata.total_pages == 0
|
||||
assert empty.metadata.page == 1
|
||||
assert not empty.metadata.has_more
|
||||
assert empty.metadata.total_spend == 0
|
||||
assert empty.metadata.total_tokens == 0
|
||||
assert empty.metadata.total_api_requests == 0
|
||||
assert empty.metadata.total_prompt_tokens == 0
|
||||
assert empty.metadata.total_completion_tokens == 0
|
||||
assert empty.metadata.total_successful_requests == 0
|
||||
assert empty.metadata.total_failed_requests == 0
|
||||
|
||||
@pytest.mark.covers("mgmt.team.daily_activity.missing_start_date_rejected")
|
||||
def test_missing_start_date_is_rejected(self, client: SpendClient) -> None:
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ from e2e_http import (
|
|||
request_with_retry,
|
||||
streaming_outcome,
|
||||
wire_body,
|
||||
without_retries,
|
||||
)
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
|
|
@ -56,6 +57,15 @@ def _issue_from(responses: Sequence[FakeResponse]) -> Callable[[], FakeResponse]
|
|||
|
||||
|
||||
class TestTransientRetryPolicy:
|
||||
def test_qualification_disables_retries_and_restores_the_default(self) -> None:
|
||||
responses: Final = (FakeResponse(529), FakeResponse(200))
|
||||
sleep: Final = SleepRecorder()
|
||||
with without_retries():
|
||||
assert request_with_retry(_issue_from(responses), sleep=sleep) is responses[0]
|
||||
assert sleep.delays == ()
|
||||
assert request_with_retry(_issue_from(responses), sleep=sleep) is responses[1]
|
||||
assert sleep.delays == (0.5,)
|
||||
|
||||
def test_transient_set_is_only_statuses_the_proxy_cannot_emit(self) -> None:
|
||||
assert TRANSIENT_STATUSES == frozenset({529})
|
||||
assert 429 not in TRANSIENT_STATUSES
|
||||
|
|
|
|||
|
|
@ -4,12 +4,20 @@ these carry no `e2e` marker and run everywhere."""
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from builtins import ExceptionGroup
|
||||
from collections.abc import Callable, Generator
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from threading import Thread
|
||||
from typing import Final
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
from e2e_http import ExternalWrite
|
||||
|
|
@ -18,6 +26,8 @@ from idp import (
|
|||
KEYCLOAK_ADMIN_USER_ENV,
|
||||
KEYCLOAK_REALM_ENV,
|
||||
KEYCLOAK_URL_ENV,
|
||||
BrowserClientBody,
|
||||
Discovery,
|
||||
Keycloak,
|
||||
PasswordCredential,
|
||||
UserCreateBody,
|
||||
|
|
@ -60,24 +70,48 @@ def _idp_server(
|
|||
) -> Generator[tuple[Keycloak, SimpleQueue[str]]]:
|
||||
"""Exercise provisioning failures through the same HTTP transport as live tests."""
|
||||
deletions: SimpleQueue[str] = SimpleQueue()
|
||||
clients: SimpleQueue[BrowserClientBody] = SimpleQueue()
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
pass
|
||||
|
||||
def do_POST(self) -> None:
|
||||
self.rfile.read(int(self.headers.get("Content-Length", "0")))
|
||||
body: Final = self.rfile.read(int(self.headers.get("Content-Length", "0")))
|
||||
if self.path.endswith("/token"):
|
||||
self.send_response(admin_status)
|
||||
self.end_headers()
|
||||
self.wfile.write(b'{"access_token":"synthetic-harness-token"}')
|
||||
else:
|
||||
if self.path.endswith("/clients"):
|
||||
clients.put(BrowserClientBody.model_validate_json(body))
|
||||
self.send_response(user_status if self.path.endswith("/users") else 201)
|
||||
self.send_header("Location", f"{self.path}/resource-1")
|
||||
self.end_headers()
|
||||
if user_status != 201 and self.path.endswith("/users"):
|
||||
self.wfile.write(b"injected create failure")
|
||||
|
||||
def do_GET(self) -> None:
|
||||
self.send_response(200)
|
||||
self.end_headers()
|
||||
if "/clients/" in self.path:
|
||||
client: Final = clients.get_nowait()
|
||||
clients.put(client)
|
||||
self.wfile.write(client.model_dump_json(by_alias=True).encode())
|
||||
else:
|
||||
issuer: Final = f"http://127.0.0.1:{server.server_port}/realms/test"
|
||||
self.wfile.write(
|
||||
Discovery(
|
||||
issuer=issuer,
|
||||
authorization_endpoint=f"{issuer}/auth",
|
||||
token_endpoint=f"{issuer}/token",
|
||||
userinfo_endpoint=f"{issuer}/userinfo",
|
||||
jwks_uri=f"{issuer}/certs",
|
||||
)
|
||||
.model_dump_json()
|
||||
.encode()
|
||||
)
|
||||
|
||||
def do_DELETE(self) -> None:
|
||||
deletions.put(self.path)
|
||||
self.send_response(delete_status)
|
||||
|
|
@ -115,6 +149,68 @@ def test_partial_provisioning_removes_the_group_when_user_creation_fails() -> No
|
|||
assert deletions.empty()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("exit_mode", "ignore_termination"), (("normal", False), ("parent", False), ("group", False), ("parent", True))
|
||||
)
|
||||
def test_oidc_launcher_removes_client_on_exit_and_termination(
|
||||
tmp_path: Path, exit_mode: Literal["normal", "parent", "group"], ignore_termination: bool
|
||||
) -> None:
|
||||
ready: Final = tmp_path / "ready"
|
||||
descendant_command: Final = (
|
||||
"import signal,socket,time; from pathlib import Path; "
|
||||
+ ("signal.signal(signal.SIGTERM, signal.SIG_IGN); " if ignore_termination else "")
|
||||
+ "listener=socket.socket(); listener.bind(('127.0.0.1',0)); listener.listen(); "
|
||||
f"Path({str(ready)!r}).write_text(str(listener.getsockname()[1])); time.sleep(120)"
|
||||
)
|
||||
child_command: Final = (
|
||||
"import os,subprocess,sys,time; from pathlib import Path; "
|
||||
'assert os.environ["GENERIC_CLIENT_SECRET"]; '
|
||||
'assert os.environ["GENERIC_CLIENT_USE_PKCE"] == "true"; '
|
||||
f"subprocess.Popen([sys.executable, '-c', {descendant_command!r}]); "
|
||||
f"ready=Path({str(ready)!r})\n"
|
||||
"while not ready.exists(): time.sleep(0.05)\n"
|
||||
+ ("raise SystemExit(7)" if exit_mode == "normal" else "time.sleep(120)")
|
||||
)
|
||||
with _idp_server() as (idp, deletions):
|
||||
with subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
str(Path(__file__).with_name("idp.py")),
|
||||
"http://127.0.0.1:9999",
|
||||
sys.executable,
|
||||
"-c",
|
||||
child_command,
|
||||
],
|
||||
env={
|
||||
**os.environ,
|
||||
KEYCLOAK_URL_ENV: idp.base_url,
|
||||
KEYCLOAK_REALM_ENV: idp.realm,
|
||||
KEYCLOAK_ADMIN_USER_ENV: idp.admin_username,
|
||||
KEYCLOAK_ADMIN_PASSWORD_ENV: idp.admin_password,
|
||||
},
|
||||
start_new_session=True,
|
||||
) as process:
|
||||
try:
|
||||
deadline: Final = time.monotonic() + 15
|
||||
while not ready.exists() and time.monotonic() < deadline and process.poll() is None:
|
||||
time.sleep(0.05)
|
||||
assert ready.exists(), "OIDC child did not start"
|
||||
if exit_mode == "parent":
|
||||
process.terminate()
|
||||
elif exit_mode == "group":
|
||||
os.killpg(process.pid, signal.SIGTERM)
|
||||
assert process.wait(timeout=15) == (7 if exit_mode == "normal" else 143)
|
||||
with socket.socket() as connection:
|
||||
connection.settimeout(1)
|
||||
assert connection.connect_ex(("127.0.0.1", int(ready.read_text()))) != 0
|
||||
finally:
|
||||
if process.poll() is None:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
process.wait(timeout=5)
|
||||
assert deletions.get(timeout=5) == "/admin/realms/test/clients/resource-1"
|
||||
assert deletions.empty()
|
||||
|
||||
|
||||
def test_successful_provisioning_cleans_up_user_before_group() -> None:
|
||||
with _idp_server() as (idp, deletions):
|
||||
with ExitStack() as cleanup:
|
||||
|
|
@ -134,6 +230,43 @@ def test_cleanup_failure_is_visible() -> None:
|
|||
idp.delete_group("group")
|
||||
|
||||
|
||||
def test_strict_cleanup_reports_each_failure_and_continues() -> None:
|
||||
from lifecycle import ResourceManager
|
||||
from proxy_client import build_proxy_client
|
||||
|
||||
with _idp_server(delete_status=500) as (idp, deletions):
|
||||
resources: Final = ResourceManager(client=build_proxy_client(), strict_cleanup=True)
|
||||
strict: Final = idp.with_strict_cleanup()
|
||||
resources.defer(lambda: strict.delete_group("group"))
|
||||
resources.defer(lambda: strict.delete_user("user"))
|
||||
with pytest.raises(ExceptionGroup, match="Resource cleanup failed") as error:
|
||||
resources.teardown()
|
||||
assert len(error.value.exceptions) == 2
|
||||
assert deletions.get_nowait() == "/admin/realms/test/users/user"
|
||||
assert deletions.get_nowait() == "/admin/realms/test/groups/group"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("groups", ((), ("one",), ("one", "two")))
|
||||
def test_provisioning_records_zero_one_or_multiple_groups(groups: tuple[str, ...]) -> None:
|
||||
with _idp_server() as (idp, deletions):
|
||||
with ExitStack() as cleanup:
|
||||
|
||||
def defer(callback: Callable[[], object]) -> None:
|
||||
cleanup.callback(callback)
|
||||
|
||||
identity: Final = idp.provision_groups(
|
||||
marker="memberships",
|
||||
groups=groups,
|
||||
defer=defer,
|
||||
)
|
||||
assert identity.groups == groups
|
||||
assert len(identity.group_ids) == len(groups)
|
||||
assert deletions.get_nowait() == "/admin/realms/test/users/resource-1"
|
||||
for _ in groups:
|
||||
assert deletions.get_nowait() == "/admin/realms/test/groups/resource-1"
|
||||
assert deletions.empty()
|
||||
|
||||
|
||||
def test_expired_admin_credentials_do_not_abort_remaining_cleanups() -> None:
|
||||
with _idp_server(admin_status=401) as (idp, _):
|
||||
cleanup: Final = ExitStack()
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ from concurrent.futures import ThreadPoolExecutor
|
|||
from contextlib import contextmanager
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
|
@ -81,9 +82,14 @@ def json_object(body: bytes) -> dict[str, object]:
|
|||
class _FakeProvider(ThreadingHTTPServer):
|
||||
daemon_threads = True
|
||||
|
||||
def __init__(self, bind: tuple[str, int]) -> None:
|
||||
def __init__(self, bind: tuple[str, int], *, echo_request: bool = True) -> None:
|
||||
super().__init__(bind, _FakeProviderHandler)
|
||||
self.hits: list[str] = []
|
||||
self.echo_request = echo_request
|
||||
self.requests: tuple[tuple[Mapping[str, str], bytes], ...] = ()
|
||||
|
||||
def capture_request(self, headers: Mapping[str, str], body: bytes) -> None:
|
||||
self.requests = (*self.requests, (MappingProxyType(dict(headers)), body))
|
||||
|
||||
|
||||
class _FakeProviderHandler(BaseHTTPRequestHandler):
|
||||
|
|
@ -101,8 +107,11 @@ class _FakeProviderHandler(BaseHTTPRequestHandler):
|
|||
length = int(self.headers.get("content-length") or "0")
|
||||
body = self.rfile.read(length) if length else b""
|
||||
provider.hits.append(f"{self.command} {self.path}")
|
||||
payload = json.dumps(
|
||||
provider.capture_request(dict(self.headers.items()), body)
|
||||
payload: Final = json.dumps(
|
||||
{"echo": body.decode("utf-8"), "path": self.path, "hit": len(provider.hits)}
|
||||
if provider.echo_request
|
||||
else {"ok": True}
|
||||
).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("content-type", "application/json")
|
||||
|
|
@ -117,8 +126,8 @@ class _FakeProviderHandler(BaseHTTPRequestHandler):
|
|||
|
||||
|
||||
@contextmanager
|
||||
def fake_provider() -> Generator[_FakeProvider]:
|
||||
server = _FakeProvider(("127.0.0.1", 0))
|
||||
def fake_provider(*, echo_request: bool = True) -> Generator[_FakeProvider]:
|
||||
server = _FakeProvider(("127.0.0.1", 0), echo_request=echo_request)
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -11,19 +11,53 @@ injected clock, so nothing here monkeypatches anything.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable, Mapping
|
||||
import json
|
||||
from builtins import ExceptionGroup
|
||||
from collections.abc import Callable, Generator, Iterable, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from itertools import chain, repeat
|
||||
from queue import SimpleQueue
|
||||
from threading import Thread
|
||||
from types import MappingProxyType
|
||||
from typing import Final, cast
|
||||
|
||||
import pytest
|
||||
from e2e_config import parse_replica_urls
|
||||
from e2e_http import Result, Success
|
||||
from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse
|
||||
from e2e_http import NoBody, Result, Success, without_retries
|
||||
from idp import Keycloak
|
||||
from lifecycle import ResourceManager
|
||||
from management.jwt_actors import ActorFactory
|
||||
from management.management_client import ManagementClient
|
||||
from models import (
|
||||
ConnectionTestBody,
|
||||
CredentialCreateBody,
|
||||
KeyGenerateBody,
|
||||
KeyInfo,
|
||||
KeyInfoResponse,
|
||||
KeyUpdateBody,
|
||||
LiteLLMParamsBody,
|
||||
McpServerCreateBody,
|
||||
McpServerUpdateBody,
|
||||
ModelListEntry,
|
||||
ModelsListResponse,
|
||||
OrgNewBody,
|
||||
OrgUpdateBody,
|
||||
SpendLogsParams,
|
||||
TagNewBody,
|
||||
TeamNewBody,
|
||||
TeamUpdateBody,
|
||||
ToolsetCreateBody,
|
||||
ToolsetUpdateBody,
|
||||
UserNewBody,
|
||||
UserUpdateBody,
|
||||
)
|
||||
from proxy_client import (
|
||||
ConvergeOutcome,
|
||||
Caller,
|
||||
Converged,
|
||||
ConvergeOutcome,
|
||||
CredentialKind,
|
||||
EverywhereConverged,
|
||||
ModelsPoller,
|
||||
NeverConvergedOn,
|
||||
|
|
@ -42,6 +76,115 @@ from proxy_client import (
|
|||
)
|
||||
from transport import Transport
|
||||
|
||||
|
||||
@contextmanager
|
||||
def caller_boundary(
|
||||
status: int = 200, bodies: SimpleQueue[bytes] | None = None, *, delete_status: int | None = None
|
||||
) -> Generator[tuple[ManagementClient, SimpleQueue[str]]]:
|
||||
received: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
pass
|
||||
|
||||
def do_GET(self) -> None:
|
||||
received.put(self.headers.get("Authorization", ""))
|
||||
self.send_response(delete_status if self.path == "/key/delete" and delete_status is not None else status)
|
||||
self.end_headers()
|
||||
self.wfile.write(
|
||||
b'{"key":"owned","info":{"key_alias":"owned"},"data":[{"id":"owned"}],"team_id":"owned","team_info":{},"model_id":"owned"}'
|
||||
)
|
||||
|
||||
def do_POST(self) -> None:
|
||||
body: Final = self.rfile.read(int(self.headers.get("Content-Length", "0")))
|
||||
if bodies is not None:
|
||||
bodies.put(body)
|
||||
self.do_GET()
|
||||
|
||||
do_PATCH = do_POST
|
||||
do_PUT = do_POST
|
||||
do_DELETE = do_POST
|
||||
|
||||
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
||||
thread: Final = Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True)
|
||||
thread.start()
|
||||
url: Final = f"http://127.0.0.1:{server.server_port}"
|
||||
proxy: Final = build_proxy_client(
|
||||
base_url=url, control_plane_base_url=url, replica_urls=(url,), master_key="bootstrap"
|
||||
)
|
||||
try:
|
||||
yield ManagementClient(proxy=proxy, master_key="bootstrap"), received
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
class TestBoundManagementCaller:
|
||||
def test_strict_key_cleanup_accepts_missing_only_when_requested(self) -> None:
|
||||
with caller_boundary(delete_status=404) as (bootstrap, received), without_retries():
|
||||
with pytest.raises(AssertionError):
|
||||
bootstrap.delete_key_strict("owned")
|
||||
bootstrap.delete_key_strict("owned", missing_ok=True)
|
||||
assert (received.get_nowait(), received.get_nowait()) == ("Bearer bootstrap", "Bearer bootstrap")
|
||||
|
||||
def test_actor_key_cleanup_reports_failure_and_continues(self) -> None:
|
||||
with caller_boundary(delete_status=500) as (bootstrap, received), without_retries():
|
||||
resources: Final = ResourceManager(client=bootstrap.proxy, strict_cleanup=True)
|
||||
remaining: SimpleQueue[str] = SimpleQueue()
|
||||
resources.defer(lambda: remaining.put("cleaned"))
|
||||
factory: Final = ActorFactory(
|
||||
bootstrap=bootstrap,
|
||||
idp=Keycloak(base_url="http://unused.test", realm="test", admin_username="test", admin_password="test"),
|
||||
resources=resources,
|
||||
)
|
||||
assert factory.key().key == "owned"
|
||||
with pytest.raises(ExceptionGroup, match="Resource cleanup failed") as failure:
|
||||
resources.teardown()
|
||||
assert len(failure.value.exceptions) == 1
|
||||
assert remaining.get_nowait() == "cleaned"
|
||||
assert (received.get_nowait(), received.get_nowait()) == ("Bearer bootstrap", "Bearer bootstrap")
|
||||
|
||||
@pytest.mark.parametrize("kind", ("direct_jwt", "virtual_key", "dashboard_session"))
|
||||
def test_direct_delegated_and_replica_reads_keep_the_bound_caller(self, kind: CredentialKind) -> None:
|
||||
with caller_boundary() as (bootstrap, received):
|
||||
caller: Final = Caller(credential="synthetic-caller", kind=kind, role="internal_user", tenant="tenant-a")
|
||||
bound: Final = bootstrap.with_caller(caller)
|
||||
bound.update_key(KeyUpdateBody(key="owned", key_alias="updated"))
|
||||
bound.proxy.key_info("owned")
|
||||
bound.proxy.read_back_everywhere(
|
||||
"/key/info",
|
||||
params=KeyUpdateBody(key="owned"),
|
||||
response_type=KeyInfoResponse,
|
||||
converged=lambda result: isinstance(result, Success),
|
||||
)
|
||||
bound.proxy.read_body_back_everywhere(
|
||||
"/key/info", KeyInfoResponse, settled=lambda result: result.info.key_alias == "owned"
|
||||
)
|
||||
assert tuple(received.get_nowait() for _ in range(4)) == ("Bearer synthetic-caller",) * 4
|
||||
assert received.empty()
|
||||
bootstrap.proxy.key_info("owned")
|
||||
assert received.get_nowait() == "Bearer bootstrap"
|
||||
|
||||
def test_explicit_override_wins_without_rebinding_or_changing_master(self) -> None:
|
||||
with caller_boundary() as (bootstrap, received):
|
||||
bound: Final = bootstrap.with_caller(Caller(credential="bound", kind="direct_jwt", role="internal_user"))
|
||||
bound.update_key(KeyUpdateBody(key="owned"), caller_key="override")
|
||||
bound.proxy.key_info("owned")
|
||||
assert received.get_nowait() == "Bearer override"
|
||||
assert received.get_nowait() == "Bearer bound"
|
||||
assert bound.master_key == "bootstrap"
|
||||
|
||||
def test_credentials_are_absent_from_binding_and_header_diagnostics(self) -> None:
|
||||
with caller_boundary() as (bootstrap, _):
|
||||
caller: Final = Caller(credential="private-value", kind="direct_jwt", role="internal_user")
|
||||
bound: Final = bootstrap.with_caller(caller)
|
||||
assert "private-value" not in repr(caller)
|
||||
assert "private-value" not in repr(bound)
|
||||
assert "private-value" not in repr(bound.proxy.management_headers())
|
||||
assert "bootstrap" not in repr(bound)
|
||||
|
||||
|
||||
MODEL: Final = "gpt-under-test"
|
||||
_NO_TRANSPORTS: Final = cast(Transport, None)
|
||||
TIMEOUT: Final = 10.0
|
||||
|
|
@ -275,3 +418,166 @@ class TestReplicasFor:
|
|||
client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={})
|
||||
with pytest.raises(AssertionError, match="no replica is configured"):
|
||||
_ = client.replicas_for("/v1/models")
|
||||
|
||||
|
||||
MANAGEMENT_OPERATIONS: Final[tuple[tuple[str, Callable[[ManagementClient], object]], ...]] = (
|
||||
("generate_key", lambda c: c.generate_key(KeyGenerateBody())),
|
||||
("llm_only_key", lambda c: c.llm_only_key()),
|
||||
("update_key", lambda c: c.update_key(KeyUpdateBody(key="owned"))),
|
||||
("update_key_models", lambda c: c.update_key_models("owned", [])),
|
||||
("key_info", lambda c: c.key_info_as("owned")),
|
||||
("delete_key_strict", lambda c: c.delete_key_strict("owned")),
|
||||
("delete_model_strict", lambda c: c.delete_model_strict("owned")),
|
||||
(
|
||||
"connection_test",
|
||||
lambda c: c.connection_test(
|
||||
ConnectionTestBody(litellm_params=LiteLLMParamsBody(model="synthetic"), mode="chat")
|
||||
),
|
||||
),
|
||||
("block_key", lambda c: c.block_key("owned")),
|
||||
("regenerate_key", lambda c: c.regenerate_key("owned")),
|
||||
("reset_key_spend", lambda c: c.reset_key_spend("owned", 0)),
|
||||
("key_list", lambda c: c.key_list("owned")),
|
||||
("key_alias_count", lambda c: c.key_alias_count("owned")),
|
||||
("create_team", lambda c: c.create_team(TeamNewBody(team_alias="owned"))),
|
||||
("update_team", lambda c: c.update_team(TeamUpdateBody(team_id="owned", team_alias="updated"))),
|
||||
("delete_team", lambda c: c.delete_team("owned")),
|
||||
("team_info", lambda c: c.team_info("owned")),
|
||||
("team_list_ids", lambda c: c.team_list_ids()),
|
||||
("team_info_status", lambda c: c.team_info_status("owned")),
|
||||
("add_team_member", lambda c: c.add_team_member("owned", "user")),
|
||||
("delete_team_member", lambda c: c.delete_team_member("owned", "user")),
|
||||
("create_user", lambda c: c.create_user(UserNewBody(user_email="actor@example.com", user_role="internal_user"))),
|
||||
("create_customer", lambda c: c.create_customer("owned")),
|
||||
("customer_info", lambda c: c.customer_info("owned")),
|
||||
("delete_customer", lambda c: c.delete_customer("owned")),
|
||||
("update_user", lambda c: c.update_user(UserUpdateBody(user_id="owned", user_role="internal_user"))),
|
||||
("delete_user", lambda c: c.delete_user("owned")),
|
||||
("delete_user_strict", lambda c: c.delete_user_strict("owned")),
|
||||
("user_info", lambda c: c.user_info("owned")),
|
||||
("user_count", lambda c: c.user_count("owned")),
|
||||
("user_list_ids", lambda c: c.user_list_ids("owned")),
|
||||
("create_org", lambda c: c.create_org(OrgNewBody(organization_alias="owned"))),
|
||||
("update_org", lambda c: c.update_org(OrgUpdateBody(organization_id="owned", organization_alias="updated"))),
|
||||
("delete_org", lambda c: c.delete_org("owned")),
|
||||
("org_info", lambda c: c.org_info("owned")),
|
||||
("org_info_status", lambda c: c.org_info_status("owned")),
|
||||
("create_tag", lambda c: c.create_tag(TagNewBody(name="owned"))),
|
||||
("delete_tag", lambda c: c.delete_tag("owned")),
|
||||
("tag_list", lambda c: c.tag_list()),
|
||||
("create_mcp_server", lambda c: c.create_mcp_server(McpServerCreateBody(alias="owned", url="http://example.test"))),
|
||||
("update_mcp_server", lambda c: c.update_mcp_server(McpServerUpdateBody(server_id="owned", alias=None))),
|
||||
("delete_mcp_server", lambda c: c.delete_mcp_server("owned")),
|
||||
("proxy.generate_key", lambda c: c.proxy.generate_key(KeyGenerateBody())),
|
||||
("proxy.delete_key", lambda c: c.proxy.delete_key("owned")),
|
||||
("proxy.delete_customers", lambda c: c.proxy.delete_customers(["owned"])),
|
||||
("proxy.key_info", lambda c: c.proxy.key_info("owned")),
|
||||
("proxy.memory_summary", lambda c: c.proxy.memory_summary_everywhere()),
|
||||
("proxy.model_info", lambda c: c.proxy.model_info()),
|
||||
("proxy.model_cost_map", lambda c: c.proxy.model_cost_map()),
|
||||
("proxy.create_model", lambda c: c.proxy.create_model("owned", LiteLLMParamsBody(model="synthetic"))),
|
||||
("proxy.update_model", lambda c: c.proxy.update_model("owned", LiteLLMParamsBody(model="synthetic"))),
|
||||
("proxy.delete_model", lambda c: c.proxy.delete_model("owned")),
|
||||
("proxy.create_toolset", lambda c: c.proxy.create_toolset(ToolsetCreateBody(toolset_name="owned", tools=[]))),
|
||||
("proxy.update_toolset", lambda c: c.proxy.update_toolset(ToolsetUpdateBody(toolset_id="owned", description=None))),
|
||||
("proxy.delete_toolset", lambda c: c.proxy.delete_toolset("owned")),
|
||||
(
|
||||
"proxy.create_credential",
|
||||
lambda c: c.proxy.create_credential(CredentialCreateBody(credential_name="owned", credential_values={})),
|
||||
),
|
||||
("proxy.delete_credential", lambda c: c.proxy.delete_credential("owned")),
|
||||
("proxy.create_team", lambda c: c.proxy.create_team(TeamNewBody(team_alias="owned"))),
|
||||
("proxy.delete_team", lambda c: c.proxy.delete_team("owned")),
|
||||
("proxy.delete_user", lambda c: c.proxy.delete_user("owned")),
|
||||
("proxy.spend_logs", lambda c: c.proxy.spend_logs(SpendLogsParams(api_key="owned"))),
|
||||
("proxy.probe", lambda c: c.proxy.probe("/user/info", params=NoBody())),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("name", "operation"), MANAGEMENT_OPERATIONS, ids=tuple(name for name, _ in MANAGEMENT_OPERATIONS)
|
||||
)
|
||||
@pytest.mark.parametrize("kind", ("master", "direct_jwt", "virtual_key", "dashboard_session"))
|
||||
def test_management_operations_send_the_selected_credential(
|
||||
name: str,
|
||||
operation: Callable[[ManagementClient], object],
|
||||
kind: CredentialKind,
|
||||
) -> None:
|
||||
with caller_boundary(status=401) as (bootstrap, received), without_retries():
|
||||
client: Final = (
|
||||
bootstrap
|
||||
if kind == "master"
|
||||
else bootstrap.with_caller(Caller(credential=f"synthetic-{kind}", kind=kind, role="internal_user"))
|
||||
)
|
||||
try:
|
||||
operation(client)
|
||||
except AssertionError:
|
||||
pass
|
||||
expected: Final = "Bearer bootstrap" if kind == "master" else f"Bearer synthetic-{kind}"
|
||||
assert received.get_nowait() == expected, name
|
||||
assert received.empty(), "an unauthorized request must not be retried"
|
||||
|
||||
|
||||
class TestSplitCallerPropagation:
|
||||
def test_control_and_data_replica_readers_keep_the_caller(self) -> None:
|
||||
with caller_boundary() as (data, data_headers), caller_boundary() as (control, control_headers):
|
||||
data_url: Final = next(iter(data.proxy.replicas))
|
||||
control_url: Final = next(iter(control.proxy.replicas))
|
||||
proxy: Final = build_proxy_client(
|
||||
base_url=data_url,
|
||||
control_plane_base_url=control_url,
|
||||
replica_urls=(data_url,),
|
||||
master_key="bootstrap",
|
||||
).with_caller(Caller(credential="tenant-token", kind="direct_jwt", role="team_member"))
|
||||
proxy.key_info("owned")
|
||||
proxy.read_body_back_everywhere(
|
||||
"/key/info", KeyInfoResponse, settled=lambda info: info.info.key_alias == "owned"
|
||||
)
|
||||
proxy.read_back_everywhere(
|
||||
"/key/info",
|
||||
params=NoBody(),
|
||||
response_type=KeyInfoResponse,
|
||||
converged=lambda result: isinstance(result, Success),
|
||||
)
|
||||
assert control_headers.get_nowait() == "Bearer tenant-token"
|
||||
assert control_headers.get_nowait() == "Bearer tenant-token"
|
||||
assert data_headers.get_nowait() == "Bearer tenant-token"
|
||||
assert control_headers.empty() and data_headers.empty()
|
||||
|
||||
def test_successful_team_and_model_polling_uses_the_bound_caller(self) -> None:
|
||||
with caller_boundary() as (bootstrap, received):
|
||||
bound: Final = bootstrap.with_caller(Caller(credential="caller", kind="direct_jwt", role="proxy_admin"))
|
||||
bound.create_team(TeamNewBody(team_alias="owned"))
|
||||
bound.proxy.create_model("owned", LiteLLMParamsBody(model="synthetic"))
|
||||
assert tuple(received.get_nowait() for _ in range(4)) == ("Bearer caller",) * 4
|
||||
assert received.empty()
|
||||
|
||||
def test_expired_shaped_token_is_sent_once_without_renewal(self) -> None:
|
||||
with caller_boundary(status=401) as (bootstrap, received):
|
||||
bound: Final = bootstrap.with_caller(
|
||||
Caller(credential="expired.payload.signature", kind="direct_jwt", role="internal_user")
|
||||
)
|
||||
result: Final = bound.key_info_as("owned")
|
||||
assert not isinstance(result, Success)
|
||||
assert received.get_nowait() == "Bearer expired.payload.signature"
|
||||
assert received.empty()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("operation", ("server", "toolset"))
|
||||
def test_partial_updates_preserve_explicit_null_at_the_http_boundary(operation: str) -> None:
|
||||
bodies: Final[SimpleQueue[bytes]] = SimpleQueue()
|
||||
with caller_boundary(status=401, bodies=bodies) as (bootstrap, _):
|
||||
try:
|
||||
if operation == "server":
|
||||
bootstrap.update_mcp_server(McpServerUpdateBody(server_id="owned", alias=None))
|
||||
else:
|
||||
bootstrap.proxy.update_toolset(ToolsetUpdateBody(toolset_id="owned", description=None))
|
||||
except AssertionError:
|
||||
pass
|
||||
expected: Final = (
|
||||
{"server_id": "owned", "alias": None}
|
||||
if operation == "server"
|
||||
else {"toolset_id": "owned", "description": None}
|
||||
)
|
||||
assert json.loads(bodies.get_nowait()) == expected
|
||||
assert bodies.empty()
|
||||
|
|
|
|||
|
|
@ -7,11 +7,9 @@ client touches requests.* or builds raw dicts; they pass pydantic models here.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Protocol
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
import e2e_http
|
||||
from e2e_http import (
|
||||
URL,
|
||||
|
|
@ -21,6 +19,7 @@ from e2e_http import (
|
|||
Result,
|
||||
StreamingResponse,
|
||||
)
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class Transport(Protocol):
|
||||
|
|
@ -85,7 +84,7 @@ class Transport(Protocol):
|
|||
self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R]
|
||||
) -> Result[R]: ...
|
||||
|
||||
def probe(self, path: str, *, params: BaseModel) -> ProbeResult: ...
|
||||
def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: ...
|
||||
|
||||
def upload[R: BaseModel](
|
||||
self,
|
||||
|
|
@ -113,7 +112,7 @@ class Transport(Protocol):
|
|||
@dataclass(frozen=True, slots=True)
|
||||
class HttpTransport:
|
||||
base_url: str
|
||||
master_key: str
|
||||
master_key: str = field(repr=False)
|
||||
request_timeout: float = 60.0
|
||||
|
||||
def _url(self, path: str) -> URL:
|
||||
|
|
@ -245,10 +244,10 @@ class HttpTransport:
|
|||
timeout=self.request_timeout,
|
||||
)
|
||||
|
||||
def probe(self, path: str, *, params: BaseModel) -> ProbeResult:
|
||||
def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult:
|
||||
return e2e_http.probe(
|
||||
self._url(path),
|
||||
headers=self.master,
|
||||
headers=self.master if headers is None else headers,
|
||||
params=params,
|
||||
timeout=self.request_timeout,
|
||||
)
|
||||
|
|
@ -434,8 +433,8 @@ class SplitTransport:
|
|||
path, headers=headers, json=json, params=params, stream=stream
|
||||
)
|
||||
|
||||
def probe(self, path: str, *, params: BaseModel) -> ProbeResult:
|
||||
return self._route(path).probe(path, params=params)
|
||||
def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult:
|
||||
return self._route(path).probe(path, params=params, headers=headers)
|
||||
|
||||
def upload[R: BaseModel](
|
||||
self,
|
||||
|
|
|
|||
30
tests/e2e/ui/oidcSetup.ts
Normal file
30
tests/e2e/ui/oidcSetup.ts
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
import { chromium, expect } from "@playwright/test";
|
||||
import * as fs from "fs";
|
||||
import * as path from "path";
|
||||
|
||||
export default async function oidcSetup() {
|
||||
const baseURL = process.env.E2E_OIDC_UI_URL;
|
||||
const issuer = process.env.JWT_ISSUER;
|
||||
const username = process.env.E2E_OIDC_USERNAME;
|
||||
const password = process.env.E2E_OIDC_PASSWORD;
|
||||
if (!baseURL || !issuer || !username || !password) {
|
||||
throw new Error("The OIDC setup requires a running stack, issuer, and provisioned actor credentials");
|
||||
}
|
||||
const artifactDir = process.env.E2E_UI_ARTIFACT_DIR || ".";
|
||||
fs.mkdirSync(artifactDir, { recursive: true });
|
||||
const browser = await chromium.launch();
|
||||
try {
|
||||
const page = await browser.newPage();
|
||||
await page.goto(`${baseURL.replace(/\/$/, "")}/sso/key/generate`);
|
||||
await expect(page).toHaveURL(new RegExp(`^${issuer.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}/`));
|
||||
await page.getByLabel("Username or email").fill(username);
|
||||
await page.getByLabel("Password", { exact: true }).fill(password);
|
||||
await page.getByRole("button", { name: "Sign In", exact: true }).click();
|
||||
await page.waitForURL((url) => url.origin === new URL(baseURL).origin && url.pathname.startsWith("/ui"));
|
||||
const statePath = path.join(artifactDir, "oidc.storageState.json");
|
||||
await page.context().storageState({ path: statePath });
|
||||
fs.chmodSync(statePath, 0o600);
|
||||
} finally {
|
||||
await browser.close();
|
||||
}
|
||||
}
|
||||
22
tests/e2e/ui/playwright.oidc.config.ts
Normal file
22
tests/e2e/ui/playwright.oidc.config.ts
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
import { defineConfig, devices } from "@playwright/test";
|
||||
import * as path from "path";
|
||||
|
||||
const baseURL = process.env.E2E_OIDC_UI_URL;
|
||||
if (!baseURL) throw new Error("E2E_OIDC_UI_URL must point to the running OIDC stack");
|
||||
|
||||
export default defineConfig({
|
||||
testDir: ".",
|
||||
testMatch: "oidc/**/*.spec.ts",
|
||||
retries: 0,
|
||||
workers: 1,
|
||||
outputDir: path.join(process.env.E2E_UI_ARTIFACT_DIR || ".", "oidc", "test-results"),
|
||||
globalSetup: require.resolve("./oidcSetup"),
|
||||
use: {
|
||||
...devices["Desktop Chrome"],
|
||||
baseURL,
|
||||
storageState: path.join(process.env.E2E_UI_ARTIFACT_DIR || ".", "oidc.storageState.json"),
|
||||
trace: "off",
|
||||
screenshot: "off",
|
||||
video: "off",
|
||||
},
|
||||
});
|
||||
|
|
@ -11,7 +11,7 @@ import litellm
|
|||
from unittest.mock import patch, MagicMock, AsyncMock
|
||||
from create_mock_standard_logging_payload import create_standard_logging_payload
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||||
from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo
|
||||
from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS
|
||||
|
||||
|
||||
|
|
@ -630,10 +630,12 @@ def test_deployment_callback_respects_cooldown_time(model_list):
|
|||
|
||||
|
||||
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
|
||||
def test_log_retry(model_list, metadata_key):
|
||||
"""log_retry appends one flat record per failed attempt and copies neither the request kwargs nor
|
||||
the request metadata into it"""
|
||||
def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) -> None:
|
||||
"""log_retry appends one flat record per failed attempt, copies neither the request kwargs nor the
|
||||
request metadata into it, counts every failed attempt of the request independently of the
|
||||
per-hop attempted_retries, and never trusts a negative count planted before the first failure"""
|
||||
router = Router(model_list=model_list)
|
||||
rate_limit_error = litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo")
|
||||
new_kwargs = router.log_retry(
|
||||
kwargs={
|
||||
"model": "gpt-3.5-turbo",
|
||||
|
|
@ -641,7 +643,7 @@ def test_log_retry(model_list, metadata_key):
|
|||
"messages": [{"role": "user", "content": "hi"}],
|
||||
metadata_key: {"model_info": {"id": "deployment-1"}, "attempted_retries": 2, "user_api_key": "sk-proxy"},
|
||||
},
|
||||
e=litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo"),
|
||||
e=rate_limit_error,
|
||||
)
|
||||
assert json.loads(json.dumps(new_kwargs[metadata_key]["previous_models"])) == [
|
||||
{
|
||||
|
|
@ -652,6 +654,10 @@ def test_log_retry(model_list, metadata_key):
|
|||
"attempted_retries": 2,
|
||||
}
|
||||
]
|
||||
assert new_kwargs[metadata_key]["request_retry_count"] == 1
|
||||
assert router.log_retry(kwargs=new_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 2
|
||||
planted_kwargs = {"model": "gpt-3.5-turbo", metadata_key: {"request_retry_count": -100}}
|
||||
assert router.log_retry(kwargs=planted_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 1
|
||||
|
||||
|
||||
def test_update_usage(model_list):
|
||||
|
|
|
|||
|
|
@ -695,6 +695,38 @@ def test_vertex_cost_and_usage_aggregation(monkeypatch):
|
|||
assert result.failed_requests == 0
|
||||
|
||||
|
||||
def test_vertex_batch_usage_preserves_modality_token_details(monkeypatch):
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"vertex_ai/gemini-embedding-2",
|
||||
{
|
||||
"input_cost_per_token_batches": 1e-7,
|
||||
"input_cost_per_audio_token_batches": 3.25e-6,
|
||||
"input_cost_per_image_token_batches": 2.25e-7,
|
||||
"input_cost_per_video_token_batches": 6e-6,
|
||||
},
|
||||
)
|
||||
responses = [
|
||||
{
|
||||
"response": {
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 84,
|
||||
"candidatesTokenCount": 0,
|
||||
"totalTokenCount": 84,
|
||||
"promptTokensDetails": [
|
||||
{"modality": "AUDIO", "tokenCount": 64},
|
||||
{"modality": "TEXT", "tokenCount": 20},
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-embedding-2")
|
||||
|
||||
assert result.prompt_cost == pytest.approx(64 * 3.25e-6 + 20 * 1e-7)
|
||||
|
||||
|
||||
def test_vertex_cost_skips_none_response_body(monkeypatch):
|
||||
import litellm.cost_calculator as cc
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,8 @@ never rewrite. It is consumed by compress() and by the Headroom guardrail, so
|
|||
the two agree on what "never compress this" means.
|
||||
"""
|
||||
|
||||
from litellm.compression.compress import get_protected_indices
|
||||
from litellm.compression.compress import compress, get_protected_indices
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
|
||||
def test_protects_system_last_user_and_last_assistant():
|
||||
|
|
@ -53,3 +54,94 @@ def test_every_system_row_is_protected():
|
|||
def test_no_user_or_assistant_rows():
|
||||
assert sorted(get_protected_indices([{"role": "system", "content": "sys"}])) == [0]
|
||||
assert get_protected_indices([]) == ()
|
||||
|
||||
|
||||
def test_mid_history_cache_control_part_is_protected():
|
||||
# A large cached tool result from a few turns back, not the last user or
|
||||
# last assistant row -- exactly the row a provider prompt-cache pins to
|
||||
# exact bytes. Rewriting it (even leaving the marker on) changes those
|
||||
# bytes and turns the next request's cache read into a cache write.
|
||||
messages = [
|
||||
{"role": "user", "content": "old question"},
|
||||
{"role": "assistant", "content": "old answer"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "a large cached tool result", "cache_control": {"type": "ephemeral"}},
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": "ack"},
|
||||
{"role": "user", "content": "live instruction"},
|
||||
]
|
||||
|
||||
# index 3 = last assistant, index 4 = last user (both protected by role
|
||||
# regardless), index 2 = the cache_control-marked row itself.
|
||||
assert sorted(get_protected_indices(messages)) == [2, 3, 4]
|
||||
|
||||
|
||||
def test_cache_control_directly_on_message_is_protected():
|
||||
messages = [
|
||||
{"role": "user", "content": "old question", "cache_control": {"type": "ephemeral"}},
|
||||
{"role": "assistant", "content": "old answer"},
|
||||
{"role": "user", "content": "live instruction"},
|
||||
]
|
||||
|
||||
assert sorted(get_protected_indices(messages)) == [0, 1, 2]
|
||||
|
||||
|
||||
def test_cache_control_protection_does_not_duplicate_already_protected_rows():
|
||||
# The last user row is already protected by role; marking it too must not
|
||||
# produce a duplicate index.
|
||||
messages = [
|
||||
{"role": "system", "content": "sys"},
|
||||
{"role": "user", "content": "live", "cache_control": {"type": "ephemeral"}},
|
||||
]
|
||||
|
||||
protected = get_protected_indices(messages)
|
||||
|
||||
assert sorted(protected) == [0, 1]
|
||||
assert len(protected) == len(set(protected))
|
||||
|
||||
|
||||
def test_content_that_is_not_a_list_of_mappings_is_not_treated_as_cache_control():
|
||||
# Defensive: a plain string content, or a list of non-dict items, must not
|
||||
# raise or be misread as carrying a breakpoint.
|
||||
messages = [
|
||||
{"role": "assistant", "content": "plain string content"},
|
||||
{"role": "user", "content": ["not", "a", "dict", "list"]},
|
||||
{"role": "user", "content": "live instruction"},
|
||||
]
|
||||
|
||||
assert sorted(get_protected_indices(messages)) == [0, 2]
|
||||
|
||||
|
||||
def test_compress_keeps_part_level_cache_control_row_verbatim():
|
||||
# compress() scores text-only copies of the rows, where a part-level marker
|
||||
# is gone; protection has to read the original rows or the pinned row is stubbed.
|
||||
stale_log = {"role": "user", "content": [{"type": "text", "text": "stale log line " * 2000}]}
|
||||
pinned = {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "cached tool result " * 2000, "cache_control": {"type": "ephemeral"}},
|
||||
],
|
||||
}
|
||||
messages = [
|
||||
stale_log,
|
||||
{"role": "assistant", "content": "old answer"},
|
||||
pinned,
|
||||
{"role": "assistant", "content": "ack"},
|
||||
{"role": "user", "content": "live instruction"},
|
||||
]
|
||||
|
||||
result = compress(
|
||||
messages,
|
||||
model="gpt-4o",
|
||||
call_type=CallTypes.anthropic_messages,
|
||||
compression_trigger=1000,
|
||||
compression_target=500,
|
||||
)
|
||||
|
||||
assert len(result["messages"]) == len(messages)
|
||||
assert result["messages"][2] == pinned
|
||||
assert result["messages"][0] != stale_log
|
||||
assert len(result["cache"]) >= 1
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import json
|
|||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
from collections.abc import Mapping
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import litellm
|
||||
|
|
@ -74,6 +75,108 @@ def test_missing_cache_read_policy_preserves_billing(prompt_tokens, read_rate, s
|
|||
assert prompt_cost == pytest.approx((prompt_tokens - 100) * billed[0] + 100 * billed[4])
|
||||
|
||||
|
||||
def test_generic_cost_per_token_prefers_audio_per_second_rate() -> None:
|
||||
model_info: ModelInfo = {
|
||||
"key": "gemini-embedding-2",
|
||||
"max_tokens": None,
|
||||
"max_input_tokens": None,
|
||||
"max_output_tokens": None,
|
||||
"input_cost_per_token": 2e-7,
|
||||
"input_cost_per_audio_token": 6.5e-6,
|
||||
"input_cost_per_audio_per_second": 0.00016,
|
||||
"output_cost_per_token": 0.0,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"mode": "embedding",
|
||||
"supported_openai_params": None,
|
||||
}
|
||||
usage = Usage(
|
||||
prompt_tokens=64,
|
||||
completion_tokens=0,
|
||||
total_tokens=64,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
audio_tokens=64,
|
||||
audio_length_seconds=2,
|
||||
),
|
||||
)
|
||||
|
||||
prompt_cost, _ = generic_cost_per_token(
|
||||
model="gemini-embedding-2",
|
||||
usage=usage,
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_info=model_info,
|
||||
)
|
||||
|
||||
assert prompt_cost == pytest.approx(2 * 0.00016)
|
||||
|
||||
|
||||
def test_generic_cost_per_token_prefers_image_per_image_rate() -> None:
|
||||
model_info: ModelInfo = {
|
||||
"key": "gemini-embedding-2",
|
||||
"max_tokens": None,
|
||||
"max_input_tokens": None,
|
||||
"max_output_tokens": None,
|
||||
"input_cost_per_token": 2e-7,
|
||||
"input_cost_per_image_token": 4.5e-7,
|
||||
"input_cost_per_image": 0.00012,
|
||||
"output_cost_per_token": 0.0,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"mode": "embedding",
|
||||
"supported_openai_params": None,
|
||||
}
|
||||
usage = Usage(
|
||||
prompt_tokens=258,
|
||||
completion_tokens=0,
|
||||
total_tokens=258,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
image_tokens=258,
|
||||
image_count=1,
|
||||
),
|
||||
)
|
||||
|
||||
prompt_cost, _ = generic_cost_per_token(
|
||||
model="gemini-embedding-2",
|
||||
usage=usage,
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_info=model_info,
|
||||
)
|
||||
|
||||
assert prompt_cost == pytest.approx(0.00012)
|
||||
|
||||
|
||||
def test_generic_cost_per_token_prefers_video_per_second_rate() -> None:
|
||||
model_info: ModelInfo = {
|
||||
"key": "gemini-embedding-2",
|
||||
"max_tokens": None,
|
||||
"max_input_tokens": None,
|
||||
"max_output_tokens": None,
|
||||
"input_cost_per_token": 2e-7,
|
||||
"input_cost_per_video_token": 1.2e-5,
|
||||
"input_cost_per_video_per_second": 0.00079,
|
||||
"output_cost_per_token": 0.0,
|
||||
"litellm_provider": "vertex_ai",
|
||||
"mode": "embedding",
|
||||
"supported_openai_params": None,
|
||||
}
|
||||
usage = Usage(
|
||||
prompt_tokens=516,
|
||||
completion_tokens=0,
|
||||
total_tokens=516,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
video_tokens=516,
|
||||
video_length_seconds=2,
|
||||
),
|
||||
)
|
||||
|
||||
prompt_cost, _ = generic_cost_per_token(
|
||||
model="gemini-embedding-2",
|
||||
usage=usage,
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_info=model_info,
|
||||
)
|
||||
|
||||
assert prompt_cost == pytest.approx(2 * 0.00079)
|
||||
|
||||
|
||||
def test_missing_cache_read_uses_off_peak_input_rate():
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
|
@ -4039,7 +4142,7 @@ def test_billed_token_rates_follow_the_token_tier_the_breakdown_bills_at(monkeyp
|
|||
cache_read_input_token_cost=6e-7,
|
||||
cache_read_input_audio_token_cost=6e-7,
|
||||
cache_creation_input_token_cost=7.5e-6,
|
||||
cache_creation_input_token_cost_above_1hr=0.0,
|
||||
cache_creation_input_token_cost_above_1hr=7.5e-6,
|
||||
output_cost_per_reasoning_token=3e-5,
|
||||
)
|
||||
assert breakdown.cache_read_cost == pytest.approx(200_000 * rates.cache_read_input_token_cost)
|
||||
|
|
@ -5334,3 +5437,72 @@ def test_realtime_models_bill_cached_text_and_audio_at_their_cache_read_rates(
|
|||
|
||||
prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=custom_llm_provider)
|
||||
assert prompt_cost == pytest.approx(expected_prompt_cost)
|
||||
|
||||
|
||||
def test_generic_cost_per_token_bills_cache_creation_at_the_input_rate_without_a_write_price():
|
||||
"""Azure and OpenAI publish no cache-write price and bill cache writes as ordinary input.
|
||||
A deployment priced with only input, output, and cache-read rates must bill the creation
|
||||
tokens the provider reports at the input rate, never at 0. The numbers are a cold 7,336-token
|
||||
prompt on a deployment that reports all but 3 of them as cache creation."""
|
||||
model_info = {
|
||||
"input_cost_per_token": 2e-7,
|
||||
"output_cost_per_token": 1.25e-6,
|
||||
"cache_read_input_token_cost": 2e-8,
|
||||
}
|
||||
usage = Usage(
|
||||
prompt_tokens=7336,
|
||||
completion_tokens=23,
|
||||
total_tokens=7359,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=0, cache_creation_tokens=7333),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model="custom-priced-deployment", usage=usage, custom_llm_provider="azure", model_info=model_info
|
||||
)
|
||||
|
||||
assert prompt_cost == pytest.approx(7336 * 2e-7)
|
||||
assert completion_cost == pytest.approx(23 * 1.25e-6)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("cache_rates", "current_time", "expected_creation", "expected_creation_1h"),
|
||||
(
|
||||
pytest.param({}, None, 2e-7, 2e-7, id="no-write-price-uses-the-input-rate"),
|
||||
pytest.param({"cache_creation_input_token_cost": 2.5e-7}, None, 2.5e-7, 2.5e-7, id="no-1h-price-uses-the-write-price"),
|
||||
pytest.param({"cache_creation_input_token_cost": 0.0}, None, 0.0, 0.0, id="explicit-zero-stays-zero"),
|
||||
pytest.param(
|
||||
{"off_peak_pricing": {"hours_utc": "00:00-23:59", "input_cost_per_token": 1e-7}},
|
||||
datetime(2026, 9, 14, 12, tzinfo=timezone.utc),
|
||||
1e-7,
|
||||
1e-7,
|
||||
id="no-write-price-uses-the-off-peak-input-rate",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
"off_peak_pricing": {
|
||||
"hours_utc": "00:00-23:59",
|
||||
"input_cost_per_token": 1e-7,
|
||||
"cache_creation_input_token_cost": 3e-7,
|
||||
}
|
||||
},
|
||||
datetime(2026, 9, 14, 12, tzinfo=timezone.utc),
|
||||
3e-7,
|
||||
3e-7,
|
||||
id="no-1h-price-uses-the-off-peak-write-price",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_get_token_base_cost_resolves_missing_cache_write_rates_like_the_tiered_path(
|
||||
cache_rates: Mapping[str, float | Mapping[str, float | str]],
|
||||
current_time: datetime | None,
|
||||
expected_creation: float,
|
||||
expected_creation_1h: float,
|
||||
) -> None:
|
||||
model_info = {"input_cost_per_token": 2e-7, "output_cost_per_token": 1.25e-6, **cache_rates}
|
||||
usage = Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11)
|
||||
|
||||
_, _, creation, creation_1h, _ = _get_token_base_cost(model_info, usage, current_time=current_time)
|
||||
|
||||
assert creation == pytest.approx(expected_creation)
|
||||
assert creation_1h == pytest.approx(expected_creation_1h)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
|
|
@ -178,9 +177,7 @@ class TestExceptionCheckers:
|
|||
]
|
||||
|
||||
for error_str in error_strings:
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(
|
||||
error_str
|
||||
)
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
|
||||
assert result is True, f"Should detect policy violation in: {error_str}"
|
||||
|
||||
def test_is_azure_content_policy_violation_error_case_insensitive(self):
|
||||
|
|
@ -194,12 +191,8 @@ class TestExceptionCheckers:
|
|||
]
|
||||
|
||||
for error_str in error_strings:
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(
|
||||
error_str
|
||||
)
|
||||
assert (
|
||||
result is True
|
||||
), f"Should detect policy violation in uppercase: {error_str}"
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
|
||||
assert result is True, f"Should detect policy violation in uppercase: {error_str}"
|
||||
|
||||
def test_is_azure_content_policy_violation_error_with_non_policy_errors(self):
|
||||
"""Test that non-policy violation errors are not detected as policy violations"""
|
||||
|
|
@ -216,12 +209,8 @@ class TestExceptionCheckers:
|
|||
]
|
||||
|
||||
for error_str in error_strings:
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(
|
||||
error_str
|
||||
)
|
||||
assert (
|
||||
result is False
|
||||
), f"Should NOT detect policy violation in: {error_str}"
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
|
||||
assert result is False, f"Should NOT detect policy violation in: {error_str}"
|
||||
|
||||
def test_is_azure_content_policy_violation_error_with_partial_matches(self):
|
||||
"""Test that partial keyword matches work correctly"""
|
||||
|
|
@ -234,9 +223,7 @@ class TestExceptionCheckers:
|
|||
]
|
||||
|
||||
for error_str in positive_cases:
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(
|
||||
error_str
|
||||
)
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
|
||||
assert result is True, f"Should detect policy violation in: {error_str}"
|
||||
|
||||
# These should not match even though they contain similar words
|
||||
|
|
@ -248,12 +235,8 @@ class TestExceptionCheckers:
|
|||
]
|
||||
|
||||
for error_str in negative_cases:
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(
|
||||
error_str
|
||||
)
|
||||
assert (
|
||||
result is False
|
||||
), f"Should NOT detect policy violation in: {error_str}"
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
|
||||
assert result is False, f"Should NOT detect policy violation in: {error_str}"
|
||||
|
||||
|
||||
gemini_context_window_test_cases = [
|
||||
|
|
@ -271,12 +254,8 @@ gemini_context_window_test_cases = [
|
|||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"error_message, should_raise_context_window", gemini_context_window_test_cases
|
||||
)
|
||||
def test_gemini_context_window_error_mapping(
|
||||
error_message, should_raise_context_window
|
||||
):
|
||||
@pytest.mark.parametrize("error_message, should_raise_context_window", gemini_context_window_test_cases)
|
||||
def test_gemini_context_window_error_mapping(error_message, should_raise_context_window):
|
||||
"""
|
||||
Tests that the exception_type function correctly maps Gemini's
|
||||
context window exceeded errors to litellm.ContextWindowExceededError.
|
||||
|
|
@ -421,9 +400,7 @@ vertex_rate_limit_test_cases = [
|
|||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"error_message, should_raise_rate_limit", vertex_rate_limit_test_cases
|
||||
)
|
||||
@pytest.mark.parametrize("error_message, should_raise_rate_limit", vertex_rate_limit_test_cases)
|
||||
def test_vertex_ai_rate_limit_error_mapping(error_message, should_raise_rate_limit):
|
||||
"""
|
||||
Tests that the exception_type function correctly maps Vertex AI's
|
||||
|
|
@ -458,10 +435,7 @@ class TestGetBodyErrorCode:
|
|||
"""Unit tests for _get_body_error_code helper."""
|
||||
|
||||
def test_parses_int_code(self):
|
||||
body = (
|
||||
'{"error":{"message":"high demand","type":"upstream_error",'
|
||||
'"param":"","code":429}}'
|
||||
)
|
||||
body = '{"error":{"message":"high demand","type":"upstream_error","param":"","code":429}}'
|
||||
assert _get_body_error_code(body) == 429
|
||||
|
||||
def test_parses_string_code(self):
|
||||
|
|
@ -498,8 +472,7 @@ gemini_body_code_429_test_cases = [
|
|||
),
|
||||
(
|
||||
503,
|
||||
'{"error":{"message":"upstream unavailable","type":"upstream_error",'
|
||||
'"param":"","code":429}}',
|
||||
'{"error":{"message":"upstream unavailable","type":"upstream_error","param":"","code":429}}',
|
||||
litellm.RateLimitError,
|
||||
"HTTP 503 envelope with body code:429 -> RateLimitError",
|
||||
),
|
||||
|
|
@ -769,9 +742,7 @@ class _UpstreamHTTPError(Exception):
|
|||
self.message = "upstream failure"
|
||||
self.status_code = status_code
|
||||
self.request = httpx.Request("POST", "https://api.example.com/v1/chat/completions")
|
||||
self.response = httpx.Response(
|
||||
status_code=status_code, request=self.request, text="upstream failure"
|
||||
)
|
||||
self.response = httpx.Response(status_code=status_code, request=self.request, text="upstream failure")
|
||||
|
||||
|
||||
UPSTREAM_STATUS_CODES = (400, 401, 403, 404, 408, 422, 429, 500, 503)
|
||||
|
|
@ -892,15 +863,13 @@ PROVIDERS_WITHOUT_A_HANDLER = tuple(
|
|||
|
||||
MINIMAX_401_BODY = (
|
||||
'{"type":"error","error":{"type":"authorized_error","message":"login fail: Please carry the API secret key '
|
||||
"in the 'Authorization' field of the request header (1004)\",\"http_code\":\"401\"},"
|
||||
'in the \'Authorization\' field of the request header (1004)","http_code":"401"},'
|
||||
'"request_id":"06ddc9ba97ee6340e38f10e09787f547"}'
|
||||
)
|
||||
|
||||
|
||||
def _expected_for(provider: str, status_code: int) -> tuple[type[Exception], int]:
|
||||
return DEVIATIONS_FROM_THE_OPENAI_SHAPE.get(provider, {}).get(
|
||||
status_code, OPENAI_SHAPED[status_code]
|
||||
)
|
||||
return DEVIATIONS_FROM_THE_OPENAI_SHAPE.get(provider, {}).get(status_code, OPENAI_SHAPED[status_code])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -910,9 +879,7 @@ def quiet_exception_mapping(monkeypatch):
|
|||
|
||||
@pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES)
|
||||
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
|
||||
def test_an_upstream_status_maps_to_one_exception_per_provider(
|
||||
provider, status_code, quiet_exception_mapping
|
||||
):
|
||||
def test_an_upstream_status_maps_to_one_exception_per_provider(provider, status_code, quiet_exception_mapping):
|
||||
expected_class, expected_status = _expected_for(provider, status_code)
|
||||
|
||||
with pytest.raises(openai.APIError) as raised:
|
||||
|
|
@ -928,9 +895,7 @@ def test_an_upstream_status_maps_to_one_exception_per_provider(
|
|||
|
||||
@pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES)
|
||||
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
|
||||
def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from(
|
||||
provider, status_code, quiet_exception_mapping
|
||||
):
|
||||
def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from(provider, status_code, quiet_exception_mapping):
|
||||
with pytest.raises(openai.APIError) as raised:
|
||||
exception_type(
|
||||
model="test-model",
|
||||
|
|
@ -943,12 +908,8 @@ def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from(
|
|||
|
||||
|
||||
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
|
||||
def test_an_already_mapped_litellm_exception_passes_through_untouched(
|
||||
provider, quiet_exception_mapping
|
||||
):
|
||||
already_mapped = litellm.RateLimitError(
|
||||
message="already mapped", llm_provider=provider, model="test-model"
|
||||
)
|
||||
def test_an_already_mapped_litellm_exception_passes_through_untouched(provider, quiet_exception_mapping):
|
||||
already_mapped = litellm.RateLimitError(message="already mapped", llm_provider=provider, model="test-model")
|
||||
|
||||
returned = exception_type(
|
||||
model="test-model",
|
||||
|
|
@ -961,9 +922,7 @@ def test_an_already_mapped_litellm_exception_passes_through_untouched(
|
|||
|
||||
@pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES)
|
||||
@pytest.mark.parametrize("provider", PROVIDERS_WITHOUT_A_HANDLER)
|
||||
def test_a_provider_without_a_handler_maps_by_the_upstream_status(
|
||||
provider, status_code, quiet_exception_mapping
|
||||
):
|
||||
def test_a_provider_without_a_handler_maps_by_the_upstream_status(provider, status_code, quiet_exception_mapping):
|
||||
expected_class, expected_status = STATUS_KEYED[status_code]
|
||||
|
||||
with pytest.raises(openai.APIError) as raised:
|
||||
|
|
@ -1015,9 +974,7 @@ def test_an_unmapped_exception_with_no_model_or_provider_is_a_connection_error(q
|
|||
assert "boom" in raised.value.message
|
||||
|
||||
|
||||
def _raise_and_map(
|
||||
model: str | None, original_exception: Exception, custom_llm_provider: str | None
|
||||
) -> None:
|
||||
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."""
|
||||
try:
|
||||
|
|
@ -1058,9 +1015,7 @@ def test_an_unmapped_exception_with_no_model_or_provider_message_keeps_traceback
|
|||
|
||||
|
||||
CONTEXT_WINDOW_MESSAGE = "This model's maximum context length is 4096 tokens."
|
||||
CONTENT_POLICY_MESSAGE = (
|
||||
'{"error": {"type": "invalid_request_error", "code": "content_policy_violation"}}'
|
||||
)
|
||||
CONTENT_POLICY_MESSAGE = '{"error": {"type": "invalid_request_error", "code": "content_policy_violation"}}'
|
||||
TIMEOUT_MESSAGE = "Request timed out."
|
||||
|
||||
PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW = (
|
||||
|
|
@ -1103,15 +1058,11 @@ class _UpstreamErrorWithMessage(_UpstreamHTTPError):
|
|||
super().__init__(status_code=status_code)
|
||||
self.args = (message,)
|
||||
self.message = message
|
||||
self.response = httpx.Response(
|
||||
status_code=status_code, request=self.request, text=message
|
||||
)
|
||||
self.response = httpx.Response(status_code=status_code, request=self.request, text=message)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
|
||||
def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it(
|
||||
provider, quiet_exception_mapping
|
||||
):
|
||||
def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it(provider, quiet_exception_mapping):
|
||||
if provider in PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW:
|
||||
expected_class, expected_status = litellm.ContextWindowExceededError, 400
|
||||
else:
|
||||
|
|
@ -1129,9 +1080,7 @@ def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it(
|
|||
|
||||
|
||||
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
|
||||
def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it(
|
||||
provider, quiet_exception_mapping
|
||||
):
|
||||
def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it(provider, quiet_exception_mapping):
|
||||
if provider in PROVIDERS_THAT_RECOGNISE_A_CONTENT_POLICY_BLOCK:
|
||||
expected_class, expected_status = litellm.ContentPolicyViolationError, 400
|
||||
else:
|
||||
|
|
@ -1149,9 +1098,7 @@ def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it(
|
|||
|
||||
|
||||
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
|
||||
def test_a_timed_out_request_is_a_timeout_for_every_provider(
|
||||
provider, quiet_exception_mapping
|
||||
):
|
||||
def test_a_timed_out_request_is_a_timeout_for_every_provider(provider, quiet_exception_mapping):
|
||||
with pytest.raises(litellm.Timeout) as raised:
|
||||
exception_type(
|
||||
model="test-model",
|
||||
|
|
@ -1409,3 +1356,97 @@ def test_bedrock_timeout_mapping_keeps_retry_after_readable(status_code):
|
|||
exception_headers = _get_response_headers(original_exception=exc_info.value)
|
||||
assert exception_headers is not None
|
||||
assert litellm.utils._get_retry_after_from_exception_header(response_headers=exception_headers) == 7
|
||||
|
||||
|
||||
_GUARDRAIL_BLOCK_ERROR = {
|
||||
"message": "Content blocked: secret_project_codename pattern detected",
|
||||
"param": "None",
|
||||
"code": "400",
|
||||
"provider_specific_fields": {
|
||||
"error": "Content blocked: secret_project_codename pattern detected",
|
||||
"pattern": "secret_project_codename",
|
||||
"guardrail_name": "block-secret-project",
|
||||
"guardrail_mode": "pre_call",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _openai_handler_error(
|
||||
error_type: str,
|
||||
headers: dict[str, str] | list[tuple[str, str]],
|
||||
status_code: int = 400,
|
||||
message: str = _GUARDRAIL_BLOCK_ERROR["message"],
|
||||
) -> OpenAIError:
|
||||
wire_error = {**_GUARDRAIL_BLOCK_ERROR, "type": error_type, "code": str(status_code), "message": message}
|
||||
return OpenAIError(
|
||||
status_code=status_code,
|
||||
message=f"Error code: {status_code} - {{'error': {wire_error}}}",
|
||||
headers=httpx.Headers(headers),
|
||||
body=wire_error,
|
||||
)
|
||||
|
||||
|
||||
_PROXY_HEADERS = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guardrails": "block-secret-project"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("error_type", "status_code"), [("None", 400), ("invalid_request_error", 400), ("None", 422)])
|
||||
def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, status_code: int):
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
exception_type(
|
||||
model="claude-haiku-4-5",
|
||||
original_exception=_openai_handler_error(error_type, _PROXY_HEADERS, status_code=status_code),
|
||||
custom_llm_provider="litellm_proxy",
|
||||
completion_kwargs={},
|
||||
extra_kwargs={},
|
||||
)
|
||||
|
||||
assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project"
|
||||
assert exc_info.value.body["type"] == error_type
|
||||
assert dict(exc_info.value.response.headers) == _PROXY_HEADERS
|
||||
|
||||
|
||||
@pytest.mark.parametrize("relayed_class", [litellm.BadRequestError, litellm.ContentPolicyViolationError])
|
||||
def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_class: type[litellm.BadRequestError]):
|
||||
message = f"litellm.{relayed_class.__name__}: {_GUARDRAIL_BLOCK_ERROR['message']}"
|
||||
|
||||
with pytest.raises(relayed_class) as exc_info:
|
||||
exception_type(
|
||||
model="claude-haiku-4-5",
|
||||
original_exception=_openai_handler_error("None", _PROXY_HEADERS, message=message),
|
||||
custom_llm_provider="litellm_proxy",
|
||||
completion_kwargs={},
|
||||
extra_kwargs={},
|
||||
)
|
||||
|
||||
assert type(exc_info.value) is relayed_class
|
||||
assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project"
|
||||
assert dict(exc_info.value.response.headers) == _PROXY_HEADERS
|
||||
|
||||
|
||||
def test_openai_compatible_vendor_400_keeps_body_but_not_headers():
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
exception_type(
|
||||
model="gpt-5.4-mini",
|
||||
original_exception=_openai_handler_error("vendor_specific_error", {"openai-organization": "org-1"}),
|
||||
custom_llm_provider="openai",
|
||||
completion_kwargs={},
|
||||
extra_kwargs={},
|
||||
)
|
||||
|
||||
assert exc_info.value.body["type"] == "vendor_specific_error"
|
||||
assert not exc_info.value.response.headers
|
||||
|
||||
|
||||
def test_litellm_proxy_repeated_response_header_keeps_each_value():
|
||||
repeated = [("x-litellm-call-id", "call-guardrail"), ("set-cookie", "a=1"), ("set-cookie", "b=2")]
|
||||
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
exception_type(
|
||||
model="claude-haiku-4-5",
|
||||
original_exception=_openai_handler_error("None", repeated),
|
||||
custom_llm_provider="litellm_proxy",
|
||||
completion_kwargs={},
|
||||
extra_kwargs={},
|
||||
)
|
||||
|
||||
assert exc_info.value.response.headers.multi_items() == repeated
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import pytest
|
|||
|
||||
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
|
||||
from litellm.llms.anthropic.chat.guardrail_translation.handler import (
|
||||
AnthropicMessagesHandler,
|
||||
|
|
@ -635,14 +636,19 @@ class TestAnthropicMessagesHandlerInputProcessing:
|
|||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert guardrail.inputs is not None
|
||||
assert guardrail.inputs["texts"] == ["safe text", "prohibited correction"]
|
||||
assert guardrail.inputs["texts"] == [
|
||||
"trusted top-level system prompt",
|
||||
"safe text",
|
||||
"prohibited correction",
|
||||
]
|
||||
structured = guardrail.inputs["structured_messages"]
|
||||
assert [m["role"] for m in structured] == ["system", "user", "system"]
|
||||
assert structured[0]["content"] == "trusted top-level system prompt"
|
||||
assert data["system"] == "trusted top-level system prompt"
|
||||
assert data["messages"][1]["content"] == "[MASKED]"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_masking_slice_is_unavailable_when_top_level_system_is_included(
|
||||
async def test_bedrock_masking_slice_lines_up_when_top_level_system_is_included(
|
||||
self,
|
||||
):
|
||||
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
|
||||
|
|
@ -668,25 +674,25 @@ class TestAnthropicMessagesHandlerInputProcessing:
|
|||
structured = guardrail.inputs["structured_messages"]
|
||||
|
||||
bedrock = BedrockGuardrail(guardrailIdentifier="gi", guardrailVersion="1")
|
||||
assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts) + 1
|
||||
assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts)
|
||||
latest_user_index = bedrock._find_latest_message_index(structured, target_role="user")
|
||||
assert (
|
||||
bedrock._locate_message_texts_slice(
|
||||
structured_messages=structured,
|
||||
target_index=latest_user_index,
|
||||
texts=texts,
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert (
|
||||
bedrock._merge_masked_texts(
|
||||
masked_texts=["{MASKED}"],
|
||||
texts=texts,
|
||||
scanned_slice=None,
|
||||
scanned_role_subset=True,
|
||||
)
|
||||
== texts
|
||||
scanned_slice = bedrock._locate_message_texts_slice(
|
||||
structured_messages=structured,
|
||||
target_index=latest_user_index,
|
||||
texts=texts,
|
||||
)
|
||||
assert scanned_slice == (3, 1)
|
||||
assert bedrock._merge_masked_texts(
|
||||
masked_texts=["{MASKED}"],
|
||||
texts=texts,
|
||||
scanned_slice=scanned_slice,
|
||||
scanned_role_subset=True,
|
||||
) == [
|
||||
"trusted top-level system prompt",
|
||||
"safe text",
|
||||
"prohibited correction",
|
||||
"{MASKED}",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("skip_system_message_in_guardrail", [True, None])
|
||||
|
|
@ -1611,7 +1617,8 @@ class TestAnthropicMessagesIncrementalScan:
|
|||
)
|
||||
assert mock_api.call_count == 1
|
||||
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [
|
||||
"What is the capital of France?"
|
||||
"You are a helpful geography assistant.",
|
||||
"What is the capital of France?",
|
||||
]
|
||||
mock_api.reset_mock()
|
||||
await handler.process_input_messages(
|
||||
|
|
@ -2150,6 +2157,213 @@ class TestAnthropicMessagesScanOnlyToolResults:
|
|||
assert guardrail.captured_inputs.get("images") == ["TOOL_IMG"]
|
||||
|
||||
|
||||
class ToolCallArgumentsMaskingGuardrail(InputsRecordingGuardrail):
|
||||
"""Masks the canary inside tool-call arguments, in place or through a fresh list of plain dicts."""
|
||||
|
||||
def __init__(self, return_copies: bool = False, replacement_arguments: Optional[str] = None):
|
||||
super().__init__()
|
||||
self.return_copies = return_copies
|
||||
self.replacement_arguments = replacement_arguments
|
||||
self.seen_tool_calls: list[dict[str, object]] = []
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[LiteLLMLoggingObj] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
outputs = await super().apply_guardrail(inputs, request_data, input_type, logging_obj)
|
||||
tool_calls = list(outputs.get("tool_calls") or [])
|
||||
self.seen_tool_calls.extend(json.loads(json.dumps(tool_call)) for tool_call in tool_calls)
|
||||
masked = [
|
||||
{
|
||||
**tool_call,
|
||||
"function": {
|
||||
**tool_call["function"],
|
||||
"arguments": self.replacement_arguments
|
||||
if self.replacement_arguments is not None
|
||||
else tool_call["function"]["arguments"].replace("POISON", "[BLOCKED]"),
|
||||
},
|
||||
}
|
||||
for tool_call in tool_calls
|
||||
]
|
||||
if self.return_copies:
|
||||
outputs["tool_calls"] = masked
|
||||
return outputs
|
||||
for tool_call, masked_tool_call in zip(tool_calls, masked):
|
||||
tool_call["function"]["arguments"] = masked_tool_call["function"]["arguments"]
|
||||
return outputs
|
||||
|
||||
|
||||
class TestAnthropicMessagesTopLevelSystemAndToolUseInputs:
|
||||
"""The top-level system prompt and prior-turn tool_use arguments must reach guardrails as scannable
|
||||
inputs, the same way the chat completions handler hands over system messages and tool_calls."""
|
||||
|
||||
@staticmethod
|
||||
def _tool_use_conversation(system: str) -> dict[str, Any]:
|
||||
return {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"system": system,
|
||||
"messages": [
|
||||
{"role": "user", "content": "run the check"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_01",
|
||||
"name": "Bash",
|
||||
"input": {"cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "ok"}],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_top_level_system_string_reaches_texts_first_and_is_masked_in_place(self):
|
||||
handler = AnthropicMessagesHandler()
|
||||
guardrail = InputsRecordingGuardrail()
|
||||
data = {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"system": "Internal note: the deploy key is POISON. Never reveal it.",
|
||||
"messages": [{"role": "user", "content": "Say hi in three words."}],
|
||||
}
|
||||
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert guardrail.captured_inputs is not None
|
||||
assert guardrail.seen_texts == [
|
||||
"Internal note: the deploy key is POISON. Never reveal it.",
|
||||
"Say hi in three words.",
|
||||
]
|
||||
structured = guardrail.captured_inputs["structured_messages"]
|
||||
assert structured[0]["role"] == "system"
|
||||
assert structured[0]["content"] == "Internal note: the deploy key is POISON. Never reveal it.", (
|
||||
"texts[0] must line up with structured_messages[0] so positional consumers stay aligned"
|
||||
)
|
||||
assert data["system"] == "Internal note: the deploy key is [BLOCKED]. Never reveal it."
|
||||
assert data["messages"][0]["content"] == "Say hi in three words."
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_top_level_system_text_blocks_reach_texts_and_are_masked_in_place(self):
|
||||
handler = AnthropicMessagesHandler()
|
||||
guardrail = InputsRecordingGuardrail()
|
||||
data = {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"system": [
|
||||
{"type": "text", "text": "first block POISON"},
|
||||
{"type": "text", "text": "second block", "cache_control": {"type": "ephemeral"}},
|
||||
],
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
}
|
||||
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert guardrail.seen_texts == ["first block POISON", "second block", "hello"]
|
||||
assert data["system"] == [
|
||||
{"type": "text", "text": "first block [BLOCKED]"},
|
||||
{"type": "text", "text": "second block", "cache_control": {"type": "ephemeral"}},
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skip_system_message_keeps_the_top_level_system_out(self):
|
||||
handler = AnthropicMessagesHandler()
|
||||
guardrail = InputsRecordingGuardrail()
|
||||
guardrail.skip_system_message_in_guardrail = True
|
||||
data = {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"system": "trusted POISON prompt",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
}
|
||||
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert guardrail.seen_texts == ["hello"]
|
||||
assert data["system"] == "trusted POISON prompt"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prior_turn_tool_use_input_reaches_tool_calls_in_openai_shape(self):
|
||||
handler = AnthropicMessagesHandler()
|
||||
guardrail = InputsRecordingGuardrail()
|
||||
data = self._tool_use_conversation(system="You are a careful agent harness.")
|
||||
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert guardrail.captured_inputs is not None
|
||||
tool_calls = guardrail.captured_inputs.get("tool_calls")
|
||||
assert tool_calls is not None and len(tool_calls) == 1
|
||||
assert tool_calls[0]["id"] == "toolu_01"
|
||||
assert tool_calls[0]["type"] == "function"
|
||||
assert tool_calls[0]["function"]["name"] == "Bash"
|
||||
assert json.loads(tool_calls[0]["function"]["arguments"]) == {
|
||||
"cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity"
|
||||
}
|
||||
assert data["messages"][1]["content"][0]["input"] == {
|
||||
"cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity"
|
||||
}, "a guardrail that leaves tool_calls alone must leave the tool_use input alone"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("return_copies", [False, True])
|
||||
async def test_masked_tool_call_arguments_write_back_into_the_tool_use_input(self, return_copies: bool):
|
||||
handler = AnthropicMessagesHandler()
|
||||
guardrail = ToolCallArgumentsMaskingGuardrail(return_copies=return_copies)
|
||||
data = self._tool_use_conversation(system="You are a careful agent harness.")
|
||||
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert [tool_call["function"]["name"] for tool_call in guardrail.seen_tool_calls] == ["Bash"]
|
||||
tool_use = data["messages"][1]["content"][0]
|
||||
assert tool_use == {
|
||||
"type": "tool_use",
|
||||
"id": "toolu_01",
|
||||
"name": "Bash",
|
||||
"input": {"cmd": "AWS_ACCESS_KEY_ID=[BLOCKED] aws sts get-caller-identity"},
|
||||
}
|
||||
assert data["messages"][2]["content"][0]["tool_use_id"] == "toolu_01"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_json_rewritten_arguments_are_rejected_by_name(self):
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
|
||||
|
||||
handler = AnthropicMessagesHandler()
|
||||
guardrail = ToolCallArgumentsMaskingGuardrail(replacement_arguments="[REDACTED]")
|
||||
data = self._tool_use_conversation(system="Internal note: the deploy key is POISON. Never reveal it.")
|
||||
data["messages"][2]["content"][0]["content"] = "fetched POISON page"
|
||||
original = json.loads(json.dumps(data))
|
||||
|
||||
with pytest.raises(UnappliableRequestRewrite) as excinfo:
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert excinfo.value.guardrail_name == "scan-only-capture"
|
||||
assert data["system"] == original["system"], "a rejected rewrite must leave the request untouched"
|
||||
assert data["messages"] == original["messages"], "a rejected rewrite must leave the request untouched"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scan_only_tool_results_keeps_system_and_tool_use_out(self):
|
||||
handler = AnthropicMessagesHandler()
|
||||
guardrail = InputsRecordingGuardrail()
|
||||
guardrail.scan_only_tool_results = True
|
||||
data = self._tool_use_conversation(system="trusted POISON prompt")
|
||||
data["messages"][2]["content"][0]["content"] = "fetched POISON page"
|
||||
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert guardrail.seen_texts == ["fetched POISON page"]
|
||||
assert guardrail.captured_inputs is not None
|
||||
assert guardrail.captured_inputs.get("tool_calls") is None
|
||||
assert data["system"] == "trusted POISON prompt"
|
||||
assert data["messages"][1]["content"][0]["input"] == {
|
||||
"cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity"
|
||||
}
|
||||
assert data["messages"][2]["content"][0]["content"] == "fetched [BLOCKED] page"
|
||||
|
||||
|
||||
class TestStructuredWriteBackKeepsToolResults:
|
||||
"""A guardrail rewrite must never leave a tool_use without its tool_result (Claude Code ToolSearch, LIT-6103)."""
|
||||
|
||||
|
|
@ -2272,6 +2486,116 @@ class TestAnthropicMessagesHandlerStreamingScanKey:
|
|||
assert ended_key != open_key
|
||||
|
||||
|
||||
class PerRowTextGuardrail(CustomGuardrail):
|
||||
"""Answers one redacted text per chat row it was shown, the way a guardrail
|
||||
that scans per message does, and hands back only texts."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(guardrail_name="per-row-redactor")
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
rows = inputs.get("structured_messages") or []
|
||||
return {**inputs, "texts": [str(row.get("content")).replace("123-45-6789", "<US_SSN>") for row in rows]}
|
||||
|
||||
|
||||
class PerSlotTextGuardrail(CustomGuardrail):
|
||||
"""Answers one redacted text per text slot of every chat row it was shown, the
|
||||
way a guardrail that counts slots per message does, and hands back only texts."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(guardrail_name="per-slot-redactor")
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts
|
||||
|
||||
rows = inputs.get("structured_messages") or []
|
||||
return {
|
||||
**inputs,
|
||||
"texts": [text.replace("123-45-6789", "<US_SSN>") for row in rows for text in message_slot_texts(row)],
|
||||
}
|
||||
|
||||
|
||||
class TestPerMessageTextWriteBack:
|
||||
"""Texts that no longer pair one-to-one with what the handler extracted must be
|
||||
rejected by name instead of sliding onto the wrong messages."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_text_per_row_over_a_system_prompt_is_applied(self):
|
||||
data = {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"system": "Reply with exactly the SSN you were given.",
|
||||
"messages": [{"role": "user", "content": "My SSN is 123-45-6789."}],
|
||||
}
|
||||
|
||||
await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerRowTextGuardrail())
|
||||
|
||||
assert data["system"] == "Reply with exactly the SSN you were given."
|
||||
assert data["messages"] == [{"role": "user", "content": "My SSN is <US_SSN>."}]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_text_per_row_over_a_multi_block_system_prompt_is_rejected_by_name(self):
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
|
||||
|
||||
data = {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"system": [
|
||||
{"type": "text", "text": "Reply with exactly the SSN you were given."},
|
||||
{"type": "text", "text": "Never apologize."},
|
||||
],
|
||||
"messages": [{"role": "user", "content": "My SSN is 123-45-6789."}],
|
||||
}
|
||||
original = json.loads(json.dumps(data))
|
||||
|
||||
with pytest.raises(UnappliableRequestRewrite) as excinfo:
|
||||
await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerRowTextGuardrail())
|
||||
|
||||
assert excinfo.value.guardrail_name == "per-row-redactor"
|
||||
assert data["system"] == original["system"], "a rejected rewrite must leave the request untouched"
|
||||
assert data["messages"] == original["messages"], "a rejected rewrite must leave the request untouched"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_text_per_slot_over_a_system_prompt_with_an_empty_block_is_applied(self):
|
||||
data = {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"system": [
|
||||
{"type": "text", "text": ""},
|
||||
{"type": "text", "text": "Reply with exactly the SSN you were given."},
|
||||
],
|
||||
"messages": [{"role": "user", "content": "My SSN is 123-45-6789."}],
|
||||
}
|
||||
|
||||
await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerSlotTextGuardrail())
|
||||
|
||||
assert data["system"] == [
|
||||
{"type": "text", "text": ""},
|
||||
{"type": "text", "text": "Reply with exactly the SSN you were given."},
|
||||
]
|
||||
assert data["messages"] == [{"role": "user", "content": "My SSN is <US_SSN>."}]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_text_per_row_without_a_system_prompt_is_applied(self):
|
||||
data = {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [{"role": "user", "content": "My SSN is 123-45-6789."}],
|
||||
}
|
||||
|
||||
await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerRowTextGuardrail())
|
||||
|
||||
assert data["messages"] == [{"role": "user", "content": "My SSN is <US_SSN>."}]
|
||||
|
||||
|
||||
class TestAnthropicMessagesHandlerPostCallHookResponse:
|
||||
def test_openai_shaped_stream_assembly_reaches_the_hook_as_a_messages_response(self):
|
||||
from litellm.types.utils import Choices, Message, ModelResponse, Usage
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ Without the fix, the AnthropicStreamWrapper silently dropped these
|
|||
arguments, causing tool_use blocks to arrive with empty input {}.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
from typing import List
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
|
|
@ -139,9 +141,7 @@ async def test_async_stream_emits_input_json_delta_for_bundled_tool_args():
|
|||
|
||||
# Verify the delta carries the tool arguments
|
||||
delta_event = events[input_json_delta_idx]
|
||||
assert delta_event["delta"][
|
||||
"partial_json"
|
||||
], "input_json_delta should have non-empty partial_json"
|
||||
assert json.loads(delta_event["delta"]["partial_json"]) == {"location": "Boston"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -300,7 +300,7 @@ def test_sync_stream_emits_input_json_delta_for_bundled_tool_args():
|
|||
assert (
|
||||
input_json_delta_idx == tool_start_idx + 1
|
||||
), "input_json_delta should immediately follow the tool_use content_block_start"
|
||||
assert events[input_json_delta_idx]["delta"]["partial_json"]
|
||||
assert json.loads(events[input_json_delta_idx]["delta"]["partial_json"]) == {"location": "Boston"}
|
||||
|
||||
|
||||
def test_sync_stream_no_extra_delta_when_tool_args_empty():
|
||||
|
|
|
|||
|
|
@ -773,6 +773,12 @@ class TestBedrockMantleCodexAdditionalTools:
|
|||
assert body["input"] == codex_agentic_items
|
||||
assert "tools" not in body
|
||||
|
||||
def test_input_without_additional_tools_sanitizes_tools_on_the_caller_params_object(self):
|
||||
params = {"tools": [{"type": "function", "name": "wait", "parameters": '{"type": "object"}'}]}
|
||||
body = self._transform(input=[self._USER_MESSAGE], params=params)
|
||||
assert body["tools"][0]["parameters"] == {"type": "object"}
|
||||
assert params["tools"][0]["parameters"] == {"type": "object"}
|
||||
|
||||
def test_malformed_additional_tools_item_without_tools_list_is_stripped(self):
|
||||
body = self._transform(
|
||||
input=[
|
||||
|
|
|
|||
|
|
@ -1893,6 +1893,183 @@ class TestScanOnlyToolResults:
|
|||
assert data["messages"][4]["content"] == "and then?"
|
||||
|
||||
|
||||
class TestNoScannableContentRecordsNotRun:
|
||||
"""LIT-6314: a guardrail whose scoping leaves nothing to scan must still persist an evaluation record"""
|
||||
|
||||
def _system_only_data(self) -> dict:
|
||||
return {"messages": [{"role": "system", "content": "SYSTEM-PROMPT"}]}
|
||||
|
||||
def _recorded_entries(self, data: dict) -> list:
|
||||
metadata = data.get("metadata") or data.get("litellm_metadata") or {}
|
||||
return metadata.get("standard_logging_guardrail_information") or []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skipped_scan_records_not_run_entry(self):
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
guardrail = MockGuardrail(guardrail_name="skip-system-guardrail")
|
||||
guardrail.skip_system_message_in_guardrail = True
|
||||
data = self._system_only_data()
|
||||
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert guardrail.last_inputs is None, "nothing survived scoping, apply_guardrail must not run"
|
||||
entries = self._recorded_entries(data)
|
||||
assert len(entries) == 1
|
||||
assert entries[0]["guardrail_name"] == "skip-system-guardrail"
|
||||
assert entries[0]["guardrail_status"] == "not_run"
|
||||
assert entries[0]["guardrail_response"] == "no scannable content after message scoping"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("skip_system", [False, True])
|
||||
async def test_empty_content_does_not_blame_scoping(self, skip_system: bool):
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
guardrail = MockGuardrail(guardrail_name="unscoped-guardrail")
|
||||
guardrail.skip_system_message_in_guardrail = skip_system
|
||||
data = {"messages": [{"role": "user", "content": None}]}
|
||||
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert guardrail.last_inputs is None
|
||||
entries = self._recorded_entries(data)
|
||||
assert len(entries) == 1
|
||||
assert entries[0]["guardrail_status"] == "not_run"
|
||||
assert entries[0]["guardrail_response"] == "no scannable content"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_self_recording_guardrail_is_left_alone(self):
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
guardrail = MockGuardrail(guardrail_name="self-recording-guardrail")
|
||||
guardrail.skip_system_message_in_guardrail = True
|
||||
guardrail.records_own_guardrail_information = True
|
||||
data = self._system_only_data()
|
||||
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert guardrail.last_inputs is None
|
||||
assert self._recorded_entries(data) == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scannable_content_records_no_extra_entry(self):
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
guardrail = MockGuardrail(guardrail_name="normal-guardrail")
|
||||
data = {"messages": [{"role": "user", "content": "hello"}]}
|
||||
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert guardrail.last_inputs is not None
|
||||
assert all(e.get("guardrail_status") != "not_run" for e in self._recorded_entries(data))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_only_content_is_not_reported_as_not_run(self):
|
||||
"""Images are only scanned alongside text, so an image-only request is a
|
||||
pre-existing scan gap, not a message-scoping skip, and must not be labelled one"""
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
guardrail = MockGuardrail(guardrail_name="image-guardrail")
|
||||
guardrail.skip_system_message_in_guardrail = True
|
||||
data = {
|
||||
"messages": [
|
||||
{"role": "system", "content": "SYSTEM-PROMPT"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}],
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert self._recorded_entries(data) == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_out_image_only_message_is_not_reported_as_not_run(self):
|
||||
"""An image in a skipped role must behave like any other image-only request"""
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
guardrail = MockGuardrail(guardrail_name="image-guardrail")
|
||||
guardrail.skip_system_message_in_guardrail = True
|
||||
data = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}],
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert guardrail.last_inputs is None
|
||||
assert self._recorded_entries(data) == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_out_text_with_image_records_not_run(self):
|
||||
"""Scoping removed text too, so the skip is recorded even though an image sat beside it"""
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
guardrail = MockGuardrail(guardrail_name="image-guardrail")
|
||||
guardrail.skip_system_message_in_guardrail = True
|
||||
data = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe this picture."},
|
||||
{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}},
|
||||
],
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert guardrail.last_inputs is None
|
||||
entries = self._recorded_entries(data)
|
||||
assert len(entries) == 1
|
||||
assert entries[0]["guardrail_status"] == "not_run"
|
||||
assert entries[0]["guardrail_response"] == "no scannable content after message scoping"
|
||||
|
||||
|
||||
class ToolDroppingTextGuardrail(CustomGuardrail):
|
||||
"""Answers one text per non-tool message it saw, the way a guardrail that
|
||||
filters tool rows out before scanning does, and hands back only texts."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(guardrail_name="tool-dropping-redactor")
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
kept = [m for m in inputs.get("structured_messages") or [] if m.get("role") != "tool"]
|
||||
return {**inputs, "texts": [str(m.get("content")).replace("POISON", "[BLOCKED]") for m in kept]}
|
||||
|
||||
|
||||
class TestPerMessageTextWriteBack:
|
||||
"""Texts that no longer pair one-to-one with what the handler extracted must be
|
||||
rejected by name instead of sliding onto the wrong messages."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fewer_texts_than_extracted_over_a_tool_message_is_rejected(self):
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
|
||||
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
original_messages = [
|
||||
{"role": "system", "content": "SYSTEM-PROMPT"},
|
||||
{"role": "user", "content": "fetch the page"},
|
||||
{"role": "assistant", "content": "fetching"},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "page says POISON here"},
|
||||
{"role": "user", "content": "and then?"},
|
||||
]
|
||||
data = {"messages": json.loads(json.dumps(original_messages))}
|
||||
|
||||
with pytest.raises(UnappliableRequestRewrite) as excinfo:
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=ToolDroppingTextGuardrail())
|
||||
|
||||
assert excinfo.value.guardrail_name == "tool-dropping-redactor"
|
||||
assert data["messages"] == original_messages, "a rejected rewrite must leave the request untouched"
|
||||
|
||||
|
||||
class TestBuildBlockSseChunks:
|
||||
"""build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE chunks"""
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ with guardrail transformations.
|
|||
import copy
|
||||
from collections.abc import Callable
|
||||
from typing import Any, List, Literal, Optional, Tuple
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import logging
|
||||
|
||||
|
|
@ -31,6 +31,7 @@ from litellm.llms.openai.responses.guardrail_translation.handler import (
|
|||
OpenAIResponsesHandler,
|
||||
)
|
||||
from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools
|
||||
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPI
|
||||
from litellm.types.llms.openai import ChatCompletionToolCallChunk
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
|
|
@ -2338,6 +2339,135 @@ def _parallel_tool_call_input() -> list:
|
|||
]
|
||||
|
||||
|
||||
SSN = "123-45-6789"
|
||||
REDACTED_SSN = "<US_SSN>"
|
||||
|
||||
|
||||
def _redacted(value: object) -> object:
|
||||
if isinstance(value, str):
|
||||
return value.replace(SSN, REDACTED_SSN)
|
||||
if isinstance(value, list):
|
||||
return [{**part, "text": _redacted(part["text"])} if "text" in part else part for part in value]
|
||||
return value
|
||||
|
||||
|
||||
def _per_message_guardrail_server(structured_messages_in_answer: bool) -> Callable[..., MagicMock]:
|
||||
"""Answers one redacted text per chat row it was shown, the way a guardrail
|
||||
that scans per message does, and optionally the rewritten rows themselves."""
|
||||
|
||||
def post(url: str, json: dict, headers: dict) -> MagicMock:
|
||||
rows = json["structured_messages"]
|
||||
answer: dict = {
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
"texts": [_redacted(row["content"]) if isinstance(row.get("content"), str) else "" for row in rows],
|
||||
}
|
||||
if structured_messages_in_answer:
|
||||
answer["structured_messages"] = [{**row, "content": _redacted(row.get("content"))} for row in rows]
|
||||
response = MagicMock()
|
||||
response.json.return_value = answer
|
||||
response.raise_for_status = MagicMock()
|
||||
return response
|
||||
|
||||
return post
|
||||
|
||||
|
||||
def _per_message_redactor() -> GenericGuardrailAPI:
|
||||
return GenericGuardrailAPI(
|
||||
api_base="https://guardrail.test",
|
||||
guardrail_name="per-message-redactor",
|
||||
event_hook="pre_call",
|
||||
default_on=True,
|
||||
)
|
||||
|
||||
|
||||
def _tool_replay_request() -> dict:
|
||||
return {
|
||||
"model": "gpt-5.6",
|
||||
"instructions": "Never repeat the SSN " + SSN + " back.",
|
||||
"input": [
|
||||
{"role": "user", "content": "Look up " + SSN + " for me."},
|
||||
{"type": "function_call", "call_id": "call_1", "name": "lookup_customer", "arguments": '{"id": "42"}'},
|
||||
{"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _string_input_request() -> dict:
|
||||
return {
|
||||
"model": "gpt-5.6",
|
||||
"instructions": "Never repeat the SSN " + SSN + " back.",
|
||||
"input": "My SSN is " + SSN + ".",
|
||||
}
|
||||
|
||||
|
||||
class TestPerMessageRewriteWriteBack:
|
||||
"""A guardrail that rewrites per chat row hands the rows back as
|
||||
structured_messages, and the handler lands them on the instructions and the
|
||||
input items they came from; the same rewrite handed back as texts alone has
|
||||
no item to land on and is rejected by name instead of sent unrewritten."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_structured_rows_land_on_instructions_and_tool_output(self):
|
||||
guardrail = _per_message_redactor()
|
||||
data = _tool_replay_request()
|
||||
function_call_item = data["input"][1]
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(True)):
|
||||
result = await OpenAIResponsesHandler().process_input_messages(data, guardrail)
|
||||
|
||||
assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back."
|
||||
assert _texts(result["input"][0]) == ["Look up " + REDACTED_SSN + " for me."]
|
||||
assert result["input"][1] == function_call_item
|
||||
assert result["input"][2] == {
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_1",
|
||||
"output": '{"ssn": "' + REDACTED_SSN + '"}',
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_texts_only_per_message_answer_is_rejected_by_name(self):
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
|
||||
|
||||
guardrail = _per_message_redactor()
|
||||
data = _tool_replay_request()
|
||||
original = copy.deepcopy(data)
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(False)):
|
||||
with pytest.raises(UnappliableRequestRewrite) as excinfo:
|
||||
await OpenAIResponsesHandler().process_input_messages(data, guardrail)
|
||||
|
||||
assert excinfo.value.guardrail_name == "per-message-redactor"
|
||||
assert data["input"] == original["input"]
|
||||
assert data["instructions"] == original["instructions"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_structured_rows_land_on_instructions_and_string_input(self):
|
||||
guardrail = _per_message_redactor()
|
||||
data = _string_input_request()
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(True)):
|
||||
result = await OpenAIResponsesHandler().process_input_messages(data, guardrail)
|
||||
|
||||
assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back."
|
||||
assert [_texts(item) for item in result["input"]] == [["My SSN is " + REDACTED_SSN + "."]]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_texts_only_per_message_answer_over_a_string_input_is_rejected_by_name(self):
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
|
||||
|
||||
guardrail = _per_message_redactor()
|
||||
data = _string_input_request()
|
||||
original = copy.deepcopy(data)
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(False)):
|
||||
with pytest.raises(UnappliableRequestRewrite) as excinfo:
|
||||
await OpenAIResponsesHandler().process_input_messages(data, guardrail)
|
||||
|
||||
assert excinfo.value.guardrail_name == "per-message-redactor"
|
||||
assert data["input"] == original["input"]
|
||||
assert data["instructions"] == original["instructions"]
|
||||
|
||||
|
||||
class TestProvenancePatching:
|
||||
"""The O(n) provenance pass must keep patching rewritten rows in place for the
|
||||
shapes real agent loops produce, and fall back safely everywhere else."""
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue