mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-16 14:17:10 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_fix-config-update-targeted-upserts
This commit is contained in:
@@ -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,7 +4,7 @@ on:
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
tag:
|
||||
description: "Release tag (e.g. v1.83.0-stable) — branch will be named release/<tag>"
|
||||
description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted) — branch will be named release/<tag>"
|
||||
required: true
|
||||
type: string
|
||||
commit_hash:
|
||||
@@ -14,7 +14,7 @@ on:
|
||||
workflow_call:
|
||||
inputs:
|
||||
tag:
|
||||
description: "Release tag"
|
||||
description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted)"
|
||||
required: true
|
||||
type: string
|
||||
commit_hash:
|
||||
@@ -40,8 +40,8 @@ jobs:
|
||||
echo "::error::commit_hash must be a full 40-character commit SHA"
|
||||
exit 1
|
||||
fi
|
||||
if ! echo "${TAG}" | grep -qE '^v[0-9]+\.[0-9]+\.[0-9]+'; then
|
||||
echo "::error::tag must start with vX.Y.Z"
|
||||
if ! echo "${TAG}" | grep -qE '^v?[0-9]+\.[0-9]+\.[0-9]+'; then
|
||||
echo "::error::tag must start with X.Y.Z (optional leading v), e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, or v1.83.10-stable"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ on:
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
tag:
|
||||
description: "Release tag (e.g. v1.83.0-stable)"
|
||||
description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted)"
|
||||
required: true
|
||||
type: string
|
||||
commit_hash:
|
||||
@@ -30,8 +30,8 @@ jobs:
|
||||
echo "::error::commit_hash must be a full 40-character commit SHA"
|
||||
exit 1
|
||||
fi
|
||||
if ! echo "${TAG}" | grep -qE '^v[0-9]+\.[0-9]+\.[0-9]+'; then
|
||||
echo "::error::tag must start with vX.Y.Z"
|
||||
if ! echo "${TAG}" | grep -qE '^v?[0-9]+\.[0-9]+\.[0-9]+'; then
|
||||
echo "::error::tag must start with X.Y.Z (optional leading v), e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, or v1.83.10-stable"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
@@ -45,6 +45,11 @@ jobs:
|
||||
const tag = process.env.TAG;
|
||||
const commitHash = process.env.COMMIT_HASH;
|
||||
|
||||
// Mark RC / dev / nightly / alpha / beta tags as GitHub pre-releases.
|
||||
// PEP 440 post-releases (e.g. `1.84.0.post1`) and legacy `-stable[.patch.N]`
|
||||
// are stable maintenance releases, not pre-releases.
|
||||
const isPrerelease = /(?:rc|nightly|alpha|beta|\.dev)/i.test(tag);
|
||||
|
||||
const cosignSection = [
|
||||
`## Verify Docker Image Signature`,
|
||||
``,
|
||||
@@ -89,7 +94,7 @@ jobs:
|
||||
target_commitish: commitHash,
|
||||
name: tag,
|
||||
owner: context.repo.owner,
|
||||
prerelease: false,
|
||||
prerelease: isPrerelease,
|
||||
repo: context.repo.repo,
|
||||
tag_name: tag,
|
||||
});
|
||||
|
||||
+75
@@ -0,0 +1,75 @@
|
||||
-- CreateTable
|
||||
CREATE TABLE "LiteLLM_WorkflowRun" (
|
||||
"run_id" TEXT NOT NULL,
|
||||
"session_id" TEXT NOT NULL,
|
||||
"workflow_type" TEXT NOT NULL,
|
||||
"status" TEXT NOT NULL DEFAULT 'pending',
|
||||
"created_by" TEXT,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL,
|
||||
"input" JSONB,
|
||||
"output" JSONB,
|
||||
"metadata" JSONB,
|
||||
|
||||
CONSTRAINT "LiteLLM_WorkflowRun_pkey" PRIMARY KEY ("run_id")
|
||||
);
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE "LiteLLM_WorkflowEvent" (
|
||||
"event_id" TEXT NOT NULL,
|
||||
"run_id" TEXT NOT NULL,
|
||||
"event_type" TEXT NOT NULL,
|
||||
"step_name" TEXT NOT NULL,
|
||||
"sequence_number" INTEGER NOT NULL,
|
||||
"data" JSONB,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
|
||||
CONSTRAINT "LiteLLM_WorkflowEvent_pkey" PRIMARY KEY ("event_id")
|
||||
);
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE "LiteLLM_WorkflowMessage" (
|
||||
"message_id" TEXT NOT NULL,
|
||||
"run_id" TEXT NOT NULL,
|
||||
"role" TEXT NOT NULL,
|
||||
"content" TEXT NOT NULL,
|
||||
"sequence_number" INTEGER NOT NULL,
|
||||
"session_id" TEXT,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
|
||||
CONSTRAINT "LiteLLM_WorkflowMessage_pkey" PRIMARY KEY ("message_id")
|
||||
);
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX "LiteLLM_WorkflowRun_session_id_key" ON "LiteLLM_WorkflowRun"("session_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_WorkflowRun_workflow_type_status_idx" ON "LiteLLM_WorkflowRun"("workflow_type", "status");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_WorkflowRun_session_id_idx" ON "LiteLLM_WorkflowRun"("session_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_WorkflowRun_created_at_idx" ON "LiteLLM_WorkflowRun"("created_at");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_WorkflowRun_created_by_idx" ON "LiteLLM_WorkflowRun"("created_by");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_WorkflowEvent_run_id_idx" ON "LiteLLM_WorkflowEvent"("run_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX "LiteLLM_WorkflowEvent_run_id_sequence_number_key" ON "LiteLLM_WorkflowEvent"("run_id", "sequence_number");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_WorkflowMessage_run_id_idx" ON "LiteLLM_WorkflowMessage"("run_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX "LiteLLM_WorkflowMessage_run_id_sequence_number_key" ON "LiteLLM_WorkflowMessage"("run_id", "sequence_number");
|
||||
|
||||
-- AddForeignKey
|
||||
ALTER TABLE "LiteLLM_WorkflowEvent" ADD CONSTRAINT "LiteLLM_WorkflowEvent_run_id_fkey" FOREIGN KEY ("run_id") REFERENCES "LiteLLM_WorkflowRun"("run_id") ON DELETE RESTRICT ON UPDATE CASCADE;
|
||||
|
||||
-- AddForeignKey
|
||||
ALTER TABLE "LiteLLM_WorkflowMessage" ADD CONSTRAINT "LiteLLM_WorkflowMessage_run_id_fkey" FOREIGN KEY ("run_id") REFERENCES "LiteLLM_WorkflowRun"("run_id") ON DELETE RESTRICT ON UPDATE CASCADE;
|
||||
|
||||
@@ -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])
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -101,5 +101,10 @@
|
||||
"param_mappings": {
|
||||
"max_completion_tokens": "max_tokens"
|
||||
}
|
||||
},
|
||||
"aihubmix": {
|
||||
"base_url": "https://aihubmix.com/v1",
|
||||
"api_key_env": "AIHUBMIX_API_KEY",
|
||||
"api_base_env": "AIHUBMIX_API_BASE"
|
||||
}
|
||||
}
|
||||
|
||||
@@ -80,6 +80,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig):
|
||||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict]:
|
||||
url = f"{api_base}/{vector_store_id}/search"
|
||||
_, request_body = super().transform_search_vector_store_request(
|
||||
@@ -89,5 +90,6 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig):
|
||||
api_base=api_base,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
return url, request_body
|
||||
|
||||
@@ -102,6 +102,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig):
|
||||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict]:
|
||||
"""RAGFlow vector stores are management-only, search is not supported."""
|
||||
raise NotImplementedError(
|
||||
|
||||
@@ -79,6 +79,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
|
||||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict]:
|
||||
"""Sync version - generates embedding synchronously."""
|
||||
# For S3 Vectors, vector_store_id should be in format: bucket_name:index_name
|
||||
@@ -140,6 +141,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
|
||||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict]:
|
||||
"""Async version - generates embedding asynchronously."""
|
||||
# For S3 Vectors, vector_store_id should be in format: bucket_name:index_name
|
||||
|
||||
@@ -4,6 +4,7 @@ from typing import Any, Coroutine, Dict, Optional, Union
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
_get_httpx_client,
|
||||
get_async_httpx_client,
|
||||
@@ -224,8 +225,14 @@ class VertexAIBatchPrediction(VertexLLM):
|
||||
},
|
||||
)
|
||||
|
||||
response = sync_handler.get(
|
||||
url=api_base,
|
||||
# ``api_base`` here can come from caller-supplied request kwargs
|
||||
# (clientside override). Wrap the fetch in ``safe_get`` so DNS
|
||||
# rebind / private / cloud-metadata targets are rejected; the
|
||||
# proxy auth gate already blocks malicious clientside ``api_base``
|
||||
# at the boundary — this is defense-in-depth for SDK callers.
|
||||
response = safe_get(
|
||||
sync_handler,
|
||||
api_base,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
@@ -270,8 +277,13 @@ class VertexAIBatchPrediction(VertexLLM):
|
||||
},
|
||||
)
|
||||
|
||||
response = await client.get(
|
||||
url=api_base,
|
||||
# Mirror the sync path: ``api_base`` may come from caller-supplied
|
||||
# request kwargs, so wrap the fetch in ``async_safe_get`` to reject
|
||||
# DNS-rebind / private / cloud-metadata targets. Defense-in-depth
|
||||
# behind the proxy auth gate's clientside ``api_base`` check.
|
||||
response = await async_safe_get(
|
||||
client,
|
||||
api_base,
|
||||
headers=headers,
|
||||
)
|
||||
if response.status_code != 200:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -2,10 +2,11 @@
|
||||
MCP Server Utilities
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, Mapping, Optional, Tuple
|
||||
from typing import Any, Dict, Iterator, Mapping, Optional, Tuple
|
||||
|
||||
import os
|
||||
import hashlib
|
||||
import importlib
|
||||
import os
|
||||
|
||||
# Constants
|
||||
LITELLM_MCP_SERVER_NAME = "litellm-mcp-server"
|
||||
@@ -14,6 +15,89 @@ LITELLM_MCP_SERVER_DESCRIPTION = "MCP Server for LiteLLM"
|
||||
MCP_TOOL_PREFIX_SEPARATOR = os.environ.get("MCP_TOOL_PREFIX_SEPARATOR", "-")
|
||||
MCP_TOOL_PREFIX_FORMAT = "{server_name}{separator}{tool_name}"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Short-ID tool prefix (opt-in)
|
||||
# ---------------------------------------------------------------------------
|
||||
# When LITELLM_USE_SHORT_MCP_TOOL_PREFIX is truthy the prefix attached to MCP
|
||||
# tool / prompt / resource / resource-template names switches from the
|
||||
# (potentially long) human-readable server name to a deterministic three
|
||||
# character ID derived from the server's ``server_id``.
|
||||
#
|
||||
# Why three characters?
|
||||
# * The first character is restricted to 52 alphabetic characters
|
||||
# ([A-Za-z]) and the remaining two characters use the full base62
|
||||
# alphabet ([0-9A-Za-z]). That guarantees the prefix never starts
|
||||
# with a digit so it remains a valid identifier for every model API
|
||||
# (some providers historically required a leading alphabetic char).
|
||||
# * 52 * 62 * 62 = 199_888 distinct IDs. The chance of a real local
|
||||
# tool name happening to begin with the exact prefix LiteLLM assigned
|
||||
# to a given MCP server is negligible in practice.
|
||||
# * The IDs are short enough that prefixed tool names stay well under
|
||||
# the 60-character upper bound enforced by some model APIs (Anthropic
|
||||
# etc.) even for long upstream tool names.
|
||||
# * The mapping is deterministic (SHA-256 of ``server_id`` → three
|
||||
# characters drawn from the alphabets above), so the prefix is stable
|
||||
# across processes, workers and restarts without any persistence
|
||||
# layer. Two servers with different ``server_id`` values can in
|
||||
# principle hash to the same three chars; that natural-hash collision
|
||||
# IS a routing-correctness issue (the second registrant would otherwise
|
||||
# have its tools misrouted to the first), so registration goes through
|
||||
# ``MCPServerManager._assign_unique_short_prefix`` which rehashes with
|
||||
# a deterministic attempt counter until it finds an unused prefix and
|
||||
# caches the result on ``MCPServer.short_prefix``. A collision is
|
||||
# logged at INFO when it happens.
|
||||
#
|
||||
# This flag is intentionally opt-in for the first release so customers can
|
||||
# migrate. It will become the default in a future release.
|
||||
SHORT_MCP_TOOL_PREFIX_LENGTH = 3
|
||||
_BASE62_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
||||
# Subset of _BASE62_ALPHABET used for the *first* character only, to
|
||||
# guarantee the prefix never starts with a digit.
|
||||
_BASE52_ALPHA_ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
||||
|
||||
|
||||
def is_short_mcp_tool_prefix_enabled() -> bool:
|
||||
"""Return True when the short-ID tool prefix mode is enabled.
|
||||
|
||||
Read at call time (not import time) so tests and runtime config changes
|
||||
take effect without reimporting the module.
|
||||
"""
|
||||
raw = os.environ.get("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "")
|
||||
return raw.strip().lower() in ("1", "true", "yes", "on")
|
||||
|
||||
|
||||
def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str:
|
||||
"""Derive the deterministic three-character prefix for a server.
|
||||
|
||||
Uses SHA-256 of ``f"{server_id}#{attempt}"`` and folds the first eight
|
||||
bytes into a fixed-length string whose first character is drawn from
|
||||
``_BASE52_ALPHA_ALPHABET`` (so the prefix never starts with a digit)
|
||||
and whose remaining characters are drawn from the full base62
|
||||
alphabet. Pass ``attempt > 0`` to rehash to a different prefix when
|
||||
the natural hash collides with a prefix already assigned to another
|
||||
server (see ``MCPServerManager._assign_unique_short_prefix``). An
|
||||
empty ``server_id`` raises ``ValueError`` — short prefixes require a
|
||||
stable identifier to be deterministic.
|
||||
"""
|
||||
if not server_id:
|
||||
raise ValueError("compute_short_server_prefix requires a non-empty server_id")
|
||||
|
||||
seed = server_id if attempt == 0 else f"{server_id}#{attempt}"
|
||||
digest = hashlib.sha256(seed.encode("utf-8")).digest()
|
||||
value = int.from_bytes(digest[:8], "big")
|
||||
|
||||
# Build chars from least-significant to most-significant; we reverse
|
||||
# at the end so the first emitted char comes from the high-order
|
||||
# bits of the digest (which is the position we constrain to be
|
||||
# alphabetic).
|
||||
chars = []
|
||||
for position in range(SHORT_MCP_TOOL_PREFIX_LENGTH):
|
||||
is_first_char = position == SHORT_MCP_TOOL_PREFIX_LENGTH - 1
|
||||
alphabet = _BASE52_ALPHA_ALPHABET if is_first_char else _BASE62_ALPHABET
|
||||
value, idx = divmod(value, len(alphabet))
|
||||
chars.append(alphabet[idx])
|
||||
return "".join(reversed(chars))
|
||||
|
||||
|
||||
def is_mcp_available() -> bool:
|
||||
"""
|
||||
@@ -82,7 +166,25 @@ def add_server_prefix_to_name(name: str, server_name: str) -> str:
|
||||
|
||||
|
||||
def get_server_prefix(server: Any) -> str:
|
||||
"""Return the prefix for a server: alias if present, else server_name, else server_id"""
|
||||
"""Return the prefix for a server.
|
||||
|
||||
When the short-prefix mode is enabled (``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``)
|
||||
a three-character base62 ID is returned. We prefer the cached
|
||||
``server.short_prefix`` value when set — that field is populated at
|
||||
registration time by ``MCPServerManager._assign_unique_short_prefix``
|
||||
and resolves natural-hash collisions deterministically — and only fall
|
||||
back to the natural hash for ad-hoc / temp-server objects without a
|
||||
cached value. In default mode the historical behaviour is preserved:
|
||||
alias if present, else server_name, else server_id.
|
||||
"""
|
||||
if is_short_mcp_tool_prefix_enabled():
|
||||
cached = getattr(server, "short_prefix", None)
|
||||
if cached:
|
||||
return cached
|
||||
server_id = getattr(server, "server_id", None)
|
||||
if server_id:
|
||||
return compute_short_server_prefix(server_id)
|
||||
|
||||
if hasattr(server, "alias") and server.alias:
|
||||
return server.alias
|
||||
if hasattr(server, "server_name") and server.server_name:
|
||||
@@ -92,6 +194,36 @@ def get_server_prefix(server: Any) -> str:
|
||||
return ""
|
||||
|
||||
|
||||
def iter_known_server_prefixes(server: Any) -> Iterator[str]:
|
||||
"""Yield every prefix form that may appear in tool names for ``server``.
|
||||
|
||||
Always includes the *current* prefix returned by ``get_server_prefix``.
|
||||
Additionally yields the historical (alias / server_name / server_id) and
|
||||
short-ID forms so the routing layer can resolve tool names regardless of
|
||||
which prefix mode was active when the client first observed them.
|
||||
"""
|
||||
seen = set()
|
||||
|
||||
def _emit(value: Optional[str]) -> Iterator[str]:
|
||||
if value and value not in seen:
|
||||
seen.add(value)
|
||||
yield value
|
||||
|
||||
yield from _emit(get_server_prefix(server))
|
||||
yield from _emit(getattr(server, "short_prefix", None))
|
||||
|
||||
server_id = getattr(server, "server_id", None)
|
||||
if server_id:
|
||||
try:
|
||||
yield from _emit(compute_short_server_prefix(server_id))
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
yield from _emit(getattr(server, "alias", None))
|
||||
yield from _emit(getattr(server, "server_name", None))
|
||||
yield from _emit(server_id)
|
||||
|
||||
|
||||
def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]:
|
||||
"""Return the unprefixed name plus the server name used as prefix."""
|
||||
if MCP_TOOL_PREFIX_SEPARATOR in prefixed_name:
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -1,22 +0,0 @@
|
||||
1:"$Sreact.fragment"
|
||||
2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
|
||||
3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
|
||||
4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
6:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"]
|
||||
7:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"]
|
||||
a:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"]
|
||||
b:"$Sreact.suspense"
|
||||
d:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"]
|
||||
f:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"]
|
||||
11:I[168027,[],"default"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}]
|
||||
0:{"P":null,"b":"zxkD4-EPlgfKHDTw8O869","c":["","chat"],"q":"","i":false,"f":[[["",{"children":["chat",{"children":["__PAGE__",{}]}]},"$undefined","$undefined",true],[["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true,"nonce":"$undefined"}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]],"forbidden":"$undefined","unauthorized":"$undefined"}]}]}]}]}]]}],{"children":[["$","$1","c",{"children":[null,["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":"$undefined","forbidden":"$undefined","unauthorized":"$undefined"}]]}],{"children":[["$","$1","c",{"children":[["$","$L6",null,{"Component":"$7","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@8","$@9"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true,"nonce":"$undefined"}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true,"nonce":"$undefined"}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true,"nonce":"$undefined"}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true,"nonce":"$undefined"}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true,"nonce":"$undefined"}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true,"nonce":"$undefined"}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true,"nonce":"$undefined"}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true,"nonce":"$undefined"}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true,"nonce":"$undefined"}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true,"nonce":"$undefined"}]],["$","$La",null,{"children":["$","$b",null,{"name":"Next.MetadataOutlet","children":"$@c"}]}]]}],{},null,false,false]},null,false,false]},null,false,false],["$","$1","h",{"children":[null,["$","$Ld",null,{"children":"$Le"}],["$","div",null,{"hidden":true,"children":["$","$Lf",null,{"children":["$","$b",null,{"name":"Next.Metadata","children":"$L10"}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],false]],"m":"$undefined","G":["$11",[]],"S":true}
|
||||
8:{}
|
||||
9:"$0:f:0:1:1:children:1:children:0:props:children:0:props:serverProvidedParams:params"
|
||||
e:[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]]
|
||||
12:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"]
|
||||
c:null
|
||||
10:[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L12","4",{}]]
|
||||
@@ -1,22 +0,0 @@
|
||||
1:"$Sreact.fragment"
|
||||
2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
|
||||
3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
|
||||
4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
6:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"]
|
||||
7:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"]
|
||||
a:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"]
|
||||
b:"$Sreact.suspense"
|
||||
d:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"]
|
||||
f:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"]
|
||||
11:I[168027,[],"default"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}]
|
||||
0:{"P":null,"b":"zxkD4-EPlgfKHDTw8O869","c":["","chat"],"q":"","i":false,"f":[[["",{"children":["chat",{"children":["__PAGE__",{}]}]},"$undefined","$undefined",true],[["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true,"nonce":"$undefined"}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]],"forbidden":"$undefined","unauthorized":"$undefined"}]}]}]}]}]]}],{"children":[["$","$1","c",{"children":[null,["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":"$undefined","forbidden":"$undefined","unauthorized":"$undefined"}]]}],{"children":[["$","$1","c",{"children":[["$","$L6",null,{"Component":"$7","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@8","$@9"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true,"nonce":"$undefined"}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true,"nonce":"$undefined"}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true,"nonce":"$undefined"}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true,"nonce":"$undefined"}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true,"nonce":"$undefined"}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true,"nonce":"$undefined"}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true,"nonce":"$undefined"}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true,"nonce":"$undefined"}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true,"nonce":"$undefined"}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true,"nonce":"$undefined"}]],["$","$La",null,{"children":["$","$b",null,{"name":"Next.MetadataOutlet","children":"$@c"}]}]]}],{},null,false,false]},null,false,false]},null,false,false],["$","$1","h",{"children":[null,["$","$Ld",null,{"children":"$Le"}],["$","div",null,{"hidden":true,"children":["$","$Lf",null,{"children":["$","$b",null,{"name":"Next.Metadata","children":"$L10"}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],false]],"m":"$undefined","G":["$11",[]],"S":true}
|
||||
8:{}
|
||||
9:"$0:f:0:1:1:children:1:children:0:props:children:0:props:serverProvidedParams:params"
|
||||
e:[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]]
|
||||
12:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"]
|
||||
c:null
|
||||
10:[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L12","4",{}]]
|
||||
@@ -1,6 +0,0 @@
|
||||
1:"$Sreact.fragment"
|
||||
2:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"]
|
||||
3:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"]
|
||||
4:"$Sreact.suspense"
|
||||
5:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"]
|
||||
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","h",{"children":[null,["$","$L2",null,{"children":[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]]}],["$","div",null,{"hidden":true,"children":["$","$L3",null,{"children":["$","$4",null,{"name":"Next.Metadata","children":[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L5","4",{}]]}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],"loading":null,"isPartial":false}
|
||||
@@ -1,8 +0,0 @@
|
||||
1:"$Sreact.fragment"
|
||||
2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
|
||||
3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
|
||||
4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
|
||||
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","template":["$","$L5",null,{}],"notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]]}]}]}]}]}]]}],"loading":null,"isPartial":false}
|
||||
@@ -1,4 +0,0 @@
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}]
|
||||
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","tree":{"name":"","paramType":null,"paramKey":"","hasRuntimePrefetch":false,"slots":{"children":{"name":"chat","paramType":null,"paramKey":"chat","hasRuntimePrefetch":false,"slots":{"children":{"name":"__PAGE__","paramType":null,"paramKey":"__PAGE__","hasRuntimePrefetch":false,"slots":null,"isRootLayout":false}},"isRootLayout":false}},"isRootLayout":true},"staleTime":300}
|
||||
@@ -1,9 +0,0 @@
|
||||
1:"$Sreact.fragment"
|
||||
2:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"]
|
||||
3:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"]
|
||||
6:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"]
|
||||
7:"$Sreact.suspense"
|
||||
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[["$","$L2",null,{"Component":"$3","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@4","$@5"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true}]],["$","$L6",null,{"children":["$","$7",null,{"name":"Next.MetadataOutlet","children":"$@8"}]}]]}],"loading":null,"isPartial":false}
|
||||
4:{}
|
||||
5:"$0:rsc:props:children:0:props:serverProvidedParams:params"
|
||||
8:null
|
||||
@@ -1,4 +0,0 @@
|
||||
1:"$Sreact.fragment"
|
||||
2:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
3:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[null,["$","$L2",null,{"parallelRouterKey":"children","template":["$","$L3",null,{}]}]]}],"loading":null,"isPartial":false}
|
||||
@@ -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",
|
||||
)
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -1,5 +1,6 @@
|
||||
from typing import Union
|
||||
from typing import Any, Awaitable, Callable, Optional, Union
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import (
|
||||
DB_CONNECTION_ERROR_TYPES,
|
||||
ProxyErrorTypes,
|
||||
@@ -123,3 +124,138 @@ class PrismaDBExceptionHandler:
|
||||
):
|
||||
return None
|
||||
raise e
|
||||
|
||||
|
||||
# Default fallback timeouts when neither the caller nor the prisma_client
|
||||
# expose `_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds`.
|
||||
# Match the auth path's existing defaults so behavior is uniform across read paths.
|
||||
_DEFAULT_RECONNECT_TIMEOUT_SECONDS = 2.0
|
||||
_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS = 0.1
|
||||
|
||||
|
||||
def _coerce_timeout(value: Any, fallback: float) -> float:
|
||||
"""Return `value` if it is a real int/float, else `fallback`. Guards
|
||||
against tests that mock `prisma_client` and leave the timeout slots as
|
||||
MagicMock instances."""
|
||||
if isinstance(value, (int, float)) and not isinstance(value, bool):
|
||||
return float(value)
|
||||
return fallback
|
||||
|
||||
|
||||
async def call_with_db_reconnect_retry(
|
||||
prisma_client: Any,
|
||||
coro_factory: Callable[[], Awaitable[Any]],
|
||||
*,
|
||||
reason: str,
|
||||
timeout_seconds: Optional[float] = None,
|
||||
lock_timeout_seconds: Optional[float] = None,
|
||||
) -> Any:
|
||||
"""Run a Prisma read coroutine with one transport-reconnect-and-retry.
|
||||
|
||||
The canonical "self-heal a transient DB transport blip" wrapper used by
|
||||
`PrismaClient.get_generic_data` and other read paths. Mirrors the inline
|
||||
pattern in `auth_checks._fetch_key_object_from_db_with_reconnect` so we
|
||||
have a single implementation rather than three drifting copies.
|
||||
|
||||
Behavior:
|
||||
1. Await `coro_factory()`. On success, return its value.
|
||||
2. On exception, if it is NOT a transport error (per
|
||||
`is_database_transport_error`), re-raise — data-layer errors like
|
||||
`UniqueViolationError` mean the DB is reachable, reconnect would be
|
||||
pointless.
|
||||
3. If `prisma_client` does not expose `attempt_db_reconnect`, re-raise.
|
||||
This guards against partial stand-ins / older clients in tests.
|
||||
4. Call `prisma_client.attempt_db_reconnect(reason=...)`. If it returns
|
||||
False (cooldown / lock contention / reconnect failure), re-raise.
|
||||
5. Otherwise await `coro_factory()` a second time and return / propagate
|
||||
its result. At-most-one retry by construction — no infinite loop.
|
||||
|
||||
`coro_factory` MUST be a zero-arg callable that returns a fresh awaitable
|
||||
on each call. Passing an already-awaited coroutine would fail on retry
|
||||
with `RuntimeError: cannot reuse already awaited coroutine`.
|
||||
|
||||
`reason` should follow `<subsystem>_<operation>_<table>_failure` so
|
||||
telemetry distinguishes between fan-out callers (e.g.
|
||||
`_update_config_from_db` issues four concurrent reads).
|
||||
|
||||
Args:
|
||||
prisma_client: The `PrismaClient` (or stand-in) that owns
|
||||
`attempt_db_reconnect` and the `_db_auth_reconnect_*` defaults.
|
||||
coro_factory: Zero-arg callable returning the read awaitable.
|
||||
reason: Telemetry tag forwarded to `attempt_db_reconnect`.
|
||||
timeout_seconds: Optional override for the reconnect cycle timeout.
|
||||
Defaults to `prisma_client._db_auth_reconnect_timeout_seconds`,
|
||||
then to 2.0s.
|
||||
lock_timeout_seconds: Optional override for how long the helper will
|
||||
wait to acquire the reconnect lock. Defaults to
|
||||
`prisma_client._db_auth_reconnect_lock_timeout_seconds`, then to
|
||||
0.1s.
|
||||
|
||||
Returns:
|
||||
Whatever `coro_factory()` returns (on first or second attempt).
|
||||
|
||||
Raises:
|
||||
Whatever `coro_factory()` raises if the failure is not a transport
|
||||
error, or if the reconnect attempt does not succeed, or if the retry
|
||||
also fails.
|
||||
"""
|
||||
try:
|
||||
return await coro_factory()
|
||||
except Exception as first_exc:
|
||||
if not PrismaDBExceptionHandler.is_database_transport_error(first_exc):
|
||||
raise
|
||||
if not hasattr(prisma_client, "attempt_db_reconnect"):
|
||||
raise
|
||||
|
||||
resolved_timeout = _coerce_timeout(
|
||||
(
|
||||
timeout_seconds
|
||||
if timeout_seconds is not None
|
||||
else getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", None)
|
||||
),
|
||||
_DEFAULT_RECONNECT_TIMEOUT_SECONDS,
|
||||
)
|
||||
resolved_lock_timeout = _coerce_timeout(
|
||||
(
|
||||
lock_timeout_seconds
|
||||
if lock_timeout_seconds is not None
|
||||
else getattr(
|
||||
prisma_client, "_db_auth_reconnect_lock_timeout_seconds", None
|
||||
)
|
||||
),
|
||||
_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.warning(
|
||||
"DB transport error on read; attempting reconnect-and-retry. reason=%s error=%s",
|
||||
reason,
|
||||
first_exc,
|
||||
)
|
||||
|
||||
# Preserve the original transport error in telemetry. If
|
||||
# `attempt_db_reconnect` itself raises (e.g. lock cancellation, timer
|
||||
# error, unexpected internal failure), surfacing that exception
|
||||
# instead of `first_exc` would mask the actual DB transport problem
|
||||
# in `failure_handler` / `db_exceptions` alerts. Chain the reconnect
|
||||
# error as the cause for debuggability without losing the original.
|
||||
try:
|
||||
did_reconnect = await prisma_client.attempt_db_reconnect(
|
||||
reason=reason,
|
||||
timeout_seconds=resolved_timeout,
|
||||
lock_timeout_seconds=resolved_lock_timeout,
|
||||
)
|
||||
except Exception as reconnect_exc:
|
||||
verbose_proxy_logger.warning(
|
||||
"DB reconnect attempt raised; preserving original transport error. "
|
||||
"reason=%s reconnect_error=%s",
|
||||
reason,
|
||||
reconnect_exc,
|
||||
)
|
||||
raise first_exc from reconnect_exc
|
||||
if not did_reconnect:
|
||||
raise
|
||||
|
||||
# At most one retry. If the retry also raises a transport error, we
|
||||
# propagate — repeated reconnect-loops are the watchdog's job, not
|
||||
# this helper's.
|
||||
return await coro_factory()
|
||||
|
||||
@@ -52,18 +52,25 @@ class PrismaWrapper:
|
||||
engine = self._original_prisma._engine
|
||||
process = getattr(engine, "process", None) if engine is not None else None
|
||||
if process is not None:
|
||||
return process.pid
|
||||
pid = process.pid
|
||||
if isinstance(pid, int):
|
||||
return pid
|
||||
except (AttributeError, TypeError):
|
||||
pass
|
||||
return 0
|
||||
|
||||
@staticmethod
|
||||
async def _kill_engine_process(pid: int) -> None:
|
||||
"""Force-kill an orphaned engine subprocess to prevent DB connection pool leaks.
|
||||
"""Force-kill the engine subprocess to prevent DB connection pool leaks.
|
||||
|
||||
Called when disconnect() fails and the old engine process may still be
|
||||
holding open connections. Sends SIGTERM for graceful shutdown, waits
|
||||
briefly, then SIGKILL as a backstop.
|
||||
Called on every reconnect (in `recreate_prisma_client`) to retire the
|
||||
old query-engine subprocess without invoking prisma-client-py's
|
||||
synchronous `disconnect()` — which blocks the asyncio event loop on
|
||||
`subprocess.Popen.wait()` for 30-120+ seconds when the engine is
|
||||
stuck on TCP close.
|
||||
|
||||
Sends SIGTERM for graceful shutdown, waits briefly, then SIGKILL as
|
||||
a backstop.
|
||||
"""
|
||||
if pid <= 0:
|
||||
return
|
||||
@@ -72,7 +79,7 @@ class PrismaWrapper:
|
||||
except (ProcessLookupError, PermissionError, OSError):
|
||||
return # Already dead or inaccessible
|
||||
verbose_proxy_logger.warning(
|
||||
"Sent SIGTERM to orphaned prisma-query-engine PID %s after failed disconnect.",
|
||||
"Sent SIGTERM to prisma-query-engine PID %s during reconnect.",
|
||||
pid,
|
||||
)
|
||||
# Brief wait for graceful shutdown, then force-kill
|
||||
@@ -217,15 +224,18 @@ class PrismaWrapper:
|
||||
async def recreate_prisma_client(
|
||||
self, new_db_url: str, http_client: Optional[Any] = None
|
||||
):
|
||||
"""Disconnect and reconnect the Prisma client with a new database URL."""
|
||||
"""Disconnect and reconnect the Prisma client with a new database URL.
|
||||
|
||||
Kills the old engine subprocess directly (SIGTERM → SIGKILL) rather than
|
||||
calling `disconnect()`. prisma-client-py's `disconnect()` calls a
|
||||
synchronous `subprocess.Popen.wait()` that can freeze the asyncio event
|
||||
loop for 30-120+ seconds when the engine is stuck on TCP close,
|
||||
breaking `/health/liveliness` and causing Kubernetes pod restarts.
|
||||
"""
|
||||
from prisma import Prisma # type: ignore
|
||||
|
||||
old_engine_pid = self._get_engine_pid()
|
||||
|
||||
try:
|
||||
await self._original_prisma.disconnect()
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(f"Failed to disconnect Prisma client: {e}")
|
||||
if old_engine_pid > 0:
|
||||
await self._kill_engine_process(old_engine_pid)
|
||||
|
||||
if http_client is not None:
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
+105
@@ -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
Reference in New Issue
Block a user