mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-07 11:13:28 +00:00
Merge pull request #26962 from BerriAI/litellm_internal_staging
[Infra] Promote Internal Staging to main
This commit is contained in:
@@ -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,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
|
||||
|
||||
|
||||
@@ -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,
|
||||
});
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+2
@@ -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[];
|
||||
+75
@@ -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])
|
||||
}
|
||||
|
||||
@@ -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==",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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())))
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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]:
|
||||
"""
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
@@ -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}
|
||||
@@ -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
@@ -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
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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",
|
||||
)
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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,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
|
||||
|
||||
@@ -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
|
||||
@@ -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()
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
Reference in New Issue
Block a user