scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-03 07:23:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 166

Merge remote-tracking branch 'origin/main' into litellm_fix_responses_stream_failure_billing

Browse source 
# Conflicts:
#	litellm/responses/streaming_iterator.py
This commit is contained in:
mateo-berri 2026-10-20 21:03:11 -07:00
commit e3e5166d82
346 changed files with 24392 additions and 4900 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.circleci/config.yml
View file

@ -1291,7 +1291,7 @@ jobs:
name: Install MCP SDK1 peer
command: |
uv venv --python 3.12 .venv-mcp-peer
uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
uv pip install --no-config --exclude-newer "3 days" --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV"
- when:
condition:

2
.circleci/scripts/classify_changes.sh
View file

@ -33,7 +33,7 @@ while IFS= read -r file || [ -n "$file" ]; do
*) has_backend=true ;;
esac
case "$file" in
model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json)
model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json)
has_cost_map=true ;;
tests/unit/proxy/*) : ;;
*) outside_cost_map_set=true ;;

8
.github/workflows/test-unit.yml vendored
View file

@ -139,6 +139,7 @@ jobs:
test-path: |-
tests/unit/batch_completion
tests/unit/decisions
tests/unit/extraction
tests/unit/litellm_core_utils
tests/unit/systemone
python-version: "3.12"
@ -433,6 +434,13 @@ jobs:
- shard: rust-bridge-harness
test-path: >-
tests/test_litellm_rust/test_dispatch.py
tests/test_litellm_rust/test_public_entries.py
tests/test_litellm_rust/chat_completions/test_dispatch.py
tests/test_litellm_rust/embeddings/test_dispatch.py
tests/test_litellm_rust/messages/test_dispatch.py
tests/test_litellm_rust/ocr/test_dispatch.py
tests/test_litellm_rust/responses/test_dispatch.py
tests/unit/compression
tests/unit/harness
tests/unit/rust_bridge

6
ci_cd/cost_map_guard.py
View file

@ -4,8 +4,7 @@ Every pull request whose diff against its merge base touches one of the two cost
checks: the files parse and the backup copy matches the root file. A pull request that leaves both untouched
skips them, since merging it keeps the base branch's copies and its head tree only carries whatever state the
branch was cut from. Pull requests from the cost map sync bot (branches named litellm_cost_map_sync_*) always
get the file checks and additionally may only touch those two files or the legacy schema
path and may only add or update models.
get the file checks and additionally may only touch those two files and may only add or update models.
"""
from __future__ import annotations
@ -21,7 +20,6 @@ from typing import Final
COST_MAP_PATH: Final = "model_prices_and_context_window.json"
BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
LEGACY_SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH)
BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_"
SPECIAL_ROOT_KEYS: Final = frozenset({"sample_spec", "fallback_generalizations"})
@ -110,7 +108,7 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str
*(
f"bot PRs may only change the cost map files, not {path}"
for path in changed_files
if path not in GUARDED_PATHS and path != LEGACY_SCHEMA_PATH
if path not in GUARDED_PATHS
),
*(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map),
*(f"bot PRs may not remove fields: {ref}" for ref in removed_fields),

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

@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tupl
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
ManagedObjectRow,
row_metadata,
)
import litellm
@ -90,11 +91,6 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li
return TeamRepository(prisma_client).table
def _row_metadata(row: object) -> dict[str, object]:
metadata: Final = getattr(row, "metadata", None)
return cast(dict[str, object], metadata) if isinstance(metadata, dict) else {} # cast-ok: prisma Json object column
class CheckBatchCost:
def __init__(
self,
@ -1063,7 +1059,9 @@ class CheckBatchCost:
if (has_key_row and key_row is None) or (team_id and team_row is None):
return True
settings: Final = get_dynamic_logging_metadata(
UserAPIKeyAuth(metadata=_row_metadata(key_row), team_metadata=_row_metadata(team_row), team_id=team_id),
UserAPIKeyAuth(
metadata=dict(row_metadata(key_row)), team_metadata=dict(row_metadata(team_row)), team_id=team_id
),
proxy_config,
)
callback_vars: Final = (settings.callback_vars if settings is not None else None) or {}

91
enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
View file

@ -7,6 +7,12 @@ attribution, priced against the deployment that served it. The poll's status che
with BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN and client reads are marked managed by the
managed files hook, so both stay free. The spend event carries the id the client holds, the row's
unified_object_id, so its spend log is found under that id with the _background_cost suffix.
The spend event reaches the creating key's or team's own logging callbacks the way that creator's
requests do. Its message content is redacted when the creator's key or team logging turns message
logging off, else as the proxy is set, the way the batch poll decides it. The row is claimed before
the creator is read, so reads that lose the race skip the lookups. A creator whose lookup fails, or
whose logging settings the logger rejects, is still billed, through the proxy's own callbacks only
and with its message content redacted.
A read that finds no row, because the create could not store one, stores it under the reading key
and bills it the same way, so a response the create left untracked is still billed once. A streamed
background create is free on its own; its stream stores the row under the creating key and bills
@ -29,6 +35,7 @@ from typing import TYPE_CHECKING, Final, cast
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
ManagedObjectRow,
ResponseCreator,
)
from pydantic import BaseModel, ConfigDict, TypeAdapter
@ -54,11 +61,13 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.proxy.proxy_server import ProxyConfig
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.prisma_protocols import TableActions
from litellm.router import Router
_STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
_CALLBACK_LIST_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...])
class _RoutedModelInfo(BaseModel):
@ -73,6 +82,8 @@ class _RoutedDeployment(BaseModel):
_ROUTED_DEPLOYMENTS: Final = TypeAdapter(tuple[_RoutedDeployment, ...])
_RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses"
_CALLBACK_LIST_FIELDS: Final = frozenset({"success_callback", "failure_callback", "litellm_disabled_callbacks"})
_REJECTED_CREATOR_LOGGING: Final[Mapping[str, object]] = MappingProxyType({"turn_off_message_logging": True})
@dataclass(frozen=True, slots=True)
@ -89,6 +100,7 @@ class BilledBackgroundResponse:
custom_llm_provider: str | None
litellm_params: Mapping[str, object]
metadata: Mapping[str, object]
creator_logging: Mapping[str, object]
@dataclass(frozen=True, slots=True)
@ -107,11 +119,51 @@ def _without_read_price(response: ResponsesAPIResponse) -> ResponsesAPIResponse:
return unpriced
async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
def _string_callbacks(callbacks: object) -> tuple[str, ...]:
if not isinstance(callbacks, list):
return ()
return tuple(
callback for callback in _CALLBACK_LIST_ADAPTER.validate_python(callbacks) if isinstance(callback, str)
)
logging_obj: Final = LiteLLMLogging(
def _running_proxy_config() -> "ProxyConfig":
from litellm.proxy.proxy_server import proxy_config
return proxy_config
def creator_logging_fields(creator: ResponseCreator, proxy_config: "ProxyConfig | None") -> Mapping[str, object]:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import key_and_team_logging_fields
if not creator.settings_read:
return _REJECTED_CREATOR_LOGGING
try:
return key_and_team_logging_fields(
user_api_key_dict=UserAPIKeyAuth(
metadata=dict(creator.key_metadata),
team_id=creator.team_id,
team_metadata=dict(creator.team_metadata),
),
proxy_config=proxy_config if proxy_config is not None else _running_proxy_config(),
)
except Exception as e:
verbose_proxy_logger.error(
f"CheckResponsesCost: could not read the creator's logging settings, so only the proxy's "
f"callbacks get the billed event: {e}"
)
return _REJECTED_CREATOR_LOGGING
def _billed_logging_obj(billed: BilledBackgroundResponse, creator_logging: Mapping[str, object]) -> "LiteLLMLogging":
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.utils import split_dynamic_success_callbacks
sync_callbacks, async_callbacks = split_dynamic_success_callbacks(
_string_callbacks(creator_logging.get("success_callback"))
)
return LiteLLMLogging(
model=billed.model,
messages=[{"role": "user", "content": "<retrieve_response>"}],
stream=False,
@ -119,7 +171,27 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi
start_time=datetime.now(),
litellm_call_id=str(uuid.uuid4()),
function_id=str(uuid.uuid4()),
dynamic_success_callbacks=list(sync_callbacks),
dynamic_async_success_callbacks=list(async_callbacks) or None,
kwargs={key: value for key, value in creator_logging.items() if key not in _CALLBACK_LIST_FIELDS},
)
def _creator_or_proxy_logging_obj(billed: BilledBackgroundResponse) -> "LiteLLMLogging":
try:
return _billed_logging_obj(billed, billed.creator_logging)
except Exception as e:
verbose_proxy_logger.error(
f"CheckResponsesCost: the creator's logging settings were rejected, so only the proxy's "
f"callbacks get the billed event: {e}"
)
return _billed_logging_obj(billed, _REJECTED_CREATOR_LOGGING)
async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge:
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
logging_obj: Final = _creator_or_proxy_logging_obj(billed)
logging_obj.update_environment_variables(
litellm_params={**billed.litellm_params, "metadata": dict(billed.metadata)},
optional_params={},
@ -187,6 +259,7 @@ class CheckResponsesCost:
prisma_client: "PrismaClient",
llm_router: "Router",
log_billed_response: BilledResponseLogger = log_billed_background_response,
proxy_config: "ProxyConfig | None" = None,
):
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
@ -196,6 +269,7 @@ class CheckResponsesCost:
self.llm_router: Router = llm_router
self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckResponsesCost")
self._log_billed_response: Final = log_billed_response
self._proxy_config: Final = proxy_config
def _resolve_deployment(self, response_id: str) -> Deployment | None:
model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
@ -307,11 +381,12 @@ class CheckResponsesCost:
data={"status": job.status or PENDING_RESPONSE_STATUSES[0]},
)
async def _billed_response(
def _billed_response(
self,
job: ManagedObjectRow,
target: _PollTarget,
response: ResponsesAPIResponse,
creator: ResponseCreator,
) -> BilledBackgroundResponse:
deployment: Final = target.deployment
return BilledBackgroundResponse(
@ -323,13 +398,14 @@ class CheckResponsesCost:
litellm_params=_deployment_litellm_params(deployment) if deployment is not None else MappingProxyType({}),
metadata=MappingProxyType(
{
**(await self._creator_attribution.response_metadata(job, job.unified_object_id)),
**creator.metadata,
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
**target.model_metadata,
**(_deployment_metadata(deployment) if deployment is not None else {}),
**self._requested_model_group(target, job.team_id),
}
),
creator_logging=creator_logging_fields(creator, self._proxy_config),
)
def _requested_model_group(self, target: _PollTarget, team_id: str | None) -> Mapping[str, object]:
@ -354,7 +430,8 @@ class CheckResponsesCost:
verbose_proxy_logger.info(f"Response {job.unified_object_id} was already billed elsewhere, skipping")
return None
try:
charge: Final = await self._log_billed_response(await self._billed_response(job, target, response))
creator: Final = await self._creator_attribution.response_creator(job, job.unified_object_id)
charge: Final = await self._log_billed_response(self._billed_response(job, target, response, creator))
except BaseException:
await self._release_billing_claim(job)
raise

185
enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py
View file

@ -10,12 +10,16 @@ from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import TYPE_CHECKING, Final, Protocol, cast
from typing import TYPE_CHECKING, Final, Protocol, TypeGuard, cast
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.constants import CLI_SESSION_KEY_PREFIX
from litellm.repositories.table_repositories import (
DeletedVerificationTokenRepository,
DeprecatedVerificationTokenRepository,
)
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
@ -70,6 +74,7 @@ class ManagedObjectRow(Protocol):
_TAG_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object])
_KEY_OWNER_SCOPE_PREFIX: Final = "key:"
_MAX_ROTATION_HOPS: Final = 16
def _string_tags(request_tags: object) -> list[str] | None:
@ -96,6 +101,20 @@ def _team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteL
return TeamRepository(prisma_client).table
def _rotated_token_table(
prisma_client: PrismaClient,
) -> TableActions[prisma_models.LiteLLM_DeprecatedVerificationToken]:
return DeprecatedVerificationTokenRepository(prisma_client).table
def _deleted_token_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_DeletedVerificationToken]:
return DeletedVerificationTokenRepository(prisma_client).table
def _deleted_team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_DeletedTeamTable]:
return TeamRepository(prisma_client).deleted_table
@dataclass(frozen=True, slots=True)
class _CreatorRows:
key: prisma_models.LiteLLM_VerificationToken | None
@ -103,6 +122,15 @@ class _CreatorRows:
user: prisma_models.LiteLLM_UserTable | None
@dataclass(frozen=True, slots=True)
class ResponseCreator:
metadata: Mapping[str, object]
key_metadata: Mapping[str, object]
team_id: str | None
team_metadata: Mapping[str, object]
settings_read: bool = True
class ManagedObjectCreatorAttribution:
"""Rebuilds the spend-tracking metadata of whoever created a managed object, so the poll
job's spend row is attributed like the creating request: key hash, team, organization,
@ -117,7 +145,12 @@ class ManagedObjectCreatorAttribution:
A background response also carries its end user, the creating key's project and its alias,
and the key, team, and user budgets and per-model budgets, the way the read that billed it
before did, and never names a key-scoped creator as its user. A key that sets no budget of
its own takes its budget tier's, the same rule auth applies.
its own takes its budget tier's, the same rule auth applies. Its creator's key and team
metadata come along too, so the billed event reaches the key's and team's logging callbacks
like that read did. A rotated key's metadata is read from the live key its rotations lead
to, and a deleted key's or team's from its deleted row. A lookup that fails leaves the
creator's settings unread and the attribution to the rows that still load, so the
response still bills and still counts toward their budgets.
"""
def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None:
@ -128,31 +161,70 @@ class ManagedObjectCreatorAttribution:
rows: Final = await self._rows(job, object_id)
return {"user_api_key_user_id": job.created_by, **_shared_metadata(job, rows)}
async def response_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]:
rows: Final = await self._rows(job, object_id)
end_user_id: Final = getattr(job, "end_user_id", None)
return {
"user_api_key_user_id": _creator_user_id(job.created_by),
**_shared_metadata(job, rows),
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
"user_api_key_project_id": getattr(rows.key, "project_id", None),
"user_api_key_project_alias": getattr(
getattr(rows.key, "litellm_project_table", None), "project_alias", None
async def response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator:
try:
return await self._read_response_creator(job, object_id)
except Exception as e:
verbose_proxy_logger.error(
f"{self._poller_name}: could not read the creator of {object_id}, so it bills without their "
f"settings: {e}"
)
return ResponseCreator(
metadata=_response_creator_metadata(job, await self._rows(job, object_id)),
key_metadata={},
team_id=getattr(job, "team_id", None),
team_metadata={},
settings_read=False,
)
async def _read_response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator:
api_key: Final = getattr(job, "api_key", None)
team_id: Final = getattr(job, "team_id", None)
rows: Final = _CreatorRows(
key=await self._live_key_row(api_key, job.created_by),
team=await self._live_team_row(team_id),
user=await self._user_row(object_id, _creator_user_id(job.created_by)),
)
return ResponseCreator(
metadata=_response_creator_metadata(job, rows),
key_metadata=row_metadata(rows.key) if rows.key is not None else await self._retired_key_metadata(api_key),
team_id=team_id,
team_metadata=(
row_metadata(rows.team) if rows.team is not None else await self._deleted_team_metadata(team_id)
),
**_key_budget_fields(rows.key),
**_budget_fields(
rows.team,
"user_api_key_team_spend",
"user_api_key_team_max_budget",
"user_api_key_team_model_max_budget",
),
**_budget_fields(
rows.user,
"user_api_key_user_spend",
"user_api_key_user_max_budget",
"user_api_key_user_model_max_budget",
),
}
)
async def _retired_key_metadata(self, api_key: str | None) -> Mapping[str, object]:
if not _names_a_key_row(api_key):
return {}
current_token: Final = await self._token_rotated_to(api_key, _MAX_ROTATION_HOPS)
current_key: Final = (
await _token_table(self._prisma_client).find_unique(where={"token": current_token})
if current_token != api_key
else None
)
if current_key is not None:
return row_metadata(current_key)
return row_metadata(
await _deleted_token_table(self._prisma_client).find_first(
where={"token": current_token}, order={"deleted_at": "desc"}
)
)
async def _token_rotated_to(self, token: str, hops_left: int) -> str:
rotation: Final = await _rotated_token_table(self._prisma_client).find_unique(where={"token": token})
if rotation is None or hops_left == 0:
return token
return await self._token_rotated_to(rotation.active_token_id, hops_left - 1)
async def _deleted_team_metadata(self, team_id: str | None) -> Mapping[str, object]:
if not team_id:
return {}
return row_metadata(
await _deleted_team_table(self._prisma_client).find_first(
where={"team_id": team_id}, order={"deleted_at": "desc"}
)
)
async def _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows:
return _CreatorRows(
@ -175,25 +247,57 @@ class ManagedObjectCreatorAttribution:
async def _key_row(
self, object_id: str, api_key: str | None, created_by: str | None
) -> prisma_models.LiteLLM_VerificationToken | None:
if not api_key or _is_cli_session_key(api_key, created_by):
return None
try:
return await _token_table(self._prisma_client).find_unique(
where={"token": api_key}, include={"litellm_budget_table": True, "litellm_project_table": True}
)
return await self._live_key_row(api_key, created_by)
except Exception as e:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up the creating key for {object_id}: {e}")
return None
async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
if not team_id:
async def _live_key_row(
self, api_key: str | None, created_by: str | None
) -> prisma_models.LiteLLM_VerificationToken | None:
if not api_key or _is_cli_session_key(api_key, created_by):
return None
return await _token_table(self._prisma_client).find_unique(
where={"token": api_key}, include={"litellm_budget_table": True, "litellm_project_table": True}
)
async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
try:
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
return await self._live_team_row(team_id)
except Exception as e:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}")
return None
async def _live_team_row(self, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
if not team_id:
return None
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
def _response_creator_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
end_user_id: Final = getattr(job, "end_user_id", None)
return {
"user_api_key_user_id": _creator_user_id(job.created_by),
**_shared_metadata(job, rows),
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
"user_api_key_project_id": getattr(rows.key, "project_id", None),
"user_api_key_project_alias": getattr(getattr(rows.key, "litellm_project_table", None), "project_alias", None),
**_key_budget_fields(rows.key),
**_budget_fields(
rows.team,
"user_api_key_team_spend",
"user_api_key_team_max_budget",
"user_api_key_team_model_max_budget",
),
**_budget_fields(
rows.user,
"user_api_key_user_spend",
"user_api_key_user_max_budget",
"user_api_key_user_model_max_budget",
),
}
def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
api_key: Final = getattr(job, "api_key", None)
@ -221,6 +325,19 @@ def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str,
}
def row_metadata(row: object) -> Mapping[str, object]:
metadata: Final = getattr(row, "metadata", None)
if not isinstance(metadata, dict):
return {}
return cast(Mapping[str, object], metadata) # cast-ok: prisma Json object column
def _names_a_key_row(api_key: object) -> TypeGuard[str]:
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
return isinstance(api_key, str) and is_valid_sha256_hash(api_key)
def _is_cli_session_key(api_key: str, created_by: str | None) -> bool:
return bool(created_by) and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}"

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.78"
version = "0.1.79"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.78"
version = "0.1.79"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

4
gateway/routes/allowlist.py
View file

@ -1,7 +1,7 @@
"""Path allowlist for the gateway component.
The gateway exposes the LLM data-plane surface: chat/completions, embeddings,
audio, batches, files, fine-tuning, rerank, decisions, ocr, rag, video, search, image,
audio, batches, files, fine-tuning, rerank, decisions, extraction, ocr, rag, video, search, image,
responses, vector stores, passthrough providers, realtime websockets, MCP
tool-call endpoints, and operational endpoints (/health, /metrics, and the
/debug/memory/summary read of the serving worker's RSS).
@ -64,6 +64,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/decisions",
"/v1/systemone",
"/systemone",
"/v1/extraction",
"/extraction",
"/v1/ocr",
"/ocr",
"/v1/rag/",

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.110"
version = "0.4.111"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -30,7 +30,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.110"
version = "0.4.111"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

2
litellm-rust/crates/python-bridge/src/cache/mod.rs vendored
View file

@ -4,7 +4,7 @@ mod selection;
pub(crate) use native::NativeCacheHandle;
pub(crate) use python::{CacheCall, PythonCache};
pub(crate) use selection::{Cached, Selection, admit_native, configure, configured_native};
pub(crate) use selection::{Cached, Selection, configure, configured_native};
use litellm_cache::Error;
use pyo3::{

2
litellm-rust/crates/python-bridge/src/cache/native/v2.rs vendored
View file

@ -279,7 +279,7 @@ pub(in crate::cache) fn configured(
)> {
let handle = native_handle(configured)?.ok_or_else(|| {
pyo3::exceptions::PyRuntimeError::new_err(
"the configured cache changed to a Python cache after native admission",
"native inference needs a v2 cache, the catalog keeps Python caches on Python",
)
})?;
let cache = handle.extract::<PyRef<'_, NativeCacheHandle>>()?;

15
litellm-rust/crates/python-bridge/src/cache/selection.rs vendored
View file

@ -76,21 +76,6 @@ fn selected_cache<'py>(
Ok(Some(configured))
}
pub(crate) fn admit_native(
py: Python<'_>,
kwargs: &Bound<'_, PyDict>,
call_type: &str,
) -> PyResult<()> {
if let Some(configured) = selected_cache(py, kwargs, call_type)?
&& native::v2::native_handle(&configured)?.is_none()
{
return Err(crate::errors::RustBridgeDeclined::new_err(
"the configured cache requires Python inference",
));
}
Ok(())
}
pub(crate) fn configured_native(
py: Python<'_>,
kwargs: &Bound<'_, PyDict>,

1
litellm-rust/crates/python-bridge/src/routes/inference.rs
View file

@ -177,7 +177,6 @@ where
} else {
R::SYNC_CALL_TYPE
};
crate::cache::admit_native(py, &call.kwargs, call_type)?;
let (arguments, hooks) = super::call_hooks(py, R::OPERATION, &call, asynchronous)?;
super::run_public_call(
py,

40
litellm-rust/crates/python-bridge/src/routes/responses/mod.rs
View file

@ -8,59 +8,19 @@ use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, ScopedCache};
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_host::{call::HostedMachine, protocol::Protocol};
use litellm_host_python::present;
use litellm_inference_responses::{ResponsesRoute, route::Responses};
use litellm_secrets::source::SecretSource;
use pyo3::prelude::*;
use serde_json::Value;
pub(crate) use websocket::ResponsesWebSocketConnection;
use super::{
NativeCall,
inference::{InferenceHost, InferenceRoute, run_inference},
};
use crate::errors::RustBridgeDeclined;
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.responses";
fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult<Py<PyAny>> {
let resolved = call.resolved()?;
if let Some(reason) = py
.import(ROUTE_HOST_MODULE)?
.getattr("decline_reason")?
.call1((&resolved,))?
.extract::<Option<String>>()?
{
return Err(RustBridgeDeclined::new_err(reason));
}
let argument = |name: &str| present(&call.kwargs, &call.base, name);
let model = argument("model")?
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("model is required"))?
.extract::<String>()?;
let provider = argument("custom_llm_provider")?
.map(|value| value.extract::<String>())
.transpose()?;
if provider
.as_deref()
.is_some_and(|provider| provider != "openai")
|| model
.strip_prefix("openai/")
.unwrap_or(&model)
.contains('/')
{
return Err(RustBridgeDeclined::new_err(
"native HTTP responses provider",
));
}
if argument("stream")?
.map(|value| litellm_host_python::from_py::<Value>(&value))
.transpose()?
.is_some_and(|value| value == Value::Bool(true))
{
return Err(RustBridgeDeclined::new_err(
"native Python responses streaming",
));
}
let host = InferenceHost::new(call.view()?, ROUTE_HOST_MODULE);
run_inference::<ResponsesRoute, _>(py, call, asynchronous, ResponsesPythonHost::new(host))
}

4
litellm/__init__.py
View file

@ -1505,6 +1505,7 @@ from .rag.main import *
from .sandbox.main import *
from .decisions.main import *
from .systemone.main import *
from .extraction.main import *
from .tool_loop import ToolLoopMaxRoundsExceeded, arun_tool_loop, run_tool_loop
from .search.main import *
from .realtime_api.main import (
@ -1711,6 +1712,9 @@ if TYPE_CHECKING:
from .llms.cloudflare.decisions.transformation import (
CloudflareDecisionsConfig as CloudflareDecisionsConfig,
)
from .llms.cloudflare.extraction.transformation import (
CloudflareExtractionConfig as CloudflareExtractionConfig,
)
from .llms.strands_decider.decisions.transformation import (
StrandsDeciderDecisionsConfig as StrandsDeciderDecisionsConfig,
)

2
litellm/_lazy_imports_registry.py
View file

@ -166,6 +166,7 @@ LLM_CONFIG_NAMES: Final = (
"OpenRouterDecisionsConfig",
"FireworksAIDecisionsConfig",
"CloudflareDecisionsConfig",
"CloudflareExtractionConfig",
"StrandsDeciderDecisionsConfig",
"DatabricksDecisionsConfig",
"HostedVLLMDecisionsConfig",
@ -735,6 +736,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
"OpenRouterDecisionsConfig": (".llms.openrouter.decisions.transformation", "OpenRouterDecisionsConfig"),
"FireworksAIDecisionsConfig": (".llms.fireworks_ai.decisions.transformation", "FireworksAIDecisionsConfig"),
"CloudflareDecisionsConfig": (".llms.cloudflare.decisions.transformation", "CloudflareDecisionsConfig"),
"CloudflareExtractionConfig": (".llms.cloudflare.extraction.transformation", "CloudflareExtractionConfig"),
"StrandsDeciderDecisionsConfig": (
".llms.strands_decider.decisions.transformation",
"StrandsDeciderDecisionsConfig",

10
litellm/caching/in_memory_cache.py
View file

@ -95,7 +95,10 @@ class InMemoryCache(BaseCache):
"""
Check if a specific key is expired
"""
return key in self.ttl_dict and self._clock() > self.ttl_dict[key]
try:
return self._clock() > self.ttl_dict[key]
except KeyError:
return False
def _cache_keys(self) -> tuple[str, ...]:
return tuple(self.cache_dict) # pyright: ignore[reportUnknownVariableType] # cache_dict is untyped legacy state
@ -227,7 +230,10 @@ class InMemoryCache(BaseCache):
if key in self.cache_dict:
if self.evict_element_if_expired(key):
return None
original_cached_response: Final = self.cache_dict[key]
try:
original_cached_response: Final = self.cache_dict[key]
except KeyError:
return None
try:
cached_response = json.loads(original_cached_response)
except Exception:

35
litellm/chat_completions/dispatch.py
View file

@ -1,15 +1,19 @@
import inspect
from collections.abc import Awaitable, Callable, Coroutine, Mapping
from types import MappingProxyType
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
from litellm import main
from litellm.rust_bridge.catalog import Route, RouteContext
from litellm.rust_bridge import detectors
from litellm.rust_bridge.catalog import Feature, Route, RouteContext
from litellm.rust_bridge.chat_completions import detectors as route_detectors
from litellm.rust_bridge.chat_completions.entrypoints import (
NATIVE_ACOMPLETION,
NATIVE_COMPLETION,
)
from litellm.rust_bridge.dispatch import PublicDispatch
from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
bind,
native_call,
@ -44,6 +48,7 @@ def _python_acompletion() -> PythonAcompletion:
_PYTHON_COMPLETION: Final = _python_completion()
_COMPLETION: Final = signature(_PYTHON_COMPLETION)
_FLAG: Final = AsyncFlag("acompletion", sync=_COMPLETION)
_PYTHON_ACOMPLETION: Final = _python_acompletion()
_ACOMPLETION: Final = signature(_PYTHON_ACOMPLETION)
@ -61,11 +66,32 @@ def _public_request(
return native_call(legacy, args, kwargs)
_FEATURES: Final[detectors.Detectors] = MappingProxyType(
{
**detectors.COMMON,
Feature.MOCK_RESPONSE: route_detectors.chat_mock_response,
Feature.AGENTIC_LOOP_HOOK: route_detectors.chat_agentic_loop_hook,
Feature.RESPONSE_CACHE: detectors.response_cache,
Feature.STREAMING: detectors.streaming,
Feature.DROP_PARAMS: detectors.drop_params,
Feature.MODIFY_PARAMS: detectors.modify_params,
Feature.EXTRA_BODY: detectors.extra_body,
Feature.RETRIES_FALLBACKS: route_detectors.chat_retries_fallbacks,
Feature.PROMPT_MANAGEMENT: detectors.prompt_management,
Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection,
Feature.MODEL_PARAM_GATING: route_detectors.chat_model_param_gating,
Feature.MODEL_LIST: route_detectors.model_list,
Feature.MODEL_ALIAS_MAP: route_detectors.model_alias_map,
}
)
def _context(request: NativeCall) -> RouteContext:
return RouteContext(
Route.CHAT_COMPLETIONS,
provider=optional_str(request.resolved.get("custom_llm_provider")),
model=str(request.resolved["model"]),
features=detectors.detected(_FEATURES, request),
)
@ -73,7 +99,6 @@ _DISPATCH: Final = PublicDispatch(
route=Route.CHAT_COMPLETIONS,
request=lambda args, kwargs: _public_request(_COMPLETION, args, kwargs),
context=_context,
bypass=lambda request: request.kwargs.get("acompletion") is True,
)
_ADISPATCH: Final = PublicDispatch(
@ -88,9 +113,11 @@ def completion(
**kwargs: object, # kwargs-ok: preserve the public chat completions call shape
) -> ChatResult | Coroutine[object, object, ChatResult]:
python: Final = _PYTHON_COMPLETION
if _FLAG.requested(kwargs):
return acompletion(**_FLAG.keywords(args, kwargs))
return _DISPATCH.run(
args,
kwargs,
_FLAG.dropped(kwargs),
python=python,
binding=NATIVE_COMPLETION,
native=native_call_hook,
@ -101,7 +128,7 @@ async def acompletion(*args: object, **kwargs: object) -> ChatResult: # kwargs-
python: Final = _PYTHON_ACOMPLETION
return await _ADISPATCH.arun(
args,
kwargs,
_FLAG.dropped(kwargs),
python=python,
binding=NATIVE_ACOMPLETION,
native=native_call_hook,

10
litellm/constants.py
View file

@ -25,6 +25,8 @@ CLAUDE_CODE_CLIENT: Final = "claude-code"
MICROSOFT_365_COPILOT_DEFAULT_TIME_ZONE: Final = "UTC"
GITHUB_COPILOT_PER_USER_AUTH_TYPE: Final = "per_user_oauth"
GITHUB_COPILOT_AUTH_TYPE_KEY: Final = "github_copilot_auth_type"
GITHUB_COPILOT_HOST_KEY: Final = "github_copilot_host"
GITHUB_COPILOT_DEFAULT_HOST: Final = "github.com"
GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS: Final = 60
GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60
GITHUB_COPILOT_DEVICE_FLOW_CACHE_PREFIX: Final = "github_copilot_device_flow"
@ -154,6 +156,9 @@ DEFAULT_REPLICATE_POLLING_RETRIES: Final = get_env_int("DEFAULT_REPLICATE_POLLIN
DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_int("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)
DEFAULT_IMAGE_TOKEN_COUNT: Final = get_env_int("DEFAULT_IMAGE_TOKEN_COUNT", 250)
HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0
HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS: Final = get_env_int_in_range(
"HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS", 300, 1, 86400
)
# Maximum wall-clock seconds a streaming response is allowed to run.
# Streams exceeding this duration are terminated with a Timeout error.
@ -256,6 +261,9 @@ MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS: Final = get_env_int("MCP_PER_USER_TOKE
# MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers.
MCP_CLIENT_TIMEOUT: Final = get_env_float("LITELLM_MCP_CLIENT_TIMEOUT", 60.0)
MCP_MAX_SSE_EVENT_SIZE_BYTES: Final = get_env_int_in_range(
"LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES", 64 * 1024 * 1024, 1, sys.maxsize
)
MCP_TOOL_LISTING_TIMEOUT: Final = get_env_float("LITELLM_MCP_TOOL_LISTING_TIMEOUT", 30.0)
MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES: Final = get_env_int("LITELLM_MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES", 3)
MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS: Final = get_env_float(
@ -1649,6 +1657,7 @@ BAGGAGE_MAPPED_TAGS_METADATA_KEY: Final = "_baggage_mapped_tags"
VECTOR_STORE_CALLER_METADATA_KWARG: Final = "litellm_vector_store_caller_metadata"
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
BILLED_RETRY_ATTEMPT_METADATA_KEY: Final = "billed_retry_attempt"
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response"
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"
@ -1812,6 +1821,7 @@ LITELLM_EXECUTED_BATCH_CONCURRENCY: Final = max(1, get_env_int("LITELLM_EXECUTED
########################### CLI SSO AUTHENTICATION CONSTANTS ###########################
LITELLM_CLI_SOURCE_IDENTIFIER: Final = "litellm-cli"
LITELLM_CLI_SESSION_TOKEN_PREFIX: Final = "litellm-session-token"
CP_RETURN_TO_COOKIE_PREFIX: Final = "return_to.v1."
CLI_SSO_SESSION_CACHE_KEY_PREFIX: Final = "cli_sso_session"
CLI_SSO_SESSION_TTL_SECONDS: Final = 600
CLI_SESSION_KEY_PREFIX: Final = "cli-session"

21
litellm/cost_calculator.py
View file

@ -112,6 +112,7 @@ from litellm.llms.xai.cost_calculator import cost_per_token as xai_cost_per_toke
from litellm.responses.utils import ResponseAPILoggingUtils
from litellm.types.agents import LiteLLMSendMessageResponse
from litellm.types.decisions import DecisionsResponse, DecisionsUsage, OpenAIDecisionResponse, OpenAIDecisionUsage
from litellm.types.extraction import ExtractionResponse, ExtractionUsage
from litellm.types.llms.base import CachedTokensDetails, LiteLLMBaseModel
from litellm.types.llms.openai import (
HttpxBinaryResponseContent,
@ -1171,6 +1172,8 @@ def get_usage_object(
if isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)):
return None if completion_response.usage is None else _decisions_usage(completion_response.usage)
if isinstance(completion_response, ExtractionResponse):
return _extraction_usage(completion_response.usage)
if isinstance(completion_response, HttpxBinaryResponseContent):
return completion_response.usage
if usage_obj is None:
@ -1227,6 +1230,14 @@ def _decisions_usage(usage: DecisionsUsage | OpenAIDecisionUsage) -> Usage:
)
def _extraction_usage(usage: ExtractionUsage) -> Usage:
return Usage(
prompt_tokens=usage.input_tokens,
completion_tokens=usage.output_tokens,
total_tokens=usage.input_tokens + usage.output_tokens,
)
def _is_known_usage_objects(usage_obj):
"""Returns True if the usage obj is a known Usage type"""
return (
@ -1234,6 +1245,7 @@ def _is_known_usage_objects(usage_obj):
or isinstance(usage_obj, ResponseAPIUsage)
or isinstance(usage_obj, DecisionsUsage)
or isinstance(usage_obj, OpenAIDecisionUsage)
or isinstance(usage_obj, ExtractionUsage)
or TranscriptionUsageObjectTransformation.is_transcription_usage_object(usage_obj)
)
@ -1263,6 +1275,8 @@ def infer_call_type(call_type: CallTypesLiteral | None, completion_response: obj
return "send_message"
elif isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)):
return "adecisions"
elif isinstance(completion_response, ExtractionResponse):
return "aextraction"
return call_type
@ -1674,6 +1688,8 @@ def completion_cost(
)
if isinstance(usage_obj, (DecisionsUsage, OpenAIDecisionUsage)):
_usage = _decisions_usage(usage_obj).model_dump()
elif isinstance(usage_obj, ExtractionUsage):
_usage = _extraction_usage(usage_obj).model_dump()
elif usage_obj is None:
_usage = {}
elif isinstance(usage_obj, BaseModel):
@ -2199,7 +2215,8 @@ def response_cost_calculator(
| Response
| SearchResponse
| DecisionsResponse
| OpenAIDecisionResponse,
| OpenAIDecisionResponse
| ExtractionResponse,
model: str,
custom_llm_provider: str | None,
call_type: Literal[
@ -2225,6 +2242,8 @@ def response_cost_calculator(
"adecisions",
"systemone",
"asystemone",
"extraction",
"aextraction",
],
optional_params: dict,
cache_hit: bool | None = None,

18
litellm/embeddings/dispatch.py
View file

@ -9,7 +9,15 @@ from litellm.rust_bridge.embeddings.entrypoints import (
NATIVE_AEMBEDDING,
NATIVE_EMBEDDING,
)
from litellm.rust_bridge.public_call import NativeCall, bind, native_call, native_call_hook, optional_str, signature
from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
bind,
native_call,
native_call_hook,
optional_str,
signature,
)
from litellm.types.utils import EmbeddingResponse
__all__ = ("aembedding", "embedding")
@ -24,6 +32,7 @@ _PYTHON_AEMBEDDING: Final = cast( # cast-ok: [LIT006] preserve the legacy publi
PythonAembedding, main.aembedding
)
_EMBEDDING_SIGNATURE: Final = signature(_PYTHON_EMBEDDING)
_FLAG: Final = AsyncFlag("aembedding", sync=_EMBEDDING_SIGNATURE)
def _public_request(
@ -50,7 +59,6 @@ _DISPATCH: Final = PublicDispatch(
route=Route.EMBEDDINGS,
request=lambda args, kwargs: _public_request(_EMBEDDING_SIGNATURE, args, kwargs),
context=_context,
bypass=lambda request: request.kwargs.get("aembedding") is True,
)
_ADISPATCH: Final = PublicDispatch(
@ -64,9 +72,11 @@ def embedding(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public embedding call shape
) -> EmbeddingResponse | Coroutine[object, object, EmbeddingResponse]:
if _FLAG.requested(kwargs):
return aembedding(**_FLAG.keywords(args, kwargs))
return _DISPATCH.run(
args,
kwargs,
_FLAG.dropped(kwargs),
python=_PYTHON_EMBEDDING,
binding=NATIVE_EMBEDDING,
native=native_call_hook,
@ -76,7 +86,7 @@ def embedding(
async def aembedding(*args: object, **kwargs: object) -> EmbeddingResponse: # kwargs-ok: preserve the public call shape
return await _ADISPATCH.arun(
args,
kwargs,
_FLAG.dropped(kwargs),
python=_PYTHON_AEMBEDDING,
binding=NATIVE_AEMBEDDING,
native=native_call_hook,

13
litellm/experimental_mcp_client/client.py
View file

@ -95,6 +95,7 @@ from pydantic import AnyUrl, TypeAdapter
from litellm._logging import verbose_logger
from litellm.constants import (
MCP_CLIENT_TIMEOUT,
MCP_MAX_SSE_EVENT_SIZE_BYTES,
MCP_NPM_CACHE_DIR,
MCP_TOOL_LISTING_MAX_PAGES,
MCP_TOOL_LISTING_TIMEOUT,
@ -292,6 +293,7 @@ class _MCPResponseStream(httpx2.AsyncByteStream):
class _MCPHTTPClient(httpx2.AsyncClient):
cleanup_scope: anyio.CancelScope | None = None
cleanup_errors: tuple[Exception, ...] = ()
post_failure: asyncio.Future[Exception] | None = None
def _record_cleanup_error(self, error: Exception) -> None:
if self.cleanup_scope is not None and self.cleanup_scope.shield:
@ -319,7 +321,12 @@ class _MCPHTTPClient(httpx2.AsyncClient):
response: Final = await super().send(request, stream=stream, auth=auth, follow_redirects=follow_redirects)
if request.method == "POST" and response.is_error and response.status_code != 404:
await response.aclose()
response.raise_for_status()
try:
response.raise_for_status()
except httpx2.HTTPStatusError as error:
if self.post_failure is not None and not self.post_failure.done():
self.post_failure.set_result(error)
raise
if stream:
response.stream = _MCPResponseStream(
cast(httpx2.AsyncByteStream, response.stream), self._record_cleanup_error
@ -641,6 +648,7 @@ class MCPClient:
transport_ctx: Final = streamable_http_client(
url=self.server_url,
http_client=http_client,
max_sse_event_size=MCP_MAX_SSE_EVENT_SIZE_BYTES,
)
return transport_ctx, http_client
@ -757,6 +765,8 @@ class MCPClient:
read_stream: Final = transport[0]
write_stream: Final = transport[1]
stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future()
if isinstance(http_client, _MCPHTTPClient):
http_client.post_failure = stream_error # rebind-ok: operation-scoped stream error
_notify_stream_error(on_stream_error, stream_error)
session_kwargs: Final = (
@ -823,6 +833,7 @@ class MCPClient:
if isinstance(http_client, _MCPHTTPClient):
http_client.cleanup_errors = ()
http_client.cleanup_scope = None
http_client.post_failure = None # rebind-ok: clear operation-scoped stream error
await anyio.lowlevel.checkpoint_if_cancelled()
if cleanup_scope.cancel_called:
raise (

3
litellm/extraction/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from litellm.extraction.main import aextraction, extraction
__all__ = ["aextraction", "extraction"]

351
litellm/extraction/main.py Normal file
View file

@ -0,0 +1,351 @@
import asyncio
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from functools import partial
from types import MappingProxyType
from typing import Final, TypeAlias
import httpx
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm.decisions.call import log_caller_request
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.tool_fields import DropPolicy, drop_policy
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.types.extraction import (
EXTRACTION_STATE_REWRITTEN_METADATA_KEY,
ExtractionAnswer,
ExtractionBase64Image,
ExtractionField,
ExtractionOptions,
ExtractionRequestBody,
ExtractionResponse,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders, is_litellm_owned_kwarg
from litellm.utils import (
ProviderConfigManager,
client, # pyright: ignore[reportUnknownVariableType] # client is an untyped decorator
get_model_cost_key,
)
ExtractionFields: TypeAlias = Mapping[str, ExtractionField | Mapping[str, object]]
ExtractionImages: TypeAlias = Sequence[str | ExtractionBase64Image | Mapping[str, object]]
ExtractionOptionsParam: TypeAlias = ExtractionOptions | Mapping[str, object]
_REQUEST_ADAPTER: Final[TypeAdapter[ExtractionRequestBody]] = TypeAdapter(ExtractionRequestBody)
_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
_HANDLER: Final = BaseLLMHTTPHandler()
_LOCAL_KWARGS: Final = frozenset(
(
*GenericLiteLLMParams.model_fields,
"extra_headers",
"base_url",
"default_headers",
"deployment_id",
"additional_drop_params",
)
)
@dataclass(frozen=True, slots=True, repr=False)
class _ExtractionCall:
model: str
requested_model: str
custom_llm_provider: str
provider_config: BaseExtractionConfig
body: Mapping[str, object]
api_base: str
api_key: str | None
logging_obj: LiteLLMLoggingObj | None
headers: Mapping[str, str]
timeout: float | httpx.Timeout | None
litellm_params: Mapping[str, object]
client: HTTPHandler | AsyncHTTPHandler | None
def supported_extraction_providers() -> tuple[str, ...]:
return tuple(
provider.value
for provider in LlmProviders
if ProviderConfigManager.get_provider_extraction_config(model="", provider=provider) is not None
)
def _provider_config(model: str, custom_llm_provider: str) -> BaseExtractionConfig:
provider: Final = next((member for member in LlmProviders if member.value == custom_llm_provider), None)
provider_config: Final = (
None if provider is None else ProviderConfigManager.get_provider_extraction_config(model, provider)
)
if provider_config is None:
raise litellm.BadRequestError(
message=(
f"Unknown Extraction provider '{custom_llm_provider}'. "
f"Supported providers: {', '.join(supported_extraction_providers())}"
),
model=model,
llm_provider=custom_llm_provider,
)
return provider_config
def _request_fields(
*,
state: str | None,
fields: ExtractionFields | None,
images: ExtractionImages | None,
reference_date: str | None,
options: ExtractionOptionsParam | None,
) -> Mapping[str, object]:
sent: Final[Mapping[str, object | None]] = MappingProxyType(
{"state": state, "fields": fields, "images": images, "reference_date": reference_date, "options": options}
)
return MappingProxyType({name: value for name, value in sent.items() if value is not None})
def _is_unsupported(name: str, value: object, policy: DropPolicy) -> bool:
return (
value is not None
and value is not False
and name not in _LOCAL_KWARGS
and not is_litellm_owned_kwarg(name)
and not policy.drops(name)
)
def _unsupported_params(kwargs: Mapping[str, object]) -> tuple[str, ...]:
policy: Final = drop_policy(kwargs)
return tuple(sorted(name for name, value in kwargs.items() if _is_unsupported(name, value, policy)))
def _prepare_call(
*,
model: str,
request_fields: Mapping[str, object],
api_key: str | None,
api_base: str | None,
timeout: float | httpx.Timeout | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, str] | None,
kwargs: Mapping[str, object],
) -> _ExtractionCall:
log_caller_request(kwargs, request_fields)
upstream_model, provider, dynamic_api_key, dynamic_api_base = litellm.get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=api_key,
)
provider_config: Final = _provider_config(model, provider)
unsupported: Final = _unsupported_params(kwargs)
if unsupported:
raise litellm.UnsupportedParamsError(
message=(
f"{provider} does not support parameters: {list(unsupported)}, for model={model}. "
"To drop these, set `litellm.drop_params=True` or for proxy:\n\n"
"`litellm_settings:\n drop_params: true`\n"
),
model=model,
llm_provider=provider,
)
try:
request: Final = _REQUEST_ADAPTER.validate_python(request_fields)
except ValidationError as error:
raise litellm.BadRequestError(
message=f"Invalid Extraction request: {error}", model=model, llm_provider=provider
) from error
canonical_model: Final = provider_config.canonical_model(upstream_model)
resolved_api_base: Final = provider_config.resolve_api_base(dynamic_api_base or api_base)
if resolved_api_base is None:
raise litellm.BadRequestError(
message=provider_config.missing_api_base_message(provider), model=model, llm_provider=provider
)
resolved_api_key: Final = provider_config.resolve_api_key(dynamic_api_key or api_key)
if resolved_api_key is None:
raise litellm.BadRequestError(
message=f"Missing API key for Extraction provider '{provider}'", model=model, llm_provider=provider
)
logging_obj: Final = kwargs.get("litellm_logging_obj")
http_client: Final = kwargs.get("client")
if isinstance(logging_obj, LiteLLMLoggingObj):
logging_obj.update_from_kwargs( # pyright: ignore[reportUnknownMemberType] # Logging takes untyped dicts
kwargs=dict(kwargs),
model=canonical_model,
litellm_params={
"litellm_call_id": kwargs.get("litellm_call_id"),
"api_base": provider_config.get_complete_url(resolved_api_base, canonical_model),
},
custom_llm_provider=provider,
)
return _ExtractionCall(
model=canonical_model,
requested_model=model,
custom_llm_provider=provider,
provider_config=provider_config,
body=provider_config.transform_extraction_request(canonical_model, request),
api_base=resolved_api_base,
api_key=resolved_api_key,
logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None,
headers=MappingProxyType(dict(extra_headers or {})),
timeout=timeout,
litellm_params=kwargs,
client=http_client if isinstance(http_client, HTTPHandler | AsyncHTTPHandler) else None,
)
def _state_rewritten_by_guardrail(litellm_params: Mapping[str, object]) -> bool:
metadata: Final = litellm_params.get(get_metadata_variable_name_from_kwargs(litellm_params))
return (
isinstance(metadata, Mapping)
and _METADATA_ADAPTER.validate_python(metadata).get(EXTRACTION_STATE_REWRITTEN_METADATA_KEY) is True
)
def _without_spans(answers: Mapping[str, ExtractionAnswer]) -> Mapping[str, ExtractionAnswer]:
return {name: answer.model_copy(update={"span": None, "spans": None}) for name, answer in answers.items()}
def _formatted_response(response: ExtractionResponse, call: _ExtractionCall) -> ExtractionResponse:
spans_index_hidden_text: Final = response.state_text is None and _state_rewritten_by_guardrail(call.litellm_params)
formatted: Final = response.model_copy(
update={
"model": call.requested_model,
**({"answers": _without_spans(response.answers)} if spans_index_hidden_text else {}),
}
)
formatted.set_hidden_params(
{
"model": f"{call.custom_llm_provider}/{call.model}",
"custom_llm_provider": call.custom_llm_provider,
"provider_response_model": (
get_model_cost_key(f"{call.custom_llm_provider}/{response.model}")
or f"{call.custom_llm_provider}/{call.model}"
),
}
)
return formatted
def _upstream_exception(error: Exception, call: _ExtractionCall) -> Exception:
if isinstance(error, BaseLLMException) and error.status_code_is_synthesized:
return litellm.APIConnectionError(
message=f"{call.custom_llm_provider.capitalize()}Exception - {error.message}",
llm_provider=call.custom_llm_provider,
model=f"{call.custom_llm_provider}/{call.model}",
)
return litellm.exception_type( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped
model=f"{call.custom_llm_provider}/{call.model}",
custom_llm_provider=call.custom_llm_provider,
original_exception=error,
)
async def _asend(call: _ExtractionCall) -> ExtractionResponse:
try:
response: Final = await _HANDLER.aextraction(
model=call.model,
custom_llm_provider=call.custom_llm_provider,
logging_obj=call.logging_obj,
provider_config=call.provider_config,
body=call.body,
api_base=call.api_base,
api_key=call.api_key,
headers=call.headers,
timeout=call.timeout,
litellm_params=call.litellm_params,
client=call.client if isinstance(call.client, AsyncHTTPHandler) else None,
)
except Exception as error:
raise _upstream_exception(error, call) from error
return _formatted_response(response, call)
def _send(call: _ExtractionCall) -> ExtractionResponse:
try:
response: Final = _HANDLER.extraction(
model=call.model,
custom_llm_provider=call.custom_llm_provider,
logging_obj=call.logging_obj,
provider_config=call.provider_config,
body=call.body,
api_base=call.api_base,
api_key=call.api_key,
headers=call.headers,
timeout=call.timeout,
litellm_params=call.litellm_params,
client=call.client if isinstance(call.client, HTTPHandler) else None,
)
except Exception as error:
raise _upstream_exception(error, call) from error
return _formatted_response(response, call)
@client
async def aextraction(
model: str,
fields: ExtractionFields,
state: str | None = None,
images: ExtractionImages | None = None,
reference_date: str | None = None,
options: ExtractionOptionsParam | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, str] | None = None,
**kwargs: object, # kwargs-ok: litellm_params forwarded by @client
) -> ExtractionResponse:
call: Final = await asyncio.to_thread(
partial(
_prepare_call,
model=model,
request_fields=_request_fields(
state=state, fields=fields, images=images, reference_date=reference_date, options=options
),
api_key=api_key,
api_base=api_base,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
kwargs=kwargs,
)
)
return await _asend(call)
@client
def extraction(
model: str,
fields: ExtractionFields,
state: str | None = None,
images: ExtractionImages | None = None,
reference_date: str | None = None,
options: ExtractionOptionsParam | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, str] | None = None,
**kwargs: object, # kwargs-ok: litellm_params forwarded by @client
) -> ExtractionResponse:
call: Final = _prepare_call(
model=model,
request_fields=_request_fields(
state=state, fields=fields, images=images, reference_date=reference_date, options=options
),
api_key=api_key,
api_base=api_base,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
kwargs=kwargs,
)
return _send(call)
__all__ = ["aextraction", "extraction"]

18
litellm/harness/endpoint.py
View file

@ -37,6 +37,7 @@ from litellm.harness.context import GatewayTarget
from litellm.harness.errors import HarnessError, HarnessInstallFailed
from litellm.harness.types import Harness, Usage
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD
from litellm.types.llms.custom_http import httpxSpecialProvider
if TYPE_CHECKING:
@ -690,18 +691,21 @@ class ModelEndpoint:
async def _invoke_sdk(
self,
route: str,
kwargs: dict[str, Any], # mutable-ok: injects stream_options into the SDK kwargs
kwargs: dict[str, Any], # mutable-ok: SDK call kwargs built by _sdk_kwargs
) -> object:
if route == ROUTE_MESSAGES:
return await litellm.anthropic.messages.acreate(**kwargs)
if route == ROUTE_CHAT:
if kwargs.get("stream"):
stream_options = kwargs.get("stream_options") or {} # mutable-ok: empty default for a JSON field
kwargs["stream_options"] = { # mutable-ok: JSON field sent to litellm.acompletion
"include_usage": True,
**stream_options,
if not kwargs.get("stream"):
return await litellm.acompletion(**kwargs)
stream_options = kwargs.get("stream_options") or {} # mutable-ok: empty default for a JSON field
return await litellm.acompletion(
**{
**kwargs,
"stream_options": {"include_usage": True, **stream_options},
STREAM_USAGE_INJECTED_FIELD: "include_usage" not in stream_options,
}
return await litellm.acompletion(**kwargs)
)
return await litellm.aresponses(**kwargs)
async def _call_sdk(self, route: str, body: Mapping[str, object]) -> Response:

1
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -77,6 +77,7 @@ OPTIONAL_KWARGS_KEYS: Final = (
"otpm",
"use_xai_oauth",
"github_copilot_auth_type",
"github_copilot_host",
"github_copilot_user_session",
"fireworks_forward_user_id",
"responses_api_version",

20
litellm/litellm_core_utils/health_check_helpers.py
View file

@ -4,6 +4,7 @@ Helper functions for health check calls.
import base64
from collections.abc import Awaitable, Callable, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal
from litellm.llms.base_llm.ocr.transformation import DocumentType
@ -25,6 +26,9 @@ IMAGE_EDIT_HEALTH_CHECK_PROMPT: Final = (
)
ANTHROPIC_MESSAGES_HEALTH_CHECK_MAX_TOKENS: Final = 16
EXTRACTION_HEALTH_CHECK_FIELDS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType(
{"reachable": MappingProxyType({"type": "text", "instructions": "Any word from the text"})}
)
def native_health_check_mode(model: str, custom_llm_provider: str | None) -> Literal["anthropic_messages"] | None:
@ -215,6 +219,7 @@ class HealthCheckHelpers:
"anthropic_messages",
"ocr",
"decisions",
"extraction",
],
Callable,
]:
@ -237,7 +242,11 @@ class HealthCheckHelpers:
from litellm.litellm_core_utils.audio_utils.utils import (
get_audio_file_for_health_check,
)
from litellm.litellm_core_utils.health_check_utils import DECISIONS_CALL_PARAMS, _filter_model_params
from litellm.litellm_core_utils.health_check_utils import (
DECISIONS_CALL_PARAMS,
EXTRACTION_CALL_PARAMS,
_filter_model_params,
)
from litellm.realtime_api.main import realtime_health_check
return {
@ -322,4 +331,13 @@ class HealthCheckHelpers:
}
)
),
"extraction": lambda: litellm.aextraction(
**EXTRACTION_CALL_PARAMS.validate_python(
{
"state": prompt or "health check",
"fields": EXTRACTION_HEALTH_CHECK_FIELDS,
**_filter_model_params(model_params=model_params),
}
)
),
}

2
litellm/litellm_core_utils/health_check_utils.py
View file

@ -8,8 +8,10 @@ from typing import Final
from pydantic import TypeAdapter
from litellm.types.decisions import DecisionsCallParams
from litellm.types.extraction import ExtractionCallParams
DECISIONS_CALL_PARAMS: Final[TypeAdapter[DecisionsCallParams]] = TypeAdapter(DecisionsCallParams)
EXTRACTION_CALL_PARAMS: Final[TypeAdapter[ExtractionCallParams]] = TypeAdapter(ExtractionCallParams)
OPTIONAL_STR: Final[TypeAdapter[str | None]] = TypeAdapter(str | None)

184
litellm/litellm_core_utils/litellm_logging.py
View file

@ -12,7 +12,7 @@ import subprocess
import sys
import time
import traceback
from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence
from collections.abc import Awaitable, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime as dt_object
from functools import lru_cache
@ -132,11 +132,13 @@ from litellm.responses.utils import ResponseAPILoggingUtils
from litellm.types.agents import LiteLLMSendMessageResponse
from litellm.types.containers.main import ContainerObject
from litellm.types.decisions import DecisionsResponse, OpenAIDecisionResponse
from litellm.types.extraction import ExtractionResponse
from litellm.types.integrations.s3_v2 import S3PartitionGranularity
from litellm.types.interactions import (
InteractionsAPIResponse,
InteractionsAPIStreamingResponse,
)
from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD
from litellm.types.llms.openai import (
AllMessageValues,
Batch,
@ -573,6 +575,27 @@ class UncachedPrice:
cost_breakdown: CostBreakdown | None
@dataclass(frozen=True, slots=True)
class BilledRoundsSpend:
usage: Usage
cost: float
cost_breakdown: CostBreakdown | None
@dataclass(frozen=True, slots=True)
class BilledRoundsDispatch:
"""A parked success log for a stream bridged through chat completions that billed several rounds, and the spend
its failure log bills when a post-call hook rejects the response, as it does for a rejected one-round bridged
stream."""
dispatch: Coroutine[object, object, None]
spend: BilledRoundsSpend
def _fixed_margin(cost_breakdown: CostBreakdown | None) -> float:
return 0.0 if cost_breakdown is None else cost_breakdown.get("margin_fixed_amount", 0.0)
def _timestamp_seconds(moment: object) -> float | None:
if isinstance(moment, datetime.datetime):
return moment.timestamp()
@ -594,6 +617,7 @@ class Logging(LiteLLMLoggingBaseClass):
prometheusLogger, \
slack_app
custom_pricing: bool = False
replaced_global_callback_keys: frozenset[str] = frozenset()
stream_options = None
litellm_request_debug: bool = False
streamed_anthropic_message_id: str | None = None
@ -648,11 +672,11 @@ class Logging(LiteLLMLoggingBaseClass):
litellm_call_id: str,
function_id: str,
litellm_trace_id: str | None = None,
dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = None,
dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None,
dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = None,
dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
applied_guardrails: list[str] | None = None,
kwargs: dict | None = None,
log_raw_request_response: bool = False,
@ -697,6 +721,7 @@ class Logging(LiteLLMLoggingBaseClass):
self._pre_call_session_id: str = session_id_var.get()
_sid: Final = kwargs.get("litellm_session_id") if kwargs else None
self.litellm_session_id: str = str(_sid) if _sid else ""
self.stream_usage_injected: Final = kwargs is not None and kwargs.get(STREAM_USAGE_INJECTED_FIELD) is True
# supports_correlation_logging is False for calls originating from the
# sync client entry point (wrapper() in utils.py): a plain OS thread
# has no per-call context isolation the way an asyncio Task does, and
@ -725,15 +750,17 @@ class Logging(LiteLLMLoggingBaseClass):
self.raw_request_only = raw_request_only
# Initialize dynamic callbacks
self.dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = dynamic_input_callbacks
self.dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = dynamic_input_callbacks
self.dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
dynamic_success_callbacks
)
self.dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
dynamic_async_success_callbacks
)
self.dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = dynamic_failure_callbacks
self.dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = (
self.dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
dynamic_failure_callbacks
)
self.dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
dynamic_async_failure_callbacks
)
@ -796,6 +823,7 @@ class Logging(LiteLLMLoggingBaseClass):
self.defer_async_logging: bool = False
self.enqueue_deferred_logging: Callable[[], None] | None = None
self._async_success_scheduled: bool = False
self._async_success_enqueued: bool = False
self._sync_success_scheduled: bool = False
self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
self._stream_failure_claimed: bool = False
@ -850,10 +878,32 @@ class Logging(LiteLLMLoggingBaseClass):
self.dynamic_async_failure_callbacks, callback
)
def replace_global_callback(self, callback: CustomLogger) -> None:
self.add_dynamic_callback(callback)
self.replaced_global_callback_keys = self.replaced_global_callback_keys | {
litellm.logging_callback_manager.get_custom_logger_key(callback)
}
def add_dynamic_function_callbacks(
self,
*,
async_success: Callable[..., object],
failure: Callable[..., object],
async_failure: Callable[..., object],
) -> None:
self.dynamic_async_success_callbacks = self._with_dynamic_callback(
self.dynamic_async_success_callbacks, async_success
)
self.dynamic_failure_callbacks = self._with_dynamic_callback(self.dynamic_failure_callbacks, failure)
self.dynamic_async_failure_callbacks = self._with_dynamic_callback(
self.dynamic_async_failure_callbacks, async_failure
)
@staticmethod
def _with_dynamic_callback(
callbacks: Sequence[str | Callable | CustomLogger] | None, callback: CustomLogger
) -> list[str | Callable | CustomLogger]:
callbacks: Sequence[str | Callable[..., object] | CustomLogger] | None,
callback: Callable[..., object] | CustomLogger,
) -> list[str | Callable[..., object] | CustomLogger]:
existing: Final = tuple(callbacks or ())
return [*existing, *(() if callback in existing else (callback,))]
@ -890,9 +940,9 @@ class Logging(LiteLLMLoggingBaseClass):
def _process_dynamic_callback_list(
self,
callback_list: list[str | Callable | CustomLogger] | None,
callback_list: list[str | Callable[..., object] | CustomLogger] | None,
dynamic_callbacks_type: Literal["input", "success", "failure", "async_success", "async_failure"],
) -> list[str | Callable | CustomLogger] | None:
) -> list[str | Callable[..., object] | CustomLogger] | None:
"""
Helper function to initialize CustomLogger compatible callbacks in self.dynamic_* callbacks
@ -904,7 +954,7 @@ class Logging(LiteLLMLoggingBaseClass):
if callback_list is None:
return None
processed_list: Final[list[str | Callable | CustomLogger]] = []
processed_list: Final[list[str | Callable[..., object] | CustomLogger]] = []
for callback in callback_list:
if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks:
for callback_instance in self._resolve_dynamic_callback_string(callback):
@ -1565,7 +1615,10 @@ class Logging(LiteLLMLoggingBaseClass):
"api_call_start_time"
]
# Input Integration Logging -> If you want to log the fact that an attempt to call the model was made
callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or [])
callbacks: Final = (
*self._without_replaced_global_callbacks(litellm.input_callback),
*(self.dynamic_input_callbacks or ()),
)
for callback in callbacks:
try:
if callback == "supabase" and supabaseClient is not None:
@ -1738,7 +1791,10 @@ class Logging(LiteLLMLoggingBaseClass):
)
# Input Integration Logging -> If you want to log the fact that an attempt to call the model was made
callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or [])
callbacks: Final = (
*self._without_replaced_global_callbacks(litellm.input_callback),
*(self.dynamic_input_callbacks or ()),
)
for callback in callbacks:
try:
if callback == "sentry" and add_breadcrumb:
@ -2209,20 +2265,23 @@ class Logging(LiteLLMLoggingBaseClass):
usage=usage,
)
def custom_pricing_for(self, result: object) -> bool:
litellm_params: Final = getattr(self, "litellm_params", None)
def spilled_over_ptu_for(self, result: object) -> bool:
result_hidden_params: Final = getattr(result, HIDDEN_PARAMS_ATTR, None) or MappingProxyType({})
additional_headers: Final = (
result_hidden_params.get("additional_headers")
if isinstance(result_hidden_params, dict)
else getattr(result_hidden_params, "additional_headers", None)
)
spilled_over: Final = is_spilled_over_ptu_request(
model_info=_deployment_model_info(litellm_params),
return is_spilled_over_ptu_request(
model_info=_deployment_model_info(getattr(self, "litellm_params", None)),
response_headers=self.model_call_details.get("response_headers"),
additional_headers=additional_headers,
)
return False if spilled_over else use_custom_pricing_for_model(litellm_params=litellm_params)
def custom_pricing_for(self, result: object) -> bool:
if self.spilled_over_ptu_for(result):
return False
return use_custom_pricing_for_model(litellm_params=getattr(self, "litellm_params", None))
def _prompt_for_cost_calculation(self) -> str:
"""
@ -2384,6 +2443,31 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["combined_usage_object"] = usage
self.model_call_details["response_cost"] = response_cost
def record_failure_after_earlier_rounds(
self, earlier: BilledRoundsSpend, failed_usage: Usage, failed_cost: float
) -> None:
"""Bill a failed round on the failure log after the rounds of its request before it were billed on the success
log.
The fixed margin is a per-request fee, so it is charged once: the failed round's margin is dropped when the
earlier rounds already paid one."""
from litellm.responses.mcp.round_billing import breakdown_less_repeated_fixed_margin
failed_breakdown: Final = None if self.cost_breakdown is earlier.cost_breakdown else self.cost_breakdown
repeated_fixed_margin: Final = (
0.0
if failed_breakdown is None
else min(_fixed_margin(earlier.cost_breakdown), _fixed_margin(failed_breakdown), earlier.cost, failed_cost)
)
self.model_call_details["combined_usage_object"] = failed_usage
self.model_call_details["response_cost"] = failed_cost - repeated_fixed_margin
self.cost_breakdown = breakdown_less_repeated_fixed_margin(failed_breakdown, repeated_fixed_margin)
def record_billed_rounds_for_failure(self, spend: BilledRoundsSpend) -> None:
"""Bill every round of a fully streamed request on the failure log when a post-call hook rejects it."""
self.record_partial_usage_for_failure(spend.usage, spend.cost)
self.cost_breakdown = spend.cost_breakdown
def attempt_scoped_copy(self) -> "Logging":
"""A copy with its own model_call_details, so logging one attempt of a request through it leaves
the state the request's other attempts share untouched."""
@ -2472,7 +2556,7 @@ class Logging(LiteLLMLoggingBaseClass):
await self.async_failure_handler(exception, traceback_exception)
finally:
if self._should_run_sync_failure_callbacks_for_async_calls():
executor.submit(self.failure_handler, exception, traceback_exception)
executor.submit(self.failure_handler, exception, traceback_exception, for_async_request=True)
def should_run_logging(
self,
@ -2487,6 +2571,16 @@ class Logging(LiteLLMLoggingBaseClass):
except Exception:
return True
@property
def async_success_log_enqueued(self) -> bool:
"""True once the success log task was actually created, not merely claimed, so a later failure raised
around it (a disconnect cancelling a still-pending guardrail task) must not bill the prompt a second
time. A claimed log parked in the deferred slot never runs once a 499 interrupts, so it stays False."""
return self._async_success_enqueued
def mark_async_success_log_enqueued(self) -> None:
self._async_success_enqueued = True
def claim_async_success_log(self) -> bool:
"""One async success log per request. The innermost @client wrapper always exits first,
since the outer one is awaiting it, so it claims the log here and the outer wrapper gets
@ -2915,6 +3009,7 @@ class Logging(LiteLLMLoggingBaseClass):
or isinstance(logging_result, SearchResponse) # Search API
or isinstance(logging_result, DecisionsResponse)
or isinstance(logging_result, OpenAIDecisionResponse)
or isinstance(logging_result, ExtractionResponse)
or (
isinstance(logging_result, InteractionsAPIResponse)
and logging_result.usage is not None
@ -4086,15 +4181,21 @@ class Logging(LiteLLMLoggingBaseClass):
traceback_exception: str,
start_time: datetime.datetime | None = None,
end_time: datetime.datetime | None = None,
*,
for_async_request: bool = False,
) -> None:
"""Restores trace_id/session_id contextvars once this attempt's own failure
logging (including any nested calls its callbacks trigger) is fully done."""
logging (including any nested calls its callbacks trigger) is fully done.
``for_async_request`` marks the sync run of a failure whose async run also logs it, so a
CustomLogger that gets ``async_log_failure_event`` from that async run skips its sync hook."""
try:
return self._failure_handler_body(
exception=exception,
traceback_exception=traceback_exception,
start_time=start_time,
end_time=end_time,
for_async_request=for_async_request,
)
finally:
self._restore_correlation_context()
@ -4105,6 +4206,7 @@ class Logging(LiteLLMLoggingBaseClass):
traceback_exception: str,
start_time: datetime.datetime | None = None,
end_time: datetime.datetime | None = None,
for_async_request: bool = False,
) -> None:
verbose_logger.debug("Logging Details LiteLLM-Failure Call: %s", litellm.failure_callback)
if not self.should_run_logging(event_type="sync_failure"): # prevent double logging
@ -4197,6 +4299,7 @@ class Logging(LiteLLMLoggingBaseClass):
isinstance(callback, CustomLogger)
and is_sync_request
and self.call_type != CallTypes.pass_through.value
and not (for_async_request and self._async_handler_logs_to(callback))
): # custom logger class
callback.log_failure_event(
start_time=start_time,
@ -4499,14 +4602,26 @@ class Logging(LiteLLMLoggingBaseClass):
_filtered_failure_callbacks = self._remove_internal_litellm_callbacks(_filtered_failure_callbacks)
return len(_filtered_failure_callbacks) > 0
def _without_replaced_global_callbacks(
self, global_callbacks: Iterable[str | Callable[..., object] | CustomLogger]
) -> tuple[str | Callable[..., object] | CustomLogger, ...]:
if not self.replaced_global_callback_keys:
return tuple(global_callbacks)
return tuple(
litellm.logging_callback_manager.without_custom_logger_keys(
global_callbacks, self.replaced_global_callback_keys
)
)
def get_combined_callback_list(
self,
dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None,
global_callbacks: list[str | Callable[..., object] | CustomLogger],
) -> list[str | Callable[..., object] | CustomLogger]:
kept_global_callbacks: Final = self._without_replaced_global_callbacks(global_callbacks)
if dynamic_success_callbacks is None:
return list(global_callbacks)
return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks))
return list(kept_global_callbacks)
return list(dict.fromkeys((*dynamic_success_callbacks, *kept_global_callbacks)))
def _remove_internal_litellm_callbacks(self, callbacks: list) -> list:
"""
@ -4559,6 +4674,12 @@ class Logging(LiteLLMLoggingBaseClass):
cb_name: Final = self._get_callback_name(cb)
return any(prefix in cb_name for prefix in INTERNAL_PREFIXES)
def _async_handler_logs_to(self, logger: CustomLogger) -> bool:
"""True when the async failure run of this call hands ``logger`` its own ``async_log_failure_event``."""
if type(logger).async_log_failure_event is CustomLogger.async_log_failure_event: # pyright: ignore[reportUnknownMemberType] # CustomLogger's hook signatures are untyped
return False
return logger in (self.dynamic_async_failure_callbacks or []) or logger in litellm._async_failure_callback # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # callback lists are untyped and have no public accessor
def _remove_internal_custom_logger_callbacks(self, callbacks: list) -> list:
"""
Removes internal custom logger callbacks from the list.
@ -7227,13 +7348,16 @@ def get_standard_logging_object_payload(
saved_cache_cost: Final = (uncached_price.cost if uncached_price is not None else None) or 0.0
## Get model cost information ##
base_model = get_base_model_from_metadata(model_call_details=kwargs)
# The router overrides completion_response.model to the model-group alias before
# this payload is built, so cost-map lookup via that alias always misses.
# Fall back to the actual deployment model set by the router in metadata.
if base_model is None:
base_model = metadata.get("deployment")
custom_pricing: Final = logging_obj.custom_pricing_for(init_response_obj)
configured_base_model: Final = get_base_model_from_metadata(model_call_details=kwargs)
spilled_over: Final = logging_obj.spilled_over_ptu_for(init_response_obj)
# The router overrides completion_response.model to the model-group alias before this payload is
# built, so cost-map lookup via that alias misses and the deployment model stands in for base_model.
# A PTU request Azure spilled onto pay-as-you-go capacity is billed at the served model's rates,
# which the deployment's zeroed PTU entry would misreport, so that fallback is skipped there.
base_model: Final = (
metadata.get("deployment") if configured_base_model is None and not spilled_over else configured_base_model
)
custom_pricing: Final = False if spilled_over else logging_obj.custom_pricing_for(init_response_obj)
raw_response_cost: Final = kwargs.get("response_cost")
llm_response_cost: Final[float] = raw_response_cost or 0.0
guardrail_cost: Final = guardrail_information_cost(metadata.get("standard_logging_guardrail_information"))

29
litellm/litellm_core_utils/llm_response_utils/response_metadata.py
View file

@ -4,6 +4,7 @@ from functools import reduce
from typing import Final
import httpx
from pydantic import TypeAdapter, ValidationError
from litellm.constants import LITELLM_DETAILED_TIMING
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, process_response_headers
@ -101,6 +102,32 @@ def response_timing_metrics(
return {"_response_ms": total_response_time_ms, "litellm_overhead_time_ms": overhead_ms}
_HIDDEN_PARAMS_DICT: Final = TypeAdapter(dict[str, object])
def _hidden_params_mapping(hidden_params: object) -> Mapping[str, object]:
try:
return _HIDDEN_PARAMS_DICT.validate_python(hidden_params)
except ValidationError:
return {}
def _existing_provider_response_model(result: object) -> str | None:
"""The pricing model a response already carries, like the decisions path sets, wins over the response model."""
hidden_params: Final[object] = getattr(result, "_hidden_params", None)
carried: Final[object] = _hidden_params_mapping(hidden_params).get("provider_response_model") or getattr(
hidden_params, "provider_response_model", None
)
return carried if isinstance(carried, str) and carried else None
def _provider_response_model(result: object) -> str | None:
"""The model the provider answered with, kept because the proxy later rewrites ``result.model`` to the alias the
client asked for and cost-map lookups for the standard logging payload would otherwise miss."""
model: Final = getattr(result, "model", None) if isinstance(result, ModelResponse) else None
return model if isinstance(model, str) and model else None
class ResponseMetadata:
"""
Handles setting and managing `_hidden_params`, `response_time_ms`, and `litellm_overhead_time_ms` for LiteLLM responses
@ -147,6 +174,8 @@ class ResponseMetadata:
preserve_litellm_internal_headers=True,
),
"litellm_model_name": model,
"provider_response_model": _existing_provider_response_model(self.result)
or _provider_response_model(self.result),
}
self._update_hidden_params(new_params)

149
litellm/litellm_core_utils/logging_callback_manager.py
View file

@ -1,4 +1,6 @@
from collections.abc import Callable
import threading
from collections.abc import Callable, Iterable, Iterator
from types import MethodType
from typing import TYPE_CHECKING, Final
import litellm
@ -15,6 +17,7 @@ else:
_custom_logger_compatible_callbacks_literal = str
_generic_api_logger_cache: Final[dict[str, GenericAPILogger]] = {}
_callback_lists_lock: Final = threading.RLock()
class LoggingCallbackManager:
@ -43,16 +46,18 @@ class LoggingCallbackManager:
Add a input callback to litellm.input_callback.
Auto-routes async callbacks to litellm._async_input_callback.
"""
if not isinstance(callback, str) and self._is_async_callable(callback):
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback)
else:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback)
with _callback_lists_lock:
if not isinstance(callback, str) and self._is_async_callable(callback):
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback)
else:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback)
def add_litellm_service_callback(self, callback: CustomLogger | str | Callable[..., object]):
"""
Add a service callback to litellm.service_callback
"""
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback)
with _callback_lists_lock:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback)
def add_litellm_callback(self, callback: CustomLogger | str | Callable[..., object]):
"""
@ -60,10 +65,11 @@ class LoggingCallbackManager:
Ensures no duplicates are added.
"""
self._safe_add_callback_to_list(
callback=callback,
parent_list=litellm.callbacks,
)
with _callback_lists_lock:
self._safe_add_callback_to_list(
callback=callback,
parent_list=litellm.callbacks,
)
def add_litellm_success_callback(self, callback: CustomLogger | str | Callable[..., object]):
"""
@ -71,37 +77,41 @@ class LoggingCallbackManager:
Auto-routes async callbacks to litellm._async_success_callback.
Special-cases 'dynamodb' and 'openmeter' as async callbacks.
"""
if (
isinstance(callback, str)
and callback in ("dynamodb", "openmeter")
or not isinstance(callback, str)
and self._is_async_callable(callback)
):
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback)
else:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback)
with _callback_lists_lock:
if (
isinstance(callback, str)
and callback in ("dynamodb", "openmeter")
or not isinstance(callback, str)
and self._is_async_callable(callback)
):
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback)
else:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback)
def add_litellm_failure_callback(self, callback: CustomLogger | str | Callable[..., object]):
"""
Add a failure callback to `litellm.failure_callback`.
Auto-routes async callbacks to litellm._async_failure_callback.
"""
if not isinstance(callback, str) and self._is_async_callable(callback):
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback)
else:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback)
with _callback_lists_lock:
if not isinstance(callback, str) and self._is_async_callable(callback):
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback)
else:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback)
def add_litellm_async_success_callback(self, callback: CustomLogger | Callable[..., object] | str):
"""
Add a success callback to litellm._async_success_callback
"""
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback)
with _callback_lists_lock:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback)
def add_litellm_async_failure_callback(self, callback: CustomLogger | Callable[..., object] | str):
"""
Add a failure callback to litellm._async_failure_callback
"""
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback)
with _callback_lists_lock:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback)
def remove_callback_from_list_by_object(self, callback_list, obj, require_self=True):
"""
@ -110,13 +120,14 @@ class LoggingCallbackManager:
if not isinstance(callback_list, list): # Not list -> do nothing
return
if require_self:
remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj]
else:
remove_list = [c for c in callback_list if c == obj]
with _callback_lists_lock:
if require_self:
remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj]
else:
remove_list = [c for c in callback_list if c == obj]
for c in remove_list:
callback_list.remove(c)
for c in remove_list:
callback_list.remove(c)
def remove_callbacks_by_type(self, callback_list, callback_type):
"""
@ -134,10 +145,11 @@ class LoggingCallbackManager:
if not isinstance(callback_list, list):
return
remove_list: Final = [c for c in callback_list if isinstance(c, callback_type)]
with _callback_lists_lock:
remove_list: Final = [c for c in callback_list if isinstance(c, callback_type)]
for c in remove_list:
callback_list.remove(c)
for c in remove_list:
callback_list.remove(c)
def _add_string_callback_to_list(
self, callback: str, parent_list: list[CustomLogger | Callable[..., object] | str]
@ -336,6 +348,55 @@ class LoggingCallbackManager:
return "-".join(key_parts)
get_custom_logger_key = _get_custom_logger_key
callback_lists_lock: Final = _callback_lists_lock
def drop_callbacks(self, callbacks: Iterable[object], method_owners: Iterable[object] = ()) -> frozenset[str]:
dropped_ids: Final = frozenset(map(id, callbacks))
owner_ids: Final = frozenset(map(id, method_owners))
def keep(callback: object) -> bool:
if id(callback) in dropped_ids:
return False
return not (isinstance(callback, MethodType) and id(callback.__self__) in owner_ids)
with _callback_lists_lock:
vacated_keys: Final = frozenset(
self._get_custom_logger_key(callback)
for callback in litellm.callbacks
if isinstance(callback, CustomLogger) and id(callback) in dropped_ids
)
litellm.callbacks = [callback for callback in litellm.callbacks if keep(callback)]
litellm.input_callback = [callback for callback in litellm.input_callback if keep(callback)]
litellm.success_callback = [callback for callback in litellm.success_callback if keep(callback)]
litellm.failure_callback = [callback for callback in litellm.failure_callback if keep(callback)]
litellm.service_callback = [callback for callback in litellm.service_callback if keep(callback)]
litellm._async_success_callback = [
callback for callback in litellm._async_success_callback if keep(callback)
]
litellm._async_failure_callback = [
callback for callback in litellm._async_failure_callback if keep(callback)
]
return vacated_keys
def with_custom_logger_keys(
self, candidates: Iterable[CustomLogger], keys: frozenset[str]
) -> tuple[CustomLogger, ...]:
if not keys:
return ()
class_names: Final = frozenset(key.split("-", 1)[0] for key in keys)
same_class: Final = (candidate for candidate in candidates if type(candidate).__name__ in class_names)
return tuple(candidate for candidate in same_class if self._get_custom_logger_key(candidate) in keys)
def without_custom_logger_keys(
self, callbacks: Iterable[str | Callable[..., object] | CustomLogger], keys: frozenset[str]
) -> Iterator[str | Callable[..., object] | CustomLogger]:
return (
callback
for callback in callbacks
if not (isinstance(callback, CustomLogger) and self._get_custom_logger_key(callback) in keys)
)
def _reset_all_callbacks(self):
"""
Reset all callbacks to an empty list
@ -363,19 +424,21 @@ class LoggingCallbackManager:
_get_all_callbacks = get_all_callbacks
def remove_callback_from_all_lists(self, obj, require_self=False) -> None:
def remove_callback_from_all_lists(self, obj: object, require_self: bool = False) -> None:
"""
Remove a callback object from every callback list it may have been
promoted into, so a re-initialized callback leaves no stale instance behind.
"""
for callback_list in (
litellm.callbacks,
litellm.success_callback,
litellm.failure_callback,
litellm._async_success_callback,
litellm._async_failure_callback,
):
self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self)
with _callback_lists_lock:
for callback_list in (
litellm.callbacks,
litellm.input_callback,
litellm.success_callback,
litellm.failure_callback,
litellm._async_success_callback,
litellm._async_failure_callback,
):
self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self)
def get_active_additional_logging_utils_from_custom_logger(
self,

84
litellm/litellm_core_utils/streaming_handler.py
View file

@ -10,7 +10,7 @@ import traceback
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Any, Final, NoReturn, Protocol, TypeVar, cast
from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeVar, cast
import anyio
import httpx
@ -60,6 +60,9 @@ from .exception_mapping_utils import exception_type, with_deployment_api_base
from .llm_response_utils.get_api_base import get_api_base
from .rules import Rules
if TYPE_CHECKING:
from .litellm_logging import BilledRoundsSpend
# Constants for special delta attribute names
AUDIO_ATTRIBUTE: Final = "audio"
IMAGE_ATTRIBUTE: Final = "images"
@ -259,6 +262,8 @@ def _provider_hidden_params(
class CustomStreamWrapper:
_earlier_rounds_spend: "BilledRoundsSpend | None" = None
@property
def hidden_params(self) -> dict[str, object]: # mutable-ok: API requires mutation
return self._hidden_params
@ -342,6 +347,10 @@ class CustomStreamWrapper:
self.messages = getattr(logging_obj, "messages", None)
self.sent_stream_usage = False
self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False
self.caller_requested_stream_usage: Final = (
self.check_send_stream_usage(getattr(logging_obj, "stream_options", None))
and getattr(logging_obj, "stream_usage_injected", False) is not True
)
self.tool_call = False
# keep track of the returned chunks - used for calculating the input/output tokens for stream options
self.chunks: list[ModelResponseStream] = []
@ -2070,7 +2079,7 @@ class CustomStreamWrapper:
except StopIteration:
if self.sent_last_chunk is True:
complete_streaming_response: Final = self._build_complete_streaming_response()
response = self.model_response_creator()
response = self._new_usage_chunk()
if complete_streaming_response is not None:
self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider)
@ -2222,27 +2231,10 @@ class CustomStreamWrapper:
)
if _has_usage:
# Store a copy ONLY when usage stripping below will mutate
# the chunk. For non-usage chunks (vast majority), store
# directly to avoid expensive model_copy() per chunk.
self.chunks.append(processed_chunk.model_copy())
# Strip usage from the outgoing chunk so it's not sent twice
# (once in the chunk, once in _hidden_params).
obj_dict = processed_chunk.model_dump()
if "usage" in obj_dict:
del obj_dict["usage"]
processed_chunk = self.model_response_creator(
chunk=obj_dict,
hidden_params=cast( # cast-ok: preserve mapping operations on dynamic storage
Mapping[str, object], getattr(processed_chunk, HIDDEN_PARAMS_ATTR)
),
)
is_empty = is_model_response_stream_empty(
model_response=cast(ModelResponseStream, processed_chunk)
)
if is_empty:
usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk)
if usage_free_chunk is None:
continue
processed_chunk = usage_free_chunk
else:
# No usage data — safe to store directly without copying
self.chunks.append(processed_chunk)
@ -2273,13 +2265,18 @@ class CustomStreamWrapper:
if processed_chunk is None:
continue
choice = processed_chunk.choices[0]
if isinstance(choice, StreamingChoices):
self._accumulate_streamed_delta(choice.delta)
if processed_chunk.choices:
choice = processed_chunk.choices[0]
if isinstance(choice, StreamingChoices):
self._accumulate_streamed_delta(choice.delta)
self.rules.post_call_rules(input=self.response_uptil_now, model=self.model)
# RETURN RESULT
self.chunks.append(processed_chunk)
return processed_chunk
if not self.send_stream_usage or getattr(processed_chunk, "usage", None) is None:
self.chunks.append(processed_chunk)
return processed_chunk
usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk)
if usage_free_chunk is None:
continue
return usage_free_chunk
except (StopAsyncIteration, StopIteration):
return await self._finalize_completed_stream(cache_hit=cache_hit)
except httpx.TimeoutException as e: # if httpx read timeout error occues
@ -2306,7 +2303,7 @@ class CustomStreamWrapper:
async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream":
if self.sent_last_chunk is True:
complete_streaming_response: Final = await self._async_build_complete_streaming_response()
response: Final = self.model_response_creator()
response: Final = self._new_usage_chunk()
if complete_streaming_response is not None:
self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider)
@ -2407,6 +2404,9 @@ class CustomStreamWrapper:
return
await asyncio.wait((self._failure_logging,), timeout=timeout)
def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None:
self._earlier_rounds_spend = earlier
def _record_partial_usage_for_failure(self) -> None:
"""
A stream that breaks mid-flight still billed the provider for the chunks
@ -2433,10 +2433,12 @@ class CustomStreamWrapper:
if self.model:
partial_response.model = self.model
backfill_missing_cache_usage_fields(usage)
self.logging_obj.model_call_details["combined_usage_object"] = usage
self.logging_obj.model_call_details["response_cost"] = (
self.logging_obj.response_cost_calculator(result=partial_response) or 0.0
)
partial_cost: Final = self.logging_obj.response_cost_calculator(result=partial_response) or 0.0
if self._earlier_rounds_spend is None:
self.logging_obj.model_call_details["combined_usage_object"] = usage
self.logging_obj.model_call_details["response_cost"] = partial_cost
return
self.logging_obj.record_failure_after_earlier_rounds(self._earlier_rounds_spend, usage, partial_cost)
except Exception as recover_error:
verbose_logger.debug(
"could not recover partial usage for interrupted stream: %s",
@ -2463,6 +2465,22 @@ class CustomStreamWrapper:
return False
return any(get(field) for field in cls._CONTINUATION_DISQUALIFYING_DELTA_FIELDS)
def _new_usage_chunk(self) -> ModelResponseStream:
usage_chunk: Final = self.model_response_creator()
if self.caller_requested_stream_usage:
usage_chunk.choices = []
return usage_chunk
def _record_and_strip_usage_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream | None:
self.chunks.append(chunk.model_copy())
usage_free_chunk: Final = self.model_response_creator(
chunk={key: value for key, value in chunk.model_dump().items() if key != "usage"},
hidden_params=cast( # cast-ok: preserve mapping operations on dynamic storage
Mapping[str, object], getattr(chunk, HIDDEN_PARAMS_ATTR)
),
)
return None if is_model_response_stream_empty(model_response=usage_free_chunk) else usage_free_chunk
def _accumulate_streamed_delta(self, delta: object) -> None:
# Shared by the sync, async, and non-aiohttp iteration sites so answer
# text and the disqualifying-content latch stay in step across all three.

4
litellm/llms/anthropic/pass_through/messages/interceptors/advisor.py
View file

@ -15,7 +15,7 @@ How it works:
"""
import uuid
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
import litellm
@ -48,7 +48,7 @@ class AdvisorOrchestrationHandler(MessagesInterceptor):
def can_handle(
self,
tools: list[dict] | None,
tools: Sequence[Mapping[str, object]] | None,
custom_llm_provider: str | None,
) -> bool:
if not tools:

4
litellm/llms/anthropic/pass_through/messages/interceptors/base.py
View file

@ -1,5 +1,5 @@
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Mapping, Sequence
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
@ -21,7 +21,7 @@ class MessagesInterceptor(ABC):
@abstractmethod
def can_handle(
self,
tools: list[dict] | None,
tools: Sequence[Mapping[str, object]] | None,
custom_llm_provider: str | None,
) -> bool:
"""Return True if this interceptor should handle the request."""

3
litellm/llms/base_llm/extraction/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .transformation import BaseExtractionConfig
__all__ = ["BaseExtractionConfig"]

86
litellm/llms/base_llm/extraction/transformation.py Normal file
View file

@ -0,0 +1,86 @@
from abc import ABC, abstractmethod
from collections.abc import Mapping
from typing import Final
import httpx
from pydantic import TypeAdapter, ValidationError
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret_str
from litellm.types.extraction import ExtractionRequestBody, ExtractionResponse
_PAYLOAD_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object)
_RESPONSE_ADAPTER: Final[TypeAdapter[ExtractionResponse]] = TypeAdapter(ExtractionResponse)
_RESERVED_HEADERS: Final[frozenset[str]] = frozenset({"authorization", "content-type"})
class BaseExtractionConfig(ABC):
api_key_env: tuple[str, ...] = ()
api_base_env: tuple[str, ...] = ()
def get_default_api_base(self) -> str | None:
return None
def missing_api_base_message(self, custom_llm_provider: str) -> str:
return f"api_base is required for Extraction provider '{custom_llm_provider}'"
def resolve_api_base(self, api_base: str | None) -> str | None:
return api_base or self._first_secret(self.api_base_env) or self.get_default_api_base()
def resolve_api_key(self, api_key: str | None) -> str | None:
return api_key or self._first_secret(self.api_key_env)
@staticmethod
def _first_secret(names: tuple[str, ...]) -> str | None:
return next((value for value in (get_secret_str(name) for name in names) if value), None)
def canonical_model(self, model: str) -> str:
return model
def request_model(self, model: str) -> str:
return model
def validate_environment(
self,
headers: Mapping[str, str],
model: str,
api_key: str | None,
litellm_params: Mapping[str, object],
) -> Mapping[str, str]:
return {
**{name: value for name, value in headers.items() if name.lower() not in _RESERVED_HEADERS},
**({"Authorization": f"Bearer {api_key}"} if api_key is not None else {}),
"Content-Type": "application/json",
}
@abstractmethod
def get_complete_url(self, api_base: str, model: str) -> str: ...
def transform_extraction_request(self, model: str, request: ExtractionRequestBody) -> Mapping[str, object]:
return {**request.model_dump(mode="json", exclude_unset=True), "model": self.request_model(model)}
def unwrap_response(self, payload: object) -> object:
return payload
def transform_extraction_response(
self,
model: str,
custom_llm_provider: str,
raw_response: httpx.Response,
) -> ExtractionResponse:
payload: Final[object] = _PAYLOAD_ADAPTER.validate_json(raw_response.content)
try:
return _RESPONSE_ADAPTER.validate_python(self.unwrap_response(payload))
except ValidationError as error:
raise BaseLLMException(
status_code=500,
message=f"Extraction provider '{custom_llm_provider}' returned an unexpected response: {error}",
) from error
def get_error_class(
self,
error_message: str,
status_code: int,
headers: Mapping[str, str] | httpx.Headers,
) -> BaseLLMException:
return BaseLLMException(status_code=status_code, message=error_message, headers=httpx.Headers(headers))

4
litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
View file

@ -60,7 +60,7 @@ def _response_header(headers: Mapping[str, str], name: str) -> str | None:
return headers.get(name, None)
def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]":
def bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]":
try:
guardrail_config: Final = _GUARDRAIL_CONFIG_VALIDATOR.validate_python(raw_guardrail_config)
except ValidationError as e:
@ -440,7 +440,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
guardrail_headers: Final = (
()
if raw_guardrail_config is None
else tuple(_bedrock_invoke_guardrail_headers(raw_guardrail_config).items())
else tuple(bedrock_invoke_guardrail_headers(raw_guardrail_config).items())
)
owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names)

38
litellm/llms/bedrock/passthrough/transformation.py
View file

@ -14,6 +14,9 @@ from litellm.constants import AWS_ERROR_TYPE_HEADER
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector
from litellm.llms.bedrock.chat.converse_to_openai import as_json_array, as_json_object
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
bedrock_invoke_guardrail_headers,
)
from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage
from litellm.types.utils import ModelResponseStream
@ -42,6 +45,9 @@ if TYPE_CHECKING:
BEDROCK_STREAMING_ACTIONS: Final = frozenset({"invoke-with-response-stream", "converse-stream"})
_BEDROCK_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"})
_BEDROCK_INVOKE_ACTIONS: Final = frozenset({"invoke", "invoke-with-response-stream"})
_BEDROCK_INVOKE_GUARDRAIL_HEADERS: Final = frozenset(
{"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"}
)
def _bedrock_endpoint_action(endpoint: str) -> str:
@ -123,6 +129,31 @@ def _with_invoke_request_metadata(
return merge_bedrock_invoke_headers(dict(headers), (), metadata_headers, owned_names), body
def _converse_guardrail_config(action: str, guardrail_config: object) -> object:
fields: Final = as_json_object(guardrail_config)
if action == "converse-stream" or fields is None:
return guardrail_config
return {key: value for key, value in fields.items() if key != "streamProcessingMode"}
def _with_deployment_guardrail(
action: str,
headers: Mapping[str, str],
body: Mapping[str, object],
guardrail_config: object,
) -> tuple[Mapping[str, str], Mapping[str, object]]:
if guardrail_config is None:
return headers, body
if action in _BEDROCK_CONVERSE_ACTIONS:
return headers, {**body, "guardrailConfig": _converse_guardrail_config(action, guardrail_config)}
if action in _BEDROCK_INVOKE_ACTIONS:
guardrail_headers: Final = tuple(bedrock_invoke_guardrail_headers(guardrail_config).items())
return merge_bedrock_invoke_headers(
dict(headers), (), guardrail_headers, _BEDROCK_INVOKE_GUARDRAIL_HEADERS
), body
return headers, body
def _is_cache_point(entry: object) -> bool:
fields: Final = as_json_object(entry)
return fields is not None and "cachePoint" in fields
@ -453,11 +484,14 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD
body=_converse_body_to_send(body) if action in _BEDROCK_CONVERSE_ACTIONS else body,
litellm_params=litellm_params,
)
guarded_headers, guarded_body = _with_deployment_guardrail(
action, signed_headers, signed_body, litellm_params.get("guardrailConfig")
)
return self._sign_request(
service_name="bedrock",
headers=dict(signed_headers),
headers=dict(guarded_headers),
optional_params=optional_params,
request_data=dict(signed_body),
request_data=dict(guarded_body),
api_base=api_base,
model=model,
api_key=self.get_bedrock_bearer_token(optional_params),

11
litellm/llms/bedrock/request_metadata.py
View file

@ -256,12 +256,11 @@ def merge_bedrock_invoke_headers(
"""
Merge the ``X-Amzn-*`` headers the Invoke paths derive from params.
``caller_owned`` (the guardrail headers) defers to a header the caller already set, which is
the long-standing behaviour for those. ``proxy_owned_names`` are dropped from the caller's
headers unconditionally and re-supplied only from ``proxy_owned``, because those names carry
proxy-authenticated identity into an AWS billing record that the caller must not be able to
write. Names are compared case-insensitively so a caller cannot leave a second spelling in
the dict and let the transport pick the winner.
``caller_owned`` defers to a header the caller already set. ``proxy_owned_names`` are dropped
from the caller's headers unconditionally and re-supplied only from ``proxy_owned``, for names
the caller must not be able to write, such as proxy-authenticated identity headed for an AWS
billing record or a deployment's guardrail. Names are compared case-insensitively so a caller
cannot leave a second spelling in the dict and let the transport pick the winner.
"""
if not caller_owned and not proxy_owned and not proxy_owned_names:
return headers

53
litellm/llms/cloudflare/common_utils.py Normal file
View file

@ -0,0 +1,53 @@
from collections.abc import Mapping
from typing import Final
from pydantic import TypeAdapter, ValidationError
from litellm.secret_managers.main import (
get_secret_str,
normalize_nonempty_secret_str,
)
CLOUDFLARE_MISSING_API_BASE_MESSAGE: Final = (
"Missing CLOUDFLARE_ACCOUNT_ID - set CLOUDFLARE_ACCOUNT_ID or pass api_base explicitly"
)
_RESPONSE_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
def workers_ai_default_api_base() -> str | None:
account_id: Final = normalize_nonempty_secret_str(get_secret_str("CLOUDFLARE_ACCOUNT_ID"))
if account_id is None:
return None
return f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run"
def workers_ai_canonical_model(model: str) -> str:
if model.startswith("@cf/"):
return model
return f"@cf/cloudflare/{model}"
def workers_ai_request_model(model: str) -> str:
return model.rsplit("/", maxsplit=1)[-1]
def workers_ai_run_url(api_base: str, model: str) -> str:
normalized_api_base: Final = api_base.rstrip("/")
if normalized_api_base.endswith("/ai/v1"):
return f"{normalized_api_base.removesuffix('/ai/v1')}/ai/run/{model}"
if normalized_api_base.endswith("/ai/run"):
return f"{normalized_api_base}/{model}"
return f"{normalized_api_base}/ai/run/{model}"
def unwrap_workers_ai_result(payload: object) -> object:
try:
response_mapping: Final = _RESPONSE_MAPPING_ADAPTER.validate_python(payload)
except ValidationError:
return payload
if "answers" in response_mapping:
return response_mapping
try:
return _RESPONSE_MAPPING_ADAPTER.validate_python(response_mapping.get("result"))
except ValidationError:
return response_mapping

48
litellm/llms/cloudflare/decisions/transformation.py
View file

@ -1,16 +1,13 @@
from collections.abc import Mapping
from typing import Final
from pydantic import TypeAdapter, ValidationError
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
from litellm.secret_managers.main import (
get_secret_str,
normalize_nonempty_secret_str,
from litellm.llms.cloudflare.common_utils import (
CLOUDFLARE_MISSING_API_BASE_MESSAGE,
unwrap_workers_ai_result,
workers_ai_canonical_model,
workers_ai_default_api_base,
workers_ai_request_model,
workers_ai_run_url,
)
_RESPONSE_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
class CloudflareDecisionsConfig(BaseDecisionsConfig):
api_key_env = ("CLOUDFLARE_API_KEY",)
@ -23,38 +20,19 @@ class CloudflareDecisionsConfig(BaseDecisionsConfig):
inlines_remote_images = True
def get_default_api_base(self) -> str | None:
account_id: Final = normalize_nonempty_secret_str(get_secret_str("CLOUDFLARE_ACCOUNT_ID"))
if account_id is None:
return None
return f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run"
return workers_ai_default_api_base()
def missing_api_base_message(self, custom_llm_provider: str) -> str:
return "Missing CLOUDFLARE_ACCOUNT_ID - set CLOUDFLARE_ACCOUNT_ID or pass api_base explicitly"
return CLOUDFLARE_MISSING_API_BASE_MESSAGE
def canonical_model(self, model: str) -> str:
if model.startswith("@cf/"):
return model
return f"@cf/cloudflare/{model}"
return workers_ai_canonical_model(model)
def request_model(self, model: str) -> str:
return model.rsplit("/", maxsplit=1)[-1]
return workers_ai_request_model(model)
def get_complete_url(self, api_base: str, model: str) -> str:
normalized_api_base: Final = api_base.rstrip("/")
if normalized_api_base.endswith("/ai/v1"):
return f"{normalized_api_base.removesuffix('/ai/v1')}/ai/run/{model}"
if normalized_api_base.endswith("/ai/run"):
return f"{normalized_api_base}/{model}"
return f"{normalized_api_base}/ai/run/{model}"
return workers_ai_run_url(api_base, model)
def unwrap_response(self, payload: object) -> object:
try:
response_mapping: Final = _RESPONSE_MAPPING_ADAPTER.validate_python(payload)
except ValidationError:
return payload
if "answers" in response_mapping:
return response_mapping
try:
return _RESPONSE_MAPPING_ADAPTER.validate_python(response_mapping.get("result"))
except ValidationError:
return response_mapping
return unwrap_workers_ai_result(payload)

11
litellm/llms/cloudflare/extraction/guardrail_translation/__init__.py Normal file
View file

@ -0,0 +1,11 @@
from typing import Final
from litellm.llms.cloudflare.extraction.guardrail_translation.handler import ExtractionHandler
from litellm.types.utils import CallTypes
guardrail_translation_mappings: Final = {
CallTypes.extraction: ExtractionHandler,
CallTypes.aextraction: ExtractionHandler,
}
__all__ = ["ExtractionHandler", "guardrail_translation_mappings"]

276
litellm/llms/cloudflare/extraction/guardrail_translation/handler.py Normal file
View file

@ -0,0 +1,276 @@
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass
from itertools import chain, groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeAlias, TypeVar
from pydantic import Field, JsonValue, TypeAdapter
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.extraction import (
EXTRACTION_STATE_REWRITTEN_METADATA_KEY,
ExtractionAnswer,
ExtractionBase64Image,
ExtractionResponse,
)
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
_FIELDS_ADAPTER: Final[TypeAdapter[Mapping[str, Mapping[str, object]]]] = TypeAdapter(
Mapping[str, Mapping[str, object]]
)
_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
_IMAGES_ADAPTER: Final[TypeAdapter[tuple[ExtractionBase64Image | JsonValue, ...]]] = TypeAdapter(
tuple[Annotated[ExtractionBase64Image | JsonValue, Field(union_mode="left_to_right")], ...]
)
_AnswerKey: TypeAlias = Literal["value", "guess", "mean", "median", "p10", "p90", "quantiles"]
_ESTIMATE_KEYS: Final[tuple[_AnswerKey, ...]] = ("value", "guess", "mean", "median", "p10", "p90", "quantiles")
_STATISTIC_KEYS: Final[frozenset[_AnswerKey]] = frozenset(("mean", "median", "p10", "p90", "quantiles"))
_NULLED_ANSWER: Final[Mapping[str, None]] = MappingProxyType(
dict.fromkeys(("value", "guess", "mean", "median", "p10", "p90", "quantiles", "steps", "span", "spans"), None)
)
@dataclass(frozen=True, slots=True)
class _AnswerText:
name: str
key: _AnswerKey
item: int | None
text: str
_SPEC_TEXT_KEYS: Final = ("instructions", "unit")
@dataclass(frozen=True, slots=True)
class _SpecText:
name: str
key: str
text: str
def _spec_texts_of(name: str, spec: Mapping[str, object]) -> tuple[_SpecText, ...]:
return tuple(
_SpecText(name=name, key=key, text=text) for key in _SPEC_TEXT_KEYS if isinstance(text := spec.get(key), str)
)
def _spec_texts(fields: Mapping[str, Mapping[str, object]]) -> tuple[_SpecText, ...]:
return tuple(chain.from_iterable(_spec_texts_of(name, spec) for name, spec in fields.items()))
_NamedText = TypeVar("_NamedText", _AnswerText, _SpecText)
def _scans_by_name(scanned: Iterable[tuple[_NamedText, str]]) -> Mapping[str, tuple[tuple[_NamedText, str], ...]]:
return {name: tuple(group) for name, group in groupby(scanned, key=lambda pair: pair[0].name)}
def _guardrailed_spec(spec: Mapping[str, object], scanned: Sequence[tuple[_SpecText, str]]) -> Mapping[str, object]:
return {**spec, **{text.key: guardrailed for text, guardrailed in scanned}}
def _metadata_key(data: Mapping[str, object]) -> Literal["metadata", "litellm_metadata"]:
return get_metadata_variable_name_from_kwargs(data)
def _metadata(data: Mapping[str, object]) -> Mapping[str, object]:
return _METADATA_ADAPTER.validate_python(data.get(_metadata_key(data)) or {})
def _is_text_answer(answer: ExtractionAnswer, spec: Mapping[str, object] | None) -> bool:
return answer.type == "text" or (answer.type == "list" and (spec is None or spec.get("items") == "text"))
def _image_urls(images: object) -> tuple[str, ...]:
if not isinstance(images, list):
return ()
return tuple(
f"data:{image.content_type};base64,{image.base64}" if isinstance(image, ExtractionBase64Image) else image
for image in _IMAGES_ADAPTER.validate_python(images)
if isinstance(image, (str, ExtractionBase64Image))
)
def _scalar_text(value: JsonValue, whole: bool) -> str | None:
if isinstance(value, bool):
return None
if isinstance(value, float) and whole:
return format(value, ".0f")
return str(value) if isinstance(value, (str, int, float)) else None
def _item_texts(items: Sequence[JsonValue], whole: bool) -> tuple[tuple[int, str], ...]:
return tuple((index, text) for index, item in enumerate(items) if (text := _scalar_text(item, whole)) is not None)
def _value_texts(name: str, key: _AnswerKey, value: JsonValue) -> tuple[_AnswerText, ...]:
whole: Final = key in _STATISTIC_KEYS
if isinstance(value, list):
return tuple(
_AnswerText(name=name, key=key, item=index, text=text) for index, text in _item_texts(value, whole)
)
text: Final = _scalar_text(value, whole)
return () if text is None else (_AnswerText(name=name, key=key, item=None, text=text),)
def _step_texts(name: str, step: JsonValue) -> tuple[_AnswerText, ...]:
if not isinstance(step, dict):
return ()
return tuple(chain.from_iterable(_value_texts(name, key, step.get(key)) for key in _ESTIMATE_KEYS))
def _answer_texts(name: str, answer: ExtractionAnswer) -> tuple[_AnswerText, ...]:
return (
*_value_texts(name, "value", answer.value),
*_value_texts(name, "guess", answer.guess),
*_value_texts(name, "mean", answer.mean),
*_value_texts(name, "median", answer.median),
*_value_texts(name, "p10", answer.p10),
*_value_texts(name, "p90", answer.p90),
*_value_texts(name, "quantiles", list(answer.quantiles or ())),
*chain.from_iterable(_step_texts(name, step) for step in answer.steps or ()),
)
def _guardrailed_value(value: JsonValue, scanned: Mapping[int | None, str]) -> JsonValue:
if isinstance(value, list):
return [scanned.get(index, item) for index, item in enumerate(value)]
return scanned.get(None, value)
def _scans_of(key: _AnswerKey, scanned: Sequence[tuple[_AnswerText, str]]) -> Mapping[int | None, str]:
return {text.item: guardrailed for text, guardrailed in scanned if text.key == key}
def _guardrailed_answer(
answer: ExtractionAnswer,
scanned: Sequence[tuple[_AnswerText, str]],
drop_spans: bool,
is_text: bool,
) -> ExtractionAnswer:
if not is_text and any(text.text != guardrailed for text, guardrailed in scanned):
return answer.model_copy(update=_NULLED_ANSWER)
value_scans: Final = _scans_of("value", scanned if is_text else ())
guess_scans: Final = _scans_of("guess", scanned if is_text else ())
update: Final = {
**({"value": _guardrailed_value(answer.value, value_scans)} if value_scans else {}),
**({"guess": _guardrailed_value(answer.guess, guess_scans)} if guess_scans else {}),
**({"span": None, "spans": None} if drop_spans else {}),
}
return answer.model_copy(update=update) if update else answer
async def _scanned_texts(
texts: Sequence[str],
model: object,
guardrail_to_apply: "CustomGuardrail",
request_data: dict[str, object], # mutable-ok: CustomGuardrail.apply_guardrail takes the request as a plain dict
input_type: Literal["request", "response"],
litellm_logging_obj: "LiteLLMLoggingObj | None",
images: Sequence[str] = (),
) -> tuple[str, ...]:
unique_texts: Final = tuple(dict.fromkeys(texts))
inputs: Final = GenericGuardrailAPIInputs(texts=list(unique_texts), model=model if isinstance(model, str) else None)
if images:
inputs["images"] = list(images)
guardrailed: Final = await guardrail_to_apply.apply_guardrail( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.apply_guardrail types request_data as a bare dict
inputs=inputs,
request_data=request_data,
input_type=input_type,
logging_obj=litellm_logging_obj,
)
returned: Final = guardrailed.get("texts", [])
guardrailed_by_text: Final = MappingProxyType(
{text: returned[index] if index < len(returned) else text for index, text in enumerate(unique_texts)}
)
return tuple(guardrailed_by_text[text] for text in texts)
class ExtractionHandler(BaseTranslation):
async def process_input_messages(
self,
data: dict[str, object], # mutable-ok: BaseTranslation hands over the proxy's request dict
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
) -> Mapping[str, object]:
state: Final = data.get("state")
states: Final = (state,) if isinstance(state, str) and state else ()
fields: Final = _FIELDS_ADAPTER.validate_python(data.get("fields") or {})
spec_texts: Final = _spec_texts(fields)
images: Final = _image_urls(data.get("images"))
if not states and not spec_texts and not images:
return data
scanned: Final = await _scanned_texts(
(*states, *(text.text for text in spec_texts)),
data.get("model"),
guardrail_to_apply,
data,
"request",
litellm_logging_obj,
images,
)
spec_scans: Final = _scans_by_name(zip(spec_texts, scanned[len(states) :]))
state_rewritten: Final = bool(states) and scanned[0] != states[0]
return {
**data,
**({"state": scanned[0]} if states else {}),
**(
{_metadata_key(data): {**_metadata(data), EXTRACTION_STATE_REWRITTEN_METADATA_KEY: True}}
if state_rewritten
else {}
),
"fields": {name: _guardrailed_spec(spec, spec_scans.get(name, ())) for name, spec in fields.items()},
}
async def process_output_response(
self,
response: object,
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict[str, object] | None = None, # mutable-ok: BaseTranslation hands over the request dict
) -> object:
if not isinstance(response, ExtractionResponse):
return response
texts: Final = tuple(
chain.from_iterable(_answer_texts(name, answer) for name, answer in response.answers.items())
)
state_texts: Final = (response.state_text,) if response.state_text else ()
request: Final = request_data or {}
scanned: Final = (
await _scanned_texts(
(*(text.text for text in texts), *state_texts),
response.model,
guardrail_to_apply,
request,
"response",
litellm_logging_obj,
)
if texts or state_texts
else ()
)
answer_scans: Final = _scans_by_name(zip(texts, scanned))
state_text: Final = next(iter(scanned[len(texts) :]), response.state_text)
specs: Final = _FIELDS_ADAPTER.validate_python(request.get("fields") or {})
return response.model_copy(
update={
"answers": {
name: _guardrailed_answer(
answer,
answer_scans.get(name, ()),
drop_spans=state_text != response.state_text,
is_text=_is_text_answer(answer, specs.get(name)),
)
for name, answer in response.answers.items()
},
**({"state_text": state_text} if state_texts else {}),
}
)

32
litellm/llms/cloudflare/extraction/transformation.py Normal file
View file

@ -0,0 +1,32 @@
from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig
from litellm.llms.cloudflare.common_utils import (
CLOUDFLARE_MISSING_API_BASE_MESSAGE,
unwrap_workers_ai_result,
workers_ai_canonical_model,
workers_ai_default_api_base,
workers_ai_request_model,
workers_ai_run_url,
)
class CloudflareExtractionConfig(BaseExtractionConfig):
api_key_env = ("CLOUDFLARE_API_KEY",)
api_base_env = ("CLOUDFLARE_API_BASE",)
def get_default_api_base(self) -> str | None:
return workers_ai_default_api_base()
def missing_api_base_message(self, custom_llm_provider: str) -> str:
return CLOUDFLARE_MISSING_API_BASE_MESSAGE
def canonical_model(self, model: str) -> str:
return workers_ai_canonical_model(model)
def request_model(self, model: str) -> str:
return workers_ai_request_model(model)
def get_complete_url(self, api_base: str, model: str) -> str:
return workers_ai_run_url(api_base, model)
def unwrap_response(self, payload: object) -> object:
return unwrap_workers_ai_result(payload)

112
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -76,6 +76,7 @@ from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
from litellm.llms.base_llm.evals.transformation import BaseEvalsAPIConfig
from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig
from litellm.llms.base_llm.files.transformation import (
BaseFilesConfig,
BaseFileUploadStream,
@ -128,6 +129,7 @@ from litellm.types.containers.main import (
DeleteContainerResult,
)
from litellm.types.decisions import DecisionsIRRequest, DecisionsIRResponse
from litellm.types.extraction import ExtractionResponse
from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig
from litellm.types.integrations.custom_logger import (
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
@ -1700,6 +1702,115 @@ class BaseLLMHTTPHandler:
model=model, custom_llm_provider=custom_llm_provider, raw_response=response, request=request
)
def _prepare_extraction_request(
self,
model: str,
logging_obj: LiteLLMLoggingObj | None,
provider_config: BaseExtractionConfig,
body: Mapping[str, object],
api_base: str,
api_key: str | None,
headers: Mapping[str, str],
litellm_params: Mapping[str, object],
) -> tuple[str, Mapping[str, str]]:
outbound_headers: Final = provider_config.validate_environment(
headers=headers,
model=model,
api_key=api_key,
litellm_params=litellm_params,
)
url: Final = provider_config.get_complete_url(api_base=api_base, model=model)
if logging_obj is not None:
logging_obj.pre_call( # pyright: ignore[reportUnknownMemberType] # Logging.pre_call is untyped
input=dict(body),
api_key=api_key,
model=model,
additional_args={"api_base": url, "complete_input_dict": dict(body), "headers": outbound_headers},
)
return url, outbound_headers
def extraction(
self,
model: str,
custom_llm_provider: str,
logging_obj: LiteLLMLoggingObj | None,
provider_config: BaseExtractionConfig,
body: Mapping[str, object],
api_base: str,
api_key: str | None,
headers: Mapping[str, str],
timeout: float | httpx.Timeout | None,
litellm_params: Mapping[str, object],
client: HTTPHandler | None = None,
) -> ExtractionResponse:
url, outbound_headers = self._prepare_extraction_request(
model=model,
logging_obj=logging_obj,
provider_config=provider_config,
body=body,
api_base=api_base,
api_key=api_key,
headers=headers,
litellm_params=litellm_params,
)
sync_httpx_client: Final = client if client is not None else get_httpx_client()
try:
response: Final = sync_httpx_client.post( # pyright: ignore[reportUnknownMemberType] # untyped
url,
json=dict(body),
headers=dict(outbound_headers),
timeout=timeout,
logging_obj=logging_obj,
)
except httpx.HTTPError as e:
raise self._handle_error(e=e, provider_config=provider_config)
return provider_config.transform_extraction_response(
model=model, custom_llm_provider=custom_llm_provider, raw_response=response
)
async def aextraction(
self,
model: str,
custom_llm_provider: str,
logging_obj: LiteLLMLoggingObj | None,
provider_config: BaseExtractionConfig,
body: Mapping[str, object],
api_base: str,
api_key: str | None,
headers: Mapping[str, str],
timeout: float | httpx.Timeout | None,
litellm_params: Mapping[str, object],
client: AsyncHTTPHandler | None = None,
) -> ExtractionResponse:
url, outbound_headers = self._prepare_extraction_request(
model=model,
logging_obj=logging_obj,
provider_config=provider_config,
body=body,
api_base=api_base,
api_key=api_key,
headers=headers,
litellm_params=litellm_params,
)
async_httpx_client: Final = (
client
if client is not None
else get_async_httpx_client(llm_provider=litellm.LlmProviders(custom_llm_provider))
)
try:
response: Final = await async_httpx_client.post( # pyright: ignore[reportUnknownMemberType] # untyped
url,
json=dict(body),
headers=dict(outbound_headers),
timeout=timeout,
logging_obj=logging_obj,
)
except httpx.HTTPError as e:
raise self._handle_error(e=e, provider_config=provider_config)
return provider_config.transform_extraction_response(
model=model, custom_llm_provider=custom_llm_provider, raw_response=response
)
def _prepare_audio_transcription_request(
self,
model: str,
@ -6462,6 +6573,7 @@ class BaseLLMHTTPHandler:
provider_config: Union[
BaseConfig,
BaseDecisionsConfig,
BaseExtractionConfig,
BaseRerankConfig,
BaseResponsesAPIConfig,
BaseImageEditConfig,

4
litellm/llms/gemini/chat/transformation.py
View file

@ -82,9 +82,9 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
return "tts" in model
def get_supported_openai_params(self, model: str) -> list[str]:
sampling_params: Final = ("temperature", "top_p") if self.supports_sampling_params(model) else ()
supported_params: Final = [
"temperature",
"top_p",
*sampling_params,
"max_tokens",
"max_completion_tokens",
"stream",

78
litellm/llms/github_copilot/github_host.py Normal file
View file

@ -0,0 +1,78 @@
import os
import re
from dataclasses import dataclass
from typing import Final
from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST
from .authenticator import (
DEFAULT_GITHUB_ACCESS_TOKEN_URL,
DEFAULT_GITHUB_API_KEY_URL,
DEFAULT_GITHUB_CLIENT_ID,
DEFAULT_GITHUB_DEVICE_CODE_URL,
)
_GHE_SUBDOMAIN_PATTERN: Final = re.compile(r"[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?")
@dataclass(frozen=True, slots=True)
class GithubHostEndpoints:
host: str
device_code_url: str
access_token_url: str
api_key_url: str
user_url: str
client_id: str
copilot_api_hosts: frozenset[str]
copilot_api_host_suffix: str | None
default_copilot_api_base: str
def normalized_github_host(raw: object) -> str | None:
if raw is None or raw == "":
return GITHUB_COPILOT_DEFAULT_HOST
if not isinstance(raw, str):
return None
host: Final = raw.lower()
if host == GITHUB_COPILOT_DEFAULT_HOST:
return host
if not host.endswith(".ghe.com"):
return None
subdomain: Final = host.removesuffix(".ghe.com")
return host if _GHE_SUBDOMAIN_PATTERN.fullmatch(subdomain) is not None else None
def github_host_endpoints(host: str) -> GithubHostEndpoints:
normalized: Final = normalized_github_host(host)
if normalized is None:
raise ValueError("host must be github.com or a single-subdomain GHE.com hostname")
client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
if normalized == GITHUB_COPILOT_DEFAULT_HOST:
return GithubHostEndpoints(
host=normalized,
device_code_url=os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL),
access_token_url=os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL),
api_key_url=os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL),
user_url="https://api.github.com/user",
client_id=client_id,
copilot_api_hosts=frozenset({"githubcopilot.com"}),
copilot_api_host_suffix=".githubcopilot.com",
default_copilot_api_base="https://api.githubcopilot.com",
)
subdomain: Final = normalized.removesuffix(".ghe.com")
api_host: Final = f"api.{subdomain}.ghe.com"
copilot_api_host: Final = f"copilot-api.{subdomain}.ghe.com"
return GithubHostEndpoints(
host=normalized,
device_code_url=f"https://{normalized}/login/device/code",
access_token_url=f"https://{normalized}/login/oauth/access_token",
api_key_url=f"https://{api_host}/copilot_internal/v2/token",
user_url=f"https://{api_host}/user",
client_id=client_id,
copilot_api_hosts=frozenset({copilot_api_host}),
copilot_api_host_suffix=None,
default_copilot_api_base=f"https://{copilot_api_host}",
)

211
litellm/llms/github_copilot/per_user_auth.py
View file

@ -5,8 +5,7 @@ touched on the per-user path.
"""
import asyncio
import hashlib
import os
import hmac
import threading
import time
from collections.abc import Callable, Mapping
@ -27,6 +26,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import (
GITHUB_COPILOT_AUTH_TYPE_KEY,
GITHUB_COPILOT_HOST_KEY,
GITHUB_COPILOT_PER_USER_AUTH_TYPE,
GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS,
)
@ -38,14 +38,8 @@ from litellm.exceptions import (
ServiceUnavailableError,
)
from .authenticator import (
DEFAULT_GITHUB_ACCESS_TOKEN_URL,
DEFAULT_GITHUB_API_KEY_URL,
DEFAULT_GITHUB_CLIENT_ID,
DEFAULT_GITHUB_DEVICE_CODE_URL,
github_api_headers,
)
from .common_utils import DEFAULT_GITHUB_COPILOT_API_BASE
from .authenticator import github_api_headers
from .github_host import GithubHostEndpoints, github_host_endpoints, normalized_github_host
GITHUB_COPILOT_USER_SESSION_KWARG_KEY: Final = "github_copilot_user_session"
_LLM_PROVIDER: Final = "github_copilot"
@ -98,7 +92,7 @@ _SESSION_CACHE: Final = cast( # cast-ok: pins the untyped module cache to _Sess
)
_EXCHANGE_LOCKS: Final = tuple(threading.Lock() for _ in range(_SYNC_LOCK_COUNT))
_IN_FLIGHT: Final[ # mutable-ok: single-flight dedup registry mutated under its own locking
dict[str, "asyncio.Future[GithubCopilotUserSession]"]
dict[str, tuple[str, "asyncio.Future[GithubCopilotUserSession]"]]
] = {}
@ -141,7 +135,7 @@ _DEVICE_CODE: Final = TypeAdapter(_DeviceCodePayload)
_ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPollPayload)
_GITHUB_USER: Final = TypeAdapter(_GithubUserPayload)
_COPILOT_ENDPOINTS: Final = TypeAdapter(_CopilotEndpointsPayload)
_CACHED_SESSION_ENTRY: Final = TypeAdapter(tuple[str, str, float])
_CACHED_SESSION_ENTRY: Final = TypeAdapter(tuple[str, str, str, float])
@dataclass(frozen=True, slots=True, repr=False)
@ -178,53 +172,72 @@ class GithubCopilotDeviceFlowPoll:
access_token: str | None = field(default=None, repr=False)
def validated_copilot_api_base(endpoints_api: object) -> str:
"""Only a genuine *.githubcopilot.com HTTPS endpoint may receive a Copilot
token; anything else falls back to the default base."""
def validated_copilot_api_base(endpoints_api: object, endpoints: GithubHostEndpoints) -> str:
if not isinstance(endpoints_api, str):
return DEFAULT_GITHUB_COPILOT_API_BASE
parsed: Final = urlparse(endpoints_api)
return endpoints.default_copilot_api_base
try:
parsed: Final = urlparse(endpoints_api)
port: Final = parsed.port
except ValueError:
return endpoints.default_copilot_api_base
hostname: Final = (parsed.hostname or "").lower()
if (
parsed.scheme.lower() != "https"
or not (hostname == "githubcopilot.com" or hostname.endswith(".githubcopilot.com"))
or not (
hostname in endpoints.copilot_api_hosts
or (endpoints.copilot_api_host_suffix is not None and hostname.endswith(endpoints.copilot_api_host_suffix))
)
or parsed.username is not None
or parsed.password is not None
or port not in (None, 443)
):
return DEFAULT_GITHUB_COPILOT_API_BASE
try:
port: Final = parsed.port
except ValueError:
return DEFAULT_GITHUB_COPILOT_API_BASE
if port not in (None, 443):
return DEFAULT_GITHUB_COPILOT_API_BASE
return endpoints.default_copilot_api_base
return endpoints_api.rstrip("/")
def _session_cache_key(user_id: str, github_token: str) -> str:
return hashlib.sha256(f"{user_id}\x00{github_token}".encode()).hexdigest()
def github_copilot_host_endpoints(credential_name: str) -> GithubHostEndpoints:
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
credential: Final = CredentialAccessor.find_credential(credential_name)
host_raw: Final[object] = (
credential.credential_values.get(GITHUB_COPILOT_HOST_KEY) if credential is not None else None
)
host: Final = normalized_github_host(host_raw)
if host is None:
raise BadRequestError(
message=f"Credential '{credential_name}' has an invalid {GITHUB_COPILOT_HOST_KEY}",
llm_provider=_LLM_PROVIDER,
model="",
)
return github_host_endpoints(host)
def evict_copilot_user_session(user_id: str, github_token: str) -> None:
_SESSION_CACHE.delete_cache(_session_cache_key(user_id, github_token))
def _session_cache_key(user_id: str, credential_name: str, github_host: str) -> str:
return f"{user_id}\x00{credential_name}\x00{github_host}"
def _cached_session(cache_key: str) -> GithubCopilotUserSession | None:
def evict_copilot_user_session(user_id: str, credential_name: str, github_host: str) -> None:
_SESSION_CACHE.delete_cache(_session_cache_key(user_id, credential_name, github_host))
def _cached_session(cache_key: str, github_token: str) -> GithubCopilotUserSession | None:
cached: Final[object] = _SESSION_CACHE.get_cache(cache_key)
try:
token, api_base, expires_at = _CACHED_SESSION_ENTRY.validate_python(cached)
cached_github_token, token, api_base, expires_at = _CACHED_SESSION_ENTRY.validate_python(cached)
except ValidationError:
return None
if not hmac.compare_digest(cached_github_token, github_token):
return None
if expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS <= time.time():
return None
return GithubCopilotUserSession(token=token, api_base=api_base)
def _cache_session(cache_key: str, session: GithubCopilotUserSession, expires_at: float) -> None:
def _cache_session(cache_key: str, github_token: str, session: GithubCopilotUserSession, expires_at: float) -> None:
ttl: Final = expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS - time.time()
if ttl <= 0:
return
_SESSION_CACHE.set_cache(cache_key, (session.token, session.api_base, expires_at), ttl=ttl)
_SESSION_CACHE.set_cache(cache_key, (github_token, session.token, session.api_base, expires_at), ttl=ttl)
def _reject_connection(credential_name: str) -> CallerCredentialAuthenticationError:
@ -238,7 +251,13 @@ def _reject_connection(credential_name: str) -> CallerCredentialAuthenticationEr
)
def _session_from_response(response: httpx.Response, credential_name: str, cache_key: str) -> GithubCopilotUserSession:
def _session_from_response(
response: httpx.Response,
github_token: str,
credential_name: str,
cache_key: str,
endpoints: GithubHostEndpoints,
) -> GithubCopilotUserSession:
if response.status_code in (401, 403, 404):
_SESSION_CACHE.delete_cache(cache_key)
raise _reject_connection(credential_name)
@ -263,15 +282,17 @@ def _session_from_response(response: httpx.Response, credential_name: str, cache
model="",
) from e
endpoints: Final[object] = payload.get("endpoints")
payload_endpoints: Final[object] = payload.get("endpoints")
endpoints_api: Final[object] = (
_COPILOT_ENDPOINTS.validate_python(endpoints).get("api") if isinstance(endpoints, Mapping) else None
_COPILOT_ENDPOINTS.validate_python(payload_endpoints).get("api")
if isinstance(payload_endpoints, Mapping)
else None
)
api_base: Final = validated_copilot_api_base(endpoints_api)
api_base: Final = validated_copilot_api_base(endpoints_api, endpoints)
session: Final = GithubCopilotUserSession(token=payload["token"], api_base=api_base)
expires_at: Final = _expires_at_seconds(payload.get("expires_at"))
if expires_at > 0:
_cache_session(cache_key, session, expires_at)
_cache_session(cache_key, github_token, session, expires_at)
return session
@ -303,7 +324,12 @@ def _expires_at_seconds(raw: object) -> float:
return 0.0
def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> GithubCopilotUserSession:
def _do_exchange(
github_token: str,
credential_name: str,
cache_key: str,
endpoints: GithubHostEndpoints,
) -> GithubCopilotUserSession:
import litellm
try:
@ -311,7 +337,7 @@ def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> Git
_SyncGetClient, litellm.module_level_client
)
response: Final = client.get(
DEFAULT_GITHUB_API_KEY_URL,
endpoints.api_key_url,
headers=github_api_headers(github_token),
)
except Exception as e:
@ -320,7 +346,7 @@ def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> Git
llm_provider=_LLM_PROVIDER,
model="",
) from e
return _session_from_response(response, credential_name, cache_key)
return _session_from_response(response, github_token, credential_name, cache_key, endpoints)
def _new_async_github_client() -> _AsyncGitHubClient:
@ -332,11 +358,16 @@ def _new_async_github_client() -> _AsyncGitHubClient:
)(_LLM_PROVIDER)
async def _do_aexchange(github_token: str, credential_name: str, cache_key: str) -> GithubCopilotUserSession:
async def _do_aexchange(
github_token: str,
credential_name: str,
cache_key: str,
endpoints: GithubHostEndpoints,
) -> GithubCopilotUserSession:
client: Final = _new_async_github_client()
try:
response: Final = await client.get(
DEFAULT_GITHUB_API_KEY_URL,
endpoints.api_key_url,
headers=github_api_headers(github_token),
)
except Exception as e:
@ -345,43 +376,50 @@ async def _do_aexchange(github_token: str, credential_name: str, cache_key: str)
llm_provider=_LLM_PROVIDER,
model="",
) from e
return _session_from_response(response, credential_name, cache_key)
return _session_from_response(response, github_token, credential_name, cache_key, endpoints)
def exchange_github_token(
user_id: str,
github_token: str,
credential_name: str,
endpoints: GithubHostEndpoints,
) -> GithubCopilotUserSession:
cache_key: Final = _session_cache_key(user_id, github_token)
cached: Final = _cached_session(cache_key)
cache_key: Final = _session_cache_key(user_id, credential_name, endpoints.host)
cached: Final = _cached_session(cache_key, github_token)
if cached is not None:
return cached
lock: Final = _EXCHANGE_LOCKS[int(cache_key, 16) % _SYNC_LOCK_COUNT]
lock: Final = _EXCHANGE_LOCKS[hash(cache_key) % _SYNC_LOCK_COUNT]
with lock:
rechecked: Final = _cached_session(cache_key)
rechecked: Final = _cached_session(cache_key, github_token)
if rechecked is not None:
return rechecked
return _do_exchange(github_token, credential_name, cache_key)
return _do_exchange(github_token, credential_name, cache_key, endpoints)
async def aexchange_github_token(
user_id: str,
github_token: str,
credential_name: str,
endpoints: GithubHostEndpoints,
) -> GithubCopilotUserSession:
cache_key: Final = _session_cache_key(user_id, github_token)
cached: Final = _cached_session(cache_key)
cache_key: Final = _session_cache_key(user_id, credential_name, endpoints.host)
cached: Final = _cached_session(cache_key, github_token)
if cached is not None:
return cached
loop: Final = asyncio.get_running_loop()
in_flight: Final = _IN_FLIGHT.get(cache_key)
if in_flight is not None and not in_flight.done() and in_flight.get_loop() is loop:
return await asyncio.shield(in_flight)
in_flight_entry: Final = _IN_FLIGHT.get(cache_key)
if in_flight_entry is not None:
in_flight_token, in_flight = in_flight_entry
if not hmac.compare_digest(in_flight_token, github_token):
return await _do_aexchange(github_token, credential_name, cache_key, endpoints)
if not in_flight.done() and in_flight.get_loop() is loop:
return await asyncio.shield(in_flight)
future: Final = loop.create_future()
_IN_FLIGHT[cache_key] = future
entry: Final = (github_token, future)
_IN_FLIGHT[cache_key] = entry
try:
session: Final = await _do_aexchange(github_token, credential_name, cache_key)
session: Final = await _do_aexchange(github_token, credential_name, cache_key, endpoints)
except BaseException as e:
if not future.done():
future.set_exception(e)
@ -392,19 +430,17 @@ async def aexchange_github_token(
future.set_result(session)
return session
finally:
if _IN_FLIGHT.get(cache_key) is future:
if _IN_FLIGHT.get(cache_key) is entry:
_IN_FLIGHT.pop(cache_key, None)
async def astart_device_flow() -> GithubCopilotDeviceFlowStart:
async def astart_device_flow(endpoints: GithubHostEndpoints) -> GithubCopilotDeviceFlowStart:
client: Final = _new_async_github_client()
device_code_url: Final = os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL)
client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
try:
response: Final = await client.post(
device_code_url,
endpoints.device_code_url,
headers=github_api_headers(),
json={"client_id": client_id, "scope": _DEVICE_FLOW_SCOPE},
json={"client_id": endpoints.client_id, "scope": _DEVICE_FLOW_SCOPE},
)
except httpx.HTTPStatusError as e:
raise ServiceUnavailableError(
@ -435,16 +471,14 @@ async def astart_device_flow() -> GithubCopilotDeviceFlowStart:
)
async def apoll_device_flow(device_code: str) -> GithubCopilotDeviceFlowPoll:
async def apoll_device_flow(device_code: str, endpoints: GithubHostEndpoints) -> GithubCopilotDeviceFlowPoll:
client: Final = _new_async_github_client()
access_token_url: Final = os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL)
client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
try:
response: Final = await client.post(
access_token_url,
endpoints.access_token_url,
headers=github_api_headers(),
json={
"client_id": client_id,
"client_id": endpoints.client_id,
"device_code": device_code,
"grant_type": "urn:ietf:params:oauth:grant-type:device_code",
},
@ -500,11 +534,11 @@ def _poll_payload_from_response(response: httpx.Response) -> GithubCopilotDevice
)
async def afetch_github_login(github_token: str) -> str:
async def afetch_github_login(github_token: str, endpoints: GithubHostEndpoints) -> str:
client: Final = _new_async_github_client()
try:
response: Final = await client.get(
"https://api.github.com/user",
endpoints.user_url,
headers=github_api_headers(github_token),
)
except Exception as e:
@ -530,11 +564,17 @@ async def afetch_github_login(github_token: str) -> str:
return payload["login"]
async def acheck_copilot_seat(user_id: str, github_token: str, credential_name: str) -> GithubCopilotUserSession:
async def acheck_copilot_seat(
user_id: str,
github_token: str,
credential_name: str,
endpoints: GithubHostEndpoints,
) -> GithubCopilotUserSession:
return await aexchange_github_token(
user_id=user_id,
github_token=github_token,
credential_name=credential_name,
endpoints=endpoints,
)
@ -567,7 +607,11 @@ def _connect_error(credential_name: str) -> CallerCredentialAuthenticationError:
)
def _caller_connection(kwargs: Mapping[str, object], credential_name: str) -> tuple[str, str]:
def _caller_connection(
kwargs: Mapping[str, object],
credential_name: str,
endpoints: GithubHostEndpoints,
) -> tuple[str, str]:
secret_fields_raw: Final[object] = kwargs.get("secret_fields")
secret_fields: Final = (
cast(Mapping[object, object], secret_fields_raw) # cast-ok: value is Mapping-checked or dict-shaped at runtime
@ -585,8 +629,21 @@ def _caller_connection(kwargs: Mapping[str, object], credential_name: str) -> tu
if isinstance(credentials_raw, Mapping)
else None
)
token: Final[object] = credentials.get(credential_name) if credentials is not None else None
if not isinstance(user_id, str) or not user_id or not isinstance(token, str) or not token:
connection_raw: Final[object] = credentials.get(credential_name) if credentials is not None else None
connection: Final = (
cast(Mapping[object, object], connection_raw) # cast-ok: runtime Mapping check validates the stored connection
if isinstance(connection_raw, Mapping)
else None
)
token: Final[object] = connection.get("access_token") if connection is not None else None
github_host: Final = normalized_github_host(connection.get("github_host") if connection is not None else None)
if (
not isinstance(user_id, str)
or not user_id
or not isinstance(token, str)
or not token
or github_host != endpoints.host
):
raise _connect_error(credential_name)
return user_id, token
@ -705,11 +762,13 @@ def attach_github_copilot_user_session(
if isinstance(kwargs.get(GITHUB_COPILOT_USER_SESSION_KWARG_KEY), GithubCopilotUserSession):
return
credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else ""
user_id, github_token = _caller_connection(kwargs, credential_name)
endpoints: Final = github_copilot_host_endpoints(credential_name)
user_id, github_token = _caller_connection(kwargs, credential_name, endpoints)
session: Final = exchange_github_token(
user_id=user_id,
github_token=github_token,
credential_name=credential_name,
endpoints=endpoints,
)
kwargs[GITHUB_COPILOT_USER_SESSION_KWARG_KEY] = session # rebind-ok: kwargs is the request mutation channel
_strip_caller_authorization(kwargs)
@ -725,11 +784,13 @@ async def aattach_github_copilot_user_session(
if isinstance(kwargs.get(GITHUB_COPILOT_USER_SESSION_KWARG_KEY), GithubCopilotUserSession):
return
credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else ""
user_id, github_token = _caller_connection(kwargs, credential_name)
endpoints: Final = github_copilot_host_endpoints(credential_name)
user_id, github_token = _caller_connection(kwargs, credential_name, endpoints)
session: Final = await aexchange_github_token(
user_id=user_id,
github_token=github_token,
credential_name=credential_name,
endpoints=endpoints,
)
kwargs[GITHUB_COPILOT_USER_SESSION_KWARG_KEY] = session # rebind-ok: kwargs is the request mutation channel
_strip_caller_authorization(kwargs)

154
litellm/llms/nvidia_riva/audio_transcription/handler.py
View file

@ -16,7 +16,9 @@ The handler is intentionally a thin orchestration layer:
3. Construct ``riva.client.Auth`` honoring NVCF (function-id metadata + TLS)
vs self-hosted (any host:port, optional TLS) modes.
4. Stream the audio through Riva's ``streaming_response_generator`` and
aggregate ``is_final`` results into a single transcript.
aggregate ``is_final`` results into a single transcript, or, when
``riva_offline`` is set, send the whole buffer through the unary
``offline_recognize`` RPC for NIMs that only ship offline engines.
5. Return a normalized ``TranscriptionResponse`` with ``duration`` exposed
on ``_hidden_params`` so cost calculation works.
@ -26,13 +28,19 @@ without the optional STT extras installed.
import asyncio
import inspect
from collections.abc import Callable, Iterable
from collections.abc import Callable, Iterable, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol
import httpx
import litellm
from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS
from litellm.litellm_core_utils.audio_utils.utils import (
get_audio_file_name,
process_audio_file,
)
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
from litellm.llms.nvidia_riva.audio_transcription.audio_utils import (
resample_to_riva_pcm,
)
@ -45,6 +53,7 @@ from litellm.llms.nvidia_riva.common_utils import (
NvidiaRivaException,
grpc_error_to_litellm_exception,
)
from litellm.secret_managers.main import str_to_bool
from litellm.types.utils import FileTypes, TranscriptionResponse
from litellm.utils import convert_to_model_response_object
@ -67,10 +76,20 @@ class _RivaAuth(Protocol):
"""Opaque ``riva.client.Auth`` handle."""
class _RecognizeFuture(Protocol):
def result(self, timeout: float | None = None) -> object: ...
def done(self) -> bool: ...
def cancel(self) -> bool: ...
class _AsrService(Protocol):
@property
def streaming_response_generator(self) -> Callable[..., Iterable[object]]: ...
def offline_recognize(self, audio_bytes: bytes, config: object, future: bool) -> _RecognizeFuture: ...
class _EndpointingConfig(Protocol):
"""Opaque ``EndpointingConfig`` protobuf message."""
@ -128,6 +147,38 @@ def _audio_encoding(riva_asr_module: _RivaAsrModule) -> _AudioEncoding:
return riva_asr_module.AudioEncoding
def _parse_riva_offline(value: object, drop_params: bool) -> bool:
match value:
case None:
return False
case bool():
return value
case str() if (parsed := str_to_bool(value)) is not None:
return parsed
case _ if drop_params:
return False
case _:
raise NvidiaRivaException(
status_code=400,
message=f"riva_offline must be true or false, got {value!r}. Set drop_params to ignore it.",
)
def _deadline_seconds(timeout: float | str | httpx.Timeout, drop_params: bool) -> float:
try:
limit: Final = CompletionTimeout.as_limit(timeout)
except ValueError:
if drop_params:
return COMPLETION_HTTP_FALLBACK_SECONDS
raise NvidiaRivaException(
status_code=400,
message=f"timeout must be a number of seconds, got {timeout!r}. Set drop_params to ignore it.",
) from None
if not isinstance(limit, httpx.Timeout):
return limit
return COMPLETION_HTTP_FALLBACK_SECONDS if limit.read is None else limit.read
class NvidiaRivaAudioTranscription:
"""Sync + async entry point for Riva ASR."""
@ -138,7 +189,7 @@ class NvidiaRivaAudioTranscription:
optional_params: dict,
litellm_params: dict,
model_response: TranscriptionResponse,
timeout: float,
timeout: float | str | httpx.Timeout | None,
logging_obj: "LiteLLMLoggingObj",
api_key: str | None,
api_base: str | None,
@ -183,7 +234,7 @@ class NvidiaRivaAudioTranscription:
optional_params: dict,
litellm_params: dict,
model_response: TranscriptionResponse,
timeout: float,
timeout: float | str | httpx.Timeout | None,
logging_obj: "LiteLLMLoggingObj",
api_key: str | None,
api_base: str | None,
@ -213,7 +264,7 @@ class NvidiaRivaAudioTranscription:
optional_params: dict,
litellm_params: dict,
model_response: TranscriptionResponse,
timeout: float,
timeout: float | str | httpx.Timeout | None,
logging_obj: "LiteLLMLoggingObj",
api_key: str | None,
api_base: str | None,
@ -259,6 +310,10 @@ class NvidiaRivaAudioTranscription:
response_format: Final = request_payload.get("response_format") or "json"
timestamp_granularities: Final = request_payload.get("timestamp_granularities")
drop_params: Final = (
"drop_params" in litellm_params and litellm_params["drop_params"] is True
) or litellm.drop_params is True
riva_offline: Final = _parse_riva_offline(optional_params.get("riva_offline"), drop_params=drop_params)
riva_module, riva_asr_module = _import_riva()
auth_obj: Final = self._construct_auth(
@ -272,10 +327,6 @@ class NvidiaRivaAudioTranscription:
riva_asr_module=riva_asr_module,
recognition_config_dict=recognition_config_dict,
)
streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig(
config=recognition_config, interim_results=False
)
logging_obj.pre_call(
input=None,
api_key=api_key,
@ -286,24 +337,31 @@ class NvidiaRivaAudioTranscription:
"recognition_config": recognition_config_dict,
"nvcf_function_id_set": bool(optional_params.get("nvcf_function_id")),
"use_ssl": optional_params.get("use_ssl"),
"riva_offline": riva_offline,
},
},
)
try:
asr_service: Final[_AsrService] = riva_module.ASRService(auth_obj)
audio_chunks: Final = self._iter_audio_chunks(resampled.pcm_bytes)
stream_kwargs: Final[dict[str, object]] = {
"audio_chunks": audio_chunks,
"streaming_config": streaming_config,
}
# Forward the deadline so the stream cannot block forever if the
# server stalls. Older riva-client versions do not accept a
# ``timeout`` kwarg, so pass it only when supported.
if timeout is not None and self._supports_timeout_kwarg(asr_service.streaming_response_generator):
stream_kwargs["timeout"] = float(timeout)
stream: Final = asr_service.streaming_response_generator(**stream_kwargs)
final_results: Final = self._collect_final_results(stream)
responses: Final = (
self._offline_responses(
asr_service=asr_service,
pcm_bytes=resampled.pcm_bytes,
recognition_config=recognition_config,
timeout=None if timeout is None else _deadline_seconds(timeout, drop_params),
)
if riva_offline
else self._streaming_responses(
asr_service=asr_service,
riva_asr_module=riva_asr_module,
pcm_bytes=resampled.pcm_bytes,
recognition_config=recognition_config,
timeout=timeout,
drop_params=drop_params,
)
)
final_results: Final = self._collect_final_results(responses, require_final=not riva_offline)
except NvidiaRivaException:
raise
except Exception as e:
@ -409,6 +467,54 @@ class NvidiaRivaAudioTranscription:
return config
def _streaming_responses(
self,
asr_service: _AsrService,
riva_asr_module: _RivaAsrModule,
pcm_bytes: bytes,
recognition_config: _RecognitionConfig,
timeout: float | str | httpx.Timeout | None,
drop_params: bool,
) -> Iterable[object]:
streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig(
config=recognition_config, interim_results=False
)
return asr_service.streaming_response_generator(
audio_chunks=self._iter_audio_chunks(pcm_bytes),
streaming_config=streaming_config,
**self._timeout_kwargs(asr_service.streaming_response_generator, timeout, drop_params),
)
def _offline_responses(
self,
asr_service: _AsrService,
pcm_bytes: bytes,
recognition_config: _RecognitionConfig,
timeout: float | None,
) -> Iterable[object]:
pending: Final = asr_service.offline_recognize(audio_bytes=pcm_bytes, config=recognition_config, future=True)
try:
return (pending.result(timeout=timeout),)
except Exception as e:
if pending.done():
raise
pending.cancel()
raise NvidiaRivaException(
status_code=408,
message=f"NVIDIA Riva offline recognition did not finish within {timeout} seconds",
) from e
@classmethod
def _timeout_kwargs(
cls, rpc: Callable[..., object], timeout: float | str | httpx.Timeout | None, drop_params: bool
) -> Mapping[str, float]:
# Forward the deadline so the call cannot block forever if the
# server stalls. Older riva-client versions do not accept a
# ``timeout`` kwarg, so pass it only when supported.
if timeout is None or not cls._supports_timeout_kwarg(rpc):
return MappingProxyType({})
return MappingProxyType({"timeout": _deadline_seconds(timeout, drop_params)})
@staticmethod
def _supports_timeout_kwarg(callable_obj: Callable[..., object]) -> bool:
try:
@ -429,9 +535,9 @@ class NvidiaRivaAudioTranscription:
yield chunk
@staticmethod
def _collect_final_results(stream) -> list[dict[str, object]]:
def _collect_final_results(stream: Iterable[object], require_final: bool = True) -> list[dict[str, object]]:
"""
Walk the gRPC stream, ignore empty / non-final chunks, and return a
Walk the gRPC responses, ignore empty / non-final chunks, and return a
list of normalized final-result dicts. Matching the user's note: the
``id`` blocks with no ``results`` are streaming heartbeats and must
be skipped.
@ -440,7 +546,7 @@ class NvidiaRivaAudioTranscription:
for response in stream:
results = getattr(response, "results", None) or []
for result in results:
if not getattr(result, "is_final", False):
if require_final and not getattr(result, "is_final", False):
continue
alternatives = getattr(result, "alternatives", None) or []
if not alternatives:

15
litellm/llms/nvidia_riva/common_utils.py
View file

@ -41,6 +41,9 @@ _GRPC_STATUS_CODE_TO_HTTP: Final[dict] = {
}
_GRPC_MESSAGE_TOO_LARGE: Final = "message larger than max"
def _extract_grpc_status_name(error: object) -> str | None:
"""
Best-effort extraction of a gRPC StatusCode name from an arbitrary error.
@ -79,8 +82,16 @@ def grpc_error_to_litellm_exception(error: Exception) -> NvidiaRivaException:
a ``NvidiaRivaException`` with an appropriate HTTP-equivalent status code.
"""
status_name: Final = _extract_grpc_status_name(error)
http_status: Final = _GRPC_STATUS_CODE_TO_HTTP.get(status_name or "", 500)
detail: Final = _extract_grpc_details(error) or str(error)
if status_name == "RESOURCE_EXHAUSTED" and _GRPC_MESSAGE_TOO_LARGE in detail:
return NvidiaRivaException(
status_code=413,
message=(
f"NVIDIA Riva gRPC error ({status_name}): {detail}. The audio is too large for one offline "
"request to this Riva server: send shorter audio, or drop riva_offline on models that support streaming."
),
)
http_status: Final = _GRPC_STATUS_CODE_TO_HTTP.get(status_name or "", 500)
message = f"NVIDIA Riva gRPC error ({status_name}): {detail}" if status_name else f"NVIDIA Riva error: {detail}"
return NvidiaRivaException(status_code=http_status, message=message)

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

@ -1107,14 +1107,18 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
"""
wrapper to call VertexGeminiConfig.map_openai_params
"""
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexGeminiConfig,
)
config: Final = VertexGeminiConfig()
_model: Final = openai_request_body.get("model", "")
raw_model: Final = openai_request_body.get("model", "")
model: Final = raw_model if isinstance(raw_model, str) else ""
sampling_params: Final = tuple(
p for p in ("temperature", "top_p", "top_k") if openai_request_body.get(p) is not None
)
if sampling_params and not VertexGeminiConfig.supports_sampling_params(model):
VertexGeminiConfig.reject_unsupported_sampling_param(
model=model, param=", ".join(sampling_params), drop_params=False
)
vertex_params: Final = config.map_openai_params(
model=_model,
model=model,
non_default_params=openai_request_body,
optional_params={},
drop_params=False,

51
litellm/llms/vertex_ai/gemini/long_decode.py
View file

@ -2,17 +2,20 @@ import json
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from itertools import accumulate, chain
from typing import Final, TypeAlias, TypeVar, cast # noqa: TID251 # sum_usage_metadata casts computed keys
from typing import Final, Literal, TypeAlias, TypeVar, cast # noqa: TID251 # sum_usage_metadata casts computed keys
import httpx
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.exceptions import Timeout
from litellm.litellm_core_utils.env_utils import get_env_int_or_none
from litellm.types.llms.vertex_ai import (
Candidates,
Citation,
CitationMetadata,
GenerateContentResponseBody,
GenerationConfig,
GroundingMetadata,
GroundingSegment,
GroundingSupport,
@ -48,8 +51,14 @@ HOP_FAILURES: Final = (VertexAIError, httpx.HTTPError, Timeout)
HOP_END_FINISH_REASON: Final = "CONTINUATION"
CUT_OFF_FINISH_REASONS: Final = frozenset({"MAX_TOKENS", HOP_END_FINISH_REASON})
DEFAULT_OUTPUT_BUDGET_ENV: Final = "LITELLM_GEMINI_DEFAULT_MAX_OUTPUT_TOKENS"
_OUTPUT_LIMIT_KEYS: Final = frozenset({"max_output_tokens", "maxOutputTokens"})
_REQUEST_KEYS_THAT_SKIP_THE_DEFAULT_BUDGET: Final = frozenset({"generation_config", "continuationToken"})
_TEXT_ONLY_MODALITIES: Final = TypeAdapter(list[Literal["TEXT"]])
_VERTEX_V1_PATH: Final = "/v1/projects/"
_VERTEX_PATH_ACCEPTING_CONTINUATION_TOKEN: Final = "/v1beta1/projects/"
_CANDIDATE_COUNT_KEYS: Final = ("candidate_count", "candidateCount")
LongDecodeHops: TypeAlias = tuple[GenerateContentResponseBody, ...]
AsyncHopPost: TypeAlias = Callable[[RequestBody], Awaitable[httpx.Response]]
@ -76,15 +85,53 @@ def _int_or_none(value: object) -> int | None:
return value if isinstance(value, int) else None
def _asks_for_one_candidate(generation_config: Mapping[str, object]) -> bool:
return all(
_int_or_none(generation_config[key] or 1) == 1 for key in _CANDIDATE_COUNT_KEYS if key in generation_config
)
def continuation_cap_of(request_body: RequestBody) -> int | None:
generation_config: Final = request_body.get("generationConfig")
if generation_config is None or "max_output_tokens" not in generation_config:
return None
if "candidate_count" in generation_config and _int_or_none(generation_config["candidate_count"]) != 1:
if not _asks_for_one_candidate(generation_config):
return None
return _int_or_none(generation_config["max_output_tokens"])
def _asks_for_text_only(generation_config: Mapping[str, object]) -> bool:
if "response_modalities" in generation_config:
return False
try:
_ = _TEXT_ONLY_MODALITIES.validate_python(generation_config.get("responseModalities", []))
except ValidationError:
return False
return True
def with_default_output_budget(
request_body: RequestBody, model_max_output_tokens: Callable[[], int | None]
) -> RequestBody:
generation_config: Final = request_body.get("generationConfig", GenerationConfig())
request_json: Final[Mapping[str, object]] = request_body
budget: Final = get_env_int_or_none(DEFAULT_OUTPUT_BUDGET_ENV)
if budget is None or budget <= 0:
return request_body
if not _REQUEST_KEYS_THAT_SKIP_THE_DEFAULT_BUDGET.isdisjoint(request_body):
return request_body
if not isinstance(request_json.get("generationConfig", {}), dict):
return request_body
if not _OUTPUT_LIMIT_KEYS.isdisjoint(generation_config):
return request_body
if not _asks_for_text_only(generation_config):
return request_body
model_max: Final = model_max_output_tokens()
if model_max is None:
return request_body
return {**request_body, "generationConfig": {**generation_config, "max_output_tokens": min(budget, model_max)}}
def client_continuation_token(candidate: Candidates) -> str | None:
if candidate.get("finishReason") not in CUT_OFF_FINISH_REASONS:
return None

130
litellm/llms/vertex_ai/gemini/transformation.py
View file

@ -7,15 +7,19 @@ Why separate file? Make it easy to see how transformation works
import os
import re
from collections.abc import Callable, Mapping
from functools import partial
from typing import TYPE_CHECKING, Any, Final, Literal, cast
from urllib.parse import quote
import httpx
from pydantic import BaseModel, TypeAdapter
from pydantic import BaseModel, PositiveInt, TypeAdapter, ValidationError
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_image_mime_type_from_url,
)
@ -64,6 +68,7 @@ from ..common_utils import (
get_supports_response_schema,
get_supports_system_message,
)
from .long_decode import with_default_output_budget
# VertexBase is imported only for type checking to avoid introducing a
# module-load-time cyclic import to vertex_llm_base. The instance is lazily
@ -868,7 +873,7 @@ def gemini_convert_messages_with_history(
"""
has_text_in_content = check_text_in_content(user_content)
if has_text_in_content is False:
verbose_logger.warning(
verbose_logger.debug(
"No text in user content. Adding a blank text to user content, to ensure Gemini doesn't fail the request. Relevant Issue - https://github.com/BerriAI/litellm/issues/5515"
)
user_content.append(
@ -1189,7 +1194,18 @@ def transform_request_body(
# labels: optional explicit param and/or metadata.requester_metadata (OpenAI metadata)
labels: Final = pop_vertex_request_labels(optional_params, litellm_params)
filtered_params = {k: v for k, v in optional_params.items() if _get_equivalent_key(k, set(config_fields))}
from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig
drops_top_k: Final = "top_k" in optional_params and not VertexGeminiConfig.supports_sampling_params(model)
if drops_top_k:
request_params: Final = cast(Mapping[str, object], litellm_params) # cast-ok: untyped shared signature
VertexGeminiConfig.reject_unsupported_sampling_param(
model=model,
param="top_k",
drop_params=drops_unsupported_params(request_params.get("drop_params")),
)
config_keys: Final = set(config_fields) - {"top_k"} if drops_top_k else set(config_fields)
filtered_params = {k: v for k, v in optional_params.items() if _get_equivalent_key(k, config_keys)}
generation_config: Final[GenerationConfig | None] = GenerationConfig(**filtered_params)
@ -1292,13 +1308,75 @@ def sync_transform_request_body(
vertex_auth_header=vertex_auth_header,
)
return transform_request_body(
messages=messages,
model=model,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
cached_content=cached_content,
optional_params=optional_params,
return _with_default_output_budget(
transform_request_body(
messages=messages,
model=model,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
cached_content=cached_content,
optional_params=optional_params,
),
model,
custom_llm_provider,
litellm_params.get("model_info"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # litellm_params is an untyped dict
)
class _CostMapOutputModalities(TypedDict):
supported_output_modalities: NotRequired[ReadOnly[tuple[str, ...]]]
_COST_MAP_OUTPUT_MODALITIES: Final = TypeAdapter(_CostMapOutputModalities)
def _answers_text(cost_map_key: str) -> bool:
try:
row: Final = _COST_MAP_OUTPUT_MODALITIES.validate_python(litellm.model_cost.get(cost_map_key)) # pyright: ignore[reportUnknownMemberType] # cost map rows are untyped JSON
except ValidationError:
return False
return "text" in row.get("supported_output_modalities", ("text",))
class _OutputLimit(TypedDict):
max_output_tokens: NotRequired[ReadOnly[PositiveInt | None]]
_OUTPUT_LIMIT: Final = TypeAdapter(_OutputLimit)
def _output_limit(fields: object) -> int | None:
try:
return _OUTPUT_LIMIT.validate_python(fields).get("max_output_tokens")
except ValidationError:
return None
def _catalog_capped_output_limit(cost_map_key: str, override: object) -> int | None:
limits: Final = (
_output_limit(GetModelCostMap.loaded_model_cost_map().get(cost_map_key)),
_output_limit(override),
)
return min((limit for limit in limits if limit is not None), default=None)
def _text_chat_model_max_output_tokens(
model: str, custom_llm_provider: str, deployment_model_info: object
) -> int | None:
info: Final = litellm.utils.lookup_model_info_for_registration(model, custom_llm_provider)
if isinstance(info, litellm.utils.ModelInfoLookupMiss) or info.get("mode") != "chat":
return None
if not _answers_text(info["key"]):
return None
return _catalog_capped_output_limit(info["key"], deployment_model_info or info)
def _with_default_output_budget(
request_body: RequestBody, model: str, custom_llm_provider: str, deployment_model_info: object
) -> RequestBody:
return with_default_output_budget(
request_body,
partial(_text_chat_model_max_output_tokens, model, custom_llm_provider, deployment_model_info),
)
@ -1368,21 +1446,8 @@ async def async_transform_request_body(
messages, should_inline=should_inline_remote_media(custom_llm_provider)
)
if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages):
# _transform_request_body may issue a sync httpx.get (up to 5s timeout)
# via _get_gcs_object_content_type to fetch GCS object metadata. Run the
# whole sync transformation on a worker thread so it does not block the
# async event loop.
return await asyncify(transform_request_body)(
messages=inlined_messages,
model=model,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
cached_content=cached_content,
optional_params=optional_params,
)
return transform_request_body(
transform: Final = partial(
transform_request_body,
messages=inlined_messages,
model=model,
custom_llm_provider=custom_llm_provider,
@ -1390,6 +1455,21 @@ async def async_transform_request_body(
cached_content=cached_content,
optional_params=optional_params,
)
# _transform_request_body may issue a sync httpx.get (up to 5s timeout)
# via _get_gcs_object_content_type to fetch GCS object metadata. Run the
# whole sync transformation on a worker thread so it does not block the
# async event loop.
request_body: Final[RequestBody] = (
await asyncify(transform)()
if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages)
else transform()
)
return _with_default_output_budget(
request_body,
model,
custom_llm_provider,
litellm_params.get("model_info"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # litellm_params is an untyped dict
)
def _default_user_message_when_system_message_passed() -> ChatCompletionUserMessage:

92
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -3,10 +3,11 @@
## Initial implementation - covers gemini + image gen calls
import json
import re
import threading
import time
from collections.abc import Awaitable, Callable, Mapping, Sequence
from copy import deepcopy
from functools import partial
from functools import lru_cache, partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args
@ -28,6 +29,7 @@ from litellm.constants import (
)
from litellm.exceptions import UnsupportedParamsError
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator
from litellm.litellm_core_utils.prompt_templates.factory import (
encode_tool_call_id_with_signature,
@ -196,6 +198,29 @@ def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsErr
)
_SAMPLING_DEPRECATION_UNMAPPED_MODELS: Final = "models outside litellm's model map"
_sampling_deprecation_lock: Final = threading.Lock()
@lru_cache(maxsize=None)
def _log_gemini_sampling_params_deprecated(model: str) -> None:
verbose_logger.warning(
"DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to function for Gemini 3+ (%s) but are planned for removal in a future release. Move sampling guidance into the `system` instructions instead.",
model,
)
def _sampling_deprecation_notice_key(model: str) -> str:
catalog: Final = GetModelCostMap.loaded_model_cost_map()
is_cataloged: Final = any(name in catalog for name in (model, f"vertex_ai/{model}", f"gemini/{model}"))
return model if is_cataloged else _SAMPLING_DEPRECATION_UNMAPPED_MODELS
def _warn_gemini_sampling_params_deprecated(model: str) -> None:
with _sampling_deprecation_lock:
_log_gemini_sampling_params_deprecated(_sampling_deprecation_notice_key(model))
def _cache_creation_tokens(logging_obj: LoggingClass) -> int:
recorded: Final[object] = getattr(logging_obj, "vertex_ai_cache_creation_input_tokens", 0)
return recorded if isinstance(recorded, int) and recorded > 0 else 0
@ -404,11 +429,43 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
return False
return True
@staticmethod
def supports_sampling_params(model: str) -> bool:
"""Whether Google still honors ``temperature`` / ``top_p`` / ``top_k`` for ``model``.
Driven by ``supports_sampling_params`` in the model map: Gemini 3.6 Flash, 3.5 Flash-Lite
and later accept the fields but ignore them, and a later generation will 400 on them.
A model the map says nothing about keeps forwarding them. ``temperature`` and ``top_p``
leave ``get_supported_openai_params`` when the flag is off, so the shared unsupported-param
check raises or drops them and ``allowed_openai_params`` still forwards them verbatim;
``top_k`` is a provider-specific kwarg that bypasses that check and is gated while the
request body is built."""
rows: Final = cast(Mapping[str, Mapping[str, object]], litellm.model_cost) # cast-ok: JSON rows by model id
bare: Final = model.split("/")[-1]
for candidate in (model, bare, f"vertex_ai/{bare}", f"gemini/{bare}"):
if isinstance(flag := rows.get(candidate, {}).get("supports_sampling_params"), bool):
return flag
return True
@staticmethod
def reject_unsupported_sampling_param(model: str, param: str, drop_params: bool) -> None:
"""Return under ``drop_params`` so the caller drops ``param``, else raise the client-side 400
that names the escape hatch."""
if litellm.drop_params or drop_params:
return
raise UnsupportedParamsError(
message=(
f"{model} does not support {param}: the model map flags it `supports_sampling_params: false` "
"because Google ignores sampling params on this model. "
"To drop unsupported params, set `litellm.drop_params = True`."
),
status_code=400,
)
def get_supported_openai_params(self, model: str) -> list[str]:
sampling_params: Final = ("temperature", "top_p", "top_k") if self.supports_sampling_params(model) else ()
supported_params: Final = [
"temperature",
"top_p",
"top_k",
*sampling_params,
"max_tokens",
"max_completion_tokens",
"stream",
@ -1266,8 +1323,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
allow_integer_enums: bool = True,
) -> dict:
self._apply_include_server_side_tool_invocations(non_default_params, optional_params)
gemini_sampling_params_warned: bool = False
supports_sampling_params: Final = self.supports_sampling_params(model)
for param, value in non_default_params.items():
if param in ("temperature", "top_p", "top_k") and not supports_sampling_params:
continue
if param == "temperature":
if VertexGeminiConfig._is_gemini_3_or_newer(model):
if value is not None and value < 1.0:
@ -1275,28 +1334,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"Warning: Setting temperature < 1.0 for Gemini 3 models (%s) can cause infinite loops, degraded reasoning performance, and failure on complex tasks. Strongly recommended to use temperature = 1.0 (default).",
model,
)
if not gemini_sampling_params_warned:
verbose_logger.warning(
"DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to function for Gemini 3+ (%s) but are planned for removal in a future release. Move sampling guidance into the `system` instructions instead.",
model,
)
gemini_sampling_params_warned = True
_warn_gemini_sampling_params_deprecated(model)
optional_params["temperature"] = value
elif param == "top_p":
if VertexGeminiConfig._is_gemini_3_or_newer(model) and not gemini_sampling_params_warned:
verbose_logger.warning(
"DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to function for Gemini 3+ (%s) but are planned for removal in a future release. Move sampling guidance into the `system` instructions instead.",
model,
)
gemini_sampling_params_warned = True
if VertexGeminiConfig._is_gemini_3_or_newer(model):
_warn_gemini_sampling_params_deprecated(model)
optional_params["top_p"] = value
elif param == "top_k":
if VertexGeminiConfig._is_gemini_3_or_newer(model) and not gemini_sampling_params_warned:
verbose_logger.warning(
"DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to function for Gemini 3+ (%s) but are planned for removal in a future release. Move sampling guidance into the `system` instructions instead.",
model,
)
gemini_sampling_params_warned = True
if VertexGeminiConfig._is_gemini_3_or_newer(model):
_warn_gemini_sampling_params_deprecated(model)
optional_params["top_k"] = value
elif (
param == "stream" and value is True

1
litellm/main.py
View file

@ -5988,6 +5988,7 @@ def completion(
*OPENAI_WIF_KWARGS_KEYS,
*OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS,
"github_copilot_auth_type",
"github_copilot_host",
"github_copilot_user_session",
PROVIDER_AFFINITY_HEADER_KWARG_KEY,
"fireworks_forward_user_id",

91
litellm/messages/dispatch.py
View file

@ -1,18 +1,21 @@
import inspect
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping
from types import MappingProxyType
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.llms.anthropic.pass_through.messages import handler as main
from litellm.llms.anthropic.pass_through.safeguards import native_route_classifies
from litellm.rust_bridge.catalog import Route, RouteContext
from litellm.rust_bridge import detectors
from litellm.rust_bridge.catalog import Feature, Route, RouteContext, Target
from litellm.rust_bridge.dispatch import PublicDispatch
from litellm.rust_bridge.messages import detectors as route_detectors
from litellm.rust_bridge.messages.entrypoints import (
NATIVE_AMESSAGES,
NATIVE_MESSAGES,
)
from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
bind,
native_call,
@ -46,6 +49,7 @@ def _python_amessages() -> PythonAmessages:
_PYTHON_MESSAGES: Final = _python_messages()
_MESSAGES: Final = signature(_PYTHON_MESSAGES)
_FLAG: Final = AsyncFlag("is_async", sync=_MESSAGES, is_set=bool)
_PYTHON_AMESSAGES: Final = _python_amessages()
_AMESSAGES: Final = signature(_PYTHON_AMESSAGES)
@ -71,45 +75,76 @@ def _public_request(
return native_call(legacy, args, _native_kwargs(fields, kwargs))
def _resolved_provider(request: NativeCall) -> str | None:
try:
return get_llm_provider(
str(request.resolved["model"]), optional_str(request.resolved.get("custom_llm_provider"))
)[1]
except BadRequestError:
return optional_str(request.resolved.get("custom_llm_provider"))
def _python_serves(request: NativeCall) -> bool:
"""Vertex AI serves only Claude through the Anthropic Messages config, as
``get_provider_anthropic_messages_config`` does; every other Vertex model keeps Python's
translation through Chat Completions. A request Python would answer with classifier
``safeguard_results`` stays on Python too."""
def _resolved(request: NativeCall) -> tuple[str, str | None]:
model: Final = str(request.resolved["model"])
vertex_without_claude: Final = _resolved_provider(request) == "vertex_ai" and "claude" not in model.lower()
return vertex_without_claude or native_route_classifies(model, request.kwargs.get("safeguards"))
provider: Final = optional_str(request.resolved.get("custom_llm_provider"))
try:
resolved_model, resolved_provider, _, _ = get_llm_provider(model, provider)
except BadRequestError:
return model, provider
return resolved_model, resolved_provider
def _context(request: NativeCall) -> RouteContext:
def _target(request: NativeCall, model: str, provider: str | None) -> Target:
requested_model: Final = str(request.resolved["model"])
model_info: Final = request.resolved.get("model_info")
api_base: Final = request.resolved.get("api_base")
native_config: Final = main._native_messages_config( # pyright: ignore[reportPrivateUsage] # the handler's own passthrough lookup
model, provider, model_info, api_base if isinstance(api_base, str) else None
)
if native_config is not None:
return Target.MESSAGES
compacting: Final = request.resolved.get("compaction") is not None
if not compacting and main._should_route_to_responses_api(provider, requested_model, model, model_info): # pyright: ignore[reportPrivateUsage] # the handler's own translation choice
return Target.RESPONSES
return Target.CHAT_COMPLETIONS
_FEATURES: Final[detectors.Detectors] = MappingProxyType(
{
**detectors.COMMON,
Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier,
Feature.RESPONSE_CACHE: detectors.response_cache,
Feature.MODIFY_PARAMS: detectors.modify_params,
Feature.WORKLOAD_IDENTITY: route_detectors.anthropic_workload_identity,
Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection,
Feature.STREAM_COST_INJECTION: route_detectors.stream_cost_injection,
Feature.THINKING_SIGNATURE_RETRY: route_detectors.thinking_signature_retry,
}
)
_ASYNC_FEATURES: Final[detectors.Detectors] = MappingProxyType(
{
**_FEATURES,
Feature.PRE_REQUEST_HOOK: route_detectors.pre_request_hook,
Feature.WEBSEARCH_INTERCEPTION: route_detectors.websearch_interception,
Feature.MESSAGES_INTERCEPTOR: route_detectors.intercepted,
}
)
def _context(request: NativeCall, features: detectors.Detectors) -> RouteContext:
requested_model: Final = str(request.resolved["model"])
model, provider = _resolved(request)
return RouteContext(
Route.MESSAGES,
provider=_resolved_provider(request),
model=str(request.resolved["model"]),
provider=provider,
model=requested_model,
target=_target(request, model, provider),
features=detectors.detected(features, request),
)
_DISPATCH: Final = PublicDispatch(
route=Route.MESSAGES,
request=lambda args, kwargs: _public_request(_MESSAGES, args, kwargs),
context=_context,
bypass=lambda request: request.kwargs.get("is_async") is True or _python_serves(request),
context=lambda request: _context(request, _FEATURES),
)
_ADISPATCH: Final = PublicDispatch(
route=Route.MESSAGES,
request=lambda args, kwargs: _public_request(_AMESSAGES, args, kwargs),
context=_context,
bypass=_python_serves,
context=lambda request: _context(request, _ASYNC_FEATURES),
)
@ -118,9 +153,11 @@ def anthropic_messages_handler(
**kwargs: object, # kwargs-ok: preserve the public Anthropic Messages call shape
) -> MessagesResult | Coroutine[object, object, MessagesResult]:
python: Final = _PYTHON_MESSAGES
if _FLAG.requested(kwargs):
return anthropic_messages(**_FLAG.keywords(args, kwargs))
return _DISPATCH.run(
args,
kwargs,
_FLAG.dropped(kwargs),
python=python,
binding=NATIVE_MESSAGES,
native=native_call_hook,
@ -131,7 +168,7 @@ async def anthropic_messages(*args: object, **kwargs: object) -> MessagesResult:
python: Final = _PYTHON_AMESSAGES
return await _ADISPATCH.arun(
args,
kwargs,
_FLAG.dropped(kwargs),
python=python,
binding=NATIVE_AMESSAGES,
native=native_call_hook,

101
litellm/model_prices_and_context_window_backup.json
View file

@ -16580,6 +16580,24 @@
"source": "https://developers.cloudflare.com/workers-ai/platform/pricing/",
"supports_vision": true
},
"cloudflare/clairic": {
"input_cost_per_token": 1e-07,
"litellm_provider": "cloudflare",
"max_input_tokens": 65536,
"mode": "extraction",
"output_cost_per_token": 0.0,
"source": "https://developers.cloudflare.com/workers-ai/models/clairic/",
"supports_vision": true
},
"cloudflare/@cf/cloudflare/clairic": {
"input_cost_per_token": 1e-07,
"litellm_provider": "cloudflare",
"max_input_tokens": 65536,
"mode": "extraction",
"output_cost_per_token": 0.0,
"source": "https://developers.cloudflare.com/workers-ai/models/clairic/",
"supports_vision": true
},
"cloudflare/@cf/meta/llama-2-7b-chat-fp16": {
"input_cost_per_token": 1.923e-06,
"litellm_provider": "cloudflare",
@ -21903,6 +21921,7 @@
],
"supports_function_calling": true,
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
},
@ -21927,6 +21946,7 @@
],
"supports_function_calling": true,
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
},
@ -21958,6 +21978,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
},
@ -22021,6 +22042,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
},
@ -29417,6 +29439,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
@ -30139,6 +30162,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
@ -30200,6 +30224,7 @@
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
@ -30268,6 +30293,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
@ -30317,6 +30343,7 @@
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_system_messages": true,
"supports_tool_choice": false,
"supports_url_context": false,
@ -31743,6 +31770,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
@ -31927,6 +31955,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
@ -31989,6 +32018,7 @@
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
@ -32051,6 +32081,7 @@
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
@ -32420,6 +32451,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
@ -32481,6 +32513,7 @@
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
@ -32549,6 +32582,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
@ -32598,6 +32632,7 @@
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_system_messages": true,
"supports_tool_choice": false,
"supports_url_context": false,
@ -33192,6 +33227,7 @@
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_vision": true
},
"github_copilot/gemini-3.8-flash": {
@ -33210,6 +33246,7 @@
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_vision": true
},
"github_copilot/gpt-3.5-turbo": {
@ -52605,6 +52642,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_url_context": true,
@ -65823,6 +65861,7 @@
"supports_tool_choice": true,
"supports_response_schema": true,
"supports_reasoning": true,
"supports_sampling_params": false,
"supports_vision": true,
"source": "https://deepinfra.com/pricing",
"supports_audio_input": true
@ -69264,6 +69303,7 @@
"perplexity/google/gemini-3.5-flash-lite": {
"litellm_provider": "perplexity",
"mode": "responses",
"supports_sampling_params": false,
"supports_web_search": true,
"supports_function_calling": true,
"input_cost_per_token": 3e-07,
@ -69274,6 +69314,7 @@
"perplexity/google/gemini-3.6-flash": {
"litellm_provider": "perplexity",
"mode": "responses",
"supports_sampling_params": false,
"supports_web_search": true,
"supports_function_calling": true,
"input_cost_per_token": 1.5e-06,
@ -69284,6 +69325,7 @@
"perplexity/google/gemini-3.7-flash": {
"litellm_provider": "perplexity",
"mode": "responses",
"supports_sampling_params": false,
"supports_web_search": true,
"supports_function_calling": true,
"input_cost_per_token": 7.5e-07,
@ -71375,6 +71417,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_vision": true,
"supports_web_search": true
},
@ -71392,6 +71435,7 @@
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
"supports_vision": true,
"supports_web_search": true
},
@ -74322,6 +74366,7 @@
"perplexity/google/gemini-3.8-flash": {
"litellm_provider": "perplexity",
"mode": "responses",
"supports_sampling_params": false,
"input_cost_per_token": 7.5e-07,
"output_cost_per_token": 3.75e-06,
"cache_read_input_token_cost": 7.5e-08,
@ -75064,12 +75109,10 @@
"gemini-nano-banana-2.1": {
"cache_read_input_token_cost": 1.5e-07,
"cache_read_input_token_cost_batches": 7.5e-08,
"cache_read_input_token_cost_flex": 7.5e-08,
"cache_read_input_token_cost_priority": 2.7e-07,
"input_cost_per_image": 0.00168,
"input_cost_per_token": 1.5e-06,
"input_cost_per_token_batches": 7.5e-07,
"input_cost_per_token_flex": 7.5e-07,
"input_cost_per_token_priority": 2.7e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 131072,
@ -75080,7 +75123,6 @@
"output_cost_per_image_token": 3e-05,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_flex": 3.75e-06,
"output_cost_per_token_priority": 1.35e-05,
"search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014},
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
@ -75113,12 +75155,10 @@
"vertex_ai/gemini-nano-banana-2.1": {
"cache_read_input_token_cost": 1.5e-07,
"cache_read_input_token_cost_batches": 7.5e-08,
"cache_read_input_token_cost_flex": 7.5e-08,
"cache_read_input_token_cost_priority": 2.7e-07,
"input_cost_per_image": 0.00168,
"input_cost_per_token": 1.5e-06,
"input_cost_per_token_batches": 7.5e-07,
"input_cost_per_token_flex": 7.5e-07,
"input_cost_per_token_priority": 2.7e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 131072,
@ -75129,7 +75169,6 @@
"output_cost_per_image_token": 3e-05,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_flex": 3.75e-06,
"output_cost_per_token_priority": 1.35e-05,
"search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014},
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
@ -76929,5 +76968,55 @@
"supported_output_modalities": [
"text"
]
},
"fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash-1017-us": {
"cache_read_input_token_cost": 9e-09,
"cache_read_input_token_cost_priority": 1.125e-08,
"input_cost_per_token": 4.5e-07,
"input_cost_per_token_priority": 5.625e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 1048576,
"max_output_tokens": 393216,
"max_tokens": 393216,
"mode": "chat",
"output_cost_per_token": 1.8e-06,
"output_cost_per_token_priority": 2.25e-06,
"source": "https://docs.fireworks.ai/serverless/us-only-serverless",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"wandb/reflection/Beam-1.0-501B-A23B": {
"cache_read_input_token_cost": 1e-07,
"input_cost_per_token": 1e-06,
"litellm_provider": "wandb",
"max_input_tokens": 1049000,
"mode": "chat",
"output_cost_per_token": 3e-06,
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": false
},
"wandb/zai-org/GLM-5.3": {
"cache_read_input_token_cost": 1.9e-07,
"input_cost_per_token": 1.2e-06,
"litellm_provider": "wandb",
"max_input_tokens": 1049000,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": false
}
}

2
litellm/models/credentials.py
View file

@ -10,6 +10,7 @@ from typing import Literal, TypeAlias
from pydantic import ConfigDict, Field, model_validator
from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST
from litellm.types.llms.base import LiteLLMBaseModel
CredentialSource: TypeAlias = Literal["db", "config"]
@ -101,6 +102,7 @@ class UserConnectionFlowHandle(LiteLLMBaseModel):
user_id: str
credential_name: str
github_host: str = GITHUB_COPILOT_DEFAULT_HOST
device_code: str = Field(repr=False)
interval: int
expires_at: float

17
litellm/ocr/dispatch.py
View file

@ -9,7 +9,14 @@ from litellm.rust_bridge import runtime
from litellm.rust_bridge.catalog import Route, RouteContext
from litellm.rust_bridge.dispatch import PublicDispatch
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
from litellm.rust_bridge.public_call import NativeCall, native_call, native_call_hook, optional_str, signature
from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
native_call,
native_call_hook,
optional_str,
signature,
)
__all__ = ("aocr", "ocr")
@ -39,6 +46,7 @@ def _bind_request(
_OCR: Final = signature(_bind_request)
_FLAG: Final = AsyncFlag("aocr", sync=_OCR)
def _public_request(name: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> NativeCall:
@ -59,7 +67,6 @@ _DISPATCH: Final = PublicDispatch(
route=Route.OCR,
request=lambda args, kwargs: _public_request("ocr", args, kwargs),
context=_context,
bypass=lambda request: request.kwargs.get("aocr") is True,
)
_ADISPATCH: Final = PublicDispatch(
@ -73,9 +80,11 @@ def ocr(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public OCR call shape
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
if _FLAG.requested(kwargs):
return aocr(**_FLAG.keywords(args, kwargs))
return _DISPATCH.run(
args,
kwargs,
_FLAG.dropped(kwargs),
python=runtime.NO_PYTHON,
binding=NATIVE_OCR,
native=native_call_hook,
@ -85,7 +94,7 @@ def ocr(
async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: preserve the public OCR call shape
return await _ADISPATCH.arun(
args,
kwargs,
_FLAG.dropped(kwargs),
python=runtime.NO_PYTHON,
binding=NATIVE_AOCR,
native=native_call_hook,

3
litellm/passthrough/main.py
View file

@ -437,6 +437,7 @@ def llm_passthrough_route(
params: QueryParamTypes | None = None,
cookies: CookieTypes | None = None,
client: HTTPHandler | AsyncHTTPHandler | None = None,
guardrailConfig: object = None,
**kwargs,
) -> (
httpx.Response
@ -547,7 +548,7 @@ def llm_passthrough_route(
_request_data: dict | None = data if isinstance(data, dict) else (json if isinstance(json, dict) else None)
headers, signed_json_body = provider_config.sign_request(
headers=headers,
litellm_params=litellm_params_dict,
litellm_params={**litellm_params_dict, "guardrailConfig": guardrailConfig},
request_data=_request_data,
api_base=str(updated_url),
model=model,

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

@ -1197,6 +1197,8 @@ def should_strip_caller_authorization(
admission credential — a virtual key, an IdP JWT, or an SSO / OIDC /
session token whose ``api_key`` is ``None`` — and must never reach
the upstream, so it is stripped regardless of the ``api_key`` value.
- **Every other auth-enabled mode**: strip when ``Authorization`` is the
credential admission validated.
"""
if _authorization_is_admission_jwt(raw_headers, user_api_key_auth):
return True
@ -1216,7 +1218,21 @@ def should_strip_caller_authorization(
mcp_server.auth_type == MCPAuth.oauth2 and mcp_server.delegate_auth_to_upstream
)
if not (mcp_server.is_oauth_passthrough or is_delegated_oauth):
return False
from litellm.proxy.auth.user_api_key_auth import is_no_auth_dev_mode
from litellm.proxy.proxy_server import (
general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict
master_key,
)
proxy_general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings)
admission_ran: Final = not is_no_auth_dev_mode(master_key, proxy_general_settings) or (
user_api_key_auth is not None and user_api_key_auth.authenticated_by_custom_auth
)
return admission_ran and (
_authorization_is_litellm_admission_credential_with_settings(
raw_headers, user_api_key_auth, proxy_general_settings
)
)
has_explicit_litellm_admission_header: Final = _has_explicit_litellm_admission_header(raw_headers)
if is_delegated_oauth:
@ -1243,16 +1259,39 @@ def _has_explicit_litellm_admission_header(raw_headers: Mapping[str, str] | None
return bool(_raw_header_value(raw_headers, "x-litellm-api-key"))
_GENERAL_SETTINGS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _authorization_is_litellm_admission_credential(
raw_headers: Mapping[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
) -> bool:
from litellm.proxy.proxy_server import (
general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict
)
proxy_general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings)
return _authorization_is_litellm_admission_credential_with_settings(
raw_headers, user_api_key_auth, proxy_general_settings
)
def _authorization_is_litellm_admission_credential_with_settings(
raw_headers: Mapping[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
proxy_general_settings: Mapping[str, object],
) -> bool:
"""True when ``Authorization`` carries the LiteLLM key admission validated.
That is the case when no usable ``x-litellm-api-key`` was sent, or when the client repeated the
same key in both headers.
That is the case when no usable admission header was sent, or when the client repeated the same key in both headers.
"""
admission_header: Final = _raw_header_value(raw_headers, "x-litellm-api-key")
custom_header_name: Final[object | None] = proxy_general_settings.get("litellm_key_header_name")
custom_admission_header: Final = (
_raw_header_value(raw_headers, custom_header_name.lower())
if isinstance(custom_header_name, str) and custom_header_name
else None
)
admission_header: Final = custom_admission_header or _raw_header_value(raw_headers, "x-litellm-api-key")
authorization: Final = _raw_header_value(raw_headers, "authorization")
if admission_header and authorization:
return strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme(admission_header, "Bearer")
@ -4322,6 +4361,7 @@ class MCPServerManager:
raw_headers: Mapping[str, str] | None = None,
client_ip: str | None = None,
protocol_version_override: MCPUpstreamProtocol | None = None,
for_tool_listing: bool = False,
) -> MCPClient:
self.catalog.assert_current(server)
record_auth_resolution(server.server_id, AuthResolution.unresolved)
@ -4355,6 +4395,7 @@ class MCPServerManager:
else None
),
upstream_transport=self._upstream_transport,
for_tool_listing=for_tool_listing,
)
_create_mcp_client = create_mcp_client
@ -4501,6 +4542,7 @@ class MCPServerManager:
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
for_tool_listing=True,
)
## HANDLE OPENAPI TOOLS

12
litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py
View file

@ -460,6 +460,8 @@ def validate_static_credential(
headers: Mapping[str, str],
upstream_token_header: str | None = None,
static_header_names: Iterable[str] = (),
*,
allow_absent: bool = False,
) -> Result[None, CredError]:
if auth_type not in _STATIC_MODES:
return Ok(None)
@ -475,17 +477,23 @@ def validate_static_credential(
)
)
values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots)
if allow_absent and not values:
return Ok(None)
if any(_usable_credential_value(auth_type, name, value) for name, value in values):
return Ok(None)
return Error(CredError.of_misconfigured(f"{auth_type} requires a usable upstream credential"))
async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient:
async def prepare_mcp_client(server: MCPServer, client: MCPClient, *, for_tool_listing: bool = False) -> MCPClient:
if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio:
return client
request: Final = await client.prepare_request_auth()
match validate_static_credential(
server.auth_type, request.headers, server.upstream_token_header, server.static_headers or ()
server.auth_type,
request.headers,
server.upstream_token_header,
server.static_headers or (),
allow_absent=for_tool_listing and server.auth_type == MCPAuth.authorization,
):
case Error(error):
raise_public(error)

3
litellm/proxy/_experimental/mcp_server/upstream.py
View file

@ -208,6 +208,7 @@ async def prepare_upstream_client(
sampling_callback: SamplingCallback | None = None,
elicitation_callback: ElicitationCallback | None = None,
upstream_transport: httpx2.AsyncBaseTransport | None = None,
for_tool_listing: bool = False,
) -> MCPClient:
transport: Final[MCPTransportType] = server.transport or MCPTransport.sse
server_spec: Final = None if transport == MCPTransport.stdio else to_server_spec_fail_closed(server)
@ -272,6 +273,7 @@ async def prepare_upstream_client(
sampling_callback=sampling_callback,
elicitation_callback=elicitation_callback,
),
for_tool_listing=for_tool_listing,
)
aws_auth: Final = (
MCPSigV4Auth(
@ -316,4 +318,5 @@ async def prepare_upstream_client(
sampling_callback=sampling_callback,
elicitation_callback=elicitation_callback,
),
for_tool_listing=for_tool_listing,
)

2
litellm/proxy/_lazy_features.py
View file

@ -275,7 +275,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
LazyFeature(
name="decisions",
module_path="litellm.proxy.decisions_endpoints.endpoints",
path_prefixes=("/v1/decisions", "/decisions", "/v1/systemone", "/systemone"),
path_prefixes=("/v1/decisions", "/decisions", "/v1/systemone", "/systemone", "/v1/extraction", "/extraction"),
),
LazyFeature(
name="claude_code_marketplace",

1037
litellm/proxy/_lazy_openapi_snapshot.json
View file

File diff suppressed because it is too large Load diff

3
litellm/proxy/_types.py
View file

@ -377,6 +377,8 @@ class LiteLLMRoutes(enum.Enum):
"/v1/decisions",
"/systemone",
"/v1/systemone",
"/extraction",
"/v1/extraction",
# OCR
"/ocr",
"/v1/ocr",
@ -3729,6 +3731,7 @@ class SpendLogsMetadata(TypedDict):
routing_decision: StandardLoggingRoutingDecision | None
kubernetes_pod_routing: ReadOnly[StandardLoggingKubernetesPodRouting | None]
internal_call_origin: InternalCallOrigin | None
billed_retry_attempt: ReadOnly[NotRequired[bool | None]]
litellm_roi_estimator: ReadOnly[NotRequired[bool | None]]
guardrail_information: list[StandardLoggingGuardrailInformation] | None
eval_information: Any | None

2
litellm/proxy/agent_endpoints/auth/managed_authorization.py
View file

@ -32,6 +32,7 @@ _MANAGED_MODEL_ROUTES: Final = frozenset(
"rerank",
"decisions",
"systemone",
"extraction",
"ocr",
),
)
@ -77,6 +78,7 @@ _MODEL_ROUTE_KINDS: Final[
"/rerank": "body",
"/decisions": "body",
"/systemone": "body",
"/extraction": "body",
"/messages/count_tokens": "body",
":countTokens": "path",
}

2
litellm/proxy/auth/auth_utils.py
View file

@ -446,6 +446,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = (
# the request away from the admin's pinned configuration.
"nvcf_function_id",
"use_ssl",
"riva_offline",
# Per-deployment opt-in that hands the whole call to the Rust core. It is a
# deployment decision, not a request one: the Rust path uses its own client
# rather than the one the deployment configured, and reports no post_call,
@ -462,6 +463,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = (
# stored credential's values and the caller's connection by the proxy's own
# secret_fields, so a body-supplied value could only spoof either.
"github_copilot_auth_type",
"github_copilot_host",
"user_provider_credentials",
"github_copilot_user_session",
# Observability credentials, hosts, and project identifiers: derived

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

@ -494,7 +494,7 @@ Each command resolves your LiteLLM key (logging in via SSO when none is stored a
The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, so the proxy lists every other group to Claude Code as `claude-router-<UTF-8 hex of the group name>` and marks a group whose input window reaches 1M with `[1m]`, and a request on such an id is served by the group. Older Claude Code versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway. Codex gets the same list as a catalog file, `$CODEX_HOME/litellm-models.json` (default `~/.codex/`), passed as `-c model_catalog_json=<path>` so `/model` lists exactly the proxy's chat models; a proxy model the installed Codex already knows (`gpt-5.5`, say) keeps that Codex's own entry, reasoning levels and prompt included, and only its place in the picker comes from the proxy, while a model Codex does not know gets the plain entry Codex uses for an unknown `-m` slug. Before launching, `lite codex` asks the installed Codex for its own list and then has it read the written file back (both through `codex debug models`), and when the fetch, either of those or the write fails (Codex releases older than 0.130 have no such command) it says so on stderr and launches with Codex's built-in catalog, leaving a rejected file in place.
pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/<id>`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. Your own flags come after the injected pin, so `lite pi --model litellm/<other-id>` wins, and inside the TUI the `/model` picker lists every synced litellm model.
pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/<id>`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. The entry also turns on pi's session header (`x-session-id`, sent by pi 0.80.7 and later), so every chat turn of one pi session shares one session id in the logs and sticks to one deployment when session affinity is on. A custom `pi --session-id` must be at least 8 letters, digits, `_`, or `-` for the proxy to read it. Your own flags come after the injected pin, so `lite pi --model litellm/<other-id>` wins, and inside the TUI the `/model` picker lists every synced litellm model.
Options (these belong to the wrapper, so put them before the agent's own flags):

4
litellm/proxy/client/cli/commands/pi.py
View file

@ -191,7 +191,7 @@ def provider_block(
Real contextWindow/maxTokens matter: pi otherwise assumes 128k/16384, which
breaks compaction thresholds and over-asks models with smaller output caps.
pi sniffs compat from the base URL, and one gateway URL fronts models with
different capabilities, so both flags are pinned off.
different capabilities, so the store and long-cache flags are pinned off.
"""
return {
"baseUrl": base_url.rstrip("/") + "/v1",
@ -199,6 +199,8 @@ def provider_block(
"compat": {
"supportsStore": False,
"supportsLongCacheRetention": False,
"sendSessionAffinityHeaders": True,
"sessionAffinityFormat": "openrouter",
},
"apiKey": f"${LITELLM_PROXY_API_KEY_ENV}",
"models": [_model_entry(model_id, limits) for model_id in model_ids],

47
litellm/proxy/common_request_processing.py
View file

@ -68,6 +68,7 @@ from litellm.litellm_core_utils.get_supported_openai_params import (
get_supported_openai_params,
)
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params
from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
from litellm.litellm_core_utils.llm_response_utils.get_headers import (
@ -134,12 +135,14 @@ from litellm.router import Router
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX
from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model
from litellm.types.extraction import ExtractionResponse
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.router import RouterRateLimitError
from litellm.types.router_weights import validate_router_weights
_LateResponseT = TypeVar("_LateResponseT", bound=Response)
_LlmCallT = TypeVar("_LlmCallT")
_ResponseT: Final = TypeVar("_ResponseT")
_ChunkT: Final = TypeVar("_ChunkT")
KNOWN_PROXY_ROUTES: Final = frozenset(
@ -191,6 +194,7 @@ ProxyRouteType: TypeAlias = Literal[
"asearch",
"adecisions",
"asystemone",
"aextraction",
"avideo_generation",
"avideo_list",
"avideo_status",
@ -263,6 +267,7 @@ from litellm.proxy.litellm_pre_call_utils import (
sanitize_for_log,
)
from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval
from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD
from litellm.types.utils import (
ModelResponse,
ModelResponseStream,
@ -549,22 +554,29 @@ def model_deployments_support_stream_options(
return all(_litellm_model_supports_stream_options(m) for m in candidate_models)
_STREAM_USAGE_FLAGS: Final = ("_litellm_strip_stream_usage", STREAM_USAGE_INJECTED_FIELD)
def stream_usage_tracking_updates(
data: Mapping[str, object],
general_settings: Mapping[str, object],
route_type: str,
supports_stream_options: Callable[[], bool],
) -> Mapping[str, object]:
scrub: Final = {"_litellm_strip_stream_usage": False} if "_litellm_strip_stream_usage" in data else {}
scrub: Final = {flag: False for flag in _STREAM_USAGE_FLAGS if flag in data}
if data.get("stream", False) is not True:
return scrub
always_include: Final = general_settings.get("always_include_stream_usage")
stream_options: Final = data.get("stream_options")
if always_include is True:
if "stream_options" not in data:
return {**scrub, "stream_options": {"include_usage": True}}
return {**scrub, "stream_options": {"include_usage": True}, STREAM_USAGE_INJECTED_FIELD: True}
if isinstance(stream_options, dict) and "include_usage" not in stream_options:
return {**scrub, "stream_options": {**stream_options, "include_usage": True}}
return {
**scrub,
"stream_options": {**stream_options, "include_usage": True},
STREAM_USAGE_INJECTED_FIELD: True,
}
return scrub
if always_include is False and route_type == "agenerate_content_stream" and "stream_options" not in data:
return {**scrub, "stream_options": {"include_usage": False}}
@ -578,6 +590,7 @@ def stream_usage_tracking_updates(
return {
"stream_options": {**merged_stream_options, "include_usage": True},
"_litellm_strip_stream_usage": True,
STREAM_USAGE_INJECTED_FIELD: True,
}
@ -1482,11 +1495,11 @@ _is_azure_model_router_request: Final = is_azure_model_router_request
def _override_openai_response_model(
*,
response_obj: object,
response_obj: _ResponseT,
requested_model: str,
log_context: str,
return_raw_model_name: bool = False,
) -> None:
) -> _ResponseT:
"""
Force the OpenAI-compatible `model` field in the response to match what the client requested.
@ -1510,7 +1523,7 @@ def _override_openai_response_model(
model group name instead of the comma-separated list the client sent.
"""
if return_raw_model_name or not requested_model:
return
return response_obj
hidden_params: Final = get_hidden_params_dict(response_obj)
if isinstance(hidden_params, dict):
@ -1523,7 +1536,7 @@ def _override_openai_response_model(
log_context,
attempted_fallbacks,
)
return
return response_obj
# For fastest_response batch completions, use the winning model's group
# name rather than the comma-separated list the client sent.
@ -1542,7 +1555,7 @@ def _override_openai_response_model(
"%s: fastest_response detected but no model group header found, preserving actual model from response.",
log_context,
)
return
return response_obj
# Check if this is an Azure Model Router request - if so, preserve the actual model used
if is_azure_model_router_request(requested_model, hidden_params):
@ -1550,11 +1563,12 @@ def _override_openai_response_model(
"%s: Azure Model Router detected - preserving actual model used from response instead of overriding to router model.",
log_context,
)
return
return response_obj
unchanged_response: Final = response_obj
if isinstance(response_obj, dict):
if "model" not in response_obj:
return
return unchanged_response
downstream_model = response_obj.get("model")
if downstream_model != requested_model:
verbose_proxy_logger.debug(
@ -1564,10 +1578,10 @@ def _override_openai_response_model(
downstream_model,
)
response_obj["model"] = requested_model
return
return unchanged_response
if not hasattr(response_obj, "model"):
return
return response_obj
downstream_model = getattr(response_obj, "model", None)
if downstream_model != requested_model:
@ -1578,6 +1592,9 @@ def _override_openai_response_model(
downstream_model,
)
if isinstance(response_obj, ExtractionResponse):
return response_obj.model_copy(update={"model": requested_model})
try:
setattr(response_obj, "model", requested_model)
except Exception as e:
@ -1589,6 +1606,7 @@ def _override_openai_response_model(
str(e),
exc_info=True,
)
return response_obj
_METADATA_BUCKET_KEYS: Final = ("metadata", "litellm_metadata")
@ -2244,6 +2262,7 @@ class ProxyBaseLLMRequestProcessing:
"asearch",
"adecisions",
"asystemone",
"aextraction",
"avideo_generation",
"avideo_list",
"avideo_status",
@ -3270,7 +3289,7 @@ class ProxyBaseLLMRequestProcessing:
# Always return the client-requested model name (not provider-prefixed internal identifiers)
# for OpenAI-compatible responses.
if requested_model_from_client:
_override_openai_response_model(
response = _override_openai_response_model( # rebind-ok: a frozen response comes back as a copy
response_obj=response,
requested_model=requested_model_from_client,
log_context=f"litellm_call_id={logging_obj.litellm_call_id}",
@ -4024,6 +4043,8 @@ class ProxyBaseLLMRequestProcessing:
async def _on_deferred_native_stream_complete(*args: object) -> None:
match args:
case (BilledRoundsDispatch(dispatch=logging_coroutine),):
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine)
case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine):
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine)
case (assembled_response, cache_hit):

108
litellm/proxy/credential_endpoints/endpoints.py
View file

@ -19,6 +19,7 @@ from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
GITHUB_COPILOT_AUTH_TYPE_KEY,
GITHUB_COPILOT_HOST_KEY,
GITHUB_COPILOT_PER_USER_AUTH_TYPE,
)
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
@ -84,6 +85,22 @@ _DISPLAY_NAME_MAX_LENGTH: Final = 255
_DISPLAY_NAME_FORBIDDEN_CATEGORIES: Final = frozenset({"Cc", "Cf", "Cs", "Zl", "Zp"})
def _normalized_github_copilot_values(
credential_values: Mapping[str, object],
credential_info: Mapping[str, object],
) -> Mapping[str, object]:
provider: Final = stored_credential_provider(credential_info.get("custom_llm_provider"))
if provider != _GITHUB_COPILOT_PROVIDER or GITHUB_COPILOT_HOST_KEY not in credential_values:
return credential_values
from litellm.llms.github_copilot.github_host import normalized_github_host
host: Final = normalized_github_host(credential_values[GITHUB_COPILOT_HOST_KEY])
if host is None:
raise HTTPException(status_code=400, detail=f"Invalid {GITHUB_COPILOT_HOST_KEY}")
return {**credential_values, GITHUB_COPILOT_HOST_KEY: host}
def _reject_non_admin_wif_fields(
wif_fields: tuple[str, ...],
user_api_key_dict: UserAPIKeyAuth,
@ -318,7 +335,14 @@ async def create_credential(
status_code=400,
detail="Credential values are required. Unable to infer credential values from model ID.",
)
_reject_non_admin_wif_fields(server_owned_wif_fields_written(credential_values), user_api_key_dict)
credential_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
credential.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field
)
normalized_values: Final = _normalized_github_copilot_values(
_CREDENTIAL_DICT_ADAPTER.validate_python(credential_values),
credential_info,
)
_reject_non_admin_wif_fields(server_owned_wif_fields_written(normalized_values), user_api_key_dict)
_reject_non_admin_wif_fields(
await named_credential_wif_fields(credential.credential_name, prisma_client), user_api_key_dict
)
@ -326,8 +350,8 @@ async def create_credential(
processed_credential: Final = CredentialItem(
credential_name=credential.credential_name,
display_name=_normalized_display_name(credential.display_name),
credential_values=_without_null_values(_CREDENTIAL_DICT_ADAPTER.validate_python(credential_values)),
credential_info=credential.credential_info,
credential_values=_without_null_values(normalized_values),
credential_info=credential_info,
)
encrypted_credential: Final = CredentialHelperUtils.encrypt_credential_values(processed_credential)
credentials_dict: Final = encrypted_credential.model_dump(exclude_none=True)
@ -614,22 +638,24 @@ async def start_user_connection(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI resolves the dependency from the default
) -> UserConnectionStartResponse:
"""Begin a GitHub device flow for the calling user's connection to a per-user credential."""
from litellm.llms.github_copilot.per_user_auth import astart_device_flow
from litellm.llms.github_copilot.per_user_auth import astart_device_flow, github_copilot_host_endpoints
from litellm.proxy.proxy_server import prisma_client
try:
user_id: Final = _authenticated_user_id(user_api_key_dict)
_per_user_credential_or_404(credential_name)
endpoints: Final = github_copilot_host_endpoints(credential_name)
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
flow: Final = await astart_device_flow()
flow: Final = await astart_device_flow(endpoints)
handle: Final = encrypt_value_helper(
UserConnectionFlowHandle(
user_id=user_id,
credential_name=credential_name,
github_host=endpoints.host,
device_code=flow.device_code,
interval=flow.interval,
expires_at=time.time() + flow.expires_in,
@ -665,12 +691,14 @@ async def poll_user_connection(
acheck_copilot_seat,
afetch_github_login,
apoll_device_flow,
github_copilot_host_endpoints,
)
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
try:
user_id: Final = _authenticated_user_id(user_api_key_dict)
_per_user_credential_or_404(credential_name)
endpoints: Final = github_copilot_host_endpoints(credential_name)
if prisma_client is None:
raise HTTPException(
status_code=500,
@ -681,11 +709,12 @@ async def poll_user_connection(
handle is None
or handle.user_id != user_id
or handle.credential_name != credential_name
or handle.github_host != endpoints.host
or handle.expires_at <= time.time()
):
raise HTTPException(status_code=400, detail="invalid or expired flow_handle")
poll: Final = await apoll_device_flow(handle.device_code)
poll: Final = await apoll_device_flow(handle.device_code, endpoints)
if poll.status in ("expired", "denied", "pending", "slow_down"):
return UserConnectionPollResponse(status=poll.status, interval=poll.interval)
@ -693,12 +722,19 @@ async def poll_user_connection(
if not github_token:
raise HTTPException(status_code=502, detail="GitHub device flow completed without a token")
try:
await acheck_copilot_seat(user_id=user_id, github_token=github_token, credential_name=credential_name)
await acheck_copilot_seat(
user_id=user_id,
github_token=github_token,
credential_name=credential_name,
endpoints=endpoints,
)
except litellm.CallerCredentialAuthenticationError:
return UserConnectionPollResponse(status="no_copilot_seat")
github_login: Final = await afetch_github_login(github_token)
github_login: Final = await afetch_github_login(github_token, endpoints)
connection_payload: Final = GithubCopilotUserConnectionPayload(
access_token=github_token, github_login=github_login
access_token=github_token,
github_login=github_login,
github_host=handle.github_host,
)
await upsert_user_provider_credential(
prisma_client,
@ -760,7 +796,7 @@ async def delete_user_connection(
)
prior: Final = await delete_user_provider_credential(prisma_client, user_id, credential_name)
if prior is not None:
evict_copilot_user_session(user_id, prior.access_token)
evict_copilot_user_session(user_id, credential_name, prior.github_host)
return UserConnectionDeleteResponse(status="disconnected")
except Exception as e: # noqa: BLE001 # endpoint boundary: every failure becomes the proxy error contract
raise handle_exception_on_proxy(e)
@ -785,16 +821,15 @@ async def _purge_user_connections_for_credential(credential_name: str) -> None:
"_purge_user_connections_for_credential: failed to drop user connections for %s", credential_name
)
return
access_tokens: Final[dict[str, str]] = {} # mutable-ok: accumulates one entry per deleted row
for row in prior_rows:
decoded = decode_user_provider_credential(row.credential_b64)
if decoded is not None:
access_tokens[row.user_id] = decoded.access_token
github_hosts: Final = {
row.user_id: decoded.github_host
for row in prior_rows
if (decoded := decode_user_provider_credential(row.credential_b64)) is not None
}
for user_id in user_ids:
await invalidate_user_provider_credential_cache(user_api_key_cache, user_id, credential_name)
token = access_tokens.get(user_id)
if token:
evict_copilot_user_session(user_id, token)
if user_id in github_hosts:
evict_copilot_user_session(user_id, credential_name, github_hosts[user_id])
@router.delete(
@ -926,7 +961,25 @@ async def update_credential(
credential_values=incoming_values,
credential_values_to_delete=credential.credential_values_to_delete,
)
merged_credential: Final = update_db_credential(db_credential, patch)
decrypted_stored_credential: Final = decrypted_credential(db_credential)
decrypted_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
decrypted_stored_credential.credential_values
)
stored_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
db_credential.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field
)
patch_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
patch.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field
)
effective_info: Final = {**stored_info, **patch_info}
normalized_patch: Final = (
patch.model_copy(
update={"credential_values": _normalized_github_copilot_values(incoming_values, effective_info)}
)
if incoming_values.get(GITHUB_COPILOT_HOST_KEY) is not None
else patch
)
merged_credential: Final = update_db_credential(db_credential, normalized_patch)
credential_object_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str
"dict[str, object]", jsonify_object(merged_credential.model_dump(exclude_none=True))
)
@ -940,13 +993,20 @@ async def update_credential(
)
# Sync in-memory credential_list (skip if not in memory - e.g., proxy restarted)
_sync_in_memory_credential(patch, credential_name)
_sync_in_memory_credential(normalized_patch, credential_name)
merged_values: Final = cast( # cast-ok: credential_values is a plain dict at runtime
Mapping[object, object], merged_credential.credential_values
from litellm.llms.github_copilot.github_host import normalized_github_host
new_decrypted_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python(
decrypted_credential(merged_credential).credential_values
)
still_per_user: Final = merged_values.get(GITHUB_COPILOT_AUTH_TYPE_KEY) == GITHUB_COPILOT_PER_USER_AUTH_TYPE
if not still_per_user:
still_per_user: Final = (
new_decrypted_values.get(GITHUB_COPILOT_AUTH_TYPE_KEY) == GITHUB_COPILOT_PER_USER_AUTH_TYPE
)
old_host: Final = normalized_github_host(decrypted_values.get(GITHUB_COPILOT_HOST_KEY))
new_host: Final = normalized_github_host(new_decrypted_values.get(GITHUB_COPILOT_HOST_KEY))
host_changed: Final = old_host is None or old_host != new_host
if not still_per_user or host_changed:
await _purge_user_connections_for_credential(credential_name)
return {"success": True, "message": "Credential updated successfully"}

30
litellm/proxy/credential_endpoints/user_provider_credentials.py
View file

@ -25,6 +25,7 @@ from litellm.caching.dual_cache import DualCache
if TYPE_CHECKING:
from litellm.caching.redis_cache import RedisCache
from litellm.constants import (
GITHUB_COPILOT_DEFAULT_HOST,
GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS,
USER_PROVIDER_CREDENTIAL_CACHE_PREFIX,
USER_PROVIDER_CREDENTIAL_NOT_CONNECTED,
@ -50,6 +51,7 @@ class GithubCopilotUserConnectionPayload(BaseModel):
access_token: str = Field(repr=False)
github_login: str
github_host: str = GITHUB_COPILOT_DEFAULT_HOST
class _UserProviderCredentialsRepository(PrismaTableRepository["prisma_models.LiteLLM_UserProviderCredentials"]):
@ -186,7 +188,7 @@ async def _try_cache_get(token_cache: "RedisCache", key: str) -> object:
try:
return await _string_cache(token_cache).async_get_cache(key)
except Exception: # noqa: BLE001 # a Redis outage must fall back to the database, never reject a connected user
verbose_proxy_logger.warning("aget_user_provider_tokens: Redis get failed; falling back to the database")
verbose_proxy_logger.warning("aget_user_provider_credentials: Redis get failed; falling back to the database")
return None
@ -196,7 +198,7 @@ async def _try_cache_set(token_cache: "RedisCache", key: str, value: str, nx: bo
key, value, nx=nx, ttl=GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS
)
except Exception: # noqa: BLE001 # caching is best-effort; a Redis outage is not worth failing the request
verbose_proxy_logger.warning("aget_user_provider_tokens: Redis set failed; skipping the cache write")
verbose_proxy_logger.warning("aget_user_provider_credentials: Redis set failed; skipping the cache write")
@with_service_target("user_provider_connections")
@ -292,16 +294,16 @@ async def set_user_provider_credential_cache(
@with_service_target("user_provider_connections")
async def aget_user_provider_tokens(
async def aget_user_provider_credentials(
prisma_client: "PrismaClient",
cache: DualCache,
user_id: str,
credential_names: Sequence[str],
) -> Mapping[str, str]:
"""Map each per-user credential name to the caller's stored GitHub token.
) -> Mapping[str, Mapping[str, str]]:
"""Map each per-user credential name to the caller's stored GitHub connection.
Cached values are the stored ciphertext or the ``_NOT_CONNECTED`` marker;
plaintext tokens only ever live in the returned dict. Redis is the only
plaintext connection data only ever lives in the returned mapping. Redis is the only
cache layer used: without it there is no shared invalidation, so a worker
serving a stale local entry after a disconnect is worse than a DB read."""
token_cache: Final = cache.redis_cache
@ -320,9 +322,9 @@ async def aget_user_provider_tokens(
misses.append(name)
if not misses:
return {
name: token
name: connection
for name, ciphertext in cached.items()
if (token := _token_from_ciphertext(ciphertext)) is not None
if (connection := _connection_from_ciphertext(ciphertext)) is not None
}
try:
@ -333,7 +335,7 @@ async def aget_user_provider_tokens(
where={"user_id": user_id},
)
except Exception: # noqa: BLE001 # a DB outage reads as not connected so the caller gets a 401, never a 500
verbose_proxy_logger.exception("aget_user_provider_tokens: DB read failed for user_id=%s", user_id)
verbose_proxy_logger.exception("aget_user_provider_credentials: DB read failed")
return {}
found: Final = {row.credential_name: row.credential_b64 for row in rows}
for name in misses:
@ -342,10 +344,14 @@ async def aget_user_provider_tokens(
if name in found:
cached[name] = found[name]
return {
name: token for name, ciphertext in cached.items() if (token := _token_from_ciphertext(ciphertext)) is not None
name: connection
for name, ciphertext in cached.items()
if (connection := _connection_from_ciphertext(ciphertext)) is not None
}
def _token_from_ciphertext(ciphertext: str) -> str | None:
def _connection_from_ciphertext(ciphertext: str) -> Mapping[str, str] | None:
payload: Final = decode_user_provider_credential(ciphertext)
return payload.access_token if payload is not None else None
if payload is None:
return None
return {"access_token": payload.access_token, "github_host": payload.github_host}

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

@ -29,6 +29,7 @@ from litellm._internal_context import service_target, with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching import RedisCache
from litellm.constants import (
BILLED_RETRY_ATTEMPT_METADATA_KEY,
DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME,
DB_SPEND_UPDATE_JOB_NAME,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
@ -312,6 +313,11 @@ def _spend_tables_left_to_send(
return cast(DBSpendUpdateTransactions, remaining) if any(remaining.values()) else None
def _is_billed_retry_attempt(metadata: SpendLogsMetadata) -> bool:
"""An earlier billed retry attempt's own spend row, not a request the caller sent."""
return BILLED_RETRY_ATTEMPT_METADATA_KEY in metadata and metadata[BILLED_RETRY_ATTEMPT_METADATA_KEY] is True
def _timed_request_duration_ms(
payload: dict | SpendLogsPayload,
request_status: Literal["success", "failure"],
@ -1651,6 +1657,8 @@ class DBSpendUpdateWriter:
window_spend_update_queue=self.window_spend_update_queue,
)
await self._flush_tool_discovery_queue(prisma_client=prisma_client)
# Only commit from redis to db if this pod is the leader
if await self.pod_lock_manager.acquire_lock(
cronjob_id=DB_SPEND_UPDATE_JOB_NAME,
@ -2803,11 +2811,12 @@ class DBSpendUpdateWriter:
if not date:
return
metadata: Final[SpendLogsMetadata] = json.loads(payload["metadata"])
internal_origin: Final = bool(metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)) # pyright: ignore[reportUnknownMemberType] # TypedDict.get overloads carry Any defaults
request_error_accumulator.record(
payload=payload,
request_status=prisma_client.get_request_status(payload), # pyright: ignore[reportUnknownMemberType] # legacy payload union; this caller supplies a typed spend payload
date=date.split("T")[0],
is_internal_call=bool(metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)), # pyright: ignore[reportUnknownMemberType] # TypedDict.get overloads carry Any defaults
is_internal_call=internal_origin or _is_billed_retry_attempt(metadata),
)
async def _common_add_spend_log_transaction_to_daily_transaction(
@ -2872,6 +2881,7 @@ class DBSpendUpdateWriter:
endpoint = ROUTE_ENDPOINT_MAPPING.get(call_type, None)
is_internal_call: Final = bool(_metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY))
caller_request: Final = not is_internal_call and not _is_billed_retry_attempt(_metadata)
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(
@ -2906,9 +2916,9 @@ class DBSpendUpdateWriter:
# 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,
api_requests=1 if caller_request else 0,
successful_requests=1 if caller_request and request_status == "success" else 0,
failed_requests=1 if caller_request 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,

21
litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
View file

@ -168,9 +168,11 @@ class RedisUpdateBuffer:
@with_service_target(SPEND_QUEUE_TARGET)
async def _store_transactions_in_redis(
self,
transactions: Mapping[str, BaseDailySpendTransaction] | None,
transactions: dict[str, BaseDailySpendTransaction] | None, # mutable-ok: add_update requires dict
redis_key: str,
service_type: ServiceTypes,
*,
source_queue: DailySpendUpdateQueue,
) -> None:
"""
Helper method to store transactions in Redis and emit an event
@ -179,6 +181,7 @@ class RedisUpdateBuffer:
transactions: The transactions to store
redis_key: The Redis key to store under
service_type: The service type for event emission
source_queue: The in-memory queue to restore transactions to if Redis push fails
"""
if transactions is None or len(transactions) == 0:
return
@ -191,6 +194,15 @@ class RedisUpdateBuffer:
key=redis_key,
values=list_of_transactions,
)
except Exception as e:
verbose_proxy_logger.error(
"Spend tracking - failed to push spend updates to Redis (redis_key=%s). "
"Restoring them to the in-memory queue for retry on next tick. Error: %s",
redis_key,
str(e),
)
await source_queue.add_update(transactions)
else:
verbose_proxy_logger.debug(
"Spend tracking - pushed spend updates to Redis buffer. redis_key=%s, buffer_size=%s",
redis_key,
@ -200,12 +212,6 @@ class RedisUpdateBuffer:
queue_size=current_redis_buffer_size,
service=service_type,
)
except Exception as e:
verbose_proxy_logger.error(
"Spend tracking - failed to push spend updates to Redis (redis_key=%s). Error: %s",
redis_key,
str(e),
)
@with_service_target(SPEND_QUEUE_TARGET)
async def store_in_memory_spend_updates_in_redis(
@ -793,6 +799,7 @@ class RedisUpdateBuffer:
transactions=daily_tag_spend_update_transactions,
redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE,
source_queue=daily_tag_spend_update_queue,
)
@with_service_target(SPEND_QUEUE_TARGET)

13
litellm/proxy/db/health_check_latest.py
View file

@ -10,8 +10,9 @@ worker on every call. ``SELECT DISTINCT ON`` keeps the transfer at one row per
from __future__ import annotations
import json
from collections.abc import Mapping, Sequence
from collections.abc import Iterable, Mapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Protocol
from pydantic import ConfigDict, JsonValue, TypeAdapter, field_validator
@ -113,3 +114,13 @@ async def fetch_latest_health_checks_for_models(
except Exception as query_err: # noqa: BLE001 # a paged model list must not fail on its health decoration
verbose_proxy_logger.error("Error getting latest health checks for models: %s", query_err)
return ()
def latest_by_model_name(rows: Iterable[LatestHealthCheckRow]) -> Mapping[str, LatestHealthCheckRow]:
"""The newest row per model name.
The queries above return one row per (model_id, model_name), so a model group with several
deployments comes back as several rows, plus one without a model_id if it was ever checked by
name. They arrive in model_id order with NULL last, which says nothing about which check is newest.
"""
return MappingProxyType({row.model_name: row for row in sorted(rows, key=lambda row: row.checked_at)})

103
litellm/proxy/decisions_endpoints/endpoints.py
View file

@ -1,6 +1,6 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import Annotated, Final, Literal
from typing import Annotated, Final, Literal, TypeAlias
from fastapi import APIRouter, Depends, Request, Response
from fastapi.responses import ORJSONResponse # pyright: ignore[reportDeprecated] # required endpoint contract
@ -18,6 +18,7 @@ from litellm.types.decisions import (
OpenAIDecisionRequestBody,
OpenAIDecisionResponse,
)
from litellm.types.extraction import ExtractionRequest, ExtractionRequestBody, ExtractionResponse
router: Final = APIRouter()
_REQUEST_DATA_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
@ -25,6 +26,7 @@ _DECISIONS_REQUEST_BODY_ADAPTER: Final[TypeAdapter[DecisionsRequestBody]] = Type
_OPENAI_DECISION_REQUEST_BODY_ADAPTER: Final[TypeAdapter[OpenAIDecisionRequestBody]] = TypeAdapter(
OpenAIDecisionRequestBody
)
_EXTRACTION_REQUEST_BODY_ADAPTER: Final[TypeAdapter[ExtractionRequestBody]] = TypeAdapter(ExtractionRequestBody)
_GENERAL_SETTINGS_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
_OPTIONAL_STRING_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None)
_OPTIONAL_FLOAT_ADAPTER: Final[TypeAdapter[float | None]] = TypeAdapter(float | None)
@ -73,16 +75,35 @@ _DECISIONS_REQUEST_BODY: Final = inline_request_body(
],
},
)
_EXTRACTION_REQUEST_BODY: Final = inline_request_body(
ExtractionRequest,
{
"model": "clairic",
"state": "Invoice #4471 from Acme Tools, dated 2026-03-03. 12 hammers. Total $174. Payment due in 30 days.",
"fields": {
"vendor": {"type": "text", "instructions": "Who sent the invoice?"},
"invoice_date": {"type": "date"},
"total": {"type": "number", "unit": "USD"},
"items": {"type": "list", "items": "text", "instructions": "What was bought?"},
"days_to_pay": {"type": "estimate", "min": 0, "max": 90},
},
},
)
_ApiName: TypeAlias = Literal["Decisions", "Extraction"]
_RouteType: TypeAlias = Literal["adecisions", "asystemone", "aextraction"]
async def _invalid_request(
raw_data: Mapping[str, object], error: ValidationError, user_api_key_dict: UserAPIKeyAuth
raw_data: Mapping[str, object],
error: ValidationError,
user_api_key_dict: UserAPIKeyAuth,
api_name: _ApiName,
) -> Exception:
from litellm.proxy.proxy_server import proxy_logging_obj, version
return await ProxyBaseLLMRequestProcessing(data=dict(raw_data)).handle_llm_api_exception(
e=BadRequestError(
message=f"Invalid Decisions request: {error}",
message=f"Invalid {api_name} request: {error}",
model=str(raw_data.get("model", "")),
llm_provider="",
),
@ -95,19 +116,21 @@ async def _invalid_request(
_FIELDS_CHECKED_AFTER_ROUTING: Final = frozenset({"safety_identifier", "images", "assets", "audio", "videos"})
def _fields_checked_before_routing(data: Mapping[str, object]) -> Mapping[str, object]:
def _fields_checked_before_routing(data: Mapping[str, object], route_type: _RouteType) -> Mapping[str, object]:
if route_type == "aextraction":
return data
return MappingProxyType({key: value for key, value in data.items() if key not in _FIELDS_CHECKED_AFTER_ROUTING})
_DEPLOYMENT_ONLY_FIELDS: Final = frozenset({"model_info"})
async def _request_data(request: Request, user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]:
async def _request_data(request: Request, user_api_key_dict: UserAPIKeyAuth, api_name: _ApiName) -> dict[str, object]:
body: Final = await request.body()
try:
data: Final = _REQUEST_DATA_ADAPTER.validate_json(body)
except ValidationError as error:
raise await _invalid_request(raw_data={}, error=error, user_api_key_dict=user_api_key_dict)
raise await _invalid_request(raw_data={}, error=error, user_api_key_dict=user_api_key_dict, api_name=api_name)
return {key: value for key, value in data.items() if key not in _DEPLOYMENT_ONLY_FIELDS}
@ -115,8 +138,10 @@ async def _process_decisions(
request: Request,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth,
body_adapter: TypeAdapter[DecisionsRequestBody] | TypeAdapter[OpenAIDecisionRequestBody],
route_type: Literal["adecisions", "asystemone"],
body_adapter: (
TypeAdapter[DecisionsRequestBody] | TypeAdapter[OpenAIDecisionRequestBody] | TypeAdapter[ExtractionRequestBody]
),
route_type: _RouteType,
) -> object:
from litellm.proxy.proxy_server import (
general_settings as proxy_general_settings,
@ -139,15 +164,17 @@ async def _process_decisions(
user_temperature as proxy_user_temperature,
)
data: Final = await _request_data(request, user_api_key_dict)
api_name: Final[_ApiName] = "Extraction" if route_type == "aextraction" else "Decisions"
data: Final = await _request_data(request, user_api_key_dict, api_name)
try:
body_adapter.validate_python(_fields_checked_before_routing(data))
body_adapter.validate_python(_fields_checked_before_routing(data, route_type))
except ValidationError as error:
raise await _invalid_request(raw_data=data, error=error, user_api_key_dict=user_api_key_dict)
raise await _invalid_request(raw_data=data, error=error, user_api_key_dict=user_api_key_dict, api_name=api_name)
general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(proxy_general_settings)
user_api_base: Final = _OPTIONAL_STRING_ADAPTER.validate_python(proxy_user_api_base)
user_model: Final = _OPTIONAL_STRING_ADAPTER.validate_python(proxy_user_model)
user_temperature: Final = _OPTIONAL_FLOAT_ADAPTER.validate_python(proxy_user_temperature)
takes_generation_settings: Final = route_type != "aextraction"
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
@ -162,9 +189,9 @@ async def _process_decisions(
select_data_generator=None,
model=None,
user_model=user_model,
user_temperature=user_temperature,
user_temperature=user_temperature if takes_generation_settings else None,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_max_tokens=user_max_tokens if takes_generation_settings else None,
user_api_base=user_api_base,
version=version,
)
@ -263,3 +290,53 @@ async def decisions(
body_adapter=_OPENAI_DECISION_REQUEST_BODY_ADAPTER,
route_type="adecisions",
)
@router.post(
"/v1/extraction",
dependencies=[Depends(user_api_key_auth)],
response_class=ORJSONResponse, # pyright: ignore[reportDeprecated] # required endpoint contract
tags=["extraction"],
summary="Extract typed fields from a state",
responses={200: {"model": ExtractionResponse}},
response_model=None,
openapi_extra={"requestBody": _EXTRACTION_REQUEST_BODY},
)
@router.post(
"/extraction",
dependencies=[Depends(user_api_key_auth)],
response_class=ORJSONResponse, # pyright: ignore[reportDeprecated] # required endpoint contract
tags=["extraction"],
summary="Extract typed fields from a state",
responses={200: {"model": ExtractionResponse}},
response_model=None,
openapi_extra={"requestBody": _EXTRACTION_REQUEST_BODY},
)
async def extraction(
request: Request,
fastapi_response: Response,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> object:
"""
Read a `state` text and up to 4 images and get one typed answer per named field in `fields`.
Field types are `text`, `number`, `date`, `time`, `datetime`, `duration`, `list` (set `items`), `estimate`
(set `min` and `max`) and `forecast` (set `min`, `max` and `horizon`). Each answer carries its `value`, which is
null when the input does not determine it, plus `confidence`, `p_null`, text `span`s into the state and, for
estimates, `mean`, `median`, `p10` and `p90`. Images are data URLs or `{content_type, base64}` objects, remote
URLs are not accepted. An unknown field type, a field name that is not 1 to 64 letters, digits, `_`, `.` or `-`
starting with a letter or `_`, a remote image URL or more than 4 images is rejected with a 400 before any
provider is called. Any other top-level key, `stream: true` and OpenAI generation settings such as `temperature`
included, is rejected with a 400 naming it unless `drop_params` or the deployment's `additional_drop_params`
drops it. Settings inside `fields` and `options` go to the provider as sent, so a field missing a setting its
type needs or an empty `state` with no images comes back with the provider's own error. `model` is any
extraction model in the proxy model list, such as Cloudflare Clairic. The request and the answers keep the
provider's format.
"""
return await _process_decisions(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
body_adapter=_EXTRACTION_REQUEST_BODY_ADAPTER,
route_type="aextraction",
)

5
litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py
View file

@ -494,7 +494,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail):
sse: Final = AnthropicSSERestorer(step)
events: Final = ResponsesStreamRestorer(step, rehydrate)
carries: Final[CarryWindows] = {}
last_chunk = None # rebind-ok: tracks the most recent chunk for the final flush.
last_chunk = None # rebind-ok: tracks the most recent chunk with choices for the final flush.
async for chunk in response:
if isinstance(chunk, (bytes, str)):
@ -506,7 +506,8 @@ class LLMShieldProxyGuardrail(CustomGuardrail):
yield event
continue
restored_chunk = detached(chunk)
last_chunk = restored_chunk
if getattr(restored_chunk, "choices", None):
last_chunk = restored_chunk
for choice in attribute_items(restored_chunk, "choices"):
await self._restore_choice(choice, carries, session_id)
yield restored_chunk

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

@ -2155,6 +2155,7 @@ async def test_model_connection(
"ocr",
"evaluation",
"decisions",
"extraction",
]
| None = fastapi.Body(
None,

56
litellm/proxy/hooks/key_management_event_hooks.py
View file

@ -254,39 +254,45 @@ class KeyManagementEventHooks:
litellm_changed_by: str | None = None,
) -> None:
from litellm.proxy.management_helpers.audit_logs import (
create_audit_log_for_update,
create_audit_logs_for_updates,
get_audit_log_changed_by,
is_audit_logging_enabled,
track_audit_task,
)
from litellm.proxy.proxy_server import litellm_proxy_admin_name
if not is_audit_logging_enabled():
return
for key in keys_being_deleted:
key_row = LiteLLM_VerificationToken.model_validate(key, from_attributes=True)
if key_row.token is None:
continue
_key_row = key_row.model_dump_json(exclude_none=True)
asyncio.create_task(
create_audit_log_for_update(
request_data=LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by=get_audit_log_changed_by(
litellm_changed_by=litellm_changed_by,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
),
changed_by_api_key=user_api_key_dict.token,
table_name=LitellmTableNames.KEY_TABLE_NAME,
object_id=key_row.token,
action="deleted",
updated_values="{}",
before_value=_key_row,
)
)
changed_by: Final = get_audit_log_changed_by(
litellm_changed_by=litellm_changed_by,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
key_rows: Final = (
LiteLLM_VerificationToken.model_validate(key, from_attributes=True) for key in keys_being_deleted
)
audit_rows: Final = tuple(
LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by=changed_by,
changed_by_api_key=user_api_key_dict.token,
table_name=LitellmTableNames.KEY_TABLE_NAME,
object_id=key_row.token,
action="deleted",
updated_values="{}",
before_value=key_row.model_dump_json(exclude_none=True),
)
for key_row in key_rows
if key_row.token is not None
)
if not audit_rows:
return
track_audit_task(
asyncio.create_task(
create_audit_logs_for_updates(audit_rows), name=f"{len(audit_rows)} deleted key audit logs"
)
)
@staticmethod
async def _store_virtual_key_in_secret_manager(secret_name: str, secret_token: str, team_id: str | None = None):

2
litellm/proxy/hooks/model_max_budget_limiter.py
View file

@ -374,6 +374,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
self._redis_increment_queue_lock = asyncio.Lock()
self._redis_increment_flush_lock = asyncio.Lock()
self._detached_increment_operations = None
self._retired = False
self._flush_after_retire = None
self.deployment_budget_config = None
@with_service_target("model_budgets")

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

@ -7,6 +7,7 @@ This is currently in development and not yet ready for production.
import asyncio
import binascii
import itertools
import json
import logging
import os
import uuid
@ -42,6 +43,7 @@ from litellm.caching.redis_batch import (
)
from litellm.caching.redis_cache import log_redis_failure
from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.cost_calculator import get_usage_object
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_str_from_messages,
@ -78,6 +80,7 @@ from litellm.router_utils.add_retry_fallback_headers import (
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.router_utils.ptu_shares import PTUTeamCeiling
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.extraction import ExtractionResponse
from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
@ -490,6 +493,7 @@ RESPONSES_API_CALL_TYPES: Final = ("aresponses", "responses")
EMBEDDING_API_CALL_TYPES: Final = ("aembedding", "embedding")
TEXT_COMPLETION_API_CALL_TYPES: Final = ("atext_completion", "text_completion")
RERANK_API_CALL_TYPES: Final = (CallTypes.rerank.value, CallTypes.arerank.value)
EXTRACTION_API_CALL_TYPES: Final = (CallTypes.extraction.value, CallTypes.aextraction.value)
GOOGLE_GENAI_NATIVE_CALL_TYPES: Final = (
CallTypes.generate_content.value,
CallTypes.agenerate_content.value,
@ -691,6 +695,7 @@ class RequestRateLimiterStash:
owner_litellm_call_id: str | None = None
rate_limit_response: RateLimitResponse | None = None
parallel_slot: ParallelSlotAcquisition | None = None
parallel_slot_reused_at: datetime | None = None
parallel_slot_release_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False)
reserved_tokens: int = 0
reserved_model: RateLimitedModel | None = None
@ -787,6 +792,19 @@ def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash:
return stash
def _ended_before_slot_reuse(ended_at: datetime | None, reused_at: datetime | None) -> bool:
if ended_at is None or reused_at is None or (ended_at.tzinfo is None) != (reused_at.tzinfo is None):
return False
return ended_at <= reused_at
def _callback_kwargs_mapping(kwargs: object) -> Mapping[str, object]:
try:
return _REQUEST_RATE_LIMIT_DATA.validate_python(kwargs)
except ValidationError:
return MappingProxyType({})
def get_request_stash_for_call(litellm_call_id: str | None) -> RequestRateLimiterStash | None:
stash: Final = _request_stash.get()
if stash is None:
@ -1094,6 +1112,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
or PROXY_MaxParallelRequestsHandler_v3._has_explicit_output_cap(data, call_type)
or is_embedding
or endpoint_type == EndpointType.DECISIONS
or call_type in EXTRACTION_API_CALL_TYPES
):
return
effective_cap: Final = max(capped_floor, configured_output_tokens or 0)
@ -1218,6 +1237,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if call_type in RESPONSES_API_CALL_TYPES or call_type in EMBEDDING_API_CALL_TYPES
else (None, estimable_data.get("prompt"), None)
if call_type in TEXT_COMPLETION_API_CALL_TYPES
else (None, self._extraction_input_text(_REQUEST_RATE_LIMIT_DATA.validate_python(estimable_data)), None)
if call_type in EXTRACTION_API_CALL_TYPES
else (estimable_data.get("messages"), None, None)
if call_type
else (
@ -1245,7 +1266,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
estimated_input_tokens: Final = max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0
explicit_max_tokens: Final = self._get_explicit_output_cap(data, call_type)
is_embedding: Final = self._is_embedding_request(data, call_type)
generates_no_output: Final = (
self._is_embedding_request(data, call_type) or call_type in EXTRACTION_API_CALL_TYPES
)
base_output_floor: Final = self.no_max_tokens_output_floor(min_configured_tpm_limit)
output_floor: Final = (
@ -1255,7 +1278,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
max_tokens_estimate: Final = (
0
if is_embedding or (explicit_max_tokens is None and total_chars == 0 and configured_output_tokens is None)
if generates_no_output
or (explicit_max_tokens is None and total_chars == 0 and configured_output_tokens is None)
else explicit_max_tokens
if explicit_max_tokens is not None
else configured_output_tokens
@ -1642,6 +1666,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
read_only: bool = False,
skip_tpm_check: bool = False,
parallel_slot_id: str | None = None,
skip_parallel_check: bool = False,
) -> RateLimitResponse:
"""
Check if any of the rate limit descriptors should be rate limited.
@ -1658,6 +1683,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
``reserve_tpm_tokens`` reservation path should set this to
avoid the +1-per-key Lua / in-memory increment double-charging
the tokens counter.
skip_parallel_check: If True, ``max_parallel_requests`` gauges are
neither checked nor acquired, for a call that runs inside a
request already holding its slot.
``max_parallel_requests`` descriptors are enforced by the dedicated
concurrency-gauge path (``_check_parallel_request_gauges``), never by
@ -1750,7 +1778,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if windowed_response["overall_code"] == "OVER_LIMIT":
return windowed_response
if not gauges:
if not gauges or skip_parallel_check:
return windowed_response
gauge_response: Final = await self._check_parallel_request_gauges(
@ -2013,12 +2041,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
parent_otel_span: Span | None,
*,
in_logging_callback: bool = False,
round_ended_at: datetime | None = None,
) -> None:
if stash is None:
return
async with stash.parallel_slot_release_lock:
acquisition: Final = stash.parallel_slot
if acquisition is None:
if acquisition is None or _ended_before_slot_reuse(round_ended_at, stash.parallel_slot_reused_at):
return
deferred: Final = in_logging_callback and await self._defer_parallel_slot_release(
acquisition, parent_otel_span
@ -3834,6 +3863,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if isinstance(part, (str, dict))
)
@staticmethod
def _extraction_input_text(data: Mapping[str, object]) -> str:
state: Final = data.get("state")
return f"{state if isinstance(state, str) else ''}\n{json.dumps(data.get('fields'), default=str)}"
def _estimate_precise_input_tokens(self, data: object, model: str | None, call_type: str | None = None) -> int:
"""
Model-aware input token estimate for the project ITPM reservation,
@ -3892,6 +3926,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if call_type in RERANK_API_CALL_TYPES
else (None, prompt, data.get("tools"), data.get("tool_choice"))
if call_type in TEXT_COMPLETION_API_CALL_TYPES
else (None, self._extraction_input_text(_REQUEST_RATE_LIMIT_DATA.validate_python(data)), None, None)
if call_type in EXTRACTION_API_CALL_TYPES
else (data.get("messages"), fallback_text, data.get("tools"), data.get("tool_choice"))
)
messages, selected_text, countable_tools, countable_tool_choice = selected_inputs
@ -3930,6 +3966,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
tpm_reservation_amount: int,
call_type: str | None = None,
endpoint_type: EndpointType = EndpointType.GENERIC,
reuses_request_slot: bool = False,
) -> None:
"""
Reserve project-scoped ITPM/OTPM tokens (Bedrock Mantle-style
@ -4009,7 +4046,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
parent_otel_span=user_api_key_dict.parent_otel_span,
)
stash.reservation_released = True
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
if not reuses_request_slot:
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
self._handle_rate_limit_error(
response=io_response,
descriptors=descriptors,
@ -4170,6 +4208,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
verbose_proxy_logger.debug("Inside Rate Limit Pre-Call Hook")
stash: Final = claim_request_stash_for_data(data)
reuses_request_slot: Final = call_type == CallTypes.call_mcp_tool.value and stash.parallel_slot is not None
if reuses_request_slot:
stash.parallel_slot_reused_at = self._get_current_time()
#########################################################
# Check if the call type has a specific rate limiter
@ -4214,11 +4255,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# then because _reserve_project_io_tokens_or_raise below charges
# them unconditionally and counting them here too would
# double-charge every request.
parallel_counter_keys: Final = [
self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests")
for d in descriptors
if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None
]
parallel_counter_keys: Final = (
[]
if reuses_request_slot
else [
self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests")
for d in descriptors
if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None
]
)
parallel_slot_id: Final = uuid.uuid4().hex if parallel_counter_keys else None
first_pass_descriptors: Final = (
@ -4233,6 +4278,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
parent_otel_span=user_api_key_dict.parent_otel_span,
skip_tpm_check=self.tpm_reservation_enabled,
parallel_slot_id=parallel_slot_id,
skip_parallel_check=reuses_request_slot,
)
if response["overall_code"] == "OVER_LIMIT":
@ -4242,7 +4288,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
requested_model=requested_model,
)
else:
stash.rate_limit_response = response
if not reuses_request_slot:
stash.rate_limit_response = response
if parallel_slot_id is not None:
stash.parallel_slot = ParallelSlotAcquisition(
slot_id=parallel_slot_id,
@ -4340,7 +4387,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
raw_estimate=estimated_tokens,
)
if stash.ptu_ceiling is not None and ptu_estimated_tokens > stash.ptu_ceiling.tpm_limit:
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
if not reuses_request_slot:
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
raise PTUShareExceededError(stash.ptu_ceiling, ptu_estimated_tokens)
tpm_response: Final = await self.reserve_tpm_tokens(
descriptors=descriptors,
@ -4350,7 +4398,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
if tpm_response["overall_code"] == "OVER_LIMIT":
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
if not reuses_request_slot:
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
self._handle_rate_limit_error(
response=tpm_response,
descriptors=descriptors,
@ -4401,6 +4450,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
tpm_reservation_amount=tpm_reservation_amount,
call_type=call_type,
endpoint_type=endpoint_type,
reuses_request_slot=reuses_request_slot,
)
def _create_pipeline_operations(
@ -4512,6 +4562,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
) -> Usage | ResponseAPIUsage | Mapping[str, object] | None:
if isinstance(response_obj, (Usage, ResponseAPIUsage)):
return response_obj
if isinstance(response_obj, ExtractionResponse):
return get_usage_object(response_obj)
if isinstance(
response_obj,
(ModelResponse, EmbeddingResponse, TextCompletionResponse, BaseLiteLLMOpenAIResponseObject),
@ -5143,21 +5195,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# than parsed out of the body) carry their usage in
# ``combined_usage_object`` instead, and would otherwise never charge
# the TPM window.
_usage: Usage | dict[str, object] | None = None
if isinstance(
response_obj,
(
ModelResponse,
EmbeddingResponse,
TextCompletionResponse,
BaseLiteLLMOpenAIResponseObject,
),
):
_usage = getattr(response_obj, "usage", None)
else:
_combined_usage: Final = kwargs.get("combined_usage_object")
if isinstance(_combined_usage, Usage):
_usage = _combined_usage
_combined_usage: Final = kwargs.get("combined_usage_object")
_usage: Final[Usage | dict[str, object] | None] = (
getattr(response_obj, "usage", None)
if isinstance(
response_obj,
(
ModelResponse,
EmbeddingResponse,
TextCompletionResponse,
BaseLiteLLMOpenAIResponseObject,
),
)
else get_usage_object(response_obj)
if isinstance(response_obj, ExtractionResponse)
else _combined_usage
if isinstance(_combined_usage, Usage)
else None
)
total_tokens = self._get_total_tokens_from_usage(usage=_usage, rate_limit_type=rate_limit_type)
if total_tokens == 0:
total_tokens = self._aggregate_only_total_tokens(usage=_usage)
@ -5310,8 +5365,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
try:
verbose_proxy_logger.debug("INSIDE parallel request limiter ASYNC SUCCESS LOGGING")
stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs))
await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span, in_logging_callback=True)
callback_kwargs: Final = _callback_kwargs_mapping(kwargs)
stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(callback_kwargs))
await self._release_stashed_parallel_slot(
stash,
litellm_parent_otel_span,
in_logging_callback=True,
round_ended_at=end_time
if isinstance(end_time, datetime) and not callback_kwargs.get("stream")
else None,
)
pipeline_operations: Final = self._build_success_event_pipeline_operations(
kwargs=kwargs,

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

@ -5,14 +5,14 @@ from contextlib import contextmanager
from contextvars import ContextVar
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Protocol, cast
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, cast
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.batches.batch_utils import batch_cost_is_final
from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED
from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED, BILLED_RETRY_ATTEMPT_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import (
budget_reservation_from_metadata,
@ -54,11 +54,13 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import (
)
from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # legacy module exports
_sanitize_error_information_for_spend_logs, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
deployment_params_provider,
get_request_model_access_groups,
sanitize_error_information_for_spend_logs,
should_store_prompts_and_responses_in_spend_logs,
)
from litellm.proxy.utils import BilledFailedAttempt, ProxyUpdateSpend
from litellm.types.router import Deployment
from litellm.types.utils import (
CallTypes,
LiteLLMBatch,
@ -187,6 +189,14 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter:
return proxy_logging_obj.db_spend_update_writer
DeploymentLookup: TypeAlias = Callable[[str], Deployment | None]
def _proxy_deployment(model_id: str) -> Deployment | None:
router: Final = get_llm_router()
return None if router is None else router.get_deployment(model_id=model_id)
def _billed_failed_attempts(request_data: Mapping[str, object]) -> tuple[BilledFailedAttempt, ...]:
"""The earlier billed attempts ``_failure_fields_to_lift`` stashed on the request, if any."""
attempts: Final = request_data.get("billed_failed_attempts")
@ -197,7 +207,10 @@ def _billed_failed_attempts(request_data: Mapping[str, object]) -> tuple[BilledF
def _failure_spend_rows(
request_data: Mapping[str, object], final_cost: float, late_guardrails_snapshot: object
request_data: Mapping[str, object],
final_cost: float,
late_guardrails_snapshot: object,
deployment_lookup: DeploymentLookup,
) -> tuple[tuple[float, Mapping[str, object]], ...]:
"""One ``(response_cost, kwargs)`` per spend row a failure writes: the final attempt under the
request id, then every earlier attempt the provider also billed under its own row."""
@ -209,25 +222,59 @@ def _failure_spend_rows(
return (
(final_cost, final_row),
*(
(attempt.response_cost, _attempt_spend_row_kwargs(request_data, attempt, position))
(
attempt.response_cost,
_attempt_spend_row_kwargs(
request_data, attempt, position, _attempt_deployment(attempt, deployment_lookup)
),
)
for position, attempt in enumerate(_billed_failed_attempts(request_data))
),
)
def _attempt_deployment(attempt: BilledFailedAttempt, deployment_lookup: DeploymentLookup) -> Deployment | None:
return None if attempt.deployment_id is None else deployment_lookup(attempt.deployment_id)
def _deployment_attribution(deployment: Deployment) -> Mapping[str, object]:
"""The spend-log attribution fields of the deployment that served an attempt: the model string it
was called with, its provider, its configured ``api_base`` and its model group."""
return MappingProxyType(
{
"model": deployment.litellm_params.model,
"custom_llm_provider": deployment_params_provider(deployment.litellm_params),
"api_base": deployment.litellm_params.api_base,
"model_group": deployment.model_name,
}
)
_NO_ATTRIBUTION: Final[Mapping[str, object]] = MappingProxyType({})
def _attempt_spend_row_kwargs(
request_data: Mapping[str, object], attempt: BilledFailedAttempt, position: int
request_data: Mapping[str, object], attempt: BilledFailedAttempt, position: int, deployment: Deployment | None
) -> Mapping[str, object]:
"""The final attempt's row kwargs re-keyed for one earlier billed attempt: its own
``request_id`` (``<id>-attempt-<position>``, since ``request_id`` is the spend-log primary key
and the retry count restarts after a fallback hop), its estimated usage and cost, its retry
count under ``attempted_retries`` and the deployment that billed it under ``model_info.id``."""
count under ``attempted_retries``, the deployment that billed it under ``model_info.id`` and,
when the router still knows that deployment, its model, provider, ``api_base`` and model group.
The row is stamped ``billed_retry_attempt`` so the daily spend tables keep its spend, tokens
and model usage without counting it as a request the caller sent."""
suffix: Final = f"-attempt-{position}"
attribution: Final = _NO_ATTRIBUTION if deployment is None else _deployment_attribution(deployment)
raw_params: Final = cast("Mapping[str, object]", request_data.get("litellm_params") or {}) # cast-ok: request dict
litellm_params: Final = dict(raw_params)
litellm_params: Final = {
**raw_params,
**({} if deployment is None else {"api_base": attribution["api_base"]}),
}
metadata: Final = {
**cast("Mapping[str, object]", litellm_params.get("metadata") or {}), # cast-ok: request dict
"attempted_retries": attempt.attempted_retries,
BILLED_RETRY_ATTEMPT_METADATA_KEY: True,
**({} if deployment is None else {"model_group": attribution["model_group"]}),
}
metadata.pop("previous_models", None)
if attempt.deployment_id is not None:
@ -238,6 +285,7 @@ def _attempt_spend_row_kwargs(
litellm_params["metadata"] = metadata
row: Final = {
**request_data,
**({} if deployment is None else {k: attribution[k] for k in ("model", "custom_llm_provider")}),
"litellm_params": litellm_params,
"combined_usage_object": attempt.usage,
"response_cost": attempt.response_cost,
@ -250,7 +298,11 @@ def _attempt_spend_row_kwargs(
if isinstance(standard_logging_object, dict):
payload: Final = cast("Mapping[str, object]", standard_logging_object) # cast-ok: StandardLoggingPayload dict
if isinstance(payload.get("id"), str):
row["standard_logging_object"] = {**payload, "id": f"{payload['id']}{suffix}"}
row["standard_logging_object"] = {
**payload,
"id": f"{payload['id']}{suffix}",
**({} if deployment is None else {**attribution, "model_id": attempt.deployment_id}),
}
return row
@ -262,10 +314,12 @@ class _ProxyDBLogger(CustomLogger):
turn_off_message_logging: bool = False,
message_logging: bool = True,
spend_writer: Callable[[], DBSpendUpdateWriter] = _proxy_spend_writer,
deployment_lookup: DeploymentLookup = _proxy_deployment,
) -> None:
super().__init__(turn_off_message_logging=turn_off_message_logging, message_logging=message_logging)
self.spend_event_producer = spend_event_producer
self._spend_writer: Final = spend_writer
self._deployment_lookup: Final = deployment_lookup
async def async_log_success_event(
self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime
@ -439,6 +493,7 @@ class _ProxyDBLogger(CustomLogger):
cast("Mapping[str, object]", request_data), # cast-ok: request dict, copied per row below
recovered_response_cost,
failure_snapshot_with_late_guardrails,
self._deployment_lookup,
),
)

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

@ -34,8 +34,9 @@ def freed_units(allowance: int, units: int, total_after: int) -> int:
return max(0, min(units, allowance - (total_after - units)))
def search_free_tier_key(allowance: SearchFreeRequests, now: datetime) -> str:
return f"litellm:search_free_tier:{allowance['shared_group']}:{now.strftime(_PERIOD_STAMP[allowance['period']])}"
def search_free_tier_key(provider: str, allowance: SearchFreeRequests, now: datetime) -> str:
period_stamp: Final = now.strftime(_PERIOD_STAMP[allowance["period"]])
return f"litellm:search_free_tier:{provider}:{allowance['shared_group']}:{period_stamp}"
def search_free_tier_ttl(allowance: SearchFreeRequests, now: datetime) -> int:
@ -107,7 +108,8 @@ async def apply_search_free_tier(
"apply_search_free_tier is on but the proxy has no Redis, grounding is billed at full price"
)
return response_cost
key: Final = search_free_tier_key(allowance, now)
provider: Final = sl_object["custom_llm_provider"] or model_info["litellm_provider"]
key: Final = search_free_tier_key(provider, allowance, now)
try:
total_after: Final = await counter.async_increment(key, units, ttl=search_free_tier_ttl(allowance, now))
except Exception: # noqa: BLE001 # any counter failure means full price, never a lost spend row

68
litellm/proxy/litellm_pre_call_utils.py
View file

@ -29,6 +29,7 @@ from litellm._service_logger import ServiceLogging
from litellm._uuid import uuid
from litellm.constants import (
BAGGAGE_MAPPED_TAGS_METADATA_KEY,
BILLED_RETRY_ATTEMPT_METADATA_KEY,
CALLER_SESSION_ID_METADATA_KEY,
CLIENT_OUTPUT_CEILING_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
@ -267,7 +268,7 @@ def sanitize_for_log(value: object) -> str:
from litellm.router import Router
from litellm.secret_managers.main import get_secret_bool
from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES
from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES, STREAM_USAGE_INJECTED_FIELD
from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS
from litellm.types.router import BaggageMappedTagsStamp, holds_api_key_secret_reference
from litellm.types.services import ServiceTypes
@ -426,6 +427,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
CONSUMED_REQUEST_TAGS_METADATA_KEY,
ROUTING_REQUEST_TAGS_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
BILLED_RETRY_ATTEMPT_METADATA_KEY,
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY,
"standard_logging_object",
"litellm_roi_estimator",
@ -1218,6 +1220,32 @@ def get_dynamic_logging_metadata(
return callback_settings_obj
def _callback_settings_fields(callback_settings_obj: TeamCallbackMetadata) -> Mapping[str, object]:
callback_vars: Final = callback_settings_obj.callback_vars
return {
"success_callback": callback_settings_obj.success_callback,
"failure_callback": callback_settings_obj.failure_callback,
**({**callback_vars, TRUSTED_CALLBACK_VARS_FIELD: callback_vars} if callback_vars is not None else {}),
}
def key_and_team_logging_fields(user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig) -> Mapping[str, object]:
callback_settings_obj: Final = get_dynamic_logging_metadata(
user_api_key_dict=user_api_key_dict, proxy_config=proxy_config
)
disabled_callbacks: Final = (user_api_key_dict.metadata or {}).get("litellm_disabled_callbacks")
return MappingProxyType(
{
**(_callback_settings_fields(callback_settings_obj) if callback_settings_obj is not None else {}),
**(
{"litellm_disabled_callbacks": disabled_callbacks}
if disabled_callbacks and isinstance(disabled_callbacks, list)
else {}
),
}
)
_get_dynamic_logging_metadata: Final = get_dynamic_logging_metadata
@ -2342,7 +2370,15 @@ def refresh_proxy_server_request_body_snapshot(
if not isinstance(proxy_server_request, dict):
return
_body_snapshot_exclude: Final = (
frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj", VECTOR_STORE_CALLER_METADATA_KWARG})
frozenset(
{
"secret_fields",
"proxy_server_request",
"litellm_logging_obj",
VECTOR_STORE_CALLER_METADATA_KWARG,
STREAM_USAGE_INJECTED_FIELD,
}
)
| _TRANSPORT_ONLY_CREDENTIAL_KEYS
| _CALLBACK_CREDENTIAL_KEYS
)
@ -2841,27 +2877,7 @@ async def add_litellm_data_to_request(
request=request, request_data=data, metadata_key=_metadata_variable_name
)
# Team Callbacks controls
callback_settings_obj: Final = get_dynamic_logging_metadata(
user_api_key_dict=user_api_key_dict, proxy_config=proxy_config
)
if callback_settings_obj is not None:
data["success_callback"] = callback_settings_obj.success_callback
data["failure_callback"] = callback_settings_obj.failure_callback
if callback_settings_obj.callback_vars is not None:
# unpack callback_vars in data
for k, v in callback_settings_obj.callback_vars.items():
data[k] = v
# Callbacks that must not honour request-supplied credentials read this
# proxy-owned field instead of the raw request kwargs.
data[TRUSTED_CALLBACK_VARS_FIELD] = callback_settings_obj.callback_vars
# Add disabled callbacks from key metadata
if user_api_key_dict.metadata and "litellm_disabled_callbacks" in user_api_key_dict.metadata:
disabled_callbacks: Final = user_api_key_dict.metadata["litellm_disabled_callbacks"]
if disabled_callbacks and isinstance(disabled_callbacks, list):
data["litellm_disabled_callbacks"] = disabled_callbacks
data.update(key_and_team_logging_fields(user_api_key_dict=user_api_key_dict, proxy_config=proxy_config))
# Guardrails from key/team metadata and policy engine
await move_guardrails_to_metadata(
@ -3321,12 +3337,12 @@ async def _resolve_user_provider_credentials_for_request(
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
from litellm.types.proxy.litellm_pre_call_utils import RedactedDict
from .credential_endpoints.user_provider_credentials import aget_user_provider_tokens
from .credential_endpoints.user_provider_credentials import aget_user_provider_credentials
if prisma_client is None:
return
try:
tokens: Final = await aget_user_provider_tokens(
connections: Final = await aget_user_provider_credentials(
prisma_client=prisma_client,
cache=user_api_key_cache,
user_id=user_id,
@ -3341,7 +3357,7 @@ async def _resolve_user_provider_credentials_for_request(
secret_fields: Final = data.get("secret_fields")
if not isinstance(secret_fields, dict):
return
secret_fields["user_provider_credentials"] = RedactedDict(dict(tokens))
secret_fields["user_provider_credentials"] = RedactedDict(dict(connections))
secret_fields["user_provider_credentials_user_id"] = user_id

116
litellm/proxy/management_endpoints/scim/scim_v2.py
View file

@ -11,7 +11,7 @@ from dataclasses import dataclass
from functools import partial, reduce
from itertools import chain, groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, NamedTuple, NoReturn, Protocol, TypeAlias, overload
from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protocol, TypeAlias, overload
from fastapi import (
APIRouter,
@ -53,7 +53,12 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
delete_cache_team_object,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
broadcast_evictions,
evict_and_broadcast,
evict_local,
evict_shared,
)
from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports
_safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
safe_get_request_headers,
@ -64,6 +69,7 @@ from litellm.proxy.management_endpoints.scim.scim_transformations import (
ScimTransformations,
)
from litellm.proxy.management_endpoints.team_endpoints import (
delete_validated_team_rows,
new_team,
team_member_add,
team_member_delete,
@ -104,6 +110,7 @@ from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.proxy.management_endpoints.scim_v2 import *
if TYPE_CHECKING:
from prisma import Prisma
from prisma.models import LiteLLM_VerificationToken as PrismaVerificationToken
@ -317,11 +324,15 @@ async def _check_team_exists(team_id: str) -> LiteLLM_TeamTable:
team: Final = await _table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id})
if not team:
raise HTTPException(status_code=404, detail={"error": f"Group not found with ID: {team_id}"})
raise _group_not_found(team_id)
return team
def _group_not_found(team_id: str) -> HTTPException:
return HTTPException(status_code=404, detail={"error": f"Group not found with ID: {team_id}"})
def _extract_scim_user_data(user: SCIMUser) -> ScimUserData:
"""Extract common data from SCIMUser object."""
user_email = None
@ -482,10 +493,7 @@ async def _recompute_scim_member_roles(prisma_client: PrismaClient, user_ids: It
Recompute and persist each user's global proxy role from their resulting team
membership. No-op unless scim_admin_group is configured, so a SCIM group write
that drops a member from the admin group demotes them just like the user
endpoints do, and the role is left untouched when the feature is off. The users'
rows are locked through ``USER_ROWS_LOCK_SQL`` before their teams are read, the
same lock a roster write takes, so a concurrent group push cannot change a
member's teams between the read and the role write.
endpoints do, and the role is left untouched when the feature is off.
"""
admin_group: Final = await _get_scim_admin_group()
if admin_group is None:
@ -494,28 +502,46 @@ async def _recompute_scim_member_roles(prisma_client: PrismaClient, user_ids: It
ids: Final = tuple(dict.fromkeys(user_ids))
if not ids:
return
default_role: Final = _default_scim_user_role()
async with prisma_client.tx(timeout=ROSTER_TX_TIMEOUT) as tx:
await tx.query_raw(USER_ROWS_LOCK_SQL, sorted(ids))
rows: Final = await find_many_in(user_tx_db(tx), "user_id", ids)
team_ids: Final = tuple(dict.fromkeys(chain.from_iterable(row.teams for row in rows)))
teams: Final = await find_many_in(team_tx_db(tx), "team_id", team_ids)
alias_of: Final = MappingProxyType({team.team_id: team.team_alias for team in teams})
resolved: Final = tuple(
(
row.user_id,
_resolve_scim_user_role(
[SCIMUserGroup(value=team_id, display=alias_of.get(team_id)) for team_id in row.teams],
admin_group,
default_role,
),
)
for row in rows
await _write_scim_member_roles_tx(tx, ids, admin_group=admin_group)
from litellm.proxy.proxy_server import user_api_key_cache
evict_local(cache_keys=ids, user_api_key_cache=user_api_key_cache)
await evict_shared(cache_keys=ids, cache=user_api_key_cache)
await broadcast_evictions(cache_keys=ids)
async def _write_scim_member_roles_tx(tx: "Prisma", user_ids: Sequence[str], *, admin_group: str) -> None:
"""
Write each user's role resolved from the teams they are in, on ``tx``. The users'
rows are locked through ``USER_ROWS_LOCK_SQL`` before their teams are read, the
same lock a roster write takes, so a concurrent group push cannot change a
member's teams between the read and the role write.
"""
if not user_ids:
return
default_role: Final = _default_scim_user_role()
await tx.query_raw(USER_ROWS_LOCK_SQL, sorted(user_ids))
rows: Final = await find_many_in(user_tx_db(tx), "user_id", user_ids)
team_ids: Final = tuple(dict.fromkeys(chain.from_iterable(row.teams for row in rows)))
teams: Final = await find_many_in(team_tx_db(tx), "team_id", team_ids)
alias_of: Final = MappingProxyType({team.team_id: team.team_alias for team in teams})
resolved: Final = tuple(
(
row.user_id,
_resolve_scim_user_role(
[SCIMUserGroup(value=team_id, display=alias_of.get(team_id)) for team_id in row.teams],
admin_group,
default_role,
),
)
for row in rows
)
for role in dict.fromkeys(role for _, role in resolved):
await tx.execute_raw(
_SET_USER_ROLE_SQL, role, sorted(user_id for user_id, user_role in resolved if user_role == role)
)
for role in dict.fromkeys(role for _, role in resolved):
await tx.execute_raw(
_SET_USER_ROLE_SQL, role, sorted(user_id for user_id, user_role in resolved if user_role == role)
)
class _ResolvedUserMember(NamedTuple):
@ -3051,7 +3077,8 @@ async def update_group(
dependencies=[Depends(user_api_key_auth)],
)
async def delete_group(
group_id: str = Path(..., title="Group ID"),
group_id: Annotated[str, Path(title="Group ID")],
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Delete a group according to SCIM v2 protocol
@ -3064,23 +3091,22 @@ async def delete_group(
existing_team: Final = await _check_team_exists(group_id)
raise_if_teams_hold_ptu_shares((group_id,), llm_router)
member_ids: Final = await _get_team_member_user_ids_from_team(existing_team)
# For each member, remove this team from their teams list
for member_id in member_ids:
user = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": member_id})
if user:
current_teams = user.teams or []
if group_id in current_teams:
new_teams = [t for t in current_teams if t != group_id]
await _table(UserRepository(prisma_client)).update(
where={"user_id": member_id}, data={"teams": new_teams}
)
await _recompute_scim_member_roles(prisma_client, member_ids)
# Delete team
await _table(TeamRepository(prisma_client)).delete(where={"team_id": group_id})
admin_group: Final = await _get_scim_admin_group()
deleted: Final = await delete_validated_team_rows(
teams=(LiteLLM_TeamTable.model_validate(existing_team.model_dump()),),
prisma_client=prisma_client,
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key=user_api_key_dict.api_key,
user_id=user_api_key_dict.user_id,
),
litellm_changed_by=None,
before_commit=(
None if admin_group is None else partial(_write_scim_member_roles_tx, admin_group=admin_group)
),
)
if group_id not in deleted.team_ids:
raise _group_not_found(group_id)
return Response(status_code=204)

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

@ -14,7 +14,7 @@ import copy
import json
import math
import traceback
from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence
from collections.abc import Set as AbstractSet
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
@ -34,7 +34,7 @@ from typing import (
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from pydantic import BaseModel, Field, JsonValue, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict, assert_never
import litellm
@ -102,6 +102,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
OrganizationNotFoundError,
_cache_team_object, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
allowed_route_check_inside_route,
await_inflight_team_membership_load,
cache_team_object,
can_org_access_model,
delete_cache_key_objects,
@ -112,6 +113,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
get_team_object,
get_user_object,
invalidate_team_member_spend_state,
jwt_key_mapping_cache_key,
)
from litellm.proxy.auth.auth_utils import (
enforce_batch_limits_are_admin_only,
@ -119,12 +121,19 @@ from litellm.proxy.auth.auth_utils import (
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
broadcast_evictions,
evict_and_broadcast,
evict_local,
evict_shared,
publish_auth_cache_invalidation,
)
from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
team_membership_auth_cache_key,
team_membership_reservation_cache_key,
)
from litellm.proxy.db.db_span import db_span
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy.hooks.model_max_budget_limiter import (
@ -207,7 +216,7 @@ from litellm.proxy.management_helpers.utils import (
)
from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.chunked_in import find_many_in
from litellm.repositories.chunked_in import delete_many_in, find_many_in
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
@ -417,6 +426,23 @@ class _TeamDeleteTx(AccessGroupSyncTx, Protocol):
@property
def litellm_teammembership(self) -> "TableActions[prisma_models.LiteLLM_TeamMembership]": ...
@property
def litellm_deletedteamtable(self) -> "TableActions[prisma_models.LiteLLM_DeletedTeamTable]": ...
@property
def litellm_verificationtoken(self) -> "TableActions[prisma_models.LiteLLM_VerificationToken]": ...
@property
def litellm_jwtkeymapping(self) -> "TableActions[prisma_models.LiteLLM_JWTKeyMapping]": ...
def _team_delete_tx(tx: "Prisma") -> _TeamDeleteTx:
return tx # pyright: ignore[reportReturnType] # TableActions widens the generated inputs to Mapping, as the repositories do
_LOCK_DELETED_TEAM_MEMBERS_SQL: Final = """
SELECT user_id, teams FROM "LiteLLM_UserTable" WHERE teams && $1::text[] ORDER BY user_id COLLATE "C" FOR UPDATE
"""
_STRIP_DELETED_TEAMS_FROM_USERS_SQL: Final = """
UPDATE "LiteLLM_UserTable"
@ -430,6 +456,22 @@ WHERE user_id IN (
)
"""
class _DetachedMember(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
user_id: str
teams: tuple[str, ...]
_DetachedMembers: Final = TypeAdapter(tuple[_DetachedMember, ...])
@dataclass(frozen=True, slots=True)
class DeletedTeams:
team_ids: tuple[str, ...]
_INCLUDE_MODEL_TABLE: Final = MappingProxyType({"litellm_model_table": True})
@ -4526,18 +4568,7 @@ async def delete_team(
}'
```
"""
from litellm.proxy.management_helpers.audit_logs import (
get_audit_log_changed_by,
is_audit_logging_enabled,
)
from litellm.proxy.proxy_server import (
create_audit_log_for_update,
litellm_proxy_admin_name,
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
from litellm.proxy.proxy_server import llm_router, prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
@ -4567,47 +4598,32 @@ async def delete_team(
raise_if_teams_hold_ptu_shares(data.team_ids, llm_router)
await _persist_deleted_team_records(
await delete_validated_teams(
teams=team_rows,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
deleted_teams: Final[_DeletedTeamsResult] = {"deleted_teams": data.team_ids}
return deleted_teams
# we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes
if is_audit_logging_enabled():
# make an audit log for each team deleted
for team_id in data.team_ids:
team_row: LiteLLM_TeamTable | None = await prisma_client.get_data(
team_id=team_id, table_name="team", query_type="find_unique"
)
if team_row is None:
continue
async def delete_validated_teams(
teams: Sequence[LiteLLM_TeamTable],
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None,
) -> DeletedTeams:
"""
What ``/team/delete`` does once the teams exist and the caller may delete them: every key
and BYOK model of the teams, then ``delete_validated_team_rows``.
"""
from litellm.proxy.proxy_server import (
proxy_logging_obj,
user_api_key_cache,
)
_team_row = team_row.json(exclude_none=True)
asyncio.create_task(
create_audit_log_for_update(
request_data=LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by=get_audit_log_changed_by(
litellm_changed_by=litellm_changed_by,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
),
changed_by_api_key=user_api_key_dict.api_key,
table_name=LitellmTableNames.TEAM_TABLE_NAME,
object_id=team_id,
action="deleted",
updated_values="{}",
before_value=_team_row,
)
)
)
# End of Audit logging
team_ids: Final = [team.team_id for team in teams]
## DELETE ASSOCIATED KEYS
# Fetch keys before deletion to persist them
@ -4615,7 +4631,7 @@ async def delete_team(
persist_deleted_verification_tokens,
)
keys_to_delete: Final = await _tokens_db(prisma_client).find_many(where={"team_id": {"in": data.team_ids}})
keys_to_delete: Final = await _tokens_db(prisma_client).find_many(where={"team_id": {"in": team_ids}})
jwt_mapping_cache_keys: Final = await get_jwt_key_mapping_cache_keys_for_tokens(
hashed_tokens=tuple(key.token for key in keys_to_delete),
prisma_client=prisma_client,
@ -4629,7 +4645,7 @@ async def delete_team(
litellm_changed_by=litellm_changed_by,
)
await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key")
await prisma_client.delete_data(team_id_list=team_ids, table_name="key")
if keys_to_delete:
KeyManagementEventHooks.create_key_deleted_audit_logs(
@ -4654,18 +4670,41 @@ async def delete_team(
from litellm.proxy.proxy_server import llm_router
await delete_team_models(
team_ids=data.team_ids,
team_ids=team_ids,
prisma_client=prisma_client,
llm_router=llm_router,
)
await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client)
member_ids_per_team: Final = await _resolve_deleted_team_member_user_ids(
teams=team_rows,
return await delete_validated_team_rows(
teams=teams,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
async def delete_validated_team_rows(
teams: Sequence[LiteLLM_TeamTable],
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None,
before_commit: Callable[["Prisma", tuple[str, ...]], Awaitable[None]] | None = None,
) -> DeletedTeams:
"""
Under every team's advisory lock, delete the teams still there with the keys their rostered
members hold on them, their membership rows and user references, and write their deleted-team
records; then evict caches and write audit rows. Keys of users off the roster and the teams'
models stay. ``team_ids`` names only the teams this call deleted. ``before_commit`` runs on the
delete's transaction with the members of the rosters read under the lock, so its writes commit
or roll back with the delete.
"""
from litellm.proxy.proxy_server import (
proxy_logging_obj,
user_api_key_cache,
)
team_ids: Final = [team.team_id for team in teams]
## DELETE TEAMS
# Both the delete and the reconcile sweep run under every team's advisory lock
# (TEAM_ADVISORY_LOCK_SQL, the same one /team/member_add takes before its own writes),
@ -4674,14 +4713,28 @@ async def delete_team(
# the lock before this transaction starts, in which case this sweep reaches what it wrote,
# or is still waiting on the lock, in which case its own re-read happens after this commits
# and sees the row gone before it writes anything.
delete_filter: Final[_TeamIdInFilter] = {"team_id": {"in": data.team_ids}}
async with prisma_client.tx(timeout=timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)) as tx:
for team_id in sorted(data.team_ids):
tx_budget: Final = timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)
async with prisma_client.tx(max_wait=tx_budget, timeout=tx_budget) as tx:
for team_id in sorted(team_ids):
await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id)
await tx.litellm_teamtable.delete_many(where=delete_filter)
await _sweep_deleted_team_references_tx(team_ids=data.team_ids, tx=tx)
team_tx: Final = _team_delete_tx(tx)
locked: Final = await _locked_teams_tx(team_ids=team_ids, tx=team_tx)
member_keys, jwt_mapping_cache_keys = await _delete_roster_member_keys_tx(
locked=locked,
tx=tx,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await _delete_locked_teams_tx(
locked=locked, tx=team_tx, user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by
)
detached_per_team: Final = await _sweep_deleted_team_references_tx(team_ids=team_ids, tx=team_tx)
if before_commit is not None:
await before_commit(tx, _locked_roster_user_ids(locked))
deleted_teams: Final[_DeletedTeamsResult] = {"deleted_teams": data.team_ids}
rostered_per_team: Final = await _resolve_deleted_team_member_user_ids(teams=locked, prisma_client=prisma_client)
member_ids_per_team: Final = _members_per_deleted_team(rostered_per_team, detached_per_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
@ -4689,7 +4742,7 @@ async def delete_team(
# 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,
teams=locked,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
@ -4697,44 +4750,167 @@ async def delete_team(
member_ids_per_team=member_ids_per_team,
user_api_key_cache=user_api_key_cache,
)
await _invalidate_deleted_key_cache(
keys=member_keys,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
for deleted_team in team_rows:
_schedule_deleted_team_audit_rows(
teams=locked, user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by
)
if member_keys:
KeyManagementEventHooks.create_key_deleted_audit_logs(
keys_being_deleted=member_keys,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
for deleted_team in locked:
_emit_team_members_metric(deleted_team.model_copy(update={"members_with_roles": ()}))
await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id)
return deleted_teams
return DeletedTeams(team_ids=tuple(team.team_id for team in locked))
async def _sweep_deleted_team_references(team_ids: Sequence[str], prisma_client: PrismaClient) -> None:
def _locked_roster_user_ids(locked: Sequence[LiteLLM_TeamTable]) -> tuple[str, ...]:
return tuple(dict.fromkeys(chain.from_iterable(_roster_user_ids(team) for team in locked)))
def _members_per_deleted_team(
rostered: Sequence[tuple[str, Sequence[str]]], detached: Sequence[tuple[str, Sequence[str]]]
) -> tuple[tuple[str, tuple[str, ...]], ...]:
detached_by_team: Final = MappingProxyType({team_id: user_ids for team_id, user_ids in detached})
return tuple(
(team_id, tuple(dict.fromkeys((*user_ids, *detached_by_team[team_id])))) for team_id, user_ids in rostered
)
def _roster_user_ids(team: LiteLLM_TeamTable) -> tuple[str, ...]:
return tuple(member.user_id for member in team.members_with_roles if member.user_id is not None)
def _is_roster_member_key(
key: "prisma_models.LiteLLM_VerificationToken", roster_by_team: Mapping[str, frozenset[str]]
) -> bool:
return key.team_id is not None and key.user_id in roster_by_team.get(key.team_id, frozenset())
async def _locked_teams_tx(team_ids: Sequence[str], tx: _TeamDeleteTx) -> tuple[LiteLLM_TeamTable, ...]:
rows: Final = await find_many_in(tx.litellm_teamtable, "team_id", team_ids)
return tuple(LiteLLM_TeamTable.model_validate(row.model_dump()) for row in rows)
async def _delete_roster_member_keys_tx(
locked: Sequence[LiteLLM_TeamTable],
tx: "Prisma",
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None,
) -> tuple[tuple["prisma_models.LiteLLM_VerificationToken", ...], tuple[str, ...]]:
from litellm.proxy.management_endpoints.key_management_endpoints import (
persist_deleted_verification_tokens,
)
team_tx: Final = _team_delete_tx(tx)
roster_by_team: Final = MappingProxyType({team.team_id: frozenset(_roster_user_ids(team)) for team in locked})
team_keys: Final = await find_many_in(team_tx.litellm_verificationtoken, "team_id", tuple(roster_by_team))
member_keys: Final = tuple(key for key in team_keys if _is_roster_member_key(key, roster_by_team))
if not member_keys:
return (), ()
tokens: Final = tuple(key.token for key in member_keys)
jwt_mappings: Final = await find_many_in(team_tx.litellm_jwtkeymapping, "token", tokens)
await persist_deleted_verification_tokens(
keys=member_keys, # pyright: ignore[reportArgumentType] # generated row model carries the same columns as LiteLLM_VerificationToken
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
tx=tx,
)
await delete_many_in(team_tx.litellm_verificationtoken, "token", tokens, atomicity="caller_transaction")
return member_keys, tuple(
jwt_key_mapping_cache_key(mapping.jwt_claim_name, mapping.jwt_claim_value, mapping.jwt_issuer)
for mapping in jwt_mappings
)
async def _delete_locked_teams_tx(
locked: Sequence[LiteLLM_TeamTable],
tx: _TeamDeleteTx,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None,
) -> None:
await _persist_deleted_team_records(
teams=locked,
deleted_teams=tx.litellm_deletedteamtable,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await delete_many_in(
tx.litellm_teamtable,
"team_id",
tuple(team.team_id for team in locked),
atomicity="caller_transaction",
)
def _schedule_deleted_team_audit_rows(
teams: Sequence[LiteLLM_TeamTable],
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None,
) -> None:
from litellm.proxy.management_helpers.audit_logs import (
get_audit_log_changed_by,
is_audit_logging_enabled,
)
from litellm.proxy.proxy_server import (
create_audit_log_for_update,
litellm_proxy_admin_name,
)
if not is_audit_logging_enabled():
return
for team in teams:
asyncio.create_task(
create_audit_log_for_update(
request_data=LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by=get_audit_log_changed_by(
litellm_changed_by=litellm_changed_by,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
),
changed_by_api_key=user_api_key_dict.api_key,
table_name=LitellmTableNames.TEAM_TABLE_NAME,
object_id=team.team_id,
action="deleted",
updated_values="{}",
before_value=team.model_dump_json(exclude_none=True),
)
)
)
async def _sweep_deleted_team_references_tx(
team_ids: Sequence[str], tx: _TeamDeleteTx
) -> tuple[tuple[str, tuple[str, ...]], ...]:
"""
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.
One statement filters the array in place 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. It locks the user rows in the order `USER_ROWS_LOCK_SQL` takes them, so it
never waits on a SCIM group push in a cycle.
Strip the deleted team ids from every user and team-membership row still referencing them, on
the transaction holding the teams' advisory locks. User rows are locked first in
``USER_ROWS_LOCK_SQL`` order so the sweep never deadlocks with a SCIM group push, and one
statement filters each array in place so a team a concurrent ``/team/member_add`` appended is
kept. Returns, per team, the users read under the lock as referencing it.
"""
_ = await prisma_client.db.execute_raw(_STRIP_DELETED_TEAMS_FROM_USERS_SQL, list(team_ids))
_ = await _team_membership_db(prisma_client).delete_many(where=_TeamIdInFilter(team_id={"in": tuple(team_ids)}))
async def _sweep_deleted_team_references_tx(team_ids: Sequence[str], tx: _TeamDeleteTx) -> None:
"""Same sweep as `_sweep_deleted_team_references`, run on the transaction that holds
every id's advisory lock and deletes the team rows, so it commits or rolls back with them."""
locked: Final = _DetachedMembers.validate_python(await tx.query_raw(_LOCK_DELETED_TEAM_MEMBERS_SQL, list(team_ids)))
_ = await tx.execute_raw(_STRIP_DELETED_TEAMS_FROM_USERS_SQL, list(team_ids))
membership_filter: Final[_TeamIdInFilter] = {"team_id": {"in": tuple(team_ids)}}
_ = await tx.litellm_teammembership.delete_many(where=membership_filter)
return tuple(
(team_id, tuple(member.user_id for member in locked if team_id in member.teams)) for team_id in team_ids
)
async def _invalidate_deleted_key_cache(
@ -4791,17 +4967,20 @@ async def _evict_deleted_team_member_cache(
member_user_ids: Sequence[str],
user_api_key_cache: UserApiKeyCache,
) -> None:
await evict_and_broadcast(cache_keys=tuple(member_user_ids), user_api_key_cache=user_api_key_cache)
await asyncio.gather(
*(
invalidate_team_member_spend_state(
user_id=user_id,
team_id=team_id,
user_api_key_cache=user_api_key_cache,
)
for user_id in member_user_ids
)
_ = await asyncio.gather(
*(await_inflight_team_membership_load(user_id=user_id, team_id=team_id) for user_id in member_user_ids)
)
cache_keys: Final = tuple(member_cache_keys(team_id, member_user_ids))
evict_local(cache_keys=cache_keys, user_api_key_cache=user_api_key_cache)
await evict_shared(cache_keys=cache_keys, cache=user_api_key_cache)
await broadcast_evictions(cache_keys=cache_keys)
def member_cache_keys(team_id: str, user_ids: Sequence[str]) -> Iterator[str]:
for user_id in user_ids:
yield user_id
yield team_membership_auth_cache_key(team_id=team_id, user_id=user_id)
yield team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
async def _resolve_deleted_team_member_user_ids(
@ -4832,7 +5011,7 @@ async def _deleted_team_member_user_ids(team: LiteLLM_TeamTable, prisma_client:
def _transform_teams_to_deleted_records(
teams: list[LiteLLM_TeamTable],
teams: Sequence[LiteLLM_TeamTable],
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None = None,
) -> list[dict[str, object]]:
@ -4881,17 +5060,17 @@ def _transform_teams_to_deleted_records(
async def _save_deleted_team_records(
records: list[dict[str, object]],
prisma_client: PrismaClient,
deleted_teams: "TableActions[prisma_models.LiteLLM_DeletedTeamTable]",
) -> None:
"""Save deleted team records to the database."""
if not records:
return
await _deleted_team_db(prisma_client).create_many(data=records)
await deleted_teams.create_many(data=records)
async def _persist_deleted_team_records(
teams: list[LiteLLM_TeamTable],
prisma_client: PrismaClient,
teams: Sequence[LiteLLM_TeamTable],
deleted_teams: "TableActions[prisma_models.LiteLLM_DeletedTeamTable]",
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None = None,
) -> None:
@ -4903,7 +5082,7 @@ async def _persist_deleted_team_records(
)
await _save_deleted_team_records(
records=records,
prisma_client=prisma_client,
deleted_teams=deleted_teams,
)

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

@ -55,6 +55,7 @@ from litellm.constants import (
CLI_SSO_CLAIM_MAX_SCALAR_LENGTH,
CLI_SSO_SESSION_CACHE_KEY_PREFIX,
CLI_SSO_SESSION_TTL_SECONDS,
CP_RETURN_TO_COOKIE_PREFIX,
LITELLM_CLI_SOURCE_IDENTIFIER,
LITELLM_UI_SESSION_DURATION,
MAX_SPENDLOG_ROWS_TO_QUERY,
@ -2285,9 +2286,7 @@ async def auth_callback(request: Request, state: str | None = None):
sso_assertion=sso_assertion,
)
# Control-plane cross-origin: read return_to from cookie.
# Starlette's cookie_parser already handles RFC 2109 unquoting.
cp_return_to: Final[str | None] = request.cookies.get("litellm_cp_return_to")
cp_return_to: Final[str | None] = read_return_to_cookie(request)
return await SSOAuthenticationHandler.get_redirect_response_from_openid(
result=result,
@ -3090,10 +3089,11 @@ def _persist_return_to_cookie(response: Response, return_to: str | None, request
safe: Final = _is_same_origin_return_path(return_to) or SSOAuthenticationHandler.validate_return_to(return_to)
except HTTPException:
return # a non-matching absolute return_to is ignored, never blocks sign-in
if safe:
sealed: Final = _seal_return_to(return_to) if safe else None
if sealed is not None:
response.set_cookie(
key="litellm_cp_return_to",
value=return_to,
value=sealed,
max_age=600,
httponly=True,
samesite="lax",
@ -3101,6 +3101,45 @@ def _persist_return_to_cookie(response: Response, return_to: str | None, request
)
def _seal_return_to(return_to: str) -> str | None:
"""``return_to`` sealed under the salt key as an opaque cookie value. The plaintext URL carries the
client's loopback ``redirect_uri`` (``http://127.0.0.1:<port>/...`` for Codex or Claude Code), which
edge firewalls read as an SSRF probe and block before the proxy sees the request. None when there
is no salt or master key yet, so the login entrypoint still renders."""
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_bearer_token
try:
return encrypt_bearer_token(return_to, prefix=CP_RETURN_TO_COOKIE_PREFIX)
except ValueError:
return None
def _clear_rejected_return_to_cookie(response: Response, request: Request) -> None:
"""Expire a ``litellm_cp_return_to`` that ``read_return_to_cookie`` rejected (plaintext, stale, or sealed
under another key) so the browser stops replaying it until it expires on its own."""
if "litellm_cp_return_to" in request.cookies:
response.delete_cookie("litellm_cp_return_to")
def _dashboard_redirect_response(request: Request, jwt_token: str, litellm_dashboard_ui: str) -> RedirectResponse:
"""The post-SSO dashboard landing, used when there is no honored ``return_to``. Sets the session
cookie and expires a ``litellm_cp_return_to`` the reader rejected."""
verbose_proxy_logger.info("Redirecting to %s", litellm_dashboard_ui)
redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
set_session_token_cookie(redirect_response, request, jwt_token)
_clear_rejected_return_to_cookie(redirect_response, request)
return redirect_response
def read_return_to_cookie(request: Request) -> str | None:
"""The ``return_to`` this proxy sealed into ``litellm_cp_return_to``, or None when the cookie is
missing, was not sealed here, or was tampered with."""
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_bearer_token
sealed: Final = request.cookies.get("litellm_cp_return_to")
return None if sealed is None else decrypt_bearer_token(sealed, prefix=CP_RETURN_TO_COOKIE_PREFIX)
def _is_blank_sso_user_id(user_id: object) -> bool:
return isinstance(user_id, str) and not user_id.strip()
@ -3976,10 +4015,7 @@ class SSOAuthenticationHandler:
if user_id is not None and isinstance(user_id, str):
litellm_dashboard_ui += "?login=success"
verbose_proxy_logger.info("Redirecting to %s", litellm_dashboard_ui)
redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
set_session_token_cookie(redirect_response, request, jwt_token)
return redirect_response
return _dashboard_redirect_response(request, jwt_token, litellm_dashboard_ui)
@staticmethod
@with_service_target(SSO_SESSIONS_TARGET)

61
litellm/proxy/management_helpers/audit_logs.py
View file

@ -4,6 +4,7 @@ Functions to create audit logs for LiteLLM Proxy
import asyncio
import json
from collections.abc import Awaitable, Sequence
from datetime import datetime, timezone
from typing import Final, TypeVar
@ -258,12 +259,19 @@ async def create_object_audit_log(
)
async def _insert_audit_log_row(repository: AuditLogRepository, request_data: LiteLLM_AuditLogs) -> None:
async def _insert_audit_logs(insert: Awaitable[object], description: str) -> None:
try:
await repository.table.create(data={**request_data.model_dump(exclude_none=True)})
await insert
except Exception as e:
# [Non-Blocking Exception. Do not allow blocking LLM API call]
verbose_proxy_logger.error("Failed Creating audit log %s", e)
verbose_proxy_logger.error("Failed Creating %s %s", description, e)
def _json_encoded(request_data: LiteLLM_AuditLogs) -> LiteLLM_AuditLogs:
values: Final = (("updated_values", request_data.updated_values), ("before_value", request_data.before_value))
return request_data.model_copy(
update={field: json.dumps(value) for field, value in values if isinstance(value, dict)}
)
async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
@ -279,15 +287,10 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
return
verbose_proxy_logger.debug("creating audit log for %s", request_data)
if isinstance(request_data.updated_values, dict):
request_data.updated_values = json.dumps(request_data.updated_values)
if isinstance(request_data.before_value, dict):
request_data.before_value = json.dumps(request_data.before_value)
encoded: Final = _json_encoded(request_data)
# Dispatch to external audit log callbacks regardless of DB availability
await _dispatch_audit_log_to_callbacks(request_data)
await _dispatch_audit_log_to_callbacks(encoded)
if prisma_client is None:
verbose_proxy_logger.error("prisma_client is None, cannot write audit log to DB")
@ -295,7 +298,41 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
await track_audit_task(
asyncio.create_task(
_insert_audit_log_row(AuditLogRepository(prisma_client), request_data),
name=_audit_task_name(request_data, "the database"),
_insert_audit_logs(
AuditLogRepository(prisma_client).table.create(data={**encoded.model_dump(exclude_none=True)}),
"audit log",
),
name=_audit_task_name(encoded, "the database"),
)
)
async def create_audit_logs_for_updates(rows: Sequence[LiteLLM_AuditLogs]) -> None:
"""Create audit logs for many objects with one insert."""
if not rows or not is_audit_logging_enabled():
return
from litellm.proxy.proxy_server import premium_user, prisma_client
if premium_user is not True:
return
encoded: Final = tuple(_json_encoded(row) for row in rows)
for row in encoded:
await _dispatch_audit_log_to_callbacks(row)
if prisma_client is None:
verbose_proxy_logger.error("prisma_client is None, cannot write audit log to DB")
return
await track_audit_task(
asyncio.create_task(
_insert_audit_logs(
AuditLogRepository(prisma_client).table.create_many(
data=tuple(row.model_dump(exclude_none=True) for row in encoded)
),
f"{len(encoded)} audit logs",
),
name=f"{len(encoded)} audit logs to the database",
)
)

68
litellm/proxy/management_helpers/model_group_custom_fields.py Normal file
View file

@ -0,0 +1,68 @@
import math
from collections.abc import Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import Final, TypeAlias
from pydantic import TypeAdapter
CustomMetadataValue: TypeAlias = str | int | float | bool
_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[object, object])
_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
def _string_tags(tags: object) -> Iterator[str]:
if isinstance(tags, list):
yield from (tag for tag in _OBJECT_LIST_ADAPTER.validate_python(tags) if isinstance(tag, str))
def _is_non_finite(value: object) -> bool:
return isinstance(value, float) and not math.isfinite(value)
def _scalar_metadata(metadata: object) -> Iterator[tuple[str, CustomMetadataValue]]:
if isinstance(metadata, dict):
yield from (
(key, value)
for key, value in _OBJECT_MAPPING_ADAPTER.validate_python(metadata).items()
if isinstance(key, str) and isinstance(value, (str, int, float, bool)) and not _is_non_finite(value)
)
def _custom_tag_values(model_infos: Sequence[Mapping[object, object]]) -> Iterator[str]:
for model_info in model_infos:
yield from _string_tags(model_info.get("custom_tags"))
def _custom_metadata_entries(
model_infos: Sequence[Mapping[object, object]],
) -> Iterator[tuple[str, CustomMetadataValue]]:
for model_info in model_infos:
yield from _scalar_metadata(model_info.get("custom_metadata"))
def _merged_custom_tags(model_infos: Sequence[Mapping[object, object]]) -> tuple[str, ...] | None:
return tuple(dict.fromkeys(_custom_tag_values(model_infos))) or None
def _merged_custom_metadata(
model_infos: Sequence[Mapping[object, object]],
) -> Mapping[str, CustomMetadataValue] | None:
entries: Final = tuple(_custom_metadata_entries(model_infos))
first_seen: Final = MappingProxyType(dict(reversed(entries)))
return MappingProxyType({key: first_seen[key] for key in dict.fromkeys(key for key, _ in entries)}) or None
def _model_info(deployment: Mapping[str, object]) -> Mapping[object, object]:
model_info: Final = deployment.get("model_info")
return _OBJECT_MAPPING_ADAPTER.validate_python(model_info) if isinstance(model_info, dict) else MappingProxyType({})
def model_group_custom_fields(deployments: Sequence[Mapping[str, object]]) -> Mapping[str, object]:
model_infos: Final = tuple(_model_info(deployment) for deployment in deployments)
return MappingProxyType(
{
"custom_tags": _merged_custom_tags(model_infos),
"custom_metadata": _merged_custom_metadata(model_infos),
}
)

14
litellm/proxy/management_helpers/team_roster_sync.py
View file

@ -1,6 +1,6 @@
import asyncio
import json
from collections.abc import Iterator, Mapping, Sequence
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import timedelta
from types import MappingProxyType
@ -23,8 +23,6 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
)
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
team_membership_auth_cache_key,
team_membership_reservation_cache_key,
)
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy.management_endpoints.key_management_endpoints import (
@ -33,6 +31,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
from litellm.proxy.management_endpoints.team_endpoints import (
_emit_team_members_metric, # pyright: ignore[reportPrivateUsage] # same gauge /team/member_add moves
_schedule_team_membership_audit_log, # pyright: ignore[reportPrivateUsage] # same audit row /team/member_add writes
member_cache_keys,
)
from litellm.proxy.management_helpers.access_group_team_sync import (
DETACH_TEAM_FROM_USERS_SQL,
@ -277,7 +276,7 @@ async def _settle_caches(
evicted: Final = (
*deleted_tokens,
*removal.jwt_mapping_cache_keys,
*_member_cache_keys(team.team_id, changed_user_ids),
*member_cache_keys(team.team_id, changed_user_ids),
)
evict_local(cache_keys=evicted, user_api_key_cache=user_api_key_cache)
await evict_shared(cache_keys=evicted, cache=user_api_key_cache)
@ -292,13 +291,6 @@ async def _settle_caches(
await broadcast_evictions(cache_keys=evicted)
def _member_cache_keys(team_id: str, user_ids: Sequence[str]) -> Iterator[str]:
for user_id in user_ids:
yield user_id
yield team_membership_auth_cache_key(team_id=team_id, user_id=user_id)
yield team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
async def sync_team_roster(
prisma_client: PrismaClient,
team_id: str,

4
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -880,8 +880,8 @@ def _build_passthrough_failure_request_payload(
@dataclass(frozen=True, slots=True)
class _TeamCallbackWiring:
success_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg
failure_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg
success_callbacks: "list[str | Callable[..., object] | CustomLogger] | None" = None # mutable-ok: Logging init arg
failure_callbacks: "list[str | Callable[..., object] | CustomLogger] | None" = None # mutable-ok: Logging init arg
logging_kwargs: dict[str, str | dict[str, str]] | None = None # mutable-ok: Logging.__init__ arg

26
litellm/proxy/proxy_server.py
View file

@ -166,6 +166,7 @@ from litellm.proxy.common_utils.token_count_settings import (
resolve_token_count_settings,
)
from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog
from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields
from litellm.router_utils.access_windows import access_windows_config_error
from litellm.router_utils.add_retry_fallback_headers import (
get_fallback_errors_from_headers,
@ -975,6 +976,7 @@ from litellm.secret_managers.main import (
)
from litellm.tracing.config import is_lens_tracing_enabled
from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingArgs
from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD
from litellm.types.llms.anthropic import (
AnthropicMessagesRequest,
AnthropicResponse,
@ -17747,10 +17749,7 @@ def get_model_group_info(
) -> list[ModelGroupInfoProxy]:
model_groups: Final[list[ModelGroupInfoProxy]] = []
unique_models: Final = []
for model in all_models_str:
if model not in unique_models:
unique_models.append(model)
unique_models: Final = tuple(dict.fromkeys(all_models_str))
for model in unique_models:
if model_group is not None and model_group != model:
@ -17759,7 +17758,18 @@ def get_model_group_info(
_model_group_info = llm_router.get_model_group_info(model_group=model, team_id=team_id)
if _model_group_info is not None:
model_groups.append(ModelGroupInfoProxy.model_validate(_model_group_info.model_dump()))
model_groups.append(
ModelGroupInfoProxy.model_validate(
{
**_model_group_info.model_dump(),
**model_group_custom_fields(
llm_router.get_model_list_of_routed_group(
resolve_model_group_alias(llm_router.model_group_alias, model) or model
)
),
}
)
)
else:
model_group_info = ModelGroupInfoProxy(
model_group=model,
@ -18210,6 +18220,7 @@ async def async_queue_request(
try:
data = await request.json()
data.pop("_litellm_strip_stream_usage", None)
data.pop(STREAM_USAGE_INJECTED_FIELD, None)
# Include original request and headers in the data
data["proxy_server_request"] = {
@ -18400,6 +18411,7 @@ async def login(request: Request):
# _is_same_origin_return_path (strictly relative path) so it can never be an open redirect, and the
# one-shot cookie is cleared after use.
from litellm.proxy.management_endpoints.ui_sso import (
read_return_to_cookie,
set_session_token_cookie,
sso_return_to_redirect,
)
@ -18411,7 +18423,7 @@ async def login(request: Request):
# return_to is honored for EVERY sign-in branch, so the write and read sets cannot diverge: it
# sets the token cookie on the same-origin arm and hands off via a one-time login code on the
# cross-origin arm, and clears the one-shot cookie in both.
cp_return_to: Final = request.cookies.get("litellm_cp_return_to")
cp_return_to: Final = read_return_to_cookie(request)
if cp_return_to:
try:
resumed = await sso_return_to_redirect( # rebind-ok: pre-existing rebinding on a rename-only line
@ -18435,7 +18447,7 @@ async def login(request: Request):
# Create redirect response with cookie
redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
set_session_token_cookie(redirect_response, request, jwt_token)
if cp_return_to:
if "litellm_cp_return_to" in request.cookies:
redirect_response.delete_cookie(key="litellm_cp_return_to")
return redirect_response

10
litellm/proxy/public_endpoints/provider_create_fields.json
View file

@ -1731,6 +1731,16 @@
"field_type": "password",
"options": null,
"default_value": null
},
{
"key": "github_copilot_host",
"label": "GitHub host",
"placeholder": "github.com",
"tooltip": "Leave empty for github.com. For GHE.com data residency enter your SUBDOMAIN.ghe.com host",
"required": false,
"field_type": "text",
"options": null,
"default_value": null
}
],
"default_model_placeholder": "gpt-3.5-turbo"

10
litellm/proxy/public_endpoints/public_endpoints.py
View file

@ -22,6 +22,7 @@ from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
from litellm.proxy._types import (
CommonProxyErrors,
)
from litellm.proxy.db.health_check_latest import latest_by_model_name
from litellm.proxy.utils import get_custom_url
from litellm.repositories.table_repositories import ClaudeCodePluginRepository
from litellm.router_strategy.complexity_router.fuse_presets import FusePresetCatalog, get_fuse_presets
@ -295,13 +296,8 @@ async def public_model_hub():
if prisma_client is not None:
try:
latest_checks: Final = await prisma_client.get_all_latest_health_checks()
for check in latest_checks:
key = check.model_id if check.model_id else check.model_name
if key:
health_check_dict = convert_health_check_to_dict(check)
health_checks_map[key] = health_check_dict
if check.model_name:
health_checks_map[check.model_name] = health_check_dict
for model_name, check in latest_by_model_name(latest_checks).items():
health_checks_map[model_name] = convert_health_check_to_dict(check)
except Exception:
pass

5
litellm/proxy/public_endpoints/public_v1/model_hub.py
View file

@ -12,6 +12,7 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.health_check_latest import latest_by_model_name
from litellm.proxy.list_api.common import PROBLEM_TYPE_BASE, ManagementProblem
from litellm.proxy.list_api.in_memory import Cells, InMemoryListExecutor
from litellm.proxy.list_api.list_framework import (
@ -59,12 +60,12 @@ class PrismaHealthSnapshotLookup:
checks: Final = await self.prisma_client.get_latest_health_checks_for_models(model_groups)
return MappingProxyType(
{
check.model_name: HealthSnapshot(
model_group: HealthSnapshot(
status=check.status,
response_time_ms=check.response_time_ms,
checked_at=check.checked_at.isoformat() if check.checked_at else None,
)
for check in checks
for model_group, check in latest_by_model_name(checks).items()
}
)

61
litellm/proxy/route_llm_request.py
View file

@ -1,8 +1,13 @@
import asyncio
from collections.abc import Mapping
import hashlib
import json
from collections import OrderedDict
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal
from uuid import uuid4
import httpx
from fastapi import HTTPException, status
@ -10,6 +15,7 @@ from pydantic import BeforeValidator, TypeAdapter, ValidationError
import litellm
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.router_utils.common_utils import is_proxy_admin_request
@ -54,7 +60,34 @@ _ROUTER_ONLY_SETTINGS: Final = frozenset(
)
def _route_user_config_request(data: dict, route_type: str):
class UserConfigRouters:
def __init__(self, max_size: int) -> None:
self._max_size: Final = max_size
self._routers: Final[OrderedDict[str, LitellmRouter]] = OrderedDict() # mutable-ok: LRU evicts in place
def get(self, config: Mapping[str, object], build: Callable[[], LitellmRouter]) -> LitellmRouter:
key: Final = hashlib.sha256(json.dumps(config, sort_keys=True, default=str).encode()).hexdigest()
cached: Final = self._routers.get(key)
if cached is not None:
self._routers.move_to_end(key)
return cached
router: Final = build()
router.serve_through_request_logging()
self._routers[key] = router
if len(self._routers) > self._max_size:
self._routers.popitem(last=False)[1].discard()
return router
def clear(self) -> None:
for router in self._routers.values():
router.discard()
self._routers.clear()
user_config_routers: Final = UserConfigRouters(max_size=32)
async def _route_user_config_request(data: dict, route_type: str):
"""Route a request using the user-provided router config."""
router_config: Final = data.pop("user_config")
@ -63,10 +96,24 @@ def _route_user_config_request(data: dict, route_type: str):
valid_args: Final = litellm.Router.get_valid_args()
filtered_config: Final = {k: v for k, v in router_config.items() if k in valid_args}
user_router: Final = litellm.Router(**filtered_config)
ret_val: Final = getattr(user_router, f"{route_type}")(**data)
user_router.discard()
return ret_val
logging_obj: Final = _request_logging_obj(JSON_OBJECT_ADAPTER.validate_python(data), route_type)
cached_router: Final = user_config_routers.get(filtered_config, lambda: litellm.Router(**filtered_config))
cached_router.bind_to_request_logging(logging_obj)
return await getattr(cached_router, f"{route_type}")(**{**data, "litellm_logging_obj": logging_obj})
def _request_logging_obj(data: Mapping[str, object], route_type: str) -> LiteLLMLoggingObj:
arrived_with: Final = data.get("litellm_logging_obj")
if isinstance(arrived_with, LiteLLMLoggingObj):
return arrived_with
logging_obj, _ = litellm.utils.function_setup(
original_function=route_type,
rules_obj=litellm.utils.Rules(),
start_time=datetime.now(), # noqa: DTZ005 # Logging subtracts it from naive end times
is_async_call=True,
**{"litellm_call_id": str(uuid4()), **data},
)
return logging_obj
def _is_a2a_agent_model(model_name: object) -> bool:
@ -111,6 +158,7 @@ ROUTE_ENDPOINT_MAPPING: Final = {
"asearch": "/search",
"adecisions": "/decisions",
"asystemone": "/systemone",
"aextraction": "/extraction",
"avideo_generation": "/videos",
"avideo_list": "/videos",
"avideo_status": "/videos/{video_id}",
@ -573,6 +621,7 @@ RouteType = Literal[
"asearch",
"adecisions",
"asystemone",
"aextraction",
"avideo_generation",
"avideo_list",
"avideo_status",

13
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -15,6 +15,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
BILLED_RETRY_ATTEMPT_METADATA_KEY,
CALLER_SESSION_ID_METADATA_KEY,
CLI_SESSION_KEY_PREFIX,
EMPTY_MAPPING,
@ -187,6 +188,7 @@ _STAMPED_METADATA_KEYS: Final = frozenset(
"used_client_oauth_token",
KUBERNETES_POD_ROUTING_KEY,
"litellm_roi_estimator",
BILLED_RETRY_ATTEMPT_METADATA_KEY,
"user_api_key_user_email",
)
)
@ -249,6 +251,7 @@ def _get_spend_logs_metadata(
guardrail_information=None,
internal_call_origin=None,
litellm_roi_estimator=False,
billed_retry_attempt=False,
eval_information=None,
cold_storage_object_key=cold_storage_object_key,
litellm_overhead_time_ms=None,
@ -270,6 +273,8 @@ def _get_spend_logs_metadata(
)
verbose_proxy_logger.debug("getting payload for SpendLogs, available keys in metadata: %s", [*metadata])
retry_key: Final = BILLED_RETRY_ATTEMPT_METADATA_KEY
billed_retry_attempt: Final = retry_key in metadata and metadata[retry_key] is True
# Filter the metadata dictionary to include only the specified keys
clean_metadata: Final = SpendLogsMetadata(
**MappingProxyType(
@ -283,6 +288,7 @@ def _get_spend_logs_metadata(
used_client_oauth_token=used_client_oauth_token,
kubernetes_pod_routing=kubernetes_pod_routing,
litellm_roi_estimator=metadata.get("litellm_roi_estimator") is True,
billed_retry_attempt=billed_retry_attempt,
user_api_key_user_email=_stored_user_email(metadata),
)
_raw_key: Final = clean_metadata.get("user_api_key")
@ -453,7 +459,12 @@ def _sl_attribution_fallback(
def _deployment_provider(deployment: DeploymentTypedDict) -> str | None:
litellm_params: Final = LiteLLM_Params.model_validate(deployment["litellm_params"])
litellm_params: Final[LiteLLM_Params] = LiteLLM_Params.model_validate(deployment["litellm_params"])
return deployment_params_provider(litellm_params)
def deployment_params_provider(litellm_params: LiteLLM_Params) -> str | None:
"""The provider a deployment's ``litellm_params`` authenticate against, None when its model cannot be resolved."""
if litellm.LiteLLMProxyChatConfig.should_use_litellm_proxy_by_default(litellm_params=litellm_params):
return LlmProviders.LITELLM_PROXY.value
declared: Final = declared_authenticating_provider(litellm_params.model, litellm_params.custom_llm_provider)

159
litellm/proxy/utils.py
View file

@ -91,7 +91,7 @@ from litellm.proxy.common_utils.openai_error_payload import (
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.model_listing import ModelInfoResponse
from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypes, CallTypesLiteral, CostPerToken, ModelInfo, Usage
from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypes, CallTypesLiteral, ModelInfo, Usage
try:
from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
@ -158,7 +158,8 @@ from litellm.litellm_core_utils.core_helpers import (
)
from litellm.litellm_core_utils.duration_parser import get_budget_window_start
from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int
from litellm.litellm_core_utils.litellm_logging import Logging, deployment_pricing_model_info
from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, Logging, deployment_pricing_model_info
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.served_output_texts import (
@ -268,6 +269,7 @@ from litellm.repositories.verification_token_repository import (
VerificationTokenRepository,
)
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.router_utils.deployment_pricing import pricing_models
from litellm.secret_managers.main import str_to_bool
from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES
from litellm.types.llms.openai import ResponsesAPIResponse
@ -1159,6 +1161,43 @@ def _timed_out_earlier_records(
return tuple(record for identity, record in by_identity.items() if identity != final_identity)
def _record_pricing(
model_call_details: Mapping[str, object], record: Mapping[str, object], base_model: str | None
) -> tuple[Mapping[str, object], str | None]:
"""The ``(model_call_details, base_model)`` to price one earlier attempt with: the record's own
deployment model and ``base_model`` when it served a different model than the final attempt
(the provider is then inferred from that model), else the final attempt's."""
record_model: Final = record.get("model")
if not isinstance(record_model, str) or record_model == model_call_details.get("model"):
return model_call_details, base_model
record_base_model: Final = record.get("base_model")
return (
{**model_call_details, "model": record_model, "custom_llm_provider": None},
record_base_model if isinstance(record_base_model, str) else None,
)
def _billed_earlier_attempt(
model_call_details: Mapping[str, object],
record: Mapping[str, object],
index: int,
estimated_usage: Usage,
base_model: str | None,
) -> BilledFailedAttempt:
pricing_details, pricing_base_model = _record_pricing(model_call_details, record, base_model)
return BilledFailedAttempt(
attempted_retries=_attempt_number(record, index),
deployment_id=_record_deployment_id(record),
usage=estimated_usage,
response_cost=_estimated_input_cost(
pricing_details,
estimated_usage.prompt_tokens,
deployment_id=_record_deployment_id(record),
base_model=pricing_base_model,
),
)
def _billed_earlier_attempts(
model_call_details: Mapping[str, object],
litellm_params: Mapping[str, object] | None,
@ -1172,54 +1211,45 @@ def _billed_earlier_attempts(
request's own deployment and retry count (``attempted_retries`` restarts at 0 after a fallback
hop, so the count alone cannot tell the groups apart). An earlier attempt that ended in a
``Timeout`` had the same prompt accepted by its deployment, so it is billed like the final
attempt, at that deployment's input rate (else the cost map's rate for ``base_model`` or the
model). Only the exception class survives on the record, so an earlier connect timeout is
billed too; any other failure class is not."""
attempt, at that deployment's input rate (else the cost map's rate for the record's own
``base_model`` or model, which a fallback hop can make a different model than the final one).
Only the exception class survives on the record, so an earlier connect timeout is billed too;
any other failure class is not."""
final_attempt: Final = _retry_metadata_bucket(litellm_params).get("attempted_retries")
final_identity: Final = (deployment_id, final_attempt if isinstance(final_attempt, int) else None)
return tuple(
BilledFailedAttempt(
attempted_retries=_attempt_number(record, index),
deployment_id=_record_deployment_id(record),
usage=estimated_usage,
response_cost=_estimated_input_cost(
model_call_details,
estimated_usage.prompt_tokens,
deployment_id=_record_deployment_id(record),
base_model=base_model,
),
)
_billed_earlier_attempt(model_call_details, record, index, estimated_usage, base_model)
for index, record in enumerate(_timed_out_earlier_records(litellm_params, final_identity))
)
def _deployment_cost_per_token(deployment_id: str | None, model: str) -> CostPerToken | None:
"""The pricing the router registered for this deployment under its ``model_info.id``, read
through ``deployment_pricing_model_info`` exactly as the success path does; None means the
public cost map. The id is ``Logging.get_router_model_id()`` (the logging object's own
``litellm_params``), never the flattened ``model_call_details`` copy a request body key can
overwrite."""
model_info: Final = deployment_pricing_model_info(model_id=deployment_id, deployment_model=model)
if model_info is None:
return None
return CostPerToken(
input_cost_per_token=model_info["input_cost_per_token"] or 0.0,
output_cost_per_token=model_info["output_cost_per_token"] or 0.0,
)
def _deployment_priced_input_cost(
model: str, custom_llm_provider: str | None, prompt_tokens: int, deployment_pricing: ModelInfo
) -> float:
"""Input price for the estimated prompt at the deployment's registered rates, read through
``deployment_pricing_model_info`` and priced by ``generic_cost_per_token`` the way the success
path prices them, so a long-context tier the deployment declares (``input_cost_per_token_above_200k_tokens``)
is honoured instead of flattened to the standard rate. An entry it cannot price keeps the row at $0."""
usage: Final = Usage(prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens)
try:
prompt_cost, _ = generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider=custom_llm_provider or str(deployment_pricing.get("litellm_provider") or ""),
model_info=deployment_pricing,
)
except Exception as exc: # noqa: BLE001 # a malformed registered entry must not drop the failure row
verbose_proxy_logger.debug("failure spend estimate: could not price deployment of %s: %s", model, exc)
return 0.0
return prompt_cost
def _priced_input_cost(
pricing_model: str,
custom_llm_provider: str | None,
prompt_tokens: int,
custom_cost_per_token: CostPerToken | None,
) -> float | None:
def _priced_input_cost(pricing_model: str, custom_llm_provider: str | None, prompt_tokens: int) -> float | None:
try:
estimated_input_cost, _ = litellm.cost_per_token(
model=pricing_model,
custom_llm_provider=custom_llm_provider,
prompt_tokens=prompt_tokens,
custom_cost_per_token=custom_cost_per_token,
)
except Exception as exc: # noqa: BLE001 # cost_per_token raises BadRequestError for an unmapped model and the row must still be written
verbose_proxy_logger.debug("failure spend estimate: could not price %s: %s", pricing_model, exc)
@ -1233,18 +1263,20 @@ def _estimated_input_cost(
deployment_id: str | None,
base_model: str | None,
) -> float:
"""Input price for the estimated prompt: the deployment's registered rate when it has one, else
the cost map's rate for ``base_model``, else the deployment model, the same candidate order
``completion_cost`` walks (an Azure deployment name is not in the cost map, its ``base_model``
is; a ``base_model`` the cost map does not know falls through to the deployment model). A model
no candidate can price keeps the row at $0 rather than dropping the failure row."""
"""Input price for the estimated prompt: the deployment's registered rates when it has any (the id
is ``Logging.get_router_model_id()``, never the flattened ``model_call_details`` copy a request body
key can overwrite), else the cost map's rate for ``base_model``, else the deployment model, the same
candidate order ``completion_cost`` walks (an Azure deployment name is not in the cost map, its
``base_model`` is; a ``base_model`` the cost map does not know falls through to the deployment
model). A model no candidate can price keeps the row at $0 rather than dropping the failure row."""
model: Final = str(model_call_details.get("model") or "")
custom_llm_provider: Final = str(model_call_details.get("custom_llm_provider") or "") or None
custom_cost_per_token: Final = _deployment_cost_per_token(deployment_id, model)
pricing_models: Final = (base_model, model) if base_model is not None else (model,)
deployment_pricing: Final = deployment_pricing_model_info(model_id=deployment_id, deployment_model=model)
if deployment_pricing is not None:
return _deployment_priced_input_cost(model, custom_llm_provider, prompt_tokens, deployment_pricing)
pricing_models_to_try: Final = pricing_models(base_model, model)
candidate_costs: Final = (
_priced_input_cost(pricing_model, custom_llm_provider, prompt_tokens, custom_cost_per_token)
for pricing_model in pricing_models
_priced_input_cost(pricing_model, custom_llm_provider, prompt_tokens) for pricing_model in pricing_models_to_try
)
return next((cost for cost in candidate_costs if cost is not None), 0.0)
@ -1383,8 +1415,9 @@ def _timed_out_reading_from_provider(exception: Exception) -> bool:
def _client_disconnected(exception: Exception) -> bool:
"""The 499 the proxy raises when the caller hung up while the provider was still working on the request."""
return getattr(exception, "status_code", None) == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
"""The 499 the proxy itself raises when the caller hung up while the provider was still working on the
request. A provider's own 499 arrives as a litellm exception and is a plain upstream failure."""
return isinstance(exception, HTTPException) and exception.status_code == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
def _provider_billed_the_failure(exception: Exception) -> bool:
@ -1406,17 +1439,19 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed:
_litellm_params: Final = getattr(_logging_obj, "litellm_params", None)
_base_model: Final = get_base_model_from_metadata(model_call_details=_model_call_details)
_deployment_id: Final = _logging_obj.get_router_model_id() if isinstance(_logging_obj, Logging) else None
_success_already_logged: Final = isinstance(_logging_obj, Logging) and _logging_obj.async_success_log_enqueued
_billed: Final = provider_billed and not _success_already_logged
_usage_to_lift: Final = _failure_usage_to_lift(
model_call_details=_model_call_details,
request_body=request_data,
dispatched=_first_handoff is not None,
provider_billed=provider_billed,
provider_billed=_billed,
deployment_id=_deployment_id,
base_model=_base_model,
)
_prompt_estimate: Final = (
_estimated_prompt_usage(_model_call_details, request_data)
if provider_billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage")
if _billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage")
else None
)
_earlier_attempts: Final = (
@ -3859,6 +3894,7 @@ class ProxyLogging:
original_exception,
traceback_str,
),
kwargs={"for_async_request": True},
daemon=True,
).start()
@ -4696,19 +4732,25 @@ class ProxyLogging:
@staticmethod
def _discard_deferred_stream_logging_for_failure(request_data: Mapping[str, object], error: Exception) -> bool:
"""Drop the parked success dispatch for an assembled chat stream that ends in an error
``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead.
"""Drop the parked success dispatch for an assembled chat stream or a stream that billed several rounds when it
ends in an error ``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead.
Returns False when the parked dispatch should still be flushed by the caller."""
logging_obj: Final = request_data.get("litellm_logging_obj")
if not isinstance(logging_obj, Logging):
return False
_args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None)
assembled: Final = _args[0] if _args else None
if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS) or not isinstance(assembled, ModelResponse):
if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS):
return False
match _args:
case (ModelResponse() as assembled, *_):
logging_obj.record_assembled_response_for_failure(assembled)
case (BilledRoundsDispatch() as billed,):
billed.dispatch.close()
logging_obj.record_billed_rounds_for_failure(billed.spend)
case _:
return False
logging_obj._on_deferred_stream_complete = None
logging_obj._deferred_stream_complete_args = None
logging_obj.record_assembled_response_for_failure(assembled)
return True
async def arelease_max_parallel_requests_on_disconnect(
@ -4796,6 +4838,7 @@ def jsonify_object(data: dict) -> dict:
_deprecated_key_cache: Final[LimitedSizeOrderedDict] = LimitedSizeOrderedDict(max_size=1000)
_DEPRECATED_KEY_CACHE_TTL_SECONDS: Final = 60
_PRISMA_DEFAULT_TX_TIMEOUT: Final = timedelta(seconds=5)
_PRISMA_DEFAULT_TX_MAX_WAIT: Final = timedelta(seconds=2)
async def _lookup_deprecated_key(
@ -5202,13 +5245,17 @@ class PrismaClient:
return self.db.read_target
return self.db
def tx(self, *, timeout: timedelta = _PRISMA_DEFAULT_TX_TIMEOUT) -> "TransactionManager":
def tx(
self, *, timeout: timedelta = _PRISMA_DEFAULT_TX_TIMEOUT, max_wait: timedelta = _PRISMA_DEFAULT_TX_MAX_WAIT
) -> "TransactionManager":
"""Open an interactive transaction on the writer.
Callers go through this instead of reaching into ``self.db`` so writer
selection and read-replica routing stay encapsulated in the wrapper.
"""
return cast("TransactionManager", self.db.tx(timeout=timeout)) # cast-ok: untyped __getattr__ delegate
return cast( # cast-ok: untyped __getattr__ delegate
"TransactionManager", self.db.tx(max_wait=max_wait, timeout=timeout)
)
def get_request_status(self, payload: dict | SpendLogsPayload) -> Literal["success", "failure"]:
"""

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