Merge pull request #26962 from BerriAI/litellm_internal_staging

[Infra] Promote Internal Staging to main
This commit is contained in:
yuneng-jiang
2026-04-30 21:59:55 -07:00
committed by GitHub
237 changed files with 48103 additions and 2209 deletions
@@ -0,0 +1,75 @@
name: Check Lazy OpenAPI Snapshot
on:
pull_request:
branches:
- main
- litellm_internal_staging
- "litellm_**"
permissions:
contents: read
checks: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
verify:
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
with:
version: "0.10.9"
- name: Cache uv dependencies
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-
- name: Install dependencies
run: uv sync --frozen --all-groups --all-extras
- name: Regenerate snapshot to /tmp
id: regen
run: |
cp litellm/proxy/_lazy_openapi_snapshot.json /tmp/snapshot.committed.json
uv run --no-sync python -m litellm.proxy._lazy_openapi_snapshot
mv litellm/proxy/_lazy_openapi_snapshot.json /tmp/snapshot.fresh.json
mv /tmp/snapshot.committed.json litellm/proxy/_lazy_openapi_snapshot.json
- name: Compare
id: diff
continue-on-error: true
run: |
diff -q /tmp/snapshot.fresh.json litellm/proxy/_lazy_openapi_snapshot.json
- name: Mark neutral if drift
if: steps.diff.outcome == 'failure'
uses: LouisBrunner/checks-action@6b626ffbad7cc56fd58627f774b9067e6118af23 # v2.0.0
with:
token: ${{ secrets.GITHUB_TOKEN }}
name: lazy-openapi-snapshot
conclusion: neutral
output: |
{
"title": "Lazy openapi snapshot is stale",
"summary": "Run `python -m litellm.proxy._lazy_openapi_snapshot` and commit the regenerated `litellm/proxy/_lazy_openapi_snapshot.json`. Not blocking — the snapshot will regenerate at release if not committed."
}
+4 -4
View File
@@ -4,7 +4,7 @@ on:
workflow_dispatch:
inputs:
tag:
description: "Release tag (e.g. v1.83.0-stable) — branch will be named release/<tag>"
description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted) — branch will be named release/<tag>"
required: true
type: string
commit_hash:
@@ -14,7 +14,7 @@ on:
workflow_call:
inputs:
tag:
description: "Release tag"
description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted)"
required: true
type: string
commit_hash:
@@ -40,8 +40,8 @@ jobs:
echo "::error::commit_hash must be a full 40-character commit SHA"
exit 1
fi
if ! echo "${TAG}" | grep -qE '^v[0-9]+\.[0-9]+\.[0-9]+'; then
echo "::error::tag must start with vX.Y.Z"
if ! echo "${TAG}" | grep -qE '^v?[0-9]+\.[0-9]+\.[0-9]+'; then
echo "::error::tag must start with X.Y.Z (optional leading v), e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, or v1.83.10-stable"
exit 1
fi
+9 -4
View File
@@ -4,7 +4,7 @@ on:
workflow_dispatch:
inputs:
tag:
description: "Release tag (e.g. v1.83.0-stable)"
description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted)"
required: true
type: string
commit_hash:
@@ -30,8 +30,8 @@ jobs:
echo "::error::commit_hash must be a full 40-character commit SHA"
exit 1
fi
if ! echo "${TAG}" | grep -qE '^v[0-9]+\.[0-9]+\.[0-9]+'; then
echo "::error::tag must start with vX.Y.Z"
if ! echo "${TAG}" | grep -qE '^v?[0-9]+\.[0-9]+\.[0-9]+'; then
echo "::error::tag must start with X.Y.Z (optional leading v), e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, or v1.83.10-stable"
exit 1
fi
@@ -45,6 +45,11 @@ jobs:
const tag = process.env.TAG;
const commitHash = process.env.COMMIT_HASH;
// Mark RC / dev / nightly / alpha / beta tags as GitHub pre-releases.
// PEP 440 post-releases (e.g. `1.84.0.post1`) and legacy `-stable[.patch.N]`
// are stable maintenance releases, not pre-releases.
const isPrerelease = /(?:rc|nightly|alpha|beta|\.dev)/i.test(tag);
const cosignSection = [
`## Verify Docker Image Signature`,
``,
@@ -89,7 +94,7 @@ jobs:
target_commitish: commitHash,
name: tag,
owner: context.repo.owner,
prerelease: false,
prerelease: isPrerelease,
repo: context.repo.repo,
tag_name: tag,
});
+1 -1
View File
@@ -2,4 +2,4 @@
# Packages needing lifecycle scripts: npm rebuild <pkg>
ignore-scripts=true
# Protects local npm install only — npm ci (used in CI) ignores this
min-release-age=3d
min-release-age=3
+1 -1
View File
@@ -2,4 +2,4 @@
# Packages needing lifecycle scripts: npm rebuild <pkg>
ignore-scripts=true
# Protects local npm install only — npm ci (used in CI) ignores this
min-release-age=3d
min-release-age=3
+1 -1
View File
@@ -2,4 +2,4 @@
# Packages needing lifecycle scripts: npm rebuild <pkg>
ignore-scripts=true
# Protects local npm install only — npm ci (used in CI) ignores this
min-release-age=3d
min-release-age=3
@@ -0,0 +1,2 @@
-- Search tool allowlists live on LiteLLM_ObjectPermissionTable (with agents, MCP, vector stores).
ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "search_tools" TEXT[] DEFAULT ARRAY[]::TEXT[];
@@ -0,0 +1,75 @@
-- CreateTable
CREATE TABLE "LiteLLM_WorkflowRun" (
"run_id" TEXT NOT NULL,
"session_id" TEXT NOT NULL,
"workflow_type" TEXT NOT NULL,
"status" TEXT NOT NULL DEFAULT 'pending',
"created_by" TEXT,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL,
"input" JSONB,
"output" JSONB,
"metadata" JSONB,
CONSTRAINT "LiteLLM_WorkflowRun_pkey" PRIMARY KEY ("run_id")
);
-- CreateTable
CREATE TABLE "LiteLLM_WorkflowEvent" (
"event_id" TEXT NOT NULL,
"run_id" TEXT NOT NULL,
"event_type" TEXT NOT NULL,
"step_name" TEXT NOT NULL,
"sequence_number" INTEGER NOT NULL,
"data" JSONB,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_WorkflowEvent_pkey" PRIMARY KEY ("event_id")
);
-- CreateTable
CREATE TABLE "LiteLLM_WorkflowMessage" (
"message_id" TEXT NOT NULL,
"run_id" TEXT NOT NULL,
"role" TEXT NOT NULL,
"content" TEXT NOT NULL,
"sequence_number" INTEGER NOT NULL,
"session_id" TEXT,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_WorkflowMessage_pkey" PRIMARY KEY ("message_id")
);
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_WorkflowRun_session_id_key" ON "LiteLLM_WorkflowRun"("session_id");
-- CreateIndex
CREATE INDEX "LiteLLM_WorkflowRun_workflow_type_status_idx" ON "LiteLLM_WorkflowRun"("workflow_type", "status");
-- CreateIndex
CREATE INDEX "LiteLLM_WorkflowRun_session_id_idx" ON "LiteLLM_WorkflowRun"("session_id");
-- CreateIndex
CREATE INDEX "LiteLLM_WorkflowRun_created_at_idx" ON "LiteLLM_WorkflowRun"("created_at");
-- CreateIndex
CREATE INDEX "LiteLLM_WorkflowRun_created_by_idx" ON "LiteLLM_WorkflowRun"("created_by");
-- CreateIndex
CREATE INDEX "LiteLLM_WorkflowEvent_run_id_idx" ON "LiteLLM_WorkflowEvent"("run_id");
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_WorkflowEvent_run_id_sequence_number_key" ON "LiteLLM_WorkflowEvent"("run_id", "sequence_number");
-- CreateIndex
CREATE INDEX "LiteLLM_WorkflowMessage_run_id_idx" ON "LiteLLM_WorkflowMessage"("run_id");
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_WorkflowMessage_run_id_sequence_number_key" ON "LiteLLM_WorkflowMessage"("run_id", "sequence_number");
-- AddForeignKey
ALTER TABLE "LiteLLM_WorkflowEvent" ADD CONSTRAINT "LiteLLM_WorkflowEvent_run_id_fkey" FOREIGN KEY ("run_id") REFERENCES "LiteLLM_WorkflowRun"("run_id") ON DELETE RESTRICT ON UPDATE CASCADE;
-- AddForeignKey
ALTER TABLE "LiteLLM_WorkflowMessage" ADD CONSTRAINT "LiteLLM_WorkflowMessage_run_id_fkey" FOREIGN KEY ("run_id") REFERENCES "LiteLLM_WorkflowRun"("run_id") ON DELETE RESTRICT ON UPDATE CASCADE;
@@ -277,6 +277,7 @@ model LiteLLM_ObjectPermissionTable {
models String[] @default([])
blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission
mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user
search_tools String[] @default([]) // search_tool_name values this key/team/user may call
teams LiteLLM_TeamTable[]
projects LiteLLM_ProjectTable[]
verification_tokens LiteLLM_VerificationToken[]
@@ -1290,3 +1291,80 @@ model LiteLLM_AdaptiveRouterSession {
@@id([session_id, router_name, model_name])
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
}
// ---------------------------------------------------------------------------
// Workflow Run Tracking
//
// Generic durable state tracking for any agent or automated workflow.
// Design: three tables — run (header + materialized status), event (append-only
// source of truth for state transitions), message (conversation inbox/outbox).
//
// Usage:
// - Set `workflow_type` to identify the owning system (e.g. "shin-builder").
// - Store domain-specific fields in `metadata` (worktree_path, pr_url, etc.).
// - `session_id` on WorkflowRun matches `x-litellm-session-id` header sent to
// the proxy — all spend logs for this run are automatically tagged.
// ---------------------------------------------------------------------------
// One instance of work being done. `status` is a materialized cache of the
// latest event; the event log is the authoritative source of truth.
model LiteLLM_WorkflowRun {
run_id String @id @default(uuid())
session_id String @unique @default(uuid())
workflow_type String
status String @default("pending")
created_by String? // user_id of the key that created this run; null = created by master key
created_at DateTime @default(now())
updated_at DateTime @updatedAt
input Json?
output Json?
metadata Json?
events LiteLLM_WorkflowEvent[]
messages LiteLLM_WorkflowMessage[]
@@index([workflow_type, status])
@@index([session_id])
@@index([created_at])
@@index([created_by])
}
// Append-only log of state transitions. Never mutate rows here.
// `step_name` and `event_type` are caller-defined strings — no hardcoded enums.
// Status auto-update rules (applied by the append endpoint):
// step.started → run.status = running
// step.failed → run.status = failed
// hook.waiting → run.status = paused
// hook.received → run.status = running
model LiteLLM_WorkflowEvent {
event_id String @id @default(uuid())
run_id String
event_type String
step_name String
sequence_number Int
data Json?
created_at DateTime @default(now())
run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id])
@@unique([run_id, sequence_number])
@@index([run_id])
}
// Conversation inbox/outbox — full message content, separate from the durable
// event log. Spend logs truncate messages; this table stores them in full.
// `session_id` here is the Claude --resume session ID (or similar).
model LiteLLM_WorkflowMessage {
message_id String @id @default(uuid())
run_id String
role String
content String
sequence_number Int
session_id String?
created_at DateTime @default(now())
run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id])
@@unique([run_id, sequence_number])
@@index([run_id])
}
+2 -2
View File
@@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.69"
version = "0.4.70"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.69"
version = "0.4.70"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",
+6 -4
View File
@@ -543,15 +543,17 @@ def _handle_retrieve_batch_providers_without_provider_config(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=(
"LiteLLM doesn't support custom_llm_provider={} for 'retrieve_batch' without a `model` kwarg. "
"Supported via this path: 'openai', 'azure', 'vertex_ai', 'anthropic'. "
"'bedrock' is supported but requires `model` to be passed so the provider config can be loaded."
).format(custom_llm_provider),
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
request=httpx.Request(method="retrieve_batch", url="https://github.com/BerriAI/litellm"), # type: ignore
),
)
return response
+5
View File
@@ -392,6 +392,7 @@ class DualCache(BaseCache):
value: float,
parent_otel_span: Optional[Span] = None,
local_only: bool = False,
refresh_ttl: bool = False,
**kwargs,
) -> Optional[float]:
"""
@@ -399,6 +400,9 @@ class DualCache(BaseCache):
Value - float - the value you want to increment by
Refresh_ttl - bool - if True, resets the Redis TTL on every write.
Default False preserves window-style semantics.
Returns - the incremented value, or None if no cache backend is
available (in_memory_cache is None and Redis failed/is absent).
"""
@@ -415,6 +419,7 @@ class DualCache(BaseCache):
value,
parent_otel_span=parent_otel_span,
ttl=kwargs.get("ttl", None),
refresh_ttl=refresh_ttl,
)
return result
+6 -4
View File
@@ -824,6 +824,7 @@ class RedisCache(BaseCache):
value: float,
ttl: Optional[int] = None,
parent_otel_span: Optional[Span] = None,
refresh_ttl: bool = False,
) -> float:
from redis.asyncio import Redis
@@ -834,11 +835,12 @@ class RedisCache(BaseCache):
try:
result = await _redis_client.incrbyfloat(name=key, amount=value)
if _used_ttl is not None:
# check if key already has ttl, if not -> set ttl
current_ttl = await _redis_client.ttl(key)
if current_ttl == -1:
# Key has no expiration
if refresh_ttl:
await _redis_client.expire(key, _used_ttl)
else:
current_ttl = await _redis_client.ttl(key)
if current_ttl == -1:
await _redis_client.expire(key, _used_ttl)
## LOGGING ##
end_time = time.time()
+15
View File
@@ -224,6 +224,16 @@ AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int(
)
AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120))
AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300))
# TCP keep-alive (SO_KEEPALIVE) — opt-in. Required when running behind NAT/LBs
# whose idle timeout is shorter than provider response timeouts (e.g. AWS NAT
# Gateway: 350s vs OpenAI/Azure: 600s). Without this, the kernel sends nothing
# during a long provider call and the NAT reaps the flow before the response
# arrives. Enabling SO_KEEPALIVE makes the kernel emit TCP probes that reset
# the NAT idle timer.
AIOHTTP_SO_KEEPALIVE = os.getenv("AIOHTTP_SO_KEEPALIVE", "False").lower() == "true"
AIOHTTP_TCP_KEEPIDLE = int(os.getenv("AIOHTTP_TCP_KEEPIDLE", 60))
AIOHTTP_TCP_KEEPINTVL = int(os.getenv("AIOHTTP_TCP_KEEPINTVL", 30))
AIOHTTP_TCP_KEEPCNT = int(os.getenv("AIOHTTP_TCP_KEEPCNT", 5))
# enable_cleanup_closed is only needed for Python versions with the SSL leak bug
# Fixed in Python 3.12.7+ and 3.13.1+ (see https://github.com/python/cpython/pull/118960)
# Reference: https://github.com/aio-libs/aiohttp/blob/master/aiohttp/connector.py#L74-L78
@@ -1383,6 +1393,10 @@ except (ValueError, TypeError):
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME = "litellm-internal-health-check"
LITTELM_CLI_SERVICE_ACCOUNT_NAME = "litellm-cli"
LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME = "litellm_internal_jobs"
# Stable identifier substituted in place of the master key on UserAPIKeyAuth
# objects so the master key (or its hash) never propagates to spend logs,
# Prometheus metrics, audit trails, or any other downstream consumer.
LITELLM_PROXY_MASTER_KEY_ALIAS = "litellm_proxy_master_key"
# Key Rotation Constants
LITELLM_KEY_ROTATION_ENABLED = os.getenv("LITELLM_KEY_ROTATION_ENABLED", "false")
@@ -1411,6 +1425,7 @@ LITELLM_PROXY_ADMIN_NAME = "default_user_id"
LITELLM_CLI_SOURCE_IDENTIFIER = "litellm-cli"
LITELLM_CLI_SESSION_TOKEN_PREFIX = "litellm-session-token"
CLI_SSO_SESSION_CACHE_KEY_PREFIX = "cli_sso_session"
CLI_SSO_SESSION_TTL_SECONDS = 600
CLI_JWT_TOKEN_NAME = "cli-jwt-token"
# Support both CLI_JWT_EXPIRATION_HOURS and LITELLM_CLI_JWT_EXPIRATION_HOURS for backwards compatibility
CLI_JWT_EXPIRATION_HOURS = int(
+7 -1
View File
@@ -1,6 +1,6 @@
import asyncio
from datetime import datetime
from typing import TYPE_CHECKING, Any, List, Optional
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy.pass_through_endpoints.success_handler import (
@@ -29,12 +29,14 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
litellm_logging_obj: LiteLLMLoggingObj,
request_body: dict,
model: str,
hidden_params: Optional[Dict[str, Any]] = None,
):
self.litellm_logging_obj = litellm_logging_obj
self.request_body = request_body
self.start_time = datetime.now()
self.collected_chunks: List[bytes] = []
self.model = model
self._hidden_params: Dict[str, Any] = hidden_params or {}
async def _handle_async_streaming_logging(
self,
@@ -76,11 +78,13 @@ class GoogleGenAIGenerateContentStreamingIterator(
litellm_metadata: dict,
custom_llm_provider: str,
request_body: Optional[dict] = None,
hidden_params: Optional[Dict[str, Any]] = None,
):
super().__init__(
litellm_logging_obj=logging_obj,
request_body=request_body or {},
model=model,
hidden_params=hidden_params,
)
self.response = response
self.model = model
@@ -130,11 +134,13 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(
litellm_metadata: dict,
custom_llm_provider: str,
request_body: Optional[dict] = None,
hidden_params: Optional[Dict[str, Any]] = None,
):
super().__init__(
litellm_logging_obj=logging_obj,
request_body=request_body or {},
model=model,
hidden_params=hidden_params,
)
self.response = response
self.model = model
@@ -87,9 +87,7 @@ class PromptManagementBase(ABC):
try:
messages = compiled_prompt_client["prompt_template"] + client_messages
except Exception as e:
raise ValueError(
f"Error compiling prompt: {e}. Prompt id={prompt_id}, prompt_variables={prompt_variables}, client_messages={client_messages}, dynamic_callback_params={dynamic_callback_params}"
)
raise ValueError(f"Error compiling prompt: {e}. Prompt id={prompt_id}")
compiled_prompt_client["completed_messages"] = messages
return compiled_prompt_client
@@ -116,9 +114,7 @@ class PromptManagementBase(ABC):
try:
messages = compiled_prompt_client["prompt_template"] + client_messages
except Exception as e:
raise ValueError(
f"Error compiling prompt: {e}. Prompt id={prompt_id}, prompt_variables={prompt_variables}, client_messages={client_messages}, dynamic_callback_params={dynamic_callback_params}"
)
raise ValueError(f"Error compiling prompt: {e}. Prompt id={prompt_id}")
compiled_prompt_client["completed_messages"] = messages
return compiled_prompt_client
@@ -1,4 +1,5 @@
from typing import Optional, Tuple
from urllib.parse import urlparse
import litellm
from litellm.constants import REPLICATE_MODEL_NAME_WITH_ID_LENGTH
@@ -8,6 +9,43 @@ from litellm.secret_managers.main import get_secret, get_secret_str
from ..types.router import LiteLLM_Params
def _endpoint_matches_api_base(endpoint: str, api_base: str) -> bool:
"""
Match a registered openai-compatible endpoint against a caller-supplied
``api_base`` using parsed-URL semantics, not unanchored substring search.
Both inputs may be a bare hostname (``api.perplexity.ai``), host+path
(``api.deepinfra.com/v1/openai``), or a full URL
(``https://api.cerebras.ai/v1``). Hostnames must match exactly
(case-insensitive); if the registered endpoint has a non-trivial path,
the api_base path must start with it on a segment boundary.
The naive ``endpoint in api_base`` shape lets a caller pass
``https://attacker.com/api.groq.com/openai/v1`` to coerce the proxy
into reading the server's GROQ_API_KEY from the environment and
forwarding it to the attacker's host as a Bearer credential.
"""
def _parse(value: str):
# Ensure urlparse sees a scheme so it populates hostname / path.
normalized = value if "://" in value else f"https://{value}"
return urlparse(normalized)
parsed_endpoint = _parse(endpoint)
parsed_url = _parse(api_base)
endpoint_host = (parsed_endpoint.hostname or "").lower()
url_host = (parsed_url.hostname or "").lower()
if not endpoint_host or endpoint_host != url_host:
return False
endpoint_path = parsed_endpoint.path.rstrip("/")
if not endpoint_path:
return True
url_path = parsed_url.path.rstrip("/")
return url_path == endpoint_path or url_path.startswith(endpoint_path + "/")
def _is_non_openai_azure_model(model: str) -> bool:
try:
model_name = model.split("/", 1)[1]
@@ -210,7 +248,7 @@ def get_llm_provider( # noqa: PLR0915
# check if api base is a known openai compatible endpoint
if api_base:
for endpoint in litellm.openai_compatible_endpoints:
if endpoint in api_base:
if _endpoint_matches_api_base(endpoint, api_base):
if endpoint == "api.perplexity.ai":
custom_llm_provider = "perplexity"
dynamic_api_key = get_secret_str("PERPLEXITYAI_API_KEY")
@@ -348,6 +386,7 @@ def get_llm_provider( # noqa: PLR0915
or "ft:gpt-3.5-turbo" in model
or "ft:gpt-4" in model # catches ft:gpt-4-0613, ft:gpt-4o
or model in litellm.openai_image_generation_models
or model.startswith("gpt-image")
or model in litellm.openai_video_generation_models
):
custom_llm_provider = "openai"
@@ -23,6 +23,13 @@ def _raise_env_reference_error(param: str, *, source: str) -> None:
)
def validate_no_callback_env_reference(
param: str, value: object, *, source: str
) -> None:
if _is_env_reference(value):
_raise_env_reference_error(param, source=source)
# Hardcoded list of supported callback params to avoid runtime inspection issues with TypedDict
_supported_callback_params = [
"langfuse_public_key",
@@ -66,8 +73,9 @@ def initialize_standard_callback_dynamic_params(
for param in _supported_callback_params:
if param in kwargs:
_param_value = kwargs.get(param)
if _is_env_reference(_param_value):
_raise_env_reference_error(param, source="request body")
validate_no_callback_env_reference(
param, _param_value, source="request body"
)
standard_callback_dynamic_params[param] = _param_value # type: ignore
# 2. Fallback: check "metadata" or "litellm_params" -> "metadata"
@@ -80,8 +88,9 @@ def initialize_standard_callback_dynamic_params(
for param in _supported_callback_params:
if param not in standard_callback_dynamic_params and param in metadata:
_param_value = metadata.get(param)
if _is_env_reference(_param_value):
_raise_env_reference_error(param, source="metadata")
validate_no_callback_env_reference(
param, _param_value, source="metadata"
)
standard_callback_dynamic_params[param] = _param_value # type: ignore
return standard_callback_dynamic_params
@@ -982,9 +982,9 @@ class CostCalculatorUtils:
image_response=completion_response,
)
elif custom_llm_provider == litellm.LlmProviders.OPENAI.value:
# Check if this is a gpt-image model (token-based pricing)
# gpt-image models use token-based pricing.
model_lower = model.lower()
if "gpt-image-1" in model_lower:
if "gpt-image" in model_lower:
from litellm.llms.openai.image_generation.cost_calculator import (
cost_calculator as openai_gpt_image_cost_calculator,
)
@@ -1004,9 +1004,9 @@ class CostCalculatorUtils:
optional_params=optional_params,
)
elif custom_llm_provider == litellm.LlmProviders.AZURE.value:
# Check if this is a gpt-image model (token-based pricing)
# gpt-image models use token-based pricing.
model_lower = model.lower()
if "gpt-image-1" in model_lower:
if "gpt-image" in model_lower:
from litellm.llms.openai.image_generation.cost_calculator import (
cost_calculator as openai_gpt_image_cost_calculator,
)
@@ -824,8 +824,6 @@ def convert_to_model_response_object( # noqa: PLR0915
stream=stream,
start_time=start_time,
end_time=end_time,
hidden_params=hidden_params,
_response_headers=_response_headers,
convert_tool_call_to_json_mode=convert_tool_call_to_json_mode,
)
raise Exception(
@@ -1661,6 +1661,20 @@ def _sanitize_anthropic_tool_use_id(tool_use_id: str) -> str:
return sanitized
_ANTHROPIC_DOCUMENT_BASE64_MEDIA_TYPES = {"application/pdf", "text/plain"}
def _is_anthropic_document_data_uri(url: str) -> bool:
# Anthropic's base64 document source accepts only application/pdf and
# text/plain (see select_anthropic_content_block_type_for_file). Routing
# other mimes here would produce a document block the API rejects, so we
# leave them on the image code path.
match = re.match(r"data:([^;,]+)", url)
if not match:
return False
return match.group(1) in _ANTHROPIC_DOCUMENT_BASE64_MEDIA_TYPES
def convert_to_anthropic_tool_result(
message: Union[ChatCompletionToolMessage, ChatCompletionFunctionMessage],
force_base64: bool = False,
@@ -1698,14 +1712,24 @@ def convert_to_anthropic_tool_result(
"""
anthropic_content: Union[
str,
List[Union[AnthropicMessagesToolResultContent, AnthropicMessagesImageParam]],
List[
Union[
AnthropicMessagesToolResultContent,
AnthropicMessagesImageParam,
AnthropicMessagesDocumentParam,
]
],
] = ""
if isinstance(message["content"], str):
anthropic_content = message["content"]
elif isinstance(message["content"], List):
content_list = message["content"]
anthropic_content_list: List[
Union[AnthropicMessagesToolResultContent, AnthropicMessagesImageParam]
Union[
AnthropicMessagesToolResultContent,
AnthropicMessagesImageParam,
AnthropicMessagesDocumentParam,
]
] = []
for content in content_list:
if content["type"] == "text":
@@ -1720,21 +1744,62 @@ def convert_to_anthropic_tool_result(
text_content["cache_control"] = cache_control_value
anthropic_content_list.append(text_content)
elif content["type"] == "image_url":
image_url_value = content["image_url"]
format = (
content["image_url"].get("format")
if isinstance(content["image_url"], dict)
image_url_value.get("format")
if isinstance(image_url_value, dict)
else None
)
_anthropic_image_param = create_anthropic_image_param(
content["image_url"], format=format, is_bedrock_invoke=force_base64
url_str = (
image_url_value.get("url")
if isinstance(image_url_value, dict)
else image_url_value
)
_anthropic_image_param = add_cache_control_to_content(
anthropic_content_element=_anthropic_image_param,
# Data URIs with non-image mime types (e.g. application/pdf) must
# translate to Anthropic document blocks, not image blocks —
# wrapping a PDF in `type: "image"` is rejected by the API.
if isinstance(url_str, str) and _is_anthropic_document_data_uri(
url_str
):
synth_file_message: ChatCompletionFileObject = {
"type": "file",
"file": {"file_data": url_str},
}
_document_block = anthropic_process_openai_file_message(
synth_file_message
)
_document_block = add_cache_control_to_content(
anthropic_content_element=cast(
AnthropicMessagesDocumentParam, _document_block
),
original_content_element=content,
)
anthropic_content_list.append(
cast(AnthropicMessagesDocumentParam, _document_block)
)
else:
_anthropic_image_param = create_anthropic_image_param(
image_url_value,
format=format,
is_bedrock_invoke=force_base64,
)
_anthropic_image_param = add_cache_control_to_content(
anthropic_content_element=_anthropic_image_param,
original_content_element=content,
)
anthropic_content_list.append(
cast(AnthropicMessagesImageParam, _anthropic_image_param)
)
elif content["type"] == "file":
file_content = cast(ChatCompletionFileObject, content)
_file_block = anthropic_process_openai_file_message(file_content)
_file_block = add_cache_control_to_content(
anthropic_content_element=cast(
AnthropicMessagesDocumentParam, _file_block
),
original_content_element=content,
)
anthropic_content_list.append(
cast(AnthropicMessagesImageParam, _anthropic_image_param)
)
anthropic_content_list.append(_file_block)
anthropic_content = anthropic_content_list
anthropic_tool_result: Optional[AnthropicMessagesToolResultParam] = None
@@ -3977,6 +4042,55 @@ def _convert_to_bedrock_tool_call_result(
tool_result_content_blocks.append(
BedrockToolResultContentBlock(image=_block["image"])
)
elif "document" in _block:
tool_result_content_blocks.append(
BedrockToolResultContentBlock(document=_block["document"])
)
else:
verbose_logger.warning(
"Bedrock Converse: unrecognized BedrockContentBlock keys "
"%s for image_url tool-result block %s; dropping.",
list(_block.keys()),
content,
)
elif content["type"] == "file":
# Match the user-message path (_process_file_message): accept
# either file_data (base64 data URI) or file_id (server-side
# reference / URL) and hand off to BedrockImageProcessor. Raise
# BadRequestError on both-None rather than silently dropping.
file_obj = content.get("file") or {}
file_data = file_obj.get("file_data")
file_id = file_obj.get("file_id")
if file_data is None and file_id is None:
raise litellm.BadRequestError(
message="file_data and file_id cannot both be None. Got={}".format(
content
),
model="",
llm_provider="bedrock",
)
file_format = file_obj.get("format")
_file_block: BedrockContentBlock = (
BedrockImageProcessor.process_image_sync(
image_url=cast(str, file_id or file_data),
format=file_format,
)
)
if "document" in _file_block:
tool_result_content_blocks.append(
BedrockToolResultContentBlock(document=_file_block["document"])
)
elif "image" in _file_block:
tool_result_content_blocks.append(
BedrockToolResultContentBlock(image=_file_block["image"])
)
else:
verbose_logger.warning(
"Bedrock Converse: unrecognized BedrockContentBlock keys "
"%s for file tool-result block %s; dropping.",
list(_file_block.keys()),
content,
)
message.get("name", "")
id = str(message.get("tool_call_id", str(uuid.uuid4())))
+59 -43
View File
@@ -60,6 +60,9 @@ def _redact_choice_content(choice):
def _redact_responses_api_output(output_items):
"""Helper to redact ResponsesAPIResponse output items."""
for output_item in output_items:
if hasattr(output_item, "text"):
output_item.text = "redacted-by-litellm"
if hasattr(output_item, "content") and isinstance(output_item.content, list):
for content_part in output_item.content:
if hasattr(content_part, "text"):
@@ -75,6 +78,28 @@ def _redact_responses_api_output(output_items):
summary_item.text = "redacted-by-litellm"
def _redact_responses_api_output_dict(output_items, redacted_str: str):
"""Helper to redact ResponsesAPIResponse output items in dict form."""
for output_item in output_items:
if not isinstance(output_item, dict):
continue
if "text" in output_item:
output_item["text"] = redacted_str
if isinstance(output_item.get("content"), list):
for content_item in output_item["content"]:
if isinstance(content_item, dict) and "text" in content_item:
content_item["text"] = redacted_str
if output_item.get("type") == "reasoning" and isinstance(
output_item.get("summary"), list
):
for summary_item in output_item["summary"]:
if isinstance(summary_item, dict) and "text" in summary_item:
summary_item["text"] = redacted_str
def _redact_standard_logging_object(model_call_details: dict):
"""Redact messages and response inside standard_logging_object if present."""
standard_logging_object = model_call_details.get("standard_logging_object")
@@ -93,28 +118,11 @@ def _redact_standard_logging_object(model_call_details: dict):
if isinstance(response, dict) and "output" in response:
# ResponsesAPIResponse format - redact content in output items
if isinstance(response.get("output"), list):
for output_item in response["output"]:
if isinstance(output_item, dict) and "content" in output_item:
if isinstance(output_item["content"], list):
for content_item in output_item["content"]:
if (
isinstance(content_item, dict)
and "text" in content_item
):
content_item["text"] = redacted_str
_redact_responses_api_output_dict(response["output"], redacted_str)
elif isinstance(response, dict) and "choices" in response:
# ModelResponse dict format - redact content in choices
if isinstance(response.get("choices"), list):
for choice in response["choices"]:
if isinstance(choice, dict):
if "message" in choice and isinstance(choice["message"], dict):
choice["message"]["content"] = redacted_str
if "audio" in choice["message"]:
choice["message"]["audio"] = None
elif "delta" in choice and isinstance(choice["delta"], dict):
choice["delta"]["content"] = redacted_str
if "audio" in choice["delta"]:
choice["delta"]["audio"] = None
_redact_model_response_dict_choices(response["choices"], redacted_str)
elif isinstance(response, str):
standard_logging_object["response"] = redacted_str
else:
@@ -122,6 +130,29 @@ def _redact_standard_logging_object(model_call_details: dict):
standard_logging_object["response"] = {"text": redacted_str}
def _redact_model_response_dict_choices(choices, redacted_str: str):
for choice in choices:
if isinstance(choice, dict):
if "message" in choice and isinstance(choice["message"], dict):
choice["message"]["content"] = redacted_str
if "reasoning_content" in choice["message"]:
choice["message"]["reasoning_content"] = redacted_str
if "thinking_blocks" in choice["message"]:
choice["message"]["thinking_blocks"] = None
if "audio" in choice["message"]:
choice["message"]["audio"] = None
elif "delta" in choice and isinstance(choice["delta"], dict):
choice["delta"]["content"] = redacted_str
if "reasoning_content" in choice["delta"]:
choice["delta"]["reasoning_content"] = redacted_str
if "thinking_blocks" in choice["delta"]:
choice["delta"]["thinking_blocks"] = None
if "audio" in choice["delta"]:
choice["delta"]["audio"] = None
else:
_redact_choice_content(choice)
def perform_redaction(model_call_details: dict, result):
"""
Performs the actual redaction on the logging object and result.
@@ -132,6 +163,7 @@ def perform_redaction(model_call_details: dict, result):
]
model_call_details["prompt"] = ""
model_call_details["input"] = ""
_redact_standard_logging_object(model_call_details)
# Redact streaming response
if (
@@ -171,30 +203,14 @@ def perform_redaction(model_call_details: dict, result):
elif isinstance(_result, dict) and "choices" in _result:
# Handle dict representation of ModelResponse (e.g., from model_dump())
if _result.get("choices") is not None:
for choice in _result["choices"]:
if isinstance(choice, dict):
if "message" in choice and isinstance(choice["message"], dict):
choice["message"]["content"] = "redacted-by-litellm"
if "reasoning_content" in choice["message"]:
choice["message"][
"reasoning_content"
] = "redacted-by-litellm"
if "thinking_blocks" in choice["message"]:
choice["message"]["thinking_blocks"] = None
if "audio" in choice["message"]:
choice["message"]["audio"] = None
elif "delta" in choice and isinstance(choice["delta"], dict):
choice["delta"]["content"] = "redacted-by-litellm"
if "reasoning_content" in choice["delta"]:
choice["delta"][
"reasoning_content"
] = "redacted-by-litellm"
if "thinking_blocks" in choice["delta"]:
choice["delta"]["thinking_blocks"] = None
if "audio" in choice["delta"]:
choice["delta"]["audio"] = None
else:
_redact_choice_content(choice)
_redact_model_response_dict_choices(
_result["choices"], "redacted-by-litellm"
)
elif isinstance(_result, dict) and "output" in _result:
if isinstance(_result.get("output"), list):
_redact_responses_api_output_dict(
_result["output"], "redacted-by-litellm"
)
elif isinstance(_result, litellm.ResponsesAPIResponse):
if hasattr(_result, "output"):
_redact_responses_api_output(_result.output)
@@ -21,6 +21,8 @@ class SensitiveDataMasker:
"auth",
"authorization",
"credential",
# Plural form: Vertex uses ``vertex_credentials``; segment-exact
# matching otherwise misses it because "credential" != "credentials".
"credentials",
"access",
"private",
+48 -22
View File
@@ -1553,25 +1553,43 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
)
data["output_config"] = output_config
def _transform_response_for_json_mode(
def _resolve_json_mode_non_streaming(
self,
json_mode: Optional[bool],
tool_calls: List[ChatCompletionToolCallChunk],
) -> Optional[LitellmMessage]:
_message: Optional[LitellmMessage] = None
if json_mode is True and len(tool_calls) == 1:
# check if tool name is the default tool name
json_mode_content_str: Optional[str] = None
if (
"name" in tool_calls[0]["function"]
and tool_calls[0]["function"]["name"] == RESPONSE_FORMAT_TOOL_NAME
):
json_mode_content_str = tool_calls[0]["function"].get("arguments")
if json_mode_content_str is not None:
_message = AnthropicConfig._convert_tool_response_to_message(
tool_calls=tool_calls,
)
return _message
) -> Tuple[
Optional[LitellmMessage],
List[ChatCompletionToolCallChunk],
Optional[str],
]:
"""Strip internal response_format tool calls; merge payload into content when mixed with user tools."""
if json_mode is not True or not tool_calls:
return None, tool_calls, None
json_indices = [
i
for i, t in enumerate(tool_calls)
if t.get("function", {}).get("name") == RESPONSE_FORMAT_TOOL_NAME
]
if not json_indices:
return None, tool_calls, None
if len(json_indices) == len(tool_calls):
json_tool = tool_calls[json_indices[0]]
if json_tool.get("function", {}).get("arguments") is None:
return None, tool_calls, None
_message = AnthropicConfig._convert_tool_response_to_message(
tool_calls=[json_tool]
)
return _message, [], None
first_json = tool_calls[json_indices[0]]
json_msg = AnthropicConfig._convert_tool_response_to_message([first_json])
extra_content: Optional[str] = (
json_msg.content if json_msg is not None else None
)
filtered_tools = [t for i, t in enumerate(tool_calls) if i not in json_indices]
return None, filtered_tools, extra_content
def extract_response_content(self, completion_response: dict) -> Tuple[
str,
@@ -1931,19 +1949,27 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
tool_calls,
)
json_mode_message, tool_calls_for_message, json_extra_content = (
self._resolve_json_mode_non_streaming(
json_mode=json_mode,
tool_calls=tool_calls,
)
)
merged_text = text_content or ""
if json_extra_content:
merged_text = (
merged_text + json_extra_content if merged_text else json_extra_content
)
_message = litellm.Message(
tool_calls=tool_calls,
content=text_content or None,
tool_calls=tool_calls_for_message,
content=merged_text or None,
provider_specific_fields=provider_specific_fields,
thinking_blocks=thinking_blocks,
reasoning_content=reasoning_content,
)
_message.provider_specific_fields = provider_specific_fields
json_mode_message = self._transform_response_for_json_mode(
json_mode=json_mode,
tool_calls=tool_calls,
)
if json_mode_message is not None:
completion_response["stop_reason"] = "stop"
_message = json_mode_message
@@ -24,6 +24,6 @@ def get_azure_image_generation_config(model: str) -> BaseImageGenerationConfig:
return AzureDallE3ImageGenerationConfig()
else:
verbose_logger.debug(
f"Using AzureGPTImageGenerationConfig for model: {model}. This follows the gpt-image-1 model format."
f"Using AzureGPTImageGenerationConfig for model: {model}. This follows the gpt-image model format."
)
return AzureGPTImageGenerationConfig()
@@ -3,7 +3,7 @@ from litellm.llms.openai.image_generation import GPTImageGenerationConfig
class AzureGPTImageGenerationConfig(GPTImageGenerationConfig):
"""
Azure gpt-image-1 image generation config
Azure gpt-image image generation config
"""
pass
@@ -95,6 +95,7 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict[str, Any]]:
"""
Transform search request for Azure AI Search API
@@ -33,6 +33,7 @@ class BaseRerankConfig(ABC):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
return {}
@@ -59,6 +59,7 @@ class BaseVectorStoreConfig:
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict]:
pass
@@ -70,6 +71,7 @@ class BaseVectorStoreConfig:
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict]:
"""
Optional async version of transform_search_vector_store_request.
@@ -84,6 +86,7 @@ class BaseVectorStoreConfig:
api_base=api_base,
litellm_logging_obj=litellm_logging_obj,
litellm_params=litellm_params,
extra_body=extra_body,
)
@abstractmethod
+30
View File
@@ -1,6 +1,7 @@
import hashlib
import json
import os
import re
import urllib.parse
from datetime import datetime
from typing import (
@@ -37,6 +38,11 @@ else:
AWSPreparedRequest = Any
# Real AWS region names are lowercase letters, digits, and hyphens
# (e.g. "us-east-1", "eu-west-2", "us-gov-west-1", "cn-north-1").
_VALID_AWS_REGION_PATTERN = re.compile(r"\A[a-z0-9-]+\Z")
class Boto3CredentialsInfo(BaseModel):
credentials: Credentials
aws_region_name: str
@@ -284,6 +290,9 @@ class BaseAWSLLM:
if not region: # Check if region is empty
return None
if not _VALID_AWS_REGION_PATTERN.match(region):
return None
return region
except Exception:
# Catch any unexpected errors and return None
@@ -481,6 +490,7 @@ class BaseAWSLLM:
str: The AWS region name
"""
aws_region_name = optional_params.get("aws_region_name", None)
self._validate_aws_region_name(aws_region_name)
### SET REGION NAME ###
if aws_region_name is None:
# check model arn #
@@ -519,8 +529,25 @@ class BaseAWSLLM:
except Exception:
aws_region_name = "us-west-2"
self._validate_aws_region_name(aws_region_name)
return aws_region_name
@staticmethod
def _validate_aws_region_name(aws_region_name: Optional[str]) -> None:
"""
Validate that an AWS region name conforms to the expected format
(lowercase alphanumerics and hyphens). Raises ValueError otherwise.
"""
if aws_region_name is None:
return
if not isinstance(aws_region_name, str) or not _VALID_AWS_REGION_PATTERN.match(
aws_region_name
):
raise ValueError(
f"Invalid AWS region format: {aws_region_name!r}. "
"Region names must contain only lowercase letters, digits, and hyphens."
)
def get_aws_region_name_for_non_llm_api_calls(
self,
aws_region_name: Optional[str] = None,
@@ -532,6 +559,7 @@ class BaseAWSLLM:
For non-llm api calls eg. Guardrails, Vector Stores we just need to check the dynamic param or env vars.
"""
self._validate_aws_region_name(aws_region_name)
if aws_region_name is None:
# check env #
litellm_aws_region_name = get_secret("AWS_REGION_NAME", None)
@@ -549,6 +577,8 @@ class BaseAWSLLM:
if aws_region_name is None:
aws_region_name = "us-west-2"
self._validate_aws_region_name(aws_region_name)
return aws_region_name
@staticmethod
@@ -1942,8 +1942,8 @@ class AmazonConverseConfig(BaseConfig):
completion_response = ConverseResponseBlock(**response.json()) # type: ignore
except Exception as e:
raise BedrockError(
message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(
response.text, str(e)
message="Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(
str(e)
),
status_code=422,
)
@@ -1,14 +1,16 @@
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
from copy import deepcopy
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
from urllib.parse import urlparse
import httpx
from litellm._logging import verbose_logger
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.types.integrations.rag.bedrock_knowledgebase import (
BedrockKBContent,
BedrockKBResponse,
BedrockKBRetrievalConfiguration,
BedrockKBResponse,
BedrockKBRetrievalQuery,
)
from litellm.types.router import GenericLiteLLMParams
@@ -202,6 +204,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict]:
if isinstance(query, list):
query = " ".join(query)
@@ -213,24 +216,46 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
}
retrieval_config: Dict[str, Any] = {}
if isinstance(extra_body, dict):
retrieval_config = deepcopy(
extra_body.get("retrievalConfiguration")
or extra_body.get("retrieval_configuration")
or {}
)
max_results = vector_store_search_optional_params.get("max_num_results")
if max_results is not None:
existing_number_of_results = retrieval_config.get(
"vectorSearchConfiguration", {}
).get("numberOfResults")
if (
existing_number_of_results is not None
and existing_number_of_results != max_results
):
verbose_logger.debug(
"Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.numberOfResults (%s) with max_num_results=%s",
existing_number_of_results,
max_results,
)
retrieval_config.setdefault("vectorSearchConfiguration", {})[
"numberOfResults"
] = max_results
filters = vector_store_search_optional_params.get("filters")
if filters is not None:
existing_filter = retrieval_config.get("vectorSearchConfiguration", {}).get(
"filter"
)
if existing_filter is not None and existing_filter != filters:
verbose_logger.debug(
"Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.filter with filters from vector_store_search_optional_params"
)
retrieval_config.setdefault("vectorSearchConfiguration", {})[
"filter"
] = filters
if retrieval_config:
# Create a properly typed retrieval configuration
typed_retrieval_config: BedrockKBRetrievalConfiguration = {}
if "vectorSearchConfiguration" in retrieval_config:
typed_retrieval_config["vectorSearchConfiguration"] = retrieval_config[
"vectorSearchConfiguration"
]
request_body["retrievalConfiguration"] = typed_retrieval_config
request_body["retrievalConfiguration"] = cast(
BedrockKBRetrievalConfiguration, retrieval_config
)
litellm_logging_obj.model_call_details["query"] = query
return url, request_body
@@ -111,6 +111,7 @@ class CohereRerankConfig(BaseRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
if "query" not in optional_rerank_params:
raise ValueError("query is required for Cohere rerank")
@@ -71,6 +71,7 @@ class CohereRerankV2Config(CohereRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
if "query" not in optional_rerank_params:
raise ValueError("query is required for Cohere rerank")
+62
View File
@@ -1,5 +1,7 @@
import asyncio
import inspect
import os
import socket
import ssl
import sys
import time
@@ -29,6 +31,10 @@ from litellm.constants import (
AIOHTTP_CONNECTOR_LIMIT_PER_HOST,
AIOHTTP_KEEPALIVE_TIMEOUT,
AIOHTTP_NEEDS_CLEANUP_CLOSED,
AIOHTTP_SO_KEEPALIVE,
AIOHTTP_TCP_KEEPCNT,
AIOHTTP_TCP_KEEPIDLE,
AIOHTTP_TCP_KEEPINTVL,
AIOHTTP_TTL_DNS_CACHE,
COMPLETION_HTTP_FALLBACK_SECONDS,
DEFAULT_SSL_CIPHERS,
@@ -54,6 +60,57 @@ except Exception:
version = "0.0.0"
# aiohttp 3.10+ exposes a `socket_factory` kwarg on TCPConnector. Older
# versions don't — detect once and skip the keep-alive wiring there.
# https://docs.aiohttp.org/en/stable/client_reference.html#aiohttp.TCPConnector
_AIOHTTP_SUPPORTS_SOCKET_FACTORY = (
"socket_factory" in inspect.signature(TCPConnector.__init__).parameters
)
def _build_aiohttp_keepalive_socket_factory() -> (
Optional[Callable[[Tuple[Any, ...]], socket.socket]]
):
"""
Build a socket_factory that enables SO_KEEPALIVE on aiohttp TCP sockets.
Why: by default, aiohttp creates sockets without SO_KEEPALIVE, so the kernel
sends nothing during a long idle TCP connection. NAT/LB hops (e.g. AWS NAT
Gateway, 350s idle timeout) reap the flow well before slow provider
responses (OpenAI/Azure: up to 600s) arrive. Enabling SO_KEEPALIVE makes
the kernel emit TCP probes that reset the NAT idle timer.
Returns None when AIOHTTP_SO_KEEPALIVE is disabled or aiohttp is too old.
"""
if not AIOHTTP_SO_KEEPALIVE or not _AIOHTTP_SUPPORTS_SOCKET_FACTORY:
return None
def factory(addr_info: Tuple[Any, ...]) -> socket.socket:
family, type_, proto = addr_info[0], addr_info[1], addr_info[2]
sock = socket.socket(family=family, type=type_, proto=proto)
sock.setblocking(False)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
# Linux: TCP_KEEPIDLE is idle-before-first-probe.
# macOS/Darwin: TCP_KEEPALIVE is the equivalent.
if hasattr(socket, "TCP_KEEPIDLE"):
sock.setsockopt(
socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, AIOHTTP_TCP_KEEPIDLE
)
elif hasattr(socket, "TCP_KEEPALIVE"):
sock.setsockopt(
socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, AIOHTTP_TCP_KEEPIDLE
)
if hasattr(socket, "TCP_KEEPINTVL"):
sock.setsockopt(
socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, AIOHTTP_TCP_KEEPINTVL
)
if hasattr(socket, "TCP_KEEPCNT"):
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, AIOHTTP_TCP_KEEPCNT)
return sock
return factory
def get_default_headers() -> dict:
"""
Get default headers for HTTP requests.
@@ -935,6 +992,11 @@ class AsyncHTTPHandler:
transport_connector_kwargs["limit_per_host"] = (
AIOHTTP_CONNECTOR_LIMIT_PER_HOST
)
# Returns None when SO_KEEPALIVE is disabled or aiohttp is too old to
# accept socket_factory — version detection lives inside the builder.
socket_factory = _build_aiohttp_keepalive_socket_factory()
if socket_factory is not None:
transport_connector_kwargs["socket_factory"] = socket_factory
return LiteLLMAiohttpTransport(
client=lambda: ClientSession(
+59 -6
View File
@@ -155,6 +155,30 @@ else:
LiteLLMLoggingObj = Any
def _google_genai_streaming_hidden_params(
*,
api_base: str,
litellm_params: GenericLiteLLMParams,
logging_obj: LiteLLMLoggingObj,
response_headers: httpx.Headers,
) -> Dict[str, Any]:
"""Pre-stream metadata for proxy response headers (mirrors CustomStreamWrapper._hidden_params)."""
from litellm.litellm_core_utils.core_helpers import process_response_headers
_model_info: Dict[str, Any] = dict(
getattr(litellm_params, "model_info", None) or {}
)
_raw_id = _model_info.get("id") or logging_obj.get_router_model_id() or ""
_model_id = _raw_id if isinstance(_raw_id, str) else str(_raw_id)
return {
"model_id": _model_id,
"api_base": api_base,
"cache_key": "",
"response_cost": "",
"additional_headers": process_response_headers(response_headers),
}
class BaseLLMHTTPHandler:
async def _make_common_async_call(
self,
@@ -983,6 +1007,7 @@ class BaseLLMHTTPHandler:
api_key: Optional[str] = None,
api_base: Optional[str] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
litellm_params: Optional[Dict[str, Any]] = None,
) -> RerankResponse:
# get config from model, custom llm provider
headers = provider_config.validate_environment(
@@ -1002,6 +1027,7 @@ class BaseLLMHTTPHandler:
model=model,
optional_rerank_params=optional_rerank_params,
headers=headers,
litellm_params=litellm_params,
)
## LOGGING
@@ -2511,10 +2537,16 @@ class BaseLLMHTTPHandler:
},
)
delete_kwargs: Dict[str, Any] = {
"url": url,
"headers": headers,
"timeout": timeout,
}
if data:
delete_kwargs["json"] = data
try:
response = await async_httpx_client.delete(
url=url, headers=headers, json=data, timeout=timeout
)
response = await async_httpx_client.delete(**delete_kwargs)
except Exception as e:
raise self._handle_error(
@@ -2595,10 +2627,16 @@ class BaseLLMHTTPHandler:
},
)
delete_kwargs: Dict[str, Any] = {
"url": url,
"headers": headers,
"timeout": timeout,
}
if data:
delete_kwargs["json"] = data
try:
response = sync_httpx_client.delete(
url=url, headers=headers, json=data, timeout=timeout
)
response = sync_httpx_client.delete(**delete_kwargs)
except Exception as e:
raise self._handle_error(
@@ -8585,6 +8623,7 @@ class BaseLLMHTTPHandler:
api_base=api_base,
litellm_logging_obj=logging_obj,
litellm_params=dict(litellm_params),
extra_body=extra_body,
)
else:
(
@@ -8597,6 +8636,7 @@ class BaseLLMHTTPHandler:
api_base=api_base,
litellm_logging_obj=logging_obj,
litellm_params=dict(litellm_params),
extra_body=extra_body,
)
all_optional_params: Dict[str, Any] = dict(litellm_params)
all_optional_params.update(vector_store_search_optional_params or {})
@@ -8697,6 +8737,7 @@ class BaseLLMHTTPHandler:
api_base=api_base,
litellm_logging_obj=logging_obj,
litellm_params=dict(litellm_params),
extra_body=extra_body,
)
all_optional_params: Dict[str, Any] = dict(litellm_params)
@@ -10425,6 +10466,12 @@ class BaseLLMHTTPHandler:
litellm_metadata=litellm_metadata or {},
custom_llm_provider=custom_llm_provider,
request_body=data,
hidden_params=_google_genai_streaming_hidden_params(
api_base=api_base,
litellm_params=litellm_params,
logging_obj=logging_obj,
response_headers=response.headers,
),
)
else:
response = sync_httpx_client.post(
@@ -10534,6 +10581,12 @@ class BaseLLMHTTPHandler:
litellm_metadata=litellm_metadata or {},
custom_llm_provider=custom_llm_provider,
request_body=data,
hidden_params=_google_genai_streaming_hidden_params(
api_base=api_base,
litellm_params=litellm_params,
logging_obj=logging_obj,
response_headers=response.headers,
),
)
else:
response = await async_httpx_client.post(
@@ -132,6 +132,7 @@ class DeepinfraRerankConfig(BaseRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
# Convert OptionalRerankParams to dict as expected by parent class
if optional_rerank_params is None:
@@ -127,6 +127,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
"""
Transform request to Fireworks AI rerank format
@@ -118,6 +118,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict]:
"""
Transform search request to Gemini's generateContent format.
@@ -121,6 +121,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
if "query" not in optional_rerank_params:
raise ValueError("query is required for Hosted VLLM rerank")
@@ -146,6 +146,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig):
model: str,
optional_rerank_params: Union[OptionalRerankParams, dict],
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
if "query" not in optional_rerank_params:
raise ValueError("query is required for HuggingFace rerank")
@@ -74,7 +74,11 @@ class JinaAIRerankConfig(BaseRerankConfig):
return cleaned_base
def transform_rerank_request(
self, model: str, optional_rerank_params: Dict, headers: Dict
self,
model: str,
optional_rerank_params: Dict,
headers: Dict,
litellm_params: Optional[dict] = None,
) -> Dict:
return {"model": model, **optional_rerank_params}
@@ -25,7 +25,6 @@ else:
LiteLLMLoggingObj = Any
MILVUS_OPTIONAL_PARAMS = {
"dbName",
"annsField",
"limit",
"filter",
@@ -33,7 +32,6 @@ MILVUS_OPTIONAL_PARAMS = {
"groupingField",
"outputFields",
"searchParams",
"partitionNames",
"consistencyLevel",
}
@@ -130,6 +128,7 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict[str, Any]]:
"""
Transform search request for Azure AI Search API
@@ -172,13 +171,21 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig):
url = f"{api_base}/v2/vectordb/entities/search"
# Build the request body for Azure AI Search with vector search
request_body = {
request_body: Dict[str, Any] = {
"collectionName": index_name,
"data": [query_vector],
"annsField": "book_intro_vector",
**vector_store_search_optional_params,
}
db_name = litellm_params.get("milvus_db_name")
if db_name:
request_body["dbName"] = db_name
partition_names = litellm_params.get("milvus_partition_names")
if partition_names:
request_body["partitionNames"] = partition_names
#########################################################
# Update logging object with details of the request
#########################################################
@@ -66,6 +66,7 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
"""
Transform request, using clean model name without 'ranking/' prefix.
@@ -75,4 +76,5 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig):
model=clean_model,
optional_rerank_params=optional_rerank_params,
headers=headers,
litellm_params=litellm_params,
)
@@ -177,6 +177,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
"""
Transform request to Nvidia NIM format.
@@ -1,5 +1,5 @@
"""
Cost calculator for OpenAI image generation models (gpt-image-1, gpt-image-1-mini)
Cost calculator for OpenAI image generation models (gpt-image family)
These models use token-based pricing instead of pixel-based pricing like DALL-E.
"""
@@ -17,13 +17,13 @@ def cost_calculator(
custom_llm_provider: Optional[str] = None,
) -> float:
"""
Calculate cost for OpenAI gpt-image-1 and gpt-image-1-mini models.
Calculate cost for OpenAI gpt-image models.
Uses the same usage format as Responses API, so we reuse the helper
to transform to chat completion format and use generic_cost_per_token.
Args:
model: The model name (e.g., "gpt-image-1", "gpt-image-1-mini")
model: The model name (e.g., "gpt-image-1", "gpt-image-2")
image_response: The ImageResponse containing usage data
custom_llm_provider: Optional provider name
@@ -15,7 +15,7 @@ if TYPE_CHECKING:
class GPTImageGenerationConfig(BaseImageGenerationConfig):
"""
OpenAI gpt-image-1 image generation config
OpenAI gpt-image image generation config
"""
def get_supported_openai_params(
@@ -106,6 +106,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict]:
url = f"{api_base}/{vector_store_id}/search"
typed_request_body = VectorStoreSearchRequest(
+5
View File
@@ -101,5 +101,10 @@
"param_mappings": {
"max_completion_tokens": "max_tokens"
}
},
"aihubmix": {
"base_url": "https://aihubmix.com/v1",
"api_key_env": "AIHUBMIX_API_KEY",
"api_base_env": "AIHUBMIX_API_BASE"
}
}
@@ -80,6 +80,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict]:
url = f"{api_base}/{vector_store_id}/search"
_, request_body = super().transform_search_vector_store_request(
@@ -89,5 +90,6 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig):
api_base=api_base,
litellm_logging_obj=litellm_logging_obj,
litellm_params=litellm_params,
extra_body=extra_body,
)
return url, request_body
@@ -102,6 +102,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict]:
"""RAGFlow vector stores are management-only, search is not supported."""
raise NotImplementedError(
@@ -79,6 +79,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict]:
"""Sync version - generates embedding synchronously."""
# For S3 Vectors, vector_store_id should be in format: bucket_name:index_name
@@ -140,6 +141,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict]:
"""Async version - generates embedding asynchronously."""
# For S3 Vectors, vector_store_id should be in format: bucket_name:index_name
+16 -4
View File
@@ -4,6 +4,7 @@ from typing import Any, Coroutine, Dict, Optional, Union
import httpx
import litellm
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
get_async_httpx_client,
@@ -224,8 +225,14 @@ class VertexAIBatchPrediction(VertexLLM):
},
)
response = sync_handler.get(
url=api_base,
# ``api_base`` here can come from caller-supplied request kwargs
# (clientside override). Wrap the fetch in ``safe_get`` so DNS
# rebind / private / cloud-metadata targets are rejected; the
# proxy auth gate already blocks malicious clientside ``api_base``
# at the boundary — this is defense-in-depth for SDK callers.
response = safe_get(
sync_handler,
api_base,
headers=headers,
)
@@ -270,8 +277,13 @@ class VertexAIBatchPrediction(VertexLLM):
},
)
response = await client.get(
url=api_base,
# Mirror the sync path: ``api_base`` may come from caller-supplied
# request kwargs, so wrap the fetch in ``async_safe_get`` to reject
# DNS-rebind / private / cloud-metadata targets. Defense-in-depth
# behind the proxy auth gate's clientside ``api_base`` check.
response = await async_safe_get(
client,
api_base,
headers=headers,
)
if response.status_code != 200:
+47
View File
@@ -27,6 +27,53 @@ class VertexAIError(BaseLLMException):
super().__init__(message=message, status_code=status_code, headers=headers)
def vertex_request_labels_from_litellm_params(
litellm_params: Optional[dict],
) -> Optional[Dict[str, str]]:
"""
Build Vertex/GCP billing labels from LiteLLM user metadata on ``litellm_params``:
``metadata`` (``completion(..., metadata=...)``) or ``litellm_metadata``,
using ``requester_metadata`` string key-value pairs (same convention as Gemini).
``metadata`` is tried first when both are present.
"""
if not litellm_params:
return None
for key in ("metadata", "litellm_metadata"):
if key not in litellm_params:
continue
metadata = litellm_params[key]
if metadata is None or not isinstance(metadata, dict):
continue
if "requester_metadata" not in metadata:
continue
rm = metadata["requester_metadata"]
if not isinstance(rm, dict):
continue
labels = {k: v for k, v in rm.items() if isinstance(v, str)}
if labels:
return labels
return None
def pop_vertex_request_labels(
optional_params: Optional[dict],
litellm_params: Optional[dict],
) -> Optional[Dict[str, str]]:
"""
Resolve labels from optional ``labels`` (Gemini-style) and/or
``litellm_params["metadata"]`` / ``litellm_params["litellm_metadata"]``
(``requester_metadata``). Pops ``labels`` from optional_params when present.
"""
labels: Optional[Dict[str, str]] = None
if optional_params is not None and "labels" in optional_params:
raw = optional_params.pop("labels")
if isinstance(raw, dict):
labels = {k: v for k, v in raw.items() if isinstance(v, str)}
if not labels:
labels = vertex_request_labels_from_litellm_params(litellm_params)
return labels if labels else None
class VertexAIModelRoute(str, Enum):
"""Enum for Vertex AI model routing"""
@@ -24,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
response_schema_prompt,
)
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
from litellm.types.files import (
get_file_mime_type_for_file_type,
get_file_type_from_extension,
@@ -714,16 +715,8 @@ def _transform_request_body( # noqa: PLR0915
optional_params.pop("output_config", None)
config_fields = GenerationConfig.__annotations__.keys()
# If the LiteLLM client sends Gemini-supported parameter "labels", add it
# as "labels" field to the request sent to the Gemini backend.
labels: Optional[dict[str, str]] = optional_params.pop("labels", None)
# If the LiteLLM client sends OpenAI-supported parameter "metadata", add it
# as "labels" field to the request sent to the Gemini backend.
if labels is None and "metadata" in litellm_params:
metadata = litellm_params["metadata"]
if metadata is not None and "requester_metadata" in metadata:
rm = metadata["requester_metadata"]
labels = {k: v for k, v in rm.items() if isinstance(v, str)}
# labels: optional explicit param and/or metadata.requester_metadata (OpenAI metadata)
labels = pop_vertex_request_labels(optional_params, litellm_params)
filtered_params = {
k: v
@@ -2395,8 +2395,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
completion_response = GenerateContentResponseBody(**raw_response.json()) # type: ignore
except Exception as e:
raise VertexAIError(
message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(
raw_response.text, str(e)
message="Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(
str(e)
),
status_code=422,
headers=raw_response.headers,
@@ -2530,8 +2530,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
except Exception as e:
raise VertexAIError(
message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(
completion_response, str(e)
message="Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(
str(e)
),
status_code=422,
headers=raw_response.headers,
@@ -7,7 +7,10 @@ import litellm
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
from litellm.llms.vertex_ai.common_utils import (
get_vertex_base_url,
pop_vertex_request_labels,
)
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
@@ -203,13 +206,16 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
"sampleCount": 1,
}
# Merge with optional params
labels = pop_vertex_request_labels(optional_params, litellm_params)
# Merge with optional params (after popping labels so they are not sent as Imagen parameters)
parameters = {**default_params, **optional_params}
request_body = {
request_body: dict = {
"instances": [{"prompt": prompt}],
"parameters": parameters,
}
if labels:
request_body["labels"] = labels
return request_body
@@ -11,12 +11,15 @@ import httpx
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.llms.vertex_ai.common_utils import (
vertex_request_labels_from_litellm_params,
)
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.secret_managers.main import get_secret_str
from litellm.types.rerank import (
RerankBilledUnits,
RerankResponse,
RerankResponseMeta,
RerankBilledUnits,
RerankResponseResult,
)
@@ -109,6 +112,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
"""
Transform the request from Cohere format to Vertex AI Discovery Engine format
@@ -145,6 +149,10 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
# When return_documents is False, we want to ignore record details (return only IDs)
request_data["ignoreRecordDetailsInResponse"] = not return_documents
user_labels = vertex_request_labels_from_litellm_params(litellm_params)
if user_labels:
request_data["userLabels"] = user_labels
return request_data
def transform_rerank_response(
@@ -100,6 +100,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict[str, Any]]:
"""
Transform search request for Vertex AI RAG API
@@ -107,6 +107,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict[str, Any]]:
"""
Transform search request for Vertex AI RAG API
@@ -1,4 +1,4 @@
from typing import Literal, Optional, Union
from typing import Dict, Literal, Optional, Union
import httpx
@@ -44,6 +44,7 @@ class VertexEmbedding(VertexBase):
vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES] = None,
gemini_api_key: Optional[str] = None,
extra_headers: Optional[dict] = None,
litellm_params: Optional[Dict] = None,
) -> EmbeddingResponse:
if aembedding is True:
return self.async_embedding( # type: ignore
@@ -61,6 +62,7 @@ class VertexEmbedding(VertexBase):
vertex_credentials=vertex_credentials,
gemini_api_key=gemini_api_key,
extra_headers=extra_headers,
litellm_params=litellm_params,
)
should_use_v1beta1_features = self.is_using_v1beta1_features(
@@ -92,7 +94,10 @@ class VertexEmbedding(VertexBase):
headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers)
vertex_request: VertexEmbeddingRequest = (
litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request(
input=input, optional_params=optional_params, model=model
input=input,
optional_params=optional_params,
model=model,
litellm_params=litellm_params,
)
)
@@ -156,6 +161,7 @@ class VertexEmbedding(VertexBase):
gemini_api_key: Optional[str] = None,
extra_headers: Optional[dict] = None,
encoding=None,
litellm_params: Optional[Dict] = None,
) -> EmbeddingResponse:
"""
Async embedding implementation
@@ -188,7 +194,10 @@ class VertexEmbedding(VertexBase):
headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers)
vertex_request: VertexEmbeddingRequest = (
litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request(
input=input, optional_params=optional_params, model=model
input=input,
optional_params=optional_params,
model=model,
litellm_params=litellm_params,
)
)
@@ -3,6 +3,7 @@ from typing import List, Literal, Optional, Union
from pydantic import BaseModel
from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
from litellm.types.utils import EmbeddingResponse, Usage
from .types import *
@@ -100,7 +101,11 @@ class VertexAITextEmbeddingConfig(BaseModel):
return optional_params
def transform_openai_request_to_vertex_embedding_request(
self, input: Union[list, str], optional_params: dict, model: str
self,
input: Union[list, str],
optional_params: dict,
model: str,
litellm_params: Optional[dict] = None,
) -> VertexEmbeddingRequest:
"""
Transforms an openai request to a vertex embedding request.
@@ -108,16 +113,26 @@ class VertexAITextEmbeddingConfig(BaseModel):
# Import here to avoid circular import issues with litellm.__init__
from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig
labels = pop_vertex_request_labels(optional_params, litellm_params)
if model.isdigit():
return self._transform_openai_request_to_fine_tuned_embedding_request(
input, optional_params, model
vertex_request = (
self._transform_openai_request_to_fine_tuned_embedding_request(
input, optional_params, model
)
)
if labels:
vertex_request["labels"] = labels
return vertex_request
if VertexBGEConfig.is_bge_model(model):
return VertexBGEConfig.transform_request(
vertex_request = VertexBGEConfig.transform_request(
input=input, optional_params=optional_params, model=model
)
if labels:
vertex_request["labels"] = labels
return vertex_request
vertex_request: VertexEmbeddingRequest = VertexEmbeddingRequest()
vertex_request = VertexEmbeddingRequest()
vertex_text_embedding_input_list: List[TextEmbeddingInput] = []
task_type: Optional[TaskType] = optional_params.get("task_type")
title = optional_params.get("title")
@@ -133,6 +148,8 @@ class VertexAITextEmbeddingConfig(BaseModel):
vertex_request["instances"] = vertex_text_embedding_input_list
vertex_request["parameters"] = EmbeddingParameters(**optional_params)
if labels:
vertex_request["labels"] = labels
return vertex_request
@@ -3,7 +3,7 @@ Types for Vertex Embeddings Requests
"""
from enum import Enum
from typing import List, Optional, Union
from typing import Dict, List, Optional, Union
from typing_extensions import TypedDict
@@ -56,6 +56,7 @@ class VertexEmbeddingRequest(TypedDict, total=False):
List[TextEmbeddingFineTunedInput],
]
parameters: Optional[Union[EmbeddingParameters, TextEmbeddingFineTunedParameters]]
labels: Optional[Dict[str, str]]
# Example usage:
+5 -1
View File
@@ -67,7 +67,11 @@ class VoyageRerankConfig(BaseRerankConfig):
return api_base
def transform_rerank_request(
self, model: str, optional_rerank_params: Dict, headers: Dict
self,
model: str,
optional_rerank_params: Dict,
headers: Dict,
litellm_params: Optional[dict] = None,
) -> Dict:
return {"model": model, **optional_rerank_params}
@@ -143,6 +143,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
"""
Transform request to IBM watsonx.ai rerank format
+1
View File
@@ -5311,6 +5311,7 @@ def embedding( # noqa: PLR0915
api_key=api_key,
api_base=api_base,
client=client,
litellm_params=litellm_params_dict,
)
elif custom_llm_provider == "oobabooga":
response = oobabooga.embedding(
@@ -712,6 +712,7 @@
},
"anthropic.claude-haiku-4-5-20251001-v1:0": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_1hr": 2e-06,
"cache_read_input_token_cost": 1e-07,
"input_cost_per_token": 1e-06,
"litellm_provider": "bedrock_converse",
@@ -735,6 +736,7 @@
},
"anthropic.claude-haiku-4-5@20251001": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_1hr": 2e-06,
"cache_read_input_token_cost": 1e-07,
"input_cost_per_token": 1e-06,
"litellm_provider": "bedrock_converse",
@@ -955,6 +957,7 @@
},
"anthropic.claude-opus-4-5-20251101-v1:0": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "bedrock_converse",
@@ -982,6 +985,7 @@
},
"anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "bedrock_converse",
@@ -1011,6 +1015,7 @@
},
"global.anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "bedrock_converse",
@@ -1040,6 +1045,7 @@
},
"us.anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 5.5e-06,
"litellm_provider": "bedrock_converse",
@@ -1127,6 +1133,7 @@
},
"anthropic.claude-opus-4-7": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "bedrock_converse",
@@ -1157,6 +1164,7 @@
},
"global.anthropic.claude-opus-4-7": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "bedrock_converse",
@@ -1187,6 +1195,7 @@
},
"us.anthropic.claude-opus-4-7": {
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 5.5e-06,
"litellm_provider": "bedrock_converse",
@@ -1277,6 +1286,7 @@
},
"anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_creation_input_token_cost_above_1hr": 6e-06,
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "bedrock_converse",
@@ -1305,6 +1315,7 @@
},
"global.anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_creation_input_token_cost_above_1hr": 6e-06,
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "bedrock_converse",
@@ -1333,6 +1344,7 @@
},
"us.anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 4.125e-06,
"cache_creation_input_token_cost_above_1hr": 6.6e-06,
"cache_read_input_token_cost": 3.3e-07,
"input_cost_per_token": 3.3e-06,
"litellm_provider": "bedrock_converse",
@@ -1447,11 +1459,13 @@
},
"anthropic.claude-sonnet-4-5-20250929-v1:0": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_creation_input_token_cost_above_1hr": 6e-06,
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 3e-06,
"input_cost_per_token_above_200k_tokens": 6e-06,
"output_cost_per_token_above_200k_tokens": 2.25e-05,
"cache_creation_input_token_cost_above_200k_tokens": 7.5e-06,
"cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05,
"cache_read_input_token_cost_above_200k_tokens": 6e-07,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 200000,
@@ -5103,6 +5117,38 @@
"/v1/images/edits"
]
},
"azure/gpt-image-2": {
"cache_read_input_image_token_cost": 2e-06,
"cache_read_input_token_cost": 1.25e-06,
"input_cost_per_token": 5e-06,
"input_cost_per_image_token": 8e-06,
"litellm_provider": "azure",
"mode": "image_generation",
"output_cost_per_token": 1e-05,
"output_cost_per_image_token": 3e-05,
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
],
"supports_vision": true,
"supports_pdf_input": true
},
"azure/gpt-image-2-2026-04-21": {
"cache_read_input_image_token_cost": 2e-06,
"cache_read_input_token_cost": 1.25e-06,
"input_cost_per_token": 5e-06,
"input_cost_per_image_token": 8e-06,
"litellm_provider": "azure",
"mode": "image_generation",
"output_cost_per_token": 1e-05,
"output_cost_per_image_token": 3e-05,
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
],
"supports_vision": true,
"supports_pdf_input": true
},
"azure/low/1024-x-1024/gpt-image-1-mini": {
"input_cost_per_pixel": 2.0751953125e-09,
"litellm_provider": "azure",
@@ -17889,11 +17935,13 @@
},
"global.anthropic.claude-sonnet-4-5-20250929-v1:0": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_creation_input_token_cost_above_1hr": 6e-06,
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 3e-06,
"input_cost_per_token_above_200k_tokens": 6e-06,
"output_cost_per_token_above_200k_tokens": 2.25e-05,
"cache_creation_input_token_cost_above_200k_tokens": 7.5e-06,
"cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05,
"cache_read_input_token_cost_above_200k_tokens": 6e-07,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 200000,
@@ -17950,6 +17998,7 @@
},
"global.anthropic.claude-haiku-4-5-20251001-v1:0": {
"cache_creation_input_token_cost": 1.25e-06,
"cache_creation_input_token_cost_above_1hr": 2e-06,
"cache_read_input_token_cost": 1e-07,
"input_cost_per_token": 1e-06,
"litellm_provider": "bedrock_converse",
@@ -19083,6 +19132,38 @@
"supports_vision": true,
"supports_pdf_input": true
},
"gpt-image-2": {
"cache_read_input_image_token_cost": 2e-06,
"cache_read_input_token_cost": 1.25e-06,
"input_cost_per_token": 5e-06,
"litellm_provider": "openai",
"mode": "image_generation",
"output_cost_per_token": 1e-05,
"input_cost_per_image_token": 8e-06,
"output_cost_per_image_token": 3e-05,
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
],
"supports_vision": true,
"supports_pdf_input": true
},
"gpt-image-2-2026-04-21": {
"cache_read_input_image_token_cost": 2e-06,
"cache_read_input_token_cost": 1.25e-06,
"input_cost_per_token": 5e-06,
"litellm_provider": "openai",
"mode": "image_generation",
"output_cost_per_token": 1e-05,
"input_cost_per_image_token": 8e-06,
"output_cost_per_image_token": 3e-05,
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
],
"supports_vision": true,
"supports_pdf_input": true
},
"low/1024-x-1024/gpt-image-1.5": {
"input_cost_per_image": 0.009,
"litellm_provider": "openai",
@@ -30052,6 +30133,7 @@
},
"us.anthropic.claude-haiku-4-5-20251001-v1:0": {
"cache_creation_input_token_cost": 1.375e-06,
"cache_creation_input_token_cost_above_1hr": 2.2e-06,
"cache_read_input_token_cost": 1.1e-07,
"input_cost_per_token": 1.1e-06,
"litellm_provider": "bedrock_converse",
@@ -30203,11 +30285,13 @@
},
"us.anthropic.claude-sonnet-4-5-20250929-v1:0": {
"cache_creation_input_token_cost": 4.125e-06,
"cache_creation_input_token_cost_above_1hr": 6.6e-06,
"cache_read_input_token_cost": 3.3e-07,
"input_cost_per_token": 3.3e-06,
"input_cost_per_token_above_200k_tokens": 6.6e-06,
"output_cost_per_token_above_200k_tokens": 2.475e-05,
"cache_creation_input_token_cost_above_200k_tokens": 8.25e-06,
"cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.32e-05,
"cache_read_input_token_cost_above_200k_tokens": 6.6e-07,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 200000,
@@ -30308,6 +30392,7 @@
},
"us.anthropic.claude-opus-4-5-20251101-v1:0": {
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 5.5e-06,
"litellm_provider": "bedrock_converse",
@@ -30335,6 +30420,7 @@
},
"global.anthropic.claude-opus-4-5-20251101-v1:0": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "bedrock_converse",
+53 -21
View File
@@ -21,6 +21,7 @@ import httpx
from httpx._types import CookieTypes, QueryParamTypes, RequestFiles
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
@@ -390,19 +391,28 @@ def _sync_streaming(
):
from litellm.utils import executor
raw_bytes: List[bytes] = []
flush_scheduled = False
try:
raw_bytes: List[bytes] = []
for chunk in response.iter_bytes(): # type: ignore
raw_bytes.append(chunk)
yield chunk
executor.submit(
litellm_logging_obj.flush_passthrough_collected_chunks,
raw_bytes=raw_bytes,
provider_config=provider_config,
)
except Exception as e:
raise e
finally:
if not flush_scheduled and raw_bytes:
flush_scheduled = True
try:
executor.submit(
litellm_logging_obj.flush_passthrough_collected_chunks,
raw_bytes=raw_bytes,
provider_config=provider_config,
)
except Exception as e:
verbose_logger.exception(
"Failed to schedule passthrough spend-tracking flush "
"in _sync_streaming; %d buffered chunks dropped: %s",
len(raw_bytes),
e,
)
async def _async_streaming(
@@ -411,23 +421,45 @@ async def _async_streaming(
provider_config: "BasePassthroughConfig",
):
iter_response = await response
try:
iter_response.raise_for_status()
raw_bytes: List[bytes] = []
async for chunk in iter_response.aiter_bytes(): # type: ignore
raw_bytes.append(chunk)
yield chunk
asyncio.create_task(
litellm_logging_obj.async_flush_passthrough_collected_chunks(
raw_bytes=raw_bytes,
provider_config=provider_config,
)
)
except Exception:
try:
await iter_response.aclose()
except Exception:
pass
raise
raw_bytes: List[bytes] = []
flush_scheduled = False
try:
async for chunk in iter_response.aiter_bytes(): # type: ignore
raw_bytes.append(chunk)
yield chunk
except Exception:
try:
await iter_response.aclose()
except Exception:
pass
raise
finally:
# GeneratorExit (raised on client disconnect) is not caught by
# `except Exception`; the finally block ensures partial usage
# still gets flushed for spend tracking. See LIT-2642.
if not flush_scheduled and raw_bytes:
flush_scheduled = True
try:
asyncio.create_task(
litellm_logging_obj.async_flush_passthrough_collected_chunks(
raw_bytes=raw_bytes,
provider_config=provider_config,
)
)
except Exception as e:
verbose_logger.exception(
"Failed to schedule passthrough spend-tracking flush "
"in _async_streaming; %d buffered chunks dropped: %s",
len(raw_bytes),
e,
)
@@ -117,7 +117,10 @@ class MCPRequestHandler:
return b"{}"
request.body = mock_body # type: ignore
if ".well-known" in str(request.url): # public routes
# Only OAuth metadata routes registered under /.well-known/ are public.
# Match on request.url.path (path-only, exact prefix) so the substring
# cannot be smuggled via query string, hostname, or a deeper URL segment.
if request.url.path.startswith("/.well-known/"):
validated_user_api_key_auth = UserAPIKeyAuth()
elif has_explicit_litellm_key:
# Explicit x-litellm-api-key provided - always validate normally
@@ -126,27 +129,37 @@ class MCPRequestHandler:
)
elif oauth2_headers:
# No x-litellm-api-key, but Authorization header present.
# Could be a LiteLLM key (backward compat) OR an OAuth2 token
# from an upstream MCP provider (e.g. Atlassian).
# Try LiteLLM auth first; on auth failure, treat as OAuth2 passthrough.
# Could be a LiteLLM key (backward compat) OR an opaque OAuth2 token
# the operator wants forwarded to an upstream OAuth2-mode MCP server.
# Try LiteLLM auth first; on auth failure, only fall back to anonymous
# passthrough when the request actually targets a server whose operator
# configured ``auth_type=oauth2``. For any other server (api_key,
# bearer_token, basic, etc.), a failed LiteLLM auth is a real failure
# and must propagate — otherwise an attacker can exchange any garbage
# bearer for an anonymous session.
try:
validated_user_api_key_auth = await user_api_key_auth(
api_key=litellm_api_key, request=request
)
except HTTPException as e:
if e.status_code in (401, 403):
except (HTTPException, ProxyException) as e:
# HTTPException.status_code is int; ProxyException.code is
# normalized to str in its __init__ but can be ``"None"`` or any
# non-numeric string when the caller didn't supply a numeric
# code, so we compare against both int and str forms rather
# than coercing (``int("None")`` would raise ValueError and
# rewrite the auth error as a 500).
status = e.status_code if isinstance(e, HTTPException) else e.code
if status in (
401,
403,
"401",
"403",
) and MCPRequestHandler._target_servers_use_oauth2(
path=request.url.path, mcp_servers=mcp_servers
):
verbose_logger.debug(
"MCP OAuth2: Authorization header is not a valid LiteLLM key, "
"treating as OAuth2 token passthrough"
)
validated_user_api_key_auth = UserAPIKeyAuth()
else:
raise
except ProxyException as e:
if str(e.code) in ("401", "403"):
verbose_logger.debug(
"MCP OAuth2: Authorization header is not a valid LiteLLM key, "
"treating as OAuth2 token passthrough"
"MCP OAuth2: target server is OAuth2-mode, treating "
"Authorization as upstream OAuth2 token passthrough"
)
validated_user_api_key_auth = UserAPIKeyAuth()
else:
@@ -165,6 +178,62 @@ class MCPRequestHandler:
dict(headers),
)
@staticmethod
def _extract_target_server_names_from_path(path: str) -> List[str]:
"""
Extract the target MCP server name from the standard MCP transport
URL patterns: ``/mcp/{server_name}[/...]`` and
``/{server_name}/mcp[/...]``. Returns ``[]`` for any other path so
callers fail closed when the target cannot be resolved.
REST/admin endpoints, OAuth2 server endpoints
(``/{server_name}/authorize``, ``/token`` etc.), and ``.well-known``
discovery routes intentionally fall through those flows do not need
OAuth2 token passthrough. Clients aggregating multiple servers should
use ``x-mcp-servers``, which takes precedence over path parsing.
"""
segments = [s for s in path.split("/") if s]
if len(segments) >= 2 and segments[0] == "mcp":
return [segments[1]]
if len(segments) >= 2 and segments[1] == "mcp":
return [segments[0]]
return []
@staticmethod
def _target_servers_use_oauth2(path: str, mcp_servers: Optional[List[str]]) -> bool:
"""
True only when EVERY MCP server the request targets is configured for
``auth_type == oauth2``. If any target is non-OAuth2 or if the target
cannot be resolved at all return False so the caller fails closed.
Used to gate the "treat Authorization as opaque OAuth2 token" fallback
in :meth:`process_mcp_request` so a failed LiteLLM-auth cannot be
exchanged for an anonymous session against a non-OAuth2 server.
"""
# Inline imports avoid a circular dependency: mcp_server_manager imports
# from this module.
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.types.mcp import MCPAuth
# Use the x-mcp-servers header verbatim when present (including the
# explicitly-empty list, which means "no targets" → fail closed).
# Only fall back to path parsing when the header was absent entirely.
target_names = (
mcp_servers
if mcp_servers is not None
else MCPRequestHandler._extract_target_server_names_from_path(path)
)
if not target_names:
return False
for name in target_names:
server = global_mcp_server_manager.get_mcp_server_by_name(name)
if server is None or server.auth_type != MCPAuth.oauth2:
return False
return True
@staticmethod
def _get_mcp_auth_header_from_headers(headers: Headers) -> Optional[str]:
"""
+99 -44
View File
@@ -1,4 +1,5 @@
import base64
import binascii
import json
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast
@@ -498,6 +499,82 @@ async def rotate_mcp_server_credentials_master_key(
)
def _decode_user_credential(stored: str) -> Optional[str]:
"""Read back a value persisted in ``LiteLLM_MCPUserCredentials.credential_b64``.
Tries nacl decryption first (current write format). Falls back to a
plain ``urlsafe_b64decode`` for rows persisted by older code that wrote
the credential without encryption. Returns ``None`` when neither path
yields a valid string.
"""
decrypted = decrypt_value_helper(
value=stored,
key="mcp_user_credential",
exception_type="debug",
return_original_value=False,
)
if decrypted is not None:
return decrypted
try:
return base64.urlsafe_b64decode(stored).decode()
except (binascii.Error, UnicodeDecodeError, ValueError, TypeError):
return None
def _decode_oauth_payload(stored: str) -> Optional[Dict[str, Any]]:
"""Return the OAuth2 payload dict if ``stored`` holds one, else ``None``.
A row is considered an OAuth2 credential iff its decoded value parses as
a JSON object with ``"type": "oauth2"``. Plain BYOK credentials (which
share the same column) decode to a non-JSON string and return ``None``.
"""
decoded = _decode_user_credential(stored)
if decoded is None:
return None
try:
parsed = json.loads(decoded)
except (ValueError, TypeError):
return None
if isinstance(parsed, dict) and parsed.get("type") == "oauth2":
return parsed
return None
async def rotate_mcp_user_credentials_master_key(
prisma_client: PrismaClient, new_master_key: str
):
"""Re-encrypt every ``LiteLLM_MCPUserCredentials`` row with ``new_master_key``.
Reads each ``credential_b64`` with the current salt key (falling back to
legacy plain base64 for unmigrated rows) and writes it back encrypted
under the new master key. Rows that are unreadable under both paths
are logged and skipped so one corrupt row does not abort the rotation.
"""
rows = await prisma_client.db.litellm_mcpusercredentials.find_many()
for row in rows:
plaintext = _decode_user_credential(row.credential_b64)
if plaintext is None:
verbose_proxy_logger.warning(
"rotate_mcp_user_credentials_master_key: could not decode "
"credential for user_id=%s server_id=%s, skipping",
row.user_id,
row.server_id,
)
continue
re_encrypted = encrypt_value_helper(
plaintext, new_encryption_key=new_master_key
)
await prisma_client.db.litellm_mcpusercredentials.update(
where={
"user_id_server_id": {
"user_id": row.user_id,
"server_id": row.server_id,
}
},
data={"credential_b64": re_encrypted},
)
async def store_user_credential(
prisma_client: PrismaClient,
user_id: str,
@@ -506,7 +583,7 @@ async def store_user_credential(
) -> None:
"""Store a user credential for a BYOK MCP server."""
encoded = base64.urlsafe_b64encode(credential.encode()).decode()
encoded = encrypt_value_helper(credential)
await prisma_client.db.litellm_mcpusercredentials.upsert(
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}},
data={
@@ -532,16 +609,7 @@ async def get_user_credential(
)
if row is None:
return None
try:
return base64.urlsafe_b64decode(row.credential_b64).decode()
except Exception:
# Fall back to nacl decryption for credentials stored by older code
return decrypt_value_helper(
value=row.credential_b64,
key="byok_credential",
exception_type="debug",
return_original_value=False,
)
return _decode_user_credential(row.credential_b64)
async def has_user_credential(
@@ -582,7 +650,7 @@ async def store_user_oauth_credential(
) -> None:
"""Persist an OAuth2 access token for a user+server pair.
The payload is JSON-serialised and stored base64-encoded in the same
The payload is JSON-serialised and stored encrypted in the same
``credential_b64`` column used by BYOK. A ``"type": "oauth2"`` key
differentiates it from plain BYOK API keys.
"""
@@ -606,29 +674,27 @@ async def store_user_oauth_credential(
payload["scopes"] = scopes
# Guard against silently overwriting a BYOK credential with an OAuth token.
# BYOK credentials lack a "type" field (or use a non-"oauth2" type).
# Skip the guard when the caller knows the row is already an OAuth2 credential
# (e.g. during token refresh), saving an extra DB round-trip.
if not skip_byok_guard:
existing = await prisma_client.db.litellm_mcpusercredentials.find_unique(
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
)
if existing is not None:
_byok_error = ValueError(
f"A non-OAuth2 credential already exists for user {user_id} "
f"and server {server_id}. Refusing to overwrite."
if (
existing is not None
and _decode_oauth_payload(existing.credential_b64) is None
):
# Existing row is either a BYOK secret or an OAuth2 row that no
# longer decrypts (e.g. after a salt-key rotation). In either
# case, refuse to overwrite — the caller would clobber data
# that may still be recoverable.
raise ValueError(
f"Existing credential for user {user_id} and server "
f"{server_id} could not be verified as an OAuth2 token. "
f"Refusing to overwrite."
)
try:
raw = json.loads(
base64.urlsafe_b64decode(existing.credential_b64).decode()
)
except Exception:
# Credential is not base64+JSON — it's a plain-text BYOK key.
raise _byok_error
if raw.get("type") != "oauth2":
raise _byok_error
encoded = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode()
encoded = encrypt_value_helper(json.dumps(payload))
await prisma_client.db.litellm_mcpusercredentials.upsert(
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}},
data={
@@ -672,15 +738,7 @@ async def get_user_oauth_credential(
)
if row is None:
return None
try:
decoded = base64.urlsafe_b64decode(row.credential_b64).decode()
parsed = json.loads(decoded)
if isinstance(parsed, dict) and parsed.get("type") == "oauth2":
return parsed
# Row exists but is a BYOK (plain string), not an OAuth token
return None
except Exception:
return None
return _decode_oauth_payload(row.credential_b64)
async def list_user_oauth_credentials(
@@ -694,14 +752,11 @@ async def list_user_oauth_credentials(
)
results: List[Dict[str, Any]] = []
for row in rows:
try:
decoded = base64.urlsafe_b64decode(row.credential_b64).decode()
parsed = json.loads(decoded)
if isinstance(parsed, dict) and parsed.get("type") == "oauth2":
parsed["server_id"] = row.server_id
results.append(parsed)
except Exception:
pass # Skip non-OAuth rows (BYOK plain strings)
payload = _decode_oauth_payload(row.credential_b64)
if payload is None:
continue
payload["server_id"] = row.server_id
results.append(payload)
return results
@@ -131,6 +131,22 @@ def decode_state_hash(encrypted_state: str) -> dict:
return state_data
def _get_validated_client_redirect_uri(state_data: Dict[str, Any]) -> str:
"""Return a loopback client redirect URI from OAuth state."""
redirect_uri = state_data.get("client_redirect_uri") or state_data.get("base_url")
if not redirect_uri or not isinstance(redirect_uri, str):
raise HTTPException(status_code=400, detail="Invalid redirect URI")
validate_loopback_redirect_uri(redirect_uri)
return redirect_uri
def _append_query_params(url: str, params: Dict[str, str]) -> str:
parsed = urlparse(url)
query_params = parse_qsl(parsed.query, keep_blank_values=True)
query_params.extend(params.items())
return urlunparse(parsed._replace(query=urlencode(query_params)))
def _resolve_oauth2_server_for_root_endpoints(
client_ip: Optional[str] = None,
) -> Optional[MCPServer]:
@@ -568,7 +584,7 @@ async def authorize(
else None
)
if mcp_server is None and mcp_server_name is None:
mcp_server = _resolve_oauth2_server_for_root_endpoints()
mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
if mcp_server is None:
raise HTTPException(status_code=404, detail="MCP server not found")
# Use server's stored client_id when caller doesn't supply one.
@@ -630,7 +646,7 @@ async def token_endpoint(
lookup_name, client_ip=client_ip
)
if mcp_server is None and mcp_server_name is None:
mcp_server = _resolve_oauth2_server_for_root_endpoints()
mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
if mcp_server is None:
raise HTTPException(status_code=404, detail="MCP server not found")
return await exchange_token_with_server(
@@ -651,7 +667,6 @@ async def token_endpoint(
async def callback(code: str, state: str):
try:
state_data = decode_state_hash(state)
base_url = state_data["base_url"]
original_state = state_data["original_state"]
# Re-validate loopback at the sink. /authorize rejects non-loopback
@@ -659,10 +674,10 @@ async def callback(code: str, state: str):
# minted before that check was added have no expiry and remain
# valid indefinitely. Validating here blocks the open-redirect +
# code-theft primitive even for pre-fix states.
validate_loopback_redirect_uri(base_url)
redirect_uri = _get_validated_client_redirect_uri(state_data)
params = {"code": code, "state": original_state}
complete_returned_url = f"{base_url}?{urlencode(params)}"
complete_returned_url = _append_query_params(redirect_uri, params)
return RedirectResponse(url=complete_returned_url, status_code=302)
except HTTPException:
@@ -719,16 +734,16 @@ def _build_oauth_protected_resource_response(
)
request_base_url = get_request_base_url(request)
client_ip = IPAddressUtils.get_mcp_client_ip(request)
# When no server name provided, try to resolve the single OAuth2 server
if mcp_server_name is None:
resolved = _resolve_oauth2_server_for_root_endpoints()
resolved = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
if resolved:
mcp_server_name = resolved.server_name or resolved.name
mcp_server: Optional[MCPServer] = None
if mcp_server_name:
client_ip = IPAddressUtils.get_mcp_client_ip(request)
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(
mcp_server_name, client_ip=client_ip
)
@@ -835,10 +850,11 @@ def _build_oauth_authorization_server_response(
)
request_base_url = get_request_base_url(request)
client_ip = IPAddressUtils.get_mcp_client_ip(request)
# When no server name provided, try to resolve the single OAuth2 server
if mcp_server_name is None:
resolved = _resolve_oauth2_server_for_root_endpoints()
resolved = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
if resolved:
mcp_server_name = resolved.server_name or resolved.name
@@ -855,7 +871,6 @@ def _build_oauth_authorization_server_response(
mcp_server: Optional[MCPServer] = None
if mcp_server_name:
client_ip = IPAddressUtils.get_mcp_client_ip(request)
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(
mcp_server_name, client_ip=client_ip
)
@@ -1007,8 +1022,9 @@ async def register_client(request: Request, mcp_server_name: Optional[str] = Non
"client_secret": "dummy",
"redirect_uris": [f"{request_base_url}/callback"],
}
client_ip = IPAddressUtils.get_mcp_client_ip(request)
if not mcp_server_name:
resolved = _resolve_oauth2_server_for_root_endpoints()
resolved = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
if resolved:
return await register_client_with_server(
request=request,
@@ -1021,7 +1037,6 @@ async def register_client(request: Request, mcp_server_name: Optional[str] = Non
)
return dummy_return
client_ip = IPAddressUtils.get_mcp_client_ip(request)
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(
mcp_server_name, client_ip=client_ip
)
@@ -41,6 +41,7 @@ from litellm.constants import (
MCP_TOOL_LISTING_TIMEOUT,
)
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get
from litellm.experimental_mcp_client.client import MCPClient, MCPSigV4Auth
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
@@ -50,8 +51,11 @@ from litellm.proxy._experimental.mcp_server.oauth2_token_cache import resolve_mc
from litellm.proxy._experimental.mcp_server.utils import (
MCP_TOOL_PREFIX_SEPARATOR,
add_server_prefix_to_name,
compute_short_server_prefix,
get_server_prefix,
is_short_mcp_tool_prefix_enabled,
is_tool_name_prefixed,
iter_known_server_prefixes,
merge_mcp_headers,
normalize_server_name,
split_server_prefix_from_name,
@@ -106,6 +110,12 @@ if not _separator_probe.is_valid:
SEP_986_URL,
)
_AZURE_ENTRA_HOSTS = {
"login.microsoftonline.com", # Global
"login.microsoftonline.us", # US Government
"login.chinacloudapi.cn", # China
}
def _warn_on_server_name_fields(
*,
@@ -364,6 +374,7 @@ class MCPServerManager:
aws_session_name=server_config.get("aws_session_name", None),
instructions=server_config.get("instructions", None),
)
self._assign_unique_short_prefix(new_server)
self.config_mcp_servers[server_id] = new_server
# Check if this is an OpenAPI-based server
@@ -726,6 +737,7 @@ class MCPServerManager:
try:
if mcp_server.server_id not in self.registry:
new_server = await self.build_mcp_server_from_table(mcp_server)
self._assign_unique_short_prefix(new_server)
self.registry[mcp_server.server_id] = new_server
await self._maybe_register_openapi_tools(new_server)
verbose_logger.debug(f"Added MCP Server: {new_server.name}")
@@ -738,6 +750,12 @@ class MCPServerManager:
try:
if mcp_server.server_id in self.registry:
new_server = await self.build_mcp_server_from_table(mcp_server)
# Carry the previously-resolved short prefix across so the
# tool names stay stable for clients holding cached lists.
existing_prefix = self.registry[mcp_server.server_id].short_prefix
if existing_prefix and not new_server.short_prefix:
new_server.short_prefix = existing_prefix
self._assign_unique_short_prefix(new_server)
self.registry[mcp_server.server_id] = new_server
await self._maybe_register_openapi_tools(new_server)
verbose_logger.debug(f"Updated MCP Server: {new_server.name}")
@@ -1236,7 +1254,11 @@ class MCPServerManager:
## HANDLE OPENAPI TOOLS
if server.spec_path:
_tools = global_mcp_tool_registry.list_tools(tool_prefix=server.name)
# OpenAPI tools were stored in the registry under the prefix
# active at registration time — fetch by that same prefix.
_tools = global_mcp_tool_registry.list_tools(
tool_prefix=get_server_prefix(server)
)
tools = global_mcp_tool_registry.convert_tools_to_mcp_sdk_tool_type(
_tools
)
@@ -1478,6 +1500,47 @@ class MCPServerManager:
)
return await client.get_prompt(get_prompt_request_params)
@staticmethod
def _is_same_authority_metadata_url(url: str, server_url: str) -> bool:
"""
Whether ``url`` shares scheme, host, and port with ``server_url``.
Same-authority metadata URLs are produced by our well-known discovery
construction and by resource servers that publish protected-resource
metadata on the resource origin. These must keep working for
administrator-configured internal MCP servers, so they are fetched
directly. Cross-origin URLs are fetched through ``async_safe_get``.
"""
try:
target = urlparse(url)
base = urlparse(server_url)
except Exception:
return False
if target.scheme not in ("http", "https") or not target.hostname:
return False
target_port = target.port or (443 if target.scheme == "https" else 80)
base_port = base.port or (443 if base.scheme == "https" else 80)
return (
base.scheme == target.scheme
and (base.hostname or "").lower() == target.hostname.lower()
and base_port == target_port
)
async def _fetch_oauth_discovery_url(self, url: str, server_url: str) -> Any:
client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.MCP,
params={"timeout": MCP_METADATA_TIMEOUT},
)
if self._is_same_authority_metadata_url(url, server_url):
# Same-authority URLs may point at administrator-configured
# internal MCP servers. Do not run them through user URL
# validation, but also do not follow redirects because the
# redirect target would not inherit the same-authority guarantee.
return await client.get(url, follow_redirects=False)
return await async_safe_get(client, url)
async def _descovery_metadata(
self,
server_url: str,
@@ -1488,11 +1551,28 @@ class MCPServerManager:
client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
response = await client.get(server_url)
response.raise_for_status()
verbose_logger.warning(
"MCP OAuth discovery unexpectedly succeeded for %s; server did not challenge",
server_url,
(
authorization_servers,
resource_scopes,
) = await self._attempt_well_known_discovery(server_url)
metadata = await self._fetch_authorization_server_metadata(
authorization_servers, server_url
)
raise RuntimeError("OAuth discovery must not succeed without a challenge")
if (
metadata is None
and not resource_scopes
and authorization_servers
and response.status_code == 200
):
verbose_logger.warning(
"MCP OAuth discovery for %s received 200 OK without RFC 9728 challenge and no discoverable authorization metadata.",
server_url,
)
if metadata is None and resource_scopes:
return MCPOAuthMetadata(scopes=resource_scopes)
if metadata is not None and resource_scopes:
metadata.scopes = resource_scopes
return metadata
except HTTPStatusError as exc:
verbose_logger.debug(
"MCP OAuth discovery for %s received status error: %s",
@@ -1510,14 +1590,14 @@ class MCPServerManager:
header_value
)
authorization_servers: List[str] = []
resource_scopes: Optional[List[str]] = None
authorization_servers = []
resource_scopes = None
if resource_metadata_url:
(
authorization_servers,
resource_scopes,
) = await self._fetch_oauth_metadata_from_resource(
resource_metadata_url
resource_metadata_url, server_url
)
else:
(
@@ -1538,7 +1618,7 @@ class MCPServerManager:
if authorization_servers:
metadata = await self._fetch_authorization_server_metadata(
authorization_servers
authorization_servers, server_url
)
preferred_scopes = scopes or resource_scopes
@@ -1578,19 +1658,26 @@ class MCPServerManager:
return resource_metadata_url, scopes
async def _fetch_oauth_metadata_from_resource(
self, resource_metadata_url: str
self, resource_metadata_url: str, server_url: str
) -> Tuple[List[str], Optional[List[str]]]:
if not resource_metadata_url:
return [], None
try:
client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.MCP,
params={"timeout": MCP_METADATA_TIMEOUT},
response = await self._fetch_oauth_discovery_url(
resource_metadata_url, server_url
)
response = await client.get(resource_metadata_url)
response.raise_for_status()
data = response.json()
except SSRFError as exc:
verbose_logger.warning(
"MCP OAuth discovery: refusing to fetch resource metadata from %s "
"(rejected by SSRF guard for server %s): %s",
resource_metadata_url,
server_url,
exc,
)
return [], None
except Exception as exc: # pragma: no cover - network issues
verbose_logger.debug(
"Failed to fetch MCP OAuth metadata from %s: %s",
@@ -1639,23 +1726,25 @@ class MCPServerManager:
(
authorization_servers,
scopes,
) = await self._fetch_oauth_metadata_from_resource(url)
) = await self._fetch_oauth_metadata_from_resource(url, server_url)
if authorization_servers:
return authorization_servers, scopes
return [], None
async def _fetch_authorization_server_metadata(
self, authorization_servers: List[str]
self, authorization_servers: List[str], server_url: str
) -> Optional[MCPOAuthMetadata]:
for issuer in authorization_servers:
metadata = await self._fetch_single_authorization_server_metadata(issuer)
metadata = await self._fetch_single_authorization_server_metadata(
issuer, server_url
)
if metadata is not None:
return metadata
return None
async def _fetch_single_authorization_server_metadata(
self, issuer_url: str
self, issuer_url: str, server_url: str
) -> Optional[MCPOAuthMetadata]:
try:
parsed = urlparse(issuer_url)
@@ -1674,19 +1763,27 @@ class MCPServerManager:
f"{base}/.well-known/oauth-authorization-server/{path}"
)
candidate_urls.append(f"{base}/.well-known/openid-configuration/{path}")
candidate_urls.append(
f"{issuer_url.rstrip('/')}/.well-known/openid-configuration"
)
candidate_urls.append(f"{base}/.well-known/oauth-authorization-server")
candidate_urls.append(f"{base}/.well-known/openid-configuration")
candidate_urls.append(issuer_url.rstrip("/"))
for url in candidate_urls:
try:
client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.MCP,
params={"timeout": MCP_METADATA_TIMEOUT},
)
response = await client.get(url)
response = await self._fetch_oauth_discovery_url(url, server_url)
response.raise_for_status()
data = response.json()
except SSRFError as exc:
verbose_logger.warning(
"MCP OAuth discovery: refusing to fetch authorization-server "
"metadata from %s (rejected by SSRF guard for server %s): %s",
url,
server_url,
exc,
)
continue
except Exception as exc: # pragma: no cover - network issues
verbose_logger.debug(
"Failed to fetch authorization metadata from %s: %s",
@@ -1713,7 +1810,28 @@ class MCPServerManager:
):
return metadata
return None
return self._build_azure_authorization_server_metadata(parsed)
@staticmethod
def _build_azure_authorization_server_metadata(
parsed_issuer_url: Any,
) -> Optional[MCPOAuthMetadata]:
path_parts = [
part for part in (parsed_issuer_url.path or "").split("/") if part
]
if (
parsed_issuer_url.netloc not in _AZURE_ENTRA_HOSTS
or len(path_parts) != 2
or path_parts[1] != "v2.0"
):
return None
tenant = path_parts[0]
base = f"{parsed_issuer_url.scheme}://{parsed_issuer_url.netloc}/{tenant}"
return MCPOAuthMetadata(
authorization_url=f"{base}/oauth2/v2.0/authorize",
token_url=f"{base}/oauth2/v2.0/token",
)
@staticmethod
def _decrypt_credential_field(
@@ -1810,6 +1928,63 @@ class MCPServerManager:
verbose_logger.warning(f"Error listing tools from {server_name}: {str(e)}")
return []
_SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024
def _assign_unique_short_prefix(self, server: MCPServer) -> None:
"""Resolve and cache a collision-free short tool prefix on ``server``.
Called at registration time for every MCP server entering the
registry. Mutates ``server.short_prefix`` in place. No-ops when
``LITELLM_USE_SHORT_MCP_TOOL_PREFIX`` is disabled, when the server
has no ``server_id`` (synthetic temp-server objects), or when a
prefix is already cached.
Collision strategy: take the natural hash; if it's already used by
a *different* server in the combined registry, rehash with an
incrementing attempt counter until we find an unused slot. The
attempt counter is folded into the hash so the resulting prefix is
still deterministic for a given (server_id, set-of-other-server-ids)
pair within one process.
"""
if not is_short_mcp_tool_prefix_enabled():
return
if server.short_prefix:
return
if not server.server_id:
return
used: Dict[str, str] = {}
for other in self.get_registry().values():
if other.server_id == server.server_id:
continue
if other.short_prefix:
used[other.short_prefix] = other.server_id
for attempt in range(self._SHORT_PREFIX_MAX_REHASH_ATTEMPTS):
candidate = compute_short_server_prefix(server.server_id, attempt=attempt)
if candidate not in used:
server.short_prefix = candidate
if attempt > 0:
verbose_logger.info(
"MCP short-prefix collision resolved for server %s: "
"natural hash collided with %s, using rehashed prefix "
"%s (attempt=%d).",
server.server_id,
used.get(
compute_short_server_prefix(server.server_id, attempt=0),
"<unknown>",
),
candidate,
attempt,
)
return
raise RuntimeError(
f"Unable to assign a unique short MCP tool prefix for server "
f"{server.server_id} after {self._SHORT_PREFIX_MAX_REHASH_ATTEMPTS} "
"attempts; the 3-character prefix space is too crowded."
)
def _create_prefixed_tools(
self, tools: List[MCPTool], server: MCPServer, add_prefix: bool = True
) -> List[MCPTool]:
@@ -1838,9 +2013,13 @@ class MCPServerManager:
tool_copy.name = name_to_use
prefixed_tools.append(tool_copy)
# Update tool to server mapping for resolution (support both forms)
# Register every known prefix form (alias, server_name, server_id,
# short ID) so call_tool can resolve regardless of which form a
# caller / cached client is using.
self.tool_name_to_mcp_server_name_mapping[original_name] = prefix
self.tool_name_to_mcp_server_name_mapping[prefixed_name] = prefix
for known_prefix in iter_known_server_prefixes(server):
qualified = add_server_prefix_to_name(original_name, known_prefix)
self.tool_name_to_mcp_server_name_mapping[qualified] = prefix
verbose_logger.info(
f"Successfully fetched {len(prefixed_tools)} tools from server {server.name}"
@@ -2601,37 +2780,43 @@ class MCPServerManager:
Returns:
MCPServer if found, None otherwise
"""
registry_servers = list(self.get_registry().values())
# Build prefix → server lookup covering every known form a tool name
# may take (alias / server_name / server_id / short ID). This is what
# makes the short-prefix mode work without breaking historical names.
prefix_to_server: Dict[str, MCPServer] = {}
for server in registry_servers:
for known_prefix in iter_known_server_prefixes(server):
normalised = normalize_server_name(known_prefix)
prefix_to_server.setdefault(normalised, server)
# First try with the original tool name
if tool_name in self.tool_name_to_mcp_server_name_mapping:
server_name = self.tool_name_to_mcp_server_name_mapping[tool_name]
for server in self.get_registry().values():
if normalize_server_name(server.name) == normalize_server_name(
server_name
):
normalised_lookup = normalize_server_name(server_name)
if normalised_lookup in prefix_to_server:
return prefix_to_server[normalised_lookup]
for server in registry_servers:
if normalize_server_name(server.name) == normalised_lookup:
return server
# If not found and tool name is prefixed, try extracting server name from prefix
known_prefixes = {
normalize_server_name(get_server_prefix(s))
for s in self.get_registry().values()
if get_server_prefix(s)
}
if is_tool_name_prefixed(tool_name, known_server_prefixes=known_prefixes):
# If not found and tool name is prefixed, extract the prefix and
# match against any known form.
if is_tool_name_prefixed(
tool_name, known_server_prefixes=set(prefix_to_server.keys())
):
(
original_tool_name,
server_name_from_prefix,
) = split_server_prefix_from_name(tool_name)
if original_tool_name in self.tool_name_to_mcp_server_name_mapping:
for server in self.get_registry().values():
if server.server_name is None:
if normalize_server_name(server.name) == normalize_server_name(
server_name_from_prefix
):
return server
elif normalize_server_name(
server.server_name
) == normalize_server_name(server_name_from_prefix):
return server
normalised_prefix = normalize_server_name(server_name_from_prefix)
matched_server = prefix_to_server.get(normalised_prefix)
if matched_server is not None and (
original_tool_name in self.tool_name_to_mcp_server_name_mapping
or tool_name in self.tool_name_to_mcp_server_name_mapping
):
return matched_server
return None
@@ -2666,6 +2851,9 @@ class MCPServerManager:
previous_registry = self.registry
new_registry: Dict[str, MCPServer] = {}
# Stage one: build every server. Stage two assigns short prefixes
# against the *full* set so dedup is deterministic regardless of
# iteration order.
for server in db_mcp_servers:
existing_server = previous_registry.get(server.server_id)
@@ -2689,10 +2877,21 @@ class MCPServerManager:
f"Building server from DB: {server.server_id} ({server.server_name})"
)
new_server = await self.build_mcp_server_from_table(server)
# Carry the cached short_prefix from the previous registry entry
# (if any) so the prefix is stable across reloads.
if existing_server is not None and existing_server.short_prefix:
new_server.short_prefix = existing_server.short_prefix
new_registry[server.server_id] = new_server
await self._maybe_register_openapi_tools(new_server)
# Swap in the new registry first so _assign_unique_short_prefix
# sees the complete set when checking for collisions.
self.registry = new_registry
for new_server in new_registry.values():
self._assign_unique_short_prefix(new_server)
# Register OpenAPI tools *after* the final short prefix is assigned
# so the tools are stored in the global registry under the same
# prefix that lookups will use.
await self._maybe_register_openapi_tools(new_server)
verbose_logger.debug(
"MCP registry refreshed (%s servers in registry)", len(new_registry)
@@ -49,6 +49,7 @@ from litellm.proxy._experimental.mcp_server.utils import (
LITELLM_MCP_SERVER_VERSION,
add_server_prefix_to_name,
get_server_prefix,
iter_known_server_prefixes,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
@@ -711,13 +712,7 @@ if MCP_AVAILABLE:
for server in allowed_mcp_servers:
if server:
match_list = [
s.lower()
for s in [
server.alias,
server.server_name,
server.server_id,
]
if s is not None
s.lower() for s in iter_known_server_prefixes(server) if s
]
if server_or_group.lower() in match_list:
@@ -2031,11 +2026,13 @@ if MCP_AVAILABLE:
# Remove prefix from tool name for logging and processing
original_tool_name, server_name = split_server_prefix_from_name(name)
# If tool name is unprefixed, resolve its server so we can enforce permissions
if not server_name:
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
if mcp_server:
server_name = mcp_server.name
# Resolve the actual MCP server up-front so the permission check uses
# the canonical server.name even when the tool name is prefixed with a
# short ID (LITELLM_USE_SHORT_MCP_TOOL_PREFIX) that doesn't match the
# server's display name directly.
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
if mcp_server is not None:
server_name = mcp_server.name
# Only enforce server-level permissions when we can resolve a server
if server_name:
+135 -3
View File
@@ -2,10 +2,11 @@
MCP Server Utilities
"""
from typing import Any, Dict, Mapping, Optional, Tuple
from typing import Any, Dict, Iterator, Mapping, Optional, Tuple
import os
import hashlib
import importlib
import os
# Constants
LITELLM_MCP_SERVER_NAME = "litellm-mcp-server"
@@ -14,6 +15,89 @@ LITELLM_MCP_SERVER_DESCRIPTION = "MCP Server for LiteLLM"
MCP_TOOL_PREFIX_SEPARATOR = os.environ.get("MCP_TOOL_PREFIX_SEPARATOR", "-")
MCP_TOOL_PREFIX_FORMAT = "{server_name}{separator}{tool_name}"
# ---------------------------------------------------------------------------
# Short-ID tool prefix (opt-in)
# ---------------------------------------------------------------------------
# When LITELLM_USE_SHORT_MCP_TOOL_PREFIX is truthy the prefix attached to MCP
# tool / prompt / resource / resource-template names switches from the
# (potentially long) human-readable server name to a deterministic three
# character ID derived from the server's ``server_id``.
#
# Why three characters?
# * The first character is restricted to 52 alphabetic characters
# ([A-Za-z]) and the remaining two characters use the full base62
# alphabet ([0-9A-Za-z]). That guarantees the prefix never starts
# with a digit so it remains a valid identifier for every model API
# (some providers historically required a leading alphabetic char).
# * 52 * 62 * 62 = 199_888 distinct IDs. The chance of a real local
# tool name happening to begin with the exact prefix LiteLLM assigned
# to a given MCP server is negligible in practice.
# * The IDs are short enough that prefixed tool names stay well under
# the 60-character upper bound enforced by some model APIs (Anthropic
# etc.) even for long upstream tool names.
# * The mapping is deterministic (SHA-256 of ``server_id`` → three
# characters drawn from the alphabets above), so the prefix is stable
# across processes, workers and restarts without any persistence
# layer. Two servers with different ``server_id`` values can in
# principle hash to the same three chars; that natural-hash collision
# IS a routing-correctness issue (the second registrant would otherwise
# have its tools misrouted to the first), so registration goes through
# ``MCPServerManager._assign_unique_short_prefix`` which rehashes with
# a deterministic attempt counter until it finds an unused prefix and
# caches the result on ``MCPServer.short_prefix``. A collision is
# logged at INFO when it happens.
#
# This flag is intentionally opt-in for the first release so customers can
# migrate. It will become the default in a future release.
SHORT_MCP_TOOL_PREFIX_LENGTH = 3
_BASE62_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
# Subset of _BASE62_ALPHABET used for the *first* character only, to
# guarantee the prefix never starts with a digit.
_BASE52_ALPHA_ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
def is_short_mcp_tool_prefix_enabled() -> bool:
"""Return True when the short-ID tool prefix mode is enabled.
Read at call time (not import time) so tests and runtime config changes
take effect without reimporting the module.
"""
raw = os.environ.get("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "")
return raw.strip().lower() in ("1", "true", "yes", "on")
def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str:
"""Derive the deterministic three-character prefix for a server.
Uses SHA-256 of ``f"{server_id}#{attempt}"`` and folds the first eight
bytes into a fixed-length string whose first character is drawn from
``_BASE52_ALPHA_ALPHABET`` (so the prefix never starts with a digit)
and whose remaining characters are drawn from the full base62
alphabet. Pass ``attempt > 0`` to rehash to a different prefix when
the natural hash collides with a prefix already assigned to another
server (see ``MCPServerManager._assign_unique_short_prefix``). An
empty ``server_id`` raises ``ValueError`` short prefixes require a
stable identifier to be deterministic.
"""
if not server_id:
raise ValueError("compute_short_server_prefix requires a non-empty server_id")
seed = server_id if attempt == 0 else f"{server_id}#{attempt}"
digest = hashlib.sha256(seed.encode("utf-8")).digest()
value = int.from_bytes(digest[:8], "big")
# Build chars from least-significant to most-significant; we reverse
# at the end so the first emitted char comes from the high-order
# bits of the digest (which is the position we constrain to be
# alphabetic).
chars = []
for position in range(SHORT_MCP_TOOL_PREFIX_LENGTH):
is_first_char = position == SHORT_MCP_TOOL_PREFIX_LENGTH - 1
alphabet = _BASE52_ALPHA_ALPHABET if is_first_char else _BASE62_ALPHABET
value, idx = divmod(value, len(alphabet))
chars.append(alphabet[idx])
return "".join(reversed(chars))
def is_mcp_available() -> bool:
"""
@@ -82,7 +166,25 @@ def add_server_prefix_to_name(name: str, server_name: str) -> str:
def get_server_prefix(server: Any) -> str:
"""Return the prefix for a server: alias if present, else server_name, else server_id"""
"""Return the prefix for a server.
When the short-prefix mode is enabled (``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``)
a three-character base62 ID is returned. We prefer the cached
``server.short_prefix`` value when set that field is populated at
registration time by ``MCPServerManager._assign_unique_short_prefix``
and resolves natural-hash collisions deterministically and only fall
back to the natural hash for ad-hoc / temp-server objects without a
cached value. In default mode the historical behaviour is preserved:
alias if present, else server_name, else server_id.
"""
if is_short_mcp_tool_prefix_enabled():
cached = getattr(server, "short_prefix", None)
if cached:
return cached
server_id = getattr(server, "server_id", None)
if server_id:
return compute_short_server_prefix(server_id)
if hasattr(server, "alias") and server.alias:
return server.alias
if hasattr(server, "server_name") and server.server_name:
@@ -92,6 +194,36 @@ def get_server_prefix(server: Any) -> str:
return ""
def iter_known_server_prefixes(server: Any) -> Iterator[str]:
"""Yield every prefix form that may appear in tool names for ``server``.
Always includes the *current* prefix returned by ``get_server_prefix``.
Additionally yields the historical (alias / server_name / server_id) and
short-ID forms so the routing layer can resolve tool names regardless of
which prefix mode was active when the client first observed them.
"""
seen = set()
def _emit(value: Optional[str]) -> Iterator[str]:
if value and value not in seen:
seen.add(value)
yield value
yield from _emit(get_server_prefix(server))
yield from _emit(getattr(server, "short_prefix", None))
server_id = getattr(server, "server_id", None)
if server_id:
try:
yield from _emit(compute_short_server_prefix(server_id))
except ValueError:
pass
yield from _emit(getattr(server, "alias", None))
yield from _emit(getattr(server, "server_name", None))
yield from _emit(server_id)
def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]:
"""Return the unprefixed name plus the server name used as prefix."""
if MCP_TOOL_PREFIX_SEPARATOR in prefixed_name:
File diff suppressed because one or more lines are too long
-22
View File
@@ -1,22 +0,0 @@
1:"$Sreact.fragment"
2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
6:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"]
7:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"]
a:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"]
b:"$Sreact.suspense"
d:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"]
f:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"]
11:I[168027,[],"default"]
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}]
0:{"P":null,"b":"zxkD4-EPlgfKHDTw8O869","c":["","chat"],"q":"","i":false,"f":[[["",{"children":["chat",{"children":["__PAGE__",{}]}]},"$undefined","$undefined",true],[["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true,"nonce":"$undefined"}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]],"forbidden":"$undefined","unauthorized":"$undefined"}]}]}]}]}]]}],{"children":[["$","$1","c",{"children":[null,["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":"$undefined","forbidden":"$undefined","unauthorized":"$undefined"}]]}],{"children":[["$","$1","c",{"children":[["$","$L6",null,{"Component":"$7","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@8","$@9"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true,"nonce":"$undefined"}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true,"nonce":"$undefined"}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true,"nonce":"$undefined"}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true,"nonce":"$undefined"}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true,"nonce":"$undefined"}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true,"nonce":"$undefined"}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true,"nonce":"$undefined"}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true,"nonce":"$undefined"}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true,"nonce":"$undefined"}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true,"nonce":"$undefined"}]],["$","$La",null,{"children":["$","$b",null,{"name":"Next.MetadataOutlet","children":"$@c"}]}]]}],{},null,false,false]},null,false,false]},null,false,false],["$","$1","h",{"children":[null,["$","$Ld",null,{"children":"$Le"}],["$","div",null,{"hidden":true,"children":["$","$Lf",null,{"children":["$","$b",null,{"name":"Next.Metadata","children":"$L10"}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],false]],"m":"$undefined","G":["$11",[]],"S":true}
8:{}
9:"$0:f:0:1:1:children:1:children:0:props:children:0:props:serverProvidedParams:params"
e:[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]]
12:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"]
c:null
10:[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L12","4",{}]]
@@ -1,22 +0,0 @@
1:"$Sreact.fragment"
2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
6:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"]
7:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"]
a:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"]
b:"$Sreact.suspense"
d:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"]
f:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"]
11:I[168027,[],"default"]
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}]
0:{"P":null,"b":"zxkD4-EPlgfKHDTw8O869","c":["","chat"],"q":"","i":false,"f":[[["",{"children":["chat",{"children":["__PAGE__",{}]}]},"$undefined","$undefined",true],[["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true,"nonce":"$undefined"}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]],"forbidden":"$undefined","unauthorized":"$undefined"}]}]}]}]}]]}],{"children":[["$","$1","c",{"children":[null,["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":"$undefined","forbidden":"$undefined","unauthorized":"$undefined"}]]}],{"children":[["$","$1","c",{"children":[["$","$L6",null,{"Component":"$7","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@8","$@9"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true,"nonce":"$undefined"}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true,"nonce":"$undefined"}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true,"nonce":"$undefined"}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true,"nonce":"$undefined"}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true,"nonce":"$undefined"}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true,"nonce":"$undefined"}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true,"nonce":"$undefined"}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true,"nonce":"$undefined"}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true,"nonce":"$undefined"}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true,"nonce":"$undefined"}]],["$","$La",null,{"children":["$","$b",null,{"name":"Next.MetadataOutlet","children":"$@c"}]}]]}],{},null,false,false]},null,false,false]},null,false,false],["$","$1","h",{"children":[null,["$","$Ld",null,{"children":"$Le"}],["$","div",null,{"hidden":true,"children":["$","$Lf",null,{"children":["$","$b",null,{"name":"Next.Metadata","children":"$L10"}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],false]],"m":"$undefined","G":["$11",[]],"S":true}
8:{}
9:"$0:f:0:1:1:children:1:children:0:props:children:0:props:serverProvidedParams:params"
e:[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]]
12:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"]
c:null
10:[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L12","4",{}]]
@@ -1,6 +0,0 @@
1:"$Sreact.fragment"
2:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"]
3:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"]
4:"$Sreact.suspense"
5:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"]
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","h",{"children":[null,["$","$L2",null,{"children":[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]]}],["$","div",null,{"hidden":true,"children":["$","$L3",null,{"children":["$","$4",null,{"name":"Next.Metadata","children":[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L5","4",{}]]}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],"loading":null,"isPartial":false}
@@ -1,8 +0,0 @@
1:"$Sreact.fragment"
2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","template":["$","$L5",null,{}],"notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]]}]}]}]}]}]]}],"loading":null,"isPartial":false}
@@ -1,4 +0,0 @@
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}]
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","tree":{"name":"","paramType":null,"paramKey":"","hasRuntimePrefetch":false,"slots":{"children":{"name":"chat","paramType":null,"paramKey":"chat","hasRuntimePrefetch":false,"slots":{"children":{"name":"__PAGE__","paramType":null,"paramKey":"__PAGE__","hasRuntimePrefetch":false,"slots":null,"isRootLayout":false}},"isRootLayout":false}},"isRootLayout":true},"staleTime":300}
@@ -1,9 +0,0 @@
1:"$Sreact.fragment"
2:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"]
3:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"]
6:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"]
7:"$Sreact.suspense"
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[["$","$L2",null,{"Component":"$3","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@4","$@5"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true}]],["$","$L6",null,{"children":["$","$7",null,{"name":"Next.MetadataOutlet","children":"$@8"}]}]]}],"loading":null,"isPartial":false}
4:{}
5:"$0:rsc:props:children:0:props:serverProvidedParams:params"
8:null
@@ -1,4 +0,0 @@
1:"$Sreact.fragment"
2:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
3:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[null,["$","$L2",null,{"parallelRouterKey":"children","template":["$","$L3",null,{}]}]]}],"loading":null,"isPartial":false}
+432
View File
@@ -0,0 +1,432 @@
"""
Lazy registration for optional feature routers. Each LAZY_FEATURES entry
imports its module only on the first request matching its path prefix,
saving ~700 MB at idle for deployments that don't use these features.
First hit pays the import cost (1-3 s for heavy modules); /openapi.json
omits each feature's routes until the feature is warmed.
"""
import asyncio
import importlib
import sys
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Callable, Dict, Tuple
from starlette.types import Receive, Scope, Send
from litellm._logging import verbose_proxy_logger
if TYPE_CHECKING:
from fastapi import APIRouter, FastAPI
def _include_router(attr_name: str = "router") -> Callable[["FastAPI", object], None]:
def _register(app: "FastAPI", module: object) -> None:
app.include_router(getattr(module, attr_name))
return _register
def _mount_app(
prefix: str, attr_name: str = "app"
) -> Callable[["FastAPI", object], None]:
def _register(app: "FastAPI", module: object) -> None:
app.mount(path=prefix, app=getattr(module, attr_name))
return _register
@dataclass(frozen=True)
class LazyFeature:
name: str
module_path: str
path_prefixes: Tuple[str, ...]
register_fn: Callable[["FastAPI", object], None] = field(
default_factory=lambda: _include_router("router")
)
# For routes whose path has a leading parameter (e.g. /{server}/authorize)
# — startswith can't match those, so the matcher also checks endswith.
path_suffixes: Tuple[str, ...] = ()
# Keep the stub injected even after load — for mounted ASGI sub-apps
# whose routes don't appear in the parent app's openapi spec.
persistent_swagger_stub: bool = False
LAZY_FEATURES: Tuple[LazyFeature, ...] = (
LazyFeature(
name="guardrails",
module_path="litellm.proxy.guardrails.guardrail_endpoints",
path_prefixes=(
"/guardrails",
"/v2/guardrails",
"/apply_guardrail",
"/policies/usage",
),
),
LazyFeature(
name="policies",
module_path="litellm.proxy.management_endpoints.policy_endpoints",
# Trailing slash to avoid matching /policies/... (policy_engine).
path_prefixes=("/policy/", "/utils/test_policies_and_guardrails"),
),
LazyFeature(
name="policy_engine",
module_path="litellm.proxy.policy_engine.policy_endpoints",
path_prefixes=("/policies",),
),
LazyFeature(
name="policy_resolve",
module_path="litellm.proxy.policy_engine.policy_resolve_endpoints",
path_prefixes=("/policies/resolve", "/policies/attachments/estimate-impact"),
),
LazyFeature(
name="agents",
module_path="litellm.proxy.agent_endpoints.endpoints",
path_prefixes=("/v1/agents", "/agents", "/agent/"),
),
LazyFeature(
name="a2a",
module_path="litellm.proxy.agent_endpoints.a2a_endpoints",
path_prefixes=("/a2a", "/v1/a2a"),
),
LazyFeature(
name="vector_stores",
module_path="litellm.proxy.vector_store_endpoints.endpoints",
path_prefixes=("/v1/vector_stores", "/vector_stores", "/v1/indexes"),
),
LazyFeature(
name="vector_store_management",
module_path="litellm.proxy.vector_store_endpoints.management_endpoints",
# Trailing slash to avoid matching /vector_stores/... (vector_stores).
path_prefixes=("/vector_store/", "/v1/vector_store/"),
),
LazyFeature(
name="vector_store_files",
# Routes appear under both /v1/vector_stores/{id}/files and the
# un-versioned form, so both prefixes must trigger the load.
module_path="litellm.proxy.vector_store_files_endpoints.endpoints",
path_prefixes=("/v1/vector_stores", "/vector_stores"),
),
LazyFeature(
name="tools",
module_path="litellm.proxy.management_endpoints.tool_management_endpoints",
path_prefixes=("/v1/tool", "/tool"),
),
LazyFeature(
name="search_tools",
module_path="litellm.proxy.search_endpoints.search_tool_management",
path_prefixes=("/search_tools",),
),
# mcp_management owns most /v1/mcp/* admin routes; mcp_app is the mounted
# streaming sub-app at /mcp.
LazyFeature(
name="mcp_management",
module_path="litellm.proxy.management_endpoints.mcp_management_endpoints",
path_prefixes=("/v1/mcp/",),
),
LazyFeature(
# Also serves /.well-known/oauth-* (OAuth metadata discovery).
# No /mcp/oauth prefix here: the mounted /mcp sub-app would
# shadow it, and there are no actual routes there anyway.
name="mcp_byok_oauth",
module_path="litellm.proxy._experimental.mcp_server.byok_oauth_endpoints",
path_prefixes=("/v1/mcp/oauth", "/.well-known/oauth-"),
),
LazyFeature(
# Serves OAuth dance endpoints (/authorize, /token, /callback,
# /register) plus several /.well-known/ discovery URLs at the proxy
# root — needed for MCP-over-OAuth flows even before /mcp is hit.
name="mcp_discoverable",
module_path="litellm.proxy._experimental.mcp_server.discoverable_endpoints",
path_prefixes=(
"/.well-known/oauth-",
"/.well-known/openid-configuration",
"/.well-known/jwks.json",
"/authorize",
"/token",
"/callback",
"/register",
),
# Catches the /{mcp_server_name}/authorize|token|register variants.
path_suffixes=("/authorize", "/token", "/register"),
),
LazyFeature(
name="mcp_rest",
module_path="litellm.proxy._experimental.mcp_server.rest_endpoints",
path_prefixes=("/mcp-rest",),
),
LazyFeature(
# Hardcoded /mcp matches BASE_MCP_ROUTE; importing the constant
# here would defeat lazy loading.
name="mcp_app",
module_path="litellm.proxy._experimental.mcp_server.server",
path_prefixes=("/mcp",),
register_fn=_mount_app("/mcp", attr_name="app"),
persistent_swagger_stub=True,
),
LazyFeature(
name="config_overrides",
module_path="litellm.proxy.management_endpoints.config_override_endpoints",
path_prefixes=("/config_overrides",),
),
LazyFeature(
name="realtime",
module_path="litellm.proxy.realtime_endpoints.endpoints",
path_prefixes=("/openai/v1/realtime", "/v1/realtime", "/realtime"),
),
LazyFeature(
name="anthropic_passthrough",
module_path="litellm.proxy.anthropic_endpoints.endpoints",
path_prefixes=("/v1/messages", "/anthropic", "/api/event_logging"),
),
LazyFeature(
name="anthropic_skills",
module_path="litellm.proxy.anthropic_endpoints.skills_endpoints",
path_prefixes=("/v1/skills", "/skills"),
),
LazyFeature(
name="langfuse_passthrough",
module_path="litellm.proxy.vertex_ai_endpoints.langfuse_endpoints",
path_prefixes=("/langfuse",),
),
LazyFeature(
name="evals",
module_path="litellm.proxy.openai_evals_endpoints.endpoints",
path_prefixes=("/v1/evals", "/evals"),
),
LazyFeature(
name="claude_code_marketplace",
module_path="litellm.proxy.anthropic_endpoints.claude_code_endpoints",
path_prefixes=("/claude-code",),
register_fn=_include_router("claude_code_marketplace_router"),
),
LazyFeature(
name="scim",
module_path="litellm.proxy.management_endpoints.scim.scim_v2",
path_prefixes=("/scim",),
register_fn=_include_router("scim_router"),
),
LazyFeature(
name="cloudzero",
module_path="litellm.proxy.spend_tracking.cloudzero_endpoints",
path_prefixes=("/cloudzero",),
),
LazyFeature(
name="vantage",
module_path="litellm.proxy.spend_tracking.vantage_endpoints",
path_prefixes=("/vantage",),
),
LazyFeature(
name="usage_ai",
module_path="litellm.proxy.management_endpoints.usage_endpoints",
path_prefixes=("/usage/ai",),
),
LazyFeature(
name="prompts",
module_path="litellm.proxy.prompts.prompt_endpoints",
path_prefixes=("/prompts", "/utils/dotprompt_json_converter"),
),
LazyFeature(
name="jwt_mappings",
module_path="litellm.proxy.management_endpoints.jwt_key_mapping_endpoints",
path_prefixes=("/jwt/key/mapping",),
),
LazyFeature(
name="compliance",
module_path="litellm.proxy.management_endpoints.compliance_endpoints",
path_prefixes=("/compliance",),
),
LazyFeature(
name="access_groups",
module_path="litellm.proxy.management_endpoints.access_group_endpoints",
path_prefixes=("/access_group", "/v1/access_group", "/v1/unified_access_group"),
),
)
class LazyFeatureMiddleware:
"""ASGI middleware that imports + registers a feature router on first
matching request. Idempotent; once loaded, subsequent requests skip."""
def __init__(
self,
app,
fastapi_app: "FastAPI",
features: Tuple[LazyFeature, ...] = LAZY_FEATURES,
):
self.app = app
self._fastapi_app = fastapi_app
self._features = features
# Loaded set / per-feature locks live on app.state so the warm endpoint
# and the middleware share them — preventing duplicate registrations
# when both paths fire for the same feature.
if not hasattr(fastapi_app.state, "lazy_loaded"):
fastapi_app.state.lazy_loaded = set()
fastapi_app.state.lazy_locks = {}
@property
def _loaded(self) -> set:
return self._fastapi_app.state.lazy_loaded
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
# Short-circuit once every feature has loaded.
if scope["type"] in ("http", "websocket") and len(self._loaded) < len(
self._features
):
path = scope.get("path", "")
for feat in self._features:
if feat.module_path in self._loaded:
continue
if any(path.startswith(p) for p in feat.path_prefixes) or any(
path.endswith(s) for s in feat.path_suffixes
):
await _force_load(self._fastapi_app, feat)
await self.app(scope, receive, send)
async def _force_load(app: "FastAPI", feat: LazyFeature) -> bool:
"""Import + register a lazy feature exactly once per (app, module).
Shared by the middleware and the /lazy/warm endpoint."""
if not hasattr(app.state, "lazy_loaded"):
app.state.lazy_loaded = set()
app.state.lazy_locks = {}
lock = app.state.lazy_locks.setdefault(feat.module_path, asyncio.Lock())
async with lock:
if feat.module_path in app.state.lazy_loaded:
return False
try:
# Import on a thread (heavy modules take 1-3 s). register_fn
# mutates app.router.routes, so it stays on the loop thread.
loop = asyncio.get_running_loop()
module = await loop.run_in_executor(
None, importlib.import_module, feat.module_path
)
feat.register_fn(app, module)
app.state.lazy_loaded.add(feat.module_path)
app.openapi_schema = None
verbose_proxy_logger.info(
"Lazy-loaded optional feature %r (module: %s)",
feat.name,
feat.module_path,
)
return True
except Exception as exc:
# Mark loaded anyway so we don't retry on every request.
app.state.lazy_loaded.add(feat.module_path)
verbose_proxy_logger.warning(
"Failed to lazy-load optional feature %r (module: %s): %s. "
"This feature's endpoints will return 404 until restart.",
feat.name,
feat.module_path,
exc,
)
return False
def attach_lazy_features(app: "FastAPI") -> None:
app.include_router(_make_warmup_router(app))
app.add_middleware(LazyFeatureMiddleware, fastapi_app=app)
def _make_warmup_router(app: "FastAPI") -> "APIRouter":
"""POST /lazy/warm/{name}: load a feature and return its partial openapi
so the Swagger plugin can merge in-place without a full /openapi.json refetch.
Requires auth anyone who can hit the proxy can already trigger the same
imports by sending a real request to a feature's prefix, but gating this
debug endpoint avoids unauthenticated callers forcing the import chain."""
from fastapi import APIRouter, Depends, HTTPException
from fastapi.openapi.utils import get_openapi
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
router = APIRouter()
@router.post(
"/lazy/warm/{name}",
include_in_schema=False,
dependencies=[Depends(user_api_key_auth)],
)
async def warm(name: str):
feat = next((f for f in LAZY_FEATURES if f.name == name), None)
if feat is None:
raise HTTPException(404, f"unknown lazy feature: {name}")
if feat.persistent_swagger_stub:
return {"stub_path": None, "paths": {}, "components": {"schemas": {}}}
await _force_load(app, feat)
feat_routes = [
r
for r in app.routes
if any(getattr(r, "path", "").startswith(p) for p in feat.path_prefixes)
]
full = get_openapi(title=app.title, version=app.version, routes=feat_routes)
# Force all operations under one tag so they group under a single Swagger
# section — many lazy modules tag routes inconsistently.
for path_ops in full.get("paths", {}).values():
for op in path_ops.values():
if isinstance(op, dict):
op["tags"] = [feat.name]
return {
"stub_path": feat.path_prefixes[0],
"paths": full.get("paths", {}),
"components": {"schemas": full.get("components", {}).get("schemas", {})},
}
return router
def inject_lazy_stubs(schema: Dict) -> Dict:
"""Inject openapi entries for unloaded features. Uses the snapshot file
when available (full route info), otherwise falls back to a single
placeholder per feature. Any failure logs and returns the schema unchanged
so /openapi.json never 500s on a cosmetic injection bug."""
try:
from litellm.proxy._lazy_openapi_snapshot import load_snapshot
snapshot = load_snapshot()
paths = schema.setdefault("paths", {})
schemas = schema.setdefault("components", {}).setdefault("schemas", {})
for feat in LAZY_FEATURES:
if feat.module_path in sys.modules and not feat.persistent_swagger_stub:
continue
fragment = (snapshot or {}).get(feat.name)
if fragment:
for p, ops in fragment.get("paths", {}).items():
paths.setdefault(p, ops)
for name, sch in (
fragment.get("components", {}).get("schemas", {}).items()
):
schemas.setdefault(name, sch)
continue
prefix = feat.path_prefixes[0]
if prefix in paths:
continue
paths[prefix] = {
"get": {
"tags": [feat.name],
"summary": feat.name,
"responses": {"200": {"description": "OK"}},
}
}
except Exception as exc:
verbose_proxy_logger.warning("inject_lazy_stubs failed: %s", exc)
return schema
def lazy_tag_to_prefix() -> Dict[str, str]:
"""feature.name -> first prefix, used by the Swagger warmup JS plugin.
Returns empty when the snapshot is loaded the plugin is unnecessary
because /openapi.json already has full route info."""
from litellm.proxy._lazy_openapi_snapshot import load_snapshot
if load_snapshot():
return {}
return {
feat.name: feat.path_prefixes[0]
for feat in LAZY_FEATURES
if not feat.persistent_swagger_stub
}
File diff suppressed because it is too large Load Diff
+70
View File
@@ -0,0 +1,70 @@
"""
Per-feature OpenAPI snapshot for lazy-loaded routers.
The committed JSON is generated by `python -m litellm.proxy._lazy_openapi_snapshot`
and consumed at runtime so /openapi.json can show full route info for unloaded
features without importing them. CI verifies the file is current and surfaces
any drift as a neutral check.
"""
import json
import sys
from pathlib import Path
from typing import Dict, Optional
SNAPSHOT_FILE = Path(__file__).parent / "_lazy_openapi_snapshot.json"
def load_snapshot() -> Optional[Dict[str, Dict]]:
if not SNAPSHOT_FILE.exists():
return None
try:
with SNAPSHOT_FILE.open() as f:
return json.load(f)
except (json.JSONDecodeError, OSError):
return None
def generate_snapshot() -> Dict[str, Dict]:
import importlib
from fastapi.openapi.utils import get_openapi
from litellm.proxy._lazy_features import LAZY_FEATURES
from litellm.proxy.proxy_server import app
for feat in LAZY_FEATURES:
if feat.module_path in sys.modules:
continue
try:
module = importlib.import_module(feat.module_path)
feat.register_fn(app, module)
except Exception as exc:
sys.stderr.write(f"warning: skip {feat.name}: {exc}\n")
fragments: Dict[str, Dict] = {}
for feat in LAZY_FEATURES:
feat_routes = [
r
for r in app.routes
if any(getattr(r, "path", "").startswith(p) for p in feat.path_prefixes)
]
if not feat_routes:
continue
full = get_openapi(title=app.title, version=app.version, routes=feat_routes)
# Group all of a feature's routes under one tag.
for path_ops in full.get("paths", {}).values():
for op in path_ops.values():
if isinstance(op, dict):
op["tags"] = [feat.name]
fragments[feat.name] = {
"paths": full.get("paths", {}),
"components": {"schemas": full.get("components", {}).get("schemas", {})},
}
return fragments
if __name__ == "__main__":
fragments = generate_snapshot()
SNAPSHOT_FILE.write_text(json.dumps(fragments, indent=2, sort_keys=True) + "\n")
sys.stdout.write(f"wrote {len(fragments)} feature fragments to {SNAPSHOT_FILE}\n")
+11 -4
View File
@@ -17,6 +17,9 @@ from typing_extensions import Required, TypedDict
from litellm._uuid import uuid
from litellm.constants import MCP_STDIO_ALLOWED_COMMANDS
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
validate_no_callback_env_reference,
)
from litellm.types.integrations.slack_alerting import AlertType
from litellm.types.llms.openai import (
AllMessageValues,
@@ -904,6 +907,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase):
agents: Optional[List[str]] = None
agent_access_groups: Optional[List[str]] = None
models: Optional[List[str]] = None
search_tools: Optional[List[str]] = None
class BudgetLimitEntry(LiteLLMPydanticObjectBase):
@@ -1868,8 +1872,10 @@ class AddTeamCallback(LiteLLMPydanticObjectBase):
raise ValueError(
f"Invalid callback variable: {key}. Must be one of {valid_keys}"
)
if not isinstance(value, str):
callback_vars[key] = str(value)
callback_vars[key] = str(value)
validate_no_callback_env_reference(
key, callback_vars[key], source="key/team callback metadata"
)
return values
@@ -1934,6 +1940,7 @@ class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase):
agent_access_groups: Optional[List[str]] = []
mcp_toolsets: Optional[List[str]] = None
blocked_tools: Optional[List[str]] = []
search_tools: Optional[List[str]] = []
class LiteLLM_TeamTable(TeamBase):
@@ -2154,8 +2161,8 @@ class PassThroughGenericEndpoint(LiteLLMPydanticObjectBase):
description="The USD cost per request to the target endpoint. This is used to calculate the cost of the request to the target endpoint.",
)
auth: bool = Field(
default=False,
description="Whether authentication is required for the pass-through endpoint. If True, requests to the endpoint will require a valid LiteLLM API key.",
default=True,
description="Whether authentication is required for the pass-through endpoint. Defaults to True so a pass-through silently created without an explicit value still requires a valid LiteLLM API key — set to False only if the endpoint is meant to be a public forwarder (e.g. an unauthenticated webhook target).",
)
guardrails: Optional[PassThroughGuardrailsConfig] = Field(
default=None,
+114 -2
View File
@@ -374,7 +374,7 @@ def _guardrail_modification_check(
coerced = _coerce_to_dict(container)
if coerced is None:
return False
return any(coerced.get(key) for key in _GUARDRAIL_MODIFICATION_KEYS)
return any(key in coerced for key in _GUARDRAIL_MODIFICATION_KEYS)
# Check both metadata keys — callers can populate either depending on the
# endpoint. Cover the top-level too so root-level injection is rejected.
@@ -915,7 +915,8 @@ async def get_team_member_default_budget(
Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"].
This budget is applied to team members whose TeamMembership row has no
linked budget. Results are cached for performance.
linked budget, or whose linked budget has max_budget=NULL. Results are
cached for performance.
Args:
budget_id: The budget_id pulled from team.metadata["team_member_budget_id"]
@@ -2962,6 +2963,116 @@ async def can_user_call_model(
)
def _search_tool_names_from_object_permission(
object_permission: Optional[LiteLLM_ObjectPermissionTable],
) -> List[str]:
"""Return allowlisted search tool names from object_permission (empty = unrestricted)."""
if object_permission is None:
return []
raw = object_permission.search_tools
if not raw:
return []
return list(raw)
def _can_object_call_search_tool(
search_tool_name: str,
allowed_search_tools: List[str],
object_type: Literal["key", "team", "project"],
) -> Literal[True]:
"""
Check if an object (key/team/project) can access a specific search tool.
Similar to _can_object_call_model but for search tools.
Args:
search_tool_name: The search tool being requested
allowed_search_tools: List of allowed search tool names for this object
object_type: Type of object for error messaging
Returns:
True if access is allowed
Raises:
ProxyException if access is denied
"""
# Empty list means all search tools are allowed
if not allowed_search_tools:
return True
# Check if the search tool is in the allowlist
if search_tool_name in allowed_search_tools:
return True
# Access denied
raise ProxyException(
message=f"{object_type.capitalize()} not allowed to access search tool: {search_tool_name}. "
f"Allowed search tools: {allowed_search_tools}",
type=ProxyErrorTypes.key_model_access_denied,
param="search_tool_name",
code=status.HTTP_403_FORBIDDEN,
)
async def can_key_call_search_tool(
search_tool_name: str,
valid_token: UserAPIKeyAuth,
) -> Literal[True]:
"""
Check if a key can access a specific search tool.
Similar to can_key_call_model but for search tools.
Args:
search_tool_name: The search tool being requested
valid_token: The authenticated key
Returns:
True if access is allowed
Raises:
ProxyException if access is denied
"""
return _can_object_call_search_tool(
search_tool_name=search_tool_name,
allowed_search_tools=_search_tool_names_from_object_permission(
valid_token.object_permission
),
object_type="key",
)
async def can_team_call_search_tool(
search_tool_name: str,
team_object: Optional[LiteLLM_TeamTable],
) -> Literal[True]:
"""
Check if a team can access a specific search tool.
Similar to can_team_access_model but for search tools.
Args:
search_tool_name: The search tool being requested
team_object: The team object
Returns:
True if access is allowed
Raises:
ProxyException if access is denied
"""
if team_object is None:
return True
return _can_object_call_search_tool(
search_tool_name=search_tool_name,
allowed_search_tools=_search_tool_names_from_object_permission(
team_object.object_permission
),
object_type="team",
)
async def is_valid_fallback_model(
model: str,
llm_router: Optional[Router],
@@ -3293,6 +3404,7 @@ async def _check_team_member_budget(
if (
team_membership is not None
and team_membership.litellm_budget_table is not None
and team_membership.litellm_budget_table.max_budget is not None
):
team_member_budget = team_membership.litellm_budget_table.max_budget
else:
+65 -10
View File
@@ -6,9 +6,11 @@ from typing import Any, List, Optional, Tuple
from fastapi import HTTPException, Request, status
import litellm
from litellm import Router, provider_list
from litellm._logging import verbose_proxy_logger
from litellm.constants import STANDARD_CUSTOMER_ID_HEADERS
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
from litellm.proxy._types import *
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS
@@ -53,6 +55,12 @@ def _check_valid_ip(
def check_complete_credentials(request_body: dict) -> bool:
"""
if 'api_base' in request body. Check if complete credentials given. Prevent malicious attacks.
Supplying an ``api_key`` is necessary but not sufficient: even with
credentials supplied, an ``api_base`` / ``base_url`` that resolves to a
private/internal/cloud-metadata address would still allow the proxy to
be used as an SSRF pivot. Validate any URL fields here so the gate
can't be bypassed with ``api_key=anything`` plus a malicious target.
"""
given_model: Optional[str] = None
@@ -70,10 +78,27 @@ def check_complete_credentials(request_body: dict) -> bool:
return False
api_key_value = request_body.get("api_key")
if api_key_value and isinstance(api_key_value, str) and api_key_value.strip():
return True
if not (api_key_value and isinstance(api_key_value, str) and api_key_value.strip()):
return False
return False
# ``validate_url`` itself doesn't consult the toggle; ``safe_get`` /
# ``async_safe_get`` do. Mirror that here so admins who explicitly
# disabled URL validation (e.g. for an internal Ollama endpoint they
# accept the SSRF risk for) aren't blocked at the proxy boundary.
if getattr(litellm, "user_url_validation", False):
for url_field in ("api_base", "base_url"):
url_value = request_body.get(url_field)
if not url_value or not isinstance(url_value, str):
continue
try:
validate_url(url_value)
except SSRFError as e:
raise ValueError(
f"Rejected request: client-side {url_field}={url_value!r} "
f"is rejected by the SSRF guard ({e})."
)
return True
def check_regex_or_str_match(request_body_value: Any, regex_str: str) -> bool:
@@ -159,15 +184,42 @@ def is_request_body_safe(
"aws_web_identity_token",
"aws_role_name",
"vertex_credentials",
# Endpoint-targeting fields that retarget the outbound request or
# an observability callback. An attacker-controlled value either
# exfiltrates the request payload (incl. messages + admin-set
# tokens) to the attacker's host, or coerces the proxy into
# authenticating against the attacker's host with admin secrets.
"aws_bedrock_runtime_endpoint",
"langsmith_base_url",
"langfuse_host",
"posthog_host",
"braintrust_host",
"slack_webhook_url",
# Provider-specific endpoint overrides that flow into the outbound
# request via ``optional_params``. Same threat as ``api_base``:
# ``s3_endpoint_url`` redirects Bedrock file uploads to attacker
# S3; ``sagemaker_base_url`` redirects all SageMaker traffic;
# ``deployment_url`` redirects SAP deployments.
"s3_endpoint_url",
"sagemaker_base_url",
"deployment_url",
]
# The blocklist is enforced unconditionally. Legitimate clientside
# credential / endpoint passthrough goes through one of the two
# explicit admin opt-ins (``general_settings.allow_client_side_credentials``
# proxy-wide or ``configurable_clientside_auth_params`` per deployment).
# Historically there was a third, *implicit*, *caller-controlled* path:
# ``check_complete_credentials`` returned True when the caller supplied
# any non-empty ``api_key``, which made the entire blocklist a no-op.
# That bypass turned every missing entry on the blocklist into an
# exploitable SSRF / credential-exfil hole — see GHSA-jh89-88fc-qrfp,
# GHSA-3frq-6r6h-7j64, and the chain of veria-admin findings (Dv_m860l,
# b_yRJeQ5, stN90yjP, LBlyOAc8, U2TD78kg). Removed: the blocklist now
# has a single, predictable failure mode for missing entries (a 400),
# not a credential leak.
for param in banned_params:
if (
param in request_body
and not check_complete_credentials( # allow client-credentials to be passed to proxy
request_body=request_body
)
):
if param in request_body:
if general_settings.get("allow_client_side_credentials") is True:
return True
elif (
@@ -182,7 +234,10 @@ def is_request_body_safe(
return True
raise ValueError(
f"Rejected Request: {param} is not allowed in request body. "
"Enable with `general_settings::allow_client_side_credentials` on proxy config.yaml. "
"Clientside passthrough requires explicit admin opt-in via "
"either `general_settings.allow_client_side_credentials = true` "
"(proxy-wide) or `configurable_clientside_auth_params` on the "
"deployment in your proxy config.yaml. "
"Relevant Issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997",
)
+12 -2
View File
@@ -21,6 +21,7 @@ import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.caching import DualCache
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
from litellm.proxy._types import *
@@ -472,7 +473,12 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints(
for endpoint in pass_through_endpoints:
if isinstance(endpoint, dict) and endpoint.get("path", "") == route:
## IF AUTH DISABLED
if endpoint.get("auth") is not True:
# Default to True: a config dict with no ``auth`` key
# otherwise produced an unauthenticated forwarder. The
# Pydantic ``PassThroughGenericEndpoint.auth`` default
# is also True, but raw config dicts skip that path —
# so this runtime check has to default to True too.
if endpoint.get("auth", True) is not True:
return UserAPIKeyAuth()
## IF AUTH ENABLED
### IF CUSTOM PARSER REQUIRED
@@ -1119,10 +1125,14 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
)
if is_master_key_valid:
# Substitute a stable alias for the raw master key so neither the
# master key nor its hash propagates into spend logs, Prometheus
# /metrics labels, audit trails, rate-limit buckets, or any other
# downstream consumer of UserAPIKeyAuth.api_key.
_user_api_key_obj = await _return_user_api_key_auth_obj(
user_obj=None,
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key=master_key,
api_key=LITELLM_PROXY_MASTER_KEY_ALIAS,
parent_otel_span=parent_otel_span,
valid_token_dict={
**end_user_params,
@@ -474,6 +474,10 @@ async def retrieve_batch( # noqa: PLR0915
)
# Fix: The helper sets "file_id" but we need "batch_id"
data["batch_id"] = data.pop("file_id", original_batch_id)
# Provider-config providers (e.g. bedrock) require `model` in kwargs
# so litellm.aretrieve_batch can load BedrockBatchesConfig. Without
# it the call falls into the legacy provider switch and 400s.
data["model"] = model_from_id
# Retrieve batch using model credentials
response = await litellm.aretrieve_batch(
+17 -16
View File
@@ -313,23 +313,24 @@ sequenceDiagram
participant Proxy as LiteLLM Proxy
participant SSO as SSO Provider
CLI->>CLI: Generate key ID (sk-uuid)
CLI->>Browser: Open /sso/key/generate?source=litellm-cli&key=sk-uuid
CLI->>Proxy: POST /sso/cli/start
Proxy->>CLI: Return login_id, poll_secret, user_code
CLI->>Browser: Open /sso/key/generate?source=litellm-cli&key=login_id
Browser->>Proxy: GET /sso/key/generate?source=litellm-cli&key=sk-uuid
Proxy->>Proxy: Set cli_state = litellm-session-token:sk-uuid
Proxy->>SSO: Redirect with state=litellm-session-token:sk-uuid
Browser->>Proxy: GET /sso/key/generate?source=litellm-cli&key=login_id
Proxy->>Proxy: Set cli_state = litellm-session-token:login_id
Proxy->>SSO: Redirect with state=litellm-session-token:login_id
SSO->>Browser: Show login page
Browser->>SSO: User authenticates
SSO->>Proxy: Redirect to /sso/callback?state=litellm-session-token:sk-uuid
SSO->>Proxy: Redirect to /sso/callback?state=litellm-session-token:login_id
Proxy->>Proxy: Check if state starts with "litellm-session-token:"
Proxy->>Proxy: Generate API key with ID=sk-uuid
Proxy->>Browser: Show success page
Proxy->>Browser: Prompt for user_code
Browser->>Proxy: POST /sso/cli/complete/login_id
CLI->>Proxy: Poll /sso/cli/poll/sk-uuid
Proxy->>CLI: Return {"status": "ready", "key": "sk-uuid"}
CLI->>Proxy: Poll /sso/cli/poll/login_id with poll_secret header
Proxy->>CLI: Return {"status": "ready", "key": "jwt"}
CLI->>CLI: Save key to ~/.litellm/token.json
```
@@ -343,13 +344,13 @@ The CLI provides three authentication commands:
### Authentication Flow Steps
1. **Generate Session ID**: CLI generates a unique key ID (`sk-{uuid}`)
2. **Open Browser**: CLI opens browser to `/sso/key/generate` with CLI source and key parameters
3. **SSO Redirect**: Proxy sets the formatted state (`litellm-session-token:sk-uuid`) as OAuth state parameter and redirects to SSO provider
1. **Start Session**: CLI creates a short-lived login session with `/sso/cli/start`
2. **Open Browser**: CLI opens browser to `/sso/key/generate` with CLI source and login ID parameters
3. **SSO Redirect**: Proxy sets the formatted state (`litellm-session-token:{login_id}`) as OAuth state parameter and redirects to SSO provider
4. **User Authentication**: User completes SSO authentication in browser
5. **Callback Processing**: SSO provider redirects back to proxy with state parameter
6. **Key Generation**: Proxy detects CLI login (state starts with "litellm-session-token:") and generates API key with pre-specified ID
7. **Polling**: CLI polls `/sso/cli/poll/{key_id}` endpoint until key is ready
6. **User Code Verification**: Browser confirms the verification code shown in the CLI
7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready
8. **Token Storage**: CLI saves the authentication token to `~/.litellm/token.json`
### Benefits of This Approach
@@ -357,7 +358,7 @@ The CLI provides three authentication commands:
- **No Local Server**: No need to run a local callback server
- **Standard OAuth**: Uses OAuth 2.0 state parameter correctly
- **Remote Compatible**: Works with remote proxy servers
- **Secure**: Uses UUID session identifiers
- **Secure**: Keeps the polling secret out of the browser handoff
- **Simple Setup**: No additional OAuth redirect URL configuration needed
### Token Storage
+39 -18
View File
@@ -5,6 +5,7 @@ import time
import webbrowser
from pathlib import Path
from typing import Any, Dict, List, Optional
from urllib.parse import urlencode
import click
import requests
@@ -241,7 +242,7 @@ def prompt_team_selection(teams: List[Dict[str, Any]]) -> Optional[Dict[str, Any
def prompt_team_selection_fallback(
teams: List[Dict[str, Any]]
teams: List[Dict[str, Any]],
) -> Optional[Dict[str, Any]]:
"""Fallback team selection for non-interactive environments"""
if not teams:
@@ -279,6 +280,7 @@ def prompt_team_selection_fallback(
def _poll_for_ready_data(
url: str,
*,
headers: Optional[Dict[str, str]] = None,
total_timeout: int = 300,
poll_interval: int = 2,
request_timeout: int = 10,
@@ -291,7 +293,10 @@ def _poll_for_ready_data(
) -> Optional[Dict[str, Any]]:
for attempt in range(total_timeout // poll_interval):
try:
response = requests.get(url, timeout=request_timeout)
request_kwargs: Dict[str, Any] = {"timeout": request_timeout}
if headers is not None:
request_kwargs["headers"] = headers
response = requests.get(url, **request_kwargs)
if response.status_code == 200:
data = response.json()
status = data.get("status")
@@ -346,7 +351,23 @@ def _normalize_teams(teams, team_details):
return []
def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
def _start_cli_sso_flow(base_url: str) -> Dict[str, Any]:
response = requests.post(f"{base_url}/sso/cli/start", timeout=10)
response.raise_for_status()
data = response.json()
required_fields = ("login_id", "poll_secret", "user_code")
if not all(isinstance(data.get(field), str) for field in required_fields):
raise ValueError("Invalid CLI SSO start response")
return data
def _get_cli_sso_poll_headers(poll_secret: str) -> Dict[str, str]:
return {"x-litellm-cli-poll-secret": poll_secret}
def _poll_for_authentication(
base_url: str, key_id: str, poll_secret: str
) -> Optional[dict]:
"""
Poll the server for authentication completion and handle team selection.
@@ -356,6 +377,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
poll_url = f"{base_url}/sso/cli/poll/{key_id}"
data = _poll_for_ready_data(
poll_url,
headers=_get_cli_sso_poll_headers(poll_secret),
pending_message="Still waiting for authentication...",
)
if not data:
@@ -373,6 +395,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
jwt_with_team = _handle_team_selection_during_polling(
base_url=base_url,
key_id=key_id,
poll_secret=poll_secret,
teams=normalized_teams,
)
@@ -410,7 +433,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
def _handle_team_selection_during_polling(
base_url: str, key_id: str, teams: List[Dict[str, Any]]
base_url: str, key_id: str, poll_secret: str, teams: List[Dict[str, Any]]
) -> Optional[str]:
"""
Handle team selection and re-poll with selected team_id.
@@ -441,6 +464,7 @@ def _handle_team_selection_during_polling(
poll_url = f"{base_url}/sso/cli/poll/{key_id}?team_id={team_id}"
data = _poll_for_ready_data(
poll_url,
headers=_get_cli_sso_poll_headers(poll_secret),
pending_message="Still waiting for team authentication...",
other_status_message="Waiting for team authentication to complete...",
http_error_log_every=10,
@@ -514,29 +538,24 @@ def _render_and_prompt_for_team_selection(teams: List[Dict[str, Any]]) -> Option
@click.pass_context
def login(ctx: click.Context):
"""Login to LiteLLM proxy using SSO authentication"""
from litellm._uuid import uuid
from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER
from litellm.proxy.client.cli.interface import show_commands
base_url = ctx.obj["base_url"]
# Check if we have an existing key to regenerate
existing_key = get_stored_api_key()
# Generate unique key ID for this login session
key_id = f"sk-{str(uuid.uuid4())}"
try:
# Construct SSO login URL with CLI source and pre-generated key
sso_url = f"{base_url}/sso/key/generate?source={LITELLM_CLI_SOURCE_IDENTIFIER}&key={key_id}"
cli_sso_flow = _start_cli_sso_flow(base_url=base_url)
key_id = cli_sso_flow["login_id"]
poll_secret = cli_sso_flow["poll_secret"]
user_code = cli_sso_flow["user_code"]
# If we have an existing key, include it as a parameter to the login endpoint
# The server will encode it in the OAuth state parameter for the SSO flow
if existing_key:
sso_url += f"&existing_key={existing_key}"
sso_url = f"{base_url}/sso/key/generate?" + urlencode(
{"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": key_id}
)
click.echo(f"Opening browser to: {sso_url}")
click.echo("Please complete the SSO authentication in your browser...")
click.echo(f"Verification code: {user_code}")
click.echo(f"Session ID: {key_id}")
# Open browser
@@ -545,7 +564,9 @@ def login(ctx: click.Context):
# Poll for authentication completion
click.echo("Waiting for authentication...")
auth_result = _poll_for_authentication(base_url=base_url, key_id=key_id)
auth_result = _poll_for_authentication(
base_url=base_url, key_id=key_id, poll_secret=poll_secret
)
if auth_result:
api_key = auth_result["api_key"]
+66 -6
View File
@@ -619,6 +619,67 @@ class ProxyBaseLLMRequestProcessing:
verbose_proxy_logger.error(f"Error setting custom headers: {e}")
return {}
@staticmethod
async def build_litellm_proxy_success_headers_from_llm_response(
*,
response: Any,
request_data: dict,
request: Request,
user_api_key_dict: UserAPIKeyAuth,
logging_obj: LiteLLMLoggingObj,
version: Optional[str],
proxy_logging_obj: ProxyLogging,
) -> Dict[str, str]:
"""
Build LiteLLM proxy response headers for routes that call the LLM directly
(e.g. Google native :generateContent) instead of base_process_llm_request.
"""
if isinstance(response, dict):
hidden_params = response.get("_hidden_params") or {}
else:
hidden_params = getattr(response, "_hidden_params", None) or {}
if not isinstance(hidden_params, dict):
hidden_params = {}
model_id = ProxyBaseLLMRequestProcessing._get_model_id_from_response(
hidden_params, request_data
)
cache_key = hidden_params.get("cache_key", None) or ""
api_base = hidden_params.get("api_base", None) or ""
response_cost = hidden_params.get("response_cost", None) or ""
fastest_response_batch_completion = hidden_params.get(
"fastest_response_batch_completion", None
)
additional_headers = hidden_params.get("additional_headers", {}) or {}
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=logging_obj.litellm_call_id,
model_id=model_id,
cache_key=cache_key,
api_base=api_base,
version=version,
response_cost=response_cost,
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
fastest_response_batch_completion=fastest_response_batch_completion,
request_data=request_data,
hidden_params=hidden_params,
litellm_logging_obj=logging_obj,
**additional_headers,
)
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
data=request_data,
user_api_key_dict=user_api_key_dict,
response=response,
request_headers=dict(request.headers),
)
if callback_headers:
custom_headers.update(callback_headers)
return custom_headers
async def common_processing_pre_call_logic(
self,
request: Request,
@@ -875,7 +936,7 @@ class ProxyBaseLLMRequestProcessing:
else:
verbose_proxy_logger.debug(
"Request received by LiteLLM:\n%s",
json.dumps(self.data, indent=4, default=str),
_payload_str,
)
async def base_process_llm_request( # noqa: PLR0915
@@ -1511,9 +1572,7 @@ class ProxyBaseLLMRequestProcessing:
_response = assembled_response
try:
from litellm.proxy.proxy_server import llm_router as _global_llm_router
from litellm.proxy.utils import (
_check_and_merge_model_level_guardrails,
)
from litellm.proxy.utils import _check_and_merge_model_level_guardrails
guardrail_data = _check_and_merge_model_level_guardrails(
data=captured_data, llm_router=_global_llm_router
@@ -1690,11 +1749,12 @@ class ProxyBaseLLMRequestProcessing:
elif isinstance(e, httpx.HTTPStatusError):
# Handle httpx.HTTPStatusError - extract actual error from response
# This matches the original behavior before the refactor in commit 511d435f6f
error_body = await e.response.aread()
http_status_error: httpx.HTTPStatusError = e
error_body = await http_status_error.response.aread()
error_text = error_body.decode("utf-8")
raise HTTPException(
status_code=e.response.status_code,
status_code=http_status_error.response.status_code,
detail={"error": error_text},
)
error_msg = f"{str(e)}"
+14 -3
View File
@@ -14,6 +14,8 @@ from litellm.types.utils import (
blue_color_code = "\033[94m"
reset_color_code = "\033[0m"
TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY = "_pillar_response_headers_trusted"
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
@@ -417,10 +419,19 @@ def get_logging_caching_headers(request_data: Dict) -> Optional[Dict]:
if "semantic-similarity" in _metadata:
headers["x-litellm-semantic-similarity"] = str(_metadata["semantic-similarity"])
is_trusted_pillar_metadata = (
_metadata.get(TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY) is True
)
pillar_headers = _metadata.get("pillar_response_headers")
if isinstance(pillar_headers, dict):
headers.update(pillar_headers)
elif "pillar_flagged" in _metadata:
if is_trusted_pillar_metadata and isinstance(pillar_headers, dict):
headers.update(
{
key: str(value)
for key, value in pillar_headers.items()
if isinstance(key, str) and key.lower().startswith("x-pillar-")
}
)
elif is_trusted_pillar_metadata and "pillar_flagged" in _metadata:
headers["x-pillar-flagged"] = str(_metadata["pillar_flagged"]).lower()
return headers
+85 -35
View File
@@ -52,6 +52,37 @@ class ResetBudgetJob:
### RESET MULTI-WINDOW BUDGETS ###
await self.reset_budget_windows()
@staticmethod
async def _invalidate_spend_counter(counter_key: str) -> None:
"""Zero a spend counter so a DB-row reset takes effect immediately.
Call AFTER the DB write commits. Clearing Redis before the DB
commit opens a window where get_current_spend reads 0 from Redis
while the DB still holds the pre-reset value, allowing bypass.
"""
try:
from litellm.proxy.proxy_server import spend_counter_cache
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key, value=0.0, ttl=60
)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key, value=0.0, ttl=60
)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to reset spend counter %s in Redis: %s. "
"Budget may be over-enforced until counter expires.",
counter_key,
redis_err,
)
except Exception as e:
verbose_proxy_logger.warning(
"Failed to reset spend counter %s: %s", counter_key, e
)
async def reset_budget_for_litellm_team_members(
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
@@ -64,46 +95,30 @@ class ResetBudgetJob:
if budget.budget_id is not None
]
# Reset spend counters for affected team members.
# Reset Redis directly so a transient failure doesn't leave stale
# counters that get_current_spend would read as authoritative.
try:
from litellm.proxy.proxy_server import spend_counter_cache
memberships = await self.prisma_client.db.litellm_teammembership.find_many(
where={"budget_id": {"in": budget_ids}}
)
for m in memberships:
counter_key = f"spend:team_member:{m.user_id}:{m.team_id}"
# Always reset in-memory
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key, value=0.0
)
# Explicitly reset Redis with warning on failure
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key, value=0.0
)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to reset team member spend counter in Redis %s: %s. "
"Budget may be over-enforced until counter expires.",
counter_key,
redis_err,
)
except Exception as e:
memberships = []
verbose_proxy_logger.warning(
"Failed to reset team member spend counters: %s", e
"Failed to fetch team memberships for counter invalidation: %s", e
)
return await self.prisma_client.db.litellm_teammembership.update_many(
update_result = await self.prisma_client.db.litellm_teammembership.update_many(
where={"budget_id": {"in": budget_ids}},
data={
"spend": 0,
},
)
for m in memberships:
await self._invalidate_spend_counter(
f"spend:team_member:{m.user_id}:{m.team_id}"
)
return update_result
async def reset_budget_for_keys_linked_to_budgets(
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
@@ -126,17 +141,36 @@ class ResetBudgetJob:
if not budget_ids:
return
return await self.prisma_client.db.litellm_verificationtoken.update_many(
where={
"budget_id": {"in": budget_ids},
"budget_duration": None, # only keys without their own reset schedule
"spend": {"gt": 0}, # only reset keys that have accumulated spend
},
data={
"spend": 0,
},
where_clause: dict = {
"budget_id": {"in": budget_ids},
"budget_duration": None, # only keys without their own reset schedule
"spend": {"gt": 0}, # only reset keys that have accumulated spend
}
try:
keys = await self.prisma_client.db.litellm_verificationtoken.find_many(
where=where_clause
)
except Exception as e:
keys = []
verbose_proxy_logger.warning(
"Failed to fetch keys for counter invalidation: %s", e
)
update_result = (
await self.prisma_client.db.litellm_verificationtoken.update_many(
where=where_clause,
data={
"spend": 0,
},
)
)
for k in keys:
await self._invalidate_spend_counter(f"spend:key:{k.token}")
return update_result
async def reset_budget_for_litellm_budget_table(self):
"""
Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired
@@ -365,6 +399,10 @@ class ResetBudgetJob:
data_list=updated_keys,
table_name="key",
)
for k in updated_keys:
token = getattr(k, "token", None)
if token:
await self._invalidate_spend_counter(f"spend:key:{token}")
end_time = time.time()
if len(failed_keys) > 0: # If any keys failed to reset
@@ -450,6 +488,12 @@ class ResetBudgetJob:
data_list=updated_users,
table_name="user",
)
for u in updated_users:
user_id = getattr(u, "user_id", None)
if user_id:
await self._invalidate_spend_counter(
f"spend:user:{user_id}"
)
end_time = time.time()
if len(failed_users) > 0: # If any users failed to reset
@@ -541,6 +585,12 @@ class ResetBudgetJob:
data_list=updated_teams,
table_name="team",
)
for t in updated_teams:
team_id = getattr(t, "team_id", None)
if team_id:
await self._invalidate_spend_counter(
f"spend:team:{team_id}"
)
end_time = time.time()
if len(failed_teams) > 0: # If any teams failed to reset
@@ -0,0 +1,52 @@
"""Helpers for unauthenticated logo / favicon endpoints."""
import os
from typing import Optional, Tuple
from litellm._logging import verbose_proxy_logger
LOCAL_IMAGE_HEADER_BYTES = 512
def detect_local_image_media_type(header: bytes) -> Optional[str]:
"""Return a browser image media type for supported local image signatures."""
if header[0:8] == b"\x89PNG\r\n\x1a\n":
return "image/png"
if header[0:4] == b"GIF8" and header[5:6] == b"a":
return "image/gif"
if header[0:3] == b"\xff\xd8\xff":
return "image/jpeg"
if header[0:4] == b"RIFF" and header[8:12] == b"WEBP":
return "image/webp"
if header[0:4] in (b"\x00\x00\x01\x00", b"\x00\x00\x02\x00"):
return "image/x-icon"
return None
def resolve_validated_local_image_path(candidate: str) -> Optional[Tuple[str, str]]:
"""Resolve ``candidate`` only when it is an existing supported image file."""
if not candidate:
return None
try:
resolved = os.path.realpath(os.path.expanduser(candidate))
except (OSError, ValueError):
return None
if not os.path.isfile(resolved):
return None
try:
with open(resolved, "rb") as f:
header = f.read(LOCAL_IMAGE_HEADER_BYTES)
except OSError as exc:
verbose_proxy_logger.debug("Could not read local asset %r: %s", candidate, exc)
return None
media_type = detect_local_image_media_type(header)
if media_type is None:
verbose_proxy_logger.warning(
"Local asset %r is not a supported image file; falling back to default.",
candidate,
)
return None
return resolved, media_type
+137 -1
View File
@@ -1,5 +1,6 @@
from typing import Union
from typing import Any, Awaitable, Callable, Optional, Union
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
ProxyErrorTypes,
@@ -123,3 +124,138 @@ class PrismaDBExceptionHandler:
):
return None
raise e
# Default fallback timeouts when neither the caller nor the prisma_client
# expose `_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds`.
# Match the auth path's existing defaults so behavior is uniform across read paths.
_DEFAULT_RECONNECT_TIMEOUT_SECONDS = 2.0
_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS = 0.1
def _coerce_timeout(value: Any, fallback: float) -> float:
"""Return `value` if it is a real int/float, else `fallback`. Guards
against tests that mock `prisma_client` and leave the timeout slots as
MagicMock instances."""
if isinstance(value, (int, float)) and not isinstance(value, bool):
return float(value)
return fallback
async def call_with_db_reconnect_retry(
prisma_client: Any,
coro_factory: Callable[[], Awaitable[Any]],
*,
reason: str,
timeout_seconds: Optional[float] = None,
lock_timeout_seconds: Optional[float] = None,
) -> Any:
"""Run a Prisma read coroutine with one transport-reconnect-and-retry.
The canonical "self-heal a transient DB transport blip" wrapper used by
`PrismaClient.get_generic_data` and other read paths. Mirrors the inline
pattern in `auth_checks._fetch_key_object_from_db_with_reconnect` so we
have a single implementation rather than three drifting copies.
Behavior:
1. Await `coro_factory()`. On success, return its value.
2. On exception, if it is NOT a transport error (per
`is_database_transport_error`), re-raise data-layer errors like
`UniqueViolationError` mean the DB is reachable, reconnect would be
pointless.
3. If `prisma_client` does not expose `attempt_db_reconnect`, re-raise.
This guards against partial stand-ins / older clients in tests.
4. Call `prisma_client.attempt_db_reconnect(reason=...)`. If it returns
False (cooldown / lock contention / reconnect failure), re-raise.
5. Otherwise await `coro_factory()` a second time and return / propagate
its result. At-most-one retry by construction no infinite loop.
`coro_factory` MUST be a zero-arg callable that returns a fresh awaitable
on each call. Passing an already-awaited coroutine would fail on retry
with `RuntimeError: cannot reuse already awaited coroutine`.
`reason` should follow `<subsystem>_<operation>_<table>_failure` so
telemetry distinguishes between fan-out callers (e.g.
`_update_config_from_db` issues four concurrent reads).
Args:
prisma_client: The `PrismaClient` (or stand-in) that owns
`attempt_db_reconnect` and the `_db_auth_reconnect_*` defaults.
coro_factory: Zero-arg callable returning the read awaitable.
reason: Telemetry tag forwarded to `attempt_db_reconnect`.
timeout_seconds: Optional override for the reconnect cycle timeout.
Defaults to `prisma_client._db_auth_reconnect_timeout_seconds`,
then to 2.0s.
lock_timeout_seconds: Optional override for how long the helper will
wait to acquire the reconnect lock. Defaults to
`prisma_client._db_auth_reconnect_lock_timeout_seconds`, then to
0.1s.
Returns:
Whatever `coro_factory()` returns (on first or second attempt).
Raises:
Whatever `coro_factory()` raises if the failure is not a transport
error, or if the reconnect attempt does not succeed, or if the retry
also fails.
"""
try:
return await coro_factory()
except Exception as first_exc:
if not PrismaDBExceptionHandler.is_database_transport_error(first_exc):
raise
if not hasattr(prisma_client, "attempt_db_reconnect"):
raise
resolved_timeout = _coerce_timeout(
(
timeout_seconds
if timeout_seconds is not None
else getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", None)
),
_DEFAULT_RECONNECT_TIMEOUT_SECONDS,
)
resolved_lock_timeout = _coerce_timeout(
(
lock_timeout_seconds
if lock_timeout_seconds is not None
else getattr(
prisma_client, "_db_auth_reconnect_lock_timeout_seconds", None
)
),
_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS,
)
verbose_proxy_logger.warning(
"DB transport error on read; attempting reconnect-and-retry. reason=%s error=%s",
reason,
first_exc,
)
# Preserve the original transport error in telemetry. If
# `attempt_db_reconnect` itself raises (e.g. lock cancellation, timer
# error, unexpected internal failure), surfacing that exception
# instead of `first_exc` would mask the actual DB transport problem
# in `failure_handler` / `db_exceptions` alerts. Chain the reconnect
# error as the cause for debuggability without losing the original.
try:
did_reconnect = await prisma_client.attempt_db_reconnect(
reason=reason,
timeout_seconds=resolved_timeout,
lock_timeout_seconds=resolved_lock_timeout,
)
except Exception as reconnect_exc:
verbose_proxy_logger.warning(
"DB reconnect attempt raised; preserving original transport error. "
"reason=%s reconnect_error=%s",
reason,
reconnect_exc,
)
raise first_exc from reconnect_exc
if not did_reconnect:
raise
# At most one retry. If the retry also raises a transport error, we
# propagate — repeated reconnect-loops are the watchdog's job, not
# this helper's.
return await coro_factory()
+22 -12
View File
@@ -52,18 +52,25 @@ class PrismaWrapper:
engine = self._original_prisma._engine
process = getattr(engine, "process", None) if engine is not None else None
if process is not None:
return process.pid
pid = process.pid
if isinstance(pid, int):
return pid
except (AttributeError, TypeError):
pass
return 0
@staticmethod
async def _kill_engine_process(pid: int) -> None:
"""Force-kill an orphaned engine subprocess to prevent DB connection pool leaks.
"""Force-kill the engine subprocess to prevent DB connection pool leaks.
Called when disconnect() fails and the old engine process may still be
holding open connections. Sends SIGTERM for graceful shutdown, waits
briefly, then SIGKILL as a backstop.
Called on every reconnect (in `recreate_prisma_client`) to retire the
old query-engine subprocess without invoking prisma-client-py's
synchronous `disconnect()` which blocks the asyncio event loop on
`subprocess.Popen.wait()` for 30-120+ seconds when the engine is
stuck on TCP close.
Sends SIGTERM for graceful shutdown, waits briefly, then SIGKILL as
a backstop.
"""
if pid <= 0:
return
@@ -72,7 +79,7 @@ class PrismaWrapper:
except (ProcessLookupError, PermissionError, OSError):
return # Already dead or inaccessible
verbose_proxy_logger.warning(
"Sent SIGTERM to orphaned prisma-query-engine PID %s after failed disconnect.",
"Sent SIGTERM to prisma-query-engine PID %s during reconnect.",
pid,
)
# Brief wait for graceful shutdown, then force-kill
@@ -217,15 +224,18 @@ class PrismaWrapper:
async def recreate_prisma_client(
self, new_db_url: str, http_client: Optional[Any] = None
):
"""Disconnect and reconnect the Prisma client with a new database URL."""
"""Disconnect and reconnect the Prisma client with a new database URL.
Kills the old engine subprocess directly (SIGTERM SIGKILL) rather than
calling `disconnect()`. prisma-client-py's `disconnect()` calls a
synchronous `subprocess.Popen.wait()` that can freeze the asyncio event
loop for 30-120+ seconds when the engine is stuck on TCP close,
breaking `/health/liveliness` and causing Kubernetes pod restarts.
"""
from prisma import Prisma # type: ignore
old_engine_pid = self._get_engine_pid()
try:
await self._original_prisma.disconnect()
except Exception as e:
verbose_proxy_logger.warning(f"Failed to disconnect Prisma client: {e}")
if old_engine_pid > 0:
await self._kill_engine_process(old_engine_pid)
if http_client is not None:
+9 -5
View File
@@ -129,7 +129,9 @@ class SpendCounterReseed:
"""
lock = await SpendCounterReseed._get_lock(counter_key)
async with lock:
# Re-check after acquiring the lock - another waiter may have warmed it.
# Re-check after acquiring the lock. Skip in-memory on a clean
# Redis miss - in-memory is per-pod-stale.
redis_clean_miss = False
if spend_counter_cache.redis_cache is not None:
try:
val = await spend_counter_cache.redis_cache.async_get_cache(
@@ -137,11 +139,13 @@ class SpendCounterReseed:
)
if val is not None:
return float(val)
redis_clean_miss = True
except Exception:
pass
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
if not redis_clean_miss:
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
db_spend = await SpendCounterReseed.from_db(prisma_client, counter_key)
if db_spend is None:
@@ -149,7 +153,7 @@ class SpendCounterReseed:
# Warm even when 0 so subsequent reads hit cache, not DB.
try:
await spend_counter_cache.async_increment_cache(
key=counter_key, value=db_spend
key=counter_key, value=db_spend, refresh_ttl=True
)
except Exception:
verbose_proxy_logger.exception(
+71 -79
View File
@@ -1,10 +1,6 @@
from datetime import datetime
from fastapi import APIRouter, Depends, Request, Response
from fastapi.responses import ORJSONResponse
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import ORJSONResponse, StreamingResponse
import litellm
from litellm._uuid import uuid
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
@@ -30,11 +26,17 @@ async def google_generate_content(
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
from litellm.proxy.proxy_server import (
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
select_data_generator,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
@@ -42,38 +44,33 @@ async def google_generate_content(
if "model" not in data:
data["model"] = model_name
# Extract generationConfig and pass it as config parameter
generation_config = data.pop("generationConfig", None)
if generation_config:
data["config"] = generation_config
# Add user authentication metadata for cost tracking
data = await add_litellm_data_to_request(
data=data,
request=request,
user_api_key_dict=user_api_key_dict,
proxy_config=proxy_config,
general_settings=general_settings,
version=version,
)
# Create logging object with full request metadata so callbacks (e.g. S3) get user/trace_id
data["litellm_call_id"] = request.headers.get(
"x-litellm-call-id", str(uuid.uuid4())
)
logging_obj, data = litellm.utils.function_setup(
original_function="agenerate_content",
rules_obj=litellm.utils.Rules(),
start_time=datetime.now(),
**data,
)
data["litellm_logging_obj"] = logging_obj
# call router
if llm_router is None:
raise HTTPException(status_code=500, detail="Router not initialized")
response = await llm_router.agenerate_content(**data)
return response
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="agenerate_content",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=model_name,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
@router.post(
@@ -90,57 +87,52 @@ async def google_stream_generate_content(
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
from litellm.proxy.proxy_server import (
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
select_data_generator,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
data = await _read_request_body(request=request)
if "model" not in data:
data["model"] = model_name
data["stream"] = True
data["stream"] = True # enforce streaming for this endpoint
# Extract generationConfig and pass it as config parameter
generation_config = data.pop("generationConfig", None)
if generation_config:
data["config"] = generation_config
# Add user authentication metadata for cost tracking
data = await add_litellm_data_to_request(
data=data,
request=request,
user_api_key_dict=user_api_key_dict,
proxy_config=proxy_config,
general_settings=general_settings,
version=version,
)
# Create logging object with full request metadata so streaming END callbacks (e.g. S3) get user/trace_id
data["litellm_call_id"] = request.headers.get(
"x-litellm-call-id", str(uuid.uuid4())
)
logging_obj, data = litellm.utils.function_setup(
original_function="agenerate_content_stream",
rules_obj=litellm.utils.Rules(),
start_time=datetime.now(),
**data,
)
data["litellm_logging_obj"] = logging_obj
# call router
if llm_router is None:
raise HTTPException(status_code=500, detail="Router not initialized")
response = await llm_router.agenerate_content_stream(**data)
# Check if response is an async iterator (streaming response)
if response is not None and hasattr(response, "__aiter__"):
return StreamingResponse(content=response, media_type="text/event-stream")
return response
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="agenerate_content_stream",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=model_name,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
@router.post(

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