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_model_access_denied_message
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> # Conflicts: # tests/test_litellm/proxy/auth/test_handle_jwt.py
This commit is contained in:
commit
5fee1c8710
238 changed files with 15495 additions and 7143 deletions
|
|
@ -9,7 +9,7 @@ commands:
|
|||
parameters:
|
||||
category:
|
||||
type: enum
|
||||
enum: ["backend", "client"]
|
||||
enum: ["backend", "client", "provider-harness"]
|
||||
default: "backend"
|
||||
steps:
|
||||
- run:
|
||||
|
|
@ -2918,19 +2918,30 @@ jobs:
|
|||
provider_replay_harness:
|
||||
docker:
|
||||
- *python312_image
|
||||
- image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f
|
||||
working_directory: ~/project
|
||||
resource_class: medium
|
||||
environment:
|
||||
E2E_CACHE_TEST_REDIS_URL: redis://127.0.0.1:6379/0
|
||||
E2E_PROVIDER_CACHE: "0"
|
||||
E2E_FIXTURE_MODE: live
|
||||
steps:
|
||||
- checkout
|
||||
- skip_if_unrelated_changes:
|
||||
category: provider-harness
|
||||
- setup_litellm_test_deps
|
||||
- wait_for_service:
|
||||
url: tcp://localhost:6379
|
||||
- run:
|
||||
name: Test provider replay harness
|
||||
name: Test provider capture and 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
|
||||
tests/code_coverage_tests/test_provider_replay_harness.py \
|
||||
tests/code_coverage_tests/test_provider_cache.py
|
||||
- store_test_results:
|
||||
path: test-results/provider-replay-harness
|
||||
|
||||
|
|
|
|||
|
|
@ -1,13 +1,19 @@
|
|||
#!/usr/bin/env bash
|
||||
set -uo pipefail
|
||||
|
||||
category="${1:?usage: classify_changes.sh <backend|client|ui>}"
|
||||
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness>}"
|
||||
|
||||
has_client=false
|
||||
has_backend=false
|
||||
has_ci=false
|
||||
has_provider_harness=false
|
||||
while IFS= read -r file || [ -n "$file" ]; do
|
||||
[ -n "$file" ] || continue
|
||||
case "$file" in
|
||||
tests/e2e/*/*.py) : ;;
|
||||
tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/test_litellm/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
|
||||
has_provider_harness=true ;;
|
||||
esac
|
||||
case "$file" in
|
||||
ui/* | tests/e2e/ui/*) has_client=true ;;
|
||||
docs/* | *.md | *.mdx) : ;;
|
||||
|
|
@ -17,6 +23,9 @@ while IFS= read -r file || [ -n "$file" ]; do
|
|||
done
|
||||
|
||||
case "$category" in
|
||||
provider-harness)
|
||||
[ "$has_provider_harness" = true ] && echo run || echo skip
|
||||
;;
|
||||
backend)
|
||||
[ "$has_backend" = true ] && echo run || echo skip
|
||||
;;
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
#!/usr/bin/env bash
|
||||
set -uo pipefail
|
||||
|
||||
category="${1:?usage: path_filter.sh <backend|client>}"
|
||||
category="${1:?usage: path_filter.sh <backend|client|provider-harness>}"
|
||||
here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
run_full() {
|
||||
|
|
@ -36,5 +36,5 @@ if [ "$decision" = run ]; then
|
|||
run_full "$category-relevant changes detected"
|
||||
fi
|
||||
|
||||
echo "path-filter[$category]: only unrelated (docs/client) changes detected; halting job as successful"
|
||||
echo "path-filter[$category]: only unrelated changes detected; halting job as successful"
|
||||
circleci-agent step halt
|
||||
|
|
|
|||
2
.github/workflows/image-scan.yml
vendored
2
.github/workflows/image-scan.yml
vendored
|
|
@ -26,6 +26,7 @@ on:
|
|||
- ui/Dockerfile
|
||||
- ui/nginx.conf
|
||||
- .github/workflows/image-scan.yml
|
||||
- .grype.yaml
|
||||
schedule:
|
||||
- cron: "41 6 * * *"
|
||||
workflow_dispatch:
|
||||
|
|
@ -93,6 +94,7 @@ jobs:
|
|||
GRYPE_MATCH_PYTHON_USING_CPES: "true"
|
||||
run: |
|
||||
"$RUNNER_TEMP/grype" litellm-image-scan:${{ github.sha }} \
|
||||
--config .grype.yaml \
|
||||
--only-fixed \
|
||||
--fail-on high \
|
||||
--output table
|
||||
|
|
|
|||
13
.grype.yaml
Normal file
13
.grype.yaml
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
# Wolfi's security database names zlib 1.3.3-r0 as the fix for CVE-2026-85091,
|
||||
# but the newest zlib published to the Wolfi apk repo is 1.3.2-r7, so every
|
||||
# wolfi-base digest reports it and no `apk upgrade` can clear it.
|
||||
# Drop this once Wolfi ships zlib >= 1.3.3-r0; expected by 2026-10-15.
|
||||
ignore:
|
||||
- vulnerability: CVE-2026-85091
|
||||
package:
|
||||
name: zlib
|
||||
type: apk
|
||||
- vulnerability: GHSA-g5fp-32jq-cfw2
|
||||
package:
|
||||
name: zlib
|
||||
type: apk
|
||||
|
|
@ -96,6 +96,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/langfuse/",
|
||||
"/vllm/",
|
||||
"/mistral/",
|
||||
"/nvidia_nim/",
|
||||
"/groq/",
|
||||
"/voyage/",
|
||||
"/cursor/",
|
||||
|
|
|
|||
|
|
@ -40,4 +40,4 @@ if not logger.handlers:
|
|||
logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
|
||||
)
|
||||
logger.addHandler(handler)
|
||||
logger.setLevel(logging.INFO)
|
||||
logger.setLevel(os.getenv("LITELLM_LOG", "INFO").upper())
|
||||
|
|
|
|||
|
|
@ -0,0 +1,23 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "total_response_time_ms" BIGINT NOT NULL DEFAULT 0;
|
||||
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "timed_requests" BIGINT NOT NULL DEFAULT 0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "total_response_time_ms" BIGINT NOT NULL DEFAULT 0;
|
||||
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "timed_requests" BIGINT NOT NULL DEFAULT 0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "total_response_time_ms" BIGINT NOT NULL DEFAULT 0;
|
||||
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "timed_requests" BIGINT NOT NULL DEFAULT 0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "total_response_time_ms" BIGINT NOT NULL DEFAULT 0;
|
||||
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "timed_requests" BIGINT NOT NULL DEFAULT 0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "total_response_time_ms" BIGINT NOT NULL DEFAULT 0;
|
||||
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "timed_requests" BIGINT NOT NULL DEFAULT 0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "total_response_time_ms" BIGINT NOT NULL DEFAULT 0;
|
||||
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "timed_requests" BIGINT NOT NULL DEFAULT 0;
|
||||
|
|
@ -801,6 +801,8 @@ model LiteLLM_DailyUserSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
|
|
@ -837,6 +839,8 @@ model LiteLLM_DailyOrganizationSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
|
|
@ -873,6 +877,8 @@ model LiteLLM_DailyEndUserSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
|
|
@ -908,6 +914,8 @@ model LiteLLM_DailyAgentSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
|
|
@ -943,6 +951,8 @@ model LiteLLM_DailyTeamSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
ptu_flat_cost Float @default(0.0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
|
@ -981,6 +991,8 @@ model LiteLLM_DailyTagSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
|
|
|
|||
|
|
@ -401,10 +401,14 @@ def _parse_json_logs_env(value: str | None) -> bool:
|
|||
return (value or "").lower() == "true"
|
||||
|
||||
|
||||
def resolve_log_level(log_level: str) -> int:
|
||||
return getattr(logging, log_level.upper())
|
||||
|
||||
|
||||
json_logs: Final = _parse_json_logs_env(os.getenv("JSON_LOGS"))
|
||||
# Create a handler for the logger (you may need to adapt this based on your needs)
|
||||
log_level: Final = os.getenv("LITELLM_LOG", "DEBUG")
|
||||
numeric_level: Final[str] = getattr(logging, log_level.upper())
|
||||
numeric_level: Final[int] = resolve_log_level(log_level)
|
||||
handler: Final = LevelRoutingStreamHandler()
|
||||
handler.setLevel(numeric_level)
|
||||
handler.addFilter(_secret_filter)
|
||||
|
|
|
|||
|
|
@ -364,6 +364,8 @@ GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS: Final = int(
|
|||
os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60)
|
||||
)
|
||||
BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS: Final = 25_000
|
||||
CONTENT_FILTER_STREAMING_HOLDBACK_CHARS: Final = 50
|
||||
CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS: Final = 512
|
||||
DEFAULT_PRESIDIO_ANALYZE_CHUNK_SIZE_BYTES: Final = 500_000
|
||||
PRESIDIO_ANALYZE_CHUNK_OVERLAP_CHARS: Final = 4096
|
||||
PRESIDIO_ANALYZE_CHUNK_CONCURRENCY: Final = 8
|
||||
|
|
|
|||
|
|
@ -1203,6 +1203,24 @@ def _without_provider_stated_cost(usage: Usage | None) -> Usage | None:
|
|||
return usage.model_copy(update=MappingProxyType({"cost": None}))
|
||||
|
||||
|
||||
def _split_responses_ws_logging_object_by_service_tier(
|
||||
completion_response: LiteLLMRealtimeStreamLoggingObject,
|
||||
) -> tuple[LiteLLMRealtimeStreamLoggingObject, ...] | None:
|
||||
partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(
|
||||
cast(Sequence[Mapping[str, object]], completion_response.results)
|
||||
)
|
||||
if len(partition) <= 1:
|
||||
return None
|
||||
return tuple(
|
||||
LiteLLMRealtimeStreamLoggingObject(
|
||||
results=cast(OpenAIRealtimeStreamList, list(group)),
|
||||
usage=ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(group),
|
||||
service_tier=tier,
|
||||
)
|
||||
for tier, group in partition.items()
|
||||
)
|
||||
|
||||
|
||||
def completion_cost(
|
||||
completion_response: object | None = None,
|
||||
model: str | None = None,
|
||||
|
|
@ -1266,6 +1284,41 @@ def completion_cost(
|
|||
try:
|
||||
call_type = _infer_call_type(call_type, completion_response) or "completion"
|
||||
|
||||
if call_type == CallTypes.aresponses_websocket.value and isinstance(
|
||||
completion_response, LiteLLMRealtimeStreamLoggingObject
|
||||
):
|
||||
ws_tier_parts: Final = _split_responses_ws_logging_object_by_service_tier(completion_response)
|
||||
if ws_tier_parts is not None:
|
||||
return sum(
|
||||
completion_cost(
|
||||
completion_response=part,
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
messages=messages,
|
||||
completion=completion,
|
||||
total_time=total_time,
|
||||
call_type=call_type,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
region_name=region_name,
|
||||
size=size,
|
||||
quality=quality,
|
||||
n=n,
|
||||
custom_cost_per_token=custom_cost_per_token,
|
||||
custom_cost_per_second=custom_cost_per_second,
|
||||
optional_params=optional_params,
|
||||
custom_pricing=custom_pricing,
|
||||
base_model=base_model,
|
||||
standard_built_in_tools_params=standard_built_in_tools_params,
|
||||
litellm_model_name=litellm_model_name,
|
||||
router_model_id=router_model_id,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
service_tier=service_tier,
|
||||
data_residency=data_residency,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
for part in ws_tier_parts
|
||||
)
|
||||
|
||||
if (
|
||||
(call_type == "aimage_generation" or call_type == "image_generation")
|
||||
and model is not None
|
||||
|
|
@ -1466,12 +1519,15 @@ def completion_cost(
|
|||
duration_seconds = usage_obj.get("duration_seconds", None)
|
||||
_vr = usage_obj.get("video_resolution", None)
|
||||
provider_reported_cost = usage_obj.get("provider_reported_cost_usd", None)
|
||||
_vc = usage_obj.get("video_count", None)
|
||||
else:
|
||||
duration_seconds = getattr(usage_obj, "duration_seconds", None)
|
||||
_vr = getattr(usage_obj, "video_resolution", None)
|
||||
provider_reported_cost = getattr(usage_obj, "provider_reported_cost_usd", None)
|
||||
_vc = getattr(usage_obj, "video_count", None)
|
||||
if _vr is not None:
|
||||
video_resolution = str(_vr).strip().lower()
|
||||
video_count = _vc if isinstance(_vc, int) and not isinstance(_vc, bool) and _vc > 1 else 1
|
||||
|
||||
if _video_model_info is None and provider_reported_cost is not None:
|
||||
return float(provider_reported_cost)
|
||||
|
|
@ -1482,12 +1538,15 @@ def completion_cost(
|
|||
video_generation_cost,
|
||||
)
|
||||
|
||||
return video_generation_cost(
|
||||
model=model,
|
||||
duration_seconds=duration_seconds,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_info=_video_model_info,
|
||||
video_resolution=video_resolution,
|
||||
return (
|
||||
video_generation_cost(
|
||||
model=model,
|
||||
duration_seconds=duration_seconds,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_info=_video_model_info,
|
||||
video_resolution=video_resolution,
|
||||
)
|
||||
* video_count
|
||||
)
|
||||
# Fallback to default video cost calculation if no duration available
|
||||
return default_video_cost_calculator(
|
||||
|
|
@ -2558,6 +2617,7 @@ _RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "re
|
|||
|
||||
class _ResponsesWsEventResponse(BaseModel):
|
||||
usage: Mapping[str, object] | None = None
|
||||
service_tier: str | None = None
|
||||
|
||||
|
||||
class _ResponsesWsEvent(BaseModel):
|
||||
|
|
@ -2565,20 +2625,39 @@ class _ResponsesWsEvent(BaseModel):
|
|||
response: _ResponsesWsEventResponse | None = None
|
||||
|
||||
|
||||
def _billable_responses_ws_events(
|
||||
results: Sequence[Mapping[str, object]],
|
||||
) -> tuple[tuple[Mapping[str, object], _ResponsesWsEventResponse], ...]:
|
||||
return tuple(
|
||||
(result, event.response)
|
||||
for result in results
|
||||
if (event := _ResponsesWsEvent.model_validate(result)).type in _RESPONSES_WS_BILLABLE_EVENT_TYPES
|
||||
and event.response is not None
|
||||
and event.response.usage is not None
|
||||
)
|
||||
|
||||
|
||||
class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor):
|
||||
@staticmethod
|
||||
def collect_usage_from_responses_ws_results(
|
||||
results: Sequence[Mapping[str, object]],
|
||||
) -> tuple[Usage, ...]:
|
||||
events: Final = tuple(_ResponsesWsEvent.model_validate(result) for result in results)
|
||||
return tuple(
|
||||
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # same shared transform the realtime processor uses
|
||||
event.response.usage
|
||||
response.usage
|
||||
)
|
||||
for event in events
|
||||
if event.type in _RESPONSES_WS_BILLABLE_EVENT_TYPES
|
||||
and event.response is not None
|
||||
and event.response.usage is not None
|
||||
for _, response in _billable_responses_ws_events(results)
|
||||
if response.usage is not None
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def partition_results_by_service_tier(
|
||||
results: Sequence[Mapping[str, object]],
|
||||
) -> Mapping[str | None, tuple[Mapping[str, object], ...]]:
|
||||
billable: Final = _billable_responses_ws_events(results)
|
||||
tiers: Final = dict.fromkeys(response.service_tier for _, response in billable)
|
||||
return MappingProxyType(
|
||||
{tier: tuple(result for result, response in billable if response.service_tier == tier) for tier in tiers}
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -2101,9 +2101,14 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
|
||||
)
|
||||
)
|
||||
ws_tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(
|
||||
results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
|
||||
)
|
||||
ws_service_tier: Final = next(iter(ws_tier_partition)) if len(ws_tier_partition) == 1 else None
|
||||
logging_result = LiteLLMRealtimeStreamLoggingObject(
|
||||
usage=combined_ws_usage,
|
||||
results=result, # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
|
||||
service_tier=ws_service_tier,
|
||||
)
|
||||
|
||||
elif (
|
||||
|
|
|
|||
|
|
@ -1167,7 +1167,9 @@ class ModelResponseIterator:
|
|||
# (matches OpenAI behavior and non-streaming Anthropic implementation)
|
||||
if self.converted_response_format_tool:
|
||||
finish_reason = "stop"
|
||||
usage: Final = self._handle_usage(anthropic_usage_chunk=message_delta["usage"])
|
||||
usage: Final = (
|
||||
self._handle_usage(anthropic_usage_chunk=message_delta["usage"]) if "usage" in message_delta else None
|
||||
)
|
||||
container: Final = message_delta["delta"].get("container")
|
||||
return finish_reason, usage, container
|
||||
|
||||
|
|
|
|||
|
|
@ -18,7 +18,9 @@ if TYPE_CHECKING:
|
|||
import litellm
|
||||
|
||||
|
||||
def cost_per_token(model: str, usage: "Usage", service_tier: str | None = None) -> tuple[float, float]:
|
||||
def cost_per_token(
|
||||
model: str, usage: "Usage", service_tier: str | None = None, model_info: "ModelInfo | None" = None
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
|
||||
|
||||
|
|
@ -27,6 +29,7 @@ def cost_per_token(model: str, usage: "Usage", service_tier: str | None = None)
|
|||
- usage: LiteLLM Usage block, containing anthropic caching information
|
||||
- service_tier: the service tier the request was served at (e.g. "priority"),
|
||||
read from the Anthropic response usage and used to select tier-specific pricing
|
||||
- model_info: effective deployment prices, when they override public rates
|
||||
|
||||
Returns:
|
||||
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
|
||||
|
|
@ -36,16 +39,23 @@ def cost_per_token(model: str, usage: "Usage", service_tier: str | None = None)
|
|||
usage=usage,
|
||||
custom_llm_provider="anthropic",
|
||||
service_tier=service_tier,
|
||||
model_info=model_info,
|
||||
)
|
||||
|
||||
# Apply provider_specific_entry multipliers for geo/speed routing
|
||||
try:
|
||||
model_info: Final = litellm.get_model_info(model=model, custom_llm_provider="anthropic")
|
||||
provider_specific_entry: Final[dict] = model_info.get("provider_specific_entry") or {}
|
||||
effective_info: Final = (
|
||||
model_info
|
||||
if model_info is not None
|
||||
else litellm.get_model_info(model=model, custom_llm_provider="anthropic")
|
||||
)
|
||||
provider_specific_entry: Final = effective_info.get("provider_specific_entry")
|
||||
|
||||
geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=model_info, usage=usage)
|
||||
geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=effective_info, usage=usage)
|
||||
speed_multiplier: Final = (
|
||||
provider_specific_entry.get("fast", 1.0) if getattr(usage, "speed", None) == "fast" else 1.0
|
||||
provider_specific_entry.get("fast", 1.0)
|
||||
if provider_specific_entry and getattr(usage, "speed", None) == "fast"
|
||||
else 1.0
|
||||
)
|
||||
|
||||
if speed_multiplier != 1.0:
|
||||
|
|
|
|||
|
|
@ -376,10 +376,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
usage_dict: UsageDelta = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
|
||||
chunk.usage
|
||||
)
|
||||
merged_chunk["usage"] = usage_dict
|
||||
if self.applied_edits and "context_management" not in merged_chunk:
|
||||
merged_chunk["context_management"] = ContextManagementResponse(applied_edits=list(self.applied_edits))
|
||||
return self._augment_message_delta_usage(merged_chunk)
|
||||
return self._augment_message_delta_usage({**merged_chunk, "usage": usage_dict})
|
||||
|
||||
def _handle_choiceless_chunk(self, chunk: "ModelResponseStream") -> bool:
|
||||
"""Consume an OpenAI-compatible chunk that carries no ``choices``.
|
||||
|
|
@ -448,8 +447,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
}
|
||||
iterations.append(message_iteration)
|
||||
augmented_usage["iterations"] = iterations
|
||||
augmented["usage"] = augmented_usage
|
||||
return augmented
|
||||
return {**augmented, "usage": augmented_usage}
|
||||
|
||||
def _next_compaction_event(self) -> dict[str, object] | None:
|
||||
"""Return the next compaction content-block SSE event, or ``None``.
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
|
|
@ -18,6 +18,8 @@ from litellm.llms.base_llm.passthrough.transformation import (
|
|||
BasePassthroughConfig,
|
||||
RelayShape,
|
||||
logged_relay_shape,
|
||||
model_group_from,
|
||||
relayed_body,
|
||||
strip_leading_model_segment,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
|
@ -35,19 +37,6 @@ if TYPE_CHECKING:
|
|||
EMPTY_QUERY: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
class PassthroughMetadata(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
model_group: str = ""
|
||||
|
||||
|
||||
def model_group_from(litellm_params: Mapping[str, object]) -> str:
|
||||
try:
|
||||
return PassthroughMetadata.model_validate(litellm_params.get("litellm_metadata")).model_group
|
||||
except ValidationError:
|
||||
return ""
|
||||
|
||||
|
||||
def api_version_from(litellm_params: Mapping[str, object]) -> str | None:
|
||||
try:
|
||||
return TypeAdapter(str | None).validate_python(litellm_params.get("api_version"))
|
||||
|
|
@ -96,14 +85,6 @@ def relay_query_params(
|
|||
return MappingProxyType({**(request_query_params or EMPTY_QUERY), "api-version": api_version})
|
||||
|
||||
|
||||
def relayed_body(httpx_response: Response) -> str | dict:
|
||||
try:
|
||||
body: Final[object] = httpx_response.json()
|
||||
except ValueError:
|
||||
return httpx_response.text
|
||||
return body if isinstance(body, dict) else httpx_response.text
|
||||
|
||||
|
||||
FOUNDRY_RELAY_SHAPES: Final = (
|
||||
RelayShape("/rerank", CallTypes.arerank, RerankResponse.model_validate),
|
||||
RelayShape("/providers/blackforestlabs/v1/flux-2-pro", CallTypes.aimage_generation, ImageResponse.model_validate),
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from collections.abc import Callable, Mapping, Sequence
|
|||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
|
|
@ -29,6 +29,19 @@ if TYPE_CHECKING:
|
|||
RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
class PassthroughMetadata(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
model_group: str = ""
|
||||
|
||||
|
||||
def model_group_from(litellm_params: Mapping[str, object]) -> str:
|
||||
try:
|
||||
return PassthroughMetadata.model_validate(litellm_params.get("litellm_metadata")).model_group
|
||||
except ValidationError:
|
||||
return ""
|
||||
|
||||
|
||||
def strip_leading_model_segment(endpoint: str, model_names: tuple[str, ...]) -> str:
|
||||
path: Final = endpoint.lstrip("/")
|
||||
for model_name in model_names:
|
||||
|
|
@ -55,6 +68,14 @@ def relayed_json_object(httpx_response: Response) -> Mapping[str, object] | None
|
|||
return None
|
||||
|
||||
|
||||
def relayed_body(httpx_response: Response) -> str | dict:
|
||||
try:
|
||||
body: Final[object] = httpx_response.json()
|
||||
except ValueError:
|
||||
return httpx_response.text
|
||||
return body if isinstance(body, dict) else httpx_response.text
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RelayShape:
|
||||
path_suffix: str
|
||||
|
|
|
|||
|
|
@ -2,9 +2,11 @@
|
|||
For calculating cost of fireworks ai serverless inference models.
|
||||
"""
|
||||
|
||||
import math
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
from typing import (
|
||||
Final,
|
||||
cast, # noqa: TID251 # the fallback entry is a dict copy of a ReadOnly TypedDict; no cast-free way to retype it
|
||||
)
|
||||
|
||||
from litellm.constants import (
|
||||
FIREWORKS_AI_4_B,
|
||||
|
|
@ -12,12 +14,10 @@ from litellm.constants import (
|
|||
FIREWORKS_AI_56_B_MOE,
|
||||
FIREWORKS_AI_176_B_MOE,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import TokenRates, apply_off_peak_pricing
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
|
||||
from litellm.types.utils import ModelInfo, Usage
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
NO_CACHE_READ_RATE: Final = float("nan")
|
||||
|
||||
|
||||
# Extract the number of billion parameters from the model name
|
||||
# only used for together_computer LLMs
|
||||
|
|
@ -67,6 +67,28 @@ def _resolve_model_info(model: str) -> ModelInfo:
|
|||
return get_model_info(model=base_model, custom_llm_provider="fireworks_ai")
|
||||
|
||||
|
||||
def _with_cache_read_fallback(model_info: ModelInfo) -> ModelInfo:
|
||||
"""Entries without a cache-read rate keep the previous calculator's input-rate fallback for cached
|
||||
reads (LIT-7845 tracks the documented discount); the shared map is never mutated, so a copy carries it."""
|
||||
input_rate: Final = model_info.get("input_cost_per_token")
|
||||
if model_info.get("cache_read_input_token_cost") is not None or input_rate is None:
|
||||
return model_info
|
||||
off_peak: Final = model_info.get("off_peak_pricing")
|
||||
if off_peak is None or "cache_read_input_token_cost" in off_peak:
|
||||
return cast(ModelInfo, {**model_info, "cache_read_input_token_cost": input_rate})
|
||||
return cast(
|
||||
ModelInfo,
|
||||
{
|
||||
**model_info,
|
||||
"cache_read_input_token_cost": input_rate,
|
||||
"off_peak_pricing": {
|
||||
**off_peak,
|
||||
"cache_read_input_token_cost": off_peak.get("input_cost_per_token", input_rate),
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def cost_per_token(model: str, usage: Usage, current_time: datetime | None = None) -> tuple[float, float]:
|
||||
"""
|
||||
Calculates the cost per token for a given model, prompt tokens, and completion tokens,
|
||||
|
|
@ -80,29 +102,11 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non
|
|||
Returns:
|
||||
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
|
||||
"""
|
||||
model_info: Final = _resolve_model_info(model)
|
||||
standard_cache_read_rate: Final = model_info.get("cache_read_input_token_cost")
|
||||
rates: Final = apply_off_peak_pricing(
|
||||
model_info,
|
||||
current_time,
|
||||
TokenRates(
|
||||
input_rate=model_info["input_cost_per_token"] or 0.0,
|
||||
output_rate=model_info["output_cost_per_token"] or 0.0,
|
||||
cache_read_rate=standard_cache_read_rate if standard_cache_read_rate is not None else NO_CACHE_READ_RATE,
|
||||
cache_creation_rate=0.0,
|
||||
reasoning_rate=None,
|
||||
),
|
||||
model_info: Final = _with_cache_read_fallback(_resolve_model_info(model))
|
||||
return generic_cost_per_token(
|
||||
model=model,
|
||||
usage=usage,
|
||||
custom_llm_provider="fireworks_ai",
|
||||
model_info=model_info,
|
||||
current_time=current_time,
|
||||
)
|
||||
cache_read_rate: Final[float] = rates.input_rate if math.isnan(rates.cache_read_rate) else rates.cache_read_rate
|
||||
|
||||
prompt_tokens_details: Final = usage.prompt_tokens_details
|
||||
cached_tokens: Final[int] = (
|
||||
prompt_tokens_details.cached_tokens
|
||||
if prompt_tokens_details is not None and prompt_tokens_details.cached_tokens is not None
|
||||
else 0
|
||||
)
|
||||
non_cached_prompt_tokens: Final[int] = max(usage.prompt_tokens - cached_tokens, 0)
|
||||
prompt_cost: Final[float] = non_cached_prompt_tokens * rates.input_rate + cached_tokens * cache_read_rate
|
||||
completion_cost: Final[float] = usage.completion_tokens * rates.output_rate
|
||||
|
||||
return prompt_cost, completion_cost
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import litellm
|
|||
from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS
|
||||
from litellm.images.utils import ImageEditRequestUtils
|
||||
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
|
||||
from litellm.llms.vertex_ai.videos.transformation import veo_video_count_from_parameters
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.gemini import (
|
||||
GeminiLongRunningOperationResponse,
|
||||
|
|
@ -354,6 +355,9 @@ class GeminiVideoConfig(BaseVideoConfig):
|
|||
video_resolution: Final = _usage_video_resolution_from_parameters(parameters)
|
||||
if video_resolution is not None:
|
||||
usage_data["video_resolution"] = video_resolution
|
||||
video_count: Final = veo_video_count_from_parameters(parameters)
|
||||
if video_count is not None:
|
||||
usage_data["video_count"] = video_count
|
||||
|
||||
video_obj.usage = usage_data
|
||||
return video_obj
|
||||
|
|
|
|||
0
litellm/llms/nvidia_nim/passthrough/__init__.py
Normal file
0
litellm/llms/nvidia_nim/passthrough/__init__.py
Normal file
139
litellm/llms/nvidia_nim/passthrough/transformation.py
Normal file
139
litellm/llms/nvidia_nim/passthrough/transformation.py
Normal file
|
|
@ -0,0 +1,139 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Collection, Iterable, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.base_llm.passthrough.transformation import (
|
||||
BasePassthroughConfig,
|
||||
model_group_from,
|
||||
relayed_body,
|
||||
strip_leading_model_segment,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.router import DeploymentTypedDict
|
||||
from litellm.types.utils import LlmProviders, StandardPassThroughResponseObject
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from httpx import URL, Response
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse
|
||||
|
||||
|
||||
API_VERSION_SEGMENT: Final = re.compile(r"^v\d+$")
|
||||
NVIDIA_NIM_MODEL_PREFIX: Final = f"{LlmProviders.NVIDIA_NIM.value}/"
|
||||
NVIDIA_NIM_ROUTE_PREFIX: Final = re.compile(rf"^/{LlmProviders.NVIDIA_NIM.value}/", re.IGNORECASE)
|
||||
|
||||
|
||||
def is_nvidia_nim_deployment(deployment: DeploymentTypedDict) -> bool:
|
||||
litellm_params: Final = deployment["litellm_params"]
|
||||
return litellm_params.get("custom_llm_provider") == LlmProviders.NVIDIA_NIM.value or litellm_params.get(
|
||||
"model", ""
|
||||
).startswith(NVIDIA_NIM_MODEL_PREFIX)
|
||||
|
||||
|
||||
def nvidia_nim_model_groups(deployments: Iterable[DeploymentTypedDict] | None) -> frozenset[str]:
|
||||
listed: Final = tuple(deployments or ())
|
||||
nim_groups: Final = frozenset(d["model_name"] for d in listed if is_nvidia_nim_deployment(d))
|
||||
other_groups: Final = frozenset(d["model_name"] for d in listed if not is_nvidia_nim_deployment(d))
|
||||
return nim_groups - other_groups
|
||||
|
||||
|
||||
def nvidia_nim_model_group_in_path(path: str, deployments: Iterable[DeploymentTypedDict] | None) -> str | None:
|
||||
return nvidia_nim_router_model_in_endpoint(
|
||||
NVIDIA_NIM_ROUTE_PREFIX.sub("", path), nvidia_nim_model_groups(deployments)
|
||||
)
|
||||
|
||||
|
||||
def nvidia_nim_router_model_in_endpoint(endpoint: str, router_models: Collection[str]) -> str | None:
|
||||
segments: Final = tuple(segment for segment in endpoint.split("/") if segment)
|
||||
return next(
|
||||
(
|
||||
"/".join(segments[:length])
|
||||
for length in range(len(segments), 0, -1)
|
||||
if "/".join(segments[:length]) in router_models
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def without_repeated_version_prefix(api_base: str, native_endpoint: str) -> str:
|
||||
url: Final = httpx.URL(api_base)
|
||||
base_segments: Final = tuple(segment for segment in url.path.split("/") if segment)
|
||||
first_native_segment: Final = native_endpoint.lstrip("/").split("/", 1)[0]
|
||||
repeated: Final = (
|
||||
bool(base_segments)
|
||||
and API_VERSION_SEGMENT.match(first_native_segment) is not None
|
||||
and base_segments[-1] == first_native_segment
|
||||
)
|
||||
kept_segments: Final = base_segments[:-1] if repeated else base_segments
|
||||
return str(url.copy_with(path="/" + "/".join(kept_segments), query=None)).rstrip("/")
|
||||
|
||||
|
||||
class NvidiaNimPassthroughConfig(BasePassthroughConfig):
|
||||
def is_streaming_request(self, endpoint: str, request_data: dict) -> bool:
|
||||
return bool(request_data.get("stream", False))
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
model: str,
|
||||
endpoint: str,
|
||||
request_query_params: dict | None,
|
||||
litellm_params: dict,
|
||||
) -> tuple[URL, str]:
|
||||
base_target_url: Final = self.get_api_base(api_base)
|
||||
if base_target_url is None:
|
||||
raise ValueError("NVIDIA NIM api base not found: set `api_base` on the deployment or NVIDIA_NIM_API_BASE")
|
||||
native_endpoint: Final = strip_leading_model_segment(endpoint, (model, model_group_from(litellm_params)))
|
||||
root: Final = without_repeated_version_prefix(base_target_url, native_endpoint)
|
||||
return (self.format_url(native_endpoint, root, request_query_params), root)
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Mapping[str, str],
|
||||
model: str,
|
||||
messages: Sequence[AllMessageValues],
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict[str, str]: # mutable-ok: base class contract returns dict for httpx
|
||||
if api_key is None:
|
||||
return dict(headers) # mutable-ok: base class contract returns dict for httpx
|
||||
return {
|
||||
**headers,
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
} # mutable-ok: base class contract returns dict for httpx
|
||||
|
||||
@staticmethod
|
||||
def get_api_base(api_base: str | None = None) -> str | None:
|
||||
return api_base or get_secret_str("NVIDIA_NIM_API_BASE")
|
||||
|
||||
@staticmethod
|
||||
def get_api_key(api_key: str | None = None) -> str | None:
|
||||
return api_key or get_secret_str("NVIDIA_NIM_API_KEY")
|
||||
|
||||
@staticmethod
|
||||
def get_base_model(model: str) -> str | None:
|
||||
return model
|
||||
|
||||
def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]:
|
||||
return []
|
||||
|
||||
def logging_non_streaming_response(
|
||||
self,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
httpx_response: Response,
|
||||
request_data: Mapping[str, object],
|
||||
logging_obj: Logging,
|
||||
endpoint: str,
|
||||
) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None:
|
||||
return StandardPassThroughResponseObject(response=relayed_body(httpx_response))
|
||||
|
|
@ -124,6 +124,12 @@ def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsErr
|
|||
)
|
||||
|
||||
|
||||
def _served_model_name(model_version: object) -> str | None:
|
||||
if not isinstance(model_version, str) or not model_version:
|
||||
return None
|
||||
return model_version.split("@", 1)[0]
|
||||
|
||||
|
||||
class VertexAIBaseConfig:
|
||||
def get_mapped_special_auth_params(self) -> dict:
|
||||
"""
|
||||
|
|
@ -1951,6 +1957,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
def _check_prompt_level_content_filter(
|
||||
processed_chunk: GenerateContentResponseBody,
|
||||
response_id: str | None,
|
||||
model: str | None = None,
|
||||
) -> Optional["ModelResponseStream"]:
|
||||
"""
|
||||
Check if prompt is blocked due to content filtering at the prompt level.
|
||||
|
|
@ -1990,7 +1997,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
enhancements=None,
|
||||
)
|
||||
|
||||
model_response: Final = ModelResponseStream(choices=[choice], id=response_id)
|
||||
model_response: Final = ModelResponseStream(choices=[choice], id=response_id, model=model)
|
||||
return model_response
|
||||
|
||||
return None
|
||||
|
|
@ -2434,7 +2441,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
completion_response = GenerateContentResponseBody(**completion_response)
|
||||
|
||||
## GET MODEL ##
|
||||
model_response.model = model
|
||||
served: Final = _served_model_name(completion_response.get("modelVersion"))
|
||||
model_response.model = served if served is not None else model
|
||||
|
||||
## CHECK IF RESPONSE FLAGGED
|
||||
if "promptFeedback" in completion_response and "blockReason" in completion_response["promptFeedback"]:
|
||||
|
|
@ -3264,12 +3272,18 @@ class ModelResponseIterator:
|
|||
|
||||
processed_chunk: Final = GenerateContentResponseBody(**chunk)
|
||||
response_id: Final = processed_chunk.get("responseId")
|
||||
model_response = ModelResponseStream(choices=[], id=response_id)
|
||||
served: Final = _served_model_name(processed_chunk.get("modelVersion"))
|
||||
model_response = ModelResponseStream(
|
||||
choices=[],
|
||||
id=response_id,
|
||||
model=served,
|
||||
)
|
||||
|
||||
# Check if prompt is blocked due to content filtering
|
||||
blocked_response: Final = VertexGeminiConfig._check_prompt_level_content_filter(
|
||||
processed_chunk=processed_chunk,
|
||||
response_id=response_id,
|
||||
model=served,
|
||||
)
|
||||
if blocked_response is not None:
|
||||
model_response = blocked_response
|
||||
|
|
|
|||
|
|
@ -70,10 +70,17 @@ def _parse_veo_operation(raw_response: httpx.Response) -> _VeoOperation:
|
|||
return operation
|
||||
|
||||
|
||||
def veo_video_count_from_parameters(parameters: Mapping[str, object]) -> int | None:
|
||||
sample_count: Final = parameters.get("sampleCount")
|
||||
if isinstance(sample_count, bool) or not isinstance(sample_count, int) or sample_count < 1:
|
||||
return None
|
||||
return sample_count
|
||||
|
||||
|
||||
def _build_vertex_video_usage_from_request_data(
|
||||
request_data: dict[str, Any] | None,
|
||||
) -> dict[str, float | str]:
|
||||
"""Build usage metadata (duration, resolution) for video cost calculation."""
|
||||
"""Build usage metadata (duration, resolution, video count) for video cost calculation."""
|
||||
usage_data: Final[dict[str, float | str]] = {}
|
||||
if not request_data:
|
||||
return usage_data
|
||||
|
|
@ -88,6 +95,9 @@ def _build_vertex_video_usage_from_request_data(
|
|||
res: Final = parameters.get("resolution")
|
||||
if res is not None and str(res).strip() != "":
|
||||
usage_data["video_resolution"] = str(res).strip().lower()
|
||||
video_count: Final = veo_video_count_from_parameters(parameters)
|
||||
if video_count is not None:
|
||||
usage_data["video_count"] = video_count
|
||||
return usage_data
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -49,7 +49,6 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
|
||||
Inherits from OpenAIResponsesAPIConfig since XAI's Responses API is largely
|
||||
compatible with OpenAI's, with a few differences:
|
||||
- Does not support the 'instructions' parameter
|
||||
- Requires code_interpreter tools to have 'container' field removed
|
||||
- Recommends store=false when sending images
|
||||
|
||||
|
|
@ -60,20 +59,6 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
def custom_llm_provider(self) -> LlmProviders:
|
||||
return LlmProviders.XAI
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
"""
|
||||
Get supported parameters for XAI Responses API.
|
||||
|
||||
XAI supports most OpenAI Responses API params except 'instructions'.
|
||||
"""
|
||||
supported_params: Final = super().get_supported_openai_params(model)
|
||||
|
||||
# Remove 'instructions' as it's not supported by XAI
|
||||
if "instructions" in supported_params:
|
||||
supported_params.remove("instructions")
|
||||
|
||||
return supported_params
|
||||
|
||||
def _transform_web_search_tool(self, tool: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""
|
||||
Transform web_search tool to XAI format.
|
||||
|
|
@ -158,19 +143,13 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
Map parameters for XAI Responses API.
|
||||
|
||||
Handles XAI-specific transformations:
|
||||
1. Drops 'instructions' parameter (not supported)
|
||||
2. Transforms code_interpreter tools to remove 'container' field
|
||||
3. Transforms web_search tools to XAI format (removes search_context_size, adds filters)
|
||||
4. Transforms x_search tools to XAI format
|
||||
5. Sets store=false when images are detected (recommended by XAI)
|
||||
1. Transforms code_interpreter tools to remove 'container' field
|
||||
2. Transforms web_search tools to XAI format (removes search_context_size, adds filters)
|
||||
3. Transforms x_search tools to XAI format
|
||||
4. Sets store=false when images are detected (recommended by XAI)
|
||||
"""
|
||||
params: Final = dict(response_api_optional_params)
|
||||
|
||||
# Drop instructions parameter (not supported by XAI)
|
||||
if "instructions" in params:
|
||||
verbose_logger.debug("XAI Responses API does not support 'instructions' parameter. Dropping it.")
|
||||
params.pop("instructions")
|
||||
|
||||
if "metadata" in params:
|
||||
verbose_logger.debug("XAI Responses API does not support 'metadata' parameter. Dropping it.")
|
||||
params.pop("metadata")
|
||||
|
|
|
|||
|
|
@ -58151,6 +58151,23 @@
|
|||
"model_info": {
|
||||
"supports_reasoning": true
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "gemini-chat-baseline",
|
||||
"pattern": "gemini-(?!.*(?:-tts|-image|-live|-audio|-embedding|-computer-use|-robotics|-transcribe|-translate))(?:2[.-][5-9]|[3-9](?:[.-]\\d{1,2})?)-(?:pro|flash)(?:-lite)?(?![a-z])",
|
||||
"description": "Any Gemini text-chat id at 2.5 or higher under any namespace, including bare ids, gemini/, vertex_ai/, openrouter/google/, deepinfra/google/, vercel_ai_gateway/google/, oci/google., and databricks-gemini-<major>-<minor>: gemini-<major>[.minor]-(pro|flash)[-lite] with any trailing preview, date or variant tag. The capability flags were verified against each of those providers' own catalogs and docs. The lookahead excludes the tts, image, live, audio, embedding, computer-use, robotics, transcribe and translate lines, which are different modes with different capabilities. Provider-specific deviations, such as Perplexity's Agent API serving these as mode responses, are carried by their exact map entries, which always win over this rule. Carries no token limits or pricing, so those stay on the standard unmapped behavior rather than a guessed number. Source check 2026-09-15: all 45 first-party 2.5+ text-chat entries in this map carry every field below, and the OpenRouter (openrouter.ai/api/v1/models), Vercel AI Gateway (ai-gateway.vercel.sh/v1/models), DeepInfra (api.deepinfra.com/models/list), OCI and Databricks model docs list reasoning, tools and image input for the same models.",
|
||||
"model_info": {
|
||||
"mode": "chat",
|
||||
"supports_reasoning": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_vision": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_web_search": true
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
"""Bridge token flow: litellm identity resolution and the DCR-bridge oauth_delegate mint/refresh pipeline."""
|
||||
|
||||
import math
|
||||
import os
|
||||
import secrets
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
|
@ -12,6 +14,9 @@ from typing_extensions import assert_never
|
|||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.proxy._experimental.mcp_server.oauth_utils import TOKEN_NO_CACHE_HEADERS
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
_V2_GCM_PREFIX, # pyright: ignore[reportPrivateUsage] # reuse the encrypted credential's format discriminator
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -24,6 +29,7 @@ if TYPE_CHECKING:
|
|||
UpstreamTokenGrant,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.handle_jwt import JWTIdentity
|
||||
|
||||
|
||||
def _litellm_key_from_request(request: Request) -> str | None:
|
||||
|
|
@ -48,6 +54,64 @@ def _litellm_key_from_request(request: Request) -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
async def oauth_authorization_uses_gateway_credential(request: Request) -> bool:
|
||||
"""Classify credentials for browser authorize; candidates still require full authorization."""
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler # noqa: PLC0415 # proxy import cycle
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the active auth configuration
|
||||
jwt_handler,
|
||||
master_key,
|
||||
user_custom_auth,
|
||||
)
|
||||
|
||||
if "x-litellm-api-key" in request.headers:
|
||||
return True
|
||||
token: Final = _litellm_key_from_request(request)
|
||||
if token is None:
|
||||
return "authorization" in request.headers
|
||||
if token.startswith("sk-") or (master_key and secrets.compare_digest(token.encode(), master_key.encode())):
|
||||
return True
|
||||
if user_custom_auth is not None or jwt_handler.litellm_jwtauth.oidc_userinfo_enabled:
|
||||
return True
|
||||
if not JWTHandler.is_jwt(token):
|
||||
return await _opaque_bearer_is_gateway_credential(token)
|
||||
claims: Final = JWTHandler.get_unverified_claims(token)
|
||||
issuer: Final = claims.get("iss") if claims is not None else None
|
||||
global_issuer: Final = os.getenv("JWT_ISSUER")
|
||||
# An unscoped global validator can accept issuers absent from the configured issuer list.
|
||||
if not isinstance(issuer, str) or not issuer or not global_issuer:
|
||||
return True
|
||||
return issuer == global_issuer or any(
|
||||
issuer == configured.issuer for configured in jwt_handler.litellm_jwtauth.issuers or ()
|
||||
)
|
||||
|
||||
|
||||
async def _opaque_bearer_is_gateway_credential(token: str) -> bool:
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
|
||||
is_envelope, # noqa: PLC0415 # envelope imports bridge types
|
||||
is_refresh_envelope,
|
||||
)
|
||||
from litellm.proxy._types import hash_token # noqa: PLC0415 # proxy import cycle
|
||||
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken # noqa: PLC0415 # proxy import cycle
|
||||
from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError # noqa: PLC0415 # proxy import cycle
|
||||
from litellm.proxy.auth.resolvers.store import IdentityStore # noqa: PLC0415 # proxy import cycle
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the identity store dependencies
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
if is_envelope(token) or is_refresh_envelope(token) or token.startswith(_V2_GCM_PREFIX):
|
||||
return True
|
||||
try:
|
||||
if ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(token) is not None:
|
||||
return True
|
||||
await IdentityStore(prisma_client, user_api_key_cache).resolve(hashed_token=hash_token(token))
|
||||
except KeyNotFoundError:
|
||||
return False
|
||||
except Exception as exc: # noqa: BLE001 # an identity lookup fault must not permit cookie fallback
|
||||
verbose_logger.debug("OAuth bearer ownership could not be checked (%s)", type(exc).__name__)
|
||||
return True
|
||||
|
||||
|
||||
def _key_is_active(key_obj: "UserAPIKeyAuth") -> bool:
|
||||
"""``True`` when the presented key is neither blocked nor past its expiry.
|
||||
|
||||
|
|
@ -243,6 +307,10 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol
|
|||
return "no_active_key"
|
||||
if user_object is None:
|
||||
return "no_active_key"
|
||||
return _active_user_record(user_object)
|
||||
|
||||
|
||||
def _active_user_record(user_object: "LiteLLM_UserTable") -> "LiteLLM_UserTable | Literal['no_active_key']":
|
||||
if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False:
|
||||
return "no_active_key"
|
||||
return user_object
|
||||
|
|
@ -301,15 +369,137 @@ async def _revalidate_active_subject(identity: "EnvelopeIdentity") -> "_KeyResol
|
|||
|
||||
|
||||
async def _extract_user_id_from_request(request: Request) -> str | None:
|
||||
"""The litellm ``user_id`` for the token request, so a per-user token is stored under the same
|
||||
identity the egress later reads it by. Storage is best-effort, so every non-resolved outcome
|
||||
(including a transient DB outage) collapses to ``None`` here and the caller simply skips the store;
|
||||
the bridge mint, which must status those outcomes differently, consumes
|
||||
:func:`_resolve_active_litellm_key` directly."""
|
||||
resolved: Final = await _resolve_active_litellm_key(request)
|
||||
if not isinstance(resolved, _ResolvedKey):
|
||||
"""Resolve the caller for identity binding without granting credential-write permission."""
|
||||
from litellm.proxy.auth.handle_jwt import JWTIdentity # noqa: PLC0415 # proxy import cycle
|
||||
|
||||
resolved: Final = await _resolve_request_auth(request)
|
||||
if isinstance(resolved, JWTIdentity):
|
||||
return resolved.user_id
|
||||
return _active_key_user_id(resolved) if resolved is not None else None
|
||||
|
||||
|
||||
async def authorize_oauth_credential_request(request: Request, server_id: str) -> str | None:
|
||||
from litellm.proxy._types import UserAPIKeyAuth # noqa: PLC0415 # proxy import cycle
|
||||
|
||||
resolved: Final = await _resolve_request_auth(request, f"/v1/mcp/server/{server_id}/oauth-user-credential")
|
||||
if not isinstance(resolved, UserAPIKeyAuth) or not _active_key_user_id(resolved):
|
||||
return None
|
||||
if not await can_store_oauth_credential(request, resolved, server_id):
|
||||
return None
|
||||
return resolved.user_id
|
||||
|
||||
|
||||
async def _resolve_request_auth(
|
||||
request: Request, write_route: str | None = None
|
||||
) -> "UserAPIKeyAuth | JWTIdentity | None":
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler # noqa: PLC0415 # proxy import cycle
|
||||
|
||||
token: Final = _litellm_key_from_request(request)
|
||||
if token is not None and JWTHandler.is_jwt(token):
|
||||
return await _resolve_jwt_auth(request, token, write_route)
|
||||
resolved: Final = await _resolve_active_litellm_key(request)
|
||||
return resolved.key if isinstance(resolved, _ResolvedKey) else None
|
||||
|
||||
|
||||
async def can_store_oauth_credential(request: Request, auth: "UserAPIKeyAuth", server_id: str) -> bool:
|
||||
"""Apply the same write policy to request credentials and verified signed-callback users."""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # registry imports auth helpers
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
|
||||
can_access_mcp_server, # noqa: PLC0415 # proxy import cycle
|
||||
)
|
||||
from litellm.proxy.auth.route_checks import RouteChecks # noqa: PLC0415 # proxy import cycle
|
||||
from litellm.proxy.auth.user_api_key_auth import ( # noqa: PLC0415 # proxy import cycle
|
||||
_run_centralized_common_checks, # pyright: ignore[reportPrivateUsage] # reuse admission policy for the credential-write action
|
||||
)
|
||||
|
||||
write_route: Final = f"/v1/mcp/server/{server_id}/oauth-user-credential"
|
||||
try:
|
||||
RouteChecks.is_virtual_key_allowed_to_call_route(route=write_route, valid_token=auth, request=request)
|
||||
await _run_centralized_common_checks(
|
||||
user_api_key_auth_obj=auth,
|
||||
request=request,
|
||||
request_data={},
|
||||
route=write_route,
|
||||
)
|
||||
return await can_access_mcp_server(auth, server_id, global_mcp_server_manager.get_allowed_mcp_servers)
|
||||
except Exception as exc: # noqa: BLE001 # authorization failure must never write credentials
|
||||
verbose_logger.debug("OAuth credential write not authorized (%s)", type(exc).__name__)
|
||||
return False
|
||||
|
||||
|
||||
async def _resolve_jwt_auth(
|
||||
request: Request,
|
||||
token: str,
|
||||
write_route: str | None,
|
||||
) -> "UserAPIKeyAuth | JWTIdentity | None":
|
||||
from litellm.proxy._types import UserAPIKeyAuth # noqa: PLC0415 # proxy import cycle
|
||||
from litellm.proxy.auth.handle_jwt import JWTAuthManager # noqa: PLC0415 # proxy import cycle
|
||||
from litellm.proxy.auth.user_api_key_auth import ( # noqa: PLC0415 # proxy import cycle
|
||||
_resolve_jwt_to_virtual_key, # pyright: ignore[reportPrivateUsage] # reuse admission mapping policy without provisioning a new key
|
||||
)
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # proxy globals initialized at startup
|
||||
general_settings,
|
||||
jwt_handler,
|
||||
premium_user,
|
||||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
if general_settings.get("enable_jwt_auth") is not True or premium_user is not True or prisma_client is None:
|
||||
return None
|
||||
try:
|
||||
if jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured():
|
||||
claims: Final = await jwt_handler.auth_jwt(token=token)
|
||||
validate: Final = jwt_handler.litellm_jwtauth.custom_validate
|
||||
if validate is not None and not validate(claims):
|
||||
return None
|
||||
mapped: Final = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims=claims,
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if isinstance(mapped, UserAPIKeyAuth):
|
||||
return None if await _key_owner_scim_deactivated(mapped) or not _active_key_user_id(mapped) else mapped
|
||||
if mapped is not None:
|
||||
return None
|
||||
if write_route is None:
|
||||
identity: Final = await JWTAuthManager.resolve_identity(
|
||||
api_key=token,
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if identity.user_object is not None and isinstance(_active_user_record(identity.user_object), str):
|
||||
return None
|
||||
return identity
|
||||
authorized: Final = await JWTAuthManager.authorize_jwt(
|
||||
api_key=token,
|
||||
jwt_handler=jwt_handler,
|
||||
request_data={},
|
||||
general_settings=general_settings,
|
||||
route=write_route,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
request_headers=dict(request.headers),
|
||||
request_method=request.method,
|
||||
)
|
||||
resolved_user: Final = authorized["user_object"]
|
||||
if resolved_user is not None and isinstance(_active_user_record(resolved_user), str):
|
||||
return None
|
||||
return JWTAuthManager.user_api_key_auth_from_result(authorized)
|
||||
except Exception as exc: # noqa: BLE001 # public OAuth exchange stays available; unvalidated identities never write credentials
|
||||
verbose_logger.debug("OAuth JWT identity could not be validated (%s)", type(exc).__name__)
|
||||
return None
|
||||
return _active_key_user_id(resolved.key)
|
||||
|
||||
|
||||
_UpstreamGrantRejection = Literal["no_access_token", "expired_lifetime"]
|
||||
|
|
|
|||
|
|
@ -32,6 +32,9 @@ from litellm.proxy._experimental.mcp_server.bridge_token_flow import (
|
|||
_prepare_bridge_mint,
|
||||
_prepare_bridge_refresh,
|
||||
_reload_active_user_by_id,
|
||||
authorize_oauth_credential_request,
|
||||
can_store_oauth_credential,
|
||||
oauth_authorization_uses_gateway_credential,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.faults import (
|
||||
CallerRejected,
|
||||
|
|
@ -836,16 +839,30 @@ async def _user_can_reach_mcp_server(user_id: str, server_id: str) -> bool:
|
|||
return server_id in await global_mcp_server_manager.get_allowed_mcp_servers(admitted)
|
||||
|
||||
|
||||
async def _bridge_authorize_access_denial(
|
||||
litellm_user_id: str,
|
||||
async def _resolve_oauth_authorization_user(
|
||||
request: Request,
|
||||
mcp_server: MCPServer,
|
||||
redirect_uri: str,
|
||||
state: str,
|
||||
) -> RedirectResponse | None:
|
||||
"""The denial redirect for a signed-in user who cannot reach the target server, or None to proceed."""
|
||||
if await _user_can_reach_mcp_server(litellm_user_id, mcp_server.server_id):
|
||||
return None
|
||||
return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
|
||||
enforce_binding: bool,
|
||||
) -> str | RedirectResponse:
|
||||
"""Resolve the authorization subject without replacing denied credentials with cookie grants."""
|
||||
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # proxy import cycle
|
||||
_user_id_from_session_cookie,
|
||||
)
|
||||
|
||||
use_gateway_credential: Final = enforce_binding and await oauth_authorization_uses_gateway_credential(request)
|
||||
request_user_id: Final = (
|
||||
await authorize_oauth_credential_request(request, mcp_server.server_id) if use_gateway_credential else None
|
||||
)
|
||||
if use_gateway_credential and request_user_id is None:
|
||||
return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
|
||||
user_id: Final = request_user_id or _user_id_from_session_cookie(request)
|
||||
if user_id is None:
|
||||
return _redirect_to_litellm_login(request)
|
||||
if not await _user_can_reach_mcp_server(user_id, mcp_server.server_id):
|
||||
return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
|
||||
return user_id
|
||||
|
||||
|
||||
async def authorize_with_server(
|
||||
|
|
@ -911,23 +928,12 @@ async def authorize_with_server(
|
|||
# Seal the authenticated caller into state so the token exchange cannot select another credential owner.
|
||||
litellm_user_id: str | None = None
|
||||
if enforce_binding or (resolved_server.is_dcr_bridge and resolved_server.is_oauth_delegate):
|
||||
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
_user_id_from_session_cookie,
|
||||
subject: Final = await _resolve_oauth_authorization_user(
|
||||
request, resolved_server, redirect_uri, state, enforce_binding
|
||||
)
|
||||
|
||||
litellm_user_id = (
|
||||
await _extract_user_id_from_request(request) if enforce_binding else None
|
||||
) or _user_id_from_session_cookie(request)
|
||||
if litellm_user_id is None:
|
||||
return _redirect_to_litellm_login(request)
|
||||
denial: Final = await _bridge_authorize_access_denial(
|
||||
litellm_user_id=litellm_user_id,
|
||||
mcp_server=resolved_server,
|
||||
redirect_uri=redirect_uri,
|
||||
state=state,
|
||||
)
|
||||
if denial is not None:
|
||||
return denial
|
||||
if isinstance(subject, RedirectResponse):
|
||||
return subject
|
||||
litellm_user_id = subject
|
||||
|
||||
oauth_nonce: Final = secrets.token_urlsafe(32) if enforce_binding else None
|
||||
encoded_state: Final = encode_state_with_base_url(
|
||||
|
|
@ -1218,12 +1224,32 @@ async def exchange_token_with_server(
|
|||
user_id: Final = resolved_user_id
|
||||
if user_id:
|
||||
try:
|
||||
await _store_per_user_token_server_side(
|
||||
server=resolved_server,
|
||||
user_id=user_id,
|
||||
token_response=token_response,
|
||||
identity_binding_proof=binding_proof,
|
||||
# Identity binding above must retain the verified caller even when a write is
|
||||
# denied. Authorize persistence separately, immediately before its side effect.
|
||||
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler
|
||||
|
||||
# A sealed code delegates a verified user for this authorized server. Raw
|
||||
# request credentials retain their own JWT/key restrictions during resolution.
|
||||
can_store: Final = (
|
||||
await can_store_oauth_credential(
|
||||
request, await MCPRequestHandler.reload_admitted_user(user_id), resolved_server.server_id
|
||||
)
|
||||
if bridge_identity is not None
|
||||
else await authorize_oauth_credential_request(request, resolved_server.server_id) == user_id
|
||||
)
|
||||
if can_store:
|
||||
await _store_per_user_token_server_side(
|
||||
server=resolved_server,
|
||||
user_id=user_id,
|
||||
token_response=token_response,
|
||||
identity_binding_proof=binding_proof,
|
||||
)
|
||||
else:
|
||||
verbose_logger.warning(
|
||||
"OAuth credential storage not authorized for user=%s server=%s",
|
||||
user_id,
|
||||
resolved_server.server_id,
|
||||
)
|
||||
except Exception as exc:
|
||||
verbose_logger.warning(
|
||||
"exchange_token_with_server: server-side storage failed for user=%s server=%s: %s",
|
||||
|
|
@ -1236,8 +1262,9 @@ async def exchange_token_with_server(
|
|||
"exchange_token_with_server: could not resolve a LiteLLM user_id for the request, "
|
||||
"so the per-user token for server=%s was NOT stored. The authorization_code egress "
|
||||
"requires the stored token, so the client will be challenged with 401 on reconnect. "
|
||||
"Ensure the request carries a valid LiteLLM key (x-litellm-api-key or Authorization), "
|
||||
"or store it via POST /mcp/server/{id}/oauth-user-credential.",
|
||||
"Ensure the request carries a valid LiteLLM key or enabled JWT identity "
|
||||
"(x-litellm-api-key or Authorization), "
|
||||
"or store it via POST /v1/mcp/server/{id}/oauth-user-credential.",
|
||||
resolved_server.server_id,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Final
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -137,3 +138,15 @@ async def build_effective_auth_contexts(
|
|||
if admitted_context is None:
|
||||
return team_contexts
|
||||
return [*team_contexts, admitted_context]
|
||||
|
||||
|
||||
async def can_access_mcp_server(
|
||||
user_api_key_auth: UserAPIKeyAuth,
|
||||
server_id: str,
|
||||
allowed_servers: Callable[[UserAPIKeyAuth], Awaitable[list[str]]],
|
||||
) -> bool:
|
||||
"""Resolve server access through the same credential contexts as MCP management."""
|
||||
for context in await build_effective_auth_contexts(user_api_key_auth):
|
||||
if server_id in await allowed_servers(context):
|
||||
return True
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -205,6 +205,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
|
|||
"/gigachat/",
|
||||
"/milvus/",
|
||||
"/mistral/",
|
||||
"/nvidia_nim/",
|
||||
"/openai/",
|
||||
"/openai_passthrough/",
|
||||
"/vertex-ai/",
|
||||
|
|
|
|||
|
|
@ -3125,6 +3125,11 @@
|
|||
"title": "Total Prompt Tokens",
|
||||
"type": "integer"
|
||||
},
|
||||
"total_response_time_ms": {
|
||||
"default": 0,
|
||||
"title": "Total Response Time Ms",
|
||||
"type": "integer"
|
||||
},
|
||||
"total_spend": {
|
||||
"default": 0.0,
|
||||
"title": "Total Spend",
|
||||
|
|
@ -3135,6 +3140,11 @@
|
|||
"title": "Total Successful Requests",
|
||||
"type": "integer"
|
||||
},
|
||||
"total_timed_requests": {
|
||||
"default": 0,
|
||||
"title": "Total Timed Requests",
|
||||
"type": "integer"
|
||||
},
|
||||
"total_tokens": {
|
||||
"default": 0,
|
||||
"title": "Total Tokens",
|
||||
|
|
@ -3643,6 +3653,16 @@
|
|||
"title": "Successful Requests",
|
||||
"type": "integer"
|
||||
},
|
||||
"timed_requests": {
|
||||
"default": 0,
|
||||
"title": "Timed Requests",
|
||||
"type": "integer"
|
||||
},
|
||||
"total_response_time_ms": {
|
||||
"default": 0,
|
||||
"title": "Total Response Time Ms",
|
||||
"type": "integer"
|
||||
},
|
||||
"total_tokens": {
|
||||
"default": 0,
|
||||
"title": "Total Tokens",
|
||||
|
|
@ -9986,7 +10006,7 @@
|
|||
},
|
||||
"unreachable_fallback": {
|
||||
"default": "fail_closed",
|
||||
"description": "Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed.",
|
||||
"description": "Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed.",
|
||||
"enum": [
|
||||
"fail_closed",
|
||||
"fail_open"
|
||||
|
|
@ -10948,6 +10968,18 @@
|
|||
"description": "Custom advisory message template used when on_flagged='inject_system_message'. Must contain a {reason} placeholder. Defaults to a generic advisory message if unset.",
|
||||
"title": "Advisory System Message"
|
||||
},
|
||||
"agent_id": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Agent identity reported to Agent 365 with every tool evaluation. When unset, the caller's key alias is used.",
|
||||
"title": "Agent Id"
|
||||
},
|
||||
"akto_account_id": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
@ -11450,6 +11482,30 @@
|
|||
"title": "Chunk Budget Chars",
|
||||
"type": "integer"
|
||||
},
|
||||
"client_id": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Client id of the gateway's Entra app registration (a confidential client). Falls back to the AGENT365_CLIENT_ID environment variable.",
|
||||
"title": "Client Id"
|
||||
},
|
||||
"client_secret": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Client secret of the gateway's Entra app registration, used to perform the On-Behalf-Of exchange. Falls back to the AGENT365_CLIENT_SECRET environment variable.",
|
||||
"title": "Client Secret"
|
||||
},
|
||||
"confidence_threshold": {
|
||||
"default": 0.5,
|
||||
"default_value": 0.5,
|
||||
|
|
@ -12496,6 +12552,18 @@
|
|||
"description": "The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set.",
|
||||
"title": "Realtime Violation Message"
|
||||
},
|
||||
"resource_app_id": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Application id of the Agent 365 resource the OBO token is minted for. Defaults to the production resource ea9ffc3e-8a23-4a7d-836d-234d7c7565c1; the Test and PreProd environments use a different id. Falls back to the AGENT365_RESOURCE_APP_ID environment variable.",
|
||||
"title": "Resource App Id"
|
||||
},
|
||||
"rules": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
@ -12733,6 +12801,18 @@
|
|||
"description": "The ID of your Model Armor template",
|
||||
"title": "Template Id"
|
||||
},
|
||||
"tenant_id": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Entra tenant id used for the On-Behalf-Of token exchange. Falls back to the AGENT365_TENANT_ID environment variable.",
|
||||
"title": "Tenant Id"
|
||||
},
|
||||
"timeout": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
@ -18945,6 +19025,228 @@
|
|||
]
|
||||
}
|
||||
},
|
||||
"/nvidia_nim/{endpoint}": {
|
||||
"delete": {
|
||||
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
|
||||
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__delete",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "endpoint",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Endpoint",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Nvidia Nim Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
]
|
||||
},
|
||||
"get": {
|
||||
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
|
||||
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__get",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "endpoint",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Endpoint",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Nvidia Nim Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
]
|
||||
},
|
||||
"patch": {
|
||||
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
|
||||
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__patch",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "endpoint",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Endpoint",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Nvidia Nim Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
]
|
||||
},
|
||||
"post": {
|
||||
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
|
||||
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__post",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "endpoint",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Endpoint",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Nvidia Nim Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
]
|
||||
},
|
||||
"put": {
|
||||
"description": "Relay a native NVIDIA NIM request through a LiteLLM model group.\n\n`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's\n`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through\nvirtual key auth, model access checks, and spend logging.",
|
||||
"operationId": "nvidia_nim_proxy_route_nvidia_nim__endpoint__put",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "endpoint",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Endpoint",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Nvidia Nim Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
]
|
||||
}
|
||||
},
|
||||
"/openai/deployments/{model}/chat/completions": {
|
||||
"post": {
|
||||
"description": "Follows the exact same API spec as `OpenAI's Chat API https://platform.openai.com/docs/api-reference/chat`\n\n```bash\ncurl -X POST http://localhost:4000/v1/chat/completions \n-H \"Content-Type: application/json\" \n-H \"Authorization: Bearer sk-1234\" \n-d '{\n \"model\": \"gpt-4o\",\n \"messages\": [\n {\n \"role\": \"user\",\n \"content\": \"Hello!\"\n }\n ]\n}'\n```",
|
||||
|
|
|
|||
|
|
@ -246,6 +246,7 @@ class Litellm_EntityType(enum.Enum):
|
|||
TEAM = "team"
|
||||
TEAM_MEMBER = "team_member"
|
||||
ORGANIZATION = "organization"
|
||||
ORGANIZATION_MEMBER = "organization_member"
|
||||
PROJECT = "project"
|
||||
TAG = "tag"
|
||||
AGENT = "agent"
|
||||
|
|
@ -485,6 +486,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/milvus",
|
||||
"/gigachat",
|
||||
"/watsonx",
|
||||
"/nvidia_nim",
|
||||
]
|
||||
|
||||
#########################################################
|
||||
|
|
@ -5234,6 +5236,8 @@ class BaseDailySpendTransaction(TypedDict):
|
|||
api_requests: int
|
||||
successful_requests: int
|
||||
failed_requests: int
|
||||
total_response_time_ms: NotRequired[int] # writable-ok: the rollup queue accumulates into this key in place
|
||||
timed_requests: NotRequired[int] # writable-ok: the rollup queue accumulates into this key in place
|
||||
|
||||
|
||||
class DailyTeamSpendTransaction(BaseDailySpendTransaction):
|
||||
|
|
@ -5272,6 +5276,7 @@ class DBSpendUpdateTransactions(TypedDict):
|
|||
team_list_transactions: dict[str, float] | None
|
||||
team_member_list_transactions: dict[str, float] | None
|
||||
org_list_transactions: dict[str, float] | None
|
||||
org_member_list_transactions: ReadOnly[dict[str, float] | None]
|
||||
tag_list_transactions: dict[str, float] | None
|
||||
agent_list_transactions: dict[str, float] | None
|
||||
model_access_group_list_transactions: ReadOnly[dict[str, float] | None]
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ from litellm.litellm_core_utils.url_utils import (
|
|||
validate_url,
|
||||
)
|
||||
from litellm.llms.azure.passthrough.transformation import azure_router_model_in_endpoint
|
||||
from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_metadata
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
|
|
@ -976,6 +977,26 @@ def _get_deployment_default_tpm_limit(model_name: str) -> int | None:
|
|||
return _get_deployment_default_limit(model_name, "default_api_key_tpm_limit")
|
||||
|
||||
|
||||
def get_key_own_model_rate_limit(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
|
||||
) -> dict[str, int] | None:
|
||||
if user_api_key_dict.metadata:
|
||||
result: Final = user_api_key_dict.metadata.get(rate_limit_key)
|
||||
if result:
|
||||
return result
|
||||
|
||||
if not user_api_key_dict.model_max_budget:
|
||||
return None
|
||||
budget_key: Final = "rpm_limit" if rate_limit_key == "model_rpm_limit" else "tpm_limit"
|
||||
model_limit: Final = {
|
||||
model: budget[budget_key]
|
||||
for model, budget in user_api_key_dict.model_max_budget.items()
|
||||
if isinstance(budget, dict) and budget.get(budget_key) is not None
|
||||
}
|
||||
return model_limit or None
|
||||
|
||||
|
||||
def get_key_model_rpm_limit(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
model_name: str | None = None,
|
||||
|
|
@ -989,20 +1010,9 @@ def get_key_model_rpm_limit(
|
|||
3. Team metadata (model_rpm_limit)
|
||||
4. Deployment default_api_key_rpm_limit (when model_name is provided)
|
||||
"""
|
||||
# 1. Check key metadata first (takes priority)
|
||||
if user_api_key_dict.metadata:
|
||||
result: Final = user_api_key_dict.metadata.get("model_rpm_limit")
|
||||
if result:
|
||||
return result
|
||||
|
||||
# 2. Check model_max_budget
|
||||
if user_api_key_dict.model_max_budget:
|
||||
model_rpm_limit: Final[dict[str, int]] = {}
|
||||
for model, budget in user_api_key_dict.model_max_budget.items():
|
||||
if isinstance(budget, dict) and budget.get("rpm_limit") is not None:
|
||||
model_rpm_limit[model] = budget["rpm_limit"]
|
||||
if model_rpm_limit:
|
||||
return model_rpm_limit
|
||||
key_own_limit: Final = get_key_own_model_rate_limit(user_api_key_dict, "model_rpm_limit")
|
||||
if key_own_limit is not None:
|
||||
return key_own_limit
|
||||
|
||||
# 3. Fallback to team metadata
|
||||
if user_api_key_dict.team_metadata:
|
||||
|
|
@ -1032,20 +1042,9 @@ def get_key_model_tpm_limit(
|
|||
3. Team metadata (model_tpm_limit)
|
||||
4. Deployment default_api_key_tpm_limit (when model_name is provided)
|
||||
"""
|
||||
# 1. Check key metadata first (takes priority)
|
||||
if user_api_key_dict.metadata:
|
||||
result: Final = user_api_key_dict.metadata.get("model_tpm_limit")
|
||||
if result:
|
||||
return result
|
||||
|
||||
# 2. Check model_max_budget (iterate per-model like RPM does)
|
||||
if user_api_key_dict.model_max_budget:
|
||||
model_tpm_limit: Final[dict[str, int]] = {}
|
||||
for model, budget in user_api_key_dict.model_max_budget.items():
|
||||
if isinstance(budget, dict) and budget.get("tpm_limit") is not None:
|
||||
model_tpm_limit[model] = budget["tpm_limit"]
|
||||
if model_tpm_limit:
|
||||
return model_tpm_limit
|
||||
key_own_limit: Final = get_key_own_model_rate_limit(user_api_key_dict, "model_tpm_limit")
|
||||
if key_own_limit is not None:
|
||||
return key_own_limit
|
||||
|
||||
# 3. Fallback to team metadata
|
||||
if user_api_key_dict.team_metadata:
|
||||
|
|
@ -2045,6 +2044,12 @@ def get_model_from_request(
|
|||
azure_model: Final = _router_model_from_azure_route(route, llm_router)
|
||||
return model if azure_model is None else azure_model
|
||||
|
||||
if route.lower().startswith("/nvidia_nim/"):
|
||||
nvidia_nim_model: Final = (
|
||||
nvidia_nim_model_group_in_path(route, llm_router.get_model_list()) if llm_router else None
|
||||
)
|
||||
return model if nvidia_nim_model is None else nvidia_nim_model
|
||||
|
||||
return model
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ import os
|
|||
import re
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Final, Literal, NoReturn, Protocol, TypeVar, cast
|
||||
|
||||
import httpx
|
||||
|
|
@ -58,7 +59,7 @@ from litellm.proxy.auth.model_access_denied import (
|
|||
)
|
||||
from litellm.proxy.auth.resolvers.grants import GrantResolver, UserLookup, canonical_user_id
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.auth.team_grants import team_model_aliases
|
||||
from litellm.proxy.auth.team_grants import team_grants, team_model_aliases
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
UserApiKeyCache,
|
||||
get_management_object_ttl,
|
||||
|
|
@ -66,6 +67,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.types.agents import AgentResponse
|
||||
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
|
||||
|
||||
from .auth_checks import (
|
||||
_allowed_routes_check,
|
||||
|
|
@ -132,6 +134,19 @@ class _UserInfoResponse(Protocol):
|
|||
def json(self) -> dict[str, object]: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class JWTIdentity:
|
||||
user_id: str | None
|
||||
user_object: LiteLLM_UserTable | None
|
||||
agent_id: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _JWTProvisioning:
|
||||
user_id_upsert: bool
|
||||
team_id_upsert: bool
|
||||
|
||||
|
||||
class AgentLookup(Protocol):
|
||||
"""The registered-agent lookups a JWT agent claim is matched against."""
|
||||
|
||||
|
|
@ -1481,6 +1496,7 @@ class JWTAuthManager:
|
|||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_id_upsert: bool | None = None,
|
||||
) -> tuple[str | None, LiteLLM_TeamTable | None]:
|
||||
"""Find and validate specific team ID from team_id_jwt_field or team_alias_jwt_field"""
|
||||
individual_team_id = jwt_handler.get_team_id(token=jwt_valid_token, default_value=None)
|
||||
|
|
@ -1508,7 +1524,9 @@ class JWTAuthManager:
|
|||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert,
|
||||
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert
|
||||
if team_id_upsert is None
|
||||
else team_id_upsert,
|
||||
)
|
||||
return individual_team_id, team_object
|
||||
except HTTPException as e:
|
||||
|
|
@ -1736,6 +1754,7 @@ class JWTAuthManager:
|
|||
proxy_logging_obj: ProxyLogging,
|
||||
route: str,
|
||||
org_alias: str | None = None,
|
||||
user_id_upsert: bool | None = None,
|
||||
) -> tuple[
|
||||
LiteLLM_UserTable | None,
|
||||
LiteLLM_OrganizationTable | None,
|
||||
|
|
@ -1799,7 +1818,11 @@ class JWTAuthManager:
|
|||
user_id=user_id,
|
||||
user_email=user_email,
|
||||
sso_user_id=user_id,
|
||||
upsert=jwt_handler.is_upsert_user_id(valid_user_email=valid_user_email),
|
||||
upsert=(
|
||||
jwt_handler.is_upsert_user_id(valid_user_email=valid_user_email)
|
||||
if user_id_upsert is None
|
||||
else user_id_upsert
|
||||
),
|
||||
),
|
||||
team_id=team_id,
|
||||
)
|
||||
|
|
@ -2020,6 +2043,7 @@ class JWTAuthManager:
|
|||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
team_id_upsert: bool | None = None,
|
||||
) -> None:
|
||||
"""Attach team context from x-litellm-team-id to an admin result.
|
||||
|
||||
|
|
@ -2037,7 +2061,7 @@ class JWTAuthManager:
|
|||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert,
|
||||
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert if team_id_upsert is None else team_id_upsert,
|
||||
)
|
||||
except Exception as e:
|
||||
# Fall back to pre-PR admin behavior: honor the admin's
|
||||
|
|
@ -2272,57 +2296,136 @@ class JWTAuthManager:
|
|||
request_headers: dict | None = None,
|
||||
request_method: str | None = None,
|
||||
) -> JWTAuthBuilderResult:
|
||||
"""Main authentication and authorization builder"""
|
||||
# Check if OIDC UserInfo endpoint is enabled, but fall back to standard
|
||||
# JWT auth if the token itself is a well-formed JWT (3-part structure).
|
||||
if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not jwt_handler.is_jwt(token=api_key):
|
||||
verbose_proxy_logger.debug("OIDC UserInfo is enabled. Fetching user info from UserInfo endpoint.")
|
||||
# Use the access token to fetch user info from OIDC UserInfo endpoint
|
||||
jwt_valid_token: dict = await jwt_handler.get_oidc_userinfo(token=api_key)
|
||||
else:
|
||||
# Default behavior: decode and validate the JWT token
|
||||
jwt_valid_token = await jwt_handler.auth_jwt(token=api_key)
|
||||
|
||||
# Check custom validate
|
||||
if jwt_handler.litellm_jwtauth.custom_validate:
|
||||
if not jwt_handler.litellm_jwtauth.custom_validate(jwt_valid_token):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="Invalid JWT token",
|
||||
)
|
||||
|
||||
# Check RBAC
|
||||
rbac_role: Final = jwt_handler.get_rbac_role(token=jwt_valid_token)
|
||||
await JWTAuthManager.check_rbac_role(
|
||||
jwt_handler,
|
||||
jwt_valid_token,
|
||||
general_settings,
|
||||
request_data,
|
||||
route,
|
||||
rbac_role,
|
||||
return await JWTAuthManager.authorize_jwt(
|
||||
api_key=api_key,
|
||||
jwt_handler=jwt_handler,
|
||||
request_data=request_data,
|
||||
general_settings=general_settings,
|
||||
route=route,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
request_headers=request_headers,
|
||||
request_method=request_method,
|
||||
provisioning=_JWTProvisioning(
|
||||
user_id_upsert=jwt_handler.litellm_jwtauth.user_id_upsert,
|
||||
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert,
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def authenticate_jwt(api_key: str, jwt_handler: JWTHandler) -> dict[str, object]:
|
||||
claims: Final = (
|
||||
await jwt_handler.get_oidc_userinfo(token=api_key)
|
||||
if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not jwt_handler.is_jwt(token=api_key)
|
||||
else await jwt_handler.auth_jwt(token=api_key)
|
||||
)
|
||||
validate: Final = jwt_handler.litellm_jwtauth.custom_validate
|
||||
if validate is not None and not validate(claims):
|
||||
raise HTTPException(status_code=403, detail="Invalid JWT token")
|
||||
return claims
|
||||
|
||||
@staticmethod
|
||||
async def resolve_identity(
|
||||
api_key: str,
|
||||
jwt_handler: JWTHandler,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> JWTIdentity:
|
||||
claims: Final = await JWTAuthManager.authenticate_jwt(api_key, jwt_handler)
|
||||
return await JWTAuthManager._resolve_claim_identity(
|
||||
claims, jwt_handler, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _resolve_claim_identity(
|
||||
claims: dict[str, object],
|
||||
jwt_handler: JWTHandler,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> JWTIdentity:
|
||||
claim_user_id, user_email, valid_user_email = await JWTAuthManager.get_user_info(jwt_handler, claims)
|
||||
user_id: Final = (
|
||||
jwt_handler.get_object_id(token=claims, default_value=None) or claim_user_id
|
||||
if jwt_handler.get_rbac_role(token=claims) == LitellmUserRoles.INTERNAL_USER
|
||||
else claim_user_id
|
||||
)
|
||||
agent_id: Final = JWTAuthManager.resolve_agent_id(jwt_handler, claims, jwt_handler.agent_lookup)
|
||||
is_admin: Final = jwt_handler.is_admin(scopes=jwt_handler.get_scopes(token=claims))
|
||||
try:
|
||||
user, _, _, _, canonical_id = await JWTAuthManager.get_objects(
|
||||
user_id=user_id,
|
||||
user_email=user_email,
|
||||
org_id=None,
|
||||
end_user_id=None,
|
||||
team_id=None,
|
||||
valid_user_email=valid_user_email,
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
route="",
|
||||
user_id_upsert=False,
|
||||
)
|
||||
except UserNotFoundError:
|
||||
if not is_admin:
|
||||
raise
|
||||
return JWTIdentity(user_id=user_id, user_object=None, agent_id=agent_id)
|
||||
return JWTIdentity(user_id=user_id if is_admin else canonical_id, user_object=user, agent_id=agent_id)
|
||||
|
||||
@staticmethod
|
||||
async def authorize_jwt(
|
||||
api_key: str,
|
||||
jwt_handler: JWTHandler,
|
||||
request_data: dict[str, object],
|
||||
general_settings: dict[str, object],
|
||||
route: str,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
request_headers: dict[str, str] | None = None,
|
||||
request_method: str | None = None,
|
||||
provisioning: _JWTProvisioning | None = None,
|
||||
) -> JWTAuthBuilderResult:
|
||||
"""Resolve and authorize JWT context; only normal admission supplies provisioning."""
|
||||
handler: Final = jwt_handler
|
||||
jwt_valid_token: Final = await JWTAuthManager.authenticate_jwt(api_key, handler)
|
||||
team_id_upsert: Final = provisioning.team_id_upsert if provisioning is not None else False
|
||||
model: Final = request_data.get("model")
|
||||
requested_model: Final = model if isinstance(model, str) else None
|
||||
|
||||
# Check RBAC
|
||||
rbac_role: Final = handler.get_rbac_role(token=jwt_valid_token)
|
||||
await JWTAuthManager.check_rbac_role(handler, jwt_valid_token, general_settings, request_data, route, rbac_role)
|
||||
|
||||
# Check Scope Based Access
|
||||
scopes: Final = jwt_handler.get_scopes(token=jwt_valid_token)
|
||||
if jwt_handler.litellm_jwtauth.enforce_scope_based_access and jwt_handler.litellm_jwtauth.scope_mappings:
|
||||
scopes: Final = handler.get_scopes(token=jwt_valid_token)
|
||||
if handler.litellm_jwtauth.enforce_scope_based_access and handler.litellm_jwtauth.scope_mappings:
|
||||
JWTAuthManager.check_scope_based_access(
|
||||
scope_mappings=jwt_handler.litellm_jwtauth.scope_mappings,
|
||||
scope_mappings=handler.litellm_jwtauth.scope_mappings,
|
||||
scopes=scopes,
|
||||
request_data=request_data,
|
||||
general_settings=general_settings,
|
||||
)
|
||||
|
||||
object_id = jwt_handler.get_object_id(token=jwt_valid_token, default_value=None)
|
||||
object_id = handler.get_object_id(token=jwt_valid_token, default_value=None)
|
||||
|
||||
# Get basic user info
|
||||
user_id, user_email, valid_user_email = await JWTAuthManager.get_user_info(jwt_handler, jwt_valid_token)
|
||||
user_id, user_email, valid_user_email = await JWTAuthManager.get_user_info(handler, jwt_valid_token)
|
||||
|
||||
# Get IDs
|
||||
org_id: Final = jwt_handler.get_org_id(token=jwt_valid_token, default_value=None)
|
||||
end_user_id: Final = jwt_handler.get_end_user_id(token=jwt_valid_token, default_value=None)
|
||||
org_id: Final = handler.get_org_id(token=jwt_valid_token, default_value=None)
|
||||
end_user_id: Final = handler.get_end_user_id(token=jwt_valid_token, default_value=None)
|
||||
team_id: str | None = None
|
||||
team_object: LiteLLM_TeamTable | None = None
|
||||
object_id = jwt_handler.get_object_id(token=jwt_valid_token, default_value=None)
|
||||
object_id = handler.get_object_id(token=jwt_valid_token, default_value=None)
|
||||
|
||||
if rbac_role and object_id:
|
||||
if rbac_role == LitellmUserRoles.TEAM:
|
||||
|
|
@ -2331,14 +2434,14 @@ class JWTAuthManager:
|
|||
user_id = object_id
|
||||
|
||||
agent_id: Final = JWTAuthManager.resolve_agent_id(
|
||||
jwt_handler=jwt_handler,
|
||||
jwt_handler=handler,
|
||||
jwt_valid_token=jwt_valid_token,
|
||||
agent_registry=jwt_handler.agent_lookup,
|
||||
agent_registry=handler.agent_lookup,
|
||||
)
|
||||
|
||||
# Check admin access
|
||||
admin_result: Final = await JWTAuthManager.check_admin_access(
|
||||
jwt_handler,
|
||||
handler,
|
||||
scopes,
|
||||
route,
|
||||
user_id,
|
||||
|
|
@ -2353,18 +2456,24 @@ class JWTAuthManager:
|
|||
admin_result=admin_result,
|
||||
route=route,
|
||||
request_headers=request_headers,
|
||||
jwt_handler=jwt_handler,
|
||||
jwt_handler=handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id_upsert=team_id_upsert,
|
||||
)
|
||||
if provisioning is None:
|
||||
identity: Final = await JWTAuthManager._resolve_claim_identity(
|
||||
jwt_valid_token, handler, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj
|
||||
)
|
||||
return {**admin_result, "user_object": identity.user_object}
|
||||
return admin_result
|
||||
|
||||
# Get team with model access
|
||||
## Check if team_id is specified via x-litellm-team-id header
|
||||
all_team_ids: Final = JWTAuthManager.get_all_team_ids(jwt_handler, jwt_valid_token)
|
||||
specific_team_id: Final = jwt_handler.get_team_id(token=jwt_valid_token, default_value=None)
|
||||
all_team_ids: Final = JWTAuthManager.get_all_team_ids(handler, jwt_valid_token)
|
||||
specific_team_id: Final = handler.get_team_id(token=jwt_valid_token, default_value=None)
|
||||
|
||||
# The DB fallback only applies when the token carries no team identity at
|
||||
# all. `get_all_jwt_team_ids` ignores `team_id_default` so a configured
|
||||
|
|
@ -2374,9 +2483,9 @@ class JWTAuthManager:
|
|||
# the RBAC team-role path (which already set `team_id`); otherwise a
|
||||
# provisional x-litellm-team-id header could override an RBAC-asserted team.
|
||||
db_team_fallback: Final = (
|
||||
jwt_handler.litellm_jwtauth.fallback_to_db_teams
|
||||
and not jwt_handler.get_all_jwt_team_ids(token=jwt_valid_token)
|
||||
and not jwt_handler.get_team_alias(token=jwt_valid_token, default_value=None)
|
||||
handler.litellm_jwtauth.fallback_to_db_teams
|
||||
and not handler.get_all_jwt_team_ids(token=jwt_valid_token)
|
||||
and not handler.get_team_alias(token=jwt_valid_token, default_value=None)
|
||||
and team_id is None
|
||||
)
|
||||
if specific_team_id and not db_team_fallback:
|
||||
|
|
@ -2401,7 +2510,7 @@ class JWTAuthManager:
|
|||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id_upsert=(jwt_handler.litellm_jwtauth.team_id_upsert and not db_team_fallback),
|
||||
team_id_upsert=(team_id_upsert and not db_team_fallback),
|
||||
)
|
||||
except HTTPException:
|
||||
if not db_team_fallback:
|
||||
|
|
@ -2413,22 +2522,23 @@ class JWTAuthManager:
|
|||
team_id,
|
||||
team_object,
|
||||
) = await JWTAuthManager.find_and_validate_specific_team_id(
|
||||
jwt_handler,
|
||||
handler,
|
||||
jwt_valid_token,
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
parent_otel_span,
|
||||
proxy_logging_obj,
|
||||
team_id_upsert=team_id_upsert,
|
||||
)
|
||||
|
||||
if not team_object and not team_id:
|
||||
## CHECK USER GROUP ACCESS
|
||||
team_id, team_object = await JWTAuthManager.find_team_with_model_access(
|
||||
team_ids=all_team_ids,
|
||||
requested_model=request_data.get("model"),
|
||||
requested_model=requested_model,
|
||||
route=route,
|
||||
request_method=request_method,
|
||||
jwt_handler=jwt_handler,
|
||||
jwt_handler=handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
|
|
@ -2452,7 +2562,7 @@ class JWTAuthManager:
|
|||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert,
|
||||
team_id_upsert=team_id_upsert,
|
||||
)
|
||||
|
||||
if team_id and not JWTAuthManager._team_has_passthrough_route_access(
|
||||
|
|
@ -2463,7 +2573,7 @@ class JWTAuthManager:
|
|||
JWTAuthManager._raise_team_passthrough_route_denial(route=route)
|
||||
|
||||
# Extract alias fields for resolution (if configured)
|
||||
org_alias: Final = jwt_handler.get_org_alias(token=jwt_valid_token, default_value=None)
|
||||
org_alias: Final = handler.get_org_alias(token=jwt_valid_token, default_value=None)
|
||||
|
||||
# get_objects returns effective_user_id for downstream spend attribution (GH #26789).
|
||||
(
|
||||
|
|
@ -2479,25 +2589,27 @@ class JWTAuthManager:
|
|||
end_user_id=end_user_id,
|
||||
team_id=team_id,
|
||||
valid_user_email=valid_user_email,
|
||||
jwt_handler=jwt_handler,
|
||||
jwt_handler=handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
route=route,
|
||||
org_alias=org_alias,
|
||||
user_id_upsert=provisioning.user_id_upsert if provisioning is not None else False,
|
||||
)
|
||||
|
||||
# Derive org_id from org_object if resolved by alias
|
||||
resolved_org_id: Final = org_object.organization_id if org_object else org_id
|
||||
|
||||
await JWTAuthManager.sync_user_role_and_teams(
|
||||
jwt_handler=jwt_handler,
|
||||
jwt_valid_token=jwt_valid_token,
|
||||
user_object=user_object,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
if provisioning is not None:
|
||||
await JWTAuthManager.sync_user_role_and_teams(
|
||||
jwt_handler=handler,
|
||||
jwt_valid_token=jwt_valid_token,
|
||||
user_object=user_object,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
# If JWT did not resolve team_id, attempt a team fallback.
|
||||
if team_id is None and db_team_fallback:
|
||||
|
|
@ -2508,11 +2620,11 @@ class JWTAuthManager:
|
|||
) = await JWTAuthManager._resolve_db_team_fallback(
|
||||
user_object=user_object,
|
||||
user_id=user_id,
|
||||
requested_model=request_data.get("model"),
|
||||
requested_model=requested_model,
|
||||
route=route,
|
||||
jwt_handler=jwt_handler,
|
||||
enforce_team_based_model_access=jwt_handler.litellm_jwtauth.enforce_team_based_model_access,
|
||||
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert,
|
||||
jwt_handler=handler,
|
||||
enforce_team_based_model_access=handler.litellm_jwtauth.enforce_team_based_model_access,
|
||||
team_id_upsert=team_id_upsert,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
|
|
@ -2540,7 +2652,7 @@ class JWTAuthManager:
|
|||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert,
|
||||
team_id_upsert=team_id_upsert,
|
||||
)
|
||||
elif db_team_fallback and team_id == header_team_id:
|
||||
JWTAuthManager._validate_header_team_in_db_membership(
|
||||
|
|
@ -2550,7 +2662,7 @@ class JWTAuthManager:
|
|||
if not JWTAuthManager._is_team_route_allowed(
|
||||
route=route,
|
||||
request_method=request_method,
|
||||
jwt_handler=jwt_handler,
|
||||
jwt_handler=handler,
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
|
|
@ -2560,16 +2672,17 @@ class JWTAuthManager:
|
|||
)
|
||||
|
||||
## MAP USER TO TEAMS
|
||||
await JWTAuthManager.map_user_to_teams(
|
||||
user_object=user_object,
|
||||
team_object=team_object,
|
||||
)
|
||||
if provisioning is not None:
|
||||
await JWTAuthManager.map_user_to_teams(
|
||||
user_object=user_object,
|
||||
team_object=team_object,
|
||||
)
|
||||
|
||||
# Validate that a valid rbac id is returned for spend tracking
|
||||
JWTAuthManager.validate_object_id(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
enforce_rbac=general_settings.get("enforce_rbac", False),
|
||||
enforce_rbac=bool(general_settings.get("enforce_rbac", False)),
|
||||
is_proxy_admin=False,
|
||||
)
|
||||
|
||||
|
|
@ -2592,3 +2705,38 @@ class JWTAuthManager:
|
|||
jwt_claims=jwt_valid_token,
|
||||
agent_id=agent_id,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def user_api_key_auth_from_result(
|
||||
result: JWTAuthBuilderResult,
|
||||
parent_otel_span: Span | None = None,
|
||||
) -> UserAPIKeyAuth:
|
||||
"""Keep JWT identity and permission attribution identical across consumers."""
|
||||
user: Final = result["user_object"]
|
||||
admin: Final = result["is_proxy_admin"]
|
||||
return UserAPIKeyAuth(
|
||||
api_key=None,
|
||||
user_role=(
|
||||
LitellmUserRoles.PROXY_ADMIN
|
||||
if admin
|
||||
else LitellmUserRoles(user.user_role)
|
||||
if user is not None and user.user_role is not None
|
||||
else LitellmUserRoles.INTERNAL_USER
|
||||
),
|
||||
user_id=result["user_id"],
|
||||
user_email=result["user_email"],
|
||||
team_id=result["team_id"],
|
||||
org_id=result["org_id"],
|
||||
end_user_id=result["end_user_id"],
|
||||
parent_otel_span=parent_otel_span,
|
||||
jwt_claims=result["jwt_claims"],
|
||||
agent_id=result.get("agent_id"),
|
||||
user_tpm_limit=user.tpm_limit if user is not None and not admin else None,
|
||||
user_rpm_limit=user.rpm_limit if user is not None and not admin else None,
|
||||
user_model_max_budget=user.model_max_budget if user is not None and not admin else None,
|
||||
**team_grants(
|
||||
team_object=result["team_object"],
|
||||
team_membership=result.get("team_membership"),
|
||||
user_id=result["user_id"],
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -155,8 +155,8 @@ class LicenseCheck:
|
|||
|
||||
def auto_router_capability_limit(self) -> int | None:
|
||||
"""
|
||||
How many auto-routers may claim each licensed capability (heuristic_v2, operator-defined
|
||||
tier_definitions): unlimited (None) only when the signed license lists the auto_router
|
||||
How many auto-routers may claim each gated classifier or customization capability:
|
||||
unlimited (None) only when the signed license lists the auto_router
|
||||
feature, otherwise one per capability. A license verified through the API carries no
|
||||
feature list, so it does not lift the limit either.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1669,13 +1669,11 @@ async def _user_api_key_auth_builder(
|
|||
|
||||
is_proxy_admin: Final = result["is_proxy_admin"]
|
||||
team_id: Final = result["team_id"]
|
||||
team_object: Final = result["team_object"]
|
||||
user_id: Final = result["user_id"]
|
||||
user_email: Final = result["user_email"]
|
||||
user_object: Final = result["user_object"]
|
||||
end_user_id = result["end_user_id"]
|
||||
org_id: Final = result["org_id"]
|
||||
team_membership: Final[LiteLLM_TeamMembership | None] = result.get("team_membership", None)
|
||||
jwt_claims = result.get("jwt_claims", None)
|
||||
agent_id: Final[str | None] = result.get("agent_id")
|
||||
|
||||
|
|
@ -1693,40 +1691,9 @@ async def _user_api_key_auth_builder(
|
|||
value=_JWT_PROXY_ADMIN_SENTINEL,
|
||||
ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl,
|
||||
)
|
||||
return UserAPIKeyAuth(
|
||||
api_key=None,
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_id=user_id,
|
||||
user_email=user_email,
|
||||
team_id=team_id,
|
||||
org_id=org_id,
|
||||
end_user_id=end_user_id,
|
||||
parent_otel_span=parent_otel_span,
|
||||
jwt_claims=jwt_claims,
|
||||
agent_id=agent_id,
|
||||
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
|
||||
)
|
||||
return JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span)
|
||||
|
||||
valid_token = UserAPIKeyAuth(
|
||||
api_key=None,
|
||||
team_id=team_id,
|
||||
user_role=(
|
||||
LitellmUserRoles(user_object.user_role)
|
||||
if user_object is not None and user_object.user_role is not None
|
||||
else LitellmUserRoles.INTERNAL_USER
|
||||
),
|
||||
user_id=user_id,
|
||||
user_email=user_email,
|
||||
org_id=org_id,
|
||||
parent_otel_span=parent_otel_span,
|
||||
end_user_id=end_user_id,
|
||||
user_tpm_limit=(user_object.tpm_limit if user_object is not None else None),
|
||||
user_rpm_limit=(user_object.rpm_limit if user_object is not None else None),
|
||||
user_model_max_budget=(user_object.model_max_budget if user_object is not None else None),
|
||||
jwt_claims=jwt_claims,
|
||||
agent_id=agent_id,
|
||||
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
|
||||
)
|
||||
valid_token = JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span)
|
||||
|
||||
# AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key.
|
||||
# JWT policy (RBAC, scope, custom_validate, email-domain)
|
||||
|
|
|
|||
|
|
@ -57,6 +57,8 @@ _COUNTER_COLUMNS: Final = (
|
|||
"cache_read_input_tokens",
|
||||
"cache_creation_input_tokens",
|
||||
"compression_saved_tokens",
|
||||
"total_response_time_ms",
|
||||
"timed_requests",
|
||||
)
|
||||
_SPEND_COLUMNS: Final = (
|
||||
"spend",
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from collections.abc import Mapping, Sequence
|
|||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload
|
||||
from urllib.parse import quote, unquote
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -85,6 +86,10 @@ else:
|
|||
RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value})
|
||||
|
||||
|
||||
def _org_member_transaction_key(org_id: str, user_id: str) -> str:
|
||||
return f"organization_id::{quote(org_id, safe='')}::user_id::{quote(user_id, safe='')}"
|
||||
|
||||
|
||||
def _is_batch_cost_row(payload: SpendLogsPayload) -> bool:
|
||||
return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success"
|
||||
|
||||
|
|
@ -110,6 +115,7 @@ class _SpendBatch(Protocol):
|
|||
litellm_teamtable: BatchTable
|
||||
litellm_teammembership: BatchTable
|
||||
litellm_organizationtable: BatchTable
|
||||
litellm_organizationmembership: BatchTable
|
||||
litellm_tagtable: BatchTable
|
||||
litellm_agentstable: BatchTable
|
||||
litellm_modelaccessgroupbudgettable: BatchTable
|
||||
|
|
@ -131,6 +137,19 @@ class _SpendTransactionManager(Protocol):
|
|||
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...
|
||||
|
||||
|
||||
def _timed_request_duration_ms(
|
||||
payload: dict | SpendLogsPayload,
|
||||
request_status: Literal["success", "failure"],
|
||||
is_internal_call: bool,
|
||||
) -> int | None:
|
||||
if is_internal_call or request_status != "success":
|
||||
return None
|
||||
duration_ms: Final = payload.get("request_duration_ms")
|
||||
if not isinstance(duration_ms, int) or duration_ms < 0:
|
||||
return None
|
||||
return duration_ms
|
||||
|
||||
|
||||
def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager:
|
||||
tx: Final[_SpendTransactionManager] = prisma_client.db.tx(timeout=timedelta(seconds=60))
|
||||
return tx
|
||||
|
|
@ -666,6 +685,7 @@ class DBSpendUpdateWriter:
|
|||
await self._update_org_db(
|
||||
response_cost=response_cost,
|
||||
org_id=org_id,
|
||||
user_id=user_id,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
except Exception:
|
||||
|
|
@ -900,6 +920,7 @@ class DBSpendUpdateWriter:
|
|||
self,
|
||||
response_cost: float | None,
|
||||
org_id: str | None,
|
||||
user_id: str | None,
|
||||
prisma_client: PrismaClient | None,
|
||||
):
|
||||
try:
|
||||
|
|
@ -916,6 +937,15 @@ class DBSpendUpdateWriter:
|
|||
response_cost=response_cost,
|
||||
)
|
||||
)
|
||||
|
||||
if user_id is not None:
|
||||
await self.spend_update_queue.add_update(
|
||||
update=SpendUpdateQueueItem(
|
||||
entity_type=Litellm_EntityType.ORGANIZATION_MEMBER,
|
||||
entity_id=_org_member_transaction_key(org_id, user_id),
|
||||
response_cost=response_cost,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
spend_log_error(
|
||||
"Spend tracking - failed to enqueue org spend update. org_id=%s, response_cost=%s - %s",
|
||||
|
|
@ -1163,14 +1193,15 @@ class DBSpendUpdateWriter:
|
|||
if db_spend_update_transactions is not None:
|
||||
verbose_proxy_logger.info(
|
||||
"Spend tracking - committing spend updates from Redis to DB: "
|
||||
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d, agents=%d, "
|
||||
"model_access_groups=%d",
|
||||
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, tags=%d, "
|
||||
"agents=%d, model_access_groups=%d",
|
||||
len(db_spend_update_transactions.get("key_list_transactions") or {}),
|
||||
len(db_spend_update_transactions.get("user_list_transactions") or {}),
|
||||
len(db_spend_update_transactions.get("team_list_transactions") or {}),
|
||||
len(db_spend_update_transactions.get("org_list_transactions") or {}),
|
||||
len(db_spend_update_transactions.get("end_user_list_transactions") or {}),
|
||||
len(db_spend_update_transactions.get("team_member_list_transactions") or {}),
|
||||
len(db_spend_update_transactions.get("org_member_list_transactions") or {}),
|
||||
len(db_spend_update_transactions.get("tag_list_transactions") or {}),
|
||||
len(db_spend_update_transactions.get("agent_list_transactions") or {}),
|
||||
len(db_spend_update_transactions.get("model_access_group_list_transactions") or {}),
|
||||
|
|
@ -1708,6 +1739,29 @@ class DBSpendUpdateWriter:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
org_member_list_transactions: Final = db_spend_update_transactions.get("org_member_list_transactions")
|
||||
verbose_proxy_logger.debug("Org Membership Spend transactions: %s", org_member_list_transactions)
|
||||
if org_member_list_transactions is not None and len(org_member_list_transactions.keys()) > 0:
|
||||
for i in range(n_retry_times + 1):
|
||||
start_time = time.time()
|
||||
try:
|
||||
async with _spend_update_tx(prisma_client) as transaction, transaction.batch_() as batcher:
|
||||
for key, response_cost in sorted(org_member_list_transactions.items()):
|
||||
_, quoted_org_id, _, quoted_user_id = key.split("::")
|
||||
batcher.litellm_organizationmembership.update_many(
|
||||
where={"organization_id": unquote(quoted_org_id), "user_id": unquote(quoted_user_id)},
|
||||
data={"spend": {"increment": response_cost}},
|
||||
)
|
||||
break
|
||||
except Exception as e:
|
||||
await self._handle_spend_update_failure(
|
||||
e=e,
|
||||
attempt=i,
|
||||
n_retry_times=n_retry_times,
|
||||
start_time=start_time,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
### UPDATE TAG TABLE ###
|
||||
tag_list_transactions: Final = db_spend_update_transactions["tag_list_transactions"]
|
||||
await DBSpendUpdateWriter._update_entity_spend_in_db(
|
||||
|
|
@ -2191,6 +2245,7 @@ class DBSpendUpdateWriter:
|
|||
recorded_autorouter_savings=_metadata.get("autorouter_savings"),
|
||||
billed_at=payload.get("endTime"),
|
||||
)
|
||||
timed_duration_ms: Final = _timed_request_duration_ms(payload, request_status, is_internal_call)
|
||||
|
||||
daily_transaction: Final = BaseDailySpendTransaction(
|
||||
date=date,
|
||||
|
|
@ -2218,6 +2273,8 @@ class DBSpendUpdateWriter:
|
|||
prompt_caching_savings_spend=savings_spend.prompt_caching,
|
||||
gateway_injected_caching_savings_spend=savings_spend.gateway_injected_caching,
|
||||
autorouter_savings_spend=0.0 if is_internal_call else savings_spend.autorouter,
|
||||
total_response_time_ms=timed_duration_ms or 0,
|
||||
timed_requests=0 if timed_duration_ms is None else 1,
|
||||
)
|
||||
return daily_transaction
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -142,6 +142,14 @@ class DailySpendUpdateQueue(BaseUpdateQueue):
|
|||
payload.get("autorouter_savings_spend", 0) or 0
|
||||
) + daily_transaction.get("autorouter_savings_spend", 0)
|
||||
|
||||
daily_transaction["total_response_time_ms"] = (
|
||||
payload.get("total_response_time_ms", 0) or 0
|
||||
) + daily_transaction.get("total_response_time_ms", 0)
|
||||
|
||||
daily_transaction["timed_requests"] = (
|
||||
payload.get("timed_requests", 0) or 0
|
||||
) + daily_transaction.get("timed_requests", 0)
|
||||
|
||||
else:
|
||||
aggregated_daily_spend_update_transactions[_key] = deepcopy(payload)
|
||||
return aggregated_daily_spend_update_transactions
|
||||
|
|
|
|||
|
|
@ -69,6 +69,7 @@ _SpendTransactionField: TypeAlias = Literal[
|
|||
"team_list_transactions",
|
||||
"team_member_list_transactions",
|
||||
"org_list_transactions",
|
||||
"org_member_list_transactions",
|
||||
"tag_list_transactions",
|
||||
"agent_list_transactions",
|
||||
"model_access_group_list_transactions",
|
||||
|
|
@ -81,6 +82,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = (
|
|||
"team_list_transactions",
|
||||
"team_member_list_transactions",
|
||||
"org_list_transactions",
|
||||
"org_member_list_transactions",
|
||||
"tag_list_transactions",
|
||||
"agent_list_transactions",
|
||||
"model_access_group_list_transactions",
|
||||
|
|
@ -412,6 +414,10 @@ class RedisUpdateBuffer:
|
|||
Litellm_EntityType.ORGANIZATION,
|
||||
db_spend_update_transactions.get("org_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.ORGANIZATION_MEMBER,
|
||||
db_spend_update_transactions.get("org_member_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.TAG,
|
||||
db_spend_update_transactions.get("tag_list_transactions"),
|
||||
|
|
@ -876,6 +882,9 @@ class RedisUpdateBuffer:
|
|||
list_of_transactions, "team_member_list_transactions"
|
||||
),
|
||||
org_list_transactions=_merged_entity_transactions(list_of_transactions, "org_list_transactions"),
|
||||
org_member_list_transactions=_merged_entity_transactions(
|
||||
list_of_transactions, "org_member_list_transactions"
|
||||
),
|
||||
tag_list_transactions=_merged_entity_transactions(list_of_transactions, "tag_list_transactions"),
|
||||
agent_list_transactions=_merged_entity_transactions(list_of_transactions, "agent_list_transactions"),
|
||||
model_access_group_list_transactions=_merged_entity_transactions(
|
||||
|
|
|
|||
|
|
@ -137,6 +137,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
|
|||
team_list_transactions={},
|
||||
team_member_list_transactions={},
|
||||
org_list_transactions={},
|
||||
org_member_list_transactions={},
|
||||
tag_list_transactions={},
|
||||
agent_list_transactions={},
|
||||
model_access_group_list_transactions={},
|
||||
|
|
@ -150,6 +151,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
|
|||
Litellm_EntityType.TEAM: "team_list_transactions",
|
||||
Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions",
|
||||
Litellm_EntityType.ORGANIZATION: "org_list_transactions",
|
||||
Litellm_EntityType.ORGANIZATION_MEMBER: "org_member_list_transactions",
|
||||
Litellm_EntityType.TAG: "tag_list_transactions",
|
||||
Litellm_EntityType.AGENT: "agent_list_transactions",
|
||||
Litellm_EntityType.MODEL_ACCESS_GROUP: "model_access_group_list_transactions",
|
||||
|
|
@ -188,6 +190,8 @@ class SpendUpdateQueue(BaseUpdateQueue):
|
|||
transactions_dict = db_spend_update_transactions["team_member_list_transactions"]
|
||||
elif dict_key == "org_list_transactions":
|
||||
transactions_dict = db_spend_update_transactions["org_list_transactions"]
|
||||
elif dict_key == "org_member_list_transactions":
|
||||
transactions_dict = db_spend_update_transactions["org_member_list_transactions"]
|
||||
elif dict_key == "tag_list_transactions":
|
||||
transactions_dict = db_spend_update_transactions["tag_list_transactions"]
|
||||
elif dict_key == "agent_list_transactions":
|
||||
|
|
|
|||
|
|
@ -0,0 +1,63 @@
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.agent_365 import (
|
||||
AGENT_365_PROD_API_BASE,
|
||||
AGENT_365_PROD_RESOURCE_APP_ID,
|
||||
)
|
||||
|
||||
from .agent_365 import Agent365Guardrail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> Agent365Guardrail:
|
||||
import litellm
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
tenant_id: Final = litellm_params.tenant_id or get_secret_str("AGENT365_TENANT_ID")
|
||||
client_id: Final = litellm_params.client_id or get_secret_str("AGENT365_CLIENT_ID")
|
||||
client_secret: Final = (
|
||||
litellm_params.client_secret or litellm_params.api_key or get_secret_str("AGENT365_CLIENT_SECRET")
|
||||
)
|
||||
api_base: Final = litellm_params.api_base or get_secret_str("AGENT365_API_BASE")
|
||||
resource_app_id: Final = litellm_params.resource_app_id or get_secret_str("AGENT365_RESOURCE_APP_ID")
|
||||
|
||||
if not tenant_id:
|
||||
raise ValueError("Microsoft Agent 365: tenant_id is required")
|
||||
if not client_id:
|
||||
raise ValueError("Microsoft Agent 365: client_id is required")
|
||||
if not client_secret:
|
||||
raise ValueError(
|
||||
"Microsoft Agent 365: client secret is required. Set client_secret, api_key, or AGENT365_CLIENT_SECRET"
|
||||
)
|
||||
|
||||
guardrail_name: Final = guardrail.get("guardrail_name")
|
||||
if not guardrail_name:
|
||||
raise ValueError("Microsoft Agent 365: guardrail_name is required")
|
||||
|
||||
agent_365_guardrail: Final = Agent365Guardrail(
|
||||
guardrail_name=guardrail_name,
|
||||
tenant_id=tenant_id,
|
||||
client_id=client_id,
|
||||
client_secret=client_secret,
|
||||
api_base=api_base or AGENT_365_PROD_API_BASE,
|
||||
resource_app_id=resource_app_id or AGENT_365_PROD_RESOURCE_APP_ID,
|
||||
agent_id=litellm_params.agent_id,
|
||||
request_timeout=litellm_params.timeout if litellm_params.timeout is not None else 10.0,
|
||||
unreachable_fallback=litellm_params.unreachable_fallback,
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback(agent_365_guardrail)
|
||||
return agent_365_guardrail
|
||||
|
||||
|
||||
guardrail_initializer_registry: Final = { # mutable-ok: registry auto-discovery requires a dict instance
|
||||
SupportedGuardrailIntegrations.AGENT_365.value: initialize_guardrail,
|
||||
}
|
||||
|
||||
guardrail_class_registry: Final = { # mutable-ok: registry auto-discovery requires a dict instance
|
||||
SupportedGuardrailIntegrations.AGENT_365.value: Agent365Guardrail,
|
||||
}
|
||||
637
litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py
Normal file
637
litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py
Normal file
|
|
@ -0,0 +1,637 @@
|
|||
"""Microsoft Agent 365 governance guardrail for MCP tool calls.
|
||||
|
||||
Before the gateway executes an MCP tool, the pending call is sent to the
|
||||
Agent 365 tool-evaluation endpoint, where Microsoft Defender scores it and
|
||||
Agent 365 records it for observability. The returned allow/block verdict is
|
||||
enforced here. Authentication is the Entra On-Behalf-Of flow: the caller's
|
||||
incoming bearer token (audienced to this gateway's app registration) is
|
||||
exchanged for a delegated Agent 365 token, so Defender evaluates and audits
|
||||
as the signed-in user.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, ClassVar, Final, Literal, NoReturn
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.exceptions import Timeout as LitellmTimeout
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.agent_365 import (
|
||||
AGENT_365_PROD_API_BASE,
|
||||
AGENT_365_PROD_RESOURCE_APP_ID,
|
||||
AGENT_365_SCOPE_NAME,
|
||||
Agent365GuardrailConfigModel,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
from litellm.types.utils import GuardrailStatus
|
||||
|
||||
TOKEN_ENDPOINT_TEMPLATE: Final = "https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token"
|
||||
EVALUATE_PATH: Final = "/agents/tool-evaluation/evaluate"
|
||||
MCP_SESSION_ID_HEADER: Final = "mcp-session-id"
|
||||
DEFENDER_STATUS_EVALUATED: Final = "Evaluated"
|
||||
_GATEWAY_OWNED_TOKEN_ERRORS: Final = frozenset(
|
||||
{"invalid_client", "unauthorized_client", "invalid_scope", "invalid_resource"}
|
||||
)
|
||||
# Entra reports a malformed or unverifiable assertion as ``invalid_client`` too; only its AADSTS50027xx
|
||||
# (InvalidJwtToken) sub-codes tell that apart from a bad gateway secret.
|
||||
_INVALID_ASSERTION_AADSTS_PREFIX: Final = "50027"
|
||||
_AADSTS_CODES_ADAPTER: Final = TypeAdapter(tuple[int, ...])
|
||||
_MCP_CALL_TYPES: Final[tuple[str, ...]] = ("mcp_call", "call_mcp_tool")
|
||||
_OBO_CACHE_MAX_ENTRIES: Final = 1000
|
||||
_DEFAULT_TOKEN_TTL_SECONDS: Final = 3599.0
|
||||
_TOKEN_EXPIRY_SLACK_SECONDS: Final = 60.0
|
||||
|
||||
|
||||
def _parse_expires_in(raw: object) -> float:
|
||||
if not isinstance(raw, (int, float, str)):
|
||||
return _DEFAULT_TOKEN_TTL_SECONDS
|
||||
try:
|
||||
return float(raw)
|
||||
except ValueError:
|
||||
return _DEFAULT_TOKEN_TTL_SECONDS
|
||||
|
||||
|
||||
def _parse_aadsts_codes(raw: object) -> tuple[int, ...]:
|
||||
try:
|
||||
return _AADSTS_CODES_ADAPTER.validate_python(raw)
|
||||
except ValidationError:
|
||||
return ()
|
||||
|
||||
|
||||
def entra_assertion(value: object) -> str | None:
|
||||
"""``value`` when it is a compact JWS, the only bearer shape the OBO exchange accepts as its assertion.
|
||||
A LiteLLM virtual key, session bearer, or opaque upstream token in ``Authorization`` yields ``None``."""
|
||||
return value if isinstance(value, str) and value.count(".") == 2 else None
|
||||
|
||||
|
||||
class _DefenderResult(TypedDict, total=False):
|
||||
status: ReadOnly[str]
|
||||
verdict: ReadOnly[str | None]
|
||||
message: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _EvaluateResponse(TypedDict, total=False):
|
||||
allowed: ReadOnly[bool]
|
||||
defender: ReadOnly[_DefenderResult]
|
||||
correlationId: ReadOnly[str]
|
||||
|
||||
|
||||
class _UnavailableDetail(TypedDict):
|
||||
error: ReadOnly[str]
|
||||
message: ReadOnly[str]
|
||||
tool: ReadOnly[str]
|
||||
|
||||
|
||||
class _BlockedDetail(TypedDict):
|
||||
error: ReadOnly[str]
|
||||
message: ReadOnly[str]
|
||||
tool: ReadOnly[str]
|
||||
correlation_id: ReadOnly[str | None]
|
||||
|
||||
|
||||
class Agent365TokenExchangeError(Exception):
|
||||
def __init__(self, status_code: int, error_code: str, description: str, aadsts_codes: tuple[int, ...] = ()) -> None:
|
||||
super().__init__(f"{error_code}: {description}")
|
||||
self.status_code = status_code
|
||||
self.error_code = error_code
|
||||
self.description = description
|
||||
self.aadsts_codes = aadsts_codes
|
||||
|
||||
@property
|
||||
def gateway_owned(self) -> bool:
|
||||
"""Whether the gateway's own client credentials, scope or resource were refused, as opposed to the
|
||||
caller's assertion. The caller cannot fix a gateway-owned rejection by signing in again."""
|
||||
if self.error_code not in _GATEWAY_OWNED_TOKEN_ERRORS:
|
||||
return False
|
||||
return not any(str(code).startswith(_INVALID_ASSERTION_AADSTS_PREFIX) for code in self.aadsts_codes)
|
||||
|
||||
|
||||
class Agent365MalformedResponseError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class Agent365ThrottledError(Exception):
|
||||
def __init__(self, status_code: int) -> None:
|
||||
super().__init__(f"HTTP {status_code}")
|
||||
self.status_code = status_code
|
||||
|
||||
|
||||
class Agent365Guardrail(CustomGuardrail):
|
||||
"""Pre-MCP-call guardrail enforcing Microsoft Agent 365 tool-evaluation verdicts.
|
||||
|
||||
Block-only: it never rewrites the call, so it runs in the post-sequential phase and judges the
|
||||
arguments the sequential guardrails hand upstream, whatever order the guardrails list uses."""
|
||||
|
||||
records_own_guardrail_information: ClassVar[bool] = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
guardrail_name: str,
|
||||
tenant_id: str,
|
||||
client_id: str,
|
||||
client_secret: str,
|
||||
api_base: str = AGENT_365_PROD_API_BASE,
|
||||
resource_app_id: str = AGENT_365_PROD_RESOURCE_APP_ID,
|
||||
agent_id: str | None = None,
|
||||
request_timeout: float = 10.0,
|
||||
unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed",
|
||||
async_handler: AsyncHTTPHandler | None = None,
|
||||
**kwargs, # noqa: ANN003 # kwargs-ok: forwarded verbatim to CustomGuardrail (event_hook, default_on)
|
||||
) -> None:
|
||||
super().__init__(
|
||||
guardrail_name=guardrail_name,
|
||||
supported_event_hooks=self.get_supported_event_hooks(),
|
||||
run_in_parallel=True,
|
||||
**kwargs,
|
||||
)
|
||||
self.guardrail_provider = "agent_365"
|
||||
self.tenant_id = tenant_id
|
||||
self.client_id = client_id
|
||||
self.client_secret = client_secret
|
||||
self.api_base = api_base.rstrip("/")
|
||||
self.resource_app_id = resource_app_id
|
||||
self.agent_id = agent_id
|
||||
self.request_timeout = request_timeout
|
||||
self.unreachable_fallback: Literal["fail_closed", "fail_open"] = (
|
||||
"fail_open" if unreachable_fallback == "fail_open" else "fail_closed"
|
||||
)
|
||||
self.async_handler = async_handler or get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.GuardrailCallback
|
||||
)
|
||||
self._obo_token_cache: OrderedDict[str, tuple[str, float]] = OrderedDict() # mutable-ok: lock-guarded LRU
|
||||
self._obo_cache_lock = threading.Lock()
|
||||
verbose_proxy_logger.info("Initialized Microsoft Agent 365 guardrail: %s", guardrail_name)
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> "type[GuardrailConfigModel] | None":
|
||||
return Agent365GuardrailConfigModel
|
||||
|
||||
@classmethod
|
||||
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: # mutable-ok: CustomGuardrail contract
|
||||
return [GuardrailEventHooks.pre_mcp_call] # mutable-ok: CustomGuardrail contract expects a list
|
||||
|
||||
@log_guardrail_information
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
cache: "DualCache",
|
||||
data: dict, # mutable-ok: hook contract; guardrail logging appends into the request metadata in place
|
||||
call_type: str,
|
||||
) -> Exception | str | dict | None: # mutable-ok: CustomGuardrail.async_pre_call_hook contract
|
||||
if call_type not in _MCP_CALL_TYPES:
|
||||
return data
|
||||
if "mcp_tool_name" not in data:
|
||||
return data
|
||||
if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_mcp_call) is not True:
|
||||
return data
|
||||
|
||||
tool_name: Final = str(data.get("mcp_tool_name") or "")
|
||||
assertion: Final = entra_assertion(data.get("incoming_bearer_token"))
|
||||
if assertion is None:
|
||||
self._handle_caller_fault(
|
||||
data=data,
|
||||
tool_name=tool_name,
|
||||
status_code=401,
|
||||
reason=(
|
||||
"the caller did not present an Entra bearer token; the Agent 365 guardrail "
|
||||
"authorizes tool calls On-Behalf-Of the signed-in user"
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
obo_token: Final = await self._get_obo_token(assertion)
|
||||
except Agent365TokenExchangeError as exc:
|
||||
if exc.gateway_owned:
|
||||
return self._handle_unavailable(
|
||||
data=data,
|
||||
tool_name=tool_name,
|
||||
reason=(
|
||||
f"Entra rejected the gateway's own Agent 365 credentials ({exc.error_code}); "
|
||||
"check the guardrail's client_id, client_secret and resource_app_id"
|
||||
),
|
||||
)
|
||||
self._handle_caller_fault(
|
||||
data=data,
|
||||
tool_name=tool_name,
|
||||
status_code=401,
|
||||
reason=f"the Entra On-Behalf-Of token exchange was rejected ({exc.error_code})",
|
||||
)
|
||||
except Agent365ThrottledError as exc:
|
||||
self._handle_throttled(
|
||||
data=data,
|
||||
tool_name=tool_name,
|
||||
reason=f"the Entra token endpoint returned HTTP {exc.status_code}",
|
||||
latency_ms=None,
|
||||
)
|
||||
except (httpx.HTTPError, LitellmTimeout, TimeoutError) as exc:
|
||||
return self._handle_unavailable(
|
||||
data=data,
|
||||
tool_name=tool_name,
|
||||
reason=f"the Entra token endpoint could not be reached ({type(exc).__name__})",
|
||||
)
|
||||
except Agent365MalformedResponseError as exc:
|
||||
return self._handle_unavailable(
|
||||
data=data,
|
||||
tool_name=tool_name,
|
||||
reason=str(exc),
|
||||
)
|
||||
|
||||
start: Final = time.perf_counter()
|
||||
try:
|
||||
response: Final = await self._post_allowing_error_status(
|
||||
url=f"{self.api_base}{EVALUATE_PATH}",
|
||||
json=self._build_evaluate_payload(data=data, user_api_key_dict=user_api_key_dict),
|
||||
headers={"Authorization": f"Bearer {obo_token}"}, # mutable-ok: httpx header dict
|
||||
)
|
||||
except (httpx.HTTPError, LitellmTimeout, TimeoutError) as exc:
|
||||
return self._handle_unavailable(
|
||||
data=data,
|
||||
tool_name=tool_name,
|
||||
reason=f"the Agent 365 endpoint could not be reached ({type(exc).__name__})",
|
||||
)
|
||||
latency_ms: Final = (time.perf_counter() - start) * 1000.0
|
||||
fallback: Final = self._handle_evaluate_error(
|
||||
data=data, tool_name=tool_name, assertion=assertion, response=response, latency_ms=latency_ms
|
||||
)
|
||||
if fallback is not None:
|
||||
return fallback
|
||||
return self._enforce_verdict(data=data, tool_name=tool_name, response=response, latency_ms=latency_ms)
|
||||
|
||||
def _handle_evaluate_error(
|
||||
self,
|
||||
data: dict, # mutable-ok: guardrail logging appends into the request metadata in place
|
||||
tool_name: str,
|
||||
assertion: str,
|
||||
response: httpx.Response,
|
||||
latency_ms: float,
|
||||
) -> dict | None: # mutable-ok: returns the request data dict per hook contract on fail_open
|
||||
if response.status_code in (408, 429):
|
||||
self._handle_throttled(
|
||||
data=data,
|
||||
tool_name=tool_name,
|
||||
reason=f"the Agent 365 endpoint returned HTTP {response.status_code}",
|
||||
latency_ms=latency_ms,
|
||||
)
|
||||
if 400 <= response.status_code < 500:
|
||||
if response.status_code == 401:
|
||||
self._evict_obo_token(assertion)
|
||||
self._record_verdict(
|
||||
data=data,
|
||||
verdict="Rejected",
|
||||
guardrail_status="guardrail_intervened",
|
||||
defender_status=None,
|
||||
correlation_id=None,
|
||||
latency_ms=latency_ms,
|
||||
reason=f"HTTP {response.status_code}: {response.text[:512]}",
|
||||
)
|
||||
rejected_detail: Final[_UnavailableDetail] = {
|
||||
"error": "Agent 365 rejected the tool evaluation request",
|
||||
"message": response.text[:512]
|
||||
if response.status_code == 400
|
||||
else f"the Agent 365 evaluation request failed with HTTP {response.status_code}",
|
||||
"tool": tool_name,
|
||||
}
|
||||
raise HTTPException(status_code=400, detail=rejected_detail)
|
||||
if response.status_code != 200:
|
||||
return self._handle_unavailable(
|
||||
data=data,
|
||||
tool_name=tool_name,
|
||||
reason=f"the Agent 365 endpoint returned HTTP {response.status_code}",
|
||||
)
|
||||
return None
|
||||
|
||||
def _enforce_verdict(
|
||||
self,
|
||||
data: dict, # mutable-ok: guardrail logging appends into the request metadata in place
|
||||
tool_name: str,
|
||||
response: httpx.Response,
|
||||
latency_ms: float,
|
||||
) -> dict: # mutable-ok: returns the request data dict per hook contract
|
||||
try:
|
||||
parsed_verdict: Final = response.json()
|
||||
except ValueError:
|
||||
return self._handle_unavailable(
|
||||
data=data,
|
||||
tool_name=tool_name,
|
||||
reason="the Agent 365 endpoint returned a non-JSON body",
|
||||
)
|
||||
if not isinstance(parsed_verdict, dict):
|
||||
return self._handle_unavailable(
|
||||
data=data,
|
||||
tool_name=tool_name,
|
||||
reason="the Agent 365 endpoint returned a non-object JSON body",
|
||||
)
|
||||
verdict: Final[_EvaluateResponse] = parsed_verdict
|
||||
allowed: Final = verdict.get("allowed")
|
||||
if not isinstance(allowed, bool):
|
||||
return self._handle_unavailable(
|
||||
data=data,
|
||||
tool_name=tool_name,
|
||||
reason="the Agent 365 endpoint returned a verdict without a boolean 'allowed' field",
|
||||
)
|
||||
raw_defender: Final = verdict.get("defender")
|
||||
defender: Final = raw_defender if isinstance(raw_defender, dict) else _DefenderResult()
|
||||
raw_correlation_id: Final = verdict.get("correlationId")
|
||||
correlation_id: Final = raw_correlation_id if isinstance(raw_correlation_id, str) else None
|
||||
defender_status: Final = defender.get("status")
|
||||
if allowed and defender_status != DEFENDER_STATUS_EVALUATED:
|
||||
return self._handle_unavailable(
|
||||
data=data,
|
||||
tool_name=tool_name,
|
||||
reason=f"Microsoft Defender did not evaluate the call (defender.status={defender_status or 'missing'})",
|
||||
defender_status=defender_status,
|
||||
correlation_id=correlation_id,
|
||||
latency_ms=latency_ms,
|
||||
)
|
||||
self._record_verdict(
|
||||
data=data,
|
||||
verdict="Allow" if allowed else "Block",
|
||||
guardrail_status="success" if allowed else "guardrail_intervened",
|
||||
defender_status=defender_status,
|
||||
correlation_id=correlation_id,
|
||||
latency_ms=latency_ms,
|
||||
)
|
||||
if not allowed:
|
||||
blocked_detail: Final[_BlockedDetail] = {
|
||||
"error": "Blocked by Microsoft Defender",
|
||||
"message": (
|
||||
defender.get("message")
|
||||
or f"Invocation of '{tool_name}' is blocked by Microsoft Threat Detection policies "
|
||||
"configured by your administrator."
|
||||
),
|
||||
"tool": tool_name,
|
||||
"correlation_id": correlation_id,
|
||||
}
|
||||
raise HTTPException(status_code=400, detail=blocked_detail)
|
||||
return data
|
||||
|
||||
def _build_evaluate_payload(
|
||||
self,
|
||||
data: Mapping[str, object],
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
) -> dict[str, object]: # mutable-ok: JSON body for AsyncHTTPHandler.post, which requires dict
|
||||
tool_name: Final = str(data.get("mcp_tool_name") or "")
|
||||
arguments: Final = data.get("mcp_arguments")
|
||||
server_name: Final = str(data.get("mcp_server_name") or "litellm")
|
||||
agent_id: Final = self.agent_id or user_api_key_dict.key_alias
|
||||
payload: Final[dict[str, object]] = { # mutable-ok: JSON body with optional fields added below
|
||||
"tool": {"name": tool_name},
|
||||
"serverName": server_name,
|
||||
"conversationId": self._resolve_conversation_id(data),
|
||||
}
|
||||
if isinstance(arguments, dict):
|
||||
payload["arguments"] = arguments
|
||||
if agent_id:
|
||||
payload["agentId"] = str(agent_id)
|
||||
return payload
|
||||
|
||||
@staticmethod
|
||||
def _resolve_conversation_id(data: Mapping[str, object]) -> str:
|
||||
"""The MCP session groups every tool call of one client conversation, so it is the conversation id
|
||||
when the transport carries one; stateless calls fall back to the per-call id."""
|
||||
raw_logging_obj: Final = data.get("litellm_logging_obj")
|
||||
logging_obj: Final = raw_logging_obj if isinstance(raw_logging_obj, LiteLLMLoggingObj) else None
|
||||
if logging_obj is not None:
|
||||
tool_call_metadata: Final = logging_obj.model_call_details.get("mcp_tool_call_metadata")
|
||||
session_from_logging: Final = (
|
||||
tool_call_metadata.get("mcp_session_id") if isinstance(tool_call_metadata, Mapping) else None
|
||||
)
|
||||
if isinstance(session_from_logging, str) and session_from_logging:
|
||||
return session_from_logging
|
||||
metadata: Final = next(
|
||||
(m for m in (data.get("metadata"), data.get("litellm_metadata")) if isinstance(m, Mapping)),
|
||||
None,
|
||||
)
|
||||
headers: Final = metadata.get("headers") if isinstance(metadata, Mapping) else None
|
||||
if isinstance(headers, Mapping):
|
||||
session_id: Final = next(
|
||||
(value for name, value in headers.items() if str(name).lower() == MCP_SESSION_ID_HEADER),
|
||||
None,
|
||||
)
|
||||
if isinstance(session_id, str) and session_id:
|
||||
return session_id
|
||||
call_id: Final = data.get("litellm_call_id") or (logging_obj.litellm_call_id if logging_obj else None)
|
||||
if isinstance(call_id, str) and call_id:
|
||||
return call_id
|
||||
return str(uuid.uuid4())
|
||||
|
||||
async def _get_obo_token(self, assertion: str) -> str:
|
||||
cache_key: Final = hashlib.sha256(assertion.encode("utf-8")).hexdigest()
|
||||
now: Final = time.time()
|
||||
with self._obo_cache_lock:
|
||||
cached: Final = self._obo_token_cache.get(cache_key)
|
||||
if cached and cached[1] > now + _TOKEN_EXPIRY_SLACK_SECONDS:
|
||||
self._obo_token_cache.move_to_end(cache_key)
|
||||
return cached[0]
|
||||
|
||||
response: Final = await self._post_allowing_error_status(
|
||||
url=TOKEN_ENDPOINT_TEMPLATE.format(tenant_id=self.tenant_id),
|
||||
data={ # mutable-ok: OAuth form body; AsyncHTTPHandler.post requires dict
|
||||
"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
|
||||
"client_id": self.client_id,
|
||||
"client_secret": self.client_secret,
|
||||
"assertion": assertion,
|
||||
"scope": f"{self.resource_app_id}/{AGENT_365_SCOPE_NAME}",
|
||||
"requested_token_use": "on_behalf_of",
|
||||
},
|
||||
headers={"Content-Type": "application/x-www-form-urlencoded"}, # mutable-ok: httpx header dict
|
||||
)
|
||||
if response.status_code in (408, 429):
|
||||
raise Agent365ThrottledError(status_code=response.status_code)
|
||||
if response.status_code >= 500:
|
||||
raise httpx.HTTPStatusError(
|
||||
f"Entra token endpoint returned {response.status_code}",
|
||||
request=response.request,
|
||||
response=response,
|
||||
)
|
||||
try:
|
||||
parsed_body: Final = response.json()
|
||||
except ValueError as exc:
|
||||
raise Agent365MalformedResponseError("the Entra token endpoint returned a non-JSON body") from exc
|
||||
if not isinstance(parsed_body, dict):
|
||||
raise Agent365MalformedResponseError("the Entra token endpoint returned a non-object JSON body")
|
||||
body: Final = parsed_body
|
||||
if response.status_code >= 400:
|
||||
raise Agent365TokenExchangeError(
|
||||
status_code=response.status_code,
|
||||
error_code=str(body.get("error", "invalid_grant")),
|
||||
description=str(body.get("error_description", ""))[:512],
|
||||
aadsts_codes=_parse_aadsts_codes(body.get("error_codes")),
|
||||
)
|
||||
if "access_token" not in body:
|
||||
raise Agent365MalformedResponseError("the Entra token endpoint returned no access_token")
|
||||
raw_access_token: Final = body.get("access_token")
|
||||
if not isinstance(raw_access_token, str) or not raw_access_token:
|
||||
raise Agent365MalformedResponseError("the Entra token endpoint returned a non-string access_token")
|
||||
access_token: Final = raw_access_token
|
||||
expires_at: Final = time.time() + _parse_expires_in(body.get("expires_in", 3599))
|
||||
with self._obo_cache_lock:
|
||||
self._obo_token_cache[cache_key] = (access_token, expires_at)
|
||||
self._obo_token_cache.move_to_end(cache_key)
|
||||
while len(self._obo_token_cache) > _OBO_CACHE_MAX_ENTRIES:
|
||||
self._obo_token_cache.popitem(last=False)
|
||||
return access_token
|
||||
|
||||
async def _post_allowing_error_status(
|
||||
self,
|
||||
url: str,
|
||||
headers: dict[str, str], # mutable-ok: AsyncHTTPHandler.post requires dict
|
||||
data: dict[str, str] | None = None, # mutable-ok: AsyncHTTPHandler.post requires dict
|
||||
json: dict[str, object] | None = None, # mutable-ok: AsyncHTTPHandler.post requires dict
|
||||
) -> httpx.Response:
|
||||
try:
|
||||
return await self.async_handler.post(
|
||||
url=url,
|
||||
data=data,
|
||||
json=json,
|
||||
headers=headers,
|
||||
timeout=self.request_timeout,
|
||||
)
|
||||
except httpx.HTTPStatusError as exc:
|
||||
return exc.response
|
||||
|
||||
def _handle_caller_fault(
|
||||
self,
|
||||
data: dict, # mutable-ok: guardrail logging appends into the request metadata in place
|
||||
tool_name: str,
|
||||
status_code: int,
|
||||
reason: str,
|
||||
) -> NoReturn:
|
||||
self._record_verdict(
|
||||
data=data,
|
||||
verdict="Rejected",
|
||||
guardrail_status="guardrail_intervened",
|
||||
defender_status=None,
|
||||
correlation_id=None,
|
||||
latency_ms=None,
|
||||
reason=reason,
|
||||
)
|
||||
caller_fault_detail: Final[_UnavailableDetail] = {
|
||||
"error": "Agent 365 guardrail rejected the tool call",
|
||||
"message": f"Tool call '{tool_name}' was blocked because {reason}.",
|
||||
"tool": tool_name,
|
||||
}
|
||||
raise HTTPException(status_code=status_code, detail=caller_fault_detail)
|
||||
|
||||
def _handle_throttled(
|
||||
self,
|
||||
data: dict, # mutable-ok: guardrail logging appends into the request metadata in place
|
||||
tool_name: str,
|
||||
reason: str,
|
||||
latency_ms: float | None,
|
||||
) -> NoReturn:
|
||||
self._record_verdict(
|
||||
data=data,
|
||||
verdict="Throttled",
|
||||
guardrail_status="guardrail_failed_to_respond",
|
||||
defender_status=None,
|
||||
correlation_id=None,
|
||||
latency_ms=latency_ms,
|
||||
reason=reason,
|
||||
)
|
||||
throttled_detail: Final[_UnavailableDetail] = {
|
||||
"error": "Agent 365 guardrail could not authorize the tool call",
|
||||
"message": f"Tool call '{tool_name}' was blocked because {reason}; "
|
||||
"throttled evaluations block regardless of unreachable_fallback.",
|
||||
"tool": tool_name,
|
||||
}
|
||||
raise HTTPException(status_code=503, detail=throttled_detail)
|
||||
|
||||
def _evict_obo_token(self, assertion: str) -> None:
|
||||
cache_key: Final = hashlib.sha256(assertion.encode("utf-8")).hexdigest()
|
||||
with self._obo_cache_lock:
|
||||
self._obo_token_cache.pop(cache_key, None)
|
||||
|
||||
def _handle_unavailable(
|
||||
self,
|
||||
data: dict, # mutable-ok: guardrail logging appends into the request metadata in place
|
||||
tool_name: str,
|
||||
reason: str,
|
||||
defender_status: str | None = None,
|
||||
correlation_id: str | None = None,
|
||||
latency_ms: float | None = None,
|
||||
) -> dict: # mutable-ok: returns the request data dict per hook contract
|
||||
if self.unreachable_fallback == "fail_open":
|
||||
verbose_proxy_logger.warning(
|
||||
"Agent 365 guardrail (%s): %s; unreachable_fallback='fail_open', allowing tool call '%s' unscanned",
|
||||
self.guardrail_name,
|
||||
reason,
|
||||
tool_name,
|
||||
)
|
||||
self._record_verdict(
|
||||
data=data,
|
||||
verdict="Unscanned",
|
||||
guardrail_status="guardrail_failed_to_respond",
|
||||
defender_status=defender_status,
|
||||
correlation_id=correlation_id,
|
||||
latency_ms=latency_ms,
|
||||
reason=reason,
|
||||
)
|
||||
return data
|
||||
self._record_verdict(
|
||||
data=data,
|
||||
verdict="Unavailable",
|
||||
guardrail_status="guardrail_failed_to_respond",
|
||||
defender_status=defender_status,
|
||||
correlation_id=correlation_id,
|
||||
latency_ms=latency_ms,
|
||||
reason=reason,
|
||||
)
|
||||
unavailable_detail: Final[_UnavailableDetail] = {
|
||||
"error": "Agent 365 guardrail could not authorize the tool call",
|
||||
"message": f"Tool call '{tool_name}' was blocked because {reason} and unreachable_fallback is "
|
||||
"'fail_closed'.",
|
||||
"tool": tool_name,
|
||||
}
|
||||
raise HTTPException(status_code=503, detail=unavailable_detail)
|
||||
|
||||
def _record_verdict(
|
||||
self,
|
||||
data: dict[str, object], # mutable-ok: standard guardrail logging appends into the request metadata in place
|
||||
verdict: str,
|
||||
guardrail_status: "GuardrailStatus",
|
||||
defender_status: str | None,
|
||||
correlation_id: str | None,
|
||||
latency_ms: float | None,
|
||||
reason: str | None = None,
|
||||
) -> None:
|
||||
payload: Final[dict[str, object]] = {"verdict": verdict} # mutable-ok: optional fields added below
|
||||
if defender_status:
|
||||
payload["defender_status"] = defender_status
|
||||
if correlation_id:
|
||||
payload["correlation_id"] = correlation_id
|
||||
if latency_ms is not None:
|
||||
payload["latency_ms"] = round(latency_ms, 1)
|
||||
if reason:
|
||||
payload["reason"] = reason
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response=payload,
|
||||
request_data=data,
|
||||
guardrail_status=guardrail_status,
|
||||
duration=(latency_ms / 1000.0) if latency_ms is not None else None,
|
||||
guardrail_provider=self.guardrail_provider,
|
||||
event_type=GuardrailEventHooks.pre_mcp_call,
|
||||
)
|
||||
|
|
@ -58,6 +58,11 @@ if TYPE_CHECKING:
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
|
||||
def _metadata_bucket(request_data: Mapping[str, object], key: str) -> Mapping[str, object]:
|
||||
bucket: Final = request_data.get(key)
|
||||
return bucket if isinstance(bucket, Mapping) else {}
|
||||
|
||||
|
||||
class CustomCodeGuardrailError(Exception):
|
||||
"""Raised when custom code guardrail execution fails."""
|
||||
|
||||
|
|
@ -280,12 +285,16 @@ class CustomCodeGuardrail(CustomGuardrail):
|
|||
Returns:
|
||||
Safe subset of request data
|
||||
"""
|
||||
metadata: Final = {
|
||||
**_metadata_bucket(request_data, "metadata"),
|
||||
**_metadata_bucket(request_data, "litellm_metadata"),
|
||||
}
|
||||
return {
|
||||
"model": request_data.get("model"),
|
||||
"user_id": request_data.get("user_api_key_user_id"),
|
||||
"team_id": request_data.get("user_api_key_team_id"),
|
||||
"end_user_id": request_data.get("user_api_key_end_user_id"),
|
||||
"metadata": request_data.get("metadata", {}),
|
||||
"user_id": metadata.get("user_api_key_user_id"),
|
||||
"team_id": metadata.get("user_api_key_team_id"),
|
||||
"end_user_id": metadata.get("user_api_key_end_user_id"),
|
||||
"metadata": metadata,
|
||||
}
|
||||
|
||||
def _process_result(
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import os
|
|||
import re
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, Coroutine, Mapping, Sequence
|
||||
from dataclasses import dataclass, replace
|
||||
from datetime import datetime
|
||||
from re import Pattern
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, cast
|
||||
|
|
@ -20,7 +21,11 @@ from fastapi import HTTPException
|
|||
|
||||
from litellm import Router
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
from litellm.constants import (
|
||||
CONTENT_FILTER_STREAMING_HOLDBACK_CHARS,
|
||||
CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS,
|
||||
DEFAULT_MAX_RECURSE_DEPTH,
|
||||
)
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -61,6 +66,7 @@ from .patterns import PATTERN_EXTRA_CONFIG, get_compiled_pattern
|
|||
|
||||
MAX_KEYWORD_VALUE_GAP_WORDS: Final = 1
|
||||
GAP_WORD_TOKENIZER: Final = re.compile(r"\b\w+\b")
|
||||
SENTENCE_TERMINATORS: Final = re.compile(r"[.!?]+")
|
||||
|
||||
|
||||
WORD_NUMBER_MAP: Final = {
|
||||
|
|
@ -112,6 +118,22 @@ class _CategoryConfigView(TypedDict):
|
|||
category_file: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _StreamedChoiceState:
|
||||
buffered_text: str = ""
|
||||
yielded_masked_text_len: int = 0
|
||||
committed_detections: tuple[ContentFilterDetection, ...] = ()
|
||||
latest_detections: tuple[ContentFilterDetection, ...] = ()
|
||||
next_trim_len: int = 0
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _StreamedScanPlan:
|
||||
context_chars: int
|
||||
exception_phrases: tuple[str, ...]
|
||||
conditional_words: tuple[str, ...]
|
||||
|
||||
|
||||
class CategoryFileData(TypedDict, total=False):
|
||||
category_name: str
|
||||
description: str
|
||||
|
|
@ -976,7 +998,7 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
|
||||
# Split text into sentences for more precise matching
|
||||
# Simple sentence splitting on common terminators
|
||||
sentences: Final = re.split(r"[.!?]+", text)
|
||||
sentences: Final = SENTENCE_TERMINATORS.split(text)
|
||||
|
||||
for category_name, config in self.conditional_categories.items():
|
||||
identifier_words = config["identifier_words"]
|
||||
|
|
@ -1950,6 +1972,81 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
exception_str=exception_str,
|
||||
)
|
||||
|
||||
def _streamed_scan_plan(self) -> _StreamedScanPlan:
|
||||
"""
|
||||
Per-stream inputs for buffer trimming: the retained tail length (the default
|
||||
context, widened to the longest configured keyword), the category exception
|
||||
phrases, which suppress matches anywhere in the scanned text, and the conditional
|
||||
category words, which only match when paired inside one sentence.
|
||||
"""
|
||||
longest_keyword: Final = max(
|
||||
map(len, (*self.blocked_words, *self.category_keywords, *self.always_block_category_keywords)),
|
||||
default=0,
|
||||
)
|
||||
return _StreamedScanPlan(
|
||||
context_chars=max(CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS, longest_keyword),
|
||||
exception_phrases=tuple(
|
||||
phrase for category in self.loaded_categories.values() for phrase in category.exceptions
|
||||
),
|
||||
conditional_words=tuple(
|
||||
word
|
||||
for config in self.conditional_categories.values()
|
||||
for word in (*config["identifier_words"], *config["block_words"])
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _cut_breaks_wider_context(buffered_text: str, head: str, tail: str, plan: _StreamedScanPlan) -> bool:
|
||||
buffered_lower: Final = buffered_text.lower()
|
||||
tail_lower: Final = tail.lower()
|
||||
if any(phrase in buffered_lower and phrase not in tail_lower for phrase in plan.exception_phrases):
|
||||
return True
|
||||
cut_sentence: Final = (
|
||||
SENTENCE_TERMINATORS.split(head.lower())[-1] + SENTENCE_TERMINATORS.split(tail_lower, maxsplit=1)[0]
|
||||
)
|
||||
return any(word in cut_sentence for word in plan.conditional_words)
|
||||
|
||||
def _trim_streamed_choice_buffer(
|
||||
self, state: _StreamedChoiceState, masked_text: str, plan: _StreamedScanPlan
|
||||
) -> _StreamedChoiceState:
|
||||
"""
|
||||
Bound the per-choice buffer rescanned on every streamed chunk.
|
||||
|
||||
Once the buffer exceeds twice the scan context, drop everything but the last
|
||||
context-sized tail, provided no exception phrase or unfinished conditional sentence
|
||||
would leave the buffer, the two halves mask to the same output as the whole (so no
|
||||
match or phrase straddles the cut), and the dropped prefix has already been yielded.
|
||||
Otherwise keep the buffer and retry once it has grown by another context length.
|
||||
|
||||
Detections found in the dropped prefix move to the state's committed detections.
|
||||
"""
|
||||
if len(state.buffered_text) <= max(2 * plan.context_chars, state.next_trim_len):
|
||||
return state
|
||||
deferred: Final = replace(state, next_trim_len=len(state.buffered_text) + plan.context_chars)
|
||||
head: Final = state.buffered_text[: -plan.context_chars]
|
||||
tail: Final = state.buffered_text[-plan.context_chars :]
|
||||
if self._cut_breaks_wider_context(state.buffered_text, head, tail, plan):
|
||||
return deferred
|
||||
head_detections: Final[list[ContentFilterDetection]] = [] # mutable-ok: filled by _filter_single_text
|
||||
try:
|
||||
masked_head: Final = self._filter_single_text(head, detections=head_detections)
|
||||
masked_tail: Final = self._filter_single_text(tail)
|
||||
except Exception:
|
||||
return deferred
|
||||
if masked_head + masked_tail != masked_text or len(masked_head) > state.yielded_masked_text_len:
|
||||
return deferred
|
||||
return replace(
|
||||
state,
|
||||
buffered_text=tail,
|
||||
yielded_masked_text_len=state.yielded_masked_text_len - len(masked_head),
|
||||
committed_detections=state.committed_detections + tuple(head_detections),
|
||||
next_trim_len=0,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _merge_detections(detections: Sequence[ContentFilterDetection]) -> tuple[ContentFilterDetection, ...]:
|
||||
return tuple(detection for index, detection in enumerate(detections) if detection not in detections[:index])
|
||||
|
||||
async def async_post_call_streaming_iterator_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -1968,10 +2065,8 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
and the UI Request Lifecycle panel. Mirrors apply_guardrail's finally-block
|
||||
contract.
|
||||
"""
|
||||
accumulated_text_by_choice: Final[dict[int, str]] = {}
|
||||
yielded_masked_text_len_by_choice: Final[dict[int, int]] = {}
|
||||
latest_detections_by_choice: Final[dict[int, list[ContentFilterDetection]]] = {}
|
||||
buffer_size: Final = 50 # Increased buffer to catch patterns split across many chunks
|
||||
state_by_choice: Final[dict[int, _StreamedChoiceState]] = {}
|
||||
plan: Final = self._streamed_scan_plan()
|
||||
|
||||
start_time: Final = datetime.now()
|
||||
scan_seconds: float = 0.0 # rebind-ok: accumulates per-chunk scan time across the stream
|
||||
|
|
@ -1997,69 +2092,60 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
|
||||
content = getattr(choice.delta, "content", None)
|
||||
is_final = bool(getattr(choice, "finish_reason", None))
|
||||
if isinstance(content, str) and content:
|
||||
accumulated_text_by_choice[choice_index] = (
|
||||
accumulated_text_by_choice.get(choice_index, "") + content
|
||||
)
|
||||
elif not is_final:
|
||||
new_content = content if isinstance(content, str) else ""
|
||||
if not new_content and not is_final:
|
||||
continue
|
||||
|
||||
text_to_check = accumulated_text_by_choice.get(choice_index, "")
|
||||
if not text_to_check:
|
||||
previous_state = state_by_choice.get(choice_index, _StreamedChoiceState())
|
||||
buffered_text = previous_state.buffered_text + new_content
|
||||
if not buffered_text:
|
||||
continue
|
||||
|
||||
# Add a space at the end if it's the final chunk to trigger word boundaries (\b)
|
||||
text_to_scan = text_to_check + (" " if is_final else "")
|
||||
text_to_scan = buffered_text + (" " if is_final else "")
|
||||
choice_detections: list[ContentFilterDetection] = []
|
||||
scan_started = time.perf_counter()
|
||||
|
||||
try:
|
||||
# _filter_single_text scans the whole accumulated
|
||||
# choice buffer every chunk, so previous-chunk
|
||||
# matches are guaranteed to be re-found. Keeping
|
||||
# only each choice's latest scan avoids duplicate
|
||||
# detections in the final log row.
|
||||
masked_text = self._filter_single_text(text_to_scan, detections=choice_detections)
|
||||
if is_final and masked_text.endswith(" "):
|
||||
masked_text = masked_text[:-1]
|
||||
latest_detections_by_choice[choice_index] = choice_detections
|
||||
latest_detections = tuple(choice_detections)
|
||||
except HTTPException:
|
||||
latest_detections_by_choice[choice_index] = choice_detections
|
||||
state_by_choice[choice_index] = replace(
|
||||
previous_state, latest_detections=tuple(choice_detections)
|
||||
)
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("ContentFilterGuardrail: Error in masking: %s", e)
|
||||
masked_text = text_to_scan # Fallback to current text
|
||||
latest_detections = previous_state.latest_detections
|
||||
finally:
|
||||
scan_seconds += time.perf_counter() - scan_started
|
||||
|
||||
# Determine how much can be safely yielded
|
||||
safe_to_yield_len = max(
|
||||
previous_state.yielded_masked_text_len,
|
||||
len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS),
|
||||
)
|
||||
choice.delta.content = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len]
|
||||
next_state = replace(
|
||||
previous_state,
|
||||
buffered_text=buffered_text,
|
||||
yielded_masked_text_len=safe_to_yield_len,
|
||||
latest_detections=latest_detections,
|
||||
)
|
||||
if is_final:
|
||||
safe_to_yield_len = len(masked_text)
|
||||
else:
|
||||
safe_to_yield_len = max(0, len(masked_text) - buffer_size)
|
||||
state_by_choice[choice_index] = next_state
|
||||
continue
|
||||
|
||||
yielded_masked_text_len = yielded_masked_text_len_by_choice.get(choice_index, 0)
|
||||
if safe_to_yield_len > yielded_masked_text_len:
|
||||
new_masked_content = masked_text[yielded_masked_text_len:safe_to_yield_len]
|
||||
choice.delta.content = new_masked_content
|
||||
yielded_masked_text_len_by_choice[choice_index] = safe_to_yield_len
|
||||
else:
|
||||
# Hold content by yielding empty content on this choice
|
||||
# while preserving chunk metadata and other choices.
|
||||
choice.delta.content = ""
|
||||
trim_started = time.perf_counter()
|
||||
state_by_choice[choice_index] = self._trim_streamed_choice_buffer(next_state, masked_text, plan)
|
||||
scan_seconds += time.perf_counter() - trim_started
|
||||
|
||||
yield item
|
||||
else:
|
||||
# Not a ModelResponseStream or no choices - yield as is
|
||||
yield item
|
||||
|
||||
# Any remaining content (should have been handled by is_final, but just in case)
|
||||
if any(
|
||||
yielded_masked_text_len_by_choice.get(choice_index, 0) < len(accumulated_text)
|
||||
for choice_index, accumulated_text in accumulated_text_by_choice.items()
|
||||
):
|
||||
# We already reached the end of the generator
|
||||
pass
|
||||
except HTTPException:
|
||||
status = "guardrail_intervened"
|
||||
raise
|
||||
|
|
@ -2070,8 +2156,8 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
finally:
|
||||
detections = [
|
||||
detection
|
||||
for choice_detections in latest_detections_by_choice.values()
|
||||
for detection in choice_detections
|
||||
for state in state_by_choice.values()
|
||||
for detection in self._merge_detections((*state.committed_detections, *state.latest_detections))
|
||||
]
|
||||
self._count_masked_entities(detections, masked_entity_count)
|
||||
self._log_guardrail_information(
|
||||
|
|
|
|||
|
|
@ -1,14 +1,17 @@
|
|||
"""LLM-as-a-Judge guardrail: uses an LLM to score responses against weighted criteria."""
|
||||
"""LLM-as-a-Judge guardrail: uses an LLM to score requests or responses against weighted criteria."""
|
||||
|
||||
from collections.abc import Callable, Sequence
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeVar
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.llm_judge import (
|
||||
default_router_provider,
|
||||
|
|
@ -16,8 +19,9 @@ from litellm.litellm_core_utils.llm_judge import (
|
|||
judge_acompletion,
|
||||
parse_json_verdict,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import get_last_user_message
|
||||
from litellm.types.guardrails import GuardrailEventHooks, Mode, SupportedGuardrailIntegrations
|
||||
from litellm.types.utils import LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN, GenericGuardrailAPIInputs, GuardrailStatus
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm import Router
|
||||
|
|
@ -26,18 +30,65 @@ if TYPE_CHECKING:
|
|||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import StandardLoggingEvalInformation
|
||||
|
||||
JUDGE_SYSTEM_PROMPT = """You are a quality judge. Evaluate the assistant's response against the criteria provided.
|
||||
For each criterion, assign a score from 0 to 100 and provide concise reasoning.
|
||||
JudgeInputType = Literal["request", "response"]
|
||||
JudgeEventHook = GuardrailEventHooks | list[GuardrailEventHooks] | Mode
|
||||
JudgeModeParam = str | list[str] | Mode | GuardrailEventHooks | list[GuardrailEventHooks] | None
|
||||
|
||||
_JUDGE_SYSTEM_PROMPT_TEMPLATE: Final = """You are a quality judge. Evaluate the {subject} against the criteria provided.
|
||||
{focus}For each criterion, assign a score from 0 to 100 and provide concise reasoning.
|
||||
Return ONLY valid JSON in this exact format:
|
||||
{
|
||||
{{
|
||||
"verdicts": [
|
||||
{"criterion_name": "<name>", "score": <0-100>, "reasoning": "<one sentence>", "passed": <true|false>, "weight": <weight>}
|
||||
{{"criterion_name": "<name>", "score": <0-100>, "reasoning": "<one sentence>", "passed": <true|false>, "weight": <weight>}}
|
||||
],
|
||||
"overall_score": <weighted average 0-100>
|
||||
}"""
|
||||
}}"""
|
||||
|
||||
JUDGE_SYSTEM_PROMPTS: Final[MappingProxyType[JudgeInputType, str]] = MappingProxyType(
|
||||
{
|
||||
"request": _JUDGE_SYSTEM_PROMPT_TEMPLATE.format(
|
||||
subject="request",
|
||||
focus="Judge the most recent user turn; treat earlier turns in the conversation only as context.\n",
|
||||
),
|
||||
"response": _JUDGE_SYSTEM_PROMPT_TEMPLATE.format(subject="assistant's response", focus=""),
|
||||
}
|
||||
)
|
||||
|
||||
_JUDGE_SUBJECT_LABELS: Final[MappingProxyType[JudgeInputType, str]] = MappingProxyType(
|
||||
{"request": "Latest request turn to evaluate", "response": "Assistant response to evaluate"}
|
||||
)
|
||||
|
||||
_LIFECYCLE_HOOKS: Final[MappingProxyType[JudgeInputType, tuple[GuardrailEventHooks, ...]]] = MappingProxyType(
|
||||
{
|
||||
"request": (GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.logging_only),
|
||||
"response": (GuardrailEventHooks.post_call, GuardrailEventHooks.logging_only),
|
||||
}
|
||||
)
|
||||
|
||||
_VALID_ON_FAILURE: Final = frozenset({"block", "log"})
|
||||
|
||||
_JUDGE_CALL_METADATA: Final = MappingProxyType(
|
||||
{INTERNAL_CALL_ORIGIN_METADATA_KEY: LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN}
|
||||
)
|
||||
|
||||
|
||||
class _LoggedCallParams(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
metadata: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
def _is_logged_judge_call(data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool:
|
||||
"""logging_only is the only event whose ``data`` is the SDK's model_call_details rather than the client body."""
|
||||
if event_type is not GuardrailEventHooks.logging_only:
|
||||
return False
|
||||
try:
|
||||
params: Final = _LoggedCallParams.model_validate(data.get("litellm_params") or {})
|
||||
except ValidationError:
|
||||
return False
|
||||
return (params.metadata or {}).get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN
|
||||
|
||||
|
||||
_default_router_provider: Final = default_router_provider
|
||||
_parse_judge_verdict: Final = parse_json_verdict
|
||||
_extract_text_from_content: Final = extract_text_from_content
|
||||
|
|
@ -86,10 +137,29 @@ def _get_litellm_param(
|
|||
return default
|
||||
|
||||
|
||||
def _coerce_event_hook(mode: JudgeModeParam) -> JudgeEventHook:
|
||||
if mode is None:
|
||||
return GuardrailEventHooks.post_call
|
||||
if isinstance(mode, Mode):
|
||||
return mode
|
||||
if isinstance(mode, list):
|
||||
return [GuardrailEventHooks(hook) for hook in mode]
|
||||
return GuardrailEventHooks(mode)
|
||||
|
||||
|
||||
def _text_under_review(inputs: GenericGuardrailAPIInputs, input_type: JudgeInputType) -> str:
|
||||
all_text: Final = "\n".join(inputs.get("texts") or [])
|
||||
if input_type == "response":
|
||||
return all_text
|
||||
latest_user_turn: Final = get_last_user_message(inputs.get("structured_messages") or [])
|
||||
return latest_user_turn if latest_user_turn is not None else all_text
|
||||
|
||||
|
||||
def _build_judge_prompt(
|
||||
criteria: Sequence[JudgeCriterion],
|
||||
messages: Sequence[JudgeMessage],
|
||||
response_text: str,
|
||||
text_under_review: str,
|
||||
input_type: JudgeInputType = "response",
|
||||
) -> str:
|
||||
criteria_block: Final = "\n".join(
|
||||
f"- {c.get('name', '')} (weight {c.get('weight', 0)}%): {c.get('description', '')}" for c in criteria
|
||||
|
|
@ -99,15 +169,16 @@ def _build_judge_prompt(
|
|||
for m in messages
|
||||
if m.get("content") is not None
|
||||
)
|
||||
conversation_block: Final = f"Conversation:\n{conversation}\n\n" if conversation or input_type == "response" else ""
|
||||
return (
|
||||
f"Criteria to evaluate:\n{criteria_block}\n\n"
|
||||
f"Conversation:\n{conversation}\n\n"
|
||||
f"Assistant response to evaluate:\n{response_text}"
|
||||
f"{conversation_block}"
|
||||
f"{_JUDGE_SUBJECT_LABELS[input_type]}:\n{text_under_review}"
|
||||
)
|
||||
|
||||
|
||||
class LLMAsAJudgeGuardrail(CustomGuardrail):
|
||||
"""Post-call guardrail that judges response quality via an LLM."""
|
||||
"""Guardrail that judges request (pre_call/during_call) or response (post_call) quality via an LLM."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -116,22 +187,15 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
criteria: Sequence[JudgeCriterion],
|
||||
overall_threshold: float = 80.0,
|
||||
on_failure: Literal["block", "log"] = "block",
|
||||
event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None,
|
||||
event_hook: JudgeModeParam = None,
|
||||
default_on: bool = False,
|
||||
router_provider: "Callable[[], Router | None] | None" = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
_event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None
|
||||
if event_hook is not None:
|
||||
if isinstance(event_hook, list):
|
||||
_event_hook = [GuardrailEventHooks(h) if isinstance(h, str) else h for h in event_hook]
|
||||
else:
|
||||
_event_hook = GuardrailEventHooks(event_hook) if isinstance(event_hook, str) else event_hook
|
||||
|
||||
super().__init__(
|
||||
guardrail_name=guardrail_name,
|
||||
supported_event_hooks=list(self.get_supported_event_hooks()),
|
||||
event_hook=_event_hook or GuardrailEventHooks.post_call,
|
||||
event_hook=_coerce_event_hook(event_hook),
|
||||
default_on=default_on,
|
||||
**kwargs,
|
||||
)
|
||||
|
|
@ -143,18 +207,24 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
|
||||
@classmethod
|
||||
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
|
||||
return [GuardrailEventHooks.post_call]
|
||||
return [GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.post_call]
|
||||
|
||||
def should_run_guardrail(self, data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool:
|
||||
if _is_logged_judge_call(data, event_type):
|
||||
return False
|
||||
return super().should_run_guardrail(data, event_type)
|
||||
|
||||
async def _run_judge(
|
||||
self,
|
||||
messages: Sequence[JudgeMessage],
|
||||
response_text: str,
|
||||
text_under_review: str,
|
||||
input_type: JudgeInputType = "response",
|
||||
) -> dict[str, object]:
|
||||
judge_messages: Final[list[AllMessageValues]] = [
|
||||
{"role": "system", "content": JUDGE_SYSTEM_PROMPT},
|
||||
{"role": "system", "content": JUDGE_SYSTEM_PROMPTS[input_type]},
|
||||
{
|
||||
"role": "user",
|
||||
"content": _build_judge_prompt(self.criteria, messages, response_text),
|
||||
"content": _build_judge_prompt(self.criteria, messages, text_under_review, input_type),
|
||||
},
|
||||
]
|
||||
response: Final = await judge_acompletion(
|
||||
|
|
@ -163,6 +233,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
judge_messages,
|
||||
response_format={"type": "json_object"},
|
||||
temperature=0,
|
||||
metadata=dict(_JUDGE_CALL_METADATA),
|
||||
)
|
||||
raw: Final = response.choices[0].message.content or "{}"
|
||||
return _parse_judge_verdict(raw)
|
||||
|
|
@ -174,13 +245,8 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
# Only evaluate post-call (response text). Fail open on pre-call.
|
||||
if input_type != "response":
|
||||
return inputs
|
||||
|
||||
texts: Final = inputs.get("texts") or []
|
||||
response_text: Final = " ".join(texts)
|
||||
if not response_text:
|
||||
text_under_review: Final = _text_under_review(inputs, input_type)
|
||||
if not text_under_review:
|
||||
return inputs
|
||||
|
||||
start_time: Final = datetime.now()
|
||||
|
|
@ -188,10 +254,12 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
judge_result: dict[str, object] = {}
|
||||
|
||||
try:
|
||||
messages: Final[Sequence[JudgeMessage]] = request_data.get("messages") or []
|
||||
messages: Final[Sequence[JudgeMessage]] = (
|
||||
inputs.get("structured_messages") or request_data.get("messages") or []
|
||||
)
|
||||
|
||||
try:
|
||||
judge_result = await self._run_judge(messages, response_text)
|
||||
judge_result = await self._run_judge(messages, text_under_review, input_type)
|
||||
except Exception as judge_err:
|
||||
verbose_logger.warning(
|
||||
"llm_as_a_judge guardrail: judge call failed, failing open. Error: %s", judge_err
|
||||
|
|
@ -230,7 +298,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
raise HTTPException(
|
||||
status_code=422,
|
||||
detail={
|
||||
"error": "LLM judge rejected response: score below threshold",
|
||||
"error": f"LLM judge rejected {input_type}: score below threshold",
|
||||
"overall_score": overall_score,
|
||||
"threshold": self.overall_threshold,
|
||||
"verdicts": judge_result.get("verdicts", []),
|
||||
|
|
@ -252,9 +320,13 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
guardrail_status=status,
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=datetime.now().timestamp(),
|
||||
event_type=GuardrailEventHooks.post_call,
|
||||
event_type=self._event_type_for(input_type),
|
||||
)
|
||||
|
||||
def _event_type_for(self, input_type: JudgeInputType) -> GuardrailEventHooks | None:
|
||||
configured: Final = tuple(hook for hook in _LIFECYCLE_HOOKS[input_type] if self._event_hook_is_event_type(hook))
|
||||
return configured[0] if len(configured) == 1 else None
|
||||
|
||||
|
||||
def initialize_guardrail(
|
||||
litellm_params: "LitellmParams",
|
||||
|
|
@ -282,10 +354,7 @@ def initialize_guardrail(
|
|||
|
||||
overall_threshold: Final = float(_get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0))
|
||||
|
||||
mode: Final[str | None] = _get_litellm_param(litellm_params, guardrail, "mode", None)
|
||||
event_hook: GuardrailEventHooks | None = None
|
||||
if isinstance(mode, str) and mode in {e.value for e in GuardrailEventHooks}:
|
||||
event_hook = GuardrailEventHooks(mode)
|
||||
mode: Final[JudgeModeParam] = _get_litellm_param(litellm_params, guardrail, "mode", None)
|
||||
|
||||
instance: Final = LLMAsAJudgeGuardrail(
|
||||
guardrail_name=guardrail_name,
|
||||
|
|
@ -293,7 +362,7 @@ def initialize_guardrail(
|
|||
criteria=criteria,
|
||||
overall_threshold=overall_threshold,
|
||||
on_failure=on_failure,
|
||||
event_hook=event_hook,
|
||||
event_hook=mode,
|
||||
default_on=bool(_get_litellm_param(litellm_params, guardrail, "default_on", False)),
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback(instance)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,7 @@
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
|
@ -19,20 +22,26 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import (
|
||||
GuardrailConfigModel,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.singulr import (
|
||||
AssistantMessage,
|
||||
SingulrGuardrailPayload,
|
||||
SingulrGuardrailRequest,
|
||||
SingulrGuardrailResponse,
|
||||
SingulrMcpGuardrailPayload,
|
||||
ToolCall,
|
||||
ToolCallFunction,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
|
||||
|
||||
_DEFAULT_API_BASE: Final = "http://localhost:8003"
|
||||
_GUARD_ENDPOINT: Final = "/api/v1/ai-gateway/litellm"
|
||||
_GUARD_ENDPOINT: Final = "/api/v1/ai-gateway/litellm-v2"
|
||||
_DEFAULT_TIMEOUT: Final = 30.0
|
||||
_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
|
||||
_MCP_MODEL_PREFIX: Final = "MCP:"
|
||||
|
||||
|
||||
class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
|
||||
|
|
@ -51,8 +60,8 @@ class SingulrGuardrail(CustomGuardrail):
|
|||
**kwargs: Unpack[_CustomGuardrailOptions],
|
||||
) -> None:
|
||||
self.singulr_api_key = singulr_api_key or os.environ.get("SINGULR_API_KEY")
|
||||
self.singulr_api_base = (singulr_api_base or os.environ.get("SINGULR_API_BASE") or _DEFAULT_API_BASE).rstrip(
|
||||
"/"
|
||||
self.singulr_api_base = (
|
||||
(singulr_api_base or os.environ.get("SINGULR_API_BASE") or _DEFAULT_API_BASE).strip().rstrip("/")
|
||||
)
|
||||
parsed: Final = urlparse(self.singulr_api_base)
|
||||
if parsed.scheme == "http" and parsed.hostname not in (
|
||||
|
|
@ -85,6 +94,9 @@ class SingulrGuardrail(CustomGuardrail):
|
|||
kwargs["supported_event_hooks"] = [
|
||||
GuardrailEventHooks.pre_call,
|
||||
GuardrailEventHooks.post_call,
|
||||
GuardrailEventHooks.logging_only,
|
||||
GuardrailEventHooks.pre_mcp_call,
|
||||
GuardrailEventHooks.post_mcp_call,
|
||||
]
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
|
@ -97,52 +109,70 @@ class SingulrGuardrail(CustomGuardrail):
|
|||
|
||||
return SingulrGuardrailConfigModel
|
||||
|
||||
def _build_payload(
|
||||
self,
|
||||
request_data: dict[str, Any],
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
input_type: str,
|
||||
) -> dict[str, object]:
|
||||
if not request_data:
|
||||
texts: Final = inputs.get("texts", [])
|
||||
|
||||
payload = SingulrGuardrailPayload(
|
||||
input_type=input_type,
|
||||
is_playground_request=True,
|
||||
playground_text=texts[0] if texts else None,
|
||||
@staticmethod
|
||||
def _metadata_containers(request_data: Mapping[str, Any]) -> tuple[Mapping[str, Any], ...]:
|
||||
litellm_params: Final = request_data.get("litellm_params") or _EMPTY_MAPPING
|
||||
return tuple(
|
||||
container
|
||||
for container in (
|
||||
request_data.get("litellm_metadata"),
|
||||
request_data.get("metadata"),
|
||||
litellm_params.get("litellm_metadata") if litellm_params else None,
|
||||
litellm_params.get("metadata") if litellm_params else None,
|
||||
)
|
||||
else:
|
||||
response: Final = request_data.get("response")
|
||||
singulr_req_object: Final = SingulrGuardrailRequest(
|
||||
model=request_data.get("model"),
|
||||
messages=request_data.get("messages"),
|
||||
tools=request_data.get("tools"),
|
||||
model_response=response.model_dump(mode="json") if input_type == "response" and response else None,
|
||||
litellm_metadata=request_data.get("litellm_metadata"),
|
||||
)
|
||||
payload = SingulrGuardrailPayload(
|
||||
litellm_call_id=request_data.get("litellm_call_id"),
|
||||
request_data=singulr_req_object,
|
||||
input_type=input_type,
|
||||
)
|
||||
|
||||
return payload.model_dump(mode="json")
|
||||
|
||||
def _build_headers(self) -> dict[str, str]:
|
||||
return dict(
|
||||
(header, value)
|
||||
for header, value in (
|
||||
("Content-Type", "application/json"),
|
||||
("X-Singulr-Gateway-Token", self.singulr_api_key),
|
||||
(
|
||||
"X-Singulr-Enforcement-Entity-Id",
|
||||
self.singulr_application_id or "",
|
||||
),
|
||||
("X-Singulr-Guardrail-Id", self.singulr_guardrail_id or ""),
|
||||
)
|
||||
if value
|
||||
if container
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _resolve_metadata_value(cls, request_data: Mapping[str, Any], key: str) -> str | None:
|
||||
for container in cls._metadata_containers(request_data=request_data):
|
||||
value = container.get(key)
|
||||
if value:
|
||||
return value
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _resolve_user_role_from_request_data(cls, request_data: Mapping[str, Any]) -> str | None:
|
||||
for container in cls._metadata_containers(request_data=request_data):
|
||||
auth = container.get("user_api_key_auth")
|
||||
if isinstance(auth, UserAPIKeyAuth) and auth.user_role:
|
||||
return auth.user_role.value
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _build_metadata(cls, request_data: Mapping[str, Any]) -> Mapping[str, str] | None:
|
||||
fields: Final = (
|
||||
"user_api_key_alias",
|
||||
"user_api_key_user_id",
|
||||
"user_api_key_user_email",
|
||||
"user_api_key_org_id",
|
||||
"user_api_key_org_alias",
|
||||
"user_api_key_team_id",
|
||||
"user_api_key_team_alias",
|
||||
)
|
||||
resolved: Final = (
|
||||
*((field, cls._resolve_metadata_value(request_data=request_data, key=field)) for field in fields),
|
||||
("user_api_key_user_role", cls._resolve_user_role_from_request_data(request_data=request_data)),
|
||||
)
|
||||
if not any(value for _, value in resolved):
|
||||
return None
|
||||
return {key: value for key, value in resolved if value} # mutable-ok: short-lived JSON payload dict
|
||||
|
||||
@staticmethod
|
||||
def _build_user_message(text: str) -> Mapping[str, Any]:
|
||||
return {"role": "user", "content": text} # mutable-ok: short-lived JSON payload dict
|
||||
|
||||
def _build_headers(self) -> Mapping[str, str]:
|
||||
all_headers: Final = MappingProxyType(
|
||||
{
|
||||
"Content-Type": "application/json",
|
||||
"X-Singulr-Gateway-Token": self.singulr_api_key,
|
||||
"X-Singulr-Enforcement-Entity-Id": self.singulr_application_id,
|
||||
"X-Singulr-Guardrail-Id": self.singulr_guardrail_id,
|
||||
}
|
||||
)
|
||||
return MappingProxyType({header: value for header, value in all_headers.items() if value})
|
||||
|
||||
async def _call_api(self, payload: dict[str, object]) -> SingulrGuardrailResponse | None:
|
||||
endpoint: Final = f"{self.singulr_api_base}{_GUARD_ENDPOINT}"
|
||||
verbose_proxy_logger.debug("Singulr: %s", endpoint)
|
||||
|
|
@ -168,7 +198,7 @@ class SingulrGuardrail(CustomGuardrail):
|
|||
if self.block_on_error:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=(f"Singulr API returned HTTP {exc.response.status_code}: {exc.response.text}"),
|
||||
message=f"Singulr API returned HTTP {exc.response.status_code}: {exc.response.text}",
|
||||
) from exc
|
||||
return None
|
||||
|
||||
|
|
@ -190,33 +220,218 @@ class SingulrGuardrail(CustomGuardrail):
|
|||
) from exc
|
||||
return None
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
async def _apply_guardrail_on_request(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: str,
|
||||
logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
texts: Sequence[str],
|
||||
structured_messages: Sequence[Any],
|
||||
request_data: Mapping[str, Any],
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
payload: Final = self._build_payload(request_data, inputs, input_type)
|
||||
if not payload:
|
||||
return inputs
|
||||
|
||||
result: Final = await self._call_api(payload)
|
||||
if result is None:
|
||||
return inputs
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Singulr: should_block=%s blocking_due_to=%s",
|
||||
result.should_block,
|
||||
result.blocking_due_to,
|
||||
messages: Final = (
|
||||
tuple(structured_messages)
|
||||
if structured_messages
|
||||
else tuple(self._build_user_message(text) for text in texts)
|
||||
)
|
||||
|
||||
if result.should_block:
|
||||
images: Final = inputs.get("images")
|
||||
tools: Final = inputs.get("tools")
|
||||
|
||||
if not messages and not images and not tools:
|
||||
verbose_proxy_logger.debug("Singulr: No messages, images, or tools to check after filtering")
|
||||
return inputs
|
||||
|
||||
metadata: Final = self._build_metadata(request_data=request_data)
|
||||
|
||||
singulr_req_obj = SingulrGuardrailPayload(
|
||||
correlation_id=request_data.get("litellm_call_id"),
|
||||
model_name=inputs.get("model"),
|
||||
guardrail_scope="request",
|
||||
messages=messages,
|
||||
images=images,
|
||||
tools=tools,
|
||||
metadata=metadata,
|
||||
)
|
||||
payload = singulr_req_obj.model_dump(mode="json")
|
||||
guardrail_resp = await self._call_api(payload)
|
||||
|
||||
if guardrail_resp is None:
|
||||
return inputs
|
||||
|
||||
if guardrail_resp.should_block:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=f"Blocked by Singulr: {result.blocking_due_to or 'unknown'}",
|
||||
status_code=400,
|
||||
message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}",
|
||||
blocked_content=True,
|
||||
)
|
||||
return inputs
|
||||
|
||||
@staticmethod
|
||||
def _mcp_tool_name(request_data: Mapping[str, Any]) -> str | None:
|
||||
return request_data.get("mcp_tool_name") or request_data.get("name")
|
||||
|
||||
@staticmethod
|
||||
def _mcp_arguments(request_data: Mapping[str, Any]) -> object:
|
||||
arguments: Final = request_data.get("mcp_arguments")
|
||||
return arguments if arguments is not None else request_data.get("arguments")
|
||||
|
||||
@staticmethod
|
||||
def _is_mcp_call(request_data: Mapping[str, Any], logging_obj: LiteLLMLoggingObj | None) -> bool:
|
||||
call_type: Final = logging_obj.call_type if logging_obj is not None else request_data.get("call_type")
|
||||
if call_type is not None:
|
||||
return call_type == CallTypes.call_mcp_tool.value
|
||||
model: Final = request_data.get("model")
|
||||
return "mcp_tool_name" in request_data or (isinstance(model, str) and model.startswith(_MCP_MODEL_PREFIX))
|
||||
|
||||
async def _apply_guardrail_on_mcp_request(self, request_data: Mapping[str, Any]) -> None:
|
||||
metadata: Final = self._build_metadata(request_data=request_data)
|
||||
|
||||
singulr_mcp_obj = SingulrMcpGuardrailPayload(
|
||||
guardrail_scope="mcp_request",
|
||||
tool_name=self._mcp_tool_name(request_data),
|
||||
tool_arguments=self._mcp_arguments(request_data),
|
||||
mcp_server_name=request_data.get("mcp_server_name"),
|
||||
metadata=metadata,
|
||||
)
|
||||
payload = singulr_mcp_obj.model_dump(mode="json")
|
||||
guardrail_resp = await self._call_api(payload)
|
||||
|
||||
if guardrail_resp is None:
|
||||
return
|
||||
|
||||
if guardrail_resp.should_block:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
status_code=400,
|
||||
message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}",
|
||||
blocked_content=True,
|
||||
)
|
||||
|
||||
async def _apply_guardrail_on_mcp_response(
|
||||
self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, Any]
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
if not texts:
|
||||
return inputs
|
||||
|
||||
metadata: Final = self._build_metadata(request_data=request_data)
|
||||
|
||||
singulr_mcp_obj = SingulrMcpGuardrailPayload(
|
||||
model_name=request_data.get("model"),
|
||||
guardrail_scope="mcp_response",
|
||||
tool_result=texts,
|
||||
metadata=metadata,
|
||||
)
|
||||
payload = singulr_mcp_obj.model_dump(mode="json")
|
||||
guardrail_resp = await self._call_api(payload)
|
||||
|
||||
if guardrail_resp is None:
|
||||
return inputs
|
||||
|
||||
if guardrail_resp.should_block:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
status_code=400,
|
||||
message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}",
|
||||
blocked_content=True,
|
||||
)
|
||||
|
||||
return inputs
|
||||
|
||||
@staticmethod
|
||||
def _build_tool_call(tool_call: Mapping[str, Any]) -> "ToolCall | None":
|
||||
tool_call_id: Final = tool_call.get("id")
|
||||
fun: Final = tool_call.get("function")
|
||||
if not tool_call_id or not fun:
|
||||
return None
|
||||
func_name: Final = fun.get("name")
|
||||
args: Final = fun.get("arguments")
|
||||
if not func_name or args is None:
|
||||
return None
|
||||
call_type: Final = tool_call.get("type")
|
||||
return ToolCall(
|
||||
id=tool_call_id,
|
||||
type=call_type if isinstance(call_type, str) and call_type else "function",
|
||||
function=ToolCallFunction(
|
||||
name=func_name,
|
||||
arguments=args if isinstance(args, str) else json.dumps(args, default=str),
|
||||
),
|
||||
)
|
||||
|
||||
async def _apply_guardrail_on_response(
|
||||
self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, Any]
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
combined_texts: Final = "\n".join(texts) if texts else None
|
||||
|
||||
tool_calls: Final = inputs.get("tool_calls", ())
|
||||
tool_calls_res: Final = tuple(
|
||||
tool_call_res
|
||||
for tool_call_res in (self._build_tool_call(tool_call) for tool_call in tool_calls)
|
||||
if tool_call_res is not None
|
||||
)
|
||||
|
||||
assistant_message: Final = AssistantMessage(
|
||||
role="assistant",
|
||||
content=combined_texts,
|
||||
tool_calls=tool_calls_res,
|
||||
)
|
||||
|
||||
metadata: Final = self._build_metadata(request_data=request_data)
|
||||
|
||||
singulr_resp_obj = SingulrGuardrailPayload(
|
||||
correlation_id=request_data.get("litellm_call_id"),
|
||||
guardrail_scope="response",
|
||||
model_name=request_data.get("model"),
|
||||
messages=request_data.get("messages"),
|
||||
images=inputs.get("images"),
|
||||
response=assistant_message,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
payload = singulr_resp_obj.model_dump(mode="json")
|
||||
guardrail_resp = await self._call_api(payload)
|
||||
|
||||
if guardrail_resp is None:
|
||||
return inputs
|
||||
|
||||
if guardrail_resp.should_block:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
status_code=400,
|
||||
message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}",
|
||||
blocked_content=True,
|
||||
)
|
||||
return inputs
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict, # mutable-ok: required by CustomGuardrail.apply_guardrail override signature
|
||||
input_type: str,
|
||||
logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
texts: Final = inputs.get("texts", ())
|
||||
structured_messages: Final = inputs.get("structured_messages", ())
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Singulr Guardrail: apply_guardrail called with input_type=%s, texts=%d, structured_messages=%d",
|
||||
input_type,
|
||||
len(texts),
|
||||
len(structured_messages),
|
||||
)
|
||||
|
||||
is_mcp_call: Final = self._is_mcp_call(request_data, logging_obj)
|
||||
if input_type == "request":
|
||||
if is_mcp_call:
|
||||
await self._apply_guardrail_on_mcp_request(request_data=request_data)
|
||||
return inputs
|
||||
return await self._apply_guardrail_on_request(
|
||||
inputs=inputs, texts=texts, structured_messages=structured_messages, request_data=request_data
|
||||
)
|
||||
elif input_type == "response":
|
||||
if is_mcp_call:
|
||||
return await self._apply_guardrail_on_mcp_response(
|
||||
inputs=inputs, texts=texts, request_data=request_data
|
||||
)
|
||||
return await self._apply_guardrail_on_response(inputs=inputs, texts=texts, request_data=request_data)
|
||||
return inputs
|
||||
|
|
|
|||
|
|
@ -41,6 +41,7 @@ from litellm.proxy._types import UserAPIKeyAuth
|
|||
from litellm.proxy.auth.auth_utils import (
|
||||
ESTIMATED_OUTPUT_TOKENS_FIELD,
|
||||
get_estimated_output_tokens,
|
||||
get_key_own_model_rate_limit,
|
||||
get_key_tag_rpm_limit,
|
||||
get_model_rate_limit_from_metadata,
|
||||
)
|
||||
|
|
@ -2892,41 +2893,67 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
return batch_limiter
|
||||
return None
|
||||
|
||||
def _key_owns_model_limit(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
requested_model: str,
|
||||
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
|
||||
) -> bool:
|
||||
key_own_limits: Final = get_key_own_model_rate_limit(user_api_key_dict, rate_limit_key)
|
||||
return key_own_limits is not None and key_own_limits.get(requested_model) is not None
|
||||
|
||||
def _inherited_team_model_limit(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
requested_model: str,
|
||||
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
|
||||
) -> int | None:
|
||||
team_limits: Final = get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", rate_limit_key)
|
||||
team_limit: Final = team_limits.get(requested_model) if team_limits else None
|
||||
if team_limit is None:
|
||||
return None
|
||||
if self._key_owns_model_limit(user_api_key_dict, requested_model, rate_limit_key):
|
||||
return None
|
||||
return team_limit
|
||||
|
||||
def _key_owns_model_tpm_limit_from_request_metadata(
|
||||
self,
|
||||
request_metadata: Mapping[str, object],
|
||||
model_group: str | None,
|
||||
) -> bool:
|
||||
if model_group is None:
|
||||
return False
|
||||
key_view: Final = UserAPIKeyAuth.model_validate(
|
||||
{
|
||||
"metadata": request_metadata.get("user_api_key_metadata") or {},
|
||||
"model_max_budget": request_metadata.get("user_api_key_model_max_budget") or {},
|
||||
}
|
||||
)
|
||||
return self._key_owns_model_limit(key_view, model_group, "model_tpm_limit")
|
||||
|
||||
def _add_team_model_rate_limit_descriptor_from_metadata(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
requested_model: str | None,
|
||||
descriptors: list[RateLimitDescriptor],
|
||||
) -> None:
|
||||
"""Add team model rate limit descriptor from team_metadata if applicable."""
|
||||
if (
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_rpm_limit") is not None
|
||||
or get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_tpm_limit") is not None
|
||||
):
|
||||
_tpm_limit_for_team_model: Final = (
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_tpm_limit") or {}
|
||||
if requested_model is None:
|
||||
return
|
||||
team_rpm_limit: Final = self._inherited_team_model_limit(user_api_key_dict, requested_model, "model_rpm_limit")
|
||||
team_tpm_limit: Final = self._inherited_team_model_limit(user_api_key_dict, requested_model, "model_tpm_limit")
|
||||
if team_rpm_limit is None and team_tpm_limit is None:
|
||||
return
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="model_per_team",
|
||||
value=f"{user_api_key_dict.team_id}:{requested_model}",
|
||||
rate_limit={
|
||||
"requests_per_unit": team_rpm_limit,
|
||||
"tokens_per_unit": team_tpm_limit,
|
||||
"window_size": self.window_size,
|
||||
},
|
||||
)
|
||||
_rpm_limit_for_team_model: Final = (
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_rpm_limit") or {}
|
||||
)
|
||||
should_check_rate_limit: Final = (
|
||||
requested_model in _tpm_limit_for_team_model or requested_model in _rpm_limit_for_team_model
|
||||
)
|
||||
|
||||
if should_check_rate_limit and requested_model is not None:
|
||||
model_specific_tpm_limit: Final = _tpm_limit_for_team_model.get(requested_model)
|
||||
model_specific_rpm_limit: Final = _rpm_limit_for_team_model.get(requested_model)
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="model_per_team",
|
||||
value=f"{user_api_key_dict.team_id}:{requested_model}",
|
||||
rate_limit={
|
||||
"requests_per_unit": model_specific_rpm_limit,
|
||||
"tokens_per_unit": model_specific_tpm_limit,
|
||||
"window_size": self.window_size,
|
||||
},
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
def _add_project_model_rate_limit_descriptor_from_metadata(
|
||||
self,
|
||||
|
|
@ -4459,6 +4486,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
kwargs=kwargs,
|
||||
model_group=reconcile_model,
|
||||
)
|
||||
charged_targets: Final = (
|
||||
[target for target in targets if target[0] != "model_per_team"]
|
||||
if self._key_owns_model_tpm_limit_from_request_metadata(request_metadata, reconcile_model)
|
||||
else targets
|
||||
)
|
||||
if reserved_tokens > 0 and total_tokens < reserved_tokens:
|
||||
verbose_proxy_logger.debug(
|
||||
"Releasing unused TPM budget on success: reserved=%s, actual=%s, release=%s",
|
||||
|
|
@ -4468,7 +4500,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
pipeline_operations.extend(
|
||||
self._build_reservation_aware_tpm_ops(
|
||||
targets=targets,
|
||||
targets=charged_targets,
|
||||
reserved_scopes=reserved_scopes,
|
||||
actual_tokens=total_tokens,
|
||||
reserved_tokens=reserved_tokens,
|
||||
|
|
|
|||
|
|
@ -114,6 +114,12 @@ class DailySpendRecord(Protocol):
|
|||
@property
|
||||
def failed_requests(self) -> int: ...
|
||||
|
||||
@property
|
||||
def total_response_time_ms(self) -> int: ...
|
||||
|
||||
@property
|
||||
def timed_requests(self) -> int: ...
|
||||
|
||||
|
||||
class _KeyMetadataDict(TypedDict, total=False):
|
||||
key_alias: ReadOnly[str | None]
|
||||
|
|
@ -162,6 +168,8 @@ class _GroupingSetsRow(SimpleNamespace):
|
|||
api_requests: int | None
|
||||
successful_requests: int | None
|
||||
failed_requests: int | None
|
||||
total_response_time_ms: int | None
|
||||
timed_requests: int | None
|
||||
|
||||
|
||||
class _EntityRollupRow(_GroupingSetsRow):
|
||||
|
|
@ -217,6 +225,8 @@ def update_metrics(existing_metrics: SpendMetrics, record: DailySpendRecord) ->
|
|||
existing_metrics.api_requests += record.api_requests or 0
|
||||
existing_metrics.successful_requests += record.successful_requests or 0
|
||||
existing_metrics.failed_requests += record.failed_requests or 0
|
||||
existing_metrics.total_response_time_ms += record.total_response_time_ms or 0
|
||||
existing_metrics.timed_requests += record.timed_requests or 0
|
||||
return existing_metrics
|
||||
|
||||
|
||||
|
|
@ -767,7 +777,9 @@ def _build_aggregated_sql_query(
|
|||
SUM(autorouter_savings_spend)::float AS autorouter_savings_spend,
|
||||
SUM(api_requests)::bigint AS api_requests,
|
||||
SUM(successful_requests)::bigint AS successful_requests,
|
||||
SUM(failed_requests)::bigint AS failed_requests
|
||||
SUM(failed_requests)::bigint AS failed_requests,
|
||||
SUM(total_response_time_ms)::bigint AS total_response_time_ms,
|
||||
SUM(timed_requests)::bigint AS timed_requests
|
||||
FROM "{pg_table}"
|
||||
WHERE {where_clause}
|
||||
GROUP BY GROUPING SETS (
|
||||
|
|
@ -846,7 +858,9 @@ def _build_entity_rollup_sql_query(
|
|||
SUM(autorouter_savings_spend)::float AS autorouter_savings_spend,
|
||||
SUM(api_requests)::bigint AS api_requests,
|
||||
SUM(successful_requests)::bigint AS successful_requests,
|
||||
SUM(failed_requests)::bigint AS failed_requests
|
||||
SUM(failed_requests)::bigint AS failed_requests,
|
||||
SUM(total_response_time_ms)::bigint AS total_response_time_ms,
|
||||
SUM(timed_requests)::bigint AS timed_requests
|
||||
FROM "{pg_table}"
|
||||
WHERE {where_clause}
|
||||
GROUP BY GROUPING SETS (
|
||||
|
|
@ -985,6 +999,8 @@ def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics:
|
|||
api_requests=record.api_requests or 0,
|
||||
successful_requests=record.successful_requests or 0,
|
||||
failed_requests=record.failed_requests or 0,
|
||||
total_response_time_ms=record.total_response_time_ms or 0,
|
||||
timed_requests=record.timed_requests or 0,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1246,6 +1262,8 @@ async def get_daily_activity(
|
|||
total_prompt_caching_savings_spend=metadata_metrics.prompt_caching_savings_spend,
|
||||
total_gateway_injected_caching_savings_spend=metadata_metrics.gateway_injected_caching_savings_spend,
|
||||
total_autorouter_savings_spend=metadata_metrics.autorouter_savings_spend,
|
||||
total_response_time_ms=metadata_metrics.total_response_time_ms,
|
||||
total_timed_requests=metadata_metrics.timed_requests,
|
||||
page=page,
|
||||
total_pages=-(-total_count // page_size), # Ceiling division
|
||||
has_more=(page * page_size) < total_count,
|
||||
|
|
@ -1423,6 +1441,8 @@ async def get_daily_activity_aggregated(
|
|||
"totals"
|
||||
].gateway_injected_caching_savings_spend,
|
||||
total_autorouter_savings_spend=aggregated["totals"].autorouter_savings_spend,
|
||||
total_response_time_ms=aggregated["totals"].total_response_time_ms,
|
||||
total_timed_requests=aggregated["totals"].timed_requests,
|
||||
page=1,
|
||||
total_pages=1,
|
||||
has_more=False,
|
||||
|
|
|
|||
|
|
@ -170,6 +170,7 @@ if MCP_AVAILABLE:
|
|||
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
|
||||
admitted_user_context,
|
||||
build_effective_auth_contexts,
|
||||
can_access_mcp_server,
|
||||
is_ui_session_credential,
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
|
|
@ -2483,10 +2484,11 @@ if MCP_AVAILABLE:
|
|||
)
|
||||
return server
|
||||
|
||||
allowed_server_ids: Final[set[str]] = set()
|
||||
for auth_context in await build_effective_auth_contexts(user_api_key_dict):
|
||||
allowed_server_ids.update(await global_mcp_server_manager.get_allowed_mcp_servers(auth_context))
|
||||
if server is None or server.server_id not in allowed_server_ids:
|
||||
if server is None or not await can_access_mcp_server(
|
||||
user_api_key_dict,
|
||||
server.server_id,
|
||||
global_mcp_server_manager.get_allowed_mcp_servers,
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail={
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
|||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
from litellm.llms.azure.passthrough.transformation import foreign_azure_deployment
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
|
||||
from litellm.proxy._types import *
|
||||
|
|
@ -1550,6 +1551,26 @@ async def _relay_azure_router_model(
|
|||
"put the model group name in the deployments segment"
|
||||
}
|
||||
raise HTTPException(status_code=400, detail=rejection)
|
||||
return await _relay_router_model(
|
||||
llm_router=llm_router,
|
||||
model=model,
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
request_body=request_body,
|
||||
is_streaming_request=is_streaming_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
async def _relay_router_model(
|
||||
llm_router: litellm.Router,
|
||||
model: str,
|
||||
endpoint: str,
|
||||
request: Request,
|
||||
request_body: Mapping[str, object],
|
||||
is_streaming_request: bool,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> Response:
|
||||
try:
|
||||
result: Final = await llm_router.allm_passthrough_route(
|
||||
model=model,
|
||||
|
|
@ -1599,6 +1620,65 @@ async def _relay_azure_router_model(
|
|||
)
|
||||
|
||||
|
||||
@router.api_route(
|
||||
"/nvidia_nim/{endpoint:path}",
|
||||
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
|
||||
tags=["NVIDIA NIM Pass-through", "pass-through"],
|
||||
)
|
||||
async def nvidia_nim_proxy_route(
|
||||
endpoint: str,
|
||||
request: Request,
|
||||
fastapi_response: Response,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
):
|
||||
"""
|
||||
Relay a native NVIDIA NIM request through a LiteLLM model group.
|
||||
|
||||
`{PROXY_BASE_URL}/nvidia_nim/{model_group}/v1/infer` forwards the body unchanged to the deployment's
|
||||
`api_base`, so object detection and OCR NIMs whose payload carries no `model` field still go through
|
||||
virtual key auth, model access checks, and spend logging.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
return await relay_nvidia_nim_request(
|
||||
llm_router=llm_router,
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
request_body=await get_request_body(request),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
async def relay_nvidia_nim_request(
|
||||
llm_router: litellm.Router | None,
|
||||
endpoint: str,
|
||||
request: Request,
|
||||
request_body: Mapping[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> Response:
|
||||
model_group: Final = nvidia_nim_model_group_in_path(endpoint, llm_router.get_model_list()) if llm_router else None
|
||||
if llm_router is None or model_group is None:
|
||||
rejection: Final[RelayRejection] = {
|
||||
"error": "no NVIDIA NIM model group in the path; call /nvidia_nim/{model_group}/v1/infer with a model "
|
||||
"group from your `model_list` whose deployments all use `nvidia_nim/` models"
|
||||
}
|
||||
raise HTTPException(status_code=400, detail=rejection)
|
||||
|
||||
is_streaming_request: Final = is_passthrough_request_streaming(request_body)
|
||||
return await open_sse_before_first_byte(
|
||||
_relay_router_model(
|
||||
llm_router=llm_router,
|
||||
model=model_group,
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
request_body=request_body,
|
||||
is_streaming_request=is_streaming_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
),
|
||||
ping_interval_seconds=(litellm.sse_keepalive_ping_interval_seconds if is_streaming_request else None),
|
||||
)
|
||||
|
||||
|
||||
@router.api_route(
|
||||
"/azure_ai/{endpoint:path}",
|
||||
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
|
||||
|
|
|
|||
|
|
@ -12,6 +12,9 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
|||
ModelResponseIterator as GeminiModelResponseIterator,
|
||||
)
|
||||
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import (
|
||||
VertexPassthroughLoggingHandler,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
ModelResponse,
|
||||
TextCompletionResponse,
|
||||
|
|
@ -40,6 +43,17 @@ class GeminiPassthroughLoggingHandler:
|
|||
request_body: dict,
|
||||
**kwargs,
|
||||
) -> PassThroughEndpointLoggingTypedDict:
|
||||
if VertexPassthroughLoggingHandler.is_interactions_route(url_route):
|
||||
return VertexPassthroughLoggingHandler.interactions_passthrough_handler(
|
||||
httpx_response=httpx_response,
|
||||
request_body=request_body,
|
||||
logging_obj=logging_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
custom_llm_provider="gemini",
|
||||
vertex_location=None,
|
||||
)
|
||||
if "predictLongRunning" in url_route:
|
||||
model = GeminiPassthroughLoggingHandler.extract_model_from_url(url_route)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,15 +1,20 @@
|
|||
import asyncio
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import VERTEX_BATCH_PREDICTION_JOBS_ROUTE
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import (
|
||||
InteractionsUsageObjectTransformation,
|
||||
)
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
get_vertex_ai_lyria_generation_cost,
|
||||
get_vertex_location_from_url,
|
||||
|
|
@ -49,8 +54,73 @@ else:
|
|||
|
||||
EndpointType = Any
|
||||
|
||||
_VERTEX_INTERACTIONS_PATH: Final = re.compile(r"/projects/[^/]+/locations/[^/]+/interactions/?$")
|
||||
_INTERACTIONS_RESPONSE_BODY: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def _interactions_model(
|
||||
response_body: Mapping[str, object],
|
||||
request_body: Mapping[str, object] | None,
|
||||
) -> str | None:
|
||||
response_model: Final = response_body.get("model")
|
||||
if isinstance(response_model, str) and response_model:
|
||||
return response_model
|
||||
request_model: Final = (request_body or {}).get("model")
|
||||
if isinstance(request_model, str) and request_model:
|
||||
return request_model
|
||||
return None
|
||||
|
||||
|
||||
class VertexPassthroughLoggingHandler:
|
||||
@staticmethod
|
||||
def is_interactions_route(url_route: str) -> bool:
|
||||
return urlparse(url_route).path.rstrip("/").endswith("/interactions")
|
||||
|
||||
@staticmethod
|
||||
def is_vertex_interactions_route(url_route: str) -> bool:
|
||||
return _VERTEX_INTERACTIONS_PATH.search(urlparse(url_route).path) is not None
|
||||
|
||||
@staticmethod
|
||||
def interactions_passthrough_handler(
|
||||
httpx_response: httpx.Response,
|
||||
request_body: Mapping[str, object] | None,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
kwargs: dict[str, object],
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
custom_llm_provider: Literal["vertex_ai", "gemini"],
|
||||
vertex_location: str | None,
|
||||
) -> PassThroughEndpointLoggingTypedDict:
|
||||
response_body: Final = _INTERACTIONS_RESPONSE_BODY.validate_python(httpx_response.json())
|
||||
usage_object: Final = response_body.get("usage")
|
||||
model: Final = _interactions_model(response_body, request_body)
|
||||
if model is None or not InteractionsUsageObjectTransformation.is_interactions_usage_object(usage_object):
|
||||
return {"result": None, "kwargs": kwargs}
|
||||
|
||||
litellm_model_response: Final = ModelResponse(
|
||||
model=model,
|
||||
usage=InteractionsUsageObjectTransformation.transform_interactions_usage_object(
|
||||
cast(Mapping[str, Any], usage_object)
|
||||
),
|
||||
)
|
||||
logging_obj.custom_llm_provider = custom_llm_provider
|
||||
logging_kwargs: Final = (
|
||||
VertexPassthroughLoggingHandler._create_vertex_response_logging_payload_for_generate_content(
|
||||
litellm_model_response=litellm_model_response,
|
||||
model=model,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
)
|
||||
return {
|
||||
"result": litellm_model_response,
|
||||
"kwargs": {**logging_kwargs, "custom_llm_provider": custom_llm_provider},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def vertex_passthrough_handler(
|
||||
httpx_response: httpx.Response,
|
||||
|
|
@ -66,6 +136,17 @@ class VertexPassthroughLoggingHandler:
|
|||
vertex_location: Final = get_vertex_location_from_url(url_route)
|
||||
if vertex_location is not None:
|
||||
logging_obj.optional_params["vertex_location"] = vertex_location
|
||||
if VertexPassthroughLoggingHandler.is_interactions_route(url_route):
|
||||
return VertexPassthroughLoggingHandler.interactions_passthrough_handler(
|
||||
httpx_response=httpx_response,
|
||||
request_body=request_body,
|
||||
logging_obj=logging_obj,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
if "predictLongRunning" in url_route:
|
||||
model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route)
|
||||
|
||||
|
|
|
|||
|
|
@ -361,7 +361,9 @@ class PassThroughEndpointLogging:
|
|||
def is_vertex_route(self, url_route: str) -> bool:
|
||||
if any(f":{method}" in url_route for method in self.TRACKED_VERTEX_METHOD_ROUTES):
|
||||
return True
|
||||
return any(resource in url_route for resource in self.TRACKED_VERTEX_RESOURCE_ROUTES)
|
||||
if any(resource in url_route for resource in self.TRACKED_VERTEX_RESOURCE_ROUTES):
|
||||
return True
|
||||
return VertexPassthroughLoggingHandler.is_vertex_interactions_route(url_route)
|
||||
|
||||
def is_anthropic_route(self, url_route: str):
|
||||
for route in self.TRACKED_ANTHROPIC_ROUTES:
|
||||
|
|
@ -434,8 +436,12 @@ class PassThroughEndpointLogging:
|
|||
|
||||
def is_gemini_route(self, url_route: str, custom_llm_provider: str | None = None):
|
||||
"""Check if the URL route is a Gemini API route."""
|
||||
if custom_llm_provider != "gemini":
|
||||
return False
|
||||
if VertexPassthroughLoggingHandler.is_interactions_route(url_route):
|
||||
return True
|
||||
for route in self.TRACKED_GEMINI_ROUTES:
|
||||
if route in url_route and custom_llm_provider == "gemini":
|
||||
if route in url_route:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
|
|
|||
|
|
@ -261,7 +261,7 @@ class ProxyInitializationHelpers:
|
|||
import uvicorn
|
||||
|
||||
import litellm
|
||||
from litellm._logging import _get_uvicorn_json_log_config
|
||||
from litellm._logging import _get_uvicorn_json_log_config, resolve_log_level
|
||||
|
||||
uvicorn_args: Final = {
|
||||
"app": "litellm.proxy.proxy_server:app",
|
||||
|
|
@ -275,6 +275,8 @@ class ProxyInitializationHelpers:
|
|||
elif litellm.json_logs:
|
||||
# Use JSON log config for uvicorn to ensure all logs (including exceptions) are JSON
|
||||
uvicorn_args["log_config"] = _get_uvicorn_json_log_config()
|
||||
elif litellm_log := os.environ.get("LITELLM_LOG"):
|
||||
uvicorn_args["log_level"] = resolve_log_level(litellm_log)
|
||||
if keepalive_timeout is not None:
|
||||
uvicorn_args["timeout_keep_alive"] = keepalive_timeout
|
||||
if timeout_worker_healthcheck is not None:
|
||||
|
|
|
|||
|
|
@ -801,6 +801,8 @@ model LiteLLM_DailyUserSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
|
|
@ -837,6 +839,8 @@ model LiteLLM_DailyOrganizationSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
|
|
@ -873,6 +877,8 @@ model LiteLLM_DailyEndUserSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
|
|
@ -908,6 +914,8 @@ model LiteLLM_DailyAgentSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
|
|
@ -943,6 +951,8 @@ model LiteLLM_DailyTeamSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
ptu_flat_cost Float @default(0.0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
|
@ -981,6 +991,8 @@ model LiteLLM_DailyTagSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
|
|
|
|||
|
|
@ -171,15 +171,25 @@ def _cost_of_usage(
|
|||
) -> float | None:
|
||||
"""What ``usage`` costs on ``model``, or ``None`` when the model has no pricing."""
|
||||
try:
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model=model.model,
|
||||
usage=usage,
|
||||
custom_llm_provider=model.provider,
|
||||
service_tier=basis.service_tier,
|
||||
data_residency=basis.data_residency,
|
||||
model_info=model_info,
|
||||
vertex_location=basis.vertex_location,
|
||||
)
|
||||
if model.provider == "anthropic":
|
||||
from litellm.llms.anthropic.cost_calculation import cost_per_token
|
||||
|
||||
prompt_cost, completion_cost = cost_per_token(
|
||||
model=model.model,
|
||||
usage=usage,
|
||||
service_tier=basis.service_tier,
|
||||
model_info=model_info,
|
||||
)
|
||||
else:
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model=model.model,
|
||||
usage=usage,
|
||||
custom_llm_provider=model.provider,
|
||||
service_tier=basis.service_tier,
|
||||
data_residency=basis.data_residency,
|
||||
model_info=model_info,
|
||||
vertex_location=basis.vertex_location,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models; degrade to zero savings
|
||||
verbose_proxy_logger.debug(
|
||||
"savings: cannot price usage for provider=%s model=%s (%s)", model.provider, model.model, e
|
||||
|
|
@ -198,11 +208,6 @@ def _cache_token_split(usage: Usage) -> tuple[int, int]:
|
|||
return int(read), int(created)
|
||||
|
||||
|
||||
_CACHE_SPLIT_FIELDS: Final = frozenset(
|
||||
("cached_tokens", "cache_creation_tokens", "cache_write_tokens", "cache_creation_token_details", "text_tokens")
|
||||
)
|
||||
|
||||
|
||||
def _baseline_cache_rate_keys(baseline_info: ModelInfo | None) -> tuple[bool, bool]:
|
||||
"""Whether the baseline model has a ``(cache read, cache write)`` rate of its own.
|
||||
|
||||
|
|
@ -274,19 +279,22 @@ def _baseline_usage(usage: Usage, conversation_continuing: bool, baseline_info:
|
|||
(getattr(details, field, 0) or 0) for field in ("audio_tokens", "image_tokens", "video_tokens")
|
||||
)
|
||||
return Usage(
|
||||
prompt_tokens=usage.prompt_tokens,
|
||||
completion_tokens=usage.completion_tokens,
|
||||
total_tokens=usage.total_tokens,
|
||||
completion_tokens_details=usage.completion_tokens_details,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
**details.model_dump(exclude=_CACHE_SPLIT_FIELDS),
|
||||
cached_tokens=reads,
|
||||
cache_creation_tokens=writes,
|
||||
cache_write_tokens=writes,
|
||||
cache_creation_token_details=details.cache_creation_token_details if writes else None,
|
||||
# Whatever no longer sits in a cache bucket is plain input on the baseline.
|
||||
text_tokens=max(usage.prompt_tokens - reads - writes - other_modalities, 0),
|
||||
),
|
||||
**{
|
||||
**usage.model_dump(),
|
||||
# Rebuild through Usage so private fallback counts agree with the public buckets.
|
||||
"cache_read_input_tokens": reads,
|
||||
"cache_creation_input_tokens": writes,
|
||||
"prompt_tokens_details": PromptTokensDetailsWrapper(
|
||||
**{
|
||||
**details.model_dump(),
|
||||
"cached_tokens": reads,
|
||||
"cache_creation_tokens": writes,
|
||||
"cache_write_tokens": writes,
|
||||
"cache_creation_token_details": details.cache_creation_token_details if writes else None,
|
||||
"text_tokens": max(usage.prompt_tokens - reads - writes - other_modalities, 0),
|
||||
}
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -176,6 +176,7 @@ class _SessionSpendRow(TypedDict):
|
|||
api_key: ReadOnly[str]
|
||||
session_total_count: ReadOnly[int]
|
||||
session_total_spend: float
|
||||
session_total_duration_ms: ReadOnly[int]
|
||||
mcp_tool_call_count: int
|
||||
mcp_tool_call_spend: float
|
||||
session_cache_hit_count: ReadOnly[int]
|
||||
|
|
@ -194,6 +195,7 @@ _SESSION_MODEL_NAME_MAX_LEN: Final = 256
|
|||
class _SessionSpendStats(NamedTuple):
|
||||
session_total_count: int
|
||||
session_total_spend: float
|
||||
session_total_duration_ms: int
|
||||
mcp_tool_call_count: int
|
||||
mcp_tool_call_spend: float
|
||||
session_cache_hit_count: int
|
||||
|
|
@ -4543,6 +4545,12 @@ async def _build_ui_spend_logs_response(
|
|||
SELECT session_id, api_key,
|
||||
COUNT(*)::int AS session_total_count,
|
||||
COALESCE(SUM(spend), 0)::double precision AS session_total_spend,
|
||||
COALESCE(SUM(
|
||||
COALESCE(
|
||||
request_duration_ms,
|
||||
(EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER
|
||||
)
|
||||
), 0)::bigint AS session_total_duration_ms,
|
||||
COUNT(*) FILTER (
|
||||
WHERE call_type IN {_MCP_CALL_TYPES_SQL}
|
||||
)::int AS mcp_tool_call_count,
|
||||
|
|
@ -4584,6 +4592,7 @@ async def _build_ui_spend_logs_response(
|
|||
(row["session_id"], row["api_key"]): _SessionSpendStats(
|
||||
session_total_count=int(row.get("session_total_count") or 0),
|
||||
session_total_spend=float(row.get("session_total_spend") or 0.0),
|
||||
session_total_duration_ms=int(row.get("session_total_duration_ms") or 0),
|
||||
mcp_tool_call_count=int(row.get("mcp_tool_call_count") or 0),
|
||||
mcp_tool_call_spend=float(row.get("mcp_tool_call_spend") or 0.0),
|
||||
session_cache_hit_count=int(row.get("session_cache_hit_count") or 0),
|
||||
|
|
@ -4615,6 +4624,7 @@ async def _build_ui_spend_logs_response(
|
|||
row_dict["session_total_count"] = session_stats.session_total_count if session_stats else 1
|
||||
if session_stats:
|
||||
row_dict["session_total_spend"] = session_stats.session_total_spend
|
||||
row_dict["session_total_duration_ms"] = session_stats.session_total_duration_ms
|
||||
if session_stats.mcp_tool_call_count:
|
||||
row_dict["mcp_tool_call_count"] = session_stats.mcp_tool_call_count
|
||||
row_dict["mcp_tool_call_spend"] = session_stats.mcp_tool_call_spend
|
||||
|
|
|
|||
|
|
@ -1269,7 +1269,12 @@ class ProxyLogging:
|
|||
# (e.g. MCPJWTSigner) to independently verify the caller's identity
|
||||
# before re-signing an outbound token (FR-5 verify+re-sign).
|
||||
"incoming_bearer_token": kwargs.get("incoming_bearer_token"),
|
||||
"metadata": {"headers": kwargs.get("headers") or {}},
|
||||
"metadata": {
|
||||
"headers": kwargs.get("headers") or {},
|
||||
"user_api_key_user_id": kwargs.get("user_api_key_user_id"),
|
||||
"user_api_key_team_id": kwargs.get("user_api_key_team_id"),
|
||||
"user_api_key_end_user_id": kwargs.get("user_api_key_end_user_id"),
|
||||
},
|
||||
}
|
||||
user_api_key_auth: Final = kwargs.get("user_api_key_auth")
|
||||
if isinstance(user_api_key_auth, UserAPIKeyAuth):
|
||||
|
|
@ -2664,34 +2669,15 @@ class ProxyLogging:
|
|||
user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(user_api_key_dict)
|
||||
else:
|
||||
user_api_key_auth_dict = user_api_key_dict
|
||||
# Add task to list for parallel execution
|
||||
if (
|
||||
"apply_guardrail" in type(callback).__dict__
|
||||
and not callback.use_native_lifecycle_hooks
|
||||
and user_api_key_dict is not None
|
||||
and not getattr(callback, "use_native_during_call_hook", False)
|
||||
):
|
||||
data["guardrail_to_apply"] = callback
|
||||
guardrail_task = self._run_guardrail_with_metrics(
|
||||
callback,
|
||||
unified_guardrail.async_moderation_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
call_type=call_type,
|
||||
),
|
||||
"during_call",
|
||||
guardrail_tasks.append(
|
||||
self._run_during_call_guardrail(
|
||||
callback=callback,
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
user_api_key_auth_dict=user_api_key_auth_dict,
|
||||
call_type=call_type,
|
||||
)
|
||||
else:
|
||||
guardrail_task = self._run_guardrail_with_metrics(
|
||||
callback,
|
||||
callback.async_moderation_hook(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_auth_dict,
|
||||
call_type=call_type,
|
||||
),
|
||||
"during_call",
|
||||
)
|
||||
guardrail_tasks.append(guardrail_task)
|
||||
)
|
||||
|
||||
# Step 2: Run all guardrail tasks in parallel
|
||||
if guardrail_tasks:
|
||||
|
|
@ -2703,6 +2689,41 @@ class ProxyLogging:
|
|||
|
||||
return data
|
||||
|
||||
async def _run_during_call_guardrail(
|
||||
self,
|
||||
callback: CustomGuardrail,
|
||||
data: dict[str, object], # mutable-ok: request payload dict, guardrail_to_apply is written in place
|
||||
user_api_key_dict: UserAPIKeyAuth | None,
|
||||
user_api_key_auth_dict: UserAPIKeyAuth | dict[str, object] | None,
|
||||
call_type: CallTypesLiteral,
|
||||
) -> None:
|
||||
if (
|
||||
"apply_guardrail" in type(callback).__dict__
|
||||
and not callback.use_native_lifecycle_hooks
|
||||
and user_api_key_dict is not None
|
||||
and not callback.use_native_during_call_hook
|
||||
):
|
||||
data["guardrail_to_apply"] = callback
|
||||
await self._run_guardrail_with_metrics(
|
||||
callback,
|
||||
unified_guardrail.async_moderation_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
call_type=call_type,
|
||||
),
|
||||
"during_call",
|
||||
)
|
||||
return
|
||||
await self._run_guardrail_with_metrics(
|
||||
callback,
|
||||
callback.async_moderation_hook(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_auth_dict,
|
||||
call_type=call_type,
|
||||
),
|
||||
"during_call",
|
||||
)
|
||||
|
||||
async def failed_tracking_alert(
|
||||
self,
|
||||
error_message: str,
|
||||
|
|
|
|||
|
|
@ -32,6 +32,9 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
|
|||
)
|
||||
from litellm.litellm_core_utils.thread_pool_executor import executor
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils
|
||||
from litellm.types.integrations.custom_logger import converted_stream_requested
|
||||
from litellm.types.llms.openai import (
|
||||
|
|
@ -257,6 +260,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self._failure_handled = False # Track if failure handler has been called
|
||||
self._yielded_first_chunk = False
|
||||
self._generated_content = ""
|
||||
self._generated_tool_arguments = ""
|
||||
self._completed_response_cached = False
|
||||
self._completed_response_logged = False
|
||||
self._completed_response_cache_hit: bool | None = None
|
||||
|
|
@ -352,6 +356,10 @@ class BaseResponsesAPIStreamingIterator:
|
|||
_delta: Final = getattr(openai_responses_api_chunk, "delta", None)
|
||||
if isinstance(_delta, str):
|
||||
self._generated_content += _delta
|
||||
elif _event_type in _TOOL_ARGUMENTS_DELTA_EVENTS:
|
||||
_args_delta: Final = getattr(openai_responses_api_chunk, "delta", None)
|
||||
if isinstance(_args_delta, str):
|
||||
self._generated_tool_arguments += _args_delta
|
||||
_stream_model_id: Final = _model_id_from_metadata(self.litellm_metadata)
|
||||
if _event_type in (
|
||||
ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
|
||||
|
|
@ -419,14 +427,41 @@ class BaseResponsesAPIStreamingIterator:
|
|||
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
|
||||
):
|
||||
self.completed_response = openai_responses_api_chunk
|
||||
_stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj)
|
||||
_response_obj: Final[object] = getattr(openai_responses_api_chunk, "response", None)
|
||||
_estimate_wanted: Final[bool] = _chunk_type in (
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
|
||||
)
|
||||
_billed_response: Final[ResponsesAPIResponse | None] = _billed_terminal_response(
|
||||
_response_obj,
|
||||
(
|
||||
lambda: (
|
||||
_estimate_usage_safely(
|
||||
self.model or "",
|
||||
self.request_data.get("input"),
|
||||
self.request_data,
|
||||
self._generated_content + self._generated_tool_arguments,
|
||||
)
|
||||
if _estimate_wanted
|
||||
else None
|
||||
)
|
||||
),
|
||||
)
|
||||
_terminal_chunk: Final = (
|
||||
openai_responses_api_chunk
|
||||
if _billed_response is None or _billed_response is _response_obj
|
||||
else openai_responses_api_chunk.model_copy(update={"response": _billed_response})
|
||||
)
|
||||
self.completed_response = _terminal_chunk
|
||||
_stamp_responses_usage_cost(_billed_response, self.logging_obj)
|
||||
|
||||
if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED:
|
||||
self._handle_logging_failed_response()
|
||||
else:
|
||||
self._handle_logging_completed_response()
|
||||
|
||||
return _terminal_chunk
|
||||
|
||||
return openai_responses_api_chunk
|
||||
|
||||
return None
|
||||
|
|
@ -655,7 +690,9 @@ class BaseResponsesAPIStreamingIterator:
|
|||
if cache is None:
|
||||
return
|
||||
|
||||
cached_response: Final = response_obj.model_dump_json()
|
||||
cached_response: Final = _dump_json_safely(response_obj)
|
||||
if cached_response is None:
|
||||
return
|
||||
if is_async:
|
||||
from litellm.caching.caching_handler import create_cache_write_task
|
||||
|
||||
|
|
@ -1301,6 +1338,31 @@ def _add_text_like_part_events(
|
|||
)
|
||||
|
||||
|
||||
def _billed_terminal_response(
|
||||
response_obj: object, estimate: Callable[[], ResponseAPIUsage | None] | None
|
||||
) -> ResponsesAPIResponse | None:
|
||||
if isinstance(response_obj, ResponsesAPIResponse):
|
||||
return (
|
||||
response_obj
|
||||
if response_obj.usage is not None or estimate is None
|
||||
else response_obj.model_copy(update={"usage": estimate()})
|
||||
)
|
||||
if not isinstance(response_obj, dict):
|
||||
return None
|
||||
usage: Final[object] = response_obj.get("usage") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # a model_constructed terminal event leaves response as an untyped dict
|
||||
return ResponsesAPIResponse.model_construct(
|
||||
**{**response_obj, "usage": usage if usage is not None or estimate is None else estimate()} # pyright: ignore[reportUnknownArgumentType, reportArgumentType] # same untyped dict spread
|
||||
)
|
||||
|
||||
|
||||
def _dump_json_safely(response: BaseModel) -> str | None:
|
||||
try:
|
||||
return response.model_dump_json()
|
||||
except Exception as exc:
|
||||
verbose_logger.debug("could not serialize completed response for cache: %s", exc)
|
||||
return None
|
||||
|
||||
|
||||
def _logging_copy(event: object) -> object:
|
||||
"""Hand logging callbacks a copy, so their usage rewrite (Responses shape to chat shape) never
|
||||
reaches the event the caller is iterating. The round trip through ``model_dump`` sidesteps the
|
||||
|
|
@ -1332,6 +1394,56 @@ def _usage_as_model(usage: object) -> ResponseAPIUsage | None:
|
|||
return None
|
||||
|
||||
|
||||
_TOOL_ARGUMENTS_DELTA_EVENTS: Final = frozenset(
|
||||
{
|
||||
ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
|
||||
ResponsesAPIStreamEvents.CUSTOM_TOOL_CALL_INPUT_DELTA,
|
||||
ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _estimate_usage_from_text(
|
||||
model: str,
|
||||
request_input: object,
|
||||
responses_api_request: Mapping[str, object],
|
||||
generated_text: str,
|
||||
) -> ResponseAPIUsage:
|
||||
messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped
|
||||
input=request_input, # pyright: ignore[reportArgumentType] # the raw Responses API input is a str or ResponseInputParam list, matching the helper's declared union
|
||||
responses_api_request=dict(responses_api_request),
|
||||
)
|
||||
input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped
|
||||
model=model, messages=messages
|
||||
)
|
||||
output_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped
|
||||
model=model, text=generated_text, count_response_tokens=True
|
||||
)
|
||||
return ResponseAPIUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=input_tokens + output_tokens,
|
||||
)
|
||||
|
||||
|
||||
def _estimate_usage_safely(
|
||||
model: str,
|
||||
request_input: object,
|
||||
responses_api_request: Mapping[str, object],
|
||||
generated_text: str,
|
||||
) -> ResponseAPIUsage | None:
|
||||
try:
|
||||
return _estimate_usage_from_text(
|
||||
model=model,
|
||||
request_input=request_input,
|
||||
responses_api_request=responses_api_request,
|
||||
generated_text=generated_text,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Could not estimate usage from stream text, billing $0: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
def _stamp_responses_usage_cost(
|
||||
response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -425,12 +425,34 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream])
|
|||
|
||||
_NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({})
|
||||
_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str])
|
||||
|
||||
|
||||
def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
|
||||
return tuple(item for item in value if isinstance(item, str)) if isinstance(value, tuple) else ()
|
||||
|
||||
|
||||
def _silent_experiment_targets(silent_model: object) -> tuple[str, ...]:
|
||||
if silent_model is None:
|
||||
return ()
|
||||
try:
|
||||
targets: Final = _SILENT_MODEL_ADAPTER.validate_python(silent_model)
|
||||
except ValidationError:
|
||||
verbose_router_logger.warning(
|
||||
"silent_model must be a model name or a list of model names, got %r; skipping shadow traffic",
|
||||
silent_model,
|
||||
)
|
||||
return ()
|
||||
return (targets,) if isinstance(targets, str) else tuple(targets)
|
||||
|
||||
|
||||
def _silent_experiment_kwargs_snapshot(kwargs: Mapping[str, object]) -> Mapping[str, object]:
|
||||
metadata: Final = kwargs.get("metadata")
|
||||
if not isinstance(metadata, Mapping):
|
||||
return MappingProxyType({**kwargs})
|
||||
return MappingProxyType({**kwargs, "metadata": dict(metadata)})
|
||||
|
||||
|
||||
def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]:
|
||||
"""
|
||||
Realtime client-secret requests carry the model inside ``session`` as well, and the caller's copy of it still
|
||||
|
|
@ -2455,18 +2477,17 @@ class Router:
|
|||
)
|
||||
silent_model: Final = litellm_params.pop("silent_model", None)
|
||||
|
||||
if silent_model is not None:
|
||||
for silent_target in _silent_experiment_targets(silent_model):
|
||||
# Mirroring traffic to a secondary model
|
||||
# Use threading.Thread (not ThreadPoolExecutor) - executor.submit()
|
||||
# requires pickling args, which fails when kwargs contain unpicklable
|
||||
# objects (e.g. _thread.RLock from OTEL spans, loggers) in deployment.
|
||||
thread: Final = threading.Thread(
|
||||
threading.Thread(
|
||||
target=self._silent_experiment_completion,
|
||||
args=(silent_model, messages),
|
||||
kwargs=kwargs,
|
||||
args=(silent_target, messages),
|
||||
kwargs=_silent_experiment_kwargs_snapshot(kwargs),
|
||||
daemon=True,
|
||||
)
|
||||
thread.start()
|
||||
).start()
|
||||
|
||||
kwargs.setdefault("messages", messages)
|
||||
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||||
|
|
@ -2567,9 +2588,6 @@ class Router:
|
|||
|
||||
silent_kwargs["metadata"]["is_silent_experiment"] = True
|
||||
|
||||
# Force stream=False so the response is fully consumed and callbacks fire
|
||||
silent_kwargs["stream"] = False
|
||||
|
||||
# Pop logging objects and call IDs to ensure a fresh logging context
|
||||
# This prevents collisions in the Proxy's database (spend_logs)
|
||||
silent_kwargs.pop("litellm_call_id", None)
|
||||
|
|
@ -2579,6 +2597,23 @@ class Router:
|
|||
|
||||
return silent_kwargs
|
||||
|
||||
async def _run_silent_experiment(
|
||||
self, silent_model: str, messages: Sequence[Mapping[str, str]], silent_kwargs: Mapping[str, object]
|
||||
) -> None:
|
||||
remaining_kwargs: Final = MappingProxyType(
|
||||
{key: value for key, value in silent_kwargs.items() if key != "stream"}
|
||||
)
|
||||
response: Final = await self.acompletion(
|
||||
model=silent_model,
|
||||
messages=cast(list[AllMessageValues], messages),
|
||||
stream=bool(silent_kwargs.get("stream", False)),
|
||||
**remaining_kwargs,
|
||||
)
|
||||
if not isinstance(response, CustomStreamWrapper):
|
||||
return
|
||||
async for _ in response:
|
||||
pass
|
||||
|
||||
def _silent_experiment_completion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs):
|
||||
"""
|
||||
Run a silent experiment in the background (thread).
|
||||
|
|
@ -2604,11 +2639,7 @@ class Router:
|
|||
try:
|
||||
|
||||
async def _run_silent_completion():
|
||||
await self.acompletion(
|
||||
model=silent_model,
|
||||
messages=cast(list[AllMessageValues], messages),
|
||||
**silent_kwargs,
|
||||
)
|
||||
await self._run_silent_experiment(silent_model, messages, silent_kwargs)
|
||||
# Drain any fire-and-forget tasks (e.g. alerting hooks)
|
||||
# scheduled via asyncio.create_task during acompletion.
|
||||
pending: Final = asyncio.all_tasks()
|
||||
|
|
@ -3500,11 +3531,7 @@ class Router:
|
|||
silent_kwargs["metadata"]["model_group"] = silent_model
|
||||
|
||||
# Trigger the silent request
|
||||
await self.acompletion(
|
||||
model=silent_model,
|
||||
messages=cast(list[AllMessageValues], messages),
|
||||
**silent_kwargs,
|
||||
)
|
||||
await self._run_silent_experiment(silent_model, messages, silent_kwargs)
|
||||
except Exception as e:
|
||||
verbose_router_logger.error("Silent experiment failed for model %s: %s", silent_model, e)
|
||||
|
||||
|
|
@ -3563,14 +3590,14 @@ class Router:
|
|||
)
|
||||
silent_model: Final = litellm_params.pop("silent_model", None)
|
||||
|
||||
if silent_model is not None:
|
||||
for silent_target in _silent_experiment_targets(silent_model):
|
||||
# Mirroring traffic to a secondary model
|
||||
# This is a silent experiment, so we don't want to block the primary request
|
||||
asyncio.create_task(
|
||||
self._silent_experiment_acompletion(
|
||||
silent_model=silent_model,
|
||||
silent_model=silent_target,
|
||||
messages=messages, # Use messages instead of *args
|
||||
**kwargs,
|
||||
**_silent_experiment_kwargs_snapshot(kwargs),
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -218,9 +218,8 @@ class GatedAutoRouterCapability:
|
|||
stored ``litellm_params`` (``{config}`` is the caller's expression for the normalized
|
||||
``complexity_router_config`` jsonb, substituted as many times as the predicate needs); they live
|
||||
on one record so they cannot drift apart. ``subject`` and ``remedy`` build the shared refusal
|
||||
message. A validated config claims at most one capability, and the validator is what makes that
|
||||
true: tier_definitions rejects every heuristic classifier_type, and it also rejects the
|
||||
classifier system_prompt, which in turn only applies to the classifier types heuristic_v2 is not.
|
||||
message. A validated config claims at most one capability: gated classifier types cannot be
|
||||
combined with operator-defined tiers or classifier prompts.
|
||||
"""
|
||||
|
||||
key: str
|
||||
|
|
@ -238,6 +237,22 @@ HEURISTIC_V2_CAPABILITY: Final = GatedAutoRouterCapability(
|
|||
sql_config_predicate="{config} ->> 'classifier_type' = 'heuristic_v2'",
|
||||
)
|
||||
|
||||
CAPABILITY_CLASSIFIER_CAPABILITY: Final = GatedAutoRouterCapability(
|
||||
key="capability",
|
||||
subject="with classifier_type 'capability' (Capability)",
|
||||
remedy="Use a different classifier or remove an existing Capability router.",
|
||||
uses=lambda config: _mapping(config).get("classifier_type") == "capability",
|
||||
sql_config_predicate="{config} ->> 'classifier_type' = 'capability'",
|
||||
)
|
||||
|
||||
LLM_V2_CAPABILITY: Final = GatedAutoRouterCapability(
|
||||
key="llm_v2",
|
||||
subject="with classifier_type 'llm_v2' (Fuse v2)",
|
||||
remedy="Use a different classifier or remove an existing Fuse v2 router.",
|
||||
uses=lambda config: _mapping(config).get("classifier_type") == "llm_v2",
|
||||
sql_config_predicate="{config} ->> 'classifier_type' = 'llm_v2'",
|
||||
)
|
||||
|
||||
_OPERATOR_PROMPT_FIELDS_SQL: Final = " OR ".join(
|
||||
f"{{config}} ->> '{field}' IS NOT NULL" for field in OPERATOR_CLASSIFIER_PROMPT_FIELDS
|
||||
)
|
||||
|
|
@ -258,7 +273,12 @@ CUSTOMIZATION_CAPABILITY: Final = GatedAutoRouterCapability(
|
|||
),
|
||||
)
|
||||
|
||||
GATED_AUTO_ROUTER_CAPABILITIES: Final = (HEURISTIC_V2_CAPABILITY, CUSTOMIZATION_CAPABILITY)
|
||||
GATED_AUTO_ROUTER_CAPABILITIES: Final = (
|
||||
HEURISTIC_V2_CAPABILITY,
|
||||
CAPABILITY_CLASSIFIER_CAPABILITY,
|
||||
LLM_V2_CAPABILITY,
|
||||
CUSTOMIZATION_CAPABILITY,
|
||||
)
|
||||
|
||||
|
||||
def claimed_capability(complexity_router_config: object) -> GatedAutoRouterCapability | None:
|
||||
|
|
|
|||
|
|
@ -8,6 +8,9 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator, model_valida
|
|||
from typing_extensions import Required, TypedDict
|
||||
|
||||
from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.agent_365 import (
|
||||
Agent365GuardrailConfigModel,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.akto import (
|
||||
AktoConfigModel,
|
||||
)
|
||||
|
|
@ -137,6 +140,7 @@ class SupportedGuardrailIntegrations(Enum):
|
|||
COMPRESR = "compresr"
|
||||
STRAIKER = "straiker"
|
||||
ALICE = "alice"
|
||||
AGENT_365 = "agent_365"
|
||||
CONDUCT = "conduct"
|
||||
|
||||
|
||||
|
|
@ -1045,7 +1049,7 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up
|
|||
default="fail_closed",
|
||||
description=(
|
||||
"Behavior when a guardrail endpoint is unreachable due to network errors. "
|
||||
"Implemented by guardrail='generic_guardrail_api', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. "
|
||||
"Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. "
|
||||
"'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed."
|
||||
),
|
||||
)
|
||||
|
|
@ -1183,6 +1187,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o
|
|||
QostodianNexusConfigModel,
|
||||
VigilGuardGuardrailConfigModel,
|
||||
SingulrGuardrailConfigModel,
|
||||
Agent365GuardrailConfigModel,
|
||||
):
|
||||
guardrail: str = Field(description="The type of guardrail integration to use")
|
||||
mode: str | list[str] | Mode = Field(
|
||||
|
|
|
|||
|
|
@ -586,7 +586,7 @@ class MessageBlockDelta(TypedDict):
|
|||
|
||||
type: Literal["message_delta"]
|
||||
delta: MessageDelta
|
||||
usage: UsageDelta
|
||||
usage: NotRequired[ReadOnly[UsageDelta]]
|
||||
context_management: NotRequired[ContextManagementResponse]
|
||||
|
||||
|
||||
|
|
|
|||
66
litellm/types/proxy/guardrails/guardrail_hooks/agent_365.py
Normal file
66
litellm/types/proxy/guardrails/guardrail_hooks/agent_365.py
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
from typing import Final
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from .base import GuardrailConfigModel
|
||||
|
||||
AGENT_365_PROD_API_BASE: Final = "https://agent365.svc.cloud.microsoft"
|
||||
AGENT_365_PROD_RESOURCE_APP_ID: Final = "ea9ffc3e-8a23-4a7d-836d-234d7c7565c1"
|
||||
AGENT_365_SCOPE_NAME: Final = "ThreatProtection.Evaluate.All"
|
||||
|
||||
|
||||
class Agent365GuardrailConfigModel(GuardrailConfigModel):
|
||||
tenant_id: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Entra tenant id used for the On-Behalf-Of token exchange. "
|
||||
"Falls back to the AGENT365_TENANT_ID environment variable."
|
||||
),
|
||||
)
|
||||
|
||||
client_id: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Client id of the gateway's Entra app registration (a confidential client). "
|
||||
"Falls back to the AGENT365_CLIENT_ID environment variable."
|
||||
),
|
||||
)
|
||||
|
||||
client_secret: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Client secret of the gateway's Entra app registration, used to perform the "
|
||||
"On-Behalf-Of exchange. Falls back to the AGENT365_CLIENT_SECRET environment variable."
|
||||
),
|
||||
)
|
||||
|
||||
api_base: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Base URL of the Microsoft Agent 365 tool-evaluation endpoint. "
|
||||
f"Defaults to the production endpoint {AGENT_365_PROD_API_BASE}. "
|
||||
"Falls back to the AGENT365_API_BASE environment variable."
|
||||
),
|
||||
)
|
||||
|
||||
resource_app_id: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Application id of the Agent 365 resource the OBO token is minted for. "
|
||||
f"Defaults to the production resource {AGENT_365_PROD_RESOURCE_APP_ID}; "
|
||||
"the Test and PreProd environments use a different id. "
|
||||
"Falls back to the AGENT365_RESOURCE_APP_ID environment variable."
|
||||
),
|
||||
)
|
||||
|
||||
agent_id: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Agent identity reported to Agent 365 with every tool evaluation. "
|
||||
"When unset, the caller's key alias is used."
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "Microsoft Agent 365"
|
||||
|
|
@ -1,24 +1,53 @@
|
|||
from typing import Any
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from .base import GuardrailConfigModel
|
||||
|
||||
|
||||
class SingulrGuardrailRequest(BaseModel):
|
||||
model: str | None = None
|
||||
messages: list[dict[str, Any]] | None = None
|
||||
tools: list[dict[str, Any]] | None = None
|
||||
model_response: dict[str, Any] | None = None
|
||||
litellm_metadata: dict[str, Any] | None = None
|
||||
class ContentBlock(BaseModel):
|
||||
type: str | None = None
|
||||
text: str | None = None
|
||||
|
||||
|
||||
class ToolCallFunction(BaseModel):
|
||||
name: str
|
||||
arguments: str
|
||||
|
||||
|
||||
class ToolCall(BaseModel):
|
||||
id: str
|
||||
type: str = "function"
|
||||
function: ToolCallFunction
|
||||
|
||||
|
||||
class AssistantMessage(BaseModel):
|
||||
role: Literal["assistant"] = "assistant"
|
||||
content: str | Sequence[ContentBlock] | None = None
|
||||
tool_calls: Sequence[ToolCall] | None = None
|
||||
|
||||
|
||||
class SingulrGuardrailPayload(BaseModel):
|
||||
litellm_call_id: str | None = None
|
||||
request_data: SingulrGuardrailRequest | None = None
|
||||
input_type: str
|
||||
is_playground_request: bool | None = None
|
||||
playground_text: str | None = None
|
||||
correlation_id: str | None = None
|
||||
model_name: str | None = None
|
||||
model_provider_name: str | None = None
|
||||
guardrail_scope: str | None = None
|
||||
messages: Sequence[Mapping[str, object]] | None = None
|
||||
images: Sequence[str] | None = None
|
||||
tools: Sequence[Mapping[str, object]] | None = None
|
||||
response: AssistantMessage | None = None
|
||||
metadata: Mapping[str, str] | None = None
|
||||
|
||||
|
||||
class SingulrMcpGuardrailPayload(BaseModel):
|
||||
model_name: str | None = None
|
||||
guardrail_scope: str | None = None
|
||||
tool_name: str | None = None
|
||||
tool_arguments: object = None
|
||||
mcp_server_name: str | None = None
|
||||
tool_result: Sequence[str] | None = None
|
||||
metadata: Mapping[str, str] | None = None
|
||||
|
||||
|
||||
class SingulrGuardrailResponse(BaseModel):
|
||||
|
|
|
|||
|
|
@ -32,6 +32,8 @@ class SpendMetrics(BaseModel):
|
|||
successful_requests: int = Field(default=0)
|
||||
failed_requests: int = Field(default=0)
|
||||
api_requests: int = Field(default=0)
|
||||
total_response_time_ms: int = Field(default=0)
|
||||
timed_requests: int = Field(default=0)
|
||||
|
||||
|
||||
class MetricBase(BaseModel):
|
||||
|
|
@ -93,6 +95,8 @@ class DailySpendMetadata(BaseModel):
|
|||
total_prompt_caching_savings_spend: float = Field(default=0.0)
|
||||
total_gateway_injected_caching_savings_spend: float = Field(default=0.0)
|
||||
total_autorouter_savings_spend: float = Field(default=0.0)
|
||||
total_response_time_ms: int = Field(default=0)
|
||||
total_timed_requests: int = Field(default=0)
|
||||
page: int = Field(default=1)
|
||||
total_pages: int = Field(default=1)
|
||||
has_more: bool = Field(default=False)
|
||||
|
|
@ -125,6 +129,8 @@ class LiteLLM_DailyUserSpend(BaseModel):
|
|||
api_requests: int = 0
|
||||
successful_requests: int = 0
|
||||
failed_requests: int = 0
|
||||
total_response_time_ms: int = 0
|
||||
timed_requests: int = 0
|
||||
|
||||
|
||||
class GroupedData(TypedDict):
|
||||
|
|
|
|||
|
|
@ -2957,6 +2957,7 @@ InternalCallOrigin = Literal[
|
|||
"autorouter_classifier",
|
||||
"shadow_eval_router",
|
||||
"shadow_eval_judge",
|
||||
"llm_as_a_judge_guardrail",
|
||||
"background_response_cost_poll",
|
||||
]
|
||||
"""Which internal litellm feature originated a billed sub-call, so a spend log row
|
||||
|
|
@ -2965,6 +2966,7 @@ records that it is not traffic the caller sent."""
|
|||
AUTOROUTER_CLASSIFIER_CALL_ORIGIN: Final[InternalCallOrigin] = "autorouter_classifier"
|
||||
SHADOW_EVAL_ROUTER_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_router"
|
||||
SHADOW_EVAL_JUDGE_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_judge"
|
||||
LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN: Final[InternalCallOrigin] = "llm_as_a_judge_guardrail"
|
||||
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN: Final[InternalCallOrigin] = "background_response_cost_poll"
|
||||
|
||||
|
||||
|
|
@ -4275,6 +4277,7 @@ class LiteLLMRealtimeStreamLoggingObject(LiteLLMPydanticObjectBase):
|
|||
# rate_limits.updated), blocks the event loop, and discards the session usage.
|
||||
results: SkipValidation[OpenAIRealtimeStreamList]
|
||||
usage: Usage
|
||||
service_tier: str | None = None
|
||||
_hidden_params: dict = {}
|
||||
|
||||
@field_serializer("results")
|
||||
|
|
|
|||
|
|
@ -8998,6 +8998,12 @@ class ProviderConfigManager:
|
|||
)
|
||||
|
||||
return WatsonxPassthroughConfig()
|
||||
elif LlmProviders.NVIDIA_NIM == provider:
|
||||
from litellm.llms.nvidia_nim.passthrough.transformation import (
|
||||
NvidiaNimPassthroughConfig,
|
||||
)
|
||||
|
||||
return NvidiaNimPassthroughConfig()
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -58151,6 +58151,23 @@
|
|||
"model_info": {
|
||||
"supports_reasoning": true
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "gemini-chat-baseline",
|
||||
"pattern": "gemini-(?!.*(?:-tts|-image|-live|-audio|-embedding|-computer-use|-robotics|-transcribe|-translate))(?:2[.-][5-9]|[3-9](?:[.-]\\d{1,2})?)-(?:pro|flash)(?:-lite)?(?![a-z])",
|
||||
"description": "Any Gemini text-chat id at 2.5 or higher under any namespace, including bare ids, gemini/, vertex_ai/, openrouter/google/, deepinfra/google/, vercel_ai_gateway/google/, oci/google., and databricks-gemini-<major>-<minor>: gemini-<major>[.minor]-(pro|flash)[-lite] with any trailing preview, date or variant tag. The capability flags were verified against each of those providers' own catalogs and docs. The lookahead excludes the tts, image, live, audio, embedding, computer-use, robotics, transcribe and translate lines, which are different modes with different capabilities. Provider-specific deviations, such as Perplexity's Agent API serving these as mode responses, are carried by their exact map entries, which always win over this rule. Carries no token limits or pricing, so those stay on the standard unmapped behavior rather than a guessed number. Source check 2026-09-15: all 45 first-party 2.5+ text-chat entries in this map carry every field below, and the OpenRouter (openrouter.ai/api/v1/models), Vercel AI Gateway (ai-gateway.vercel.sh/v1/models), DeepInfra (api.deepinfra.com/models/list), OCI and Databricks model docs list reasoning, tools and image input for the same models.",
|
||||
"model_info": {
|
||||
"mode": "chat",
|
||||
"supports_reasoning": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_vision": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_web_search": true
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
|
|
|
|||
|
|
@ -801,6 +801,8 @@ model LiteLLM_DailyUserSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
|
|
@ -837,6 +839,8 @@ model LiteLLM_DailyOrganizationSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
|
|
@ -873,6 +877,8 @@ model LiteLLM_DailyEndUserSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
|
|
@ -908,6 +914,8 @@ model LiteLLM_DailyAgentSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
|
|
@ -943,6 +951,8 @@ model LiteLLM_DailyTeamSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
ptu_flat_cost Float @default(0.0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
|
@ -981,6 +991,8 @@ model LiteLLM_DailyTagSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
total_response_time_ms BigInt @default(0)
|
||||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
|
|
|
|||
|
|
@ -67,6 +67,7 @@ IGNORE_FUNCTIONS = [
|
|||
"_redact_agent_params_tree", # max depth set (default 10), same shape as _redact_sensitive_litellm_params.
|
||||
"_restore_redacted_nested_value", # max depth set (default 10), mirrors _redact_agent_params_tree on the write side.
|
||||
"_unqualified", # bounded by the qualifier depth of a static TypedDict annotation (Annotated, Required/NotRequired, ReadOnly around one type, no cycles possible).
|
||||
"completion_cost", # max depth 1: recursion only fires for mixed-tier Responses WS logging objects, and each split part carries a single service_tier so _split_responses_ws_logging_object_by_service_tier returns None.
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
472
tests/code_coverage_tests/test_provider_cache.py
Normal file
472
tests/code_coverage_tests/test_provider_cache.py
Normal file
|
|
@ -0,0 +1,472 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import socket
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Generator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, replace
|
||||
from http.client import HTTPConnection
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import pytest
|
||||
from e2e_http import NetworkError, PreparedForward, RawResponse, StreamChunk, StreamHead, forward, prepare_forward
|
||||
from models import LiteLLMParamsBody
|
||||
from provider_cache import CacheEdge, CacheHit, CaptureLease, exact_key, successful_response
|
||||
from provider_cache_redis import PUBLISH, RedisCommands, RedisResponseStore, configured_cache, redis_store
|
||||
from provider_cache_routing import LIVE_PROVIDER_REQUIRED, route_cache_model
|
||||
from provider_edge import configured_cache_backend, start_provider_edge
|
||||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
|
||||
SECRET: Final = b"synthetic-cache-hmac-key-for-tests"
|
||||
BODY: Final = b'{"model":"test","messages":[{"role":"user","content":"hello"}]}'
|
||||
SUCCESS: Final = b'{"id":"provider-fixed-id","choices":[{"message":{"content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}'
|
||||
HEADERS: Final = {"content-type": "application/json", "authorization": "Bearer synthetic-account-one"}
|
||||
|
||||
|
||||
class Provider(ThreadingHTTPServer):
|
||||
hits: tuple[tuple[str, bytes], ...] = ()
|
||||
response: bytes = SUCCESS
|
||||
status: int = 200
|
||||
delay: float = 0
|
||||
stream: bool = False
|
||||
truncated: bool = False
|
||||
cookie: str = ""
|
||||
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
def do_POST(self) -> None:
|
||||
server: Final = self.server
|
||||
assert isinstance(server, Provider)
|
||||
body: Final = self.rfile.read(int(self.headers.get("content-length", "0")))
|
||||
server.hits += ((self.path, body),)
|
||||
time.sleep(server.delay)
|
||||
self.send_response(server.status)
|
||||
if server.stream:
|
||||
self.send_header("content-type", "text/event-stream")
|
||||
self.send_header("transfer-encoding", "chunked")
|
||||
self.end_headers()
|
||||
self.wfile.write(b"%x\r\n%s\r\n" % (len(server.response), server.response))
|
||||
if server.truncated:
|
||||
self.close_connection = True
|
||||
return
|
||||
self.wfile.write(b"0\r\n\r\n")
|
||||
return
|
||||
self.send_header("content-type", "application/json")
|
||||
self.send_header("content-length", str(len(server.response)))
|
||||
if server.cookie:
|
||||
self.send_header("set-cookie", server.cookie)
|
||||
self.end_headers()
|
||||
self.wfile.write(server.response)
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def provider() -> Generator[Provider, None, None]:
|
||||
server: Final = Provider(("127.0.0.1", 0), Handler)
|
||||
thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
yield server
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def redis_url(tmp_path_factory: pytest.TempPathFactory) -> Generator[str, None, None]:
|
||||
configured: Final = os.environ.get("E2E_CACHE_TEST_REDIS_URL")
|
||||
if configured:
|
||||
yield configured
|
||||
return
|
||||
binary: Final = shutil.which("redis-server")
|
||||
assert binary is not None, "Set E2E_CACHE_TEST_REDIS_URL or install Redis for cache integration checks"
|
||||
root: Final = tmp_path_factory.mktemp("provider-cache-redis")
|
||||
with socket.socket() as probe:
|
||||
probe.bind(("127.0.0.1", 0))
|
||||
port: Final = probe.getsockname()[1]
|
||||
with (root / "redis.log").open("wb") as log:
|
||||
process: Final = subprocess.Popen(
|
||||
[binary, "--bind", "127.0.0.1", "--port", str(port), "--save", "", "--appendonly", "no", "--dir", str(root)],
|
||||
stdout=log, stderr=subprocess.STDOUT,
|
||||
)
|
||||
try:
|
||||
deadline: Final = time.monotonic() + 5
|
||||
while True:
|
||||
try:
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=0.1):
|
||||
break
|
||||
except OSError:
|
||||
assert process.poll() is None and time.monotonic() < deadline
|
||||
time.sleep(0.02)
|
||||
yield f"redis://127.0.0.1:{port}/0"
|
||||
finally:
|
||||
process.terminate()
|
||||
process.wait(timeout=5)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(redis_url: str) -> RedisResponseStore:
|
||||
return redis_store(redis_url, "test-" + uuid.uuid4().hex)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def edge(cache: CacheEdge, provider: Provider) -> Generator[str, None, None]:
|
||||
upstream: Final = f"http://127.0.0.1:{provider.server_port}"
|
||||
running: Final = start_provider_edge(cache, mounts={"openai": upstream})
|
||||
try:
|
||||
yield running.edge.api_base("openai") + "/v1/chat/completions"
|
||||
finally:
|
||||
running.shutdown()
|
||||
|
||||
|
||||
def call(url: str, body: bytes = BODY, headers: dict[str, str] = HEADERS) -> RawResponse:
|
||||
result: Final = forward("POST", url, headers=headers, body=body, timeout=5)
|
||||
assert isinstance(result, RawResponse), result
|
||||
return result
|
||||
|
||||
|
||||
def test_success_is_reusable_across_fresh_edges(store: RedisResponseStore, provider: Provider) -> None:
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
assert call(url).body == SUCCESS
|
||||
assert call(url).body == SUCCESS
|
||||
with edge(CacheEdge(store, SECRET), provider) as other:
|
||||
assert call(other).body == SUCCESS
|
||||
assert len(provider.hits) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("body", [BODY + b" ", BODY.replace(b"hello", b"Hello"), BODY.replace(b"test", b"test2")])
|
||||
def test_any_body_change_calls_live(store: RedisResponseStore, provider: Provider, body: bytes) -> None:
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
call(url)
|
||||
call(url, body)
|
||||
call(url, body)
|
||||
assert len(provider.hits) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name,value", [("authorization", "Bearer another-account"), ("x-request-id", "one"), ("anthropic-version", "new")])
|
||||
def test_changed_header_cannot_reuse(store: RedisResponseStore, provider: Provider, name: str, value: str) -> None:
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
call(url)
|
||||
call(url, headers=HEADERS | {name: value})
|
||||
call(url + "?x=1")
|
||||
assert len(provider.hits) == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status,response", [(429, b'{"error":"rate limited"}'), (500, b'failed'), (200, b'{"error":"bad"}'), (200, b'not json')])
|
||||
def test_failed_provider_responses_never_enter_cache(store: RedisResponseStore, provider: Provider, status: int, response: bytes) -> None:
|
||||
provider.status = status
|
||||
provider.response = response
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
assert call(url).status_code == status
|
||||
assert call(url).body == response
|
||||
assert len(provider.hits) == 2
|
||||
|
||||
|
||||
def test_cookie_setting_success_is_reused_without_the_cookie(store: RedisResponseStore, provider: Provider) -> None:
|
||||
provider.cookie = "__cf_bm=synthetic-bot-management; Path=/; HttpOnly; Secure"
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
replies: Final = tuple(call(url) for _ in range(2))
|
||||
assert len(provider.hits) == 1
|
||||
assert all(reply.body == SUCCESS and "set-cookie" not in reply.headers for reply in replies)
|
||||
|
||||
|
||||
def test_expiry_does_not_slide(store: RedisResponseStore, provider: Provider) -> None:
|
||||
short: Final = replace(store, lifetime_ms=250)
|
||||
with edge(CacheEdge(short, SECRET), provider) as url:
|
||||
call(url)
|
||||
call(url)
|
||||
time.sleep(0.3)
|
||||
call(url)
|
||||
call(url)
|
||||
assert len(provider.hits) == 2
|
||||
|
||||
|
||||
def test_concurrent_requests_publish_atomically(store: RedisResponseStore, provider: Provider) -> None:
|
||||
provider.delay = 0.15
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
with ThreadPoolExecutor(max_workers=5) as executor:
|
||||
replies: Final = tuple(executor.map(lambda _: call(url).body, range(5)))
|
||||
assert replies == (SUCCESS,) * 5
|
||||
assert len(provider.hits) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("age_past_expiry_ms", [0, 1])
|
||||
def test_expired_response_is_rejected_without_physical_eviction(
|
||||
store: RedisResponseStore, age_past_expiry_ms: int,
|
||||
) -> None:
|
||||
response_key: Final = store.keys("expired")[0]
|
||||
retained: Final = store.client.eval(
|
||||
"""
|
||||
local clock = redis.call('TIME')
|
||||
local expires = clock[1] * 1000 + math.floor(clock[2] / 1000) - tonumber(ARGV[1])
|
||||
redis.call('HSET', KEYS[1], 'captured', expires - 86400000, 'expires', expires, 'payload', 'old-response')
|
||||
return redis.call('PTTL', KEYS[1])
|
||||
""",
|
||||
1, response_key, age_past_expiry_ms,
|
||||
)
|
||||
assert retained == -1
|
||||
replacement: Final = store.lookup("expired")
|
||||
assert isinstance(replacement, CaptureLease)
|
||||
assert replacement.expires_at_ms - replacement.captured_at_ms == 86_400_000
|
||||
assert store.publish("expired", replacement, b"fresh-response")
|
||||
hit: Final = store.lookup("expired")
|
||||
assert isinstance(hit, CacheHit) and hit.payload == b"fresh-response"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("truncated", [False, True])
|
||||
def test_stream_completion_controls_publication(store: RedisResponseStore, provider: Provider, truncated: bool) -> None:
|
||||
provider.stream = True
|
||||
provider.truncated = truncated
|
||||
provider.response = b'data: {"choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n'
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
for _ in range(2):
|
||||
result: Final = forward("POST", url, headers=HEADERS, body=BODY, timeout=5)
|
||||
if truncated:
|
||||
assert isinstance(result, NetworkError)
|
||||
else:
|
||||
assert isinstance(result, RawResponse) and result.body == provider.response
|
||||
assert len(provider.hits) == (2 if truncated else 1)
|
||||
|
||||
|
||||
def test_store_outage_preserves_provider_success(provider: Provider) -> None:
|
||||
with socket.socket() as probe:
|
||||
probe.bind(("127.0.0.1", 0))
|
||||
port: Final = probe.getsockname()[1]
|
||||
unavailable: Final = redis_store(f"redis://127.0.0.1:{port}/0", "unavailable")
|
||||
with edge(CacheEdge(unavailable, SECRET), provider) as url:
|
||||
assert call(url).body == SUCCESS
|
||||
assert call(url).body == SUCCESS
|
||||
assert len(provider.hits) == 2
|
||||
|
||||
|
||||
def test_old_lease_cannot_overwrite_new_owner(store: RedisResponseStore) -> None:
|
||||
short: Final = replace(store, lease_ms=50)
|
||||
old: Final = short.lookup("key")
|
||||
assert isinstance(old, CaptureLease)
|
||||
time.sleep(0.08)
|
||||
current: Final = short.lookup("key")
|
||||
assert isinstance(current, CaptureLease)
|
||||
assert not short.publish("key", old, b"old")
|
||||
assert short.publish("key", current, b"new")
|
||||
hit: Final = short.lookup("key")
|
||||
assert isinstance(hit, CacheHit) and hit.payload == b"new"
|
||||
|
||||
|
||||
def test_identity_preserves_values_and_never_contains_credentials() -> None:
|
||||
variants: Final = (b'{}', b'{"a":null}', b'{"a":false}', b'{"a":0}', b'{"a":0.0}', b'{"a":"0"}', b' { }', None, b'')
|
||||
keys: Final = tuple(exact_key(SECRET, "POST", "https://example.invalid/v1/chat/completions", HEADERS, body) for body in variants)
|
||||
assert len(set(keys)) == len(variants)
|
||||
assert all(len(key) == 64 and "synthetic-account" not in key for key in keys)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", [b"corrupt response", '{"response":"{}","signature":"é"}'.encode()])
|
||||
def test_corrupt_entry_is_replaced_by_same_successful_request(store: RedisResponseStore, provider: Provider, payload: bytes) -> None:
|
||||
upstream: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions"
|
||||
prepared: Final = prepare_forward("POST", upstream, HEADERS, BODY)
|
||||
assert isinstance(prepared, PreparedForward)
|
||||
key: Final = exact_key(SECRET, "POST", upstream, prepared.headers, BODY)
|
||||
lease: Final = store.lookup(key)
|
||||
assert isinstance(lease, CaptureLease)
|
||||
assert store.publish(key, lease, payload)
|
||||
cache: Final = CacheEdge(store, SECRET)
|
||||
for _ in range(2):
|
||||
head = cache.forward("POST", upstream, HEADERS, BODY, 5)
|
||||
assert isinstance(head, StreamHead)
|
||||
assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS
|
||||
assert len(provider.hits) == 1
|
||||
assert dict(cache.counters.counts) == {
|
||||
"corrupt": 1, "misses": 1, "upstream_attempts": 1, "writes": 1, "hits": 1,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", [
|
||||
b'data: {}\n\ndata: [DONE]\n\n',
|
||||
b'data: {"choices":[{"index":0,"delta":{}}]}\n\ndata: [DONE]\n\n',
|
||||
b'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\ndata: [DONE]',
|
||||
b'data: {"error":{"message":"failed"}}\n\ndata: [DONE]\n\n',
|
||||
])
|
||||
def test_malformed_success_stream_is_never_cached(store: RedisResponseStore, provider: Provider, payload: bytes) -> None:
|
||||
provider.stream = True
|
||||
provider.response = payload
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
assert call(url).body == payload
|
||||
assert call(url).body == payload
|
||||
assert len(provider.hits) == 2
|
||||
|
||||
|
||||
def test_anthropic_stream_requires_start_finish_and_stop() -> None:
|
||||
start: Final = b'data: {"type":"message_start","message":{}}\n\n'
|
||||
finish: Final = b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"}}\n\n'
|
||||
stop: Final = b'data: {"type":"message_stop"}\n\n'
|
||||
url: Final = "https://example.invalid/v1/messages"
|
||||
headers: Final = {"content-type": "text/event-stream"}
|
||||
assert successful_response(url, 200, headers, start + finish + stop)
|
||||
assert not successful_response(url, 200, headers, start + stop)
|
||||
assert not successful_response(url, 200, headers, finish + stop)
|
||||
assert not successful_response(url, 200, headers, start + finish)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", "")])
|
||||
def test_normal_registration_routes_supported_providers(provider: str, suffix: str) -> None:
|
||||
params: Final = LiteLLMParamsBody(model=f"{provider}/test", api_key="os.environ/SYNTHETIC_KEY", timeout=12)
|
||||
routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True)
|
||||
assert routed.api_base == f"http://edge.invalid/{provider}{suffix}"
|
||||
assert routed.model_dump(exclude={"api_base"}) == params.model_dump(exclude={"api_base"})
|
||||
assert params.api_base is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("params", [
|
||||
LiteLLMParamsBody(model="bedrock/test"),
|
||||
LiteLLMParamsBody(model="azure/test"),
|
||||
LiteLLMParamsBody(model="openai/test", api_base="https://custom.invalid/v1"),
|
||||
LiteLLMParamsBody(model="openai/test", api_base=""),
|
||||
LiteLLMParamsBody(model="openai/test", litellm_credential_name="named-credential"),
|
||||
LiteLLMParamsBody(model="openai/test", mock_response="synthetic"),
|
||||
])
|
||||
def test_registration_preserves_unsupported_or_explicit_routes(params: LiteLLMParamsBody) -> None:
|
||||
def unexpected_edge(mount: str) -> str:
|
||||
pytest.fail(f"should not start edge for {mount}")
|
||||
assert route_cache_model(params, unexpected_edge, enabled=True) is params
|
||||
|
||||
|
||||
def test_rollback_and_live_only_policy_keep_direct_provider_route() -> None:
|
||||
params: Final = LiteLLMParamsBody(model="openai/test")
|
||||
assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=False) is params
|
||||
assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=True, mode="realtime") is params
|
||||
token: Final = LIVE_PROVIDER_REQUIRED.set(True)
|
||||
try:
|
||||
assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=True) is params
|
||||
finally:
|
||||
LIVE_PROVIDER_REQUIRED.reset(token)
|
||||
assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=True).api_base == "http://edge.invalid/v1"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PublishOutage:
|
||||
client: RedisCommands
|
||||
|
||||
def eval(self, script: str, numkeys: int, *args: str | bytes | int) -> object:
|
||||
if script == PUBLISH:
|
||||
raise RedisConnectionError("synthetic publication outage")
|
||||
return self.client.eval(script, numkeys, *args)
|
||||
|
||||
|
||||
def test_write_outage_preserves_success_without_hidden_retry(store: RedisResponseStore, provider: Provider) -> None:
|
||||
unavailable: Final = replace(store, client=PublishOutage(store.client))
|
||||
cache: Final = CacheEdge(unavailable, SECRET)
|
||||
with edge(cache, provider) as url:
|
||||
assert call(url).body == SUCCESS
|
||||
assert call(url).body == SUCCESS
|
||||
assert len(provider.hits) == 2
|
||||
assert dict(cache.counters.counts)["write_failures"] == 2
|
||||
with edge(CacheEdge(store, SECRET), provider) as url:
|
||||
assert call(url).body == SUCCESS
|
||||
assert call(url).body == SUCCESS
|
||||
assert len(provider.hits) == 3
|
||||
|
||||
|
||||
def test_connection_failure_releases_capture_lease(store: RedisResponseStore) -> None:
|
||||
with socket.socket() as unavailable:
|
||||
unavailable.bind(("127.0.0.1", 0))
|
||||
url: Final = f"http://127.0.0.1:{unavailable.getsockname()[1]}/v1/chat/completions"
|
||||
cache: Final = CacheEdge(store, SECRET)
|
||||
assert isinstance(cache.forward("POST", url, HEADERS, BODY, 0.2), NetworkError)
|
||||
prepared: Final = prepare_forward("POST", url, HEADERS, BODY)
|
||||
assert isinstance(prepared, PreparedForward)
|
||||
key: Final = exact_key(SECRET, "POST", url, prepared.headers, BODY)
|
||||
slot: Final = store.lookup(key)
|
||||
assert isinstance(slot, CaptureLease)
|
||||
assert store.release(key, slot)
|
||||
assert dict(cache.counters.counts)["rejected"] == 1
|
||||
|
||||
|
||||
def test_close_before_first_chunk_releases_lease(store: RedisResponseStore, provider: Provider) -> None:
|
||||
url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions"
|
||||
cache: Final = CacheEdge(store, SECRET)
|
||||
head: Final = cache.forward("POST", url, HEADERS, BODY, 5)
|
||||
assert isinstance(head, StreamHead)
|
||||
head.steps.close()
|
||||
prepared: Final = prepare_forward("POST", url, HEADERS, BODY)
|
||||
assert isinstance(prepared, PreparedForward)
|
||||
key: Final = exact_key(SECRET, "POST", url, prepared.headers, BODY)
|
||||
slot: Final = store.lookup(key)
|
||||
assert isinstance(slot, CaptureLease)
|
||||
assert store.release(key, slot)
|
||||
|
||||
|
||||
def test_effective_account_change_cannot_reuse_cache(
|
||||
store: RedisResponseStore, provider: Provider, monkeypatch: pytest.MonkeyPatch, tmp_path,
|
||||
) -> None:
|
||||
url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions"
|
||||
cache: Final = CacheEdge(store, SECRET)
|
||||
for account in ("account-a", "account-b", "account-b"):
|
||||
netrc = tmp_path / account
|
||||
netrc.write_text(f"machine 127.0.0.1 login {account} password synthetic\n")
|
||||
monkeypatch.setenv("NETRC", str(netrc))
|
||||
head = cache.forward("POST", url, HEADERS, BODY, 5)
|
||||
assert isinstance(head, StreamHead)
|
||||
assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS
|
||||
assert len(provider.hits) == 2
|
||||
assert dict(cache.counters.counts)["hits"] == 1
|
||||
|
||||
|
||||
def test_enabled_environment_reuses_store_across_fresh_backends(
|
||||
redis_url: str, provider: Provider, monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("E2E_PROVIDER_CACHE", "1")
|
||||
monkeypatch.setenv("E2E_PROVIDER_CACHE_REDIS_URL", redis_url)
|
||||
monkeypatch.setenv("E2E_PROVIDER_CACHE_HMAC_KEY", SECRET.decode())
|
||||
monkeypatch.setenv("E2E_PROVIDER_CACHE_NAMESPACE", "environment-" + uuid.uuid4().hex)
|
||||
configured_cache.cache_clear()
|
||||
try:
|
||||
for _ in range(2):
|
||||
backend = configured_cache_backend()
|
||||
assert isinstance(backend, CacheEdge)
|
||||
with edge(backend, provider) as url:
|
||||
assert call(url).body == SUCCESS
|
||||
configured_cache.cache_clear()
|
||||
assert len(provider.hits) == 1
|
||||
monkeypatch.setenv("E2E_PROVIDER_CACHE", "0")
|
||||
assert configured_cache_backend() is None
|
||||
finally:
|
||||
configured_cache.cache_clear()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("known_mount", (True, False))
|
||||
def test_duplicate_headers_bypass_cache_and_count_live_calls(
|
||||
store: RedisResponseStore, provider: Provider, known_mount: bool,
|
||||
) -> None:
|
||||
cache: Final = CacheEdge(store, SECRET)
|
||||
with edge(cache, provider) as url:
|
||||
parsed: Final = urlsplit(url)
|
||||
for _ in range(2):
|
||||
connection = HTTPConnection(str(parsed.hostname), parsed.port, timeout=5)
|
||||
try:
|
||||
connection.putrequest("POST", parsed.path if known_mount else "/unknown/v1/chat/completions")
|
||||
connection.putheader("content-length", str(len(BODY)))
|
||||
connection.putheader("content-type", "application/json")
|
||||
connection.putheader("x-duplicate", "first")
|
||||
connection.putheader("x-duplicate", "second")
|
||||
connection.endheaders(BODY)
|
||||
response = connection.getresponse()
|
||||
assert response.status == (200 if known_mount else 404)
|
||||
payload = response.read()
|
||||
assert payload == SUCCESS if known_mount else b"unknown provider mount" in payload
|
||||
finally:
|
||||
connection.close()
|
||||
assert len(provider.hits) == (2 if known_mount else 0)
|
||||
assert dict(cache.counters.counts)["duplicate_header_bypass"] == 2
|
||||
assert dict(cache.counters.counts).get("upstream_attempts", 0) == (2 if known_mount else 0)
|
||||
33
tests/e2e/PROVIDER_CACHE.md
Normal file
33
tests/e2e/PROVIDER_CACHE.md
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
# Shared provider-response cache
|
||||
|
||||
`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live
|
||||
|
||||
The edge caches complete successful POST responses for `/v1/chat/completions` and `/v1/messages`, including streams. Unsupported endpoints pass through. It matches the method, original URL, effective outbound headers (including authentication and HTTP-library defaults), body presence and exact body bytes using a full keyed digest. It sends the same prepared request used for matching. No prompts, random markers, JSON values or credentials are normalized away. Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies
|
||||
|
||||
An eligible miss calls the provider. A complete successful response is stored immediately even if a later test assertion fails. Provider errors, malformed responses, truncated streams and cancelled captures are not stored. Cache reads, writes and lease failures fall through to normal provider behavior; they introduce no provider retry. An already-started response cannot be restarted after a delivery failure
|
||||
|
||||
Recordings are shared across workers and builds through dedicated Redis, separate from the candidate's own cache. They expire 86,400 seconds after capture starts, based on Redis time. Reads never extend expiry. There is no scheduled recapture: the next miss calls the provider again. Bounded coordination reduces duplicate concurrent calls, but slow or failed captures may lead to extra live calls after the wait expires
|
||||
|
||||
## Configuration
|
||||
|
||||
The trusted runner receives:
|
||||
|
||||
- `E2E_PROVIDER_CACHE`: `1` to enable, `0` to use the normal live path
|
||||
- `E2E_PROVIDER_CACHE_REDIS_URL`: authenticated dedicated Redis URL
|
||||
- `E2E_PROVIDER_CACHE_HMAC_KEY`: dedicated secret containing at least 32 bytes
|
||||
- `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision
|
||||
- `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory
|
||||
|
||||
Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits
|
||||
|
||||
Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay
|
||||
|
||||
## Recorded response semantics
|
||||
|
||||
Replay preserves the original response ID, usage and end-to-end headers. The proxy can therefore deduplicate repeated provider IDs when storing spend-log rows, just as it does when a live upstream returns the same ID twice. One spend-log row per invocation is not guaranteed for identical recorded responses. Existing spend reconciliation requests use distinct prompt markers and retain their distinct-ID and row-count assertions; accounting tests are not automatically excluded from caching
|
||||
|
||||
Provider remaining-quota headers describe the captured response. Metrics derived from them are historical on a cache hit, not a measurement of current provider capacity. Gateway-generated API-key quota headers are a separate contract. A test of fresh provider quota or timing must use the live-provider policy; replay can still exercise how the proxy processes the recorded headers
|
||||
|
||||
## Qualification
|
||||
|
||||
`tests/code_coverage_tests/test_provider_cache.py` exercises local HTTP providers and disposable real Redis. CI runs these checks with the existing provider-edge and replay harness tests. These component checks do not establish Buildkite deployment, full-suite cross-build reuse or a genuine 24-hour expiry observation; those require separate runtime evidence
|
||||
|
|
@ -22,7 +22,6 @@ from typing import Final
|
|||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from e2e_config import (
|
||||
CONTROL_PLANE_BASE_URL,
|
||||
FIXTURE_DIR,
|
||||
|
|
@ -41,6 +40,7 @@ from idp import Identity, Keycloak, keycloak_from_env
|
|||
from junit_properties import attach_result_properties
|
||||
from lifecycle import ProxyClientProvider, ResourceManager
|
||||
from models import TeamNewBody, UserNewBody, UserNewResponse
|
||||
from provider_cache_routing import LIVE_PROVIDER_REQUIRED
|
||||
from provider_edge import replay_leftover_error
|
||||
from proxy_client import ProxyClient, build_proxy_client
|
||||
|
||||
|
|
@ -85,6 +85,7 @@ def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient)
|
|||
|
||||
|
||||
def pytest_configure(config: pytest.Config) -> None:
|
||||
config.addinivalue_line("markers", "provider_live: requires actual provider timing, limits or state; bypass shared cache")
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"e2e: live test that requires a running proxy and real provider keys",
|
||||
|
|
@ -192,11 +193,13 @@ def _proxy_fail_reason() -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
@pytest.hookimpl(tryfirst=True)
|
||||
def pytest_runtest_setup(item: pytest.Item) -> None:
|
||||
"""Hard-fail `e2e`-marked tests unless a proxy answers its liveness probe.
|
||||
Unmarked tests (unit coverage of the harness) don't touch the proxy, so they
|
||||
run even when none is up. Never skip for a missing proxy. Replay mode needs
|
||||
the proxy too: only provider-bound traffic replays from the bundle."""
|
||||
LIVE_PROVIDER_REQUIRED.set(item.get_closest_marker("provider_live") is not None)
|
||||
if item.get_closest_marker("e2e") is None:
|
||||
return
|
||||
reason = _proxy_fail_reason()
|
||||
|
|
@ -235,6 +238,7 @@ def pytest_runtest_teardown(item: pytest.Item) -> Generator[None, None, None]:
|
|||
yield so fixture finalizers replay their recorded calls first. Failed tests
|
||||
are left alone - their own failure already explains any unconsumed tail."""
|
||||
result = yield
|
||||
LIVE_PROVIDER_REQUIRED.set(False)
|
||||
if not item.stash.get(_CALL_PASSED, False):
|
||||
return result
|
||||
reason = replay_leftover_error(
|
||||
|
|
|
|||
|
|
@ -853,6 +853,7 @@ def _stream_steps(resp: requests.Response) -> Generator[StreamStep, None, None]:
|
|||
the chunks already delivered are exactly what makes a mid-stream failure
|
||||
different from a request that never streamed at all."""
|
||||
try:
|
||||
yield StreamChunk(b"")
|
||||
for piece in cast("Iterator[bytes]", resp.iter_content(chunk_size=None)):
|
||||
if piece:
|
||||
yield StreamChunk(data=piece)
|
||||
|
|
@ -862,6 +863,44 @@ def _stream_steps(resp: requests.Response) -> Generator[StreamStep, None, None]:
|
|||
resp.close()
|
||||
|
||||
|
||||
def primed_steps(steps: Generator[StreamStep, None, None]) -> Generator[StreamStep, None, None]:
|
||||
first: Final = next(steps)
|
||||
assert isinstance(first, StreamChunk) and first.data == b""
|
||||
return steps
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, repr=False)
|
||||
class PreparedForward:
|
||||
request: requests.PreparedRequest
|
||||
url: str
|
||||
headers: dict[str, str]
|
||||
|
||||
|
||||
def prepare_forward(
|
||||
method: str, url: str, headers: dict[str, str], body: bytes | None,
|
||||
) -> PreparedForward | NetworkError:
|
||||
try:
|
||||
with requests.Session() as session:
|
||||
request: Final = session.prepare_request(requests.Request(method, url, headers=headers, data=body))
|
||||
except requests.RequestException as exc:
|
||||
return NetworkError(message=str(exc))
|
||||
assert request.url is not None
|
||||
return PreparedForward(request, request.url, dict(request.headers))
|
||||
|
||||
|
||||
def forward_prepared_stream(prepared: PreparedForward, timeout: float) -> StreamHead | NetworkError:
|
||||
try:
|
||||
with requests.Session() as session:
|
||||
settings: Final = session.merge_environment_settings(prepared.url, {}, True, None, None)
|
||||
resp: Final = session.send(prepared.request, timeout=timeout, allow_redirects=False, **settings)
|
||||
except requests.RequestException as exc:
|
||||
return NetworkError(message=str(exc))
|
||||
return StreamHead(
|
||||
resp.status_code, {name.lower(): value for name, value in resp.headers.items()},
|
||||
primed_steps(_stream_steps(resp)),
|
||||
)
|
||||
|
||||
|
||||
def open_stream(url: URL, *, headers: BaseModel, json: BaseModel, timeout: float = 60.0) -> StreamHead | NetworkError:
|
||||
"""POST a streaming request and return the moment its response head arrives,
|
||||
leaving the body unread behind ``StreamHead.steps``. For a test that must keep
|
||||
|
|
@ -907,5 +946,5 @@ def forward_stream(
|
|||
return StreamHead(
|
||||
status_code=resp.status_code,
|
||||
headers={name.lower(): value for name, value in resp.headers.items()},
|
||||
steps=_stream_steps(resp),
|
||||
steps=primed_steps(_stream_steps(resp)),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -44,7 +44,7 @@ from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody, Usage
|
|||
from passthrough_client import PassthroughClient
|
||||
import os
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
pytestmark = [pytest.mark.e2e, pytest.mark.provider_live]
|
||||
|
||||
BEDROCK_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
VERTEX_MODEL = "vertex_ai/gemini-2.5-flash"
|
||||
|
|
|
|||
|
|
@ -24,10 +24,10 @@ from models import (
|
|||
AnthropicAssistantTurn,
|
||||
AnthropicContentBlock,
|
||||
AnthropicCustomTool,
|
||||
AnthropicMessagesBody,
|
||||
AnthropicToolChoice,
|
||||
AnthropicToolResultBlock,
|
||||
AnthropicToolResultTurn,
|
||||
AnthropicMessagesBody,
|
||||
ChatMessage,
|
||||
JsonSchemaProperty,
|
||||
LiteLLMParamsBody,
|
||||
|
|
@ -165,6 +165,7 @@ class TestAnthropicMessages:
|
|||
)
|
||||
|
||||
@pytest.mark.covers("llm.messages.anthropic.basic.stream.works")
|
||||
@pytest.mark.provider_live
|
||||
def test_messages_streams_completion(
|
||||
self, endpoints_client: EndpointsClient, resources: ResourceManager
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -744,6 +744,7 @@ class TestTogetherMessages:
|
|||
assert "22" in text, f"the model never saw the tool result: {response.content}"
|
||||
|
||||
@pytest.mark.covers("llm.messages.together_ai.basic.stream.works")
|
||||
@pytest.mark.provider_live
|
||||
def test_streams_text_deltas(
|
||||
self, client: PassthroughClient, resources: ResourceManager, reasoning_tool_backend: str
|
||||
) -> None:
|
||||
|
|
|
|||
299
tests/e2e/provider_cache.py
Normal file
299
tests/e2e/provider_cache.py
Normal file
|
|
@ -0,0 +1,299 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import io
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Generator, Mapping
|
||||
from contextlib import closing
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Final, Literal, Protocol
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from e2e_http import (
|
||||
NetworkError,
|
||||
StreamChunk,
|
||||
StreamHead,
|
||||
StreamStep,
|
||||
StreamTruncation,
|
||||
forward_prepared_stream,
|
||||
forward_stream,
|
||||
prepare_forward,
|
||||
primed_steps,
|
||||
)
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
LIFETIME_SECONDS: Final = 86_400
|
||||
MAX_REQUEST_BYTES: Final = 256 * 1024
|
||||
MAX_RESPONSE_BYTES: Final = 8 * 1024 * 1024
|
||||
UNRECORDED_RESPONSE_HEADERS: Final = frozenset({"set-cookie"})
|
||||
JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CacheHit:
|
||||
payload: bytes
|
||||
valid_until: float
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CaptureLease:
|
||||
token: str
|
||||
captured_at_ms: int
|
||||
expires_at_ms: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CacheBusy:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CacheUnavailable:
|
||||
pass
|
||||
|
||||
|
||||
type CacheLookup = CacheHit | CaptureLease | CacheBusy | CacheUnavailable
|
||||
|
||||
|
||||
class ResponseStore(Protocol):
|
||||
def lookup(self, key: str) -> CacheLookup: ...
|
||||
|
||||
def publish(self, key: str, lease: CaptureLease, payload: bytes) -> bool: ...
|
||||
|
||||
def release(self, key: str, lease: CaptureLease) -> bool: ...
|
||||
|
||||
def discard(self, key: str, payload: bytes) -> bool: ...
|
||||
|
||||
|
||||
class CachedResponse(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", strict=True)
|
||||
format_version: Literal[1] = 1
|
||||
request_key: str
|
||||
status_code: int
|
||||
headers: dict[str, str]
|
||||
chunks: tuple[str, ...]
|
||||
|
||||
|
||||
class SignedResponse(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid", strict=True)
|
||||
response: str
|
||||
signature: str
|
||||
|
||||
|
||||
def exact_key(secret: bytes, method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> str:
|
||||
fields: Final = (
|
||||
b"provider-cache-exact-v1", method.encode(), url.encode(),
|
||||
*(part.encode() for pair in sorted(headers.items()) for part in pair),
|
||||
b"no-body" if body is None else b"body", b"" if body is None else body,
|
||||
)
|
||||
encoded: Final = b"".join(len(part).to_bytes(8, "big") + part for part in fields)
|
||||
return hmac.new(secret, encoded, hashlib.sha256).hexdigest()
|
||||
|
||||
|
||||
def cacheable_endpoint(method: str, url: str, body: bytes | None) -> bool:
|
||||
return (
|
||||
method == "POST"
|
||||
and urlsplit(url).path in {"/v1/chat/completions", "/v1/messages"}
|
||||
and body is not None
|
||||
and len(body) <= MAX_REQUEST_BYTES
|
||||
)
|
||||
|
||||
|
||||
def successful_response(url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool:
|
||||
if not 200 <= status < 300 or len(body) > MAX_RESPONSE_BYTES:
|
||||
return False
|
||||
streaming: Final = "text/event-stream" in headers.get("content-type", "").lower()
|
||||
if streaming:
|
||||
try:
|
||||
text: Final = body.decode("utf-8").replace("\r\n", "\n")
|
||||
if not text.endswith("\n\n"):
|
||||
return False
|
||||
events: Final = tuple(
|
||||
"\n".join(line[5:].removeprefix(" ") for line in event.split("\n") if line.startswith("data:"))
|
||||
for event in text.split("\n\n") if any(line.startswith("data:") for line in event.split("\n"))
|
||||
)
|
||||
values: Final = tuple(JSON_VALUE.validate_json(event) for event in events if event != "[DONE]")
|
||||
except (UnicodeDecodeError, ValidationError):
|
||||
return False
|
||||
if not values or any(not isinstance(value, dict) or "error" in value or value.get("type") == "error" for value in values):
|
||||
return False
|
||||
if urlsplit(url).path == "/v1/chat/completions":
|
||||
return events[-1] == "[DONE]" and "[DONE]" not in events[:-1] and complete_chat_stream(values)
|
||||
return (
|
||||
"[DONE]" not in events
|
||||
and isinstance(values[0], dict) and values[0].get("type") == "message_start"
|
||||
and isinstance(values[-1], dict) and values[-1].get("type") == "message_stop"
|
||||
and any(
|
||||
isinstance(value, dict) and value.get("type") == "message_delta"
|
||||
and isinstance(delta := value.get("delta"), dict) and isinstance(delta.get("stop_reason"), str)
|
||||
for value in values
|
||||
)
|
||||
)
|
||||
try:
|
||||
value: Final = JSON_VALUE.validate_json(body)
|
||||
except ValidationError:
|
||||
return False
|
||||
if not isinstance(value, dict) or "error" in value:
|
||||
return False
|
||||
if urlsplit(url).path == "/v1/messages":
|
||||
return value.get("type") == "message" and isinstance(value.get("content"), list) and isinstance(value.get("stop_reason"), str)
|
||||
choices: Final = value.get("choices")
|
||||
return isinstance(choices, list) and bool(choices) and all(
|
||||
isinstance(choice, dict) and isinstance(choice.get("message"), dict) and isinstance(choice.get("finish_reason"), str)
|
||||
for choice in choices
|
||||
)
|
||||
|
||||
|
||||
def complete_chat_stream(values: tuple[JsonValue, ...]) -> bool:
|
||||
if any(not isinstance(value, dict) or not isinstance(value.get("choices"), list) for value in values):
|
||||
return False
|
||||
choices: Final = tuple(
|
||||
choice for value in values if isinstance(value, dict)
|
||||
if isinstance(items := value.get("choices"), list) for choice in items
|
||||
)
|
||||
if not choices or any(
|
||||
not isinstance(choice, dict) or type(choice.get("index")) is not int
|
||||
or not isinstance(choice.get("delta"), dict)
|
||||
for choice in choices
|
||||
):
|
||||
return False
|
||||
indices: Final = frozenset(choice["index"] for choice in choices if isinstance(choice, dict))
|
||||
return all(
|
||||
isinstance(tuple(choice for choice in choices if isinstance(choice, dict) and choice["index"] == index)[-1].get("finish_reason"), str)
|
||||
for index in indices
|
||||
)
|
||||
|
||||
|
||||
def encode_response(secret: bytes, response: CachedResponse) -> bytes:
|
||||
raw: Final = response.model_dump_json()
|
||||
return SignedResponse(response=raw, signature=hmac.new(secret, raw.encode(), hashlib.sha256).hexdigest()).model_dump_json().encode()
|
||||
|
||||
|
||||
def decode_response(secret: bytes, key: str, payload: bytes, url: str) -> CachedResponse | None:
|
||||
if len(payload) > 2 * MAX_RESPONSE_BYTES:
|
||||
return None
|
||||
try:
|
||||
signed: Final = SignedResponse.model_validate_json(payload)
|
||||
if not hmac.compare_digest(signed.signature.encode(), hmac.new(secret, signed.response.encode(), hashlib.sha256).hexdigest().encode()):
|
||||
return None
|
||||
response: Final = CachedResponse.model_validate_json(signed.response)
|
||||
chunks: Final = tuple(base64.b64decode(chunk, validate=True) for chunk in response.chunks)
|
||||
except (ValidationError, ValueError):
|
||||
return None
|
||||
if response.request_key != key or not successful_response(url, response.status_code, response.headers, b"".join(chunks)):
|
||||
return None
|
||||
return response
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class CacheCounters:
|
||||
counts: tuple[tuple[str, int], ...] = ()
|
||||
lock: threading.Lock = field(default_factory=threading.Lock)
|
||||
|
||||
def increment(self, name: str) -> None:
|
||||
with self.lock:
|
||||
current: Final = dict(self.counts)
|
||||
self.counts = tuple((current | {name: current.get(name, 0) + 1}).items())
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ResponseCapture:
|
||||
buffer: io.BytesIO = field(default_factory=io.BytesIO)
|
||||
size: int = 0
|
||||
eligible: bool = True
|
||||
|
||||
def observe(self, step: StreamStep) -> None:
|
||||
if not self.eligible:
|
||||
return
|
||||
if isinstance(step, StreamTruncation) or self.size + len(step.data) + 8 > MAX_RESPONSE_BYTES:
|
||||
self.eligible = False
|
||||
self.buffer.close()
|
||||
return
|
||||
self.buffer.write(len(step.data).to_bytes(8, "big"))
|
||||
self.buffer.write(step.data)
|
||||
self.size += len(step.data) + 8
|
||||
|
||||
def chunks(self) -> tuple[bytes, ...]:
|
||||
self.buffer.seek(0)
|
||||
return tuple(self.buffer.read(int.from_bytes(size, "big")) for size in iter(lambda: self.buffer.read(8), b""))
|
||||
|
||||
|
||||
def response_steps(response: CachedResponse) -> Generator[StreamStep, None, None]:
|
||||
for chunk in response.chunks:
|
||||
yield StreamChunk(base64.b64decode(chunk, validate=True))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CacheEdge:
|
||||
store: ResponseStore
|
||||
secret: bytes = field(repr=False)
|
||||
counters: CacheCounters = field(default_factory=CacheCounters)
|
||||
wait_seconds: float = 2.0
|
||||
clock: Callable[[], float] = time.monotonic
|
||||
sleep: Callable[[float], None] = time.sleep
|
||||
|
||||
def lookup(self, key: str) -> CacheLookup:
|
||||
deadline: Final = self.clock() + self.wait_seconds
|
||||
while isinstance(result := self.store.lookup(key), CacheBusy) and self.clock() < deadline:
|
||||
self.sleep(min(0.05, max(0, deadline - self.clock())))
|
||||
return result
|
||||
|
||||
def forward(self, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float) -> StreamHead | NetworkError:
|
||||
if not cacheable_endpoint(method, url, body):
|
||||
self.counters.increment("bypass")
|
||||
self.counters.increment("upstream_attempts")
|
||||
return forward_stream(method, url, headers=headers, body=body, timeout=timeout)
|
||||
prepared: Final = prepare_forward(method, url, headers, body)
|
||||
if isinstance(prepared, NetworkError):
|
||||
self.counters.increment("rejected")
|
||||
return prepared
|
||||
key: Final = exact_key(self.secret, method, url, prepared.headers, body)
|
||||
found: Final = self.lookup(key)
|
||||
if isinstance(found, CacheHit):
|
||||
response: Final = decode_response(self.secret, key, found.payload, url)
|
||||
if response is not None and self.clock() < found.valid_until:
|
||||
self.counters.increment("hits")
|
||||
return StreamHead(response.status_code, response.headers, response_steps(response))
|
||||
self.counters.increment("corrupt" if response is None else "expired")
|
||||
self.store.discard(key, found.payload)
|
||||
capture_slot: Final = self.lookup(key) if isinstance(found, CacheHit) else found
|
||||
self.counters.increment("misses")
|
||||
if isinstance(capture_slot, CacheUnavailable):
|
||||
self.counters.increment("cache_errors")
|
||||
self.counters.increment("upstream_attempts")
|
||||
head: Final = forward_prepared_stream(prepared, timeout)
|
||||
if not isinstance(capture_slot, CaptureLease):
|
||||
return head
|
||||
if isinstance(head, NetworkError):
|
||||
self.store.release(key, capture_slot)
|
||||
self.counters.increment("rejected")
|
||||
return head
|
||||
return StreamHead(head.status_code, head.headers, primed_steps(self.capture(key, capture_slot, url, head)))
|
||||
|
||||
def capture(self, key: str, lease: CaptureLease, url: str, head: StreamHead) -> Generator[StreamStep, None, None]:
|
||||
capture: Final = ResponseCapture()
|
||||
try:
|
||||
with closing(head.steps):
|
||||
yield StreamChunk(b"")
|
||||
for step in head.steps:
|
||||
yield step
|
||||
capture.observe(step)
|
||||
chunks: Final = capture.chunks() if capture.eligible else ()
|
||||
headers: Final = {
|
||||
name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS
|
||||
}
|
||||
if not capture.eligible or not successful_response(url, head.status_code, headers, b"".join(chunks)):
|
||||
self.counters.increment("rejected")
|
||||
return
|
||||
response: Final = CachedResponse(
|
||||
request_key=key, status_code=head.status_code, headers=headers,
|
||||
chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks),
|
||||
)
|
||||
published: Final = self.store.publish(key, lease, encode_response(self.secret, response))
|
||||
self.counters.increment("writes" if published else "write_failures")
|
||||
finally:
|
||||
self.store.release(key, lease)
|
||||
capture.buffer.close()
|
||||
154
tests/e2e/provider_cache_redis.py
Normal file
154
tests/e2e/provider_cache_redis.py
Normal file
|
|
@ -0,0 +1,154 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import functools
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final, Protocol, cast
|
||||
|
||||
from provider_cache import LIFETIME_SECONDS, CacheBusy, CacheEdge, CacheHit, CacheLookup, CacheUnavailable, CaptureLease
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from redis import Redis
|
||||
from redis.exceptions import RedisError
|
||||
|
||||
REDIS_ARRAY: Final[TypeAdapter[list[bytes]]] = TypeAdapter(list[bytes])
|
||||
|
||||
LOOKUP: Final = """
|
||||
local clock = redis.call('TIME')
|
||||
local now = clock[1] * 1000 + math.floor(clock[2] / 1000)
|
||||
local row = redis.call('HMGET', KEYS[1], 'captured', 'expires', 'payload')
|
||||
if row[3] then
|
||||
local captured = tonumber(row[1])
|
||||
local expires = tonumber(row[2])
|
||||
if captured and expires and captured <= now and expires > now
|
||||
and expires - captured == tonumber(ARGV[2]) then
|
||||
return {'hit', row[3], tostring(expires - now)}
|
||||
end
|
||||
redis.call('DEL', KEYS[1])
|
||||
end
|
||||
if redis.call('SET', KEYS[2], ARGV[1], 'NX', 'PX', ARGV[3]) then
|
||||
return {'lease', tostring(now), tostring(now + tonumber(ARGV[2]))}
|
||||
end
|
||||
return {'busy'}
|
||||
"""
|
||||
|
||||
PUBLISH: Final = """
|
||||
if redis.call('GET', KEYS[2]) ~= ARGV[1] then return 0 end
|
||||
local clock = redis.call('TIME')
|
||||
local now = clock[1] * 1000 + math.floor(clock[2] / 1000)
|
||||
local captured = tonumber(ARGV[2])
|
||||
local expires = tonumber(ARGV[3])
|
||||
if captured > now or expires <= now or expires - captured ~= tonumber(ARGV[5]) then return 0 end
|
||||
if redis.call('EXISTS', KEYS[1]) == 1 then return 0 end
|
||||
redis.call('HSET', KEYS[1], 'captured', ARGV[2], 'expires', ARGV[3], 'payload', ARGV[4])
|
||||
redis.call('PEXPIREAT', KEYS[1], expires)
|
||||
redis.call('DEL', KEYS[2])
|
||||
return 1
|
||||
"""
|
||||
|
||||
RELEASE: Final = """
|
||||
if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end
|
||||
return redis.call('DEL', KEYS[1])
|
||||
"""
|
||||
|
||||
DISCARD: Final = """
|
||||
if redis.call('HGET', KEYS[1], 'payload') ~= ARGV[1] then return 0 end
|
||||
return redis.call('DEL', KEYS[1])
|
||||
"""
|
||||
|
||||
|
||||
class RedisCommands(Protocol):
|
||||
def eval(self, script: str, numkeys: int, *args: str | bytes | int) -> object: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RedisResponseStore:
|
||||
client: RedisCommands
|
||||
namespace: str
|
||||
lifetime_ms: int = LIFETIME_SECONDS * 1000
|
||||
lease_ms: int = 120_000
|
||||
|
||||
def keys(self, key: str) -> tuple[str, str]:
|
||||
prefix: Final = f"e2e-provider-cache:v1:{self.namespace}:{{{key}}}"
|
||||
return prefix + ":response", prefix + ":lease"
|
||||
|
||||
def lookup(self, key: str) -> CacheLookup:
|
||||
token: Final = uuid.uuid4().hex
|
||||
started: Final = time.monotonic()
|
||||
try:
|
||||
result: Final = self.client.eval(LOOKUP, 2, *self.keys(key), token, self.lifetime_ms, self.lease_ms)
|
||||
except (RedisError, OSError):
|
||||
return CacheUnavailable()
|
||||
try:
|
||||
parts: Final = tuple(REDIS_ARRAY.validate_python(result, strict=True))
|
||||
except ValidationError:
|
||||
return CacheUnavailable()
|
||||
if len(parts) == 3 and parts[0] == b"hit" and parts[2].isdigit():
|
||||
return CacheHit(parts[1], started + int(parts[2]) / 1000)
|
||||
if len(parts) == 3 and parts[0] == b"lease" and parts[1].isdigit() and parts[2].isdigit():
|
||||
return CaptureLease(token, int(parts[1]), int(parts[2]))
|
||||
if parts == (b"busy",):
|
||||
return CacheBusy()
|
||||
return CacheUnavailable()
|
||||
|
||||
def publish(self, key: str, lease: CaptureLease, payload: bytes) -> bool:
|
||||
try:
|
||||
result: Final = self.client.eval(
|
||||
PUBLISH, 2, *self.keys(key), lease.token, lease.captured_at_ms, lease.expires_at_ms, payload, self.lifetime_ms,
|
||||
)
|
||||
except (RedisError, OSError):
|
||||
return False
|
||||
return result == 1
|
||||
|
||||
def release(self, key: str, lease: CaptureLease) -> bool:
|
||||
try:
|
||||
result: Final = self.client.eval(RELEASE, 1, self.keys(key)[1], lease.token)
|
||||
except (RedisError, OSError):
|
||||
return False
|
||||
return result == 1
|
||||
|
||||
def discard(self, key: str, payload: bytes) -> bool:
|
||||
try:
|
||||
result: Final = self.client.eval(DISCARD, 1, self.keys(key)[0], payload)
|
||||
except (RedisError, OSError):
|
||||
return False
|
||||
return result == 1
|
||||
|
||||
|
||||
def redis_store(url: str, namespace: str) -> RedisResponseStore:
|
||||
client: Final = Redis.from_url(url, socket_timeout=0.25, socket_connect_timeout=0.25, decode_responses=False)
|
||||
return RedisResponseStore(cast(RedisCommands, client), namespace)
|
||||
|
||||
|
||||
def write_metrics(cache: CacheEdge) -> None:
|
||||
report: Final = json.dumps({"provider_cache": dict(cache.counters.counts)})
|
||||
directory: Final = os.environ.get("E2E_PROVIDER_CACHE_METRICS_DIR")
|
||||
if directory:
|
||||
try:
|
||||
root: Final = Path(directory)
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
(root / f"{os.getpid()}.json").write_text(report + "\n")
|
||||
except OSError:
|
||||
logging.getLogger(__name__).warning("provider cache metrics artifact unavailable")
|
||||
logging.getLogger(__name__).info("%s", report)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def configured_cache() -> CacheEdge | None:
|
||||
if os.environ.get("E2E_PROVIDER_CACHE", "0") == "0":
|
||||
return None
|
||||
if os.environ.get("E2E_PROVIDER_CACHE") != "1":
|
||||
raise ValueError("E2E_PROVIDER_CACHE must be 0 or 1")
|
||||
secret: Final = os.environ.get("E2E_PROVIDER_CACHE_HMAC_KEY", "").encode()
|
||||
namespace: Final = os.environ.get("E2E_PROVIDER_CACHE_NAMESPACE", "")
|
||||
if len(secret) < 32 or re.fullmatch(r"[a-zA-Z0-9_-]{1,64}", namespace) is None:
|
||||
raise ValueError("provider cache requires a dedicated key and namespace")
|
||||
cache: Final = CacheEdge(redis_store(os.environ["E2E_PROVIDER_CACHE_REDIS_URL"], namespace), secret)
|
||||
atexit.register(write_metrics, cache)
|
||||
return cache
|
||||
23
tests/e2e/provider_cache_routing.py
Normal file
23
tests/e2e/provider_cache_routing.py
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from contextvars import ContextVar
|
||||
from typing import Final
|
||||
|
||||
from models import LiteLLMParamsBody, ModelMode
|
||||
|
||||
LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_required", default=False)
|
||||
|
||||
|
||||
def route_cache_model(
|
||||
params: LiteLLMParamsBody, base_for: Callable[[str], str | None], *, enabled: bool, mode: ModelMode | None = None,
|
||||
) -> LiteLLMParamsBody:
|
||||
if not enabled or mode == "realtime" or LIVE_PROVIDER_REQUIRED.get() or params.api_base is not None or params.mock_response is not None:
|
||||
return params
|
||||
provider: Final = params.model.partition("/")[0]
|
||||
if provider not in {"openai", "anthropic"} or params.litellm_credential_name is not None:
|
||||
return params
|
||||
base: Final = base_for(provider)
|
||||
if base is None:
|
||||
return params
|
||||
return params.model_copy(update={"api_base": f"{base}/v1" if provider == "openai" else base})
|
||||
|
|
@ -42,6 +42,7 @@ import base64
|
|||
import difflib
|
||||
import functools
|
||||
import hashlib
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
from collections import deque
|
||||
|
|
@ -93,6 +94,8 @@ from fixture_mode import (
|
|||
parse_fixture_mode,
|
||||
)
|
||||
from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity
|
||||
from provider_cache import CacheEdge
|
||||
from provider_cache_routing import LIVE_PROVIDER_REQUIRED
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType(
|
||||
|
|
@ -506,7 +509,7 @@ class LiveEdge:
|
|||
pass
|
||||
|
||||
|
||||
type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge
|
||||
type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge | CacheEdge
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
|
|
@ -750,12 +753,16 @@ def _handle_record(
|
|||
|
||||
|
||||
def _handle_live(
|
||||
method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float
|
||||
method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float,
|
||||
cache: CacheEdge | None = None,
|
||||
) -> EdgeOutcome:
|
||||
forwarded: Final = {
|
||||
name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS
|
||||
}
|
||||
head: Final = forward_stream(method, url, headers=forwarded, body=body, timeout=timeout)
|
||||
head: Final = (
|
||||
forward_stream(method, url, headers=forwarded, body=body, timeout=timeout)
|
||||
if cache is None else cache.forward(method, url, forwarded, body, timeout)
|
||||
)
|
||||
match head:
|
||||
case NetworkError(message=message):
|
||||
return _recorded_outcome(_network_error_response(message))
|
||||
|
|
@ -821,6 +828,10 @@ def handle_edge_request(
|
|||
else edge_request(method, split.path, split.query, body, _header_value(headers, "content-type"))
|
||||
)
|
||||
match backend:
|
||||
case CacheEdge():
|
||||
return _handle_live(
|
||||
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, backend,
|
||||
)
|
||||
case LiveEdge():
|
||||
return _handle_live(
|
||||
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout
|
||||
|
|
@ -871,11 +882,19 @@ class _EdgeHandler(BaseHTTPRequestHandler):
|
|||
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):
|
||||
if strict and len({name.lower() for name in self.headers}) != len(self.headers):
|
||||
self._write_reply(_text_reply(REPLAY_MISS_STATUS, "stateless_v1 eligibility error: duplicate headers"))
|
||||
return
|
||||
duplicate_headers: Final = len({name.lower() for name in self.headers}) != len(self.headers)
|
||||
selected_backend: Final = (
|
||||
LiveEdge() if isinstance(edge_server.backend, CacheEdge) and duplicate_headers else edge_server.backend
|
||||
)
|
||||
if isinstance(edge_server.backend, CacheEdge) and duplicate_headers:
|
||||
edge_server.backend.counters.increment("duplicate_header_bypass")
|
||||
if urlsplit(self.path).path.lstrip("/").partition("/")[0] in edge_server.mounts:
|
||||
edge_server.backend.counters.increment("upstream_attempts")
|
||||
outcome: Final = handle_edge_request(
|
||||
edge_server.backend,
|
||||
selected_backend,
|
||||
edge_server.mounts,
|
||||
self.command,
|
||||
self.path,
|
||||
|
|
@ -908,12 +927,12 @@ class _EdgeHandler(BaseHTTPRequestHandler):
|
|||
shuts down write-side first: the proxy sees a graceful close mid-message,
|
||||
which is the incomplete chunked read a provider hanging up produces, and not
|
||||
the reset that could discard the chunks already in flight."""
|
||||
self.send_response(stream.status_code)
|
||||
for name, value in stream.headers.items():
|
||||
self.send_header(name, value)
|
||||
self.send_header("transfer-encoding", "chunked")
|
||||
self.end_headers()
|
||||
with closing(stream.steps) as steps:
|
||||
self.send_response(stream.status_code)
|
||||
for name, value in stream.headers.items():
|
||||
self.send_header(name, value)
|
||||
self.send_header("transfer-encoding", "chunked")
|
||||
self.end_headers()
|
||||
for step in steps:
|
||||
match step:
|
||||
case StreamChunk(data=data):
|
||||
|
|
@ -923,7 +942,7 @@ class _EdgeHandler(BaseHTTPRequestHandler):
|
|||
return
|
||||
case _:
|
||||
assert_never(step)
|
||||
self.wfile.write(b"0\r\n\r\n")
|
||||
self.wfile.write(b"0\r\n\r\n")
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
"""Silence the per-request stderr line BaseHTTPRequestHandler emits."""
|
||||
|
|
@ -1056,6 +1075,8 @@ def provider_edge_api_base(
|
|||
case InvalidFixtureMode(value=value):
|
||||
raise ValueError(f"E2E_FIXTURE_MODE={value!r} is not one of {', '.join(FIXTURE_MODES)}")
|
||||
case "live":
|
||||
if configured_cache_backend() is not None:
|
||||
return _shared_cache_edge(bind_host, advertise_host, forward_timeout).api_base(mount)
|
||||
return None
|
||||
case "record" | "replay":
|
||||
if mount not in EDGE_MOUNTS:
|
||||
|
|
@ -1073,7 +1094,7 @@ def _observed_backend(mode_raw: str, bundle_dir: Path) -> EdgeBackend:
|
|||
case InvalidFixtureMode(value=value):
|
||||
raise ValueError(f"E2E_FIXTURE_MODE={value!r} is not one of {', '.join(FIXTURE_MODES)}")
|
||||
case "live":
|
||||
return LiveEdge()
|
||||
return configured_cache_backend() or LiveEdge()
|
||||
case "record":
|
||||
return RecordEdge(_shared_recorder(bundle_dir, match_profile()), threading.Lock())
|
||||
case "replay":
|
||||
|
|
@ -1082,6 +1103,24 @@ def _observed_backend(mode_raw: str, bundle_dir: Path) -> EdgeBackend:
|
|||
assert_never(mode)
|
||||
|
||||
|
||||
def configured_cache_backend() -> CacheEdge | None:
|
||||
if LIVE_PROVIDER_REQUIRED.get() or os.environ.get("E2E_PROVIDER_CACHE", "0") == "0":
|
||||
return None
|
||||
from provider_cache_redis import configured_cache
|
||||
|
||||
return configured_cache()
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=8)
|
||||
def _shared_cache_edge(bind_host: str, advertise_host: str, forward_timeout: float) -> ProviderEdge:
|
||||
backend: Final = configured_cache_backend()
|
||||
assert backend is not None
|
||||
return start_provider_edge(
|
||||
backend, mounts=EDGE_MOUNTS, bind_host=bind_host,
|
||||
advertise_host=advertise_host, forward_timeout=forward_timeout,
|
||||
).edge
|
||||
|
||||
|
||||
@contextmanager
|
||||
def observed_provider_edge(
|
||||
observation: ProviderRequestObservation,
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ ProxyClient's key/customer methods for cleanup. Read-backs are eventually consis
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Callable, Mapping
|
||||
|
|
@ -26,6 +27,7 @@ from e2e_config import (
|
|||
PROXY_REPLICA_URLS,
|
||||
REQUEST_TIMEOUT,
|
||||
SLOW_PROVIDER_TIMEOUT_SECONDS,
|
||||
provider_edge_base,
|
||||
settle_propagation,
|
||||
)
|
||||
from e2e_http import (
|
||||
|
|
@ -93,6 +95,7 @@ from models import (
|
|||
UserDeleteBody,
|
||||
UserDeleteResponse,
|
||||
)
|
||||
from provider_cache_routing import route_cache_model
|
||||
from pydantic import BaseModel
|
||||
from transport import HttpTransport, SplitTransport, Transport, is_control_plane_path
|
||||
|
||||
|
|
@ -645,7 +648,10 @@ class ProxyClient:
|
|||
self.transport.post(
|
||||
"/model/new",
|
||||
headers=self.management_headers(),
|
||||
json=body,
|
||||
json=body.model_copy(update={"litellm_params": route_cache_model(
|
||||
body.litellm_params, provider_edge_base,
|
||||
enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1", mode=body.model_info.mode,
|
||||
)}),
|
||||
response_type=ModelNewResponse,
|
||||
)
|
||||
).model_id
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@ from models import (
|
|||
)
|
||||
from quota_client import QuotaClient
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
pytestmark = [pytest.mark.e2e, pytest.mark.provider_live]
|
||||
|
||||
# Anthropic prompt caching (host has ANTHROPIC_API_KEY; Bedrock was "Operation not allowed").
|
||||
ANTHROPIC_MODEL = "anthropic/claude-haiku-4-5-20251001"
|
||||
|
|
|
|||
|
|
@ -1254,7 +1254,8 @@ class TestHandleEdgeRequestPure:
|
|||
|
||||
|
||||
class TestApiBaseSeam:
|
||||
def test_live_mode_returns_none(self, tmp_path: Path) -> None:
|
||||
def test_live_mode_returns_none(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("E2E_PROVIDER_CACHE", raising=False)
|
||||
for mode_raw in ("live", ""):
|
||||
assert (
|
||||
provider_edge_api_base(
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@
|
|||
-- Idempotent: deletes all e2e-* rows then re-inserts deterministic data.
|
||||
|
||||
-- 1. Clean up in dependency order
|
||||
DELETE FROM "LiteLLM_InvitationLink"
|
||||
WHERE "user_id" LIKE 'e2e-%' OR "created_by" LIKE 'e2e-%' OR "updated_by" LIKE 'e2e-%';
|
||||
DELETE FROM "LiteLLM_TeamMembership" WHERE "user_id" LIKE 'e2e-%';
|
||||
DELETE FROM "LiteLLM_VerificationToken" WHERE token LIKE 'e2e-%';
|
||||
DELETE FROM "LiteLLM_TeamTable" WHERE "team_id" LIKE 'e2e-%';
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import { chromium, expect, request } from "@playwright/test";
|
||||
import { users, Role, STORAGE_PATHS } from "./fixtures/users";
|
||||
import { ARTIFACT_DIR, UI_BASE_URL } from "./constants";
|
||||
import { expectUnrestrictedDashboard, setInvitedUserPassword } from "./helpers/userOnboarding";
|
||||
import * as fs from "fs";
|
||||
import * as path from "path";
|
||||
|
||||
|
|
@ -30,32 +31,37 @@ async function globalSetup() {
|
|||
throw new Error(`Enabling enable_projects_ui failed (${settingsRes.status()}): ${await settingsRes.text()}`);
|
||||
}
|
||||
|
||||
for (const { email, password, seedApiRole } of Object.values(users)) {
|
||||
if (!seedApiRole) {
|
||||
continue;
|
||||
}
|
||||
const createRes = await api.post(`${UI_BASE_URL}${rootPath}/user/new`, {
|
||||
headers: { Authorization: `Bearer ${masterKey}` },
|
||||
data: { user_email: email, user_role: seedApiRole, auto_create_key: false },
|
||||
});
|
||||
if (!createRes.ok() && createRes.status() !== 409) {
|
||||
throw new Error(`Seeding user ${email} failed (${createRes.status()}): ${await createRes.text()}`);
|
||||
}
|
||||
const passwordRes = await api.post(`${UI_BASE_URL}${rootPath}/user/update`, {
|
||||
headers: { Authorization: `Bearer ${masterKey}` },
|
||||
data: { user_email: email, password },
|
||||
});
|
||||
if (!passwordRes.ok()) {
|
||||
throw new Error(`Setting password for ${email} failed (${passwordRes.status()}): ${await passwordRes.text()}`);
|
||||
}
|
||||
}
|
||||
await api.dispose();
|
||||
|
||||
for (const role of Object.values(Role)) {
|
||||
const { email, password } = users[role];
|
||||
const roles = [Role.ProxyAdmin, ...Object.values(Role).filter((role) => role !== Role.ProxyAdmin)];
|
||||
for (const role of roles) {
|
||||
const { email, password, seedApiRole } = users[role];
|
||||
const storagePath = STORAGE_PATHS[role];
|
||||
const page = await browser.newPage();
|
||||
try {
|
||||
if (seedApiRole) {
|
||||
const createRes = await api.post(`${UI_BASE_URL}${rootPath}/user/new`, {
|
||||
headers: { Authorization: `Bearer ${masterKey}` },
|
||||
data: { user_email: email, user_role: seedApiRole, auto_create_key: false },
|
||||
});
|
||||
if (!createRes.ok() && createRes.status() !== 409) {
|
||||
throw new Error(`Seeding user ${email} failed (${createRes.status()}): ${await createRes.text()}`);
|
||||
}
|
||||
const userId = createRes.ok()
|
||||
? (await createRes.json()).user_id
|
||||
: await (async () => {
|
||||
const existing = await api.get(`${UI_BASE_URL}${rootPath}/user/list`, {
|
||||
headers: { Authorization: `Bearer ${masterKey}` },
|
||||
params: { user_email: email },
|
||||
});
|
||||
expect(existing.ok(), `Find seeded user ${email}: HTTP ${existing.status()}`).toBe(true);
|
||||
const matches = (await existing.json()).users.filter(
|
||||
(user: { user_email: string }) => user.user_email === email,
|
||||
);
|
||||
expect(matches, `Exactly one seeded user for ${email}`).toHaveLength(1);
|
||||
return matches[0].user_id;
|
||||
})();
|
||||
expect(typeof userId, `User ID for ${email}`).toBe("string");
|
||||
await setInvitedUserPassword(api, userId, password);
|
||||
}
|
||||
await page.goto(`${UI_BASE_URL}${rootPath}/ui/login`);
|
||||
await page.getByPlaceholder("Enter your username").fill(email);
|
||||
await page.getByPlaceholder("Enter your password").fill(password);
|
||||
|
|
@ -63,7 +69,7 @@ async function globalSetup() {
|
|||
await page.waitForURL((url) => url.pathname.startsWith(`${rootPath}/ui`) && !url.pathname.includes("/login"), {
|
||||
timeout: 30_000,
|
||||
});
|
||||
await expect(page.locator("a", { hasText: "Virtual Keys" })).toBeVisible({ timeout: 30_000 });
|
||||
await expectUnrestrictedDashboard(page);
|
||||
// Dismiss feedback popup if present
|
||||
const dismiss = page.getByText("Don't ask me again");
|
||||
if (await dismiss.isVisible({ timeout: 1_500 }).catch(() => false)) {
|
||||
|
|
@ -100,6 +106,7 @@ async function globalSetup() {
|
|||
}
|
||||
}
|
||||
|
||||
await api.dispose();
|
||||
await browser.close();
|
||||
}
|
||||
|
||||
|
|
|
|||
59
tests/e2e/ui/helpers/userOnboarding.ts
Normal file
59
tests/e2e/ui/helpers/userOnboarding.ts
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
import { expect, type APIRequestContext, type Page } from "@playwright/test";
|
||||
import { UI_BASE_URL } from "../constants";
|
||||
import { masterKey, rootPath } from "./traffic";
|
||||
|
||||
const endpoint = (route: string): string => `${UI_BASE_URL}${rootPath()}${route}`;
|
||||
|
||||
export async function setInvitedUserPassword(
|
||||
request: APIRequestContext,
|
||||
userId: string,
|
||||
password: string,
|
||||
): Promise<void> {
|
||||
const invitation = await request.post(endpoint("/invitation/new"), {
|
||||
headers: { Authorization: `Bearer ${masterKey()}` },
|
||||
data: { user_id: userId },
|
||||
});
|
||||
expect(invitation.ok(), `Create invitation for ${userId}: HTTP ${invitation.status()}`).toBe(true);
|
||||
const { id } = await invitation.json();
|
||||
expect(typeof id, "invitation ID").toBe("string");
|
||||
|
||||
const onboarding = await request.get(endpoint("/onboarding/get_token"), {
|
||||
params: { invite_link: id },
|
||||
});
|
||||
expect(onboarding.ok(), `Get onboarding session for ${userId}: HTTP ${onboarding.status()}`).toBe(true);
|
||||
const { token } = await onboarding.json();
|
||||
const payload = JSON.parse(Buffer.from(token.split(".")[1], "base64url").toString("utf-8"));
|
||||
expect(typeof payload.key, "onboarding credential").toBe("string");
|
||||
const claimed = await request.post(endpoint("/onboarding/claim_token"), {
|
||||
headers: { Authorization: `Bearer ${payload.key}` },
|
||||
data: { invitation_link: id, user_id: userId, password },
|
||||
});
|
||||
expect(claimed.ok(), `Claim invitation for ${userId}: HTTP ${claimed.status()}`).toBe(true);
|
||||
}
|
||||
|
||||
export async function readDashboardSession(page: Page): Promise<{
|
||||
key: string;
|
||||
user_id: string;
|
||||
password_reset_required?: boolean;
|
||||
}> {
|
||||
await expect.poll(async () => (await page.context().cookies()).some((cookie) => cookie.name === "token")).toBe(true);
|
||||
const cookie = (await page.context().cookies()).find((candidate) => candidate.name === "token")!;
|
||||
return JSON.parse(Buffer.from(cookie.value.split(".")[1], "base64url").toString("utf-8"));
|
||||
}
|
||||
|
||||
export async function expectUnrestrictedDashboard(page: Page): Promise<void> {
|
||||
const virtualKeys = page.getByRole("complementary").getByRole("link", { name: "Virtual Keys", exact: true });
|
||||
await expect(virtualKeys).toBeVisible({ timeout: 30_000 });
|
||||
const session = await readDashboardSession(page);
|
||||
expect(session.password_reset_required === true, "login must not require a password reset").toBe(false);
|
||||
await virtualKeys.click();
|
||||
await expect(page.getByRole("main").getByRole("heading", { name: "Virtual Keys", exact: true })).toBeVisible({
|
||||
timeout: 30_000,
|
||||
});
|
||||
const info = await page.request.get(endpoint("/user/info"), {
|
||||
headers: { Authorization: `Bearer ${session.key}` },
|
||||
params: { user_id: session.user_id },
|
||||
});
|
||||
expect(info.ok(), `Read own user with dashboard session: HTTP ${info.status()}`).toBe(true);
|
||||
expect((await info.json()).user_id).toBe(session.user_id);
|
||||
}
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
import { expectUnrestrictedDashboard, setInvitedUserPassword } from "../../helpers/userOnboarding";
|
||||
import { test, expect, type APIRequestContext } from "@playwright/test";
|
||||
import { Page } from "../../fixtures/pages";
|
||||
import {
|
||||
|
|
@ -76,10 +77,7 @@ test.describe("Internal User - own team key model scope", () => {
|
|||
user_role: "internal_user",
|
||||
auto_create_key: false,
|
||||
});
|
||||
await postAsMaster(request, "/user/update", {
|
||||
user_id: userId,
|
||||
password: MEMBER_PASSWORD,
|
||||
});
|
||||
await setInvitedUserPassword(request, userId, MEMBER_PASSWORD);
|
||||
await postAsMaster(request, "/team/member_add", {
|
||||
team_id: teamId,
|
||||
member: { role: "user", user_id: userId },
|
||||
|
|
@ -99,10 +97,7 @@ test.describe("Internal User - own team key model scope", () => {
|
|||
.getByPlaceholder("Enter your password")
|
||||
.fill(MEMBER_PASSWORD);
|
||||
await page.getByRole("button", { name: "Login", exact: true }).click();
|
||||
await expect(
|
||||
page.locator("a", { hasText: "Virtual Keys" }),
|
||||
`${email} never reached the dashboard`,
|
||||
).toBeVisible({ timeout: 30_000 });
|
||||
await expectUnrestrictedDashboard(page);
|
||||
await dismissFeedbackPopup(page);
|
||||
|
||||
await navigateToPage(page, Page.ApiKeys);
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import {
|
|||
import {
|
||||
E2E_TEAM_CRUD_ALIAS,
|
||||
E2E_TEAM_ORG_ALIAS,
|
||||
E2E_TEAM_ORG_ID,
|
||||
INTERNAL_USER_STORAGE_PATH,
|
||||
} from "../../constants";
|
||||
import { Page } from "../../fixtures/pages";
|
||||
|
|
@ -179,18 +180,28 @@ test.describe("Models and Endpoints for an internal user", () => {
|
|||
`switching to ${ALL_MODELS_VIEW} leaves the table populated rather than blanking it`,
|
||||
).toHaveCount(1, { timeout: 15_000 });
|
||||
|
||||
await expect(page).toHaveURL((url) =>
|
||||
url.searchParams.get("filter_team") === E2E_TEAM_ORG_ID &&
|
||||
url.searchParams.get("view_mode") === "all",
|
||||
);
|
||||
await page.reload();
|
||||
await expect(
|
||||
teamSelector(page),
|
||||
"the team selection is not persisted across a reload, so the table returns to the personal view",
|
||||
).toContainText(PERSONAL_TEAM, { timeout: 15_000 });
|
||||
"the selected team is restored from the URL after a reload",
|
||||
).toContainText(E2E_TEAM_ORG_ALIAS, { timeout: 15_000 });
|
||||
await expect(
|
||||
viewSelector(page),
|
||||
"the view selection is not persisted across a reload either",
|
||||
).toContainText(CURRENT_TEAM_VIEW, { timeout: 15_000 });
|
||||
"the selected view is restored from the URL after a reload",
|
||||
).toContainText(ALL_MODELS_VIEW, { timeout: 15_000 });
|
||||
await expect(modelRow(page, CHAT_MODEL_A)).toHaveCount(1, { timeout: 15_000 });
|
||||
await expect(page.getByTestId("pagination-range")).toHaveText("Showing 1-1 of 1");
|
||||
await expect(modelRow(page, CHAT_MODEL_B)).toHaveCount(0);
|
||||
await expect(modelRow(page, ungrantedModelName)).toHaveCount(0);
|
||||
|
||||
await chooseOption(page, teamSelector(page), PERSONAL_TEAM);
|
||||
await expect(
|
||||
modelRow(page, ungrantedModelName),
|
||||
"the personal view still renders models after a reload rather than coming back empty",
|
||||
"switching back to the personal team restores models outside the selected team",
|
||||
).toHaveCount(1, { timeout: 30_000 });
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import { expectUnrestrictedDashboard, setInvitedUserPassword } from "../../helpers/userOnboarding";
|
||||
import { test, expect } from "@playwright/test";
|
||||
import { ADMIN_STORAGE_PATH } from "../../constants";
|
||||
import { Page } from "../../fixtures/pages";
|
||||
|
|
@ -46,19 +47,13 @@ test.describe("Second proxy admin", () => {
|
|||
|
||||
const userId = await inviteAdminUser();
|
||||
try {
|
||||
const passwordRes = await request.post("/user/update", {
|
||||
headers: auth,
|
||||
data: { user_email: email, password },
|
||||
});
|
||||
expect(passwordRes.ok(), `setting password failed (${passwordRes.status()}): ${await passwordRes.text()}`).toBe(
|
||||
true,
|
||||
);
|
||||
await setInvitedUserPassword(request, userId, password);
|
||||
|
||||
await page.goto("/ui/login");
|
||||
await page.getByPlaceholder("Enter your username").fill(email);
|
||||
await page.getByPlaceholder("Enter your password").fill(password);
|
||||
await page.getByRole("button", { name: "Login", exact: true }).click();
|
||||
await expect(page.locator("a", { hasText: "Virtual Keys" })).toBeVisible({ timeout: 30_000 });
|
||||
await expectUnrestrictedDashboard(page);
|
||||
await dismissFeedbackPopup(page);
|
||||
|
||||
await navigateToPage(page, Page.ApiKeys);
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import { expectUnrestrictedDashboard, setInvitedUserPassword } from "../../helpers/userOnboarding";
|
||||
import { test, expect, type Browser, type BrowserContext, type Page as PlaywrightPage } from "@playwright/test";
|
||||
import { Page } from "../../fixtures/pages";
|
||||
import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation";
|
||||
|
|
@ -24,7 +25,7 @@ async function signIn(browser: Browser, email: string): Promise<BrowserContext>
|
|||
await page.getByPlaceholder("Enter your username").fill(email);
|
||||
await page.getByPlaceholder("Enter your password").fill(PASSWORD);
|
||||
await page.getByRole("button", { name: "Login", exact: true }).click();
|
||||
await expect(page.locator("a", { hasText: "Virtual Keys" })).toBeVisible({ timeout: 30_000 });
|
||||
await expectUnrestrictedDashboard(page);
|
||||
await dismissFeedbackPopup(page);
|
||||
return context;
|
||||
}
|
||||
|
|
@ -49,11 +50,7 @@ test.describe("Team Admin - Member permissions", () => {
|
|||
data: { user_id: userId, user_email: email, user_role: "internal_user", auto_create_key: false },
|
||||
});
|
||||
expect(created.ok(), `POST /user/new for ${userId} (${created.status()}): ${await created.text()}`).toBe(true);
|
||||
const password = await request.post("/user/update", {
|
||||
headers: auth(),
|
||||
data: { user_id: userId, password: PASSWORD },
|
||||
});
|
||||
expect(password.ok(), `POST /user/update for ${userId} (${password.status()})`).toBe(true);
|
||||
await setInvitedUserPassword(request, userId, PASSWORD);
|
||||
};
|
||||
|
||||
let teamId = "";
|
||||
|
|
|
|||
|
|
@ -0,0 +1,37 @@
|
|||
import importlib
|
||||
import logging
|
||||
from collections.abc import Iterator
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm_proxy_extras._logging as extras_logging
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fresh_extras_logger() -> Iterator[logging.Logger]:
|
||||
logger = logging.getLogger("litellm_proxy_extras")
|
||||
saved_handlers = logger.handlers[:]
|
||||
saved_level = logger.level
|
||||
logger.handlers[:] = []
|
||||
try:
|
||||
yield logger
|
||||
finally:
|
||||
logger.handlers[:] = saved_handlers
|
||||
logger.setLevel(saved_level)
|
||||
|
||||
|
||||
def test_litellm_log_error_silences_extras_info_lines(monkeypatch, fresh_extras_logger):
|
||||
monkeypatch.setenv("LITELLM_LOG", "ERROR")
|
||||
reloaded = importlib.reload(extras_logging).logger
|
||||
assert reloaded is fresh_extras_logger
|
||||
assert reloaded.isEnabledFor(logging.INFO) is False
|
||||
assert reloaded.isEnabledFor(logging.ERROR) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("litellm_log", [None, "info", "DEBUG"])
|
||||
def test_unset_or_verbose_litellm_log_keeps_extras_info_lines(monkeypatch, fresh_extras_logger, litellm_log):
|
||||
if litellm_log is None:
|
||||
monkeypatch.delenv("LITELLM_LOG", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("LITELLM_LOG", litellm_log)
|
||||
assert importlib.reload(extras_logging).logger.isEnabledFor(logging.INFO) is True
|
||||
|
|
@ -163,6 +163,9 @@ async def test_reset_budget_keys_partial_failure():
|
|||
key1, key2, key3, key4, key5, key6 = (
|
||||
_attrify(k) for k in [key1, key2, key3, key4, key5, key6]
|
||||
)
|
||||
pre_reset_spend = {
|
||||
k["token"]: k["spend"] for k in [key2, key3, key4, key5, key6]
|
||||
}
|
||||
prisma_client.get_data = AsyncMock(
|
||||
return_value=[key1, key2, key3, key4, key5, key6]
|
||||
)
|
||||
|
|
@ -201,7 +204,7 @@ async def test_reset_budget_keys_partial_failure():
|
|||
# And every write must carry only {spend, budget_reset_at} — never the full row.
|
||||
for c in key_writes:
|
||||
assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
assert c["data"]["spend"] == 0
|
||||
assert c["data"]["spend"] == {"decrement": pre_reset_spend[c["where"]["token"]]}
|
||||
|
||||
# Verify that the failure logging hook was scheduled (due to the failure for key1)
|
||||
failure_hook_calls = (
|
||||
|
|
@ -252,6 +255,9 @@ async def test_reset_budget_users_partial_failure():
|
|||
user1, user2, user3, user4, user5, user6 = (
|
||||
_attrify(u) for u in [user1, user2, user3, user4, user5, user6]
|
||||
)
|
||||
pre_reset_spend = {
|
||||
u["user_id"]: u["spend"] for u in [user2, user3, user4, user5, user6]
|
||||
}
|
||||
prisma_client.get_data = AsyncMock(
|
||||
return_value=[user1, user2, user3, user4, user5, user6]
|
||||
)
|
||||
|
|
@ -280,7 +286,9 @@ async def test_reset_budget_users_partial_failure():
|
|||
assert written_ids == ["user2", "user3", "user4", "user5", "user6"]
|
||||
for c in user_writes:
|
||||
assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
assert c["data"]["spend"] == 0
|
||||
assert c["data"]["spend"] == {
|
||||
"decrement": pre_reset_spend[c["where"]["user_id"]]
|
||||
}
|
||||
|
||||
failure_hook_calls = (
|
||||
proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
|
||||
|
|
@ -441,6 +449,7 @@ async def test_reset_budget_teams_partial_failure():
|
|||
for t in [team1, team2]:
|
||||
t.setdefault("team_id", t["id"])
|
||||
team1, team2 = _attrify(team1), _attrify(team2)
|
||||
pre_reset_spend = team2["spend"]
|
||||
prisma_client.get_data = AsyncMock(return_value=[team1, team2])
|
||||
|
||||
async def fake_reset_team(team, current_time, reset_settings=None):
|
||||
|
|
@ -465,7 +474,7 @@ async def test_reset_budget_teams_partial_failure():
|
|||
assert len(team_writes) == 1
|
||||
assert team_writes[0]["where"] == {"team_id": "team2"}
|
||||
assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
assert team_writes[0]["data"]["spend"] == 0
|
||||
assert team_writes[0]["data"]["spend"] == {"decrement": pre_reset_spend}
|
||||
|
||||
failure_hook_calls = (
|
||||
proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
|
||||
|
|
@ -542,6 +551,11 @@ async def test_reset_budget_continues_other_categories_on_failure():
|
|||
user1, user2 = _attrify(user1), _attrify(user2)
|
||||
team1, team2 = _attrify(team1), _attrify(team2)
|
||||
enduser1 = _attrify(enduser1)
|
||||
pre_reset_spend = {
|
||||
**{k["token"]: k["spend"] for k in [key1, key2]},
|
||||
**{u["user_id"]: u["spend"] for u in [user2]},
|
||||
**{t["team_id"]: t["spend"] for t in [team1, team2]},
|
||||
}
|
||||
_wire_cascade_reads_for_test(prisma_client)
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
|
|
@ -618,7 +632,9 @@ async def test_reset_budget_continues_other_categories_on_failure():
|
|||
# Every batched write must carry only the two reset fields, never the full row.
|
||||
for c in key_writes + user_writes + team_writes:
|
||||
assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
assert c["data"]["spend"] == 0
|
||||
assert c["data"]["spend"] == {
|
||||
"decrement": pre_reset_spend[next(iter(c["where"].values()))]
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -86,6 +86,7 @@ def test_async_fallbacks(caplog):
|
|||
if "Task exception was never retrieved" not in log
|
||||
and "Task was destroyed but it is pending" not in log
|
||||
and "get_available_deployment" not in log
|
||||
and "Selected deployment for model" not in log
|
||||
and "in the Langfuse queue" not in log
|
||||
and "Unclosed client session" not in log
|
||||
and "Unclosed connector" not in log
|
||||
|
|
|
|||
|
|
@ -1069,6 +1069,7 @@ async def test_jwt_non_admin_team_route_access(monkeypatch):
|
|||
|
||||
mock_jwt_response = {
|
||||
"is_proxy_admin": False,
|
||||
"jwt_claims": {},
|
||||
"team_id": None,
|
||||
"team_object": None,
|
||||
"user_id": None,
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -135,9 +135,7 @@ def test_fill_missing_requires_per_rule_opt_in(restore_generalizations):
|
|||
"supports_vision": True,
|
||||
}
|
||||
|
||||
restore_generalizations(
|
||||
[{"name": "base", "pattern": r"^acme-", "model_info": {"supports_reasoning": True}}]
|
||||
)
|
||||
restore_generalizations([{"name": "base", "pattern": r"^acme-", "model_info": {"supports_reasoning": True}}])
|
||||
assert match_fill_missing_generalizations("acme-1", "openai") is None
|
||||
|
||||
restore_generalizations(
|
||||
|
|
@ -451,6 +449,94 @@ def shipped_cost_map(monkeypatch):
|
|||
set_fallback_generalizations(previous_rules)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model,provider",
|
||||
[
|
||||
("gemini-4-pro", "gemini"),
|
||||
("gemini/gemini-4-pro", None),
|
||||
("gemini-3.9-flash-lite-preview-09-2026", "vertex_ai"),
|
||||
("vertex_ai/gemini-4-pro", None),
|
||||
("gemini-4-pro-preview-customtools", "gemini"),
|
||||
("google/gemini-4-pro", "openrouter"),
|
||||
("google/gemini-4-pro", "deepinfra"),
|
||||
("google/gemini-4-pro", "vercel_ai_gateway"),
|
||||
("google.gemini-4-pro", "oci"),
|
||||
("databricks-gemini-4-1-pro", "databricks"),
|
||||
],
|
||||
)
|
||||
def test_shipped_gemini_chat_baseline_resolves_unmapped_ids(shipped_cost_map, model, provider):
|
||||
assert model not in litellm.model_cost
|
||||
if provider == "gemini":
|
||||
assert f"gemini/{model}" not in litellm.model_cost
|
||||
elif provider in {"openrouter", "deepinfra", "vercel_ai_gateway", "oci", "databricks"}:
|
||||
assert f"{provider}/{model}" not in litellm.model_cost
|
||||
|
||||
info = litellm.get_model_info(model, custom_llm_provider=provider)
|
||||
assert info["litellm_provider"] == (provider or model.split("/")[0])
|
||||
assert info["mode"] == "chat"
|
||||
assert not info.get("max_input_tokens")
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_system_messages"] is True
|
||||
assert info["supports_vision"] is True
|
||||
assert info["supports_response_schema"] is True
|
||||
assert info["supports_pdf_input"] is True
|
||||
assert info["supports_prompt_caching"] is True
|
||||
assert info["supports_web_search"] is True
|
||||
assert not info.get("input_cost_per_token")
|
||||
assert not info.get("output_cost_per_token")
|
||||
|
||||
|
||||
def test_shipped_gemini_chat_baseline_loses_to_perplexity_exact_entries(shipped_cost_map):
|
||||
info = litellm.get_model_info("google/gemini-2.5-pro", custom_llm_provider="perplexity")
|
||||
entry = litellm.model_cost["perplexity/google/gemini-2.5-pro"]
|
||||
assert info["mode"] == "responses"
|
||||
assert entry["supports_reasoning"] is False
|
||||
|
||||
|
||||
def test_shipped_gemini_chat_baseline_skips_non_chat_and_pre_2_5_ids(shipped_cost_map):
|
||||
for model in (
|
||||
"gemini/gemini-4-flash-image",
|
||||
"gemini/gemini-3.9-flash-preview-tts",
|
||||
"gemini/gemini-4-flash-live-preview",
|
||||
"gemini/gemini-4-flash-native-audio",
|
||||
"gemini/gemini-embedding-4",
|
||||
"gemini/gemini-2.5-computer-use-preview-12-2026",
|
||||
"gemini/gemini-2.0-flash-new",
|
||||
"gemini/gemini-1.5-pro-new",
|
||||
"gemini/gemini-4-flashy",
|
||||
"gemini/gemini-4-flash-transcribe",
|
||||
"gemini/gemini-4-flash-live-translate-preview",
|
||||
"databricks-gemini-3-1-flash-image",
|
||||
"openrouter/google/gemini-2.0-flash-001",
|
||||
):
|
||||
assert match_capability_generalizations(model) is None, model
|
||||
|
||||
|
||||
def test_shipped_gemini_chat_baseline_keeps_reasoning_effort_on_unmapped_model(shipped_cost_map):
|
||||
assert litellm.supports_reasoning(model="gemini-4-pro", custom_llm_provider="gemini") is True
|
||||
|
||||
optional_params = litellm.utils.get_optional_params(
|
||||
model="gemini-4-pro",
|
||||
custom_llm_provider="gemini",
|
||||
reasoning_effort="medium",
|
||||
drop_params=False,
|
||||
)
|
||||
assert isinstance(optional_params, dict)
|
||||
assert optional_params["thinkingConfig"]["thinkingBudget"] > 0
|
||||
assert optional_params["thinkingConfig"]["includeThoughts"] is True
|
||||
|
||||
|
||||
def test_shipped_gemini_chat_baseline_loses_to_exact_entries(shipped_cost_map):
|
||||
model = "gemini-2.5-flash-lite"
|
||||
info = litellm.get_model_info(model, custom_llm_provider="gemini")
|
||||
entry = litellm.model_cost["gemini/gemini-2.5-flash-lite"]
|
||||
assert info["max_tokens"] == entry["max_tokens"]
|
||||
assert info["input_cost_per_token"] == entry["input_cost_per_token"]
|
||||
assert entry["input_cost_per_token"] > 0
|
||||
|
||||
|
||||
def test_shipped_bare_claude_id_routes_to_anthropic(shipped_cost_map):
|
||||
_, provider, _, _ = litellm.get_llm_provider(model="claude-haiku-4-6")
|
||||
assert provider == "anthropic"
|
||||
|
|
|
|||
|
|
@ -1,7 +1,5 @@
|
|||
|
||||
import pytest
|
||||
|
||||
|
||||
from litellm.litellm_core_utils.get_supported_openai_params import (
|
||||
get_supported_openai_params,
|
||||
)
|
||||
|
|
@ -33,9 +31,7 @@ def test_base_model_label_alone_lacks_bedrock_tools():
|
|||
"""The label by itself does not advertise tools; this is what made the union
|
||||
necessary. Guards against the discrepancy disappearing (and the regression test
|
||||
above silently passing for the wrong reason)."""
|
||||
params = get_supported_openai_params(
|
||||
model=BEDROCK_LABEL, custom_llm_provider="bedrock"
|
||||
)
|
||||
params = get_supported_openai_params(model=BEDROCK_LABEL, custom_llm_provider="bedrock")
|
||||
|
||||
assert params is not None
|
||||
assert "tools" not in params
|
||||
|
|
@ -46,14 +42,8 @@ def test_base_model_is_additive_not_replacement():
|
|||
|
||||
Bedrock: real id supports ``tools`` but not the label's reasoning hint; the union
|
||||
must contain the real model's ``tools`` regardless of the label being a subset."""
|
||||
real_only = set(
|
||||
get_supported_openai_params(
|
||||
model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock"
|
||||
)
|
||||
)
|
||||
label_only = set(
|
||||
get_supported_openai_params(model=BEDROCK_LABEL, custom_llm_provider="bedrock")
|
||||
)
|
||||
real_only = set(get_supported_openai_params(model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock"))
|
||||
label_only = set(get_supported_openai_params(model=BEDROCK_LABEL, custom_llm_provider="bedrock"))
|
||||
combined = set(
|
||||
get_supported_openai_params(
|
||||
model=BEDROCK_REAL_MODEL,
|
||||
|
|
@ -70,19 +60,15 @@ def test_base_model_is_additive_not_replacement():
|
|||
def test_base_model_adds_capabilities_the_real_model_lacks():
|
||||
"""Regression for #27717 (the behavior the union must preserve).
|
||||
|
||||
``gemini-3.1-pro`` isn't in the cost map so it advertises no reasoning support,
|
||||
``gemini-exp-9999`` isn't in the cost map so it advertises no reasoning support,
|
||||
but the registered ``gemini-3.1-pro-preview`` base_model does. The hint must add
|
||||
``reasoning_effort``/``thinking`` without the call erroring."""
|
||||
real_only = set(
|
||||
get_supported_openai_params(
|
||||
model="gemini-3.1-pro", custom_llm_provider="gemini"
|
||||
)
|
||||
)
|
||||
real_only = set(get_supported_openai_params(model="gemini-exp-9999", custom_llm_provider="gemini"))
|
||||
assert "reasoning_effort" not in real_only
|
||||
|
||||
combined = set(
|
||||
get_supported_openai_params(
|
||||
model="gemini-3.1-pro",
|
||||
model="gemini-exp-9999",
|
||||
custom_llm_provider="gemini",
|
||||
base_model="gemini-3.1-pro-preview",
|
||||
)
|
||||
|
|
@ -93,21 +79,15 @@ def test_base_model_adds_capabilities_the_real_model_lacks():
|
|||
|
||||
def test_no_base_model_is_unchanged():
|
||||
"""Omitting ``base_model`` must resolve purely from ``model``."""
|
||||
with_none = get_supported_openai_params(
|
||||
model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock", base_model=None
|
||||
)
|
||||
plain = get_supported_openai_params(
|
||||
model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock"
|
||||
)
|
||||
with_none = get_supported_openai_params(model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock", base_model=None)
|
||||
plain = get_supported_openai_params(model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock")
|
||||
|
||||
assert with_none == plain
|
||||
|
||||
|
||||
def test_base_model_equal_to_model_is_unchanged():
|
||||
"""A ``base_model`` identical to ``model`` must not double-resolve or reorder."""
|
||||
plain = get_supported_openai_params(
|
||||
model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock"
|
||||
)
|
||||
plain = get_supported_openai_params(model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock")
|
||||
same = get_supported_openai_params(
|
||||
model=BEDROCK_REAL_MODEL,
|
||||
custom_llm_provider="bedrock",
|
||||
|
|
@ -152,14 +132,10 @@ def test_bedrock_converse_alias_resolves_like_bedrock():
|
|||
params saw no Bedrock capabilities for a Converse model invoked via the alias."""
|
||||
anthropic_model = "bedrock/converse/us.anthropic.claude-sonnet-4-6"
|
||||
|
||||
via_alias = get_supported_openai_params(
|
||||
model=anthropic_model, custom_llm_provider="bedrock_converse"
|
||||
)
|
||||
via_alias = get_supported_openai_params(model=anthropic_model, custom_llm_provider="bedrock_converse")
|
||||
|
||||
assert via_alias is not None
|
||||
assert via_alias == get_supported_openai_params(
|
||||
model=anthropic_model, custom_llm_provider="bedrock"
|
||||
)
|
||||
assert via_alias == get_supported_openai_params(model=anthropic_model, custom_llm_provider="bedrock")
|
||||
assert "web_search_options" not in via_alias
|
||||
assert "tools" in via_alias
|
||||
|
||||
|
|
@ -167,9 +143,7 @@ def test_bedrock_converse_alias_resolves_like_bedrock():
|
|||
def test_bedrock_converse_alias_keeps_nova_web_search_options():
|
||||
"""Nova on the ``bedrock_converse`` alias still advertises web_search_options, proving the
|
||||
alias routes through the model-aware config rather than a blanket Bedrock default."""
|
||||
nova_params = get_supported_openai_params(
|
||||
model="amazon.nova-pro-v1:0", custom_llm_provider="bedrock_converse"
|
||||
)
|
||||
nova_params = get_supported_openai_params(model="amazon.nova-pro-v1:0", custom_llm_provider="bedrock_converse")
|
||||
|
||||
assert nova_params is not None
|
||||
assert "web_search_options" in nova_params
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue