Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_fix-config-update-targeted-upserts

This commit is contained in:
Yuneng Jiang
2026-04-29 18:26:19 -07:00
134 changed files with 7968 additions and 1253 deletions
+4
View File
@@ -2,6 +2,10 @@
<!-- e.g. "Fixes #000" -->
## Linear ticket
<!-- if you are an internal contributor, add the Linear ticket e.g. "Resolves LIT-1234" to magically link the Linear ticket to the GitHub PR -->
## Pre-Submission checklist
**Please complete all items before asking a LiteLLM maintainer to review your PR**
+4 -4
View File
@@ -4,7 +4,7 @@ on:
workflow_dispatch:
inputs:
tag:
description: "Release tag (e.g. v1.83.0-stable) — branch will be named release/<tag>"
description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted) — branch will be named release/<tag>"
required: true
type: string
commit_hash:
@@ -14,7 +14,7 @@ on:
workflow_call:
inputs:
tag:
description: "Release tag"
description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted)"
required: true
type: string
commit_hash:
@@ -40,8 +40,8 @@ jobs:
echo "::error::commit_hash must be a full 40-character commit SHA"
exit 1
fi
if ! echo "${TAG}" | grep -qE '^v[0-9]+\.[0-9]+\.[0-9]+'; then
echo "::error::tag must start with vX.Y.Z"
if ! echo "${TAG}" | grep -qE '^v?[0-9]+\.[0-9]+\.[0-9]+'; then
echo "::error::tag must start with X.Y.Z (optional leading v), e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, or v1.83.10-stable"
exit 1
fi
+9 -4
View File
@@ -4,7 +4,7 @@ on:
workflow_dispatch:
inputs:
tag:
description: "Release tag (e.g. v1.83.0-stable)"
description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted)"
required: true
type: string
commit_hash:
@@ -30,8 +30,8 @@ jobs:
echo "::error::commit_hash must be a full 40-character commit SHA"
exit 1
fi
if ! echo "${TAG}" | grep -qE '^v[0-9]+\.[0-9]+\.[0-9]+'; then
echo "::error::tag must start with vX.Y.Z"
if ! echo "${TAG}" | grep -qE '^v?[0-9]+\.[0-9]+\.[0-9]+'; then
echo "::error::tag must start with X.Y.Z (optional leading v), e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, or v1.83.10-stable"
exit 1
fi
@@ -45,6 +45,11 @@ jobs:
const tag = process.env.TAG;
const commitHash = process.env.COMMIT_HASH;
// Mark RC / dev / nightly / alpha / beta tags as GitHub pre-releases.
// PEP 440 post-releases (e.g. `1.84.0.post1`) and legacy `-stable[.patch.N]`
// are stable maintenance releases, not pre-releases.
const isPrerelease = /(?:rc|nightly|alpha|beta|\.dev)/i.test(tag);
const cosignSection = [
`## Verify Docker Image Signature`,
``,
@@ -89,7 +94,7 @@ jobs:
target_commitish: commitHash,
name: tag,
owner: context.repo.owner,
prerelease: false,
prerelease: isPrerelease,
repo: context.repo.repo,
tag_name: tag,
});
@@ -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;
@@ -1290,3 +1290,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])
}
+59 -2
View File
@@ -650,7 +650,10 @@ class Cache:
verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}")
def _convert_to_cached_embedding(
self, embedding_response: Any, model: Optional[str]
self,
embedding_response: Any,
model: Optional[str],
prompt_tokens_details: Optional[dict] = None,
) -> CachedEmbedding:
"""
Convert any embedding response into the standardized CachedEmbedding TypedDict format.
@@ -662,6 +665,7 @@ class Cache:
"index": embedding_response.get("index"),
"object": embedding_response.get("object"),
"model": model,
"prompt_tokens_details": prompt_tokens_details,
}
elif hasattr(embedding_response, "model_dump"):
data = embedding_response.model_dump()
@@ -670,6 +674,7 @@ class Cache:
"index": data.get("index"),
"object": data.get("object"),
"model": model,
"prompt_tokens_details": prompt_tokens_details,
}
else:
data = vars(embedding_response)
@@ -678,10 +683,54 @@ class Cache:
"index": data.get("index"),
"object": data.get("object"),
"model": model,
"prompt_tokens_details": prompt_tokens_details,
}
except KeyError as e:
raise ValueError(f"Missing expected key in embedding response: {e}")
def _get_per_item_prompt_tokens_details(
self,
result: EmbeddingResponse,
idx_in_result_data: int,
) -> Optional[dict]:
"""
Extract per-item prompt_tokens_details from a response for caching.
For single-item responses (common for multimodal providers like Bedrock Titan,
Nova, Vertex AI), returns the full prompt_tokens_details.
For multi-item responses, distributes integer fields evenly across items
so that summing all per-item details reconstructs the original totals.
"""
if result.usage is None or result.usage.prompt_tokens_details is None:
return None
details = result.usage.prompt_tokens_details
if hasattr(details, "model_dump"):
details_dict = details.model_dump(exclude_none=True)
elif isinstance(details, dict):
details_dict = {k: v for k, v in details.items() if v is not None}
else:
return None
if not details_dict:
return None
num_items = len(result.data)
if num_items <= 1:
return details_dict
# Distribute integer/float fields evenly across items
per_item: dict = {}
for key, value in details_dict.items():
if isinstance(value, int):
quotient, remainder = divmod(value, num_items)
per_item[key] = quotient + (1 if idx_in_result_data < remainder else 0)
elif isinstance(value, float):
per_item[key] = value / num_items
else:
per_item[key] = value
return per_item if per_item else None
def add_embedding_response_to_cache(
self,
result: EmbeddingResponse,
@@ -693,10 +742,18 @@ class Cache:
kwargs["cache_key"] = preset_cache_key
embedding_response = result.data[idx_in_result_data]
# Extract per-item prompt_tokens_details from response usage
prompt_tokens_details = self._get_per_item_prompt_tokens_details(
result=result,
idx_in_result_data=idx_in_result_data,
)
# Always convert to properly typed CachedEmbedding
model_name = result.model
embedding_dict: CachedEmbedding = self._convert_to_cached_embedding(
embedding_response, model_name
embedding_response,
model_name,
prompt_tokens_details=prompt_tokens_details,
)
cache_key, cached_data, kwargs = self._add_cache_logic(
+88
View File
@@ -59,6 +59,7 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.utils import PromptTokensDetailsWrapper
else:
LiteLLMLoggingObj = Any
@@ -415,6 +416,7 @@ class LLMCachingHandler:
final_embedding_cached_response._hidden_params["cache_hit"] = True
prompt_tokens = 0
aggregated_details: Optional[dict] = None
for val in non_null_list:
idx, cr = val # (idx, cr) tuple
if cr is not None:
@@ -431,11 +433,35 @@ class LLMCachingHandler:
prompt_tokens += token_counter(
text=kwargs_input_as_list[idx], count_response_tokens=True
)
# Aggregate prompt_tokens_details from cached items
item_details = cr.get("prompt_tokens_details")
if item_details:
if aggregated_details is None:
aggregated_details = {}
for key, value in item_details.items():
if isinstance(value, (int, float)):
aggregated_details[key] = (
aggregated_details.get(key, 0) + value
)
else:
aggregated_details[key] = value
## USAGE
prompt_tokens_details: Optional["PromptTokensDetailsWrapper"] = None
if aggregated_details:
from litellm.types.utils import PromptTokensDetailsWrapper
try:
prompt_tokens_details = PromptTokensDetailsWrapper(
**aggregated_details
)
except Exception:
prompt_tokens_details = None
usage = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=0,
total_tokens=prompt_tokens,
prompt_tokens_details=prompt_tokens_details,
)
final_embedding_cached_response.usage = usage
if len(remaining_list) == 0:
@@ -478,8 +504,70 @@ class LLMCachingHandler:
prompt_tokens=usage1.prompt_tokens + usage2.prompt_tokens,
completion_tokens=usage1.completion_tokens + usage2.completion_tokens,
total_tokens=usage1.total_tokens + usage2.total_tokens,
prompt_tokens_details=self._merge_prompt_tokens_details(
usage1.prompt_tokens_details,
usage2.prompt_tokens_details,
),
)
def _merge_prompt_tokens_details(
self,
details1: Optional["PromptTokensDetailsWrapper"],
details2: Optional["PromptTokensDetailsWrapper"],
) -> Optional["PromptTokensDetailsWrapper"]:
"""Merge two PromptTokensDetailsWrapper objects by summing numeric fields."""
if details1 is None and details2 is None:
return None
if details1 is None:
return details2
if details2 is None:
return details1
dict1 = (
details1.model_dump(exclude_none=True)
if hasattr(details1, "model_dump")
else {}
)
dict2 = (
details2.model_dump(exclude_none=True)
if hasattr(details2, "model_dump")
else {}
)
merged: dict = {}
for key in set(dict1.keys()) | set(dict2.keys()):
v1 = dict1.get(key, 0)
v2 = dict2.get(key, 0)
if isinstance(v1, (int, float)) and isinstance(v2, (int, float)):
merged[key] = v1 + v2
elif isinstance(v1, dict) and isinstance(v2, dict):
# Recursively merge nested dicts (e.g. cache_creation_token_details)
nested: dict = {}
for nk in set(v1.keys()) | set(v2.keys()):
nv1 = v1.get(nk, 0)
nv2 = v2.get(nk, 0)
if isinstance(nv1, (int, float)) and isinstance(nv2, (int, float)):
nested[nk] = nv1 + nv2
elif nv1:
nested[nk] = nv1
else:
nested[nk] = nv2
merged[key] = nested
elif v1:
merged[key] = v1
else:
merged[key] = v2
if not merged:
return None
from litellm.types.utils import PromptTokensDetailsWrapper
try:
return PromptTokensDetailsWrapper(**merged)
except Exception:
return None
def _combine_cached_embedding_response_with_api_result(
self,
_caching_handler_response: CachingHandlerResponse,
+20
View File
@@ -224,6 +224,16 @@ AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int(
)
AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120))
AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300))
# TCP keep-alive (SO_KEEPALIVE) — opt-in. Required when running behind NAT/LBs
# whose idle timeout is shorter than provider response timeouts (e.g. AWS NAT
# Gateway: 350s vs OpenAI/Azure: 600s). Without this, the kernel sends nothing
# during a long provider call and the NAT reaps the flow before the response
# arrives. Enabling SO_KEEPALIVE makes the kernel emit TCP probes that reset
# the NAT idle timer.
AIOHTTP_SO_KEEPALIVE = os.getenv("AIOHTTP_SO_KEEPALIVE", "False").lower() == "true"
AIOHTTP_TCP_KEEPIDLE = int(os.getenv("AIOHTTP_TCP_KEEPIDLE", 60))
AIOHTTP_TCP_KEEPINTVL = int(os.getenv("AIOHTTP_TCP_KEEPINTVL", 30))
AIOHTTP_TCP_KEEPCNT = int(os.getenv("AIOHTTP_TCP_KEEPCNT", 5))
# enable_cleanup_closed is only needed for Python versions with the SSL leak bug
# Fixed in Python 3.12.7+ and 3.13.1+ (see https://github.com/python/cpython/pull/118960)
# Reference: https://github.com/aio-libs/aiohttp/blob/master/aiohttp/connector.py#L74-L78
@@ -1396,6 +1406,15 @@ LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS = int(
os.getenv("LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS", 600)
) # 10 minutes default — caps the deadlock window if a pod crashes mid-rotation
UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard"
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED = os.getenv(
"LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED", "false"
)
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS = int(
os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS", 86400)
) # 24 hours default
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE = int(
os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE", 1000)
)
LITELLM_PROXY_ADMIN_NAME = "default_user_id"
########################### CLI SSO AUTHENTICATION CONSTANTS ###########################
@@ -1425,6 +1444,7 @@ CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int(
)
SPEND_LOG_CLEANUP_JOB_NAME = "spend_log_cleanup"
KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job"
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME = "litellm_expired_ui_session_key_cleanup_job"
SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500))
SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000))
SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100))
+7 -1
View File
@@ -1,6 +1,6 @@
import asyncio
from datetime import datetime
from typing import TYPE_CHECKING, Any, List, Optional
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy.pass_through_endpoints.success_handler import (
@@ -29,12 +29,14 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
litellm_logging_obj: LiteLLMLoggingObj,
request_body: dict,
model: str,
hidden_params: Optional[Dict[str, Any]] = None,
):
self.litellm_logging_obj = litellm_logging_obj
self.request_body = request_body
self.start_time = datetime.now()
self.collected_chunks: List[bytes] = []
self.model = model
self._hidden_params: Dict[str, Any] = hidden_params or {}
async def _handle_async_streaming_logging(
self,
@@ -76,11 +78,13 @@ class GoogleGenAIGenerateContentStreamingIterator(
litellm_metadata: dict,
custom_llm_provider: str,
request_body: Optional[dict] = None,
hidden_params: Optional[Dict[str, Any]] = None,
):
super().__init__(
litellm_logging_obj=logging_obj,
request_body=request_body or {},
model=model,
hidden_params=hidden_params,
)
self.response = response
self.model = model
@@ -130,11 +134,13 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(
litellm_metadata: dict,
custom_llm_provider: str,
request_body: Optional[dict] = None,
hidden_params: Optional[Dict[str, Any]] = None,
):
super().__init__(
litellm_logging_obj=logging_obj,
request_body=request_body or {},
model=model,
hidden_params=hidden_params,
)
self.response = response
self.model = model
@@ -11,8 +11,9 @@ import json
import os
import re
import traceback
from typing import Dict, List, Literal, Optional, Union
from typing import Any, Dict, List, Literal, Optional, Union
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm._uuid import uuid
@@ -103,6 +104,9 @@ class GenericAPILogger(CustomBatchLogger):
event_types: Optional[List[API_EVENT_TYPES]] = None,
callback_name: Optional[str] = None,
log_format: Optional[LOG_FORMAT_TYPES] = None,
max_retries: int = 0,
retry_delay: float = 1.0,
timeout: Optional[Union[float, httpx.Timeout]] = None,
**kwargs,
):
"""
@@ -114,6 +118,9 @@ class GenericAPILogger(CustomBatchLogger):
event_types: Optional[List[API_EVENT_TYPES]] = None,
callback_name: Optional[str] = None - If provided, loads config from generic_api_compatible_callbacks.json
log_format: Optional[LOG_FORMAT_TYPES] = None - Format for log output: "json_array" (default), "ndjson", or "single"
max_retries: Number of retry attempts after the initial request fails. Defaults to 0.
retry_delay: Initial retry delay in seconds. Retries use exponential backoff.
timeout: Optional timeout to use for Generic API callback requests.
"""
#########################################################
# Check if callback_name is provided and load config
@@ -162,6 +169,10 @@ class GenericAPILogger(CustomBatchLogger):
self.endpoint: str = endpoint
self.event_types: Optional[List[API_EVENT_TYPES]] = event_types
self.callback_name: Optional[str] = callback_name
self.max_retries = max(0, int(max_retries or 0))
retry_delay_value = 0.0 if retry_delay is None else retry_delay
self.retry_delay = max(0.0, float(retry_delay_value))
self.timeout = timeout
# Validate and store log_format
if log_format is not None and log_format not in [
@@ -226,6 +237,53 @@ class GenericAPILogger(CustomBatchLogger):
return headers_dict
def _should_retry_exception(self, exception: Exception) -> bool:
if isinstance(exception, (litellm.Timeout, httpx.TransportError)):
return True
if isinstance(exception, httpx.HTTPStatusError):
return exception.response.status_code >= 500
return False
async def _sleep_before_retry(self, attempt: int) -> None:
if self.retry_delay <= 0:
return
delay = self.retry_delay * (2**attempt)
await asyncio.sleep(delay)
async def _post_with_retries(self, data: str) -> httpx.Response:
post_kwargs: Dict[str, Any] = {
"url": self.endpoint,
"headers": self.headers,
"data": data,
}
if self.timeout is not None:
post_kwargs["timeout"] = self.timeout
total_attempts = self.max_retries + 1
for attempt in range(total_attempts):
try:
return await self.async_httpx_client.post(**post_kwargs)
except Exception as e:
is_last_attempt = attempt == self.max_retries
should_retry = self._should_retry_exception(e)
if is_last_attempt or not should_retry:
raise
verbose_logger.warning(
"Generic API Logger - retrying request to %s after error: %s "
"(attempt %s/%s)",
self.endpoint,
str(e),
attempt + 1,
total_attempts,
)
await self._sleep_before_retry(attempt)
raise RuntimeError("Generic API Logger retry loop exited unexpectedly")
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
"""
Async Log success events to Generic API Endpoint
@@ -325,11 +383,7 @@ class GenericAPILogger(CustomBatchLogger):
# Send each log as individual HTTP request in parallel
tasks = []
for log_entry in self.log_queue:
task = self.async_httpx_client.post(
url=self.endpoint,
headers=self.headers,
data=safe_dumps(log_entry),
)
task = self._post_with_retries(data=safe_dumps(log_entry))
tasks.append(task)
# Execute all requests in parallel
@@ -356,11 +410,7 @@ class GenericAPILogger(CustomBatchLogger):
raise ValueError(f"Unknown log_format: {self.log_format}")
# Make POST request
response = await self.async_httpx_client.post(
url=self.endpoint,
headers=self.headers,
data=data,
)
response = await self._post_with_retries(data=data)
verbose_logger.debug(
f"Generic API Logger - sent batch to {self.endpoint}, "
@@ -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"
+28 -25
View File
@@ -1467,6 +1467,8 @@ class Logging(LiteLLMLoggingBaseClass):
LiteLLMRealtimeStreamLoggingObject,
OpenAIModerationResponse,
"SearchResponse",
dict,
list,
],
cache_hit: Optional[bool] = None,
litellm_model_name: Optional[str] = None,
@@ -1725,12 +1727,18 @@ class Logging(LiteLLMLoggingBaseClass):
return
if self.model_call_details.get("litellm_params") is None:
return
self.model_call_details["litellm_params"].setdefault("metadata", {})
if self.model_call_details["litellm_params"]["metadata"] is None:
self.model_call_details["litellm_params"]["metadata"] = {}
self.model_call_details["litellm_params"]["metadata"]["hidden_params"] = (
getattr(logging_result, "_hidden_params", {})
)
metadata_hidden_params = hidden_params.copy()
response_cost = self.model_call_details.get("response_cost")
if (
metadata_hidden_params.get("response_cost") is None
and response_cost is not None
):
metadata_hidden_params["response_cost"] = response_cost
litellm_params = self.model_call_details["litellm_params"]
metadata = litellm_params.get("metadata") or {}
litellm_params["metadata"] = metadata
metadata["hidden_params"] = metadata_hidden_params
def _process_hidden_params_and_response_cost(
self,
@@ -1738,6 +1746,7 @@ class Logging(LiteLLMLoggingBaseClass):
start_time,
end_time,
):
"""Resolve hidden params, compute response cost, and emit the standard logging payload."""
hidden_params = getattr(logging_result, "_hidden_params", {})
if hidden_params:
if self.model_call_details.get("litellm_params") is not None:
@@ -1871,24 +1880,12 @@ class Logging(LiteLLMLoggingBaseClass):
):
if self._is_recognized_call_type_for_logging(
logging_result=logging_result
):
) or isinstance(logging_result, (dict, list)):
self._process_hidden_params_and_response_cost(
logging_result=logging_result,
start_time=start_time,
end_time=end_time,
)
elif isinstance(result, dict) or isinstance(result, list):
self.model_call_details["standard_logging_object"] = (
self._build_standard_logging_payload(
result, start_time, end_time
)
)
if (
standard_logging_payload := self.model_call_details.get(
"standard_logging_object"
)
) is not None:
emit_standard_logging_payload(standard_logging_payload)
elif standard_logging_object is not None:
self.model_call_details["standard_logging_object"] = (
standard_logging_object
@@ -5438,11 +5435,6 @@ def get_standard_logging_object_payload(
completion_start_time_float=completion_start_time_float,
stream=kwargs.get("stream", False),
)
# clean up litellm hidden params
clean_hidden_params = StandardLoggingPayloadSetup.get_hidden_params(
hidden_params
)
# clean up litellm metadata
clean_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata(
metadata=metadata,
@@ -5476,6 +5468,18 @@ def get_standard_logging_object_payload(
## Get model cost information ##
base_model = _get_base_model_from_metadata(model_call_details=kwargs)
custom_pricing = use_custom_pricing_for_model(litellm_params=litellm_params)
raw_response_cost = kwargs.get("response_cost")
response_cost: float = raw_response_cost or 0.0
# clean up litellm hidden params
clean_hidden_params = StandardLoggingPayloadSetup.get_hidden_params(
hidden_params
)
if (
clean_hidden_params["response_cost"] is None
and raw_response_cost is not None
):
clean_hidden_params["response_cost"] = response_cost
model_cost_information = StandardLoggingPayloadSetup.get_model_cost_information(
base_model=base_model,
@@ -5484,7 +5488,6 @@ def get_standard_logging_object_payload(
init_response_obj=init_response_obj,
api_base=litellm_params.get("api_base"),
)
response_cost: float = kwargs.get("response_cost", 0) or 0.0
error_information = StandardLoggingPayloadSetup.get_error_information(
original_exception=original_exception,
@@ -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(
@@ -221,6 +221,13 @@ class LoggingCallbackManager:
headers = callback_config.get("headers")
event_types = callback_config.get("event_types")
log_format = callback_config.get("log_format")
max_retries = max(0, int(callback_config.get("max_retries", 0) or 0))
retry_delay_value = callback_config.get("retry_delay")
retry_delay = max(
0.0,
float(0.0 if retry_delay_value is None else retry_delay_value),
)
timeout = callback_config.get("timeout")
if endpoint is None or headers is None:
verbose_logger.warning(
@@ -236,6 +243,9 @@ class LoggingCallbackManager:
and cached_logger.headers == headers
and cached_logger.event_types == event_types
and cached_logger.log_format == log_format
and cached_logger.max_retries == max_retries
and cached_logger.retry_delay == retry_delay
and cached_logger.timeout == timeout
):
return cached_logger
@@ -244,6 +254,9 @@ class LoggingCallbackManager:
headers=headers,
event_types=event_types,
log_format=log_format,
max_retries=max_retries,
retry_delay=retry_delay,
timeout=timeout,
)
_generic_api_logger_cache[callback] = new_logger
return new_logger
@@ -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
@@ -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
@@ -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
+62
View File
@@ -1,5 +1,7 @@
import asyncio
import inspect
import os
import socket
import ssl
import sys
import time
@@ -29,6 +31,10 @@ from litellm.constants import (
AIOHTTP_CONNECTOR_LIMIT_PER_HOST,
AIOHTTP_KEEPALIVE_TIMEOUT,
AIOHTTP_NEEDS_CLEANUP_CLOSED,
AIOHTTP_SO_KEEPALIVE,
AIOHTTP_TCP_KEEPCNT,
AIOHTTP_TCP_KEEPIDLE,
AIOHTTP_TCP_KEEPINTVL,
AIOHTTP_TTL_DNS_CACHE,
COMPLETION_HTTP_FALLBACK_SECONDS,
DEFAULT_SSL_CIPHERS,
@@ -54,6 +60,57 @@ except Exception:
version = "0.0.0"
# aiohttp 3.10+ exposes a `socket_factory` kwarg on TCPConnector. Older
# versions don't — detect once and skip the keep-alive wiring there.
# https://docs.aiohttp.org/en/stable/client_reference.html#aiohttp.TCPConnector
_AIOHTTP_SUPPORTS_SOCKET_FACTORY = (
"socket_factory" in inspect.signature(TCPConnector.__init__).parameters
)
def _build_aiohttp_keepalive_socket_factory() -> (
Optional[Callable[[Tuple[Any, ...]], socket.socket]]
):
"""
Build a socket_factory that enables SO_KEEPALIVE on aiohttp TCP sockets.
Why: by default, aiohttp creates sockets without SO_KEEPALIVE, so the kernel
sends nothing during a long idle TCP connection. NAT/LB hops (e.g. AWS NAT
Gateway, 350s idle timeout) reap the flow well before slow provider
responses (OpenAI/Azure: up to 600s) arrive. Enabling SO_KEEPALIVE makes
the kernel emit TCP probes that reset the NAT idle timer.
Returns None when AIOHTTP_SO_KEEPALIVE is disabled or aiohttp is too old.
"""
if not AIOHTTP_SO_KEEPALIVE or not _AIOHTTP_SUPPORTS_SOCKET_FACTORY:
return None
def factory(addr_info: Tuple[Any, ...]) -> socket.socket:
family, type_, proto = addr_info[0], addr_info[1], addr_info[2]
sock = socket.socket(family=family, type=type_, proto=proto)
sock.setblocking(False)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
# Linux: TCP_KEEPIDLE is idle-before-first-probe.
# macOS/Darwin: TCP_KEEPALIVE is the equivalent.
if hasattr(socket, "TCP_KEEPIDLE"):
sock.setsockopt(
socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, AIOHTTP_TCP_KEEPIDLE
)
elif hasattr(socket, "TCP_KEEPALIVE"):
sock.setsockopt(
socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, AIOHTTP_TCP_KEEPIDLE
)
if hasattr(socket, "TCP_KEEPINTVL"):
sock.setsockopt(
socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, AIOHTTP_TCP_KEEPINTVL
)
if hasattr(socket, "TCP_KEEPCNT"):
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, AIOHTTP_TCP_KEEPCNT)
return sock
return factory
def get_default_headers() -> dict:
"""
Get default headers for HTTP requests.
@@ -935,6 +992,11 @@ class AsyncHTTPHandler:
transport_connector_kwargs["limit_per_host"] = (
AIOHTTP_CONNECTOR_LIMIT_PER_HOST
)
# Returns None when SO_KEEPALIVE is disabled or aiohttp is too old to
# accept socket_factory — version detection lives inside the builder.
socket_factory = _build_aiohttp_keepalive_socket_factory()
if socket_factory is not None:
transport_connector_kwargs["socket_factory"] = socket_factory
return LiteLLMAiohttpTransport(
client=lambda: ClientSession(
@@ -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,
@@ -8585,6 +8609,7 @@ class BaseLLMHTTPHandler:
api_base=api_base,
litellm_logging_obj=logging_obj,
litellm_params=dict(litellm_params),
extra_body=extra_body,
)
else:
(
@@ -8597,6 +8622,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 +8723,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 +10452,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 +10567,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(
@@ -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.
@@ -130,6 +130,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
@@ -1,5 +1,5 @@
"""
Cost calculator for OpenAI image generation models (gpt-image-1, gpt-image-1-mini)
Cost calculator for OpenAI image generation models (gpt-image family)
These models use token-based pricing instead of pixel-based pricing like DALL-E.
"""
@@ -17,13 +17,13 @@ def cost_calculator(
custom_llm_provider: Optional[str] = None,
) -> float:
"""
Calculate cost for OpenAI gpt-image-1 and gpt-image-1-mini models.
Calculate cost for OpenAI gpt-image models.
Uses the same usage format as Responses API, so we reuse the helper
to transform to chat completion format and use generic_cost_per_token.
Args:
model: The model name (e.g., "gpt-image-1", "gpt-image-1-mini")
model: The model name (e.g., "gpt-image-1", "gpt-image-2")
image_response: The ImageResponse containing usage data
custom_llm_provider: Optional provider name
@@ -15,7 +15,7 @@ if TYPE_CHECKING:
class GPTImageGenerationConfig(BaseImageGenerationConfig):
"""
OpenAI gpt-image-1 image generation config
OpenAI gpt-image image generation config
"""
def get_supported_openai_params(
@@ -106,6 +106,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict]:
url = f"{api_base}/{vector_store_id}/search"
typed_request_body = VectorStoreSearchRequest(
+5
View File
@@ -101,5 +101,10 @@
"param_mappings": {
"max_completion_tokens": "max_tokens"
}
},
"aihubmix": {
"base_url": "https://aihubmix.com/v1",
"api_key_env": "AIHUBMIX_API_KEY",
"api_base_env": "AIHUBMIX_API_BASE"
}
}
@@ -80,6 +80,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict]:
url = f"{api_base}/{vector_store_id}/search"
_, request_body = super().transform_search_vector_store_request(
@@ -89,5 +90,6 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig):
api_base=api_base,
litellm_logging_obj=litellm_logging_obj,
litellm_params=litellm_params,
extra_body=extra_body,
)
return url, request_body
@@ -102,6 +102,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict]:
"""RAGFlow vector stores are management-only, search is not supported."""
raise NotImplementedError(
@@ -79,6 +79,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict]:
"""Sync version - generates embedding synchronously."""
# For S3 Vectors, vector_store_id should be in format: bucket_name:index_name
@@ -140,6 +141,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, Dict]:
"""Async version - generates embedding asynchronously."""
# For S3 Vectors, vector_store_id should be in format: bucket_name:index_name
+16 -4
View File
@@ -4,6 +4,7 @@ from typing import Any, Coroutine, Dict, Optional, Union
import httpx
import litellm
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
get_async_httpx_client,
@@ -224,8 +225,14 @@ class VertexAIBatchPrediction(VertexLLM):
},
)
response = sync_handler.get(
url=api_base,
# ``api_base`` here can come from caller-supplied request kwargs
# (clientside override). Wrap the fetch in ``safe_get`` so DNS
# rebind / private / cloud-metadata targets are rejected; the
# proxy auth gate already blocks malicious clientside ``api_base``
# at the boundary — this is defense-in-depth for SDK callers.
response = safe_get(
sync_handler,
api_base,
headers=headers,
)
@@ -270,8 +277,13 @@ class VertexAIBatchPrediction(VertexLLM):
},
)
response = await client.get(
url=api_base,
# Mirror the sync path: ``api_base`` may come from caller-supplied
# request kwargs, so wrap the fetch in ``async_safe_get`` to reject
# DNS-rebind / private / cloud-metadata targets. Defense-in-depth
# behind the proxy auth gate's clientside ``api_base`` check.
response = await async_safe_get(
client,
api_base,
headers=headers,
)
if response.status_code != 200:
+11 -8
View File
@@ -597,7 +597,14 @@ def process_items(schema, depth=0):
f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema. Please check the schema for excessive nesting."
)
if isinstance(schema, dict):
if "items" in schema and schema["items"] == {}:
# Vertex requires `items` whenever `type == "array"` (even inside anyOf).
# Normalize: empty `items: {}` and missing-items both become {"type": "object"}.
type_val = schema.get("type")
if (
isinstance(type_val, str)
and type_val.lower() == "array"
and ("items" not in schema or schema.get("items") == {})
):
schema["items"] = {"type": "object"}
for key, value in schema.items():
if isinstance(value, dict):
@@ -710,14 +717,10 @@ def convert_anyof_null_to_nullable(schema, depth=0):
if contains_null:
# set all types to nullable following guidance found here: https://cloud.google.com/vertex-ai/generative-ai/docs/samples/generativeaionvertexai-gemini-controlled-generation-response-schema-3#generativeaionvertexai_gemini_controlled_generation_response_schema_3-python
# Empty `items: {}` on array branches is left in place; downstream
# process_items() converts it to {"type": "object"}, which Vertex
# requires whenever type == "array" (even inside anyOf).
for atype in anyof:
# Remove items field if type is array and items is empty
if (
atype.get("type") == "array"
and "items" in atype
and not atype["items"]
):
atype.pop("items")
atype["nullable"] = True
properties = schema.get("properties", None)
@@ -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,
@@ -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
@@ -4735,17 +4735,17 @@
"supports_web_search": true
},
"azure/gpt-5.5-pro": {
"cache_read_input_token_cost": 6e-06,
"cache_read_input_token_cost_above_272k_tokens": 1.2e-05,
"input_cost_per_token": 6e-05,
"input_cost_per_token_above_272k_tokens": 0.00012,
"cache_read_input_token_cost": 3e-06,
"cache_read_input_token_cost_above_272k_tokens": 6e-06,
"input_cost_per_token": 3e-05,
"input_cost_per_token_above_272k_tokens": 6e-05,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 0.00036,
"output_cost_per_token_above_272k_tokens": 0.00054,
"output_cost_per_token": 0.00018,
"output_cost_per_token_above_272k_tokens": 0.00027,
"supported_endpoints": [
"/v1/batch",
"/v1/responses"
@@ -4774,17 +4774,17 @@
"supports_low_reasoning_effort": false
},
"azure/gpt-5.5-pro-2026-04-23": {
"cache_read_input_token_cost": 6e-06,
"cache_read_input_token_cost_above_272k_tokens": 1.2e-05,
"input_cost_per_token": 6e-05,
"input_cost_per_token_above_272k_tokens": 0.00012,
"cache_read_input_token_cost": 3e-06,
"cache_read_input_token_cost_above_272k_tokens": 6e-06,
"input_cost_per_token": 3e-05,
"input_cost_per_token_above_272k_tokens": 6e-05,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 0.00036,
"output_cost_per_token_above_272k_tokens": 0.00054,
"output_cost_per_token": 0.00018,
"output_cost_per_token_above_272k_tokens": 0.00027,
"supported_endpoints": [
"/v1/batch",
"/v1/responses"
@@ -5103,6 +5103,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",
@@ -19083,6 +19115,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",
@@ -19898,21 +19962,21 @@
"supports_minimal_reasoning_effort": true
},
"gpt-5.5-pro": {
"cache_read_input_token_cost": 6e-06,
"cache_read_input_token_cost_above_272k_tokens": 1.2e-05,
"input_cost_per_token": 6e-05,
"input_cost_per_token_above_272k_tokens": 0.00012,
"input_cost_per_token_flex": 3e-05,
"input_cost_per_token_batches": 3e-05,
"cache_read_input_token_cost": 3e-06,
"cache_read_input_token_cost_above_272k_tokens": 6e-06,
"input_cost_per_token": 3e-05,
"input_cost_per_token_above_272k_tokens": 6e-05,
"input_cost_per_token_flex": 1.5e-05,
"input_cost_per_token_batches": 1.5e-05,
"litellm_provider": "openai",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 0.00036,
"output_cost_per_token_above_272k_tokens": 0.00054,
"output_cost_per_token_flex": 0.00018,
"output_cost_per_token_batches": 0.00018,
"output_cost_per_token": 0.00018,
"output_cost_per_token_above_272k_tokens": 0.00027,
"output_cost_per_token_flex": 9e-05,
"output_cost_per_token_batches": 9e-05,
"supported_endpoints": [
"/v1/responses",
"/v1/batch"
@@ -19941,21 +20005,21 @@
"supports_minimal_reasoning_effort": true
},
"gpt-5.5-pro-2026-04-23": {
"cache_read_input_token_cost": 6e-06,
"cache_read_input_token_cost_above_272k_tokens": 1.2e-05,
"input_cost_per_token": 6e-05,
"input_cost_per_token_above_272k_tokens": 0.00012,
"input_cost_per_token_flex": 3e-05,
"input_cost_per_token_batches": 3e-05,
"cache_read_input_token_cost": 3e-06,
"cache_read_input_token_cost_above_272k_tokens": 6e-06,
"input_cost_per_token": 3e-05,
"input_cost_per_token_above_272k_tokens": 6e-05,
"input_cost_per_token_flex": 1.5e-05,
"input_cost_per_token_batches": 1.5e-05,
"litellm_provider": "openai",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 0.00036,
"output_cost_per_token_above_272k_tokens": 0.00054,
"output_cost_per_token_flex": 0.00018,
"output_cost_per_token_batches": 0.00018,
"output_cost_per_token": 0.00018,
"output_cost_per_token_above_272k_tokens": 0.00027,
"output_cost_per_token_flex": 9e-05,
"output_cost_per_token_batches": 9e-05,
"supported_endpoints": [
"/v1/responses",
"/v1/batch"
@@ -50,8 +50,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 +109,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 +373,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 +736,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 +749,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 +1253,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
)
@@ -1488,11 +1509,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
)
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,8 +1548,8 @@ 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,
@@ -1674,6 +1712,9 @@ 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("/"))
@@ -1713,7 +1754,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 +1872,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 +1957,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 +2724,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 +2795,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 +2821,21 @@ class MCPServerManager:
f"Building server from DB: {server.server_id} ({server.server_name})"
)
new_server = await self.build_mcp_server_from_table(server)
# Carry the cached short_prefix from the previous registry entry
# (if any) so the prefix is stable across reloads.
if existing_server is not None and existing_server.short_prefix:
new_server.short_prefix = existing_server.short_prefix
new_registry[server.server_id] = new_server
await self._maybe_register_openapi_tools(new_server)
# Swap in the new registry first so _assign_unique_short_prefix
# sees the complete set when checking for collisions.
self.registry = new_registry
for new_server in new_registry.values():
self._assign_unique_short_prefix(new_server)
# Register OpenAPI tools *after* the final short prefix is assigned
# so the tools are stored in the global registry under the same
# prefix that lookups will use.
await self._maybe_register_openapi_tools(new_server)
verbose_logger.debug(
"MCP registry refreshed (%s servers in registry)", len(new_registry)
@@ -49,6 +49,7 @@ from litellm.proxy._experimental.mcp_server.utils import (
LITELLM_MCP_SERVER_VERSION,
add_server_prefix_to_name,
get_server_prefix,
iter_known_server_prefixes,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
@@ -711,13 +712,7 @@ if MCP_AVAILABLE:
for server in allowed_mcp_servers:
if server:
match_list = [
s.lower()
for s in [
server.alias,
server.server_name,
server.server_id,
]
if s is not None
s.lower() for s in iter_known_server_prefixes(server) if s
]
if server_or_group.lower() in match_list:
@@ -2031,11 +2026,13 @@ if MCP_AVAILABLE:
# Remove prefix from tool name for logging and processing
original_tool_name, server_name = split_server_prefix_from_name(name)
# If tool name is unprefixed, resolve its server so we can enforce permissions
if not server_name:
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
if mcp_server:
server_name = mcp_server.name
# Resolve the actual MCP server up-front so the permission check uses
# the canonical server.name even when the tool name is prefixed with a
# short ID (LITELLM_USE_SHORT_MCP_TOOL_PREFIX) that doesn't match the
# server's display name directly.
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
if mcp_server is not None:
server_name = mcp_server.name
# Only enforce server-level permissions when we can resolve a server
if server_name:
+135 -3
View File
@@ -2,10 +2,11 @@
MCP Server Utilities
"""
from typing import Any, Dict, Mapping, Optional, Tuple
from typing import Any, Dict, Iterator, Mapping, Optional, Tuple
import os
import hashlib
import importlib
import os
# Constants
LITELLM_MCP_SERVER_NAME = "litellm-mcp-server"
@@ -14,6 +15,89 @@ LITELLM_MCP_SERVER_DESCRIPTION = "MCP Server for LiteLLM"
MCP_TOOL_PREFIX_SEPARATOR = os.environ.get("MCP_TOOL_PREFIX_SEPARATOR", "-")
MCP_TOOL_PREFIX_FORMAT = "{server_name}{separator}{tool_name}"
# ---------------------------------------------------------------------------
# Short-ID tool prefix (opt-in)
# ---------------------------------------------------------------------------
# When LITELLM_USE_SHORT_MCP_TOOL_PREFIX is truthy the prefix attached to MCP
# tool / prompt / resource / resource-template names switches from the
# (potentially long) human-readable server name to a deterministic three
# character ID derived from the server's ``server_id``.
#
# Why three characters?
# * The first character is restricted to 52 alphabetic characters
# ([A-Za-z]) and the remaining two characters use the full base62
# alphabet ([0-9A-Za-z]). That guarantees the prefix never starts
# with a digit so it remains a valid identifier for every model API
# (some providers historically required a leading alphabetic char).
# * 52 * 62 * 62 = 199_888 distinct IDs. The chance of a real local
# tool name happening to begin with the exact prefix LiteLLM assigned
# to a given MCP server is negligible in practice.
# * The IDs are short enough that prefixed tool names stay well under
# the 60-character upper bound enforced by some model APIs (Anthropic
# etc.) even for long upstream tool names.
# * The mapping is deterministic (SHA-256 of ``server_id`` → three
# characters drawn from the alphabets above), so the prefix is stable
# across processes, workers and restarts without any persistence
# layer. Two servers with different ``server_id`` values can in
# principle hash to the same three chars; that natural-hash collision
# IS a routing-correctness issue (the second registrant would otherwise
# have its tools misrouted to the first), so registration goes through
# ``MCPServerManager._assign_unique_short_prefix`` which rehashes with
# a deterministic attempt counter until it finds an unused prefix and
# caches the result on ``MCPServer.short_prefix``. A collision is
# logged at INFO when it happens.
#
# This flag is intentionally opt-in for the first release so customers can
# migrate. It will become the default in a future release.
SHORT_MCP_TOOL_PREFIX_LENGTH = 3
_BASE62_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
# Subset of _BASE62_ALPHABET used for the *first* character only, to
# guarantee the prefix never starts with a digit.
_BASE52_ALPHA_ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
def is_short_mcp_tool_prefix_enabled() -> bool:
"""Return True when the short-ID tool prefix mode is enabled.
Read at call time (not import time) so tests and runtime config changes
take effect without reimporting the module.
"""
raw = os.environ.get("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "")
return raw.strip().lower() in ("1", "true", "yes", "on")
def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str:
"""Derive the deterministic three-character prefix for a server.
Uses SHA-256 of ``f"{server_id}#{attempt}"`` and folds the first eight
bytes into a fixed-length string whose first character is drawn from
``_BASE52_ALPHA_ALPHABET`` (so the prefix never starts with a digit)
and whose remaining characters are drawn from the full base62
alphabet. Pass ``attempt > 0`` to rehash to a different prefix when
the natural hash collides with a prefix already assigned to another
server (see ``MCPServerManager._assign_unique_short_prefix``). An
empty ``server_id`` raises ``ValueError`` short prefixes require a
stable identifier to be deterministic.
"""
if not server_id:
raise ValueError("compute_short_server_prefix requires a non-empty server_id")
seed = server_id if attempt == 0 else f"{server_id}#{attempt}"
digest = hashlib.sha256(seed.encode("utf-8")).digest()
value = int.from_bytes(digest[:8], "big")
# Build chars from least-significant to most-significant; we reverse
# at the end so the first emitted char comes from the high-order
# bits of the digest (which is the position we constrain to be
# alphabetic).
chars = []
for position in range(SHORT_MCP_TOOL_PREFIX_LENGTH):
is_first_char = position == SHORT_MCP_TOOL_PREFIX_LENGTH - 1
alphabet = _BASE52_ALPHA_ALPHABET if is_first_char else _BASE62_ALPHABET
value, idx = divmod(value, len(alphabet))
chars.append(alphabet[idx])
return "".join(reversed(chars))
def is_mcp_available() -> bool:
"""
@@ -82,7 +166,25 @@ def add_server_prefix_to_name(name: str, server_name: str) -> str:
def get_server_prefix(server: Any) -> str:
"""Return the prefix for a server: alias if present, else server_name, else server_id"""
"""Return the prefix for a server.
When the short-prefix mode is enabled (``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``)
a three-character base62 ID is returned. We prefer the cached
``server.short_prefix`` value when set that field is populated at
registration time by ``MCPServerManager._assign_unique_short_prefix``
and resolves natural-hash collisions deterministically and only fall
back to the natural hash for ad-hoc / temp-server objects without a
cached value. In default mode the historical behaviour is preserved:
alias if present, else server_name, else server_id.
"""
if is_short_mcp_tool_prefix_enabled():
cached = getattr(server, "short_prefix", None)
if cached:
return cached
server_id = getattr(server, "server_id", None)
if server_id:
return compute_short_server_prefix(server_id)
if hasattr(server, "alias") and server.alias:
return server.alias
if hasattr(server, "server_name") and server.server_name:
@@ -92,6 +194,36 @@ def get_server_prefix(server: Any) -> str:
return ""
def iter_known_server_prefixes(server: Any) -> Iterator[str]:
"""Yield every prefix form that may appear in tool names for ``server``.
Always includes the *current* prefix returned by ``get_server_prefix``.
Additionally yields the historical (alias / server_name / server_id) and
short-ID forms so the routing layer can resolve tool names regardless of
which prefix mode was active when the client first observed them.
"""
seen = set()
def _emit(value: Optional[str]) -> Iterator[str]:
if value and value not in seen:
seen.add(value)
yield value
yield from _emit(get_server_prefix(server))
yield from _emit(getattr(server, "short_prefix", None))
server_id = getattr(server, "server_id", None)
if server_id:
try:
yield from _emit(compute_short_server_prefix(server_id))
except ValueError:
pass
yield from _emit(getattr(server, "alias", None))
yield from _emit(getattr(server, "server_name", None))
yield from _emit(server_id)
def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]:
"""Return the unprefixed name plus the server name used as prefix."""
if MCP_TOOL_PREFIX_SEPARATOR in prefixed_name:
File diff suppressed because one or more lines are too long
-22
View File
@@ -1,22 +0,0 @@
1:"$Sreact.fragment"
2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
6:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"]
7:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"]
a:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"]
b:"$Sreact.suspense"
d:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"]
f:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"]
11:I[168027,[],"default"]
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}]
0:{"P":null,"b":"zxkD4-EPlgfKHDTw8O869","c":["","chat"],"q":"","i":false,"f":[[["",{"children":["chat",{"children":["__PAGE__",{}]}]},"$undefined","$undefined",true],[["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true,"nonce":"$undefined"}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]],"forbidden":"$undefined","unauthorized":"$undefined"}]}]}]}]}]]}],{"children":[["$","$1","c",{"children":[null,["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":"$undefined","forbidden":"$undefined","unauthorized":"$undefined"}]]}],{"children":[["$","$1","c",{"children":[["$","$L6",null,{"Component":"$7","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@8","$@9"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true,"nonce":"$undefined"}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true,"nonce":"$undefined"}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true,"nonce":"$undefined"}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true,"nonce":"$undefined"}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true,"nonce":"$undefined"}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true,"nonce":"$undefined"}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true,"nonce":"$undefined"}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true,"nonce":"$undefined"}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true,"nonce":"$undefined"}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true,"nonce":"$undefined"}]],["$","$La",null,{"children":["$","$b",null,{"name":"Next.MetadataOutlet","children":"$@c"}]}]]}],{},null,false,false]},null,false,false]},null,false,false],["$","$1","h",{"children":[null,["$","$Ld",null,{"children":"$Le"}],["$","div",null,{"hidden":true,"children":["$","$Lf",null,{"children":["$","$b",null,{"name":"Next.Metadata","children":"$L10"}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],false]],"m":"$undefined","G":["$11",[]],"S":true}
8:{}
9:"$0:f:0:1:1:children:1:children:0:props:children:0:props:serverProvidedParams:params"
e:[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]]
12:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"]
c:null
10:[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L12","4",{}]]
@@ -1,22 +0,0 @@
1:"$Sreact.fragment"
2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
6:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"]
7:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"]
a:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"]
b:"$Sreact.suspense"
d:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"]
f:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"]
11:I[168027,[],"default"]
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}]
0:{"P":null,"b":"zxkD4-EPlgfKHDTw8O869","c":["","chat"],"q":"","i":false,"f":[[["",{"children":["chat",{"children":["__PAGE__",{}]}]},"$undefined","$undefined",true],[["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true,"nonce":"$undefined"}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]],"forbidden":"$undefined","unauthorized":"$undefined"}]}]}]}]}]]}],{"children":[["$","$1","c",{"children":[null,["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":"$undefined","forbidden":"$undefined","unauthorized":"$undefined"}]]}],{"children":[["$","$1","c",{"children":[["$","$L6",null,{"Component":"$7","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@8","$@9"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true,"nonce":"$undefined"}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true,"nonce":"$undefined"}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true,"nonce":"$undefined"}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true,"nonce":"$undefined"}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true,"nonce":"$undefined"}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true,"nonce":"$undefined"}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true,"nonce":"$undefined"}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true,"nonce":"$undefined"}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true,"nonce":"$undefined"}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true,"nonce":"$undefined"}]],["$","$La",null,{"children":["$","$b",null,{"name":"Next.MetadataOutlet","children":"$@c"}]}]]}],{},null,false,false]},null,false,false]},null,false,false],["$","$1","h",{"children":[null,["$","$Ld",null,{"children":"$Le"}],["$","div",null,{"hidden":true,"children":["$","$Lf",null,{"children":["$","$b",null,{"name":"Next.Metadata","children":"$L10"}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],false]],"m":"$undefined","G":["$11",[]],"S":true}
8:{}
9:"$0:f:0:1:1:children:1:children:0:props:children:0:props:serverProvidedParams:params"
e:[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]]
12:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"]
c:null
10:[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L12","4",{}]]
@@ -1,6 +0,0 @@
1:"$Sreact.fragment"
2:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"]
3:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"]
4:"$Sreact.suspense"
5:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"]
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","h",{"children":[null,["$","$L2",null,{"children":[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]]}],["$","div",null,{"hidden":true,"children":["$","$L3",null,{"children":["$","$4",null,{"name":"Next.Metadata","children":[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L5","4",{}]]}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],"loading":null,"isPartial":false}
@@ -1,8 +0,0 @@
1:"$Sreact.fragment"
2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","template":["$","$L5",null,{}],"notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]]}]}]}]}]}]]}],"loading":null,"isPartial":false}
@@ -1,4 +0,0 @@
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}]
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","tree":{"name":"","paramType":null,"paramKey":"","hasRuntimePrefetch":false,"slots":{"children":{"name":"chat","paramType":null,"paramKey":"chat","hasRuntimePrefetch":false,"slots":{"children":{"name":"__PAGE__","paramType":null,"paramKey":"__PAGE__","hasRuntimePrefetch":false,"slots":null,"isRootLayout":false}},"isRootLayout":false}},"isRootLayout":true},"staleTime":300}
@@ -1,9 +0,0 @@
1:"$Sreact.fragment"
2:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"]
3:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"]
6:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"]
7:"$Sreact.suspense"
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[["$","$L2",null,{"Component":"$3","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@4","$@5"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true}]],["$","$L6",null,{"children":["$","$7",null,{"name":"Next.MetadataOutlet","children":"$@8"}]}]]}],"loading":null,"isPartial":false}
4:{}
5:"$0:rsc:props:children:0:props:serverProvidedParams:params"
8:null
@@ -1,4 +0,0 @@
1:"$Sreact.fragment"
2:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
3:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[null,["$","$L2",null,{"parallelRouterKey":"children","template":["$","$L3",null,{}]}]]}],"loading":null,"isPartial":false}
+65 -10
View File
@@ -6,9 +6,11 @@ from typing import Any, List, Optional, Tuple
from fastapi import HTTPException, Request, status
import litellm
from litellm import Router, provider_list
from litellm._logging import verbose_proxy_logger
from litellm.constants import STANDARD_CUSTOMER_ID_HEADERS
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
from litellm.proxy._types import *
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS
@@ -53,6 +55,12 @@ def _check_valid_ip(
def check_complete_credentials(request_body: dict) -> bool:
"""
if 'api_base' in request body. Check if complete credentials given. Prevent malicious attacks.
Supplying an ``api_key`` is necessary but not sufficient: even with
credentials supplied, an ``api_base`` / ``base_url`` that resolves to a
private/internal/cloud-metadata address would still allow the proxy to
be used as an SSRF pivot. Validate any URL fields here so the gate
can't be bypassed with ``api_key=anything`` plus a malicious target.
"""
given_model: Optional[str] = None
@@ -70,10 +78,27 @@ def check_complete_credentials(request_body: dict) -> bool:
return False
api_key_value = request_body.get("api_key")
if api_key_value and isinstance(api_key_value, str) and api_key_value.strip():
return True
if not (api_key_value and isinstance(api_key_value, str) and api_key_value.strip()):
return False
return False
# ``validate_url`` itself doesn't consult the toggle; ``safe_get`` /
# ``async_safe_get`` do. Mirror that here so admins who explicitly
# disabled URL validation (e.g. for an internal Ollama endpoint they
# accept the SSRF risk for) aren't blocked at the proxy boundary.
if getattr(litellm, "user_url_validation", False):
for url_field in ("api_base", "base_url"):
url_value = request_body.get(url_field)
if not url_value or not isinstance(url_value, str):
continue
try:
validate_url(url_value)
except SSRFError as e:
raise ValueError(
f"Rejected request: client-side {url_field}={url_value!r} "
f"is rejected by the SSRF guard ({e})."
)
return True
def check_regex_or_str_match(request_body_value: Any, regex_str: str) -> bool:
@@ -159,15 +184,42 @@ def is_request_body_safe(
"aws_web_identity_token",
"aws_role_name",
"vertex_credentials",
# Endpoint-targeting fields that retarget the outbound request or
# an observability callback. An attacker-controlled value either
# exfiltrates the request payload (incl. messages + admin-set
# tokens) to the attacker's host, or coerces the proxy into
# authenticating against the attacker's host with admin secrets.
"aws_bedrock_runtime_endpoint",
"langsmith_base_url",
"langfuse_host",
"posthog_host",
"braintrust_host",
"slack_webhook_url",
# Provider-specific endpoint overrides that flow into the outbound
# request via ``optional_params``. Same threat as ``api_base``:
# ``s3_endpoint_url`` redirects Bedrock file uploads to attacker
# S3; ``sagemaker_base_url`` redirects all SageMaker traffic;
# ``deployment_url`` redirects SAP deployments.
"s3_endpoint_url",
"sagemaker_base_url",
"deployment_url",
]
# The blocklist is enforced unconditionally. Legitimate clientside
# credential / endpoint passthrough goes through one of the two
# explicit admin opt-ins (``general_settings.allow_client_side_credentials``
# proxy-wide or ``configurable_clientside_auth_params`` per deployment).
# Historically there was a third, *implicit*, *caller-controlled* path:
# ``check_complete_credentials`` returned True when the caller supplied
# any non-empty ``api_key``, which made the entire blocklist a no-op.
# That bypass turned every missing entry on the blocklist into an
# exploitable SSRF / credential-exfil hole — see GHSA-jh89-88fc-qrfp,
# GHSA-3frq-6r6h-7j64, and the chain of veria-admin findings (Dv_m860l,
# b_yRJeQ5, stN90yjP, LBlyOAc8, U2TD78kg). Removed: the blocklist now
# has a single, predictable failure mode for missing entries (a 400),
# not a credential leak.
for param in banned_params:
if (
param in request_body
and not check_complete_credentials( # allow client-credentials to be passed to proxy
request_body=request_body
)
):
if param in request_body:
if general_settings.get("allow_client_side_credentials") is True:
return True
elif (
@@ -182,7 +234,10 @@ def is_request_body_safe(
return True
raise ValueError(
f"Rejected Request: {param} is not allowed in request body. "
"Enable with `general_settings::allow_client_side_credentials` on proxy config.yaml. "
"Clientside passthrough requires explicit admin opt-in via "
"either `general_settings.allow_client_side_credentials = true` "
"(proxy-wide) or `configurable_clientside_auth_params` on the "
"deployment in your proxy config.yaml. "
"Relevant Issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997",
)
+66 -6
View File
@@ -619,6 +619,67 @@ class ProxyBaseLLMRequestProcessing:
verbose_proxy_logger.error(f"Error setting custom headers: {e}")
return {}
@staticmethod
async def build_litellm_proxy_success_headers_from_llm_response(
*,
response: Any,
request_data: dict,
request: Request,
user_api_key_dict: UserAPIKeyAuth,
logging_obj: LiteLLMLoggingObj,
version: Optional[str],
proxy_logging_obj: ProxyLogging,
) -> Dict[str, str]:
"""
Build LiteLLM proxy response headers for routes that call the LLM directly
(e.g. Google native :generateContent) instead of base_process_llm_request.
"""
if isinstance(response, dict):
hidden_params = response.get("_hidden_params") or {}
else:
hidden_params = getattr(response, "_hidden_params", None) or {}
if not isinstance(hidden_params, dict):
hidden_params = {}
model_id = ProxyBaseLLMRequestProcessing._get_model_id_from_response(
hidden_params, request_data
)
cache_key = hidden_params.get("cache_key", None) or ""
api_base = hidden_params.get("api_base", None) or ""
response_cost = hidden_params.get("response_cost", None) or ""
fastest_response_batch_completion = hidden_params.get(
"fastest_response_batch_completion", None
)
additional_headers = hidden_params.get("additional_headers", {}) or {}
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=logging_obj.litellm_call_id,
model_id=model_id,
cache_key=cache_key,
api_base=api_base,
version=version,
response_cost=response_cost,
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
fastest_response_batch_completion=fastest_response_batch_completion,
request_data=request_data,
hidden_params=hidden_params,
litellm_logging_obj=logging_obj,
**additional_headers,
)
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
data=request_data,
user_api_key_dict=user_api_key_dict,
response=response,
request_headers=dict(request.headers),
)
if callback_headers:
custom_headers.update(callback_headers)
return custom_headers
async def common_processing_pre_call_logic(
self,
request: Request,
@@ -875,7 +936,7 @@ class ProxyBaseLLMRequestProcessing:
else:
verbose_proxy_logger.debug(
"Request received by LiteLLM:\n%s",
json.dumps(self.data, indent=4, default=str),
_payload_str,
)
async def base_process_llm_request( # noqa: PLR0915
@@ -1511,9 +1572,7 @@ class ProxyBaseLLMRequestProcessing:
_response = assembled_response
try:
from litellm.proxy.proxy_server import llm_router as _global_llm_router
from litellm.proxy.utils import (
_check_and_merge_model_level_guardrails,
)
from litellm.proxy.utils import _check_and_merge_model_level_guardrails
guardrail_data = _check_and_merge_model_level_guardrails(
data=captured_data, llm_router=_global_llm_router
@@ -1690,11 +1749,12 @@ class ProxyBaseLLMRequestProcessing:
elif isinstance(e, httpx.HTTPStatusError):
# Handle httpx.HTTPStatusError - extract actual error from response
# This matches the original behavior before the refactor in commit 511d435f6f
error_body = await e.response.aread()
http_status_error: httpx.HTTPStatusError = e
error_body = await http_status_error.response.aread()
error_text = error_body.decode("utf-8")
raise HTTPException(
status_code=e.response.status_code,
status_code=http_status_error.response.status_code,
detail={"error": error_text},
)
error_msg = f"{str(e)}"
@@ -0,0 +1,156 @@
"""
Expired UI session key cleanup manager.
Deletes expired virtual keys created for LiteLLM dashboard sessions.
"""
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache
from litellm.constants import (
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE,
LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
UI_SESSION_TOKEN_TEAM_ID,
)
from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken, UserAPIKeyAuth
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy.management_endpoints.key_management_endpoints import (
delete_verification_tokens,
)
from litellm.proxy.utils import PrismaClient
class ExpiredUISessionKeyCleanupManager:
"""
Cleans up expired UI session keys.
"""
def __init__(
self,
prisma_client: PrismaClient,
user_api_key_cache: DualCache,
pod_lock_manager=None,
):
self.prisma_client = prisma_client
self.user_api_key_cache = user_api_key_cache
self.pod_lock_manager = pod_lock_manager
async def cleanup_expired_keys(self) -> int:
"""
Main entry point for deleting expired UI session keys.
Uses PodLockManager to ensure only one pod runs cleanup in multi-pod deployments.
"""
lock_acquired = False
try:
if self.pod_lock_manager and self.pod_lock_manager.redis_cache:
lock_acquired = (
await self.pod_lock_manager.acquire_lock(
cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
)
or False
)
if not lock_acquired:
verbose_proxy_logger.debug(
"Expired UI session key cleanup: another pod is already "
"running cleanup or Redis lock acquisition failed - "
"skipping this cycle."
)
return 0
verbose_proxy_logger.info("Starting expired UI session key cleanup...")
expired_keys = await self._find_expired_ui_session_keys()
if not expired_keys:
verbose_proxy_logger.debug("No expired UI session keys found")
return 0
tokens = [key.token for key in expired_keys if key.token is not None]
if not tokens:
return 0
system_user = UserAPIKeyAuth.get_litellm_internal_jobs_user_api_key_auth()
response, keys_being_deleted = await delete_verification_tokens(
tokens=tokens,
user_api_key_cache=self.user_api_key_cache,
user_api_key_dict=system_user,
litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
)
await KeyManagementEventHooks.async_key_deleted_hook(
data=KeyRequest(keys=tokens),
keys_being_deleted=keys_being_deleted,
response=response or {},
user_api_key_dict=system_user,
litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
)
deleted_count = self._get_deleted_token_count(
tokens=tokens,
response=response,
)
verbose_proxy_logger.info(
"Deleted %s expired UI session key(s)", deleted_count
)
return deleted_count
except Exception as e:
if getattr(e, "status_code", None) == 404:
verbose_proxy_logger.debug(
"Expired UI session key cleanup skipped because selected keys "
"were already deleted: %s",
e,
)
return 0
verbose_proxy_logger.error(f"Expired UI session key cleanup failed: {e}")
return 0
finally:
if (
lock_acquired
and self.pod_lock_manager
and self.pod_lock_manager.redis_cache
):
await self.pod_lock_manager.release_lock(
cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
)
@staticmethod
def _get_deleted_token_count(
tokens: List[str],
response: Optional[Dict[str, Any]],
) -> int:
"""
Return the number of tokens actually deleted from the delete helper response.
"""
if response is None:
return len(tokens)
deleted_keys = response.get("deleted_keys")
if isinstance(deleted_keys, list):
return len(deleted_keys)
if isinstance(deleted_keys, int):
return deleted_keys
if isinstance(deleted_keys, dict):
nested_deleted_keys = deleted_keys.get("deleted_keys")
if isinstance(nested_deleted_keys, list):
return len(nested_deleted_keys)
if isinstance(nested_deleted_keys, int):
return nested_deleted_keys
failed_tokens = response.get("failed_tokens") or []
if failed_tokens:
return max(len(tokens) - len(set(failed_tokens)), 0)
return len(tokens)
async def _find_expired_ui_session_keys(self) -> List[LiteLLM_VerificationToken]:
"""
Find expired LiteLLM dashboard session keys.
"""
now = datetime.now(timezone.utc)
return await self.prisma_client.db.litellm_verificationtoken.find_many(
where={
"team_id": UI_SESSION_TOKEN_TEAM_ID,
"expires": {"lt": now},
},
take=LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE,
)
+137 -1
View File
@@ -1,5 +1,6 @@
from typing import Union
from typing import Any, Awaitable, Callable, Optional, Union
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
ProxyErrorTypes,
@@ -123,3 +124,138 @@ class PrismaDBExceptionHandler:
):
return None
raise e
# Default fallback timeouts when neither the caller nor the prisma_client
# expose `_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds`.
# Match the auth path's existing defaults so behavior is uniform across read paths.
_DEFAULT_RECONNECT_TIMEOUT_SECONDS = 2.0
_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS = 0.1
def _coerce_timeout(value: Any, fallback: float) -> float:
"""Return `value` if it is a real int/float, else `fallback`. Guards
against tests that mock `prisma_client` and leave the timeout slots as
MagicMock instances."""
if isinstance(value, (int, float)) and not isinstance(value, bool):
return float(value)
return fallback
async def call_with_db_reconnect_retry(
prisma_client: Any,
coro_factory: Callable[[], Awaitable[Any]],
*,
reason: str,
timeout_seconds: Optional[float] = None,
lock_timeout_seconds: Optional[float] = None,
) -> Any:
"""Run a Prisma read coroutine with one transport-reconnect-and-retry.
The canonical "self-heal a transient DB transport blip" wrapper used by
`PrismaClient.get_generic_data` and other read paths. Mirrors the inline
pattern in `auth_checks._fetch_key_object_from_db_with_reconnect` so we
have a single implementation rather than three drifting copies.
Behavior:
1. Await `coro_factory()`. On success, return its value.
2. On exception, if it is NOT a transport error (per
`is_database_transport_error`), re-raise data-layer errors like
`UniqueViolationError` mean the DB is reachable, reconnect would be
pointless.
3. If `prisma_client` does not expose `attempt_db_reconnect`, re-raise.
This guards against partial stand-ins / older clients in tests.
4. Call `prisma_client.attempt_db_reconnect(reason=...)`. If it returns
False (cooldown / lock contention / reconnect failure), re-raise.
5. Otherwise await `coro_factory()` a second time and return / propagate
its result. At-most-one retry by construction no infinite loop.
`coro_factory` MUST be a zero-arg callable that returns a fresh awaitable
on each call. Passing an already-awaited coroutine would fail on retry
with `RuntimeError: cannot reuse already awaited coroutine`.
`reason` should follow `<subsystem>_<operation>_<table>_failure` so
telemetry distinguishes between fan-out callers (e.g.
`_update_config_from_db` issues four concurrent reads).
Args:
prisma_client: The `PrismaClient` (or stand-in) that owns
`attempt_db_reconnect` and the `_db_auth_reconnect_*` defaults.
coro_factory: Zero-arg callable returning the read awaitable.
reason: Telemetry tag forwarded to `attempt_db_reconnect`.
timeout_seconds: Optional override for the reconnect cycle timeout.
Defaults to `prisma_client._db_auth_reconnect_timeout_seconds`,
then to 2.0s.
lock_timeout_seconds: Optional override for how long the helper will
wait to acquire the reconnect lock. Defaults to
`prisma_client._db_auth_reconnect_lock_timeout_seconds`, then to
0.1s.
Returns:
Whatever `coro_factory()` returns (on first or second attempt).
Raises:
Whatever `coro_factory()` raises if the failure is not a transport
error, or if the reconnect attempt does not succeed, or if the retry
also fails.
"""
try:
return await coro_factory()
except Exception as first_exc:
if not PrismaDBExceptionHandler.is_database_transport_error(first_exc):
raise
if not hasattr(prisma_client, "attempt_db_reconnect"):
raise
resolved_timeout = _coerce_timeout(
(
timeout_seconds
if timeout_seconds is not None
else getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", None)
),
_DEFAULT_RECONNECT_TIMEOUT_SECONDS,
)
resolved_lock_timeout = _coerce_timeout(
(
lock_timeout_seconds
if lock_timeout_seconds is not None
else getattr(
prisma_client, "_db_auth_reconnect_lock_timeout_seconds", None
)
),
_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS,
)
verbose_proxy_logger.warning(
"DB transport error on read; attempting reconnect-and-retry. reason=%s error=%s",
reason,
first_exc,
)
# Preserve the original transport error in telemetry. If
# `attempt_db_reconnect` itself raises (e.g. lock cancellation, timer
# error, unexpected internal failure), surfacing that exception
# instead of `first_exc` would mask the actual DB transport problem
# in `failure_handler` / `db_exceptions` alerts. Chain the reconnect
# error as the cause for debuggability without losing the original.
try:
did_reconnect = await prisma_client.attempt_db_reconnect(
reason=reason,
timeout_seconds=resolved_timeout,
lock_timeout_seconds=resolved_lock_timeout,
)
except Exception as reconnect_exc:
verbose_proxy_logger.warning(
"DB reconnect attempt raised; preserving original transport error. "
"reason=%s reconnect_error=%s",
reason,
reconnect_exc,
)
raise first_exc from reconnect_exc
if not did_reconnect:
raise
# At most one retry. If the retry also raises a transport error, we
# propagate — repeated reconnect-loops are the watchdog's job, not
# this helper's.
return await coro_factory()
+22 -12
View File
@@ -52,18 +52,25 @@ class PrismaWrapper:
engine = self._original_prisma._engine
process = getattr(engine, "process", None) if engine is not None else None
if process is not None:
return process.pid
pid = process.pid
if isinstance(pid, int):
return pid
except (AttributeError, TypeError):
pass
return 0
@staticmethod
async def _kill_engine_process(pid: int) -> None:
"""Force-kill an orphaned engine subprocess to prevent DB connection pool leaks.
"""Force-kill the engine subprocess to prevent DB connection pool leaks.
Called when disconnect() fails and the old engine process may still be
holding open connections. Sends SIGTERM for graceful shutdown, waits
briefly, then SIGKILL as a backstop.
Called on every reconnect (in `recreate_prisma_client`) to retire the
old query-engine subprocess without invoking prisma-client-py's
synchronous `disconnect()` which blocks the asyncio event loop on
`subprocess.Popen.wait()` for 30-120+ seconds when the engine is
stuck on TCP close.
Sends SIGTERM for graceful shutdown, waits briefly, then SIGKILL as
a backstop.
"""
if pid <= 0:
return
@@ -72,7 +79,7 @@ class PrismaWrapper:
except (ProcessLookupError, PermissionError, OSError):
return # Already dead or inaccessible
verbose_proxy_logger.warning(
"Sent SIGTERM to orphaned prisma-query-engine PID %s after failed disconnect.",
"Sent SIGTERM to prisma-query-engine PID %s during reconnect.",
pid,
)
# Brief wait for graceful shutdown, then force-kill
@@ -217,15 +224,18 @@ class PrismaWrapper:
async def recreate_prisma_client(
self, new_db_url: str, http_client: Optional[Any] = None
):
"""Disconnect and reconnect the Prisma client with a new database URL."""
"""Disconnect and reconnect the Prisma client with a new database URL.
Kills the old engine subprocess directly (SIGTERM SIGKILL) rather than
calling `disconnect()`. prisma-client-py's `disconnect()` calls a
synchronous `subprocess.Popen.wait()` that can freeze the asyncio event
loop for 30-120+ seconds when the engine is stuck on TCP close,
breaking `/health/liveliness` and causing Kubernetes pod restarts.
"""
from prisma import Prisma # type: ignore
old_engine_pid = self._get_engine_pid()
try:
await self._original_prisma.disconnect()
except Exception as e:
verbose_proxy_logger.warning(f"Failed to disconnect Prisma client: {e}")
if old_engine_pid > 0:
await self._kill_engine_process(old_engine_pid)
if http_client is not None:
+28 -1
View File
@@ -35,6 +35,7 @@ async def google_generate_content(
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
version,
)
@@ -73,6 +74,16 @@ async def google_generate_content(
if llm_router is None:
raise HTTPException(status_code=500, detail="Router not initialized")
response = await llm_router.agenerate_content(**data)
success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
response=response,
request_data=data,
request=request,
user_api_key_dict=user_api_key_dict,
logging_obj=logging_obj,
version=version,
proxy_logging_obj=proxy_logging_obj,
)
fastapi_response.headers.update(success_headers)
return response
@@ -95,6 +106,7 @@ async def google_stream_generate_content(
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
version,
)
@@ -137,9 +149,24 @@ async def google_stream_generate_content(
raise HTTPException(status_code=500, detail="Router not initialized")
response = await llm_router.agenerate_content_stream(**data)
success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
response=response,
request_data=data,
request=request,
user_api_key_dict=user_api_key_dict,
logging_obj=logging_obj,
version=version,
proxy_logging_obj=proxy_logging_obj,
)
# 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 StreamingResponse(
content=response,
media_type="text/event-stream",
headers=success_headers,
)
fastapi_response.headers.update(success_headers)
return response
@@ -0,0 +1,492 @@
"""
WORKFLOW RUN MANAGEMENT
Generic durable state tracking for agents and automated workflows.
POST /v1/workflows/runs - Create a workflow run
GET /v1/workflows/runs - List runs (filter by type, status)
GET /v1/workflows/runs/{run_id} - Get run with latest event
PATCH /v1/workflows/runs/{run_id} - Update status, metadata, output
POST /v1/workflows/runs/{run_id}/events - Append event (updates run status)
GET /v1/workflows/runs/{run_id}/events - Full event log
POST /v1/workflows/runs/{run_id}/messages - Append conversation message
GET /v1/workflows/runs/{run_id}/messages - Fetch conversation history
"""
import json
from typing import Any, Dict, Literal, Optional
from fastapi import APIRouter, Depends, HTTPException, Query
try:
from prisma.errors import UniqueViolationError
except ImportError:
UniqueViolationError = None # type: ignore
from pydantic import BaseModel
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
router = APIRouter()
_MAX_SEQUENCE_RETRIES = 5
def _json(value: Any) -> str:
"""Serialize a Python value for prisma-client-py Json fields (must be a string)."""
return json.dumps(value)
def _is_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
def _caller_key(user_api_key_dict: UserAPIKeyAuth) -> Optional[str]:
"""Return the hashed key token that identifies this caller, or None for master key."""
return user_api_key_dict.token
# Status transitions driven by event_type
_EVENT_STATUS_MAP: Dict[str, str] = {
"step.started": "running",
"step.failed": "failed",
"hook.waiting": "paused",
"hook.received": "running",
}
# ---------------------------------------------------------------------------
# Request / Response models
# ---------------------------------------------------------------------------
class WorkflowRunCreateRequest(BaseModel):
workflow_type: str
input: Optional[Dict[str, Any]] = None
metadata: Optional[Dict[str, Any]] = None
WorkflowRunStatus = Literal["pending", "running", "paused", "completed", "failed"]
class WorkflowRunUpdateRequest(BaseModel):
status: Optional[WorkflowRunStatus] = None
output: Optional[Dict[str, Any]] = None
metadata: Optional[Dict[str, Any]] = None
class WorkflowEventCreateRequest(BaseModel):
event_type: str
step_name: str
data: Optional[Dict[str, Any]] = None
class WorkflowMessageCreateRequest(BaseModel):
role: str
content: str
session_id: Optional[str] = None
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
async def _get_next_sequence_number(prisma_client: Any, run_id: str, table: str) -> int:
"""Return MAX(sequence_number) + 1 for the given run, for either events or messages."""
if table == "events":
rows = await prisma_client.db.litellm_workflowevent.find_many(
where={"run_id": run_id},
order={"sequence_number": "desc"},
take=1,
)
else:
rows = await prisma_client.db.litellm_workflowmessage.find_many(
where={"run_id": run_id},
order={"sequence_number": "desc"},
take=1,
)
return (rows[0].sequence_number + 1) if rows else 0
async def _require_run(
prisma_client: Any,
run_id: str,
user_api_key_dict: Optional[UserAPIKeyAuth] = None,
) -> Any:
"""Return the run or raise 404. For non-admin callers, also enforce key ownership."""
run = await prisma_client.db.litellm_workflowrun.find_unique(
where={"run_id": run_id}
)
if run is None:
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
if user_api_key_dict is not None and not _is_admin(user_api_key_dict):
caller = _caller_key(user_api_key_dict)
if not caller or run.created_by != caller:
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
return run
# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------
@router.post(
"/v1/workflows/runs",
tags=["workflow management"],
dependencies=[Depends(user_api_key_auth)],
)
async def create_workflow_run(
data: WorkflowRunCreateRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""Create a new workflow run. Returns run_id and session_id.
The caller's API key token is stored as created_by so that non-admin keys
can only see and modify their own runs.
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
try:
create_data: Dict[str, Any] = {
"workflow_type": data.workflow_type,
"created_by": _caller_key(user_api_key_dict),
}
if data.input is not None:
create_data["input"] = _json(data.input)
if data.metadata is not None:
create_data["metadata"] = _json(data.metadata)
run = await prisma_client.db.litellm_workflowrun.create(data=create_data)
return run
except Exception as e:
verbose_proxy_logger.exception("Error creating workflow run: %s", e)
raise HTTPException(status_code=500, detail=str(e))
@router.get(
"/v1/workflows/runs",
tags=["workflow management"],
dependencies=[Depends(user_api_key_auth)],
)
async def list_workflow_runs(
workflow_type: Optional[str] = Query(None),
status: Optional[str] = Query(None),
limit: int = Query(50, ge=1, le=250),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""List workflow runs. Filter by workflow_type and/or status.
Non-admin callers only see runs created by their own API key.
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
where: Dict[str, Any] = {}
if workflow_type:
where["workflow_type"] = workflow_type
if status:
statuses = [s.strip() for s in status.split(",")]
where["status"] = {"in": statuses} if len(statuses) > 1 else statuses[0]
# Non-admin callers are scoped to their own key.
if not _is_admin(user_api_key_dict):
caller = _caller_key(user_api_key_dict)
if caller:
where["created_by"] = caller
try:
runs = await prisma_client.db.litellm_workflowrun.find_many(
where=where,
order={"created_at": "desc"},
take=limit,
)
return {"runs": runs, "count": len(runs)}
except Exception as e:
verbose_proxy_logger.exception("Error listing workflow runs: %s", e)
raise HTTPException(status_code=500, detail=str(e))
@router.get(
"/v1/workflows/runs/{run_id}",
tags=["workflow management"],
dependencies=[Depends(user_api_key_auth)],
)
async def get_workflow_run(
run_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""Get a workflow run with its most recent event."""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
try:
run = await prisma_client.db.litellm_workflowrun.find_unique(
where={"run_id": run_id},
include={"events": {"order_by": {"sequence_number": "desc"}, "take": 1}},
)
if run is None:
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
if not _is_admin(user_api_key_dict):
caller = _caller_key(user_api_key_dict)
if not caller or run.created_by != caller:
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
return run
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception("Error getting workflow run: %s", e)
raise HTTPException(status_code=500, detail=str(e))
@router.patch(
"/v1/workflows/runs/{run_id}",
tags=["workflow management"],
dependencies=[Depends(user_api_key_auth)],
)
async def update_workflow_run(
run_id: str,
data: WorkflowRunUpdateRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""Update status, metadata, or output on a workflow run."""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
update: Dict[str, Any] = {}
if data.status is not None:
update["status"] = data.status
if data.output is not None:
update["output"] = _json(data.output)
if data.metadata is not None:
update["metadata"] = _json(data.metadata)
if not update:
raise HTTPException(status_code=400, detail="No fields to update")
# Enforce ownership before writing.
await _require_run(prisma_client, run_id, user_api_key_dict)
try:
run = await prisma_client.db.litellm_workflowrun.update(
where={"run_id": run_id},
data=update,
)
if run is None:
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
return run
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception("Error updating workflow run: %s", e)
raise HTTPException(status_code=500, detail=str(e))
@router.post(
"/v1/workflows/runs/{run_id}/events",
tags=["workflow management"],
dependencies=[Depends(user_api_key_auth)],
)
async def append_workflow_event(
run_id: str,
data: WorkflowEventCreateRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""Append an event to the run's event log. Also updates run.status if event_type maps to a status.
Sequence numbers use optimistic concurrency: on a unique-constraint collision
(concurrent append), retries up to _MAX_SEQUENCE_RETRIES times with a fresh MAX+1.
The event+status update is atomic in a single DB transaction.
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
await _require_run(prisma_client, run_id, user_api_key_dict)
new_status = _EVENT_STATUS_MAP.get(data.event_type)
for attempt in range(_MAX_SEQUENCE_RETRIES):
try:
seq = await _get_next_sequence_number(prisma_client, run_id, "events")
event_data: Dict[str, Any] = {
"run_id": run_id,
"event_type": data.event_type,
"step_name": data.step_name,
"sequence_number": seq,
}
if data.data is not None:
event_data["data"] = _json(data.data)
async with prisma_client.db.tx() as tx:
event = await tx.litellm_workflowevent.create(data=event_data)
if new_status:
await tx.litellm_workflowrun.update(
where={"run_id": run_id},
data={"status": new_status},
)
return event
except Exception as e:
if UniqueViolationError is not None and isinstance(e, UniqueViolationError):
if attempt == _MAX_SEQUENCE_RETRIES - 1:
verbose_proxy_logger.exception(
"Sequence number collision after %d retries for run %s",
_MAX_SEQUENCE_RETRIES,
run_id,
)
raise HTTPException(
status_code=409,
detail="Concurrent write conflict — please retry",
)
continue
verbose_proxy_logger.exception("Error appending workflow event: %s", e)
raise HTTPException(status_code=500, detail=str(e))
raise HTTPException(
status_code=500, detail="Failed to append event"
) # pragma: no cover
@router.get(
"/v1/workflows/runs/{run_id}/events",
tags=["workflow management"],
dependencies=[Depends(user_api_key_auth)],
)
async def list_workflow_events(
run_id: str,
limit: int = Query(100, ge=1, le=500),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""Fetch event log for a run, ordered by sequence_number. Default limit 100, max 500."""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
await _require_run(prisma_client, run_id, user_api_key_dict)
try:
events = await prisma_client.db.litellm_workflowevent.find_many(
where={"run_id": run_id},
order={"sequence_number": "asc"},
take=limit,
)
return {"events": events, "count": len(events)}
except Exception as e:
verbose_proxy_logger.exception("Error listing workflow events: %s", e)
raise HTTPException(status_code=500, detail=str(e))
@router.post(
"/v1/workflows/runs/{run_id}/messages",
tags=["workflow management"],
dependencies=[Depends(user_api_key_auth)],
)
async def append_workflow_message(
run_id: str,
data: WorkflowMessageCreateRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""Append a conversation message. Stores full content (not truncated).
Uses optimistic concurrency for sequence numbers.
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
await _require_run(prisma_client, run_id, user_api_key_dict)
for attempt in range(_MAX_SEQUENCE_RETRIES):
try:
seq = await _get_next_sequence_number(prisma_client, run_id, "messages")
msg_data: Dict[str, Any] = {
"run_id": run_id,
"role": data.role,
"content": data.content,
"sequence_number": seq,
}
if data.session_id is not None:
msg_data["session_id"] = data.session_id
msg = await prisma_client.db.litellm_workflowmessage.create(data=msg_data)
return msg
except Exception as e:
if UniqueViolationError is not None and isinstance(e, UniqueViolationError):
if attempt == _MAX_SEQUENCE_RETRIES - 1:
verbose_proxy_logger.exception(
"Sequence number collision after %d retries for run %s",
_MAX_SEQUENCE_RETRIES,
run_id,
)
raise HTTPException(
status_code=409,
detail="Concurrent write conflict — please retry",
)
continue
verbose_proxy_logger.exception("Error appending workflow message: %s", e)
raise HTTPException(status_code=500, detail=str(e))
raise HTTPException(
status_code=500, detail="Failed to append message"
) # pragma: no cover
@router.get(
"/v1/workflows/runs/{run_id}/messages",
tags=["workflow management"],
dependencies=[Depends(user_api_key_auth)],
)
async def list_workflow_messages(
run_id: str,
limit: int = Query(100, ge=1, le=500),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""Fetch conversation history for a run, ordered by sequence_number. Default limit 100, max 500."""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
await _require_run(prisma_client, run_id, user_api_key_dict)
try:
messages = await prisma_client.db.litellm_workflowmessage.find_many(
where={"run_id": run_id},
order={"sequence_number": "asc"},
take=limit,
)
return {"messages": messages, "count": len(messages)}
except Exception as e:
verbose_proxy_logger.exception("Error listing workflow messages: %s", e)
raise HTTPException(status_code=500, detail=str(e))
+133 -16
View File
@@ -426,6 +426,9 @@ from litellm.proxy.management_endpoints.team_endpoints import (
from litellm.proxy.management_endpoints.tool_management_endpoints import (
router as tool_management_router,
)
from litellm.proxy.management_endpoints.workflow_management_endpoints import (
router as workflow_management_router,
)
from litellm.proxy.memory.memory_endpoints import router as memory_router
from litellm.proxy.management_endpoints.ui_sso import (
get_disabled_non_admin_personal_key_creation,
@@ -497,14 +500,18 @@ from litellm.proxy.utils import (
_get_redoc_url,
_is_projected_spend_over_limit,
_is_valid_team_configs,
get_config_param,
get_custom_url,
get_error_message_str,
get_server_root_path,
handle_exception_on_proxy,
hash_password,
hash_token,
invalidate_config_param,
litellm_config_cache,
migrate_passwords_to_scrypt_async,
model_dump_with_preserved_fields,
prefetch_config_params,
update_spend,
)
from litellm.proxy.vector_store_endpoints.endpoints import router as vector_store_router
@@ -738,6 +745,10 @@ async def _initialize_shared_aiohttp_session():
try:
from aiohttp import ClientSession, TCPConnector
from litellm.llms.custom_httpx.http_handler import (
_build_aiohttp_keepalive_socket_factory,
)
connector_kwargs: Dict[str, Any] = {
"keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT,
"ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE,
@@ -748,6 +759,9 @@ async def _initialize_shared_aiohttp_session():
connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT
if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0:
connector_kwargs["limit_per_host"] = AIOHTTP_CONNECTOR_LIMIT_PER_HOST
socket_factory = _build_aiohttp_keepalive_socket_factory()
if socket_factory is not None:
connector_kwargs["socket_factory"] = socket_factory
connector = TCPConnector(**connector_kwargs)
session = ClientSession(connector=connector)
@@ -2929,8 +2943,13 @@ class ProxyConfig:
## INIT PROXY REDIS USAGE CLIENT ##
redis_usage_cache = litellm.cache.cache
spend_counter_cache.redis_cache = redis_usage_cache
litellm_config_cache.redis_cache = redis_usage_cache
# Note: PKCE verifier storage uses redis_usage_cache directly (not
# user_api_key_cache) to avoid routing all API-key lookups through Redis.
elif litellm_config_cache.redis_cache is None:
verbose_proxy_logger.info(
"litellm_config_cache: no Redis configured; cluster-wide cache sharing disabled."
)
def switch_on_llm_response_caching(self):
"""
@@ -4846,10 +4865,7 @@ class ProxyConfig:
"environment_variables",
]
for k in keys:
response = prisma_client.get_generic_data(
key="param_name", value=k, table_name="config"
)
_tasks.append(response)
_tasks.append(get_config_param(prisma_client, k))
responses = await asyncio.gather(*_tasks)
for response in responses:
@@ -4931,6 +4947,19 @@ class ProxyConfig:
global llm_router, llm_model_list, master_key, general_settings
try:
# warm the config cache so the per-param reads below all hit
await prefetch_config_params(
prisma_client,
[
"general_settings",
"router_settings",
"litellm_settings",
"environment_variables",
"model_cost_map_reload_config",
"anthropic_beta_headers_reload_config",
],
)
# Only load models from DB if "models" is in supported_db_objects (or if supported_db_objects is not set)
if self._should_load_db_object(object_type="models"):
new_models = await self._get_models_from_db(prisma_client=prisma_client)
@@ -4940,8 +4969,8 @@ class ProxyConfig:
new_models=new_models, proxy_logging_obj=proxy_logging_obj
)
db_general_settings = await prisma_client.db.litellm_config.find_first(
where={"param_name": "general_settings"}
db_general_settings = await get_config_param(
prisma_client, "general_settings"
)
# update general settings
@@ -5034,10 +5063,7 @@ class ProxyConfig:
from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook
try:
# Load litellm_settings from DB
config_record = await prisma_client.db.litellm_config.find_unique(
where={"param_name": "litellm_settings"}
)
config_record = await get_config_param(prisma_client, "litellm_settings")
if config_record is None or config_record.param_value is None:
return
@@ -5192,8 +5218,8 @@ class ProxyConfig:
"""
try:
# Get model cost map reload configuration from database
config_record = await prisma_client.db.litellm_config.find_unique(
where={"param_name": "model_cost_map_reload_config"}
config_record = await get_config_param(
prisma_client, "model_cost_map_reload_config"
)
if config_record is None or config_record.param_value is None:
@@ -5288,6 +5314,7 @@ class ProxyConfig:
},
},
)
await invalidate_config_param("model_cost_map_reload_config")
verbose_proxy_logger.info(
f"Model cost map reloaded successfully. Models count: {len(new_model_cost_map) if new_model_cost_map else 0}"
@@ -5307,8 +5334,8 @@ class ProxyConfig:
"""
try:
# Get anthropic beta headers reload configuration from database
config_record = await prisma_client.db.litellm_config.find_unique(
where={"param_name": "anthropic_beta_headers_reload_config"}
config_record = await get_config_param(
prisma_client, "anthropic_beta_headers_reload_config"
)
if config_record is None or config_record.param_value is None:
@@ -5396,6 +5423,7 @@ class ProxyConfig:
},
},
)
await invalidate_config_param("anthropic_beta_headers_reload_config")
# Count providers in config
provider_count = sum(
@@ -6688,6 +6716,10 @@ class ProxyStartupEvent:
Args:
scheduler: The scheduler to add the background jobs to
"""
global prisma_client
global proxy_logging_obj
global user_api_key_cache
########################################################
# CloudZero Background Job
########################################################
@@ -6761,8 +6793,6 @@ class ProxyStartupEvent:
)
# Get prisma_client and proxy_logging_obj from global scope
global prisma_client
global proxy_logging_obj
if prisma_client is not None:
# Reuse the PodLockManager from db_spend_update_writer
pod_lock_manager = (
@@ -6792,6 +6822,83 @@ class ProxyStartupEvent:
"Key rotation disabled (set LITELLM_KEY_ROTATION_ENABLED=true to enable)"
)
await cls._initialize_expired_ui_session_key_cleanup_background_job(
scheduler=scheduler
)
@classmethod
async def _initialize_expired_ui_session_key_cleanup_background_job(
cls, scheduler: AsyncIOScheduler
):
"""
Initialize the expired UI session key cleanup background job.
"""
global prisma_client
global proxy_logging_obj
global user_api_key_cache
########################################################
# Expired UI Session Key Cleanup Background Job
########################################################
from litellm.constants import (
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED,
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS,
)
expired_ui_session_key_cleanup_enabled: Optional[bool] = str_to_bool(
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED
)
verbose_proxy_logger.debug(
"expired_ui_session_key_cleanup_enabled: "
f"{expired_ui_session_key_cleanup_enabled}"
)
if expired_ui_session_key_cleanup_enabled is True:
try:
from litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager import (
ExpiredUISessionKeyCleanupManager,
)
if prisma_client is not None:
pod_lock_manager = (
proxy_logging_obj.db_spend_update_writer.pod_lock_manager
)
expired_ui_session_key_cleanup_manager = (
ExpiredUISessionKeyCleanupManager(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
pod_lock_manager=pod_lock_manager,
)
)
verbose_proxy_logger.debug(
"Expired UI session key cleanup background job scheduled "
"every "
f"{LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS} "
"seconds "
"(LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true)"
)
scheduler.add_job(
expired_ui_session_key_cleanup_manager.cleanup_expired_keys,
"interval",
seconds=LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS,
id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
)
else:
verbose_proxy_logger.warning(
"Expired UI session key cleanup enabled but prisma_client "
"not available"
)
except Exception as e:
verbose_proxy_logger.warning(
f"Failed to setup expired UI session key cleanup job: {e}"
)
else:
verbose_proxy_logger.debug(
"Expired UI session key cleanup disabled (set "
"LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true to enable)"
)
@classmethod
async def _initialize_slack_alerting_jobs(
cls,
@@ -12633,6 +12740,7 @@ async def update_config( # noqa: PLR0915
updated_litellm_settings["success_callback"] = normalize_callback_names(
incoming_cb
)
await invalidate_config_param(k)
merged = {**existing, **updated_litellm_settings}
@@ -12770,6 +12878,7 @@ async def update_config_general_settings(
"update": {"param_value": json.dumps(general_settings)}, # type: ignore
},
)
await invalidate_config_param("general_settings")
return response
@@ -13053,6 +13162,7 @@ async def delete_config_general_settings(
"update": {"param_value": json.dumps(general_settings)}, # type: ignore
},
)
await invalidate_config_param("general_settings")
return response
@@ -13418,6 +13528,7 @@ async def reload_model_cost_map(
},
},
)
await invalidate_config_param("model_cost_map_reload_config")
models_count = len(new_model_cost_map) if new_model_cost_map else 0
verbose_proxy_logger.info(
@@ -13487,6 +13598,7 @@ async def schedule_model_cost_map_reload(
},
},
)
await invalidate_config_param("model_cost_map_reload_config")
verbose_proxy_logger.info(
f"Model cost map reload scheduled for every {hours} hours"
@@ -13540,6 +13652,7 @@ async def cancel_model_cost_map_reload(
await prisma_client.db.litellm_config.delete(
where={"param_name": "model_cost_map_reload_config"}
)
await invalidate_config_param("model_cost_map_reload_config")
verbose_proxy_logger.info("Model cost map reload schedule cancelled")
@@ -13770,6 +13883,7 @@ async def reload_anthropic_beta_headers(
},
},
)
await invalidate_config_param("anthropic_beta_headers_reload_config")
provider_count = sum(
1 for k in new_config.keys() if k not in ["provider_aliases", "description"]
@@ -13843,6 +13957,7 @@ async def schedule_anthropic_beta_headers_reload(
},
},
)
await invalidate_config_param("anthropic_beta_headers_reload_config")
verbose_proxy_logger.info(
f"Anthropic beta headers reload scheduled for every {hours} hours"
@@ -13896,6 +14011,7 @@ async def cancel_anthropic_beta_headers_reload(
await prisma_client.db.litellm_config.delete(
where={"param_name": "anthropic_beta_headers_reload_config"}
)
await invalidate_config_param("anthropic_beta_headers_reload_config")
verbose_proxy_logger.info("Anthropic beta headers reload schedule cancelled")
@@ -14158,6 +14274,7 @@ app.include_router(model_management_router)
app.include_router(model_access_group_management_router)
app.include_router(tag_management_router)
app.include_router(tool_management_router)
app.include_router(workflow_management_router)
app.include_router(memory_router)
app.include_router(cost_tracking_settings_router)
app.include_router(router_settings_router)
+77
View File
@@ -1290,3 +1290,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])
}
+138 -23
View File
@@ -106,7 +106,10 @@ from litellm.proxy.db.create_views import (
should_create_missing_views,
)
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.db.exception_handler import (
PrismaDBExceptionHandler,
call_with_db_reconnect_retry,
)
from litellm.proxy.db.log_db_metrics import log_db_metrics
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
@@ -2442,6 +2445,92 @@ async def _lookup_deprecated_key(
return None
# DualCache for LiteLLM_Config param_name reads.
# Redis layer is attached in proxy_server._init_cache.
LITELLM_CONFIG_CACHE_TTL_SECONDS: int = int(
os.environ.get("LITELLM_CONFIG_PARAM_CACHE_TTL_SECONDS", "60")
)
_CONFIG_CACHE_MISS: str = "__litellm_config_param_miss__"
litellm_config_cache: DualCache = DualCache(
default_in_memory_ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS,
default_redis_ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS,
)
class _ConfigRow:
"""Mimics the Prisma litellm_config row shape for cached entries."""
__slots__ = ("param_name", "param_value")
def __init__(self, param_name: str, param_value: Any) -> None:
self.param_name = param_name
self.param_value = param_value
def _config_cache_key(param_name: str) -> str:
return f"litellm_config:param:{param_name}"
def _pack_config_row(row: Any) -> Dict[str, Any]:
return {"param_name": row.param_name, "param_value": row.param_value}
def _unpack_config_row(cached: Any) -> Optional[_ConfigRow]:
if cached is None or cached == _CONFIG_CACHE_MISS:
return None
if isinstance(cached, dict):
return _ConfigRow(cached["param_name"], cached["param_value"])
return None
async def get_config_param(prisma_client: Any, param_name: str) -> Optional[Any]:
"""Cached read of a LiteLLM_Config row; returns row, _ConfigRow shim, or None."""
cache_key = _config_cache_key(param_name)
cached = await litellm_config_cache.async_get_cache(cache_key)
if cached is not None:
return _unpack_config_row(cached)
row = await prisma_client.get_generic_data(
key="param_name", value=param_name, table_name="config"
)
cache_value: Any = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS
await litellm_config_cache.async_set_cache(
cache_key, cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS
)
return row
async def invalidate_config_param(param_name: str) -> None:
"""Evict from both cache layers; call after every LiteLLM_Config write."""
await litellm_config_cache.async_delete_cache(_config_cache_key(param_name))
async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> None:
"""Batch-load LiteLLM_Config rows into the cache with one find_many."""
if not param_names:
return
try:
rows = await prisma_client.db.litellm_config.find_many(
where={"param_name": {"in": param_names}} # type: ignore
)
except Exception as e:
verbose_proxy_logger.debug(
"prefetch_config_params failed, falling through to per-param queries: %s",
e,
)
return
by_name = {row.param_name: row for row in rows}
for name in param_names:
row = by_name.get(name)
cache_value: Any = (
_pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS
)
await litellm_config_cache.async_set_cache(
_config_cache_key(name), cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS
)
class PrismaClient:
spend_log_transactions: List = []
_spend_log_transactions_lock = asyncio.Lock()
@@ -2693,30 +2782,42 @@ class PrismaClient:
table_name: Literal["users", "keys", "config", "spend"],
):
"""
Generic implementation of get data
Generic implementation of get data.
Self-heals across a single transient transport blip via
`call_with_db_reconnect_retry`: on `httpx.ReadError` /
`ClientNotConnectedError` / similar, attempt one DB reconnect and
retry once before surfacing the failure. Restores the 1.82.6 behavior
that was lost in 1.83.x see issue #25143.
"""
start_time = time.time()
try:
async def _do_query():
if table_name == "users":
response = await self.db.litellm_usertable.find_first(
return await self.db.litellm_usertable.find_first(
where={key: value} # type: ignore
)
elif table_name == "keys":
response = await self.db.litellm_verificationtoken.find_first( # type: ignore
return await self.db.litellm_verificationtoken.find_first( # type: ignore
where={key: value} # type: ignore
)
elif table_name == "config":
response = await self.db.litellm_config.find_first( # type: ignore
return await self.db.litellm_config.find_first( # type: ignore
where={key: value} # type: ignore
)
elif table_name == "spend":
response = await self.db.l.find_first( # type: ignore
return await self.db.l.find_first( # type: ignore
where={key: value} # type: ignore
)
return response
except Exception as e:
import traceback
return None
try:
return await call_with_db_reconnect_retry(
self,
_do_query,
reason=f"prisma_get_generic_data_{table_name}_lookup_failure",
)
except Exception as e:
error_msg = f"LiteLLM Prisma Client Exception get_generic_data: {str(e)}"
verbose_proxy_logger.error(error_msg)
error_msg = error_msg + "\nException Type: {}".format(type(e))
@@ -3310,6 +3411,9 @@ class PrismaClient:
tasks.append(updated_table_row)
await asyncio.gather(*tasks)
# invalidate cache so other pods see writes from save_config
for k in data.keys():
await invalidate_config_param(k)
verbose_proxy_logger.info("Data Inserted into Config Table")
elif table_name == "spend":
db_data = self.jsonify_object(data=data)
@@ -4094,8 +4198,11 @@ class PrismaClient:
Uses the _engine_confirmed_dead flag (set by waitpid thread / pidfd / poll
handlers) to choose between heavy reconnect (engine dead -- recreate
Prisma client, re-arm watcher) and lightweight reconnect (network
blip -- disconnect, connect, SELECT 1).
Prisma client, re-arm watcher) and direct reconnect (network blip --
recreate Prisma client, re-arm watcher, SELECT 1). Both paths recreate
the client via the non-blocking kill-then-construct flow rather than
calling disconnect(), which blocks the event loop on the synchronous
subprocess.Popen.wait() inside prisma-client-py (see issue #26191).
"""
effective_timeout = (
timeout_seconds
@@ -4115,7 +4222,6 @@ class PrismaClient:
)
self._reap_all_zombies()
self._cleanup_engine_watcher()
self._engine_confirmed_dead = False
async def _do_heavy_reconnect() -> None:
db_url = os.getenv("DATABASE_URL", "")
@@ -4128,23 +4234,32 @@ class PrismaClient:
await self._start_engine_watcher()
await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout)
# Only clear the "dead engine" flag after the heavy reconnect
# actually completed. If `_do_heavy_reconnect()` raises (timeout,
# missing DATABASE_URL, recreate failure), the flag stays True so
# the next attempt re-enters the heavy branch instead of silently
# demoting to the lightweight path.
self._engine_confirmed_dead = False
else:
verbose_proxy_logger.debug(
"Performing Prisma DB reconnect (engine alive or unknown)."
)
async def _do_direct_reconnect() -> None:
old_pid = self._get_engine_pid()
try:
await self.db.disconnect()
except Exception as disconnect_err:
verbose_proxy_logger.warning(
"Prisma DB disconnect before reconnect failed: %s",
disconnect_err,
db_url = os.getenv("DATABASE_URL", "")
if not db_url:
verbose_proxy_logger.error(
"DATABASE_URL not set; cannot reconnect Prisma client."
)
await PrismaWrapper._kill_engine_process(old_pid)
await self.db.connect()
raise RuntimeError("DATABASE_URL not set")
# Fresh Prisma client + new engine subprocess. The previous
# "lightweight" path called `disconnect()` which blocks the
# event loop on `subprocess.Popen.wait()`; since that call
# ends up killing the engine anyway, we do it non-blockingly
# via `_kill_engine_process` inside `recreate_prisma_client`.
self._cleanup_engine_watcher()
await self.db.recreate_prisma_client(db_url)
await self._start_engine_watcher()
await self.db.query_raw("SELECT 1")
await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout)
+150
View File
@@ -0,0 +1,150 @@
# Workflow Run Tracking
Generic durable state tracking for agents and automated workflows built on the LiteLLM proxy.
## The Problem
Agents like [shin-builder](https://github.com/BerriAI/shin-builder) run multi-stage pipelines (triage → plan → implement → PR). Their task state and conversation history lived in memory — a process restart lost everything.
## Three-Table Design
```
WorkflowRun one instance of work (header + materialized status)
WorkflowEvent append-only state transitions (source of truth for replay)
WorkflowMessage conversation inbox/outbox (full content, not truncated)
```
**WorkflowEvent is the source of truth.** `WorkflowRun.status` is a materialized cache updated automatically when events are appended. If you need to debug a run, replay its events.
## API
All endpoints require a valid LiteLLM API key (`Authorization: Bearer sk-...`).
### Runs
```
POST /v1/workflows/runs Create a run
GET /v1/workflows/runs List runs (?workflow_type=&status=)
GET /v1/workflows/runs/{run_id} Get run + latest event
PATCH /v1/workflows/runs/{run_id} Update status / metadata / output
```
### Events
```
POST /v1/workflows/runs/{run_id}/events Append event (auto-updates run status)
GET /v1/workflows/runs/{run_id}/events Full event log (ordered by sequence)
```
### Messages
```
POST /v1/workflows/runs/{run_id}/messages Append message
GET /v1/workflows/runs/{run_id}/messages Conversation history (ordered by sequence)
```
## Quick Start
```bash
# Create a run
curl -X POST http://localhost:4000/v1/workflows/runs \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{"workflow_type": "shin-builder", "metadata": {"title": "Fix login bug"}}'
# {"run_id": "abc-123", "session_id": "xyz-456", "status": "pending", ...}
# Mark step started (sets status → running)
curl -X POST http://localhost:4000/v1/workflows/runs/abc-123/events \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{"event_type": "step.started", "step_name": "grill", "data": {"claude_session_id": "sess-789"}}'
# Store a conversation message
curl -X POST http://localhost:4000/v1/workflows/runs/abc-123/messages \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{"role": "user", "content": "What is the expected behavior?", "session_id": "sess-789"}'
# Restart recovery: fetch active runs and resume from last event's data.claude_session_id
curl "http://localhost:4000/v1/workflows/runs?status=running,paused&workflow_type=shin-builder" \
-H "Authorization: Bearer sk-1234"
```
## Status Auto-Update Rules
When you append an event, the run's status is updated automatically:
| event_type | run.status |
|-----------------|------------|
| `step.started` | `running` |
| `step.failed` | `failed` |
| `hook.waiting` | `paused` |
| `hook.received` | `running` |
Set `status = completed` explicitly via PATCH when the workflow finishes.
## Linking to Spend Logs
`WorkflowRun.session_id` is generated automatically (UUID). Pass it as the `x-litellm-session-id` header when making completions through the proxy:
```python
headers = {"x-litellm-session-id": run.session_id}
```
All spend log entries for this run are then tagged automatically. Query cost per run:
```
POST /ui/spend_logs/view_session_spend_logs?session_id={run.session_id}
```
## Sequence Numbers
Sequence numbers on events and messages are assigned server-side (`MAX + 1` per run). Callers never supply them. This guarantees ordering even under concurrent writes.
## Using from shin-builder
Replace the in-memory `tasks.py` dict with calls to these endpoints:
```python
import httpx
class WorkflowRunClient:
def __init__(self, base_url: str, api_key: str):
self._client = httpx.AsyncClient(
base_url=base_url,
headers={"Authorization": f"Bearer {api_key}"},
)
async def create_task(self, title: str, **metadata) -> dict:
r = await self._client.post("/v1/workflows/runs", json={
"workflow_type": "shin-builder",
"metadata": {"title": title, **metadata},
})
r.raise_for_status()
return r.json()
async def list_active_tasks(self) -> list:
r = await self._client.get(
"/v1/workflows/runs",
params={"workflow_type": "shin-builder", "status": "running,paused"},
)
r.raise_for_status()
return r.json()["runs"]
async def transition(self, run_id: str, step_name: str, event_type: str, data: dict = None):
r = await self._client.post(f"/v1/workflows/runs/{run_id}/events", json={
"event_type": event_type,
"step_name": step_name,
"data": data or {},
})
r.raise_for_status()
async def append_message(self, run_id: str, role: str, content: str, session_id: str = None):
r = await self._client.post(f"/v1/workflows/runs/{run_id}/messages", json={
"role": role, "content": content, "session_id": session_id,
})
r.raise_for_status()
```
On startup, call `list_active_tasks()` to restore in-flight runs. The last `step.started` event's `data.claude_session_id` gives you the `--resume` ID.
@@ -11,9 +11,60 @@ If given, generate a unique model_id for the deployment.
Ensures cooldowns are applied correctly.
"""
from typing import List
clientside_credential_keys = ["api_key", "api_base", "base_url"]
def _admin_config_fields_to_clear_on_base_override() -> List[str]:
"""
Provider-specific credential / endpoint-targeting fields that must NOT
flow through to a client-redirected upstream.
Built dynamically from ``CredentialLiteLLMParams.model_fields`` so any
new provider field added there (Bedrock endpoint, Watsonx region, etc.)
is gated automatically plus a fixed list of kwargs-only fields that
aren't declared on the typed model.
"""
from litellm.types.router import CredentialLiteLLMParams
typed_fields = [
f
for f in CredentialLiteLLMParams.model_fields
if f not in clientside_credential_keys
]
kwargs_only_fields = [
# Caller-supplied via **kwargs, not declared on CredentialLiteLLMParams.
"organization",
"extra_body",
"extra_headers",
"default_headers",
"api_type",
"azure_ad_token",
"azure_ad_token_provider",
"aws_session_token",
"aws_sts_endpoint",
"aws_web_identity_token",
"aws_role_name",
# OCI provider — consumed by litellm/llms/oci/* via optional_params
# and not declared on CredentialLiteLLMParams. Without these here,
# an admin's OCI signing key / tenancy / fingerprint would flow
# through to an attacker-redirected upstream.
"oci_signer",
"oci_user",
"oci_fingerprint",
"oci_tenancy",
"oci_key",
"oci_key_file",
]
return typed_fields + kwargs_only_fields
_ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE = (
_admin_config_fields_to_clear_on_base_override()
)
def is_clientside_credential(request_kwargs: dict) -> bool:
"""
Check if the credential is a clientside credential.
@@ -34,4 +85,20 @@ def get_dynamic_litellm_params(litellm_params: dict, request_kwargs: dict) -> di
for key in clientside_credential_keys:
if key in request_kwargs:
litellm_params[key] = request_kwargs[key]
# If the caller redirected api_base/base_url to a client-controlled value,
# don't forward the admin's organization / extra_body / region / token /
# vertex / aws fields — those were meant for the original upstream.
# Always drop the admin's value first, then write the caller's value back
# if they resupplied the field. The naive
# ``if field not in request_kwargs: pop`` shape lets a caller *echo* a
# field name (with any value, including an empty string) to keep the
# admin's value in ``litellm_params`` and have it forwarded to the
# redirected upstream.
if "api_base" in request_kwargs or "base_url" in request_kwargs:
for field in _ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE:
litellm_params.pop(field, None)
if field in request_kwargs:
litellm_params[field] = request_kwargs[field]
return litellm_params
+1
View File
@@ -118,3 +118,4 @@ class CachedEmbedding(TypedDict):
index: Optional[int]
object: Optional[str]
model: Optional[str]
prompt_tokens_details: Optional[dict]
@@ -81,6 +81,12 @@ class MCPServer(BaseModel):
# Defaults to the token's expires_in minus the expiry buffer, or
# MCP_PER_USER_TOKEN_DEFAULT_TTL when expires_in is absent.
token_storage_ttl_seconds: Optional[int] = None
# Resolved short-ID tool prefix when LITELLM_USE_SHORT_MCP_TOOL_PREFIX is
# enabled. Set by ``MCPServerManager._assign_unique_short_prefix`` at
# registration time so that natural-hash collisions between two
# different ``server_id`` values are bumped deterministically. Left
# ``None`` in default-prefix mode.
short_prefix: Optional[str] = None
model_config = ConfigDict(arbitrary_types_allowed=True)
@property
+1 -1
View File
@@ -2659,7 +2659,7 @@ class StandardLoggingHiddenParams(TypedDict):
] # id of the model in the router, separates multiple models with the same name but different credentials
cache_key: Optional[str]
api_base: Optional[str]
response_cost: Optional[str]
response_cost: Optional[Union[str, float]]
litellm_overhead_time_ms: Optional[float]
additional_headers: Optional[StandardLoggingAdditionalHeaders]
batch_models: Optional[List[str]]
+14 -2
View File
@@ -6526,6 +6526,7 @@ def validate_environment( # noqa: PLR0915
or model in litellm.open_ai_text_completion_models
or model in litellm.open_ai_embedding_models
or model in litellm.openai_image_generation_models
or model.startswith("gpt-image")
):
if "OPENAI_API_KEY" in os.environ:
keys_in_environment = True
@@ -8410,6 +8411,17 @@ class ProviderConfigManager:
model: str,
provider: LlmProviders,
) -> Optional[BaseAnthropicMessagesConfig]:
return ProviderConfigManager._get_provider_anthropic_messages_config_cached(
model=model, provider=provider
)
@staticmethod
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
def _get_provider_anthropic_messages_config_cached(
model: str,
provider: LlmProviders,
) -> Optional[BaseAnthropicMessagesConfig]:
model_lower = model.lower()
if litellm.LlmProviders.ANTHROPIC == provider:
return litellm.AnthropicMessagesConfig()
# The 'BEDROCK' provider corresponds to Amazon's implementation of Anthropic Claude v3.
@@ -8419,14 +8431,14 @@ class ProviderConfigManager:
return BedrockModelInfo.get_bedrock_provider_config_for_messages_api(model)
elif litellm.LlmProviders.VERTEX_AI == provider:
if "claude" in model.lower():
if "claude" in model_lower:
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import (
VertexAIPartnerModelsAnthropicMessagesConfig,
)
return VertexAIPartnerModelsAnthropicMessagesConfig()
elif litellm.LlmProviders.AZURE_AI == provider:
if "claude" in model.lower():
if "claude" in model_lower:
from litellm.llms.azure_ai.anthropic.messages_transformation import (
AzureAnthropicMessagesConfig,
)
+96 -32
View File
@@ -4749,17 +4749,17 @@
"supports_web_search": true
},
"azure/gpt-5.5-pro": {
"cache_read_input_token_cost": 6e-06,
"cache_read_input_token_cost_above_272k_tokens": 1.2e-05,
"input_cost_per_token": 6e-05,
"input_cost_per_token_above_272k_tokens": 0.00012,
"cache_read_input_token_cost": 3e-06,
"cache_read_input_token_cost_above_272k_tokens": 6e-06,
"input_cost_per_token": 3e-05,
"input_cost_per_token_above_272k_tokens": 6e-05,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 0.00036,
"output_cost_per_token_above_272k_tokens": 0.00054,
"output_cost_per_token": 0.00018,
"output_cost_per_token_above_272k_tokens": 0.00027,
"supported_endpoints": [
"/v1/batch",
"/v1/responses"
@@ -4788,17 +4788,17 @@
"supports_low_reasoning_effort": false
},
"azure/gpt-5.5-pro-2026-04-23": {
"cache_read_input_token_cost": 6e-06,
"cache_read_input_token_cost_above_272k_tokens": 1.2e-05,
"input_cost_per_token": 6e-05,
"input_cost_per_token_above_272k_tokens": 0.00012,
"cache_read_input_token_cost": 3e-06,
"cache_read_input_token_cost_above_272k_tokens": 6e-06,
"input_cost_per_token": 3e-05,
"input_cost_per_token_above_272k_tokens": 6e-05,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 0.00036,
"output_cost_per_token_above_272k_tokens": 0.00054,
"output_cost_per_token": 0.00018,
"output_cost_per_token_above_272k_tokens": 0.00027,
"supported_endpoints": [
"/v1/batch",
"/v1/responses"
@@ -5117,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",
@@ -19097,6 +19129,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",
@@ -19912,21 +19976,21 @@
"supports_minimal_reasoning_effort": true
},
"gpt-5.5-pro": {
"cache_read_input_token_cost": 6e-06,
"cache_read_input_token_cost_above_272k_tokens": 1.2e-05,
"input_cost_per_token": 6e-05,
"input_cost_per_token_above_272k_tokens": 0.00012,
"input_cost_per_token_flex": 3e-05,
"input_cost_per_token_batches": 3e-05,
"cache_read_input_token_cost": 3e-06,
"cache_read_input_token_cost_above_272k_tokens": 6e-06,
"input_cost_per_token": 3e-05,
"input_cost_per_token_above_272k_tokens": 6e-05,
"input_cost_per_token_flex": 1.5e-05,
"input_cost_per_token_batches": 1.5e-05,
"litellm_provider": "openai",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 0.00036,
"output_cost_per_token_above_272k_tokens": 0.00054,
"output_cost_per_token_flex": 0.00018,
"output_cost_per_token_batches": 0.00018,
"output_cost_per_token": 0.00018,
"output_cost_per_token_above_272k_tokens": 0.00027,
"output_cost_per_token_flex": 9e-05,
"output_cost_per_token_batches": 9e-05,
"supported_endpoints": [
"/v1/responses",
"/v1/batch"
@@ -19955,21 +20019,21 @@
"supports_minimal_reasoning_effort": true
},
"gpt-5.5-pro-2026-04-23": {
"cache_read_input_token_cost": 6e-06,
"cache_read_input_token_cost_above_272k_tokens": 1.2e-05,
"input_cost_per_token": 6e-05,
"input_cost_per_token_above_272k_tokens": 0.00012,
"input_cost_per_token_flex": 3e-05,
"input_cost_per_token_batches": 3e-05,
"cache_read_input_token_cost": 3e-06,
"cache_read_input_token_cost_above_272k_tokens": 6e-06,
"input_cost_per_token": 3e-05,
"input_cost_per_token_above_272k_tokens": 6e-05,
"input_cost_per_token_flex": 1.5e-05,
"input_cost_per_token_batches": 1.5e-05,
"litellm_provider": "openai",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 0.00036,
"output_cost_per_token_above_272k_tokens": 0.00054,
"output_cost_per_token_flex": 0.00018,
"output_cost_per_token_batches": 0.00018,
"output_cost_per_token": 0.00018,
"output_cost_per_token_above_272k_tokens": 0.00027,
"output_cost_per_token_flex": 9e-05,
"output_cost_per_token_batches": 9e-05,
"supported_endpoints": [
"/v1/responses",
"/v1/batch"
+17
View File
@@ -193,6 +193,23 @@
"a2a": false
}
},
"aihubmix": {
"display_name": "AIHubMix (`aihubmix`)",
"url": "https://docs.litellm.ai/docs/providers/aihubmix",
"endpoints": {
"chat_completions": true,
"messages": true,
"responses": true,
"embeddings": true,
"image_generations": true,
"audio_transcriptions": true,
"audio_speech": true,
"moderations": true,
"batches": false,
"rerank": true,
"a2a": false
}
},
"assemblyai": {
"display_name": "AssemblyAI (`assemblyai`)",
"url": "https://docs.litellm.ai/docs/pass_through/assembly_ai",
+2 -2
View File
@@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.83.14"
version = "1.84.0"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.14"
@@ -236,7 +236,7 @@ source-exclude = [
profile = "black"
[tool.commitizen]
version = "1.83.14"
version = "1.84.0"
version_files = [
"pyproject.toml:^version",
]
+77
View File
@@ -1290,3 +1290,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])
}
@@ -254,32 +254,48 @@ async def test_run_reconnect_cycle_uses_heavy_path_when_confirmed_dead(
@pytest.mark.asyncio
async def test_run_reconnect_cycle_uses_lightweight_path_when_engine_alive(
async def test_run_reconnect_cycle_uses_direct_path_when_engine_alive(
engine_client,
) -> None:
"""_run_reconnect_cycle uses disconnect/connect when engine is alive."""
engine_client._engine_pid = 1234
"""Direct reconnect (engine alive) calls recreate_prisma_client + SELECT 1.
with patch.object(engine_client, "_is_engine_alive", return_value=True):
The old "lightweight" path called `disconnect()` + `connect()`, which
blocks the event loop on the sync `process.wait()` inside aclose().
The fix routes both engine-alive and engine-dead paths through
`recreate_prisma_client`, which non-blockingly kills the old engine.
"""
engine_client._engine_pid = 1234
engine_client._start_engine_watcher = AsyncMock()
with (
patch.object(engine_client, "_is_engine_alive", return_value=True),
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
):
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
engine_client.db.connect.assert_awaited_once()
engine_client.db.recreate_prisma_client.assert_awaited_once_with(
"postgresql://test"
)
engine_client.db.query_raw.assert_awaited_once_with("SELECT 1")
engine_client.db.recreate_prisma_client.assert_not_awaited()
engine_client.db.disconnect.assert_not_awaited()
@pytest.mark.asyncio
async def test_run_reconnect_cycle_uses_lightweight_path_when_pid_unknown(
async def test_run_reconnect_cycle_uses_direct_path_when_pid_unknown(
engine_client,
) -> None:
"""_run_reconnect_cycle uses lightweight path when engine PID is not tracked."""
"""When the engine PID is not tracked, direct reconnect still runs."""
engine_client._engine_pid = 0
engine_client._start_engine_watcher = AsyncMock()
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
engine_client.db.connect.assert_awaited_once()
engine_client.db.recreate_prisma_client.assert_awaited_once_with(
"postgresql://test"
)
engine_client.db.query_raw.assert_awaited_once_with("SELECT 1")
engine_client.db.recreate_prisma_client.assert_not_awaited()
engine_client.db.disconnect.assert_not_awaited()
@pytest.mark.asyncio
@@ -473,36 +489,38 @@ def test_on_engine_death_from_thread_ignores_stale_pid(engine_client):
@pytest.mark.asyncio
async def test_escalation_after_consecutive_lightweight_failures(engine_client):
"""After N consecutive lightweight reconnect failures, _engine_confirmed_dead
async def test_escalation_after_consecutive_direct_reconnect_failures(engine_client):
"""After N consecutive direct reconnect failures, _engine_confirmed_dead
is set to True so _run_reconnect_cycle takes the heavy reconnect path."""
engine_client._reconnect_escalation_threshold = 3
engine_client._consecutive_reconnect_failures = 0
engine_client._db_reconnect_cooldown_seconds = 0 # disable cooldown for test
engine_client._start_engine_watcher = AsyncMock(return_value=None)
# Make lightweight reconnect fail every time
engine_client.db.disconnect = AsyncMock(return_value=None)
engine_client.db.connect = AsyncMock(side_effect=Exception("connect failed"))
# Make direct reconnect fail every time
engine_client.db.recreate_prisma_client = AsyncMock(
side_effect=Exception("recreate failed")
)
# Run 3 failed reconnect attempts
for i in range(3):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=5.0
)
assert result is False
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
for _ in range(3):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=5.0
)
assert result is False
assert engine_client._consecutive_reconnect_failures == 3
# Next attempt should escalate: _engine_confirmed_dead set to True before _run_reconnect_cycle
# Next attempt should escalate to the heavy path (recreate_prisma_client still
# the call, but via the _engine_confirmed_dead branch that also re-arms the watcher).
engine_client.db.recreate_prisma_client = AsyncMock(return_value=None)
engine_client._start_engine_watcher = AsyncMock(return_value=None)
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test_escalation", timeout_seconds=5.0
)
# Heavy reconnect should have been attempted (recreate_prisma_client called)
engine_client.db.recreate_prisma_client.assert_awaited_once()
@@ -511,15 +529,16 @@ async def test_successful_reconnect_resets_failure_counter(engine_client):
"""A successful reconnect resets _consecutive_reconnect_failures to 0."""
engine_client._consecutive_reconnect_failures = 2
engine_client._db_reconnect_cooldown_seconds = 0
engine_client._start_engine_watcher = AsyncMock()
# Make reconnect succeed
engine_client.db.disconnect = AsyncMock(return_value=None)
engine_client.db.connect = AsyncMock(return_value=None)
engine_client.db.recreate_prisma_client = AsyncMock(return_value=None)
engine_client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=5.0
)
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=5.0
)
assert result is True
assert engine_client._consecutive_reconnect_failures == 0
@@ -314,11 +314,11 @@ def test_update_litellm_params_for_health_check():
# Issue #15807: Fixes health checks sending "region/model" as model ID to AWS
model_info = {}
litellm_params = {
"model": "bedrock/us-gov-west-1/anthropic.claude-3-7-sonnet-20250219-v1:0",
"model": "bedrock/us-gov-west-1/anthropic.claude-sonnet-4-5-20250929-v1:0",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert updated_params["model"] == "anthropic.claude-3-7-sonnet-20250219-v1:0"
assert updated_params["model"] == "anthropic.claude-sonnet-4-5-20250929-v1:0"
# Test with Bedrock cross-region inference profile - should preserve the inference profile prefix
# AWS requires inference profile IDs like "us.anthropic.claude..." for cross-region routing
@@ -366,3 +366,40 @@ def test_generic_api_compatible_callbacks_json_unknown_callback():
# Should return the string unchanged
assert result == "unknown_callback", "Unknown callback should be returned as-is"
assert isinstance(result, str), "Unknown callback should remain a string"
@pytest.mark.asyncio
async def test_generic_api_callback_settings_retry_config():
"""
Test that generic_api callback_settings are passed to GenericAPILogger.
"""
from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger
from litellm.litellm_core_utils.logging_callback_manager import (
_generic_api_logger_cache,
)
callback_name = "test_generic_api_retry_config"
_generic_api_logger_cache.pop(callback_name, None)
litellm.callback_settings[callback_name] = {
"callback_type": "generic_api",
"endpoint": "https://example.com/api/logs",
"headers": {"Content-Type": "application/json"},
"max_retries": 2,
"retry_delay": 0.5,
"timeout": 3,
}
try:
result = LoggingCallbackManager._add_custom_callback_generic_api_str(
callback_name
)
assert isinstance(result, GenericAPILogger)
assert result.endpoint == "https://example.com/api/logs"
assert result.headers == {"Content-Type": "application/json"}
assert result.max_retries == 2
assert result.retry_delay == 0.5
assert result.timeout == 3
finally:
litellm.callback_settings.pop(callback_name, None)
_generic_api_logger_cache.pop(callback_name, None)
+3 -3
View File
@@ -2309,11 +2309,11 @@ def test_get_provider_audio_transcription_config():
@pytest.mark.parametrize(
"model, expected_bool",
[
("anthropic.claude-3-7-sonnet-20250219-v1:0", True),
("us.anthropic.claude-3-7-sonnet-20250219-v1:0", True),
("anthropic.claude-sonnet-4-5-20250929-v1:0", True),
("us.anthropic.claude-sonnet-4-5-20250929-v1:0", True),
],
)
def test_claude_3_7_sonnet_supports_pdf_input(model, expected_bool):
def test_claude_sonnet_4_5_supports_pdf_input(model, expected_bool):
from litellm.utils import supports_pdf_input
assert supports_pdf_input(model) == expected_bool
@@ -134,7 +134,7 @@ class TestBedrockAnthropicPromptCachingRegression:
if "converse" in model_prefix:
config = AmazonConverseConfig()
result = config.transform_request(
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
optional_params={},
litellm_params={},
@@ -162,7 +162,7 @@ class TestBedrockAnthropicPromptCachingRegression:
else:
config = AmazonAnthropicClaudeConfig()
result = config.transform_request(
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
optional_params={},
litellm_params={},
@@ -227,7 +227,7 @@ class TestBedrockAnthropicPromptCachingRegression:
if "converse" in model_prefix:
config = AmazonConverseConfig()
result = config._transform_request_helper(
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
system_content_blocks=[],
optional_params={},
messages=messages,
@@ -236,7 +236,7 @@ class TestBedrockAnthropicPromptCachingRegression:
else:
config = AmazonAnthropicClaudeConfig()
result = config.transform_request(
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
optional_params={},
litellm_params={},
@@ -498,7 +498,7 @@ class TestBedrockAnthropicCombinedRegressions:
if "converse" in model_prefix:
config = AmazonConverseConfig()
result = config._transform_request_helper(
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
system_content_blocks=[],
optional_params={},
messages=messages,
@@ -518,7 +518,7 @@ class TestBedrockAnthropicCombinedRegressions:
else:
config = AmazonAnthropicClaudeConfig()
result = config.transform_request(
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
optional_params={},
litellm_params={},
@@ -1323,7 +1323,7 @@ def test_base_aws_llm_get_credentials():
def test_bedrock_completion_test_2():
litellm.set_verbose = True
data = {
"model": "bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0",
"model": "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0",
"messages": [
{
"role": "system",
@@ -1630,7 +1630,7 @@ def test_bedrock_completion_test_4(modify_params):
litellm.modify_params = modify_params
data = {
"model": "anthropic.claude-3-7-sonnet-20250219-v1:0",
"model": "anthropic.claude-sonnet-4-5-20250929-v1:0",
"messages": [
{
"role": "user",
@@ -2115,7 +2115,7 @@ class TestBedrockConverseAnthropicUnitTests(BaseAnthropicChatTest):
def get_base_completion_call_args_with_thinking(self) -> dict:
return {
"model": "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
"model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
"thinking": {"type": "enabled", "budget_tokens": 16000},
}
@@ -2828,7 +2828,7 @@ async def test_bedrock_thinking_in_assistant_message(sync_mode):
client = AsyncHTTPHandler()
params = {
"model": "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
"model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
"messages": [
{
"role": "assistant",
@@ -2887,7 +2887,7 @@ async def test_bedrock_stream_thinking_content_openwebui():
```
"""
response = await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[{"role": "user", "content": "Hello who is this?"}],
stream=True,
max_tokens=1080,
@@ -580,7 +580,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers():
def test_litellm_gateway_from_sdk_with_thinking_param():
try:
response = litellm.completion(
model="litellm_proxy/anthropic.claude-3-7-sonnet-20250219-v1:0",
model="litellm_proxy/anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[{"role": "user", "content": "Hello world"}],
api_base="http://0.0.0.0:4000",
api_key="sk-PIp1h0RekR",
@@ -1828,7 +1828,7 @@ def test_azure_response_format_param():
"model, provider",
[
("claude-3-7-sonnet-20240620-v1:0", "anthropic"),
("anthropic.claude-3-7-sonnet-20250219-v1:0", "bedrock"),
("anthropic.claude-sonnet-4-5-20250929-v1:0", "bedrock"),
("invoke/anthropic.claude-3-7-sonnet-20240620-v1:0", "bedrock"),
("claude-3-7-sonnet@20250219", "vertex_ai"),
],
@@ -3493,8 +3493,14 @@ def test_litellm_api_base(monkeypatch, provider, route):
def test_gemini_tool_calling_working_demo():
load_vertex_ai_credentials()
litellm._turn_on_debug()
"""
Regression test: tool params with anyOf containing a `{"type": "array"}`
branch (no items field at all) must synthesize items before the request
is sent to Vertex (Vertex rejects array types missing items).
"""
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
args = {
"messages": [
{
@@ -3564,13 +3570,75 @@ def test_gemini_tool_calling_working_demo():
],
"vertex_location": "global",
}
response = completion(model="vertex_ai/gemini-3-flash-preview", **args)
print(response)
client = HTTPHandler()
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/json"}
mock_response.json.return_value = {
"candidates": [
{
"content": {
"role": "model",
"parts": [{"text": "Hello!"}],
},
"finishReason": "STOP",
}
],
"usageMetadata": {
"promptTokenCount": 10,
"candidatesTokenCount": 5,
"totalTokenCount": 15,
},
}
with (
patch.object(client, "post", return_value=mock_response) as mock_post,
patch.object(
VertexBase,
"_ensure_access_token",
return_value=("fake-token", "fake-project"),
),
):
completion(
model="vertex_ai/gemini-3-flash-preview",
client=client,
**args,
)
sent_body = mock_post.call_args.kwargs.get(
"json"
) or mock_post.call_args.kwargs.get("data")
assert sent_body is not None, "expected request body to be sent"
if isinstance(sent_body, str):
sent_body = json.loads(sent_body)
function_decl = sent_body["tools"][0]["function_declarations"][0]
callbacks_schema = function_decl["parameters"]["properties"]["config"][
"properties"
]["callbacks"]
array_branches = [
branch
for branch in callbacks_schema["anyOf"]
if branch.get("type", "").lower() == "array"
]
assert array_branches, "expected an array branch in callbacks anyOf"
for branch in array_branches:
assert "items" in branch and branch["items"], (
f"array branch in callbacks.anyOf must include non-empty items "
f"(Vertex rejects array types missing items). Got: {branch}"
)
def test_gemini_tool_calling_not_working():
load_vertex_ai_credentials()
litellm._turn_on_debug()
"""
Regression test: tool params with anyOf containing both an empty-items
array branch and a null branch must serialize with items present on the
array branch (Vertex rejects array types missing `items`).
"""
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
args = {
"messages": [
{
@@ -3637,8 +3705,64 @@ def test_gemini_tool_calling_not_working():
],
"vertex_location": "global",
}
response = completion(model="vertex_ai/gemini-3-flash-preview", **args)
print(response)
client = HTTPHandler()
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/json"}
mock_response.json.return_value = {
"candidates": [
{
"content": {
"role": "model",
"parts": [{"text": "Hello!"}],
},
"finishReason": "STOP",
}
],
"usageMetadata": {
"promptTokenCount": 10,
"candidatesTokenCount": 5,
"totalTokenCount": 15,
},
}
with (
patch.object(client, "post", return_value=mock_response) as mock_post,
patch.object(
VertexBase,
"_ensure_access_token",
return_value=("fake-token", "fake-project"),
),
):
completion(
model="vertex_ai/gemini-3-flash-preview",
client=client,
**args,
)
sent_body = mock_post.call_args.kwargs.get(
"json"
) or mock_post.call_args.kwargs.get("data")
assert sent_body is not None, "expected request body to be sent"
if isinstance(sent_body, str):
sent_body = json.loads(sent_body)
function_decl = sent_body["tools"][0]["function_declarations"][0]
callbacks_schema = function_decl["parameters"]["properties"]["config"][
"properties"
]["callbacks"]
array_branches = [
branch
for branch in callbacks_schema["anyOf"]
if branch.get("type", "").lower() == "array"
]
assert array_branches, "expected an array branch in callbacks anyOf"
for branch in array_branches:
assert "items" in branch and branch["items"], (
f"array branch in callbacks.anyOf must include non-empty items "
f"(Vertex rejects array types missing items). Got: {branch}"
)
def test_vertex_ai_llama_tool_calling():
+1 -1
View File
@@ -159,7 +159,7 @@ def test_aaparallel_function_call(model):
"model",
[
"anthropic/claude-4-sonnet-20250514",
"bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
"bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
],
)
@pytest.mark.flaky(retries=3, delay=1)
+3 -4
View File
@@ -30,10 +30,9 @@ def test_model_alias_map(caplog):
)
print(response.model)
captured_logs = [rec.levelname for rec in caplog.records]
for log in captured_logs:
assert "ERROR" not in log
for rec in caplog.records:
if rec.levelname == "ERROR" and rec.name.startswith("LiteLLM"):
pytest.fail(f"Unexpected litellm ERROR log: {rec.getMessage()}")
assert "llama-3.1-8b-instant" in response.model
except litellm.ServiceUnavailableError:
@@ -354,6 +354,7 @@ async def test_bedrock_kb_request_body_has_transformed_filters(
api_base=api_base,
litellm_logging_obj=logging_obj,
litellm_params=litellm_params_dict,
extra_body=None,
)
)
captured_request_body["url"] = url
@@ -8,6 +8,7 @@ sys.path.insert(0, os.path.abspath("../.."))
import asyncio
import litellm
import gzip
import httpx
import json
import logging
import time
@@ -470,3 +471,96 @@ async def test_generic_api_callback_invalid_log_format():
endpoint=test_endpoint,
log_format="invalid_format", # type: ignore # Intentionally invalid for testing
)
@pytest.mark.asyncio
async def test_generic_api_callback_retries_timeout_then_succeeds():
"""
Test that GenericAPILogger retries LiteLLM timeout errors when configured.
"""
test_endpoint = "https://example.com/api/logs"
generic_logger = GenericAPILogger(
endpoint=test_endpoint,
max_retries=1,
retry_delay=0,
timeout=0.2,
)
mock_post = AsyncMock()
mock_post.side_effect = [
litellm.Timeout(
message="Connection timed out",
model="default-model-name",
llm_provider="litellm-httpx-handler",
),
type("Response", (), {"status_code": 200})(),
]
generic_logger.async_httpx_client.post = mock_post
generic_logger.log_queue = [{"event": "timeout-retry"}]
await generic_logger.async_send_batch()
assert mock_post.call_count == 2
first_call = mock_post.call_args_list[0][1]
assert first_call["url"] == test_endpoint
assert first_call["timeout"] == 0.2
assert json.loads(first_call["data"]) == [{"event": "timeout-retry"}]
@pytest.mark.asyncio
async def test_generic_api_callback_retries_5xx_then_succeeds():
"""
Test that GenericAPILogger retries transient HTTP 5xx errors when configured.
"""
test_endpoint = "https://example.com/api/logs"
generic_logger = GenericAPILogger(
endpoint=test_endpoint,
max_retries=1,
retry_delay=0,
)
request = httpx.Request("POST", test_endpoint)
response = httpx.Response(status_code=503, request=request)
mock_post = AsyncMock()
mock_post.side_effect = [
httpx.HTTPStatusError(
"Server error",
request=request,
response=response,
),
type("Response", (), {"status_code": 200})(),
]
generic_logger.async_httpx_client.post = mock_post
generic_logger.log_queue = [{"event": "5xx-retry"}]
await generic_logger.async_send_batch()
assert mock_post.call_count == 2
@pytest.mark.asyncio
async def test_generic_api_callback_does_not_retry_4xx():
"""
Test that GenericAPILogger does not retry non-transient HTTP 4xx errors.
"""
test_endpoint = "https://example.com/api/logs"
generic_logger = GenericAPILogger(
endpoint=test_endpoint,
max_retries=2,
retry_delay=0,
)
request = httpx.Request("POST", test_endpoint)
response = httpx.Response(status_code=401, request=request)
mock_post = AsyncMock()
mock_post.side_effect = httpx.HTTPStatusError(
"Unauthorized",
request=request,
response=response,
)
generic_logger.async_httpx_client.post = mock_post
generic_logger.log_queue = [{"event": "4xx-no-retry"}]
await generic_logger.async_send_batch()
mock_post.assert_called_once()
@@ -96,8 +96,8 @@ class BaseAnthropicMessagesPromptCachingTest(ABC):
Returns the model string to use for tests.
Examples:
- "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0"
- "bedrock/invoke/anthropic.claude-3-7-sonnet-20250219-v1:0"
- "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0"
- "bedrock/invoke/anthropic.claude-sonnet-4-5-20250929-v1:0"
"""
pass
@@ -31,7 +31,7 @@ class TestBedrockConversePromptCaching(BaseAnthropicMessagesPromptCachingTest):
"""
def get_model(self) -> str:
return "bedrock/converse/us.anthropic.claude-3-7-sonnet-20250219-v1:0"
return "bedrock/converse/us.anthropic.claude-sonnet-4-5-20250929-v1:0"
class TestBedrockInvokePromptCaching(BaseAnthropicMessagesPromptCachingTest):
@@ -43,4 +43,4 @@ class TestBedrockInvokePromptCaching(BaseAnthropicMessagesPromptCachingTest):
"""
def get_model(self) -> str:
return "bedrock/invoke/us.anthropic.claude-3-7-sonnet-20250219-v1:0"
return "bedrock/invoke/us.anthropic.claude-sonnet-4-5-20250929-v1:0"
@@ -52,3 +52,183 @@ async def test_process_async_embedding_cached_response():
print(f"response: {response}")
assert len(response.data) == 1
@pytest.mark.asyncio
async def test_embedding_cache_preserves_prompt_tokens_details():
"""Test that prompt_tokens_details (including image_count) survives a full cache hit."""
llm_caching_handler = LLMCachingHandler(
original_function=MagicMock(),
request_kwargs={},
start_time=datetime.now(),
)
cached_result = [
{
"embedding": [-0.025, -0.019],
"index": 0,
"object": "embedding",
"model": "amazon.titan-embed-image-v1",
"prompt_tokens_details": {"image_count": 1},
}
]
mock_logging_obj = MagicMock()
mock_logging_obj.async_success_handler = AsyncMock()
response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
final_embedding_cached_response=None,
cached_result=cached_result,
kwargs={"model": "amazon.titan-embed-image-v1", "input": "base64imagedata"},
logging_obj=mock_logging_obj,
start_time=datetime.now(),
model="amazon.titan-embed-image-v1",
)
assert cache_hit
assert response.usage is not None
assert response.usage.prompt_tokens_details is not None
assert response.usage.prompt_tokens_details.image_count == 1
@pytest.mark.asyncio
async def test_embedding_cache_backward_compat_no_prompt_tokens_details():
"""Test that old cached items without prompt_tokens_details still work."""
llm_caching_handler = LLMCachingHandler(
original_function=MagicMock(),
request_kwargs={},
start_time=datetime.now(),
)
# Old-format cached item — no prompt_tokens_details field
cached_result = [
{
"embedding": [-0.025, -0.019],
"index": 0,
"object": "embedding",
"model": "text-embedding-ada-002",
}
]
mock_logging_obj = MagicMock()
mock_logging_obj.async_success_handler = AsyncMock()
response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
final_embedding_cached_response=None,
cached_result=cached_result,
kwargs={"model": "text-embedding-ada-002", "input": "test"},
logging_obj=mock_logging_obj,
start_time=datetime.now(),
model="text-embedding-ada-002",
)
assert cache_hit
assert response.usage is not None
assert response.usage.prompt_tokens_details is None
@pytest.mark.asyncio
async def test_embedding_cache_aggregates_multiple_image_counts():
"""Test that image_count is summed correctly across multiple cached items."""
llm_caching_handler = LLMCachingHandler(
original_function=MagicMock(),
request_kwargs={},
start_time=datetime.now(),
)
cached_result = [
{
"embedding": [-0.025, -0.019],
"index": 0,
"object": "embedding",
"model": "amazon.titan-embed-image-v1",
"prompt_tokens_details": {"image_count": 1},
},
{
"embedding": [0.031, 0.042],
"index": 1,
"object": "embedding",
"model": "amazon.titan-embed-image-v1",
"prompt_tokens_details": {"image_count": 1},
},
]
mock_logging_obj = MagicMock()
mock_logging_obj.async_success_handler = AsyncMock()
response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
final_embedding_cached_response=None,
cached_result=cached_result,
kwargs={
"model": "amazon.titan-embed-image-v1",
"input": ["img1", "img2"],
},
logging_obj=mock_logging_obj,
start_time=datetime.now(),
model="amazon.titan-embed-image-v1",
)
assert cache_hit
assert response.usage.prompt_tokens_details is not None
assert response.usage.prompt_tokens_details.image_count == 2
def test_combine_usage_merges_prompt_tokens_details():
"""Test that combine_usage merges prompt_tokens_details from both Usage objects."""
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
llm_caching_handler = LLMCachingHandler(
original_function=MagicMock(),
request_kwargs={},
start_time=datetime.now(),
)
usage1 = Usage(
prompt_tokens=10,
completion_tokens=0,
total_tokens=10,
prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1),
)
usage2 = Usage(
prompt_tokens=20,
completion_tokens=0,
total_tokens=20,
prompt_tokens_details=PromptTokensDetailsWrapper(image_count=2),
)
combined = llm_caching_handler.combine_usage(usage1, usage2)
assert combined.prompt_tokens == 30
assert combined.total_tokens == 30
assert combined.prompt_tokens_details is not None
assert combined.prompt_tokens_details.image_count == 3
def test_combine_usage_handles_none_details():
"""Test that combine_usage works when one or both sides have null prompt_tokens_details."""
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
llm_caching_handler = LLMCachingHandler(
original_function=MagicMock(),
request_kwargs={},
start_time=datetime.now(),
)
# Both null
usage_a = Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10)
usage_b = Usage(prompt_tokens=20, completion_tokens=0, total_tokens=20)
combined = llm_caching_handler.combine_usage(usage_a, usage_b)
assert combined.prompt_tokens_details is None
# Only first has details
usage_c = Usage(
prompt_tokens=10,
completion_tokens=0,
total_tokens=10,
prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1),
)
combined = llm_caching_handler.combine_usage(usage_c, usage_b)
assert combined.prompt_tokens_details is not None
assert combined.prompt_tokens_details.image_count == 1
# Only second has details
combined = llm_caching_handler.combine_usage(usage_a, usage_c)
assert combined.prompt_tokens_details is not None
assert combined.prompt_tokens_details.image_count == 1
File diff suppressed because one or more lines are too long
@@ -1 +1 @@
{"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}
{"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}
@@ -220,7 +220,7 @@ async def test_anthropic_cache_control_hook_negative_indices():
with patch.object(client, "post", return_value=mock_response) as mock_post:
# Test with multiple messages and negative indices
response = await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[
{
"role": "system",
@@ -352,7 +352,7 @@ async def test_anthropic_cache_control_hook_out_of_bounds_logging():
]
await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
cache_control_injection_points=[
{"location": "message", "index": 10}
@@ -420,7 +420,7 @@ async def test_anthropic_cache_control_hook_negative_out_of_bounds_logging():
]
await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
cache_control_injection_points=[
{
@@ -486,7 +486,7 @@ async def test_anthropic_cache_control_hook_multiple_user_messages():
with patch.object(client, "post", return_value=mock_response) as mock_post:
# Test with multiple user messages and negative indices
response = await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[
{
"role": "user",
@@ -586,7 +586,7 @@ async def test_anthropic_cache_control_hook_out_of_bounds(bad_index):
]
await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
cache_control_injection_points=[
{"location": "message", "index": bad_index}
@@ -651,7 +651,7 @@ async def test_anthropic_cache_control_hook_single_message(message_list):
client = AsyncHTTPHandler()
with patch.object(client, "post", return_value=mock_response) as mock_post:
await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=message_list,
cache_control_injection_points=[{"location": "message", "index": -1}],
client=client,
@@ -691,7 +691,7 @@ async def test_anthropic_cache_control_hook_empty_message_list():
match="bedrock requires at least one non-system message",
):
await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[],
cache_control_injection_points=[
{"location": "message", "index": -1}
@@ -742,7 +742,7 @@ async def test_anthropic_cache_control_hook_no_op():
]
await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
# No cache_control_injection_points parameter
client=client,
@@ -799,7 +799,7 @@ async def test_anthropic_cache_control_hook_multiple_content_items_last_only():
client = AsyncHTTPHandler()
with patch.object(client, "post", return_value=mock_response) as mock_post:
response = await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[
{
"role": "user",
@@ -874,7 +874,7 @@ async def test_anthropic_cache_control_hook_document_analysis_multiple_pages():
client = AsyncHTTPHandler()
with patch.object(client, "post", return_value=mock_response) as mock_post:
response = await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[
{
"role": "user",
@@ -1057,7 +1057,7 @@ async def test_anthropic_cache_control_hook_string_negative_index():
client = AsyncHTTPHandler()
with patch.object(client, "post", return_value=mock_response) as mock_post:
await litellm.acompletion(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[
{"role": "user", "content": "First message"},
{"role": "assistant", "content": "First response"},
@@ -262,7 +262,7 @@ class TestOpenTelemetryProviderInitialization(unittest.TestCase):
class TestOpenTelemetry(unittest.TestCase):
POLL_INTERVAL = 0.05
POLL_TIMEOUT = 2.0
MODEL = "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0"
MODEL = "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0"
HERE = os.path.dirname(__file__)
@patch.dict(os.environ, {}, clear=True)
@@ -369,7 +369,7 @@ def test_generic_cost_per_token_gpt55():
def test_generic_cost_per_token_gpt55_pro():
"""gpt-5.5-pro: responses-only model — $60/1M input, $360/1M output, $6/1M cached input."""
"""gpt-5.5-pro: responses-only model — $30/1M input, $180/1M output, $3/1M cached input."""
model = "gpt-5.5-pro"
custom_llm_provider = "openai"
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
@@ -378,18 +378,18 @@ def test_generic_cost_per_token_gpt55_pro():
model_cost_map = litellm.model_cost[model]
# Sanity-check the map values match OpenAI's published pricing.
assert model_cost_map["input_cost_per_token"] == 6e-5
assert model_cost_map["output_cost_per_token"] == 3.6e-4
assert model_cost_map["cache_read_input_token_cost"] == 6e-6
assert model_cost_map["input_cost_per_token"] == 3e-5
assert model_cost_map["output_cost_per_token"] == 1.8e-4
assert model_cost_map["cache_read_input_token_cost"] == 3e-6
assert model_cost_map["litellm_provider"] == "openai"
# gpt-5.5-pro is a responses-only model (no /v1/chat/completions endpoint).
assert model_cost_map["mode"] == "responses"
assert "/v1/chat/completions" not in model_cost_map["supported_endpoints"]
assert "/v1/responses" in model_cost_map["supported_endpoints"]
# Inherits GPT-5.4-pro's long-context window + tiered pricing (scaled 2x).
# Inherits GPT-5.4-pro's long-context window + tiered pricing.
assert model_cost_map["max_input_tokens"] == 1050000
assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 1.2e-4
assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 5.4e-4
assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 6e-5
assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 2.7e-4
prompt_tokens = 1000
completion_tokens = 500
@@ -454,8 +454,8 @@ def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities(
[
("azure/gpt-5.5", "chat", 5e-6, 3e-5, 5e-7),
("azure/gpt-5.5-2026-04-23", "chat", 5e-6, 3e-5, 5e-7),
("azure/gpt-5.5-pro", "responses", 6e-5, 3.6e-4, 6e-6),
("azure/gpt-5.5-pro-2026-04-23", "responses", 6e-5, 3.6e-4, 6e-6),
("azure/gpt-5.5-pro", "responses", 3e-5, 1.8e-4, 3e-6),
("azure/gpt-5.5-pro-2026-04-23", "responses", 3e-5, 1.8e-4, 3e-6),
],
)
def test_azure_gpt55_entries_present_with_correct_pricing(
@@ -464,7 +464,7 @@ def test_azure_gpt55_entries_present_with_correct_pricing(
"""Day-0 Azure entries for GPT-5.5 mirror the OpenAI pricing structure.
Pricing parity with openai/gpt-5.5* (verified against OpenAI's pricing page
on 2026-04-24): $5/$30 input/output per 1M for chat, $60/$360 for pro.
on 2026-04-24): $5/$30 input/output per 1M for chat, $30/$180 for pro.
Cache discount is 10% of input.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
@@ -77,7 +77,7 @@ async def test_anthropic_bedrock_thinking_blocks_with_none_content():
# test _bedrock_converse_messages_pt_async
result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
messages=messages,
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
llm_provider="bedrock",
)
@@ -0,0 +1,138 @@
"""
Regression tests for the parsed-URL hostname match used to identify a
caller-supplied ``api_base`` as a known openai-compatible provider.
The previous shape (``if endpoint in api_base:``) used unanchored
substring search, which let a caller pass
``https://attacker.com/api.groq.com/openai/v1`` and have the proxy
return ``GROQ_API_KEY`` as the dynamic credential exfiltrating the
server's real provider key to an attacker-controlled host on the
outbound request.
"""
import os
import sys
from unittest.mock import patch
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
from litellm.litellm_core_utils.get_llm_provider_logic import (
_endpoint_matches_api_base,
get_llm_provider,
)
class TestEndpointMatchesApiBase:
"""Direct unit tests on the parsed-URL matcher."""
@pytest.mark.parametrize(
"endpoint, api_base",
[
# Bare hostname endpoint, exact host match.
("api.perplexity.ai", "https://api.perplexity.ai/v1"),
# Endpoint includes a path; api_base path starts with it.
("api.groq.com/openai/v1", "https://api.groq.com/openai/v1"),
# Endpoint with full URL scheme.
("https://api.cerebras.ai/v1", "https://api.cerebras.ai/v1/chat"),
# Trailing-slash on registered endpoint must not break match.
("https://llm.chutes.ai/v1/", "https://llm.chutes.ai/v1/chat"),
# Case-insensitive on hostname.
("api.groq.com/openai/v1", "https://API.GROQ.COM/openai/v1"),
],
)
def test_legitimate_provider_urls_match(self, endpoint, api_base):
assert _endpoint_matches_api_base(endpoint, api_base) is True
@pytest.mark.parametrize(
"endpoint, api_base",
[
# Attacker host, registered endpoint smuggled into path.
(
"api.groq.com/openai/v1",
"https://attacker.com/api.groq.com/openai/v1",
),
# Attacker host, registered endpoint smuggled into a path segment.
(
"api.groq.com/openai/v1",
"https://attacker.com/foo/api.groq.com/openai/v1",
),
# Lookalike host that contains the registered host as a suffix label.
(
"api.groq.com/openai/v1",
"https://api.groq.com.attacker.com/openai/v1",
),
# Lookalike host with the registered host as a prefix.
(
"api.groq.com/openai/v1",
"https://api.groq.com.evil.example/openai/v1",
),
# Right host, wrong path — endpoint requires ``/openai/v1`` prefix.
("api.groq.com/openai/v1", "https://api.groq.com/v1"),
# Path-segment lookalike: ``/openai/v10`` must not match ``/openai/v1``.
("api.groq.com/openai/v1", "https://api.groq.com/openai/v10"),
# Userinfo / @-injection trick — the ``hostname`` after ``@`` is
# what httpx connects to.
(
"api.groq.com/openai/v1",
"https://api.groq.com@attacker.com/openai/v1",
),
],
)
def test_attacker_smuggling_does_not_match(self, endpoint, api_base):
assert _endpoint_matches_api_base(endpoint, api_base) is False
class TestGetLlmProviderRejectsAttackerSmuggledApiBase:
"""
End-to-end: ``get_llm_provider`` must NOT return the server's stored
secret (e.g. ``GROQ_API_KEY``) for an api_base whose hostname is
attacker-controlled, even when the registered endpoint string appears
elsewhere in the URL.
"""
def test_attacker_host_does_not_yield_groq_secret(self):
# The function may either fall through (different provider) or
# raise BadRequestError because the model can't be identified.
# The invariant under test is that ``GROQ_API_KEY`` is never
# looked up against an attacker-controlled hostname.
import litellm
with patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_secret_str",
return_value="server-real-groq-key",
) as mocked_secret:
try:
_, _, dynamic_api_key, _ = get_llm_provider(
model="some-model",
api_base="https://attacker.com/api.groq.com/openai/v1",
)
# If it returned, the dynamic key must not be the secret.
assert dynamic_api_key != "server-real-groq-key"
except litellm.exceptions.BadRequestError:
# Acceptable outcome: provider unidentifiable, no secret
# was returned.
pass
# Regardless of return / raise, the secret must never have been
# read against this attacker-controlled api_base.
groq_lookups = [
call
for call in mocked_secret.call_args_list
if call.args and call.args[0] == "GROQ_API_KEY"
]
assert groq_lookups == []
def test_legitimate_groq_api_base_still_resolves(self):
with patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_secret_str",
return_value="server-real-groq-key",
):
_, provider, dynamic_api_key, _ = get_llm_provider(
model="some-model",
api_base="https://api.groq.com/openai/v1",
)
assert provider == "groq"
assert dynamic_api_key == "server-real-groq-key"
@@ -2337,6 +2337,104 @@ def test_merge_hidden_params_from_response_into_metadata_populates_metadata():
assert meta["hidden_params"]["model_id"] == "mid-test"
def test_merge_hidden_params_from_response_into_metadata_backfills_response_cost():
"""Streaming metadata should include the already-calculated response cost."""
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
logging_obj = LiteLLMLoggingObj(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="acompletion",
start_time=time.time(),
litellm_call_id="merge-hp-cost-test",
function_id="merge-hp-cost-fn",
)
logging_obj.model_call_details = {
"litellm_params": {"metadata": {}},
"response_cost": 0.002,
}
class _Resp:
_hidden_params = {"response_cost": None, "model_id": "mid-test"}
response = _Resp()
logging_obj._merge_hidden_params_from_response_into_metadata(response)
meta = logging_obj.model_call_details["litellm_params"]["metadata"]
assert meta["hidden_params"]["response_cost"] == 0.002
assert meta["hidden_params"]["model_id"] == "mid-test"
assert response._hidden_params["response_cost"] is None
def test_standard_logging_hidden_params_backfills_response_cost_without_mutating_response():
"""Streaming standard logging payload should expose the calculated response cost."""
from datetime import datetime
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.utils import Usage
logging_obj = LiteLLMLoggingObj(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="acompletion",
start_time=time.time(),
litellm_call_id="standard-hp-cost-test",
function_id="standard-hp-cost-fn",
)
logging_obj.model_call_details = {
"litellm_params": {"metadata": {}, "proxy_server_request": {}},
"litellm_call_id": "standard-hp-cost-test",
"call_type": "acompletion",
"stream": True,
"model": "gpt-4o-mini",
"custom_llm_provider": "openai",
"optional_params": {"stream": True},
"response_cost": 0.002,
}
response = ModelResponse(
id="standard-hp-cost-response",
model="gpt-4o-mini",
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
)
response._hidden_params = {"response_cost": None, "model_id": "mid-test"}
payload = logging_obj._build_standard_logging_payload(
response, datetime.now(), datetime.now()
)
assert payload is not None
assert payload["hidden_params"]["response_cost"] == 0.002
assert response._hidden_params["response_cost"] is None
def test_merge_hidden_params_from_response_into_metadata_preserves_response_cost():
"""Do not overwrite provider-supplied response cost when it already exists."""
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
logging_obj = LiteLLMLoggingObj(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="acompletion",
start_time=time.time(),
litellm_call_id="merge-hp-preserve-cost-test",
function_id="merge-hp-preserve-cost-fn",
)
logging_obj.model_call_details = {
"litellm_params": {"metadata": {}},
"response_cost": 0.002,
}
class _Resp:
_hidden_params = {"response_cost": 0.001, "model_id": "mid-test"}
logging_obj._merge_hidden_params_from_response_into_metadata(_Resp())
meta = logging_obj.model_call_details["litellm_params"]["metadata"]
assert meta["hidden_params"]["response_cost"] == 0.001
assert meta["hidden_params"]["model_id"] == "mid-test"
def test_merge_hidden_params_from_response_into_metadata_no_op_when_empty():
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@@ -2436,3 +2534,141 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob
assert payload is not None
assert payload["litellm_call_id"] == call_id
def _make_dict_logging_obj():
"""Build a Logging instance configured for a non-streaming dict result."""
obj = LitellmLogging(
model="claude-haiku-4-5@20251001",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type="acompletion",
litellm_call_id="test-call-id",
start_time=time.time(),
function_id="test-fn",
)
obj.model_call_details = {
"model": "claude-haiku-4-5@20251001",
"custom_llm_provider": "vertex_ai",
"litellm_params": {"metadata": {}},
"response_cost": None,
}
return obj
def test_success_handler_computes_cost_for_dict_response():
"""Non-streaming dict responses run through the cost calculator."""
logging_obj = _make_dict_logging_obj()
expected_cost = 0.42
with (
patch.object(
logging_obj,
"_response_cost_calculator",
return_value=expected_cost,
) as mock_calc,
patch.object(
logging_obj,
"_build_standard_logging_payload",
return_value={"response_cost": expected_cost},
),
patch(
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
),
patch.object(
logging_obj,
"_is_recognized_call_type_for_logging",
return_value=False,
),
patch.object(
logging_obj,
"_transform_usage_objects",
side_effect=lambda result: result,
),
):
logging_obj.success_handler(
result={"id": "msg_1"},
start_time=time.time(),
end_time=time.time(),
)
mock_calc.assert_called_once()
assert logging_obj.model_call_details["response_cost"] == expected_cost
def test_success_handler_preserves_precomputed_cost_for_dict_response():
"""Precomputed response_cost on model_call_details must not be overwritten."""
logging_obj = _make_dict_logging_obj()
precomputed_cost = 1.23
logging_obj.model_call_details["response_cost"] = precomputed_cost
with (
patch.object(
logging_obj,
"_response_cost_calculator",
return_value=9.99,
) as mock_calc,
patch.object(
logging_obj,
"_build_standard_logging_payload",
return_value={"response_cost": precomputed_cost},
),
patch(
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
),
patch.object(
logging_obj,
"_is_recognized_call_type_for_logging",
return_value=False,
),
patch.object(
logging_obj,
"_transform_usage_objects",
side_effect=lambda result: result,
),
):
logging_obj.success_handler(
result={"id": "msg_2"},
start_time=time.time(),
end_time=time.time(),
)
mock_calc.assert_not_called()
assert logging_obj.model_call_details["response_cost"] == precomputed_cost
def test_success_handler_unified_helper_runs_for_typed_results():
"""Recognized typed responses still flow through the unified helper."""
logging_obj = _make_dict_logging_obj()
expected_cost = 0.10
typed_result = MagicMock()
typed_result._hidden_params = {}
with (
patch.object(
logging_obj,
"_response_cost_calculator",
return_value=expected_cost,
) as mock_calc,
patch.object(
logging_obj,
"_build_standard_logging_payload",
return_value={"response_cost": expected_cost},
),
patch(
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
),
patch.object(
logging_obj,
"_is_recognized_call_type_for_logging",
return_value=True,
),
patch.object(
logging_obj,
"_transform_usage_objects",
side_effect=lambda result: result,
),
):
logging_obj.success_handler(
result=typed_result,
start_time=time.time(),
end_time=time.time(),
)
mock_calc.assert_called_once()
assert logging_obj.model_call_details["response_cost"] == expected_cost
@@ -279,7 +279,7 @@ def test_reasoning_with_forced_tool_choice_switches_to_auto():
}
optional_params = config.map_openai_params(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
non_default_params=non_default_params,
optional_params={},
drop_params=False,
@@ -2797,7 +2797,7 @@ def test_thinking_with_max_completion_tokens():
result = config.map_openai_params(
non_default_params=non_default_params_with_max_completion,
optional_params=optional_params,
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
drop_params=False,
)
@@ -2819,7 +2819,7 @@ def test_thinking_with_max_completion_tokens():
result = config.map_openai_params(
non_default_params=non_default_params_with_max_tokens,
optional_params=optional_params,
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
drop_params=False,
)
@@ -2842,7 +2842,7 @@ def test_thinking_with_max_completion_tokens():
result = config.map_openai_params(
non_default_params=non_default_params_without_max,
optional_params=optional_params,
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
drop_params=False,
)
@@ -3617,7 +3617,7 @@ class TestBedrockMinThinkingBudgetTokens:
"""Test that thinking.budget_tokens is clamped to the Bedrock minimum (1024)."""
def _map_params(
self, thinking_value, model="anthropic.claude-3-7-sonnet-20250219-v1:0"
self, thinking_value, model="anthropic.claude-sonnet-4-5-20250929-v1:0"
):
"""Helper to call map_openai_params with the given thinking value."""
config = AmazonConverseConfig()
@@ -3651,7 +3651,7 @@ class TestBedrockMinThinkingBudgetTokens:
result = config.map_openai_params(
non_default_params={},
optional_params={},
model="anthropic.claude-3-7-sonnet-20250219-v1:0",
model="anthropic.claude-sonnet-4-5-20250929-v1:0",
drop_params=False,
)
assert "thinking" not in result or result.get("thinking") is None
@@ -4146,3 +4146,39 @@ def test_transform_response_finish_reason_stop_when_json_mode_filters_all_tools(
# finish_reason must be "stop", not "tool_calls"
assert result.choices[0].finish_reason == "stop"
def test_transform_response_does_not_leak_body_on_parse_failure():
from litellm.llms.bedrock.common_utils import BedrockError
leaky_body = {"output": {"message": {"content": [{"text": "secret content"}]}}}
class MockResponse:
def json(self):
return leaky_body
@property
def text(self):
return json.dumps(leaky_body)
with patch(
"litellm.llms.bedrock.chat.converse_transformation.ConverseResponseBlock",
side_effect=KeyError("missing required field"),
):
with pytest.raises(BedrockError) as exc_info:
AmazonConverseConfig()._transform_response(
model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
response=MockResponse(),
model_response=ModelResponse(),
stream=False,
logging_obj=None,
optional_params={},
api_key=None,
data=None,
messages=[],
encoding=None,
)
msg = str(exc_info.value)
assert "secret content" not in msg
assert "Error converting to valid response block" in msg
@@ -21,7 +21,112 @@ def test_transform_search_request():
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
litellm_logging_obj=mock_log,
litellm_params={},
extra_body=None,
)
assert url.endswith("/kb123/retrieve")
assert body["retrievalQuery"].get("text") == "hello"
def test_transform_search_request_uses_only_retrieval_config_from_extra_body():
config = BedrockVectorStoreConfig()
mock_log = MagicMock()
mock_log.model_call_details = {}
url, body = config.transform_search_vector_store_request(
vector_store_id="kb123",
query="hello",
vector_store_search_optional_params={},
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
litellm_logging_obj=mock_log,
litellm_params={},
extra_body={
"retrievalConfiguration": {
"vectorSearchConfiguration": {
"overrideSearchType": "HYBRID",
"numberOfResults": 8,
}
},
"unrelatedField": {"should_not": "be_forwarded"},
},
)
assert url.endswith("/kb123/retrieve")
assert body["retrievalQuery"].get("text") == "hello"
assert (
body["retrievalConfiguration"]["vectorSearchConfiguration"][
"overrideSearchType"
]
== "HYBRID"
)
assert "unrelatedField" not in body
def test_transform_search_request_does_not_mutate_extra_body_and_overrides_number_of_results():
config = BedrockVectorStoreConfig()
mock_log = MagicMock()
mock_log.model_call_details = {}
extra_body = {
"retrievalConfiguration": {
"vectorSearchConfiguration": {
"overrideSearchType": "HYBRID",
"numberOfResults": 8,
}
}
}
_, body = config.transform_search_vector_store_request(
vector_store_id="kb123",
query="hello",
vector_store_search_optional_params={"max_num_results": 10},
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
litellm_logging_obj=mock_log,
litellm_params={},
extra_body=extra_body,
)
assert (
body["retrievalConfiguration"]["vectorSearchConfiguration"]["numberOfResults"]
== 10
)
assert (
extra_body["retrievalConfiguration"]["vectorSearchConfiguration"][
"numberOfResults"
]
== 8
)
def test_transform_search_request_overrides_filter_without_mutating_extra_body():
config = BedrockVectorStoreConfig()
mock_log = MagicMock()
mock_log.model_call_details = {}
extra_body = {
"retrievalConfiguration": {
"vectorSearchConfiguration": {
"filter": {"equals": {"key": "tenant", "value": "a"}}
}
}
}
new_filter = {"equals": {"key": "tenant", "value": "b"}}
_, body = config.transform_search_vector_store_request(
vector_store_id="kb123",
query="hello",
vector_store_search_optional_params={"filters": new_filter},
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
litellm_logging_obj=mock_log,
litellm_params={},
extra_body=extra_body,
)
assert (
body["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"]
== new_filter
)
assert (
extra_body["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"][
"equals"
]["value"]
== "a"
)
@@ -0,0 +1,161 @@
import socket
from unittest.mock import MagicMock, patch
def _invoke_connector_factory(http_handler_module):
"""
Drive the lambda factory installed on the transport so TCPConnector is
actually constructed. _create_aiohttp_transport returns a transport whose
_client_factory is the lambda that builds (TCPConnector ClientSession);
invoking it directly avoids relying on _get_valid_client_session's internal
branching to trigger connector construction.
"""
transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(
shared_session=None
)
transport._client_factory()
return transport
def test_socket_factory_omitted_when_disabled(monkeypatch):
from litellm.llms.custom_httpx import http_handler as http_handler_module
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", False)
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
connector_mock = MagicMock(name="connector")
session_mock = MagicMock(name="session")
with patch.object(
http_handler_module, "TCPConnector", return_value=connector_mock
) as mock_tcp_connector:
with patch.object(
http_handler_module, "ClientSession", return_value=session_mock
):
_invoke_connector_factory(http_handler_module)
assert mock_tcp_connector.call_count >= 1
assert "socket_factory" not in mock_tcp_connector.call_args.kwargs
def test_socket_factory_attached_when_enabled(monkeypatch):
from litellm.llms.custom_httpx import http_handler as http_handler_module
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
connector_mock = MagicMock(name="connector")
session_mock = MagicMock(name="session")
with patch.object(
http_handler_module, "TCPConnector", return_value=connector_mock
) as mock_tcp_connector:
with patch.object(
http_handler_module, "ClientSession", return_value=session_mock
):
_invoke_connector_factory(http_handler_module)
assert mock_tcp_connector.call_count >= 1
factory = mock_tcp_connector.call_args.kwargs.get("socket_factory")
assert callable(factory)
def test_socket_factory_skipped_on_old_aiohttp(monkeypatch):
from litellm.llms.custom_httpx import http_handler as http_handler_module
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", False)
connector_mock = MagicMock(name="connector")
session_mock = MagicMock(name="session")
with patch.object(
http_handler_module, "TCPConnector", return_value=connector_mock
) as mock_tcp_connector:
with patch.object(
http_handler_module, "ClientSession", return_value=session_mock
):
_invoke_connector_factory(http_handler_module)
assert mock_tcp_connector.call_count >= 1
assert "socket_factory" not in mock_tcp_connector.call_args.kwargs
def test_socket_factory_sets_keepalive_options(monkeypatch):
from litellm.llms.custom_httpx import http_handler as http_handler_module
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPIDLE", 45)
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPINTVL", 15)
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPCNT", 4)
factory = http_handler_module._build_aiohttp_keepalive_socket_factory()
assert factory is not None
addr_info = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", ("", 0))
fake_sock = MagicMock(spec=socket.socket)
with patch("socket.socket", return_value=fake_sock) as sock_ctor:
returned = factory(addr_info)
sock_ctor.assert_called_once_with(
family=socket.AF_INET, type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP
)
assert returned is fake_sock
fake_sock.setblocking.assert_called_once_with(False)
setsockopt_calls = {
(call.args[0], call.args[1]): call.args[2]
for call in fake_sock.setsockopt.call_args_list
}
assert setsockopt_calls[(socket.SOL_SOCKET, socket.SO_KEEPALIVE)] == 1
if hasattr(socket, "TCP_KEEPIDLE"):
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE)] == 45
elif hasattr(socket, "TCP_KEEPALIVE"):
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE)] == 45
if hasattr(socket, "TCP_KEEPINTVL"):
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL)] == 15
if hasattr(socket, "TCP_KEEPCNT"):
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPCNT)] == 4
def test_socket_factory_uses_tcp_keepalive_when_keepidle_unavailable(monkeypatch):
"""
Cover the macOS/Darwin branch: when TCP_KEEPIDLE is missing but TCP_KEEPALIVE
is present, the factory should fall back to TCP_KEEPALIVE for the idle timer.
Linux CI runners always have TCP_KEEPIDLE, so we patch socket itself to
simulate the BSD-derived environment.
"""
from litellm.llms.custom_httpx import http_handler as http_handler_module
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPIDLE", 60)
factory = http_handler_module._build_aiohttp_keepalive_socket_factory()
assert factory is not None
fake_socket_module = MagicMock(spec=[])
fake_socket_module.SOL_SOCKET = socket.SOL_SOCKET
fake_socket_module.SO_KEEPALIVE = socket.SO_KEEPALIVE
fake_socket_module.IPPROTO_TCP = socket.IPPROTO_TCP
fake_socket_module.TCP_KEEPALIVE = getattr(socket, "TCP_KEEPALIVE", 0x10)
fake_sock = MagicMock(spec=socket.socket)
fake_socket_module.socket = MagicMock(return_value=fake_sock)
addr_info = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", ("", 0))
with patch.object(http_handler_module, "socket", fake_socket_module):
factory(addr_info)
setsockopt_calls = {
(call.args[0], call.args[1]): call.args[2]
for call in fake_sock.setsockopt.call_args_list
}
assert setsockopt_calls[(socket.SOL_SOCKET, socket.SO_KEEPALIVE)] == 1
assert (
setsockopt_calls[(socket.IPPROTO_TCP, fake_socket_module.TCP_KEEPALIVE)] == 60
)
assert (socket.IPPROTO_TCP, getattr(socket, "TCP_KEEPIDLE", -1)) not in setsockopt_calls
@@ -2,12 +2,16 @@ import os
import sys
from unittest.mock import AsyncMock, Mock, patch
import httpx
import pytest
sys.path.insert(
0, os.path.abspath("../../../..")
) # Adds the parent directory to the system path
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import (
BaseLLMHTTPHandler,
_google_genai_streaming_hidden_params,
)
from litellm.types.router import GenericLiteLLMParams
@@ -320,3 +324,29 @@ async def test_async_anthropic_messages_handler_header_priority():
assert captured_headers["X-Forwarded-Only"] == "keep"
assert captured_headers["X-Extra-Only"] == "also-keep"
assert captured_headers["X-Provider-Only"] == "keep-this-too"
def test_google_genai_streaming_hidden_params_model_info_and_router_fallback():
logging_obj = Mock()
logging_obj.get_router_model_id = Mock(return_value="router-model-id")
from_model_info = _google_genai_streaming_hidden_params(
api_base="https://generativelanguage.googleapis.com/v1beta",
litellm_params=GenericLiteLLMParams(model_info={"id": "info-id"}),
logging_obj=logging_obj,
response_headers=httpx.Headers({"x-ratelimit-remaining": "10"}),
)
assert from_model_info["model_id"] == "info-id"
assert (
from_model_info["api_base"]
== "https://generativelanguage.googleapis.com/v1beta"
)
assert isinstance(from_model_info["additional_headers"], dict)
from_router = _google_genai_streaming_hidden_params(
api_base="https://x",
litellm_params=GenericLiteLLMParams(),
logging_obj=logging_obj,
response_headers=httpx.Headers({}),
)
assert from_router["model_id"] == "router-model-id"
@@ -59,6 +59,7 @@ class TestS3VectorsVectorStoreConfig:
api_base="https://s3vectors.us-west-2.api.aws",
litellm_logging_obj=mock_logging_obj,
litellm_params={},
extra_body=None,
)
def test_transform_search_response(self):
@@ -4261,3 +4261,32 @@ def test_sync_streaming_uses_custom_client():
# Verify that gemini_client is in the partial's keywords
assert "gemini_client" in partial_make_sync_call.keywords
assert partial_make_sync_call.keywords["gemini_client"] is mock_client
def test_transform_response_does_not_leak_body_on_parse_failure():
leaky_body = {"candidates": [{"content": {"parts": [{"text": "secret content"}]}}]}
raw_response = MagicMock()
raw_response.json.return_value = leaky_body
raw_response.text = json.dumps(leaky_body)
raw_response.headers = {}
with patch(
"litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.GenerateContentResponseBody",
side_effect=KeyError("missing required field"),
):
with pytest.raises(VertexAIError) as exc_info:
VertexGeminiConfig().transform_response(
model="gemini-pro",
raw_response=raw_response,
model_response=ModelResponse(),
logging_obj=MagicMock(),
request_data={},
messages=[],
optional_params={},
litellm_params={},
encoding=None,
)
msg = str(exc_info.value)
assert "secret content" not in msg
assert "Error converting to valid response block" in msg

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