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

chore: sync SSO concurrency fix with main

Browse source
This commit is contained in:
Yuneng Jiang 2026-10-10 15:09:53 -07:00
commit 146be5bd34
No known key found for this signature in database
131 changed files with 11836 additions and 1130 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

BIN
.github/assets/scaledown/provider-after.jpg vendored Normal file
View file

Binary file not shown.
Side by side

After

Width:  |  Height:  |  Size: 108 KiB

BIN
.github/assets/scaledown/provider-before.jpg vendored Normal file
View file

Binary file not shown.
Side by side

After

Width:  |  Height:  |  Size: 109 KiB

16
.github/workflows/publish-lint-base-counts.yml vendored
View file

@ -35,6 +35,8 @@ jobs:
script: scripts/type_discipline_gate.py
- checker: test-quality
script: scripts/test_quality_gate.py
- checker: openapi-docs
script: scripts/openapi_docs_gate.py
- checker: basedpyright
script: scripts/type_check_gate.py
@ -60,16 +62,24 @@ jobs:
uses: ./.github/actions/cache-cargo-build
- name: Cache Prisma binaries
if: matrix.checker == 'basedpyright'
if: matrix.checker == 'basedpyright' || matrix.checker == 'openapi-docs'
uses: ./.github/actions/cache-prisma-binaries
# The three source scanners only need the pinned dev tools (ruff and the
# The source scanners only need the pinned dev tools (ruff and the
# stdlib checkers), the same versions test-linting.yml's python job runs.
- name: Install the dev tools
if: matrix.checker != 'basedpyright'
if: matrix.checker != 'basedpyright' && matrix.checker != 'openapi-docs'
run: |
uv sync --frozen --only-group dev --no-install-project
# The OpenAPI gate imports the proxy app, so it needs the same install the
# test-linting.yml python job uses, Prisma client included.
- name: Install the proxy
if: matrix.checker == 'openapi-docs'
run: |
uv sync --frozen --group proxy-dev --group e2e-dev --extra proxy
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Emit ${{ matrix.checker }} counts for HEAD
if: matrix.checker != 'basedpyright'
run: |

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

@ -156,6 +156,14 @@ jobs:
run: |
uv run --no-sync python scripts/test_quality_gate.py --base "$GATE_BASE_SHA"
- name: Check OpenAPI docs (POST routes without a request schema, description or 200 schema, delta vs merge-base counts)
if: steps.changes.outputs.decision != 'skip'
env:
GH_TOKEN: ${{ github.token }}
run: |
uv sync --frozen --group proxy-dev --group e2e-dev --extra proxy
uv run --no-sync python scripts/openapi_docs_gate.py --base "$GATE_BASE_SHA"
- name: Print OpenAI version
if: steps.changes.outputs.decision != 'skip'
run: |

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

@ -327,6 +327,7 @@ jobs:
tests/unit/llms/sagemaker
tests/unit/llms/sail
tests/unit/llms/sambanova
tests/unit/llms/scaledown
tests/unit/llms/sap
tests/unit/llms/scaleway
tests/unit/llms/searchapi

2
AGENTS.md
View file

@ -54,7 +54,7 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a
Python max line length is 120, not 88
The four lint gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gate.py`, `scripts/test_quality_gate.py`, `scripts/type_check_gate.py`) compare each rule's codebase count on your branch against the count at its merge-base with the default branch, and a rule may not grow. There are no budget files to edit or ratchet: when a gate fails, fix the violations the branch introduced or remove at least as many of that rule elsewhere in the tree. The one exception is `reportAny` / `reportExplicitAny`, which share a fixed codebase-wide cap in `ANY_CAPS` in `scripts/type_check_gate.py` because Any spreads past the lines you touch. A branch may add Anys while the total stays under that cap. Only lower `ANY_CAPS`, and only in its own PR to the default branch, never raise it on a feature branch
The five lint gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gate.py`, `scripts/test_quality_gate.py`, `scripts/openapi_docs_gate.py`, `scripts/type_check_gate.py`) compare each rule's codebase count on your branch against the count at its merge-base with the default branch, and a rule may not grow. There are no budget files to edit or ratchet: when a gate fails, fix the violations the branch introduced or remove at least as many of that rule elsewhere in the tree. The one exception is `reportAny` / `reportExplicitAny`, which share a fixed codebase-wide cap in `ANY_CAPS` in `scripts/type_check_gate.py` because Any spreads past the lines you touch. A branch may add Anys while the total stays under that cap. Only lower `ANY_CAPS`, and only in its own PR to the default branch, never raise it on a feature branch
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice

11
Makefile
View file

@ -7,7 +7,7 @@
test-rust-extension rust-sqlx-prepare \
info lint lint-inner lint-dev lint-checks format \
lint-basedpyright lint-e2e-basedpyright lint-type-discipline \
lint-ruff-strict lint-gate lint-test-quality \
lint-ruff-strict lint-gate lint-test-quality lint-openapi-docs \
install-dev install-proxy-dev install-test-deps install-hooks \
install-helm-unittest check-circular-imports check-import-safety check check-inner pre-commit \
lint-install lint-fetch-base bootstrap
@ -35,6 +35,7 @@ help:
@echo " make lint-ruff-strict - Gate each strict ruff rule's codebase total against its merge-base count"
@echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches the default branch, simulates the merge)"
@echo " make lint-test-quality - Gate the test suite's TQ counts against their merge-base counts"
@echo " make lint-openapi-docs - Gate undocumented POST routes in the proxy OpenAPI document against their merge-base counts"
@echo " make check-circular-imports - Check for circular imports"
@echo " make check-import-safety - Check import safety"
@echo " make test - Run all tests"
@ -219,6 +220,12 @@ lint-type-discipline: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
lint-test-quality: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/test_quality_gate.py --base "$(BASE_REF)"
# OpenAPI docs gate (POST routes without a request body schema, a description or a
# 200 schema in app.openapi()), counted the same delta-vs-base way so a new endpoint
# has to show up documented in Swagger UI.
lint-openapi-docs: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/openapi_docs_gate.py --base "$(BASE_REF)"
lint-format: format-check
lint-ruff-strict: install-dev
@ -250,7 +257,7 @@ lint-inner: lint-install
@base_ref=$$($(RESOLVE_BASE)) && \
$(MAKE) BASE_REF="$$base_ref" -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-test-quality lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-test-quality lint-openapi-docs lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
# Faster linting for local development (only checks changed code)
lint-dev: lint-format-changed check-circular-imports check-import-safety

1
README.md
View file

@ -407,6 +407,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
| [Sagemaker Chat (`sagemaker_chat`)](https://docs.litellm.ai/docs/providers/aws_sagemaker) | ✅ | ✅ | ✅ | | | | | | | |
| [Sail (`sail`)](https://docs.litellm.ai/docs/providers/sail) | ✅ | ✅ | ✅ | | | | | | | |
| [Sambanova (`sambanova`)](https://docs.litellm.ai/docs/providers/sambanova) | ✅ | ✅ | ✅ | | | | | | | |
| [ScaleDown (`scaledown`)](https://docs.litellm.ai/docs/providers/scaledown) | ✅ | | | | | | | | | |
| [Snowflake (`snowflake`)](https://docs.litellm.ai/docs/providers/snowflake) | ✅ | ✅ | ✅ | | | | | | | |
| [Strands Decider (`strands_decider`)](https://docs.litellm.ai/docs/providers) | | | | | | | | | | |
| [Text Completion Codestral (`text-completion-codestral`)](https://docs.litellm.ai/docs/providers/codestral) | ✅ | ✅ | ✅ | | | | | | | |

10
litellm/__init__.py
View file

@ -674,6 +674,7 @@ nvidia_nim_models: Set = set()
nvidia_riva_models: Set = set()
soniox_models: Set = set()
sambanova_models: Set = set()
scaledown_models: Final[Set[str]] = set() # Price reloads update this registry through existing references.
sambanova_embedding_models: Set = set()
novita_models: Set = set()
assemblyai_models: Set = set()
@ -756,7 +757,7 @@ def is_openai_finetune_model(key: str) -> bool:
return key.startswith("ft:") and not key.count(":") > 1
def _populate_provider_model_sets(model_cost_map: Dict) -> None:
def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, object]]) -> None:
for key, value in model_cost_map.items():
if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key):
open_ai_chat_completion_models.add(key)
@ -918,6 +919,8 @@ def _populate_provider_model_sets(model_cost_map: Dict) -> None:
soniox_models.add(key)
elif value.get("litellm_provider") == "sambanova":
sambanova_models.add(key)
elif value.get("litellm_provider") == "scaledown":
scaledown_models.add(key)
elif value.get("litellm_provider") == "sambanova-embedding-models":
sambanova_embedding_models.add(key)
elif value.get("litellm_provider") == "novita":
@ -1106,6 +1109,7 @@ model_list = list(
| nvidia_riva_models
| soniox_models
| sambanova_models
| scaledown_models
| azure_text_models
| novita_models
| assemblyai_models
@ -1215,6 +1219,7 @@ def _build_models_by_provider() -> dict:
"nvidia_riva": nvidia_riva_models,
"soniox": soniox_models,
"sambanova": sambanova_models | sambanova_embedding_models,
"scaledown": scaledown_models,
"novita": novita_models,
"nebius": nebius_models | nebius_embedding_models,
"aiml": aiml_models,
@ -2054,6 +2059,9 @@ if TYPE_CHECKING:
from .llms.featherless_ai.chat.transformation import (
FeatherlessAIConfig as FeatherlessAIConfig,
)
from .llms.scaledown.chat.transformation import (
ScaleDownChatConfig as ScaleDownChatConfig,
)
from .llms.cerebras.chat import CerebrasConfig as CerebrasConfig
from .llms.nadir.chat.transformation import NadirConfig as NadirConfig
from .llms.baseten.chat import BasetenConfig as BasetenConfig

5
litellm/_lazy_imports_registry.py
View file

@ -277,6 +277,7 @@ LLM_CONFIG_NAMES: Final = (
"NvidiaNimConfig",
"NvidiaNimEmbeddingConfig",
"FeatherlessAIConfig",
"ScaleDownChatConfig",
"CerebrasConfig",
"NadirConfig",
"BasetenConfig",
@ -1103,6 +1104,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.featherless_ai.chat.transformation",
"FeatherlessAIConfig",
),
"ScaleDownChatConfig": (
".llms.scaledown.chat.transformation",
"ScaleDownChatConfig",
),
"CerebrasConfig": (".llms.cerebras.chat", "CerebrasConfig"),
"NadirConfig": (".llms.nadir.chat.transformation", "NadirConfig"),
"BasetenConfig": (".llms.baseten.chat", "BasetenConfig"),

2
litellm/_service_logger.py
View file

@ -16,7 +16,7 @@ from .types.services import ServiceLoggerPayload, ServiceTypes
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
Span = _Span | Any
OTELClass = OpenTelemetry

2
litellm/constants.py
View file

@ -234,6 +234,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0"
MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0"))
MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
MCP_UPSTREAM_SESSION_SHUTDOWN_TIMEOUT_SECONDS: Final = 5.0
MCP_TOOL_LISTING_MAX_PAGES: Final = 1000
MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH: Final = 8
MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60
@ -825,6 +826,7 @@ LITELLM_CHAT_PROVIDERS: Final = [
"lemonade",
"docker_model_runner",
"amazon_nova",
"scaledown",
]
# Resolving these providers runs an OAuth device flow (their provider info IS the login), so any

407
litellm/experimental_mcp_client/client.py
View file

@ -8,12 +8,12 @@ import hashlib
import json
import os
import time
from collections.abc import AsyncIterator, Awaitable, Callable, Generator, Sequence
from contextlib import AbstractAsyncContextManager
from collections.abc import AsyncIterator, Awaitable, Callable, Generator, Mapping, Sequence
from contextlib import AbstractAsyncContextManager, AsyncExitStack, nullcontext
from functools import partial
from importlib.metadata import version
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeAlias, TypeVar, cast
from typing import TYPE_CHECKING, Final, Literal, TypeAlias, TypeVar, cast, overload
import anyio
import httpx2
@ -21,26 +21,42 @@ from httpx2._client import UseClientDefault
from httpx2._types import AuthTypes
from mcp import ClientSession, MCPError, ReadResourceResult, Resource, StdioServerParameters
from mcp.client._input_required import run_input_required_driver
from mcp.client.session import ClientRequestContext, ElicitationFnT, LoggingFnT, SamplingFnT
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import streamable_http_client
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.message import SessionMessage
from typing_extensions import ReadOnly, TypedDict, assert_never
_TransportStreams: TypeAlias = tuple[
ReadStream[SessionMessage | Exception],
WriteStream[SessionMessage],
]
_TransportContext: TypeAlias = AbstractAsyncContextManager[_TransportStreams]
_SessionCallbackName: TypeAlias = Literal["sampling_callback", "elicitation_callback", "logging_callback"]
class _SessionCallbackMap(TypedDict):
sampling_callback: ReadOnly[SamplingFnT | None]
elicitation_callback: ReadOnly[ElicitationFnT | None]
logging_callback: ReadOnly[LoggingFnT | None]
from mcp.types import (
CONNECTION_CLOSED,
METHOD_NOT_FOUND,
REQUEST_TIMEOUT,
CacheableResult,
ClientCapabilities,
CreateMessageRequestParams,
CreateMessageResult,
CreateMessageResultWithTools,
DiscoverResult,
ElicitationCapability,
ElicitRequestParams,
ElicitResult,
ErrorData,
FormElicitationCapability,
GetPromptRequestParams,
GetPromptResult,
@ -58,6 +74,7 @@ from mcp.types import (
ListResourceTemplatesResult,
ListToolsRequest,
ListToolsResult,
LoggingMessageNotificationParams,
PaginatedRequestParams,
Prompt,
ResourceTemplate,
@ -302,6 +319,12 @@ class MCPSigV4Auth(httpx2.Auth):
raise ImportError("Missing botocore to use AWS SigV4 authentication. Run 'pip install boto3'.")
self.service_name = aws_service_name or "bedrock-agentcore"
self.region_name = aws_region_name or "us-east-1"
self.identity: Final[tuple[str, str, str | None, str | None]] = (
self.service_name,
self.region_name,
aws_role_name,
aws_access_key_id,
)
# Note: os.environ/ prefixed values are already resolved by
# ProxyConfig._check_for_os_environ_vars() at config load time.
# Values arrive here as plain strings.
@ -346,7 +369,7 @@ class MCPSigV4Auth(httpx2.Auth):
import time
import boto3
from botocore.credentials import Credentials
from botocore.credentials import RefreshableCredentials
session_name: Final = aws_session_name or f"litellm-mcp-{int(time.time())}"
sts_kwargs: Final[dict] = {"region_name": aws_region_name}
@ -356,15 +379,23 @@ class MCPSigV4Auth(httpx2.Auth):
if aws_session_token:
sts_kwargs["aws_session_token"] = aws_session_token
sts_client: Final = boto3.client("sts", **sts_kwargs)
sts_response: Final = sts_client.assume_role(
RoleArn=aws_role_name,
RoleSessionName=session_name,
)
sts_creds: Final = sts_response["Credentials"]
return Credentials(
access_key=sts_creds["AccessKeyId"],
secret_key=sts_creds["SecretAccessKey"],
token=sts_creds["SessionToken"],
def assume() -> Mapping[str, str]:
sts_creds: Final = sts_client.assume_role(
RoleArn=aws_role_name,
RoleSessionName=session_name,
)["Credentials"]
return MappingProxyType(
{
"access_key": sts_creds["AccessKeyId"],
"secret_key": sts_creds["SecretAccessKey"],
"token": sts_creds["SessionToken"],
"expiry_time": sts_creds["Expiration"].isoformat(),
}
)
return RefreshableCredentials.create_from_metadata(
metadata=dict(assume()), refresh_using=assume, method="sts-assume-role"
)
def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]:
@ -390,6 +421,19 @@ class MCPSigV4Auth(httpx2.Auth):
yield request
def _notify_stream_error(
callback: Callable[[asyncio.Future[Exception]], None] | None,
stream_error: asyncio.Future[Exception],
) -> None:
if callback is not None:
callback(stream_error)
def _notify_cleanup(callback: Callable[[], None] | None) -> None:
if callback is not None:
callback()
class MCPClient:
"""
MCP Client supporting:
@ -443,12 +487,37 @@ class MCPClient:
self._sampling_callback: Callable | None = sampling_callback
self._elicitation_callback: ElicitationCallback | None = elicitation_callback
self._logging_callback: Callable | None = logging_callback
self._session_callbacks_override: _SessionCallbackMap | None = None
# handle the basic auth value if provided
if auth_value:
self.update_auth_value(auth_value)
def session_callbacks(self) -> _SessionCallbackMap:
return {
"sampling_callback": self._sampling_callback,
"elicitation_callback": self._elicitation_callback,
"logging_callback": self._logging_callback,
}
def session_settings(self) -> tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool]:
callbacks: Final = self.session_callbacks()
return (
self.transport_type,
self.protocol_version,
self.timeout,
callbacks["sampling_callback"] is not None,
callbacks["elicitation_callback"] is not None,
callbacks["logging_callback"] is not None,
)
def route_session_callbacks(self, callbacks: _SessionCallbackMap) -> None:
self._session_callbacks_override = callbacks
async def discovery_auth_fingerprint(self) -> str:
return self._hash_discovery_auth(await self.prepare_request_auth())
return self._hash_discovery_auth(await self.prepare_request_auth(), self._aws_identity())
def _aws_identity(self) -> tuple[str, str, str | None, str | None] | None:
return self._aws_auth.identity if isinstance(self._aws_auth, MCPSigV4Auth) else None
async def prepare_request_auth(self) -> httpx2.Request:
"""Preview the authenticated request without sending it, closing the auth flow afterwards."""
@ -465,8 +534,10 @@ class MCPClient:
await flow.aclose()
@staticmethod
def _hash_discovery_auth(request: httpx2.Request) -> str:
material: Final = json.dumps((str(request.url), tuple(sorted(request.headers.multi_items()))))
def _hash_discovery_auth(
request: httpx2.Request, aws_identity: tuple[str, str, str | None, str | None] | None
) -> str:
material: Final = json.dumps((str(request.url), tuple(sorted(request.headers.multi_items())), aws_identity))
return hashlib.sha256(material.encode()).hexdigest()
def _create_transport_context(
@ -605,6 +676,9 @@ class MCPClient:
transport_ctx: _TransportContext,
operation: Callable[[ClientSession], Awaitable[TSessionResult]],
http_client: httpx2.AsyncClient | None = None,
*,
on_stream_error: Callable[[asyncio.Future[Exception]], None] | None = None,
on_cleanup: Callable[[], None] | None = None,
) -> TSessionResult:
"""
Execute an operation within a transport and session context.
@ -623,6 +697,7 @@ class MCPClient:
read_stream: Final = transport[0]
write_stream: Final = transport[1]
stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future()
_notify_stream_error(on_stream_error, stream_error)
async def receive_message(
message: ServerNotification | Exception,
@ -634,6 +709,11 @@ class MCPClient:
# The SDK closes pending requests when its message handler raises.
raise RuntimeError("MCP response stream failed")
session_kwargs: Final = (
self._session_callbacks_override
if self._session_callbacks_override is not None
else self.session_callbacks()
)
# The SDK drops a response stream that ends without a JSON-RPC reply, so nothing else
# ever fails the request.
session_ctx: Final = ClientSession(
@ -641,9 +721,7 @@ class MCPClient:
write_stream,
read_timeout_seconds=self.timeout,
message_handler=receive_message,
sampling_callback=self._sampling_callback,
elicitation_callback=self._elicitation_callback,
logging_callback=self._logging_callback,
**session_kwargs,
)
session: Final = await session_ctx.__aenter__()
try:
@ -659,6 +737,7 @@ class MCPClient:
raise stream_error.result()
raise
finally:
_notify_cleanup(on_cleanup)
cleanup_scope.shield = True
cleanup_scope.deadline = anyio.current_time() + 5
try:
@ -671,6 +750,7 @@ class MCPClient:
in_flight_error = e
raise
finally:
_notify_cleanup(on_cleanup)
cleanup_scope.shield = True
cleanup_scope.deadline = min(cleanup_scope.deadline, anyio.current_time() + 5)
try:
@ -703,6 +783,8 @@ class MCPClient:
operation: Callable[[ClientSession], Awaitable[TSessionResult]],
*,
quiet_on_error: bool = False,
on_stream_error: Callable[[asyncio.Future[Exception]], None] | None = None,
on_cleanup: Callable[[], None] | None = None,
) -> TSessionResult:
"""Open a session, run the provided coroutine, and clean up.
@ -714,7 +796,18 @@ class MCPClient:
try:
self._last_initialize_instructions = None
transport_ctx, http_client = self._create_transport_context()
result: Final = await self._execute_session_operation(transport_ctx, operation, http_client=http_client)
execution: Final = (
self._execute_session_operation(
transport_ctx,
operation,
http_client=http_client,
on_stream_error=on_stream_error,
on_cleanup=on_cleanup,
)
if on_stream_error is not None or on_cleanup is not None
else self._execute_session_operation(transport_ctx, operation, http_client=http_client)
)
result: Final = await execution
except Exception as e:
read_timeout: Final = as_mcp_read_timeout(e)
if read_timeout is not None:
@ -741,6 +834,12 @@ class MCPClient:
await anyio.lowlevel.checkpoint_if_cancelled()
return result
def open_persistent_session(
self,
admission: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> "PersistentMCPSession":
return PersistentMCPSession(self, admission)
def update_auth_value(self, mcp_auth_value: str | dict[str, str]) -> None:
"""
Set the authentication header for the MCP client.
@ -972,10 +1071,15 @@ class MCPClient:
host_progress_callback: Callable | None = None,
raise_on_error: bool = False,
allow_input_required: bool = False,
persistent_session: "PersistentMCPSession | None" = None,
) -> MCPCallToolResult | InputRequiredResult:
"""
Call an MCP Tool.
persistent_session runs the call inside an already-open upstream session (one
upstream mcp-session-id shared by every call of a stateful gateway session) instead of
the per-call initialize + teardown that run_with_session performs.
Args:
raise_on_error: When True, re-raise the underlying exception instead of returning an
``isError=True`` result. The token-exchange (OBO) tool-call path uses this to detect
@ -1030,7 +1134,15 @@ class MCPClient:
)
try:
tool_result: Final = await self.run_with_session(_call_tool_operation, quiet_on_error=raise_on_error)
tool_result: Final = (
await self.run_with_session(_call_tool_operation, quiet_on_error=raise_on_error)
if persistent_session is None
else await persistent_session.run(
_call_tool_operation,
quiet_on_error=raise_on_error,
callbacks=self.session_callbacks(),
)
)
verbose_logger.info("MCP client tool call '%s' completed successfully", call_tool_request_params.name)
return tool_result
except asyncio.CancelledError:
@ -1395,3 +1507,256 @@ class MCPClient:
"the MCP server may have crashed, disconnected, or timed out."
)
raise
_PendingOperation: TypeAlias = tuple[
Callable[[ClientSession], Awaitable[object]],
asyncio.Future[object],
_SessionCallbackMap | None,
]
_MAX_PENDING_OPERATIONS: Final = 64
def _ends_session(outcome: BaseException) -> bool:
if isinstance(
outcome,
(httpx2.HTTPError, OSError, anyio.BrokenResourceError, anyio.ClosedResourceError, anyio.EndOfStream),
):
return True
return isinstance(outcome, MCPError) and outcome.error.code == CONNECTION_CLOSED
def _deliver_persistent_outcome(
future: asyncio.Future[object],
outcome: object,
stream_error: asyncio.Future[Exception] | None,
) -> bool:
mapped_outcome: Final = (
stream_error.result()
if isinstance(outcome, MCPError) and stream_error is not None and stream_error.done()
else outcome
)
if not future.done():
if isinstance(mapped_outcome, Exception):
future.set_exception(mapped_outcome)
elif isinstance(mapped_outcome, BaseException):
future.set_exception(RuntimeError("upstream MCP operation was cancelled"))
else:
future.set_result(mapped_outcome)
return (
isinstance(outcome, BaseException)
and _ends_session(outcome)
or isinstance(mapped_outcome, BaseException)
and _ends_session(mapped_outcome)
or stream_error is not None
and stream_error.done()
)
class UpstreamSessionClosedError(RuntimeError):
def __init__(self) -> None:
super().__init__("upstream MCP session closed")
class PersistentMCPSession:
"""One upstream MCP session kept open across operations.
The SDK transport owns anyio cancel scopes that must be entered and exited by the same
task, so a dedicated task holds the session open and serves queued operations in order.
"""
def __init__(
self,
client: MCPClient,
admission: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> None:
loop: Final = asyncio.get_running_loop()
self._client: Final = client
self._admission: Final = admission
self._queue: Final[asyncio.Queue[_PendingOperation]] = asyncio.Queue(maxsize=_MAX_PENDING_OPERATIONS)
self._ready: Final[asyncio.Future[None]] = loop.create_future()
self._active: asyncio.Future[object] | None = None
self._stream_error: asyncio.Future[Exception] | None = None
self._retiring: bool = False
self._ended: bool = False
self._initial_callbacks: Final[_SessionCallbackMap] = client.session_callbacks()
self._current_callbacks: _SessionCallbackMap = self._initial_callbacks
delegated_callbacks: Final[_SessionCallbackMap] = {
"sampling_callback": (
self._delegate("sampling_callback")
if self._initial_callbacks["sampling_callback"] is not None
else None
),
"elicitation_callback": (
self._delegate("elicitation_callback")
if self._initial_callbacks["elicitation_callback"] is not None
else None
),
"logging_callback": (
self._delegate("logging_callback") if self._initial_callbacks["logging_callback"] is not None else None
),
}
client.route_session_callbacks(delegated_callbacks)
self._task: Final = loop.create_task(self._serve())
@overload
def _delegate(self, name: Literal["sampling_callback"]) -> SamplingFnT: ...
@overload
def _delegate(self, name: Literal["elicitation_callback"]) -> ElicitationFnT: ...
@overload
def _delegate(self, name: Literal["logging_callback"]) -> LoggingFnT: ...
def _delegate(self, name: _SessionCallbackName) -> SamplingFnT | ElicitationFnT | LoggingFnT:
match name:
case "sampling_callback":
async def sampling_callback(
context: ClientRequestContext, params: CreateMessageRequestParams
) -> CreateMessageResult | CreateMessageResultWithTools | ErrorData:
current_callback: Final = self._current_callbacks.get(name)
initial_callback: Final = self._initial_callbacks.get(name)
callback: Final = current_callback if current_callback is not None else initial_callback
if callback is None:
raise RuntimeError("sampling callback is unavailable")
return await callback(context, params)
return sampling_callback
case "elicitation_callback":
async def elicitation_callback(
context: ClientRequestContext, params: ElicitRequestParams
) -> ElicitResult | ErrorData:
current_callback: Final = self._current_callbacks.get(name)
initial_callback: Final = self._initial_callbacks.get(name)
callback: Final = current_callback if current_callback is not None else initial_callback
if callback is None:
raise RuntimeError("elicitation callback is unavailable")
return await callback(context, params)
return elicitation_callback
case "logging_callback":
async def logging_callback(params: LoggingMessageNotificationParams) -> None:
current_callback: Final = self._current_callbacks.get(name)
initial_callback: Final = self._initial_callbacks.get(name)
callback: Final = current_callback if current_callback is not None else initial_callback
if callback is None:
raise RuntimeError("logging callback is unavailable")
await callback(params)
return logging_callback
assert_never(name)
@property
def closed(self) -> bool:
return self._task.done()
@property
def reusable(self) -> bool:
return not self._ended and not self._retiring and not self._task.done()
def _record_stream_error(self, stream_error: asyncio.Future[Exception]) -> None:
self._stream_error = stream_error
def _mark_ended(self) -> None:
self._ended = True
async def _serve(self) -> None:
admission_stack: Final = AsyncExitStack()
async def drain(session: ClientSession) -> None:
await admission_stack.aclose()
self._ready.set_result(None)
while True:
operation, future, callbacks = await self._queue.get()
if not future.done():
self._active = future
if callbacks is not None:
self._current_callbacks = callbacks
async with self._admission() if self._admission is not None else nullcontext():
if not future.done():
(outcome,) = await asyncio.gather(operation(session), return_exceptions=True)
if _deliver_persistent_outcome(future, outcome, self._stream_error):
self._active = None
self._ended = True
return
self._active = None
if self._retiring and self._queue.empty():
return
try:
if self._admission is not None:
await admission_stack.enter_async_context(self._admission())
(ended,) = await asyncio.gather(
self._client.run_with_session(
drain,
quiet_on_error=True,
on_stream_error=self._record_stream_error,
on_cleanup=self._mark_ended,
),
return_exceptions=True,
)
self._ended = True
except asyncio.CancelledError:
self._ended = True
self._fail_waiters(None)
raise
finally:
await admission_stack.aclose()
self._fail_waiters(ended if isinstance(ended, Exception) else None)
def _fail_waiters(self, cause: Exception | None) -> None:
pending: Final = (self._ready, self._active, *(future for _, future, _ in self._drained()))
for future in pending:
if future is not None and not future.done():
future.set_exception(UpstreamSessionClosedError() if cause is None else cause)
def _drained(self) -> tuple[_PendingOperation, ...]:
return tuple(self._queue.get_nowait() for _ in range(self._queue.qsize()))
async def run(
self,
operation: Callable[[ClientSession], Awaitable[TSessionResult]],
*,
quiet_on_error: bool = False,
callbacks: _SessionCallbackMap | None = None,
) -> TSessionResult:
if self.closed:
await asyncio.shield(self._ready)
raise UpstreamSessionClosedError()
future: Final[asyncio.Future[object]] = asyncio.get_running_loop().create_future()
await self._queue.put((operation, future, callbacks))
try:
_ = await asyncio.wait((future, self._task), return_when=asyncio.FIRST_COMPLETED)
except asyncio.CancelledError:
_ = future.cancel()
raise
if not future.done():
raise UpstreamSessionClosedError()
return cast(TSessionResult, future.result()) # cast-ok: one queue serves operations of every result type
def close(self) -> None:
_ = self._task.cancel()
def retire(self) -> None:
if self._retiring:
return
self._retiring = True
if self.closed:
return
async def retired_operation(_: ClientSession) -> None:
return None
loop: Final = asyncio.get_running_loop()
future: Final[asyncio.Future[object]] = loop.create_future()
future.set_result(None)
try:
self._queue.put_nowait((retired_operation, future, None))
except asyncio.QueueFull:
return
async def wait_closed(self) -> None:
await asyncio.gather(self._task, return_exceptions=True)

2
litellm/integrations/SlackAlerting/slack_alerting.py
View file

@ -38,7 +38,6 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.repositories.table_repositories import InvitationLinkRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
@ -52,6 +51,7 @@ from litellm.types.integrations.slack_alerting import (
VirtualKeyEvent,
WebhookEvent,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.proxy.model_deprecation import (
DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS,
DEPRECATION_IDLE_POLL_SECONDS,

65
litellm/integrations/azure_storage/azure_storage.py
View file

@ -2,7 +2,7 @@ import asyncio
import os
import time
from collections.abc import Callable
from datetime import datetime, timedelta
from datetime import datetime, timedelta, timezone
from functools import cache
from typing import Final
from urllib.parse import unquote
@ -28,7 +28,7 @@ from litellm.secret_managers.get_azure_ad_token_provider import (
from litellm.types.secret_managers.get_azure_ad_token_provider import (
AzureCredentialType,
)
from litellm.types.utils import StandardLoggingPayload
from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
AZURE_STORAGE_TOKEN_SCOPE: Final = "https://storage.azure.com/.default"
_ADLS_SAFE_NAME: Final = str.maketrans("/", "_", "=")
@ -160,6 +160,24 @@ class AzureBlobStorageLogger(CustomBatchLogger):
except Exception as e:
verbose_logger.exception("AzureBlobStorageLogger Layer Error - %s", e)
async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> None:
now: Final = datetime.now(timezone.utc)
file_path: Final = f"audit_logs/{now.strftime('%Y-%m-%d')}/{now.strftime('%H-%M-%S')}_{audit_log['id']}.json"
await self._upload_json_to_file_path(file_path=file_path, json_payload=safe_dumps(audit_log))
async def _upload_json_to_file_path(self, file_path: str, json_payload: str) -> None:
payload_bytes: Final = json_payload.encode("utf-8")
if self.azure_storage_account_key:
await self._upload_bytes_with_account_key(file_path=file_path, content=payload_bytes)
return
await self.set_valid_azure_ad_token()
async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
base_url: Final = f"{self.azure_storage_dfs_endpoint}/{self.azure_storage_file_system}/{file_path}"
await self._create_file(async_client, base_url)
await self._append_data(async_client, base_url, json_payload)
await self._flush_data(async_client, base_url, len(payload_bytes))
async def async_send_batch(self):
"""
Sends the in memory logs queue to Azure Blob Storage
@ -366,41 +384,22 @@ class AzureBlobStorageLogger(CustomBatchLogger):
This is used when Azure Storage Account Key is set - Azure Storage Account Key does not work directly with Azure Rest API
"""
# Create an async service client
service_client: Final = await self.get_service_client()
# Get file system client
file_system_client: Final = service_client.get_file_system_client(file_system=self.azure_storage_file_system)
try:
# Create directory with today's date
from datetime import datetime
today: Final = datetime.now().strftime("%Y-%m-%d")
directory_client: Final = file_system_client.get_directory_client(today)
# check if the directory exists
if not await directory_client.exists():
await directory_client.create_directory()
verbose_logger.debug("Created directory: %s", today)
# Create a file client
file_name: Final = adls_safe_file_name(payload.get("id"))
file_client: Final = directory_client.get_file_client(file_name)
# Create the file
await file_client.create_file()
# Content to append
content: Final = safe_dumps(payload).encode("utf-8")
# Append content to the file
await file_client.append_data(data=content, offset=0, length=len(content))
# Flush the content to finalize the file
await file_client.flush_data(position=len(content), offset=0)
await self._upload_bytes_with_account_key(
file_path=f"{today}/{file_name}", content=safe_dumps(payload).encode("utf-8")
)
verbose_logger.debug("Successfully uploaded and wrote to %s/%s", today, file_name)
except Exception as e:
verbose_logger.exception("Error occurred: %s", e)
async def _upload_bytes_with_account_key(self, file_path: str, content: bytes) -> None:
service_client: Final = await self.get_service_client()
file_client: Final = service_client.get_file_system_client(
file_system=self.azure_storage_file_system
).get_file_client(file_path)
await file_client.create_file()
await file_client.append_data(data=content, offset=0, length=len(content))
await file_client.flush_data(position=len(content), offset=0)

4
litellm/integrations/custom_guardrail.py
View file

@ -44,7 +44,7 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
dc: Final = DualCache()
@ -119,7 +119,7 @@ def is_guardrail_intervention(e: Exception) -> bool:
def _user_api_key_auth_from_request(request_data: Mapping[str, object]) -> "UserAPIKeyAuth":
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
metadata: Final = request_data.get(get_metadata_variable_name_from_kwargs(request_data))
stamped: Final[Mapping[str, object]] = metadata if isinstance(metadata, dict) else {}

2
litellm/integrations/custom_logger.py
View file

@ -35,12 +35,12 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp import (
MCPPostCallResponseObject,
MCPPreCallRequestObject,
MCPPreCallResponseObject,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.router import PreRoutingHookResponse
Span = _Span

2
litellm/integrations/datadog/datadog.py
View file

@ -66,7 +66,7 @@ from ..additional_logging_utils import AdditionalLoggingUtils
if TYPE_CHECKING:
from fastapi import HTTPException
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class _DatadogLoggingKwargs(TypedDict, total=False):

4
litellm/integrations/otel/README.md
View file

@ -295,7 +295,9 @@ nothing here imports outside it:
- [`config.py`](./model/config.py) — `OpenTelemetryV2Config`, a pydantic-settings
model that reads `OTEL_*` / `LITELLM_OTEL_*` env vars, plus the feature gate.
`capture_span_content` gates whether prompt/response bodies may be written as
span attributes; it defaults **off** (`no_content`). The Baggage allowlists are
span attributes; it defaults **off** (`no_content`). A team or key can override
the global mode per OTel v2 destination with `capture_message_content`; omitted
values follow the proxy's own callback for that backend, else the global mode. The Baggage allowlists are
configurable, not hard-coded: set `LITELLM_OTEL_BAGGAGE_PROMOTED_KEYS` /
`LITELLM_OTEL_BAGGAGE_METADATA_KEYS` /
`LITELLM_OTEL_BAGGAGE_TEAM_METADATA_KEYS` (comma-separated) as env vars, or

8
litellm/integrations/otel/langfuse_logger.py
View file

@ -13,7 +13,7 @@ from litellm.integrations.otel.model.trace_controls import caller_trace_controls
from litellm.integrations.otel.plumbing.context import request_root_span
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.utils import ModelResponseStream
@ -51,6 +51,10 @@ class LangfuseContentOpenTelemetryV2(LangfuseOpenTelemetryV2):
response: "AsyncIterator[ModelResponseStream]",
request_data: Mapping[str, object],
) -> "AsyncGenerator[ModelResponseStream, None]":
if not self._capture_span_content():
async for chunk in response:
yield chunk
return
relayed: Final[list[ModelResponseStream]] = [] # mutable-ok: relayed as they arrive, assembled at end of stream
async for chunk in response:
relayed.append(chunk)
@ -58,6 +62,8 @@ class LangfuseContentOpenTelemetryV2(LangfuseOpenTelemetryV2):
self._stamp_root_io(request_data, lambda: stream_output(tuple(relayed), request_data))
def _stamp_root_io(self, data: Mapping[str, object], render_output: Callable[[], str | None]) -> None:
if not self._capture_span_content():
return
root: Final = request_root_span()
if root is None or not root.is_recording():
return

19
litellm/integrations/otel/logger.py
View file

@ -58,6 +58,7 @@ from litellm.integrations.otel.plumbing.context import (
is_recordable_span,
mcp_message_transport_span,
post_response_root,
request_destinations,
request_root_http_route,
request_root_span,
resolve_internal_call_span_context,
@ -82,12 +83,13 @@ from litellm.integrations.otel.plumbing.providers import (
resolve_meter_provider,
)
from litellm.integrations.otel.plumbing.routing import TenantTracerCache
from litellm.types.utils import captures_span_content
if TYPE_CHECKING:
from opentelemetry.metrics import MeterProvider
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.services import ServiceLoggerPayload
from litellm.types.utils import (
CallTypesLiteral,
@ -254,6 +256,11 @@ class OpenTelemetryV2(CustomLogger):
return None
return GenAIEventRecorder(get_event_logger(provider, LITELLM_TRACER_NAME), provider.resource)
def _capture_span_content(self) -> bool:
return self.config.capture_span_content or any(
captures_span_content(destination.capture_message_content) for destination in request_destinations()
)
# ====================================================================== #
# Proxy global registration
# ====================================================================== #
@ -433,7 +440,7 @@ class OpenTelemetryV2(CustomLogger):
return False
payload: Final = cast("StandardLoggingPayload", raw_payload)
data: Final = MCPToolCallSpanData.from_standard_logging_payload(
payload, capture_content=self.config.capture_span_content
payload, capture_content=self._capture_span_content()
)
# A stray LLM carrier from a ``pre_call`` that mis-fired for this id would
# otherwise linger until evicted; drop it so it's neither leaked nor closed
@ -477,7 +484,7 @@ class OpenTelemetryV2(CustomLogger):
return False
payload: Final = cast("StandardLoggingPayload", raw_payload)
data: Final = MCPListToolsSpanData.from_standard_logging_payload(
payload, capture_content=self.config.capture_span_content
payload, capture_content=self._capture_span_content()
)
if data.identity.call_id:
self._release_carrier(self._open_llm_calls.pop(data.identity.call_id, None))
@ -569,7 +576,7 @@ class OpenTelemetryV2(CustomLogger):
return None
data: Final = LLMCallSpanData.from_standard_logging_payload(
payload,
capture_content=self.config.capture_span_content,
capture_content=self._capture_span_content(),
time_to_first_chunk_seconds=call.time_to_first_chunk_seconds,
request_route=request_root_http_route(),
request_purpose=call.purpose,
@ -1076,6 +1083,6 @@ def build_otel_v2_logger(
def _logger_class(config: OpenTelemetryV2Config) -> type[OpenTelemetryV2]:
if "langfuse" not in config.mapper_names:
return OpenTelemetryV2
from litellm.integrations.otel.langfuse_logger import LangfuseContentOpenTelemetryV2, LangfuseOpenTelemetryV2
from litellm.integrations.otel.langfuse_logger import LangfuseContentOpenTelemetryV2
return LangfuseContentOpenTelemetryV2 if config.capture_span_content else LangfuseOpenTelemetryV2
return LangfuseContentOpenTelemetryV2

59
litellm/integrations/otel/model/config.py
View file

@ -16,19 +16,17 @@ from litellm.integrations.otel.model.baggage import (
)
from litellm.integrations.otel.model.spans import POSTGRESQL, db_system
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.utils import OtelSpanScope
from litellm.types.utils import (
CAPTURE_MESSAGE_CONTENT_VALUES,
CaptureMessageContent,
OtelSpanScope,
captures_span_content,
)
#: Master feature-flag env var. The logger is inert until this is truthy.
OTEL_V2_ENV: Final = "LITELLM_OTEL_V2"
class CaptureMessageContent(str):
NO_CONTENT = "no_content"
SPAN_ONLY = "span_only"
EVENT_ONLY = "event_only"
SPAN_AND_EVENT = "span_and_event"
class ExporterOwner(str, Enum):
"""The preset that contributed an exporter. Values match the callback names
in ``presets.PRESET_BY_CALLBACK`` so per-request dynamic-credential routing
@ -297,9 +295,7 @@ class OpenTelemetryV2Config(BaseSettings):
here are lower_snake_case; normalizing at the boundary keeps both
spellings working and lets every downstream comparison stay exact.
"""
if isinstance(value, str):
return value.lower()
return value
return _normalized_capture_message_content(value)
@field_validator("langfuse_span_scope", mode="before")
@classmethod
@ -371,22 +367,39 @@ class OpenTelemetryV2Config(BaseSettings):
@property
def capture_span_content(self) -> bool:
"""Whether prompt/response content may be stamped as span attributes.
Defaults off (``no_content``): an operator must opt in before message
bodies leave the process, so a user request can never force its prompt
or completion into the configured backend while capture is disabled.
"""
return self.capture_message_content in (
CaptureMessageContent.SPAN_ONLY,
CaptureMessageContent.SPAN_AND_EVENT,
)
"""Whether the global setting permits prompt/response span attributes."""
return captures_span_content(self.capture_message_content)
@classmethod
def from_env(cls) -> "OpenTelemetryV2Config":
return cls()
def _normalized_capture_message_content(value: object) -> object:
if isinstance(value, str):
return value.lower()
return value
class _CaptureMessageContentSettings(BaseSettings):
model_config = SettingsConfigDict(extra="ignore", frozen=True)
capture_message_content: str = Field(
default=CaptureMessageContent.NO_CONTENT,
validation_alias=AliasChoices("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"),
)
@field_validator("capture_message_content", mode="before")
@classmethod
def _normalize_capture_message_content(cls, value: object) -> object:
return _normalized_capture_message_content(value)
def capture_message_content_from_env() -> str:
"""The global capture mode, read without building a settings model that rereads every other env var"""
return _CaptureMessageContentSettings().capture_message_content
_EXCLUDED_SERVICES_INPUT: Final[TypeAdapter[str | tuple[object, ...]]] = TypeAdapter(str | tuple[object, ...])
@ -395,6 +408,10 @@ def excluded_db_systems_from(value: object) -> frozenset[str]:
return _normalize_excluded_services(excluded_service_names(value))
def parse_capture_message_content(value: object) -> str | None:
return value if isinstance(value, str) and value in CAPTURE_MESSAGE_CONTENT_VALUES else None
def excluded_service_names(value: object) -> frozenset[str]:
"""Read a YAML list or comma-separated string of service names, logging and dropping unusable input
so a malformed value cannot stop the OTel logger from being built"""

19
litellm/integrations/otel/model/destination.py
View file

@ -11,7 +11,7 @@ from urllib.parse import quote
from pydantic import ConfigDict, Field
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.utils import OtelSpanScope
from litellm.types.utils import OtelSpanScope, captures_span_content
class OtelDestination(LiteLLMBaseModel):
@ -50,6 +50,17 @@ class OtelDestination(LiteLLMBaseModel):
default=None,
description="Same, for the requests with a failed span in their tree; ``None`` forwards every one.",
)
capture_message_content: str | None = Field(
default=None,
description=(
"An explicit mode overrides the global capture policy for this destination. "
"Omitted follows the global setting."
),
)
def captures_content(self, default: str) -> bool:
setting: Final = self.capture_message_content
return captures_span_content(default if setting is None else setting)
def header_string(self) -> str:
"""Render headers as the ``k=v,k2=v2`` form an ``ExporterSpec`` expects.
@ -64,9 +75,9 @@ class OtelDestination(LiteLLMBaseModel):
def cache_key(self) -> tuple[str, tuple[tuple[str, str], ...], tuple[tuple[str, str], ...], str | None]:
"""Identity for processor reuse, so one destination means one exporter.
``span_scope`` and the sampling rates are left out on purpose: they decide which
spans reach the processor, not how the processor exports them, so a full and an
``llm_only`` view of the same account share one exporter.
``span_scope``, the sampling rates and ``capture_message_content`` are left out on
purpose: they decide which spans reach the processor and what they carry, not how
the processor exports them, so two views of the same account share one exporter.
"""
return (
self.endpoint,

187
litellm/integrations/otel/plumbing/providers.py
View file

@ -5,7 +5,7 @@ import random
import threading
import time
from collections import OrderedDict
from collections.abc import Callable, Iterable, Mapping, Sequence
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal
@ -40,8 +40,19 @@ from opentelemetry.util.types import Attributes, AttributeValue
from litellm._logging import verbose_logger
from litellm._version import version as litellm_version
from litellm.integrations.otel.mappers.langfuse import LANGFUSE_TRACE_NAME
from litellm.integrations.otel.model.config import ExporterOwner, ExporterSpec, OpenTelemetryV2Config
from litellm.integrations.otel.mappers.langfuse import (
LANGFUSE_OBSERVATION_INPUT,
LANGFUSE_OBSERVATION_OUTPUT,
LANGFUSE_TRACE_NAME,
)
from litellm.integrations.otel.model.config import (
CaptureMessageContent,
ExporterOwner,
ExporterSpec,
OpenTelemetryV2Config,
capture_message_content_from_env,
captures_span_content,
)
from litellm.integrations.otel.model.semconv import (
DB,
MCP,
@ -385,6 +396,26 @@ _CAPTURED_HEADER_PREFIXES: Final = ("http.request.header.", "http.response.heade
# ``?key=`` query parameter.
_URL_KEYS: Final = frozenset({"http.url", "http.target", "url.full"})
_URL_QUERY_KEY: Final = "url.query"
# Prompt and response content, in every vocabulary the mappers write it in once the
# global capture policy collects it: GenAI semconv (MCP tool-call arguments and results
# included), Langfuse, OpenInference's blobs and per-index messages, Langtrace, Weave.
_MESSAGE_CONTENT_KEYS: Final = frozenset(
{
GenAI.INPUT_MESSAGES,
GenAI.OUTPUT_MESSAGES,
GenAI.SYSTEM_INSTRUCTIONS,
GenAI.TOOL_CALL_ARGUMENTS,
GenAI.TOOL_CALL_RESULT,
LANGFUSE_OBSERVATION_INPUT,
LANGFUSE_OBSERVATION_OUTPUT,
"input.value",
"output.value",
"llm.prompts",
"llm.completions",
"weave.output",
}
)
_INDEXED_MESSAGE_PREFIXES: Final = ("llm.input_messages.", "llm.output_messages.")
class _SpanView(ReadableSpan):
@ -415,6 +446,19 @@ class _SpanView(ReadableSpan):
)
def _is_message_content(key: str) -> bool:
return key in _MESSAGE_CONTENT_KEYS or key.startswith(_INDEXED_MESSAGE_PREFIXES)
def without_message_content(span: ReadableSpan) -> ReadableSpan:
"""``span`` without its prompt and response content, everything else left as it was."""
attributes: Final = span.attributes or _NO_ATTRIBUTES
if not any(_is_message_content(key) for key in attributes):
return span
kept: Final = MappingProxyType({key: value for key, value in attributes.items() if not _is_message_content(key)})
return _SpanView(span, span.resource, kept, span.events, span.status, parent=span.parent)
def _is_database_span(attributes: Mapping[str, AttributeValue]) -> bool:
return any(key in attributes for key in _DB_SYSTEM_KEYS)
@ -553,7 +597,11 @@ def _tenant_resource(resource: Resource, defaults: Mapping[str, str], extra: Map
return filled.merge(Resource(extra)) if extra else filled
def _for_destination(span: ReadableSpan, destination: "OtelDestination") -> ReadableSpan:
def _for_destination(
span: ReadableSpan,
destination: "OtelDestination",
default_capture: str,
) -> ReadableSpan:
"""The view of ``span`` a tenant destination receives.
A span the tenant's own call produced keeps its error text. Every other span is
@ -564,8 +612,10 @@ def _for_destination(span: ReadableSpan, destination: "OtelDestination") -> Read
exception it raised and names the operator's guardrail endpoint. Stack traces walk
the operator's install and come off every span, as do the headers the operator
captures on the server span, whose request side holds the caller's bearer token,
and the query string of the request URL, which can hold the same key. The span
itself stays, so the tenant still gets the whole trace tree.
and the query string of the request URL, which can hold the same key. A
destination whose effective setting disables span content loses the prompt and
response content too. The span itself stays, so the tenant still gets the whole
trace tree.
"""
extra: Final = destination.resource_attributes
defaults: Final = destination.resource_defaults
@ -573,11 +623,12 @@ def _for_destination(span: ReadableSpan, destination: "OtelDestination") -> Read
database: Final = _is_database_span(attributes)
owned: Final = _is_tenant_owned_span(attributes)
unreachable: Final = _guardrail_unreachable(attributes)
redacted: Final = not destination.captures_content(default_capture)
kept: Final = MappingProxyType(
{
key: _without_query(key, value)
for key, value in attributes.items()
if _tenant_visible(key, database, owned, unreachable)
if _tenant_visible(key, database, owned, unreachable) and not (redacted and _is_message_content(key))
}
)
recorded: Final = span.events
@ -612,6 +663,8 @@ class TenantFanOutSpanProcessor(SpanProcessor):
shutdown_drain_seconds: float = _SHUTDOWN_DRAIN_SECONDS,
operator_sinks: 'Mapping[_SinkKey, "OtelSpanScope"]' = MappingProxyType({}),
excluded_db_systems: frozenset[str] = frozenset(),
default_capture: str = CaptureMessageContent.NO_CONTENT,
default_capture_by_backend: Mapping[str, str] = MappingProxyType({}),
pending_drains: int = _MAX_PENDING_DRAINS,
drain_pool: _DrainPool | None = None,
sampling_draw: Callable[[], float] | None = None,
@ -619,6 +672,8 @@ class TenantFanOutSpanProcessor(SpanProcessor):
self._operator_sinks: Final = operator_sinks
self._draw: Final = sampling_draw if sampling_draw is not None else random.random
self._excluded_db_systems: Final = excluded_db_systems
self._default_capture: Final = default_capture
self._default_capture_by_backend: Final = default_capture_by_backend
self._drain_seconds: Final = shutdown_drain_seconds
self._lock: Final = threading.Condition()
self._closed = False # guarded by ``_lock``: an unlocked read races the teardown it gates
@ -661,12 +716,22 @@ class TenantFanOutSpanProcessor(SpanProcessor):
for held_span, held_destination in self._settle(trace_id, failed):
self._forward(held_span, held_destination)
def _default_capture_for(self, destination: "OtelDestination") -> str:
backend: Final = destination.callback_name
return (
self._default_capture_by_backend.get(backend, self._default_capture) if backend else self._default_capture
)
def _forward(self, span: ReadableSpan, destination: "OtelDestination") -> None:
processor: Final = self._acquire(destination)
if processor is None:
return
try:
processor.on_end(_scoped(_for_destination(span, destination), destination.span_scope))
processor.on_end(
_scoped(
_for_destination(span, destination, self._default_capture_for(destination)), destination.span_scope
)
)
except Exception as exc: # noqa: BLE001 # one destination's failure must not cost the others their span
verbose_logger.debug("OTel V2 fan-out: forwarding to %s failed: %s", destination.endpoint, exc)
finally:
@ -963,14 +1028,15 @@ class _OverriddenBackendFilter(SpanProcessor):
Wrapping is the only place this works: ``SynchronousMultiSpanProcessor.on_end``
ignores return values, so a sibling processor can never veto the export.
Under ``additive`` mode nothing is suppressed, so the wrapper passes every span
straight through and the operator keeps its copy.
Under ``additive`` mode spans are not suppressed, though global content policy
still strips message content when disabled.
``scope`` narrows what the exporter receives independently of that: under
``llm_only`` the model-call spans go through as trace roots and the rest of the
tree is held back, unless a destination of the request names ``sink``, the account
this exporter writes to, with a wider scope: the fan-out then delivers the rest of
the tree there and the model call keeps its place in it.
the tree there and the model call keeps its place in it. A destination naming
``sink`` sets this exporter's content policy; otherwise it follows global capture.
"""
def __init__(
@ -979,11 +1045,15 @@ class _OverriddenBackendFilter(SpanProcessor):
owner: str | None,
scope: "OtelSpanScope" = "full",
sink: _SinkKey | None = None,
global_captures: bool = True,
filter_content: bool = True,
) -> None:
self._inner: Final = inner
self._owner: Final = owner
self._scope: Final = scope
self._sink: Final = sink
self._global_captures: Final = global_captures
self._filter_content: Final = filter_content
def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None:
self._inner.on_start(span, parent_context)
@ -991,7 +1061,28 @@ class _OverriddenBackendFilter(SpanProcessor):
def on_end(self, span: ReadableSpan) -> None:
if self._owner in suppressed_backends() or not _in_scope(span, self._scope):
return
self._inner.on_end(_scoped(span, self._account_scope()))
scoped: Final = _scoped(span, self._account_scope())
self._inner.on_end(
without_message_content(scoped) if self._filter_content and not self._account_captures_content() else scoped
)
def _account_captures_content(self) -> bool:
"""Whether a destination of this request that names ``sink`` captures content.
The fan-out skips a destination this exporter already writes to, so this copy is
the one that account gets and has to honor the destination's content mode.
"""
explicit: Final = next(
(
destination.capture_message_content
for destination in request_destinations()
if self._sink is not None
and destination.capture_message_content is not None
and _sink_key(destination.endpoint, destination.headers) == self._sink
),
None,
)
return captures_span_content(explicit) if explicit is not None else self._global_captures
def _account_scope(self) -> "OtelSpanScope":
if self._scope == "full" or self._sink is None:
@ -1010,6 +1101,25 @@ class _OverriddenBackendFilter(SpanProcessor):
return self._inner.force_flush(timeout_millis)
class _MessageContentFilter(SpanProcessor):
"""Strip prompt and response content from every span before ``inner`` exports it."""
def __init__(self, inner: SpanProcessor) -> None:
self._inner: Final = inner
def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None:
self._inner.on_start(span, parent_context)
def on_end(self, span: ReadableSpan) -> None:
self._inner.on_end(without_message_content(span))
def shutdown(self) -> None:
self._inner.shutdown()
def force_flush(self, timeout_millis: int = 30000) -> bool:
return self._inner.force_flush(timeout_millis)
def build_span_exporter(config: OpenTelemetryV2Config) -> SpanExporter:
"""Build a single exporter from the top-level config fields.
@ -1261,12 +1371,23 @@ def build_resource(config: OpenTelemetryV2Config) -> Resource:
return Resource.create(attributes)
def _spec_processor(
spec: ExporterSpec, use_simple_processor: bool | None, strip_message_content: bool
) -> SpanProcessor:
exporting: Final = _processor_for(
_exporter_from_spec(spec),
(spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor),
)
return _MessageContentFilter(exporting) if strip_message_content else exporting
def build_tracer_provider(
config: OpenTelemetryV2Config,
exporter: SpanExporter | None = None,
baggage_processor: SpanProcessor | None = None,
use_simple_processor: bool | None = None,
tenant_overrides: bool = False,
strip_message_content: bool = False,
) -> TracerProvider:
"""Build the shared :class:`TracerProvider`.
@ -1285,6 +1406,9 @@ def build_tracer_provider(
``config.langfuse_span_scope`` narrows the exporter owned by ``langfuse_otel``
alone; a collector or any other backend in the same config keeps the full tree.
``strip_message_content`` removes prompt and response content before every
exporter, for a routed provider whose spans can carry content another
destination asked for.
"""
provider: Final = TracerProvider(resource=build_resource(config))
if baggage_processor is None:
@ -1300,17 +1424,23 @@ def build_tracer_provider(
for spec in config.exporters:
if spec.requires_headers and not spec.headers:
continue
exp = _exporter_from_spec(spec)
processor = _processor_for(
exp,
(spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor),
)
processor = _spec_processor(spec, use_simple_processor, strip_message_content)
owner = spec.owner.value if tenant_overrides and spec.owner is not None else None
scope = _operator_scope(config, spec)
sink = _sink_key(spec.endpoint, parse_headers(spec.headers)) if _exports_to_the_wire(spec) else None
provider.add_span_processor(
_OverriddenBackendFilter(processor, owner, scope, sink)
if owner is not None or scope != "full"
_OverriddenBackendFilter(
processor,
owner,
scope,
sink,
config.capture_span_content,
tenant_overrides,
)
if owner is not None
or scope != "full"
or sink is not None
or (tenant_overrides and not config.capture_span_content)
else processor
)
return provider
@ -1340,10 +1470,29 @@ def attach_tenant_fan_out(
TenantFanOutSpanProcessor(
operator_sinks=operator_sink_scopes(*configs),
excluded_db_systems=excluded_db_systems,
default_capture=_operator_capture_mode(configs),
default_capture_by_backend=MappingProxyType(dict(_owned_capture_modes(configs))),
)
)
def _operator_capture_mode(configs: Sequence[OpenTelemetryV2Config]) -> str:
"""The mode of the operator's generic ``otel`` callback, so ``callback_settings.otel`` counts as the global default"""
return next(
(config.capture_message_content for config in configs if any(spec.owner is None for spec in config.exporters)),
capture_message_content_from_env(),
)
def _owned_capture_modes(
configs: Sequence[OpenTelemetryV2Config],
) -> Iterator[tuple[str, str]]:
for config in configs:
for spec in config.exporters:
if spec.owner is not None:
yield spec.owner.value, config.capture_message_content
def deliverable_destinations(
destinations: Iterable["OtelDestination"],
provider: trace.TracerProvider | None = None,

3
litellm/integrations/otel/plumbing/routing.py
View file

@ -282,7 +282,8 @@ class TenantTracerCache:
self._providers.move_to_end(cache_key)
return cached
built: Final = build_tracer_provider(
self._routed_config(credential_headers, project_headers, endpoint, service_name)
self._routed_config(credential_headers, project_headers, endpoint, service_name),
strip_message_content=not self._config.capture_span_content,
)
self._providers[cache_key] = built
return built

2
litellm/integrations/otel/presets/destinations.py
View file

@ -187,6 +187,7 @@ def destination_for(
callback_name: str,
params: StandardCallbackDynamicParams,
service_name: str | None = None,
capture_message_content: str | None = None,
) -> OtelDestination | None:
"""The destination ``params`` names for ``callback_name``, or ``None``.
@ -217,6 +218,7 @@ def destination_for(
callback_name=callback_name,
protocol=protocol,
span_scope=_span_scope(callback_name, params),
capture_message_content=capture_message_content,
success_sampling_rate=success_sampling_rate,
error_sampling_rate=error_sampling_rate,
)

7
litellm/integrations/otel/presets/newrelic.py
View file

@ -9,6 +9,7 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
from litellm._logging import verbose_logger
from litellm.integrations.otel.model.config import (
CaptureMessageContent,
ExporterOwner,
ExporterSpec,
OpenTelemetryV2Config,
@ -66,7 +67,11 @@ def newrelic_preset(
# New Relic ingests the OTLP GenAI semantic conventions natively.
"mapper_names": ensure_mappers(base.mapper_names, "genai"),
**(
{"capture_message_content": ("span_only" if settings.record_content else "no_content")}
{
"capture_message_content": (
CaptureMessageContent.SPAN_ONLY if settings.record_content else CaptureMessageContent.NO_CONTENT
)
}
if settings.record_content is not None
else {}
),

2
litellm/integrations/prometheus.py
View file

@ -53,7 +53,6 @@ from litellm.models.end_user import LiteLLM_EndUserTable
from litellm.models.team import LiteLLM_TeamTable
from litellm.models.user import LiteLLM_UserTable
from litellm.models.verification_token import LiteLLM_DeletedVerificationToken
from litellm.proxy._types import UserAPIKeyAuth
from litellm.repositories.base_repository import BaseRepository
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.organization_repository import OrganizationRepository
@ -67,6 +66,7 @@ from litellm.types.integrations.prometheus import (
sanitize_prometheus_label_value,
validate_prometheus_deployment_and_latency_caller_identity,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.proxy.carried_budget_state import (
KeyBudgetSnapshot,
OrgBudgetSnapshot,

2
litellm/integrations/rubrik.py
View file

@ -44,7 +44,7 @@ if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
_WEBHOOK_PATH_RESPONSE_MODERATION: Final = "/v1/after_completion/openai/v1"
_WEBHOOK_PATH_PROMPT_MODERATION: Final = "/v1/before_prompt/openai/v1"

2
litellm/integrations/shadow_eval_logger.py
View file

@ -631,13 +631,13 @@ async def _key_or_team_is_over_budget(metadata: Mapping[str, object]) -> bool:
"""
try:
from litellm.exceptions import BudgetExceededError
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import (
get_team_object,
team_max_budget_check,
virtual_key_max_budget_check,
)
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
except ImportError:
return False

2
litellm/integrations/websearch_interception/handler.py
View file

@ -77,10 +77,10 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper

2
litellm/litellm_core_utils/health_check_helpers.py
View file

@ -122,8 +122,8 @@ class HealthCheckHelpers:
2. `user_api_key_auth`: This helps identify health check calls in the DB.
We need this since the DB requires an API Key to track a log in the SpendLogs Table
"""
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
_metadata_variable_name: Final = "litellm_metadata"
litellm_metadata: Final = HealthCheckHelpers._get_metadata_for_health_check_call()

8
litellm/litellm_core_utils/initialize_dynamic_callback_params.py
View file

@ -6,6 +6,7 @@ from typing import Any, Final
from litellm.types.utils import (
ARIZE_OTLP_PROTOCOLS,
CAPTURE_MESSAGE_CONTENT_VALUES,
OTEL_SPAN_SCOPES,
TRUSTED_CALLBACK_VARS_FIELD,
StandardCallbackDynamicParams,
@ -91,6 +92,13 @@ def validate_arize_otlp_protocol_value(value: str) -> None:
raise ValueError(f"Invalid arize_otlp_protocol {value!r}: must be one of {sorted(ARIZE_OTLP_PROTOCOLS)}")
def validate_capture_message_content_value(value: str) -> None:
if value not in CAPTURE_MESSAGE_CONTENT_VALUES:
raise ValueError(
f"Invalid capture_message_content {value!r}: must be one of {sorted(CAPTURE_MESSAGE_CONTENT_VALUES)}"
)
# Hardcoded list of supported callback params to avoid runtime inspection issues with TypedDict
_supported_callback_params: Final[tuple[str, ...]] = (
"langfuse_public_key",

93
litellm/litellm_core_utils/resource_ownership.py Normal file
View file

@ -0,0 +1,93 @@
from typing import Final
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.proxy.auth.user_roles import LitellmUserRoles
def is_proxy_admin(user_api_key_dict: UserAPIKeyAuth | None) -> bool:
if user_api_key_dict is None:
return False
return (
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
)
def get_resource_owner_scopes(
user_api_key_dict: UserAPIKeyAuth | None,
) -> list[str]:
"""
Return ownership scopes that may access a user-created proxy resource.
Raw user_id is included for rows created before scope prefixes existed.
Prefixes avoid collisions when falling back to team/org/key ownership
for keys that do not have a user_id.
Identity-less callers (no user_id, team_id, org_id, api_key, or token)
return ``[]`` — they share no scope with any other caller, so access
checks against an existing owner always fail and creates that depend
on a primary scope must reject up front. Returning a shared sentinel
here would let any two identity-less callers see each other's data.
"""
if user_api_key_dict is None:
return []
scopes: Final[list[str]] = []
def _add(scope: str | None) -> None:
if scope and scope not in scopes:
scopes.append(scope)
if user_api_key_dict.user_id:
_add(user_api_key_dict.user_id)
_add(f"user:{user_api_key_dict.user_id}")
if user_api_key_dict.team_id:
_add(f"team:{user_api_key_dict.team_id}")
if user_api_key_dict.org_id:
_add(f"org:{user_api_key_dict.org_id}")
if user_api_key_dict.api_key:
_add(f"key:{user_api_key_dict.api_key}")
if user_api_key_dict.token:
_add(f"key:{user_api_key_dict.token}")
return scopes
def get_primary_resource_owner_scope(
user_api_key_dict: UserAPIKeyAuth | None,
) -> str | None:
"""Return the canonical owner scope to stamp on newly-created rows.
``None`` for identity-less callers — callers that depend on a primary
scope to record ownership must surface that as a hard error rather
than fall back to a shared sentinel (which would collapse every
identity-less caller into the same logical owner).
"""
if user_api_key_dict is None:
return None
if user_api_key_dict.user_id:
return user_api_key_dict.user_id
if user_api_key_dict.team_id:
return f"team:{user_api_key_dict.team_id}"
if user_api_key_dict.org_id:
return f"org:{user_api_key_dict.org_id}"
if user_api_key_dict.api_key:
return f"key:{user_api_key_dict.api_key}"
if user_api_key_dict.token:
return f"key:{user_api_key_dict.token}"
return None
def user_can_access_resource_owner(
owner: str | None,
user_api_key_dict: UserAPIKeyAuth | None,
) -> bool:
if user_api_key_dict is None:
return True
if is_proxy_admin(user_api_key_dict):
return True
if owner is None:
return False
return owner in get_resource_owner_scopes(user_api_key_dict)

2
litellm/llms/a2a/chat/guardrail_translation/handler.py
View file

@ -26,7 +26,7 @@ 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.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class _A2ATextPart(TypedDict, total=False):

2
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -76,10 +76,10 @@ if TYPE_CHECKING:
ModifyResponseException,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
@dataclass(frozen=True, slots=True)

2
litellm/llms/anthropic/pass_through/adapters/handler.py
View file

@ -33,8 +33,8 @@ from litellm.types.utils import ModelResponse
from litellm.utils import get_model_info
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.router import Router
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
# Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge.
ANTHROPIC_ONLY_REQUEST_KEYS: Final[frozenset[str]] = frozenset({"output_config", "safeguards"})

2
litellm/llms/anthropic/pass_through/context_management/dispatcher.py
View file

@ -13,8 +13,8 @@ from .editors import apply_clear_tool_uses_20250919, apply_compact_20260112
from .result import PolyfillResult
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.router import Router
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
EditorResult: TypeAlias = "PolyfillResult | tuple[list[dict[str, object]], AppliedEdit | None]"

2
litellm/llms/anthropic/pass_through/context_management/editors/compact.py
View file

@ -29,7 +29,6 @@ from litellm.types.llms.anthropic import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
RateLimitDescriptor,
RateLimitDescriptorRateLimitObject,
@ -41,6 +40,7 @@ if TYPE_CHECKING:
AllAnthropicToolsValues,
)
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.utils import ModelResponse
from ..constants import (

2
litellm/llms/base_llm/files/transformation.py
View file

@ -22,9 +22,9 @@ from ..chat.transformation import BaseConfig
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
from litellm.proxy._types import UserAPIKeyAuth as _UserAPIKeyAuth
from litellm.router import Router as _Router
from litellm.types.llms.openai import HttpxBinaryResponseContent
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth as _UserAPIKeyAuth
LiteLLMLoggingObj = _LiteLLMLoggingObj
Span = Any

2
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -11,8 +11,8 @@ if TYPE_CHECKING:
ModifyResponseException,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import AllMessageValues
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
@dataclass(slots=True)

2
litellm/llms/base_llm/managed_resources/base_managed_resource.py
View file

@ -16,7 +16,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
can_access_resource,
resolve_resource_owner_id,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.utils import SpecialEnums
MANAGED_RESOURCES_TARGET: Final = "managed_resources"

6
litellm/llms/base_llm/managed_resources/isolation.py
View file

@ -12,10 +12,10 @@ are denied so an empty user_id can never select an unscoped query.
from typing import Any, Final
from litellm.proxy._types import (
from litellm.types.proxy.auth.user_api_key_auth import (
UserAPIKeyAuth,
)
from litellm.proxy._types import (
from litellm.types.proxy.auth.user_api_key_auth import (
user_api_key_has_admin_view as _user_has_admin_view,
)
@ -29,7 +29,7 @@ def resolve_resource_owner_id(
A key with neither a user_id nor a team_id would otherwise stamp
``created_by=None`` and be locked out of its own resources, so it owns
them under its hashed token instead, using the ``key:`` scope prefix
already used by ``proxy/common_utils/resource_ownership.py``. ``None``
already used by ``litellm_core_utils/resource_ownership.py``. ``None``
means the caller has no usable identity of its own and must fall back
to team scoping, or be denied.
"""

18
litellm/llms/bedrock/passthrough/guardrail_translation/handler.py
View file

@ -17,8 +17,8 @@ if TYPE_CHECKING:
from litellm.llms.pass_through.guardrail_translation.handler import (
PassThroughEndpointHandler,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"})
_EVENT_STREAM_CONTENT_TYPE: Final = "vnd.amazon.eventstream"
@ -206,12 +206,14 @@ class _TextContainer(Protocol):
_DeltaHolder = tuple[_GroupKey, _TextContainer, str]
_StreamGroupKey: TypeAlias = tuple[int | None, _GroupKey]
class _StreamFrame(TypedDict):
"""One raw event-stream frame plus the guardrailable texts it carries."""
raw: ReadOnly[bytes]
texts: ReadOnly[Sequence[tuple[_GroupKey, str]]]
texts: ReadOnly[Sequence[tuple[_StreamGroupKey, str]]]
def _unpack_uint32(buffer: bytes) -> int:
@ -302,12 +304,14 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
frames.append({"raw": frame_raw, "texts": []})
continue
texts: list[tuple[_GroupKey, str]] = []
texts: list[tuple[_StreamGroupKey, str]] = []
if event_type == "contentBlockDelta":
try:
payload_dict: dict[str, object] = _json.loads(payload_bytes)
block_index = payload_dict.get("contentBlockIndex")
# Group per content block: a guardrail edit that changes length must not move text between blocks.
texts = [
(group_key, container[key])
((block_index if isinstance(block_index, int) else None, group_key), container[key])
for group_key, container, key in _collect_stream_delta_text_holders(payload_dict.get("delta"))
]
except Exception as e:
@ -321,9 +325,9 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
trailing_bytes: Final = body_bytes[offset:]
group_order: Final[list[_GroupKey]] = []
group_members: Final[dict[_GroupKey, list[tuple[int, int]]]] = {}
group_texts: Final[dict[_GroupKey, list[str]]] = {}
group_order: Final[list[_StreamGroupKey]] = []
group_members: Final[dict[_StreamGroupKey, list[tuple[int, int]]]] = {}
group_texts: Final[dict[_StreamGroupKey, list[str]]] = {}
for frame_idx, frame in enumerate(frames):
for local_idx, (group_key, text) in enumerate(frame["texts"]):
if group_key not in group_members:

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

@ -213,7 +213,6 @@ if TYPE_CHECKING:
FakeAnthropicMessagesStreamIterator,
)
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.google_genai.main import GenerateContentResponse
from litellm.types.llms.openai_evals import (
CancelEvalResponse,
@ -225,6 +224,7 @@ if TYPE_CHECKING:
Run,
RunDeleteResponse,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:

2
litellm/llms/gemini/google_genai/guardrail_translation/handler.py
View file

@ -25,7 +25,7 @@ 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.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
_EMPTY_REQUEST_DATA: Final[Mapping[str, object]] = MappingProxyType({})

14
litellm/llms/litellm_proxy/skills/handler.py
View file

@ -11,19 +11,19 @@ from typing import Final
from litellm._logging import verbose_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.llms.litellm_proxy.skills.constants import (
LITELLM_SKILL_ID_PREFIX,
MAX_SKILLS_PER_SEARCH,
)
from litellm.models.skills import LiteLLM_SkillsTable, NewSkillRequest
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.resource_ownership import (
from litellm.litellm_core_utils.resource_ownership import (
get_primary_resource_owner_scope,
get_resource_owner_scopes,
is_proxy_admin,
user_can_access_resource_owner,
)
from litellm.llms.litellm_proxy.skills.constants import (
LITELLM_SKILL_ID_PREFIX,
MAX_SKILLS_PER_SEARCH,
)
from litellm.models.skills import LiteLLM_SkillsTable, NewSkillRequest
from litellm.repositories.table_repositories import SkillsRepository
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
# Skills are looked up on every chat completion that has skills enabled
# (`SkillsInjectionHook` calls ``fetch_skill_from_db``). 60s LRU/TTL cache

2
litellm/llms/litellm_proxy/skills/skill_search.py
View file

@ -21,9 +21,9 @@ from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.models.skills import LiteLLM_SkillsTable
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
DEFAULT_SKILL_SEARCH_TOP_K: Final = 5
MAX_SKILL_SEARCH_TOP_K: Final = 100

2
litellm/llms/litellm_proxy/skills/transformation.py
View file

@ -22,7 +22,7 @@ from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.models.skills import LiteLLM_SkillsTable
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class LiteLLMSkillsTransformationHandler:

2
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -69,7 +69,7 @@ if TYPE_CHECKING:
ModifyResponseException,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class OpenAIChatCompletionsHandler(BaseTranslation):

2
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -102,8 +102,8 @@ if TYPE_CHECKING:
ModifyResponseException,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import ResponseInputParam
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class _ToolCallShape(NamedTuple):

2
litellm/llms/openai/videos/guardrail_translation/handler.py
View file

@ -6,7 +6,7 @@ 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.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class OpenAIVideoGenerationHandler(BaseTranslation):

2
litellm/llms/pass_through/guardrail_translation/handler.py
View file

@ -18,8 +18,8 @@ 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.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class PassThroughEndpointHandler(BaseTranslation):

3
litellm/llms/scaledown/__init__.py Normal file
View file

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

3
litellm/llms/scaledown/chat/__init__.py Normal file
View file

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

648
litellm/llms/scaledown/chat/transformation.py Normal file
View file

@ -0,0 +1,648 @@
import json
import time
import uuid
from collections.abc import Mapping, Sequence
from functools import reduce
from typing import TYPE_CHECKING, Final, NoReturn
from urllib.parse import unquote
import httpx
from pydantic import JsonValue, TypeAdapter
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import Choices, Message, ModelResponse, Usage
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
JSON_MESSAGES: Final = TypeAdapter(list[dict[str, JsonValue]])
DEFAULT_API_BASE: Final = "https://api.scaledown.xyz"
DOMAIN_MODELS: Final = frozenset({"classify", "extract", "summarize", "compress"})
DECISIONS_MODELS: Final = frozenset({"decisions"})
DECISIONS_UPSTREAM_MODEL: Final = "classify-1"
NATIVE_PATHS: Final[Mapping[str, str]] = {
"classify": "/classify",
"extract": "/extract",
"summarize": "/summarization/abstractive",
"compress": "/compress/raw/",
}
DECISION_QUESTION_TYPES: Final = frozenset({"choice", "noul", "score"})
EXTRA_BODY_KEYS: Final[Mapping[str, frozenset[str]]] = {
"extract": frozenset({"threshold", "top_n"}),
"summarize": frozenset(),
"compress": frozenset({"compression_rate"}),
"classify": frozenset({"labels"}),
"decisions": frozenset({"questions"}),
}
# Set by the LiteLLM router and proxy on every call; handled by LiteLLM, never sent upstream.
ROUTER_PARAMS: Final = frozenset({"max_retries", "stream_options"})
SCORE_MIN_LEVELS: Final = 2
SCORE_MAX_LEVELS: Final = 10
MAX_ENTITIES: Final = 1000
MAX_REF_CHAIN: Final = 32
MAX_SCHEMA_DEPTH: Final = 32
class ScaleDownError(BaseLLMException):
pass
def _reject(message: str, status_code: int = 400) -> NoReturn:
"""The single place this adapter raises its own errors."""
raise ScaleDownError(status_code=status_code, message=message)
class ScaleDownChatConfig(BaseConfig):
@property
def custom_llm_provider(self) -> str | None:
return "scaledown"
def validate_environment(
self,
headers: dict[str, str], # mutable-ok: BaseConfig signature
model: str,
messages: list[AllMessageValues], # mutable-ok: BaseConfig signature
optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
api_key: str | None = None,
api_base: str | None = None,
) -> dict[str, JsonValue]: # mutable-ok: BaseConfig signature
resolved_key: Final = api_key or get_secret_str("SCALEDOWN_API_KEY")
if resolved_key is None:
_reject("Missing ScaleDown API key. Set SCALEDOWN_API_KEY or pass api_key to the call.", 401)
if not api_key and api_base is not None and _root(api_base) != _trusted_root():
_reject(
"A custom api_base needs its own api_key: the SCALEDOWN_API_KEY from the environment is only "
"sent to the default host or the host in SCALEDOWN_API_BASE."
)
return {**headers, "x-api-key": resolved_key, "content-type": "application/json"}
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
stream: bool | None = None,
) -> str:
base: Final = _root(api_base or _trusted_root())
operation: Final = _operation(model)
if operation in DECISIONS_MODELS:
return f"{base}/v1/scaledown"
return f"{base}{NATIVE_PATHS[operation]}"
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseConfig signature
operation: Final = _operation(model)
if operation == "extract":
return ["response_format", "stream"]
if operation == "summarize":
return ["max_tokens", "max_completion_tokens", "stream"]
return ["stream"]
def map_openai_params(
self,
non_default_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
model: str,
drop_params: bool,
) -> dict[str, JsonValue]: # mutable-ok: BaseConfig signature
supported: Final = frozenset(self.get_supported_openai_params(model))
unsupported: Final = sorted(key for key in non_default_params if key not in supported | ROUTER_PARAMS)
if unsupported and not drop_params:
_reject(
f"ScaleDown model '{model}' does not support {unsupported}. "
f"Supported: {sorted(supported)}. Set litellm.drop_params=True to ignore."
)
mapped: Final = {
("max_tokens" if key == "max_completion_tokens" else key): value
for key, value in non_default_params.items()
if key in supported
}
return {**optional_params, **mapped}
def should_fake_stream(
self, model: str | None, stream: bool | None, custom_llm_provider: str | None = None
) -> bool:
return bool(stream)
def transform_extra_body(
self,
extra_body: Mapping[str, object],
request: Mapping[str, object],
model: str,
litellm_params: Mapping[str, object],
) -> Mapping[str, JsonValue]:
# extra_body is merged after guardrails ran on the messages, so it may only add options.
# Anything that carries prompt text or picks the model is refused.
options: Final = JSON_OBJECT.validate_python(extra_body)
operation: Final = _operation(model)
allowed: Final = EXTRA_BODY_KEYS[operation]
unexpected: Final = sorted(set(options) - allowed)
if unexpected:
allowance: Final = f"may only set {sorted(allowed)}" if allowed else "is not accepted"
_reject(
f"extra_body for scaledown/{operation} {allowance}, got {unexpected}. "
"Text and instructions must be passed as messages."
)
if operation == "compress":
return {"scaledown": {"rate": options["compression_rate"]}} if "compression_rate" in options else {}
if "questions" in options:
_validate_questions(options["questions"])
if "labels" in options:
_validate_labels(options["labels"])
return dict(options)
def sign_request(
self,
headers: dict[str, str], # mutable-ok: BaseConfig signature
optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
request_data: dict[str, JsonValue], # mutable-ok: BaseConfig signature
api_base: str,
api_key: str | None = None,
model: str | None = None,
stream: bool | None = None,
fake_stream: bool | None = None,
) -> tuple[dict[str, str], bytes | None]: # mutable-ok: BaseConfig signature
# Runs after extra_body is merged, so this is where a decisions body is known to be complete.
if model is not None and _operation(model) in DECISIONS_MODELS:
_require_complete_decisions(request_data)
if model is not None and _operation(model) == "classify":
_validate_labels(request_data.get("labels"))
return headers, None
def transform_request(
self,
model: str,
messages: list[AllMessageValues], # mutable-ok: BaseConfig signature
optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
headers: dict[str, str], # mutable-ok: BaseConfig signature
) -> dict[str, JsonValue]: # mutable-ok: BaseConfig signature
checked_messages: Final = JSON_MESSAGES.validate_python(messages)
operation: Final = _operation(model)
_validate_text_input(operation, checked_messages, optional_params)
if operation == "classify":
return {
"text": _last_user_text(operation, checked_messages),
**({"labels": optional_params["labels"]} if "labels" in optional_params else {}),
}
if operation in DECISIONS_MODELS:
return dict(_decisions_request(checked_messages, optional_params))
if operation == "extract":
return dict(_extract_request(checked_messages, optional_params))
if operation == "summarize":
return dict(_summarize_request(checked_messages, optional_params))
return dict(_compress_request(checked_messages, optional_params))
def transform_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ModelResponse,
logging_obj: "LiteLLMLoggingObj",
request_data: dict[str, JsonValue], # mutable-ok: BaseConfig signature
messages: list[AllMessageValues], # mutable-ok: BaseConfig signature
optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
encoding: "Tokenizer | None",
api_key: str | None = None,
json_mode: bool | None = None,
) -> ModelResponse:
try:
raw: Final = JSON_OBJECT.validate_json(raw_response.content)
except ValueError:
_reject(
f"ScaleDown returned a non-JSON response: {raw_response.text[:500]}",
raw_response.status_code if raw_response.is_error else 502,
)
operation: Final = _operation(model)
if operation in DECISIONS_MODELS:
return _decisions_response(operation, raw, model_response)
return _native_response(operation, raw, model_response, request_data.get("entities"))
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, str] | httpx.Headers, # mutable-ok: BaseConfig signature
) -> BaseLLMException:
return ScaleDownError(status_code=status_code, message=error_message, headers=headers)
def _root(api_base: str) -> str:
return api_base.rstrip("/").removesuffix("/v1")
def _trusted_root() -> str:
return _root(get_secret_str("SCALEDOWN_API_BASE") or DEFAULT_API_BASE)
def _operation(model: str) -> str:
operation: Final = model.split("/", 1)[1] if "/" in model else model
if operation not in DOMAIN_MODELS | DECISIONS_MODELS:
_reject(f"Unknown ScaleDown model '{model}'. Choose from {sorted(DOMAIN_MODELS | DECISIONS_MODELS)}.")
return operation
def _text_of(message: Mapping[str, JsonValue]) -> str | None:
content: Final = message.get("content")
if isinstance(content, str):
return content
if isinstance(content, list):
parts: Final = [
str(part.get("text", "")) for part in content if isinstance(part, dict) and part.get("type") == "text"
]
return "\n".join(parts) if parts else None
return None
def _validate_text_input(
operation: str, messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue]
) -> None:
if any(key in optional_params for key in ("state", "document", "document_mime_type", "text")):
_reject("ScaleDown accepts text through messages only. Native state and document inputs are not supported.")
for message in messages:
if operation in {"classify", "decisions", "extract"} and message.get("role") in {"system", "developer"}:
_reject(
f"ScaleDown {operation} does not support system messages. Put instructions in the field definitions."
)
if _has_non_text_content(message):
_reject("ScaleDown currently supports text only. Image and document inputs are not supported.")
def _has_non_text_content(message: Mapping[str, JsonValue]) -> bool:
content: Final = message.get("content")
return isinstance(content, list) and any(
not isinstance(part, dict) or part.get("type") != "text" for part in content
)
def _last_user_text(operation: str, messages: Sequence[Mapping[str, JsonValue]]) -> str:
message: Final = next((message for message in reversed(messages) if message.get("role") == "user"), None)
text: Final = _text_of(message) if message is not None else None
if not text:
_reject(f"ScaleDown {operation} requires text in the last user message.")
return text
def _validate_labels(labels: JsonValue) -> None:
if not isinstance(labels, list) or not labels:
_reject("ScaleDown classify requires a non-empty 'labels' list of objects with 'name' and 'rubric' strings.")
for label in labels:
if not isinstance(label, Mapping) or not all(
isinstance(label.get(key), str) and label[key] for key in ("name", "rubric")
):
_reject("Each classify label needs non-empty 'name' and 'rubric' strings.")
def _child(node: JsonValue, key: str) -> JsonValue:
return node.get(key.replace("~1", "/").replace("~0", "~")) if isinstance(node, Mapping) else None
def _extend_ref_path(path: tuple[str, ...], ref: str) -> tuple[str, ...]:
return (*path, ref)
def _nullable_branch(schema: Mapping[str, JsonValue]) -> Mapping[str, JsonValue] | None:
if "allOf" in schema:
_reject("ScaleDown extraction does not support allOf. Supply the object's properties directly.")
variants: Final = schema.get("anyOf", schema.get("oneOf"))
if variants is None:
return None
if any(key in schema for key in ("properties", "items", "$ref", "type")) or (
"anyOf" in schema and "oneOf" in schema
):
_reject("ScaleDown extraction cannot combine nullable branches with sibling schema definitions.")
if not isinstance(variants, list) or len(variants) != 2:
_reject("ScaleDown extraction supports anyOf/oneOf only for one schema plus null.")
branches: Final = tuple(branch for branch in variants if isinstance(branch, dict))
non_null: Final = tuple(branch for branch in branches if branch.get("type") != "null")
if len(branches) != 2 or len(non_null) != 1:
_reject("ScaleDown extraction supports anyOf/oneOf only for one schema plus null.")
return {
**non_null[0],
**{key: value for key, value in schema.items() if key not in {"anyOf", "oneOf"}},
}
def _deref(
prop: Mapping[str, JsonValue], root: Mapping[str, JsonValue], path: tuple[str, ...]
) -> tuple[Mapping[str, JsonValue], tuple[str, ...], bool]:
"""Follow nullable branches and local `$ref`s (`#/$defs/X`, `#/definitions/X`).
Returns the schema, the refs followed so far, and whether a ref repeated (a cycle).
The walk is a bounded loop, so a long chain of distinct refs cannot exhaust the stack.
"""
current: Mapping[str, JsonValue] = prop # rebind-ok: bounded walk along a ref chain
followed: tuple[str, ...] = path # rebind-ok: bounded walk along a ref chain
for _ in range(MAX_REF_CHAIN):
nullable: Mapping[str, JsonValue] | None = _nullable_branch(current)
ref: JsonValue = current.get("$ref")
if nullable is not None:
current = nullable
continue
if "$ref" not in current:
return current, followed, False
if not isinstance(ref, str) or not ref.startswith("#"):
_reject("ScaleDown extraction supports only local JSON Pointer $refs.")
pointer: str = unquote(ref[1:])
if not pointer.startswith("/"):
_reject("ScaleDown extraction supports only local JSON Pointer $refs beginning with #/.")
if any(key in current for key in ("properties", "items", "type")):
_reject("ScaleDown extraction cannot combine $ref with sibling schema definitions.")
if ref in followed:
return current, followed, True
target: JsonValue = reduce(_child, pointer[1:].split("/"), dict(root))
if not isinstance(target, Mapping):
_reject(f"ScaleDown extraction cannot resolve $ref '{ref}' to a schema object.")
current, followed = target, _extend_ref_path(followed, ref)
_reject(f"The response_format schema follows more than {MAX_REF_CHAIN} chained $refs or nullable branches.")
def _build_properties(
properties: Mapping[str, JsonValue],
root: Mapping[str, JsonValue],
path: tuple[str, ...],
budget: int,
depth: int = 0,
) -> tuple[dict[str, JsonValue], int]: # mutable-ok: serialized JSON entity map
"""Turn JSON-schema properties into the /extract entity map, spending one budget unit per entity.
Each property becomes an entity whose value is its description (or its name when there is
none). Nested objects stay nested, arrays of objects become a one-element list holding the
item's entity map, local `$ref`s are followed, and a ref that points back into itself is cut
off at that point instead of expanded.
"""
if depth > MAX_SCHEMA_DEPTH:
_reject(f"The response_format schema nests deeper than {MAX_SCHEMA_DEPTH} levels.")
built: Final[dict[str, JsonValue]] = {} # mutable-ok: filled once while threading the budget through
remaining = budget # rebind-ok: the budget is threaded through the loop
for name, prop in properties.items():
if remaining <= 0:
_reject(
f"The response_format schema expands to more than {MAX_ENTITIES} entities; "
"remove recursive or heavily repeated $refs."
)
built[name], remaining = _build_entity(name, prop, root, path, remaining - 1, depth)
return built, remaining
def _build_entity(
name: str, prop: JsonValue, root: Mapping[str, JsonValue], path: tuple[str, ...], budget: int, depth: int
) -> tuple[JsonValue, int]:
schema, followed, cyclic = _deref(prop if isinstance(prop, Mapping) else {}, root, path)
leaf: Final = schema.get("description") or name
nested: Final = schema.get("properties")
if cyclic:
return leaf, budget
if isinstance(nested, Mapping):
return _build_properties(nested, root, followed, budget, depth + 1)
raw_items: Final = schema.get("items")
schema_type: Final = schema.get("type")
is_array: Final = schema_type == "array" or (isinstance(schema_type, list) and "array" in schema_type)
if is_array and isinstance(raw_items, Mapping):
items, item_path, item_cyclic = _deref(raw_items, root, followed)
item_properties: Final = items.get("properties")
if not item_cyclic and isinstance(item_properties, Mapping):
built, remaining = _build_properties(item_properties, root, item_path, budget, depth + 1)
return [built], remaining
return leaf, budget
def _extract_request(
messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue]
) -> Mapping[str, JsonValue]:
response_format: Final = optional_params.get("response_format")
json_schema: Final = response_format.get("json_schema") if isinstance(response_format, Mapping) else None
if isinstance(json_schema, Mapping) and json_schema.get("strict") is True:
_reject(
"ScaleDown extract does not support strict JSON schemas. Use strict=false and validate the returned fields."
)
schema: Final = json_schema.get("schema") if isinstance(json_schema, Mapping) else None
properties: Final = _deref(schema, schema, ())[0].get("properties") if isinstance(schema, Mapping) else None
if not isinstance(schema, dict) or not isinstance(properties, dict) or not properties:
_reject(
"ScaleDown extract needs response_format={'type': 'json_schema', ...} with a non-empty "
"'properties' map; each property becomes an entity and its description the extraction hint."
)
return {
"text": _last_user_text("extract", messages),
"entities": _build_properties(properties, schema, (), MAX_ENTITIES)[0],
**{key: optional_params[key] for key in ("threshold", "top_n") if key in optional_params},
}
def _summarize_request(
messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue]
) -> Mapping[str, JsonValue]:
instructions: Final = [
text for message in messages if message.get("role") in {"system", "developer"} and (text := _text_of(message))
]
return {
"text": _last_user_text("summarize", messages),
**({"instructions": "\n".join(instructions)} if instructions else {}),
**({"max_tokens": optional_params["max_tokens"]} if "max_tokens" in optional_params else {}),
}
def _compress_request(
messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue]
) -> Mapping[str, JsonValue]:
prompt: Final = _last_user_text("compress", messages)
last_user_index: Final = max(i for i, m in enumerate(messages) if m.get("role") == "user" and _text_of(m))
context: Final = "\n\n".join(text for i, m in enumerate(messages) if i != last_user_index and (text := _text_of(m)))
return {
"context": context,
"prompt": prompt,
"scaledown": {"rate": optional_params.get("compression_rate", "auto")},
}
def _decisions_request(
messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue]
) -> Mapping[str, JsonValue]:
questions: Final = optional_params.get("questions")
if questions is not None:
_validate_questions(questions)
return {
"model": DECISIONS_UPSTREAM_MODEL,
"state": {"text": _last_user_text("decisions", messages)},
**({"questions": questions} if questions is not None else {}),
}
def _require_complete_decisions(body: Mapping[str, JsonValue]) -> None:
if not body.get("questions"):
_reject(
"ScaleDown decisions requires a non-empty 'questions' map, passed as an extra parameter, "
'e.g. {"questions": {"category": {"type": "choice", "criteria": {"billing": "..."}}}}. '
"Question types are 'choice', 'noul', and 'score'."
)
def _validate_questions(questions: JsonValue) -> None:
if not isinstance(questions, Mapping):
_reject(f"'questions' must be a map of name to question, got {type(questions).__name__}.")
for name, question in questions.items():
if not isinstance(question, Mapping):
_reject(f"Question '{name}' must be an object.")
question_type = question.get("type")
if not isinstance(question_type, str) or question_type not in DECISION_QUESTION_TYPES:
_reject(f"Question '{name}' has type {question_type!r}; expected one of {sorted(DECISION_QUESTION_TYPES)}.")
criteria = question.get("criteria")
if question_type == "choice" and (not isinstance(criteria, Mapping) or not criteria):
_reject(f"Choice question '{name}' needs a non-empty 'criteria' map of option key to description.")
if question_type == "score" and (
not isinstance(criteria, list) or not SCORE_MIN_LEVELS <= len(criteria) <= SCORE_MAX_LEVELS
):
_reject(
f"Score question '{name}' needs 'criteria' as an ordered list of "
f"{SCORE_MIN_LEVELS} to {SCORE_MAX_LEVELS} level descriptions, lowest to highest."
)
def _unwrap(value: JsonValue) -> JsonValue:
wrapper_keys: Final = {"_value", "_span_anchor"}
if isinstance(value, Mapping) and "_value" in value and set(value) <= wrapper_keys:
return value["_value"]
return value
def _clean_extraction(value: JsonValue, requested: JsonValue) -> JsonValue:
"""Reduce an extraction result to the fields the caller's schema asked for.
The request's entity map says which keys are real fields, so provider-added keys such as
`<field>_span_anchor` and `{"_value": ..., "_span_anchor": ...}` wrappers are dropped without
touching a requested field that happens to share one of those names. The untouched payload
stays on `_hidden_params["scaledown_response"]`.
"""
if isinstance(requested, Mapping) and isinstance(value, Mapping):
return {key: _clean_extraction(value[key], requested[key]) for key in requested if key in value}
if isinstance(requested, list) and requested and isinstance(value, list):
return [_clean_extraction(item, requested[0]) for item in value]
return _unwrap(value)
def _extracted_fields(raw: Mapping[str, JsonValue], requested: JsonValue) -> JsonValue:
entities: Final = raw.get("entities")
matches: Final = (
tuple(entity for entity in entities if isinstance(entity, Mapping)) if isinstance(entities, list) else ()
)
scalar_names: Final = (
tuple(name for name, hint in requested.items() if isinstance(hint, str))
if isinstance(requested, Mapping)
else ()
)
scalars: Final = {
name: next(entity["text"] for entity in matches if entity.get("type") == name and "text" in entity)
for name in scalar_names
if any(entity.get("type") == name and "text" in entity for entity in matches)
}
structured: Final = raw.get("structured_result")
return _clean_extraction({**scalars, **(structured if isinstance(structured, Mapping) else {})}, requested)
def _native_response(
operation: str, raw: JsonValue, model_response: ModelResponse, requested: JsonValue = None
) -> ModelResponse:
"""Wrap a native /extract, /summarization/abstractive or /compress/raw/ payload.
Summarize and compress return the upstream payload as JSON on choices[0].message.content.
Extract returns the fields requested by response_format without enforcing schema types.
The native APIs report input tokens only; they return no output token
count, so completion_tokens stays 0 because it is unmeasured, not because nothing was generated.
"""
if not isinstance(raw, Mapping) or not raw:
_reject(
f"ScaleDown '{operation}' returned an empty or malformed response: {json.dumps(raw)[:500]}",
500,
)
results: Final = raw.get("results")
input_tokens: Final = raw.get("input_tokens") or (
(results.get("original_prompt_tokens") if isinstance(results, Mapping) else None)
or raw.get("original_prompt_tokens")
if operation == "compress"
else None
)
content: Final = _extracted_fields(raw, requested) if operation == "extract" else raw
model_response.id = f"chatcmpl-{uuid.uuid4().hex}" # rebind-ok: LiteLLM fills the response object it passes in
model_response.created = int(time.time()) # rebind-ok: LiteLLM fills the response object it passes in
model_response.model = f"scaledown/{operation}" # rebind-ok: LiteLLM fills the response object it passes in
model_response.object = "chat.completion" # rebind-ok: LiteLLM fills the response object it passes in
model_response.choices = [ # rebind-ok: LiteLLM fills the response object it passes in
Choices(
index=0,
message=Message(role="assistant", content=json.dumps(content, separators=(",", ":"))),
finish_reason="stop",
)
]
_set_usage(model_response, {"prompt_tokens": input_tokens or 0})
model_response.hidden_params["scaledown_response"] = raw # pyright: ignore[reportUnknownMemberType] # Inherited hidden_params has untyped values. # rebind-ok: supplied response
return model_response
def _decisions_response(operation: str, raw: JsonValue, model_response: ModelResponse) -> ModelResponse:
answers: Final = raw.get("answers") if isinstance(raw, Mapping) else None
if answers is None:
_reject(f"ScaleDown decisions response contained no answers: {json.dumps(raw)[:500]}", 500)
model_response.id = f"chatcmpl-{uuid.uuid4().hex}" # rebind-ok: LiteLLM fills the response object it passes in
model_response.created = int(time.time()) # rebind-ok: LiteLLM fills the response object it passes in
model_response.model = f"scaledown/{operation}" # rebind-ok: LiteLLM fills the response object it passes in
model_response.object = "chat.completion" # rebind-ok: LiteLLM fills the response object it passes in
model_response.choices = [ # rebind-ok: LiteLLM fills the response object it passes in
Choices(
index=0,
message=Message(role="assistant", content=json.dumps(answers, separators=(",", ":"))),
finish_reason="stop",
)
]
usage: Final = raw.get("usage") if isinstance(raw, dict) else None
counts: Final = usage if isinstance(usage, dict) else {}
_set_usage(
model_response,
{
"prompt_tokens": counts.get("input_tokens", 0),
"completion_tokens": counts.get("output_tokens", 0),
},
)
model_response.hidden_params["scaledown_response"] = raw # pyright: ignore[reportUnknownMemberType] # Inherited hidden_params has untyped values. # rebind-ok: supplied response
return model_response
def _set_usage(model_response: ModelResponse, usage: Mapping[str, JsonValue]) -> None:
prompt_tokens: Final = _token_count(usage.get("prompt_tokens"))
completion_tokens: Final = _token_count(usage.get("completion_tokens"))
total_tokens: Final = usage.get("total_tokens")
model_response.usage = Usage( # pyright: ignore[reportAttributeAccessIssue] # ModelResponse stores usage as a dynamic Pydantic field. # rebind-ok: fills the supplied response
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=_token_count(total_tokens) if total_tokens is not None else prompt_tokens + completion_tokens,
)
def _token_count(value: JsonValue) -> int:
if value is None:
return 0
if isinstance(value, int) and value >= 0:
return value
_reject("ScaleDown returned an invalid token count.", 500)

35
litellm/main.py
View file

@ -1938,6 +1938,39 @@ def _complete_ragflow(ctx: CompletionDispatchContext) -> _CompletionDispatchResu
return response
def _complete_scaledown(ctx: CompletionDispatchContext) -> _CompletionDispatchResult:
try:
return base_llm_http_handler.completion(
model=ctx.model,
messages=ctx.messages,
headers=ctx.headers,
model_response=ctx.model_response,
api_key=ctx.api_key,
api_base=ctx.api_base,
acompletion=ctx.acompletion,
logging_obj=ctx.logging,
optional_params=ctx.optional_params,
litellm_params=ctx.litellm_params,
shared_session=ctx.shared_session,
timeout=cast( # cast-ok: CompletionTimeout.resolve normalizes this value
float | httpx.Timeout, ctx.timeout
),
client=_dispatch_client_http(ctx),
custom_llm_provider=ctx.custom_llm_provider,
encoding=_get_encoding(),
stream=ctx.stream,
provider_config=ctx.provider_config,
)
except Exception as exc:
ctx.logging.post_call(
input=ctx.messages,
api_key=ctx.api_key,
original_response=str(exc),
additional_args={"headers": ctx.headers},
)
raise
def _complete_xai(ctx: CompletionDispatchContext) -> _CompletionDispatchResult:
acompletion: Final = ctx.acompletion
api_base: Final = ctx.api_base
@ -5859,6 +5892,8 @@ def completion(
elif custom_llm_provider == "ragflow":
## COMPLETION CALL - RAGFlow uses HTTP handler to support custom URL paths
response = _complete_ragflow(_dispatch_ctx)
elif custom_llm_provider == "scaledown":
response = _complete_scaledown(_dispatch_ctx) # rebind-ok: shared provider return
elif custom_llm_provider == "xai":
## COMPLETION CALL
response = _complete_xai(_dispatch_ctx)

48
litellm/model_prices_and_context_window_backup.json
View file

@ -66876,7 +66876,7 @@
"supports_audio_input": true
},
"azure/gpt-realtime-translate": {
"deprecation_date": "2027-05-06",
"deprecation_date": "2027-11-09",
"input_cost_per_second": 0.000566666666667,
"litellm_provider": "azure",
"max_input_tokens": 32000,
@ -71842,5 +71842,51 @@
"image"
],
"supports_embedding_image_input": true
},
"scaledown/extract": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": false,
"supports_function_calling": false
},
"scaledown/summarize": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_system_messages": true,
"supports_vision": false,
"supports_function_calling": false
},
"scaledown/compress": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_system_messages": true,
"supports_function_calling": false,
"supports_vision": false
},
"scaledown/classify": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_system_messages": false,
"supports_vision": false,
"supports_function_calling": false
},
"scaledown/decisions": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_system_messages": false,
"supports_vision": false,
"supports_function_calling": false
}
}

18
litellm/provider_endpoints_support_backup.json
View file

@ -2064,6 +2064,24 @@
"search": true
}
},
"scaledown": {
"display_name": "ScaleDown (`scaledown`)",
"url": "https://docs.litellm.ai/docs/providers/scaledown",
"endpoints": {
"chat_completions": true,
"messages": false,
"responses": false,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,
"batches": false,
"rerank": false,
"a2a": false,
"interactions": false
}
},
"sambanova": {
"display_name": "Sambanova (`sambanova`)",
"url": "https://docs.litellm.ai/docs/providers/sambanova",

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

@ -95,9 +95,7 @@ def _normalize_caller_admission_credential(value: str) -> str:
def _admission_request(scope: Scope) -> Request:
"""Request whose ``body()`` serves the routing layer's lazily peeked JSON-RPC
bytes (``b"{}"`` when no peek callable was stashed) instead of the ASGI
receive channel."""
"""Request whose ``body()`` serves the peeked JSON-RPC method (``b"{}"`` when nothing is stashed)."""
request: Final = Request(scope=scope)
peeked_body: Final[Callable[[], Awaitable[bytes]] | None] = scope.get(MCP_PEEKED_BODY_SCOPE_KEY)

4
litellm/proxy/_experimental/mcp_server/legacy_callbacks.py
View file

@ -27,6 +27,10 @@ class ElicitationCallback(Protocol):
async def __call__(self, context: object, params: ElicitRequestParams) -> ElicitResult | ErrorData: ...
def has_active_mcp_request_context() -> bool:
return get_active_mcp_request_ctx() is not None
def create_sampling_callback(
user_api_key_auth: UserAPIKeyAuth | None = None,
raw_headers: Mapping[str, str] | None = None,

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

@ -24,7 +24,7 @@ from collections.abc import (
MutableMapping,
Sequence,
)
from contextlib import ExitStack, asynccontextmanager
from contextlib import AbstractAsyncContextManager, ExitStack, asynccontextmanager, nullcontext
from dataclasses import dataclass, replace
from functools import lru_cache
from itertools import chain, groupby
@ -69,9 +69,16 @@ from litellm.constants import (
MCP_HEALTH_CHECK_TIMEOUT,
MCP_METADATA_TIMEOUT,
MCP_TOOL_LISTING_TIMEOUT,
MCP_UPSTREAM_SESSION_SHUTDOWN_TIMEOUT_SECONDS,
)
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme, to_basic_credentials
from litellm.experimental_mcp_client.client import (
MCPClient,
PersistentMCPSession,
UpstreamSessionClosedError,
strip_auth_scheme,
to_basic_credentials,
)
from litellm.integrations.custom_guardrail import (
_sync_guardrail_info_to_logging_obj, # pyright: ignore[reportPrivateUsage] - the same bridge @log_guardrail_information uses; reimplementing it here would fork the metadata-key logic
)
@ -98,6 +105,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
upstream_auth_challenge,
)
from litellm.proxy._experimental.mcp_server.interactions import bind_target
from litellm.proxy._experimental.mcp_server.legacy_callbacks import has_active_mcp_request_context
from litellm.proxy._experimental.mcp_server.mcp_debug import describe_upstream_http_failure, record_auth_resolution
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
MCPPerUserTokenCache,
@ -343,6 +351,12 @@ _ToolParamMap: TypeAlias = dict[str, list[str]]
_EnvVarList: TypeAlias = list[dict[str, object]]
_InMemoryCacheDict: TypeAlias = dict[str, object]
_ToolArguments: TypeAlias = dict[str, object]
_UpstreamSessionKey: TypeAlias = tuple[
str,
str,
str,
tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool],
]
@dataclass(frozen=True, slots=True)
@ -1927,6 +1941,9 @@ class MCPServerManager:
self.catalog = CatalogSnapshots(self)
self.registry: dict[str, MCPServer] = {}
self._upstream_sessions: dict[_UpstreamSessionKey, PersistentMCPSession] = {} # mutable-ok: session registry
self._retiring_upstream_sessions: list[tuple[str, PersistentMCPSession]] = [] # mutable-ok: session registry
self._live_gateway_sessions: frozenset[str] = frozenset()
self._openapi_health_probes: Callable[[str], _OpenAPIHealthProbe] = lru_cache(maxsize=128)(_OpenAPIHealthProbe)
self.config_mcp_servers: dict[str, MCPServer] = {}
"""
@ -6123,6 +6140,108 @@ class MCPServerManager:
async with semaphore:
yield
def _outbound_limit(
self,
mcp_server: MCPServer,
persistent_session: PersistentMCPSession | None,
) -> AbstractAsyncContextManager[None]:
if persistent_session is not None:
return nullcontext()
return self._limit_outbound_concurrency(mcp_server)
async def _upstream_session_for(
self,
client: MCPClient,
mcp_server: MCPServer,
raw_headers: Mapping[str, str] | None,
) -> PersistentMCPSession | None:
gateway_session_id: Final = next(
(
value
for key, value in (raw_headers.items() if raw_headers else ())
if key.lower() == "mcp-session-id" and value
),
None,
)
if (
not has_active_mcp_request_context()
or gateway_session_id is None
or gateway_session_id not in self._live_gateway_sessions
or mcp_server.transport == MCPTransport.stdio
):
return None
key: Final[_UpstreamSessionKey] = (
gateway_session_id,
mcp_server.server_id,
await client.discovery_auth_fingerprint(),
client.session_settings(),
)
if gateway_session_id not in self._live_gateway_sessions:
return None
existing: Final = self._upstream_sessions.get(key)
if existing is not None and existing.reusable:
return existing
opened: Final = client.open_persistent_session(
lambda: self._limit_outbound_concurrency(self.get_mcp_server_by_id(mcp_server.server_id) or mcp_server)
)
self._retiring_upstream_sessions = [
(gateway_session_id, session)
for gateway_session_id, session in self._retiring_upstream_sessions
if not session.closed
]
previous_keys: Final = tuple(
registered_key for registered_key in self._upstream_sessions if registered_key[:2] == key[:2]
)
previous_sessions: Final = tuple(
(key[0], self._upstream_sessions.pop(previous_key)) for previous_key in previous_keys
)
for retired_gateway_session_id, previous_session in previous_sessions:
previous_session.retire()
self._retiring_upstream_sessions.append((retired_gateway_session_id, previous_session))
self._upstream_sessions[key] = opened
return opened
def track_gateway_session(self, gateway_session_id: str) -> None:
self._live_gateway_sessions = self._live_gateway_sessions | frozenset((gateway_session_id,))
def release_upstream_sessions(self, gateway_session_id: str) -> None:
self._live_gateway_sessions = self._live_gateway_sessions - frozenset((gateway_session_id,))
for key in tuple(key for key in self._upstream_sessions if key[0] == gateway_session_id):
self._upstream_sessions.pop(key).close()
for retiring_gateway_session_id, session in self._retiring_upstream_sessions:
if retiring_gateway_session_id == gateway_session_id:
session.close()
self._retiring_upstream_sessions = [
(retiring_gateway_session_id, session)
for retiring_gateway_session_id, session in self._retiring_upstream_sessions
if retiring_gateway_session_id != gateway_session_id and not session.closed
]
async def close_all_upstream_sessions(self) -> None:
sessions: Final = (
*self._upstream_sessions.values(),
*(session for _, session in self._retiring_upstream_sessions),
)
self._upstream_sessions.clear()
self._retiring_upstream_sessions.clear()
self._live_gateway_sessions = frozenset()
for session in sessions:
session.close()
try:
await asyncio.wait_for(
asyncio.gather(*(session.wait_closed() for session in sessions)),
MCP_UPSTREAM_SESSION_SHUTDOWN_TIMEOUT_SECONDS,
)
except asyncio.TimeoutError:
verbose_logger.warning("Timed out waiting for persistent MCP sessions to close during shutdown")
def _drop_upstream_session(self, session: PersistentMCPSession | None) -> None:
matching_keys: Final = tuple(key for key, value in self._upstream_sessions.items() if value is session)
dropped_sessions: Final = tuple((key[0], self._upstream_sessions.pop(key)) for key in matching_keys)
for gateway_session_id, dropped in dropped_sessions:
dropped.close()
self._retiring_upstream_sessions.append((gateway_session_id, dropped))
async def _obo_call_tool_with_retry(
self,
*,
@ -6138,27 +6257,33 @@ class MCPServerManager:
raw_headers: Mapping[str, str] | None = None,
client_ip: str | None = None,
allow_input_required: bool = False,
persistent_session: PersistentMCPSession | None = None,
) -> CallToolResult | InputRequiredResult:
"""Call a token_exchange (OBO) tool; on an upstream 401/403 re-mint the token once and retry.
The exchanged token is baked into the client at build time, so the retry invalidates the
cached exchange and rebuilds the client (which re-exchanges). One retry only: a non-auth
failure or a second auth failure degrades to the normal ``isError`` result, and a re-exchange
that now fails surfaces its own 401 challenge from ``_create_mcp_client``.
that now fails surfaces its own 401 challenge from ``create_mcp_client``.
"""
try:
return await client.call_tool(
call_tool_params,
host_progress_callback=host_progress_callback,
raise_on_error=True,
allow_input_required=allow_input_required,
)
async with self._outbound_limit(mcp_server, persistent_session):
return await client.call_tool(
call_tool_params,
host_progress_callback=host_progress_callback,
raise_on_error=True,
allow_input_required=allow_input_required,
persistent_session=persistent_session,
)
except Exception as exc:
if _extract_upstream_auth_failure(exc) is None:
auth_failure: Final = _extract_upstream_auth_failure(exc)
if auth_failure is None and not isinstance(exc, UpstreamSessionClosedError):
return MCPClient.error_tool_result(exc)
spec: Final = to_server_spec(mcp_server)
if spec is not None:
await self._cred_provider.invalidate_credentials(to_subject(user_api_key_auth, subject_token), spec)
if auth_failure is not None:
spec: Final = to_server_spec(mcp_server)
if spec is not None:
await self._cred_provider.invalidate_credentials(to_subject(user_api_key_auth, subject_token), spec)
self._drop_upstream_session(persistent_session)
retry_client: Final = await self.create_mcp_client(
server=mcp_server,
mcp_auth_header=server_auth_header,
@ -6169,11 +6294,14 @@ class MCPServerManager:
raw_headers=raw_headers,
client_ip=client_ip,
)
return await retry_client.call_tool(
call_tool_params,
host_progress_callback=host_progress_callback,
allow_input_required=allow_input_required,
)
retry_persistent_session: Final = await self._upstream_session_for(retry_client, mcp_server, raw_headers)
async with self._outbound_limit(mcp_server, retry_persistent_session):
return await retry_client.call_tool(
call_tool_params,
host_progress_callback=host_progress_callback,
allow_input_required=allow_input_required,
persistent_session=retry_persistent_session,
)
async def _call_regular_mcp_tool(
self,
@ -6330,6 +6458,7 @@ class MCPServerManager:
raw_headers=raw_headers,
client_ip=client_ip,
)
persistent_session: Final = await self._upstream_session_for(client, mcp_server, raw_headers)
call_tool_params: Final = MCPCallToolRequestParams(
name=original_tool_name,
@ -6342,24 +6471,21 @@ class MCPServerManager:
# OBO / ID-JAG: the exchanged token may have been revoked/rotated upstream since it was
# cached, so an upstream 401 gets one invalidate + re-mint + retry. Gated to these modes;
# all others keep the plain single call below.
async def _obo_call_tool_limited():
async with self._limit_outbound_concurrency(mcp_server):
return await self._obo_call_tool_with_retry(
client=client,
call_tool_params=call_tool_params,
host_progress_callback=host_progress_callback,
mcp_server=mcp_server,
server_auth_header=server_auth_header,
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
allow_input_required=allow_input_required,
)
tool_call_coro = _obo_call_tool_limited()
tool_call_coro = self._obo_call_tool_with_retry(
client=client,
call_tool_params=call_tool_params,
host_progress_callback=host_progress_callback,
mcp_server=mcp_server,
server_auth_header=server_auth_header,
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
allow_input_required=allow_input_required,
persistent_session=persistent_session,
)
else:
# Scoped to the two client-forwarded token modes this stack introduced; legacy
# oauth2 + delegate_auth_to_upstream (is_oauth_passthrough) is being removed, so it is not
@ -6368,12 +6494,13 @@ class MCPServerManager:
server_label: Final = mcp_server.name or mcp_server.server_name or mcp_server.alias or ""
async def _call_tool_via_client(client, params):
async with self._limit_outbound_concurrency(mcp_server):
async with self._outbound_limit(mcp_server, persistent_session):
if not relays_upstream_auth:
return await client.call_tool(
params,
host_progress_callback=host_progress_callback,
allow_input_required=allow_input_required,
persistent_session=persistent_session,
)
# The client-forwarded modes carry the caller's own upstream token, so an upstream
# 401 (expired/invalid token) is the caller's to resolve: relay it as
@ -6390,6 +6517,7 @@ class MCPServerManager:
host_progress_callback=host_progress_callback,
raise_on_error=True,
allow_input_required=allow_input_required,
persistent_session=persistent_session,
)
except Exception as e:
auth_info: Final = _extract_upstream_auth_failure(e)

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

@ -658,6 +658,7 @@ if MCP_AVAILABLE:
_stateful_session_locks.pop(session_id, None)
_stateful_session_active_request_counts.pop(session_id, None)
_stateful_session_client_info.pop(session_id, None)
operations.global_mcp_server_manager.release_upstream_sessions(session_id)
# Keep this alias so existing references to session_manager still work
session_manager: Final = session_manager_stateless
@ -793,6 +794,7 @@ if MCP_AVAILABLE:
_stateful_auth_context_cleanup_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await _stateful_auth_context_cleanup_task
await operations.global_mcp_server_manager.close_all_upstream_sessions()
if _session_manager_stateful_cm:
await _session_manager_stateful_cm.__aexit__(None, None, None)
if _session_manager_cm:
@ -1264,13 +1266,16 @@ if MCP_AVAILABLE:
return f"ip:{hashlib.sha256(client_ip.encode('utf-8')).hexdigest()}"
return "anonymous"
def _peeked_json_object_body(body: bytes) -> bytes | None:
"""The peeked bytes when they form a complete JSON-RPC object, else None so a
truncated prefix or a batch array fails closed and stays budget-enforced."""
def _admission_auth_body(body: bytes) -> bytes:
"""Return only the JSON-RPC method for admission auth."""
try:
return body if isinstance(json.loads(body), dict) else None
except (json.JSONDecodeError, TypeError):
return None
data: Final = TypeAdapter(dict[str, object]).validate_json(body)
except (ValueError, TypeError):
return b"{}"
method: Final = data.get("method")
if not isinstance(method, str):
return b"{}"
return json.dumps({"method": method}).encode()
def _is_initialize_request(body: bytes) -> bool:
"""
@ -2193,10 +2198,10 @@ if MCP_AVAILABLE:
if peek is not None:
receive = peek.receive # rebind-ok: replay peeked ASGI messages to the downstream handler
async def peeked_json_object_body() -> bytes:
return _peeked_json_object_body(await peek.admission_body()) or b"{}"
async def admission_auth_body() -> bytes:
return _admission_auth_body(await peek.admission_body())
scope[MCP_PEEKED_BODY_SCOPE_KEY] = peeked_json_object_body
scope[MCP_PEEKED_BODY_SCOPE_KEY] = admission_auth_body
(
user_api_key_auth,
mcp_auth_header,
@ -2856,6 +2861,7 @@ if MCP_AVAILABLE:
_stateful_session_auth_contexts[session_id] = auth_user
_stateful_session_auth_context_last_seen[session_id] = time.monotonic()
_stateful_session_owners[session_id] = owner_fingerprint
operations.global_mcp_server_manager.track_gateway_session(session_id)
if client_info is not None:
_stateful_session_client_info[session_id] = client_info
break

1805
litellm/proxy/_lazy_openapi_snapshot.json
View file

File diff suppressed because it is too large Load diff

583
litellm/proxy/_types.py
View file

@ -3,7 +3,7 @@ import json
import os
from collections.abc import Callable, Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from types import MappingProxyType as MappingProxyType # noqa: PLC0414 # public re-export
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple, TypeAlias
import httpx
@ -16,7 +16,6 @@ from pydantic import (
Json,
JsonValue,
PositiveInt,
PrivateAttr,
TypeAdapter,
field_validator,
model_validator,
@ -26,19 +25,24 @@ from typing_extensions import (
ReadOnly,
Required,
TypedDict,
TypeIs, # noqa: TID251 # narrows the before-validator input without copying it
)
from litellm._uuid import uuid
from litellm.constants import DEFAULT_STAGGER_WINDOW_SECONDS, MCP_STDIO_ALLOWED_COMMANDS
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
validate_arize_otlp_protocol_value,
validate_capture_message_content_value,
validate_langfuse_environment_value,
validate_langfuse_span_scope_value,
validate_no_callback_env_reference,
)
from litellm.proxy._experimental.mcp_server.stdio_gate import MCP_STDIO_DISABLED_MESSAGE, is_mcp_stdio_enabled
from litellm.types.agents import AgentCaller, AgentResponse
from litellm.types.agents import (
AgentCaller as AgentCaller, # noqa: PLC0414 # public re-export
)
from litellm.types.agents import (
AgentResponse as AgentResponse, # noqa: PLC0414 # public re-export
)
from litellm.types.integrations.compression_interception import (
CompressionSavingsMetadata,
)
@ -67,11 +71,17 @@ from litellm.types.proxy.agent_identity import ManagedAgentContext
from litellm.types.proxy.auth.special_headers import (
SpecialHeaders as SpecialHeaders, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.auth.user_api_key_auth import is_jwt
from litellm.types.proxy.auth.user_api_key_auth import (
is_jwt as is_jwt, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.carried_budget_state import (
OrgBudgetSnapshot,
TeamBudgetSnapshot,
UserBudgetSnapshot,
OrgBudgetSnapshot as OrgBudgetSnapshot, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.carried_budget_state import (
TeamBudgetSnapshot as TeamBudgetSnapshot, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.carried_budget_state import (
UserBudgetSnapshot as UserBudgetSnapshot, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
from litellm.types.proxy.spend_capture_rate import SpendCaptureRateCheckSettings
@ -79,6 +89,7 @@ from litellm.types.router import AllowedModelRegion, RouterErrors, UpdateRouterC
from litellm.types.router_weights import validate_router_settings_dict
from litellm.types.secret_managers.main import KeyManagementSystem
from litellm.types.utils import (
CAPTURE_MESSAGE_CONTENT_VAR,
AzureSpillover,
CallTypes,
CostBreakdown,
@ -196,17 +207,9 @@ class LitellmTableNames(str, enum.Enum):
from litellm.types.integrations.slack_alerting import (
Litellm_EntityType as Litellm_EntityType, # noqa: E402, PLC0414 # public re-export
)
def hash_token(token: str) -> str:
import hashlib
# Hash the string using SHA-256
hashed_token: Final = hashlib.sha256(token.encode()).hexdigest()
return hashed_token
from litellm.types.proxy.auth.user_api_key_auth import ( # noqa: E402 # re-export after the definitions above
hash_token as hash_token, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.auth.user_roles import ( # noqa: E402 # re-export after the definitions above
KeyManagementRoutes as KeyManagementRoutes, # noqa: PLC0414 # public re-export
)
@ -1163,6 +1166,18 @@ from litellm.models.team import BudgetLimitEntry as BudgetLimitEntry # noqa: E4
from litellm.types.object_permission import ( # noqa: E402
ObjectPermissionDict as ObjectPermissionDict,
)
from litellm.types.proxy.management_endpoints.key_requests import ( # noqa: E402 # re-export after the definitions above
AllowedVectorStoreIndexItem as AllowedVectorStoreIndexItem, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.key_requests import ( # noqa: E402 # re-export after the definitions above
GenerateKeyRequest as GenerateKeyRequest, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.key_requests import ( # noqa: E402 # re-export after the definitions above
KeyRequestBase as KeyRequestBase, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.key_requests import ( # noqa: E402 # re-export after the definitions above
LiteLLMKeyType as LiteLLMKeyType, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.request_base import ( # noqa: E402 # re-export after the definitions above
GenerateRequestBase as GenerateRequestBase, # noqa: PLC0414 # public re-export
)
@ -1171,71 +1186,6 @@ from litellm.types.proxy.management_endpoints.request_base import ( # noqa: E40
)
class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase):
index_name: str
index_permissions: list[Literal["read", "write"]]
class KeyRequestBase(GenerateRequestBase):
key: str | None = None
tpd_limit: int | None = None
default_estimated_output_tokens: PositiveInt | None = None
default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None
budget_id: str | None = None
end_user_budget_id: str | None = None
tags: list[str] | None = None
disable_global_guardrails: bool | None = None
enable_prompt_caching: bool | None = None
throttle_on_budget_exceeded: bool | None = None
enforced_params: list[str] | None = None
allowed_routes: list[Any] | None = []
allowed_passthrough_routes: list[object] | None = None
denied_passthrough_routes: list[str] | None = None
allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None
rpm_limit_type: Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] | None = (
None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm
)
tpm_limit_type: Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] | None = (
None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm
)
router_settings: UpdateRouterConfig | None = None
access_group_ids: list[str] | None = None
class LiteLLMKeyType(str, enum.Enum):
"""
Enum for key types that determine what routes a key can access
"""
LLM_API = "llm_api" # Can call LLM API routes (chat/completions, embeddings, etc.)
MANAGEMENT = "management" # Can call management routes (user/team/key management)
READ_ONLY = "read_only" # Can only call info/read routes
DEFAULT = "default" # Uses default allowed routes
class GenerateKeyRequest(KeyRequestBase):
soft_budget: float | None = None
send_invite_email: bool | None = None
key_type: LiteLLMKeyType | None = Field(
default=LiteLLMKeyType.DEFAULT,
description="Type of key that determines default allowed routes.",
)
auto_rotate: bool | None = Field(default=False, description="Whether this key should be automatically rotated")
rotation_interval: str | None = Field(
default=None,
description="How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True",
)
organization_id: str | None = None
project_id: str | None = None
@field_validator("team_id", "organization_id", "project_id", mode="before")
@classmethod
def treat_cleared_id_as_unset(cls, v: object) -> object:
if v == "":
return None
return v
class GenerateKeyResponse(KeyRequestBase):
key: str
key_name: str | None = None
@ -1278,65 +1228,12 @@ class GenerateKeyResponse(KeyRequestBase):
return values
def _is_mapping(value: object) -> TypeIs[Mapping[object, object]]: # guard-ok: raw before-validator payload
return isinstance(value, Mapping)
class UpdateKeyRequest(KeyRequestBase):
# Note: the defaults of all Params here MUST BE NONE
# else they will get overwritten
duration: str | None = None
spend: float | None = None
soft_budget: float | None = None
metadata: dict[Any, Any] | None = None
temp_budget_increase: float | None = None
temp_budget_expiry: datetime | None = None
auto_rotate: bool | None = None
rotation_interval: str | None = None
organization_id: str | None = None
project_id: str | None = Field(
default=None,
description="Omit to retain the project, or send null to detach. Assigning a different project is not supported.",
)
@model_validator(mode="before")
@classmethod
def drop_blank_team_id(cls, values: object) -> object:
if _is_mapping(values) and values.get("team_id") == "":
return MappingProxyType({k: v for k, v in values.items() if k != "team_id"})
return values
@field_validator("organization_id", mode="before")
@classmethod
def treat_cleared_organization_id_as_unset(cls, v: object) -> object:
if v == "":
return None
return v
@model_validator(mode="after")
def validate_temp_budget(self) -> "UpdateKeyRequest":
if self.temp_budget_increase is not None or self.temp_budget_expiry is not None:
if self.temp_budget_increase is None or self.temp_budget_expiry is None:
raise ValueError("temp_budget_increase and temp_budget_expiry must be set together")
return self
@model_validator(mode="after")
def validate_key_identifier(self) -> "UpdateKeyRequest":
if self.key is None and self.key_alias is None:
raise ValueError("either key or key_alias must be provided")
return self
class RegenerateKeyRequest(GenerateKeyRequest):
# This needs to be different from UpdateKeyRequest, because "key" is optional for this
key: str | None = None
new_key: str | None = None
duration: str | None = None
spend: float | None = None
metadata: dict[Any, Any] | None = None
new_master_key: str | None = None
grace_period: str | None = None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke
from litellm.types.proxy.management_endpoints.key_requests import ( # noqa: E402 # re-export after the definitions above
RegenerateKeyRequest as RegenerateKeyRequest, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.key_requests import ( # noqa: E402 # re-export after the definitions above
UpdateKeyRequest as UpdateKeyRequest, # noqa: PLC0414 # public re-export
)
class ResetSpendRequest(LiteLLMPydanticObjectBase):
@ -1788,41 +1685,12 @@ class ListSkillsRequest(LiteLLMPydanticObjectBase):
offset: int | None = 0
class NewUserRequestTeam(LiteLLMPydanticObjectBase):
team_id: str
max_budget_in_team: float | None = None
user_role: Literal["user", "admin"] = "user"
class NewUserRequest(GenerateRequestBase):
max_budget: float | None = None
user_email: str | None = None
user_alias: str | None = None
user_role: (
Literal[
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
]
| None
) = None
teams: list[str] | list[NewUserRequestTeam] | None = None
auto_create_key: bool = True # flag used for returning a key as part of the /user/new response
send_invite_email: bool | None = None
sso_user_id: str | None = None
organizations: list[str] | None = None
password: str | None = None
@field_validator("password")
@classmethod
def password_not_supported(cls, value: str | None) -> str | None:
if value is not None:
raise ValueError(
"password cannot be set via /user/new. Users set their own password through an "
"invitation link (POST /invitation/new)."
)
return value
from litellm.types.proxy.management_endpoints.user_requests import ( # noqa: E402 # re-export after the definitions above
NewUserRequest as NewUserRequest, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.user_requests import ( # noqa: E402 # re-export after the definitions above
NewUserRequestTeam as NewUserRequestTeam, # noqa: PLC0414 # public re-export
)
class NewUserResponse(GenerateKeyResponse):
@ -1844,36 +1712,12 @@ class NewUserResponse(GenerateKeyResponse):
updated_at: datetime | None = None
class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with BulkUpdateUserRequest
# repr=False keeps the plaintext out of management-endpoint alerts, which str() the request model
password: str | None = Field(default=None, repr=False)
spend: float | None = None
metadata: dict[object, object] | None = None
user_alias: str | None = None
user_role: (
Literal[
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
]
| None
) = None
max_budget: float | None = None
class UpdateUserRequest(UpdateUserRequestNoUserIDorEmail):
# Note: the defaults of all Params here MUST BE NONE
# else they will get overwritten
user_id: str | None = None
user_email: str | None = None
@model_validator(mode="before")
@classmethod
def check_user_info(cls, values: Mapping[str, object]) -> Mapping[str, object]:
if values.get("user_id") is None and values.get("user_email") is None:
raise ValueError("Either user id or user email must be provided")
return values
from litellm.types.proxy.management_endpoints.user_requests import ( # noqa: E402 # re-export after the definitions above
UpdateUserRequest as UpdateUserRequest, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.user_requests import ( # noqa: E402 # re-export after the definitions above
UpdateUserRequestNoUserIDorEmail as UpdateUserRequestNoUserIDorEmail, # noqa: PLC0414 # public re-export
)
class ChangePasswordRequest(LiteLLMPydanticObjectBase):
@ -2202,7 +2046,7 @@ class AddTeamCallback(LiteLLMPydanticObjectBase):
@classmethod
def validate_callback_vars(cls, values):
callback_vars: Final = values.get("callback_vars", {})
valid_keys: Final = set(StandardCallbackDynamicParams.__annotations__.keys())
valid_keys: Final = {*StandardCallbackDynamicParams.__annotations__, CAPTURE_MESSAGE_CONTENT_VAR}
for key, value in callback_vars.items():
if key not in valid_keys:
raise ValueError(f"Invalid callback variable: {key}. Must be one of {valid_keys}")
@ -2214,6 +2058,8 @@ class AddTeamCallback(LiteLLMPydanticObjectBase):
validate_langfuse_span_scope_value(callback_vars[key])
if key == "arize_otlp_protocol":
validate_arize_otlp_protocol_value(callback_vars[key])
if key == CAPTURE_MESSAGE_CONTENT_VAR:
validate_capture_message_content_value(callback_vars[key])
return values
@ -3213,296 +3059,15 @@ from litellm.models.verification_token import ( # noqa: E402
from litellm.models.verification_token import ( # noqa: E402
LiteLLM_VerificationToken as LiteLLM_VerificationToken,
)
_UntypedDict: TypeAlias = dict[Any, Any]
def _is_str_keyed_dict(value: object) -> TypeIs[dict[str, object]]: # guard-ok: pydantic input and JSON keys are str
return isinstance(value, dict)
class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
"""
Combined view of litellm verification token + litellm team table (select values)
"""
team_spend: float | None = None
team_alias: str | None = None
team_tpm_limit: int | None = None
team_rpm_limit: int | None = None
team_tpd_limit: int | None = None
team_max_budget: float | None = None
team_soft_budget: float | None = None
team_model_max_budget: dict[str, object] | None = None
team_models: list[Any] = []
team_blocked: bool = False
soft_budget: float | None = None
team_model_aliases: _UntypedDict | None = None
team_member: Member | None = None
team_metadata: _UntypedDict | None = None
team_object_permission_id: str | None = None
# Team Member Specific Params
team_member_spend: float | None = None
team_member_tpm_limit: int | None = None
team_member_rpm_limit: int | None = None
# End User Params
end_user_id: str | None = None
end_user_tpm_limit: int | None = None
end_user_rpm_limit: int | None = None
end_user_tpd_limit: int | None = None
end_user_max_budget: float | None = None
end_user_model_max_budget: _UntypedDict | None = None
# Organization Params
organization_alias: str | None = None
organization_max_budget: float | None = None
organization_tpm_limit: int | None = None
organization_rpm_limit: int | None = None
organization_metadata: _UntypedDict | None = None
# Project Params
project_alias: str | None = None
project_metadata: _UntypedDict | None = None
# Time stamps
last_refreshed_at: float | None = None # last time joint view was pulled from db
def __init__(self, **kwargs: object) -> None:
# Handle litellm_budget_table_* keys (budget table overrides when key value is None or empty)
for key, value in list(kwargs.items()):
if key.startswith("litellm_budget_table_") and value is not None:
# Extract the corresponding attribute name
attr_name = key.replace("litellm_budget_table_", "")
# Use key's value from kwargs (from DB view), not class default
current = kwargs.get(attr_name)
if current is None:
current = getattr(self, attr_name, None)
# Apply budget value when key has no value, or for model_max_budget when key has empty dict
should_apply = current is None or (
attr_name == "model_max_budget" and _is_str_keyed_dict(current) and len(current) == 0
)
if should_apply:
kwargs[attr_name] = value
if key == "end_user_id" and value is not None and isinstance(value, int):
kwargs[key] = str(value)
if kwargs.get("organization_id") is not None:
kwargs["org_id"] = kwargs.pop("organization_id")
# Initialize the superclass
super().__init__(**kwargs) # pyright: ignore[reportArgumentType, reportUnknownMemberType] # pydantic validates raw values
class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response object for user api key auth
"""
Return the row in the db
"""
api_key: str | None = None
user_role: LitellmUserRoles | None = None
allowed_model_region: AllowedModelRegion | None = None
parent_otel_span: Span | None = None
rpm_limit_per_model: dict[str, int] | None = None
tpm_limit_per_model: dict[str, int] | None = None
user_tpm_limit: int | None = None
user_rpm_limit: int | None = None
user_email: str | None = None
user_spend: float | None = None
user_max_budget: float | None = None
# Values stay `object` rather than BudgetConfig: this is the raw JSON column,
# and validating it here would make one malformed row fail auth outright.
# resolve_model_budget validates the single entry a request actually needs.
user_model_max_budget: Mapping[str, object] | None = None
team_member_model_max_budget: Mapping[str, object] | None = Field(default=None, exclude=True)
request_route: str | None = None
is_session_token: bool = False
# Server-only marker set exclusively by the MCP gateway admission path
# (reload_admitted_user) for a keyless user-subject admitted via a gateway DCR session
# bearer or bridge envelope. Not a DB column and never populated from caller-controlled key
# metadata or JWT claims, so it cannot be forged to gain the team-inherited MCP grant union
# or to escape the caller-Authorization egress scrub. exclude=True keeps it out of serialization.
mcp_admitted_user_subject: bool = Field(default=False, exclude=True)
requires_fresh_policy: bool = Field(default=False, exclude=True)
mcp_explicit_grants_only: bool = Field(default=False, exclude=True)
# team_id -> that team's mcp_rpm_limit map, for a keyless admitted subject that reaches MCP
# servers through several teams at once and therefore has no single team_id for the limiter to
# key off. Server-only and stripped from validated input for the same reason as the marker
# above: a forged entry would let a caller pick which team's rpm bucket it is charged against.
mcp_source_team_rpm_limits: dict[str, dict[str, int]] | None = Field(default=None, exclude=True)
# The single MCP server_id a gateway session bearer was scoped to at authorize time (RFC 8707
# resource), or None for an aggregate-scope session. A RESTRICTION intersected against the live
# grant resolution, never a grant. Server-only, set exclusively by the MCP gateway admission
# path via post-construction assignment and stripped from validated input like the markers
# above; a forged value could at most narrow, but the stripping keeps the field's provenance
# single-owner so its meaning stays trustworthy.
mcp_session_resource_server_id: str | None = Field(default=None, exclude=True)
mcp_toolset_id: str | None = Field(default=None, exclude=True)
authenticated_by_custom_auth: bool = Field(default=False, exclude=True)
via_virtual_key: bool = Field(
default=False,
exclude=True,
description=(
"Server-only marker set exclusively by the DB virtual-key and master-key auth paths via "
"post-construction assignment. Stripped from validated input so custom auth handlers, JWT "
"claims, or key metadata cannot forge it. Gates overwrite_user_with_key_hash stamping: only "
"a credential the proxy itself validated as a key may be forwarded as the provider-facing "
"user id."
),
)
invoked_agent_id: str | None = Field(default=None, exclude=True)
invoked_agent_policy: AgentResponse | None = Field(default=None, exclude=True)
agent_invocation_cost: float | None = Field(default=None, exclude=True)
billing_agent_policy: AgentResponse | None = Field(default=None, exclude=True)
_managed_delegation_verified: bool = PrivateAttr(default=False)
managed_agent_policy: AgentResponse | None = Field(default=None, exclude=True)
managed_agent_context: ManagedAgentContext | None = Field(default=None, exclude=True)
agent_caller: AgentCaller | None = Field(
default=None,
exclude=True,
description=(
"Set per request from the x-litellm-user-id / x-litellm-team-id headers an agent echoes back on "
"calls made with its own key. Every check treats it as a ceiling, so a forged value can only "
"narrow the agent's access."
),
)
budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True)
team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True)
user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True)
org_budget_snapshot: OrgBudgetSnapshot | None = Field(default=None, exclude=True)
matched_model_access_groups: list[str] | None = Field(default=None, exclude=True)
budget_throttle_pct: float | None = Field(default=None, exclude=True)
user: Any | None = None # Expanded user object when expand=user is used
created_by_user: Any | None = None # Expanded created_by user when expand=user is used
end_user_object_permission: LiteLLM_ObjectPermissionTable | None = None
# Team object_permission preloaded in auth (e.g. get_team_object) to avoid
# per-request object_permission fetches in downstream checks (vector stores, etc.)
team_object_permission: LiteLLM_ObjectPermissionTable | None = None
# Decoded upstream IdP claims (groups, roles, etc.) propagated by JWT auth machinery
# and forwarded into outbound tokens by guardrails such as MCPJWTSigner.
jwt_claims: _UntypedDict | None = None
model_config = ConfigDict(arbitrary_types_allowed=True)
@model_validator(mode="before")
@classmethod
def check_api_key(cls, values: object) -> object:
# If values is already an instance (not a dict), return it as-is
if not _is_str_keyed_dict(values):
return values
# mcp_admitted_user_subject is a server-only marker, set ONLY by the MCP gateway admission
# path via post-construction assignment. Strip it from any validated input (constructor
# kwargs, model_validate, a JWT/key claim splat) so it can never be forged from caller data.
values.pop("mcp_admitted_user_subject", None)
values.pop("requires_fresh_policy", None)
values.pop("mcp_explicit_grants_only", None)
values.pop("mcp_source_team_rpm_limits", None)
values.pop("mcp_session_resource_server_id", None)
values.pop("mcp_toolset_id", None)
values.pop("via_virtual_key", None)
values.pop("authenticated_by_custom_auth", None)
values.pop("agent_caller", None)
values.pop("managed_agent_context", None)
values.pop("managed_agent_policy", None)
values.pop("invoked_agent_id", None)
values.pop("invoked_agent_policy", None)
values.pop("agent_invocation_cost", None)
values.pop("billing_agent_policy", None)
api_key: Final = values.get("api_key")
if api_key is not None:
values.update({"token": cls._safe_hash_litellm_api_key(api_key)}) # pyright: ignore[reportArgumentType] # non-str fails in the hasher
if isinstance(api_key, str):
values.update({"api_key": cls._safe_hash_litellm_api_key(api_key)})
return values
@classmethod
def _safe_hash_litellm_api_key(cls, api_key: str) -> str:
"""
Helper to ensure all logged keys are hashed
Covers:
1. Regular API keys from LiteLLM DB
2. JWT tokens used for connecting to LiteLLM API
"""
normalized = api_key
if normalized[:7].lower() == "bearer ":
normalized = normalized[7:]
if normalized.startswith("sk-"):
return hash_token(normalized)
if is_jwt(token=normalized):
return f"hashed-jwt-{hash_token(token=normalized)}"
return normalized
@classmethod
def get_litellm_internal_health_check_user_api_key_auth(cls) -> "UserAPIKeyAuth":
"""
Returns a `UserAPIKeyAuth` object for the litellm internal health check service account.
This is used to track number of requests/spend for health check calls.
"""
from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
team_id=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
key_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
team_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
)
@classmethod
def get_litellm_cli_user_api_key_auth(cls) -> "UserAPIKeyAuth":
"""
Returns a `UserAPIKeyAuth` object for the litellm internal health check service account.
This is used to track number of requests/spend for health check calls.
"""
from litellm.constants import LITTELM_CLI_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
team_id=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
key_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
team_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
)
@classmethod
def get_litellm_internal_jobs_user_api_key_auth(cls) -> "UserAPIKeyAuth":
"""
Returns a `UserAPIKeyAuth` object for internal LiteLLM jobs like key rotation.
This is used to track actions performed by automated system jobs.
"""
from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
team_id="system",
key_alias=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
team_alias="system",
user_id="system",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
@property
def is_team_service_account(self) -> bool:
return (
self.user_id is None
and self.team_id is not None
and bool(self.metadata) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped base field
and self.metadata.get("service_account_id") is not None # pyright: ignore[reportUnknownMemberType] # untyped base field
)
def user_api_key_has_admin_view(user_api_key_dict: UserAPIKeyAuth) -> bool:
"""Return True if the caller's role grants unscoped read access to all
tenant resources (managed files, batches, vector stores, spend rows, etc).
Lives on _types.py so leaf modules (e.g. litellm.llms.base_llm.managed_resources)
can use it without pulling in litellm.proxy.utils via management_endpoints.
"""
return user_api_key_dict.user_role in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
)
from litellm.types.proxy.auth.user_api_key_auth import ( # noqa: E402 # re-export after the definitions above
LiteLLM_VerificationTokenView as LiteLLM_VerificationTokenView, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.auth.user_api_key_auth import (
UserAPIKeyAuth as UserAPIKeyAuth, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.auth.user_api_key_auth import (
user_api_key_has_admin_view as user_api_key_has_admin_view, # noqa: PLC0414 # public re-export
)
class UserInfoResponse(LiteLLMPydanticObjectBase):
@ -3734,10 +3299,6 @@ class LiteLLM_UserTableFiltered(LiteLLMBaseModel): # done to avoid exposing sen
user_email: str | None = None
class LiteLLM_UserTableWithKeyCount(LiteLLM_UserTable):
key_count: int = 0
from litellm.models.access_group import ( # noqa: E402
LiteLLM_AccessGroupTable as LiteLLM_AccessGroupTable,
)
@ -3751,6 +3312,9 @@ from litellm.models.spend_logs import ( # noqa: E402
LiteLLM_SpendLogs as LiteLLM_SpendLogs,
)
from litellm.models.tag import LiteLLM_TagTable as LiteLLM_TagTable # noqa: E402
from litellm.types.proxy.management_endpoints.user_requests import ( # noqa: E402 # re-export after the definitions above
LiteLLM_UserTableWithKeyCount as LiteLLM_UserTableWithKeyCount, # noqa: PLC0414 # public re-export
)
AUDIT_ACTIONS = Literal["created", "updated", "deleted", "blocked", "unblocked", "rotated", "kill_switch_fired"]
@ -4526,16 +4090,9 @@ class OrganizationAddMemberResponse(LiteLLMPydanticObjectBase):
updated_organization_memberships: list[LiteLLM_OrganizationMembershipTable]
class MemberDeleteRequest(LiteLLMPydanticObjectBase):
user_id: str | None = None
user_email: str | None = None
@model_validator(mode="before")
@classmethod
def check_user_info(cls, values: Mapping[str, object]) -> Mapping[str, object]:
if values.get("user_id") is None and values.get("user_email") is None:
raise ValueError("Either user id or user email must be provided")
return values
from litellm.types.proxy.management_endpoints.team_endpoints import ( # noqa: E402 # re-export after the definitions above
MemberDeleteRequest as MemberDeleteRequest, # noqa: PLC0414 # public re-export
)
class MemberUpdateResponse(LiteLLMPydanticObjectBase):

81
litellm/proxy/common_utils/callback_config_validation.py
View file

@ -10,6 +10,8 @@ from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final
from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VAR
_NEWRELIC_CALLBACK: Final = "newrelic"
_NEWRELIC_VAR_PREFIX: Final = "newrelic_"
_LANGFUSE_OTEL_CALLBACK: Final = "langfuse_otel"
@ -36,6 +38,9 @@ def callback_config_error(
)
if langfuse_error is not None:
return langfuse_error
capture_error: Final = _capture_message_content_error(callback_name, callback_vars, callback_type)
if capture_error is not None:
return capture_error
if callback_name != _NEWRELIC_CALLBACK:
return None
return _newrelic_config_error(callback_vars)
@ -80,6 +85,32 @@ def _langfuse_span_scope_error(callback_name: str | None, callback_vars: Mapping
return None
def _capture_message_content_error(
callback_name: str | None, callback_vars: Mapping[str, str], callback_type: str | None
) -> str | None:
value: Final = callback_vars.get(CAPTURE_MESSAGE_CONTENT_VAR)
if value is None:
return None
from litellm.integrations.otel.model.config import is_otel_v2_enabled
from litellm.integrations.otel.presets.destinations import destination_capable_backends
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
validate_capture_message_content_value,
)
supported: Final = sorted(destination_capable_backends())
if callback_name not in supported:
return f"{CAPTURE_MESSAGE_CONTENT_VAR} applies to the OTel v2 callbacks {supported} only, not {callback_name!r}"
if not is_otel_v2_enabled():
return f"Per-destination {CAPTURE_MESSAGE_CONTENT_VAR} requires the proxy to run with LITELLM_OTEL_V2=true."
try:
validate_capture_message_content_value(value)
except ValueError as e:
return str(e)
if callback_type == "failure":
return f"{CAPTURE_MESSAGE_CONTENT_VAR} needs callback_type 'success' or 'success_and_failure'"
return None
# Which credential family a dynamic variable belongs to. The families are the
# integrations that share one account: every langfuse_* variable configures the
# same Langfuse project whether it rides the classic callback or the OTel one,
@ -190,20 +221,41 @@ def conflicting_span_scope_error(
)
def conflicting_capture_error(
callback_name: str,
callback_vars: Mapping[str, str] | None,
stored_entries: Sequence[tuple[str | None, Mapping[str, str]]],
) -> str | None:
incoming: Final = None if callback_vars is None else callback_vars.get(CAPTURE_MESSAGE_CONTENT_VAR)
if incoming is None:
return None
return next(
(
f"{CAPTURE_MESSAGE_CONTENT_VAR} is already set to {stored!r} by another {callback_name} entry. "
f"Every {callback_name} entry shares one value: remove that entry or send the same value."
for name, entry_vars in stored_entries
if name == callback_name and (stored := entry_vars.get(CAPTURE_MESSAGE_CONTENT_VAR)) not in (None, incoming)
),
None,
)
def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str | None:
"""Validate every ``logging`` entry of a team/key metadata payload."""
if not metadata:
return None
entries: Final = metadata.get("logging")
if not isinstance(entries, Sequence) or isinstance(entries, (str, bytes)):
entries: Final = _object_sequence(metadata.get("logging"))
if entries is None:
return None
entry_vars: Final = tuple(_entry_callback_vars(entry) for entry in entries)
named_vars: Final = tuple(zip((_entry_callback_name(entry) for entry in entries), entry_vars))
return next(
(
error
for error in (
*(_logging_entry_error(entry) for entry in entries),
*(conflicting_span_scope_error(entry_vars[i], entry_vars[:i]) for i in range(len(entry_vars))),
*(_conflicting_entry_capture_error(named_vars[i], named_vars[:i]) for i in range(len(named_vars))),
)
if error is not None
),
@ -211,6 +263,31 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str
)
def stored_capture_entries(logging_entries: object) -> tuple[tuple[str | None, Mapping[str, str]], ...]:
"""The callback name and vars of every stored ``logging`` entry, for ``conflicting_capture_error``."""
entries: Final = _object_sequence(logging_entries) or ()
return tuple((_entry_callback_name(entry), _entry_callback_vars(entry)) for entry in entries)
def _object_sequence(value: object) -> Sequence[object] | None:
return value if isinstance(value, Sequence) and not isinstance(value, (str, bytes)) else None
def _entry_callback_name(entry: object) -> str | None:
match entry:
case {"callback_name": str() as callback_name}:
return callback_name
case _:
return None
def _conflicting_entry_capture_error(
entry: tuple[str | None, Mapping[str, str]], earlier: Sequence[tuple[str | None, Mapping[str, str]]]
) -> str | None:
callback_name, callback_vars = entry
return None if callback_name is None else conflicting_capture_error(callback_name, callback_vars, earlier)
def _entry_callback_vars(entry: object) -> Mapping[str, str]:
callback_vars: Final = entry.get("callback_vars") if isinstance(entry, Mapping) else None
if not isinstance(callback_vars, Mapping):

38
litellm/proxy/common_utils/custom_openapi_spec.py
View file

@ -437,3 +437,41 @@ class CustomOpenAPISpec:
# Add responses API request schema
return CustomOpenAPISpec.add_responses_api_request_schema(with_embeddings)
_MAX_INLINE_DEPTH: Final = 32
def _inlined(node: JsonValue, defs: Mapping[str, JsonValue], depth: int = 0) -> JsonValue:
if depth >= _MAX_INLINE_DEPTH or not isinstance(node, (list, dict)):
return node
if isinstance(node, list):
return [_inlined(item, defs, depth + 1) for item in node]
ref: Final = node.get("$ref")
if isinstance(ref, str) and ref.startswith("#/$defs/"):
target: Final = _inlined(defs[ref.removeprefix("#/$defs/")], defs, depth + 1)
siblings: Final = {key: _inlined(value, defs, depth + 1) for key, value in node.items() if key != "$ref"}
return {**target, **siblings} if isinstance(target, dict) else siblings
if "propertyName" in node and "mapping" in node:
return {"propertyName": node["propertyName"]}
return {key: _inlined(value, defs, depth + 1) for key, value in node.items() if key != "$defs"}
def inline_request_body(model_class: type, example: Mapping[str, object]) -> JsonObject:
"""
Build an ``openapi_extra["requestBody"]`` for a route that reads its body with ``request.body()``.
Swagger UI resolves ``$ref`` against the whole document, so a schema placed inline in an operation
cannot keep Pydantic's ``#/$defs/...`` references. Every reference is expanded in place instead,
which also keeps the operation self-contained inside the lazy OpenAPI snapshot.
"""
schema: Final = CustomOpenAPISpec.get_pydantic_schema(model_class)
media: Final[JsonObject] = {"example": cast(JsonValue, dict(example))} # cast-ok: JSON example literal
if schema is None:
return {"required": True, "content": {"application/json": media}}
raw_defs: Final = schema.get("$defs")
defs: Final[Mapping[str, JsonValue]] = raw_defs if isinstance(raw_defs, dict) else MappingProxyType({})
return {
"required": True,
"content": {"application/json": {"schema": _inlined(schema, defs), **media}},
}

106
litellm/proxy/common_utils/resource_ownership.py
View file

@ -1,92 +1,20 @@
from typing import Final
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.litellm_core_utils.resource_ownership import (
get_primary_resource_owner_scope,
get_resource_owner_scopes,
is_proxy_admin,
user_can_access_resource_owner,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.proxy.auth.user_roles import LitellmUserRoles
def is_proxy_admin(user_api_key_dict: UserAPIKeyAuth | None) -> bool:
if user_api_key_dict is None:
return False
return (
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
)
def get_resource_owner_scopes(
user_api_key_dict: UserAPIKeyAuth | None,
) -> list[str]:
"""
Return ownership scopes that may access a user-created proxy resource.
Raw user_id is included for rows created before scope prefixes existed.
Prefixes avoid collisions when falling back to team/org/key ownership
for keys that do not have a user_id.
Identity-less callers (no user_id, team_id, org_id, api_key, or token)
return ``[]`` — they share no scope with any other caller, so access
checks against an existing owner always fail and creates that depend
on a primary scope must reject up front. Returning a shared sentinel
here would let any two identity-less callers see each other's data.
"""
if user_api_key_dict is None:
return []
scopes: Final[list[str]] = []
def _add(scope: str | None) -> None:
if scope and scope not in scopes:
scopes.append(scope)
if user_api_key_dict.user_id:
_add(user_api_key_dict.user_id)
_add(f"user:{user_api_key_dict.user_id}")
if user_api_key_dict.team_id:
_add(f"team:{user_api_key_dict.team_id}")
if user_api_key_dict.org_id:
_add(f"org:{user_api_key_dict.org_id}")
if user_api_key_dict.api_key:
_add(f"key:{user_api_key_dict.api_key}")
if user_api_key_dict.token:
_add(f"key:{user_api_key_dict.token}")
return scopes
def get_primary_resource_owner_scope(
user_api_key_dict: UserAPIKeyAuth | None,
) -> str | None:
"""Return the canonical owner scope to stamp on newly-created rows.
``None`` for identity-less callers — callers that depend on a primary
scope to record ownership must surface that as a hard error rather
than fall back to a shared sentinel (which would collapse every
identity-less caller into the same logical owner).
"""
if user_api_key_dict is None:
return None
if user_api_key_dict.user_id:
return user_api_key_dict.user_id
if user_api_key_dict.team_id:
return f"team:{user_api_key_dict.team_id}"
if user_api_key_dict.org_id:
return f"org:{user_api_key_dict.org_id}"
if user_api_key_dict.api_key:
return f"key:{user_api_key_dict.api_key}"
if user_api_key_dict.token:
return f"key:{user_api_key_dict.token}"
return None
def user_can_access_resource_owner(
owner: str | None,
user_api_key_dict: UserAPIKeyAuth | None,
) -> bool:
if user_api_key_dict is None:
return True
if is_proxy_admin(user_api_key_dict):
return True
if owner is None:
return False
return owner in get_resource_owner_scopes(user_api_key_dict)
__all__ = [
"Final",
"LitellmUserRoles",
"UserAPIKeyAuth",
"get_primary_resource_owner_scope",
"get_resource_owner_scopes",
"is_proxy_admin",
"user_can_access_resource_owner",
]

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

@ -9,7 +9,15 @@ from pydantic import TypeAdapter, ValidationError
from litellm.exceptions import BadRequestError
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.types.decisions import DecisionsRequestBody, OpenAIDecisionRequestBody
from litellm.proxy.common_utils.custom_openapi_spec import inline_request_body
from litellm.types.decisions import (
DecisionsRequest,
DecisionsRequestBody,
DecisionsResponse,
OpenAIDecisionRequest,
OpenAIDecisionRequestBody,
OpenAIDecisionResponse,
)
router: Final = APIRouter()
_REQUEST_DATA_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
@ -20,6 +28,51 @@ _OPENAI_DECISION_REQUEST_BODY_ADAPTER: Final[TypeAdapter[OpenAIDecisionRequestBo
_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)
_SYSTEMONE_REQUEST_BODY: Final = inline_request_body(
DecisionsRequest,
{
"model": "jev",
"state": "Customer wrote: I was charged twice for order #4411 and want one charge refunded today.",
"questions": {
"is_refund_request": {"type": "noul", "instructions": "Is the customer asking for a refund?"},
"urgency": {
"type": "choice",
"instructions": "How urgent is this?",
"criteria": {"low": "can wait a week", "high": "needs action today"},
},
"frustration": {
"type": "score",
"instructions": "How frustrated is the customer?",
"criteria": ["calm", "mildly annoyed", "angry"],
},
},
},
)
_DECISIONS_REQUEST_BODY: Final = inline_request_body(
OpenAIDecisionRequest,
{
"model": "luna",
"input": "Customer wrote: I was charged twice for order #4411 and want one charge refunded today.",
"questions": [
{"type": "predicate", "name": "is_refund_request", "instructions": "Is the customer asking for a refund?"},
{
"type": "choice",
"name": "urgency",
"instructions": "How urgent is this?",
"choices": [
{"value": "low", "description": "can wait a week"},
{"value": "high", "description": "needs action today"},
],
},
{
"type": "score",
"name": "frustration",
"instructions": "How frustrated is the customer?",
"levels": [{"label": "calm"}, {"label": "mildly annoyed"}, {"label": "angry"}],
},
],
},
)
async def _invalid_request(
@ -121,18 +174,35 @@ async def _process_decisions(
dependencies=[Depends(user_api_key_auth)],
response_class=ORJSONResponse, # pyright: ignore[reportDeprecated] # required endpoint contract
tags=["decisions"],
summary="Ask named questions about a state (System One format)",
responses={200: {"model": DecisionsResponse}},
openapi_extra={"requestBody": _SYSTEMONE_REQUEST_BODY},
)
@router.post(
"/systemone",
dependencies=[Depends(user_api_key_auth)],
response_class=ORJSONResponse, # pyright: ignore[reportDeprecated] # required endpoint contract
tags=["decisions"],
summary="Ask named questions about a state (System One format)",
responses={200: {"model": DecisionsResponse}},
openapi_extra={"requestBody": _SYSTEMONE_REQUEST_BODY},
)
async def systemone(
request: Request,
fastapi_response: Response,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Judge a `state` against 1 to 128 named `questions` and get one calibrated answer per question name.
Question types are `noul` (probability the answer is yes), `choice` (one label out of `criteria`, with
`confidence` and `probabilities`) and `score` (an index into the `criteria` levels, with `confidence`,
`legend` and `probabilities`). Requests that break these rules are rejected with a 400 before any provider
is called. `model` is any decision model in the proxy model list. OpenAI decision models accept this format
too, LiteLLM translates it and keeps the answers keyed by question name. Streaming is not supported.
[Docs](https://docs.litellm.ai/docs/decisions)
"""
return await _process_decisions(
request=request,
fastapi_response=fastapi_response,
@ -146,18 +216,35 @@ async def systemone(
dependencies=[Depends(user_api_key_auth)],
response_class=ORJSONResponse, # pyright: ignore[reportDeprecated] # required endpoint contract
tags=["decisions"],
summary="Ask questions about an input (OpenAI Decisions format)",
responses={200: {"model": OpenAIDecisionResponse}},
openapi_extra={"requestBody": _DECISIONS_REQUEST_BODY},
)
@router.post(
"/decisions",
dependencies=[Depends(user_api_key_auth)],
response_class=ORJSONResponse, # pyright: ignore[reportDeprecated] # required endpoint contract
tags=["decisions"],
summary="Ask questions about an input (OpenAI Decisions format)",
responses={200: {"model": OpenAIDecisionResponse}},
openapi_extra={"requestBody": _DECISIONS_REQUEST_BODY},
)
async def decisions(
request: Request,
fastapi_response: Response,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Judge an `input` against 1 to 128 `questions` and get the answers back in the same order.
Question types are `predicate` (yes/no with `probability`), `choice` (one of the `choices`, with
`probabilities`) and `score` (an index into `levels`, with `probabilities`). Requests that break these
rules are rejected with a 400 before any provider is called. `model` is any decision model in the proxy
model list. System One decision models accept this format too, LiteLLM translates the request and the
answers. Streaming is not supported.
[Docs](https://docs.litellm.ai/docs/decisions)
"""
return await _process_decisions(
request=request,
fastapi_response=fastapi_response,

39
litellm/proxy/litellm_pre_call_utils.py
View file

@ -227,6 +227,7 @@ from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES
from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS
from litellm.types.services import ServiceTypes
from litellm.types.utils import (
CAPTURE_MESSAGE_CONTENT_VAR,
CustomPricingLiteLLMParams,
LlmProviders,
ProviderSpecificHeader,
@ -975,6 +976,8 @@ def convert_key_logging_metadata_to_callback(
team_callback_settings_obj.callbacks.append(data.callback_name)
for var, value in data.callback_vars.items():
if var == CAPTURE_MESSAGE_CONTENT_VAR:
continue
# New Relic routing reads these from the trusted-vars overlay with no
# callback-name check, so scope them to the newrelic entry: a team that
# put newrelic_* under a different callback never asked for New Relic and
@ -1162,7 +1165,6 @@ def resolve_tenant_otel_destinations(
span still routes to the tenant's credentials the way it did then.
"""
from litellm.integrations.otel.model.config import is_otel_v2_enabled
from litellm.integrations.otel.presets.destinations import destination_for
if not is_otel_v2_enabled():
return ()
@ -1186,23 +1188,28 @@ def resolve_tenant_otel_destinations(
for name in dict.fromkeys(
callback.callback_name for callback in callbacks if callback.callback_type != "failure"
)
if (
destination := destination_for(
name,
_tenant_otel_params(
MappingProxyType(
{
var: value
for callback in callbacks
if callback.callback_name == name
for var, value in callback.callback_vars.items()
}
)
),
_tenant_service_name(user_api_key_dict),
if (destination := _tenant_destination(name, callbacks, _tenant_service_name(user_api_key_dict))) is not None
)
def _tenant_destination(
name: str, callbacks: Sequence[AddTeamCallback], service_name: str | None
) -> "OtelDestination | None":
from litellm.integrations.otel.model.config import parse_capture_message_content
from litellm.integrations.otel.presets.destinations import destination_for
callback_vars: Final = MappingProxyType(
dict(
itertools.chain.from_iterable(
callback.callback_vars.items() for callback in callbacks if callback.callback_name == name
)
)
is not None
)
return destination_for(
name,
_tenant_otel_params(callback_vars),
service_name,
parse_capture_message_content(callback_vars.get(CAPTURE_MESSAGE_CONTENT_VAR)),
)

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

@ -8040,7 +8040,10 @@ def _validate_key_alias_format(key_alias: str | None) -> None:
raise_if_unsafe_secret_name(key_alias)
except ValueError:
raise ProxyException(
message="Invalid key_alias",
message=(
"Invalid key_alias: must not contain control characters (such as tab or newline) "
+ 'or a ".." path segment'
),
type=ProxyErrorTypes.bad_request_error,
param="key_alias",
code=400,

13
litellm/proxy/management_endpoints/team_callback_endpoints.py
View file

@ -31,8 +31,10 @@ from litellm.proxy._types import (
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.callback_config_validation import (
callback_config_error,
conflicting_capture_error,
conflicting_span_scope_error,
cross_entry_family_error,
stored_capture_entries,
)
from litellm.proxy.common_utils.callback_utils import ( # noqa: F401 # legacy module exports
_CALLBACK_VAR_ENCRYPTED_PREFIX, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
@ -288,6 +290,9 @@ async def add_team_callbacks(
- langfuse_host: The host for the Langfuse callback
- langfuse_environment: The tracing environment for the Langfuse callback (lowercase; falls back to LANGFUSE_TRACING_ENVIRONMENT)
- langfuse_span_scope: For langfuse_otel, "full" (default) sends the whole request trace, "llm_only" sends only the model-call spans
- capture_message_content: For OTel v2 callbacks, "no_content" or "span_only".
An explicit value overrides the global setting for this destination; omitted follows the proxy's own callback for that backend, else the global setting.
"span_only" puts prompt and response content on the destination's spans, "no_content" leaves it out
- gcs_bucket_name: The name of the GCS bucket
- gcs_path_service_account: The path to the GCS service account
- langsmith_api_key: The API key for the Langsmith callback
@ -351,11 +356,19 @@ async def add_team_callbacks(
# Decrypted, because the checks compare the incoming values against
# the stored ones and the credentials are encrypted at rest.
decrypted_logging: Final = decrypt_callback_vars(team_metadata).get("logging")
stored_capture: Final = stored_capture_entries(decrypted_logging)
stored_entries: Final = decrypted_logging if isinstance(decrypted_logging, list) else ()
stored_entry_vars: Final = [entry.get("callback_vars") or {} for entry in stored_entries]
scope_error: Final = conflicting_span_scope_error(data.callback_vars, stored_entry_vars)
if scope_error is not None:
raise _callback_config_error(scope_error)
capture_error: Final = conflicting_capture_error(
data.callback_name,
data.callback_vars,
stored_capture,
)
if capture_error is not None:
raise _callback_config_error(capture_error)
# One entry has to own a credential family end to end. The entries are
# flattened into one dict before a request reads them, so an entry
# naming only a destination would pair with a key written on another

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

@ -1,4 +1,22 @@
[
{
"provider": "SCALEDOWN",
"provider_display_name": "ScaleDown",
"litellm_provider": "scaledown",
"credential_fields": [
{
"key": "api_key",
"label": "API Key",
"placeholder": null,
"tooltip": null,
"required": true,
"field_type": "password",
"options": null,
"default_value": null
}
],
"default_model_placeholder": "scaledown/summarize"
},
{
"provider": "AIML",
"provider_display_name": "AI/ML API",

2
litellm/responses/mcp/chat_completions_handler.py
View file

@ -15,7 +15,7 @@ from litellm.types.utils import Message, ModelResponse
from litellm.utils import CustomStreamWrapper
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class _MCPCompletionKwargs(TypedDict, total=False, extra_items=object):

2
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -42,8 +42,8 @@ if TYPE_CHECKING:
from mcp.types import Tool as MCPTool
from litellm.proxy._experimental.mcp_server.ui_session_utils import GrantedToolsetIds
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
else:
MCPTool = Any

2
litellm/responses/mcp/mcp_streaming_iterator.py
View file

@ -27,8 +27,8 @@ if TYPE_CHECKING:
from mcp.types import Tool as MCPTool
from litellm.proxy._types import UserAPIKeyAuth
from litellm.responses.mcp.litellm_proxy_mcp_handler import MCPToolResult
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
else:
MCPTool = Any

2
litellm/responses/mcp/request_context.py
View file

@ -18,7 +18,7 @@ from pydantic import TypeAdapter
from typing_extensions import NotRequired, ReadOnly, TypedDict
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class _AuthCarryingMetadata(TypedDict):

2
litellm/responses/streaming_iterator.py
View file

@ -50,7 +50,7 @@ from litellm.utils import async_post_call_success_deployment_hook
if TYPE_CHECKING:
from litellm.caching.caching_handler import LLMCachingHandler
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.responses.streaming_websocket import (
PresidioGuardrailCallback,
ResponsesBackendWebSocket,

2
litellm/router_strategy/adaptive_router/hooks.py
View file

@ -28,7 +28,7 @@ from litellm.router_strategy.adaptive_router.config import (
from litellm.router_strategy.adaptive_router.signals import Turn
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
# Identity fields hashed into a derived session key so the same conversation
# from the same caller produces a stable key, while different keys/teams/users

86
litellm/types/decisions.py
View file

@ -15,8 +15,13 @@ MAX_DECISION_QUESTIONS: Final = 128
class NoulQuestion(LiteLLMPydanticObjectBase):
type: Literal["noul"]
instructions: DecisionsJSON | None = None
criteria: NoulCriteria | None = None
instructions: Annotated[DecisionsJSON | None, Field(description="The yes/no question to ask about the state")] = (
None
)
criteria: Annotated[
NoulCriteria | None,
Field(description="Optional descriptions of what makes the answer true and what makes it false"),
] = None
model_config = ConfigDict(extra="allow", frozen=True)
@ -29,16 +34,24 @@ class NoulQuestion(LiteLLMPydanticObjectBase):
class ChoiceQuestion(LiteLLMPydanticObjectBase):
type: Literal["choice"]
instructions: DecisionsJSON | None = None
criteria: Annotated[Mapping[str, DecisionsJSON | None], Field(min_length=1, max_length=255)]
instructions: Annotated[DecisionsJSON | None, Field(description="The question to ask about the state")] = None
criteria: Annotated[
Mapping[str, DecisionsJSON | None],
Field(min_length=1, max_length=255, description="Candidate labels mapped to an optional description of each"),
]
model_config = ConfigDict(extra="allow", frozen=True)
class ScoreQuestion(LiteLLMPydanticObjectBase):
type: Literal["score"]
instructions: DecisionsJSON | None = None
criteria: Annotated[Sequence[DecisionsJSON], Field(min_length=1, max_length=10)]
instructions: Annotated[DecisionsJSON | None, Field(description="The question to ask about the state")] = None
criteria: Annotated[
Sequence[DecisionsJSON],
Field(
min_length=1, max_length=10, description="Scale levels from lowest to highest. The score is a level index"
),
]
model_config = ConfigDict(extra="allow", frozen=True)
@ -55,14 +68,20 @@ DecisionQuestionMap: TypeAlias = Annotated[
class DecisionsRequestBody(LiteLLMPydanticObjectBase):
state: DecisionsJSON
questions: DecisionQuestionMap
state: Annotated[
DecisionsJSON,
Field(description="The thing being judged, such as a support ticket, a document or a chat transcript"),
]
questions: Annotated[
DecisionQuestionMap,
Field(description="Named noul, choice and score questions. Each key becomes a key in the answers"),
]
model_config = ConfigDict(extra="allow", frozen=True)
class DecisionsRequest(DecisionsRequestBody):
model: str
model: Annotated[str, Field(description="A decision model from the proxy model_list")]
@with_config(ConfigDict(extra="allow"))
@ -171,8 +190,8 @@ OpenAIDecisionInput: TypeAlias = str | Sequence[OpenAIDecisionInputMessage]
class OpenAIPredicateQuestion(LiteLLMPydanticObjectBase):
type: Literal["predicate"]
name: str | None = None
instructions: str
name: Annotated[str | None, Field(description="Echoed in the matching answer")] = None
instructions: Annotated[str, Field(description="The yes/no question to ask about the input")]
model_config = ConfigDict(extra="forbid", frozen=True)
@ -192,9 +211,12 @@ def systemone_choice_key(value: str | bool) -> str:
class OpenAIChoiceQuestion(LiteLLMPydanticObjectBase):
type: Literal["choice"]
name: str | None = None
instructions: str
choices: Annotated[Sequence[OpenAIChoiceOption], Field(min_length=2, max_length=255)]
name: Annotated[str | None, Field(description="Echoed in the matching answer")] = None
instructions: Annotated[str, Field(description="The question to ask about the input")]
choices: Annotated[
Sequence[OpenAIChoiceOption],
Field(min_length=2, max_length=255, description="The options the model picks from. Values must be unique"),
]
model_config = ConfigDict(extra="forbid", frozen=True)
@ -215,9 +237,14 @@ class OpenAIScoreLevel(LiteLLMPydanticObjectBase):
class OpenAIScoreQuestion(LiteLLMPydanticObjectBase):
type: Literal["score"]
name: str | None = None
instructions: str
levels: Annotated[Sequence[OpenAIScoreLevel], Field(min_length=2, max_length=10)]
name: Annotated[str | None, Field(description="Echoed in the matching answer")] = None
instructions: Annotated[str, Field(description="The question to ask about the input")]
levels: Annotated[
Sequence[OpenAIScoreLevel],
Field(
min_length=2, max_length=10, description="Scale levels from lowest to highest. The score is a level index"
),
]
model_config = ConfigDict(extra="forbid", frozen=True)
@ -229,13 +256,32 @@ OpenAIDecisionQuestion: TypeAlias = Annotated[
class OpenAIDecisionRequestBody(LiteLLMPydanticObjectBase):
input: OpenAIDecisionInput
questions: Annotated[Sequence[OpenAIDecisionQuestion], Field(min_length=1, max_length=MAX_DECISION_QUESTIONS)]
safety_identifier: str | None = None
input: Annotated[
OpenAIDecisionInput,
Field(
description="The text being judged, or user messages whose content mixes input_text and input_image parts"
),
]
questions: Annotated[
Sequence[OpenAIDecisionQuestion],
Field(
min_length=1,
max_length=MAX_DECISION_QUESTIONS,
description="Predicate, choice and score questions. Answers come back in the same order",
),
]
safety_identifier: Annotated[
str | None,
Field(description="A stable id for your end user. Sent to OpenAI and dropped for other providers"),
] = None
model_config = ConfigDict(extra="allow", frozen=True)
class OpenAIDecisionRequest(OpenAIDecisionRequestBody):
model: Annotated[str, Field(description="A decision model from the proxy model_list")]
class OpenAIPredicateAnswer(LiteLLMPydanticObjectBase):
type: Literal["predicate"] = "predicate"
name: str | None

2
litellm/types/passthrough_endpoints/managed_id_rewriter.py
View file

@ -24,8 +24,8 @@ from pydantic import JsonValue
if TYPE_CHECKING:
from litellm.models.managed_files import LiteLLM_ManagedFileTable
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import OpenAIFileObject
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
SortOrder: TypeAlias = Literal["asc", "desc"]
ResourceKind: TypeAlias = Literal["files", "batches"]

331
litellm/types/proxy/auth/user_api_key_auth.py
View file

@ -1,4 +1,33 @@
from typing import Final
from collections.abc import Mapping
from typing import (
TYPE_CHECKING,
Any, # noqa: TID251 # moved UserAPIKeyAuth fields and the Span alias are typed Any
Final,
TypeAlias,
)
from pydantic import ConfigDict, Field, PrivateAttr, model_validator
from typing_extensions import TypeIs # noqa: TID251 # narrows the before-validator input without copying it
from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
from litellm.models.team import Member
from litellm.models.verification_token import LiteLLM_VerificationToken
from litellm.types.agents import AgentCaller, AgentResponse
from litellm.types.proxy.agent_identity import ManagedAgentContext
from litellm.types.proxy.auth.user_roles import LitellmUserRoles
from litellm.types.proxy.carried_budget_state import (
OrgBudgetSnapshot,
TeamBudgetSnapshot,
UserBudgetSnapshot,
)
from litellm.types.router import AllowedModelRegion
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
Span: TypeAlias = _Span | Any
else:
Span: TypeAlias = Any
def is_jwt(token: str | None) -> bool:
@ -6,3 +35,303 @@ def is_jwt(token: str | None) -> bool:
return False
parts: Final = token.split(".")
return len(parts) == 3
def hash_token(token: str) -> str:
import hashlib
# Hash the string using SHA-256
hashed_token: Final = hashlib.sha256(token.encode()).hexdigest()
return hashed_token
_UntypedDict: TypeAlias = dict[Any, Any]
def _is_str_keyed_dict(value: object) -> TypeIs[dict[str, object]]: # guard-ok: pydantic input and JSON keys are str
return isinstance(value, dict)
class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
"""
Combined view of litellm verification token + litellm team table (select values)
"""
team_spend: float | None = None
team_alias: str | None = None
team_tpm_limit: int | None = None
team_rpm_limit: int | None = None
team_tpd_limit: int | None = None
team_max_budget: float | None = None
team_soft_budget: float | None = None
team_model_max_budget: dict[str, object] | None = None
team_models: list[Any] = []
team_blocked: bool = False
soft_budget: float | None = None
team_model_aliases: _UntypedDict | None = None
team_member: Member | None = None
team_metadata: _UntypedDict | None = None
team_object_permission_id: str | None = None
# Team Member Specific Params
team_member_spend: float | None = None
team_member_tpm_limit: int | None = None
team_member_rpm_limit: int | None = None
# End User Params
end_user_id: str | None = None
end_user_tpm_limit: int | None = None
end_user_rpm_limit: int | None = None
end_user_tpd_limit: int | None = None
end_user_max_budget: float | None = None
end_user_model_max_budget: _UntypedDict | None = None
# Organization Params
organization_alias: str | None = None
organization_max_budget: float | None = None
organization_tpm_limit: int | None = None
organization_rpm_limit: int | None = None
organization_metadata: _UntypedDict | None = None
# Project Params
project_alias: str | None = None
project_metadata: _UntypedDict | None = None
# Time stamps
last_refreshed_at: float | None = None # last time joint view was pulled from db
def __init__(self, **kwargs: object) -> None:
# Handle litellm_budget_table_* keys (budget table overrides when key value is None or empty)
for key, value in list(kwargs.items()):
if key.startswith("litellm_budget_table_") and value is not None:
# Extract the corresponding attribute name
attr_name = key.replace("litellm_budget_table_", "")
# Use key's value from kwargs (from DB view), not class default
current = kwargs.get(attr_name)
if current is None:
current = getattr(self, attr_name, None)
# Apply budget value when key has no value, or for model_max_budget when key has empty dict
should_apply = current is None or (
attr_name == "model_max_budget" and _is_str_keyed_dict(current) and len(current) == 0
)
if should_apply:
kwargs[attr_name] = value
if key == "end_user_id" and value is not None and isinstance(value, int):
kwargs[key] = str(value)
if kwargs.get("organization_id") is not None:
kwargs["org_id"] = kwargs.pop("organization_id")
# Initialize the superclass
super().__init__(**kwargs) # pyright: ignore[reportArgumentType, reportUnknownMemberType] # pydantic validates raw values
class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response object for user api key auth
"""
Return the row in the db
"""
api_key: str | None = None
user_role: LitellmUserRoles | None = None
allowed_model_region: AllowedModelRegion | None = None
parent_otel_span: Span | None = None
rpm_limit_per_model: dict[str, int] | None = None
tpm_limit_per_model: dict[str, int] | None = None
user_tpm_limit: int | None = None
user_rpm_limit: int | None = None
user_email: str | None = None
user_spend: float | None = None
user_max_budget: float | None = None
# Values stay `object` rather than BudgetConfig: this is the raw JSON column,
# and validating it here would make one malformed row fail auth outright.
# resolve_model_budget validates the single entry a request actually needs.
user_model_max_budget: Mapping[str, object] | None = None
team_member_model_max_budget: Mapping[str, object] | None = Field(default=None, exclude=True)
request_route: str | None = None
is_session_token: bool = False
# Server-only marker set exclusively by the MCP gateway admission path
# (reload_admitted_user) for a keyless user-subject admitted via a gateway DCR session
# bearer or bridge envelope. Not a DB column and never populated from caller-controlled key
# metadata or JWT claims, so it cannot be forged to gain the team-inherited MCP grant union
# or to escape the caller-Authorization egress scrub. exclude=True keeps it out of serialization.
mcp_admitted_user_subject: bool = Field(default=False, exclude=True)
requires_fresh_policy: bool = Field(default=False, exclude=True)
mcp_explicit_grants_only: bool = Field(default=False, exclude=True)
# team_id -> that team's mcp_rpm_limit map, for a keyless admitted subject that reaches MCP
# servers through several teams at once and therefore has no single team_id for the limiter to
# key off. Server-only and stripped from validated input for the same reason as the marker
# above: a forged entry would let a caller pick which team's rpm bucket it is charged against.
mcp_source_team_rpm_limits: dict[str, dict[str, int]] | None = Field(default=None, exclude=True)
# The single MCP server_id a gateway session bearer was scoped to at authorize time (RFC 8707
# resource), or None for an aggregate-scope session. A RESTRICTION intersected against the live
# grant resolution, never a grant. Server-only, set exclusively by the MCP gateway admission
# path via post-construction assignment and stripped from validated input like the markers
# above; a forged value could at most narrow, but the stripping keeps the field's provenance
# single-owner so its meaning stays trustworthy.
mcp_session_resource_server_id: str | None = Field(default=None, exclude=True)
mcp_toolset_id: str | None = Field(default=None, exclude=True)
authenticated_by_custom_auth: bool = Field(default=False, exclude=True)
via_virtual_key: bool = Field(
default=False,
exclude=True,
description=(
"Server-only marker set exclusively by the DB virtual-key and master-key auth paths via "
"post-construction assignment. Stripped from validated input so custom auth handlers, JWT "
"claims, or key metadata cannot forge it. Gates overwrite_user_with_key_hash stamping: only "
"a credential the proxy itself validated as a key may be forwarded as the provider-facing "
"user id."
),
)
invoked_agent_id: str | None = Field(default=None, exclude=True)
invoked_agent_policy: AgentResponse | None = Field(default=None, exclude=True)
agent_invocation_cost: float | None = Field(default=None, exclude=True)
billing_agent_policy: AgentResponse | None = Field(default=None, exclude=True)
_managed_delegation_verified: bool = PrivateAttr(default=False)
managed_agent_policy: AgentResponse | None = Field(default=None, exclude=True)
managed_agent_context: ManagedAgentContext | None = Field(default=None, exclude=True)
agent_caller: AgentCaller | None = Field(
default=None,
exclude=True,
description=(
"Set per request from the x-litellm-user-id / x-litellm-team-id headers an agent echoes back on "
"calls made with its own key. Every check treats it as a ceiling, so a forged value can only "
"narrow the agent's access."
),
)
budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True)
team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True)
user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True)
org_budget_snapshot: OrgBudgetSnapshot | None = Field(default=None, exclude=True)
matched_model_access_groups: list[str] | None = Field(default=None, exclude=True)
budget_throttle_pct: float | None = Field(default=None, exclude=True)
user: Any | None = None # Expanded user object when expand=user is used
created_by_user: Any | None = None # Expanded created_by user when expand=user is used
end_user_object_permission: LiteLLM_ObjectPermissionTable | None = None
# Team object_permission preloaded in auth (e.g. get_team_object) to avoid
# per-request object_permission fetches in downstream checks (vector stores, etc.)
team_object_permission: LiteLLM_ObjectPermissionTable | None = None
# Decoded upstream IdP claims (groups, roles, etc.) propagated by JWT auth machinery
# and forwarded into outbound tokens by guardrails such as MCPJWTSigner.
jwt_claims: _UntypedDict | None = None
model_config = ConfigDict(arbitrary_types_allowed=True)
@model_validator(mode="before")
@classmethod
def check_api_key(cls, values: object) -> object:
# If values is already an instance (not a dict), return it as-is
if not _is_str_keyed_dict(values):
return values
# mcp_admitted_user_subject is a server-only marker, set ONLY by the MCP gateway admission
# path via post-construction assignment. Strip it from any validated input (constructor
# kwargs, model_validate, a JWT/key claim splat) so it can never be forged from caller data.
values.pop("mcp_admitted_user_subject", None)
values.pop("requires_fresh_policy", None)
values.pop("mcp_explicit_grants_only", None)
values.pop("mcp_source_team_rpm_limits", None)
values.pop("mcp_session_resource_server_id", None)
values.pop("mcp_toolset_id", None)
values.pop("via_virtual_key", None)
values.pop("authenticated_by_custom_auth", None)
values.pop("agent_caller", None)
values.pop("managed_agent_context", None)
values.pop("managed_agent_policy", None)
values.pop("invoked_agent_id", None)
values.pop("invoked_agent_policy", None)
values.pop("agent_invocation_cost", None)
values.pop("billing_agent_policy", None)
api_key: Final = values.get("api_key")
if api_key is not None:
values.update({"token": cls._safe_hash_litellm_api_key(api_key)}) # pyright: ignore[reportArgumentType] # non-str fails in the hasher
if isinstance(api_key, str):
values.update({"api_key": cls._safe_hash_litellm_api_key(api_key)})
return values
@classmethod
def _safe_hash_litellm_api_key(cls, api_key: str) -> str:
"""
Helper to ensure all logged keys are hashed
Covers:
1. Regular API keys from LiteLLM DB
2. JWT tokens used for connecting to LiteLLM API
"""
normalized = api_key
if normalized[:7].lower() == "bearer ":
normalized = normalized[7:]
if normalized.startswith("sk-"):
return hash_token(normalized)
if is_jwt(token=normalized):
return f"hashed-jwt-{hash_token(token=normalized)}"
return normalized
@classmethod
def get_litellm_internal_health_check_user_api_key_auth(cls) -> "UserAPIKeyAuth":
"""
Returns a `UserAPIKeyAuth` object for the litellm internal health check service account.
This is used to track number of requests/spend for health check calls.
"""
from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
team_id=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
key_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
team_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
)
@classmethod
def get_litellm_cli_user_api_key_auth(cls) -> "UserAPIKeyAuth":
"""
Returns a `UserAPIKeyAuth` object for the litellm internal health check service account.
This is used to track number of requests/spend for health check calls.
"""
from litellm.constants import LITTELM_CLI_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
team_id=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
key_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
team_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
)
@classmethod
def get_litellm_internal_jobs_user_api_key_auth(cls) -> "UserAPIKeyAuth":
"""
Returns a `UserAPIKeyAuth` object for internal LiteLLM jobs like key rotation.
This is used to track actions performed by automated system jobs.
"""
from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
team_id="system",
key_alias=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
team_alias="system",
user_id="system",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
@property
def is_team_service_account(self) -> bool:
return (
self.user_id is None
and self.team_id is not None
and bool(self.metadata) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped base field
and self.metadata.get("service_account_id") is not None # pyright: ignore[reportUnknownMemberType] # untyped base field
)
def user_api_key_has_admin_view(user_api_key_dict: UserAPIKeyAuth) -> bool:
"""Return True if the caller's role grants unscoped read access to all
tenant resources (managed files, batches, vector stores, spend rows, etc).
Lives on _types.py so leaf modules (e.g. litellm.llms.base_llm.managed_resources)
can use it without pulling in litellm.proxy.utils via management_endpoints.
"""
return user_api_key_dict.user_role in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
)

6
litellm/types/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -4,14 +4,14 @@ from typing import Any, Final, Literal
from pydantic import ConfigDict, Field, field_validator
from typing_extensions import ReadOnly, TypedDict
from litellm.proxy._types import (
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.proxy.management_endpoints.management_v1 import ResourceResponse
from litellm.types.proxy.management_endpoints.user_requests import (
LiteLLM_UserTableWithKeyCount,
NewUserRequest,
UpdateUserRequest,
UpdateUserRequestNoUserIDorEmail,
)
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.proxy.management_endpoints.management_v1 import ResourceResponse
MAX_BULK_DELETE_USERS: Final = 500

8
litellm/types/proxy/management_endpoints/key_management_endpoints.py
View file

@ -5,14 +5,14 @@ from pydantic import ConfigDict, model_validator
from typing_extensions import ReadOnly, TypedDict
from litellm.models.verification_token import LiteLLM_VerificationToken
from litellm.proxy._types import (
from litellm.types.llms.base import LiteLLMBaseModel, LiteLLMPydanticObjectBase
from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains
from litellm.types.proxy.management_endpoints.key_requests import (
GenerateKeyRequest,
LiteLLM_ObjectPermissionBase,
RegenerateKeyRequest,
UpdateKeyRequest,
)
from litellm.types.llms.base import LiteLLMBaseModel, LiteLLMPydanticObjectBase
from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains
from litellm.types.proxy.management_endpoints.request_base import LiteLLM_ObjectPermissionBase
class KeyTokenWhere(TypedDict):

138
litellm/types/proxy/management_endpoints/key_requests.py Normal file
View file

@ -0,0 +1,138 @@
import enum
from collections.abc import Mapping
from datetime import datetime
from types import MappingProxyType
from typing import Any, Literal # noqa: TID251 # allowed_routes and metadata stay Any for existing callers
from pydantic import Field, PositiveInt, field_validator, model_validator
from typing_extensions import TypeIs # noqa: TID251 # narrows the raw payload a before-validator receives
from litellm.types.proxy.management_endpoints.request_base import GenerateRequestBase
from litellm.types.router import UpdateRouterConfig
from litellm.types.utils import LiteLLMPydanticObjectBase
class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase):
index_name: str
index_permissions: list[Literal["read", "write"]]
class KeyRequestBase(GenerateRequestBase):
key: str | None = None
tpd_limit: int | None = None
default_estimated_output_tokens: PositiveInt | None = None
default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None
budget_id: str | None = None
end_user_budget_id: str | None = None
tags: list[str] | None = None
disable_global_guardrails: bool | None = None
enable_prompt_caching: bool | None = None
throttle_on_budget_exceeded: bool | None = None
enforced_params: list[str] | None = None
allowed_routes: list[Any] | None = []
allowed_passthrough_routes: list[object] | None = None
denied_passthrough_routes: list[str] | None = None
allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None
rpm_limit_type: Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] | None = (
None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm
)
tpm_limit_type: Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] | None = (
None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm
)
router_settings: UpdateRouterConfig | None = None
access_group_ids: list[str] | None = None
class LiteLLMKeyType(str, enum.Enum):
"""
Enum for key types that determine what routes a key can access
"""
LLM_API = "llm_api" # Can call LLM API routes (chat/completions, embeddings, etc.)
MANAGEMENT = "management" # Can call management routes (user/team/key management)
READ_ONLY = "read_only" # Can only call info/read routes
DEFAULT = "default" # Uses default allowed routes
class GenerateKeyRequest(KeyRequestBase):
soft_budget: float | None = None
send_invite_email: bool | None = None
key_type: LiteLLMKeyType | None = Field(
default=LiteLLMKeyType.DEFAULT,
description="Type of key that determines default allowed routes.",
)
auto_rotate: bool | None = Field(default=False, description="Whether this key should be automatically rotated")
rotation_interval: str | None = Field(
default=None,
description="How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True",
)
organization_id: str | None = None
project_id: str | None = None
@field_validator("team_id", "organization_id", "project_id", mode="before")
@classmethod
def treat_cleared_id_as_unset(cls, v: object) -> object:
if v == "":
return None
return v
def _is_mapping(value: object) -> TypeIs[Mapping[object, object]]: # guard-ok: raw before-validator payload
return isinstance(value, Mapping)
class UpdateKeyRequest(KeyRequestBase):
# Note: the defaults of all Params here MUST BE NONE
# else they will get overwritten
duration: str | None = None
spend: float | None = None
soft_budget: float | None = None
metadata: dict[Any, Any] | None = None
temp_budget_increase: float | None = None
temp_budget_expiry: datetime | None = None
auto_rotate: bool | None = None
rotation_interval: str | None = None
organization_id: str | None = None
project_id: str | None = Field(
default=None,
description="Omit to retain the project, or send null to detach. Assigning a different project is not supported.",
)
@model_validator(mode="before")
@classmethod
def drop_blank_team_id(cls, values: object) -> object:
if _is_mapping(values) and values.get("team_id") == "":
return MappingProxyType({k: v for k, v in values.items() if k != "team_id"})
return values
@field_validator("organization_id", mode="before")
@classmethod
def treat_cleared_organization_id_as_unset(cls, v: object) -> object:
if v == "":
return None
return v
@model_validator(mode="after")
def validate_temp_budget(self) -> "UpdateKeyRequest":
if self.temp_budget_increase is not None or self.temp_budget_expiry is not None:
if self.temp_budget_increase is None or self.temp_budget_expiry is None:
raise ValueError("temp_budget_increase and temp_budget_expiry must be set together")
return self
@model_validator(mode="after")
def validate_key_identifier(self) -> "UpdateKeyRequest":
if self.key is None and self.key_alias is None:
raise ValueError("either key or key_alias must be provided")
return self
class RegenerateKeyRequest(GenerateKeyRequest):
# This needs to be different from UpdateKeyRequest, because "key" is optional for this
key: str | None = None
new_key: str | None = None
duration: str | None = None
spend: float | None = None
metadata: dict[Any, Any] | None = None
new_master_key: str | None = None
grace_period: str | None = None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke

19
litellm/types/proxy/management_endpoints/team_endpoints.py
View file

@ -1,16 +1,15 @@
from collections.abc import Mapping
from typing import Any, Final, Literal
from pydantic import ConfigDict, Field, field_validator, model_validator
from litellm.models.team import LiteLLM_DeletedTeamTable, LiteLLM_TeamTable, Member
from litellm.models.team_membership import LiteLLM_TeamMembership
from litellm.proxy._types import (
KeyManagementRoutes,
MemberDeleteRequest,
)
from litellm.proxy.common_utils.timezone_utils import budget_duration_error
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.proxy.auth.user_roles import KeyManagementRoutes
from litellm.types.proxy.management_endpoints.management_v1 import ResourceResponse
from litellm.types.utils import LiteLLMPydanticObjectBase
TeamIdSearchMatch = Literal["exact", "prefix"]
@ -19,6 +18,18 @@ MAX_BULK_TEAM_MEMBER_DELETES: Final = 500
MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES: Final = 500
class MemberDeleteRequest(LiteLLMPydanticObjectBase):
user_id: str | None = None
user_email: str | None = None
@model_validator(mode="before")
@classmethod
def check_user_info(cls, values: Mapping[str, object]) -> Mapping[str, object]:
if values.get("user_id") is None and values.get("user_email") is None:
raise ValueError("Either user id or user email must be provided")
return values
class GetTeamMemberPermissionsRequest(LiteLLMBaseModel):
"""Request to get the team member permissions for a team"""

82
litellm/types/proxy/management_endpoints/user_requests.py Normal file
View file

@ -0,0 +1,82 @@
from collections.abc import Mapping
from typing import Literal
from pydantic import Field, field_validator, model_validator
from litellm.models.user import LiteLLM_UserTable
from litellm.types.proxy.auth.user_roles import LitellmUserRoles
from litellm.types.proxy.management_endpoints.request_base import GenerateRequestBase
from litellm.types.utils import LiteLLMPydanticObjectBase
class NewUserRequestTeam(LiteLLMPydanticObjectBase):
team_id: str
max_budget_in_team: float | None = None
user_role: Literal["user", "admin"] = "user"
class NewUserRequest(GenerateRequestBase):
max_budget: float | None = None
user_email: str | None = None
user_alias: str | None = None
user_role: (
Literal[
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
]
| None
) = None
teams: list[str] | list[NewUserRequestTeam] | None = None
auto_create_key: bool = True # flag used for returning a key as part of the /user/new response
send_invite_email: bool | None = None
sso_user_id: str | None = None
organizations: list[str] | None = None
password: str | None = None
@field_validator("password")
@classmethod
def password_not_supported(cls, value: str | None) -> str | None:
if value is not None:
raise ValueError(
"password cannot be set via /user/new. Users set their own password through an "
"invitation link (POST /invitation/new)."
)
return value
class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with BulkUpdateUserRequest
# repr=False keeps the plaintext out of management-endpoint alerts, which str() the request model
password: str | None = Field(default=None, repr=False)
spend: float | None = None
metadata: dict[object, object] | None = None
user_alias: str | None = None
user_role: (
Literal[
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
]
| None
) = None
max_budget: float | None = None
class UpdateUserRequest(UpdateUserRequestNoUserIDorEmail):
# Note: the defaults of all Params here MUST BE NONE
# else they will get overwritten
user_id: str | None = None
user_email: str | None = None
@model_validator(mode="before")
@classmethod
def check_user_info(cls, values: Mapping[str, object]) -> Mapping[str, object]:
if values.get("user_id") is None and values.get("user_email") is None:
raise ValueError("Either user id or user email must be provided")
return values
class LiteLLM_UserTableWithKeyCount(LiteLLM_UserTable):
key_count: int = 0

23
litellm/types/utils.py
View file

@ -3726,6 +3726,28 @@ ArizeOtlpProtocol = Literal["grpc", "http/protobuf"]
ARIZE_OTLP_PROTOCOLS: Final[frozenset[str]] = frozenset(get_args(ArizeOtlpProtocol))
CAPTURE_MESSAGE_CONTENT_VAR: Final = "capture_message_content"
class CaptureMessageContent(str):
NO_CONTENT = "no_content"
SPAN_ONLY = "span_only"
EVENT_ONLY = "event_only"
SPAN_AND_EVENT = "span_and_event"
# The values a team or key destination may set. OTel v2 writes content only to span attributes,
# so the event modes stay global-only settings
CAPTURE_MESSAGE_CONTENT_VALUES: Final[frozenset[str]] = frozenset(
{CaptureMessageContent.NO_CONTENT, CaptureMessageContent.SPAN_ONLY}
)
def captures_span_content(mode: str | None) -> bool:
"""Whether a capture mode puts prompt and response content on spans"""
return mode in (CaptureMessageContent.SPAN_ONLY, CaptureMessageContent.SPAN_AND_EVENT)
class StandardCallbackDynamicParams(TypedDict, total=False):
# Langfuse dynamic params
langfuse_public_key: str | None
@ -4202,6 +4224,7 @@ class LlmProviders(str, Enum):
FRIENDLIAI = "friendliai"
FEATHERLESS_AI = "featherless_ai"
WATSONX = "watsonx"
SCALEDOWN = "scaledown"
WATSONX_TEXT = "watsonx_text"
TRITON = "triton"
PREDIBASE = "predibase"

1
litellm/utils.py
View file

@ -8655,6 +8655,7 @@ class ProviderConfigManager:
False,
),
LlmProviders.FEATHERLESS_AI: (lambda: litellm.FeatherlessAIConfig(), False),
LlmProviders.SCALEDOWN: (lambda: litellm.ScaleDownChatConfig(), False),
LlmProviders.NOVITA: (lambda: litellm.NovitaConfig(), False),
LlmProviders.NEBIUS: (lambda: litellm.NebiusConfig(), False),
LlmProviders.WANDB: (lambda: litellm.WandbConfig(), False),

48
model_prices_and_context_window.json
View file

@ -66876,7 +66876,7 @@
"supports_audio_input": true
},
"azure/gpt-realtime-translate": {
"deprecation_date": "2027-05-06",
"deprecation_date": "2027-11-09",
"input_cost_per_second": 0.000566666666667,
"litellm_provider": "azure",
"max_input_tokens": 32000,
@ -71842,5 +71842,51 @@
"image"
],
"supports_embedding_image_input": true
},
"scaledown/extract": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": false,
"supports_function_calling": false
},
"scaledown/summarize": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_system_messages": true,
"supports_vision": false,
"supports_function_calling": false
},
"scaledown/compress": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_system_messages": true,
"supports_function_calling": false,
"supports_vision": false
},
"scaledown/classify": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_system_messages": false,
"supports_vision": false,
"supports_function_calling": false
},
"scaledown/decisions": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_system_messages": false,
"supports_vision": false,
"supports_function_calling": false
}
}

18
provider_endpoints_support.json
View file

@ -2349,6 +2349,24 @@
"search": true
}
},
"scaledown": {
"display_name": "ScaleDown (`scaledown`)",
"url": "https://docs.litellm.ai/docs/providers/scaledown",
"endpoints": {
"chat_completions": true,
"messages": false,
"responses": false,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,
"batches": false,
"rerank": false,
"a2a": false,
"interactions": false
}
},
"sambanova": {
"display_name": "Sambanova (`sambanova`)",
"url": "https://docs.litellm.ai/docs/providers/sambanova",

73
scripts/check_openapi_docs.py Executable file
View file

@ -0,0 +1,73 @@
#!/usr/bin/env python3
"""List the POST operations in the proxy's OpenAPI document that Swagger UI cannot explain.
Every finding is one line, ``<METHOD> <path> <code>``, so scripts/openapi_docs_gate.py can
count them per rule the same delta-vs-base way as the other lint gates. The rules:
OA001 no request body schema (the handler reads ``request.body()`` and never told FastAPI
what is in it, see ``inline_request_body`` in custom_openapi_spec.py for the fix)
OA002 no description (the handler has no docstring and the route no ``description=``)
OA003 the 200 response has no schema (no ``response_model`` and no ``responses={200: ...}``)
"""
from __future__ import annotations
import json
import sys
from collections.abc import Iterator, Mapping
from typing import Final, NamedTuple
GATED_METHODS: Final = frozenset({"post"})
class Violation(NamedTuple):
method: str
path: str
code: str
def _as_mapping(node: object) -> Mapping[str, object]:
return node if isinstance(node, Mapping) else {}
def _json_schema(container: object) -> object:
content: Final = _as_mapping(_as_mapping(container).get("content"))
return _as_mapping(content.get("application/json")).get("schema")
def _operation_codes(operation: Mapping[str, object]) -> Iterator[str]:
if not _json_schema(operation.get("requestBody")):
yield "OA001"
description: Final = operation.get("description")
if not isinstance(description, str) or not description.strip():
yield "OA002"
if not _json_schema(_as_mapping(operation.get("responses")).get("200")):
yield "OA003"
def violations(spec: Mapping[str, object]) -> tuple[Violation, ...]:
paths: Final = _as_mapping(spec.get("paths"))
return tuple(
Violation(method.upper(), path, code)
for path, operations in sorted(paths.items())
for method, operation in sorted(_as_mapping(operations).items())
if method in GATED_METHODS
for code in _operation_codes(_as_mapping(operation))
) # comprehension-ok: one pass over path, method and code is the whole scan
def proxy_spec() -> Mapping[str, object]:
from litellm.proxy.proxy_server import app
return _as_mapping(app.openapi())
def main(argv: list[str] = sys.argv[1:]) -> int:
spec: Final = json.loads(open(argv[0]).read()) if argv else proxy_spec()
for violation in violations(spec):
print(f"{violation.method} {violation.path} {violation.code}")
return 0
if __name__ == "__main__":
raise SystemExit(main())

178
scripts/openapi_docs_gate.py Executable file
View file

@ -0,0 +1,178 @@
#!/usr/bin/env python3
"""Delta-vs-base gate for the OA* rules in scripts/check_openapi_docs.py.
Sibling of scripts/test_quality_gate.py, pointed at the proxy's OpenAPI document
instead of a source tree. Each rule is counted over every POST operation the proxy
serves at HEAD and at the merge-base with the branch this change merges into, and the
gate fails only when a rule grew past the merge-base count, so a new endpoint has to
ship with a request schema, a description and a 200 schema while the undocumented
routes that already exist are only ever allowed to shrink.
The merge-base counts come from scripts/lint_base_counts.py: the disk cache, then the
CI artifact published for that commit, then a pass of the current checker over a
detached worktree at the merge-base with that tree first on PYTHONPATH.
"""
from __future__ import annotations
import argparse
import os
import re
import shutil
import signal
import subprocess
import sys
import tempfile
from collections import Counter
from collections.abc import Mapping, Sequence
from pathlib import Path
from types import FrameType, MappingProxyType
from typing import Final, NamedTuple
from lint_base_counts import (
Checker,
base_counts_cached,
emit_counts,
evaluate,
head_sha,
resolve_base_point,
sha256_of,
)
REPO_ROOT: Final = Path(__file__).resolve().parent.parent
CHECKER: Final = REPO_ROOT / "scripts" / "check_openapi_docs.py"
TARGET: Final = "litellm"
TERMINATION_SIGNALS: Final = (signal.SIGTERM, signal.SIGHUP)
_LINE: Final = re.compile(r"^(?P<method>[A-Z]+) (?P<path>\S+) (?P<code>OA\d+)$")
_ADDED_LINE: Final = re.compile(r"^\+(?!\+\+)(?P<text>.*)$", re.MULTILINE)
class Violation(NamedTuple):
method: str
path: str
code: str
def _run(cmd: Sequence[str], cwd: Path = REPO_ROOT, env: Mapping[str, str] | None = None) -> str:
proc: Final = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True, env=env)
if proc.returncode != 0:
sys.stderr.write(proc.stderr)
raise SystemExit(f"{cmd[0]} exited {proc.returncode}")
return proc.stdout
def checker_identity(checker: Path = CHECKER) -> Checker:
return Checker("openapi-docs", (sha256_of(checker),))
def parse_violations(out: str) -> tuple[Violation, ...]:
return tuple(
Violation(match.group("method"), match.group("path"), match.group("code"))
for line in out.splitlines()
if (match := _LINE.match(line)) is not None
)
def _check(root: Path, checker: Path) -> tuple[Violation, ...]:
resolved: Final = root.resolve()
env: Final = {**os.environ, "PYTHONPATH": os.pathsep.join((str(resolved), os.environ.get("PYTHONPATH", "")))}
return parse_violations(_run([sys.executable, str(checker)], cwd=resolved, env=env))
def head_violations() -> tuple[Violation, ...]:
return _check(REPO_ROOT, CHECKER)
def count_by_rule(violations: Sequence[Violation]) -> Mapping[str, int]:
return MappingProxyType(dict(Counter(v.code for v in violations)))
def _exit_on_termination(signum: int, _frame: FrameType | None) -> None:
raise SystemExit(128 + signum)
def _install_termination_handlers() -> None:
for termination in TERMINATION_SIGNALS:
if signal.getsignal(termination) == signal.SIG_DFL:
signal.signal(termination, _exit_on_termination)
def base_counts(ref: str, repo_root: Path = REPO_ROOT, checker: Path = CHECKER) -> Mapping[str, int]:
"""Rule counts at `ref`, measured with the *current* checker over the base tree."""
_install_termination_handlers()
parent: Final = Path(tempfile.mkdtemp(prefix="oa_base_"))
worktree: Final = parent / "wt"
try:
_run(["git", "worktree", "add", "--detach", str(worktree), ref], cwd=repo_root)
(worktree / "scripts").mkdir(parents=True, exist_ok=True)
base_checker: Final = worktree / "scripts" / "check_openapi_docs.py"
shutil.copy(checker, base_checker)
return count_by_rule(_check(worktree, base_checker))
finally:
# Teardown must never raise, or it masks the real error when the body failed.
subprocess.run(
["git", "worktree", "remove", "--force", str(worktree)],
cwd=repo_root, capture_output=True, text=True,
)
shutil.rmtree(parent, ignore_errors=True)
def added_text(diff_text: str) -> str:
return "\n".join(match.group("text") for match in _ADDED_LINE.finditer(diff_text))
def introduced(violations: Sequence[Violation], added: str) -> tuple[Violation, ...]:
"""The violations whose route path is quoted on a line this change added. A new
route is registered with its path as a string literal, so this names the endpoint
the change brought in without needing the base document's own violation list."""
return tuple(v for v in violations if f'"{v.path}"' in added or f"'{v.path}'" in added)
def cmd_check(base: str) -> None:
head: Final = head_violations()
base_point: Final = resolve_base_point(base)
breaches: Final = evaluate(count_by_rule(head), base_counts_cached(checker_identity(), base_point, base_counts))
if not breaches:
print(f"OK: no OA rule grew past its merge-base count (base {base})")
return
diff: Final = _run(["git", "diff", base_point, "--unified=0", "--no-color", "--", TARGET])
new: Final = introduced(head, added_text(diff))
print(f"FAIL: OA-rule totals grew past their merge-base count (base {base}):")
for breach in breaches:
print(f" {breach.rule}: total {breach.total} over ceiling {breach.ceiling} (this change added {breach.added})")
for violation in sorted(v for v in new if v.code == breach.rule):
print(f" {violation.method} {violation.path}")
print(
"Give every new POST route a request body schema (a typed body parameter, or "
"`openapi_extra={'requestBody': inline_request_body(Model, example)}` when the handler reads "
"`request.body()`), a docstring, and a 200 schema (`response_model=` or `responses={200: {'model': ...}}`), "
"or document an equal number of existing routes; the ceiling is the merge-base count. Run "
"`python scripts/check_openapi_docs.py` to see every finding."
)
raise SystemExit(1)
def main() -> None:
parser: Final = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)")
parser.add_argument(
"--emit-counts-dir",
type=Path,
help="Write HEAD's per-rule counts to this directory as a base-counts artifact instead of gating",
)
args: Final = parser.parse_args()
from default_branch import resolve_base_ref
from gate_slot_lock import held_slot
if args.emit_counts_dir is not None:
with held_slot():
emit_counts(checker_identity(), count_by_rule(head_violations()), args.emit_counts_dir, head_sha())
return
base_ref: Final = resolve_base_ref(args.base, REPO_ROOT)
with held_slot():
cmd_check(base_ref)
if __name__ == "__main__":
main()

2
tests/code_coverage_tests/recursive_detector.py
View file

@ -58,6 +58,7 @@ IGNORE_FUNCTIONS = [
"sanitize_oci_schema", # OCI: bounded by JSON-schema tree depth.
"_freeze_for_dedupe", # OTEL: max depth set (default 16, _FREEZE_MAX_DEPTH); fails closed by returning repr(value) at the cap.
"apply_json_merge_patch", # max depth set (_MAX_MERGE_DEPTH=64); fails closed by raising ValueError at the cap.
"_inlined", # max depth set (_MAX_INLINE_DEPTH=32); passes the schema node through untouched at the cap. Walks a Pydantic JSON schema once at import time.
"_filter_argument_value", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by blocking the tool call at the cap.
"_redact_scanned_content", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by returning "[REDACTED]" at the cap.
"replace_ciphertexts", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); walks stored JSON, which has no cycles, and leaves values below the cap untouched.
@ -67,6 +68,7 @@ IGNORE_FUNCTIONS = [
"strict_json_schema", # harness: max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by raising ValueError at the cap.
"toml_value", # harness/codex: max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by raising OptionsMismatch at the cap.
"with_json_string_leaves", # transitively bounded: only runs on a tree json_string_leaves already walked under the cap.
"_clean_extraction", # ScaleDown: traverses the entity map built under MAX_SCHEMA_DEPTH=32 and MAX_ENTITIES=1000.
"json_unrewritable_labels", # max depth set (MAX_STRUCTURED_CONTENT_SCAN_DEPTH); returns the None sentinel at the cap so the caller blocks.
"_flatten_form_field", # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible).
"_flatten_form_data_field", # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible).

51
tests/integration/_support/mcp.py
View file

@ -4,7 +4,7 @@ import os
import queue
import sys
import time
from collections.abc import Callable, Iterator, Mapping
from collections.abc import Callable, Generator, Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
@ -383,6 +383,55 @@ def peer_of(kind: PeerKind, *, rich: bool = False) -> Iterator[McpPeer]:
yield candidate
@contextmanager
def stateful_mcp_peer() -> Generator[McpPeer]:
service: Final = MCPServer("integration-stateful")
selected: Final[dict[str, str]] = {} # mutable-ok: per-session state the upstream keeps across tool calls
def upstream_session(ctx: Context) -> str:
assert ctx.headers is not None, "stateful upstream requires HTTP request headers"
return ctx.headers["mcp-session-id"]
@service.tool()
def select_project(name: str, ctx: Context) -> str:
selected[upstream_session(ctx)] = name
return f"selected {name}"
@service.tool()
def create_feature(title: str, ctx: Context) -> str:
project: Final = selected.get(upstream_session(ctx))
if project is None:
raise ValueError("no project selected in this session")
return f"{project}/{title}"
app: Final = service.streamable_http_app(
stateless_http=False,
json_response=True,
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
)
observed: Final[queue.Queue[dict[str, object]]] = queue.Queue()
async def capture(scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] == "http" and scope["method"] == "GET":
await Response(status_code=405)(scope, receive, send)
return
if scope["type"] != "http" or scope["method"] != "POST":
await app(scope, receive, send)
return
body: Final = await StarletteRequest(scope, receive).body()
observed.put({"body": json.loads(body) if body else None, "headers": dict(scope["headers"])})
message: Final[Message] = {"type": "http.request", "body": body, "more_body": False}
pending: Final = iter((message,))
async def replay() -> Message:
return next(pending, {"type": "http.disconnect"})
await app(scope, replay, send)
with asgi_server(capture) as url:
yield McpPeer(url + "/mcp", observed)
def register_mcp(scenario: Scenario, peer: McpPeer, alias: str, **fields: object) -> str:
response: Final = scenario.gateway.request(
"POST", "/v1/mcp/server", {"server_name": alias, "alias": alias, **peer.registration(), **fields}

5
tests/integration/_support/otlp_sink.py
View file

@ -186,6 +186,7 @@ class _Handler(BaseHTTPRequestHandler):
"count": len(recorded),
"host": self.headers.get("host", ""),
"headers": dict(self.headers),
"span_ids": [span["span_id"] for span in recorded],
}
)
self._send_json({"recorded": len(recorded)}, status=self.state.status)
@ -466,9 +467,7 @@ def _await_sink(url: str) -> None:
def owned_grpc_sink(directory: Path) -> Iterator[GrpcSink]:
http_port: Final = _free_port()
grpc_port: Final = _free_port()
with _spawn_sink(
directory, "otlp-grpc-sink.log", ["--port", str(http_port), "--grpc-port", str(grpc_port)]
):
with _spawn_sink(directory, "otlp-grpc-sink.log", ["--port", str(http_port), "--grpc-port", str(grpc_port)]):
control_url: Final = f"http://127.0.0.1:{http_port}"
_await_sink(control_url)
yield GrpcSink(url=f"http://127.0.0.1:{grpc_port}", control_url=control_url)

45
tests/integration/mcp/test_mcp_lifecycle.py
View file

@ -1,3 +1,4 @@
import asyncio
import functools
import json
import uuid
@ -5,13 +6,18 @@ from collections.abc import Mapping
from contextlib import ExitStack
from hashlib import sha256
from pathlib import Path
from typing import Final
from typing import Final, cast
import httpx
import httpx2
import pytest
import yaml
from hypothesis import settings
from hypothesis import strategies as st
from hypothesis.stateful import RuleBasedStateMachine, invariant, rule, run_state_machine_as_test
from mcp import ClientSession
from mcp.client.streamable_http import streamable_http_client
from mcp.types import CallToolResult, TextContent
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests
@ -25,6 +31,7 @@ from integration._support.mcp import (
register_mcp,
scripted_peer,
text_result,
stateful_mcp_peer,
tool_calls,
tool_names,
)
@ -466,6 +473,42 @@ def test_key_grant_added_by_key_update_is_visible_to_mcp_tool_listing_before_the
assert set(seen.tools) == {f"{alias}-add", f"{alias}-multiply", f"{alias}-fail"}, seen.raw
async def _select_then_create(url: str, key: str) -> tuple[CallToolResult, CallToolResult]:
async with httpx2.AsyncClient(headers={"x-litellm-api-key": key}, timeout=30, trust_env=False) as http:
async with streamable_http_client(url, http_client=http) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
names: Final = {tool.name.rsplit("-", 1)[-1]: tool.name for tool in (await session.list_tools()).tools}
selected: Final = await session.call_tool(names["select_project"], {"name": "alpha"})
created: Final = await session.call_tool(names["create_feature"], {"title": "login"})
return selected, created
def _upstream_session_ids(observed: tuple[dict[str, object], ...]) -> set[str | None]:
headers: Final = tuple(cast(dict[bytes, bytes], item["headers"]) for item in observed)
return {raw.decode() if isinstance(raw := item.get(b"mcp-session-id"), bytes) else None for item in headers}
def test_stateful_upstream_keeps_selection_across_tool_calls_through_gateway(gateway: Gateway) -> None:
with stateful_mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "integration" + uuid.uuid4().hex
identity: Final = register_mcp(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
peer.drain()
selected, created = asyncio.run(_select_then_create(f"{gateway.client.base_url}/{alias}/mcp", key))
observed: Final = peer.drain()
for session_id in _upstream_session_ids(observed):
if session_id is not None:
httpx.delete(peer.url, headers={"mcp-session-id": session_id}, trust_env=False)
call_sessions: Final = _upstream_session_ids(
tuple(item for item in observed if cast(dict[str, object], item["body"]).get("method") == "tools/call")
)
assert selected.is_error is False, selected.model_dump_json()
assert created.is_error is False, created.model_dump_json()
assert isinstance(created.content[0], TextContent) and created.content[0].text == "alpha/login"
assert len(call_sessions) == 1 and None not in call_sessions, call_sessions
def _update_tool_permissions(
gateway: Gateway, key: str, identity: str, permissions: dict[str, list[str]] | None
) -> None:

720
tests/integration/observability/test_otel_team_capture_message_content.py Normal file
View file

@ -0,0 +1,720 @@
import json
import re
import uuid
from collections.abc import Generator, Iterator, Mapping
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Final, Literal
from urllib.parse import urlparse
import httpx
import pytest
import yaml
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value
from integration._support.otlp_sink import (
ConnectSink,
Span,
SpanSinks,
configure_sink,
owned_connect_sink,
recorded_requests,
recorded_spans,
)
from integration._support.process import OwnedProxy, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
MARKER: Final = re.compile(rb"cmc-(?:tool-|fail-)?[0-9a-f]{32}")
JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
SERVER: Final = 2
CALL_ID: Final = "litellm.call_id"
OPERATION: Final = "gen_ai.operation.name"
INPUT_TOKENS: Final = "gen_ai.usage.input_tokens"
Endpoint = Literal["chat", "responses", "messages"]
Capture = Literal["no_content", "span_only"]
def _marker(kind: str = "") -> str:
return f"cmc-{kind}{uuid.uuid4().hex}"
def _usage() -> dict[str, JsonValue]:
return {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9}
def _chat_reply(marker: str, stream: bool) -> Reply:
identity: Final = f"chatcmpl-{uuid.uuid4().hex}"
if marker.startswith("cmc-tool-"):
call: Final[dict[str, JsonValue]] = {
"id": "call_1",
"type": "function",
"function": {"name": "lookup", "arguments": json.dumps({"query": marker})},
}
message: Final[dict[str, JsonValue]] = {"role": "assistant", "content": None, "tool_calls": [call]}
return Reply(
body=json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": message, "finish_reason": "tool_calls"}],
"usage": _usage(),
}
).encode()
)
if not stream:
return Reply(
body=json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": marker}, "finish_reason": "stop"}
],
"usage": _usage(),
}
).encode()
)
chunk: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"}
deltas: Final[tuple[dict[str, JsonValue], ...]] = (
{**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": marker}}]},
{**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
{**chunk, "choices": [], "usage": _usage()},
)
return Reply(
content_type="text/event-stream",
chunks=(*(b"data: " + json.dumps(delta).encode() + b"\n\n" for delta in deltas), b"data: [DONE]\n\n"),
)
def _responses_reply(marker: str, stream: bool) -> Reply:
identity: Final = uuid.uuid4().hex
response: Final[dict[str, JsonValue]] = {
"id": f"resp_{identity}",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o-mini",
"output": [
{
"id": f"msg_{identity}",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": marker, "annotations": []}],
}
],
"usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9},
}
if not stream:
return Reply(body=json.dumps(response).encode())
events: Final[tuple[dict[str, JsonValue], ...]] = (
{
"type": "response.created",
"sequence_number": 0,
"response": {**response, "status": "in_progress", "output": []},
},
{
"type": "response.output_text.delta",
"sequence_number": 1,
"item_id": f"msg_{identity}",
"output_index": 0,
"content_index": 0,
"delta": marker,
},
{"type": "response.completed", "sequence_number": 2, "response": response},
)
return Reply(
content_type="text/event-stream",
chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events),
)
def _upstream(request: Request) -> Reply:
found: Final = MARKER.search(request.body)
if found is None:
return Reply(status=404, body=b'{"error":"no marker"}')
marker: Final = found.group(0).decode()
if marker.startswith("cmc-fail-"):
return Reply(status=500, body=b'{"error":{"message":"scripted upstream failure","type":"server_error"}}')
stream: Final = object_value(JSON.validate_json(request.body)).get("stream") is True
if request.target.endswith("/responses"):
return _responses_reply(marker, stream)
return _chat_reply(marker, stream)
def _body(model: str, endpoint: Endpoint, marker: str, stream: bool) -> tuple[str, dict[str, JsonValue]]:
if endpoint == "responses":
return "/v1/responses", {"model": model, "input": marker, "stream": stream}
messages: Final[list[JsonValue]] = [{"role": "user", "content": marker}]
if endpoint == "messages":
return "/v1/messages", {"model": model, "max_tokens": 16, "messages": messages, "stream": stream}
return "/v1/chat/completions", {"model": model, "messages": messages, "stream": stream}
@dataclass(frozen=True, slots=True)
class Sent:
marker: str
trace_id: str
call_id: str
status: int
text: str
@dataclass(frozen=True, slots=True)
class Cursors:
operator: int
tenant: int
arize: int
newrelic: int
@dataclass(frozen=True, slots=True)
class Rig:
proxy: Gateway
owned: OwnedProxy
scenario: Scenario
model: str
upstream: Wire
sinks: SpanSinks
newrelic: ConnectSink
def cursors(self) -> Cursors:
self.upstream.drain()
return Cursors(
recorded_spans(self.sinks.operator)[0],
recorded_spans(self.sinks.tenant)[0],
recorded_spans(self.sinks.arize)[0],
recorded_spans(self.newrelic.control_url)[0],
)
def upstream_hits(self, marker: str) -> int:
return sum(1 for request in self.upstream.drain() if marker.encode() in request.body)
def langfuse(self, name: str, host: str | None = None) -> dict[str, str]:
return {
"langfuse_public_key": f"pk-lf-{name}",
"langfuse_secret_key": f"sk-lf-{name}",
"langfuse_host": host or self.sinks.tenant,
}
def attach(self, team: str, callback: str, variables: Mapping[str, str], kind: str | None = None) -> httpx.Response:
body: Final[dict[str, JsonValue]] = {"callback_name": callback, "callback_vars": dict(variables)}
return self.proxy.request(
"POST", f"/team/{team}/callback", body if kind is None else {**body, "callback_type": kind}
)
def team_key(
self,
capture: Capture | None,
name: str,
*,
callback: str = "langfuse_otel",
variables: Mapping[str, str] | None = None,
kind: str | None = None,
) -> tuple[str, str]:
team: Final = self.scenario.team()
base: Final = dict(variables) if variables is not None else self.langfuse(name)
attached: Final = self.attach(
team, callback, base if capture is None else {**base, "capture_message_content": capture}, kind
)
assert attached.status_code == 200, attached.text
return team, self.scenario.key(team_id=team)
def send(
self,
key: str,
endpoint: Endpoint = "chat",
stream: bool = False,
*,
marker: str | None = None,
extra: Mapping[str, JsonValue] | None = None,
) -> Sent:
tagged: Final = marker or _marker()
path, body = _body(self.model, endpoint, tagged, stream)
trace_id: Final = uuid.uuid4().hex
call_id: Final = str(uuid.uuid4())
headers: Final = {
"Authorization": f"Bearer {key}",
"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01",
"x-litellm-call-id": call_id,
}
with self.proxy.client.stream("POST", path, json={**body, **(extra or {})}, headers=headers) as response:
text: Final = response.read().decode()
return Sent(tagged, trace_id, call_id, response.status_code, text)
def _served(sent: Sent) -> Sent:
assert sent.status == 200, sent.text
assert sent.marker in sent.text, sent.text
return sent
Credential = tuple[str, str]
def _sent_with(url: str, credential: Credential) -> frozenset[str]:
"""Span ids ``url`` received in export requests whose ``credential`` header matched."""
header, value = credential
return frozenset(
span_id
for request in recorded_requests(url)
if {name.lower(): seen for name, seen in object_value(request.get("headers") or {}).items()}.get(header)
== value
for span_id in request.get("span_ids") or ()
if isinstance(span_id, str)
)
def _of_request(url: str, since: int, sent: Sent, credential: Credential | None = None) -> tuple[Span, ...]:
"""Spans of ``sent``'s trace at ``url``, narrowed to one account's export requests when ``credential`` is set."""
spans: Final = recorded_spans(url, since)[1]
traces: Final = {sent.trace_id} | {
span["trace_id"] for span in spans if span["attributes"].get(CALL_ID) == sent.call_id
}
owned: Final = None if credential is None else _sent_with(url, credential)
return tuple(span for span in spans if span["trace_id"] in traces and (owned is None or span["span_id"] in owned))
def _carries(span: Span, marker: str) -> bool:
return marker in json.dumps(span["attributes"])
def _model_spans(spans: tuple[Span, ...]) -> tuple[Span, ...]:
return tuple(span for span in spans if OPERATION in span["attributes"])
def _names(spans: tuple[Span, ...]) -> list[str]:
return sorted(span["name"] for span in spans)
def _with_content(url: str, since: int, sent: Sent) -> tuple[Span, ...]:
"""The copy of ``sent`` at ``url`` once its model span arrived and some span carries the marker."""
return eventually(
lambda: _of_request(url, since, sent),
lambda spans: bool(_model_spans(spans)) and any(_carries(span, sent.marker) for span in spans),
seconds=40,
)
def _arrived(
url: str, since: int, sent: Sent, expected: frozenset[str], credential: Credential | None = None
) -> tuple[Span, ...]:
"""The copy of ``sent`` at ``url`` once every span name in ``expected`` and a usage-bearing model span landed."""
spans: Final = eventually(
lambda: _of_request(url, since, sent, credential),
lambda found: (
expected <= {span["name"] for span in found}
and any(INPUT_TOKENS in span["attributes"] for span in _model_spans(found))
),
seconds=40,
)
assert all(
span["attributes"].get(INPUT_TOKENS) == 7 for span in _model_spans(spans) if INPUT_TOKENS in span["attributes"]
)
return spans
def _assert_redacted_twin(
rig: Rig, sent: Sent, cursors: Cursors, sink: str, since: int, credential: Credential | None = None
) -> tuple[Span, ...]:
operator: Final = _with_content(rig.sinks.operator, cursors.operator, sent)
carried: Final = frozenset(span["name"] for span in operator if _carries(span, sent.marker))
tenant: Final = _arrived(sink, since, sent, carried, credential)
leaked: Final = [span["name"] for span in tenant if _carries(span, sent.marker)]
assert leaked == [], f"{leaked} carried content; tenant saw {_names(tenant)}"
model: Final = _model_spans(tenant)[0]["attributes"]
assert "gen_ai.input.messages" not in model and "gen_ai.output.messages" not in model, sorted(model)
assert model.get("gen_ai.request.model"), sorted(model)
return tenant
def _assert_content_twin(rig: Rig, sent: Sent, cursors: Cursors) -> None:
operator: Final = _with_content(rig.sinks.operator, cursors.operator, sent)
carried: Final = frozenset(span["name"] for span in operator if _carries(span, sent.marker))
tenant: Final = _arrived(rig.sinks.tenant, cursors.tenant, sent, carried)
kept: Final = frozenset(span["name"] for span in tenant if _carries(span, sent.marker))
assert carried <= kept, f"content lost on {sorted(carried - kept)}"
@pytest.mark.parametrize(
("endpoint", "stream"),
[("chat", False), ("chat", True), ("responses", False), ("messages", True)],
ids=["chat-unary", "chat-stream", "responses-unary", "messages-stream"],
)
def test_no_content_team_keeps_its_trace_and_metadata_without_content(
rig: Rig, endpoint: Endpoint, stream: bool
) -> None:
_, key = rig.team_key("no_content", "a")
cursors: Final = rig.cursors()
sent: Final = _served(rig.send(key, endpoint, stream))
assert rig.upstream_hits(sent.marker) == 1
_assert_redacted_twin(rig, sent, cursors, rig.sinks.tenant, cursors.tenant)
@pytest.mark.parametrize("capture", ["span_only", None], ids=["span_only", "omitted"])
def test_span_only_and_omitted_teams_keep_the_globally_captured_content(rig: Rig, capture: Capture | None) -> None:
_, key = rig.team_key(capture, "b")
cursors: Final = rig.cursors()
sent: Final = _served(rig.send(key))
_assert_content_twin(rig, sent, cursors)
def test_tool_call_arguments_never_reach_a_no_content_team(rig: Rig) -> None:
_, key = rig.team_key("no_content", "a")
cursors: Final = rig.cursors()
tool: Final[dict[str, JsonValue]] = {
"type": "function",
"function": {"name": "lookup", "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}},
}
sent: Final = rig.send(key, marker=_marker("tool-"), extra={"tools": [tool]})
assert sent.status == 200 and sent.marker in sent.text, sent.text
operator: Final = _with_content(rig.sinks.operator, cursors.operator, sent)
assert any(sent.marker in str(span["attributes"].get("gen_ai.output.messages", "")) for span in operator), (
"operator copy lacks the tool-call arguments"
)
tenant: Final = _assert_redacted_twin(rig, sent, cursors, rig.sinks.tenant, cursors.tenant)
assert "lookup" in json.dumps([span["attributes"] for span in tenant]), "tool definitions are metadata and stay"
def test_cache_hit_twin_of_a_no_content_team_is_redacted(rig: Rig) -> None:
_, key = rig.team_key("no_content", "a")
marker: Final = _marker()
_served(rig.send(key, marker=marker))
assert rig.upstream_hits(marker) == 1
cursors: Final = rig.cursors()
hit: Final = eventually(
lambda: _served(rig.send(key, marker=marker)), lambda _: rig.upstream_hits(marker) == 0, seconds=20
)
operator: Final = eventually(
lambda: _of_request(rig.sinks.operator, cursors.operator, hit),
lambda spans: any(span["kind"] == SERVER for span in spans),
seconds=40,
)
expected: Final = frozenset(span["name"] for span in operator if span["kind"] == SERVER)
tenant: Final = eventually(
lambda: _of_request(rig.sinks.tenant, cursors.tenant, hit),
lambda spans: expected <= {span["name"] for span in spans},
seconds=40,
)
assert [span["name"] for span in tenant if _carries(span, marker)] == [], _names(tenant)
def test_a_failure_only_entry_rejects_the_setting(rig: Rig) -> None:
team: Final = rig.scenario.team()
response: Final = rig.attach(
team, "langfuse_otel", {**rig.langfuse("fail"), "capture_message_content": "no_content"}, "failure"
)
assert response.status_code == 400, response.text
assert "success_and_failure" in response.text, response.text
def test_key_level_logging_with_no_content_is_redacted(rig: Rig) -> None:
logging: Final[list[JsonValue]] = [
{
"callback_name": "langfuse_otel",
"callback_type": "success",
"callback_vars": {**rig.langfuse("key"), "capture_message_content": "no_content"},
}
]
key: Final = rig.scenario.key(metadata={"logging": logging})
cursors: Final = rig.cursors()
sent: Final = _served(rig.send(key))
_assert_redacted_twin(rig, sent, cursors, rig.sinks.tenant, cursors.tenant)
def test_same_account_destination_redacts_the_shared_copy_only_for_that_team(rig: Rig) -> None:
_, restricted = rig.team_key("no_content", "operator", variables=rig.langfuse("operator", host=rig.sinks.operator))
_, other = rig.team_key("span_only", "b")
cursors: Final = rig.cursors()
shared: Final = _served(rig.send(restricted))
assert rig.upstream_hits(shared.marker) == 1
copy: Final = eventually(
lambda: _of_request(rig.sinks.operator, cursors.operator, shared),
lambda spans: any(INPUT_TOKENS in span["attributes"] for span in _model_spans(spans)),
seconds=40,
)
assert [span["name"] for span in copy if _carries(span, shared.marker)] == [], _names(copy)
assert len([span for span in copy if span["kind"] == SERVER and span["name"].startswith("POST ")]) == 1, _names(
copy
)
_assert_content_twin(rig, _served(rig.send(other)), cursors)
def test_arize_destination_drops_openinference_content(rig: Rig) -> None:
_, key = rig.team_key(
"no_content", "arize", callback="arize", variables={"arize_space_key": "space-a", "arize_api_key": "arize-a"}
)
cursors: Final = rig.cursors()
sent: Final = _served(rig.send(key))
operator: Final = _with_content(rig.sinks.operator, cursors.operator, sent)
assert any(sent.marker in str(span["attributes"].get("input.value", "")) for span in operator), (
"operator copy has no OpenInference content to strip"
)
tenant: Final = _assert_redacted_twin(
rig, sent, cursors, rig.sinks.arize, cursors.arize, ("arize-space-id", "space-a")
)
model: Final = _model_spans(tenant)[0]["attributes"]
assert not any(name.startswith(("llm.input_messages.", "llm.output_messages.")) for name in model), sorted(model)
assert "input.value" not in model and "output.value" not in model, sorted(model)
assert model.get("llm.model_name"), sorted(model)
def test_weave_destination_drops_its_content(rig: Rig) -> None:
_, key = rig.team_key(
"no_content",
"weave",
callback="weave_otel",
variables={"wandb_api_key": "wandb-a", "weave_project_id": "team/a"},
)
cursors: Final = rig.cursors()
sent: Final = _served(rig.send(key))
operator: Final = eventually(
lambda: _of_request(rig.sinks.arize, cursors.arize, sent, ("project_id", "operator/weave")),
lambda spans: any(sent.marker in str(span["attributes"].get("weave.output", "")) for span in spans),
seconds=40,
)
assert operator, "operator weave copy has no weave.output to strip"
tenant: Final = _assert_redacted_twin(rig, sent, cursors, rig.sinks.arize, cursors.arize, ("project_id", "team/a"))
assert not any("weave.output" in span["attributes"] for span in tenant), _names(tenant)
def test_newrelic_destination_drops_genai_content(rig: Rig) -> None:
_, key = rig.team_key("no_content", "newrelic", callback="newrelic", variables={"newrelic_api_key": "nr-team-a"})
cursors: Final = rig.cursors()
sent: Final = _served(rig.send(key))
tenant: Final = _assert_redacted_twin(
rig, sent, cursors, rig.newrelic.control_url, cursors.newrelic, ("api-key", "nr-team-a")
)
hosts: Final = {
object_value(request.get("headers") or {}).get("Host")
for request in recorded_requests(rig.newrelic.control_url)
if object_value(request.get("headers") or {}).get("api-key") == "nr-team-a"
}
assert hosts == {"otlp.nr-data.net"}, hosts
assert tenant
def test_span_only_team_lifts_global_no_content_only_for_its_destination(dark_rig: Rig) -> None:
"""The sibling Arize entry gets a fan-out destination and the SigNoz entry rides the routed tracer."""
team, key = dark_rig.team_key("span_only", "b")
siblings: Final = (
dark_rig.attach(team, "arize", {"arize_space_key": "space-sibling", "arize_api_key": "arize-sibling"}),
dark_rig.attach(
team,
"signoz",
{"signoz_ingestion_endpoint": dark_rig.sinks.arize + "/v1/traces", "signoz_ingestion_key": "signoz-team"},
),
)
assert [response.status_code for response in siblings] == [200, 200], [response.text for response in siblings]
cursors: Final = dark_rig.cursors()
sent: Final = _served(dark_rig.send(key))
assert dark_rig.upstream_hits(sent.marker) == 1
operator: Final = _arrived(dark_rig.sinks.operator, cursors.operator, sent, frozenset())
tenant: Final = _with_content(dark_rig.sinks.tenant, cursors.tenant, sent)
sibling: Final = _arrived(
dark_rig.sinks.arize,
cursors.arize,
sent,
frozenset(),
("arize-space-id", "space-sibling"),
)
signoz: Final = _arrived(
dark_rig.sinks.arize,
cursors.arize,
sent,
frozenset(),
("signoz-ingestion-key", "signoz-team"),
)
assert any(_carries(span, sent.marker) for span in tenant), _names(tenant)
assert not any(_carries(span, sent.marker) for span in operator), _names(operator)
assert not any(_carries(span, sent.marker) for span in sibling), _names(sibling)
assert not any(_carries(span, sent.marker) for span in signoz), _names(signoz)
def test_an_unsupported_value_fails_registration(rig: Rig) -> None:
team: Final = rig.scenario.team()
response: Final = rig.attach(team, "langfuse_otel", {**rig.langfuse("a"), "capture_message_content": "bogus"})
assert response.status_code == 422, response.text
assert "Invalid capture_message_content 'bogus'" in response.text, response.text
def test_classic_langfuse_rejects_the_setting(rig: Rig) -> None:
team: Final = rig.scenario.team()
response: Final = rig.attach(team, "langfuse", {**rig.langfuse("a"), "capture_message_content": "no_content"})
assert response.status_code == 400, response.text
assert "capture_message_content" in response.text, response.text
def test_a_second_entry_for_the_same_backend_cannot_disagree(rig: Rig) -> None:
team, _ = rig.team_key("no_content", "a")
response: Final = rig.attach(team, "langfuse_otel", {"capture_message_content": "span_only"}, "success_and_failure")
assert response.status_code == 400, response.text
assert "already set to 'no_content' by another langfuse_otel entry" in response.text, response.text
def test_replacing_the_registration_switches_the_team_to_content(rig: Rig) -> None:
team, key = rig.team_key("no_content", "a")
cursors: Final = rig.cursors()
_assert_redacted_twin(rig, _served(rig.send(key)), cursors, rig.sinks.tenant, cursors.tenant)
removed: Final = rig.proxy.request("DELETE", f"/team/{team}/callback/langfuse_otel")
assert removed.status_code == 200, removed.text
replaced: Final = rig.attach(team, "langfuse_otel", {**rig.langfuse("a"), "capture_message_content": "span_only"})
assert replaced.status_code == 200, replaced.text
def attempt() -> bool:
since: Final = rig.cursors()
sent: Final = _served(rig.send(key))
operator: Final = _with_content(rig.sinks.operator, since.operator, sent)
carried: Final = frozenset(span["name"] for span in operator if _carries(span, sent.marker))
tenant: Final = _arrived(rig.sinks.tenant, since.tenant, sent, carried)
return any(_carries(span, sent.marker) for span in tenant)
assert eventually(attempt, bool, seconds=60)
def test_request_metadata_cannot_lift_the_team_restriction(rig: Rig) -> None:
_, key = rig.team_key("no_content", "a")
cursors: Final = rig.cursors()
sent: Final = _served(rig.send(key, extra={"metadata": {"capture_message_content": "span_only"}}))
assert rig.upstream_hits(sent.marker) == 1
_assert_redacted_twin(rig, sent, cursors, rig.sinks.tenant, cursors.tenant)
def test_stalled_tenant_during_a_mixed_burst_keeps_serving_and_each_teams_policy(rig: Rig) -> None:
_, restricted = rig.team_key("no_content", "a")
_, open_key = rig.team_key("span_only", "b")
shapes: Final[tuple[tuple[Endpoint, bool], ...]] = (
("chat", False),
("chat", True),
("responses", False),
("messages", True),
)
plan: Final = [(restricted if index % 2 else open_key, *shapes[index % len(shapes)]) for index in range(30)]
cursors: Final = rig.cursors()
configure_sink(rig.sinks.tenant, paused=True)
try:
with ThreadPoolExecutor(max_workers=10) as pool:
sent: Final = list(pool.map(lambda job: rig.send(job[0], job[1], job[2]), plan))
finally:
configure_sink(rig.sinks.tenant, paused=False)
assert [result.status for result in sent] == [200] * 30, [result.text for result in sent if result.status != 200]
for (key, _, _), result in zip(plan, sent, strict=True):
roots = [
span
for span in _with_content(rig.sinks.operator, cursors.operator, result)
if span["kind"] == SERVER and span["name"].startswith("POST ")
]
assert len(roots) == 1, roots
if key == restricted:
_assert_redacted_twin(rig, result, cursors, rig.sinks.tenant, cursors.tenant)
else:
_assert_content_twin(rig, result, cursors)
def _write_config(directory: Path, sinks: SpanSinks, name: str, callbacks: tuple[str, ...]) -> Path:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["litellm_settings"] = {
**config.get("litellm_settings", {}),
"callbacks": list(callbacks),
"otel_tenant_destination_mode": "additive",
"provider_url_destination_allowed_hosts": [
urlparse(sinks.tenant).netloc,
urlparse(sinks.operator).netloc,
urlparse(sinks.arize).netloc,
],
}
config["general_settings"] = {**config.get("general_settings", {}), "user_api_key_cache_ttl": 2}
path: Final = directory / f"{name}.yaml"
path.write_text(yaml.safe_dump(config))
return path
@contextmanager
def _started(
provider: Wire,
sinks: SpanSinks,
newrelic: ConnectSink,
directory: Path,
capture: str | None,
callbacks: tuple[str, ...] = ("langfuse_otel",),
) -> Generator[Rig]:
environment: Final = {
"LITELLM_OTEL_V2": "1",
"OTEL_BSP_SCHEDULE_DELAY": "300",
"MAPPER_NAMES": "genai,openinference",
"LANGFUSE_HOST": sinks.operator,
"LANGFUSE_PUBLIC_KEY": "pk-lf-operator",
"LANGFUSE_SECRET_KEY": "sk-lf-operator",
"ARIZE_SPACE_KEY": "space-operator",
"ARIZE_API_KEY": "arize-operator",
"ARIZE_HTTP_ENDPOINT": sinks.arize + "/v1/traces",
"WANDB_API_KEY": "wandb-operator",
"WANDB_PROJECT_ID": "operator/weave",
"WANDB_HOST": sinks.arize,
"SIGNOZ_INGESTION_ENDPOINT": sinks.arize + "/v1/traces",
"SIGNOZ_INGESTION_KEY": "signoz-operator",
"HTTPS_PROXY": newrelic.proxy_url,
"NO_PROXY": "127.0.0.1,localhost",
"REQUESTS_CA_BUNDLE": newrelic.ca_pem,
**({} if capture is None else {"OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT": capture}),
}
with (
gateway_from_environment() as gateway,
owned_proxy_process(
gateway,
directory,
environment,
config=_write_config(directory, sinks, "capture", callbacks),
remove_environment=(
"OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT",
"LITELLM_OTEL_TENANT_DESTINATION_MODE",
),
workers=2,
) as owned,
owned.gateway.scenario() as scenario,
):
yield Rig(
owned.gateway, owned, scenario, scenario.model(api_base=provider.url + "/v1"), provider, sinks, newrelic
)
@pytest.fixture(scope="module")
def provider() -> Iterator[Wire]:
with wire_server(_upstream) as wire:
yield wire
@pytest.fixture(scope="module")
def newrelic_sink(tmp_path_factory: pytest.TempPathFactory) -> Iterator[ConnectSink]:
with owned_connect_sink(tmp_path_factory.mktemp("newrelic-sink")) as sink:
yield sink
@pytest.fixture(scope="module")
def rig(
provider: Wire, audit_sinks: SpanSinks, newrelic_sink: ConnectSink, tmp_path_factory: pytest.TempPathFactory
) -> Iterator[Rig]:
directory: Final = tmp_path_factory.mktemp("capture-span-only")
with _started(provider, audit_sinks, newrelic_sink, directory, "span_only") as started:
yield started
@pytest.fixture(scope="module")
def dark_rig(
provider: Wire, audit_sinks: SpanSinks, newrelic_sink: ConnectSink, tmp_path_factory: pytest.TempPathFactory
) -> Iterator[Rig]:
with _started(
provider,
audit_sinks,
newrelic_sink,
tmp_path_factory.mktemp("capture-no-content"),
None,
("langfuse_otel", "signoz"),
) as started:
yield started

27
tests/integration/providers/test_decisions_openai_format_wire.py
View file

@ -105,19 +105,25 @@ class _Provider:
wraps_result: bool
cost_map_key: str | None
speaks_openai: bool = False
provider_reported_cost: float | None = None
def upstream_body(self) -> dict[str, JsonValue]:
if self.speaks_openai:
return {"model": self.body_model, "input": _INPUT, "questions": _QUESTIONS}
return {"model": self.body_model, "state": _INPUT, "questions": _SYSTEM_ONE_QUESTIONS}
def _with_reported_cost(self, usage: dict[str, JsonValue]) -> dict[str, JsonValue]:
if self.provider_reported_cost is None:
return usage
return {**usage, "cost": self.provider_reported_cost}
def upstream_reply(self) -> dict[str, JsonValue]:
if self.speaks_openai:
return {"model": self.body_model, "answers": _ANSWERS, "usage": _USAGE}
return {"model": self.body_model, "answers": _ANSWERS, "usage": self._with_reported_cost(_USAGE)}
answer: Final[dict[str, JsonValue]] = {
"model": self.body_model,
"answers": _SYSTEM_ONE_ANSWERS,
"usage": _SYSTEM_ONE_USAGE,
"usage": self._with_reported_cost(_SYSTEM_ONE_USAGE),
}
return {"result": answer, "success": True} if self.wraps_result else answer
@ -143,7 +149,8 @@ _PROVIDERS: Final = (
"typesafe/jev-1.13",
_API_KEY,
False,
"openrouter/typesafe/jev-1.13",
None,
provider_reported_cost=1.5834e-5,
),
_Provider(
"strands_decider", "strands_decider/systemone-decider", "/v1/systemone", "systemone-decider", None, False, None
@ -235,11 +242,13 @@ def _number(value: JsonValue) -> float:
return float(value)
def _expected_spend(cost_map_key: str | None) -> float:
if cost_map_key is None:
def _expected_spend(provider: _Provider) -> float:
if provider.provider_reported_cost is not None:
return provider.provider_reported_cost
if provider.cost_map_key is None:
return 0.0
cost_map: Final = _JSON_OBJECT.validate_json(Path("model_prices_and_context_window.json").read_bytes())
prices: Final = object_value(cost_map[cost_map_key])
prices: Final = object_value(cost_map[provider.cost_map_key])
return _INPUT_TOKENS * _number(prices["input_cost_per_token"]) + _OUTPUT_TOKENS * _number(
prices["output_cost_per_token"]
)
@ -320,10 +329,8 @@ def _drop_params_config(directory: Path, api_base: str) -> Path:
@pytest.mark.parametrize("provider", _PROVIDERS, ids=_provider_id)
def test_each_provider_gets_its_own_path_key_and_body_and_is_billed_from_the_cost_map(
gateway: Gateway, provider: _Provider
) -> None:
expected_spend: Final = _expected_spend(provider.cost_map_key)
def test_each_provider_gets_its_own_path_key_and_body_and_is_billed(gateway: Gateway, provider: _Provider) -> None:
expected_spend: Final = _expected_spend(provider)
with gateway.scenario() as scenario:
handle: Final = _register(scenario, provider.upstream_reply())
model: Final = _deployment(scenario, handle, provider)

1174
tests/unit/experimental_mcp_client/test_mcp_client.py
View file

File diff suppressed because it is too large Load diff

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