scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

Merge branch 'litellm_internal_staging' of https://github.com/BerriAI/litellm into litellm_presidio_new_entities

Browse source 
# Conflicts:
#	type-discipline-budget.json
This commit is contained in:
Michael van den Berg 2026-08-14 14:59:53 +02:00
commit 8cd641f426
677 changed files with 48240 additions and 21254 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

84
.circleci/config.yml
View file

@ -2744,84 +2744,6 @@ jobs:
file: ./coverage.xml
flags: circleci
ui_build:
docker:
- image: cimg/node:24.19@sha256:8966565f07189a67d64d6808a2b127f31dafae566508e3547f55640e1070bfad
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
resource_class: medium+
working_directory: ~/project
steps:
- checkout
- skip_if_unrelated_changes:
category: client
- setup_google_dns
- restore_cache:
keys:
- ui-build-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
- ui-build-deps-v1-
- restore_cache:
keys:
- ui-nextjs-cache-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
- ui-nextjs-cache-v1-
- run:
name: Install dependencies
command: |
cd ui/litellm-dashboard
npm ci
- save_cache:
key: ui-build-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- run:
name: Build UI
command: |
cd ui/litellm-dashboard
source ./build_ui.sh
- save_cache:
key: ui-nextjs-cache-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
paths:
- ui/litellm-dashboard/.next/cache
- persist_to_workspace:
root: .
paths:
- litellm/proxy/_experimental/out
ui_unit_tests:
docker:
- image: cimg/node:24.19@sha256:8966565f07189a67d64d6808a2b127f31dafae566508e3547f55640e1070bfad
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
resource_class: xlarge
working_directory: ~/project
steps:
- checkout
- skip_if_unrelated_changes:
category: client
- setup_google_dns
- restore_cache:
keys:
- ui-unit-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
- ui-unit-deps-v1-
- run:
name: Install dependencies
command: |
cd ui/litellm-dashboard
npm ci
- save_cache:
key: ui-unit-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- run:
name: Run UI unit tests (Vitest)
command: |
cd ui/litellm-dashboard
CI=true npm run test -- --run \
--pool forks --poolOptions.forks.maxForks=6
e2e_ui_testing:
docker:
- image: cimg/python:3.12-browsers@sha256:b432899af01c9a311bf74f4f22e9ada2e5306d4b1b4383f8d29e1228a5844ef2
@ -3181,12 +3103,6 @@ workflows:
filters: *main_branches
- litellm_router_unit_testing:
filters: *main_branches
- ui_build:
filters: *main_branches
- ui_unit_tests:
requires:
- ui_build
filters: *main_branches
- auth_ui_unit_tests:
filters: *main_branches
- proxy_behavior_tests:

54
.github/workflows/test-terraform-modules.yml vendored Normal file
View file

@ -0,0 +1,54 @@
name: Terraform Modules
on:
push:
paths:
- "terraform/litellm/aws/**"
- ".github/workflows/test-terraform-modules.yml"
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
paths:
- "terraform/litellm/aws/**"
- ".github/workflows/test-terraform-modules.yml"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
aws-module:
name: fmt, validate, test (aws)
runs-on: ubuntu-latest
timeout-minutes: 15
defaults:
run:
working-directory: terraform/litellm/aws
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- uses: hashicorp/setup-terraform@b9cd54a3c349d3f38e8881555d616ced269862dd # v3.1.2
with:
terraform_version: 1.13.3
terraform_wrapper: false
- name: fmt
run: terraform fmt -recursive -check -diff
- name: init
run: terraform init -backend=false -input=false
- name: validate
run: terraform validate
# Plan-only, mock_provider-backed: no AWS credentials, no API calls.
- name: test
run: terraform test

2
.github/workflows/test-unit-proxy-db.yml vendored
View file

@ -135,8 +135,6 @@ jobs:
test-path: >-
tests/proxy_unit_tests/test_proxy_server.py
tests/proxy_unit_tests/test_proxy_server_keys.py
tests/proxy_unit_tests/test_proxy_server_caching.py
tests/proxy_unit_tests/test_proxy_server_langfuse.py
tests/proxy_unit_tests/test_proxy_server_spend.py
tests/proxy_unit_tests/test_aproxy_startup.py
workers: 4

2
CLAUDE.md
View file

@ -29,7 +29,7 @@ End-to-end tests belong in `tests/e2e/` and must follow the harness conventions
When creating PRs, don't set base to `main`. `litellm_internal_staging` is the default base branch and serves that purpose for both internal and external / OSS contributions
When writing a PR body, treat the comments and imperative instructions inside @.github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
Same applies for filing bug reports and feature requests, with .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml, respectively

20
basedpyright-code-budget.json
View file

@ -1,9 +1,9 @@
{
"reportAny": {
"limit": 23919
"limit": 22947
},
"reportArgumentType": {
"limit": 2580
"limit": 2579
},
"reportAssignmentType": {
"limit": 323
@ -24,7 +24,7 @@
"limit": 19
},
"reportExplicitAny": {
"limit": 7573
"limit": 7312
},
"reportFunctionMemberAccess": {
"limit": 7
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5719
"limit": 5707
},
"reportMissingTypeArgument": {
"limit": 15657
"limit": 15642
},
"reportMissingTypeStubs": {
"limit": 40
@ -99,22 +99,22 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44832
"limit": 44776
},
"reportUnknownLambdaType": {
"limit": 113
},
"reportUnknownMemberType": {
"limit": 39269
"limit": 39237
},
"reportUnknownParameterType": {
"limit": 19988
"limit": 19969
},
"reportUnknownVariableType": {
"limit": 30923
"limit": 30881
},
"reportUnnecessaryCast": {
"limit": 118
"limit": 117
},
"reportUnnecessaryComparison": {
"limit": 699

103
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -3,7 +3,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
"""
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Tuple
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -23,6 +23,15 @@ if TYPE_CHECKING:
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
"completed",
"complete",
"failed",
"expired",
"cancelled",
"stale_expired",
)
class CheckBatchCost:
def __init__(
@ -132,11 +141,11 @@ class CheckBatchCost:
in non-terminal states as 'stale_expired'. These will never complete and
should not be polled.
"""
cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
cutoff: Final = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
result: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
where={
"file_purpose": "batch",
"status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
"status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)},
"created_at": {"lt": cutoff},
},
data={"status": "stale_expired"},
@ -147,6 +156,26 @@ class CheckBatchCost:
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
)
if not self._has_batch_processed_column:
return
# A row already in a terminal status is never rewritten by the sweep above, so
# without this it keeps a poll-page slot forever and starves newer batches.
retired: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
where={
"file_purpose": "batch",
"batch_processed": False,
"status": {"in": ["complete", "completed"]},
"created_at": {"lt": cutoff},
},
data={"batch_processed": True},
)
if retired > 0:
verbose_proxy_logger.warning(
f"CheckBatchCost: gave up on {retired} completed managed objects older than "
f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
)
async def _fallback_find_jobs(self) -> list:
"""Query batch jobs without the batch_processed filter (for older schemas)."""
return await self.prisma_client.db.litellm_managedobjecttable.find_many(
@ -167,6 +196,68 @@ class CheckBatchCost:
order={"created_at": "asc"},
)
async def _retire_job(self, job: "LiteLLM_ManagedObjectTable", reason: str) -> None:
"""
Take a row that can never be costed out of the poll page. Leaving it selectable
would burn one of the MAX_OBJECTS_PER_POLL_CYCLE slots on every future cycle, and
once enough such rows accumulate no newer batch is ever reached. Older schemas
without batch_processed can only be excluded through the status filter.
"""
data: Final = (
{"batch_processed": True}
if self._has_batch_processed_column
else {"status": "stale_expired"}
)
try:
await self.prisma_client.db.litellm_managedobjecttable.update(
where={"id": job.id},
data=data,
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}"
)
return
verbose_proxy_logger.warning(
f"CheckBatchCost: job {job.id} can never be costed ({reason}), "
"so it will no longer be polled"
)
@staticmethod
def _has_unified_id_without_model(job: "LiteLLM_ManagedObjectTable") -> bool:
"""A unified id that decodes but carries no model_id can never be routed."""
from litellm.proxy.openai_files_endpoints.common_utils import (
convert_b64_uid_to_unified_uid,
get_model_id_from_unified_batch_id,
)
decoded: Final = convert_b64_uid_to_unified_uid(job.unified_object_id)
return (
decoded != job.unified_object_id
and get_model_id_from_unified_batch_id(decoded) is None
)
@staticmethod
def _is_batch_gone_at_provider(error: Exception, batch_id: str) -> bool:
"""
A 404 naming the batch means the provider dropped its record of it, so no later
retrieve can ever succeed. A 404 about anything else, a renamed Azure deployment
or a fallback deployment that never saw this batch, is still fixable in config, so
it keeps retrying.
"""
import openai
from litellm.exceptions import NotFoundError
return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error)
def _batch_deployment_exists(self, model_id: str) -> bool:
"""A 404 only proves the batch is gone when it came from the batch's own
deployment. Once that deployment leaves the router, default fallbacks can
silently send the retrieve to a provider that never saw the batch, so its
404 must not retire the row; the staleness sweep bounds it instead."""
return self.llm_router.get_deployment(model_id=model_id) is not None
@staticmethod
def _record_error(
prom_logger: Optional["PrometheusLogger"], error_type: str
@ -645,6 +736,8 @@ class CheckBatchCost:
for job in jobs:
routing = self._resolve_job_routing(job, prom_logger)
if routing is None:
if self._has_unified_id_without_model(job):
await self._retire_job(job, "unified object id has no model id")
continue
model_id, batch_id = routing
@ -667,6 +760,8 @@ class CheckBatchCost:
)
if prom_logger:
prom_logger.record_check_batch_cost_error("provider_retrieval_error")
if self._is_batch_gone_at_provider(e, batch_id) and self._batch_deployment_exists(model_id):
await self._retire_job(job, f"batch {batch_id} no longer exists at the provider")
continue
## RETRIEVE THE BATCH JOB OUTPUT FILE

4
helm/litellm-helm/templates/migrations-job.yaml
View file

@ -105,6 +105,10 @@ spec:
{{- toYaml . | nindent 8 }}
{{- end }}
restartPolicy: OnFailure
{{- with .Values.nodeSelector }}
nodeSelector:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.affinity }}
affinity:
{{- toYaml . | nindent 8 }}

24
helm/litellm-helm/tests/migrations-job_tests.yaml
View file

@ -290,3 +290,27 @@ tests:
value:
allowPrivilegeEscalation: false
readOnlyRootFilesystem: true
- it: should schedule onto the same nodes as the gateway
template: migrations-job.yaml
set:
migrationJob:
enabled: true
nodeSelector:
karpenter.sh/nodepool: litellm-e2e
tolerations:
- key: workload
operator: Equal
value: litellm-e2e
effect: NoSchedule
asserts:
- equal:
path: spec.template.spec.nodeSelector
value:
karpenter.sh/nodepool: litellm-e2e
- equal:
path: spec.template.spec.tolerations
value:
- key: workload
operator: Equal
value: litellm-e2e
effect: NoSchedule

49
litellm-proxy-extras/litellm_proxy_extras/migrations/20260811172448_add_shadow_eval/migration.sql Normal file
View file

@ -0,0 +1,49 @@
-- CreateTable
CREATE TABLE "LiteLLM_ShadowEvalJob" (
"id" TEXT NOT NULL,
"api_key_id" TEXT NOT NULL,
"router_name" TEXT NOT NULL,
"judge_model" TEXT NOT NULL,
"shadow_percentage" DOUBLE PRECISION NOT NULL,
"max_turns" INTEGER NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"created_by" TEXT,
"ends_at" TIMESTAMP(3) NOT NULL,
"stopped_at" TIMESTAMP(3),
CONSTRAINT "LiteLLM_ShadowEvalJob_pkey" PRIMARY KEY ("id")
);
-- CreateTable
CREATE TABLE "LiteLLM_ShadowEvalAttempt" (
"id" TEXT NOT NULL,
"job_id" TEXT NOT NULL,
"request_id" TEXT NOT NULL,
"outcome" TEXT NOT NULL,
"tier" TEXT,
"real_model" TEXT,
"shadow_model" TEXT,
"confidence" DOUBLE PRECISION,
"judge_cost" DOUBLE PRECISION NOT NULL DEFAULT 0,
"error" TEXT,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_ShadowEvalAttempt_pkey" PRIMARY KEY ("id")
);
-- CreateIndex
CREATE INDEX "LiteLLM_ShadowEvalJob_api_key_id_idx" ON "LiteLLM_ShadowEvalJob"("api_key_id");
-- CreateIndex
CREATE INDEX "LiteLLM_ShadowEvalJob_created_at_idx" ON "LiteLLM_ShadowEvalJob"("created_at");
-- CreateIndex
CREATE INDEX "LiteLLM_ShadowEvalAttempt_job_id_idx" ON "LiteLLM_ShadowEvalAttempt"("job_id");
-- One active job per key, enforced by the database rather than a read-then-create in the
-- start endpoint, which races against a concurrent start on another pod. Partial indexes
-- are not expressible in schema.prisma, so this lives here only. Active means not yet
-- stopped; the start endpoint stamps stopped_at on expired jobs before creating.
CREATE UNIQUE INDEX "LiteLLM_ShadowEvalJob_one_active_per_key"
ON "LiteLLM_ShadowEvalJob"("api_key_id") WHERE "stopped_at" IS NULL;

38
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1450,6 +1450,44 @@ model LiteLLM_AutoRouterSession {
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
}
// Shadow eval: pre-adoption evaluation of an auto-router against a key's live traffic.
// A sampled slice of requests is duplicated through the router in a detached task and an
// LLM judge compares real vs shadow responses blind. The job row is immutable config plus
// stopped_at; every count, status, and spend figure is derived from the append-only
// attempt rows, so nothing can disagree across pods or stop races.
model LiteLLM_ShadowEvalJob {
id String @id @default(cuid())
api_key_id String // hashed virtual key whose traffic is shadowed
router_name String
judge_model String
shadow_percentage Float
max_turns Int // sample budget: judge at most this many turns
created_at DateTime @default(now())
created_by String?
ends_at DateTime
stopped_at DateTime?
@@index([api_key_id])
@@index([created_at])
}
// One row per sampled pipeline: a blind verdict (real | shadow | tie) or an error.
model LiteLLM_ShadowEvalAttempt {
id String @id @default(cuid())
job_id String
request_id String // the judged real request
outcome String // real | shadow | tie | error
tier String? // router's tier for the prompt, when classified
real_model String?
shadow_model String?
confidence Float?
judge_cost Float @default(0)
error String?
created_at DateTime @default(now())
@@index([job_id])
}
// ---------------------------------------------------------------------------
// Workflow Run Tracking
//

1
litellm/__init__.py
View file

@ -246,6 +246,7 @@ use_chat_completions_url_for_anthropic_messages: bool = bool(
# config.yaml.
strip_anthropic_total_tokens: bool = False
anthropic_sse_ping_interval_seconds: float = 15.0
sse_keepalive_ping_interval_seconds: float | None = None
route_all_chat_openai_to_responses: bool = (
os.getenv("LITELLM_ROUTE_ALL_CHAT_OPENAI_TO_RESPONSES", "false").lower() == "true"
) # When True, routes all OpenAI /chat/completions requests through the Responses API bridge

23
litellm/a2a_protocol/litellm_completion_bridge/handler.py
View file

@ -10,7 +10,7 @@ A2A Streaming Events (in order):
4. Status update (kind: "status-update") - Final status "completed" with final=true
"""
from collections.abc import AsyncIterator, Mapping
from collections.abc import AsyncIterator, Callable, Coroutine, Mapping
from typing import Any, Final
import litellm
@ -54,7 +54,7 @@ class A2ACompletionBridgeHandler:
agent_extra_headers: Mapping[str, str] | None,
*,
stream: bool,
) -> Mapping[str, Any]:
) -> Mapping[str, object]:
# Extract message from params
message: Final = params.get("message", {})
@ -63,7 +63,7 @@ class A2ACompletionBridgeHandler:
# Get completion params
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
model: Final = litellm_params.get("model", "agent")
model: Final[str] = litellm_params.get("model", "agent")
# Build full model string if provider specified
# Skip prepending if model already starts with the provider prefix
@ -109,13 +109,16 @@ class A2ACompletionBridgeHandler:
return completion_params
@staticmethod
async def _acompletion(completion_params: Mapping[str, Any]) -> ModelResponse | CustomStreamWrapper:
return await litellm.acompletion(**completion_params)
async def _acompletion(completion_params: Mapping[str, object]) -> ModelResponse | CustomStreamWrapper:
acompletion_fn: Final[Callable[..., Coroutine[object, object, ModelResponse | CustomStreamWrapper]]] = vars(
litellm
)["acompletion"]
return await acompletion_fn(**completion_params)
@staticmethod
async def handle_non_streaming(
request_id: str,
params: dict[str, Any],
params: dict[str, object],
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
@ -296,8 +299,8 @@ class A2ACompletionBridgeHandler:
# Convenience functions that delegate to the class methods
async def handle_a2a_completion(
request_id: str,
params: dict[str, Any],
litellm_params: dict[str, Any],
params: dict[str, object],
litellm_params: dict[str, object],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, object]:
@ -313,8 +316,8 @@ async def handle_a2a_completion(
async def handle_a2a_completion_streaming(
request_id: str,
params: dict[str, Any],
litellm_params: dict[str, Any],
params: dict[str, object],
litellm_params: dict[str, object],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
) -> AsyncIterator[dict[str, object]]:

66
litellm/a2a_protocol/main.py
View file

@ -12,7 +12,8 @@ Provides standalone functions with @client decorator for LiteLLM logging integra
import asyncio
import datetime
import uuid
from collections.abc import AsyncIterator, Coroutine
from collections.abc import AsyncIterator, Coroutine, Mapping
from types import ModuleType
from typing import TYPE_CHECKING, Any, Final, Optional, cast
import litellm
@ -38,12 +39,15 @@ if TYPE_CHECKING:
SendMessageResponse,
SendStreamingMessageRequest,
SendStreamingMessageResponse,
SendStreamingMessageSuccessResponse,
Task,
)
from a2a.types.a2a_pb2 import SendMessageRequest as CoreSendMessageRequest
from a2a.types.a2a_pb2 import StreamResponse as CoreStreamResponse
# Runtime imports — requires a2a-sdk>=1.1.0
A2A_SDK_AVAILABLE = False
_a2a_conversions: Any = None
_a2a_conversions: ModuleType | None = None
try:
from a2a.client import Client, ClientCallContext, ClientConfig, create_client
@ -128,7 +132,7 @@ _A2A_COST_PARAM_KEYS: Final = ("cost_per_query", "input_cost_per_token", "output
def _set_litellm_params_on_logging_obj(
kwargs: dict[str, Any],
litellm_params: dict[str, Any],
litellm_params: Mapping[str, object],
) -> None:
"""
Merge the agent's pricing params into model_call_details["litellm_params"]
@ -150,7 +154,7 @@ def _set_litellm_params_on_logging_obj(
logging_obj.model_call_details["litellm_params"] = {**existing, **cost_params}
def _get_a2a_model_info(a2a_client: Any, kwargs: dict[str, Any]) -> str:
def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: dict[str, Any]) -> str:
"""
Extract agent info and set model/custom_llm_provider for cost tracking.
@ -179,7 +183,7 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: dict[str, Any]) -> str:
return agent_name
def _get_a2a_client_agent_card(a2a_client: Any) -> Optional["AgentCard"]:
def _get_a2a_client_agent_card(a2a_client: "A2AClientType") -> Optional["AgentCard"]:
agent_card = cast(Optional["AgentCard"], getattr(a2a_client, "_litellm_agent_card", None))
if agent_card is not None:
return agent_card
@ -191,9 +195,9 @@ def _get_a2a_client_agent_card(a2a_client: Any) -> Optional["AgentCard"]:
async def _send_message_via_completion_bridge(
request: "SendMessageRequest",
custom_llm_provider: str,
custom_llm_provider: object,
api_base: str | None,
litellm_params: dict[str, Any],
litellm_params: dict[str, object],
agent_extra_headers: dict[str, str] | None = None,
) -> LiteLLMSendMessageResponse:
"""
@ -224,6 +228,20 @@ def _get_a2a_call_context(a2a_client: "A2AClientType") -> Optional["A2ACallConte
return getattr(a2a_client, "_litellm_call_context", None)
def _to_core_send_message_request(request: "SendMessageRequest") -> "CoreSendMessageRequest":
from a2a.compat.v0_3 import conversions
return conversions.to_core_send_message_request(request)
def _to_compat_stream_response(
event: "CoreStreamResponse", request_id: str | int
) -> "SendStreamingMessageSuccessResponse":
from a2a.compat.v0_3 import conversions
return conversions.to_compat_stream_response(event, request_id=request_id)
async def _send_message(a2a_client: "A2AClientType", request: "SendMessageRequest") -> "SendMessageResponse":
"""Send a non-streaming message via a2a-sdk 1.x and return JSON-RPC response."""
if _a2a_conversions is None:
@ -231,17 +249,14 @@ async def _send_message(a2a_client: "A2AClientType", request: "SendMessageReques
"The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk"
)
pb_request: Final = _a2a_conversions.to_core_send_message_request(request)
pb_request: Final = _to_core_send_message_request(request)
last_event = None
async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)):
last_event = event
if last_event is None:
raise RuntimeError("A2A send_message failed: no response received from agent.")
stream_compat: Final = _a2a_conversions.to_compat_stream_response(
last_event,
request_id=request.id,
)
stream_compat: Final = _to_compat_stream_response(last_event, request_id=request.id)
result: Final = stream_compat.result
if not isinstance(result, (Message, Task)):
raise RuntimeError(
@ -306,12 +321,9 @@ async def _stream_messages(
"The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk"
)
pb_request: Final = _a2a_conversions.to_core_send_message_request(request)
pb_request: Final[CoreSendMessageRequest] = _a2a_conversions.to_core_send_message_request(request)
async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)):
compat_chunk = _a2a_conversions.to_compat_stream_response(
event,
request_id=request.id,
)
compat_chunk = _to_compat_stream_response(event, request_id=request.id)
yield SendStreamingMessageResponse(root=compat_chunk)
@ -368,10 +380,10 @@ async def asend_message(
a2a_client: Optional["A2AClientType"] = None,
request: Optional["SendMessageRequest"] = None,
api_base: str | None = None,
litellm_params: dict[str, Any] | None = None,
litellm_params: dict[str, object] | None = None,
agent_id: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
**kwargs: Any,
**kwargs: object,
) -> LiteLLMSendMessageResponse:
"""
Async: Send a message to an A2A agent.
@ -485,7 +497,7 @@ async def asend_message(
response: Final = LiteLLMSendMessageResponse.from_a2a_response(a2a_response, request_id=str(request.id))
# Calculate token usage from request and response
response_dict: Final = a2a_response.model_dump(mode="json", exclude_none=True)
response_dict: Final[dict[str, object]] = a2a_response.model_dump(mode="json", exclude_none=True)
(
prompt_tokens,
completion_tokens,
@ -516,7 +528,7 @@ def send_message(
a2a_client: "A2AClientType",
request: "SendMessageRequest",
**kwargs: Any,
) -> LiteLLMSendMessageResponse | Coroutine[Any, Any, LiteLLMSendMessageResponse]:
) -> LiteLLMSendMessageResponse | Coroutine[object, object, LiteLLMSendMessageResponse]:
"""
Sync: Send a message to an A2A agent.
@ -545,9 +557,9 @@ def _build_streaming_logging_obj(
request: "SendStreamingMessageRequest",
agent_name: str,
agent_id: str | None,
litellm_params: dict[str, Any] | None,
metadata: dict[str, Any] | None,
proxy_server_request: dict[str, Any] | None,
litellm_params: dict[str, object] | None,
metadata: dict[str, object] | None,
proxy_server_request: dict[str, object] | None,
) -> Logging:
"""Build logging object for streaming A2A requests."""
start_time: Final = datetime.datetime.now()
@ -588,10 +600,10 @@ async def asend_message_streaming(
a2a_client: Optional["A2AClientType"] = None,
request: Optional["SendStreamingMessageRequest"] = None,
api_base: str | None = None,
litellm_params: dict[str, Any] | None = None,
litellm_params: dict[str, object] | None = None,
agent_id: str | None = None,
metadata: dict[str, Any] | None = None,
proxy_server_request: dict[str, Any] | None = None,
metadata: dict[str, object] | None = None,
proxy_server_request: dict[str, object] | None = None,
agent_extra_headers: dict[str, str] | None = None,
**kwargs: object,
) -> AsyncIterator[Any]:

13
litellm/constants.py
View file

@ -472,6 +472,8 @@ EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE: Final = float(
### ANTHROPIC CONSTANTS ###
ANTHROPIC_TOKEN_COUNTING_BETA_VERSION = os.getenv("ANTHROPIC_TOKEN_COUNTING_BETA_VERSION", "token-counting-2024-11-01")
ANTHROPIC_SKILLS_API_BETA_VERSION: Final = "skills-2025-10-02"
ANTHROPIC_BATCHES_ROUTE: Final = "/v1/messages/batches"
VERTEX_BATCH_PREDICTION_JOBS_ROUTE: Final = "batchPredictionJobs"
ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES: Final = {
"low": 1,
"medium": 5,
@ -1323,6 +1325,7 @@ LITELLM_METADATA_FIELD: Final = "litellm_metadata"
OLD_LITELLM_METADATA_FIELD: Final = "metadata"
RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name"
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (
@ -1488,6 +1491,9 @@ SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEA
SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = float(
os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5)
)
SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", "300"))
SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", "30"))
SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = int(os.getenv("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", "100000"))
TOOL_SPEND_TOP_TOOLS: Final = 100
SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day")
SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7))
@ -1527,6 +1533,10 @@ APSCHEDULER_REPLACE_EXISTING: Final = os.getenv("APSCHEDULER_REPLACE_EXISTING",
"1",
] # always replace existing jobs
# Width of the window scheduled background jobs are spread across, so they do not all fire
# on one instant on every replica. Tunable per deployment via general_settings.
DEFAULT_STAGGER_WINDOW_SECONDS: Final = 300
# The number of tag entries are higher than number of user, team entries. This leads to a higher QPS.
# This will run tag spcific tasks at a later time to smooth QPS
DAILY_TAG_SPEND_BATCH_MULTIPLIER: Final = 2.3
@ -1735,6 +1745,9 @@ PTU_ROLLUP_LOCK_TTL_SECONDS: Final[int] = 900
# Furthest back the catch-up pass looks for unpriced PTU days when a deployment
# declares no ptu_effective_from, bounding the scan for an open-ended window.
PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90
# Deployments named in the lapsed-window alert before it is truncated, so a fleet-wide
# expiry cannot produce an alert too large for the channel delivering it.
PTU_LAPSED_ALERT_LIMIT: Final[int] = 10
# Slack allowed when deciding a sentinel row is stale. The row's updated_at and the
# run's cutoff are stamped by different hosts, so clock skew between them must not let
# one run delete a charge another just wrote. A stale row is hours old and a concurrent

46
litellm/experimental_mcp_client/client.py
View file

@ -6,10 +6,11 @@ import asyncio
import base64
import os
from collections.abc import Awaitable, Callable, Generator
from datetime import timedelta
from typing import Any, Final, TypeVar
import httpx
from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters
from mcp import ClientSession, McpError, ReadResourceResult, Resource, StdioServerParameters
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
@ -69,6 +70,29 @@ def _first_non_cancelled_cause(exc: BaseException) -> BaseException | None:
return None
_SDK_READ_TIMEOUT_CODE: Final = int(httpx.codes.REQUEST_TIMEOUT)
"""The code the MCP SDK puts on its own elapsed read timeout, an HTTP status in a field that
otherwise carries JSON-RPC error codes."""
def _as_read_timeout(exc: BaseException) -> TimeoutError | None:
"""The session read timeout elapsing, re-expressed as a ``TimeoutError``, or ``None``.
The SDK reports its own elapsed read timeout as ``McpError`` carrying an HTTP status code in a
field that otherwise holds JSON-RPC error codes, and it relays an upstream's JSON-RPC error
through that same class and field. The numeric code alone therefore cannot separate the two, and
an upstream answering with application code 408 would be reported as a gateway timeout it never
caused. The SDK raises its own from inside an ``except TimeoutError``, so the elapsed timeout is
on the context chain, while a relayed error is built from a received message and has no such
chain; that is the discriminator.
"""
if not isinstance(exc, McpError) or exc.error.code != _SDK_READ_TIMEOUT_CODE:
return None
if not isinstance(exc.__context__, TimeoutError):
return None
return TimeoutError(exc.error.message)
TSessionResult = TypeVar("TSessionResult")
@ -347,7 +371,14 @@ class MCPClient:
session_kwargs["elicitation_callback"] = self._elicitation_callback
if self._logging_callback is not None:
session_kwargs["logging_callback"] = self._logging_callback
session_ctx: Final = ClientSession(read_stream, write_stream, **session_kwargs)
# The SDK drops a response stream that ends without a JSON-RPC reply, so nothing else
# ever fails the request.
session_ctx: Final = ClientSession(
read_stream,
write_stream,
read_timeout_seconds=timedelta(seconds=self.timeout),
**session_kwargs,
)
session: Final = await session_ctx.__aenter__()
try:
init_result: Final = await session.initialize()
@ -390,7 +421,16 @@ class MCPClient:
self._last_initialize_instructions = None
transport_ctx, http_client = self._create_transport_context()
return await self._execute_session_operation(transport_ctx, operation)
except Exception:
except Exception as e:
read_timeout: Final = _as_read_timeout(e)
if read_timeout is not None:
verbose_logger.warning(
"MCP client timed out after %ss waiting for %s to answer; the server accepted the "
"request and ended its response stream without a JSON-RPC reply",
self.timeout,
self.server_url or "stdio",
)
raise read_timeout from e
_log: Final = verbose_logger.debug if quiet_on_error else verbose_logger.warning
_log("MCP client run_with_session failed for %s", self.server_url or "stdio")
raise

61
litellm/google_genai/adapters/transformation.py
View file

@ -1,6 +1,8 @@
import json
from collections.abc import AsyncIterator, Iterator
from typing import Any, Final, cast
from typing import Any, Final, TypedDict, cast
from typing_extensions import ReadOnly
from litellm import verbose_logger
from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema
@ -28,6 +30,19 @@ from litellm.types.utils import (
)
class _GenAITextPart(TypedDict, total=False):
text: ReadOnly[str]
class _GenAISystemInstruction(TypedDict, total=False):
parts: ReadOnly[list[_GenAITextPart]]
class _GenAIPart(TypedDict, total=False):
text: ReadOnly[str]
functionCall: ReadOnly[dict[str, object]]
class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
"""
Wrapper for streaming Google GenAI generate_content responses.
@ -36,9 +51,9 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
sent_first_chunk: bool = False
# State tracking for accumulating partial tool calls
accumulated_tool_calls: dict[str, dict[str, Any]]
accumulated_tool_calls: dict[str, dict[str, str]]
def __init__(self, completion_stream: Any):
def __init__(self, completion_stream: object):
self.sent_first_chunk = False
self.accumulated_tool_calls = {}
self._returned_response = False
@ -85,7 +100,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
# After the stream is exhausted, check for any remaining accumulated tool calls
if self.accumulated_tool_calls:
try:
parts: Final = []
parts: Final[list[_GenAIPart]] = []
for (
tool_call_index,
tool_call_data,
@ -94,7 +109,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
# For tool calls with no arguments, accumulated_args will be "", which is not valid JSON.
# We default to an empty JSON object in this case.
parsed_args = json.loads(tool_call_data["arguments"] or "{}")
function_call_part = {
function_call_part: _GenAIPart = {
"functionCall": {
"name": tool_call_data["name"] or "undefined_tool_name",
"args": parsed_args,
@ -110,7 +125,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
tool_call_data["arguments"],
)
if parts:
final_chunk: Final = {
final_chunk: Final[dict[str, object]] = {
"candidates": [
{
"content": {"parts": parts, "role": "model"},
@ -273,9 +288,9 @@ class GoogleGenAIAdapter:
def _add_generic_litellm_params_to_request(
self,
completion_request_dict: dict[str, Any],
completion_request_dict: dict[str, object],
litellm_params: GenericLiteLLMParams | None = None,
) -> dict:
) -> dict[str, object]:
"""Add generic litellm params to request. e.g add api_base, api_key, api_version, etc.
Args:
@ -295,7 +310,7 @@ class GoogleGenAIAdapter:
def translate_completion_output_params_streaming(
self,
completion_stream: Any,
completion_stream: object,
) -> AsyncIterator[bytes] | None:
"""Transform streaming completion output to Google GenAI format"""
google_genai_wrapper: Final = GoogleGenAIStreamWrapper(completion_stream=completion_stream)
@ -307,12 +322,12 @@ class GoogleGenAIAdapter:
tools: list[dict[str, Any]],
) -> list[ChatCompletionToolParam]:
"""Transform Google GenAI tools to OpenAI tools format"""
openai_tools: Final[list[dict[str, Any]]] = []
openai_tools: Final[list[dict[str, object]]] = []
for tool in tools:
if "functionDeclarations" in tool:
for func_decl in tool["functionDeclarations"]:
function_chunk: dict[str, Any] = {
function_chunk: dict[str, object] = {
"name": func_decl.get("name", ""),
}
@ -321,7 +336,7 @@ class GoogleGenAIAdapter:
if "parametersJsonSchema" in func_decl:
function_chunk["parameters"] = func_decl["parametersJsonSchema"]
openai_tool = {"type": "function", "function": function_chunk}
openai_tool: dict[str, object] = {"type": "function", "function": function_chunk}
openai_tools.append(openai_tool)
# normalize the tool schemas
@ -345,7 +360,7 @@ class GoogleGenAIAdapter:
def _transform_contents_to_messages(
self,
contents: list[dict[str, Any]],
system_instruction: dict[str, Any] | None = None,
system_instruction: _GenAISystemInstruction | None = None,
) -> list[AllMessageValues]:
"""Transform Google GenAI contents to OpenAI messages format"""
messages: Final[list[AllMessageValues]] = []
@ -461,7 +476,7 @@ class GoogleGenAIAdapter:
def translate_completion_to_generate_content(
self,
response: ModelResponse,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Transform litellm completion response to Google GenAI generate_content format
@ -490,7 +505,7 @@ class GoogleGenAIAdapter:
parts = [{"text": message_content}] if message_content else []
# Create Google GenAI format response
generate_content_response: Final[dict[str, Any]] = {
generate_content_response: Final[dict[str, object]] = {
"candidates": [
{
"content": {"parts": parts, "role": "model"},
@ -524,7 +539,7 @@ class GoogleGenAIAdapter:
self,
response: ModelResponse | ModelResponseStream,
wrapper: GoogleGenAIStreamWrapper,
) -> dict[str, Any] | None:
) -> dict[str, object] | None:
"""
Transform streaming litellm completion chunk to Google GenAI generate_content format
@ -560,7 +575,7 @@ class GoogleGenAIAdapter:
return None
# Create Google GenAI streaming format response
streaming_chunk: Final[dict[str, Any]] = {
streaming_chunk: Final[dict[str, object]] = {
"candidates": [
{
"content": {"parts": parts, "role": "model"},
@ -597,9 +612,9 @@ class GoogleGenAIAdapter:
def _transform_openai_message_to_google_genai_parts(
self,
message: Any,
) -> list[dict[str, Any]]:
) -> list[_GenAIPart]:
"""Transform OpenAI message to Google GenAI parts format"""
parts: Final[list[dict[str, Any]]] = []
parts: Final[list[_GenAIPart]] = []
# Add text content if present
if hasattr(message, "content") and message.content:
@ -614,7 +629,7 @@ class GoogleGenAIAdapter:
except json.JSONDecodeError:
args = {}
function_call_part = {
function_call_part: _GenAIPart = {
"functionCall": {
"name": tool_call.function.name or "undefined_tool_name",
"args": args,
@ -626,14 +641,14 @@ class GoogleGenAIAdapter:
def _transform_openai_delta_to_google_genai_parts_with_accumulation(
self, delta: Any, wrapper: GoogleGenAIStreamWrapper
) -> list[dict[str, Any]]:
) -> list[_GenAIPart]:
"""Transforms OpenAI delta to Google GenAI parts, accumulating streaming tool calls."""
# 1. Initialize wrapper state if it doesn't exist
if not hasattr(wrapper, "accumulated_tool_calls"):
wrapper.accumulated_tool_calls = {}
parts: Final[list[dict[str, Any]]] = []
parts: Final[list[_GenAIPart]] = []
if hasattr(delta, "content") and delta.content:
parts.append({"text": delta.content})
@ -686,7 +701,7 @@ class GoogleGenAIAdapter:
# The part will be created by a later chunk that brings the name.
if accumulated_name:
# If successful, create the part and clean up
function_call_part = {"functionCall": {"name": accumulated_name, "args": parsed_args}}
function_call_part: _GenAIPart = {"functionCall": {"name": accumulated_name, "args": parsed_args}}
parts.append(function_call_part)
# Remove the completed tool call from the accumulator

112
litellm/integrations/langfuse/langfuse.py
View file

@ -2,8 +2,9 @@
# On success, logs events to Langfuse
import os
import traceback
from collections.abc import Callable
from collections.abc import Callable, Iterable, Mapping
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, cast
from packaging.version import Version
@ -30,6 +31,7 @@ from litellm.types.utils import (
ImageResponse,
ModelResponse,
RerankResponse,
StandardLoggingMetadata,
StandardLoggingPayload,
StandardLoggingPromptManagementMetadata,
TextCompletionResponse,
@ -46,6 +48,11 @@ else:
Langfuse = Any
_DENIED_STEERING_KEYS: Final = frozenset({"headers", "endpoint", "caching_groups", "previous_models"})
_NO_METADATA: Final[Mapping[str, Any]] = MappingProxyType({})
_REDACTED_PROXY_HEADERS: Final[frozenset[str]] = frozenset({"authorization", "cookie", "referer"})
def _extract_cache_read_input_tokens(usage_obj) -> int:
"""
Extract cache_read_input_tokens from usage object.
@ -75,6 +82,22 @@ def _extract_cache_read_input_tokens(usage_obj) -> int:
return cache_read_input_tokens
def _as_steering_flag(value: object) -> bool:
"""A string ``str_to_bool`` does not recognise falls back to its truthiness."""
if isinstance(value, str):
parsed: Final = str_to_bool(value)
return bool(value) if parsed is None else parsed
return bool(value)
def _as_steering_key_sequence(value: object) -> tuple[str, ...]:
if isinstance(value, str):
return tuple(key.strip() for key in value.split(",") if key.strip())
if isinstance(value, Iterable):
return tuple(str(key) for key in value)
return ()
def resolve_langfuse_credentials(
langfuse_public_key=None,
langfuse_secret=None,
@ -496,16 +519,14 @@ class LangFuseLogger:
else []
)
if standard_logging_object is None:
end_user_id = None
prompt_management_metadata: StandardLoggingPromptManagementMetadata | None = None
else:
end_user_id = standard_logging_object["metadata"].get("user_api_key_end_user_id", None)
prompt_management_metadata = cast(
StandardLoggingPromptManagementMetadata | None,
standard_logging_object["metadata"].get("prompt_management_metadata", None),
)
allowlisted_metadata: Final[StandardLoggingMetadata | dict[str, Any]] = (
standard_logging_object["metadata"] if standard_logging_object is not None else _NO_METADATA
)
end_user_id: Final = allowlisted_metadata.get("user_api_key_end_user_id", None)
prompt_management_metadata: Final[StandardLoggingPromptManagementMetadata | None] = cast(
StandardLoggingPromptManagementMetadata | None,
allowlisted_metadata.get("prompt_management_metadata", None),
)
# Clean Metadata before logging - never log raw metadata
# the raw metadata can contain circular references which leads to infinite recursion
@ -524,12 +545,7 @@ class LangFuseLogger:
tags.append(f"{key}:{value}")
# clean litellm metadata before logging
if key in [
"headers",
"endpoint",
"caching_groups",
"previous_models",
]:
if key in _DENIED_STEERING_KEYS:
continue
else:
clean_metadata[key] = value
@ -552,10 +568,10 @@ class LangFuseLogger:
# This allows continuing an existing trace while still returning the correct trace_id
if existing_trace_id is not None:
trace_id = existing_trace_id
update_trace_keys: Final = cast(list, clean_metadata.pop("update_trace_keys", []))
update_trace_keys: Final = _as_steering_key_sequence(clean_metadata.pop("update_trace_keys", ()))
debug: Final = clean_metadata.pop("debug_langfuse", None)
mask_input: Final = clean_metadata.pop("mask_input", False)
mask_output: Final = clean_metadata.pop("mask_output", False)
mask_input: Final = _as_steering_flag(clean_metadata.pop("mask_input", False))
mask_output: Final = _as_steering_flag(clean_metadata.pop("mask_output", False))
# Look for masking function in the dedicated location first (set by scrub_sensitive_keys_in_metadata)
# Fall back to metadata for backwards compatibility
masking_function: Final = litellm_params.get("_langfuse_masking_function") or clean_metadata.pop(
@ -614,19 +630,18 @@ class LangFuseLogger:
trace_params["output"] = output if not mask_output else "redacted-by-litellm"
if debug is True or (isinstance(debug, str) and debug.lower() == "true"):
if "metadata" in trace_params:
# log the raw_metadata in the trace
trace_params["metadata"]["metadata_passed_to_litellm"] = metadata
else:
trace_params["metadata"] = {"metadata_passed_to_litellm": metadata}
debug_metadata: Final = {
key: value for key, value in metadata.items() if isinstance(value, (str, int, float, bool))
}
trace_params["metadata"] = {
**(trace_params.get("metadata") or _NO_METADATA),
"metadata_passed_to_litellm": debug_metadata,
}
cost: Final = kwargs.get("response_cost", None)
verbose_logger.debug("trace: %s", cost)
clean_metadata["litellm_response_cost"] = cost
if standard_logging_object is not None:
hidden_params: Final = standard_logging_object.get("hidden_params", {})
clean_metadata["hidden_params"] = filter_exceptions_from_params(hidden_params)
hidden_params: Final = standard_logging_object.get("hidden_params") if standard_logging_object else None
if (
litellm.langfuse_default_tags is not None
@ -638,22 +653,24 @@ class LangFuseLogger:
tags.append(f"proxy_base_url:{proxy_base_url}")
api_base: Final = litellm_params.get("api_base", None)
if api_base:
clean_metadata["api_base"] = api_base
vertex_location: Final = kwargs.get("vertex_location", None)
if vertex_location:
clean_metadata["vertex_location"] = vertex_location
aws_region_name: Final = kwargs.get("aws_region_name", None)
if aws_region_name:
clean_metadata["aws_region_name"] = aws_region_name
candidate_enrichments: Final = (
("litellm_response_cost", cost, True),
("hidden_params", filter_exceptions_from_params(hidden_params), hidden_params is not None),
("api_base", api_base, bool(api_base)),
("vertex_location", vertex_location, bool(vertex_location)),
("aws_region_name", aws_region_name, bool(aws_region_name)),
("cache_hit", kwargs.get("cache_hit") or False, self._supports_tags() and "cache_hit" in kwargs),
)
enrichments: Final[Mapping[str, Any]] = {
key: value for key, value, include in candidate_enrichments if include
}
if self._supports_tags():
if "cache_hit" in kwargs:
if kwargs["cache_hit"] is None:
kwargs["cache_hit"] = False
clean_metadata["cache_hit"] = kwargs["cache_hit"]
if "cache_hit" in kwargs and kwargs["cache_hit"] is None:
kwargs["cache_hit"] = False # rebind-ok: pre-existing normalization other integrations rely on
if existing_trace_id is None:
trace_params.update({"tags": tags})
@ -666,13 +683,13 @@ class LangFuseLogger:
if headers:
for key, value in headers.items():
# these headers can leak our API keys and/or JWT tokens
if key.lower() not in ["authorization", "cookie", "referer"]:
if key.lower() not in _REDACTED_PROXY_HEADERS:
clean_headers[key] = value
trace: Final[StatefulTraceClient] = self.Langfuse.trace(**trace_params)
# Log provider specific information as a span
log_provider_specific_information_as_span(trace, clean_metadata)
log_provider_specific_information_as_span(trace, enrichments)
# Log guardrail information as a span
self._log_guardrail_information_as_span(
@ -745,7 +762,10 @@ class LangFuseLogger:
"output": output if not mask_output else "redacted-by-litellm",
"usage": usage,
"usage_details": usage_details,
"metadata": log_requester_metadata(clean_metadata),
"metadata": {
**log_requester_metadata(redact_user_api_key_info(metadata=allowlisted_metadata)),
**enrichments,
},
"level": level,
"version": clean_metadata.pop("version", None),
}
@ -1042,7 +1062,7 @@ def _add_prompt_to_generation_params(
def log_provider_specific_information_as_span(
trace,
clean_metadata,
clean_metadata: Mapping[str, Any],
):
"""
Logs provider-specific information as spans.
@ -1082,7 +1102,7 @@ def log_provider_specific_information_as_span(
)
def log_requester_metadata(clean_metadata: dict):
def log_requester_metadata(clean_metadata: Mapping[str, Any]):
returned_metadata: Final = {}
requester_metadata: Final = clean_metadata.get("requester_metadata") or {}
for k, v in clean_metadata.items():

10
litellm/integrations/langfuse/langfuse_otel.py
View file

@ -90,7 +90,6 @@ class LangfuseOtelLogger(OpenTelemetry):
"generation_name": LangfuseSpanAttributes.GENERATION_NAME,
"generation_id": LangfuseSpanAttributes.GENERATION_ID,
"parent_observation_id": LangfuseSpanAttributes.PARENT_OBSERVATION_ID,
"version": LangfuseSpanAttributes.GENERATION_VERSION,
"mask_input": LangfuseSpanAttributes.MASK_INPUT,
"mask_output": LangfuseSpanAttributes.MASK_OUTPUT,
"trace_user_id": LangfuseSpanAttributes.TRACE_USER_ID,
@ -99,13 +98,18 @@ class LangfuseOtelLogger(OpenTelemetry):
"trace_name": LangfuseSpanAttributes.TRACE_NAME,
"trace_id": LangfuseSpanAttributes.TRACE_ID,
"trace_metadata": LangfuseSpanAttributes.TRACE_METADATA,
"trace_version": LangfuseSpanAttributes.TRACE_VERSION,
"trace_release": LangfuseSpanAttributes.TRACE_RELEASE,
"trace_release": LangfuseSpanAttributes.RELEASE,
"existing_trace_id": LangfuseSpanAttributes.EXISTING_TRACE_ID,
"update_trace_keys": LangfuseSpanAttributes.UPDATE_TRACE_KEYS,
"debug_langfuse": LangfuseSpanAttributes.DEBUG_LANGFUSE,
}
version: Final = (
metadata.get("trace_version") if metadata.get("trace_version") is not None else metadata.get("version")
)
if version is not None:
safe_set_attribute(span, LangfuseSpanAttributes.VERSION.value, version)
for key, enum_attr in mapping.items():
if key in metadata and metadata[key] is not None:
value = metadata[key]

8
litellm/integrations/otel/presets/langfuse.py
View file

@ -42,9 +42,7 @@ def langfuse_dynamic_headers(params: StandardCallbackDynamicParams) -> dict[str,
public_key: Final = params.get("langfuse_public_key")
secret_key: Final = params.get("langfuse_secret_key")
if public_key and secret_key:
return {
"Authorization": _V1Langfuse._get_langfuse_authorization_header(
public_key=public_key, secret_key=secret_key
)
}
return _V1Langfuse._build_langfuse_otel_headers(
_V1Langfuse._get_langfuse_authorization_header(public_key=public_key, secret_key=secret_key)
)
return {}

85
litellm/integrations/rubrik.py
View file

@ -6,12 +6,13 @@ import random
import time
import uuid
from collections import Counter
from collections.abc import Mapping, Sequence
from collections.abc import Awaitable, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict
import httpx
from typing_extensions import Never, ReadOnly
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
@ -48,7 +49,20 @@ _WEBHOOK_PATH_PROMPT_MODERATION: Final = "/v1/before_prompt/openai/v1"
_WEBHOOK_PATH_LOGGING_BATCH: Final = "/v1/litellm/batch"
_MAX_QUEUE_SIZE: Final = 10_000
_DROP_WARNING_INTERVAL_SECONDS: Final = 60.0
_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
_EMPTY_MAPPING: Final[Mapping[str, Never]] = MappingProxyType({})
class _ServiceToolCall(TypedDict):
id: ReadOnly[str]
class _ServiceMessage(TypedDict, total=False):
content: ReadOnly[str]
tool_calls: ReadOnly[Sequence[_ServiceToolCall]]
class _ServiceChoice(TypedDict, total=False):
message: ReadOnly[_ServiceMessage]
class _MalformedToolBlockingResponseError(Exception):
@ -143,7 +157,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
else {"Content-Type": "application/json"}
)
self._periodic_flush_task: asyncio.Task[Any] | None = self._start_periodic_flush_task()
self._periodic_flush_task: asyncio.Task[None] | None = self._start_periodic_flush_task()
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
@ -191,7 +205,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
params={"timeout": httpx.Timeout(5.0, connect=2.0)},
)
def _start_periodic_flush_task(self) -> asyncio.Task[Any] | None:
def _start_periodic_flush_task(self) -> asyncio.Task[None] | None:
"""Start the periodic flush task only when an event loop is already running."""
try:
loop: Final = asyncio.get_running_loop()
@ -212,7 +226,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
Closing them here would close the shared connection pool for every
other logger instance; let LiteLLM manage their lifecycle instead.
"""
task: Final = getattr(self, "_periodic_flush_task", None)
task: Final[asyncio.Task[None] | None] = getattr(self, "_periodic_flush_task", None)
if task is not None:
task.cancel()
@ -253,7 +267,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
@staticmethod
async def _guarded(
coro: Any,
coro: Awaitable[GenericGuardrailAPIInputs],
inputs: GenericGuardrailAPIInputs,
label: str,
) -> GenericGuardrailAPIInputs:
@ -400,7 +414,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
request_data["_rubrik_logging_obj"] = logging_obj
@staticmethod
def _normalize_tool_calls(tool_calls: Any) -> tuple[ChatCompletionMessageToolCall, ...]:
def _normalize_tool_calls(tool_calls: Sequence[object]) -> tuple[ChatCompletionMessageToolCall, ...]:
"""Convert tool_calls from inputs to ChatCompletionMessageToolCall objects."""
return tuple(RubrikLogger._normalize_tool_call(tc) for tc in tool_calls)
@ -427,7 +441,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
raise TypeError(f"Cannot normalize tool_call of type {type(tc).__name__}: {tc!r}")
@staticmethod
def _join_texts(texts: Any) -> str:
def _join_texts(texts: Sequence[str] | None) -> str:
"""Join response text segments into the single content string the
webhook evaluates. Empty when there is no assistant text."""
if not texts:
@ -439,14 +453,14 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
tool_calls: Sequence[ChatCompletionMessageToolCall],
content: str,
request_id: str | None,
) -> Mapping[str, Any]:
) -> Mapping[str, object]:
"""Build an OpenAI ChatCompletion-format dict (assistant text + tool
calls) for the after_completion webhook.
``content`` is sent so the webhook can moderate the response text;
``None`` when the assistant produced no text (tool-call-only response).
"""
message: Final[dict[str, Any]] = {
message: Final[dict[str, object]] = {
"role": "assistant",
"content": content or None,
}
@ -467,7 +481,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
}
@staticmethod
def _flatten_messages_for_moderation(messages: Any) -> tuple[Mapping[str, Any], ...]:
def _flatten_messages_for_moderation(messages: Sequence[object] | None) -> tuple[Mapping[str, Any], ...]:
"""Collapse each message's content to a plain string for the webhook.
litellm normalizes Anthropic ``/v1/messages`` requests to OpenAI shape,
@ -506,8 +520,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
@staticmethod
def _build_prompt_moderation_payload(
inputs: GenericGuardrailAPIInputs,
request_data: Mapping[str, Any],
) -> Mapping[str, Any]:
request_data: Mapping[str, object],
) -> Mapping[str, object]:
"""Build the bare OpenAI request the before_prompt webhook consumes.
Unlike the after_completion envelope, this endpoint takes a raw OpenAI
@ -516,7 +530,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
``/v1/messages`` requests too. Optional fields are sent only when
present so the payload stays clean.
"""
payload: Final[dict[str, Any]] = {
payload: Final[dict[str, object]] = {
"model": inputs.get("model") or request_data.get("model") or "",
"messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")),
}
@ -540,8 +554,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
@staticmethod
def _extract_request_data(
call_details: Mapping[str, Any],
request_data: Mapping[str, Any] | None,
) -> Mapping[str, Any]:
request_data: Mapping[str, object] | None,
) -> Mapping[str, object]:
"""Extract original request data from model_call_details for the
response moderation service envelope.
@ -576,7 +590,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
}
@staticmethod
def _sanitize_proxy_server_request(proxy_server_request: Any) -> Any:
def _sanitize_proxy_server_request(proxy_server_request: object) -> object:
"""Allowlist only routing fields (``url``, ``method``) when forwarding
``proxy_server_request`` to an external webhook, dropping inbound
``headers`` (Authorization, Cookie, x-api-key, ...) and the raw
@ -586,17 +600,18 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
return {key: proxy_server_request[key] for key in ("url", "method") if key in proxy_server_request}
@staticmethod
def _resolve_model(request_data: Mapping[str, Any], call_details: Mapping[str, Any]) -> str:
def _resolve_model(request_data: Mapping[str, object], call_details: Mapping[str, str]) -> str:
"""Get the model name for the ModifyResponseException."""
response: Final = request_data.get("response")
if response and hasattr(response, "model"):
return response.model or "unknown"
response_model: Final[str | None] = getattr(response, "model", None)
return response_model or "unknown"
return call_details.get("model", "unknown")
# -- Logging hooks ---------------------------------------------------------
@staticmethod
def _correlation_id(call_details: Mapping[str, Any], request_data: Mapping[str, Any] | None = None) -> str | None:
def _correlation_id(call_details: Mapping[str, str], request_data: Mapping[str, str] | None = None) -> str | None:
"""The id that joins a blocked request's two S3 logs by filename: the
moderation (``_blocking``) log and the failure (response) log.
@ -610,7 +625,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
return call_details.get("litellm_call_id") or (request_data or _EMPTY_MAPPING).get("litellm_call_id")
@classmethod
def _apply_correlation_id(cls, payload: dict[str, Any], source: Mapping[str, Any]) -> None:
def _apply_correlation_id(cls, payload: dict[str, object], source: Mapping[str, str]) -> None:
"""Pin ``payload["id"]`` to ``litellm_call_id`` in place so this log
shares its S3 filename id with the moderation (``_blocking``) and
failure logs for the same request -- for every provider.
@ -630,7 +645,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
payload["id"] = correlated
@staticmethod
def _prepend_system_prompt(payload: dict[str, Any], source: Mapping[str, Any]) -> None:
def _prepend_system_prompt(payload: dict[str, object], source: Mapping[str, object]) -> None:
"""Prepend ``source["system"]`` onto ``payload["messages"]``.
Builds a NEW messages list rather than mutating ``payload["messages"]``
@ -658,7 +673,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
exc_info=True,
)
async def _prepare_log_payload(self, kwargs: Mapping[str, Any], event_type: str) -> StandardLoggingPayload | None:
async def _prepare_log_payload(
self, kwargs: Mapping[str, object], event_type: str
) -> StandardLoggingPayload | None:
"""Shared logic for success logging (sampled)."""
if random.random() > self.sampling_rate:
verbose_logger.debug("Skipping Rubrik %s logging (sampling_rate=%s)", event_type, self.sampling_rate)
@ -697,7 +714,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
self._dropped_since_warning = 0
self._last_drop_warning_time = now
async def _enqueue_log_event(self, kwargs: Mapping[str, Any], event_type: str):
async def _enqueue_log_event(self, kwargs: Mapping[str, object], event_type: str):
try:
payload: Final = await self._prepare_log_payload(kwargs, event_type)
if payload is None:
@ -862,7 +879,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
base: Final = call_details.get("standard_logging_object")
if base is not None:
payload: dict = safe_deep_copy(base)
payload: dict[str, object] = safe_deep_copy(base)
else:
verbose_logger.debug(
"Rubrik: standard_logging_object not yet on model_call_details "
@ -908,7 +925,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
cls,
call_details: Mapping[str, Any],
user_api_key_dict: "UserAPIKeyAuth",
) -> dict[str, Any]:
) -> dict[str, object]:
# Convert datetime to a Unix float so json.dumps can serialize it.
# httpx's json= parameter uses stdlib json.dumps with no custom encoder.
_raw_start: Final = call_details.get("start_time")
@ -996,7 +1013,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
# -- Webhook services ------------------------------------------------------
async def _post_json(self, endpoint: str, payload: Mapping[str, Any], service_name: str) -> Mapping[str, Any]:
async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> Mapping[str, Any]:
"""POST ``payload`` to a Rubrik webhook and return its dict response.
Raises:
@ -1010,7 +1027,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
headers=self._headers,
)
http_response.raise_for_status()
result: Final = http_response.json()
result: Final[object] = http_response.json()
if not isinstance(result, dict):
raise TypeError(
f"{service_name} returned non-dict JSON "
@ -1021,8 +1038,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
async def _post_to_response_moderation_endpoint(
self,
response_data: Mapping[str, Any],
request_data: Mapping[str, Any],
response_data: Mapping[str, object],
request_data: Mapping[str, object],
) -> Mapping[str, Any]:
"""Post the ``{request, response}`` envelope to the after_completion
webhook and return its (possibly rewritten) response.
@ -1039,7 +1056,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
"Response moderation service",
)
async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, Any]) -> Mapping[str, Any]:
async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> Mapping[str, Any]:
"""Post a bare OpenAI request to the before_prompt webhook.
Returns ``{}`` (passthrough) or a synthetic chat.completion (block).
@ -1054,7 +1071,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
chat.completion whose ``choices[0].message.content`` is the refusal
explanation.
"""
choices: Final = service_response.get("choices")
choices: Final[Sequence[_ServiceChoice] | None] = service_response.get("choices")
if not choices:
return None
message: Final = choices[0].get("message") or _EMPTY_MAPPING
@ -1086,7 +1103,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
Expects service_response in OpenAI chat completion format:
{"choices": [{"message": {"tool_calls": [...], "content": "..."}}]}
"""
choices: Final = service_response.get("choices") or ()
choices: Final[Sequence[_ServiceChoice]] = service_response.get("choices") or ()
if not choices:
raise _MalformedToolBlockingResponseError("Response moderation service returned empty response")

563
litellm/integrations/shadow_eval_logger.py Normal file
View file

@ -0,0 +1,563 @@
"""Shadow Eval Logger: samples a shadowed key's successful chat requests, duplicates each
through the auto-router in a detached task, blind-judges real vs shadow, and appends one
``LiteLLM_ShadowEvalAttempt`` row (verdict or error) as the feature's only hot-path write.
Counts, status, and spend derive from those rows at read time, so nothing can disagree
across pods or stop races; the hook reads active jobs through a short-TTL cache."""
import asyncio
import hashlib
import random
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from pydantic import BaseModel
from litellm._logging import verbose_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs
from litellm.litellm_core_utils.internal_call_metadata import sanitized_forwardable_call_metadata
from litellm.litellm_core_utils.llm_judge import (
default_router_provider,
extract_text_from_content,
judge_acompletion,
parse_json_verdict,
)
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
from litellm.types.utils import SHADOW_EVAL_JUDGE_CALL_ORIGIN, SHADOW_EVAL_ROUTER_CALL_ORIGIN
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
from litellm.router import Router
from litellm.types.utils import StandardLoggingPayload
# A job starting, stopping, or hitting its turn budget propagates to sampling within one
# TTL; the turn budget can overshoot by at most one TTL of in-flight samples per pod.
_JOBS_CACHE_TTL_SECONDS: Final = 10
# Concurrent shadow+judge pipelines per pod: a traffic spike turns into skipped samples
# rather than an unbounded task pileup.
_MAX_CONCURRENT_SHADOW_TASKS: Final = 16
# Total character budget for the judge's user prompt, however long the conversation and
# the two responses are, so the prompt can never overflow a judge model's context window.
_MAX_JUDGE_RESPONSE_CHARS: Final = 8_000
_MAX_JUDGE_PROMPT_CHARS: Final = 24_000
# The judge answers with a small JSON object; a tighter budget truncates the JSON
# mid-object and the attempt is lost to an error row.
JUDGE_MAX_OUTPUT_TOKENS: Final = 500
_MAX_ERROR_CHARS: Final = 500
_EMPTY_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
_SAMPLED_CALL_TYPES: Final = frozenset({"completion", "acompletion"})
PAIRWISE_JUDGE_SYSTEM_PROMPT: Final = """You are an impartial quality judge comparing two responses to the same conversation.
The responses are labeled A and B in random order. You do not know which system produced which.
Criteria: correctness, completeness, clarity, conciseness.
Return ONLY valid JSON in this exact format, no other text:
{
"preference": "A" | "B" | "tie",
"confidence": <0.0 to 1.0>,
"reasoning": "<one sentence>"
}"""
class PairwiseVerdict(BaseModel):
"""The judge's blind A/B verdict, validated at the parse boundary."""
preference: str = "tie"
confidence: float = 0.0
def _sample_hits(request_id: str, job_id: str, percentage: float) -> bool:
"""Deterministically decide whether a request falls in the shadowed slice: hash-based
rather than random so retries sample the same way and pods agree without coordination."""
digest: Final = hashlib.sha256(f"{job_id}:{request_id}".encode()).digest()
bucket: Final = int.from_bytes(digest[:8], "big") / float(2**64)
return bucket * 100.0 < percentage
def _judge_call_cost(response: object) -> float:
"""Price a judge call, treating an unmapped judge model as free rather than fatal."""
import litellm
try:
return litellm.completion_cost(completion_response=response) or 0.0
except Exception: # noqa: BLE001 # unmapped judge model: the verdict still counts, cost stays 0
return 0.0
def _unmask_preference(raw_preference: str, real_is_a: bool) -> str:
"""Map the judge's blind A/B/tie verdict back to real/shadow/tie."""
normalized: Final = raw_preference.strip().lower()
if normalized == "a":
return "real" if real_is_a else "shadow"
if normalized == "b":
return "shadow" if real_is_a else "real"
return "tie"
def _judge_user_prompt(conversation: str, response_a: str, response_b: str) -> str:
"""The judge prompt under one total character budget: each response is capped, and
the conversation tail gets whatever budget the responses left over."""
a: Final = response_a[:_MAX_JUDGE_RESPONSE_CHARS]
b: Final = response_b[:_MAX_JUDGE_RESPONSE_CHARS]
conversation_budget: Final = _MAX_JUDGE_PROMPT_CHARS - len(a) - len(b)
return (
f"Conversation:\n{conversation[-conversation_budget:]}\n\n"
f"Response A:\n{a}\n\n"
f"Response B:\n{b}\n\n"
"Which response is better?"
)
async def _key_or_team_is_over_budget(metadata: Mapping[str, object]) -> bool:
"""Whether the shadowed key or its team is over budget, decided by the same owners
the request path uses, so counter keys and thresholds can never drift from auth's.
Advisory and fail-open: real traffic on an over-budget key is already rejected at
auth (so nothing reaches the success hook), and this gate only closes the race
where the key crosses its budget while a request is in flight.
"""
try:
from litellm.exceptions import BudgetExceededError
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import (
_team_max_budget_check,
_virtual_key_max_budget_check,
get_team_object,
)
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
except ImportError:
return False
auth: Final = metadata.get("user_api_key_auth")
if not isinstance(auth, UserAPIKeyAuth):
return False
try:
await _virtual_key_max_budget_check(valid_token=auth, proxy_logging_obj=proxy_logging_obj)
if auth.team_id:
team: Final = await get_team_object(
team_id=auth.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
check_cache_only=True,
)
await _team_max_budget_check(team_object=team, valid_token=auth, proxy_logging_obj=proxy_logging_obj)
except BudgetExceededError:
return True
except Exception as e: # noqa: BLE001 # advisory gate: a failed read must not block sampling
verbose_logger.debug("shadow_eval: budget read failed: %s", e)
return False
def _request_was_routed_by(request_metadata: Mapping[str, object], router_name: str) -> bool:
"""Duplicating a request the shadowed router already served compares the router to
itself: guaranteed ties, judge spend for zero information."""
decision: Final = request_metadata.get("routing_decision")
if not isinstance(decision, Mapping):
return False
return decision.get("router_model_name") == router_name
@dataclass(frozen=True, slots=True)
class _CallFailure:
"""A shadow or judge call that produced no usable response. cost carries any judge
spend the failed attempt still billed, so job-level judge_spend never undercounts."""
error: str
cost: float = 0.0
@dataclass(frozen=True, slots=True)
class _ShadowResponse:
"""A successful shadow call, with what the attempt row records."""
text: str
model: str
tier: str | None
@dataclass(frozen=True, slots=True)
class _JudgeVerdict:
"""A parsed judge verdict, unmasked back to real/shadow/tie."""
preference: str
confidence: float
cost: float
@dataclass(frozen=True, slots=True)
class ActiveShadowEvalJob:
"""One active job as the sampling path needs it: immutable config plus the attempt
count as of the cache fill (the turn budget's staleness is bounded by the cache TTL)."""
id: str
router_name: str
shadow_percentage: float
judge_model: str
max_turns: int
ends_at: datetime
attempts: int
def _as_utc(value: datetime) -> datetime:
return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value
_jobs_cache: Final = InMemoryCache(max_size_in_memory=4, default_ttl=_JOBS_CACHE_TTL_SECONDS)
_JOBS_CACHE_KEY: Final = "shadow_eval:active_jobs"
class ShadowEvalLogger(CustomLogger):
"""Fires blind pairwise shadow evaluations for keys with an active shadow-eval job."""
def __init__(
self,
router_provider: Callable[[], "Router | None"] | None = None,
prisma_provider: Callable[[], "PrismaClient | None"] | None = None,
jobs_cache: InMemoryCache | None = None,
) -> None:
"""Providers are callables so the proxy's lazily-initialized globals are resolved
at call time, not at logger construction."""
self._router_provider = router_provider or default_router_provider
self._prisma_provider = prisma_provider or _default_prisma_provider
self._jobs_cache = jobs_cache or _jobs_cache
self._inflight_shadow_tasks: int = 0
# Starts per job since the last cache fill, never decremented within a
# generation; the refill absorbs written rows and resets.
self._job_starts: dict[str, int] = {} # mutable-ok: per-generation counter
async def _active_jobs(self) -> Mapping[str, ActiveShadowEvalJob]:
"""Active jobs by api_key_id, cache-first. A DB fault returns empty without
caching, so sampling pauses for that request and the next one retries."""
cached: Final = await self._jobs_cache.async_get_cache(_JOBS_CACHE_KEY)
if cached is not None:
return cached # pyright: ignore[reportReturnType] # cache stores exactly this mapping shape
prisma: Final = self._prisma_provider()
if prisma is None:
return _EMPTY_JOBS
try:
records: Final = await prisma.db.litellm_shadowevaljob.find_many(
where={ # mutable-ok: Prisma filter
"stopped_at": None,
"ends_at": {"gt": datetime.now(timezone.utc)}, # mutable-ok: Prisma filter
},
)
grouped: Final = (
await prisma.db.litellm_shadowevalattempt.group_by(
by=["job_id"],
count=True,
where={"job_id": {"in": [str(record.id) for record in records]}}, # mutable-ok: Prisma filter
)
if records
else ()
)
attempt_counts: Final = {str(row["job_id"]): int(row["_count"]["_all"]) for row in grouped or []}
jobs: Final = {
str(record.api_key_id): ActiveShadowEvalJob(
id=str(record.id),
router_name=str(record.router_name),
shadow_percentage=float(record.shadow_percentage),
judge_model=str(record.judge_model),
max_turns=int(record.max_turns),
ends_at=_as_utc(record.ends_at),
attempts=attempt_counts.get(str(record.id), 0),
)
for record in records or []
}
await self._jobs_cache.async_set_cache(_JOBS_CACHE_KEY, jobs)
self._job_starts = {} # rebind-ok: new generation, counts absorbed into the fill
return jobs
except Exception as e: # noqa: BLE001 # a DB blip must never break request logging
verbose_logger.debug("shadow_eval: active-job read failed: %s", e)
return _EMPTY_JOBS
#### hook ####
async def async_log_success_event(
self,
kwargs: Mapping[str, object],
response_obj: object,
start_time: object,
end_time: object,
) -> None:
try:
payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") # pyright: ignore[reportAssignmentType] # untyped callback kwargs
if payload is None:
return
raw_meta: Final = get_litellm_metadata_from_kwargs(dict(kwargs)) # mutable-ok: helper needs dict
request_metadata: Final = raw_meta if isinstance(raw_meta, Mapping) else _EMPTY_METADATA
if request_metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY):
return # internal sub-call (our own shadow/judge, a classifier), not user traffic
# redaction rewrites logged content before callbacks run, so this hook
# only ever sees placeholders for a redacted request
if should_redact_message_logging(dict(kwargs)): # mutable-ok: predicate takes a plain dict
return
metadata: Final = payload.get("metadata") or _EMPTY_METADATA
api_key_hash: Final = metadata.get("user_api_key_hash")
if not api_key_hash:
return
job: Final = (await self._active_jobs()).get(str(api_key_hash))
if job is None:
return
if datetime.now(timezone.utc) >= job.ends_at:
return
if job.attempts + self._job_starts.get(job.id, 0) >= job.max_turns:
return
request_id: Final = payload.get("id") or ""
if not request_id:
return
if not _sample_hits(request_id, job.id, job.shadow_percentage):
return
if payload.get("call_type") not in _SAMPLED_CALL_TYPES:
return # only known chat-shaped traffic is comparable; unknown or missing types fail closed
if _request_was_routed_by(request_metadata, job.router_name):
return
if self._inflight_shadow_tasks >= _MAX_CONCURRENT_SHADOW_TASKS:
return
raw_messages: Final = kwargs.get("messages")
self._job_starts[job.id] = self._job_starts.get(job.id, 0) + 1
self._inflight_shadow_tasks += 1
task: Final = asyncio.create_task(
self._run_shadow_eval(
job=job,
request_id=request_id,
messages=tuple(m for m in raw_messages if isinstance(m, Mapping))
if isinstance(raw_messages, Sequence)
else (),
response_obj=response_obj,
real_model=payload.get("model") or "",
model_parameters=MappingProxyType(
dict(payload.get("model_parameters") or {}) # mutable-ok: frozen snapshot
),
parent_metadata=MappingProxyType(dict(request_metadata)), # mutable-ok: frozen snapshot
)
)
task.add_done_callback(self._release_shadow_slot)
except Exception as e: # noqa: BLE001 # logging hooks must never fail the request
verbose_logger.debug("shadow_eval: failed to schedule task: %s", e)
def _release_shadow_slot(self, _task: "asyncio.Task[None]") -> None:
self._inflight_shadow_tasks -= 1
#### the detached pipeline: one attempt row per sampled request, verdict or error ####
async def _run_shadow_eval(
self,
job: ActiveShadowEvalJob,
request_id: str,
messages: Sequence[Mapping[str, object]],
response_obj: object,
real_model: str,
model_parameters: Mapping[str, object],
parent_metadata: Mapping[str, object],
) -> None:
"""Budget gate -> shadow call -> blind judge -> one attempt row. The prisma gate
sits above the dispatch so no provider spend happens without a place to record
the outcome, and the budget read lives here rather than in the success hook."""
prisma: Final = self._prisma_provider()
try:
if prisma is None:
return
real_text: Final = self._extract_response_text(response_obj)
if not real_text or not messages:
return
if await _key_or_team_is_over_budget(parent_metadata):
return
shadow: Final = await self._call_router_shadow(job.router_name, messages, model_parameters, parent_metadata)
if isinstance(shadow, _CallFailure):
await self._record_attempt(prisma, job, request_id, outcome="error", error=shadow.error)
return
verdict: Final = await self._call_judge(
judge_model=job.judge_model,
messages=messages,
real_text=real_text,
shadow_text=shadow.text,
parent_metadata=parent_metadata,
)
if isinstance(verdict, _CallFailure):
await self._record_attempt(
prisma,
job,
request_id,
outcome="error",
error=verdict.error,
shadow=shadow,
judge_cost=verdict.cost,
)
return
await self._record_attempt(
prisma,
job,
request_id,
outcome=verdict.preference,
shadow=shadow,
real_model=real_model,
confidence=verdict.confidence,
judge_cost=verdict.cost,
)
except Exception as e: # noqa: BLE001 # detached task: record what happened, never raise
verbose_logger.debug("shadow_eval: pipeline failed for %s: %s", request_id, e)
await self._record_attempt(prisma, job, request_id, outcome="error", error=f"pipeline error: {e}")
@staticmethod
async def _record_attempt(
prisma: "PrismaClient | None",
job: ActiveShadowEvalJob,
request_id: str,
*,
outcome: str,
shadow: _ShadowResponse | None = None,
real_model: str = "",
confidence: float | None = None,
judge_cost: float = 0.0,
error: str | None = None,
) -> None:
if prisma is None:
return
try:
await prisma.db.litellm_shadowevalattempt.create(
data={ # mutable-ok: Prisma payload
"job_id": job.id,
"request_id": request_id,
"outcome": outcome,
"tier": shadow.tier if shadow else None,
"real_model": real_model or None,
"shadow_model": shadow.model if shadow else None,
"confidence": confidence,
"judge_cost": judge_cost,
"error": error[:_MAX_ERROR_CHARS] if error else None,
}
)
except Exception as e: # noqa: BLE001 # a lost row degrades sample size, nothing can disagree with it
verbose_logger.debug("shadow_eval: attempt write failed for %s: %s", request_id, e)
async def _call_router_shadow(
self,
router_name: str,
messages: Sequence[Mapping[str, object]],
model_parameters: Mapping[str, object],
parent_metadata: Mapping[str, object],
) -> "_ShadowResponse | _CallFailure":
"""Send the prompt through the auto-router being evaluated. The metadata carries
the shadowed key's identity (spend attribution) and receives the router's routing
decision write-back, read back for tier attribution."""
router: Final = self._router_provider()
if router is None:
return _CallFailure("no router configured on this pod")
shadow_metadata: Final[dict[str, object]] = ( # mutable-ok: router writes its routing decision back
sanitized_forwardable_call_metadata(parent_metadata, SHADOW_EVAL_ROUTER_CALL_ORIGIN)
)
shadow_params: Final = { # mutable-ok: splatted as kwargs
k: v for k, v in model_parameters.items() if k not in ("stream", "metadata")
}
try:
response: Final = await router.acompletion(
model=router_name,
messages=messages, # pyright: ignore[reportArgumentType] # snapshot of the SDK's own message dicts
metadata=shadow_metadata,
num_retries=0,
fallbacks=[], # mutable-ok: SDK kwarg; a failed shadow is a recorded error, never a spend multiplier
**shadow_params,
)
except Exception as e: # noqa: BLE001 # provider errors become error rows, not crashes
verbose_logger.debug("shadow_eval: router call failed: %s", e)
return _CallFailure(f"shadow router call failed: {e}")
text: Final = self._extract_response_text(response)
if not text:
return _CallFailure("shadow router returned an empty response")
raw_decision: Final = shadow_metadata.get("routing_decision")
routing_decision: Final = raw_decision if isinstance(raw_decision, Mapping) else _EMPTY_METADATA
raw_tier: Final = routing_decision.get("tier_label") or routing_decision.get("tier")
return _ShadowResponse(
text=text,
model=str(getattr(response, "model", None) or routing_decision.get("routed_model") or ""),
tier=str(raw_tier) if raw_tier is not None else None,
)
async def _call_judge(
self,
judge_model: str,
messages: Sequence[Mapping[str, object]],
real_text: str,
shadow_text: str,
parent_metadata: Mapping[str, object],
) -> "_JudgeVerdict | _CallFailure":
"""Blind pairwise judge with A/B labels randomized to cancel position bias."""
real_is_a: Final = random.random() < 0.5
response_a: Final = real_text if real_is_a else shadow_text
response_b: Final = shadow_text if real_is_a else real_text
conversation: Final = "\n".join(
f"{str(m.get('role', 'user')).upper()}: {extract_text_from_content(m.get('content'))}"
for m in messages
if m.get("content") is not None
)
judge_metadata: Final = sanitized_forwardable_call_metadata(parent_metadata, SHADOW_EVAL_JUDGE_CALL_ORIGIN)
judge_messages: Final = [ # mutable-ok: SDK takes a list
{"role": "system", "content": PAIRWISE_JUDGE_SYSTEM_PROMPT}, # mutable-ok: SDK message
{
"role": "user",
"content": _judge_user_prompt(conversation, response_a, response_b),
}, # mutable-ok: SDK message
]
try:
response: Final = await judge_acompletion(
self._router_provider(),
judge_model,
judge_messages, # pyright: ignore[reportArgumentType] # plain SDK message dicts
temperature=0,
max_tokens=JUDGE_MAX_OUTPUT_TOKENS,
metadata=judge_metadata,
)
except Exception as e: # noqa: BLE001 # judge outages become error rows, not crashes
verbose_logger.debug("shadow_eval: judge call failed: %s", e)
return _CallFailure(f"judge call failed: {e}")
try:
raw: Final = response["choices"][0]["message"]["content"] or ""
verdict: Final = PairwiseVerdict.model_validate(parse_json_verdict(raw))
except Exception as e: # noqa: BLE001 # malformed verdicts become error rows
verbose_logger.debug("shadow_eval: unparseable judge verdict: %s", e)
return _CallFailure(f"unparseable judge verdict: {e}", cost=_judge_call_cost(response))
return _JudgeVerdict(
preference=_unmask_preference(verdict.preference, real_is_a),
confidence=max(0.0, min(1.0, verdict.confidence)),
cost=_judge_call_cost(response),
)
@staticmethod
def _extract_response_text(response_obj: object) -> str:
"""Extract the assistant's text from a ModelResponse-shaped object or dict."""
try:
content: Final = (
response_obj["choices"][0]["message"]["content"]
if isinstance(response_obj, Mapping)
else response_obj.choices[0].message.content # pyright: ignore[reportAttributeAccessIssue] # duck-typed ModelResponse
)
except (AttributeError, KeyError, IndexError, TypeError):
return ""
return extract_text_from_content(content)
_EMPTY_JOBS: Final[Mapping[str, ActiveShadowEvalJob]] = MappingProxyType({})
def _default_prisma_provider() -> "PrismaClient | None":
try:
from litellm.proxy.proxy_server import prisma_client
except ImportError:
return None
return prisma_client

38
litellm/interactions/litellm_responses_transformation/handler.py
View file

@ -2,8 +2,8 @@
Handler for transforming interactions API requests to litellm.responses requests.
"""
from collections.abc import AsyncIterator, Coroutine, Iterator
from typing import Any, Final, cast
from collections.abc import AsyncIterator, Callable, Coroutine, Iterator
from typing import Any, Final
import litellm
from litellm.interactions.litellm_responses_transformation.streaming_iterator import (
@ -37,7 +37,7 @@ class LiteLLMResponsesInteractionsHandler:
) -> (
InteractionsAPIResponse
| Iterator[InteractionsAPIStreamingResponse]
| Coroutine[Any, Any, InteractionsAPIResponse | AsyncIterator[InteractionsAPIStreamingResponse]]
| Coroutine[object, object, InteractionsAPIResponse | AsyncIterator[InteractionsAPIStreamingResponse]]
):
"""
Handle Interactions API request by calling litellm.responses().
@ -55,13 +55,15 @@ class LiteLLMResponsesInteractionsHandler:
InteractionsAPIResponse or streaming iterator
"""
# Transform interactions request to responses request
responses_request = LiteLLMResponsesInteractionsConfig.transform_interactions_request_to_responses_request(
model=model,
input=input,
optional_params=optional_params,
custom_llm_provider=custom_llm_provider,
stream=stream,
**kwargs,
responses_request: Final = (
LiteLLMResponsesInteractionsConfig.transform_interactions_request_to_responses_request(
model=model,
input=input,
optional_params=optional_params,
custom_llm_provider=custom_llm_provider,
stream=stream,
**kwargs,
)
)
if _is_async:
@ -76,7 +78,10 @@ class LiteLLMResponsesInteractionsHandler:
# Call litellm.responses()
# Note: litellm.responses() returns Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]
# but the type checker may see it as a coroutine in some contexts
responses_response: Final = litellm.responses(
responses_fn: Final[Callable[..., ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]] = vars(litellm)[
"responses"
]
responses_response: Final = responses_fn(
**responses_request,
)
@ -92,8 +97,7 @@ class LiteLLMResponsesInteractionsHandler:
)
# At this point, responses_response must be ResponsesAPIResponse (not streaming)
# Cast to satisfy type checker since we've already checked it's not a streaming iterator
responses_api_response: Final = cast(ResponsesAPIResponse, responses_response)
responses_api_response: Final = responses_response
# Transform responses response to interactions response
return LiteLLMResponsesInteractionsConfig.transform_responses_response_to_interactions_response(
@ -112,7 +116,10 @@ class LiteLLMResponsesInteractionsHandler:
"""Async handler for interactions API requests."""
# Call litellm.aresponses()
# Note: litellm.aresponses() returns Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]
responses_response: Final = await litellm.aresponses(
aresponses_fn: Final[
Callable[..., Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]]
] = vars(litellm)["aresponses"]
responses_response: Final = await aresponses_fn(
**responses_request,
)
@ -128,8 +135,7 @@ class LiteLLMResponsesInteractionsHandler:
)
# At this point, responses_response must be ResponsesAPIResponse (not streaming)
# Cast to satisfy type checker since we've already checked it's not a streaming iterator
responses_api_response: Final = cast(ResponsesAPIResponse, responses_response)
responses_api_response: Final = responses_response
# Transform responses response to interactions response
return LiteLLMResponsesInteractionsConfig.transform_responses_response_to_interactions_response(

154
litellm/interactions/litellm_responses_transformation/transformation.py
View file

@ -2,12 +2,16 @@
Transformation utilities for bridging Interactions API to Responses API.
This module handles transforming between:
- Interactions API format (Google's format with Turn[], system_instruction, etc.)
- Interactions API format (Google's format with Step[]/Turn[], system_instruction, etc.)
- Responses API format (OpenAI's format with input[], instructions, etc.)
"""
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Any, Final, cast
from pydantic import BaseModel
from litellm.types.interactions import (
InteractionInput,
InteractionsAPIOptionalRequestParams,
@ -19,6 +23,8 @@ from litellm.types.llms.openai import (
ResponsesAPIResponse,
)
_STEP_TYPE_ROLES: Final = MappingProxyType({"user_input": "user", "model_output": "assistant"})
class LiteLLMResponsesInteractionsConfig:
"""Configuration class for transforming between Interactions API and Responses API."""
@ -91,112 +97,94 @@ class LiteLLMResponsesInteractionsConfig:
Interactions API input can be:
- string: "Hello"
- Turn[]: [{"role": "user", "content": [...]}]
- Content object
- Step[]: [{"type": "user_input", "content": [...]}, {"type": "model_output", "content": [...]}]
- Turn[] (legacy): [{"role": "user", "content": [...]}]
- Content | Content[]: one user message worth of content parts
Responses API input is:
- string: "Hello"
- Message[]: [{"role": "user", "content": [...]}]
- Message[]: [{"role": "user", "content": [{"type": "input_text", ...}]}]
"""
if isinstance(input, str):
# ResponseInputParam accepts str
return cast(ResponseInputParam, input)
if isinstance(input, list):
# Turn[] format - convert to Responses API Message[] format
messages: Final = []
for turn in input:
if isinstance(turn, dict):
role = turn.get("role", "user")
content = turn.get("content", [])
transformed: Final = (
[
LiteLLMResponsesInteractionsConfig._transform_history_item(item)
for item in input
if LiteLLMResponsesInteractionsConfig._is_history_item(item)
]
if any(LiteLLMResponsesInteractionsConfig._is_history_item(item) for item in input)
else [
{
"role": "user",
"content": LiteLLMResponsesInteractionsConfig._transform_content_array(input, "user"),
}
]
)
return cast(ResponseInputParam, transformed)
# Transform content array
transformed_content = LiteLLMResponsesInteractionsConfig._transform_content_array(content)
messages.append(
{
"role": role,
"content": transformed_content,
}
)
elif isinstance(turn, Turn):
# Pydantic model
role = turn.role if hasattr(turn, "role") else "user"
content = turn.content if hasattr(turn, "content") else []
# Ensure content is a list for _transform_content_array
# Cast to List[Any] to handle various content types
if isinstance(content, list):
content_list: list[Any] = list(content)
elif content is not None:
content_list = [content]
else:
content_list = []
transformed_content = LiteLLMResponsesInteractionsConfig._transform_content_array(content_list)
messages.append(
{
"role": role,
"content": transformed_content,
}
)
return cast(ResponseInputParam, messages)
# Single content object - wrap in message
if isinstance(input, dict):
raw_content: Final = input.get("content")
content_items: Final = raw_content if isinstance(raw_content, list) else [input]
return cast(
ResponseInputParam,
[
{
"role": "user",
"content": LiteLLMResponsesInteractionsConfig._transform_content_array(
input.get("content", []) if isinstance(input.get("content"), list) else [input]
),
"content": LiteLLMResponsesInteractionsConfig._transform_content_array(content_items, "user"),
}
],
)
# Fallback: convert to string
return cast(ResponseInputParam, str(input))
@staticmethod
def _transform_content_array(content: list[Any]) -> list[dict[str, Any]]:
"""Transform Interactions API content array to Responses API format."""
if not isinstance(content, list):
# Single content item - wrap in array
content = [content]
def _is_history_item(item: object) -> bool:
if isinstance(item, Turn):
return True
return isinstance(item, dict) and ("role" in item or item.get("type") in _STEP_TYPE_ROLES)
transformed: Final[list[dict[str, Any]]] = []
for item in content:
if isinstance(item, dict):
# Already in dict format, pass through
transformed.append(item)
elif isinstance(item, str):
# Plain string - wrap in text format
transformed.append({"type": "text", "text": item})
else:
# Pydantic model or other - convert to dict
if hasattr(item, "model_dump"):
dumped = item.model_dump()
if isinstance(dumped, dict):
transformed.append(dumped)
else:
# Fallback: wrap in text format
transformed.append({"type": "text", "text": str(dumped)})
elif hasattr(item, "dict"):
dumped = item.dict()
if isinstance(dumped, dict):
transformed.append(dumped)
else:
# Fallback: wrap in text format
transformed.append({"type": "text", "text": str(dumped)})
else:
# Fallback: wrap in text format
transformed.append({"type": "text", "text": str(item)})
@staticmethod
def _transform_history_item(item: object) -> Mapping[str, object]:
raw: Final = item.model_dump(exclude_none=True) if isinstance(item, Turn) else item
fields: Final = raw if isinstance(raw, Mapping) else {}
role: Final = LiteLLMResponsesInteractionsConfig._responses_role(fields)
raw_content: Final = fields.get("content")
content_items: Final = (
raw_content if isinstance(raw_content, list) else [] if raw_content is None else [raw_content]
)
return {
"role": role,
"content": LiteLLMResponsesInteractionsConfig._transform_content_array(content_items, role),
}
return transformed
@staticmethod
def _responses_role(item: Mapping[str, object]) -> str:
step_role: Final = _STEP_TYPE_ROLES.get(str(item.get("type", "")))
if step_role is not None:
return step_role
raw_role: Final = str(item.get("role") or "user")
return "assistant" if raw_role == "model" else raw_role
@staticmethod
def _transform_content_array(content: Sequence[object], role: str) -> Sequence[Mapping[str, object]]:
"""Transform Interactions API content parts to Responses API parts for the given role."""
return [LiteLLMResponsesInteractionsConfig._transform_content_item(item, role) for item in content]
@staticmethod
def _transform_content_item(item: object, role: str) -> Mapping[str, object]:
text_type: Final = "output_text" if role == "assistant" else "input_text"
if isinstance(item, str):
return {"type": text_type, "text": item}
if isinstance(item, Mapping):
if item.get("type") == "text":
return {"type": text_type, "text": str(item.get("text", ""))}
return item
if isinstance(item, BaseModel):
return LiteLLMResponsesInteractionsConfig._transform_content_item(item.model_dump(exclude_none=True), role)
return {"type": text_type, "text": str(item)}
@staticmethod
def transform_responses_response_to_interactions_response(

4
litellm/litellm_core_utils/core_helpers.py
View file

@ -1,7 +1,7 @@
# What is this?
## Helper utilities
import copy
from collections.abc import Iterable
from collections.abc import Iterable, Mapping
from typing import TYPE_CHECKING, Any, Final, Literal
import httpx
@ -181,7 +181,7 @@ def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metad
def get_metadata_variable_name_from_kwargs(
kwargs: dict,
kwargs: Mapping[str, object],
) -> Literal["metadata", "litellm_metadata"]:
"""
Helper to return what the "metadata" field should be called in the request data

94
litellm/litellm_core_utils/internal_call_metadata.py Normal file
View file

@ -0,0 +1,94 @@
"""Metadata a request forwards to the internal LLM sub-calls it triggers.
Internal features (the auto-router's classifier and embeddings, shadow eval's shadow and
judge calls) bill real provider spend that nobody typed a prompt for. That spend must land
on the same key/team/org/user as the request that caused it, so the sub-call carries the
caller's identity metadata, minus two things that must never be forwarded as-is:
* ``user_api_key_budget_reservation`` (and the reservation nested inside
``user_api_key_auth``) belongs to the parent completion. If a sub-call's cost callback
sees it, that callback finalizes the reservation and the parent's own callback then
skips incrementing the key/team budget counters, losing the parent's spend.
``user_api_key_auth`` itself is kept, sanitized, because model access-group filtering
needs it.
* The sub-call is stamped with ``INTERNAL_CALL_ORIGIN_METADATA_KEY`` so its spend log row
records that it is not traffic the caller sent.
"""
from __future__ import annotations
from collections.abc import Mapping
from typing import Final
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.types.utils import InternalCallOrigin
BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"})
_USER_API_KEY_AUTH_KEY: Final = "user_api_key_auth"
FORWARDABLE_IDENTITY_METADATA_KEYS: Final = frozenset(
{
"user_api_key",
"user_api_key_hash",
"user_api_key_alias",
"user_api_key_team_id",
"user_api_key_org_id",
"user_api_key_user_id",
"user_api_key_end_user_id",
_USER_API_KEY_AUTH_KEY,
}
)
"""The caller-identity subset a detached sub-call needs to be attributed and
budget-checked like the request that spawned it. Everything else on the parent's metadata
(routing decision, guardrail state, logging payload) describes the parent call and would
be a lie on a sub-call that runs after it returned."""
def sanitize_user_api_key_auth(auth: object) -> object:
"""Copy of the auth object with its budget reservation removed; the cost callback
falls back to reading the reservation from inside the auth object."""
if isinstance(auth, dict):
return {k: v for k, v in auth.items() if k != "budget_reservation"} # mutable-ok: SDK metadata value
reservation: Final[object] = getattr(auth, "budget_reservation", None)
model_copy: Final[object] = getattr(auth, "model_copy", None)
if reservation is not None and callable(model_copy):
return model_copy(update={"budget_reservation": None}) # mutable-ok: pydantic update payload
return auth
def _sanitized(parent_metadata: Mapping[str, object]) -> dict[str, object]: # mutable-ok: SDK metadata kwarg
return { # mutable-ok: SDK metadata kwarg
k: sanitize_user_api_key_auth(v) if k == _USER_API_KEY_AUTH_KEY else v
for k, v in parent_metadata.items()
if k not in BUDGET_RESERVATION_METADATA_KEYS
}
def forwarded_internal_call_metadata(
parent_metadata: Mapping[str, object] | None,
call_origin: InternalCallOrigin,
) -> dict[str, object]: # mutable-ok: SDK metadata kwarg
"""Parent metadata, minus its budget reservation, stamped with the sub-call's origin.
For sub-calls made inside the parent request (classifier, embeddings), where the
parent's full context still describes the call being made.
"""
if not parent_metadata:
return {} # mutable-ok: SDK metadata kwarg
return _sanitized(parent_metadata) | { # mutable-ok: SDK metadata kwarg
INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin
}
def sanitized_forwardable_call_metadata(
parent_metadata: Mapping[str, object],
call_origin: InternalCallOrigin,
) -> dict[str, object]: # mutable-ok: SDK metadata kwarg
"""Just the caller's identity, stamped with the sub-call's origin.
For sub-calls detached from the parent request (shadow eval), which outlive it and
must not inherit per-request state such as its routing decision or logging payload.
"""
identity: Final = {k: v for k, v in parent_metadata.items() if k in FORWARDABLE_IDENTITY_METADATA_KEYS}
return _sanitized(identity) | {INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin} # mutable-ok: SDK metadata kwarg

3
litellm/litellm_core_utils/litellm_logging.py
View file

@ -3462,7 +3462,7 @@ class Logging(LiteLLMLoggingBaseClass):
model=self.model,
messages=[],
logging_obj=self,
optional_params={},
optional_params=self.optional_params or {},
api_key="",
request_data={},
encoding=litellm.encoding,
@ -3483,6 +3483,7 @@ class Logging(LiteLLMLoggingBaseClass):
),
model_response=litellm.ModelResponse(),
json_mode=None,
speed=self.optional_params.get("speed") if self.optional_params else None,
)
return result

19
litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py
View file

@ -2,6 +2,7 @@
Helper utilities for tracking the cost of built-in tools.
"""
from collections.abc import Mapping
from typing import Any, Final, Literal
import litellm
@ -23,6 +24,14 @@ from litellm.types.utils import (
)
def _usage_reports_server_side_web_search_calls(usage: Usage) -> bool:
details: Final = getattr(usage, "server_side_tool_usage_details", None)
if not isinstance(details, Mapping):
return False
calls: Final = details.get("web_search_calls")
return isinstance(calls, int) and calls > 0
class StandardBuiltInToolCostTracking:
"""
Helper class for tracking the cost of built-in tools
@ -351,6 +360,10 @@ class StandardBuiltInToolCostTracking:
# and _handle_web_search_cost() is never called.
if hasattr(usage, "server_tool_use") and _get_web_search_requests(usage.server_tool_use) is not None:
return True
# xAI reports usage.server_side_tool_usage_details.web_search_calls; a searched
# answer with no url_citation annotations has no other chat-path signal
if _usage_reports_server_side_web_search_calls(usage):
return True
return False
elif isinstance(response_object, ResponsesAPIResponse):
# response api explicitly includes web_search_call in the output
@ -370,6 +383,8 @@ class StandardBuiltInToolCostTracking:
)
):
return True
if _usage_reports_server_side_web_search_calls(usage):
return True
return False
@ -432,7 +447,9 @@ class StandardBuiltInToolCostTracking:
"""
output: Final = response_object.output
for output_item in output:
_output_type: str | None = getattr(output_item, "type", None)
_output_type: str | None = (
output_item.get("type") if isinstance(output_item, dict) else getattr(output_item, "type", None)
)
if _output_type == output_type:
return True
return False

87
litellm/litellm_core_utils/llm_judge.py Normal file
View file

@ -0,0 +1,87 @@
"""Shared primitives for LLM-judge features (llm_as_a_judge guardrail, shadow eval)."""
from __future__ import annotations
import json
import re
from typing import TYPE_CHECKING, Final
import litellm
if TYPE_CHECKING:
from litellm import Router
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponse
JSON_FENCE_RE: Final = re.compile(r"```(?:json)?\s*(.*?)\s*```", re.DOTALL | re.IGNORECASE)
def default_router_provider() -> Router | None:
try:
from litellm.proxy.proxy_server import llm_router
except ImportError:
return None
return llm_router
def parse_json_verdict(raw: str) -> dict[str, object]: # mutable-ok: plain parsed-JSON payload
"""Parse a judge's JSON verdict, tolerating markdown fences and surrounding prose."""
text = raw.strip() # rebind-ok: progressively narrowed to the JSON payload
fenced: Final = JSON_FENCE_RE.search(text)
if fenced is not None:
text = fenced.group(1).strip() # rebind-ok: progressively narrowed to the JSON payload
parsed: object
try:
parsed = json.loads(text)
except json.JSONDecodeError:
start: Final = text.find("{")
end: Final = text.rfind("}")
if start == -1 or end <= start:
raise
parsed = json.loads(text[start : end + 1])
if not isinstance(parsed, dict):
raise ValueError("judge response is not a JSON object")
return {str(k): v for k, v in parsed.items()} # mutable-ok: plain parsed-JSON payload
def extract_text_from_content(content: object) -> str:
"""Return plain text from a message content field (str or multimodal list)."""
if isinstance(content, str):
return content
if isinstance(content, list):
return " ".join(
str(part.get("text", "")) for part in content if isinstance(part, dict) and part.get("type") == "text"
)
return ""
def router_resolves_model(router: Router | None, model: str) -> bool:
"""Whether the model name resolves through the proxy's router (configured deployment
or model-group alias), the same check the judge dispatch itself makes, so start-time
validation cannot accept a name the call path then fails on."""
return router is not None and bool(model in router.model_group_alias or router.get_model_list(model_name=model))
async def judge_acompletion(
router: Router | None,
judge_model: str,
messages: list[AllMessageValues], # mutable-ok: the SDK acompletion signature takes a list
**params: object,
) -> ModelResponse:
"""Dispatch a judge call through the proxy's router when the judge model is a
configured deployment (DB-stored credentials work), through the SDK for
provider-qualified public names. The router path never retries or falls back:
a failed judge call is the caller's counted failure, not a spend multiplier.
Sampling preferences are advisory: models that removed sampling params (e.g.
claude-sonnet-5) drop them instead of rejecting the judge call."""
if router_resolves_model(router, judge_model):
return await router.acompletion( # pyright: ignore[reportOptionalMemberAccess] # router_resolves_model implies router is not None
model=judge_model,
messages=messages,
num_retries=0,
fallbacks=[],
drop_params=True,
**params,
)
return await litellm.acompletion(model=judge_model, messages=messages, num_retries=0, drop_params=True, **params)

31
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -803,8 +803,28 @@ class ChunkProcessor:
completion_tokens_details=completion_tokens_details,
prompt_tokens_details=prompt_tokens_details,
cost=cost,
inference_geo=self._last_provider_pricing_field(chunks, "inference_geo"),
speed=self._last_provider_pricing_field(chunks, "speed"),
)
def _last_provider_pricing_field(
self,
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
field: str,
) -> str | None:
"""
Last value of a provider-specific usage field that changes pricing but is not a
declared ``Usage`` field, e.g. Anthropic's ``speed`` (fast mode multiplies
non-cache token cost) and ``inference_geo``.
"""
values: Final = [
value
for chunk in chunks
if (usage_chunk := self._extract_usage_chunk(chunk)) is not None
and isinstance(value := getattr(usage_chunk, field, None), str)
]
return values[-1] if values else None
@staticmethod
def _reset_anthropic_cursor_completion_tokens(
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
@ -934,7 +954,16 @@ class ChunkProcessor:
# Return a new usage object with the new values
returned_usage = Usage(**returned_usage.model_dump())
provider_pricing_fields: Final = {
field: value
for field, value in (
("inference_geo", calculated_usage_per_chunk["inference_geo"]),
("speed", calculated_usage_per_chunk["speed"]),
)
if value is not None
}
returned_usage = Usage(**returned_usage.model_dump(), **provider_pricing_fields)
return returned_usage

66
litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py
View file

@ -13,8 +13,10 @@ Mirrors Anthropic's native ``compact_20260112`` for non-Anthropic providers:
"""
import re
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, Optional, TypedDict, Union, cast
from typing_extensions import ReadOnly
import litellm
from litellm._logging import verbose_logger
@ -29,9 +31,8 @@ if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.router import Router
from litellm.types.llms.anthropic import (
AllAnthropicPassThroughMessageValues,
AllAnthropicToolsValues,
AnthopicMessagesAssistantMessageParam,
AnthropicMessagesUserMessageParam,
)
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.utils import ModelResponse
@ -534,7 +535,7 @@ def _augment_system_with_summary(
return [{"type": "text", "text": prefix.rstrip()}, *system]
def _resolve_trigger_tokens(edit_spec: dict[str, object]) -> tuple[int, list[str]]:
def _resolve_trigger_tokens(edit_spec: Mapping[str, object]) -> tuple[int, list[str]]:
"""Validate and resolve ``trigger.value``.
Raises ``AnthropicContextManagementError`` if the explicitly-supplied value
@ -568,7 +569,7 @@ def _resolve_trigger_tokens(edit_spec: dict[str, object]) -> tuple[int, list[str
return value, warnings
def _build_summary_prompt(edit_spec: dict[str, object], tools: list[dict[str, object]] | None) -> str:
def _build_summary_prompt(edit_spec: Mapping[str, object], tools: Sequence[Mapping[str, object]] | None) -> str:
custom: Final = edit_spec.get("instructions")
if isinstance(custom, str) and custom.strip():
return custom
@ -623,7 +624,7 @@ def _count_effective_tokens(
try:
openai_shape = adapter.translate_anthropic_messages_to_openai(
messages=cast(
"list[AnthropicMessagesUserMessageParam | AnthopicMessagesAssistantMessageParam]",
"list[AllAnthropicPassThroughMessageValues]",
messages_without_compaction,
)
)
@ -736,7 +737,7 @@ def _extract_summary_text(raw: str | None) -> str | None:
def _system_to_openai_message(
system: str | list[dict[str, Any]] | None,
) -> dict[str, Any] | None:
) -> dict[str, object] | None:
"""Translate Anthropic-shaped ``system`` to an OpenAI system message.
Accepts a bare string or a list of Anthropic content blocks; returns
@ -773,7 +774,7 @@ def _build_summary_messages(
try:
openai_messages = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(
messages=cast(
"list[AnthropicMessagesUserMessageParam | AnthopicMessagesAssistantMessageParam]",
"list[AllAnthropicPassThroughMessageValues]",
stripped,
)
)
@ -809,7 +810,7 @@ def _is_user_message(msg: object) -> bool:
return isinstance(msg, dict) and msg.get("role") == "user"
def _append_text_to_content(content: Any, extra_text: str) -> Any:
def _append_text_to_content(content: object, extra_text: str) -> object:
"""Append ``extra_text`` to an OpenAI-shape message ``content`` field.
Handles the two common shapes: ``str`` and ``list`` of content parts.
@ -820,10 +821,29 @@ def _append_text_to_content(content: Any, extra_text: str) -> Any:
if isinstance(content, str):
return f"{content}\n\n{extra_text}"
if isinstance(content, list):
return [*content, {"type": "text", "text": extra_text}]
appended: Final[list[object]] = [*content, {"type": "text", "text": extra_text}]
return appended
return [content, {"type": "text", "text": extra_text}]
class _SummaryCallUserKwarg(TypedDict, total=False):
user: ReadOnly[object]
class _SummaryCallRegionKwarg(TypedDict, total=False):
allowed_model_region: ReadOnly[str]
class _SummaryCallKwargs(TypedDict):
model: ReadOnly[str]
messages: ReadOnly[list[dict[str, object]]]
max_tokens: ReadOnly[int]
timeout: ReadOnly[float]
litellm_metadata: ReadOnly[Mapping[str, object]]
user: NotRequired[ReadOnly[object]]
allowed_model_region: NotRequired[ReadOnly[str]]
async def _call_summary_model(
*,
summary_model: str,
@ -860,22 +880,24 @@ async def _call_summary_model(
# the parent ``/v1/messages`` request. On timeout the caller catches the
# exception and surfaces ``applied_edits[0].error = "summary_call_failed"``,
# forwarding the request without compaction rather than hanging.
call_kwargs: Final[dict[str, Any]] = {
"model": summary_model,
"messages": summary_messages,
"max_tokens": max_tokens,
"timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS,
"litellm_metadata": metadata,
}
# The end-user id must also travel as the top-level ``user`` kwarg: legacy
# limiter hooks and prometheus end-user tracking read it from there rather
# than from ``litellm_metadata``, so without it the summary tokens would not
# debit the caller's end-user counters.
end_user_id: Final = metadata.get("user_api_key_end_user_id")
if end_user_id:
call_kwargs["user"] = end_user_id
if allowed_model_region is not None:
call_kwargs["allowed_model_region"] = allowed_model_region
call_kwargs: Final[_SummaryCallKwargs] = {
"model": summary_model,
"messages": summary_messages,
"max_tokens": max_tokens,
"timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS,
"litellm_metadata": metadata,
**(_SummaryCallUserKwarg(user=end_user_id) if end_user_id else _SummaryCallUserKwarg()),
**(
_SummaryCallRegionKwarg(allowed_model_region=allowed_model_region)
if allowed_model_region is not None
else _SummaryCallRegionKwarg()
),
}
if llm_router is not None and hasattr(llm_router, "acompletion"):
return await llm_router.acompletion(**call_kwargs)
return await litellm.acompletion(**call_kwargs)

31
litellm/llms/azure/common_utils.py
View file

@ -2,11 +2,12 @@ import asyncio
import hashlib
import json
import os
from collections.abc import Callable
from collections.abc import Callable, Mapping
from typing import Any, Final, Literal, NamedTuple, cast
import httpx
from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
@ -23,6 +24,22 @@ from litellm.utils import _add_path_to_api_base
azure_ad_cache: Final = DualCache()
class _AzureAdTokenJson(TypedDict, total=False):
access_token: ReadOnly[str]
expires_in: ReadOnly[int]
class _AzureV1ClientParams(TypedDict, total=False, extra_items=object):
base_url: ReadOnly[str]
class _AzureGatewayClientParams(TypedDict, total=False, extra_items=object):
api_version: ReadOnly[str]
base_url: ReadOnly[str]
max_retries: ReadOnly[int]
timeout: ReadOnly[float | httpx.Timeout]
class AzureOpenAIError(BaseLLMException):
def __init__(
self,
@ -220,7 +237,7 @@ def get_azure_ad_token_from_oidc(
message=req_token.text,
)
azure_ad_token_json: Final = req_token.json()
azure_ad_token_json: Final[_AzureAdTokenJson] = req_token.json()
azure_ad_token_access_token = azure_ad_token_json.get("access_token", None)
azure_ad_token_expires_in: Final = azure_ad_token_json.get("expires_in", None)
@ -486,7 +503,7 @@ class BaseAzureLLM(BaseOpenAILLM):
v1_api_key = _async_v1_api_key
v1_params: Final[dict[str, Any]] = {
v1_params: Final[_AzureV1ClientParams] = {
"api_key": v1_api_key,
"base_url": f"{api_base}/openai/v1/",
}
@ -643,7 +660,7 @@ class BaseAzureLLM(BaseOpenAILLM):
api_base += "/"
api_base += f"{model}"
azure_client_params: Final[dict[str, Any]] = {
azure_client_params: Final[_AzureGatewayClientParams] = {
"api_version": api_version,
"base_url": f"{api_base}",
"http_client": litellm.client_session,
@ -702,7 +719,7 @@ class BaseAzureLLM(BaseOpenAILLM):
@staticmethod
def _get_base_azure_url(
api_base: str | None,
litellm_params: GenericLiteLLMParams | dict[str, Any] | None,
litellm_params: GenericLiteLLMParams | Mapping[str, object] | None,
route: Literal["/openai/responses", "/openai/vector_stores"] | str,
default_api_version: str | Literal["latest", "preview"] | None = None,
) -> str:
@ -757,7 +774,9 @@ class BaseAzureLLM(BaseOpenAILLM):
return False
return api_version in {"preview", "latest", "v1"}
def _resolve_env_var(self, litellm_params: dict[str, Any], param_key: str, env_var_key: str) -> str | None:
def _resolve_env_var(
self, litellm_params: Mapping[str, str | None], param_key: str, env_var_key: str
) -> str | None:
"""Resolve the environment variable for a given parameter key.
The logic here is different from `params.get(key, os.getenv(env_var))` because

6
litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py
View file

@ -19,9 +19,9 @@ from litellm.llms.bedrock.common_utils import (
convert_bedrock_invoke_output_format_to_inline_schema,
get_anthropic_beta_from_headers,
normalize_bedrock_opus_output_config_effort,
normalize_custom_field_on_tools,
normalize_tool_input_schema_types_for_bedrock_invoke,
pop_bedrock_invoke_output_config_format,
remove_custom_field_from_tools,
)
from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER
from litellm.types.llms.openai import AllMessageValues
@ -243,8 +243,8 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
if "anthropic_version" not in anthropic_request:
anthropic_request["anthropic_version"] = self.anthropic_version
# Remove `custom` field from tools (Bedrock doesn't support it)
remove_custom_field_from_tools(anthropic_request)
# Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it)
normalize_custom_field_on_tools(anthropic_request)
normalize_tool_input_schema_types_for_bedrock_invoke(anthropic_request)
return anthropic_request

21
litellm/llms/bedrock/common_utils.py
View file

@ -176,13 +176,14 @@ def convert_bedrock_invoke_output_format_to_inline_schema(
request_body["messages"] = new_messages
def remove_custom_field_from_tools(request_body: dict) -> None:
def normalize_custom_field_on_tools(request_body: dict) -> None:
"""
Remove ``custom`` field from each tool in the request body.
Drop the ``custom`` field from each tool, first hoisting a boolean
``custom.defer_loading`` onto the top-level ``defer_loading`` flag that
Bedrock and Anthropic actually document, unless the tool already carries one.
Claude Code (v2.1.69+) sends ``custom: {defer_loading: true}`` on tool
definitions, which Anthropic's API accepts but Bedrock rejects with
``"Extra inputs are not permitted"``.
Claude Code (v2.1.69+) is reported to send ``custom: {defer_loading: true}`` on
tool definitions, which Bedrock rejects with ``"Extra inputs are not permitted"``.
Args:
request_body: The request dictionary to modify in-place.
@ -193,8 +194,14 @@ def remove_custom_field_from_tools(request_body: dict) -> None:
if not tools or not isinstance(tools, list):
return
for tool in tools:
if isinstance(tool, dict):
tool.pop("custom", None)
if not isinstance(tool, dict):
continue
custom: dict[str, object] | None = tool.pop("custom", None)
if not isinstance(custom, dict) or "defer_loading" in tool:
continue
deferred: object = custom.get("defer_loading")
if isinstance(deferred, bool):
tool["defer_loading"] = deferred
def normalize_json_schema_custom_types_to_object(schema: dict) -> None:

94
litellm/llms/bedrock/files/transformation.py
View file

@ -2,17 +2,18 @@ import base64
import json
import os
import time
from collections.abc import Iterable, Mapping, MutableMapping
from collections.abc import Iterable, Mapping, MutableMapping, Sequence
from functools import cache
from itertools import chain
from types import MappingProxyType
from typing import Any, Final
from typing import Any, Final, TypeAlias, TypedDict
from urllib.parse import unquote
import httpx
from httpx import Headers, Response
from openai.types.file_deleted import FileDeleted
from pydantic import BaseModel, ConfigDict, TypeAdapter
from typing_extensions import ReadOnly
from litellm._logging import verbose_logger
from litellm._uuid import uuid
@ -63,10 +64,39 @@ from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resol
S3_SIGNED_GET_HEADERS_PARAM: Final = "_s3_signed_get_headers"
def _frozen_mapping(items: Iterable[tuple[str, Any]]) -> Mapping[str, Any]:
def _frozen_mapping(items: Iterable[tuple[str, object]]) -> Mapping[str, object]:
return MappingProxyType(dict(items))
_EmbeddingBatchInput: TypeAlias = (
str | int | float | Sequence[str] | Sequence[int] | Sequence[Sequence[int]] | Mapping[str, object]
)
class _OpenAIBatchRecordBody(TypedDict, total=False):
model: ReadOnly[str]
prompt: ReadOnly[str | Sequence[str] | Sequence[int] | Sequence[Sequence[int]]]
input: ReadOnly[_EmbeddingBatchInput]
metadata: ReadOnly[Mapping[str, object]]
class _OpenAIBatchRecord(TypedDict, total=False):
custom_id: ReadOnly[str]
url: ReadOnly[str]
body: ReadOnly[_OpenAIBatchRecordBody]
class _BedrockBatchRecord(TypedDict):
recordId: ReadOnly[str]
modelInput: ReadOnly[Mapping[str, object]]
class _S3UploadResponse(TypedDict, total=False):
Key: ReadOnly[str]
Bucket: ReadOnly[str]
ContentLength: ReadOnly[int]
# JSONL batch records are untyped json, so the `/v1/responses` fields are
# validated into their concrete Responses API types before being handed to the
# Responses-to-Chat bridge. Both adapters drop keys the Responses API doesn't
@ -231,7 +261,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
def _get_s3_object_name_from_batch_jsonl(
self,
openai_jsonl_content: list[dict[str, Any]],
openai_jsonl_content: Sequence[_OpenAIBatchRecord],
) -> str:
"""
Gets a unique S3 object name for the Bedrock batch processing job
@ -341,7 +371,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
OPENAI_RESPONSES_URL = "/v1/responses"
@staticmethod
def _classify_batch_record(openai_jsonl_record: Mapping[str, Any]) -> BedrockBatchRecordKind:
def _classify_batch_record(openai_jsonl_record: _OpenAIBatchRecord) -> BedrockBatchRecordKind:
"""
Decide which OpenAI endpoint shape an OpenAI batch JSONL line carries.
@ -484,7 +514,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
return value if isinstance(value, str) and value else None
@staticmethod
def _coerce_embedding_input_to_string(raw_input: Any, model: str = "") -> str:
def _coerce_embedding_input_to_string(raw_input: _EmbeddingBatchInput | None, model: str = "") -> str:
"""
Normalize an OpenAI /v1/embeddings `input` field into the single
string that Bedrock Titan v2 InvokeModel expects in `inputText`.
@ -541,8 +571,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
def _map_openai_embedding_to_bedrock_params(
self,
openai_request_body: dict[str, Any],
) -> dict[str, Any]:
openai_request_body: _OpenAIBatchRecordBody,
) -> dict[str, object]:
"""
Transform an OpenAI /v1/embeddings request body into the
Bedrock InvokeModel `modelInput` for embedding models that AWS
@ -588,7 +618,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
return dict(titan_config._transform_request(input=input_text, inference_params=inference_params))
@staticmethod
def _transform_text_completion_body_to_chat_body(openai_request_body: Mapping[str, Any]) -> Mapping[str, Any]:
def _transform_text_completion_body_to_chat_body(
openai_request_body: _OpenAIBatchRecordBody,
) -> Mapping[str, object]:
"""
Rewrite an OpenAI `/v1/completions` batch body as a Chat Completions body.
@ -610,7 +642,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
)
@staticmethod
def _transform_responses_body_to_chat_body(openai_request_body: Mapping[str, Any]) -> Mapping[str, Any]:
def _transform_responses_body_to_chat_body(openai_request_body: _OpenAIBatchRecordBody) -> Mapping[str, object]:
"""
Rewrite an OpenAI `/v1/responses` batch body as a Chat Completions body.
@ -631,23 +663,25 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
"Batch record for /v1/responses is missing required `input` field: "
f"model={openai_request_body.get('model', '')}"
)
chat_body: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
model=openai_request_body.get("model", ""),
input=_responses_input_adapter().validate_python(responses_input),
responses_api_request=_responses_request_adapter().validate_python(
_frozen_mapping(
(key, value) for key, value in openai_request_body.items() if key not in ("model", "input")
)
),
metadata=openai_request_body.get("metadata"),
chat_body: Final[Mapping[str, object]] = (
LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
model=openai_request_body.get("model", ""),
input=_responses_input_adapter().validate_python(responses_input),
responses_api_request=_responses_request_adapter().validate_python(
_frozen_mapping(
(key, value) for key, value in openai_request_body.items() if key not in ("model", "input")
)
),
metadata=openai_request_body.get("metadata"),
)
)
return _frozen_mapping((key, value) for key, value in chat_body.items() if key != "tools" or value)
@staticmethod
def _transform_batch_body_to_chat_body(
openai_request_body: Mapping[str, Any],
openai_request_body: _OpenAIBatchRecordBody,
record_kind: BedrockBatchRecordKind,
) -> Mapping[str, Any]:
) -> Mapping[str, object]:
"""
Normalize a non-embedding batch body to the Chat Completions shape the
per-provider Bedrock transformations expect.
@ -666,7 +700,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
self,
openai_request_body: Mapping[str, Any],
provider: str | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Transform OpenAI request body to Bedrock-compatible modelInput
parameters using existing transformation logic.
@ -677,7 +711,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
"""
from litellm.types.utils import LlmProviders
_model: Final = openai_request_body.get("model", "")
_model: Final[str] = openai_request_body.get("model", "")
messages: Final = openai_request_body.get("messages", [])
optional_params: Final = {k: v for k, v in openai_request_body.items() if k not in ["model", "messages"]}
@ -733,8 +767,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
}
def _transform_openai_jsonl_content_to_bedrock_jsonl_content(
self, openai_jsonl_content: list[dict[str, Any]]
) -> list[dict[str, Any]]:
self, openai_jsonl_content: Sequence[_OpenAIBatchRecord]
) -> list[_BedrockBatchRecord]:
"""
Transforms OpenAI JSONL content to Bedrock batch format
@ -1026,7 +1060,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
response_headers: Final = raw_response.headers
# Extract S3 object information from the response
# S3 PUT object returns ETag and other metadata in headers
content_length: Final = response_headers.get("Content-Length", "0")
content_length: Final[str] = response_headers.get("Content-Length", "0")
# Use the actual upload URL that was used for the S3 upload
upload_url: Final = litellm_params.get("upload_url")
@ -1224,7 +1258,9 @@ class BedrockJsonlFilesTransformation:
object_name: Final = self._get_s3_object_name(openai_jsonl_content=openai_jsonl_content)
return bedrock_jsonl_string, object_name
def _transform_openai_jsonl_content_to_bedrock_jsonl_content(self, openai_jsonl_content: list[dict[str, Any]]):
def _transform_openai_jsonl_content_to_bedrock_jsonl_content(
self, openai_jsonl_content: Sequence[_OpenAIBatchRecord]
):
"""
Delegate to the main BedrockFilesConfig transformation method
"""
@ -1233,7 +1269,7 @@ class BedrockJsonlFilesTransformation:
def _get_s3_object_name(
self,
openai_jsonl_content: list[dict[str, Any]],
openai_jsonl_content: Sequence[_OpenAIBatchRecord],
) -> str:
"""
Gets a unique S3 object name for the Bedrock batch processing job
@ -1285,7 +1321,7 @@ class BedrockJsonlFilesTransformation:
return content
def transform_s3_bucket_response_to_openai_file_object(
self, create_file_data: CreateFileRequest, s3_upload_response: dict[str, Any]
self, create_file_data: CreateFileRequest, s3_upload_response: _S3UploadResponse
) -> OpenAIFileObject:
"""
Transforms S3 Bucket upload file response to OpenAI FileObject

8
litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
View file

@ -33,9 +33,9 @@ from litellm.llms.bedrock.common_utils import (
get_anthropic_beta_from_headers,
is_claude_4_5_on_bedrock,
normalize_bedrock_opus_output_config_effort,
normalize_custom_field_on_tools,
normalize_tool_input_schema_types_for_bedrock_invoke,
pop_bedrock_invoke_output_config_format,
remove_custom_field_from_tools,
)
from litellm.types.llms.anthropic import (
ANTHROPIC_BETA_HEADER_VALUES,
@ -749,11 +749,9 @@ class AmazonAnthropicClaudeMessagesConfig(
model,
)
# 5b. Remove `custom` field from tools (Bedrock doesn't support it)
# Claude Code sends `custom: {defer_loading: true}` on tool definitions,
# which causes Bedrock to reject the request with "Extra inputs are not permitted"
# 5b. Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it)
# Ref: https://github.com/BerriAI/litellm/issues/22847
remove_custom_field_from_tools(anthropic_messages_request)
normalize_custom_field_on_tools(anthropic_messages_request)
normalize_tool_input_schema_types_for_bedrock_invoke(anthropic_messages_request)
ensure_bedrock_anthropic_messages_tool_names(anthropic_messages_request)

120
litellm/llms/nvidia_nim/rerank/ranking_transformation.py
View file

@ -8,13 +8,23 @@ Reference: https://build.nvidia.com/nvidia/llama-3_2-nv-rerankqa-1b-v2/deploy
from typing import Final
import httpx
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig
from litellm.types.rerank import RerankResponse
class NvidiaNimRankingConfig(NvidiaNimRerankConfig):
"""
Configuration for NVIDIA NIM models that use the /v1/ranking endpoint.
The native /v1/ranking request schema accepts only 'model', 'query',
'passages', and 'truncate' -- requests containing 'top_k' are rejected
with a 400 validation error. Cohere-compatible 'top_n' is therefore
applied client-side by truncating the converted response instead of
being forwarded to the endpoint.
Example:
curl -X "POST" 'https://ai.api.nvidia.com/v1/ranking' \
-H 'Accept: application/json' \
@ -27,6 +37,16 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig):
}'
"""
SUPPORTED_PASSAGE_FIELDS: tuple[str, ...] = ("text", "image")
def __init__(self) -> None:
super().__init__()
# top_n captured in transform_rerank_request and applied in
# transform_rerank_response. The provider config is instantiated
# per-request (see ProviderConfigManager.get_provider_rerank_config),
# so this does not leak across requests.
self._client_side_top_n: int | None = None
def _get_clean_model_name(self, model: str) -> str:
"""Strip 'nvidia_nim/' and 'ranking/' prefixes from model name."""
# First strip nvidia_nim/ prefix if present
@ -58,6 +78,47 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig):
return f"{api_base}/v1/ranking"
def map_cohere_rerank_params(
self,
non_default_params: dict | None, # mutable-ok: matches BaseRerankConfig's request contract
model: str,
drop_params: bool,
query: str,
documents: list[str | dict[str, object]], # mutable-ok: matches BaseRerankConfig's document contract
custom_llm_provider: str | None = None,
top_n: int | None = None,
rank_fields: list[str] | None = None, # mutable-ok: matches BaseRerankConfig's field contract
return_documents: bool | None = True,
max_chunks_per_doc: int | None = None,
max_tokens_per_doc: int | None = None,
instruction: str | None = None,
) -> dict: # mutable-ok: LiteLLM provider transforms return mutable request dictionaries
"""
Keep Cohere's top_n as-is instead of mapping it to top_k.
The native /v1/ranking endpoint rejects top_k, so top_n is applied
client-side after the response is converted.
"""
optional_params: Final = super().map_cohere_rerank_params(
non_default_params=non_default_params,
model=model,
drop_params=drop_params,
query=query,
documents=documents,
custom_llm_provider=custom_llm_provider,
top_n=None, # do not map top_n -> top_k for /v1/ranking
rank_fields=rank_fields,
return_documents=return_documents,
max_chunks_per_doc=max_chunks_per_doc,
max_tokens_per_doc=max_tokens_per_doc,
instruction=instruction,
)
# /v1/ranking rejects top_k even when passed as a provider-specific param
optional_params.pop("top_k", None)
if top_n is not None:
optional_params["top_n"] = top_n
return optional_params
def transform_rerank_request(
self,
model: str,
@ -67,11 +128,66 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig):
) -> dict:
"""
Transform request, using clean model name without 'ranking/' prefix.
top_n / top_k are stripped from the outgoing request: the native
/v1/ranking endpoint accepts only model, query, passages, and
truncate. top_n is stashed and applied client-side in
transform_rerank_response.
"""
top_n: Final = optional_rerank_params.get("top_n")
if top_n is not None:
if isinstance(top_n, bool) or not isinstance(top_n, int) or top_n < 1:
raise ValueError(f"top_n must be a positive integer, got: {top_n!r}")
self._client_side_top_n = top_n
clean_model: Final = self._get_clean_model_name(model)
filtered_params: Final = { # mutable-ok: the base transformer requires a mutable request dictionary
k: v for k, v in optional_rerank_params.items() if k not in ("top_n", "top_k")
}
return super().transform_rerank_request(
model=clean_model,
optional_rerank_params=optional_rerank_params,
optional_rerank_params=filtered_params,
headers=headers,
litellm_params=litellm_params,
)
def transform_rerank_response(
self,
model: str,
raw_response: httpx.Response,
model_response: RerankResponse,
logging_obj: LiteLLMLoggingObj,
api_key: str | None = None,
request_data: dict | None = None, # mutable-ok: matches BaseRerankConfig's response contract
optional_params: dict | None = None, # mutable-ok: matches BaseRerankConfig's response contract
litellm_params: dict | None = None, # mutable-ok: matches BaseRerankConfig's response contract
) -> RerankResponse:
"""
Convert the native ranking response, then apply top_n client-side.
/v1/ranking returns rankings sorted by relevance, but sort before
truncating in case a server returns them unsorted.
"""
resolved_request_data: Final = request_data or {} # mutable-ok: the base transformer requires a dictionary
resolved_optional_params: Final = optional_params or {} # mutable-ok: response options are keyed lookups
resolved_litellm_params: Final = litellm_params or {} # mutable-ok: the base transformer requires a dictionary
response: Final = super().transform_rerank_response(
model=model,
raw_response=raw_response,
model_response=model_response,
logging_obj=logging_obj,
api_key=api_key,
request_data=resolved_request_data,
optional_params=resolved_optional_params,
litellm_params=resolved_litellm_params,
)
top_n: Final = resolved_optional_params.get("top_n") or self._client_side_top_n
if top_n is not None and response.results is not None and len(response.results) > top_n:
response.results = sorted(
response.results,
key=lambda result: result["relevance_score"],
reverse=True,
)[:top_n]
return response

28
litellm/llms/nvidia_nim/rerank/transformation.py
View file

@ -21,8 +21,9 @@ class NvidiaNimQueryObject(TypedDict):
text: Required[str]
class NvidiaNimPassageObject(TypedDict):
text: Required[str]
class NvidiaNimPassageObject(TypedDict, total=False):
text: str
image: str
class NvidiaNimRerankRequest(TypedDict, total=False):
@ -53,6 +54,10 @@ class NvidiaNimRerankConfig(BaseRerankConfig):
DEFAULT_NIM_RERANK_API_BASE = "https://ai.api.nvidia.com"
# The legacy retrieval rerank route accepts text passages only. The native
# ranking subclass expands this tuple for VL models that accept images.
SUPPORTED_PASSAGE_FIELDS: tuple[str, ...] = ("text",)
def __init__(self) -> None:
pass
@ -206,11 +211,17 @@ class NvidiaNimRerankConfig(BaseRerankConfig):
if isinstance(doc, str):
passages.append({"text": doc})
elif isinstance(doc, dict):
# If document is already a dict, check if it has 'text' field
if "text" in doc:
passages.append({"text": doc["text"]})
# Preserve only the structured passage fields supported by the
# selected rerank route.
supported_fields: NvidiaNimPassageObject = {} # mutable-ok: assembling a request TypedDict
if "text" in self.SUPPORTED_PASSAGE_FIELDS and "text" in doc:
supported_fields["text"] = doc["text"]
if "image" in self.SUPPORTED_PASSAGE_FIELDS and "image" in doc:
supported_fields["image"] = doc["image"]
if supported_fields:
passages.append(supported_fields)
else:
# Otherwise, stringify the dict
# No supported fields - stringify the dict
import json
passages.append({"text": json.dumps(doc)})
@ -304,9 +315,10 @@ class NvidiaNimRerankConfig(BaseRerankConfig):
"relevance_score": ranking["logit"],
}
# Include document if it was in the original request
# Include document if it was in the original request.
# Image-only passages carry no 'text' field, so guard the lookup.
index: int = ranking["index"]
if index < len(original_passages):
if index < len(original_passages) and "text" in original_passages[index]:
result_item["document"] = {"text": original_passages[index]["text"]}
results.append(result_item)

10
litellm/llms/openai/responses/transformation.py
View file

@ -353,6 +353,16 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
)
return event_pydantic_model.model_construct(**parsed_chunk)
@staticmethod
def parse_terminal_response_from_stream_chunks(all_chunks: list[str]) -> ResponsesAPIResponse | None:
for chunk_str in reversed(all_chunks):
for event_model in (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent):
try:
return event_model.model_validate_json(chunk_str.removeprefix("data: ")).response
except ValueError:
continue
return None
@staticmethod
def get_event_model_class(event_type: str) -> Any:
"""

55
litellm/llms/runwayml/videos/transformation.py
View file

@ -1,8 +1,10 @@
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Any, Final, Literal
import httpx
from httpx._types import RequestFiles
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm.constants import RUNWAYML_DEFAULT_API_VERSION
@ -31,6 +33,29 @@ else:
LiteLLMLoggingObj = Any
class _RunwayTaskResponse(TypedDict, total=False):
id: ReadOnly[str]
status: ReadOnly[str]
createdAt: ReadOnly[str]
completedAt: ReadOnly[str]
output: ReadOnly[Sequence[str] | str]
failureCode: ReadOnly[str]
failure: ReadOnly[str]
progress: ReadOnly[int]
class _VideoObjectData(TypedDict, extra_items=object):
id: ReadOnly[str]
object: ReadOnly[Literal["video"]]
status: ReadOnly[str]
created_at: ReadOnly[int]
def _parse_runway_task_response(raw_response: httpx.Response) -> _RunwayTaskResponse:
response_data: Final[_RunwayTaskResponse] = raw_response.json()
return response_data
class RunwayMLVideoConfig(BaseVideoConfig):
"""
Configuration class for RunwayML video generation.
@ -78,7 +103,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
- size -> ratio (convert "WIDTHxHEIGHT" to "WIDTH:HEIGHT")
- seconds -> duration (convert to integer)
"""
mapped_params: Final[dict[str, Any]] = {}
mapped_params: Final[dict[str, object]] = {}
# Handle input_reference parameter - map to promptImage
if "input_reference" in video_create_optional_params:
@ -180,7 +205,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
}
"""
# Build the request data
request_data: Final[dict[str, Any]] = {
request_data: Final[dict[str, object]] = {
"model": model,
"promptText": prompt,
}
@ -189,7 +214,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
request_data.update(video_create_optional_request_params)
# RunwayML uses JSON body, no files multipart
files_list: Final[list[tuple[str, Any]]] = []
files_list: Final[RequestFiles] = []
# Append the specific endpoint for video generation
full_api_base: Final = f"{api_base}/image_to_video"
@ -216,10 +241,10 @@ class RunwayMLVideoConfig(BaseVideoConfig):
We map this to OpenAI VideoObject format.
"""
response_data: Final = raw_response.json()
response_data: Final = _parse_runway_task_response(raw_response)
# Map RunwayML task response to VideoObject format
video_data: Final[dict[str, Any]] = {
video_data: Final[_VideoObjectData] = {
"id": response_data.get("id", ""),
"object": "video",
"status": self._map_runway_status(response_data.get("status", "pending")),
@ -326,7 +351,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
# Get task status to retrieve video URL
url: Final = f"{api_base}/tasks/{encoded_video_id}"
params: Final[dict[str, Any]] = {}
params: Final[dict[str, str]] = {}
return url, params
@ -421,7 +446,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict,
extra_body: dict[str, Any] | None = None,
extra_body: Mapping[str, object] | None = None,
) -> tuple[str, dict]:
"""
Transform the video remix request for RunwayML API.
@ -448,7 +473,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
after: str | None = None,
limit: int | None = None,
order: str | None = None,
extra_query: dict[str, Any] | None = None,
extra_query: Mapping[str, object] | None = None,
) -> tuple[str, dict]:
"""
Transform the video list request for RunwayML API.
@ -484,7 +509,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
# Construct the URL for task cancellation
url: Final = f"{api_base}/tasks/{encoded_video_id}/cancel"
data: Final[dict[str, Any]] = {}
data: Final[dict[str, str]] = {}
return url, data
@ -494,7 +519,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
logging_obj: LiteLLMLoggingObj,
) -> VideoObject:
"""Transform the RunwayML video delete/cancel response."""
response_data: Final = raw_response.json()
response_data: Final = _parse_runway_task_response(raw_response)
video_obj: Final = VideoObject(
id=response_data.get("id", ""),
@ -524,7 +549,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
url: Final = f"{api_base}/tasks/{encoded_video_id}"
# Empty dict for GET request (no body)
data: Final[dict[str, Any]] = {}
data: Final[dict[str, str]] = {}
return url, data
@ -537,10 +562,10 @@ class RunwayMLVideoConfig(BaseVideoConfig):
"""
Transform the RunwayML video status retrieve response.
"""
response_data: Final = raw_response.json()
response_data: Final = _parse_runway_task_response(raw_response)
# Map RunwayML task response to VideoObject format
video_data: Final[dict[str, Any]] = {
video_data: Final[_VideoObjectData] = {
"id": response_data.get("id", ""),
"object": "video",
"status": self._map_runway_status(response_data.get("status", "pending")),
@ -572,7 +597,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
return video_obj
def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers):
def transform_video_create_character_request(self, name, video: object, api_base, litellm_params, headers):
raise NotImplementedError("video create character is not supported for RunwayML")
def transform_video_create_character_response(self, raw_response, logging_obj):

71
litellm/llms/vertex_ai/files/transformation.py
View file

@ -5,12 +5,13 @@ import json
import os
import re
import time
from collections.abc import Callable, Iterable, Iterator
from typing import Any, Final
from collections.abc import Callable, Iterable, Iterator, Mapping
from typing import Any, Final, TypedDict
import httpx
from httpx import Headers, Response
from openai.types.file_deleted import FileDeleted
from typing_extensions import ReadOnly
import litellm
from litellm._uuid import uuid
@ -50,6 +51,7 @@ from litellm.types.llms.openai import (
HttpxBinaryResponseContent,
OpenAICreateFileRequestOptionalParams,
OpenAIFileObject,
OpenAIFilesPurpose,
PathLike,
)
from litellm.types.llms.vertex_ai import GcsBucketResponse
@ -62,6 +64,46 @@ _GCP_LABEL_VALUE_MAX_LEN: Final = 63
_CUSTOM_ID_RAW_LABEL_PREFIX: Final = "b32_"
class _GcsObjectMetadataJson(TypedDict, total=False):
purpose: ReadOnly[OpenAIFilesPurpose]
class _GcsObjectJson(TypedDict, total=False):
id: ReadOnly[str]
name: ReadOnly[str]
size: ReadOnly[str]
timeCreated: ReadOnly[str]
metadata: ReadOnly[_GcsObjectMetadataJson]
class _VertexBatchRowRequest(TypedDict, total=False):
labels: ReadOnly[Mapping[str, object]]
class _VertexBatchRow(TypedDict, total=False):
request: ReadOnly[_VertexBatchRowRequest]
status: ReadOnly[str]
processed_time: ReadOnly[str]
class _OpenAIBatchOutputError(TypedDict):
code: ReadOnly[str]
message: ReadOnly[str]
class _OpenAIBatchOutputResponse(TypedDict):
status_code: ReadOnly[int]
request_id: ReadOnly[str]
body: ReadOnly[Mapping[str, object]]
class _OpenAIBatchOutputRow(TypedDict):
id: ReadOnly[str]
custom_id: ReadOnly[str]
response: ReadOnly[_OpenAIBatchOutputResponse | None]
error: ReadOnly[_OpenAIBatchOutputError | None]
def _sanitize_gcp_label_value(value: str) -> str:
"""
Sanitize a string to meet GCP label value constraints.
@ -106,7 +148,7 @@ def _decode_gcp_label_value_chunks(values: list[str]) -> str | None:
return None
def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any) -> None:
def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: object) -> None:
"""
Store OpenAI batch custom_id for Vertex batch correlation.
@ -122,7 +164,7 @@ def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any)
labels[f"litellm_custom_id_raw_{index}"] = raw_label_chunk
def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str:
def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, object]) -> str:
"""Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels)."""
raw: Final = labels.get("litellm_custom_id_raw")
if raw:
@ -186,7 +228,7 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]:
``str.splitlines()`` + ``line.strip()`` for ``\\n`` / ``\\r\\n`` delimited
JSONL.
"""
content: Any = openai_file_content
content: FileTypes | str = openai_file_content
if isinstance(content, tuple):
content = content[1]
@ -246,6 +288,11 @@ def _iter_openai_jsonl_entries(
yield json.loads(line)
def _parse_vertex_batch_output_row(line: str) -> _VertexBatchRow:
row: Final[_VertexBatchRow] = json.loads(line)
return row
class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream):
"""Streams an OpenAI batch JSONL upload as Vertex-wrapped JSONL one row at a
time, so the transformed payload is never held in full.
@ -463,7 +510,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
"""
Transform VertexAI File upload response into OpenAI-style FileObject
"""
response_json: Final = raw_response.json()
response_json: Final[GcsBucketResponse] = raw_response.json()
try:
response_object: Final = GcsBucketResponse(**response_json)
@ -523,7 +570,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> OpenAIFileObject:
response_json: Final = raw_response.json()
response_json: Final[_GcsObjectJson] = raw_response.json()
gcs_id = response_json.get("id", "")
gcs_id = "/".join(gcs_id.split("/")[:-1]) if gcs_id else ""
return OpenAIFileObject(
@ -682,7 +729,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
# discriminating fields. Anything else (e.g. a binary file whose
# first line is not valid UTF-8/JSON) raises and falls through to the
# passthrough below, leaving the content untouched.
first_row: Final = json.loads(first_line)
first_row: Final = _parse_vertex_batch_output_row(first_line)
is_vertex_batch_output: Final = (
"request" in first_row
and "response" in first_row
@ -723,7 +770,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
for line in itertools.chain([first_line], lines):
try:
openai_output = self._transform_single_vertex_batch_output_to_openai(
vertex_output=json.loads(line),
vertex_output=_parse_vertex_batch_output_row(line),
vertex_gemini_config=vertex_gemini_config,
logging_obj=batch_transform_logging_obj,
mock_httpx_response=mock_httpx_response,
@ -742,18 +789,18 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
def _transform_single_vertex_batch_output_to_openai(
self,
vertex_output: dict[str, Any],
vertex_output: _VertexBatchRow,
vertex_gemini_config: VertexGeminiConfig,
logging_obj: Logging,
mock_httpx_response: httpx.Response,
) -> dict[str, Any]:
) -> _OpenAIBatchOutputRow:
"""
Transform a single Vertex AI batch output line to OpenAI format.
Uses the existing VertexGeminiConfig transformation for the response.
"""
# Extract custom_id from request labels (prefer raw for OpenAI round-trip)
request_data: Final = vertex_output.get("request", {})
labels: Final = request_data.get("labels", {}) or {}
labels: Final[Mapping[str, object]] = request_data.get("labels", {}) or {}
custom_id: Final = _get_litellm_batch_custom_id_from_labels(labels)
# Check if there's an error

115
litellm/llms/vertex_ai/vertex_gemma_models/transformation.py
View file

@ -14,6 +14,11 @@ from typing import Any, Final, cast
import httpx
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
_get_httpx_client,
)
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponse
@ -123,6 +128,79 @@ class VertexGemmaConfig(OpenAIGPTConfig):
return response_json["predictions"]
@staticmethod
def _sync_post(
client: HTTPHandler | httpx.Client | None,
api_base: str,
headers: dict[str, str], # mutable-ok: forwarded to post(headers: dict | None)
request_data: dict[str, Any], # mutable-ok: forwarded to post(json: dict | ...)
timeout: float | httpx.Timeout | None,
) -> httpx.Response:
if isinstance(client, HTTPHandler):
return client.post(
url=api_base,
headers=headers,
json=request_data,
timeout=timeout,
)
if isinstance(client, httpx.Client):
if timeout is None:
return client.post(
url=api_base,
headers=headers,
json=request_data,
)
return client.post(
url=api_base,
headers=headers,
json=request_data,
timeout=timeout,
)
return _get_httpx_client().post(
url=api_base,
headers=headers,
json=request_data,
timeout=timeout,
)
@staticmethod
async def _async_post(
client: AsyncHTTPHandler | httpx.AsyncClient | None,
api_base: str,
headers: dict[str, str], # mutable-ok: forwarded to post(headers: dict | None)
request_data: dict[str, Any], # mutable-ok: forwarded to post(json: dict | ...)
timeout: float | httpx.Timeout | None,
) -> httpx.Response:
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.utils import LlmProviders
if isinstance(client, AsyncHTTPHandler):
return await client.post(
url=api_base,
headers=headers,
json=request_data,
timeout=timeout,
)
if isinstance(client, httpx.AsyncClient):
if timeout is None:
return await client.post(
url=api_base,
headers=headers,
json=request_data,
)
return await client.post(
url=api_base,
headers=headers,
json=request_data,
timeout=timeout,
)
return await get_async_httpx_client(llm_provider=LlmProviders.VERTEX_AI).post(
url=api_base,
headers=headers,
json=request_data,
timeout=timeout,
)
def completion(
self,
model: str,
@ -137,7 +215,7 @@ class VertexGemmaConfig(OpenAIGPTConfig):
acompletion: bool,
litellm_params: dict,
logger_fn: Callable | None = None,
client: httpx.Client | None = None,
client: HTTPHandler | AsyncHTTPHandler | httpx.Client | httpx.AsyncClient | None = None,
timeout: float | httpx.Timeout | None = None,
encoding=None,
custom_llm_provider: str = "vertex_ai",
@ -147,6 +225,7 @@ class VertexGemmaConfig(OpenAIGPTConfig):
Supports both sync and async requests with fake streaming.
"""
if acompletion:
async_client = client if isinstance(client, (AsyncHTTPHandler, httpx.AsyncClient)) else None
return self._async_completion(
model=model,
messages=messages,
@ -157,10 +236,12 @@ class VertexGemmaConfig(OpenAIGPTConfig):
logging_obj=logging_obj,
optional_params=optional_params,
litellm_params=litellm_params,
client=async_client,
timeout=timeout,
encoding=encoding,
)
else:
sync_client = client if isinstance(client, (HTTPHandler, httpx.Client)) else None
return self._sync_completion(
model=model,
messages=messages,
@ -171,6 +252,7 @@ class VertexGemmaConfig(OpenAIGPTConfig):
logging_obj=logging_obj,
optional_params=optional_params,
litellm_params=litellm_params,
client=sync_client,
timeout=timeout,
encoding=encoding,
)
@ -186,11 +268,11 @@ class VertexGemmaConfig(OpenAIGPTConfig):
logging_obj: Any,
optional_params: dict,
litellm_params: dict,
timeout: float | httpx.Timeout | None,
encoding: Any,
client: HTTPHandler | httpx.Client | None = None,
timeout: float | httpx.Timeout | None = None,
encoding: Any = None,
):
"""Synchronous completion request"""
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.utils import convert_to_model_response_object
# Check if streaming is requested (will be faked)
@ -222,11 +304,11 @@ class VertexGemmaConfig(OpenAIGPTConfig):
)
# Make the HTTP request
http_handler: Final = HTTPHandler(concurrent_limit=1)
response: Final = http_handler.post(
url=api_base,
response: Final = self._sync_post(
client=client,
api_base=api_base,
headers=headers,
json=request_data,
request_data=request_data,
timeout=timeout,
)
@ -276,12 +358,11 @@ class VertexGemmaConfig(OpenAIGPTConfig):
logging_obj: Any,
optional_params: dict,
litellm_params: dict,
timeout: float | httpx.Timeout | None,
encoding: Any,
client: AsyncHTTPHandler | httpx.AsyncClient | None = None,
timeout: float | httpx.Timeout | None = None,
encoding: Any = None,
):
"""Asynchronous completion request"""
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.utils import LlmProviders
from litellm.utils import convert_to_model_response_object
# Check if streaming is requested (will be faked)
@ -313,13 +394,11 @@ class VertexGemmaConfig(OpenAIGPTConfig):
)
# Make the HTTP request
http_handler: Final = get_async_httpx_client(
llm_provider=LlmProviders.VERTEX_AI,
)
response: Final = await http_handler.post(
url=api_base,
response: Final = await self._async_post(
client=client,
api_base=api_base,
headers=headers,
json=request_data,
request_data=request_data,
timeout=timeout,
)

71
litellm/llms/vertex_ai/videos/transformation.py
View file

@ -7,10 +7,12 @@ Based on: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-refer
import base64
import time
from typing import TYPE_CHECKING, Any, Final, cast
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
import httpx
from httpx._types import RequestFiles
from typing_extensions import ReadOnly
from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS
from litellm.images.utils import ImageEditRequestUtils
@ -40,11 +42,37 @@ else:
BaseLLMException = Any
class _VeoVideo(TypedDict, total=False):
gcsUri: ReadOnly[str]
bytesBase64Encoded: ReadOnly[str]
mimeType: ReadOnly[str]
class _VeoOperationResponse(TypedDict, total=False):
videos: ReadOnly[Sequence[_VeoVideo]]
class _VeoOperationMetadata(TypedDict, total=False):
createTime: ReadOnly[str]
class _VeoOperation(TypedDict, total=False):
name: ReadOnly[str]
done: ReadOnly[bool]
metadata: ReadOnly[_VeoOperationMetadata]
response: ReadOnly[_VeoOperationResponse]
def _parse_veo_operation(raw_response: httpx.Response) -> _VeoOperation:
operation: Final[_VeoOperation] = raw_response.json()
return operation
def _build_vertex_video_usage_from_request_data(
request_data: dict[str, Any] | None,
) -> dict[str, Any]:
) -> dict[str, float | str]:
"""Build usage metadata (duration, resolution) for video cost calculation."""
usage_data: Final[dict[str, Any]] = {}
usage_data: Final[dict[str, float | str]] = {}
if not request_data:
return usage_data
@ -125,7 +153,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
video_create_optional_params: VideoCreateOptionalRequestParams,
model: str,
drop_params: bool,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Map OpenAI-style parameters to Veo format.
@ -135,7 +163,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
- size → aspectRatio (e.g., "1280x720" → "16:9")
- seconds → durationSeconds (defaults to 4 seconds if not provided)
"""
mapped_params: Final[dict[str, Any]] = {}
mapped_params: Final[dict[str, object]] = {}
# Map input_reference to image (will be processed in transform_video_create_request)
if "input_reference" in video_create_optional_params:
@ -289,7 +317,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
}
"""
# Build instance with prompt
instance_dict: Final[dict[str, Any]] = {"prompt": prompt}
instance_dict: Final[dict[str, object]] = {"prompt": prompt}
params_copy: Final = video_create_optional_request_params.copy()
# Check if user wants to provide full instance dict
@ -324,13 +352,13 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
# {"parameters": {"parameters": {...}}} ← wrong
# {"parameters": {...}} ← correct
nested_params: Final = params_copy.pop("parameters", None)
vertex_params: Final[dict[str, Any]] = {}
vertex_params: Final[dict[str, object]] = {}
if isinstance(nested_params, dict):
vertex_params.update(nested_params)
vertex_params.update(params_copy)
# Build request data directly (TypedDict doesn't have model_dump)
request_data: Final[dict[str, Any]] = {"instances": [instance_dict]}
request_data: Final[dict[str, object]] = {"instances": [instance_dict]}
# Only add parameters if there are any
if vertex_params:
@ -363,7 +391,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
- status: "processing"
- usage: includes duration_seconds and optional video_resolution for cost calculation
"""
response_data: Final = raw_response.json()
response_data: Final = _parse_veo_operation(raw_response)
operation_name: Final = response_data.get("name")
if not operation_name:
@ -441,7 +469,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
}
}
"""
response_data: Final = raw_response.json()
response_data: Final = _parse_veo_operation(raw_response)
operation_name: Final = response_data.get("name", "")
is_done: Final = response_data.get("done", False)
@ -513,7 +541,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
Extracts the base64 encoded video from the response and decodes it to bytes.
"""
response_data: Final = raw_response.json()
response_data: Final = _parse_veo_operation(raw_response)
if not response_data.get("done", False):
raise ValueError(
@ -548,7 +576,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict,
extra_body: dict[str, Any] | None = None,
extra_body: dict[str, object] | None = None,
) -> tuple[str, dict]:
"""
Video remix is not supported by Veo API.
@ -574,7 +602,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
after: str | None = None,
limit: int | None = None,
order: str | None = None,
extra_query: dict[str, Any] | None = None,
extra_query: dict[str, object] | None = None,
) -> tuple[str, dict]:
"""
Video list is not supported by Veo API.
@ -615,7 +643,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
"""Video delete is not supported."""
raise NotImplementedError("Video delete is not supported by Vertex AI Veo.")
def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers):
def transform_video_create_character_request(self, name, video: object, api_base, litellm_params, headers):
raise NotImplementedError("video create character is not supported for Vertex AI")
def transform_video_create_character_response(self, raw_response, logging_obj):
@ -649,7 +677,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict,
extra_body: dict[str, Any] | None = None,
extra_body: dict[str, object] | None = None,
prefetched_source_data: dict[str, Any] | None = None,
) -> tuple[str, dict]:
"""
@ -667,12 +695,13 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
if not prefetched_source_data.get("done", False):
raise ValueError("Source video generation is not complete yet. Check the video status before editing.")
videos: Final = prefetched_source_data.get("response", {}).get("videos", [])
source_response: Final[_VeoOperationResponse] = prefetched_source_data.get("response", {})
videos: Final = source_response.get("videos", [])
if not videos:
raise ValueError("No videos found in the completed operation. Cannot edit.")
source_video: Final = videos[0]
video_input: Final[dict[str, Any]] = {}
video_input: Final[dict[str, str]] = {}
if "gcsUri" in source_video:
video_input["gcsUri"] = source_video["gcsUri"]
elif "bytesBase64Encoded" in source_video:
@ -684,13 +713,13 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
operation_name: Final = extract_original_video_id(video_id)
model: Final = self.extract_model_from_operation_name(operation_name) or ""
instance_dict: Final[dict[str, Any]] = {"prompt": prompt, "video": video_input}
request_data: Final[dict[str, Any]] = {"instances": [instance_dict]}
instance_dict: Final[dict[str, object]] = {"prompt": prompt, "video": video_input}
request_data: Final[dict[str, object]] = {"instances": [instance_dict]}
if extra_body:
extra_body_copy: Final = dict(extra_body)
nested_params: Final = extra_body_copy.pop("parameters", None)
vertex_params: Final[dict[str, Any]] = {}
vertex_params: Final[dict[str, object]] = {}
if isinstance(nested_params, dict):
vertex_params.update(nested_params)
vertex_params.update(extra_body_copy)
@ -716,7 +745,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
usage includes duration_seconds and optional video_resolution from the
edit request parameters for cost calculation.
"""
response_data: Final = raw_response.json()
response_data: Final = _parse_veo_operation(raw_response)
operation_name: Final = response_data.get("name")
if not operation_name:

31
litellm/llms/xai/chat/transformation.py
View file

@ -1,4 +1,4 @@
from collections.abc import AsyncIterator, Iterator
from collections.abc import AsyncIterator, Iterator, Mapping
from typing import Any, Final
import httpx
@ -12,13 +12,15 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
strip_name_from_messages,
)
from litellm.llms.xai.common_utils import XAIModelInfo
from litellm.llms.xai.cost_calculator import (
apply_server_side_tool_usage_details_to_usage,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import (
Choices,
ModelResponse,
ModelResponseStream,
PromptTokensDetailsWrapper,
Usage,
)
@ -248,7 +250,7 @@ class XAIChatConfig(OpenAIGPTConfig):
XAI API returns empty string for finish_reason when using tools,
so we need to fix this after the standard OpenAI transformation.
Also handles X.AI web search usage tracking by extracting num_sources_used.
Also handles X.AI web search usage tracking.
"""
# First, let the parent class handle the standard transformation
@ -351,25 +353,20 @@ class XAIChatConfig(OpenAIGPTConfig):
def _enhance_usage_with_xai_web_search_fields(self, model_response: ModelResponse, raw_response_json: dict) -> None:
"""
Extract num_sources_used from X.AI response and map it to web_search_requests.
Copy usage.server_side_tool_usage_details from the provider usage block
onto model_response.usage for tool cost calculation.
"""
if not hasattr(model_response, "usage") or model_response.usage is None:
return
usage: Final[Usage] = model_response.usage
num_sources_used = None
response_usage: Final = raw_response_json.get("usage", {})
if isinstance(response_usage, dict) and "num_sources_used" in response_usage:
num_sources_used = response_usage.get("num_sources_used")
# Map num_sources_used to web_search_requests for cost detection
if num_sources_used is not None and num_sources_used > 0:
if usage.prompt_tokens_details is None:
usage.prompt_tokens_details = PromptTokensDetailsWrapper()
usage.prompt_tokens_details.web_search_requests = int(num_sources_used)
setattr(usage, "num_sources_used", int(num_sources_used))
verbose_logger.debug("X.AI web search sources used: %s", num_sources_used)
response_usage: Final = raw_response_json.get("usage")
if not isinstance(response_usage, dict):
return
details: Final = response_usage.get("server_side_tool_usage_details")
if isinstance(details, Mapping):
apply_server_side_tool_usage_details_to_usage(usage, details)
verbose_logger.debug("X.AI server_side_tool_usage_details: %s", details)
@staticmethod
def _normalize_openai_compatible_usage_totals(

98
litellm/llms/xai/cost_calculator.py
View file

@ -4,14 +4,37 @@ Helper util for handling XAI-specific cost calculation
- Handles XAI-specific reasoning token billing (billed as part of completion tokens)
"""
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
from litellm.types.utils import Usage
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
if TYPE_CHECKING:
from litellm.types.utils import ModelInfo
# https://docs.x.ai/developers/pricing#tools-pricing — default when unset in model map
_DEFAULT_WEB_SEARCH_COST_PER_CALL: Final = 5.0 / 1000.0
def apply_server_side_tool_usage_details_to_usage(usage: Usage, details: Mapping[str, object] | None) -> None:
"""
Attach server_side_tool_usage_details and mirror web_search_calls onto
prompt_tokens_details.web_search_requests for built-in tool cost gating.
"""
if details is None:
return
usage.server_side_tool_usage_details = details # pyright: ignore[reportAttributeAccessIssue] # extra # rebind-ok: extras
try:
web_search_calls: Final = int(details.get("web_search_calls") or 0)
except (TypeError, ValueError):
return
if web_search_calls <= 0:
return
prompt_tokens_details: Final = usage.prompt_tokens_details or PromptTokensDetailsWrapper()
prompt_tokens_details.web_search_requests = web_search_calls
usage.prompt_tokens_details = prompt_tokens_details # rebind-ok: write details onto caller usage
def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
"""
@ -32,9 +55,11 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
prompt_tokens: Final = int(getattr(usage, "prompt_tokens", 0) or 0)
completion_tokens: Final = int(getattr(usage, "completion_tokens", 0) or 0)
total_tokens: Final = int(getattr(usage, "total_tokens", 0) or 0)
reasoning_tokens = 0
if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details:
reasoning_tokens = int(getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0)
reasoning_tokens: Final = (
int(getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0)
if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details
else 0
)
already_normalised: Final = total_tokens == prompt_tokens + completion_tokens
total_completion_tokens: Final = completion_tokens if already_normalised else completion_tokens + reasoning_tokens
@ -52,33 +77,48 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
return prompt_cost, completion_cost
def _web_search_cost_per_call_from_model_info(model_info: "ModelInfo") -> float:
"""
Per-invocation web_search price from model_info when configured.
Prefer ``search_context_cost_per_query`` (same shape as Gemini/Anthropic web
search pricing in the model cost map). Fall back to current xAI list pricing.
"""
search_costs: Final = model_info.get("search_context_cost_per_query")
if not isinstance(search_costs, Mapping):
return _DEFAULT_WEB_SEARCH_COST_PER_CALL
for key in (
"search_context_size_medium",
"search_context_size_low",
"search_context_size_high",
):
value = search_costs.get(key)
if value is None:
continue
try:
cost = float(value)
except (TypeError, ValueError):
continue
if cost > 0:
return cost
return _DEFAULT_WEB_SEARCH_COST_PER_CALL
def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> float:
"""
Calculate the cost of web search requests for X.AI models.
X.AI Live Search costs $25 per 1,000 sources used.
Each source costs $0.025.
The number of sources is stored in prompt_tokens_details.web_search_requests
by the transformation layer to be compatible with the existing detection system.
Counts invocations from usage.server_side_tool_usage_details.web_search_calls.
Per-call rate comes from model_info.search_context_cost_per_query when set,
otherwise the default xAI tools rate ($5 / 1k calls).
"""
# Cost per source used: $25 per 1,000 sources = $0.025 per source
cost_per_source: Final = 25.0 / 1000.0 # $0.025
num_sources_used = 0
if (
hasattr(usage, "prompt_tokens_details")
and usage.prompt_tokens_details is not None
and hasattr(usage.prompt_tokens_details, "web_search_requests")
and usage.prompt_tokens_details.web_search_requests is not None
):
num_sources_used = int(usage.prompt_tokens_details.web_search_requests)
# Fallback: try to get from num_sources_used if set directly
elif hasattr(usage, "num_sources_used") and usage.num_sources_used is not None:
num_sources_used = int(usage.num_sources_used)
total_cost: Final = cost_per_source * num_sources_used
return total_cost
details: Final = getattr(usage, "server_side_tool_usage_details", None)
if not isinstance(details, Mapping):
return 0.0
try:
web_search_calls: Final = int(details.get("web_search_calls") or 0)
except (TypeError, ValueError):
return 0.0
if web_search_calls <= 0:
return 0.0
return _web_search_cost_per_call_from_model_info(model_info) * web_search_calls

9
litellm/llms/xai/responses/transformation.py
View file

@ -1,4 +1,4 @@
from typing import TYPE_CHECKING, Any, Final
from typing import Any, Final
import litellm
from litellm._logging import verbose_logger
@ -12,13 +12,6 @@ from litellm.types.llms.xai import XAIWebSearchTool, XAIXSearchTool
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
"""

848
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

4
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -1190,7 +1190,7 @@ class MCPRequestHandler:
DEPRECATED: This method is deprecated in favor of server-specific auth headers using the format x-mcp-{{server_alias}}-{{header_name}} instead.
"""
mcp_client_side_auth_header_name: Final[str] = MCPRequestHandler._get_mcp_client_side_auth_header_name()
mcp_client_side_auth_header_name: Final[str] = MCPRequestHandler.get_mcp_client_side_auth_header_name()
auth_header: Final = headers.get(mcp_client_side_auth_header_name)
if auth_header:
verbose_logger.warning(
@ -1265,7 +1265,7 @@ class MCPRequestHandler:
return oauth2_headers
@staticmethod
def _get_mcp_client_side_auth_header_name() -> str:
def get_mcp_client_side_auth_header_name() -> str:
"""
Get the header name used to pass the MCP auth header to the MCP server

107
litellm/proxy/_experimental/mcp_server/db.py
View file

@ -552,13 +552,20 @@ async def get_all_mcp_servers(
) -> list[LiteLLM_MCPServerTable]:
"""
Returns mcp servers from the db, optionally filtered by approval_status.
Pass approval_status=None to return all servers regardless of approval state.
Pass approval_status=None to return every server except drafts, which back the admin OAuth
session flow, are addressable only by their own server_id, and must never appear in a listing.
NULL approval_status predates the approval workflow, so those rows are kept explicitly rather
than dropped by a bare inequality, which SQL evaluates as NULL and would silently hide them.
"""
try:
where: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = {}
if approval_status is not None:
where["approval_status"] = approval_status
mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where if where else {})
where: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = (
{"approval_status": approval_status}
if approval_status is not None
# mutable-ok: prisma where-inputs must be plain dicts, and both `NOT` and `not` drop
# NULL rows (measured), so the OR is the only NULL-preserving way to exclude drafts
else {"OR": [{"approval_status": None}, {"approval_status": {"not": MCPApprovalStatus.draft}}]}
)
mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where)
tables: Final = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers]
for table in tables:
@ -814,6 +821,96 @@ async def create_mcp_server(
return new_mcp_server
async def create_draft_mcp_server(
prisma_client: PrismaClient,
data: NewMCPServerRequest,
touched_by: str,
ttl_seconds: int,
server_id: str | None = None,
) -> LiteLLM_MCPServerTable:
"""
Persist a short-lived draft row backing the admin OAuth "Authorize & Fetch Token" flow.
The draft lives in the database rather than in process memory so that the /register,
/authorize and /token legs resolve it whichever worker or replica accepts each request.
Writing is strictly create-if-absent. Any existing row for the id is returned untouched, which
covers both a live draft for this same session and a real server the edit form is
re-authorizing against its own id, where writing a draft would collide on the primary key.
Each click of Authorize mints a fresh id, so nothing is lost by never overwriting, and it is
what makes concurrent callers sharing one id safe rather than mutually destructive.
"""
draft_id: Final = server_id or data.server_id or str(uuid.uuid4())
await _prune_expired_draft_mcp_servers(prisma_client, ttl_seconds)
existing: Final = await _db_find_mcp_server_row(prisma_client, draft_id)
if existing is not None:
# Already usable by every worker, whether it is a live draft for this same session or a
# real server the edit form is re-authorizing. Either way there is nothing to write, and
# not writing is what keeps concurrent callers for one server_id from racing each other.
return LiteLLM_MCPServerTable.model_validate(existing.model_dump())
draft_payload: Final = data.model_copy(update={"server_id": draft_id, "approval_status": MCPApprovalStatus.draft})
try:
return await create_mcp_server(prisma_client, draft_payload, touched_by)
except Exception:
# Lost the create race: the read above and this create are two statements, not one. The
# winner wrote a draft for this same session, so adopt it rather than failing a caller
# whose session is in fact ready. Anything else still raises.
raced: Final = await _db_find_mcp_server_row(prisma_client, draft_id)
if raced is None or raced.approval_status != MCPApprovalStatus.draft:
raise
return LiteLLM_MCPServerTable.model_validate(raced.model_dump())
async def _prune_expired_draft_mcp_servers(prisma_client: PrismaClient, ttl_seconds: int) -> None:
"""Drop drafts already past ``ttl_seconds``, so abandoned OAuth sessions do not accumulate.
Runs on each draft write rather than on a schedule, mirroring the in-memory cache this
replaces, which pruned on every store. Expired drafts are unreadable by then anyway, so the
only thing at stake is row count, and the work is bounded by how often admins authorize.
"""
cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=max(1, ttl_seconds))
# Age is filtered here rather than in the query: the draft set is bounded by how many OAuth
# authorizations are in flight, so it is a handful of rows even on a busy proxy.
drafts: Final = await _db_find_mcp_server_rows(
prisma_client,
where={"approval_status": MCPApprovalStatus.draft},
)
for row in drafts:
# A row without a timestamp has no age to judge, so leave it rather than guess it is stale.
# Two workers sweeping the same row is harmless: prisma's delete returns None for a row
# that is already gone rather than raising, so the loser of that race is a no-op.
if row.updated_at is not None and row.updated_at < cutoff:
await delete_mcp_server(prisma_client, row.server_id)
async def get_draft_mcp_server(
prisma_client: PrismaClient, server_id: str, ttl_seconds: int
) -> LiteLLM_MCPServerTable | None:
"""
Return the draft row for ``server_id`` if it has not yet aged past ``ttl_seconds``, else None.
Age is enforced in the query rather than by a sweeper so an expired draft is unreadable the
moment it lapses, regardless of which process last ran a cleanup.
"""
cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=max(1, ttl_seconds))
draft_rows: Final = await _db_find_mcp_server_rows(
prisma_client,
where={
"server_id": server_id,
"approval_status": MCPApprovalStatus.draft,
"updated_at": {"gte": cutoff},
},
)
if not draft_rows:
return None
table: Final = LiteLLM_MCPServerTable.model_validate(draft_rows[0].model_dump())
decrypt_global_env_var_values(table.env_vars)
return table
async def update_mcp_server(
prisma_client: PrismaClient,
data: UpdateMCPServerRequest,

2
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -118,6 +118,7 @@ from litellm.proxy._experimental.mcp_server.utils import (
is_short_mcp_tool_prefix_enabled,
iter_known_server_prefixes,
iter_known_tool_name_spellings,
logging_safe_mcp_headers,
match_known_server_prefix,
match_known_tool_name,
merge_mcp_headers,
@ -4603,6 +4604,7 @@ class MCPServerManager:
),
"user_api_key_hash": (getattr(user_api_key_auth, "api_key_hash", None) if user_api_key_auth else None),
"incoming_bearer_token": incoming_bearer_token,
"headers": logging_safe_mcp_headers(raw_headers),
}
# Create MCP request object for processing

99
litellm/proxy/_experimental/mcp_server/sampling_handler.py
View file

@ -1042,100 +1042,15 @@ def _build_sampling_request(
raw_headers: dict[str, str] | None = None,
client_ip: str | None = None,
) -> "Request":
"""Build a synthetic FastAPI Request for sampling sub-calls.
"""The synthetic FastAPI Request for sampling sub-calls, carrying the original
MCP connection's headers and client IP."""
from litellm.proxy._experimental.mcp_server.utils import build_synthetic_mcp_request
Converts the original MCP connection's HTTP headers into ASGI
scope format so that ``add_litellm_data_to_request`` can apply
header-dependent guardrails, tag-based routing, trace correlation,
and ``forward_llm_provider_auth_headers``.
Key fields populated:
- **headers**: All original HTTP headers are forwarded (except
hop-by-hop: content-length, transfer-encoding). This ensures
``traceparent``, ``authorization``, ``user-agent``, and
``x-litellm-api-key`` are visible to pre-call utils.
- **client**: The ASGI ``(host, port)`` tuple so that
``request.client.host`` returns the real client IP for
IP-based routing and guardrails.
- **server**: Derived from the running proxy's ``server_host``
/ ``server_port`` when available, avoiding the misleading
``127.0.0.1:0`` placeholder.
- **x-forwarded-for**: Injected from ``client_ip`` if the
original headers don't already carry it, as a fallback for
IP attribution.
"""
from fastapi import Request
# --- Build ASGI headers ---
_scope_headers: Final[list[tuple[bytes, bytes]]] = [(b"content-type", b"application/json")]
# Hop-by-hop headers that must NOT be forwarded into the
# synthetic request (they describe the original HTTP framing,
# not the logical request).
_HOP_BY_HOP: Final = frozenset(
{
"content-length",
"transfer-encoding",
"connection",
"keep-alive",
"upgrade",
"te",
"trailer",
}
return build_synthetic_mcp_request(
path="/mcp/sampling/createMessage",
raw_headers=raw_headers,
client_ip=client_ip,
)
if raw_headers:
for hdr_name, hdr_value in raw_headers.items():
_key = hdr_name.lower()
# Skip content-type (already set), x-forwarded-for (use resolved
# client_ip instead to prevent spoofing), and hop-by-hop headers
if _key in {"content-type", "x-forwarded-for"} or _key in _HOP_BY_HOP:
continue
_scope_headers.append(
(
_key.encode("latin-1", errors="replace"),
hdr_value.encode("utf-8"),
)
)
# Inject x-forwarded-for from captured client_ip if the
# original headers don't already carry it
if client_ip and not any(h[0] == b"x-forwarded-for" for h in _scope_headers):
_scope_headers.append((b"x-forwarded-for", client_ip.encode("utf-8")))
# --- Derive server (host, port) from the running proxy ---
_server_host = "127.0.0.1"
_server_port = 4000 # LiteLLM default
try:
from litellm.proxy import proxy_server
_proxy_host: Final[str | None] = getattr(proxy_server, "server_host", None)
_proxy_port: Final[str | int | None] = getattr(proxy_server, "server_port", None)
if _proxy_host:
_server_host = str(_proxy_host)
if _proxy_port:
_server_port = int(_proxy_port)
except (ImportError, AttributeError, TypeError, ValueError):
pass
# --- Build ASGI client tuple for request.client.host ---
_client_tuple = None
if client_ip:
_client_tuple = (client_ip, 0)
scope: Final[dict[str, object]] = {
"type": "http",
"method": "POST",
"path": "/mcp/sampling/createMessage",
"scheme": "http",
"server": (_server_host, _server_port),
"query_string": b"",
"root_path": "",
"headers": _scope_headers,
}
if _client_tuple is not None:
scope["client"] = _client_tuple
return Request(scope=scope)
async def _build_completion_kwargs(

36
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -58,9 +58,11 @@ from litellm.proxy._experimental.mcp_server.utils import (
LITELLM_MCP_SERVER_VERSION,
MCPMissingUserEnvVarsError,
add_server_prefix_to_name,
build_synthetic_mcp_request,
extract_mcp_tool_result_error_message,
get_server_prefix,
iter_known_server_prefixes,
logging_safe_mcp_headers,
match_known_tool_name,
)
from litellm.proxy._types import (
@ -860,11 +862,11 @@ if MCP_AVAILABLE:
name: str,
arguments: dict[str, object],
user_api_key_auth: UserAPIKeyAuth,
raw_headers: Mapping[str, str] | None = None,
client_ip: str | None = None,
) -> LiteLLMLoggingObj | None:
"""Run the pre-call pipeline (guardrails + logging setup) for a virtual
mcp_tool_call so the SSE path spend-logs like the REST path."""
from fastapi import Request
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
)
@ -874,13 +876,10 @@ if MCP_AVAILABLE:
proxy_logging_obj,
)
request: Final = Request(
scope={
"type": "http",
"method": "POST",
"path": "/mcp/tools/call",
"headers": [(b"content-type", b"application/json")],
}
request: Final = build_synthetic_mcp_request(
path="/mcp/tools/call",
raw_headers=raw_headers,
client_ip=client_ip,
)
_, virtual_logging_obj = await ProxyBaseLLMRequestProcessing(
data={"name": name, "arguments": arguments}
@ -952,7 +951,11 @@ if MCP_AVAILABLE:
assert user_api_key_auth is not None # guaranteed by the flag check above
virtual_logging_obj: Final = await _build_virtual_call_logging_obj(
name=name, arguments=args, user_api_key_auth=user_api_key_auth
name=name,
arguments=args,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
)
return await handle_mcp_tool_call(
tool_name=args.get("tool_name", ""),
@ -979,7 +982,6 @@ if MCP_AVAILABLE:
Raises:
HTTPException: If tool not found or arguments missing
"""
from fastapi import Request
from mcp.server.lowlevel.server import request_ctx
from mcp.types import CallToolResult
@ -1041,13 +1043,10 @@ if MCP_AVAILABLE:
body_data["litellm_trace_id"] = chain_id
body_data["litellm_session_id"] = chain_id
request: Final = Request(
scope={
"type": "http",
"method": "POST",
"path": "/mcp/tools/call",
"headers": [(b"content-type", b"application/json")],
}
request: Final = build_synthetic_mcp_request(
path="/mcp/tools/call",
raw_headers=raw_headers,
client_ip=_client_ip,
)
if user_api_key_auth is not None:
data = await add_litellm_data_to_request(
@ -1905,6 +1904,7 @@ if MCP_AVAILABLE:
"litellm_trace_id": effective_litellm_trace_id,
"metadata": {
"spend_logs_metadata": spend_logs_metadata,
"headers": logging_safe_mcp_headers(raw_headers),
**({"tags": request_tags} if request_tags else {}),
},
# Provide a small input payload for standard logging

218
litellm/proxy/_experimental/mcp_server/utils.py
View file

@ -7,12 +7,38 @@ import importlib
import json
import os
import re
import typing
from collections.abc import Iterable, Iterator, Mapping, MutableMapping, MutableSequence
from typing import Any, Final
from collections.abc import Set as AbstractSet
from typing import Any, Final, Protocol
from urllib.parse import quote
from litellm.types.mcp_server.mcp_server_manager import MCPServer
if typing.TYPE_CHECKING:
from fastapi import Request
class _McpServerLike(Protocol):
@property
def server_id(self) -> str: ...
@property
def server_name(self) -> str | None: ...
@property
def alias(self) -> str | None: ...
@property
def short_prefix(self) -> str | None: ...
class McpServerPayloadLike(Protocol):
alias: str | None
@property
def server_name(self) -> str | None: ...
@property
def tool_name_to_display_name(self) -> Mapping[str, str] | None: ...
# Constants
#
# NOTE: The environment-backed values below are read once, when this module is
@ -102,7 +128,7 @@ def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str:
# at the end so the first emitted char comes from the high-order
# bits of the digest (which is the position we constrain to be
# alphabetic).
chars: Final = []
chars: Final[list[str]] = []
for position in range(SHORT_MCP_TOOL_PREFIX_LENGTH):
is_first_char = position == SHORT_MCP_TOOL_PREFIX_LENGTH - 1
alphabet = _BASE52_ALPHA_ALPHABET if is_first_char else _BASE62_ALPHABET
@ -176,34 +202,34 @@ def lookup_mcp_server_auth_in_headers(
MCP_TOOL_ALLOWLIST_ENFORCED_KEY: Final = "tool_allowlist_enforced"
def _parse_mcp_info_dict(mcp_info: Any) -> dict[str, Any] | None:
def _parse_mcp_info_dict(mcp_info: object) -> Mapping[str, object] | None:
if mcp_info is None:
return None
if isinstance(mcp_info, dict):
return mcp_info
if isinstance(mcp_info, str):
try:
parsed: Final = json.loads(mcp_info)
parsed: Final[object] = json.loads(mcp_info)
except (ValueError, TypeError):
return None
return parsed if isinstance(parsed, dict) else None
return None
def is_server_tool_allowlist_enforced(mcp_server: Any) -> bool:
def is_server_tool_allowlist_enforced(mcp_server: object) -> bool:
mcp_info: Final = _parse_mcp_info_dict(getattr(mcp_server, "mcp_info", None))
if not mcp_info:
return False
return bool(mcp_info.get(MCP_TOOL_ALLOWLIST_ENFORCED_KEY))
def server_applies_tool_allowlist(mcp_server: Any) -> bool:
def server_applies_tool_allowlist(mcp_server: object) -> bool:
"""Whether server-level allowed_tools whitelist filtering is active."""
allowed_tools: Final = getattr(mcp_server, "allowed_tools", None) or []
allowed_tools: Final[object] = getattr(mcp_server, "allowed_tools", None) or []
return is_server_tool_allowlist_enforced(mcp_server) or bool(allowed_tools)
def validate_and_normalize_mcp_server_payload(payload: Any) -> None:
def validate_and_normalize_mcp_server_payload(payload: McpServerPayloadLike) -> None:
"""
Validate and normalize MCP server payload fields (server_name, alias, and
tool_name_to_display_name).
@ -233,8 +259,8 @@ def validate_and_normalize_mcp_server_payload(payload: Any) -> None:
validate_tool_display_names(payload.tool_name_to_display_name)
# Alias normalization and defaulting
alias = getattr(payload, "alias", None)
server_name: Final = getattr(payload, "server_name", None)
alias: str | None = getattr(payload, "alias", None)
server_name: Final[str | None] = getattr(payload, "server_name", None)
if not alias and server_name:
alias = normalize_server_name(server_name)
@ -257,7 +283,7 @@ def add_server_prefix_to_name(name: str, server_name: str) -> str:
)
def get_server_prefix(server: Any) -> str:
def get_server_prefix(server: object) -> str:
"""Return the prefix for a server.
When the short-prefix mode is enabled (``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``)
@ -270,23 +296,26 @@ def get_server_prefix(server: Any) -> str:
alias if present, else server_name, else server_id.
"""
if is_short_mcp_tool_prefix_enabled():
cached: Final = getattr(server, "short_prefix", None)
cached: Final[str | None] = getattr(server, "short_prefix", None)
if cached:
return cached
server_id: Final = getattr(server, "server_id", None)
server_id: Final[str | None] = getattr(server, "server_id", None)
if server_id:
return compute_short_server_prefix(server_id)
if hasattr(server, "alias") and server.alias:
return server.alias
if hasattr(server, "server_name") and server.server_name:
return server.server_name
alias: Final[str | None] = getattr(server, "alias", None)
if alias:
return alias
server_name: Final[str | None] = getattr(server, "server_name", None)
if server_name:
return server_name
if hasattr(server, "server_id"):
return server.server_id
fallback_server_id: Final[str] = getattr(server, "server_id", "")
return fallback_server_id
return ""
def iter_known_server_prefixes(server: Any) -> Iterator[str]:
def iter_known_server_prefixes(server: _McpServerLike) -> Iterator[str]:
"""Yield every prefix form that may appear in tool names for ``server``.
Always includes the *current* prefix returned by ``get_server_prefix``.
@ -304,7 +333,7 @@ def iter_known_server_prefixes(server: Any) -> Iterator[str]:
yield from _emit(get_server_prefix(server))
yield from _emit(getattr(server, "short_prefix", None))
server_id: Final = getattr(server, "server_id", None)
server_id: Final[str | None] = getattr(server, "server_id", None)
if server_id:
try:
yield from _emit(compute_short_server_prefix(server_id))
@ -397,7 +426,7 @@ def match_known_server_prefix(name: str, known_prefixes: Iterable[str]) -> tuple
return None
def strip_known_server_prefix(name: str, server: Any | None) -> str:
def strip_known_server_prefix(name: str, server: _McpServerLike | None) -> str:
"""Strip ``server``'s registered prefix from a prefixed tool/resource name.
Unlike :func:`split_server_prefix_from_name`, which guesses the boundary at
@ -420,7 +449,7 @@ def strip_known_server_prefix(name: str, server: Any | None) -> str:
def is_tool_name_prefixed(
tool_name: str,
known_server_prefixes: set | None = None,
known_server_prefixes: AbstractSet[str] | None = None,
) -> bool:
"""
Check if tool name has a known MCP server prefix.
@ -640,7 +669,7 @@ def parse_admin_env_vars(
if raw is None:
continue
if hasattr(raw, "model_dump"):
entry = raw.model_dump()
entry: Mapping[str, object] = raw.model_dump()
elif isinstance(raw, dict):
entry = raw
else:
@ -837,3 +866,146 @@ def set_mcp_tool_result_structured_content(result: object, value: object) -> boo
return True
except (AttributeError, TypeError, ValueError):
return False
_HOP_BY_HOP_HEADERS: Final = frozenset(
{
"content-length",
"transfer-encoding",
"connection",
"keep-alive",
"upgrade",
"te",
"trailer",
}
)
_SYNTHETIC_REQUEST_EXCLUDED_HEADERS: Final = _HOP_BY_HOP_HEADERS | frozenset({"content-type", "x-forwarded-for"})
_SYNTHETIC_REQUEST_SERVER: Final = ("127.0.0.1", 4000)
_MCP_SERVER_AUTH_HEADER_PREFIX: Final = "x-mcp-"
def _custom_litellm_key_header_name() -> str | None:
"""``general_settings.litellm_key_header_name``, the deployment's custom header name for
the proxy virtual key, so it is stripped from observability copies like the standard ones."""
try:
from litellm.proxy.proxy_server import general_settings
except ImportError:
return None
return general_settings.get("litellm_key_header_name") if general_settings else None
def _mcp_client_side_auth_header_name() -> str:
"""The header name the client passes the upstream MCP credential in, falling back to the
default when ``general_settings`` is unavailable (the SDK, outside a running proxy)."""
from .auth.user_api_key_auth_mcp import MCPRequestHandler
try:
return MCPRequestHandler.get_mcp_client_side_auth_header_name()
except ImportError:
return MCPRequestHandler.LITELLM_MCP_AUTH_HEADER_NAME
def _upstream_credential_headers(header_names: Iterable[str]) -> frozenset[str]:
"""Lowercased names of the headers in ``header_names`` that carry an upstream MCP
credential rather than request context: the configured client side auth header and
the per-server ``x-mcp-{alias}-{header}`` family. ``clean_headers`` only knows the
credential headers of the chat completions path, so these are dropped on top of it.
"""
from .auth.user_api_key_auth_mcp import MCPRequestHandler
non_credential: Final = frozenset(
{
MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME.lower(),
MCPRequestHandler.LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME.lower(),
}
)
client_side_auth: Final = _mcp_client_side_auth_header_name().lower()
return frozenset(
name
for name in (raw_name.lower() for raw_name in header_names)
if name == client_side_auth or (name.startswith(_MCP_SERVER_AUTH_HEADER_PREFIX) and name not in non_credential)
)
def build_synthetic_mcp_request(
*,
path: str,
raw_headers: Mapping[str, str] | None = None,
client_ip: str | None = None,
) -> "Request":
"""A synthetic FastAPI ``Request`` carrying the MCP connection's HTTP headers.
The MCP protocol transports do not hand a per-call ``Request`` to the tool
handlers, so one is reconstructed from the connection's ``raw_headers``. That
lets ``add_litellm_data_to_request`` derive ``metadata.headers``,
``proxy_server_request``, header-based tags, guardrails and trace correlation
exactly as on the chat completions path. Hop-by-hop headers describe the
original HTTP framing rather than the logical request, so they are dropped, and
``x-forwarded-for`` comes from the resolved ``client_ip`` to avoid spoofing. Upstream
MCP credentials and the deployment's proxy key header, including a custom
``litellm_key_header_name``, are dropped so they cannot reach a callback or a guardrail
through the derived metadata even when a caller omits ``general_settings``.
"""
from fastapi import Request
custom_key_header: Final = _custom_litellm_key_header_name()
excluded: Final = (
_SYNTHETIC_REQUEST_EXCLUDED_HEADERS
| _upstream_credential_headers(raw_headers.keys() if raw_headers else ())
| (frozenset({custom_key_header.lower()}) if custom_key_header else frozenset())
)
forwarded: Final = tuple(
(
name.lower().encode("latin-1", errors="replace"),
value.encode("utf-8", errors="replace"),
)
for name, value in (raw_headers.items() if raw_headers else ())
if name.lower() not in excluded
)
xff: Final = ((b"x-forwarded-for", client_ip.encode("utf-8")),) if client_ip else ()
return Request(
scope={
"type": "http",
"method": "POST",
"path": path,
"scheme": "http",
"server": _SYNTHETIC_REQUEST_SERVER,
"query_string": b"",
"root_path": "",
"headers": ((b"content-type", b"application/json"), *forwarded, *xff),
**({"client": (client_ip, 0)} if client_ip else {}),
}
)
def logging_safe_mcp_headers(raw_headers: Mapping[str, str] | None) -> Mapping[str, str]:
"""The MCP request's client headers, sanitized the way the chat completions path
sanitizes them before they reach a logging callback or a guardrail: proxy key
headers stripped, including the custom key header name the deployment configured,
upstream MCP credentials dropped, and credential-bearing values masked.
Client-controlled behaviour flags (``litellm-disable-message-redaction``) are dropped
too: these headers are read back out of the metadata to change proxy behaviour, so
leaving one in place would let any MCP client turn off the redaction an admin
configured. This path carries no key or team object to authorize an opt-out with, so
it always strips them."""
from starlette.datastructures import Headers
from litellm.proxy.litellm_pre_call_utils import (
UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS,
clean_headers,
redact_credential_headers,
)
excluded: Final = (
_upstream_credential_headers(raw_headers.keys() if raw_headers else ())
| UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS
)
cleaned: Final = clean_headers(
Headers(raw_headers),
litellm_key_header_name=_custom_litellm_key_header_name(),
)
return redact_credential_headers({name: value for name, value in cleaned.items() if name.lower() not in excluded})

85
litellm/proxy/_types.py
View file

@ -3,7 +3,7 @@ import json
import os
from collections.abc import Callable, Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple
import httpx
from pydantic import (
@ -18,7 +18,7 @@ from pydantic import (
from typing_extensions import NotRequired, Required, TypedDict
from litellm._uuid import uuid
from litellm.constants import MCP_STDIO_ALLOWED_COMMANDS
from litellm.constants import DEFAULT_STAGGER_WINDOW_SECONDS, MCP_STDIO_ALLOWED_COMMANDS
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
validate_no_callback_env_reference,
)
@ -73,6 +73,27 @@ else:
Span = Any
class ReconcileOutcome(NamedTuple):
"""What a model reconcile observed, captured while it still held the reconcile
lock.
Both fields have to be read under that lock to be worth anything. ``live_after``
in particular is the router's serving state the instant this reconcile finished,
which is NOT the same as what a later snapshot would see: any other model write
admitted in between briefly un-serves every db model (see ``clear_cache``), so a
caller that re-snapshots at verdict time can observe that hole and blame its own
reload for it.
- ``still_desired``: the db + config ids the reconcile reconciled against, or None
when no reconcile ran and the desired set is therefore unknown.
- ``live_after``: the ids the router served immediately after the reconcile, or
None when no reconcile ran.
"""
still_desired: frozenset[str] | None
live_after: frozenset[str] | None
class SupportedDBObjectType(str, enum.Enum):
"""
Supported database object types for fine-grained DB storage control.
@ -1262,6 +1283,9 @@ class MCPApprovalStatus(str, enum.Enum):
pending_review = "pending_review"
active = "active"
rejected = "rejected"
# Short-lived row backing the admin OAuth "Authorize & Fetch Token" flow. Never served: the
# registry loader and every listing exclude it, so it is reachable only by its own server_id.
draft = "draft"
from litellm.models.mcp_server import ( # noqa: E402
@ -2251,6 +2275,39 @@ class CoordinationRedisParams(LiteLLMPydanticObjectBase):
return any(value is not None for value in (self.host, self.url, self.startup_nodes, self.sentinel_nodes))
class ScheduledJobStaggerSettings(LiteLLMPydanticObjectBase):
"""
Spreads the proxy's scheduled background jobs across a window instead of firing them
all on one instant, on every replica, forever.
"""
model_config = ConfigDict(frozen=True, extra="forbid", protected_namespaces=())
enabled: bool = Field(default=True, description="apply deterministic phase offsets to scheduled background jobs")
window_seconds: int = Field(
default=DEFAULT_STAGGER_WINDOW_SECONDS,
ge=0,
description=(
"width of the window jobs are spread over. An interval job is never offset by "
"more than one of its own periods, so it is not delayed past the wait it already has"
),
)
identity: str | None = Field(
default=None,
description=(
"replaces the POD_NAME/HOSTNAME-derived component of the offset hash. Set this "
"when replicas share a hostname and would otherwise land on the same offset"
),
)
offsets: Mapping[str, int] = Field(
default_factory=dict,
description=(
"explicit offset in seconds per scheduler job id, overriding the derived value. "
"0 pins a job to its unshifted schedule"
),
)
class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
"""
Documents all the fields supported by `general_settings` in config.yaml
@ -2437,6 +2494,14 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="By default, the user calling /team/new is automatically added to the new team as a team admin. If True, proxy admins are no longer auto-added; members explicitly listed in members_with_roles are unaffected. Default is False.",
)
scheduled_job_stagger: ScheduledJobStaggerSettings | None = Field(
None,
description=(
"Spreads the proxy's scheduled background jobs (spend flushes, budget resets, "
"config reloads, exports) across a window instead of firing them together on "
"every replica. On by default; set to tune the window, pin a job, or turn it off."
),
)
maximum_spend_logs_retention_period: str | None = Field(
None,
description="Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted.",
@ -2449,6 +2514,22 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="If True and LiteLLM_SpendLogs has been converted to a range-partitioned table (db_scripts/partition_spend_logs.sql), retention cleanup drops expired partitions instead of deleting rows, and pre-creates upcoming partitions. Default is False.",
)
maximum_spend_logs_cleanup_batch_size: int | None = Field(
None,
description="Rows deleted per DELETE statement by the spend log cleanup job. Defaults to 1000.",
)
maximum_spend_logs_cleanup_max_batches: int | None = Field(
None,
description="Maximum DELETE statements the spend log cleanup job issues per table per run. Defaults to 500.",
)
maximum_spend_logs_cleanup_run_budget: str | None = Field(
None,
description="Wall-clock budget for one spend log cleanup run (e.g. '5m'), shared across every table it prunes. A run that hits the budget stops and the next run resumes from where it left off. Defaults to '5m'.",
)
maximum_spend_logs_cleanup_batch_timeout: str | None = Field(
None,
description="Postgres statement_timeout and lock_timeout applied to each spend log cleanup delete batch (e.g. '30s'), so cleanup cannot hold row locks or a connection indefinitely. Defaults to '30s'.",
)
mcp_internal_ip_ranges: list[str] | None = Field(
None,
description="Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8).",

137
litellm/proxy/auth/auth_checks.py
View file

@ -14,6 +14,7 @@ import math
import re
import time
from collections.abc import Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast
from fastapi import HTTPException, Request, status
@ -65,6 +66,7 @@ from litellm.proxy.auth.budget_throttle import (
should_throttle_budget_exceeded,
)
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
from litellm.proxy.common_utils.http_parsing_utils import (
_safe_get_request_headers,
@ -375,6 +377,16 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None
zero_cost_cache[model_name] = False
return False
if _has_ptu_flat_cost(model_name, llm_router):
verbose_proxy_logger.debug(
"Model %s prices reserved PTU capacity as a flat cost, so its zero per-token "
"rate is not a free model (enforce budget)",
safe_name,
)
if zero_cost_cache is not None:
zero_cost_cache[model_name] = False
return False
verbose_proxy_logger.debug(
"Model %s has zero cost explicitly configured (input: %s, output: %s)",
safe_name,
@ -393,6 +405,24 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None
return True
_NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({})
def _has_ptu_flat_cost(model: str, llm_router: "Router") -> bool:
"""Whether any deployment in the model group bills reserved PTU capacity as a flat cost.
Such a deployment carries an explicit zero per-token price so the flat cost is not charged
twice, which otherwise reads here as a free model and waives every budget check for it.
"""
for deployment in llm_router.model_list:
if deployment.get("model_name") != model:
continue
model_info = deployment.get("model_info") or _NO_MODEL_INFO
if model_info.get("ptu_count") is not None and model_info.get("cost_per_ptu_per_hour") is not None:
return True
return False
def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool:
"""
Check if any deployment in the model group has cost fields explicitly
@ -2016,6 +2046,44 @@ async def _cache_team_object(
)
async def delete_cache_team_object(
team_id: str,
team_alias: str | None,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging | None,
) -> None:
"""
Evict both keys `_cache_team_object` writes.
`get_team_object` reads the id key and the JWT `team_alias_jwt_field` path reads the alias key,
so leaving either behind keeps a deleted team resolvable for auth until its TTL expires.
Mirrors `delete_cached_project_object`: evicting locally only reaches the worker handling the
delete, so every key is also broadcast to drop the other workers' in-memory copies.
Eviction is best-effort, matching `_cache_team_object`. `delete_team` calls this after the team
rows are already gone, so letting an unreachable cache backend raise here would fail a request
whose delete has committed.
"""
keys: Final = (f"team_id:{team_id}", *((f"team_alias:{team_alias}",) if team_alias else ()))
for key in keys:
try:
user_api_key_cache.delete_cache(key=key)
## UPDATE REDIS CACHE ##
if proxy_logging_obj is not None:
await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key)
except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not abort the delete
verbose_proxy_logger.warning(
"Failed to invalidate cached team entry %s on delete; "
"a deleted team may be served until its TTL expires: %s",
key,
e,
)
await publish_auth_cache_invalidation(cache_key=key)
async def _cache_key_object(
hashed_token: str,
user_api_key_obj: UserAPIKeyAuth,
@ -2051,6 +2119,61 @@ async def _delete_cache_key_object(
await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key)
class TeamNotFoundError(HTTPException):
"""The team row is provably absent, as opposed to merely unreadable.
``get_team_object`` reports every failure as a 404, so a deleted team and a
database that would not answer are indistinguishable to its callers. Callers
that must not treat a degraded read as a definitive answer, such as the
authorization fallback in ``user_api_key_auth``, key on this subclass. It
stays a 404 carrying the same detail, so every other caller is unaffected.
"""
def __init__(self, team_id: str) -> None:
super().__init__(
status_code=404,
detail={"error": f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call."},
)
async def delete_cache_key_objects(
hashed_tokens: Sequence[str],
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging | None,
) -> None:
"""
Evict a batch of key objects, for callers that delete keys in bulk rather than through
`/key/delete`. Auth resolves a cached key object without re-reading its team, so a key left
cached after its row is gone keeps buying access until its TTL expires.
Evicting locally only reaches this worker, so each token is also broadcast: a deleted key left
in a peer worker's in-memory cache still authenticates there until its TTL expires.
Best-effort per key: the rows are already deleted by the time this runs, so an unreachable
cache backend must not abort the caller partway through its own cascade.
"""
results: Final = await asyncio.gather(
*(
_delete_cache_key_object(
hashed_token=hashed_token,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
for hashed_token in hashed_tokens
),
return_exceptions=True,
)
for hashed_token, result in zip(hashed_tokens, results):
if isinstance(result, BaseException):
verbose_proxy_logger.warning(
"Failed to evict cached key entry for %s; a deleted key may authenticate until its TTL expires: %s",
hashed_token,
result,
)
await publish_auth_cache_invalidation(cache_key=hashed_token)
@log_db_metrics
async def _get_team_db_check(
team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None
@ -2096,6 +2219,10 @@ async def _get_team_object_from_user_api_key_cache(
)
if should_check_db:
response = await _get_team_db_check(team_id=team_id, prisma_client=prisma_client, team_id_upsert=team_id_upsert)
# The database answered and the row is not there. Distinct from every
# other failure here, which leaves the team's grant unknown.
if response is None:
raise TeamNotFoundError(team_id=team_id)
else:
response = None
@ -2217,6 +2344,8 @@ async def get_team_object(
key=key,
team_id_upsert=team_id_upsert,
)
except TeamNotFoundError:
raise
except Exception:
raise HTTPException(
status_code=404,
@ -2556,6 +2685,8 @@ class ExperimentalUIJWTToken:
user_info: LiteLLM_UserTable,
team_id: str | None = None,
team_alias: str | None = None,
team_models: Sequence[str] | None = None,
team_model_aliases: Mapping[str, str] | None = None,
max_budget: float | None = None,
) -> str:
"""
@ -2568,6 +2699,8 @@ class ExperimentalUIJWTToken:
user_info: User information from the database
team_id: Team ID for the user (optional, uses user's team if available)
team_alias: Team alias for the selected team, if available
team_models: Model allowlist granted by the selected team
team_model_aliases: Team model aliases for the selected team
Returns:
Encrypted JWT token string
@ -2606,7 +2739,9 @@ class ExperimentalUIJWTToken:
user_id=user_info.user_id,
team_id=_team_id,
team_alias=team_alias,
models=user_info.models,
team_models=list(team_models) if team_models is not None else [],
team_model_aliases=dict(team_model_aliases) if team_model_aliases is not None else None,
models=[] if _team_id is not None else user_info.models,
max_parallel_requests=None,
user_role=LitellmUserRoles(user_info.user_role),
is_session_token=True,

28
litellm/proxy/auth/model_checks.py
View file

@ -52,20 +52,20 @@ def _get_models_from_access_groups(
model_access_groups: dict[str, list[str]],
all_models: list[str],
include_model_access_groups: bool | None = False,
proxy_model_list: Sequence[str] | None = None,
) -> list[str]:
idx_to_remove: Final = []
new_models: Final = []
for idx, model in enumerate(all_models):
if model in model_access_groups:
if not include_model_access_groups: # remove access group, unless requested - e.g. when creating a key
idx_to_remove.append(idx)
new_models.extend(model_access_groups[model])
for idx in sorted(idx_to_remove, reverse=True):
all_models.pop(idx)
all_models.extend(new_models)
return all_models
# a grant naming both a deployed model and an access group means both at runtime
# (_check_model_access_helper unions them), so listings must keep the literal too
deployed_model_names: Final = frozenset(proxy_model_list or ())
kept_models: Final = [
model
for model in all_models
if model not in model_access_groups or include_model_access_groups or model in deployed_model_names
]
member_models: Final = [
member for model in all_models if model in model_access_groups for member in model_access_groups[model]
]
return kept_models + member_models
async def get_mcp_server_ids(
@ -128,6 +128,7 @@ def get_key_models(
model_access_groups=model_access_groups,
all_models=all_models,
include_model_access_groups=include_model_access_groups,
proxy_model_list=proxy_model_list,
)
# deduplicate while preserving order
@ -169,6 +170,7 @@ def get_team_models(
model_access_groups=model_access_groups,
all_models=list(all_models_set),
include_model_access_groups=include_model_access_groups,
proxy_model_list=proxy_model_list,
)
# deduplicate while preserving order

61
litellm/proxy/auth/user_api_key_auth.py
View file

@ -34,6 +34,7 @@ from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import (
ExperimentalUIJWTToken,
TeamNotFoundError,
_cache_key_object,
_can_object_call_model,
_check_end_user_budget,
@ -85,6 +86,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
)
from litellm.proxy.common_utils.realtime_utils import _realtime_request_body
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.utils import (
PrismaClient,
@ -1060,6 +1062,31 @@ async def _read_request_body_deferring_parse_failure(
return populate_request_with_path_params(request_data=parsed_body, request=request), None
async def _record_unparsable_body_failure(
user_api_key_dict: UserAPIKeyAuth,
body_parse_exception: ProxyException,
route: str,
) -> None:
"""Record the 400 an unparsable body earns as a failed request log.
The endpoint never runs for these, so no downstream failure hook writes the
spend log row the Admin UI reads. Logging must not change what the caller
sees, so a failure here is swallowed and the 400 is raised either way.
"""
from litellm.proxy.proxy_server import proxy_logging_obj
try:
await proxy_logging_obj.post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # bare dict in sig
request_data={}, # mutable-ok: the failure hook seeds the call id and metadata onto this dict
original_exception=body_parse_exception,
user_api_key_dict=user_api_key_dict,
error_type=ProxyErrorTypes.bad_request_error,
route=route,
)
except Exception as e: # noqa: BLE001 # any logging failure must leave the caller's 400 untouched
verbose_proxy_logger.exception("Failed to log the request rejected for an unparsable body: %s", e)
async def _user_api_key_auth_builder(
request: Request,
api_key: str,
@ -2136,6 +2163,28 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached
)
def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseException) -> bool:
"""Whether the token's own team fields may stand in for a team that failed to
resolve, without widening access.
A team that is provably gone is a definitive answer, not a degraded read, so
nothing may stand in for it and no setting may override that.
Otherwise the team's grant is merely unknown. A token carrying one may vouch,
since replaying a recorded grant cannot widen it and denying every team key
while the row is briefly unreadable would trade the widening for an outage. A
token carrying none may not: ``team_models=[]`` reads as every model and
``team_blocked=False`` as unblocked. ``allow_requests_on_db_unavailable`` opts
back out, and is only consulted here because the failure is known by this
point to be a degraded read.
"""
if isinstance(lookup_error, TeamNotFoundError):
return False
if valid_token.team_models:
return True
return PrismaDBExceptionHandler.should_allow_request_on_db_unavailable()
@tracer.wrap()
async def _run_centralized_common_checks(
user_api_key_auth_obj: UserAPIKeyAuth,
@ -2339,7 +2388,12 @@ async def _run_centralized_common_checks(
if isinstance(team_result, BaseException):
# Token-derived fallback only valid when a team_id is set;
# _team_obj_from_token asserts that precondition.
team_object = _team_obj_from_token(user_api_key_auth_obj) if user_api_key_auth_obj.team_id is not None else None
if user_api_key_auth_obj.team_id is None:
team_object = None
elif _token_can_vouch_for_team(user_api_key_auth_obj, team_result):
team_object = _team_obj_from_token(user_api_key_auth_obj)
else:
raise team_result
else:
team_object = team_result
@ -2673,6 +2727,11 @@ async def user_api_key_auth(
user_api_key_auth_obj.request_route = normalize_request_route(route)
if body_parse_exception is not None:
await _record_unparsable_body_failure(
user_api_key_dict=user_api_key_auth_obj,
body_parse_exception=body_parse_exception,
route=route,
)
raise body_parse_exception
# Resolve caller identity once, here at the seam, into a single per-request

5
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -715,7 +715,7 @@ async def list_batches(
operation_context="batch listing",
)
data.update(credentials)
prepare_data_with_credentials(data=data, credentials=credentials)
response = await litellm.alist_batches(
custom_llm_provider=credentials["custom_llm_provider"],
@ -948,9 +948,10 @@ async def cancel_batch(
# SCENARIO 3: Fallback to custom_llm_provider (uses env variables)
else:
body_custom_llm_provider = data.pop("custom_llm_provider", None)
custom_llm_provider: Final = (
provider
or data.pop("custom_llm_provider", None)
or body_custom_llm_provider
or get_custom_llm_provider_from_request_headers(request=request)
or get_custom_llm_provider_from_request_query(request=request)
or "openai"

11
litellm/proxy/client/cli/README.md
View file

@ -36,6 +36,17 @@ The base URL is resolved in this order of precedence:
3. `base_url` from `~/.litellm/config.json`
4. `http://localhost:4000`
### Hiding commands from the listings
Deployments that hand `lite` to end users often want to advertise only part of it. Store the commands to keep out of the listings, comma separated:
```bash
lite config set hidden_commands codex,opencode
lite config unset hidden_commands # list everything again
```
Hidden commands drop out of both `lite --help` and the interactive shell's "Available commands" block, and stay runnable so existing scripts keep working
## Global Options
- `--version`, `-v`: Print the LiteLLM Proxy client and server version and exit.

103
litellm/proxy/client/cli/commands/agents.py
View file

@ -1,5 +1,6 @@
import os
import shutil
import subprocess
import sys
from collections.abc import Callable, Mapping, Sequence
from typing import Final
@ -142,8 +143,95 @@ def verify_proxy_key(
)
def _exec(path: str, args: Sequence[str], env: Mapping[str, str]) -> None:
os.execvpe(path, list(args), dict(env))
_WINDOWS_SHIM_SUFFIXES: Final[frozenset[str]] = frozenset({".cmd", ".bat"})
_CMD_PERCENT_GUARD: Final = "%%cd:~,%"
_CMD_LINE_BREAKS: Final = ("\r", "\n")
def _double_trailing_backslashes(segment: str) -> str:
bare: Final = segment.rstrip("\\")
return bare + "\\" * 2 * (len(segment) - len(bare))
def _quote_for_cmd(token: str) -> str:
"""Quote one token so both parsers that read it see the original text.
Follows the algorithm the Rust standard library settled on for batch files
after CVE-2024-24576. Two parsers see this token: cmd.exe, which ends a
quoted string on a lone `"` and so wants an embedded one doubled, and the
shim's own interpreter, which re-splits `%*` under C runtime rules where a
backslash escapes the quote that follows it, so every backslash run standing
before a quote is doubled. Quoting cannot stop cmd expanding `%VAR%`, so each
`%` is prefixed with `%%cd:~,`: the zero-length substring of the always
defined `cd` expands to nothing and leaves no `%` pair for cmd to match.
"""
escaped: Final = '""'.join(_double_trailing_backslashes(part) for part in token.split('"'))
return '"' + escaped.replace("%", _CMD_PERCENT_GUARD) + '"'
def _windows_command(path: str, args: Sequence[str]) -> str | tuple[str, ...]:
"""Build what CreateProcess runs, routing batch shims through cmd.exe.
npm installs Claude Code as `claude.cmd`, which PATHEXT lets shutil.which
resolve but CreateProcess refuses to run (WinError 193), so a shim has to go
through the command processor. cmd.exe does not follow the C runtime quoting
that subprocess would apply to an argument list, and it would split on `&` or
`|` in a forwarded argument, so the shim case is emitted as one verbatim
command line with every token quoted. Every switch is load-bearing: `/s`
makes cmd strip only the outer pair, leaving each token quoted and its
metacharacters inert, `/e:on` keeps the command extensions that the percent
guard is built out of, `/v:off` keeps `!` from expanding, and `/d` keeps a
machine's AutoRun commands out of the launch. argv[0] carries the
caller-facing name on POSIX; Windows needs the resolved path there.
Raises AgentRunError for an argument holding a line break, which cmd would
read as the end of the command line and silently drop the rest of.
"""
rest: Final = tuple(args[1:])
if os.path.splitext(path)[1].lower() not in _WINDOWS_SHIM_SUFFIXES:
return (path, *rest)
if any(brk in token for token in rest for brk in _CMD_LINE_BREAKS):
raise AgentRunError(
f"Cannot pass an argument containing a line break to `{os.path.basename(path)}` on "
"Windows: cmd.exe ends the command line there, so the agent would silently lose it."
)
inner: Final = " ".join(_quote_for_cmd(token) for token in (path, *rest))
return f'cmd.exe /d /e:on /v:off /s /c "{inner}"'
def _spawn_and_wait(command: str | Sequence[str], env: Mapping[str, str]) -> int:
return subprocess.run(command, env=dict(env), check=False).returncode
def _replace_process(
path: str,
args: Sequence[str],
env: Mapping[str, str],
*,
execvpe: Callable[..., None] = os.execvpe,
) -> None:
execvpe(path, list(args), dict(env))
def _hand_off(
path: str,
args: Sequence[str],
env: Mapping[str, str],
*,
platform: str = sys.platform,
replace: Callable[[str, Sequence[str], Mapping[str, str]], None] = _replace_process,
spawn: Callable[[str | Sequence[str], Mapping[str, str]], int] = _spawn_and_wait,
) -> None:
"""Replace this process with the agent; on Windows, run it as a child instead.
os.exec* has no process-replacement semantics on Windows: the C runtime
spawns a detached child and terminates the parent, so the shell reclaims the
console and the agent's TUI never gets one. Windows therefore waits on the
child and exits with its status.
"""
if platform.startswith("win"):
raise SystemExit(spawn(_windows_command(path, args), env))
replace(path, list(args), dict(env))
def _restore_controlling_terminal() -> None:
@ -175,13 +263,14 @@ def run_agent(
base_env: Mapping[str, str] | None = None,
which: Callable[[str], str | None] = shutil.which,
verify: Callable[[str, str], None] = verify_proxy_key,
launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _exec,
launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _hand_off,
reattach_terminal: Callable[[], None] | None = None,
) -> None:
"""Validate, wire the environment, and hand off to the agent.
On success this replaces the current process and never returns. Raises
AgentRunError for missing binaries, an unreachable proxy, or a rejected key.
On success this never returns: POSIX replaces the current process, Windows
waits on the agent and exits with its status. Raises AgentRunError for
missing binaries, an unreachable proxy, or a rejected key.
reattach_terminal, when given, runs just before handoff to restore stdin.
"""
if not command:
@ -277,9 +366,9 @@ def _make_agent_command(binary: str, display_name: str) -> click.Command:
return _command
def agent_commands() -> list[click.Command]:
def agent_commands() -> tuple[click.Command, ...]:
"""Build one top-level command per known agent, e.g. `lite claude`."""
return [_make_agent_command(binary, name) for binary, (name, _profiles) in _KNOWN_AGENTS.items()]
return tuple(_make_agent_command(binary, name) for binary, (name, _profiles) in _KNOWN_AGENTS.items())
__all__ = [

59
litellm/proxy/client/cli/commands/config.py
View file

@ -1,8 +1,9 @@
import json
import os
import sys
from collections.abc import Mapping
from collections.abc import Callable, Mapping
from pathlib import Path
from types import MappingProxyType
from typing import Final
from urllib.parse import urlparse
@ -11,7 +12,7 @@ from pydantic import TypeAdapter
from .private_json import write_private_json
ALLOWED_CONFIG_KEYS: Final[tuple[str, ...]] = ("base_url",)
HIDDEN_COMMANDS_KEY: Final = "hidden_commands"
_config_adapter: Final[TypeAdapter[Mapping[str, str]]] = TypeAdapter(Mapping[str, str])
@ -49,6 +50,48 @@ def get_config_value(key: str) -> str | None:
return load_config().get(key)
def parse_hidden_commands(raw: str | None) -> frozenset[str]:
"""Split a stored `hidden_commands` value, e.g. "codex, opencode"."""
return frozenset(name.strip() for name in (raw or "").split(",") if name.strip())
def hidden_command_names() -> frozenset[str]:
"""Top-level commands the operator chose to keep out of `lite`'s listings."""
return parse_hidden_commands(get_config_value(HIDDEN_COMMANDS_KEY))
def _normalize_base_url(value: str) -> str:
parsed: Final = urlparse(value)
if parsed.scheme not in ("http", "https") or not parsed.netloc:
raise click.UsageError("base_url must be a full http:// or https:// URL including a host")
if "?" in value or "#" in value:
raise click.UsageError("base_url must not include a query string or fragment")
return value.rstrip("/")
def _normalize_hidden_commands(value: str) -> str:
names: Final = parse_hidden_commands(value)
if not names:
raise click.UsageError(
f"{HIDDEN_COMMANDS_KEY} must be a comma-separated list of command names, e.g. "
f"`lite config set {HIDDEN_COMMANDS_KEY} codex,opencode`. To list everything again, "
f"run `lite config unset {HIDDEN_COMMANDS_KEY}`"
)
if any(" " in name for name in names):
raise click.UsageError(f"{HIDDEN_COMMANDS_KEY} entries must be single command names, without spaces")
return ",".join(sorted(names))
_NORMALIZERS: Final[Mapping[str, Callable[[str], str]]] = MappingProxyType(
{
"base_url": _normalize_base_url,
HIDDEN_COMMANDS_KEY: _normalize_hidden_commands,
}
)
ALLOWED_CONFIG_KEYS: Final[tuple[str, ...]] = tuple(_NORMALIZERS)
@click.group(name="config")
def config_commands() -> None:
"""Manage persistent CLI configuration (~/.litellm/config.json)"""
@ -59,17 +102,11 @@ def config_commands() -> None:
@click.argument("value")
def set_config(key: str, value: str) -> None:
"""Set a config KEY to VALUE (e.g. `lite config set base_url https://your-proxy.example.com`)"""
if key not in ALLOWED_CONFIG_KEYS:
normalizer: Final = _NORMALIZERS.get(key)
if normalizer is None:
raise click.UsageError(f"Unknown config key '{key}'. Allowed keys: {', '.join(ALLOWED_CONFIG_KEYS)}")
if key == "base_url":
parsed: Final = urlparse(value)
if parsed.scheme not in ("http", "https") or not parsed.netloc:
raise click.UsageError("base_url must be a full http:// or https:// URL including a host")
if "?" in value or "#" in value:
raise click.UsageError("base_url must not include a query string or fragment")
normalized_value: Final = value.rstrip("/")
normalized_value: Final = normalizer(value)
save_config({**load_config(), key: normalized_value})
click.echo(f"Set {key} = {normalized_value} in {get_config_file_path()}")

8
litellm/proxy/client/cli/interface.py
View file

@ -74,8 +74,9 @@ def styled_prompt():
def show_commands():
"""Display available commands."""
"""Display available commands, minus any the operator chose to hide."""
from .commands.agents import agent_commands
from .commands.config import hidden_command_names
commands = [
("login", "Authenticate with the LiteLLM proxy server"),
@ -96,9 +97,12 @@ def show_commands():
("quit", "Exit the interactive session"),
]
hidden: Final = hidden_command_names()
click.echo("Available commands:")
for cmd, description in commands:
click.echo(f" {cmd:<20} {description}")
if cmd not in hidden:
click.echo(f" {cmd:<20} {description}")
click.echo()

18
litellm/proxy/client/cli/main.py
View file

@ -12,7 +12,7 @@ from .commands.agents import agent_commands
from .commands.auth import auth_group, get_stored_api_key, login, logout, whoami
from .commands.autoroute.commands import autoroute_group
from .commands.chat import chat
from .commands.config import config_commands, get_config_value
from .commands.config import config_commands, get_config_value, hidden_command_names
from .commands.credentials import credentials
from .commands.encryption import encryption
from .commands.http import http
@ -43,7 +43,21 @@ def print_version(base_url: str, api_key: str | None):
click.echo(f"Could not retrieve server version: {e}")
@click.group(invoke_without_command=True)
class HideConfiguredCommandsGroup(click.Group):
"""Group that omits operator-hidden commands from listings, still running them.
Deployments hand `lite` to users who should only see a curated subset of
commands (`lite config set hidden_commands codex,opencode`). Filtering the
listing rather than dropping the commands keeps anyone's existing scripts
working.
"""
def list_commands(self, ctx: click.Context) -> list[str]:
hidden: Final = hidden_command_names()
return [name for name in super().list_commands(ctx) if name not in hidden]
@click.group(cls=HideConfiguredCommandsGroup, invoke_without_command=True)
@click.option(
"--version",
"-v",

90
litellm/proxy/common_utils/callback_utils.py
View file

@ -1,12 +1,19 @@
import copy
import os
from collections.abc import Callable, Iterable
from typing import TYPE_CHECKING, Any, Final, Optional
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias
from typing_extensions import assert_never
import litellm
from litellm import get_secret
from litellm._logging import verbose_proxy_logger
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
from litellm.constants import (
CONSUMED_REQUEST_TAGS_METADATA_KEY,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
@ -46,6 +53,66 @@ if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
@dataclass(frozen=True, slots=True)
class _CallbackResolvedToClass:
entry: str
loaded: type
tag: Literal["resolved_to_class"] = "resolved_to_class"
@dataclass(frozen=True, slots=True)
class _CallbackNotDispatchable:
entry: str
loaded: object
tag: Literal["not_dispatchable"] = "not_dispatchable"
_CallbackLoadError: TypeAlias = _CallbackResolvedToClass | _CallbackNotDispatchable
def _classify_loaded_callback(entry: str, loaded: object) -> CustomLogger | Callable[..., object] | _CallbackLoadError:
"""
Decide whether what a ``litellm_settings.callbacks`` dotted path resolved to can be dispatched.
A dotted path only ever runs as a ``CustomLogger`` instance or as a callback function. Anything
else (most commonly a class instead of an instance) used to load without complaint and then be
skipped on every request, with no log line and no error.
"""
if isinstance(loaded, CustomLogger) or (callable(loaded) and not isinstance(loaded, type)):
return loaded
if isinstance(loaded, type):
return _CallbackResolvedToClass(entry=entry, loaded=loaded)
return _CallbackNotDispatchable(entry=entry, loaded=loaded)
def _raise_callback_load_error(error: _CallbackLoadError) -> NoReturn:
"""The one edge that raises: map a load error onto config load's failure contract."""
match error:
case _CallbackResolvedToClass():
module_path: Final = error.entry.rsplit(".", 1)[0] if "." in error.entry else error.entry
raise ValueError(
f"litellm_settings.callbacks entry '{error.entry}' resolved to the class "
f"{error.loaded.__module__}.{error.loaded.__qualname__}, which is neither a "
"CustomLogger instance nor a callable, so the proxy would never run it."
f" Point it at an instance instead, e.g. add `proxy_handler_instance = {error.loaded.__name__}()` to "
f'{module_path} and set `callbacks: ["{module_path}.proxy_handler_instance"]`.'
)
case _CallbackNotDispatchable():
raise ValueError(
f"litellm_settings.callbacks entry '{error.entry}' resolved to "
f"{type(error.loaded).__name__} {error.loaded!r}, which is neither a "
"CustomLogger instance nor a callable, so the proxy would never run it."
)
assert_never(error)
def _loaded_callback_or_raise(entry: str, loaded: object) -> CustomLogger | Callable[..., object]:
resolved: Final = _classify_loaded_callback(entry=entry, loaded=loaded)
if isinstance(resolved, _CallbackResolvedToClass | _CallbackNotDispatchable):
_raise_callback_load_error(resolved)
return resolved
def initialize_callbacks_on_proxy(
value: Any,
premium_user: bool,
@ -301,9 +368,12 @@ def initialize_callbacks_on_proxy(
"%s attempting to import custom calback=%s %s", blue_color_code, callback, reset_color_code
)
imported_list.append(
get_instance_fn(
value=callback,
config_file_path=config_file_path,
_loaded_callback_or_raise(
entry=callback,
loaded=get_instance_fn(
value=callback,
config_file_path=config_file_path,
),
)
)
if isinstance(litellm.callbacks, list):
@ -317,9 +387,12 @@ def initialize_callbacks_on_proxy(
PrometheusLogger._mount_metrics_endpoint()
else:
litellm.callbacks = [
get_instance_fn(
value=value,
config_file_path=config_file_path,
_loaded_callback_or_raise(
entry=value,
loaded=get_instance_fn(
value=value,
config_file_path=config_file_path,
),
)
]
verbose_proxy_logger.debug("%s Initialized Callbacks - %s %s", blue_color_code, litellm.callbacks, reset_color_code)
@ -426,6 +499,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset(
"_pipeline_managed_guardrails",
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
"disable_global_guardrails",
"disable_global_guardrail",
"opted_out_global_guardrails",

347
litellm/proxy/common_utils/scheduled_job_stagger.py Normal file
View file

@ -0,0 +1,347 @@
"""
Deterministic phase offsets for the proxy's scheduled background jobs.
APScheduler anchors an ``interval`` job at ``now + interval``, so every job registered in
the same startup shares one firing instant for the life of the process, and every replica
brought up by the same rollout shares it too. The result is a burst: each tick, every job
on every replica queries Postgres at the same moment, competing with the request path for
the connection pool. The product's own daily/monthly crons are worse still, since they name
a wall-clock instant that is identical on every replica by construction.
The fix is a phase offset derived from ``sha256(job_id, identity)``, where ``identity``
covers the pod and the worker process. Different jobs get different offsets, different
replicas get different offsets for the same job, and nothing collapses back onto a shared
instant after a restart. Hashing rather than randomising keeps a given process's schedule
stable for its whole life and lets the applied offsets be logged once and reasoned about
later.
The offset lives in the trigger rather than in a one-off ``next_run_time`` because a cron
trigger recomputes each fire from the wall clock and would otherwise snap straight back
onto the shared instant after its first shifted run.
Only schedules LiteLLM itself chose are shifted. Interval jobs are always eligible; cron
jobs only when their id is one of the product's own defaults, so an operator-supplied
crontab keeps the exact instant it asks for. A job whose call site passed an explicit
``next_run_time`` already anchors itself and is left alone.
"""
# apscheduler ships no type information, so its imports have no stubs. The Protocols below
# narrow everything it hands back, which is why this is the only diagnostic left to silence.
# pyright: reportMissingTypeStubs=false
import hashlib
import os
import socket
from collections.abc import Callable, Mapping, Sequence
from datetime import datetime, timedelta
from types import MappingProxyType
from typing import Final, Protocol
from apscheduler.events import EVENT_JOB_SUBMITTED
from apscheduler.triggers.base import BaseTrigger
from apscheduler.triggers.interval import IntervalTrigger
from pydantic import ValidationError
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import (
MONTHLY_SPEND_REPORT_JOB_ID,
PROMETHEUS_FALLBACK_STATS_JOB_ID,
PTU_ROLLUP_JOB_ID,
PTU_ROLLUP_LOCK_TTL_SECONDS,
)
from litellm.proxy._types import ScheduledJobStaggerSettings
GENERAL_SETTINGS_KEY: Final = "scheduled_job_stagger"
#: Cron schedules LiteLLM picks on the operator's behalf, so shifting them changes nothing the
#: operator asked for. Every other cron trigger is an operator-supplied crontab, preserved exactly.
#:
#: The value is the span over which a second firing would redo work the first already did, which
#: is how long each job's leader-election lock stays held. Two replicas further apart than that
#: both find the key free and both run, which for the spend report means the customer gets it
#: twice. Offsets for these jobs are bounded by it, so widening the window cannot resurrect the
#: duplicate-work failure this feature exists to avoid.
DEFAULT_CRON_DEDUPE_SECONDS: Final = MappingProxyType(
{
MONTHLY_SPEND_REPORT_JOB_ID: 3600,
PROMETHEUS_FALLBACK_STATS_JOB_ID: 3600,
PTU_ROLLUP_JOB_ID: PTU_ROLLUP_LOCK_TTL_SECONDS,
}
)
class Trigger(Protocol):
"""The one method APScheduler asks a trigger for"""
def get_next_fire_time(self, previous_fire_time: datetime | None, now: datetime) -> datetime | None: ...
class ScheduledJob(Protocol):
@property
def id(self) -> str: ...
@property
def trigger(self) -> Trigger: ...
class JobScheduler(Protocol):
"""The slice of ``AsyncIOScheduler`` this module uses, which ships no type information"""
@property
def running(self) -> bool: ...
def get_jobs(self) -> Sequence[ScheduledJob]: ...
def modify_job(self, job_id: str, *, trigger: Trigger) -> object: ...
def add_listener(self, callback: Callable[["JobSubmission"], None], mask: int = ...) -> None: ...
class JobSubmission(Protocol):
"""An ``EVENT_JOB_SUBMITTED`` event"""
@property
def job_id(self) -> str: ...
@property
def scheduled_run_times(self) -> Sequence[datetime]: ...
class _OffsetTrigger:
"""
Delegates to ``base`` on a clock rolled back by ``offset``, then rolls the answer
forward again, so every fire lands exactly ``offset`` later than it otherwise would
while the underlying schedule keeps its own semantics.
Composed rather than derived from ``BaseTrigger``: APScheduler only ever asks a trigger
for its next fire time, and it accepts this by virtual registration below.
"""
__slots__ = ("base", "offset")
def __init__(self, base: Trigger, offset: timedelta) -> None:
self.base = base
self.offset = offset
def get_next_fire_time(self, previous_fire_time: datetime | None, now: datetime) -> datetime | None:
shifted_previous: Final = None if previous_fire_time is None else previous_fire_time - self.offset
next_fire_time: Final = self.base.get_next_fire_time(shifted_previous, now - self.offset)
return None if next_fire_time is None else next_fire_time + self.offset
def __str__(self) -> str:
return f"{self.base}[+{int(self.offset.total_seconds())}s]"
# APScheduler type-checks assigned triggers with isinstance, so it has to accept this one
BaseTrigger.register(_OffsetTrigger)
def parse_stagger_settings(general_settings: Mapping[str, object]) -> ScheduledJobStaggerSettings:
raw: Final = general_settings.get(GENERAL_SETTINGS_KEY)
if raw is None:
return ScheduledJobStaggerSettings()
try:
return ScheduledJobStaggerSettings.model_validate(raw)
except ValidationError as exc:
verbose_proxy_logger.warning(
"Ignoring invalid general_settings.%s, falling back to defaults: %s",
GENERAL_SETTINGS_KEY,
exc,
)
return ScheduledJobStaggerSettings()
def resolve_stagger_identity(configured: str | None) -> str:
"""
The value hashed alongside a job id to place this process in the stagger window.
The process id is part of it because a pod runs one scheduler per uvicorn worker, and
workers sharing a hostname would otherwise all land on the same offset. That makes the
offsets change across restarts, which is what stops a simultaneous rollout from
reconverging; the applied values are logged so a given run stays explainable.
"""
host: Final = configured or os.getenv("POD_NAME") or os.getenv("HOSTNAME") or _hostname()
return f"{host}:{os.getpid()}"
def _hostname() -> str:
try:
return socket.gethostname()
except OSError:
return str(uuid.uuid4())
def offset_seconds(*, job_id: str, identity: str, window_seconds: int) -> int:
"""A stable point in ``[0, window_seconds)`` for this job on this process"""
if window_seconds <= 0:
return 0
digest: Final = hashlib.sha256(f"{job_id}\x00{identity}".encode()).digest()
return int.from_bytes(digest[:8], "big") % window_seconds
def _interval_seconds(job: ScheduledJob) -> int | None:
if not isinstance(job.trigger, IntervalTrigger):
return None
interval: Final = getattr(job.trigger, "interval", None)
return int(interval.total_seconds()) if isinstance(interval, timedelta) else None
def _is_staggerable(job: ScheduledJob) -> bool:
if hasattr(job, "next_run_time"):
# the call site anchored the first fire itself
return False
if _interval_seconds(job) is not None:
return True
return job.id in DEFAULT_CRON_DEDUPE_SECONDS
def _window_for(*, job_id: str, period_seconds: int | None, settings: ScheduledJobStaggerSettings) -> int:
"""
Exclusive upper bound on this job's offset. An interval job is never offset by more than
one of its own periods, so it is not delayed past the wait it already had, and a
leader-elected cron is never offset past the span in which a second replica would redo
its work.
"""
limits: Final = (settings.window_seconds, period_seconds, DEFAULT_CRON_DEDUPE_SECONDS.get(job_id))
return min(limit for limit in limits if limit is not None)
def _clamped_override(*, job_id: str, requested: int) -> int:
horizon: Final = DEFAULT_CRON_DEDUPE_SECONDS.get(job_id)
if horizon is None or requested < horizon:
return requested
verbose_proxy_logger.warning(
"general_settings.%s.offsets[%s]=%ss would place replicas more than %ss apart, "
"which is long enough for a second replica to redo the run; using %ss instead",
GENERAL_SETTINGS_KEY,
job_id,
requested,
horizon,
horizon - 1,
)
return horizon - 1
def _offset_for(
*,
job_id: str,
period_seconds: int | None,
staggerable: bool,
settings: ScheduledJobStaggerSettings,
identity: str,
) -> int:
override: Final = settings.offsets.get(job_id)
if override is not None:
return _clamped_override(job_id=job_id, requested=max(0, override))
if not staggerable:
return 0
return offset_seconds(
job_id=job_id,
identity=identity,
window_seconds=_window_for(job_id=job_id, period_seconds=period_seconds, settings=settings),
)
def stagger_trigger(
*,
job_id: str,
trigger: Trigger,
period_seconds: int | None,
settings: ScheduledJobStaggerSettings,
identity: str | None = None,
) -> Trigger:
"""
The trigger a job should carry, shifted by its own share of the window.
For a job registered against an already-running scheduler, which the startup sweep cannot
reach: every job carries a ``next_run_time`` by then, so re-running the sweep would treat
them all as self-anchored and change nothing.
"""
offset: Final = _offset_for(
job_id=job_id,
period_seconds=period_seconds,
staggerable=True,
settings=settings,
identity=identity or resolve_stagger_identity(settings.identity),
)
return trigger if offset == 0 else _OffsetTrigger(trigger, timedelta(seconds=offset))
def apply_scheduled_job_stagger(
*,
scheduler: JobScheduler,
settings: ScheduledJobStaggerSettings,
identity: str | None = None,
) -> Mapping[str, int]:
"""
Shift each eligible job's schedule by its own offset. Call this once, after every job is
registered and before the scheduler starts, so the offset is folded into the first fire
rather than applied to a schedule already running.
``identity`` is resolved from the environment when the caller does not supply one.
Returns the offset applied to every registered job, including the zeroes, so the caller
and the logs describe the same thing.
"""
resolved_identity: Final = identity or resolve_stagger_identity(settings.identity)
if scheduler.running:
# every job already carries a next_run_time by now, so the sweep would skip all of
# them and report success while changing nothing
verbose_proxy_logger.warning(
"Scheduled job stagger skipped: the scheduler is already running, so offsets must be "
"applied before it starts"
)
return MappingProxyType({job.id: 0 for job in scheduler.get_jobs()})
if not settings.enabled:
verbose_proxy_logger.info(
"Scheduled job stagger disabled via general_settings.%s; all jobs keep their unshifted schedule",
GENERAL_SETTINGS_KEY,
)
return MappingProxyType({job.id: 0 for job in scheduler.get_jobs()})
offsets: Final = MappingProxyType(
{
job.id: _offset_for(
job_id=job.id,
period_seconds=_interval_seconds(job),
staggerable=_is_staggerable(job),
settings=settings,
identity=resolved_identity,
)
for job in scheduler.get_jobs()
}
)
for job in scheduler.get_jobs():
if offsets[job.id] > 0:
scheduler.modify_job(
job.id,
trigger=_OffsetTrigger(job.trigger, timedelta(seconds=offsets[job.id])),
)
verbose_proxy_logger.info(
"Scheduled job stagger applied (identity=%s, window=%ss): %s",
resolved_identity,
settings.window_seconds,
", ".join(f"{job_id}=+{seconds}s" for job_id, seconds in sorted(offsets.items())),
)
return offsets
def attach_job_timing_logger(scheduler: JobScheduler) -> None:
"""Log each fire's scheduled instant against the instant it actually started"""
scheduler.add_listener(_log_job_submitted, EVENT_JOB_SUBMITTED)
def _log_job_submitted(event: JobSubmission) -> None:
if not event.scheduled_run_times:
return
scheduled: Final = event.scheduled_run_times[0]
started: Final = datetime.now(scheduled.tzinfo)
verbose_proxy_logger.debug(
"Scheduled job %s started: scheduled_run_time=%s actual_start_time=%s delay=%.3fs",
event.job_id,
scheduled.isoformat(),
started.isoformat(),
(started - scheduled).total_seconds(),
)

11
litellm/proxy/common_utils/sse_keepalive.py
View file

@ -21,6 +21,17 @@ def _coerce_interval(ping_interval_seconds: float | str | None) -> float | None:
return interval
def keepalive_ping_has_fired(elapsed_seconds: float, ping_interval_seconds: float | str | None) -> bool:
"""Whether a keepalive ping has already gone out, which flushes the response headers.
A caller that discovers a failure after that point cannot raise its way to the client, since
the status line is already on the wire. With pings disabled nothing flushes early, so a raise
still carries its real status.
"""
interval: Final = _coerce_interval(ping_interval_seconds)
return interval is not None and elapsed_seconds >= interval
def wrap_sse_stream_with_keepalive_pings(
stream: AsyncGenerator[str, None],
ping_interval_seconds: float | str | None,

12
litellm/proxy/db/autorouter_session_rollup.py
View file

@ -24,6 +24,7 @@ from itertools import groupby
from typing import TYPE_CHECKING, Final, NamedTuple
from litellm._logging import verbose_proxy_logger
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES
if TYPE_CHECKING:
@ -180,12 +181,17 @@ def build_autorouter_turn_transaction(
The routing_decision record is what says a request was auto-routed at all, so a
request without one (including the auto-router's own classifier sub-calls) never
reaches the rollup. Failed requests served nothing and are excluded. Cache facts
are derived from the payload's own usage record through the savings owner, never
handed in beside it.
reaches the rollup. Internal sub-calls that DO carry one (a shadow eval's duplicate
of a request through the router) are excluded by their internal_call_origin stamp:
they are not traffic a user sent, so counting them would manufacture sessions and
savings in the adoption metrics. Failed requests served nothing and are excluded.
Cache facts are derived from the payload's own usage record through the savings
owner, never handed in beside it.
"""
if payload.get("status") != "success":
return None
if metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY):
return None
routing_decision: Final = metadata.get("routing_decision")
if not isinstance(routing_decision, Mapping) or not routing_decision:
return None

75
litellm/proxy/db/create_views.py
View file

@ -1,15 +1,50 @@
from typing import Any, Final
from typing import Any, Final, Protocol
from litellm import verbose_logger
_db = Any
class SupportsExecuteRaw(Protocol):
"""The one database operation create_view_tolerating_race needs.
Narrower than the `_db = Any` the rest of this module still uses, so the
helper's contract is checkable at its call sites without retyping every
function here.
"""
async def execute_raw(self, query: str, *args: object) -> int: ...
# Markers that indicate a view/relation does not yet exist in the database.
# Keeping these in one place avoids repeating the check across all view blocks
# and prevents overly broad matches (e.g. bare 'undefined' would also match
# 'undefined function' or 'column undefined_col referenced in query').
_VIEW_NOT_FOUND_MARKERS: Final = ("does not exist", "no such table", "undefined table")
# Markers for the inverse condition: another replica created the view between
# our existence probe and our CREATE.
_VIEW_ALREADY_EXISTS_MARKERS: Final = ("already exists", "duplicate object", "duplicate table")
async def create_view_tolerating_race(db: SupportsExecuteRaw, view_name: str, ddl: str) -> None:
"""
Create a view, treating "a concurrent creator won" as success.
Every replica booting against the same fresh database observes the view as
absent and issues the CREATE; Postgres fails all but one with a
duplicate-object error. The desired end state is still reached, so losing
that race is success. Without this, the loser's exception propagates out of
a detached startup task and the remaining views are never created.
"""
try:
await db.execute_raw(ddl)
verbose_logger.debug("%s Created!", view_name)
except Exception as e:
if not any(marker in str(e).lower() for marker in _VIEW_ALREADY_EXISTS_MARKERS):
raise
verbose_logger.debug("%s already created by a concurrent replica", view_name)
async def create_missing_views(db: _db):
"""
@ -34,7 +69,10 @@ async def create_missing_views(db: _db):
if not any(marker in error_msg for marker in _VIEW_NOT_FOUND_MARKERS):
raise
# If an error occurs, the view does not exist, so create it
await db.execute_raw("""
await create_view_tolerating_race(
db,
"LiteLLM_VerificationTokenView",
"""
CREATE VIEW "LiteLLM_VerificationTokenView" AS
SELECT
v.*,
@ -46,9 +84,8 @@ async def create_missing_views(db: _db):
FROM "LiteLLM_VerificationToken" v
LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id
LEFT JOIN "LiteLLM_ProjectTable" p ON v.project_id = p.project_id;
""")
verbose_logger.debug("LiteLLM_VerificationTokenView Created!")
""",
)
try:
await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpend" LIMIT 1""")
@ -69,9 +106,7 @@ async def create_missing_views(db: _db):
GROUP BY
DATE("startTime");
"""
await db.execute_raw(query=sql_query)
verbose_logger.debug("MonthlyGlobalSpend Created!")
await create_view_tolerating_race(db, "MonthlyGlobalSpend", sql_query)
try:
await db.query_raw("""SELECT 1 FROM "Last30dKeysBySpend" LIMIT 1""")
@ -100,9 +135,7 @@ async def create_missing_views(db: _db):
ORDER BY
total_spend DESC;
"""
await db.execute_raw(query=sql_query)
verbose_logger.debug("Last30dKeysBySpend Created!")
await create_view_tolerating_race(db, "Last30dKeysBySpend", sql_query)
try:
await db.query_raw("""SELECT 1 FROM "Last30dModelsBySpend" LIMIT 1""")
@ -126,9 +159,7 @@ async def create_missing_views(db: _db):
ORDER BY
total_spend DESC;
"""
await db.execute_raw(query=sql_query)
verbose_logger.debug("Last30dModelsBySpend Created!")
await create_view_tolerating_race(db, "Last30dModelsBySpend", sql_query)
try:
await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpendPerKey" LIMIT 1""")
verbose_logger.debug("MonthlyGlobalSpendPerKey Exists!")
@ -150,9 +181,7 @@ async def create_missing_views(db: _db):
DATE("startTime"),
api_key;
"""
await db.execute_raw(query=sql_query)
verbose_logger.debug("MonthlyGlobalSpendPerKey Created!")
await create_view_tolerating_race(db, "MonthlyGlobalSpendPerKey", sql_query)
try:
await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpendPerUserPerKey" LIMIT 1""")
verbose_logger.debug("MonthlyGlobalSpendPerUserPerKey Exists!")
@ -176,9 +205,7 @@ async def create_missing_views(db: _db):
"user",
api_key;
"""
await db.execute_raw(query=sql_query)
verbose_logger.debug("MonthlyGlobalSpendPerUserPerKey Created!")
await create_view_tolerating_race(db, "MonthlyGlobalSpendPerUserPerKey", sql_query)
try:
await db.query_raw("""SELECT 1 FROM "DailyTagSpend" LIMIT 1""")
@ -197,9 +224,7 @@ async def create_missing_views(db: _db):
FROM "LiteLLM_SpendLogs" s
GROUP BY individual_request_tag, DATE(s."startTime");
"""
await db.execute_raw(query=sql_query)
verbose_logger.debug("DailyTagSpend Created!")
await create_view_tolerating_race(db, "DailyTagSpend", sql_query)
try:
await db.query_raw("""SELECT 1 FROM "Last30dTopEndUsersSpend" LIMIT 1""")
@ -218,9 +243,7 @@ async def create_missing_views(db: _db):
ORDER BY total_spend DESC
LIMIT 100;
"""
await db.execute_raw(query=sql_query)
verbose_logger.debug("Last30dTopEndUsersSpend Created!")
await create_view_tolerating_race(db, "Last30dTopEndUsersSpend", sql_query)
async def should_create_missing_views(db: _db) -> bool:

15
litellm/proxy/db/db_spend_update_writer.py
View file

@ -21,6 +21,7 @@ from litellm.caching import RedisCache
from litellm.constants import (
DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME,
DB_SPEND_UPDATE_JOB_NAME,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
)
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.proxy._types import (
@ -1794,6 +1795,7 @@ class DBSpendUpdateWriter:
if call_type:
endpoint = ROUTE_ENDPOINT_MAPPING.get(call_type, None)
is_internal_call: Final = bool(_metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY))
cache_read_input_tokens: Final = extract_cache_read_tokens(usage_obj)
compression_saved_tokens: Final = extract_compression_saved_tokens(_metadata)
savings_spend: Final = compute_savings_spend(
@ -1818,15 +1820,20 @@ class DBSpendUpdateWriter:
prompt_tokens=payload["prompt_tokens"],
completion_tokens=payload["completion_tokens"],
spend=payload["spend"],
api_requests=1,
successful_requests=1 if request_status == "success" else 0,
failed_requests=1 if request_status != "success" else 0,
# Internal sub-calls (auto-router classifier, shadow eval's shadow and
# judge) bill real spend and tokens to the key, but they are not
# requests the caller made: counting them inflates request-volume
# readers, and an auto-router savings figure computed on a shadow
# duplicate credits savings for traffic no user sent.
api_requests=0 if is_internal_call else 1,
successful_requests=1 if not is_internal_call and request_status == "success" else 0,
failed_requests=1 if not is_internal_call and request_status != "success" else 0,
cache_read_input_tokens=cache_read_input_tokens,
cache_creation_input_tokens=extract_cache_creation_tokens(usage_obj),
compression_saved_tokens=compression_saved_tokens,
compression_savings_spend=savings_spend.compression,
prompt_caching_savings_spend=savings_spend.prompt_caching,
autorouter_savings_spend=savings_spend.autorouter,
autorouter_savings_spend=0.0 if is_internal_call else savings_spend.autorouter,
)
return daily_transaction
except Exception as e:

453
litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
View file

@ -1,22 +1,60 @@
import asyncio
import time
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Final
from typing import Final, Literal, TypeAlias
from pydantic import BaseModel, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.caching import RedisCache
from litellm.constants import (
SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS,
SPEND_LOG_CLEANUP_BATCH_SIZE,
SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS,
SPEND_LOG_CLEANUP_JOB_NAME,
SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES,
SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP,
SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS,
SPEND_LOG_RUN_LOOPS,
)
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import (
RunOutcome,
SpendLogCleanupMetrics,
)
from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import (
RemainingTimeoutMs,
SpendLogsPartitionManager,
)
from litellm.proxy.utils import PrismaClient
StopReason: TypeAlias = Literal["exhausted", "budget_exhausted", "batch_cap_reached", "aborted"]
@dataclass(frozen=True, slots=True)
class TableCleanupResult:
"""Outcome of pruning one table, so the caller can report why a run ended."""
rows_deleted: int
stop_reason: StopReason
class _RemainingRow(BaseModel):
"""One row of the capped outstanding-rows probe, validated out of prisma's untyped result."""
remaining: int
_REMAINING_ROWS: Final = TypeAdapter(list[_RemainingRow])
SPEND_LOG_CLEANUP_BOUND_SETTINGS: Final = (
"maximum_spend_logs_cleanup_batch_size",
"maximum_spend_logs_cleanup_max_batches",
"maximum_spend_logs_cleanup_run_budget",
"maximum_spend_logs_cleanup_batch_timeout",
)
class SpendLogCleanup:
"""
@ -26,6 +64,24 @@ class SpendLogCleanup:
dropping whole partitions (instant, frees disk immediately). Otherwise it
falls back to deleting logs in batches.
Uses PodLockManager to ensure only one pod runs cleanup in multi-pod deployments.
Every run is bounded so it can never monopolise the database: a wall-clock
budget shared across all tables, a per-table batch cap, and a Postgres
statement/lock timeout on every statement the job issues, deletes and the
outstanding-rows probe alike. A run that hits a bound stops cleanly and the
next run resumes from where it left off, because the cutoff is recomputed
and deleted rows are gone.
The budget is a hard wall clock, not an advisory one. Every statement this
job issues, deletes, the outstanding-rows probe and partition DDL alike, is
issued with a timeout clamped to the budget that is still left, so one
started just under the deadline is cancelled by Postgres at the deadline
rather than running a further batch timeout past it. No statement is issued
at all once the budget is spent, which is why the probe is skipped on that
path. Partition DDL additionally carries a lock_timeout, because it takes an
ACCESS EXCLUSIVE lock and would otherwise queue behind a long-running reader
for as long as that reader lives; a partition this run cannot get is left
for the next one.
"""
def __init__(
@ -34,17 +90,88 @@ class SpendLogCleanup:
redis_cache: RedisCache | None = None,
partition_manager: SpendLogsPartitionManager | None = None,
):
self.batch_size = SPEND_LOG_CLEANUP_BATCH_SIZE
self.retention_seconds: int | None = None
self.partition_manager = partition_manager or SpendLogsPartitionManager()
from litellm.proxy.proxy_server import general_settings as default_settings
self.general_settings = general_settings or default_settings
self._refresh_bounds()
from litellm.proxy.proxy_server import proxy_logging_obj
pod_lock_manager: Final = proxy_logging_obj.db_spend_update_writer.pod_lock_manager
self.pod_lock_manager = pod_lock_manager
verbose_proxy_logger.info("SpendLogCleanup initialized with batch size: %s", self.batch_size)
verbose_proxy_logger.info(
"SpendLogCleanup initialized: batch_size=%s max_batches=%s run_budget=%ss batch_timeout=%ss",
self.batch_size,
self.max_batches,
self.run_budget_seconds,
self.batch_timeout_seconds,
)
def _refresh_bounds(self) -> None:
"""
Re-read every bound in SPEND_LOG_CLEANUP_BOUND_SETTINGS from settings.
The scheduler holds one long-lived instance, so a bound captured at
construction would never reflect a dashboard change. general_settings is
the same dict the periodic config reload mutates in place, so reading it
per run is what makes these knobs live. Every bound falls back to its
shipped default, so clearing a field restores that default.
"""
self.batch_size: int = self._positive_int_setting(
"maximum_spend_logs_cleanup_batch_size", SPEND_LOG_CLEANUP_BATCH_SIZE
)
self.max_batches: int = self._positive_int_setting(
"maximum_spend_logs_cleanup_max_batches", SPEND_LOG_RUN_LOOPS
)
self.run_budget_seconds: float = self._duration_setting(
"maximum_spend_logs_cleanup_run_budget", SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS
)
self.batch_timeout_seconds: float = self._duration_setting(
"maximum_spend_logs_cleanup_batch_timeout", SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS
)
def _positive_int_setting(self, setting_name: str, default: int) -> int:
"""
Read a positive-integer knob, falling back to the default when unset or unusable.
"""
raw: Final = self.general_settings.get(setting_name)
if raw is None:
return default
try:
parsed: Final = int(raw)
except (TypeError, ValueError):
verbose_proxy_logger.warning("Invalid %s value: %s, using default %s", setting_name, raw, default)
return default
if parsed <= 0:
verbose_proxy_logger.warning("%s must be positive, got %s, using default %s", setting_name, parsed, default)
return default
return parsed
def _duration_setting(self, setting_name: str, default_seconds: float) -> float:
"""
Read a duration knob (e.g. '5m'), falling back to the default when unset or unusable.
The knob must never be able to remove the bound it exists to enforce, so
anything the parser rejects (including the non-finite spellings 'inf' and
'nan') and anything non-positive falls back rather than being honoured.
"""
raw: Final = self.general_settings.get(setting_name)
if raw is None:
return default_seconds
try:
parsed: Final = float(duration_in_seconds(str(raw)))
except (ValueError, TypeError) as e:
verbose_proxy_logger.warning(
"Invalid %s value: %s (%s), using default %ss", setting_name, raw, e, default_seconds
)
return default_seconds
if parsed <= 0:
verbose_proxy_logger.warning(
"%s must be a positive duration, got %s, using default %ss", setting_name, raw, default_seconds
)
return default_seconds
return parsed
def _retention_seconds_for(self, setting_name: str) -> int | None:
"""
@ -78,6 +205,91 @@ class SpendLogCleanup:
self.retention_seconds = self._retention_seconds_for("maximum_spend_logs_retention_period")
return self.retention_seconds is not None
def _timeout_ms(self, deadline: float) -> int:
"""
The per-statement bound in milliseconds: the batch timeout, or whatever
is left of the run budget, whichever is smaller.
Clamping to the remaining budget is what makes the budget a real
wall-clock bound rather than an advisory one. Postgres offers no "stop
at time T", only a per-statement duration, so a statement issued just
under the deadline would otherwise run a full batch timeout past it, and
with several tables those overruns stack.
Interpolating this into SQL is safe by construction: an int cannot carry
SQL, and SET does not accept a bind parameter.
"""
remaining_ms: Final = int((deadline - time.monotonic()) * 1000)
return max(1, min(int(self.batch_timeout_seconds * 1000), remaining_ms))
def _remaining_timeout_ms(self, deadline: float) -> RemainingTimeoutMs:
"""
The per-statement bound for work this job delegates, as a callable.
Partition maintenance issues one statement per partition, so handing it a
number would bound each statement by the budget that was left before the
FIRST one and never by what remains. Re-evaluating per statement is what
makes the loop itself bounded, and None tells the callee to stop rather
than issue a statement it has no budget for.
"""
def remaining() -> int | None:
return None if time.monotonic() >= deadline else self._timeout_ms(deadline)
return remaining
async def _execute_delete_batch(
self, prisma_client: PrismaClient, delete_sql: str, cutoff_date: datetime, deadline: float
) -> int | None:
"""
Run one delete batch under a Postgres statement and lock timeout.
The timeouts are what actually bound the work: a Prisma transaction
timeout cannot interrupt a statement that is already executing, so
without these a single batch blocked behind a lock would hold its
connection, and the row locks it already took, indefinitely. SET LOCAL
scopes both to this transaction so the pooled connection is unaffected.
Returns the row count, or None when the driver returned something that
is not a row count. That is a contract violation rather than a transient
fault, so the caller stops instead of retrying.
"""
timeout_ms: Final = self._timeout_ms(deadline)
async with prisma_client.db.tx() as tx:
await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}")
deleted_result: Final = await tx.execute_raw(delete_sql, cutoff_date, self.batch_size)
return deleted_result if isinstance(deleted_result, int) else None
async def _count_remaining(
self, prisma_client: PrismaClient, cutoff_date: datetime, table_name: str, time_column: str, deadline: float
) -> int | None:
"""
Count expired rows still outstanding, stopping at a cap.
An uncapped COUNT(*) over an expired backlog would itself be the kind of
long scan this job exists to avoid, so the probe reads at most
SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP index entries. A result equal to
the cap means "at least this many".
"""
count_sql: Final = f"""
SELECT count(*)::int AS remaining FROM (
SELECT 1 FROM "{table_name}"
WHERE "{time_column}" < $1::timestamptz
LIMIT $2
) capped
"""
try:
async with prisma_client.db.tx() as tx:
await tx.execute_raw(f"SET LOCAL statement_timeout = {self._timeout_ms(deadline)}")
rows: Final = _REMAINING_ROWS.validate_python(
await tx.query_raw(count_sql, cutoff_date, SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP)
)
except Exception as e: # noqa: BLE001 - an observability probe must never fail the cleanup run
verbose_proxy_logger.warning("Could not count remaining %s rows: %s", table_name, e)
return None
return rows[0].remaining if rows else None
async def _delete_old_rows_batched(
self,
prisma_client: PrismaClient,
@ -85,10 +297,14 @@ class SpendLogCleanup:
table_name: str,
key_columns: tuple[str, ...],
time_column: str,
) -> int:
deadline: float,
) -> TableCleanupResult:
"""
Helper method to delete a table's rows older than the cutoff in batches.
Returns the total number of rows deleted.
Delete a table's rows older than the cutoff in batches.
Stops at whichever bound is reached first: the backlog running out, the
shared wall-clock deadline, the per-table batch cap, or too many
consecutive batch failures.
"""
key_list: Final = ", ".join(f'"{col}"' for col in key_columns)
delete_sql: Final = f"""
@ -103,23 +319,46 @@ class SpendLogCleanup:
run_count = 0
consecutive_failures = 0
while True:
if run_count > SPEND_LOG_RUN_LOOPS:
if time.monotonic() >= deadline:
verbose_proxy_logger.info(
"Run budget exhausted during %s cleanup after %d rows; the next run resumes from here",
table_name,
total_deleted,
)
return await self._finish_table(
prisma_client, cutoff_date, table_name, time_column, total_deleted, "budget_exhausted", deadline
)
if run_count >= self.max_batches:
verbose_proxy_logger.info(
"Max batches reached for %s cleanup, remaining rows will be deleted in next run", table_name
)
break
# Step 1: Find rows and delete them in one go without fetching to application
# Delete in batches, limited by self.batch_size
try:
deleted_result = await prisma_client.db.execute_raw(
delete_sql,
cutoff_date,
self.batch_size,
return await self._finish_table(
prisma_client, cutoff_date, table_name, time_column, total_deleted, "batch_cap_reached", deadline
)
# Find rows and delete them in one go without fetching to application
batch_started_at = time.monotonic()
try:
batch_result = await self._execute_delete_batch(prisma_client, delete_sql, cutoff_date, deadline)
except Exception as batch_exc:
if time.monotonic() >= deadline:
# The statement timeout was clamped to the budget that was
# left, so this batch was cancelled by the deadline itself.
# That is the bound working, not a database fault, and
# counting it would both inflate the failure metric and push
# every budget-exhausted run toward the abort threshold.
verbose_proxy_logger.info(
"Run budget exhausted mid-batch during %s cleanup after %d rows; "
"the next run resumes from here",
table_name,
total_deleted,
)
return await self._finish_table(
prisma_client, cutoff_date, table_name, time_column, total_deleted, "budget_exhausted", deadline
)
# A single batch failure (e.g. Prisma/DB timeout) must not abort
# the whole run — subsequent batches may still succeed.
consecutive_failures += 1
SpendLogCleanupMetrics.record_batch_failure(table_name)
verbose_proxy_logger.exception(
"%s cleanup batch failed "
"(run_count=%d, consecutive_failures=%d, batch_size=%d, "
@ -140,28 +379,31 @@ class SpendLogCleanup:
consecutive_failures,
total_deleted,
)
break
return await self._finish_table(
prisma_client, cutoff_date, table_name, time_column, total_deleted, "aborted", deadline
)
await asyncio.sleep(SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS)
continue
consecutive_failures = 0
deleted_count = 0
if isinstance(deleted_result, int):
deleted_count = deleted_result
else:
if batch_result is None:
verbose_proxy_logger.error(
"Unexpected execute_raw return type for %s cleanup: %s; aborting cleanup to avoid infinite loop",
"Unexpected execute_raw return type for %s cleanup; aborting cleanup to avoid infinite loop",
table_name,
type(deleted_result),
)
break
return await self._finish_table(
prisma_client, cutoff_date, table_name, time_column, total_deleted, "aborted", deadline
)
consecutive_failures = 0
deleted_count = batch_result
SpendLogCleanupMetrics.record_batch(table_name, deleted_count, time.monotonic() - batch_started_at)
verbose_proxy_logger.info("Deleted %s %s rows in this batch", deleted_count, table_name)
if deleted_count == 0:
verbose_proxy_logger.info("No more %s rows to delete. Total deleted: %s", table_name, total_deleted)
break
return await self._finish_table(
prisma_client, cutoff_date, table_name, time_column, total_deleted, "exhausted", deadline
)
total_deleted += deleted_count
run_count += 1
@ -169,18 +411,49 @@ class SpendLogCleanup:
# Add a small sleep to prevent overwhelming the database
await asyncio.sleep(0.1)
return total_deleted
async def _finish_table(
self,
prisma_client: PrismaClient,
cutoff_date: datetime,
table_name: str,
time_column: str,
rows_deleted: int,
stop_reason: StopReason,
deadline: float,
) -> TableCleanupResult:
"""
Publish how much of this table is still outstanding, then report the run's result.
async def _delete_old_logs(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int:
The probe is skipped once the budget is spent. It is the one piece of
work that would otherwise be ISSUED after the deadline, and every table
exits through here, including the ones a spent run never started, so
keeping it would put one more statement per table past the bound. A run
that ends this way already reports "budget_exhausted", which tells an
operator the backlog was not drained; the gauge simply keeps its value
from the last run that finished inside its budget.
"""
if time.monotonic() >= deadline:
return TableCleanupResult(rows_deleted=rows_deleted, stop_reason=stop_reason)
remaining: Final = await self._count_remaining(prisma_client, cutoff_date, table_name, time_column, deadline)
if remaining is not None:
SpendLogCleanupMetrics.set_rows_remaining(table_name, remaining)
return TableCleanupResult(rows_deleted=rows_deleted, stop_reason=stop_reason)
async def _delete_old_logs(
self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float
) -> TableCleanupResult:
return await self._delete_old_rows_batched(
prisma_client,
cutoff_date,
table_name="LiteLLM_SpendLogs",
key_columns=("request_id", "startTime"),
time_column="startTime",
deadline=deadline,
)
async def _delete_old_tool_index_rows(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int:
async def _delete_old_tool_index_rows(
self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float
) -> TableCleanupResult:
# SpendLogToolIndex rows are derived from spend logs, so they expire on the
# same cutoff; rows older than retention point at already-deleted logs.
return await self._delete_old_rows_batched(
@ -189,17 +462,87 @@ class SpendLogCleanup:
table_name="LiteLLM_SpendLogToolIndex",
key_columns=("request_id", "tool_name"),
time_column="start_time",
deadline=deadline,
)
async def _delete_old_autorouter_session_rows(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int:
async def _delete_old_autorouter_session_rows(
self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float
) -> TableCleanupResult:
return await self._delete_old_rows_batched(
prisma_client,
cutoff_date,
table_name="LiteLLM_AutoRouterSession",
key_columns=("api_key", "session_id", "router_name"),
time_column="last_turn_at",
deadline=deadline,
)
async def _clean_spend_log_tables(
self, prisma_client: PrismaClient, deadline: float
) -> tuple[TableCleanupResult, ...]:
"""
Prune the spend logs and the tool index rows derived from them.
When the table is range-partitioned, whole expired partitions are dropped
first because that reclaims disk immediately. Expired rows can still sit in
the DEFAULT partition (backfill, coverage gaps) or in a partition that spans
the cutoff, so retention still deletes those stragglers row-wise.
"""
cutoff_date: Final = datetime.now(timezone.utc) - timedelta(seconds=float(self.retention_seconds or 0))
verbose_proxy_logger.info("Removing logs older than %s", cutoff_date.isoformat())
# Partition maintenance is DDL taking an ACCESS EXCLUSIVE lock, so it is
# only STARTED while the run still has budget, and each statement carries
# the same timeouts the batches do. Without those, a DROP would queue
# behind any long-running reader for as long as that reader lives, which
# is the one way this job could still outlast its budget without bound.
remaining_timeout_ms: Final = self._remaining_timeout_ms(deadline)
if time.monotonic() >= deadline:
verbose_proxy_logger.info("Run budget already spent, skipping partition maintenance this run")
elif self.general_settings.get(
"use_spend_logs_partitioning", False
) and await self.partition_manager.is_partitioned(prisma_client, remaining_timeout_ms):
await self.partition_manager.ensure_partitions(prisma_client, remaining_timeout_ms)
dropped: Final = await self.partition_manager.drop_partitions_older_than(
prisma_client, cutoff_date, remaining_timeout_ms
)
verbose_proxy_logger.info("Dropped %d expired spend-log partitions: %s", len(dropped), dropped)
logs_result: Final = await self._delete_old_logs(prisma_client, cutoff_date, deadline)
verbose_proxy_logger.info("Deleted %s logs", logs_result.rows_deleted)
index_result: Final = await self._delete_old_tool_index_rows(prisma_client, cutoff_date, deadline)
verbose_proxy_logger.info("Deleted %s expired tool index rows", index_result.rows_deleted)
return (logs_result, index_result)
async def _clean_session_rollup(
self, prisma_client: PrismaClient, retention_seconds: int, deadline: float
) -> tuple[TableCleanupResult, ...]:
"""
Prune auto-router session rollup rows, which carry their own retention horizon.
"""
session_cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=float(retention_seconds))
sessions_result: Final = await self._delete_old_autorouter_session_rows(prisma_client, session_cutoff, deadline)
verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_result.rows_deleted)
return (sessions_result,)
@staticmethod
def _run_outcome(results: tuple[TableCleanupResult, ...]) -> RunOutcome:
"""
Report the most operationally significant reason the run stopped.
A bound that was hit matters more than a table that simply ran dry, so
those win over "completed", and an abort wins over everything.
"""
reasons: Final = frozenset(result.stop_reason for result in results)
if "aborted" in reasons:
return "aborted"
if "budget_exhausted" in reasons:
return "budget_exhausted"
if "batch_cap_reached" in reasons:
return "batch_cap_reached"
return "completed"
async def cleanup_old_spend_logs(self, prisma_client: PrismaClient) -> None:
"""
Main cleanup function. Deletes old spend logs in batches.
@ -209,16 +552,19 @@ class SpendLogCleanup:
lock_acquired = False
try:
verbose_proxy_logger.info("Cleanup job triggered at %s", datetime.now())
self._refresh_bounds()
delete_spend_logs: Final = self._should_delete_spend_logs()
autorouter_retention_seconds: Final = self._retention_seconds_for(
"maximum_autorouter_session_retention_period"
)
if not delete_spend_logs and autorouter_retention_seconds is None:
SpendLogCleanupMetrics.record_run("skipped_disabled")
return
if delete_spend_logs and self.retention_seconds is None:
verbose_proxy_logger.error("Retention seconds is None, cannot proceed with cleanup")
SpendLogCleanupMetrics.record_run("skipped_disabled")
return
# If we have a pod lock manager, try to acquire the lock
@ -235,43 +581,23 @@ class SpendLogCleanup:
if not lock_acquired:
verbose_proxy_logger.info("Another pod is already running cleanup")
SpendLogCleanupMetrics.record_run("skipped_locked")
return
if delete_spend_logs and self.retention_seconds is not None:
cutoff_date: Final = datetime.now(timezone.utc) - timedelta(seconds=float(self.retention_seconds))
verbose_proxy_logger.info("Removing logs older than %s", cutoff_date.isoformat())
deadline: Final = time.monotonic() + self.run_budget_seconds
if self.general_settings.get(
"use_spend_logs_partitioning", False
) and await self.partition_manager.is_partitioned(prisma_client):
await self.partition_manager.ensure_partitions(prisma_client)
dropped: Final = await self.partition_manager.drop_partitions_older_than(prisma_client, cutoff_date)
verbose_proxy_logger.info(
"Dropped %d expired spend-log partitions: %s",
len(dropped),
dropped,
)
# DROP only reclaims whole expired partitions. Expired rows can
# still sit in the DEFAULT partition (backfill, coverage gaps)
# or in a partition that spans the cutoff, so retention must
# also delete those stragglers row-wise.
total_deleted = await self._delete_old_logs(prisma_client, cutoff_date)
verbose_proxy_logger.info(
"Deleted %s expired logs not covered by dropped partitions", total_deleted
)
else:
total_deleted = await self._delete_old_logs(prisma_client, cutoff_date)
verbose_proxy_logger.info("Deleted %s logs", total_deleted)
spend_log_results: Final = (
await self._clean_spend_log_tables(prisma_client, deadline)
if delete_spend_logs and self.retention_seconds is not None
else ()
)
session_results: Final = (
await self._clean_session_rollup(prisma_client, autorouter_retention_seconds, deadline)
if autorouter_retention_seconds is not None
else ()
)
index_deleted: Final = await self._delete_old_tool_index_rows(prisma_client, cutoff_date)
verbose_proxy_logger.info("Deleted %s expired tool index rows", index_deleted)
if autorouter_retention_seconds is not None:
session_cutoff: Final = datetime.now(timezone.utc) - timedelta(
seconds=float(autorouter_retention_seconds)
)
sessions_deleted: Final = await self._delete_old_autorouter_session_rows(prisma_client, session_cutoff)
verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_deleted)
SpendLogCleanupMetrics.record_run(self._run_outcome(spend_log_results + session_results))
except Exception as e:
# .exception() captures the traceback; str(e) alone on a Prisma/DB
@ -281,6 +607,7 @@ class SpendLogCleanup:
type(e).__name__,
e,
)
SpendLogCleanupMetrics.record_run("aborted")
return # Return after error handling
finally:
# Only release the lock if it was actually acquired

122
litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py Normal file
View file

@ -0,0 +1,122 @@
"""
Prometheus metrics for the spend-log retention cleanup job.
The job runs in the background on a single elected pod, so its cost is invisible
from request-path metrics. These instruments make a run's database footprint
observable: how much it deleted, how long each batch took, how much work is
still outstanding, and why a run stopped.
``prometheus_client`` is an optional dependency, so every recorder degrades to a
no-op when it is absent.
"""
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from litellm._logging import verbose_proxy_logger
if TYPE_CHECKING:
# aliased so the annotations below cannot be mistaken for collections.Counter
from prometheus_client import Counter as PrometheusCounter
from prometheus_client import Gauge as PrometheusGauge
from prometheus_client import Histogram as PrometheusHistogram
RunOutcome: TypeAlias = Literal[
"completed",
"budget_exhausted",
"batch_cap_reached",
"skipped_locked",
"skipped_disabled",
"aborted",
]
_BATCH_DURATION_BUCKETS: Final = (0.005, 0.025, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0)
_TABLE_LABEL: Final = ("table",)
_OUTCOME_LABEL: Final = ("outcome",)
class SpendLogCleanupMetrics:
"""
Lazily-registered Prometheus instruments for the retention cleanup job.
Registration is deferred to first use so that importing this module never
touches the Prometheus registry, which keeps it safe to import from the
proxy regardless of whether Prometheus is a configured callback.
"""
_initialized: bool = False
rows_deleted: "PrometheusCounter | None" = None
batch_duration: "PrometheusHistogram | None" = None
rows_remaining: "PrometheusGauge | None" = None
batch_failures: "PrometheusCounter | None" = None
runs: "PrometheusCounter | None" = None
@classmethod
def _ensure_initialized(cls) -> None:
if cls._initialized:
return
cls._initialized = True
try:
# prometheus_client is an optional extra, so it is resolved here rather
# than at module import: this module is reachable from proxy startup
# regardless of whether Prometheus is a configured callback.
from prometheus_client import Counter, Gauge, Histogram
cls.rows_deleted = Counter(
"litellm_spend_log_cleanup_rows_deleted_total",
"Rows deleted by the spend-log retention cleanup job",
labelnames=_TABLE_LABEL,
)
cls.batch_duration = Histogram(
"litellm_spend_log_cleanup_batch_duration_seconds",
"Wall-clock duration of one retention cleanup delete batch",
labelnames=_TABLE_LABEL,
buckets=_BATCH_DURATION_BUCKETS,
)
cls.rows_remaining = Gauge(
"litellm_spend_log_cleanup_rows_remaining",
"Expired rows still awaiting deletion, counted only up to "
"SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP so the probe itself cannot scan a "
"large table; a value equal to that cap means at least that many remain",
labelnames=_TABLE_LABEL,
multiprocess_mode="livemax",
)
cls.batch_failures = Counter(
"litellm_spend_log_cleanup_batch_failures_total",
"Retention cleanup delete batches that raised",
labelnames=_TABLE_LABEL,
)
cls.runs = Counter(
"litellm_spend_log_cleanup_runs_total",
"Retention cleanup runs, labelled by why the run ended",
labelnames=_OUTCOME_LABEL,
)
except Exception as e: # noqa: BLE001 - a metrics problem must never fail the cleanup run
# Covers the extra being absent, a duplicate registration (repeated
# imports under a test runner), and registry misconfiguration alike.
verbose_proxy_logger.warning("Could not register spend-log cleanup metrics: %s", e)
@classmethod
def record_batch(cls, table_name: str, rows_deleted: int, duration_seconds: float) -> None:
cls._ensure_initialized()
if cls.rows_deleted is not None:
cls.rows_deleted.labels(table=table_name).inc(rows_deleted)
if cls.batch_duration is not None:
cls.batch_duration.labels(table=table_name).observe(duration_seconds)
@classmethod
def record_batch_failure(cls, table_name: str) -> None:
cls._ensure_initialized()
if cls.batch_failures is not None:
cls.batch_failures.labels(table=table_name).inc()
@classmethod
def set_rows_remaining(cls, table_name: str, remaining: int) -> None:
cls._ensure_initialized()
if cls.rows_remaining is not None:
cls.rows_remaining.labels(table=table_name).set(remaining)
@classmethod
def record_run(cls, outcome: RunOutcome) -> None:
cls._ensure_initialized()
if cls.runs is not None:
cls.runs.labels(outcome=outcome).inc()

128
litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py
View file

@ -14,8 +14,9 @@ keeps the batched-DELETE path, so existing deployments are untouched.
"""
import re
from collections.abc import Callable
from datetime import date, datetime, timedelta, timezone
from typing import Final
from typing import TYPE_CHECKING, Final, TypeAlias
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
@ -23,8 +24,23 @@ from litellm.constants import (
SPEND_LOG_PARTITION_PRECREATE_AHEAD,
)
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
SPEND_LOGS_TABLE: Final = "LiteLLM_SpendLogs"
RemainingTimeoutMs: TypeAlias = Callable[[], "int | None"]
"""
The per-statement bound in milliseconds, or None once the caller's budget is
spent.
Injected rather than passed as a number so it is re-evaluated before EVERY
statement: a value read once at entry would let a loop issue N statements each
bounded by the budget that was left before the first of them, which is not a
bound on the loop at all. The caller owns the policy; this module only asks how
much time it may still use.
"""
PartitionInterval = str # "day" | "week" | "month"
VALID_PARTITION_INTERVALS: Final = {"day", "week", "month"}
@ -116,21 +132,26 @@ class SpendLogsPartitionManager:
self.interval = interval
self.precreate_ahead = precreate_ahead
async def is_partitioned(self, prisma_client) -> bool:
async def is_partitioned(self, prisma_client: "PrismaClient", remaining_timeout_ms: RemainingTimeoutMs) -> bool:
budget_ms: Final = remaining_timeout_ms()
if budget_ms is None:
return False
try:
rows: Final = await prisma_client.db.query_raw(
"""
SELECT EXISTS (
SELECT 1
FROM pg_partitioned_table pt
JOIN pg_class c ON c.oid = pt.partrelid
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relname = $1
AND n.nspname = current_schema()
) AS partitioned
""",
SPEND_LOGS_TABLE,
)
async with prisma_client.db.tx() as tx:
await tx.execute_raw(f"SET LOCAL statement_timeout = {budget_ms}")
rows: Final = await tx.query_raw(
"""
SELECT EXISTS (
SELECT 1
FROM pg_partitioned_table pt
JOIN pg_class c ON c.oid = pt.partrelid
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relname = $1
AND n.nspname = current_schema()
) AS partitioned
""",
SPEND_LOGS_TABLE,
)
except Exception as e:
verbose_proxy_logger.warning(
"Could not determine if %s is partitioned, assuming it is not: %s",
@ -140,7 +161,25 @@ class SpendLogsPartitionManager:
return False
return bool(rows and rows[0].get("partitioned"))
async def ensure_partitions(self, prisma_client) -> list[str]:
@staticmethod
async def _execute_bounded_ddl(prisma_client: "PrismaClient", statement: str, timeout_ms: int) -> None:
"""
Run one DDL statement under a Postgres statement and lock timeout.
Partition DDL takes an ACCESS EXCLUSIVE lock, so an unbounded statement
queues behind any long-running reader for as long as that reader lives,
and the caller's run budget cannot cut it short. lock_timeout bounds the
wait for the lock and statement_timeout bounds the work itself, so a
partition this run cannot get is simply left for the next one.
"""
async with prisma_client.db.tx() as tx:
await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}")
await tx.execute_raw(statement)
async def ensure_partitions(
self, prisma_client: "PrismaClient", remaining_timeout_ms: RemainingTimeoutMs
) -> list[str]:
"""
Ensure the current and upcoming partitions exist, returning the names
now present. CREATE TABLE IF NOT EXISTS is a no-op for partitions that
@ -150,42 +189,61 @@ class SpendLogsPartitionManager:
for name, lower, upper in upcoming_partitions(
datetime.now(timezone.utc).date(), self.interval, self.precreate_ahead
):
budget_ms = remaining_timeout_ms()
if budget_ms is None:
verbose_proxy_logger.info("Run budget spent, leaving the remaining partitions for the next run")
break
try:
await prisma_client.db.execute_raw(
await self._execute_bounded_ddl(
prisma_client,
f'CREATE TABLE IF NOT EXISTS "{name}" '
f'PARTITION OF "{SPEND_LOGS_TABLE}" '
f"FOR VALUES FROM ('{lower.isoformat()}') TO ('{upper.isoformat()}')"
f"FOR VALUES FROM ('{lower.isoformat()}') TO ('{upper.isoformat()}')",
budget_ms,
)
ensured.append(name)
except Exception as e:
verbose_proxy_logger.warning("Failed to ensure spend-log partition %s: %s", name, e)
return ensured
async def _list_partitions(self, prisma_client) -> list[tuple[str, datetime | None]]:
rows: Final = await prisma_client.db.query_raw(
"""
SELECT c.relname AS name,
pg_get_expr(c.relpartbound, c.oid) AS bound
FROM pg_inherits i
JOIN pg_class c ON c.oid = i.inhrelid
JOIN pg_class p ON p.oid = i.inhparent
JOIN pg_namespace n ON n.oid = p.relnamespace
WHERE p.relname = $1
AND n.nspname = current_schema()
""",
SPEND_LOGS_TABLE,
)
async def _list_partitions(
self, prisma_client: "PrismaClient", timeout_ms: int
) -> list[tuple[str, datetime | None]]:
async with prisma_client.db.tx() as tx:
await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
rows: Final = await tx.query_raw(
"""
SELECT c.relname AS name,
pg_get_expr(c.relpartbound, c.oid) AS bound
FROM pg_inherits i
JOIN pg_class c ON c.oid = i.inhrelid
JOIN pg_class p ON p.oid = i.inhparent
JOIN pg_namespace n ON n.oid = p.relnamespace
WHERE p.relname = $1
AND n.nspname = current_schema()
""",
SPEND_LOGS_TABLE,
)
return [(row["name"], parse_partition_upper_bound(row.get("bound") or "")) for row in rows]
async def drop_partitions_older_than(self, prisma_client, cutoff: datetime) -> list[str]:
async def drop_partitions_older_than(
self, prisma_client: "PrismaClient", cutoff: datetime, remaining_timeout_ms: RemainingTimeoutMs
) -> list[str]:
"""DROP every partition whose whole range is older than `cutoff`."""
list_budget_ms: Final = remaining_timeout_ms()
if list_budget_ms is None:
return []
cutoff_naive: Final = cutoff.astimezone(timezone.utc).replace(tzinfo=None)
partitions: Final = await self._list_partitions(prisma_client)
partitions: Final = await self._list_partitions(prisma_client, list_budget_ms)
to_drop: Final = select_partitions_to_drop(partitions, cutoff_naive)
dropped: Final[list[str]] = []
for name in to_drop:
budget_ms = remaining_timeout_ms()
if budget_ms is None:
verbose_proxy_logger.info("Run budget spent, leaving the remaining partitions for the next run")
break
try:
await prisma_client.db.execute_raw(f'DROP TABLE IF EXISTS "{name}"')
await self._execute_bounded_ddl(prisma_client, f'DROP TABLE IF EXISTS "{name}"', budget_ms)
dropped.append(name)
except Exception as e:
verbose_proxy_logger.warning("Failed to drop spend-log partition %s: %s", name, e)

125
litellm/proxy/guardrails/anthropic_sse.py Normal file
View file

@ -0,0 +1,125 @@
"""Anthropic SSE <-> ModelResponse conversion for guardrail streaming hooks.
`/v1/messages` streams reach a guardrail's `async_post_call_streaming_iterator_hook` as raw SSE
frames rather than chunk objects, which `stream_chunk_builder` cannot assemble. These helpers let a
hook scan such a stream, and re-emit it when the guardrail rewrote the response.
"""
from __future__ import annotations
import json
from collections.abc import Mapping, Sequence
from typing import Final
from litellm.types.utils import Choices, ModelResponse
def is_raw_sse_stream(all_chunks: Sequence[object]) -> bool:
return any(isinstance(chunk, (str, bytes)) for chunk in all_chunks)
def _joined_sse_stream(all_chunks: Sequence[object]) -> str | None:
raw: Final = b"".join(
chunk if isinstance(chunk, bytes) else chunk.encode("utf-8")
for chunk in all_chunks
if isinstance(chunk, (str, bytes))
)
try:
return raw.decode("utf-8")
except UnicodeDecodeError:
return None
def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None:
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
)
return next(
(
message
for event in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(sse_stream) # pyright: ignore[reportPrivateUsage] # same parser the assembler uses
if (event_data := AnthropicPassthroughLoggingHandler._extract_sse_data(event)) is not None # pyright: ignore[reportPrivateUsage] # same parser the assembler uses; a private import beats forking SSE parsing
and event_data.get("type") == "message_start"
and isinstance(message := event_data.get("message"), dict)
),
None,
)
def assemble_anthropic_sse_stream(
all_chunks: Sequence[object], *, restore_identity: bool = False
) -> ModelResponse | None:
"""Assemble raw Anthropic SSE frames into a ModelResponse.
``restore_identity`` stamps the upstream message id and model onto the result, which the
assembler does not carry through. It is off by default so callers that re-emit the assembled
response keep the wire shape they had before this helper was shared. The writes land on a
freshly built object that is unreachable from caller state until returned.
"""
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
)
sse_stream: Final = _joined_sse_stream(all_chunks)
if sse_stream is None:
return None
message_start: Final = _anthropic_message_start(sse_stream)
if message_start is None:
return None
model: Final = message_start.get("model") if restore_identity else None
try:
assembled: Final = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only SSE-to-ModelResponse assembler; reimplementing it here would fork the parser
all_chunks=(sse_stream,),
litellm_logging_obj=None, # pyright: ignore[reportArgumentType] # only forwarded to stream_chunk_builder, which accepts None
model=model if isinstance(model, str) else "",
)
except Exception: # noqa: BLE001 # stream_chunk_builder re-raises every assembly failure as litellm.APIError
return None
if not isinstance(assembled, ModelResponse):
return None
if not restore_identity:
return assembled
message_id: Final = message_start.get("id")
if isinstance(message_id, str):
assembled.id = message_id
if isinstance(model, str) and model:
assembled.model = model
return assembled
def model_response_text(response: ModelResponse) -> str:
"""Assistant text of a response, used to detect whether a guardrail rewrote it."""
return "".join(
choice.message.content
for choice in response.choices
if isinstance(choice, Choices) # pyright: ignore[reportUnnecessaryIsInstance] # runtime choices can be StreamingChoices
and isinstance(choice.message.content, str)
)
def anthropic_sse_error_frames(message: str) -> tuple[bytes, ...]:
"""Anthropic error event, for a failure discovered after the response headers were flushed.
Once a keepalive ping has been sent a raise cannot reach the client, so the failure has to
travel as a frame.
"""
body: Final = json.dumps(message)
return (
f'event: error\ndata: {{"type": "error", "error": {{"type": "guardrail_error", '
f'"message": {body}}}}}\n\n'.encode(),
)
def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, ...]:
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import (
LiteLLMAnthropicMessagesAdapter,
)
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
anthropic_response: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(
response=assembled
)
return tuple(FakeAnthropicMessagesStreamIterator(response=anthropic_response).chunks)

65
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -14,6 +14,7 @@ import copy
import json
import re
import sys
import time
from collections.abc import AsyncGenerator, Mapping, Sequence
from datetime import datetime, timezone
from itertools import accumulate, groupby
@ -30,6 +31,7 @@ from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
from litellm.exceptions import ModifyResponseException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_scan_only_tool_results_for_guardrail,
)
@ -39,6 +41,15 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_request_processing import _serialize_http_exception_detail
from litellm.proxy.common_utils.sse_keepalive import keepalive_ping_has_fired
from litellm.proxy.guardrails.anthropic_sse import (
anthropic_sse_chunks_from_response,
anthropic_sse_error_frames,
assemble_anthropic_sse_stream,
is_raw_sse_stream,
model_response_text,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.guardrails import BedrockChecksConfigModel, GuardrailEventHooks
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
@ -2578,14 +2589,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
from litellm.types.utils import TextCompletionResponse
# Collect all chunks to process them together
started_at: Final = time.monotonic()
all_chunks: Final[list[ModelResponseStream]] = []
async for chunk in response:
all_chunks.append(chunk)
assembled_model_response: ModelResponse | TextCompletionResponse | None = stream_chunk_builder(
chunks=all_chunks,
# /v1/messages arrives as SSE frames, which stream_chunk_builder cannot assemble
raw_sse: Final = is_raw_sse_stream(all_chunks)
assembled_model_response: ModelResponse | TextCompletionResponse | None = (
assemble_anthropic_sse_stream(all_chunks, restore_identity=True)
if raw_sse
else stream_chunk_builder(chunks=all_chunks)
)
if isinstance(assembled_model_response, ModelResponse):
pre_guardrail_text: Final = model_response_text(assembled_model_response)
_pre_block_response: Final = assembled_model_response
####################################################################
########## 1. Make Bedrock Apply Guardrail API request ##########
#
@ -2609,7 +2627,32 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data=request_data,
logging_event_type=GuardrailEventHooks.post_call,
)
except HTTPException as block_exc:
block_detail: Final = block_exc.detail
# A policy block is the only 400 carrying a structured detail; a service failure
# either details a plain string or reports a non-400 status. Re-raising a service
# failure keeps its real status, but only while the headers are unflushed: past the
# first keepalive ping the raise reaches nobody, so it has to travel as a frame too
is_block: Final = raw_sse and block_exc.status_code == 400 and isinstance(block_detail, Mapping)
headers_flushed: Final = keepalive_ping_has_fired(
time.monotonic() - started_at, litellm.anthropic_sse_ping_interval_seconds
)
if not raw_sse or (not is_block and not headers_flushed):
raise
block_message, _ = _serialize_http_exception_detail(block_detail)
for error_frame in anthropic_sse_error_frames(
block_message if is_block else f"{block_exc.status_code}: {block_message}"
):
yield error_frame
return
except ModifyResponseException as e:
if raw_sse:
e.model = _pre_block_response.model or e.model # rebind-ok: exc.model defaults to the guardrail
if e.original_response is None:
e.original_response = _pre_block_response # rebind-ok: the block builder reads usage off this
for block_chunk in AnthropicMessagesHandler().build_block_sse_chunks(e, stream_started=False):
yield block_chunk
return
# Preserve upstream usage from the LLM call we already
# consumed. Non-streaming blocks carry it via
# ModifyResponseException.original_response +
@ -2642,11 +2685,29 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
#########################################################################
########## 3. Return the (potentially masked) chunks ##########
#########################################################################
if raw_sse:
for sse_chunk in (
anthropic_sse_chunks_from_response(assembled_model_response)
if model_response_text(assembled_model_response) != pre_guardrail_text
else all_chunks
):
yield sse_chunk
return
mock_response: Final = MockResponseIterator(model_response=assembled_model_response)
# Return the reconstructed stream
async for chunk in mock_response:
yield chunk
elif raw_sse:
# Forwarding an unscannable stream would silently disable the guardrail, so fail closed.
# A raise cannot reach the client once a keepalive ping has flushed the headers, so the
# refusal travels as a frame, matching how a block is delivered above
for error_frame in anthropic_sse_error_frames(
f"{self.guardrail_name}: streamed response could not be assembled for scanning, blocking it"
):
yield error_frame
return
else:
for chunk in all_chunks:
yield chunk

32
litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py
View file

@ -11,6 +11,7 @@ import requests
from fastapi import HTTPException
from httpx import HTTPStatusError
from requests.auth import HTTPBasicAuth
from typing_extensions import ReadOnly
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
@ -55,6 +56,26 @@ class _HiddenlayerResponse(TypedDict, total=False):
modified_data: Mapping[str, _HiddenlayerModifiedSide]
class _LoggedCallMetadata(TypedDict, total=False):
headers: ReadOnly[Mapping[str, str]]
class _LoggedCallLitellmParams(TypedDict, total=False):
metadata: ReadOnly[_LoggedCallMetadata]
class _HiddenlayerOutputMessage(TypedDict, total=False):
content: ReadOnly[str | Sequence[Mapping[str, str]]]
class _HiddenlayerChoiceMessage(TypedDict, total=False):
content: ReadOnly[str]
class _HiddenlayerChoice(TypedDict, total=False):
message: ReadOnly[_HiddenlayerChoiceMessage]
def is_saas(host: str) -> bool:
"""Checks whether the connection is to the SaaS platform"""
@ -155,7 +176,10 @@ class HiddenlayerGuardrail(CustomGuardrail):
# from the logger object on the response from the model.
headers = request_data.get("proxy_server_request", {}).get("headers", {})
if not headers and logging_obj and logging_obj.model_call_details:
headers = logging_obj.model_call_details.get("litellm_params", {}).get("metadata", {}).get("headers", {})
logged_litellm_params: Final[_LoggedCallLitellmParams] = logging_obj.model_call_details.get(
"litellm_params", {}
)
headers = logged_litellm_params.get("metadata", {}).get("headers", {})
hl_request_metadata["requester_id"] = headers.get("hl-requester-id") or "LiteLLM"
project_id: Final = headers.get("hl-project-id")
@ -408,7 +432,8 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
if input_type == "request":
inputs["structured_messages"] = output
for message in output.get("messages", []):
modified_messages: Final[Sequence[_HiddenlayerOutputMessage]] = output.get("messages", [])
for message in modified_messages:
content = message.get("content", "")
if isinstance(content, list):
text_parts = [
@ -422,7 +447,8 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
inputs["texts"] = new_texts
elif input_type == "response" and inputs.get("texts"):
inputs["texts"] = [output.get("choices", [{}])[-1].get("message", {}).get("content", "")]
redacted_choices: Final[Sequence[_HiddenlayerChoice]] = output.get("choices", [{}])
inputs["texts"] = [redacted_choices[-1].get("message", {}).get("content", "")]
elif input_type == "response" and inputs.get("tool_calls"):
inputs["tool_calls"] = output

83
litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py
View file

@ -1,16 +1,20 @@
"""LLM-as-a-Judge guardrail: uses an LLM to score responses against weighted criteria."""
import json
import re
from collections.abc import Callable
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
from fastapi import HTTPException
import litellm
from litellm._logging import verbose_logger
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.llm_judge import (
default_router_provider,
extract_text_from_content,
judge_acompletion,
parse_json_verdict,
)
from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus
@ -32,50 +36,9 @@ Return ONLY valid JSON in this exact format:
_VALID_ON_FAILURE: Final = frozenset({"block", "log"})
def _default_router_provider() -> "Router | None":
try:
from litellm.proxy.proxy_server import llm_router
except ImportError:
return None
return llm_router
_JSON_FENCE_RE: Final = re.compile(r"```(?:json)?\s*(.*?)\s*```", re.DOTALL | re.IGNORECASE)
def _parse_judge_verdict(raw: str) -> dict[str, Any]:
"""Parse the judge's JSON verdict, tolerating markdown fences and surrounding prose."""
text = raw.strip()
fenced: Final = _JSON_FENCE_RE.search(text)
if fenced is not None:
text = fenced.group(1).strip()
parsed: object
try:
parsed = json.loads(text)
except json.JSONDecodeError:
start: Final = text.find("{")
end: Final = text.rfind("}")
if start == -1 or end <= start:
raise
parsed = json.loads(text[start : end + 1])
if not isinstance(parsed, dict):
raise ValueError("judge response is not a JSON object")
return cast(dict[str, Any], parsed) # cast-ok: narrowed to dict by the isinstance guard above
def _extract_text_from_content(content: Any) -> str:
"""Return plain text from a message content field (str or multimodal list)."""
if isinstance(content, str):
return content
if isinstance(content, list):
parts: Final = []
for part in content:
if isinstance(part, dict) and part.get("type") == "text":
parts.append(part.get("text", ""))
return " ".join(parts)
return ""
_default_router_provider: Final = default_router_provider
_parse_judge_verdict: Final = parse_json_verdict
_extract_text_from_content: Final = extract_text_from_content
def _get_litellm_param(
@ -168,25 +131,13 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
"content": _build_judge_prompt(self.criteria, messages, response_text),
},
]
router: Final = self._router_provider()
if router is not None and (
self.judge_model in router.model_group_alias or router.get_model_list(model_name=self.judge_model)
):
response = await router.acompletion(
model=self.judge_model,
messages=judge_messages,
response_format={"type": "json_object"},
temperature=0,
num_retries=0,
fallbacks=[],
)
else:
response = await litellm.acompletion(
model=self.judge_model,
messages=judge_messages,
response_format={"type": "json_object"},
temperature=0,
)
response: Final = await judge_acompletion(
self._router_provider(),
self.judge_model,
judge_messages,
response_format={"type": "json_object"},
temperature=0,
)
raw: Final = response.choices[0].message.content or "{}"
return _parse_judge_verdict(raw)

74
litellm/proxy/guardrails/guardrail_hooks/tool_permission.py
View file

@ -17,6 +17,11 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
)
from litellm.proxy.guardrails.anthropic_sse import (
anthropic_sse_chunks_from_response,
assemble_anthropic_sse_stream,
is_raw_sse_stream,
)
from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
PermissionError,
@ -870,7 +875,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
all_chunks.append(chunk)
assembled_model_response: Final[ModelResponse | TextCompletionResponse | None] = (
stream_chunk_builder(chunks=all_chunks) if not self._is_raw_sse_stream(all_chunks) else None
stream_chunk_builder(chunks=all_chunks) if not is_raw_sse_stream(all_chunks) else None
)
if isinstance(assembled_model_response, ModelResponse):
denied_tools = self._check_assembled_stream(assembled_model_response)
@ -883,9 +888,9 @@ class ToolPermissionGuardrail(CustomGuardrail):
yield chunk
return
anthropic_response: Final = self._assemble_anthropic_stream(all_chunks)
anthropic_response: Final = assemble_anthropic_sse_stream(all_chunks)
if anthropic_response is None:
if self._is_raw_sse_stream(all_chunks):
if is_raw_sse_stream(all_chunks):
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
message=(
@ -904,13 +909,9 @@ class ToolPermissionGuardrail(CustomGuardrail):
return
self._modify_response_with_permission_errors(anthropic_response, anthropic_denials)
for sse_chunk in self._rewritten_anthropic_sse_chunks(anthropic_response):
for sse_chunk in anthropic_sse_chunks_from_response(anthropic_response):
yield sse_chunk
@staticmethod
def _is_raw_sse_stream(all_chunks: Sequence[Any]) -> bool:
return any(isinstance(chunk, (str, bytes)) for chunk in all_chunks)
def _check_assembled_stream(
self, assembled: ModelResponse
) -> tuple[tuple[ChatCompletionMessageToolCall, PermissionError], ...]:
@ -924,60 +925,3 @@ class ToolPermissionGuardrail(CustomGuardrail):
if not denied_tools:
verbose_proxy_logger.debug("Tool Permission Guardrail Post-Call Hook: All tools allowed")
return denied_tools
@staticmethod
def _joined_sse_stream(all_chunks: Sequence[Any]) -> str | None:
raw: Final = b"".join(
chunk if isinstance(chunk, bytes) else chunk.encode("utf-8")
for chunk in all_chunks
if isinstance(chunk, (str, bytes))
)
try:
return raw.decode("utf-8")
except UnicodeDecodeError:
return None
@staticmethod
def _has_anthropic_message_start(sse_stream: str) -> bool:
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
)
return any(
(event_data := AnthropicPassthroughLoggingHandler._extract_sse_data(event)) is not None # pyright: ignore[reportPrivateUsage] # same parser the assembler uses; a private import beats forking SSE parsing
and event_data.get("type") == "message_start"
for event in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(sse_stream) # pyright: ignore[reportPrivateUsage] # same parser the assembler uses
)
@staticmethod
def _assemble_anthropic_stream(all_chunks: Sequence[Any]) -> ModelResponse | None:
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
)
sse_stream: Final = ToolPermissionGuardrail._joined_sse_stream(all_chunks)
if sse_stream is None or not ToolPermissionGuardrail._has_anthropic_message_start(sse_stream):
return None
try:
assembled = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only SSE-to-ModelResponse assembler; reimplementing it here would fork the parser
all_chunks=(sse_stream,),
litellm_logging_obj=None, # pyright: ignore[reportArgumentType] # only forwarded to stream_chunk_builder, which accepts None
model="",
)
except (AttributeError, TypeError, ValueError, json.JSONDecodeError):
return None
return assembled if isinstance(assembled, ModelResponse) else None
@staticmethod
def _rewritten_anthropic_sse_chunks(assembled: ModelResponse) -> tuple[bytes, ...]:
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import (
LiteLLMAnthropicMessagesAdapter,
)
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
anthropic_response: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(
response=assembled
)
return tuple(FakeAnthropicMessagesStreamIterator(response=anthropic_response).chunks)

38
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -2,9 +2,10 @@
import importlib
import os
from collections.abc import Callable, Iterator, Mapping
from datetime import datetime, timezone
from itertools import chain, count
from typing import Any, Final, Literal, Optional, cast
from typing import Any, Final, Literal, Optional, Protocol, cast
from pydantic import ValidationError
@ -59,6 +60,13 @@ from .guardrail_initializers import (
initialize_tool_permission,
)
class _GuardrailRowLike(Protocol):
@property
def guardrail_id(self) -> str: ...
def __iter__(self) -> Iterator[tuple[str, object]]: ...
guardrail_initializer_registry: Final = {
SupportedGuardrailIntegrations.BEDROCK.value: initialize_bedrock,
SupportedGuardrailIntegrations.LAKERA.value: initialize_lakera,
@ -125,7 +133,9 @@ def get_guardrail_initializer_from_hooks():
# Check for guardrail_initializer_registry dictionary
if hasattr(module, "guardrail_initializer_registry"):
registry = getattr(module, "guardrail_initializer_registry")
registry: Mapping[str, Callable[..., CustomGuardrail]] | None = getattr(
module, "guardrail_initializer_registry", None
)
if isinstance(registry, dict):
discovered_initializers.update(registry)
verbose_proxy_logger.debug(
@ -135,7 +145,7 @@ def get_guardrail_initializer_from_hooks():
# Check for standalone initialize_guardrail function (fallback for directory-based guardrails)
elif hasattr(module, "initialize_guardrail"):
# For directories with just initialize_guardrail, use the directory name as the key
initialize_fn = getattr(module, "initialize_guardrail")
initialize_fn: Callable[..., CustomGuardrail] | None = getattr(module, "initialize_guardrail", None)
discovered_initializers[item] = initialize_fn
verbose_proxy_logger.debug("Found initialize_guardrail function in %s", module_path)
@ -206,7 +216,9 @@ def get_guardrail_class_from_hooks():
# Check for guardrail_initializer_registry dictionary
if hasattr(module, "guardrail_class_registry"):
registry = getattr(module, "guardrail_class_registry")
registry: Mapping[str, type[CustomGuardrail]] | None = getattr(
module, "guardrail_class_registry", None
)
if isinstance(registry, dict):
discovered_classes.update(registry)
@ -275,7 +287,7 @@ class GuardrailRegistry:
guardrail_info: Final[str] = safe_dumps(guardrail.get("guardrail_info", {}))
# Create guardrail in DB
created_guardrail: Final = await GuardrailsRepository(prisma_client).table.create(
created_guardrail: Final[_GuardrailRowLike] = await GuardrailsRepository(prisma_client).table.create(
data={
"guardrail_name": guardrail_name,
"litellm_params": litellm_params,
@ -321,7 +333,7 @@ class GuardrailRegistry:
guardrail_info: Final[str] = safe_dumps(guardrail.get("guardrail_info", {}))
# Update in DB
updated_guardrail: Final = await GuardrailsRepository(prisma_client).table.update(
updated_guardrail: Final[_GuardrailRowLike] = await GuardrailsRepository(prisma_client).table.update(
where={"guardrail_id": guardrail_id},
data={
"guardrail_name": guardrail_name,
@ -482,7 +494,7 @@ class InMemoryGuardrailHandler:
custom_guardrail_callback = initializer(litellm_params, guardrail)
elif isinstance(guardrail_type, str) and "." in guardrail_type:
custom_guardrail_callback = self.initialize_custom_guardrail(
guardrail=cast(dict, guardrail),
guardrail=guardrail,
guardrail_type=guardrail_type,
litellm_params=litellm_params,
config_file_path=config_file_path,
@ -512,7 +524,7 @@ class InMemoryGuardrailHandler:
"skip_tool_message_in_guardrail are enabled together, which excludes every message from "
"scanning, so no request content would ever be scanned. Remove one of the two."
)
configured_run_in_parallel: Final = getattr(litellm_params, "run_in_parallel", None)
configured_run_in_parallel: Final[bool | None] = getattr(litellm_params, "run_in_parallel", None)
if configured_run_in_parallel is not None:
custom_guardrail_callback.run_in_parallel = bool(configured_run_in_parallel)
@ -532,7 +544,7 @@ class InMemoryGuardrailHandler:
def initialize_custom_guardrail(
self,
guardrail: dict,
guardrail: Guardrail,
guardrail_type: str,
litellm_params: LitellmParams,
config_file_path: str | None = None,
@ -550,7 +562,9 @@ class InMemoryGuardrailHandler:
guardrail_type,
)
_guardrail_class: Final = get_instance_fn(guardrail_type, config_file_path=config_file_path)
_guardrail_class: Final[Callable[..., CustomGuardrail]] = get_instance_fn(
guardrail_type, config_file_path=config_file_path
)
mode: Final = litellm_params.mode
if mode is None:
@ -683,8 +697,8 @@ class InMemoryGuardrailHandler:
@staticmethod
def _normalize_litellm_params_for_comparison(
params: Any | None,
) -> dict[str, Any] | None:
params: LitellmParams | Mapping[str, object] | None,
) -> Mapping[str, object] | None:
"""
Render litellm_params to a canonical dict so an in-memory LitellmParams and
the raw dict loaded from the DB compare equal when they describe the same

29
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -50,6 +50,7 @@ from litellm.router_utils.clientside_credential_handler import (
_ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the router path
clientside_credential_keys,
)
from litellm.secret_managers.main import get_secret_bool
#### Health ENDPOINTS ####
@ -1447,6 +1448,31 @@ def callback_name(callback):
return str(callback)
DISABLE_NO_REDIS_WARNING_ENV_VAR: Final = "LITELLM_DISABLE_NO_REDIS_WARNING"
def _show_no_redis_warning() -> bool:
"""
Whether the UI should warn that no Redis is configured.
Redis is what makes rate limits, budgets, router state, and cache
invalidation consistent across workers, so a proxy running without it is
only safe as a single worker. Both places a Redis can land count: the
coordination cache (from a Redis response cache, general_settings.
coordination_redis, or the REDIS_* env fallback) and the router's own
Redis (router_settings.redis_host), which backs cooldowns and usage-based
routing on its own. Operators who know they run one worker can silence the
warning with LITELLM_DISABLE_NO_REDIS_WARNING=true.
"""
from litellm.proxy.proxy_server import llm_router, redis_usage_cache
if redis_usage_cache is not None:
return False
if llm_router is not None and llm_router.cache.redis_cache is not None:
return False
return get_secret_bool(DISABLE_NO_REDIS_WARNING_ENV_VAR, False) is not True
async def _get_health_readiness_details(
response: Response | None = None,
) -> dict[str, Any]:
@ -1487,6 +1513,7 @@ async def _get_health_readiness_details(
# check log level
log_level_name: Final = logging.getLevelName(verbose_logger.getEffectiveLevel())
is_detailed_debug: Final = verbose_logger.isEnabledFor(logging.DEBUG)
show_no_redis_warning: Final = _show_no_redis_warning()
# check DB
if prisma_client is not None: # if db passed in, check if it's connected
@ -1506,6 +1533,7 @@ async def _get_health_readiness_details(
"use_aiohttp_transport": AsyncHTTPHandler._should_use_aiohttp_transport(),
"log_level": log_level_name,
"is_detailed_debug": is_detailed_debug,
"show_no_redis_warning": show_no_redis_warning,
}
else:
return {
@ -1517,6 +1545,7 @@ async def _get_health_readiness_details(
"use_aiohttp_transport": AsyncHTTPHandler._should_use_aiohttp_transport(),
"log_level": log_level_name,
"is_detailed_debug": is_detailed_debug,
"show_no_redis_warning": show_no_redis_warning,
}
except Exception as e:
raise HTTPException(status_code=503, detail=f"Service Unhealthy ({e})")

8
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -24,7 +24,7 @@ from typing import (
from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE
from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_str_from_messages,
@ -2991,6 +2991,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
rate_limit_type: Literal["output", "input", "total"],
) -> list[RedisPipelineIncrementOperation]:
"""Build Redis pipeline increment ops for TPM / parallel-request counters."""
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs
from litellm.proxy.common_utils.callback_utils import (
get_model_group_from_litellm_kwargs,
)
@ -2998,6 +2999,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# Get metadata from standard_logging_object - this correctly handles both
# 'metadata' and 'litellm_metadata' fields from litellm_params
standard_logging_object: Final = kwargs.get("standard_logging_object") or {}
request_metadata: Final = get_litellm_metadata_from_kwargs(kwargs)
if request_metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY):
# Internal sub-calls bill spend to the caller but are not the caller's
# traffic; charging them here would let background evals eat TPM headroom.
return []
standard_logging_metadata: Final = standard_logging_object.get("metadata") or {}
model_group: Final = get_model_group_from_litellm_kwargs(kwargs)

9
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -39,11 +39,10 @@ _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset(
CallTypes.pass_through.value,
CallTypes.llm_passthrough_route.value,
CallTypes.allm_passthrough_route.value,
# CheckBatchCost's synthetic logging_obj for a completed managed batch only ever
# carries user_api_key_user_id (from LiteLLM_ManagedObjectTable.created_by) and
# user_api_key_team_id (from .team_id) -- both are None for batches created with
# the master key or a team-less key, since the table never stores the raw key
# hash. The batch already incurred real provider cost, so track it regardless.
# CheckBatchCost's synthetic logging_obj for a completed managed batch carries
# whatever LiteLLM_ManagedObjectTable stored at create time, and all of it is
# None for a batch created before those columns were persisted, or by the master
# key. The batch already incurred real provider cost, so track it regardless.
CallTypes.aretrieve_batch.value,
}
)

6
litellm/proxy/litellm_pre_call_utils.py
View file

@ -16,6 +16,7 @@ import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.constants import (
CONSUMED_REQUEST_TAGS_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
LITELLM_PROXY_MASTER_KEY_ALIAS,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
@ -261,6 +262,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
"policy_sources",
"routing_decision",
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
"standard_logging_object",
"proxy_server_request",
@ -272,7 +274,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
)
_UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: Final = frozenset(
UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: Final = frozenset(
{
"litellm-disable-message-redaction",
}
@ -353,7 +355,7 @@ def _strip_untrusted_request_header_controls(
return
for header_name in list(headers.keys()):
if isinstance(header_name, str) and header_name.lower() in _UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS:
if isinstance(header_name, str) and header_name.lower() in UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS:
if allow_client_message_redaction_opt_out:
continue
headers.pop(header_name, None)

20
litellm/proxy/management_endpoints/access_group_endpoints.py
View file

@ -14,11 +14,11 @@ from litellm.proxy.auth.auth_checks import (
_cache_access_object,
_cache_key_object,
_cache_team_object,
_delete_cache_access_object,
_get_team_object_from_cache,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.management_helpers.access_group_team_sync import invalidate_access_group_cache
from litellm.proxy.utils import get_prisma_client_or_throw
from litellm.repositories.table_repositories import AccessGroupRepository
from litellm.types.access_group import (
@ -146,22 +146,6 @@ async def _cache_access_group_record(record: _AccessGroupRecord) -> None:
)
async def _invalidate_cache_access_group(access_group_id: str) -> None:
"""
Invalidate (delete) an access group entry from both in-memory and Redis caches.
Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server
to avoid circular imports, following the same pattern as key_management_endpoints.
"""
from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
await _delete_cache_access_object(
access_group_id=access_group_id,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
# ---------------------------------------------------------------------------
# DB sync helpers (called inside a Prisma transaction)
# ---------------------------------------------------------------------------
@ -595,7 +579,7 @@ async def delete_access_group(
from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
await _invalidate_cache_access_group(access_group_id)
await invalidate_access_group_cache(access_group_id)
await _patch_team_caches_remove_access_group(
affected_team_ids, access_group_id, user_api_key_cache, proxy_logging_obj
)

347
litellm/proxy/management_endpoints/auto_router_endpoints.py
View file

@ -13,6 +13,7 @@ from pydantic import BaseModel, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.exceptions import BudgetExceededError
from litellm.litellm_core_utils.llm_judge import router_resolves_model
from litellm.proxy._types import (
CommonProxyErrors,
LiteLLM_TeamTable,
@ -39,11 +40,16 @@ from litellm.types.management_endpoints.auto_router_endpoints import (
AutoRouterRoutingTestRequest,
AutoRouterRoutingTestResponse,
RequestComplexityRouterConfig,
ShadowEvalJobResponse,
ShadowEvalResult,
ShadowEvalSlice,
StartShadowEvalRequest,
)
if TYPE_CHECKING:
from fastapi import APIRouter, Depends, HTTPException, Query, status
from litellm.proxy.utils import PrismaClient
from litellm.router import Router
else:
try:
@ -388,14 +394,7 @@ async def get_auto_router_benchmarks(
"""
from litellm.proxy.proxy_server import prisma_client
if user_api_key_dict.user_role not in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
):
raise HTTPException(
status_code=403,
detail="Only proxy admin roles can view auto-router benchmarks across the deployment",
)
_require_admin_viewer(user_api_key_dict, "view auto-router benchmarks across the deployment")
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
@ -430,3 +429,335 @@ async def get_auto_router_benchmarks(
totals=_benchmark_totals(_summed_agg_row(rows)),
groups=groups,
)
# ---------------------------------------------------------------------------
# Shadow eval: pre-adoption evaluation of an auto-router against live traffic.
# The job row is immutable config plus stopped_at; status, counts, spend, and errors
# are derived from the append-only attempt rows, so reads here are aggregations
# bounded by each job's max_turns through the attempt table's job_id index.
# ---------------------------------------------------------------------------
def _require_admin_viewer(user_api_key_dict: UserAPIKeyAuth, action: str) -> None:
if user_api_key_dict.user_role not in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
):
raise HTTPException(status_code=403, detail=f"Only proxy admin roles can {action}")
def _require_admin_writer(user_api_key_dict: UserAPIKeyAuth, action: str) -> None:
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(status_code=403, detail=f"Only a proxy admin can {action}")
def _is_configured_pre_routing_strategy(llm_router: "Router", router_name: str) -> bool:
return any(
router_name in registry
for registry in (
llm_router.auto_routers,
llm_router.complexity_routers,
llm_router.adaptive_routers,
llm_router.quality_routers,
)
)
def _validate_judge_model(llm_router: "Router | None", judge_model: str) -> None:
"""Reject a judge model the dispatch path cannot resolve, at start rather than as a
silently growing error count once the job is already sampling and billing."""
if llm_router is not None and _is_configured_pre_routing_strategy(llm_router, judge_model):
raise HTTPException(
status_code=400,
detail=f"judge_model '{judge_model}' is an auto-router; the judge must be a plain model",
)
if router_resolves_model(llm_router, judge_model):
return
import litellm
try:
litellm.get_llm_provider(model=judge_model)
except Exception as e:
raise HTTPException(
status_code=400,
detail=(
f"judge_model '{judge_model}' is neither a model configured on this proxy nor a "
"provider-qualified public model name (e.g. 'anthropic/claude-sonnet-5')"
),
) from e
def _is_unique_violation(error: Exception) -> bool:
"""Whether a Prisma create failed on a unique index. One active job per key lives in
a partial unique index (raw SQL in the migration; schema.prisma cannot express partial
indexes), so the read-then-create check above it is advisory: two concurrent starts
pass the read, and the loser must surface as the same 409 rather than a 500."""
try:
from prisma.errors import UniqueViolationError
except ImportError:
return "unique constraint" in str(error).lower() or "P2002" in str(error)
return isinstance(error, UniqueViolationError)
class _AttemptAggRow(BaseModel):
grp: str
turn_count: int
real_wins: int
shadow_wins: int
ties: int
avg_confidence: float | None
_ATTEMPT_AGG_ROWS: Final = TypeAdapter(list[_AttemptAggRow])
_ATTEMPT_AGG_SELECT: Final = """
COUNT(*)::int AS turn_count,
COUNT(*) FILTER (WHERE outcome = 'real')::int AS real_wins,
COUNT(*) FILTER (WHERE outcome = 'shadow')::int AS shadow_wins,
COUNT(*) FILTER (WHERE outcome = 'tie')::int AS ties,
AVG(confidence)::float AS avg_confidence
FROM "LiteLLM_ShadowEvalAttempt"
WHERE job_id = $1 AND outcome != 'error'
GROUP BY 1
"""
_ATTEMPT_AGG_BY_TIER_SQL: Final = "SELECT COALESCE(tier, 'UNCLASSIFIED') AS grp," + _ATTEMPT_AGG_SELECT
_ATTEMPT_AGG_BY_MODEL_SQL: Final = "SELECT COALESCE(real_model, 'unknown') AS grp," + _ATTEMPT_AGG_SELECT
_SWEEP_FINISHED_JOBS_SQL: Final = """
UPDATE "LiteLLM_ShadowEvalJob" j SET stopped_at = NOW()
WHERE j.api_key_id = $1 AND j.stopped_at IS NULL
AND (
j.ends_at <= NOW()
OR (SELECT COUNT(*) FROM "LiteLLM_ShadowEvalAttempt" a WHERE a.job_id = j.id) >= j.max_turns
)
"""
_ATTEMPT_TOTALS_SQL: Final = """
SELECT
COUNT(*) FILTER (WHERE outcome != 'error')::int AS judged_count,
COUNT(*) FILTER (WHERE outcome = 'error')::int AS error_count,
COALESCE(SUM(judge_cost), 0)::float AS judge_spend
FROM "LiteLLM_ShadowEvalAttempt"
WHERE job_id = $1
"""
class _AttemptTotalsRow(BaseModel):
judged_count: int
error_count: int
judge_spend: float
_ATTEMPT_TOTALS_ROWS: Final = TypeAdapter(list[_AttemptTotalsRow])
def _pct_of(numerator: int, denominator: int) -> float:
return _pct(numerator, denominator)
def _slices(rows: Sequence[_AttemptAggRow]) -> tuple[ShadowEvalSlice, ...]:
return tuple(
ShadowEvalSlice(
group=row.grp,
turn_count=row.turn_count,
real_win_rate_pct=_pct_of(row.real_wins, row.turn_count),
shadow_win_rate_pct=_pct_of(row.shadow_wins, row.turn_count),
tie_rate_pct=_pct_of(row.ties, row.turn_count),
avg_judge_confidence=round(row.avg_confidence or 0.0, 3),
)
for row in sorted(rows, key=lambda r: r.turn_count, reverse=True)
)
async def _shadow_eval_results(prisma_client: "PrismaClient", job_id: str) -> ShadowEvalResult | None:
"""Both stratifications of one job's verdicts. Tier answers "where does the router do
well"; current-model answers "which of the models this key uses today would the router
beat". Reads are bounded by the job's own attempts (<= max_turns) via the job_id index."""
by_tier: Final = _ATTEMPT_AGG_ROWS.validate_python(
await prisma_client.db.query_raw(_ATTEMPT_AGG_BY_TIER_SQL, job_id) or ()
)
if not by_tier:
return None
by_model: Final = _ATTEMPT_AGG_ROWS.validate_python(
await prisma_client.db.query_raw(_ATTEMPT_AGG_BY_MODEL_SQL, job_id) or ()
)
total_turns: Final = sum(r.turn_count for r in by_tier)
return ShadowEvalResult(
by_tier=_slices(by_tier),
by_current_model=_slices(by_model),
overall_shadow_win_rate_pct=_pct_of(sum(r.shadow_wins for r in by_tier), total_turns),
overall_tie_rate_pct=_pct_of(sum(r.ties for r in by_tier), total_turns),
)
@router.post(
"/auto_router/shadow_eval/start",
tags=("auto router",),
dependencies=(Depends(user_api_key_auth),),
response_model=ShadowEvalJobResponse,
status_code=status.HTTP_201_CREATED,
)
async def start_shadow_eval(
data: StartShadowEvalRequest,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> ShadowEvalJobResponse:
"""
Start a pre-adoption shadow eval: duplicate a sampled slice of a key's live traffic
through an auto-router, judge real vs. shadow responses blind, and stratify win rates
by the router's tier classification and by the incumbent model.
Shadow responses are never served to users. The job samples until it has judged
max_turns turns, reaches the end of its window, or is stopped; sampling changes
propagate to pods within about 10 seconds. Shadow and judge calls bill to the
shadowed key but are excluded from request counts and auto-router adoption metrics.
"""
from litellm.proxy.proxy_server import llm_router, prisma_client
_require_admin_writer(user_api_key_dict, "start a shadow eval")
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
if llm_router is None or not _is_configured_pre_routing_strategy(llm_router, data.router_name):
raise HTTPException(status_code=400, detail=f"'{data.router_name}' is not a configured auto-router")
_validate_judge_model(llm_router, data.judge_model)
key_row: Final = await prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": data.api_key_id} # mutable-ok: Prisma filter
)
if key_row is None:
raise HTTPException(
status_code=400,
detail=(
f"api_key_id '{data.api_key_id}' is not a key on this proxy; pass the key's token hash, "
"the value the key list and key info endpoints report"
),
)
# A job that expired or exhausted its turn budget stopped sampling on its own, but
# still holds the one-active-per-key partial unique index until stamped; free it so
# a new eval can start.
await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, data.api_key_id)
active: Final = await prisma_client.db.litellm_shadowevaljob.find_first(
where={"api_key_id": data.api_key_id, "stopped_at": None}, # mutable-ok: Prisma filter
)
if active is not None:
raise HTTPException(
status_code=409,
detail=f"Key already has an active shadow eval job ({active.id}). Stop it first.",
)
now: Final = datetime.now(timezone.utc)
try:
job: Final = await prisma_client.db.litellm_shadowevaljob.create(
data={ # mutable-ok: Prisma payload
"api_key_id": data.api_key_id,
"router_name": data.router_name,
"judge_model": data.judge_model,
"shadow_percentage": data.shadow_percentage,
"max_turns": data.max_turns,
"created_by": user_api_key_dict.user_id,
"ends_at": now + timedelta(days=data.duration_days),
}
)
except Exception as e:
if not _is_unique_violation(e):
raise
raise HTTPException(
status_code=409,
detail="Key already has an active shadow eval job (started concurrently). Stop it first.",
) from e
return ShadowEvalJobResponse.model_validate(job, from_attributes=True)
@router.get(
"/auto_router/shadow_eval",
tags=("auto router",),
dependencies=(Depends(user_api_key_auth),),
response_model=list[ShadowEvalJobResponse],
)
async def list_shadow_eval_jobs(
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
api_key_id: Annotated[str | None, Query(description="Filter to jobs shadowing this key")] = None,
limit: Annotated[int, Query(ge=1, le=200, description="Newest jobs to return")] = 50,
) -> tuple[ShadowEvalJobResponse, ...]:
"""List shadow eval jobs, newest first. Counts and results ride the detail endpoint only."""
from litellm.proxy.proxy_server import prisma_client
_require_admin_viewer(user_api_key_dict, "view shadow evals")
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
records: Final = await prisma_client.db.litellm_shadowevaljob.find_many(
where={"api_key_id": api_key_id} if api_key_id else {}, # mutable-ok: Prisma filter
order={"created_at": "desc"}, # mutable-ok: Prisma order
take=limit,
)
return tuple(ShadowEvalJobResponse.model_validate(record, from_attributes=True) for record in records or ())
@router.get(
"/auto_router/shadow_eval/{job_id}",
tags=("auto router",),
dependencies=(Depends(user_api_key_auth),),
response_model=ShadowEvalJobResponse,
)
async def get_shadow_eval_job(
job_id: str,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> ShadowEvalJobResponse:
"""One job with derived counts, judge spend, latest error, and stratified results."""
from litellm.proxy.proxy_server import prisma_client
_require_admin_viewer(user_api_key_dict, "view shadow evals")
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
record: Final = await prisma_client.db.litellm_shadowevaljob.find_unique(
where={"id": job_id} # mutable-ok: Prisma filter
)
if record is None:
raise HTTPException(status_code=404, detail=f"No shadow eval job {job_id}")
totals: Final = _ATTEMPT_TOTALS_ROWS.validate_python(
await prisma_client.db.query_raw(_ATTEMPT_TOTALS_SQL, job_id) or ()
)
latest_error: Final = await prisma_client.db.litellm_shadowevalattempt.find_first(
where={"job_id": job_id, "outcome": "error"}, # mutable-ok: Prisma filter
order={"created_at": "desc"}, # mutable-ok: Prisma order
)
return ShadowEvalJobResponse.model_validate(record, from_attributes=True).model_copy(
update={ # mutable-ok: pydantic update payload
"judged_count": totals[0].judged_count if totals else 0,
"error_count": totals[0].error_count if totals else 0,
"judge_spend": round(totals[0].judge_spend, 6) if totals else 0.0,
"last_error": latest_error.error if latest_error else None,
"results": await _shadow_eval_results(prisma_client, job_id),
}
)
@router.post(
"/auto_router/shadow_eval/{job_id}/stop",
tags=("auto router",),
dependencies=(Depends(user_api_key_auth),),
response_model=ShadowEvalJobResponse,
)
async def stop_shadow_eval_job(
job_id: str,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> ShadowEvalJobResponse:
"""Stop an active shadow eval job. Attempts are kept; sampling halts within ~10s."""
from litellm.proxy.proxy_server import prisma_client
_require_admin_writer(user_api_key_dict, "stop a shadow eval")
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
record: Final = await prisma_client.db.litellm_shadowevaljob.find_unique(
where={"id": job_id} # mutable-ok: Prisma filter
)
if record is None:
raise HTTPException(status_code=404, detail=f"No shadow eval job {job_id}")
current: Final = ShadowEvalJobResponse.model_validate(record, from_attributes=True)
if current.status != "running":
raise HTTPException(status_code=400, detail=f"Job {job_id} is already {current.status}")
updated: Final = await prisma_client.db.litellm_shadowevaljob.update(
where={"id": job_id}, # mutable-ok: Prisma filter
data={"stopped_at": datetime.now(timezone.utc)}, # mutable-ok: Prisma payload
)
return ShadowEvalJobResponse.model_validate(updated, from_attributes=True)

110
litellm/proxy/management_endpoints/cost_tracking_settings.py
View file

@ -10,6 +10,8 @@ PATCH /config/cost_margin_config - Update cost margin configuration
POST /cost/estimate - Estimate cost for a given model and token counts
"""
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final
from fastapi import APIRouter, Depends, HTTPException
@ -24,29 +26,65 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.types.utils import LlmProvidersSet
from litellm.types.utils import CostPerToken, LlmProvidersSet, ModelInfo
router: Final = APIRouter()
def _resolve_model_for_cost_lookup(model: str) -> tuple[str, str | None]:
@dataclass(frozen=True, slots=True)
class ResolvedCostModel:
model: str
provider: str | None
custom_cost_per_token: CostPerToken | None
def _configured_price(key: str, sources: tuple[Mapping[str, object], ...]) -> float | None:
values: Final = (source.get(key) for source in sources)
numeric: Final = (float(value) for value in values if isinstance(value, (int, float)))
return next(numeric, None)
def _extract_custom_pricing(
litellm_params: Mapping[str, object], model_info: Mapping[str, object]
) -> CostPerToken | None:
"""
Pull per-token pricing configured on a deployment so on-prem / self-hosted
models (absent from the public cost map) still estimate a real cost.
Pricing may live on ``litellm_params`` or ``model_info``; ``litellm_params``
wins, matching the router's cost-map registration precedence.
"""
sources: Final = (litellm_params, model_info)
input_price: Final = _configured_price("input_cost_per_token", sources)
output_price: Final = _configured_price("output_cost_per_token", sources)
if input_price is None and output_price is None:
return None
return CostPerToken(
input_cost_per_token=input_price or 0.0,
output_cost_per_token=output_price or 0.0,
)
def _lookup_model_info(model: str) -> ModelInfo | None:
try:
return litellm.get_model_info(model=model)
except Exception:
return None
def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel:
"""
Resolve a model name (which may be a router alias/model_group) to the
underlying litellm model name for cost lookup.
underlying litellm model name, provider, and any deployment-configured
pricing used for cost lookup.
Args:
model: The model name from the request (could be a router alias like 'e-model-router'
or an actual model name like 'azure_ai/gpt-4')
Returns:
Tuple of (resolved_model_name, custom_llm_provider)
- resolved_model_name: The actual model name to use for cost lookup
- custom_llm_provider: The provider if resolved from router, None otherwise
"""
from litellm.proxy.proxy_server import llm_router
custom_llm_provider: str | None = None
# Try to resolve from router if available
if llm_router is not None:
try:
@ -57,31 +95,25 @@ def _resolve_model_for_cost_lookup(model: str) -> tuple[str, str | None]:
first_deployment: Final = deployments[0]
litellm_params: Final = first_deployment.get("litellm_params", {})
model_info: Final = first_deployment.get("model_info", {})
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
provider: Final = str(custom_llm_provider) if custom_llm_provider is not None else None
custom_cost_per_token: Final = _extract_custom_pricing(litellm_params, model_info)
# Check base_model first (needed for Azure custom deployment names)
base_model: Final = model_info.get("base_model") or litellm_params.get("base_model")
if base_model:
verbose_proxy_logger.debug("Resolved model '%s' to base_model '%s' from router", model, base_model)
custom_llm_provider = litellm_params.get("custom_llm_provider")
return (
str(base_model),
(str(custom_llm_provider) if custom_llm_provider is not None else None),
)
return ResolvedCostModel(str(base_model), provider, custom_cost_per_token)
resolved_model: Final = litellm_params.get("model")
if resolved_model:
verbose_proxy_logger.debug("Resolved model '%s' to '%s' from router", model, resolved_model)
custom_llm_provider = litellm_params.get("custom_llm_provider")
return (
str(resolved_model),
(str(custom_llm_provider) if custom_llm_provider is not None else None),
)
return ResolvedCostModel(str(resolved_model), provider, custom_cost_per_token)
except Exception as e:
verbose_proxy_logger.debug("Could not resolve model '%s' from router: %s", model, e)
# Return original model if not resolved
return model, custom_llm_provider
return ResolvedCostModel(model, None, None)
def _calculate_period_costs(num_requests, cost_per_request, input_cost, output_cost, margin_cost):
@ -450,7 +482,9 @@ async def estimate_cost(
from litellm.types.utils import ModelResponse, Usage
# Resolve model name (handles router aliases like 'e-model-router' -> 'azure_ai/gpt-4')
resolved_model, resolved_provider = _resolve_model_for_cost_lookup(request.model)
resolved: Final = _resolve_model_for_cost_lookup(request.model)
resolved_model: Final = resolved.model
resolved_provider: Final = resolved.provider
verbose_proxy_logger.debug("Cost estimate: request.model='%s' resolved to '%s'", request.model, resolved_model)
@ -480,6 +514,8 @@ async def estimate_cost(
cost_per_request: Final = completion_cost(
completion_response=mock_response,
model=resolved_model,
custom_llm_provider=resolved_provider,
custom_cost_per_token=resolved.custom_cost_per_token,
litellm_logging_obj=litellm_logging_obj,
)
except Exception as e:
@ -497,20 +533,22 @@ async def estimate_cost(
output_cost: Final = cost_breakdown.get("output_cost", 0.0) if cost_breakdown else 0.0
margin_cost: Final = cost_breakdown.get("margin_total_amount", 0.0) if cost_breakdown else 0.0
# Get model info for per-token pricing display
try:
model_info: Final = litellm.get_model_info(model=resolved_model)
input_cost_per_token = model_info.get("input_cost_per_token")
output_cost_per_token = model_info.get("output_cost_per_token")
custom_llm_provider = model_info.get("litellm_provider")
except Exception:
input_cost_per_token = None
output_cost_per_token = None
custom_llm_provider = None
model_info: Final = _lookup_model_info(resolved_model)
mapped_input_price: Final = model_info.get("input_cost_per_token") if model_info is not None else None
mapped_output_price: Final = model_info.get("output_cost_per_token") if model_info is not None else None
mapped_provider: Final = model_info.get("litellm_provider") if model_info is not None else None
# Use provider from router resolution if not found in model_info
if custom_llm_provider is None and resolved_provider is not None:
custom_llm_provider = resolved_provider
input_cost_per_token: Final = (
resolved.custom_cost_per_token["input_cost_per_token"]
if resolved.custom_cost_per_token is not None
else mapped_input_price
)
output_cost_per_token: Final = (
resolved.custom_cost_per_token["output_cost_per_token"]
if resolved.custom_cost_per_token is not None
else mapped_output_price
)
custom_llm_provider: Final = mapped_provider if mapped_provider is not None else resolved_provider
# Calculate daily and monthly costs
(

118
litellm/proxy/management_endpoints/customer_endpoints.py
View file

@ -10,12 +10,19 @@ All /customer management endpoints
"""
#### END-USER/CUSTOMER MANAGEMENT ####
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta
from typing import Final
from typing import TYPE_CHECKING, Final, Protocol, TypeVar, overload
import fastapi
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
if TYPE_CHECKING:
from prisma.models import LiteLLM_BudgetTable as PrismaBudgetRow
from prisma.models import LiteLLM_EndUserTable as PrismaEndUserRow
from litellm.proxy.utils import PrismaClient
import litellm
from litellm._logging import verbose_proxy_logger
@ -41,6 +48,54 @@ from litellm.types.proxy.management_endpoints.customer_endpoints import (
UnblockUsersResponse,
)
_RowT_co: Final = TypeVar("_RowT_co", covariant=True)
_STR_OBJECT_DICT: Final = TypeAdapter(dict[str, object])
if TYPE_CHECKING:
class _TableOps(Protocol[_RowT_co]):
async def find_first(
self,
where: Mapping[str, object] | None = None,
include: Mapping[str, bool] | None = None,
) -> _RowT_co | None: ...
async def find_many(
self,
where: Mapping[str, object] | None = None,
include: Mapping[str, bool] | None = None,
) -> Sequence[_RowT_co]: ...
async def create(
self,
data: Mapping[str, object],
include: Mapping[str, bool] | None = None,
) -> _RowT_co: ...
async def update(
self,
where: Mapping[str, object],
data: Mapping[str, object],
include: Mapping[str, bool] | None = None,
) -> _RowT_co | None: ...
async def upsert(
self,
where: Mapping[str, object],
data: Mapping[str, Mapping[str, object]],
) -> _RowT_co: ...
async def delete_many(self, where: Mapping[str, object]) -> int: ...
@overload
def _typed_table(repo: EndUserRepository) -> "_TableOps[PrismaEndUserRow]": ...
@overload
def _typed_table(repo: BudgetRepository) -> "_TableOps[PrismaBudgetRow]": ...
def _typed_table(repo: EndUserRepository | BudgetRepository) -> object:
return repo.table
router: Final = APIRouter()
@ -89,7 +144,7 @@ async def block_user(data: BlockUsers):
records: Final = []
if prisma_client is not None:
for id in data.user_ids:
record = await EndUserRepository(prisma_client).table.upsert(
record = await _typed_table(EndUserRepository(prisma_client)).upsert(
where={"user_id": id},
data={
"create": {"user_id": id, "blocked": True},
@ -184,7 +239,7 @@ def new_budget_request(data: NewCustomerRequest) -> BudgetNewRequest | None:
budget_kv_pairs[field_name] = value
if budget_kv_pairs:
budget_request: Final = BudgetNewRequest(**budget_kv_pairs)
budget_request: Final = BudgetNewRequest.model_validate(budget_kv_pairs)
validate_budget_duration(budget_request.budget_duration)
if budget_request.budget_reset_at is None and budget_request.budget_duration is not None:
budget_request.budget_reset_at = datetime.utcnow() + timedelta(
@ -195,10 +250,10 @@ def new_budget_request(data: NewCustomerRequest) -> BudgetNewRequest | None:
async def _handle_customer_object_permission_update(
non_default_values: dict,
non_default_values: dict[str, object],
end_user_table_data_typed: LiteLLM_EndUserTable | None,
update_end_user_table_data: dict,
prisma_client,
update_end_user_table_data: dict[str, object],
prisma_client: "PrismaClient",
) -> None:
"""
Handle object permission updates for customer endpoints.
@ -344,13 +399,13 @@ async def new_end_user(
},
)
new_end_user_obj: dict = {}
new_end_user_obj: dict[str, object] = {}
## CREATE BUDGET ## if set
_new_budget: Final = new_budget_request(data)
if _new_budget is not None:
try:
budget_record: Final = await BudgetRepository(prisma_client).table.create(
budget_record: Final = await _typed_table(BudgetRepository(prisma_client)).create(
data={
**_new_budget.model_dump(exclude_unset=True),
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
@ -364,16 +419,18 @@ async def new_end_user(
elif data.budget_id is not None:
new_end_user_obj["budget_id"] = data.budget_id
_user_data: Final = data.dict(exclude_none=True)
_user_data: Final = _STR_OBJECT_DICT.validate_python(data.dict(exclude_none=True))
for k, v in _user_data.items():
if k not in BudgetNewRequest.model_fields:
new_end_user_obj[k] = v
## Handle Object Permission - MCP Servers, Vector Stores etc.
new_end_user_obj = await _set_object_permission(
data_json=new_end_user_obj,
prisma_client=prisma_client,
new_end_user_obj = _STR_OBJECT_DICT.validate_python(
await _set_object_permission(
data_json=new_end_user_obj,
prisma_client=prisma_client,
)
)
# Ensure object_permission is not in the data being sent to create
@ -386,7 +443,7 @@ async def new_end_user(
new_end_user_obj.pop("object_permission", None)
## WRITE TO DB ##
end_user_record: Final = await EndUserRepository(prisma_client).table.create(
end_user_record: Final = await _typed_table(EndUserRepository(prisma_client)).create(
data=new_end_user_obj,
include={"litellm_budget_table": True, "object_permission": True},
)
@ -442,7 +499,7 @@ async def end_user_info(
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
user_info: Final = await EndUserRepository(prisma_client).table.find_first(
user_info: Final = await _typed_table(EndUserRepository(prisma_client)).find_first(
where={"user_id": end_user_id},
include={"litellm_budget_table": True, "object_permission": True},
)
@ -535,13 +592,13 @@ async def update_end_user(
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client
try:
data_json: Final[dict] = data.json()
data_json: Final = _STR_OBJECT_DICT.validate_python(data.json())
# get the row from db
if prisma_client is None:
raise Exception("Not connected to DB!")
# get non default values for key
non_default_values: Final = {}
non_default_values: Final = dict[str, object]()
for k, v in data_json.items():
if v is not None and v not in (
[],
@ -551,7 +608,7 @@ async def update_end_user(
non_default_values[k] = v
## Get end user table data ##
end_user_table_data: Final = await EndUserRepository(prisma_client).table.find_first(
end_user_table_data: Final = await _typed_table(EndUserRepository(prisma_client)).find_first(
where={"user_id": data.user_id}, include={"litellm_budget_table": True}
)
@ -563,14 +620,14 @@ async def update_end_user(
param="user_id",
)
end_user_table_data_typed: Final = LiteLLM_EndUserTable(**end_user_table_data.model_dump())
end_user_table_data_typed: Final = LiteLLM_EndUserTable.model_validate(end_user_table_data.model_dump())
## Get budget table data ##
end_user_budget_table: Final = end_user_table_data_typed.litellm_budget_table
## Get all params for budget table ##
budget_table_data: Final = {}
update_end_user_table_data: Final = {}
budget_table_data: Final = dict[str, object]()
update_end_user_table_data: Final = dict[str, object]()
for k, v in non_default_values.items():
# budget_id is for linking to existing budget, not for creating new budget
if k == "budget_id":
@ -593,7 +650,7 @@ async def update_end_user(
if budget_table_data:
if end_user_budget_table is None:
## Create new budget ##
budget_table_data_record = await BudgetRepository(prisma_client).table.create(
budget_table_data_record = await _typed_table(BudgetRepository(prisma_client)).create(
data={
**budget_table_data,
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
@ -605,7 +662,7 @@ async def update_end_user(
update_end_user_table_data["budget_id"] = budget_table_data_record.budget_id
else:
## Update existing budget ##
budget_table_data_record = await BudgetRepository(prisma_client).table.update(
budget_table_data_record = await _typed_table(BudgetRepository(prisma_client)).update(
where={"budget_id": end_user_budget_table.budget_id},
data=budget_table_data,
)
@ -625,7 +682,7 @@ async def update_end_user(
if data.user_id is not None and len(data.user_id) > 0:
update_end_user_table_data["user_id"] = data.user_id
verbose_proxy_logger.debug("In update customer, user_id condition block.")
response: Final = await EndUserRepository(prisma_client).table.update(
response: Final = await _typed_table(EndUserRepository(prisma_client)).update(
where={"user_id": data.user_id},
data=update_end_user_table_data,
include={"litellm_budget_table": True, "object_permission": True},
@ -688,7 +745,7 @@ async def delete_end_user(
verbose_proxy_logger.debug("/customer/delete: Received data = %s", data)
if data.user_ids is not None and isinstance(data.user_ids, list) and len(data.user_ids) > 0:
# First check if all users exist
existing_users: Final = await EndUserRepository(prisma_client).table.find_many(
existing_users: Final = await _typed_table(EndUserRepository(prisma_client)).find_many(
where={"user_id": {"in": data.user_ids}}
)
existing_user_ids: Final = {user.user_id for user in existing_users}
@ -703,7 +760,7 @@ async def delete_end_user(
)
# All users exist, proceed with deletion
response: Final = await EndUserRepository(prisma_client).table.delete_many(
response: Final = await _typed_table(EndUserRepository(prisma_client)).delete_many(
where={"user_id": {"in": data.user_ids}}
)
verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response)
@ -764,7 +821,7 @@ async def list_end_user(
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
response: Final = await EndUserRepository(prisma_client).table.find_many(
response: Final = await _typed_table(EndUserRepository(prisma_client)).find_many(
include={"litellm_budget_table": True, "object_permission": True}
)
@ -827,11 +884,10 @@ async def get_customer_daily_activity(
exclude_end_user_ids_list = exclude_end_user_ids.split(",") if exclude_end_user_ids else None
# Fetch organization aliases for metadata
where_condition: Final = {}
where_condition: Final = dict[str, object]()
if end_user_ids_list:
where_condition["user_id"] = {"in": list(end_user_ids_list)}
end_user_aliases: Final = await EndUserRepository(prisma_client).table.find_many(where=where_condition)
end_user_alias_metadata: Final = {e.user_id: {"alias": e.alias} for e in end_user_aliases}
end_user_aliases: Final = await _typed_table(EndUserRepository(prisma_client)).find_many(where=where_condition)
# Query daily activity for organizations
return await get_daily_activity(
@ -839,7 +895,7 @@ async def get_customer_daily_activity(
table_name="litellm_dailyenduserspend",
entity_id_field="end_user_id",
entity_id=end_user_ids_list,
entity_metadata_field=end_user_alias_metadata,
entity_metadata_field={e.user_id: {"alias": e.alias} for e in end_user_aliases},
exclude_entity_ids=exclude_end_user_ids_list,
start_date=start_date,
end_date=end_date,

4
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -2311,7 +2311,7 @@ async def delete_user(
fetch_all_teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_row.teams}})
teams_to_update = []
for team in fetch_all_teams:
is_member_in_team, new_team_members = _cleanup_members_with_roles(
removed_team_members, new_team_members = _cleanup_members_with_roles(
existing_team_row=LiteLLM_TeamTable.model_validate(team.model_dump()),
data=TeamMemberDeleteRequest(
team_id=team.team_id,
@ -2319,7 +2319,7 @@ async def delete_user(
user_email=user_row.user_email,
),
)
if is_member_in_team:
if removed_team_members:
_db_new_team_members: list[dict] = [m.model_dump() for m in new_team_members]
team.members_with_roles = json.dumps(_db_new_team_members)
teams_to_update.append(team)

84
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -88,6 +88,11 @@ from litellm.proxy.management_endpoints.common_utils import (
from litellm.proxy.management_endpoints.model_management_endpoints import (
_add_model_to_db,
)
from litellm.proxy.management_helpers.access_group_key_sync import (
sync_key_access_group_membership,
sync_key_regeneration_access_group_membership,
sync_key_update_access_group_membership,
)
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
from litellm.proxy.management_helpers.object_permission_utils import (
_set_object_permission,
@ -888,17 +893,24 @@ async def _common_key_generation_helper(
if litellm.default_key_generate_params is not None:
for elem in data:
key, value = elem
if value is None and key in [
"max_budget",
"user_id",
"team_id",
"max_parallel_requests",
"tpm_limit",
"rpm_limit",
"budget_duration",
"duration",
]:
setattr(data, key, litellm.default_key_generate_params.get(key, None))
if (
value is None
and (key != "budget_duration" or key not in data.model_fields_set)
and key
in [
"max_budget",
"user_id",
"team_id",
"max_parallel_requests",
"tpm_limit",
"rpm_limit",
"budget_duration",
"duration",
]
):
default_value = litellm.default_key_generate_params.get(key)
if default_value is not None:
setattr(data, key, default_value)
elif key == "models" and value == []:
setattr(data, key, litellm.default_key_generate_params.get(key, []))
elif key == "metadata" and value == {}:
@ -2340,6 +2352,17 @@ async def _process_single_key_update(
proxy_logging_obj=proxy_logging_obj,
)
# After the key's own cache entry is dropped, so a failure here cannot leave the key
# authenticating against the access groups it just lost.
await sync_key_update_access_group_membership(
prisma_client=prisma_client,
key_token=_hash_token_if_needed(
_resolve_token_to_update(data=update_key_request, existing_key_row=existing_key_row)
),
data=update_key_request,
existing_key_row=existing_key_row,
)
# Trigger async hook
asyncio.create_task(
KeyManagementEventHooks.async_key_updated_hook(
@ -2821,6 +2844,15 @@ async def update_key_fn(
proxy_logging_obj=proxy_logging_obj,
)
# After the key's own cache entry is dropped, so a failure here cannot leave the key
# authenticating against the access groups it just lost.
await sync_key_update_access_group_membership(
prisma_client=prisma_client,
key_token=_hash_token_if_needed(key),
data=data,
existing_key_row=existing_key_row,
)
if data.spend is not None:
from litellm.proxy.proxy_server import spend_counter_cache
@ -3764,7 +3796,7 @@ async def generate_key_helper_fn(
auto_rotate: bool | None = None,
rotation_interval: str | None = None,
router_settings: dict | None = None,
access_group_ids: list | None = None,
access_group_ids: list[str] | None = None,
budget_limits: list | None = None, # multiple concurrent budget windows
):
from litellm.proxy.proxy_server import premium_user, prisma_client
@ -3972,6 +4004,14 @@ async def generate_key_helper_fn(
create_key_response: Final = await prisma_client.insert_data(data=key_data, table_name="key")
key_data["token_id"] = getattr(create_key_response, "token", None)
created_token_hash: Final = getattr(create_key_response, "token", None)
if isinstance(created_token_hash, str):
await sync_key_access_group_membership(
prisma_client=prisma_client,
key_token=created_token_hash,
previous_access_group_ids=None,
updated_access_group_ids=access_group_ids,
)
key_data["litellm_budget_table"] = getattr(create_key_response, "litellm_budget_table", None)
key_data["created_at"] = getattr(create_key_response, "created_at", None)
key_data["updated_at"] = getattr(create_key_response, "updated_at", None)
@ -4189,6 +4229,7 @@ async def delete_verification_tokens(
deleted_tokens = [key.token for key in authorized_keys]
if len(deleted_tokens) != len(tokens):
failed_tokens = [token for token in tokens if token not in deleted_tokens]
else:
raise Exception("DB not connected. prisma_client is None")
except Exception as e:
@ -4204,6 +4245,16 @@ async def delete_verification_tokens(
hashed_token = hash_token(cast(str, key))
user_api_key_cache.delete_cache(hashed_token)
# After credential invalidation, so a failure here can never keep a deleted key alive.
for deleted_key in authorized_keys:
if deleted_key.token is not None:
await sync_key_access_group_membership(
prisma_client=prisma_client,
key_token=deleted_key.token,
previous_access_group_ids=deleted_key.access_group_ids,
updated_access_group_ids=None,
)
return {
"deleted_keys": deleted_tokens,
"failed_tokens": failed_tokens,
@ -4719,6 +4770,15 @@ async def _execute_virtual_key_regeneration(
proxy_logging_obj=proxy_logging_obj,
)
# After credential invalidation, so a failure here can never keep the old key alive.
await sync_key_regeneration_access_group_membership(
prisma_client=prisma_client,
previous_key_token=hashed_api_key,
new_key_token=new_token_hash,
data=data,
existing_key_row=key_in_db,
)
response: Final = GenerateKeyResponse.model_validate(updated_token_dict)
asyncio.create_task(
KeyManagementEventHooks.async_key_rotated_hook(

109
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -22,7 +22,7 @@ import os
from collections.abc import Iterable
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any, Final, Literal
from typing import TYPE_CHECKING, Any, Final, Literal
from fastapi import (
APIRouter,
@ -50,6 +50,7 @@ from litellm.constants import LITELLM_PROXY_ADMIN_NAME
from litellm.proxy._experimental.mcp_server.utils import (
LITELLM_MCP_SERVER_DESCRIPTION,
LITELLM_MCP_SERVER_NAME,
McpServerPayloadLike,
build_env_var_setup_url,
collect_env_var_references,
get_server_prefix,
@ -91,6 +92,9 @@ def does_mcp_server_exist(mcp_server_records: Iterable[Any], mcp_server_id: str)
DEFAULT_MCP_REGISTRY_VERSION: Final = "1.0.0"
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
try:
importlib.import_module("mcp")
except ImportError as e:
@ -114,11 +118,13 @@ if MCP_AVAILABLE:
from litellm.proxy._experimental.mcp_server.db import (
approve_mcp_server,
create_draft_mcp_server,
create_mcp_server,
delete_mcp_server,
delete_user_credential,
delete_user_env_vars,
get_all_mcp_servers_for_user,
get_draft_mcp_server,
get_mcp_server,
get_mcp_servers,
get_mcp_submissions,
@ -196,7 +202,7 @@ if MCP_AVAILABLE:
server: MCPServer
expires_at: datetime
def _validate_mcp_server_name_fields(payload: Any) -> None:
def _validate_mcp_server_name_fields(payload: McpServerPayloadLike) -> None:
candidates: Final[list[tuple[str, str | None]]] = []
server_name: Final = getattr(payload, "server_name", None)
@ -223,7 +229,7 @@ if MCP_AVAILABLE:
detail={"error": error_messages_text},
)
def validate_and_normalize_mcp_server_payload(payload: Any) -> None:
def validate_and_normalize_mcp_server_payload(payload: McpServerPayloadLike) -> None:
_base_validate_and_normalize_mcp_server_payload(payload)
_validate_mcp_server_name_fields(payload)
@ -466,19 +472,68 @@ if MCP_AVAILABLE:
verbose_proxy_logger.debug("Invalid temporary MCP server payload in Redis cache: %s", e)
return None
def _get_prisma_client_or_none() -> "PrismaClient | None":
"""Non-throwing counterpart to ``get_prisma_client_or_throw`` for paths that degrade
gracefully: a proxy configured without a database keeps the in-memory OAuth session."""
from litellm.proxy.proxy_server import prisma_client
return prisma_client
async def _persist_draft_mcp_server(
payload: NewMCPServerRequest,
server_id: str,
created_by: str,
) -> None:
"""Write the draft row that makes the OAuth session resolvable from any worker.
A failure here is raised, not swallowed: without the shared row the flow degrades to
the per-process cache and fails intermittently, which is the defect being fixed.
"""
prisma_client: Final = _get_prisma_client_or_none()
if prisma_client is None:
return
await create_draft_mcp_server(
prisma_client,
payload,
created_by,
ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS,
server_id=server_id,
)
async def _get_draft_mcp_server_as_mcp_server(server_id: str) -> MCPServer | None:
"""Resolve a database-backed draft, which is the only lookup that works across workers."""
prisma_client: Final = _get_prisma_client_or_none()
if prisma_client is None:
return None
draft: Final = await get_draft_mcp_server(
prisma_client, server_id, ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS
)
if draft is None:
return None
return await global_mcp_server_manager.build_mcp_server_from_table(draft)
async def get_cached_temporary_mcp_server(
server_id: str,
) -> MCPServer | None:
_prune_expired_temporary_mcp_servers()
entry: Final = _temporary_mcp_servers.get(server_id)
if entry is None:
redis_server: Final = await _get_temporary_mcp_server_from_redis(server_id)
if redis_server is None:
return None
# Intentionally avoid repopulating local cache from Redis to prevent
# extending effective lifetime beyond the remaining Redis TTL.
return redis_server
return entry.server
if entry is not None:
return entry.server
# A miss here means either an expired session or, on a multi-worker or multi-replica
# proxy, that a different process served /session. The draft row is shared, so it
# resolves the second case; the in-memory hit above still serves single-process
# deployments with no database configured.
draft_server: Final = await _get_draft_mcp_server_as_mcp_server(server_id)
if draft_server is not None:
return draft_server
redis_server: Final = await _get_temporary_mcp_server_from_redis(server_id)
if redis_server is None:
return None
# Intentionally avoid repopulating local cache from Redis to prevent
# extending effective lifetime beyond the remaining Redis TTL.
return redis_server
def _redact_mcp_credentials(
mcp_server: LiteLLM_MCPServerTable,
@ -708,12 +763,36 @@ if MCP_AVAILABLE:
payload_dict["credentials"] = inherited_credentials
return NewMCPServerRequest.model_validate(payload_dict)
async def _resolve_session_server_id(payload: NewMCPServerRequest) -> str:
"""Decide the id an OAuth session runs under.
A caller-supplied id is honoured only when it names a server that really exists, which is
the edit form re-authorizing a saved server against its own id. Anything else gets a fresh
id, so two concurrent sessions can never land on one id and silently adopt each other's
URL or client credentials. Without a database there is nothing shared to collide over, so
the supplied id is kept and behaviour is unchanged.
"""
supplied: Final = payload.server_id
if not supplied:
return str(uuid.uuid4())
if global_mcp_server_manager.get_mcp_server_by_id(supplied) is not None:
return supplied
prisma_client: Final = _get_prisma_client_or_none()
if prisma_client is None:
return supplied
# A draft is another session's row, not a saved server, so re-supplying an id this
# endpoint previously handed back must not let a later session adopt its configuration.
existing: Final = await get_mcp_server(prisma_client, supplied)
if existing is None or existing.approval_status == MCPApprovalStatus.draft:
return str(uuid.uuid4())
return supplied
def _build_temporary_mcp_server_record(
payload: NewMCPServerRequest,
created_by: str | None,
server_id: str,
) -> LiteLLM_MCPServerTable:
now: Final = datetime.utcnow()
server_id: Final = payload.server_id or str(uuid.uuid4())
server_name: Final = payload.server_name or payload.alias or server_id
return LiteLLM_MCPServerTable(
server_id=server_id,
@ -1543,6 +1622,7 @@ if MCP_AVAILABLE:
temp_record: Final = _build_temporary_mcp_server_record(
payload_with_credentials,
created_by,
await _resolve_session_server_id(payload_with_credentials),
)
try:
@ -1554,6 +1634,11 @@ if MCP_AVAILABLE:
temporary_server,
ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS,
)
await _persist_draft_mcp_server(
payload_with_credentials,
temp_record.server_id,
created_by,
)
await _cache_temporary_mcp_server_in_redis(
temporary_server,
ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS,

389
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -35,6 +35,7 @@ from litellm.proxy._types import (
PrismaCompatibleUpdateDBModel,
ProxyErrorTypes,
ProxyException,
ReconcileOutcome,
TeamModelAddRequest,
TeamModelDeleteRequest,
UserAPIKeyAuth,
@ -67,6 +68,7 @@ from litellm.repositories.team_repository import TeamRepository
from litellm.router import Router
from litellm.router_strategy.complexity_router import (
DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE,
ClassificationRubric,
ComplexityRouterConfig,
ComplexityTier,
classification_system_prompt,
@ -87,6 +89,7 @@ from litellm.types.router import (
ModelInfo,
updateDeployment,
)
from litellm.types.utils import CustomPricingLiteLLMParams
from litellm.utils import get_utc_datetime
router: Final = APIRouter()
@ -240,6 +243,7 @@ def _raise_on_strategy_router_write_violation(
_PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effective_from", "ptu_effective_to")
_PTU_PRICED_PAIR: Final = frozenset({"ptu_count", "cost_per_ptu_per_hour"})
def _explicitly_cleared_ptu_fields(model_info: ModelInfo | None) -> frozenset[str]:
@ -263,9 +267,10 @@ def _merged_ptu_model_info(*, db_model: Deployment, patch_data: updateDeployment
A PTU invariant holds over the deployment as it will exist, not over whichever subset
of fields a caller happened to send.
"""
empty: Final[Mapping[str, object]] = MappingProxyType({})
stored: Final = db_model.model_info.model_dump(exclude_none=True) if db_model.model_info else empty
incoming: Final = patch_data.model_info.model_dump(exclude_none=True) if patch_data.model_info else empty
stored: Final = db_model.model_info.model_dump(exclude_none=True) if db_model.model_info else _EMPTY_MODEL_INFO
incoming: Final = (
patch_data.model_info.model_dump(exclude_none=True) if patch_data.model_info else _EMPTY_MODEL_INFO
)
cleared: Final = _explicitly_cleared_ptu_fields(patch_data.model_info)
return MappingProxyType({k: v for k, v in {**stored, **incoming}.items() if k not in cleared})
@ -337,6 +342,140 @@ def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None:
)
# The six mirrored pricing fields plus the three remaining fields
# Router._inherit_builtin_cache_pricing back-fills from the public cost map. An unset field is
# what that back-fill targets, so a field left out here is one a PTU deployment still bills.
_PTU_ZEROED_PRICING_FIELDS: Final = SPECIAL_MODEL_INFO_PARAMS + (
"cache_creation_input_token_cost_above_1hr",
"cache_creation_input_token_cost_above_200k_tokens",
"cache_read_input_token_cost_above_200k_tokens",
)
_PTU_ZEROED_PRICING: Final[Mapping[str, float]] = MappingProxyType(dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0))
_NO_PRICING_OVERRIDE: Final[Mapping[str, float]] = MappingProxyType({})
_EMPTY_MODEL_INFO: Final[Mapping[str, object]] = _NO_PRICING_OVERRIDE
# Rate fields only. CustomPricingLiteLLMParams also carries settings that are not charges
# (an embedding's output_vector_size, the regional uplift multipliers), and zeroing one of
# those would destroy the deployment's configuration rather than stop a charge.
_CUSTOM_PRICING_FIELDS: Final = frozenset(f for f in CustomPricingLiteLLMParams.model_fields if "cost" in f)
def _is_nonzero_price(value: object) -> bool:
return isinstance(value, (int, float)) and not isinstance(value, bool) and value != 0
def _is_zero_price(value: object) -> bool:
return isinstance(value, (int, float)) and not isinstance(value, bool) and value == 0
def _raise_if_ptu_deployment_is_priced(*, model_info: Mapping[str, object], supplied: Mapping[str, object]) -> None:
"""Refuse a rate the caller supplies for a deployment that bills reserved capacity.
Separate from the zeroing so the team-model path can run it before it touches the team, whose
ACL write autocommits: a refusal raised after it would leave the team changed and the
deployment row never written.
"""
if not is_ptu_cost_attribution_enabled():
return
if model_info.get("ptu_count") is None or model_info.get("cost_per_ptu_per_hour") is None:
return
priced: Final = tuple(sorted(field for field in _CUSTOM_PRICING_FIELDS if _is_nonzero_price(supplied.get(field))))
if not priced:
return
raise HTTPException(
status_code=400,
detail=(
f"A PTU deployment bills by reserved capacity, so {', '.join(priced)} cannot be charged on "
"top of it. Send 0 or no value, or remove ptu_count and cost_per_ptu_per_hour to bill per token."
),
)
def _ptu_zeroed_pricing(
*,
model_info: Mapping[str, object],
litellm_params: Mapping[str, object],
supplied: Mapping[str, object],
) -> Mapping[str, float]:
"""The pricing a PTU deployment must carry, empty unless one is being stored.
Reserved capacity is already billed by the flat cost the rollup writes, so charging the
traffic it serves bills the same tokens twice. Left unset the rate falls back to the public
cost map, which makes the double charge the default rather than an opt-in.
Only a price the caller supplies is refused. A non-zero price already on the row is zeroed
instead, so a deployment priced through a path this rule does not cover heals on its next
save rather than rejecting every later edit of a field that has nothing to do with pricing.
``supplied`` is the caller's litellm_params alone, because that is the blob a price is
authored on. model_info's copy is written by the server, both by the mirror in
``Deployment.__init__`` and by the cost-map defaults /model/info fills in, so a client that
round-trips a model_info blob sends back prices it never chose.
"""
if not is_ptu_cost_attribution_enabled():
return _NO_PRICING_OVERRIDE
if model_info.get("ptu_count") is None or model_info.get("cost_per_ptu_per_hour") is None:
return _NO_PRICING_OVERRIDE
_raise_if_ptu_deployment_is_priced(model_info=model_info, supplied=supplied)
stored: Final = frozenset(
field
for field in _CUSTOM_PRICING_FIELDS
if _is_nonzero_price(model_info.get(field)) or _is_nonzero_price(litellm_params.get(field))
)
if not stored:
return _PTU_ZEROED_PRICING
return MappingProxyType({**_PTU_ZEROED_PRICING, **dict.fromkeys(stored, 0.0)})
def _ptu_pricing_delta(
*,
stored_model_info: Mapping[str, object],
model_info: Mapping[str, object],
litellm_params: Mapping[str, object],
patch: updateDeployment,
) -> tuple[Mapping[str, float], frozenset[str]]:
"""The pricing a patch must write into both blobs, and the pricing it must drop from them.
A patch that takes the deployment off PTU takes the zeroed pricing with it, since the zeros
exist only to stop the double charge. Left behind they would serve the deployment for free.
Reading the stored row rather than the patch alone keeps that release off a deployment that
never carried PTU config, whose zero price is a rate its operator chose. A zero the patch
itself carries is released with the rest, because the dashboard echoes the whole stored
blob on every save, so a supplied zero cannot be told apart from the one this rule wrote.
The release spans every field the zeroing could have written, not just the mirrored ones, or
a rate zeroed on the way in (per-second, per-character tiers) would bill nothing forever.
"""
supplied: Final = patch.litellm_params.model_dump(exclude_none=True) if patch.litellm_params else _EMPTY_MODEL_INFO
zeroed: Final = _ptu_zeroed_pricing(model_info=model_info, litellm_params=litellm_params, supplied=supplied)
if zeroed:
return zeroed, frozenset()
was_ptu: Final = any(stored_model_info.get(field) is not None for field in _PTU_PRICED_PAIR)
if not was_ptu or not _explicitly_cleared_ptu_fields(patch.model_info) & _PTU_PRICED_PAIR:
return _NO_PRICING_OVERRIDE, frozenset()
return _NO_PRICING_OVERRIDE, frozenset(
field
for field in _CUSTOM_PRICING_FIELDS.union(_PTU_ZEROED_PRICING_FIELDS)
if _is_zero_price(model_info.get(field)) or _is_zero_price(litellm_params.get(field))
)
def _ptu_priced_deployment(model_params: Deployment) -> Deployment:
"""``model_params`` with PTU pricing applied, or itself when it configures no PTU."""
model_info: Final = model_params.model_info.model_dump(exclude_none=True)
litellm_params: Final = model_params.litellm_params.model_dump(exclude_none=True)
override: Final = _ptu_zeroed_pricing(model_info=model_info, litellm_params=litellm_params, supplied=litellm_params)
if not override:
return model_params
return model_params.model_copy(
update=MappingProxyType(
{
"litellm_params": model_params.litellm_params.model_copy(update=override),
"model_info": model_params.model_info.model_copy(update=override),
}
)
)
def _parse_ptu_datetime(value: object) -> datetime.datetime | None:
"""``value`` as a datetime, parsing an ISO string, else None."""
if isinstance(value, datetime.datetime):
@ -402,6 +541,19 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
merged_model_info.pop(field, None)
_validate_ptu_model_info(merged_model_info)
ptu_pricing, ptu_released = _ptu_pricing_delta(
stored_model_info=db_model.model_info.model_dump(exclude_none=True)
if db_model.model_info
else _EMPTY_MODEL_INFO,
model_info=merged_model_info,
litellm_params=merged_litellm_params,
patch=updated_patch,
)
merged_model_info.update(ptu_pricing)
merged_litellm_params.update(ptu_pricing)
for field in ptu_released:
merged_model_info.pop(field, None)
merged_litellm_params.pop(field, None)
# convert to prisma compatible format
@ -534,7 +686,7 @@ async def patch_model(
# Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates)
live_before_reload: Final = live_model_ids_snapshot()
still_desired_ids: Final = await clear_cache()
reload_outcome: Final = await clear_cache()
## CREATE AUDIT LOG ##
asyncio.create_task(
@ -554,7 +706,8 @@ async def patch_model(
before=live_before_reload,
written_models=[(model_id, getattr(updated_model, "model_info", None))],
action="update",
still_desired=still_desired_ids,
still_desired=reload_outcome.still_desired,
live_after=reload_outcome.live_after,
)
return updated_model
@ -640,7 +793,7 @@ async def _set_model_blocked_status(
)
live_before_reload: Final = live_model_ids_snapshot()
still_desired_ids: Final = await clear_cache()
reload_outcome: Final = await clear_cache()
asyncio.create_task(
create_object_audit_log(
@ -661,7 +814,8 @@ async def _set_model_blocked_status(
before=live_before_reload,
written_models=[(data.model_id, getattr(updated_model, "model_info", None))],
action=action,
still_desired=still_desired_ids,
still_desired=reload_outcome.still_desired,
live_after=reload_outcome.live_after,
)
return updated_model
@ -859,6 +1013,12 @@ async def _update_team_model_in_db(
if patch_data.model_info is not None:
_raise_if_ptu_cost_attribution_disabled(patch_data.model_info.model_dump(exclude_none=True))
_validate_ptu_model_info(_merged_ptu_model_info(db_model=db_model, patch_data=patch_data))
_raise_if_ptu_deployment_is_priced(
model_info=_merged_ptu_model_info(db_model=db_model, patch_data=patch_data),
supplied=(
patch_data.litellm_params.model_dump(exclude_none=True) if patch_data.litellm_params else _EMPTY_MODEL_INFO
),
)
patch_team_id: Final = patch_data.model_info.team_id if patch_data.model_info else None
@ -1033,9 +1193,15 @@ async def delete_team_models(
if deleted_model_ids:
await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable")
# Under MODEL_RECONCILE_LOCK, for the same reason as delete_model: the rows are
# gone, but a reconcile holding a pre-delete snapshot would upsert these ids back
# onto this pod. The lock orders the eviction after any in-flight reconcile.
if llm_router is not None:
for model_id in deleted_model_ids:
llm_router.delete_deployment(id=model_id)
from litellm.proxy.proxy_server import MODEL_RECONCILE_LOCK
async with MODEL_RECONCILE_LOCK:
for model_id in deleted_model_ids:
llm_router.delete_deployment(id=model_id)
return deleted_model_ids
@ -1355,6 +1521,7 @@ async def delete_model(
"""
from litellm.proxy.proxy_server import (
MODEL_RECONCILE_LOCK,
llm_router,
premium_user,
prisma_client,
@ -1403,8 +1570,15 @@ async def delete_model(
)
## DELETE FROM ROUTER ##
# Under MODEL_RECONCILE_LOCK. The db row is already gone, but a reconcile
# that snapshotted the db BEFORE that delete still lists this id as desired,
# and its _add_deployment upserts the deployment straight back -- leaving
# this pod serving a model the database no longer has, until the next
# reconcile. Taking the lock orders this eviction after any such in-flight
# reconcile's re-add, so the eviction is the last word.
if llm_router is not None:
llm_router.delete_deployment(id=model_info.id)
async with MODEL_RECONCILE_LOCK:
llm_router.delete_deployment(id=model_info.id)
# Runs after the row delete so the sibling check sees post-delete state.
if model_params.model_info.team_id is not None:
@ -1571,6 +1745,7 @@ async def add_new_model(
incoming_model_info: Final = model_params.model_info.model_dump(exclude_none=True)
_raise_if_ptu_cost_attribution_disabled(incoming_model_info)
_validate_ptu_model_info(incoming_model_info)
priced_model_params: Final = _ptu_priced_deployment(model_params)
if store_model_in_db is True:
"""
@ -1579,22 +1754,22 @@ async def add_new_model(
"""
live_before_reload: Final = live_model_ids_snapshot()
still_desired_ids: frozenset[str] | None = None
reload_outcome: ReconcileOutcome = ReconcileOutcome(still_desired=None, live_after=None)
try:
_original_litellm_model_name: Final = model_params.model_name
if model_params.model_info.team_id is None:
model_response = await _add_model_to_db(
model_params=model_params,
model_params=priced_model_params,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
else:
model_response = await _add_team_model_to_db(
model_params=model_params,
model_params=priced_model_params,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
still_desired_ids = await proxy_config.add_deployment(
reload_outcome = await proxy_config.add_deployment(
prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
)
# don't let failed slack alert block the /model/new response
@ -1602,9 +1777,9 @@ async def add_new_model(
if "slack" in _alerting:
# send notification - new model added
await proxy_logging_obj.slack_alerting_instance.model_added_alert(
model_name=model_params.model_name,
model_name=priced_model_params.model_name,
litellm_model_name=_original_litellm_model_name,
passed_model_info=model_params.model_info,
passed_model_info=priced_model_params.model_info,
)
except Exception as e:
verbose_proxy_logger.exception("Exception in add_new_model: %s", e)
@ -1641,7 +1816,8 @@ async def add_new_model(
before=live_before_reload,
written_models=[(model_response.model_id, getattr(model_response, "model_info", None))],
action="create",
still_desired=still_desired_ids,
still_desired=reload_outcome.still_desired,
live_after=reload_outcome.live_after,
)
return model_response
@ -1768,7 +1944,7 @@ async def update_model(
# Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates)
live_before_reload: Final = live_model_ids_snapshot()
still_desired_ids: Final = await clear_cache()
reload_outcome: Final = await clear_cache()
## CREATE AUDIT LOG ##
asyncio.create_task(
create_object_audit_log(
@ -1795,7 +1971,8 @@ async def update_model(
before=live_before_reload,
written_models=[(_model_id, getattr(model_response, "model_info", None))],
action="update",
still_desired=still_desired_ids,
still_desired=reload_outcome.still_desired,
live_after=reload_outcome.live_after,
)
return model_response
@ -2006,19 +2183,23 @@ def _labeled_tiers_from_query(tier_labels: str | None) -> tuple[tuple[Complexity
async def get_auto_router_classifier_default_prompt(
context_window_size: int = DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE,
tier_labels: str | None = None,
classification_rubric: ClassificationRubric | None = None,
) -> AutoRouterClassifierDefaultPromptResponse:
"""
Get the default classifier system prompt, so the dashboard's prompt editor can prefill it.
The prompt's closing line depends on whether prior conversation turns are quoted to the
classifier, and its tier bullets are named by the router's tier_labels, so the caller passes both
to get the text that router would actually send rather than a rubric it does not use.
classifier, its tier bullets are named by the router's tier_labels, and its calibration examples
come from the router's classification rubric, so the caller passes all three to get the text that router
would actually send rather than a rubric it does not use.
Parameters:
- context_window_size: int - The router's classifier_context_window_size. Defaults to the
built-in default.
- tier_labels: str | None - The router's tier_labels as a JSON object of canonical tier name to
display name, e.g. `{"SIMPLE": "Cheap"}`. Omit or pass an empty object for the default names.
- classification_rubric: ClassificationRubric | None - The router's
classifier_llm_config.classification_rubric. Omit for the default.
"""
if context_window_size < 0:
raise ProxyException(
@ -2031,9 +2212,11 @@ async def get_auto_router_classifier_default_prompt(
labeled_tiers: Final = _labeled_tiers_from_query(tier_labels)
return AutoRouterClassifierDefaultPromptResponse(
system_prompt=(
classification_system_prompt(context_window_size)
classification_system_prompt(context_window_size, classification_rubric=classification_rubric)
if labeled_tiers is None
else classification_system_prompt(context_window_size, labeled_tiers=labeled_tiers)
else classification_system_prompt(
context_window_size, labeled_tiers=labeled_tiers, classification_rubric=classification_rubric
)
)
)
@ -2100,6 +2283,7 @@ def reload_serving_verdict(
written_models: Sequence[tuple[str, object]],
written_must_serve: bool,
still_desired: frozenset[str] | None = None,
live_after: frozenset[str] | None = None,
) -> tuple[tuple[str, ...], tuple[str, ...]]:
"""Judge a write-triggered reload by diffing the router's serving state instead of
trusting any layer of the reload stack to report its own failure.
@ -2121,9 +2305,16 @@ def reload_serving_verdict(
yet polled, so the reload dropping it is the reconcile working rather than damage.
Without it (no reconcile ran) every drop is reported, which is the safe direction.
``live_after`` is the router's serving state captured by the reload itself, while it
still held MODEL_RECONCILE_LOCK. Pass it whenever the caller has it: re-reading the
router here instead means sampling it after the lock was released, where the NEXT
reconcile's leading wipe (clear_cache un-serves every db model before reloading
them) shows up as this reload having dropped them. Falling back to a fresh read is
only correct when no reconcile ran and there is nothing to be concurrent with.
Returns (written ids violating their obligation, collateral ids no longer served).
"""
now: Final = live_model_ids_snapshot()
now: Final = live_model_ids_snapshot() if live_after is None else live_after
written_ids: Final = frozenset(model_id for model_id, _ in written_models)
if written_must_serve:
missing = tuple(
@ -2143,16 +2334,23 @@ def raise_if_reload_degraded_serving(
written_models: Sequence[tuple[str, object]],
action: str,
still_desired: frozenset[str] | None = None,
live_after: frozenset[str] | None = None,
) -> None:
"""The caller-visible error this pod's model-write endpoints owe their caller when
the model they wrote is not being served after the reload they triggered. The DB
write is durable either way and every other pod reloads on its own interval; this
speaks only for the handling pod."""
speaks only for the handling pod.
Callers hold a ReconcileOutcome from the reload; pass BOTH of its fields. Supplying
still_desired without live_after mixes a snapshot taken under the reconcile lock
with one taken after it was released, which is what makes a concurrent model write
look like collateral damage."""
missing, collateral = reload_serving_verdict(
before=before,
written_models=written_models,
written_must_serve=True,
still_desired=still_desired,
live_after=live_after,
)
if not missing and not collateral:
return
@ -2179,14 +2377,20 @@ def raise_if_reload_degraded_serving(
)
async def clear_cache() -> frozenset[str] | None:
async def clear_cache() -> ReconcileOutcome:
"""
Clear router caches and reload models.
Returns the db + config id set the reload reconciled against, or None when no
reload ran, so callers can pass it to raise_if_reload_degraded_serving.
Returns what the reload saw (see ReconcileOutcome) so callers can pass it to
raise_if_reload_degraded_serving.
Runs under MODEL_RECONCILE_LOCK for its whole extent, not just the reload at the
end, so the auto-router reset and the reload that rebuilds those routers are atomic
to any other reconcile. The inner call is _add_deployment_locked because
add_deployment would re-acquire the same non-reentrant lock and deadlock.
"""
from litellm.proxy.proxy_server import (
MODEL_RECONCILE_LOCK,
llm_router,
prisma_client,
proxy_config,
@ -2196,61 +2400,88 @@ async def clear_cache() -> frozenset[str] | None:
if llm_router is None or prisma_client is None:
verbose_proxy_logger.debug("llm_router or prisma_client is None, skipping cache clear")
return None
return ReconcileOutcome(still_desired=None, live_after=None)
try:
# Only clear DB models, preserve config models
verbose_proxy_logger.debug("Clearing only DB models, preserving config models")
async with MODEL_RECONCILE_LOCK:
try:
# Only clear DB models, preserve config models
verbose_proxy_logger.debug("Clearing only DB models, preserving config models")
# Get current models and filter out DB models
current_models: Final = llm_router.model_list.copy()
config_models: Final = []
db_model_ids: Final = []
# Get current models and filter out DB models
current_models: Final = llm_router.model_list.copy()
config_models: Final = []
db_model_ids: Final = []
for model in current_models:
model_info = model.get("model_info", {})
if model_info.get("db_model", False):
# This is a DB model, mark for deletion
db_model_ids.append(model_info.get("id"))
else:
# This is a config model, preserve it
config_models.append(model)
db_router_names: Final = set()
# Clear only DB models
for model_id in db_model_ids:
llm_router.delete_deployment(id=model_id)
for model in current_models:
model_info = model.get("model_info", {})
if model_info.get("db_model", False):
db_model_ids.append(model_info.get("id"))
# Auto-router deployments (and only those) are wiped here, in the
# same pass, so the reload rebuilds them -- see the comment below.
model_name = model.get("model_name")
if model_name is not None and str(model.get("litellm_params", {}).get("model", "")).startswith(
"auto_router/"
):
db_router_names.add(model_name)
router_model_id = model_info.get("id")
if router_model_id is not None:
llm_router.delete_deployment(id=router_model_id)
else:
# This is a config model, preserved by the reconcile below
config_models.append(model)
# Clear only DB-backed auto-router-family entries, keyed by model_name, so the
# reload below rebuilds them fresh. A blanket .clear() would also drop config-defined
# routers, which are never re-added below (add_deployment only reloads DB models),
# leaving them permanently unroutable until a full proxy restart for every tenant.
# Restrict to deployments whose model is actually an auto_router/* so a config
# router that merely shares a model_name with a regular DB model isn't evicted. The
# auto_router/ prefix also covers quality_router/ and adaptive_router/, so pop the
# name from every router registry (no-op where absent); missing quality/adaptive
# entries would otherwise make init raise "already exists" on reload and abort it.
db_router_names: Final = {
model.get("model_name")
for model in current_models
if model.get("model_name") is not None
and model.get("model_info", {}).get("db_model", False)
and str(model.get("litellm_params", {}).get("model", "")).startswith("auto_router/")
}
for model_name in db_router_names:
llm_router.auto_routers.pop(model_name, None)
llm_router.complexity_routers.pop(model_name, None)
llm_router.adaptive_routers.pop(model_name, None)
llm_router.quality_routers.pop(model_name, None)
# ORDINARY db deployments are deliberately NOT wiped. This used to
# delete_deployment() every db model before the reload put them back, which
# left the router serving ZERO db models for the whole width of the reload
# -- a real data-plane hole that every inference request landing in it fell
# into. It was also redundant for them: the reload's _delete_deployment
# evicts exactly the ids the db no longer lists, and upsert_deployment
# pops-and-re-adds a deployment whose params changed while no-opping one
# that did not, so the reconcile converges on its own. Every mutation is
# visible to that comparison -- `blocked` and (for premium) `updated_at`
# are written into model_info.
#
# AUTO-ROUTER db deployments are the exception and ARE wiped -- in the
# classification pass above, together with the strategy entries popped
# just below. Their strategy registries are keyed
# by model_name, which no deployment-id reconcile touches, so they have to
# be popped and rebuilt here. But the rebuild only happens on the ADD path:
# Router.upsert_deployment returns early when a deployment is unchanged and
# never reaches add_deployment -> _add_deployment ->
# init_auto_router_deployment, which is what repopulates the registries.
# Popping without deleting would therefore strip every db-backed auto,
# complexity, adaptive and quality router on this pod and never put it back,
# so ANY unrelated model write would leave them unroutable until a restart.
# Deleting the deployment forces upsert down the add path, which rebuilds
# both the deployment and its strategy entry.
#
# That pass restricts the wipe to deployments whose model is actually an
# auto_router/* so a config router that merely shares a model_name with a
# regular db model isn't evicted -- config routers are never re-added by the
# reload (it only reloads db models) and would be permanently unroutable.
# The auto_router/ prefix also covers quality_router/ and adaptive_router/,
# so pop the name from every registry (no-op where absent); a missing
# quality/adaptive entry would otherwise make init raise "already exists"
# on reload and abort it.
for model_name in db_router_names:
llm_router.auto_routers.pop(model_name, None)
llm_router.complexity_routers.pop(model_name, None)
llm_router.adaptive_routers.pop(model_name, None)
llm_router.quality_routers.pop(model_name, None)
# Reload only DB models
still_desired_ids: Final = await proxy_config.add_deployment(
prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
)
# Reload only DB models. _add_deployment_locked, not add_deployment: this
# coroutine already holds MODEL_RECONCILE_LOCK and asyncio.Lock is not
# reentrant, so the public wrapper would deadlock against itself.
outcome: Final = await proxy_config._add_deployment_locked(
prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
)
verbose_proxy_logger.debug(
"Cleared %s DB models, preserved %s config models", len(db_model_ids), len(config_models)
)
return still_desired_ids
except Exception as e:
verbose_proxy_logger.exception("Failed to clear cache and reload models. Due to error - %s", e)
return None
verbose_proxy_logger.debug(
"Reconciled %s DB models, preserved %s config models", len(db_model_ids), len(config_models)
)
return outcome
except Exception as e:
verbose_proxy_logger.exception("Failed to clear cache and reload models. Due to error - %s", e)
return ReconcileOutcome(still_desired=None, live_after=None)

267
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -15,6 +15,7 @@ import math
import traceback
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Annotated, Final, Protocol, TypedDict, TypeVar, cast
import fastapi
@ -77,6 +78,8 @@ from litellm.proxy.auth.auth_checks import (
_cache_team_object,
allowed_route_check_inside_route,
can_org_access_model,
delete_cache_key_objects,
delete_cache_team_object,
get_org_object,
get_team_membership,
get_team_object,
@ -104,6 +107,12 @@ from litellm.proxy.management_endpoints.organization_endpoints import (
from litellm.proxy.management_endpoints.tag_management_endpoints import (
get_daily_activity,
)
from litellm.proxy.management_helpers.access_group_team_sync import (
AccessGroupSyncTx,
invalidate_access_group_caches,
reconcile_team_access_group_membership,
sync_team_access_group_membership,
)
from litellm.proxy.management_helpers.object_permission_utils import (
_set_object_permission,
enforce_all_proxy_mcp_servers_grant_is_admin_only,
@ -313,6 +322,18 @@ class _TeamIdInFilter(TypedDict, total=False):
team_id: Mapping[str, Sequence[str]]
class _TeamCreateTx(AccessGroupSyncTx, Protocol):
@property
def litellm_teamtable(self) -> "_PrismaTableActions[LiteLLM_TeamTable]": ...
_STRIP_DELETED_TEAM_FROM_USERS_SQL: Final = """
UPDATE "LiteLLM_UserTable" SET teams = array_remove(teams, $1) WHERE $1 = ANY(teams)
"""
_INCLUDE_MODEL_TABLE: Final = MappingProxyType({"litellm_model_table": True})
def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]":
return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable)
@ -1313,8 +1334,9 @@ async def new_team(
if isinstance(default_organization_id, str):
data.organization_id = default_organization_id
# Apply defaults from litellm.default_team_params for any fields
# not explicitly provided in the request.
# Apply defaults from litellm.default_team_params to null fields.
# budget_duration alone distinguishes explicit null (a deliberate
# never-resetting budget, which the default must not override) from omitted.
for field in (
"max_budget",
"budget_duration",
@ -1322,7 +1344,9 @@ async def new_team(
"rpm_limit",
"team_member_permissions",
):
if getattr(data, field, None) is None:
if getattr(data, field, None) is None and (
field != "budget_duration" or field not in data.model_fields_set
):
default_value = _get_default_team_param(field)
if default_value is not None:
setattr(data, field, default_value)
@ -1501,10 +1525,15 @@ async def new_team(
complete_team_data_dict = prisma_client.jsonify_team_object(db_data=complete_team_data_dict)
team_creation_data: Final[Mapping[str, object]] = complete_team_data_dict
team_row: Final[LiteLLM_TeamTable] = await _team_db(prisma_client).create(
data=team_creation_data,
include={"litellm_model_table": True},
)
tx: _TeamCreateTx
async with prisma_client.db.tx() as tx:
team_row: Final[LiteLLM_TeamTable] = await tx.litellm_teamtable.create(
data=team_creation_data,
include=_INCLUDE_MODEL_TABLE,
)
affected_access_groups: Final = await reconcile_team_access_group_membership(tx, team_row.team_id)
await invalidate_access_group_caches(affected_access_groups)
## ADD TEAM ID TO USER TABLE ##
team_member_add_request: Final = TeamMemberAddRequest(
@ -2207,6 +2236,7 @@ async def update_team(
)
verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id)
await sync_team_access_group_membership(prisma_client=prisma_client, team_id=team_row.team_id)
await _refresh_cached_team(
team_row=team_row,
user_api_key_cache=user_api_key_cache,
@ -2654,6 +2684,11 @@ async def _add_team_members_to_team(
serialize on the row lock and each appends onto the other's committed
result, instead of both rewriting the whole JSON array from a stale
snapshot (which silently drops one member on the losing write).
The same lock serializes this against /team/delete: the delete cannot remove
the row while the reconcile holds it, and a reconcile that finds the row
already gone cleans up after itself rather than leaving the member pointing
at a deleted team id.
"""
# Process and add new members
updated_users, updated_team_memberships = await _process_team_members(
@ -2664,11 +2699,42 @@ async def _add_team_members_to_team(
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
async with prisma_client.tx() as tx:
complete_team_data.members_with_roles = await TeamRepository(prisma_client).get_members_with_roles_locked(
tx, data.team_id
updated_team: Final = await _write_members_with_roles_locked(
data=data,
complete_team_data=complete_team_data,
prisma_client=prisma_client,
updated_users=updated_users,
)
if updated_team is None:
await _sweep_deleted_team_references(team_ids=(data.team_id,), prisma_client=prisma_client)
raise HTTPException(
status_code=404,
detail={"error": f"Team={data.team_id} was deleted while this member add was running"},
)
return updated_team, updated_users, updated_team_memberships
async def _write_members_with_roles_locked(
data: TeamMemberAddRequest,
complete_team_data: LiteLLM_TeamTable,
prisma_client: PrismaClient,
updated_users: list[LiteLLM_UserTable],
) -> LiteLLM_TeamTable | None:
"""Reconcile members_with_roles under the team row lock. None when the team row is gone.
That read is at least as recent as the user and membership writes the caller
already made, so a missing row means /team/delete committed after them. Its
post-delete sweep can have run before those writes landed, which is why the
caller sweeps this team id again rather than only reporting the 404.
"""
async with prisma_client.tx() as tx:
locked_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, data.team_id)
if locked_members is None:
return None
complete_team_data.members_with_roles = locked_members
await _update_team_members_list(
data=data,
complete_team_data=complete_team_data,
@ -2676,13 +2742,11 @@ async def _add_team_members_to_team(
)
_db_team_members: Final = [m.model_dump() for m in complete_team_data.members_with_roles]
updated_team: Final = await tx.litellm_teamtable.update(
return await tx.litellm_teamtable.update(
where={"team_id": data.team_id},
data={"members_with_roles": json.dumps(_db_team_members)},
)
return updated_team, updated_users, updated_team_memberships
def _emit_team_members_metric(team: LiteLLM_TeamTable) -> None:
"""Update the Prometheus team members gauge after a membership change.
@ -3088,26 +3152,27 @@ async def team_member_add(
)
def _is_member_addressed_by(member: Member, data: TeamMemberDeleteRequest) -> bool:
return (data.user_id is not None and member.user_id is not None and data.user_id == member.user_id) or (
data.user_email is not None and member.user_email is not None and data.user_email == member.user_email
)
def _cleanup_members_with_roles(
existing_team_row: LiteLLM_TeamTable,
data: TeamMemberDeleteRequest,
) -> tuple[bool, list[Member]]:
"""Cleanup members_with_roles list for a team."""
is_member_in_team = False
new_team_members: Final[list[Member]] = []
for m in existing_team_row.members_with_roles:
if (
data.user_id is not None
and m.user_id is not None
and data.user_id == m.user_id
or data.user_email is not None
and m.user_email is not None
and data.user_email == m.user_email
):
is_member_in_team = True
continue
new_team_members.append(m)
return is_member_in_team, new_team_members
) -> tuple[tuple[Member, ...], list[Member]]:
"""Split a team's members_with_roles into the entries the request addresses and the ones that stay.
The addressed entries are returned rather than a bare found/not-found flag because they carry the
user_id the request may not have supplied, and every cleanup that keys off the user rather than
off the roster has to run against that id.
"""
removed_team_members: Final = tuple(
m for m in existing_team_row.members_with_roles if _is_member_addressed_by(m, data)
)
new_team_members: Final = [m for m in existing_team_row.members_with_roles if not _is_member_addressed_by(m, data)]
return removed_team_members, new_team_members
@router.post(
@ -3179,12 +3244,12 @@ async def team_member_delete(
)
## DELETE MEMBER FROM TEAM
is_member_in_team, new_team_members = _cleanup_members_with_roles(
removed_team_members, new_team_members = _cleanup_members_with_roles(
existing_team_row=existing_team_row,
data=data,
)
if not is_member_in_team:
if not removed_team_members:
raise HTTPException(status_code=400, detail={"error": "User not found in team"})
existing_team_row.members_with_roles = new_team_members
@ -3202,38 +3267,28 @@ async def team_member_delete(
## DELETE TEAM ID from USER ROW, IF EXISTS ##
# get user row
key_val: Final = {}
if data.user_id is not None:
key_val["user_id"] = data.user_id
elif data.user_email is not None:
key_val["user_email"] = data.user_email
existing_user_rows: Final[Sequence[LiteLLM_UserTable] | None] = await UserRepository(prisma_client).table.find_many(
where=key_val
removed_user_ids: Final = frozenset(m.user_id for m in removed_team_members if m.user_id is not None)
key_val: Final[Mapping[str, object]] = (
{"user_id": {"in": sorted(removed_user_ids)}} if removed_user_ids else {"user_email": data.user_email}
)
existing_user_rows: Final[Sequence[LiteLLM_UserTable]] = await _user_db(prisma_client).find_many(where=key_val)
if existing_user_rows is not None and (isinstance(existing_user_rows, list) and len(existing_user_rows) > 0):
for existing_user in existing_user_rows:
team_list = []
if data.team_id in existing_user.teams:
team_list = existing_user.teams
team_list.remove(data.team_id)
await _user_db(prisma_client).update(
where={
"user_id": existing_user.user_id,
},
data={"teams": {"set": team_list}},
)
for existing_user in existing_user_rows:
if data.team_id in existing_user.teams:
await _user_db(prisma_client).update(
where={
"user_id": existing_user.user_id,
},
data={"teams": {"set": [team for team in existing_user.teams if team != data.team_id]}},
)
# Also clean up any existing team membership rows for this user and team
user_ids_to_delete: Final = set[str]()
if data.user_id is not None:
user_ids_to_delete.add(data.user_id)
if existing_user_rows is not None and isinstance(existing_user_rows, list):
for existing_user in existing_user_rows:
if getattr(existing_user, "user_id", None):
user_ids_to_delete.add(existing_user.user_id)
user_ids_to_delete: Final = removed_user_ids.union(
(data.user_id,) if data.user_id is not None else (),
(user.user_id for user in existing_user_rows if user.user_id),
)
for _uid in user_ids_to_delete:
for _uid in sorted(user_ids_to_delete):
await _team_membership_db(prisma_client).delete_many(where={"team_id": data.team_id, "user_id": _uid})
## DELETE KEYS CREATED BY USER FOR THIS TEAM
@ -3245,7 +3300,7 @@ async def team_member_delete(
# Fetch keys before deletion to persist them
keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many(
where={
"user_id": {"in": list(user_ids_to_delete)},
"user_id": {"in": sorted(user_ids_to_delete)},
"team_id": data.team_id,
}
)
@ -3260,7 +3315,7 @@ async def team_member_delete(
await _tokens_db(prisma_client).delete_many(
where={
"user_id": {"in": list(user_ids_to_delete)},
"user_id": {"in": sorted(user_ids_to_delete)},
"team_id": data.team_id,
}
)
@ -3659,6 +3714,8 @@ async def delete_team(
create_audit_log_for_update,
litellm_proxy_admin_name,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if prisma_client is None:
@ -3752,6 +3809,12 @@ async def delete_team(
await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key")
await _invalidate_deleted_key_cache(
keys=keys_to_delete,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
## DELETE ASSOCIATED BYOK MODELS
# Runs before the team rows are deleted so a mid-flight failure never leaves
# the team gone with its models orphaned.
@ -3785,11 +3848,93 @@ async def delete_team(
)
await asyncio.gather(*tasks)
await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client)
## DELETE TEAMS
deleted_teams: Final = await prisma_client.delete_data(team_id_list=data.team_ids, table_name="team")
# Evict AFTER the rows are gone. Both writers of these keys (`_cache_team_object` and
# `get_team_object_by_alias`) hydrate from the db, so evicting first leaves a window where a
# concurrent auth lookup re-caches the still-present team and the delete looks like it never
# invalidated anything. Nothing fallible runs between the delete and this, or a failure there
# would strand the deleted team in cache.
await _invalidate_deleted_team_cache(
teams=team_rows,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
# Sweep again now the team is gone. A `/team/member_add` that landed between the first sweep
# and the delete would have re-appended the reference; an add still in flight sees the row
# missing under its own row lock and sweeps what it wrote. Both passes are idempotent, and
# keeping the first one means a failure here still leaves a team the admin can retry deleting.
await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client)
for deleted_team in team_rows:
await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id)
return deleted_teams
async def _sweep_deleted_team_references(team_ids: Sequence[str], prisma_client: PrismaClient) -> None:
"""
Strip the deleted team ids from every user row and team-membership row that still references them.
The per-member `team_member_delete` pass above only reaches users listed in the team's
`members_with_roles`, so a user row that outlived its roster entry is invisible to it and keeps
surfacing the team on `/user/info` after the team is gone.
#36839 closed the route that created that drift, by resolving member removal off the roster
entry's `user_id` rather than the identifier the caller happened to pass. It does not backfill
rows that already drifted, which is the state this was reported against, so the sweep still has
to run on delete.
`array_remove` rather than read-filter-write: rewriting the whole array from a snapshot read
outside a transaction drops any team a concurrent `/team/member_add` appended in between.
"""
for team_id in team_ids:
_ = await prisma_client.db.execute_raw(_STRIP_DELETED_TEAM_FROM_USERS_SQL, team_id)
_ = await _team_membership_db(prisma_client).delete_many(where=_TeamIdInFilter(team_id={"in": tuple(team_ids)}))
async def _invalidate_deleted_key_cache(
keys: Sequence[LiteLLM_VerificationToken],
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> None:
"""
Evict the auth cache entry for every key deleted along with the team.
`/key/delete` evicts as it goes, but the bulk delete above writes straight to the db. Auth
resolves a cached key object without re-reading the team, so a key belonging to a deleted team
keeps buying access until its TTL expires.
"""
await delete_cache_key_objects(
hashed_tokens=tuple(key.token for key in keys),
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
async def _invalidate_deleted_team_cache(
teams: Sequence[LiteLLM_TeamTable],
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> None:
_ = await asyncio.gather(
*(
delete_cache_team_object(
team_id=team.team_id,
team_alias=team.team_alias,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
for team in teams
)
)
def _transform_teams_to_deleted_records(
teams: list[LiteLLM_TeamTable],
user_api_key_dict: UserAPIKeyAuth,

149
litellm/proxy/management_endpoints/tool_management_endpoints.py
View file

@ -10,13 +10,21 @@ POST /v1/tool/policy - Update the input_policy / output_policy for a
"""
import uuid
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Annotated, Any, Final
from typing import TYPE_CHECKING, Annotated, Final, Protocol, TypeAlias, TypeVar, overload
from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import BaseModel, Field, TypeAdapter
if TYPE_CHECKING:
from prisma.models import LiteLLM_DailyToolSpend as PrismaDailyToolSpendRow
from prisma.models import LiteLLM_ObjectPermissionTable as PrismaObjectPermissionRow
from prisma.models import LiteLLM_SpendLogs as PrismaSpendLogRow
from prisma.models import LiteLLM_SpendLogToolIndex as PrismaSpendLogToolIndexRow
from prisma.models import LiteLLM_TeamTable as PrismaTeamRow
from prisma.models import LiteLLM_VerificationToken as PrismaVerificationTokenRow
from litellm.proxy.utils import PrismaClient
from litellm._logging import verbose_proxy_logger
@ -49,6 +57,72 @@ from litellm.types.tool_management import (
ToolUsageLogsResponse,
)
_RowT_co: Final = TypeVar("_RowT_co", covariant=True)
if TYPE_CHECKING:
class _TableOps(Protocol[_RowT_co]):
async def find_many(
self,
where: Mapping[str, object] | None = None,
order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
skip: int | None = None,
take: int | None = None,
) -> Sequence[_RowT_co]: ...
async def find_unique(self, where: Mapping[str, object]) -> _RowT_co | None: ...
async def count(self, where: Mapping[str, object] | None = None) -> int: ...
async def create(self, data: Mapping[str, object]) -> _RowT_co: ...
async def update_many(
self,
where: Mapping[str, object],
data: Mapping[str, object],
) -> int: ...
async def delete(self, where: Mapping[str, object]) -> _RowT_co | None: ...
async def group_by(
self,
by: Sequence[str],
sum: Mapping[str, bool] | None = None,
where: Mapping[str, object] | None = None,
order: Mapping[str, object] | None = None,
take: int | None = None,
) -> Sequence[Mapping[str, object]]: ...
class _SpendLogRow(Protocol):
@property
def messages(self) -> object: ...
@property
def proxy_server_request(self) -> str | Mapping[str, object] | None: ...
@overload
def _typed_table(repo: DailyToolSpendRepository) -> "_TableOps[PrismaDailyToolSpendRow]": ...
@overload
def _typed_table(repo: SpendLogToolIndexRepository) -> "_TableOps[PrismaSpendLogToolIndexRow]": ...
@overload
def _typed_table(repo: SpendLogsRepository) -> "_TableOps[PrismaSpendLogRow]": ...
@overload
def _typed_table(repo: VerificationTokenRepository) -> "_TableOps[PrismaVerificationTokenRow]": ...
@overload
def _typed_table(repo: TeamRepository) -> "_TableOps[PrismaTeamRow]": ...
@overload
def _typed_table(repo: ObjectPermissionRepository) -> "_TableOps[PrismaObjectPermissionRow]": ...
def _typed_table(
repo: DailyToolSpendRepository
| SpendLogToolIndexRepository
| SpendLogsRepository
| VerificationTokenRepository
| TeamRepository
| ObjectPermissionRepository,
) -> object:
return repo.table
router: Final = APIRouter()
TOOL_POLICY_OPTIONS: Final = ToolPolicyOptionsResponse(
@ -201,7 +275,7 @@ async def get_tool_spend(
end_str: Final = end_day.strftime("%Y-%m-%d")
date_window: Final = {"date": {"gte": start_str, "lte": end_str}}
table: Final = DailyToolSpendRepository(prisma_client).table
table: Final = _typed_table(DailyToolSpendRepository(prisma_client))
top_tools: Final = _TOP_TOOL_ROWS.validate_python(
await table.group_by(
by=["tool_name"],
@ -222,7 +296,7 @@ async def get_tool_spend(
for row in top_tools
]
daily_rows: Final = (
daily_rows: Final[Sequence[PrismaDailyToolSpendRow]] = (
await table.find_many(
where={**date_window, "tool_name": {"in": [row.tool_name for row in top_tools]}},
order=[{"date": "asc"}, {"spend": "desc"}],
@ -270,36 +344,43 @@ async def get_tool_detail(
raise HTTPException(status_code=500, detail=str(e))
def _input_snippet_for_tool_log(sl: Any, max_len: int = 200) -> str | None:
_ParsedJson: TypeAlias = dict[str, object] | list[object] | str | int | float | bool | None
_PARSED_JSON: Final[TypeAdapter[_ParsedJson]] = TypeAdapter(_ParsedJson)
_STR_OBJECT_DICT: Final = TypeAdapter(dict[str, object])
def _input_snippet_for_tool_log(sl: "_SpendLogRow | None", max_len: int = 200) -> str | None:
"""Short snippet from messages or proxy_server_request for tool usage log row."""
if sl is None:
return None
messages: Final = getattr(sl, "messages", None)
messages: Final = sl.messages
if messages is not None:
s = _snippet_str(messages, max_len)
if s:
return s
psr = getattr(sl, "proxy_server_request", None)
psr = sl.proxy_server_request
if not psr:
return None
if isinstance(psr, str):
import json
try:
psr = json.loads(psr)
psr = _PARSED_JSON.validate_python(json.loads(psr))
except Exception:
return _snippet_str(psr, max_len)
if isinstance(psr, dict):
msgs = psr.get("messages")
if msgs is None and isinstance(psr.get("body"), dict):
msgs = psr["body"].get("messages")
if msgs is None:
body: Final = psr.get("body")
if isinstance(body, dict):
msgs = _STR_OBJECT_DICT.validate_python(body).get("messages")
s = _snippet_str(msgs, max_len)
if s:
return s
return _snippet_str(psr, max_len)
def _snippet_str(text: Any, max_len: int = 200) -> str | None:
def _snippet_str(text: object, max_len: int = 200) -> str | None:
if text is None:
return None
if isinstance(text, str):
@ -344,7 +425,7 @@ async def get_tool_usage_logs(
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
try:
where: Final[dict] = {"tool_name": tool_name}
where: Final[dict[str, object]] = {"tool_name": tool_name}
if start_date or end_date:
start_time_filter: datetime | None = None
end_time_filter: datetime | None = None
@ -363,14 +444,14 @@ async def get_tool_usage_logs(
except ValueError:
pass
if start_time_filter is not None or end_time_filter is not None:
where["start_time"] = {}
if start_time_filter is not None:
where["start_time"]["gte"] = start_time_filter
if end_time_filter is not None:
where["start_time"]["lte"] = end_time_filter
where["start_time"] = {
key: value
for key, value in (("gte", start_time_filter), ("lte", end_time_filter))
if value is not None
}
total: Final = await SpendLogToolIndexRepository(prisma_client).table.count(where=where)
index_rows: Final = await SpendLogToolIndexRepository(prisma_client).table.find_many(
total: Final = await _typed_table(SpendLogToolIndexRepository(prisma_client)).count(where=where)
index_rows: Final = await _typed_table(SpendLogToolIndexRepository(prisma_client)).find_many(
where=where,
order={"start_time": "desc"},
skip=(page - 1) * page_size,
@ -380,7 +461,9 @@ async def get_tool_usage_logs(
if not request_ids:
return ToolUsageLogsResponse(logs=[], total=total, page=page, page_size=page_size)
spend_logs = await SpendLogsRepository(prisma_client).table.find_many(where={"request_id": {"in": request_ids}})
spend_logs = await _typed_table(SpendLogsRepository(prisma_client)).find_many(
where={"request_id": {"in": request_ids}}
)
log_by_id: Final = {s.request_id: s for s in spend_logs}
logs_out: Final[list[ToolUsageLogEntry]] = []
@ -449,24 +532,24 @@ async def _resolve_key_hash_to_object_permission_id(
hashed: Final = key_hash if "sk-" not in (key_hash or "") else hash_token(key_hash)
if not hashed:
return None
row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed})
row = await _typed_table(VerificationTokenRepository(prisma_client)).find_unique(where={"token": hashed})
if row is None:
return None
op_id: Final = getattr(row, "object_permission_id", None)
op_id: Final = row.object_permission_id
if op_id:
return op_id
new_id: Final = str(uuid.uuid4())
await ObjectPermissionRepository(prisma_client).table.create(
await _typed_table(ObjectPermissionRepository(prisma_client)).create(
data={"object_permission_id": new_id, "blocked_tools": []}
)
updated_count: Final = await VerificationTokenRepository(prisma_client).table.update_many(
updated_count: Final = await _typed_table(VerificationTokenRepository(prisma_client)).update_many(
where={"token": hashed, "object_permission_id": None},
data={"object_permission_id": new_id},
)
if updated_count == 0:
await ObjectPermissionRepository(prisma_client).table.delete(where={"object_permission_id": new_id})
row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed})
return getattr(row, "object_permission_id", None) if row else None
await _typed_table(ObjectPermissionRepository(prisma_client)).delete(where={"object_permission_id": new_id})
row = await _typed_table(VerificationTokenRepository(prisma_client)).find_unique(where={"token": hashed})
return row.object_permission_id if row else None
return new_id
@ -478,24 +561,24 @@ async def _resolve_team_id_to_object_permission_id(
if not team_id or not team_id.strip():
return None
team_id_clean: Final = team_id.strip()
row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id_clean})
row = await _typed_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id_clean})
if row is None:
return None
op_id: Final = getattr(row, "object_permission_id", None)
op_id: Final = row.object_permission_id
if op_id:
return op_id
new_id: Final = str(uuid.uuid4())
await ObjectPermissionRepository(prisma_client).table.create(
await _typed_table(ObjectPermissionRepository(prisma_client)).create(
data={"object_permission_id": new_id, "blocked_tools": []}
)
updated_count: Final = await TeamRepository(prisma_client).table.update_many(
updated_count: Final = await _typed_table(TeamRepository(prisma_client)).update_many(
where={"team_id": team_id_clean, "object_permission_id": None},
data={"object_permission_id": new_id},
)
if updated_count == 0:
await ObjectPermissionRepository(prisma_client).table.delete(where={"object_permission_id": new_id})
row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id_clean})
return getattr(row, "object_permission_id", None) if row else None
await _typed_table(ObjectPermissionRepository(prisma_client)).delete(where={"object_permission_id": new_id})
row = await _typed_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id_clean})
return row.object_permission_id if row else None
return new_id

132
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -21,6 +21,7 @@ from copy import deepcopy
from html import escape
from typing import (
TYPE_CHECKING,
Annotated,
Any,
Final,
Literal,
@ -40,6 +41,7 @@ if TYPE_CHECKING:
import jwt
from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status
from fastapi.responses import RedirectResponse
from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
@ -185,6 +187,7 @@ class _PrismaTableActions(Protocol[_DbRecordT]):
async def find_many(
self,
where: Mapping[str, object] | None = None,
include: Mapping[str, bool] | None = None,
) -> Sequence[_DbRecordT]: ...
async def update(
@ -241,6 +244,45 @@ def _team_detail_db(repo: "_HasTeamDetailTable") -> "_PrismaTableActions[_TeamDe
return repo.table
_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str])
def _decode_model_aliases(value: object) -> object:
"""``/team/new`` stores team model aliases as a JSON-encoded string in the Json column."""
if not isinstance(value, str):
return value
try:
return _MODEL_ALIASES_ADAPTER.validate_json(value)
except ValidationError:
return None
class _TeamModelAliasTable(BaseModel):
model_config = ConfigDict(protected_namespaces=())
model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None
class _TeamRowGrants(BaseModel):
team_id: str
team_alias: str | None = None
models: tuple[str, ...] = ()
litellm_model_table: _TeamModelAliasTable | None = None
class _CliSsoTeamDetail(BaseModel):
"""The per-team snapshot cached in the CLI SSO flow and echoed to the CLI on poll."""
team_id: str | None = None
team_alias: str | None = None
team_models: tuple[str, ...]
team_model_aliases: Mapping[str, str] | None = None
_CLI_SSO_TEAM_DETAILS_ADAPTER: Final = TypeAdapter(tuple[_CliSsoTeamDetail, ...])
_TEAMLESS_CLI_SSO_TEAM_DETAIL: Final = _CliSsoTeamDetail(team_models=())
class _CustomSsoCall(Protocol):
async def __call__(self, sso_response: object) -> SSOUserDefinedValues | None: ...
@ -2147,27 +2189,55 @@ async def _build_cli_sso_user_defined_values(
)
def _cli_sso_team_detail(team_row: Mapping[str, object]) -> _CliSsoTeamDetail:
team: Final = _TeamRowGrants.model_validate(team_row)
alias_table: Final = team.litellm_model_table
return _CliSsoTeamDetail(
team_id=team.team_id,
team_alias=team.team_alias,
team_models=team.models,
team_model_aliases=alias_table.model_aliases if alias_table is not None else None,
)
async def _fetch_cli_sso_team_details(
prisma_client: PrismaClient,
teams: Sequence[str],
) -> list[dict[str, object]]:
team_details: Final[list[dict[str, object]]] = []
) -> tuple[_CliSsoTeamDetail, ...] | None:
"""``None`` means the lookup itself failed, which is not the same as the user having no teams."""
if not teams:
return ()
try:
if teams:
prisma_teams: Final = await _team_detail_db(TeamRepository(prisma_client)).find_many(
where={"team_id": {"in": teams}}
)
for team_row in prisma_teams:
team_dict = team_row.model_dump()
team_details.append(
{
"team_id": team_dict.get("team_id"),
"team_alias": team_dict.get("team_alias"),
}
)
prisma_teams: Final = await _team_detail_db(TeamRepository(prisma_client)).find_many(
where={"team_id": {"in": teams}},
include={"litellm_model_table": True},
)
except Exception as e:
verbose_proxy_logger.error("Error fetching team details for CLI SSO session: %s", e)
return team_details
return None
return tuple(_cli_sso_team_detail(team_row.model_dump()) for team_row in prisma_teams)
def _cli_sso_session_teams(team_details: Sequence[_CliSsoTeamDetail]) -> list[str]:
"""The teams a login may bind to: only those whose row still exists.
A team deleted out from under a membership, which is what deleting an organization
leaves behind, can never resolve its grants, so offering it would refuse every
future login for that user with nothing they could do to recover.
"""
return [detail.team_id for detail in team_details if detail.team_id is not None]
def _selected_cli_sso_team_detail(team_details: object, team_id: str | None) -> _CliSsoTeamDetail | None:
"""``None`` means the team's grants are unknown. An empty grant is a real value meaning unrestricted,
so an unknown one must not be minted as empty."""
if team_id is None:
return _TEAMLESS_CLI_SSO_TEAM_DETAIL
try:
details: Final = _CLI_SSO_TEAM_DETAILS_ADAPTER.validate_python(team_details)
except ValidationError:
return None
return next((detail for detail in details if detail.team_id == team_id), None)
async def _complete_cli_sso_callback_session(
@ -2210,6 +2280,12 @@ async def _complete_cli_sso_callback_session(
teams = user_info.teams if isinstance(user_info.teams, list) else []
team_details: Final = await _fetch_cli_sso_team_details(prisma_client=prisma_client, teams=teams)
if team_details is None:
raise HTTPException(
status_code=500,
detail="Could not resolve team model grants for this login. Please try again",
)
resolved_teams: Final = _cli_sso_session_teams(team_details)
attribution_metadata: Final = build_cli_sso_attribution_metadata(result=result)
if attribution_metadata:
await _persist_cli_sso_user_metadata(
@ -2223,8 +2299,8 @@ async def _complete_cli_sso_callback_session(
"user_role": user_info.user_role,
"models": user_info.models if hasattr(user_info, "models") else [],
"user_email": user_email,
"teams": teams,
"team_details": team_details,
"teams": resolved_teams,
"team_details": [detail.model_dump() for detail in team_details],
"attribution_metadata": attribution_metadata,
}
flow["sso_complete"] = True
@ -2233,7 +2309,10 @@ async def _complete_cli_sso_callback_session(
_set_cli_sso_flow(login_id=key, cache=cli_sso_session_cache, flow=flow)
verbose_proxy_logger.info(
"Stored CLI SSO session for user: %s, teams: %s, num_teams: %s", user_info.user_id, teams, len(teams)
"Stored CLI SSO session for user: %s, teams: %s, num_teams: %s",
user_info.user_id,
resolved_teams,
len(resolved_teams),
)
verify_url: Final = get_custom_url(
request_base_url=str(request.base_url),
@ -2401,11 +2480,14 @@ async def cli_poll_key(
# If no team_id provided and user has 0 or 1 team, use first team (or None)
team_id = user_teams[0] if len(user_teams) > 0 else None
team_alias = None
if team_id and isinstance(user_team_details, list):
team_alias = next(
(team.get("team_alias") for team in user_team_details if team.get("team_id") == team_id),
None,
selected_team: Final = _selected_cli_sso_team_detail(
team_details=user_team_details,
team_id=team_id,
)
if selected_team is None:
raise HTTPException(
status_code=500,
detail=f"Could not resolve the model grants for team: {team_id}. Please run `lite login` again",
)
user_info: Final = LiteLLM_UserTable(
@ -2417,7 +2499,9 @@ async def cli_poll_key(
jwt_token: Final = ExperimentalUIJWTToken.get_cli_jwt_auth_token(
user_info=user_info,
team_id=team_id,
team_alias=team_alias,
team_alias=selected_team.team_alias,
team_models=selected_team.team_models,
team_model_aliases=selected_team.team_model_aliases,
max_budget=None,
)

75
litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py
View file

@ -4,11 +4,11 @@ usage/spend data by querying the aggregated daily activity endpoints.
"""
import json
from collections.abc import AsyncIterator, Callable
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence
from datetime import date
from typing import Any, Final, Literal, cast
from typing import Any, Final, Literal, Protocol, cast, overload
from typing_extensions import TypedDict
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
@ -73,9 +73,36 @@ class SSEErrorEvent(TypedDict):
SSEEvent = SSEStatusEvent | SSEToolCallEvent | SSEChunkEvent | SSEDoneEvent | SSEErrorEvent
class _EntityEntry(TypedDict, total=False):
metrics: ReadOnly[Mapping[str, float]]
metadata: ReadOnly[Mapping[str, str]]
class _DayDump(TypedDict, total=False):
breakdown: ReadOnly[Mapping[str, Mapping[str, _EntityEntry]]]
class _UsageDump(Protocol):
@overload
def get(self, key: Literal["metadata"], default: Mapping[str, float], /) -> Mapping[str, float]: ...
@overload
def get(self, key: Literal["results"], default: Sequence[_DayDump], /) -> Sequence[_DayDump]: ...
class _ToolFunctionDef(TypedDict):
name: ReadOnly[str]
description: ReadOnly[str]
parameters: ReadOnly[Mapping[str, object]]
class _ToolDef(TypedDict):
type: ReadOnly[str]
function: ReadOnly[_ToolFunctionDef]
class ToolHandler(TypedDict):
fetch: Callable[..., Any]
summarise: Callable[[dict[str, Any]], str]
fetch: Callable[..., Awaitable[_UsageDump]]
summarise: Callable[[_UsageDump], str]
label: str
@ -88,7 +115,7 @@ _DATE_PARAMS: Final = {
"end_date": {"type": "string", "description": "End date in YYYY-MM-DD format"},
}
_TOOL_USAGE: Final = {
_TOOL_USAGE: Final[_ToolDef] = {
"type": "function",
"function": {
"name": "get_usage_data",
@ -111,7 +138,7 @@ _TOOL_USAGE: Final = {
},
}
_TOOL_TEAM: Final = {
_TOOL_TEAM: Final[_ToolDef] = {
"type": "function",
"function": {
"name": "get_team_usage_data",
@ -133,7 +160,7 @@ _TOOL_TEAM: Final = {
},
}
_TOOL_TAG: Final = {
_TOOL_TAG: Final[_ToolDef] = {
"type": "function",
"function": {
"name": "get_tag_usage_data",
@ -159,7 +186,7 @@ TOOLS_BASE: Final = [_TOOL_USAGE]
TOOLS_ADMIN: Final = [_TOOL_USAGE, _TOOL_TEAM, _TOOL_TAG]
def get_tools_for_role(is_admin: bool) -> list[dict[str, Any]]:
def get_tools_for_role(is_admin: bool) -> list[_ToolDef]:
"""Return the tool list appropriate for the user's role."""
return TOOLS_ADMIN if is_admin else TOOLS_BASE
@ -254,7 +281,7 @@ async def _query_activity(
)
async def _fetch_usage_data(start_date: str, end_date: str, user_id: str | None = None) -> dict[str, Any]:
async def _fetch_usage_data(start_date: str, end_date: str, user_id: str | None = None) -> _UsageDump:
resp: Final = await _query_activity(
TABLE_DAILY_USER_SPEND,
ENTITY_FIELD_USER,
@ -266,7 +293,7 @@ async def _fetch_usage_data(start_date: str, end_date: str, user_id: str | None
return resp.model_dump(mode="json")
async def _fetch_team_usage_data(start_date: str, end_date: str, team_ids: str | None = None) -> dict[str, Any]:
async def _fetch_team_usage_data(start_date: str, end_date: str, team_ids: str | None = None) -> _UsageDump:
resp: Final = await _query_activity(
TABLE_DAILY_TEAM_SPEND,
ENTITY_FIELD_TEAM,
@ -277,7 +304,7 @@ async def _fetch_team_usage_data(start_date: str, end_date: str, team_ids: str |
return resp.model_dump(mode="json")
async def _fetch_tag_usage_data(start_date: str, end_date: str, tags: str | None = None) -> dict[str, Any]:
async def _fetch_tag_usage_data(start_date: str, end_date: str, tags: str | None = None) -> _UsageDump:
resp: Final = await _query_activity(
TABLE_DAILY_TAG_SPEND,
ENTITY_FIELD_TAG,
@ -294,7 +321,7 @@ async def _fetch_tag_usage_data(start_date: str, end_date: str, tags: str | None
def _accumulate_breakdown(
results: list[dict[str, Any]], dimension: str, fields: list[str]
results: Sequence[_DayDump], dimension: str, fields: Sequence[str]
) -> dict[str, dict[str, float]]:
"""Aggregate a single breakdown dimension across days."""
totals: Final[dict[str, dict[str, float]]] = {}
@ -317,7 +344,7 @@ def _ranked_lines(
return [fmt(name, vals) for name, vals in sorted(totals.items(), key=lambda x: -x[1].get("spend", 0))[:limit]]
def _summarise_usage_data(data: dict[str, Any]) -> str:
def _summarise_usage_data(data: _UsageDump) -> str:
meta: Final = data.get("metadata", {})
results: Final = data.get("results", [])
@ -349,7 +376,7 @@ def _summarise_usage_data(data: dict[str, Any]) -> str:
return "\n".join(sections)
def _summarise_entity_data(data: dict[str, Any], entity_label: str) -> str:
def _summarise_entity_data(data: _UsageDump, entity_label: str) -> str:
"""Summarise team/tag entity usage data."""
results: Final = data.get("results", [])
if not results:
@ -409,16 +436,16 @@ def _sse(event: SSEEvent) -> str:
def _resolve_fetch_kwargs(
fn_name: str,
fn_args: dict[str, str],
fn_args: Mapping[str, str],
user_id: str | None,
is_admin: bool,
) -> dict[str, Any]:
) -> dict[str, str]:
"""Build keyword arguments for a tool's fetch function."""
start_date: Final = fn_args.get("start_date", "")
end_date: Final = fn_args.get("end_date", "")
if not start_date or not end_date:
raise ValueError("Missing required start_date or end_date from tool arguments")
kwargs: Final[dict[str, Any]] = {"start_date": start_date, "end_date": end_date}
kwargs: Final[dict[str, str]] = {"start_date": start_date, "end_date": end_date}
if fn_name == "get_usage_data":
if not is_admin:
if user_id is None:
@ -443,7 +470,7 @@ def _resolve_fetch_kwargs(
async def _execute_tool_call(
handler: ToolHandler,
fn_name: str,
fn_args: dict[str, str],
fn_args: Mapping[str, str],
user_id: str | None,
is_admin: bool,
) -> str:
@ -455,13 +482,13 @@ async def _execute_tool_call(
async def _process_tool_call(
tc: Any,
chat_messages: list[dict[str, Any]],
chat_messages: list[Mapping[str, object]],
user_id: str | None,
is_admin: bool,
) -> AsyncIterator[str]:
"""Execute a single tool call, yielding SSE events for status."""
fn_name: Final = tc.function.name
fn_args: Final = json.loads(tc.function.arguments)
fn_name: Final[str] = tc.function.name
fn_args: Final[Mapping[str, str]] = json.loads(tc.function.arguments)
allowed_names: Final = {t["function"]["name"] for t in get_tools_for_role(is_admin)}
handler: Final = TOOL_HANDLERS.get(fn_name)
@ -495,7 +522,7 @@ async def _process_tool_call(
chat_messages.append({"role": "tool", "tool_call_id": tc.id, "content": tool_result})
async def _stream_final_response(model: str, chat_messages: list[dict[str, Any]]) -> AsyncIterator[str]:
async def _stream_final_response(model: str, chat_messages: list[Mapping[str, object]]) -> AsyncIterator[str]:
"""Stream the final LLM response after tool results are appended."""
yield _sse({"type": "status", "message": "Analyzing results..."})
@ -520,7 +547,7 @@ async def stream_usage_ai_chat(
"""Stream SSE events: status → tool_call → chunk → done."""
resolved_model: Final = (model or "").strip() or DEFAULT_COMPETITOR_DISCOVERY_MODEL
truncated: Final = messages[-MAX_CHAT_MESSAGES:] if len(messages) > MAX_CHAT_MESSAGES else messages
chat_messages: Final[list[dict[str, Any]]] = [
chat_messages: Final[list[Mapping[str, object]]] = [
{"role": "system", "content": _build_system_prompt(is_admin)},
*truncated,
]

119
litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py
View file

@ -11,11 +11,19 @@ These endpoints use optimized single SQL queries with joins to efficiently calcu
user metrics from tag activity data and return time series for dashboard visualization.
"""
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta
from typing import Any, Final
from typing import TYPE_CHECKING, Final, Protocol, TypeVar, overload
from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
if TYPE_CHECKING:
from prisma.models import LiteLLM_DailyTagSpend as PrismaDailyTagSpendRow
from prisma.models import LiteLLM_UserTable as PrismaUserRow
from prisma.models import LiteLLM_VerificationToken as PrismaVerificationTokenRow
from litellm.proxy.utils import PrismaClient
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -103,6 +111,54 @@ class PerUserAnalyticsResponse(BaseModel):
total_pages: int
class _DistinctTagRow(BaseModel):
tag: str
class _ActiveUsersRow(BaseModel):
tag: str
active_users: int
date: str
period_start: str | None = None
period_end: str | None = None
class _TagSummaryRow(BaseModel):
tag: str
unique_users: int | None = None
total_requests: float | int | str | None = None
successful_requests: float | int | str | None = None
failed_requests: float | int | str | None = None
total_tokens: float | int | str | None = None
total_spend: float | int | str | None = None
_DISTINCT_TAG_ROWS: Final = TypeAdapter(list[_DistinctTagRow])
_ACTIVE_USERS_ROWS: Final = TypeAdapter(list[_ActiveUsersRow])
_TAG_SUMMARY_ROWS: Final = TypeAdapter(list[_TagSummaryRow])
_RowT_co: Final = TypeVar("_RowT_co", covariant=True)
if TYPE_CHECKING:
class _TableOps(Protocol[_RowT_co]):
async def find_many(self, where: Mapping[str, object] | None = None) -> Sequence[_RowT_co]: ...
@overload
def _typed_table(repo: DailyTagSpendRepository) -> "_TableOps[PrismaDailyTagSpendRow]": ...
@overload
def _typed_table(repo: VerificationTokenRepository) -> "_TableOps[PrismaVerificationTokenRow]": ...
@overload
def _typed_table(repo: UserRepository) -> "_TableOps[PrismaUserRow]": ...
def _typed_table(repo: DailyTagSpendRepository | VerificationTokenRepository | UserRepository) -> object:
return repo.table
async def _query_raw(prisma_client: "PrismaClient", sql_query: str, *params: object) -> object:
return await prisma_client.db.query_raw(sql_query, *params)
@router.get(
"/tag/distinct",
response_model=DistinctTagsResponse,
@ -141,9 +197,9 @@ async def get_distinct_user_agent_tags(
LIMIT {MAX_TAGS}
"""
db_response: Final = await prisma_client.db.query_raw(sql_query)
db_response: Final = _DISTINCT_TAG_ROWS.validate_python(await _query_raw(prisma_client, sql_query))
results: Final = [DistinctTagResponse(tag=row["tag"]) for row in db_response]
results: Final = [DistinctTagResponse(tag=row.tag) for row in db_response]
return DistinctTagsResponse(results=results)
@ -231,11 +287,10 @@ async def get_daily_active_users(
ORDER BY dts.date DESC, active_users DESC
"""
db_response: Final = await prisma_client.db.query_raw(sql_query, *params)
db_response: Final = _ACTIVE_USERS_ROWS.validate_python(await _query_raw(prisma_client, sql_query, *params))
results: Final = [
TagActiveUsersResponse(tag=row["tag"], active_users=row["active_users"], date=row["date"])
for row in db_response
TagActiveUsersResponse(tag=row.tag, active_users=row.active_users, date=row.date) for row in db_response
]
return ActiveUsersAnalyticsResponse(results=results)
@ -346,15 +401,15 @@ async def get_weekly_active_users(
ORDER BY week_offset DESC, active_users DESC
"""
db_response: Final = await prisma_client.db.query_raw(sql_query, *params)
db_response: Final = _ACTIVE_USERS_ROWS.validate_python(await _query_raw(prisma_client, sql_query, *params))
results: Final = [
TagActiveUsersResponse(
tag=row["tag"],
active_users=row["active_users"],
date=row["date"], # This will be "Week 1 (Jan 15)", "Week 2 (Jan 8)", etc.
period_start=row["period_start"],
period_end=row["period_end"],
tag=row.tag,
active_users=row.active_users,
date=row.date, # This will be "Week 1 (Jan 15)", "Week 2 (Jan 8)", etc.
period_start=row.period_start,
period_end=row.period_end,
)
for row in db_response
]
@ -467,15 +522,15 @@ async def get_monthly_active_users(
ORDER BY month_offset DESC, active_users DESC
"""
db_response: Final = await prisma_client.db.query_raw(sql_query, *params)
db_response: Final = _ACTIVE_USERS_ROWS.validate_python(await _query_raw(prisma_client, sql_query, *params))
results: Final = [
TagActiveUsersResponse(
tag=row["tag"],
active_users=row["active_users"],
date=row["date"], # This will be "Month 1 (Jan)", "Month 2 (Dec)", etc.
period_start=row["period_start"],
period_end=row["period_end"],
tag=row.tag,
active_users=row.active_users,
date=row.date, # This will be "Month 1 (Jan)", "Month 2 (Dec)", etc.
period_start=row.period_start,
period_end=row.period_end,
)
for row in db_response
]
@ -565,17 +620,17 @@ async def get_tag_summary(
ORDER BY total_requests DESC
"""
db_response: Final = await prisma_client.db.query_raw(sql_query, *params)
db_response: Final = _TAG_SUMMARY_ROWS.validate_python(await _query_raw(prisma_client, sql_query, *params))
results: Final = [
TagSummaryMetrics(
tag=row["tag"],
unique_users=row["unique_users"] or 0,
total_requests=int(row["total_requests"] or 0),
successful_requests=int(row["successful_requests"] or 0),
failed_requests=int(row["failed_requests"] or 0),
total_tokens=int(row["total_tokens"] or 0),
total_spend=float(row["total_spend"] or 0.0),
tag=row.tag,
unique_users=row.unique_users or 0,
total_requests=int(row.total_requests or 0),
successful_requests=int(row.successful_requests or 0),
failed_requests=int(row.failed_requests or 0),
total_tokens=int(row.total_tokens or 0),
total_spend=float(row.total_spend or 0.0),
)
for row in db_response
]
@ -648,7 +703,7 @@ async def get_per_user_analytics(
start_date: Final = start_dt.strftime("%Y-%m-%d")
# Build where clause with date range
where_clause: Final[dict[str, Any]] = {"date": {"gte": start_date, "lte": end_date}}
where_clause: Final[dict[str, object]] = {"date": {"gte": start_date, "lte": end_date}}
# Add tag filtering if provided
if tag_filters and len(tag_filters) > 0:
@ -657,7 +712,7 @@ async def get_per_user_analytics(
where_clause["tag"] = {"contains": tag_filter}
# Get all tag records in the date range with optional tag filtering
tag_records: Final = await DailyTagSpendRepository(prisma_client).table.find_many(where=where_clause)
tag_records: Final = await _typed_table(DailyTagSpendRepository(prisma_client)).find_many(where=where_clause)
# Get unique api_keys
api_keys: Final = set(record.api_key for record in tag_records if record.api_key)
@ -672,7 +727,7 @@ async def get_per_user_analytics(
)
# Lookup user_id for each api_key
api_key_records: Final = await VerificationTokenRepository(prisma_client).table.find_many(
api_key_records: Final = await _typed_table(VerificationTokenRepository(prisma_client)).find_many(
where={"token": {"in": list(api_keys)}}
)
@ -681,7 +736,9 @@ async def get_per_user_analytics(
# Get user emails for the user_ids
user_ids: Final = list(set(api_key_to_user_id.values()))
user_records: Final = await UserRepository(prisma_client).table.find_many(where={"user_id": {"in": user_ids}})
user_records: Final = await _typed_table(UserRepository(prisma_client)).find_many(
where={"user_id": {"in": user_ids}}
)
# Create mapping from user_id to user_email
user_id_to_email: Final = {record.user_id: record.user_email for record in user_records}

173
litellm/proxy/management_helpers/access_group_key_sync.py Normal file
View file

@ -0,0 +1,173 @@
"""
Reverse sync for the key side of the key <-> access group relationship.
`litellm_accessgrouptable.assigned_key_ids` and `litellm_verificationtoken.access_group_ids`
are the two halves of one relationship and BOTH are read: the access group's
attached-keys view reads the former, and so does the grant check in
`auth_checks.get_authorized_resources_from_key_access_groups`, which authorizes a
key only when the group lists the key's token (or the key's team). The access-group
endpoints maintain both halves already; this module is what the key write paths call
so an edit from that side is mirrored back.
Every write is a single guarded statement rather than a read-modify-write. Prisma has no
atomic scalar-list removal (see `TeamRepository.remove_member`), and the read-modify-write
it otherwise forces is not safe here: a lost update would put an already revoked token back
into a group and restore its grants, or drop a grant an admin just made. The guards also
make each statement idempotent, so a retry cannot duplicate an entry. Each statement covers
every group the request touches at once, so the size of the caller's id list does not turn
into a matching number of round trips, and returns the ids it actually moved so only those
groups are dropped from cache.
It deliberately lives outside `access_group_endpoints`, which is a lazily
registered feature router (see `_lazy_features.LAZY_FEATURES`). Importing that
module eagerly from `key_management_endpoints` would put it in `sys.modules`
without its router ever being included, which drops its routes from the OpenAPI
schema.
"""
from collections.abc import Sequence
from typing import Final, Protocol
from pydantic import BaseModel
from litellm.proxy._types import (
LiteLLM_VerificationToken,
RegenerateKeyRequest,
UpdateKeyRequest,
)
from litellm.proxy.auth.auth_checks import (
_delete_cache_access_object, # pyright: ignore[reportPrivateUsage] # the access-group endpoints reach for this same cache primitive
)
from litellm.repositories.table_repositories import AccessGroupRepository
class _MovedGroupRow(BaseModel):
access_group_id: str
class _RawExecutor(Protocol):
async def query_raw(self, query: str, *args: str | Sequence[str]) -> Sequence[object]: ...
_ATTACH_KEY_SQL: Final = (
'UPDATE "LiteLLM_AccessGroupTable" '
'SET "assigned_key_ids" = array_append("assigned_key_ids", $1) '
'WHERE "access_group_id" = ANY($2::text[]) AND NOT ($1 = ANY("assigned_key_ids")) '
'RETURNING "access_group_id"'
)
_DETACH_KEY_SQL: Final = (
'UPDATE "LiteLLM_AccessGroupTable" '
'SET "assigned_key_ids" = array_remove("assigned_key_ids", $1) '
'WHERE "access_group_id" = ANY($2::text[]) AND $1 = ANY("assigned_key_ids") '
'RETURNING "access_group_id"'
)
_REPOINT_KEY_SQL: Final = (
'UPDATE "LiteLLM_AccessGroupTable" '
'SET "assigned_key_ids" = array_append(array_remove(array_remove("assigned_key_ids", $1), $2), $2) '
'WHERE $1 = ANY("assigned_key_ids") '
'RETURNING "access_group_id"'
)
def _raw_executor(prisma_client: object) -> _RawExecutor:
"""Narrow the untyped Prisma client down to the raw-query call this module makes."""
return AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client
async def _invalidate_access_group_cache(access_group_id: str) -> None:
"""
Drop an access group entry from both the in-memory and Redis caches.
Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server
to avoid circular imports, following the same pattern as key_management_endpoints.
"""
from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
await _delete_cache_access_object(
access_group_id=access_group_id,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
async def _invalidate_moved_groups(moved_rows: Sequence[object]) -> None:
for row in moved_rows:
await _invalidate_access_group_cache(_MovedGroupRow.model_validate(row).access_group_id)
async def _write_membership(prisma_client: object, sql: str, access_group_ids: frozenset[str], key_token: str) -> None:
"""Run one guarded membership statement for every listed group, dropping the cache of those it moved."""
if not access_group_ids:
return
await _invalidate_moved_groups(
await _raw_executor(prisma_client).query_raw(sql, key_token, sorted(access_group_ids))
)
async def sync_key_access_group_membership(
prisma_client: object,
key_token: str,
previous_access_group_ids: Sequence[str] | None,
updated_access_group_ids: Sequence[str] | None,
) -> None:
"""Mirror a key-side change to `access_group_ids` onto each access group's `assigned_key_ids`."""
previous: Final = frozenset(previous_access_group_ids or ())
updated: Final = frozenset(updated_access_group_ids or ())
await _write_membership(prisma_client, _ATTACH_KEY_SQL, updated - previous, key_token)
await _write_membership(prisma_client, _DETACH_KEY_SQL, previous - updated, key_token)
async def sync_key_update_access_group_membership(
prisma_client: object,
key_token: str,
data: UpdateKeyRequest | RegenerateKeyRequest,
existing_key_row: LiteLLM_VerificationToken,
) -> None:
"""
Mirror a key UPDATE onto the group side, honouring `exclude_unset` semantics.
The key row is written from `model_dump(exclude_unset=True)`, so a request that never
mentions `access_group_ids` leaves the key's own list alone and must leave the group's
copy alone too. Reading the attribute instead of `model_fields_set` would see None on
every unrelated edit and withdraw the token from every group it belongs to.
"""
if "access_group_ids" not in data.model_fields_set:
return
await sync_key_access_group_membership(
prisma_client=prisma_client,
key_token=key_token,
previous_access_group_ids=existing_key_row.access_group_ids,
updated_access_group_ids=data.access_group_ids,
)
async def sync_key_regeneration_access_group_membership(
prisma_client: object,
previous_key_token: str,
new_key_token: str,
data: RegenerateKeyRequest | None,
existing_key_row: LiteLLM_VerificationToken,
) -> None:
"""
Re-point every group's copy from the old token to the regenerated one.
Regeneration replaces the token, which is the identity `assigned_key_ids` stores, so
leaving the old hash behind both points the group at a row that no longer exists and
denies the regenerated key the group's grants. The swap is driven by the groups that
hold the old token when the statement runs, not by the key row read earlier, so a group
edited in between is neither resurrected nor skipped. Removing the new token before
appending it keeps a re-run from duplicating it.
"""
await _invalidate_moved_groups(
await _raw_executor(prisma_client).query_raw(_REPOINT_KEY_SQL, previous_key_token, new_key_token)
)
if data is not None:
await sync_key_update_access_group_membership(
prisma_client=prisma_client,
key_token=new_key_token,
data=data,
existing_key_row=existing_key_row,
)

155
litellm/proxy/management_helpers/access_group_team_sync.py Normal file
View file

@ -0,0 +1,155 @@
"""
Reverse sync for the team side of the team <-> access group relationship.
`litellm_accessgrouptable.assigned_team_ids` and `litellm_teamtable.access_group_ids`
are two copies of the same relationship, and both are read: the access group's
attached-teams view reads the former, and so does the key-side grant check in
`auth_checks.get_authorized_resources_from_key_access_groups`. The access-group
endpoints maintain both copies already; this module is what the team write paths
call so an edit from that side is mirrored back.
It deliberately lives outside `access_group_endpoints`, which is a lazily
registered feature router (see `_lazy_features.LAZY_FEATURES`). Importing that
module eagerly from `team_endpoints` would put it in `sys.modules` without its
router ever being included, which drops its routes from the OpenAPI schema.
"""
import asyncio
from collections.abc import Mapping, Sequence
from typing import Final, Protocol
from pydantic import BaseModel, TypeAdapter
from litellm.proxy.auth.auth_checks import _delete_cache_access_object
# hashtext collisions only cost two unrelated teams a little serialization, and the
# lock is never taken by the access-group endpoints, so it cannot join their
# access-group-then-team lock order to form a cycle.
_LOCK_TEAM_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked"
_READ_TEAM_SQL: Final = 'SELECT access_group_ids FROM "LiteLLM_TeamTable" WHERE team_id = $1'
# The groups the team is on either side of the reconcile, so the cache step is driven by
# desired state rather than by which rows this attempt happened to change. A retry after a
# failed invalidation finds the same set even though its statements are already no-ops.
_AFFECTED_SQL: Final = """
SELECT access_group_id FROM "LiteLLM_AccessGroupTable"
WHERE access_group_id = ANY($2::TEXT[])
OR $1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]))
"""
_ATTACH_SQL: Final = """
UPDATE "LiteLLM_AccessGroupTable"
SET assigned_team_ids = array_append(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]), $1)
WHERE access_group_id = ANY($2::TEXT[])
AND NOT ($1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[])))
RETURNING access_group_id
"""
_DETACH_SQL: Final = """
UPDATE "LiteLLM_AccessGroupTable"
SET assigned_team_ids = array_remove(assigned_team_ids, $1)
WHERE $1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]))
AND NOT (access_group_id = ANY($2::TEXT[]))
RETURNING access_group_id
"""
class _AffectedGroup(BaseModel):
access_group_id: str
class _TeamGroups(BaseModel):
access_group_ids: tuple[str, ...] | None = None
_AffectedGroups: Final = TypeAdapter(tuple[_AffectedGroup, ...])
_TeamRows: Final = TypeAdapter(tuple[_TeamGroups, ...])
class AccessGroupSyncTx(Protocol):
async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ...
class _Transaction(Protocol):
async def __aenter__(self) -> AccessGroupSyncTx: ...
async def __aexit__(self, *exc_info: object) -> None: ...
class _PrismaDb(Protocol):
def tx(self) -> _Transaction: ...
class _PrismaClient(Protocol):
@property
def db(self) -> _PrismaDb: ...
async def invalidate_access_group_cache(access_group_id: str) -> None:
"""
Drop an access group entry from both the in-memory and Redis caches.
Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server
to avoid circular imports, following the same pattern as key_management_endpoints.
"""
from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
await _delete_cache_access_object(
access_group_id=access_group_id,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
async def invalidate_access_group_caches(access_group_ids: Sequence[str]) -> None:
"""
Drop every given access group from the caches, then raise if any drop failed.
Every entry is attempted even when one raises, so a single unreachable cache cannot
leave the rest of the reconciled groups serving a grant the admin revoked.
"""
outcomes: Final = await asyncio.gather(
*(invalidate_access_group_cache(access_group_id) for access_group_id in access_group_ids),
return_exceptions=True,
)
for outcome in outcomes:
if isinstance(outcome, BaseException):
raise outcome
async def reconcile_team_access_group_membership(tx: AccessGroupSyncTx, team_id: str) -> tuple[str, ...]:
"""
Reconcile every access group's `assigned_team_ids` against the team's own
`access_group_ids`, and return the groups whose cache the caller has to drop once the
transaction commits.
Call this inside the transaction that writes the team row, or after that row is
written or deleted: a team with no row reconciles to an empty set, which detaches it
from every group.
The team row is read here rather than passed in, under an advisory lock held for the
rest of the transaction. That is what makes concurrent writes to the same team
converge, since each mirror reconciles against the row as the transaction sees it
instead of against the snapshot its own caller happened to see. It also means a retry
heals a sync that failed partway, where a before/after delta would compute nothing.
Both mirror statements are set-based and mutate the array inside the statement, so a
concurrent write for a different team cannot be lost the way a read-modify-write of
the whole array can, and the pair commits together or not at all.
"""
await tx.query_raw(_LOCK_TEAM_SQL, team_id)
team_rows: Final = _TeamRows.validate_python(await tx.query_raw(_READ_TEAM_SQL, team_id))
desired: Final = (team_rows[0].access_group_ids or ()) if team_rows else ()
affected: Final = _AffectedGroups.validate_python(await tx.query_raw(_AFFECTED_SQL, team_id, desired))
await tx.query_raw(_ATTACH_SQL, team_id, desired)
await tx.query_raw(_DETACH_SQL, team_id, desired)
return tuple(group.access_group_id for group in affected)
async def sync_team_access_group_membership(prisma_client: _PrismaClient, team_id: str) -> None:
"""Reconcile the mirror for an already committed team write, in its own transaction."""
async with prisma_client.db.tx() as tx:
affected: Final = await reconcile_team_access_group_membership(tx, team_id)
await invalidate_access_group_caches(affected)

2
litellm/proxy/openai_files_endpoints/files_endpoints.py
View file

@ -1352,7 +1352,7 @@ async def list_files(
if should_route and credentials is not None:
# Use model-based routing with credentials from config
data.update(credentials)
prepare_data_with_credentials(data=data, credentials=credentials)
response = await litellm.afile_list(
custom_llm_provider=credentials["custom_llm_provider"],
purpose=purpose,

75
litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
View file

@ -1,5 +1,6 @@
import asyncio
import json
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, cast
@ -7,6 +8,7 @@ import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import ANTHROPIC_BATCHES_ROUTE
from litellm.litellm_core_utils.core_helpers import map_finish_reason
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model
@ -20,6 +22,12 @@ from litellm.llms.anthropic.chat.handler import (
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
is_collection_route,
log_batch_registration_result,
optional_str,
request_tags_from_metadata,
)
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
PassthroughStandardLoggingPayload,
)
@ -74,6 +82,9 @@ class AnthropicPassthroughLoggingHandler:
)
model: Final = response_body.get("model", "")
speed: Final = AnthropicPassthroughLoggingHandler._cost_relevant_speed(
request_body or kwargs.get("request_body")
)
anthropic_config: Final = get_anthropic_config(url_route)
litellm_model_response: Final[ModelResponse] = anthropic_config().transform_response(
raw_response=httpx_response,
@ -81,7 +92,7 @@ class AnthropicPassthroughLoggingHandler:
model=model,
messages=[],
logging_obj=logging_obj,
optional_params={},
optional_params={"speed": speed} if speed else {},
api_key="",
request_data={},
encoding=litellm.encoding,
@ -103,6 +114,15 @@ class AnthropicPassthroughLoggingHandler:
"kwargs": kwargs,
}
@staticmethod
def _cost_relevant_speed(request_body: Mapping[str, object] | None) -> str | None:
"""
Anthropic's ``speed=fast`` multiplies non-cache token cost, and only the request
carries it, so it has to reach the usage-building paths for spend to be right.
"""
speed: Final = (request_body or {}).get("speed")
return speed if isinstance(speed, str) else None
@staticmethod
def _get_user_from_metadata(
passthrough_logging_payload: PassthroughStandardLoggingPayload,
@ -316,6 +336,7 @@ class AnthropicPassthroughLoggingHandler:
- Logs in litellm callbacks
"""
speed: Final = AnthropicPassthroughLoggingHandler._cost_relevant_speed(request_body)
model = request_body.get("model", "")
# Check if it's available in the logging object
if (
@ -335,6 +356,7 @@ class AnthropicPassthroughLoggingHandler:
all_chunks=all_chunks,
litellm_logging_obj=litellm_logging_obj,
model=model,
speed=speed,
)
except Exception as e:
# stream_chunk_builder re-raises assembly failures (as litellm.APIError)
@ -356,6 +378,7 @@ class AnthropicPassthroughLoggingHandler:
complete_streaming_response = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks(
all_chunks=all_chunks,
model=model,
speed=speed,
)
except Exception as e:
verbose_proxy_logger.warning(
@ -420,6 +443,7 @@ class AnthropicPassthroughLoggingHandler:
all_chunks: Sequence[str | bytes],
litellm_logging_obj: LiteLLMLoggingObj,
model: str,
speed: str | None = None,
) -> ModelResponse | TextCompletionResponse | None:
"""
Builds complete response from raw Anthropic chunks.
@ -444,11 +468,13 @@ class AnthropicPassthroughLoggingHandler:
all_chunks=collapsed,
litellm_logging_obj=litellm_logging_obj,
model=model,
speed=speed,
)
return AnthropicPassthroughLoggingHandler._build_complete_streaming_response_legacy(
all_chunks=all_chunks,
litellm_logging_obj=litellm_logging_obj,
model=model,
speed=speed,
)
# Anthropic SSE block/delta types that the fast path is NOT allowed to
@ -576,6 +602,7 @@ class AnthropicPassthroughLoggingHandler:
all_chunks: Sequence[str | bytes],
litellm_logging_obj: LiteLLMLoggingObj,
model: str,
speed: str | None = None,
) -> ModelResponse | TextCompletionResponse | None:
"""
Original reconstruction: convert every SSE event to a generic chunk
@ -591,6 +618,7 @@ class AnthropicPassthroughLoggingHandler:
anthropic_model_response_iterator: Final = AnthropicModelResponseIterator(
streaming_response=None,
sync_stream=False,
speed=speed,
)
all_openai_chunks: Final = []
@ -650,6 +678,7 @@ class AnthropicPassthroughLoggingHandler:
def _build_usage_only_response_from_chunks(
all_chunks: Sequence[str | bytes],
model: str,
speed: str | None = None,
) -> ModelResponse | None:
"""
Build a usage-bearing ModelResponse from Anthropic SSE token-usage events, for
@ -743,7 +772,9 @@ class AnthropicPassthroughLoggingHandler:
usage_object["server_tool_use"] = _server_tool_use
if inference_geo is not None:
usage_object["inference_geo"] = inference_geo
usage_obj: Final = AnthropicConfig().calculate_usage(usage_object=usage_object, reasoning_content=None)
usage_obj: Final = AnthropicConfig().calculate_usage(
usage_object=usage_object, reasoning_content=None, speed=speed
)
return ModelResponse(
model=resolved_model,
choices=[
@ -833,13 +864,14 @@ class AnthropicPassthroughLoggingHandler:
# Store the managed object for cost tracking
# This will be picked up by check_batch_cost polling mechanism
AnthropicPassthroughLoggingHandler._store_batch_managed_object(
unified_object_id=unified_object_id,
batch_object=litellm_batch_response,
model_object_id=batch_id,
logging_obj=logging_obj,
**kwargs,
)
if is_collection_route(url_route, ANTHROPIC_BATCHES_ROUTE):
AnthropicPassthroughLoggingHandler._store_batch_managed_object(
unified_object_id=unified_object_id,
batch_object=litellm_batch_response,
model_object_id=batch_id,
logging_obj=logging_obj,
**kwargs,
)
# Create a batch job response for logging
litellm_model_response = ModelResponse()
@ -964,8 +996,12 @@ class AnthropicPassthroughLoggingHandler:
**kwargs,
) -> None:
"""
Store batch managed object for cost tracking.
Register a newly created batch for cost tracking.
This will be picked up by the check_batch_cost polling mechanism.
Only the create reaches here, so the row records the creating key and its tags.
An id-scoped route cannot rebuild the unified object id anyway: the model comes
from the create's request body, which a retrieve does not have.
"""
try:
# Get the managed files hook from the logging object
@ -981,7 +1017,7 @@ class AnthropicPassthroughLoggingHandler:
user_api_key_dict: Final = UserAPIKeyAuth(
user_id=_request_metadata.get("user_api_key_user_id", "default-user"),
api_key="",
api_key=optional_str(_request_metadata.get("user_api_key")),
team_id=_request_metadata.get("user_api_key_team_id"),
team_alias=None,
user_role=LitellmUserRoles.CUSTOMER, # Use proper enum value
@ -1003,9 +1039,7 @@ class AnthropicPassthroughLoggingHandler:
)
# Store the unified object for batch cost tracking
import asyncio
asyncio.create_task(
task: Final = asyncio.create_task(
managed_files_hook.store_unified_object_id(
unified_object_id=unified_object_id,
file_object=batch_object,
@ -1013,13 +1047,14 @@ class AnthropicPassthroughLoggingHandler:
model_object_id=model_object_id,
file_purpose="batch",
user_api_key_dict=user_api_key_dict,
request_tags=request_tags_from_metadata(_request_metadata),
persist_attribution=True,
)
)
verbose_proxy_logger.info(
"Stored Anthropic batch managed object with unified_object_id=%s, batch_id=%s",
unified_object_id,
model_object_id,
task.add_done_callback(
lambda finished: log_batch_registration_result(
finished, "Anthropic", unified_object_id, model_object_id, is_batch_create=True
)
)
else:
verbose_proxy_logger.warning(

79
litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py Normal file
View file

@ -0,0 +1,79 @@
"""Spend attribution for batches created through a passthrough endpoint.
The creating key and its tags are read off the passthrough request's metadata and
persisted on the managed object row, because the batch cost lands hours later in a
background poll that has no request to read them from.
"""
import asyncio
from collections.abc import Mapping, Sequence
from typing import Final
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.safe_json_dumps import strip_null_bytes
def optional_str(value: object) -> str | None:
return value if isinstance(value, str) else None
def _sanitized_str_tuple(value: object) -> tuple[str, ...] | None:
if not isinstance(value, list):
return None
items: Final[Sequence[object]] = value
return tuple(strip_null_bytes(tag) for tag in items if isinstance(tag, str))
def is_collection_route(url_route: str, collection_suffix: str) -> bool:
"""Whether the route addresses the batch collection itself rather than one batch.
A POST to the collection is the create; every id-scoped route is a retrieve,
results or cancel.
"""
return url_route.split("?")[0].rstrip("/").endswith(collection_suffix)
def request_tags_from_metadata(request_metadata: Mapping[str, object]) -> tuple[str, ...] | None:
"""Tags for the batch-cost spend row: the request's own tags when it sent any,
otherwise the key's tags, which auth exposes as user_api_key_auth_metadata (a
tagged key does not put its tags in the top-level metadata "tags" on the
passthrough path)
"""
tags: Final = _sanitized_str_tuple(request_metadata.get("tags"))
if tags:
return tags
key_auth_metadata: Final = request_metadata.get("user_api_key_auth_metadata")
if isinstance(key_auth_metadata, dict):
return _sanitized_str_tuple(key_auth_metadata.get("tags"))
return None
def log_batch_registration_result(
finished: asyncio.Task[None],
provider: str,
unified_object_id: str,
model_object_id: str,
is_batch_create: bool,
) -> None:
"""Report the outcome of the fire-and-forget managed object write. A create that
fails is not retried by a later poll, so its cost is never tracked at all.
"""
error: Final = finished.exception() if not finished.cancelled() else None
if finished.cancelled() or error is not None:
consequence: Final = (
"its cost will not be tracked" if is_batch_create else "its status and output file may be stale"
)
verbose_proxy_logger.error(
"Failed to store %s batch managed object with unified_object_id=%s, batch_id=%s; %s: %s",
provider,
unified_object_id,
model_object_id,
consequence,
error,
)
return
verbose_proxy_logger.info(
"Stored %s batch managed object with unified_object_id=%s, batch_id=%s",
provider,
unified_object_id,
model_object_id,
)

2
litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py
View file

@ -82,7 +82,7 @@ class CoherePassthroughLoggingHandler(BasePassthroughLoggingHandler):
Handle Cohere passthrough logging with route detection and cost tracking.
"""
# Check if this is an embed endpoint
if "/v1/embed" in url_route:
if "/v1/embed" in url_route and "/v1/embeddings" not in url_route:
model: Final = request_body.get("model", response_body.get("model", ""))
try:
cohere_embed_config: Final = CohereEmbeddingConfig()

Some files were not shown because too many files have changed in this diff Show more