mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
Merge pull request #40943 from BerriAI/litellm_internal_staging
Some checks are pending
ai-gateway image / ai-gateway release image (push) Waiting to run
CI Coverage / assert-ci-coverage (push) Waiting to run
CodeQL / Analyze (actions) (push) Waiting to run
CodeQL / Analyze (javascript-typescript) (push) Waiting to run
Code Quality Checks / python-310-import-smoke (push) Waiting to run
Unit Tests: Proxy DB Operations / auth-checks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / budgets (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / custom-logging (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / db-and-spend (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / key-generation (push) Blocked by required conditions
CodeQL / Analyze (python) (push) Waiting to run
CodSpeed Benchmarks / benchmarks (push) Waiting to run
Helm unit test / unit-test (push) Waiting to run
Publish basedpyright base counts / publish (push) Waiting to run
Scorecard supply-chain security / Scorecard analysis (push) Waiting to run
Code Quality Checks / code-quality (push) Waiting to run
UI Unit Tests / ui-unit-tests (push) Waiting to run
Postgres Tests / proxy-security (push) Waiting to run
Postgres Tests / schema-migration (push) Waiting to run
Postgres Tests / proxy-behavior (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
Terraform Provider / gofmt, vet, build, test (push) Waiting to run
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Waiting to run
Unit Tests: Documentation Validation / documentation (push) Waiting to run
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Waiting to run
Unit Tests / integrations (push) Waiting to run
Unit Tests / All Other Providers (push) Waiting to run
Unit Tests / Vertex AI (push) Waiting to run
Unit Tests / misc (push) Waiting to run
Unit Tests / proxy-auth (push) Waiting to run
Unit Tests / proxy-endpoints (push) Waiting to run
Unit Tests / proxy-extras (push) Waiting to run
Unit Tests / proxy-infra (push) Waiting to run
Unit Tests / proxy-server (push) Waiting to run
Unit Tests / responses-caching-types (push) Waiting to run
GitHub Actions Security Analysis / zizmor (push) Waiting to run
Unit Tests: Proxy DB Operations / logging-misc (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-runtime (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-server-core (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-utils (push) Blocked by required conditions
Unit Tests / caching-local (push) Waiting to run
Unit Tests / core-utils (push) Waiting to run
Unit Tests / enterprise-package (push) Waiting to run
Unit Tests / enterprise-routing (push) Waiting to run
Some checks are pending
ai-gateway image / ai-gateway release image (push) Waiting to run
CI Coverage / assert-ci-coverage (push) Waiting to run
CodeQL / Analyze (actions) (push) Waiting to run
CodeQL / Analyze (javascript-typescript) (push) Waiting to run
Code Quality Checks / python-310-import-smoke (push) Waiting to run
Unit Tests: Proxy DB Operations / auth-checks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / budgets (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / custom-logging (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / db-and-spend (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / key-generation (push) Blocked by required conditions
CodeQL / Analyze (python) (push) Waiting to run
CodSpeed Benchmarks / benchmarks (push) Waiting to run
Helm unit test / unit-test (push) Waiting to run
Publish basedpyright base counts / publish (push) Waiting to run
Scorecard supply-chain security / Scorecard analysis (push) Waiting to run
Code Quality Checks / code-quality (push) Waiting to run
UI Unit Tests / ui-unit-tests (push) Waiting to run
Postgres Tests / proxy-security (push) Waiting to run
Postgres Tests / schema-migration (push) Waiting to run
Postgres Tests / proxy-behavior (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
Terraform Provider / gofmt, vet, build, test (push) Waiting to run
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Waiting to run
Unit Tests: Documentation Validation / documentation (push) Waiting to run
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Waiting to run
Unit Tests / integrations (push) Waiting to run
Unit Tests / All Other Providers (push) Waiting to run
Unit Tests / Vertex AI (push) Waiting to run
Unit Tests / misc (push) Waiting to run
Unit Tests / proxy-auth (push) Waiting to run
Unit Tests / proxy-endpoints (push) Waiting to run
Unit Tests / proxy-extras (push) Waiting to run
Unit Tests / proxy-infra (push) Waiting to run
Unit Tests / proxy-server (push) Waiting to run
Unit Tests / responses-caching-types (push) Waiting to run
GitHub Actions Security Analysis / zizmor (push) Waiting to run
Unit Tests: Proxy DB Operations / logging-misc (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-runtime (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-server-core (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-utils (push) Blocked by required conditions
Unit Tests / caching-local (push) Waiting to run
Unit Tests / core-utils (push) Waiting to run
Unit Tests / enterprise-package (push) Waiting to run
Unit Tests / enterprise-routing (push) Waiting to run
chore(ci): remerge internal staging
This commit is contained in:
commit
30f33a949b
67 changed files with 5454 additions and 357 deletions
5
.github/pull_request_template.md
vendored
5
.github/pull_request_template.md
vendored
|
|
@ -47,6 +47,10 @@ After: the same request comes back with real token counts, so the dashboard show
|
|||
|
||||
<!-- e.g., "Fixes #000" -->
|
||||
|
||||
## Affected release
|
||||
|
||||
<!-- Only for a fix to a regression in a released or rc version (perf, memory, crash, or behavior): name the version it regressed in, e.g. "regression in v1.100.0" or "since v1.101.0-rc.1", and add the `backport-stable` label so the fix is cherry-picked onto the rc line before the stable is tagged. Leave the section blank otherwise -->
|
||||
|
||||
## Linear ticket
|
||||
|
||||
<!-- if you are an internal contributor, add "Resolves " followed by the Linear ticket e.g., "Resolves LIT-1234" to link the Linear ticket to the GitHub PR. If you don't have one, leave the section blank rather than guessing -->
|
||||
|
|
@ -154,3 +158,4 @@ Example checklists:
|
|||
## Final Attestation
|
||||
|
||||
- [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
import asyncio
|
||||
import aiohttp
|
||||
import json
|
||||
import math
|
||||
from typing import Any
|
||||
|
||||
# Asynchronously fetch data from a given URL
|
||||
async def fetch_data(url):
|
||||
|
|
@ -21,11 +23,157 @@ async def fetch_data(url):
|
|||
print("Error fetching data from URL:", e)
|
||||
return None
|
||||
|
||||
|
||||
FRIENDLI_API_URL = "https://api.friendli.ai/serverless/v1/models"
|
||||
FRIENDLI_PROVIDER = "friendliai"
|
||||
|
||||
INHERITABLE_BASE_KEYS = (
|
||||
"supports_pdf_input",
|
||||
"supports_assistant_prefill",
|
||||
"supports_adaptive_thinking",
|
||||
"supports_output_config",
|
||||
)
|
||||
|
||||
REASONING_EFFORT_LEVEL_ORDER = ("none", "minimal", "low", "medium", "high", "xhigh", "max")
|
||||
|
||||
|
||||
def _find_base_model_entry(base_model: str, local_data: dict) -> str | None:
|
||||
if not base_model:
|
||||
return None
|
||||
bm_tail = base_model.split("/")[-1].lower()
|
||||
if base_model in local_data:
|
||||
return base_model
|
||||
for key in local_data:
|
||||
if key.startswith("sample_spec") or key == "fallback_generalizations":
|
||||
continue
|
||||
if key.split("/")[-1].lower() == bm_tail:
|
||||
return key
|
||||
return None
|
||||
|
||||
|
||||
def _reasoning_effort_levels(reasoning_options: list) -> list:
|
||||
offered = {
|
||||
val
|
||||
for opt in reasoning_options or []
|
||||
if opt.get("type") == "effort"
|
||||
for val in opt.get("values", [])
|
||||
}
|
||||
return [level for level in REASONING_EFFORT_LEVEL_ORDER if level in offered]
|
||||
|
||||
|
||||
def _valid_token_price(value: object) -> bool:
|
||||
try:
|
||||
price = float(value) # pyright: ignore[reportArgumentType] # non-numeric values are rejected via the except
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
return math.isfinite(price) and price >= 0
|
||||
|
||||
|
||||
def _has_valid_token_prices(pricing: dict | None) -> bool:
|
||||
prices = pricing or {}
|
||||
return _valid_token_price(prices.get("input")) and _valid_token_price(prices.get("output"))
|
||||
|
||||
|
||||
def _pricing(pricing: dict) -> dict:
|
||||
out: dict[str, Any] = {}
|
||||
if not pricing:
|
||||
return out
|
||||
if "input" in pricing:
|
||||
out["input_cost_per_token"] = float(pricing["input"])
|
||||
if "output" in pricing:
|
||||
out["output_cost_per_token"] = float(pricing["output"])
|
||||
if "input_cache_read" in pricing and pricing["input_cache_read"] is not None:
|
||||
out["cache_read_input_token_cost"] = float(pricing["input_cache_read"])
|
||||
return out
|
||||
|
||||
|
||||
def _modality_flags(input_mods: list) -> dict:
|
||||
mods = input_mods or []
|
||||
has_image = "image" in mods
|
||||
return {
|
||||
"supports_vision": has_image,
|
||||
"supports_image_input": has_image,
|
||||
"supports_video_input": "video" in mods,
|
||||
}
|
||||
|
||||
|
||||
def transform_friendli_data(data: list, local_data: dict) -> dict:
|
||||
transformed: dict[str, dict] = {}
|
||||
if not data:
|
||||
return transformed
|
||||
for model in data:
|
||||
# An unpriced row must never wholesale-replace an already priced local entry:
|
||||
# missing prices cost-calculate as zero, silently zeroing tracked spend
|
||||
if not _has_valid_token_prices(model.get("pricing")):
|
||||
continue
|
||||
model_id = model["id"]
|
||||
base_model = model.get("base_model") or ""
|
||||
entry: dict[str, Any] = {
|
||||
"litellm_provider": FRIENDLI_PROVIDER,
|
||||
}
|
||||
|
||||
base_key = _find_base_model_entry(base_model, local_data)
|
||||
if base_key:
|
||||
base_entry = local_data[base_key]
|
||||
for k in INHERITABLE_BASE_KEYS:
|
||||
if k in base_entry:
|
||||
entry[k] = base_entry[k]
|
||||
|
||||
ctx = model.get("context_length")
|
||||
if ctx is not None:
|
||||
entry["max_input_tokens"] = int(ctx)
|
||||
max_out = model.get("max_completion_tokens")
|
||||
if max_out is not None:
|
||||
entry["max_output_tokens"] = int(max_out)
|
||||
entry["max_tokens"] = int(max_out)
|
||||
|
||||
pricing = _pricing(model.get("pricing", {}))
|
||||
entry.update(pricing)
|
||||
entry["supports_prompt_caching"] = "cache_read_input_token_cost" in pricing
|
||||
|
||||
reasoning = model.get("reasoning") is True
|
||||
entry["supports_reasoning"] = reasoning
|
||||
if reasoning:
|
||||
entry["reasoning_effort_levels"] = _reasoning_effort_levels(
|
||||
model.get("reasoning_options", [])
|
||||
)
|
||||
|
||||
func = model.get("functionality", {})
|
||||
entry["supports_function_calling"] = func.get("tool_call") is True
|
||||
entry["supports_parallel_function_calling"] = func.get("parallel_tool_call") is True
|
||||
is_struct = func.get("structured_output") is True
|
||||
entry["supports_response_schema"] = is_struct
|
||||
entry["supports_native_structured_output"] = is_struct
|
||||
entry["supports_system_messages"] = func.get("system_messages") is True
|
||||
entry["supports_tool_choice"] = func.get("tool_choice") is True
|
||||
|
||||
entry.update(_modality_flags(model.get("input_modalities", [])))
|
||||
|
||||
entry["mode"] = model.get("mode", "chat")
|
||||
|
||||
desc = model.get("description")
|
||||
if desc:
|
||||
entry["comment"] = desc
|
||||
|
||||
dep = model.get("deprecation_date")
|
||||
if dep:
|
||||
entry["deprecation_date"] = dep.split("T")[0]
|
||||
|
||||
entry["source"] = FRIENDLI_API_URL
|
||||
|
||||
transformed[f"{FRIENDLI_PROVIDER}/{model_id}"] = entry
|
||||
return transformed
|
||||
|
||||
# Synchronize local data with remote data
|
||||
def sync_local_data_with_remote(local_data, remote_data):
|
||||
def sync_local_data_with_remote(local_data, remote_data, replace_keys=frozenset()):
|
||||
# Update existing keys in local_data with values from remote_data
|
||||
# (replace_keys entries are swapped wholesale so a field the remote catalog
|
||||
# dropped, e.g. cache pricing, cannot survive as a stale value)
|
||||
for key in (set(local_data) & set(remote_data)):
|
||||
local_data[key].update(remote_data[key])
|
||||
if key in replace_keys:
|
||||
local_data[key] = remote_data[key]
|
||||
else:
|
||||
local_data[key].update(remote_data[key])
|
||||
|
||||
# Add new keys from remote_data to local_data
|
||||
for key in (set(remote_data) - set(local_data)):
|
||||
|
|
@ -46,6 +194,8 @@ def write_to_file(file_path, data):
|
|||
# Update the existing models and add the missing models for OpenRouter
|
||||
def transform_openrouter_data(data):
|
||||
transformed = {}
|
||||
if not data:
|
||||
return transformed
|
||||
for row in data:
|
||||
# Add the fields 'max_tokens' and 'input_cost_per_token'
|
||||
obj = {
|
||||
|
|
@ -84,7 +234,14 @@ def transform_openrouter_data(data):
|
|||
# Update the existing models and add the missing models for Vercel AI Gateway
|
||||
def transform_vercel_ai_gateway_data(data):
|
||||
transformed = {}
|
||||
if not data:
|
||||
return transformed
|
||||
for row in data:
|
||||
# Rows without token pricing or token limits (video/embedding models) previously KeyError'd the whole sync
|
||||
if any(row.get(k) is None for k in ("context_window", "max_tokens")) or any(
|
||||
row.get("pricing", {}).get(k) is None for k in ("input", "output")
|
||||
):
|
||||
continue
|
||||
obj = {
|
||||
"max_tokens": row["context_window"],
|
||||
"input_cost_per_token": float(row["pricing"]["input"]),
|
||||
|
|
@ -143,13 +300,16 @@ def main():
|
|||
vercel_data = asyncio.run(fetch_data(vercel_ai_gateway_url))
|
||||
# Transform the fetched Vercel AI Gateway data
|
||||
vercel_data = transform_vercel_ai_gateway_data(vercel_data)
|
||||
|
||||
friendli_data = asyncio.run(fetch_data(FRIENDLI_API_URL))
|
||||
friendli_data = transform_friendli_data(friendli_data, local_data)
|
||||
|
||||
# Combine both datasets
|
||||
all_remote_data = {**openrouter_data, **vercel_data}
|
||||
all_remote_data = {**openrouter_data, **vercel_data, **friendli_data}
|
||||
|
||||
# If both local and openrouter data are available, synchronize and save
|
||||
if local_data and all_remote_data:
|
||||
sync_local_data_with_remote(local_data, all_remote_data)
|
||||
sync_local_data_with_remote(local_data, all_remote_data, replace_keys=frozenset(friendli_data))
|
||||
write_to_file(local_file_path, local_data)
|
||||
else:
|
||||
print("Failed to fetch model data from either local file or URL.")
|
||||
|
|
|
|||
|
|
@ -27,10 +27,13 @@ import litellm
|
|||
from litellm import Router, verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.constants import MAX_FILE_LIST_LIMIT
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
extract_file_metadata,
|
||||
)
|
||||
from openai.types.file_deleted import FileDeleted
|
||||
|
||||
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
|
||||
from litellm.llms.base_llm.managed_resources.isolation import (
|
||||
build_list_page,
|
||||
|
|
@ -48,7 +51,6 @@ from litellm.proxy._types import (
|
|||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
BATCH_CREATE_HIDDEN_PARAM,
|
||||
FILE_LIST_CONTINUATION_CHUNK_SIZE,
|
||||
MAX_FILE_LIST_LIMIT,
|
||||
_is_base64_encoded_unified_file_id,
|
||||
apply_unified_file_ids,
|
||||
decode_model_from_file_id,
|
||||
|
|
@ -1787,7 +1789,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
litellm_parent_otel_span: Optional[Span],
|
||||
llm_router: Router,
|
||||
**data: Dict,
|
||||
) -> OpenAIFileObject:
|
||||
) -> FileDeleted:
|
||||
|
||||
# Check if file deletion should be blocked due to batch references
|
||||
await self._check_file_deletion_allowed(file_id)
|
||||
|
|
@ -1795,7 +1797,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
# file_id = convert_b64_uid_to_unified_uid(file_id)
|
||||
model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span)
|
||||
|
||||
delete_response = None
|
||||
specific_model_file_id_mapping = model_file_id_mapping.get(file_id)
|
||||
if specific_model_file_id_mapping:
|
||||
# Remove conflicting keys from data to avoid duplicate keyword arguments
|
||||
|
|
@ -1810,23 +1811,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
else {}
|
||||
),
|
||||
}
|
||||
delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
|
||||
await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
|
||||
|
||||
stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span)
|
||||
await self.delete_unified_file_id(file_id, litellm_parent_otel_span)
|
||||
|
||||
# Record successful deletion metric only on actual success
|
||||
if stored_file_object or delete_response:
|
||||
prom_logger = self._get_prometheus_logger()
|
||||
if prom_logger:
|
||||
prom_logger.record_managed_file_deleted(result="success")
|
||||
|
||||
if stored_file_object:
|
||||
return OpenAIFileObject.model_validate(stored_file_object).model_copy(update={"id": file_id})
|
||||
elif delete_response:
|
||||
delete_response.id = file_id
|
||||
return delete_response
|
||||
else:
|
||||
raise Exception(f"LiteLLM Managed File object with id={file_id} not found")
|
||||
prom_logger = self._get_prometheus_logger()
|
||||
if prom_logger:
|
||||
prom_logger.record_managed_file_deleted(result="success")
|
||||
return FileDeleted(id=file_id, object="file", deleted=True)
|
||||
|
||||
async def afile_content(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "litellm_call_id" TEXT;
|
||||
|
|
@ -0,0 +1,12 @@
|
|||
-- CreateIndex (CONCURRENTLY)
|
||||
--
|
||||
-- Disclaimer:
|
||||
-- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a
|
||||
-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction.
|
||||
-- - Builds are slower and use more I/O than a blocking CREATE INDEX; if the build is
|
||||
-- interrupted, Postgres may leave an INVALID index that must be dropped and recreated.
|
||||
-- - Do not edit this file after it has been applied to any database: Prisma checksums
|
||||
-- migrations; add a new migration instead.
|
||||
-- - Requires PostgreSQL that supports CONCURRENTLY with IF NOT EXISTS (use a new migration
|
||||
-- without IF NOT EXISTS if you must support older versions).
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id");
|
||||
|
|
@ -659,12 +659,14 @@ model LiteLLM_SpendLogs {
|
|||
mcp_namespaced_tool_name String?
|
||||
agent_id String?
|
||||
proxy_server_request Json? @default("{}")
|
||||
litellm_call_id String?
|
||||
created_at DateTime @default(now()) @map("created_at")
|
||||
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
|
||||
@@index([startTime])
|
||||
@@index([startTime, request_id])
|
||||
@@index([end_user])
|
||||
@@index([session_id])
|
||||
@@index([litellm_call_id])
|
||||
}
|
||||
|
||||
model LiteLLM_BudgetWindowSpend {
|
||||
|
|
|
|||
|
|
@ -53,6 +53,7 @@ S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64
|
|||
S3_PREFIX_DIGEST_CHARS: Final = 16
|
||||
# s3 allows 2048 bytes of combined metadata headers, which Content-Disposition counts against
|
||||
MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
|
||||
MAX_FILE_LIST_LIMIT: Final = 10000
|
||||
DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
|
||||
DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
|
||||
budget_reservation_disabled_info_emitted = False
|
||||
|
|
@ -143,6 +144,7 @@ DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD: Final = float(
|
|||
os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD", 0.3)
|
||||
)
|
||||
MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH: Final = int(os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150))
|
||||
MAX_LITELLM_CALL_ID_LENGTH: Final = 256
|
||||
MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH: Final = 2048
|
||||
|
||||
DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS: Final = 2000
|
||||
|
|
|
|||
|
|
@ -481,7 +481,8 @@ def cost_per_token(
|
|||
else:
|
||||
model_with_provider = f"{custom_llm_provider}/{model}"
|
||||
if region_name is not None:
|
||||
model_with_provider_and_region: Final = f"{custom_llm_provider}/{region_name}/{model}"
|
||||
bare_model: Final = model[len(_prov_prefix) :] if model_is_str and model.startswith(_prov_prefix) else model
|
||||
model_with_provider_and_region: Final = f"{custom_llm_provider}/{region_name}/{bare_model}"
|
||||
if model_with_provider_and_region in model_cost_ref: # use region based pricing, if it's available
|
||||
model_with_provider = model_with_provider_and_region
|
||||
else:
|
||||
|
|
@ -778,6 +779,7 @@ def _select_model_name_for_cost_calc(
|
|||
custom_pricing: bool | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
router_model_id: str | None = None,
|
||||
region_name: str | None = None,
|
||||
) -> str | None:
|
||||
"""
|
||||
1. If custom pricing is true, return received model name
|
||||
|
|
@ -799,8 +801,8 @@ def _select_model_name_for_cost_calc(
|
|||
provider_response_model: Final = _get_hidden_str_for_cost_calc(hidden_params, "provider_response_model")
|
||||
explicit_pricing: Final = custom_pricing is True or base_model is not None
|
||||
priced_from_response: Final = provider_response_model is not None or completion_response_model is not None
|
||||
region_name: Final = (
|
||||
_get_hidden_str_for_cost_calc(hidden_params, "region_name")
|
||||
priced_region: Final = (
|
||||
_get_hidden_str_for_cost_calc(hidden_params, "region_name") or region_name
|
||||
if not explicit_pricing and priced_from_response
|
||||
else None
|
||||
)
|
||||
|
|
@ -837,8 +839,10 @@ def _select_model_name_for_cost_calc(
|
|||
and custom_llm_provider is not None
|
||||
and not _model_contains_known_llm_provider(return_model)
|
||||
): # add provider prefix if not already present, to match model_cost
|
||||
provider_prefix: Final = custom_llm_provider if region_name is None else f"{custom_llm_provider}/{region_name}"
|
||||
return_model = _strip_unregistered_leading_segments(f"{provider_prefix}/{return_model}", region_name)
|
||||
provider_prefix: Final = (
|
||||
custom_llm_provider if priced_region is None else f"{custom_llm_provider}/{priced_region}"
|
||||
)
|
||||
return_model = _strip_unregistered_leading_segments(f"{provider_prefix}/{return_model}", priced_region)
|
||||
|
||||
return return_model
|
||||
|
||||
|
|
@ -1300,6 +1304,7 @@ def completion_cost(
|
|||
|
||||
service_tier = _normalize_service_tier(service_tier)
|
||||
|
||||
explicit_pricing: Final = custom_pricing is True or base_model is not None
|
||||
selected_model: Final = _select_model_name_for_cost_calc(
|
||||
model=model,
|
||||
completion_response=completion_response,
|
||||
|
|
@ -1307,6 +1312,7 @@ def completion_cost(
|
|||
custom_pricing=custom_pricing,
|
||||
base_model=base_model,
|
||||
router_model_id=router_model_id,
|
||||
region_name=region_name,
|
||||
)
|
||||
|
||||
potential_model_names: Final = [
|
||||
|
|
@ -1651,7 +1657,7 @@ def completion_cost(
|
|||
completion_tokens=completion_tokens or 0,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
response_time_ms=total_time,
|
||||
region_name=region_name,
|
||||
region_name=None if explicit_pricing else region_name,
|
||||
custom_cost_per_second=custom_cost_per_second,
|
||||
custom_cost_per_token=custom_cost_per_token,
|
||||
prompt_characters=prompt_characters,
|
||||
|
|
@ -1861,6 +1867,7 @@ def response_cost_calculator(
|
|||
data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us")
|
||||
### VERTEX LOCATION ###
|
||||
vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global")
|
||||
region_name: str | None = None,
|
||||
) -> float:
|
||||
"""
|
||||
Returns
|
||||
|
|
@ -1894,6 +1901,7 @@ def response_cost_calculator(
|
|||
service_tier=service_tier,
|
||||
data_residency=data_residency,
|
||||
vertex_location=vertex_location,
|
||||
region_name=region_name,
|
||||
)
|
||||
return response_cost
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -682,6 +682,10 @@ def file_list(
|
|||
)
|
||||
if provider_config is not None:
|
||||
litellm_params_dict: Final = get_litellm_params(**kwargs)
|
||||
add_trusted_model_credentials_to_litellm_params(
|
||||
litellm_params_dict=litellm_params_dict,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
litellm_params_dict["api_key"] = optional_params.api_key
|
||||
litellm_params_dict["api_base"] = optional_params.api_base
|
||||
|
||||
|
|
|
|||
|
|
@ -601,12 +601,15 @@ def _get_openai_compatible_provider_info(
|
|||
dynamic_api_key,
|
||||
) = litellm.GroqChatConfig()._get_openai_compatible_provider_info(api_base, api_key)
|
||||
elif custom_llm_provider == "bedrock_mantle":
|
||||
from litellm.llms.bedrock_mantle.common_utils import split_mantle_region_prefix
|
||||
|
||||
(
|
||||
api_base,
|
||||
dynamic_api_key,
|
||||
) = litellm.BedrockMantleChatConfig()._get_openai_compatible_provider_info(
|
||||
api_base, api_key, litellm_params=litellm_params, model=model
|
||||
)
|
||||
model = split_mantle_region_prefix(model)[1] # rebind-ok: the prefix is routing only, not a Mantle model id
|
||||
elif custom_llm_provider == "nvidia_nim":
|
||||
# nvidia_nim is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.endpoints.anyscale.com/v1
|
||||
api_base = api_base or get_secret("NVIDIA_NIM_API_BASE") or "https://integrate.api.nvidia.com/v1"
|
||||
|
|
|
|||
|
|
@ -454,6 +454,17 @@ def _resolve_vertex_location_for_cost(
|
|||
return VertexBase.get_vertex_region(configured_location, model)
|
||||
|
||||
|
||||
def _resolve_mantle_region_for_cost(
|
||||
custom_llm_provider: str | None,
|
||||
litellm_params: Mapping[str, object] | None,
|
||||
) -> str | None:
|
||||
if custom_llm_provider != "bedrock_mantle":
|
||||
return None
|
||||
from litellm.llms.bedrock_mantle.common_utils import resolve_mantle_region
|
||||
|
||||
return resolve_mantle_region(litellm_params or MappingProxyType({}))
|
||||
|
||||
|
||||
def _provider_response_id(source: object) -> str | None:
|
||||
candidate: Final = source.get("id") if isinstance(source, dict) else getattr(source, "id", None)
|
||||
return candidate if isinstance(candidate, str) and candidate else None
|
||||
|
|
@ -1768,6 +1779,10 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
optional_params=self.optional_params,
|
||||
model=litellm_model_name or self.model,
|
||||
),
|
||||
"region_name": _resolve_mantle_region_for_cost(
|
||||
custom_llm_provider=self.model_call_details.get("custom_llm_provider", None),
|
||||
litellm_params=self.model_call_details.get("litellm_params"),
|
||||
),
|
||||
}
|
||||
except Exception as e: # error creating kwargs for cost calculation
|
||||
debug_info = StandardLoggingModelCostFailureDebugInformation(
|
||||
|
|
|
|||
|
|
@ -19,6 +19,8 @@ from typing_extensions import assert_never
|
|||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.types.llms.anthropic import (
|
||||
AppliedEdit,
|
||||
CompactionBlock,
|
||||
|
|
@ -58,6 +60,25 @@ def _optional_attr_sequence(obj: object, name: str) -> Sequence[object]:
|
|||
return value if value else ()
|
||||
|
||||
|
||||
def _error_status_and_message(exc: Exception) -> tuple[int, str]:
|
||||
if isinstance(exc, (BaseLLMException, MidStreamFallbackError)):
|
||||
return exc.status_code, exc.message
|
||||
return 500, str(exc) or "Upstream stream ended before completion"
|
||||
|
||||
|
||||
def _mid_stream_error_sse_event(exc: Exception) -> bytes:
|
||||
from litellm.anthropic_interface.exceptions.exception_mapping_utils import (
|
||||
AnthropicExceptionMapping,
|
||||
)
|
||||
|
||||
status_code, message = _error_status_and_message(exc)
|
||||
error_response = AnthropicExceptionMapping.transform_to_anthropic_error(
|
||||
status_code=status_code,
|
||||
raw_message=message,
|
||||
)
|
||||
return f"event: error\ndata: {json.dumps(error_response)}\n\n".encode()
|
||||
|
||||
|
||||
def _delta_payload_field(delta_type: StreamingContentBlockDeltaType) -> str:
|
||||
match delta_type:
|
||||
case "text_delta":
|
||||
|
|
@ -990,14 +1011,17 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
Async version of anthropic_sse_wrapper.
|
||||
Convert AnthropicStreamWrapper dict chunks to Server-Sent Events format.
|
||||
"""
|
||||
async for chunk in self:
|
||||
if isinstance(chunk, dict):
|
||||
event_type: str = str(chunk.get("type", "message"))
|
||||
payload = f"event: {event_type}\ndata: {json.dumps(chunk)}\n\n"
|
||||
yield payload.encode()
|
||||
else:
|
||||
# For non-dict chunks, forward the original value unchanged
|
||||
yield chunk
|
||||
try:
|
||||
async for chunk in self:
|
||||
if isinstance(chunk, dict):
|
||||
event_type: str = str(chunk.get("type", "message"))
|
||||
payload = f"event: {event_type}\ndata: {json.dumps(chunk)}\n\n"
|
||||
yield payload.encode()
|
||||
else:
|
||||
yield chunk
|
||||
except Exception as e: # noqa: BLE001 # boundary before the socket: any upstream failure becomes an Anthropic error event
|
||||
verbose_logger.exception("Anthropic Adapter - mid-stream error, emitting Anthropic error event: %s", e)
|
||||
yield _mid_stream_error_sse_event(e)
|
||||
|
||||
def _increment_content_block_index(self):
|
||||
self.current_content_block_index += 1
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Iterator
|
||||
from collections.abc import Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Union
|
||||
|
||||
import httpx
|
||||
|
|
@ -160,6 +160,15 @@ class BaseFilesConfig(BaseConfig):
|
|||
) -> tuple[str, dict]:
|
||||
"""Transform file list request into provider-specific format."""
|
||||
|
||||
def transform_list_files_next_request(
|
||||
self,
|
||||
raw_response: httpx.Response,
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: dict, # mutable-ok: carries provider stashes from the request transform to the response one
|
||||
) -> tuple[str, dict[str, str]] | None:
|
||||
"""Request for the page after `raw_response`, or None once the listing is complete."""
|
||||
return None
|
||||
|
||||
@abstractmethod
|
||||
def transform_list_files_response(
|
||||
self,
|
||||
|
|
@ -258,7 +267,7 @@ class BaseFileEndpoints(ABC):
|
|||
litellm_parent_otel_span: Span | None,
|
||||
llm_router: Router,
|
||||
**data: dict,
|
||||
) -> OpenAIFileObject:
|
||||
) -> FileDeleted:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
|
|
|
|||
|
|
@ -1,14 +1,18 @@
|
|||
import base64
|
||||
import json
|
||||
import os
|
||||
import posixpath
|
||||
import time
|
||||
import xml.etree.ElementTree as ET
|
||||
from collections.abc import Iterable, Mapping, MutableMapping, Sequence
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from functools import cache
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Literal, TypeAlias, TypedDict
|
||||
from urllib.parse import unquote
|
||||
from urllib.parse import quote, unquote, urlencode
|
||||
|
||||
import httpx
|
||||
from httpx import Headers, Response
|
||||
|
|
@ -23,6 +27,7 @@ from litellm.files.utils import FilesAPIUtils
|
|||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.litellm_core_utils.cloud_storage_security import (
|
||||
BEDROCK_MANAGED_S3_BATCH_PREFIX,
|
||||
BEDROCK_MANAGED_S3_OUTPUT_PREFIX,
|
||||
BEDROCK_MANAGED_S3_PREFIXES,
|
||||
BEDROCK_MANAGED_S3_UPLOAD_PREFIX,
|
||||
build_managed_cloud_object_name,
|
||||
|
|
@ -62,6 +67,10 @@ from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resol
|
|||
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers"
|
||||
|
||||
LIST_FILES_PURPOSE_PARAM: Final = "_s3_list_files_purpose"
|
||||
|
||||
LIST_FILES_LOCATION_PARAM: Final = "_s3_list_files_location"
|
||||
|
||||
|
||||
class _S3DeleteContext(BaseModel):
|
||||
file_id: str = Field(min_length=1)
|
||||
|
|
@ -152,6 +161,13 @@ class _BedrockS3RequestParams(BaseModel):
|
|||
s3_endpoint_url: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _S3RequestTarget:
|
||||
endpoint_url: str
|
||||
aws_region_name: str
|
||||
request_params: _BedrockS3RequestParams
|
||||
|
||||
|
||||
class _TrustedS3ModelCredentials(BaseModel):
|
||||
"""The S3 buckets the server trusts file ids against, from the deployment snapshot."""
|
||||
|
||||
|
|
@ -248,6 +264,128 @@ def _validate_file_id_against_configured_buckets(
|
|||
return validate_against(configured_bucket_names[-1])
|
||||
|
||||
|
||||
_REJECTED_FILE_ID_REQUEST_URL: Final = "https://litellm.ai"
|
||||
|
||||
|
||||
def _rejected_file_id(reason: ValueError) -> BedrockError:
|
||||
message: Final = str(reason)
|
||||
return BedrockError(
|
||||
status_code=400,
|
||||
message=message,
|
||||
response=httpx.Response(
|
||||
status_code=400,
|
||||
text=message,
|
||||
request=httpx.Request(method="GET", url=_REJECTED_FILE_ID_REQUEST_URL),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _resolve_managed_s3_object(file_id: str, litellm_params: Mapping[str, object]) -> tuple[str, str]:
|
||||
configured_bucket_names: Final = get_configured_s3_bucket_names(litellm_params)
|
||||
allow_legacy_cloud_file_ids: Final = should_allow_legacy_cloud_file_ids(litellm_params)
|
||||
try:
|
||||
return _validate_file_id_against_configured_buckets(
|
||||
s3_uri=extract_s3_uri_from_file_id(file_id),
|
||||
configured_bucket_names=configured_bucket_names,
|
||||
allow_legacy_cloud_file_ids=allow_legacy_cloud_file_ids,
|
||||
)
|
||||
except ValueError as reason:
|
||||
raise _rejected_file_id(reason) from reason
|
||||
|
||||
|
||||
_ANY_MANAGED_LISTING_PREFIX: Final = os.path.commonprefix(BEDROCK_MANAGED_S3_PREFIXES)
|
||||
_MANAGED_LISTING_PREFIX_BY_PURPOSE: Final = MappingProxyType(
|
||||
{
|
||||
"batch": os.path.commonprefix((BEDROCK_MANAGED_S3_BATCH_PREFIX, BEDROCK_MANAGED_S3_UPLOAD_PREFIX)),
|
||||
"batch_output": BEDROCK_MANAGED_S3_OUTPUT_PREFIX,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
_EMPTY_LISTING_QUERY: Final = (("list-type", "2"), ("max-keys", "0"))
|
||||
|
||||
|
||||
def _managed_listing_prefix(configured_prefix: str, purpose: str | None) -> str | None:
|
||||
managed_prefix: Final = _MANAGED_LISTING_PREFIX_BY_PURPOSE.get(purpose) if purpose else _ANY_MANAGED_LISTING_PREFIX
|
||||
if managed_prefix is None:
|
||||
return None
|
||||
return f"{configured_prefix}/{managed_prefix}" if configured_prefix else managed_prefix
|
||||
|
||||
|
||||
def _listing_query(configured_prefix: str, purpose: str | None) -> tuple[tuple[str, str], ...]:
|
||||
listing_prefix: Final = _managed_listing_prefix(configured_prefix, purpose)
|
||||
if listing_prefix is None:
|
||||
return _EMPTY_LISTING_QUERY
|
||||
return (("list-type", "2"), ("prefix", listing_prefix))
|
||||
|
||||
|
||||
def _requested_listing_purpose(litellm_params: Mapping[str, object]) -> str | None:
|
||||
requested_purpose: Final = litellm_params.get(LIST_FILES_PURPOSE_PARAM)
|
||||
return requested_purpose if isinstance(requested_purpose, str) else None
|
||||
|
||||
|
||||
def _walked_listing_purpose(litellm_params: Mapping[str, object]) -> str | None:
|
||||
walked_purpose: Final = litellm_params.get(LIST_FILES_LOCATION_PARAM)
|
||||
return walked_purpose if isinstance(walked_purpose, str) else _requested_listing_purpose(litellm_params)
|
||||
|
||||
|
||||
def _output_location_still_unlisted(litellm_params: Mapping[str, object]) -> bool:
|
||||
if _walked_listing_purpose(litellm_params) is not None:
|
||||
return False
|
||||
return _listing_bucket_name(litellm_params, "batch_output") != _listing_bucket_name(litellm_params, None)
|
||||
|
||||
|
||||
def _listing_bucket_name(litellm_params: Mapping[str, object], purpose: str | None) -> str:
|
||||
if purpose != "batch_output":
|
||||
return get_configured_s3_bucket_name(litellm_params)
|
||||
trusted: Final = _trusted_s3_model_credentials(litellm_params)
|
||||
return (
|
||||
trusted.s3_output_bucket_name
|
||||
or os.getenv("AWS_S3_OUTPUT_BUCKET_NAME")
|
||||
or get_configured_s3_bucket_name(litellm_params)
|
||||
)
|
||||
|
||||
|
||||
def _listed_object_created_at(entry: ET.Element) -> int:
|
||||
last_modified: Final = entry.findtext("{*}LastModified")
|
||||
if not last_modified:
|
||||
return 0
|
||||
return int(datetime.fromisoformat(last_modified.replace("Z", "+00:00")).timestamp())
|
||||
|
||||
|
||||
def _listed_managed_file(
|
||||
entry: ET.Element,
|
||||
bucket_name: str,
|
||||
configured_bucket_name: str,
|
||||
allow_legacy_cloud_file_ids: bool,
|
||||
) -> OpenAIFileObject | None:
|
||||
object_key: Final = entry.findtext("{*}Key")
|
||||
if not object_key:
|
||||
return None
|
||||
file_id: Final = f"s3://{bucket_name}/{object_key}"
|
||||
try:
|
||||
validate_managed_cloud_file_id(
|
||||
file_id=file_id,
|
||||
scheme="s3://",
|
||||
configured_bucket_name=configured_bucket_name,
|
||||
allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES,
|
||||
allow_legacy_cloud_file_ids=allow_legacy_cloud_file_ids,
|
||||
)
|
||||
except ValueError:
|
||||
return None
|
||||
_, configured_prefix = split_configured_cloud_bucket_name(configured_bucket_name)
|
||||
relative_key: Final = object_key[len(configured_prefix) + 1 :] if configured_prefix else object_key
|
||||
return OpenAIFileObject(
|
||||
id=file_id,
|
||||
bytes=int(entry.findtext("{*}Size") or 0),
|
||||
created_at=_listed_object_created_at(entry),
|
||||
filename=posixpath.basename(object_key),
|
||||
object="file",
|
||||
purpose="batch_output" if relative_key.startswith(BEDROCK_MANAGED_S3_OUTPUT_PREFIX) else "batch",
|
||||
status="uploaded",
|
||||
)
|
||||
|
||||
|
||||
def _uploaded_object_size(litellm_params: Mapping[str, object], raw_response: Response) -> int:
|
||||
"""
|
||||
S3 answers PutObject with an empty body, so the stored object size comes from the
|
||||
|
|
@ -1213,18 +1351,86 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
def transform_list_files_request(
|
||||
self,
|
||||
purpose: str | None,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> tuple[str, dict]:
|
||||
raise NotImplementedError("BedrockFilesConfig does not support file listing")
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: MutableMapping[str, object],
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
litellm_params[LIST_FILES_PURPOSE_PARAM] = purpose # rebind-ok: handed to the response transform
|
||||
litellm_params[LIST_FILES_LOCATION_PARAM] = purpose # rebind-ok: names the location the next page walks
|
||||
return self._signed_listing_request(purpose, optional_params, litellm_params, continuation_token=None)
|
||||
|
||||
def transform_list_files_next_request(
|
||||
self,
|
||||
raw_response: httpx.Response,
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: MutableMapping[str, object],
|
||||
) -> tuple[str, dict[str, str]] | None:
|
||||
if raw_response.status_code >= 400:
|
||||
return None
|
||||
continuation_token: Final = ET.fromstring(raw_response.content).findtext("{*}NextContinuationToken")
|
||||
if continuation_token:
|
||||
return self._signed_listing_request(
|
||||
_walked_listing_purpose(litellm_params), optional_params, litellm_params, continuation_token
|
||||
)
|
||||
if not _output_location_still_unlisted(litellm_params):
|
||||
return None
|
||||
litellm_params[LIST_FILES_LOCATION_PARAM] = "batch_output" # rebind-ok: the input location is fully listed
|
||||
return self._signed_listing_request("batch_output", optional_params, litellm_params, continuation_token=None)
|
||||
|
||||
def _signed_listing_request(
|
||||
self,
|
||||
purpose: str | None,
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: MutableMapping[str, object],
|
||||
continuation_token: str | None,
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
bucket_name, configured_prefix = split_configured_cloud_bucket_name(
|
||||
_listing_bucket_name(litellm_params, purpose)
|
||||
)
|
||||
target: Final = self._s3_request_target(optional_params=optional_params, litellm_params=litellm_params)
|
||||
url: Final = f"{target.endpoint_url}/{bucket_name}/"
|
||||
listing_query: Final = _listing_query(configured_prefix, purpose)
|
||||
continuation_query: Final = (("continuation-token", continuation_token),) if continuation_token else ()
|
||||
query: Final[dict[str, str]] = dict( # mutable-ok: the base files contract returns the query as a dict
|
||||
listing_query + continuation_query
|
||||
)
|
||||
signed_headers: Final = self._sign_s3_request_without_body(
|
||||
method="GET",
|
||||
api_base=f"{url}?{urlencode(query, quote_via=quote, safe='')}",
|
||||
aws_region_name=target.aws_region_name,
|
||||
request_params=target.request_params,
|
||||
)
|
||||
litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = signed_headers # rebind-ok: handed to validate_environment
|
||||
return url, query
|
||||
|
||||
def transform_list_files_response(
|
||||
self,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> list[OpenAIFileObject]:
|
||||
raise NotImplementedError("BedrockFilesConfig does not support file listing")
|
||||
if raw_response.status_code >= 400:
|
||||
raise BedrockError(
|
||||
status_code=raw_response.status_code,
|
||||
message=raw_response.text,
|
||||
headers=raw_response.headers,
|
||||
response=raw_response,
|
||||
)
|
||||
purpose: Final = _requested_listing_purpose(litellm_params)
|
||||
configured_bucket_name: Final = _listing_bucket_name(litellm_params, _walked_listing_purpose(litellm_params))
|
||||
allow_legacy_cloud_file_ids: Final = should_allow_legacy_cloud_file_ids(litellm_params)
|
||||
listing: Final = ET.fromstring(raw_response.content)
|
||||
bucket_name: Final = (
|
||||
listing.findtext("{*}Name") or split_configured_cloud_bucket_name(configured_bucket_name)[0]
|
||||
)
|
||||
listed_files: Final = (
|
||||
_listed_managed_file(entry, bucket_name, configured_bucket_name, allow_legacy_cloud_file_ids)
|
||||
for entry in listing.iterfind("{*}Contents")
|
||||
)
|
||||
return [ # mutable-ok: the base files contract returns a list
|
||||
listed_file
|
||||
for listed_file in listed_files
|
||||
if listed_file is not None and (purpose is None or listed_file.purpose == purpose)
|
||||
]
|
||||
|
||||
def transform_file_content_request(
|
||||
self,
|
||||
|
|
@ -1255,39 +1461,54 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
optional_params: Mapping[str, object],
|
||||
litellm_params: MutableMapping[str, object],
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
s3_uri: Final = extract_s3_uri_from_file_id(file_id)
|
||||
bucket_name, object_key = _validate_file_id_against_configured_buckets(
|
||||
s3_uri=s3_uri,
|
||||
configured_bucket_names=get_configured_s3_bucket_names(litellm_params),
|
||||
allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params),
|
||||
bucket_name, object_key = _resolve_managed_s3_object(file_id=file_id, litellm_params=litellm_params)
|
||||
target: Final = self._s3_request_target(optional_params=optional_params, litellm_params=litellm_params)
|
||||
url: Final = f"{target.endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}"
|
||||
signed_headers: Final = self._sign_s3_request_without_body(
|
||||
method=method,
|
||||
api_base=url,
|
||||
aws_region_name=target.aws_region_name,
|
||||
request_params=target.request_params,
|
||||
)
|
||||
litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = signed_headers # rebind-ok: handed to validate_environment
|
||||
return url, {} # mutable-ok: the base files contract returns the query as a dict
|
||||
|
||||
request_params: Final = _BedrockS3RequestParams.model_validate({**litellm_params, **optional_params})
|
||||
|
||||
def _s3_request_target(
|
||||
self,
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> _S3RequestTarget:
|
||||
"""
|
||||
The shared files handler passes optional_params={}, so AWS credentials and
|
||||
region arrive via litellm_params here (unlike the upload path).
|
||||
s3_region_name wins over aws_region_name, same priority as get_complete_file_url.
|
||||
"""
|
||||
request_params: Final = _BedrockS3RequestParams.model_validate(
|
||||
MappingProxyType({**litellm_params, **optional_params})
|
||||
)
|
||||
region_preference: Final = request_params.s3_region_name or request_params.aws_region_name
|
||||
region_params: Final[dict[str, str | None]] = {"aws_region_name": region_preference}
|
||||
aws_region_name: Final = self._get_aws_region_name(optional_params=region_params, model="")
|
||||
|
||||
s3_endpoint_url: Final = (
|
||||
aws_region_name: Final = self._get_aws_region_name(
|
||||
optional_params={"aws_region_name": region_preference}, # mutable-ok: BaseAWSLLM takes a dict
|
||||
model="",
|
||||
)
|
||||
endpoint_url: Final = (
|
||||
request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}"
|
||||
).rstrip("/")
|
||||
url: Final = f"{s3_endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}"
|
||||
|
||||
litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = self._sign_s3_request_without_body(
|
||||
api_base=url,
|
||||
aws_region_name=aws_region_name,
|
||||
request_params=request_params,
|
||||
method=method,
|
||||
return _S3RequestTarget(
|
||||
endpoint_url=endpoint_url, aws_region_name=aws_region_name, request_params=request_params
|
||||
)
|
||||
return url, {}
|
||||
|
||||
def _sign_s3_request_without_body(
|
||||
self,
|
||||
method: Literal["GET", "DELETE"],
|
||||
api_base: str,
|
||||
aws_region_name: str,
|
||||
request_params: _BedrockS3RequestParams,
|
||||
method: Literal["GET", "DELETE"] = "GET",
|
||||
) -> dict[str, str]:
|
||||
) -> Mapping[str, str]:
|
||||
"""
|
||||
SigV4-sign a bodiless S3 request (GetObject, DeleteObject, ListObjectsV2),
|
||||
mirroring `_sign_s3_request` (PUT).
|
||||
"""
|
||||
try:
|
||||
import hashlib
|
||||
|
||||
|
|
@ -1313,11 +1534,11 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
aws_request: Final = AWSRequest( # any-ok: botocore AWSRequest is untyped
|
||||
method=method,
|
||||
url=api_base,
|
||||
headers={"x-amz-content-sha256": empty_body_hash},
|
||||
headers={"x-amz-content-sha256": empty_body_hash}, # mutable-ok: botocore AWSRequest takes a dict
|
||||
)
|
||||
auth: Final = S3SigV4Auth(credentials, "s3", aws_region_name) # any-ok: botocore untyped
|
||||
auth.add_auth(aws_request) # any-ok: botocore request mutation is untyped
|
||||
return dict(aws_request.headers) # any-ok: botocore headers are untyped
|
||||
return MappingProxyType(dict(aws_request.headers)) # any-ok: botocore headers are untyped
|
||||
|
||||
def transform_file_content_response(
|
||||
self,
|
||||
|
|
@ -1330,6 +1551,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
status_code=raw_response.status_code,
|
||||
message=raw_response.text,
|
||||
headers=raw_response.headers,
|
||||
response=raw_response,
|
||||
)
|
||||
return HttpxBinaryResponseContent(response=raw_response)
|
||||
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ from litellm.types.router import GenericLiteLLMParams
|
|||
from ...base_llm.chat.transformation import BaseLLMException
|
||||
from ...bedrock.common_utils import BedrockError
|
||||
from ...openai_like.chat.transformation import OpenAILikeChatConfig
|
||||
from ..common_utils import mantle_base_segment
|
||||
from ..common_utils import mantle_base_segment, split_mantle_region_prefix
|
||||
|
||||
|
||||
class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
|
||||
|
|
@ -61,8 +61,10 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
|
|||
litellm_params: GenericLiteLLMParams | None = None,
|
||||
model: str | None = None,
|
||||
) -> tuple[str | None, str | None]:
|
||||
prefix_region, base_model = split_mantle_region_prefix(model) if model else (None, None)
|
||||
region: Final = (
|
||||
(litellm_params.aws_region_name if litellm_params else None)
|
||||
or prefix_region
|
||||
or get_secret_str("BEDROCK_MANTLE_REGION")
|
||||
or get_secret_str("AWS_REGION_NAME")
|
||||
or get_secret_str("AWS_REGION")
|
||||
|
|
@ -75,7 +77,7 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
|
|||
api_base = (
|
||||
api_base
|
||||
or get_secret_str("BEDROCK_MANTLE_API_BASE")
|
||||
or f"https://bedrock-mantle.{region}.api.aws/{mantle_base_segment(model, litellm.model_cost)}"
|
||||
or f"https://bedrock-mantle.{region}.api.aws/{mantle_base_segment(base_model, litellm.model_cost)}"
|
||||
)
|
||||
dynamic_api_key: Final = self._resolve_bearer_token(api_key)
|
||||
return api_base, dynamic_api_key
|
||||
|
|
|
|||
|
|
@ -24,9 +24,11 @@ from botocore.exceptions import (
|
|||
)
|
||||
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS
|
||||
from litellm.llms.bedrock.common_utils import AmazonBedrockGlobalConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1"
|
||||
BEDROCK_REGIONS: Final = frozenset(AmazonBedrockGlobalConfig().get_all_regions())
|
||||
|
||||
# Standard Mantle host: https://bedrock-mantle.<region>.api.aws (group 1 = region).
|
||||
MANTLE_HOST_RE: Final = re.compile(r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws(?=/|$)", re.IGNORECASE)
|
||||
|
|
@ -36,6 +38,13 @@ def resolve_mantle_bearer_token(api_key: str | None) -> str | None:
|
|||
return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
|
||||
|
||||
|
||||
def split_mantle_region_prefix(model: str) -> tuple[str | None, str]:
|
||||
head, sep, tail = model.partition("/")
|
||||
if sep and head in BEDROCK_REGIONS:
|
||||
return head, tail
|
||||
return None, model
|
||||
|
||||
|
||||
def resolve_mantle_region(params: Mapping[str, object]) -> str:
|
||||
region: Final = params.get("aws_region_name")
|
||||
if isinstance(region, str) and region:
|
||||
|
|
@ -130,7 +139,7 @@ def mantle_supports_responses(model: str | None, model_cost: dict) -> bool:
|
|||
gpt-oss substring), so a substring gate would be wrong. A model absent from
|
||||
model_cost simply has no signal and returns False (chat-completions emulation).
|
||||
"""
|
||||
entry: Final = model_cost.get(f"bedrock_mantle/{model}", {})
|
||||
entry: Final = model_cost.get(f"bedrock_mantle/{split_mantle_region_prefix(model)[1]}", {}) if model else {}
|
||||
if "/v1/responses" in (entry.get("supported_endpoints") or []):
|
||||
return True
|
||||
return entry.get("mode") == "responses"
|
||||
|
|
@ -147,5 +156,5 @@ def mantle_base_segment(model: str | None, model_cost: dict) -> str:
|
|||
the base for the model's whole OpenAI-compatible surface, so both the chat and
|
||||
responses configs derive from it -- there is no separate model-name rule.
|
||||
"""
|
||||
entry: Final = model_cost.get(f"bedrock_mantle/{model}", {})
|
||||
entry: Final = model_cost.get(f"bedrock_mantle/{split_mantle_region_prefix(model)[1]}", {}) if model else {}
|
||||
return "openai/v1" if entry.get("use_openai_responses_path") is True else "v1"
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ import litellm.types
|
|||
import litellm.types.utils
|
||||
from litellm._logging import _redact_string, verbose_logger
|
||||
from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta
|
||||
from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES
|
||||
from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES
|
||||
from litellm.litellm_core_utils.agentic_loop_settings import (
|
||||
DEFAULT_MAX_AGENTIC_LOOPS,
|
||||
validated_max_agentic_loops,
|
||||
|
|
@ -4981,15 +4981,16 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
try:
|
||||
response: Final = sync_httpx_client.get(url=url, headers=headers, params=params)
|
||||
response: Final = sync_httpx_client.get(url=url, headers=headers, params=params, timeout=timeout)
|
||||
except Exception as e:
|
||||
raise self._handle_error(e=e, provider_config=provider_config)
|
||||
|
||||
return provider_config.transform_list_files_response(
|
||||
raw_response=response,
|
||||
logging_obj=logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
files_per_page: Final = self._files_per_listing_page(
|
||||
response, provider_config, logging_obj, litellm_params, headers, sync_httpx_client, timeout
|
||||
)
|
||||
return [ # mutable-ok: the files contract returns the listing as a list
|
||||
listed_file for page_files in files_per_page for listed_file in page_files
|
||||
]
|
||||
|
||||
async def async_list_files(
|
||||
self,
|
||||
|
|
@ -5037,16 +5038,101 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
try:
|
||||
response: Final = await async_httpx_client.get(url=url, headers=headers, params=params)
|
||||
response: Final = await async_httpx_client.get(url=url, headers=headers, params=params, timeout=timeout)
|
||||
except Exception as e:
|
||||
raise self._handle_error(e=e, provider_config=provider_config)
|
||||
|
||||
return provider_config.transform_list_files_response(
|
||||
raw_response=response,
|
||||
logging_obj=logging_obj,
|
||||
files_per_page: Final = self._files_per_async_listing_page(
|
||||
response, provider_config, logging_obj, litellm_params, headers, async_httpx_client, timeout
|
||||
)
|
||||
return [ # mutable-ok: the files contract returns the listing as a list
|
||||
listed_file async for page_files in files_per_page for listed_file in page_files
|
||||
]
|
||||
|
||||
def _files_per_listing_page(
|
||||
self,
|
||||
first_page: httpx.Response,
|
||||
provider_config: BaseFilesConfig,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict, # mutable-ok: handed to the files contract, which types it as a dict
|
||||
headers: dict, # mutable-ok: handed to validate_environment, which types it as a dict
|
||||
client: HTTPHandler,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
) -> Iterator[list[OpenAIFileObject]]: # mutable-ok: each page arrives as the list the files contract returns
|
||||
latest_page = first_page # rebind-ok: advances one page per loop turn
|
||||
listed_count = 0 # rebind-ok: grows per page so the listing stops at MAX_FILE_LIST_LIMIT, OpenAI's ceiling
|
||||
while True:
|
||||
page_files = provider_config.transform_list_files_response(
|
||||
raw_response=latest_page, logging_obj=logging_obj, litellm_params=litellm_params
|
||||
)
|
||||
yield page_files[: MAX_FILE_LIST_LIMIT - listed_count]
|
||||
listed_count += len(page_files)
|
||||
next_request = self._next_listing_request(latest_page, provider_config, litellm_params, listed_count)
|
||||
if next_request is None:
|
||||
return
|
||||
url, params = next_request
|
||||
next_headers = self._next_listing_page_headers(provider_config, headers, litellm_params)
|
||||
try:
|
||||
latest_page = client.get(url=url, headers=next_headers, params=params, timeout=timeout)
|
||||
except Exception as e: # noqa: BLE001 # _handle_error maps every failure kind, like the first page's fetch
|
||||
raise self._handle_error(e=e, provider_config=provider_config)
|
||||
|
||||
async def _files_per_async_listing_page(
|
||||
self,
|
||||
first_page: httpx.Response,
|
||||
provider_config: BaseFilesConfig,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict, # mutable-ok: handed to the files contract, which types it as a dict
|
||||
headers: dict, # mutable-ok: handed to validate_environment, which types it as a dict
|
||||
client: AsyncHTTPHandler,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
) -> AsyncIterator[list[OpenAIFileObject]]: # mutable-ok: each page arrives as the list the files contract returns
|
||||
latest_page = first_page # rebind-ok: advances one page per loop turn
|
||||
listed_count = 0 # rebind-ok: grows per page so the listing stops at MAX_FILE_LIST_LIMIT, OpenAI's ceiling
|
||||
while True:
|
||||
page_files = provider_config.transform_list_files_response(
|
||||
raw_response=latest_page, logging_obj=logging_obj, litellm_params=litellm_params
|
||||
)
|
||||
yield page_files[: MAX_FILE_LIST_LIMIT - listed_count]
|
||||
listed_count += len(page_files)
|
||||
next_request = self._next_listing_request(latest_page, provider_config, litellm_params, listed_count)
|
||||
if next_request is None:
|
||||
return
|
||||
url, params = next_request
|
||||
next_headers = self._next_listing_page_headers(provider_config, headers, litellm_params)
|
||||
try:
|
||||
latest_page = await client.get(url=url, headers=next_headers, params=params, timeout=timeout)
|
||||
except Exception as e: # noqa: BLE001 # _handle_error maps every failure kind, like the first page's fetch
|
||||
raise self._handle_error(e=e, provider_config=provider_config)
|
||||
|
||||
def _next_listing_page_headers(
|
||||
self,
|
||||
provider_config: BaseFilesConfig,
|
||||
headers: dict, # mutable-ok: handed to validate_environment, which types it as a dict
|
||||
litellm_params: dict, # mutable-ok: handed to the files contract, which types it as a dict
|
||||
) -> dict: # mutable-ok: validate_environment returns the header dict the files contract types
|
||||
return provider_config.validate_environment(
|
||||
api_key=litellm_params.get("api_key"),
|
||||
headers=headers,
|
||||
model="",
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
def _next_listing_request(
|
||||
self,
|
||||
latest_page: httpx.Response,
|
||||
provider_config: BaseFilesConfig,
|
||||
litellm_params: dict, # mutable-ok: handed to the files contract, which types it as a dict
|
||||
listed_count: int,
|
||||
) -> tuple[str, dict[str, str]] | None: # mutable-ok: the base files contract returns the query as a dict
|
||||
if listed_count >= MAX_FILE_LIST_LIMIT:
|
||||
return None
|
||||
return provider_config.transform_list_files_next_request(
|
||||
raw_response=latest_page, optional_params={}, litellm_params=litellm_params
|
||||
)
|
||||
|
||||
def retrieve_file_content(
|
||||
self,
|
||||
file_content_request: "FileContentRequest",
|
||||
|
|
|
|||
|
|
@ -23083,58 +23083,206 @@
|
|||
},
|
||||
"friendliai/zai-org/GLM-5.3-Flash": {
|
||||
"litellm_provider": "friendliai",
|
||||
"supports_reasoning": true,
|
||||
"supports_function_calling": true,
|
||||
"max_input_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"max_output_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"output_cost_per_token": 5e-07,
|
||||
"cache_read_input_token_cost": 3e-08,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
"max"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"mode": "chat",
|
||||
"comment": "Native multimodal GLM model for efficient coding and long-horizon agent tasks",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models",
|
||||
"supports_vision": true,
|
||||
"supports_image_input": true,
|
||||
"supports_video_input": true
|
||||
"supports_video_input": true,
|
||||
"mode": "chat",
|
||||
"comment": "Native multimodal GLM model for efficient coding and long-horizon agent tasks",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"friendliai/zai-org/GLM-5.3": {
|
||||
"litellm_provider": "friendliai",
|
||||
"supports_reasoning": true,
|
||||
"supports_function_calling": true,
|
||||
"max_input_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"max_output_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"input_cost_per_token": 1.26e-06,
|
||||
"output_cost_per_token": 3.96e-06,
|
||||
"cache_read_input_token_cost": 2.34e-07,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
"max"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"supports_image_input": false,
|
||||
"supports_video_input": false,
|
||||
"mode": "chat",
|
||||
"comment": "Flagship GLM model for long-horizon coding, agents, and complex project delivery",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"friendliai/google/gemma-4-31B-it": {
|
||||
"litellm_provider": "friendliai",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"input_cost_per_token": 1.4e-07,
|
||||
"output_cost_per_token": 4e-07,
|
||||
"supports_prompt_caching": false,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_image_input": true,
|
||||
"supports_video_input": false,
|
||||
"mode": "chat",
|
||||
"comment": "Largest Gemma 4 instruction model for open, self-hosted chat and reasoning",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"friendliai/zai-org/GLM-5.2": {
|
||||
"litellm_provider": "friendliai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"input_cost_per_token": 1.4e-06,
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"cache_read_input_token_cost": 2.6e-07,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [
|
||||
"high",
|
||||
"max"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"supports_image_input": false
|
||||
"supports_image_input": false,
|
||||
"supports_video_input": false,
|
||||
"mode": "chat",
|
||||
"comment": "Open flagship GLM for long-horizon coding agents and million-token context work",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"friendliai/LGAI-EXAONE/K-EXAONE-2.0-750B-A37B": {
|
||||
"litellm_provider": "friendliai",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"input_cost_per_token": 6e-07,
|
||||
"output_cost_per_token": 2.4e-06,
|
||||
"cache_read_input_token_cost": 1.2e-07,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"supports_image_input": false,
|
||||
"supports_video_input": false,
|
||||
"mode": "chat",
|
||||
"comment": "Frontier-scale multilingual language model developed by LG AI Research",
|
||||
"deprecation_date": "2026-09-06",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"friendliai/deepseek-ai/DeepSeek-V3.2": {
|
||||
"litellm_provider": "friendliai",
|
||||
"max_input_tokens": 163840,
|
||||
"max_output_tokens": 163840,
|
||||
"max_tokens": 163840,
|
||||
"input_cost_per_token": 5e-07,
|
||||
"output_cost_per_token": 1.5e-06,
|
||||
"cache_read_input_token_cost": 2.5e-07,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"supports_image_input": false,
|
||||
"supports_video_input": false,
|
||||
"mode": "chat",
|
||||
"comment": "DeepSeek chat model for instruction following, coding, and analysis",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"friendliai/MiniMaxAI/MiniMax-M2.5": {
|
||||
"litellm_provider": "friendliai",
|
||||
"max_input_tokens": 196608,
|
||||
"max_output_tokens": 196608,
|
||||
"max_tokens": 196608,
|
||||
"input_cost_per_token": 3e-07,
|
||||
"output_cost_per_token": 1.2e-06,
|
||||
"cache_read_input_token_cost": 6e-08,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"supports_image_input": false,
|
||||
"supports_video_input": false,
|
||||
"mode": "chat",
|
||||
"comment": "Prior MiniMax coding model for agent workflows, office edits, and automation",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"friendliai/zai-org/GLM-5.1": {
|
||||
"litellm_provider": "friendliai",
|
||||
"max_input_tokens": 202752,
|
||||
"max_output_tokens": 202752,
|
||||
"max_tokens": 202752,
|
||||
"input_cost_per_token": 1.4e-06,
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"cache_read_input_token_cost": 2.6e-07,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"supports_image_input": false,
|
||||
"supports_video_input": false,
|
||||
"mode": "chat",
|
||||
"comment": "Strong GLM coding model for agentic engineering, terminals, and repository generation",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"ft:babbage-002": {
|
||||
"deprecation_date": "2026-10-23",
|
||||
|
|
|
|||
|
|
@ -3882,6 +3882,7 @@ class SpendLogsPayload(TypedDict):
|
|||
session_id: str | None
|
||||
request_duration_ms: int | None
|
||||
status: Literal["success", "failure"]
|
||||
litellm_call_id: ReadOnly[str | None]
|
||||
|
||||
|
||||
class SpanAttributes(str, enum.Enum):
|
||||
|
|
@ -4580,7 +4581,6 @@ class UserManagementEndpointParamDocStringEnums(str, enum.Enum):
|
|||
)
|
||||
metadata_doc_str = """Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }"""
|
||||
max_parallel_requests_doc_str = """Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x."""
|
||||
soft_budget_doc_str = """Optional[float] - Get alerts when user crosses given budget, doesn't block requests."""
|
||||
model_max_budget_doc_str = """Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)"""
|
||||
model_rpm_limit_doc_str = """Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)"""
|
||||
model_tpm_limit_doc_str = """Optional[float] - Model-specific tpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)"""
|
||||
|
|
|
|||
|
|
@ -72,6 +72,7 @@ from litellm.proxy.auth.budget_throttle import (
|
|||
)
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation
|
||||
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
|
||||
from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
_safe_get_request_headers,
|
||||
_safe_get_request_query_params,
|
||||
|
|
@ -80,6 +81,7 @@ from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
|||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL,
|
||||
MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL,
|
||||
NO_TEAM_MEMBERSHIP_SENTINEL,
|
||||
TAG_REGISTRY_OVERFLOW_SENTINEL,
|
||||
UserApiKeyCache,
|
||||
end_user_cache_key,
|
||||
|
|
@ -2164,10 +2166,10 @@ async def get_team_membership(
|
|||
_key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
|
||||
|
||||
# check if in cache
|
||||
cached_membership_obj: Final = await user_api_key_cache.async_get_cache(
|
||||
key=_key,
|
||||
model_type=LiteLLM_TeamMembership,
|
||||
)
|
||||
cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key)
|
||||
if cached == NO_TEAM_MEMBERSHIP_SENTINEL:
|
||||
return None
|
||||
cached_membership_obj: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership)
|
||||
if cached_membership_obj is not None:
|
||||
return cached_membership_obj
|
||||
|
||||
|
|
@ -2179,6 +2181,11 @@ async def get_team_membership(
|
|||
)
|
||||
|
||||
if response is None:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=_key,
|
||||
value=NO_TEAM_MEMBERSHIP_SENTINEL,
|
||||
ttl=get_management_object_ttl(user_api_key_cache),
|
||||
)
|
||||
return None
|
||||
|
||||
_response: Final = LiteLLM_TeamMembership.model_validate(response.dict())
|
||||
|
|
|
|||
|
|
@ -52,6 +52,7 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import can_team_access_model
|
||||
from litellm.proxy.auth.resolvers.grants import GrantResolver, UserLookup, canonical_user_id
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.auth.team_grants import team_model_aliases
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
|
|
@ -1656,9 +1657,7 @@ class JWTAuthManager:
|
|||
``get_user_object`` resolved a legacy row with a different ``user_id``,
|
||||
use that row's id; otherwise keep the claim. GH #26789.
|
||||
"""
|
||||
if user_object is not None and user_object.user_id:
|
||||
return user_object.user_id
|
||||
return user_id
|
||||
return canonical_user_id(user_id=user_id, user_object=user_object)
|
||||
|
||||
@staticmethod
|
||||
async def get_objects(
|
||||
|
|
@ -1725,22 +1724,23 @@ class JWTAuthManager:
|
|||
code=403,
|
||||
)
|
||||
|
||||
user_object: LiteLLM_UserTable | None = None
|
||||
if user_id:
|
||||
user_object = (
|
||||
await get_user_object(
|
||||
user_id=user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=jwt_handler.is_upsert_user_id(valid_user_email=valid_user_email),
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
user_email=user_email,
|
||||
sso_user_id=user_id,
|
||||
)
|
||||
if user_id
|
||||
else None
|
||||
)
|
||||
user_object, team_membership_object, effective_user_id = await GrantResolver(
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
load_user=get_user_object,
|
||||
load_team=get_team_object,
|
||||
load_membership=get_team_membership,
|
||||
).resolve_identity(
|
||||
UserLookup(
|
||||
user_id=user_id,
|
||||
user_email=user_email,
|
||||
sso_user_id=user_id,
|
||||
upsert=jwt_handler.is_upsert_user_id(valid_user_email=valid_user_email),
|
||||
),
|
||||
team_id=team_id,
|
||||
)
|
||||
|
||||
end_user_object: LiteLLM_EndUserTable | None = None
|
||||
if end_user_id:
|
||||
|
|
@ -1757,37 +1757,12 @@ class JWTAuthManager:
|
|||
else None
|
||||
)
|
||||
|
||||
# Rebind to resolved DB user_id for team_membership + auth_builder (GH #26789).
|
||||
effective_user_id: Final = JWTAuthManager._canonical_user_id_from_db(user_id=user_id, user_object=user_object)
|
||||
if effective_user_id != user_id:
|
||||
verbose_proxy_logger.debug(
|
||||
"JWT Auth: rebinding user_id %r -> DB user_id %r (email/sso match)",
|
||||
user_id,
|
||||
effective_user_id,
|
||||
)
|
||||
user_id = effective_user_id
|
||||
|
||||
team_membership_object: LiteLLM_TeamMembership | None = None
|
||||
if user_id and team_id:
|
||||
team_membership_object = (
|
||||
await get_team_membership(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if user_id and team_id
|
||||
else None
|
||||
)
|
||||
|
||||
return (
|
||||
user_object,
|
||||
org_object,
|
||||
end_user_object,
|
||||
team_membership_object,
|
||||
user_id,
|
||||
effective_user_id,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
267
litellm/proxy/auth/resolvers/grants.py
Normal file
267
litellm/proxy/auth/resolvers/grants.py
Normal file
|
|
@ -0,0 +1,267 @@
|
|||
"""Load a caller's user row, team row, and team membership from the database and validate them together.
|
||||
|
||||
The virtual-key path reads these off the combined-view SQL join. Every other credential (an IdP JWT, a
|
||||
``lite login`` session token) carries only identifiers, or a snapshot of grants taken when it was minted, so
|
||||
it has to read the live rows on each request. Both of those paths resolve the same rows with the same
|
||||
membership rule, and ``GrantResolver`` is the one place that rule lives.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Coroutine, Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, NoReturn, Protocol, TypeAlias
|
||||
|
||||
from fastapi import HTTPException, status
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from pydantic.main import IncEx
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_TeamMembership,
|
||||
LiteLLM_TeamTableCachedObj,
|
||||
LiteLLM_UserTable,
|
||||
ProxyErrorTypes,
|
||||
ProxyException,
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
TeamNotFoundError,
|
||||
UserNotFoundError,
|
||||
get_team_membership,
|
||||
get_team_object,
|
||||
get_user_object,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy._types import Span
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
|
||||
|
||||
class UserLoader(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
user_id: str | None,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
user_id_upsert: bool,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging | None,
|
||||
sso_user_id: str | None,
|
||||
user_email: str | None,
|
||||
) -> Coroutine[object, object, LiteLLM_UserTable | None]: ...
|
||||
|
||||
|
||||
class TeamLoader(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
team_id: str,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging | None,
|
||||
) -> Coroutine[object, object, LiteLLM_TeamTableCachedObj]: ...
|
||||
|
||||
|
||||
class MembershipLoader(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
user_id: str,
|
||||
team_id: str,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging | None,
|
||||
) -> Coroutine[object, object, LiteLLM_TeamMembership | None]: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UserLookup:
|
||||
"""The user a credential names, plus the hints ``get_user_object`` may fall back to when the id alone
|
||||
matches no row."""
|
||||
|
||||
user_id: str | None
|
||||
user_email: str | None = None
|
||||
sso_user_id: str | None = None
|
||||
upsert: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ResolvedGrants:
|
||||
"""The live rows behind a credential. ``effective_user_id`` is the DB row's id when a fuzzy match found a
|
||||
legacy row under a different id (GH #26789), otherwise the id the credential named."""
|
||||
|
||||
user_object: LiteLLM_UserTable | None
|
||||
team_object: LiteLLM_TeamTableCachedObj | None
|
||||
team_membership: LiteLLM_TeamMembership | None
|
||||
effective_user_id: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UserGone:
|
||||
user_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TeamGone:
|
||||
team_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NotAMember:
|
||||
user_id: str
|
||||
team_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class LookupDegraded:
|
||||
"""A row could not be read for a reason that says nothing about the caller: the database is down or a
|
||||
loader failed. The caller decides whether a grant it already holds may stand in."""
|
||||
|
||||
error: Exception
|
||||
|
||||
|
||||
GrantDenial: TypeAlias = UserGone | TeamGone | NotAMember
|
||||
GrantOutcome: TypeAlias = ResolvedGrants | GrantDenial | LookupDegraded
|
||||
|
||||
|
||||
_MODELS_COLUMN: Final[Mapping[str, IncEx | bool]] = MappingProxyType({"models": True})
|
||||
|
||||
|
||||
class _UserModelColumn(BaseModel):
|
||||
"""``LiteLLM_UserTable.models`` is a bare ``list``; re-read it with the shape a token's ``models`` takes."""
|
||||
|
||||
models: tuple[str, ...] = ()
|
||||
|
||||
|
||||
def user_models(user_object: LiteLLM_UserTable) -> tuple[str, ...]:
|
||||
try:
|
||||
return _UserModelColumn.model_validate(user_object.model_dump(include=_MODELS_COLUMN)).models
|
||||
except ValidationError:
|
||||
return ()
|
||||
|
||||
|
||||
def canonical_user_id(user_id: str | None, user_object: LiteLLM_UserTable | None) -> str | None:
|
||||
if user_object is not None and user_object.user_id:
|
||||
return user_object.user_id
|
||||
return user_id
|
||||
|
||||
|
||||
def raise_public(denial: GrantDenial) -> NoReturn:
|
||||
match denial:
|
||||
case UserGone(user_id=user_id):
|
||||
raise ProxyException(
|
||||
message=f"Authentication Error, user '{user_id}' no longer exists.",
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="user_id",
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
case TeamGone(team_id=team_id):
|
||||
raise TeamNotFoundError(team_id=team_id)
|
||||
case NotAMember(team_id=team_id):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=f"Team '{team_id}' is not in your team memberships.",
|
||||
)
|
||||
case _:
|
||||
assert_never(denial)
|
||||
|
||||
|
||||
class GrantResolver:
|
||||
"""Reads the user, membership, and team rows for a credential through injected loaders.
|
||||
|
||||
The loaders default to the shared ``auth_checks`` readers. A caller passes its own module's names for them
|
||||
so the reads stay interceptable where that module's callers already intercept them. ``resolve_identity``
|
||||
is the JWT half: user and membership only, since the JWT builder selects the team itself and lets loader
|
||||
errors surface as they are. ``resolve`` also reads the team row and applies the membership rule, which is
|
||||
what a credential carrying a grant snapshot needs to refresh it.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
prisma_client: PrismaClient | None,
|
||||
cache: UserApiKeyCache,
|
||||
*,
|
||||
parent_otel_span: Span | None = None,
|
||||
proxy_logging_obj: ProxyLogging | None = None,
|
||||
load_user: UserLoader = get_user_object,
|
||||
load_team: TeamLoader = get_team_object,
|
||||
load_membership: MembershipLoader = get_team_membership,
|
||||
) -> None:
|
||||
self._prisma = prisma_client
|
||||
self._cache = cache
|
||||
self._parent_otel_span = parent_otel_span
|
||||
self._proxy_logging_obj = proxy_logging_obj
|
||||
self._load_user = load_user
|
||||
self._load_team = load_team
|
||||
self._load_membership = load_membership
|
||||
|
||||
async def resolve_identity(
|
||||
self, lookup: UserLookup, team_id: str | None
|
||||
) -> tuple[LiteLLM_UserTable | None, LiteLLM_TeamMembership | None, str | None]:
|
||||
user_object: Final = await self._user(lookup) if lookup.user_id else None
|
||||
effective_user_id: Final = canonical_user_id(lookup.user_id, user_object)
|
||||
if effective_user_id != lookup.user_id:
|
||||
verbose_proxy_logger.debug(
|
||||
"Auth: rebinding user_id %r -> DB user_id %r (email/sso match)",
|
||||
lookup.user_id,
|
||||
effective_user_id,
|
||||
)
|
||||
membership: Final = (
|
||||
await self._membership(user_id=effective_user_id, team_id=team_id)
|
||||
if effective_user_id and team_id
|
||||
else None
|
||||
)
|
||||
return user_object, membership, effective_user_id
|
||||
|
||||
async def resolve(self, lookup: UserLookup, team_id: str | None) -> GrantOutcome:
|
||||
try:
|
||||
user_object, membership, effective_user_id = await self.resolve_identity(lookup, team_id)
|
||||
except UserNotFoundError:
|
||||
return UserGone(user_id=lookup.user_id or "")
|
||||
except Exception as error:
|
||||
return LookupDegraded(error=error)
|
||||
if team_id is None:
|
||||
return ResolvedGrants(user_object, None, membership, effective_user_id)
|
||||
if user_object is not None and team_id not in user_object.teams:
|
||||
return NotAMember(user_id=user_object.user_id, team_id=team_id)
|
||||
try:
|
||||
team_object: Final = await self._load_team(
|
||||
team_id=team_id,
|
||||
prisma_client=self._prisma,
|
||||
user_api_key_cache=self._cache,
|
||||
parent_otel_span=self._parent_otel_span,
|
||||
proxy_logging_obj=self._proxy_logging_obj,
|
||||
)
|
||||
except TeamNotFoundError:
|
||||
return TeamGone(team_id=team_id)
|
||||
except Exception as error:
|
||||
return LookupDegraded(error=error)
|
||||
return ResolvedGrants(user_object, team_object, membership, effective_user_id)
|
||||
|
||||
async def _user(self, lookup: UserLookup) -> LiteLLM_UserTable | None:
|
||||
return await self._load_user(
|
||||
user_id=lookup.user_id,
|
||||
prisma_client=self._prisma,
|
||||
user_api_key_cache=self._cache,
|
||||
user_id_upsert=lookup.upsert,
|
||||
parent_otel_span=self._parent_otel_span,
|
||||
proxy_logging_obj=self._proxy_logging_obj,
|
||||
user_email=lookup.user_email,
|
||||
sso_user_id=lookup.sso_user_id,
|
||||
)
|
||||
|
||||
async def _membership(self, user_id: str, team_id: str) -> LiteLLM_TeamMembership | None:
|
||||
return await self._load_membership(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
prisma_client=self._prisma,
|
||||
user_api_key_cache=self._cache,
|
||||
parent_otel_span=self._parent_otel_span,
|
||||
proxy_logging_obj=self._proxy_logging_obj,
|
||||
)
|
||||
|
|
@ -56,6 +56,7 @@ from litellm.proxy.auth.auth_checks import (
|
|||
get_jwt_key_mapping_object,
|
||||
get_object_permission,
|
||||
get_project_object,
|
||||
get_team_membership,
|
||||
get_team_object,
|
||||
get_user_object,
|
||||
is_valid_fallback_model,
|
||||
|
|
@ -82,6 +83,14 @@ from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_conte
|
|||
from litellm.proxy.auth.oauth2_check import Oauth2Handler
|
||||
from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request
|
||||
from litellm.proxy.auth.resolvers import CredentialRef, Principal
|
||||
from litellm.proxy.auth.resolvers.grants import (
|
||||
GrantResolver,
|
||||
LookupDegraded,
|
||||
ResolvedGrants,
|
||||
UserLookup,
|
||||
raise_public,
|
||||
user_models,
|
||||
)
|
||||
from litellm.proxy.auth.resolvers.store import IdentityStore
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.auth.team_grants import team_grants
|
||||
|
|
@ -1230,6 +1239,52 @@ async def _record_unparsable_body_failure(
|
|||
verbose_proxy_logger.exception("Failed to log the request rejected for an unparsable body: %s", e)
|
||||
|
||||
|
||||
async def _refresh_session_token_grants(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> UserAPIKeyAuth:
|
||||
"""Rebuild a ``lite login`` session token's grants from the live user and team rows.
|
||||
|
||||
The blob only proves who logged in and which team they picked. Team models, aliases, the user's own model
|
||||
list, and their role are re-read every request, so a `/team/update` or a demotion shows up without a
|
||||
re-login, and a user removed from the team or deleted outright is refused. When a row cannot be read for
|
||||
a reason unrelated to the caller, the minted grants stand in exactly as they did before this refresh.
|
||||
"""
|
||||
outcome: Final = await GrantResolver(
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
load_user=get_user_object,
|
||||
load_team=get_team_object,
|
||||
load_membership=get_team_membership,
|
||||
).resolve(UserLookup(user_id=valid_token.user_id), team_id=valid_token.team_id)
|
||||
match outcome:
|
||||
case ResolvedGrants(
|
||||
user_object=LiteLLM_UserTable() as user_object, team_object=team_object, team_membership=team_membership
|
||||
):
|
||||
return UserAPIKeyAuth.model_validate(
|
||||
MappingProxyType(
|
||||
{
|
||||
**valid_token.model_dump(exclude_none=True),
|
||||
**team_grants(team_object, team_membership, user_object.user_id),
|
||||
"user_role": _get_user_role(user_object),
|
||||
"models": () if team_object is not None else user_models(user_object),
|
||||
}
|
||||
)
|
||||
)
|
||||
case ResolvedGrants():
|
||||
return valid_token
|
||||
case LookupDegraded(error=error):
|
||||
verbose_proxy_logger.debug("Session token grants not refreshed, keeping minted grants: %s", error)
|
||||
return valid_token
|
||||
case _:
|
||||
raise_public(outcome)
|
||||
|
||||
|
||||
async def _resolve_object_permission_for_unresolvable_team(
|
||||
object_permission_id: str | None,
|
||||
prisma_client: PrismaClient | None,
|
||||
|
|
@ -1774,6 +1829,15 @@ async def _user_api_key_auth_builder(
|
|||
):
|
||||
valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(api_key)
|
||||
|
||||
if valid_token is not None and valid_token.is_session_token and prisma_client is not None:
|
||||
valid_token = await _refresh_session_token_grants( # rebind-ok: later checks read this name
|
||||
valid_token=valid_token,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
if (
|
||||
valid_token is not None
|
||||
and isinstance(valid_token, UserAPIKeyAuth)
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ from litellm.constants import (
|
|||
DEFAULT_MAX_RECURSE_DEPTH,
|
||||
LITELLM_DETAILED_TIMING,
|
||||
LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED,
|
||||
MAX_LITELLM_CALL_ID_LENGTH,
|
||||
MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG,
|
||||
NON_INFERENCE_CALL_TYPES,
|
||||
RETURN_RAW_MODEL_NAME_METADATA_KEY,
|
||||
|
|
@ -217,6 +218,12 @@ def _withheld_provider_output(response: object) -> bool:
|
|||
return getattr(response, "has_buffered_provider_output", False) is True
|
||||
|
||||
|
||||
def resolve_litellm_call_id(client_call_id: str | None) -> str:
|
||||
if client_call_id is not None and 0 < len(client_call_id) <= MAX_LITELLM_CALL_ID_LENGTH:
|
||||
return client_call_id
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
def _should_return_raw_model_name(request_data: dict[str, object]) -> bool:
|
||||
return any(
|
||||
isinstance(metadata, dict) and metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY) is True
|
||||
|
|
@ -1938,7 +1945,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
if alias_target is not None:
|
||||
self.data["model"] = alias_target
|
||||
|
||||
self.data["litellm_call_id"] = request.headers.get("x-litellm-call-id", str(uuid.uuid4()))
|
||||
self.data["litellm_call_id"] = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
|
||||
DDSpanTagger.tag_call_id(self.data.get("litellm_call_id"))
|
||||
DDSpanTagger.tag_request(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
|
|||
|
|
@ -336,6 +336,14 @@ def team_membership_reservation_cache_key(user_id: str, team_id: str) -> str:
|
|||
return f"team_membership:{user_id}:{team_id}"
|
||||
|
||||
|
||||
#: Cached under ``team_membership_reservation_cache_key`` when a member has no ``LiteLLM_TeamMembership``
|
||||
#: row, so a session-token member without a per-member budget costs no DB read per request. Lives beside
|
||||
#: the key builder because it is part of the same cache protocol: every reader of the key must know that
|
||||
#: a plain string here means "no row", distinct from a serialized membership. The two budget readers
|
||||
#: already treat a non-model value as "no row", so they need no change to stay correct.
|
||||
NO_TEAM_MEMBERSHIP_SENTINEL: Final = "__no_team_membership__"
|
||||
|
||||
|
||||
def get_management_object_ttl(cache: DualCache) -> float:
|
||||
"""
|
||||
In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...).
|
||||
|
|
|
|||
|
|
@ -443,6 +443,11 @@ async def health_services_endpoint(
|
|||
}
|
||||
return pointfive_health
|
||||
if service == "webhook":
|
||||
if not _is_proxy_admin(user_api_key_dict):
|
||||
webhook_non_admin_detail: Final[_ServiceTestErrorDetail] = {
|
||||
"error": "Only proxy admins can trigger the webhook test alert."
|
||||
}
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=webhook_non_admin_detail)
|
||||
user_info: Final = CallInfo(
|
||||
token=user_api_key_dict.token or "",
|
||||
spend=1,
|
||||
|
|
|
|||
|
|
@ -478,7 +478,6 @@ async def new_user(
|
|||
- permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.
|
||||
- metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
||||
- max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||||
- soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests.
|
||||
- model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
|
|
@ -1651,7 +1650,6 @@ async def user_update(
|
|||
- permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.
|
||||
- metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
||||
- max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||||
- soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests.
|
||||
- model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ import copy
|
|||
import json
|
||||
import math
|
||||
import traceback
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from collections.abc import Set as AbstractSet
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
|
|
@ -1903,6 +1903,39 @@ def validate_team_org_change(
|
|||
return True
|
||||
|
||||
|
||||
def _member_user_ids(members_with_roles: Sequence[dict[str, object]]) -> tuple[str, ...]:
|
||||
"""Extract the string ``user_id`` of each team member, dropping rows without one.
|
||||
|
||||
``members_with_roles`` is a Prisma-deserialized JSON column, so its ``user_id`` is typed
|
||||
``object``; the ``isinstance`` narrows it to the ``str`` ``invalidate_team_member_spend_state`` needs.
|
||||
"""
|
||||
return tuple(user_id for member in members_with_roles if isinstance((user_id := member.get("user_id")), str))
|
||||
|
||||
|
||||
async def _evict_created_membership_caches(
|
||||
user_ids: Iterable[str],
|
||||
team_id: str,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
) -> None:
|
||||
"""Evict the ``get_team_membership`` negative-cache sentinel for members whose row was just created.
|
||||
|
||||
A session-token request caches ``NO_TEAM_MEMBERSHIP_SENTINEL`` for a member with no
|
||||
``LiteLLM_TeamMembership`` row. When a create path (``/team/member_add`` or the ``/team/update``
|
||||
budget backfill) later writes that row with a per-member budget, the stale sentinel keeps the
|
||||
member's budget unenforced until the membership cache TTL expires, so it must be evicted here.
|
||||
"""
|
||||
await asyncio.gather(
|
||||
*(
|
||||
invalidate_team_member_spend_state(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
for user_id in user_ids
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.post("/team/update", tags=["team management"], dependencies=[Depends(user_api_key_auth)])
|
||||
@management_endpoint_wrapper
|
||||
async def update_team(
|
||||
|
|
@ -2239,6 +2272,11 @@ async def update_team(
|
|||
team_member_budget_id=_backfill_budget_id,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
await _evict_created_membership_caches(
|
||||
user_ids=_member_user_ids(existing_team_row.members_with_roles),
|
||||
team_id=data.team_id,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
elif _team_member_fields_in_request:
|
||||
updated_kv = await TeamMemberBudgetHandler.clear_team_member_budget_fields(
|
||||
team_table=existing_team_row,
|
||||
|
|
@ -3191,6 +3229,12 @@ async def team_member_add(
|
|||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
)
|
||||
|
||||
await _evict_created_membership_caches(
|
||||
user_ids=(tm.user_id for tm in updated_team_memberships),
|
||||
team_id=data.team_id,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
_emit_team_members_metric(complete_team_data)
|
||||
|
||||
await _create_team_member_add_audit_logs(
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from typing import (
|
|||
)
|
||||
|
||||
from litellm.batches.batch_utils import batch_cost_is_final
|
||||
from litellm.constants import MAX_FILE_LIST_LIMIT
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.repositories.table_repositories import (
|
||||
ManagedFileRepository,
|
||||
|
|
@ -34,8 +35,6 @@ if TYPE_CHECKING:
|
|||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
|
||||
MAX_FILE_LIST_LIMIT: Final = 10000
|
||||
|
||||
FILE_LIST_CONTINUATION_CHUNK_SIZE: Final = 500
|
||||
|
||||
BATCH_CREATE_HIDDEN_PARAM: Final = "batch_create"
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ from litellm.litellm_core_utils.cloud_storage_security import (
|
|||
)
|
||||
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
|
||||
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
|
||||
from litellm.llms.base_llm.managed_resources.isolation import build_list_page
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
|
@ -90,6 +91,7 @@ from litellm.router import Router
|
|||
from litellm.types.llms.openai import (
|
||||
CREATE_FILE_REQUESTS_PURPOSE,
|
||||
FileExpiresAfter,
|
||||
FileListPage,
|
||||
OpenAIFileObject,
|
||||
OpenAIFilesPurpose,
|
||||
)
|
||||
|
|
@ -97,6 +99,7 @@ from litellm.types.llms.openai import (
|
|||
router: Final = APIRouter()
|
||||
|
||||
_MAX_BATCH_FILE_SIZE_MB_ADAPTER: Final = TypeAdapter(int | None)
|
||||
_LISTED_FILES_ADAPTER: Final = TypeAdapter(list[OpenAIFileObject])
|
||||
|
||||
|
||||
class UploadedFileInfo(TypedDict):
|
||||
|
|
@ -1287,6 +1290,11 @@ async def delete_file(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
|
||||
)
|
||||
if is_managed_cloud_storage_uri(file_id) and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="Raw cloud storage file ids can only be deleted by a proxy admin key. Use the LiteLLM managed file id returned when the file was created.",
|
||||
)
|
||||
|
||||
custom_llm_provider: Final = (
|
||||
provider
|
||||
|
|
@ -1446,6 +1454,12 @@ async def delete_file(
|
|||
)
|
||||
|
||||
|
||||
def _as_file_list_page(response: object) -> object:
|
||||
if not isinstance(response, list):
|
||||
return response
|
||||
return FileListPage(**build_list_page(_LISTED_FILES_ADAPTER.validate_python(response)))
|
||||
|
||||
|
||||
@router.get(
|
||||
"/{provider}/v1/files",
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
|
|
@ -1524,7 +1538,7 @@ async def list_files(
|
|||
|
||||
if should_route and credentials is not None:
|
||||
# Use model-based routing with credentials from config
|
||||
prepare_data_with_credentials(data=data, credentials=credentials)
|
||||
prepare_data_with_credentials(data=data, credentials=credentials, include_internal_credentials=True)
|
||||
response = await litellm.afile_list(
|
||||
custom_llm_provider=credentials["custom_llm_provider"],
|
||||
purpose=purpose,
|
||||
|
|
@ -1550,7 +1564,7 @@ async def list_files(
|
|||
model_id=target_model_names_list[0],
|
||||
operation_context="file list",
|
||||
)
|
||||
prepare_data_with_credentials(data=data, credentials=credentials)
|
||||
prepare_data_with_credentials(data=data, credentials=credentials, include_internal_credentials=True)
|
||||
response = await litellm.afile_list(
|
||||
custom_llm_provider=credentials["custom_llm_provider"],
|
||||
purpose=purpose,
|
||||
|
|
@ -1592,6 +1606,7 @@ async def list_files(
|
|||
status_code=500,
|
||||
detail="Either 'provider' or 'target_model_names' must be provided e.g. `?target_model_names=gpt-4o`",
|
||||
)
|
||||
response = _as_file_list_page(response) # rebind-ok: each dispatch branch above binds response
|
||||
|
||||
## POST CALL HOOKS ###
|
||||
_response: Final = await proxy_logging_obj.post_call_success_hook(
|
||||
|
|
|
|||
|
|
@ -659,12 +659,14 @@ model LiteLLM_SpendLogs {
|
|||
mcp_namespaced_tool_name String?
|
||||
agent_id String?
|
||||
proxy_server_request Json? @default("{}")
|
||||
litellm_call_id String?
|
||||
created_at DateTime @default(now()) @map("created_at")
|
||||
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
|
||||
@@index([startTime])
|
||||
@@index([startTime, request_id])
|
||||
@@index([end_user])
|
||||
@@index([session_id])
|
||||
@@index([litellm_call_id])
|
||||
}
|
||||
|
||||
model LiteLLM_BudgetWindowSpend {
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import collections
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from itertools import groupby
|
||||
from types import MappingProxyType
|
||||
|
|
@ -17,6 +18,7 @@ from typing import (
|
|||
TypeAlias,
|
||||
TypedDict,
|
||||
TypeVar,
|
||||
cast, # noqa: TID251 # custom-logger and cold-storage payloads are untyped JSON
|
||||
)
|
||||
|
||||
import fastapi
|
||||
|
|
@ -75,6 +77,7 @@ _SPEND_LOG_LIST_COLUMNS: Final = """
|
|||
cache_hit, cache_key, request_tags, team_id,
|
||||
organization_id, end_user, requester_ip_address,
|
||||
session_id, status, mcp_namespaced_tool_name, agent_id,
|
||||
litellm_call_id,
|
||||
COALESCE(request_duration_ms,
|
||||
(EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms
|
||||
"""
|
||||
|
|
@ -91,9 +94,9 @@ class _SupportsModelDump(Protocol):
|
|||
def model_dump(self) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
class _SpendLogOwnershipRow(Protocol):
|
||||
user: str | None
|
||||
team_id: str | None
|
||||
class _SpendLogOwnerRow(TypedDict):
|
||||
user: ReadOnly[str | None]
|
||||
team_id: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _ActivityRow(TypedDict):
|
||||
|
|
@ -330,12 +333,37 @@ async def _find_spend_logs(
|
|||
return rows
|
||||
|
||||
|
||||
async def _find_spend_log_row(prisma_client: PrismaClient, request_id: str) -> _SpendLogOwnershipRow | None:
|
||||
"""Read the single spend log row identified by ``request_id``."""
|
||||
return await _spend_logs_table(prisma_client).find_unique(
|
||||
where={"request_id": request_id},
|
||||
include=None,
|
||||
)
|
||||
class _RequestIdEquals(TypedDict):
|
||||
request_id: ReadOnly[str]
|
||||
|
||||
|
||||
class _LitellmCallIdEquals(TypedDict):
|
||||
litellm_call_id: ReadOnly[str]
|
||||
|
||||
|
||||
def _request_id_or_call_id_clause(request_id: str) -> tuple[_RequestIdEquals, _LitellmCallIdEquals]:
|
||||
request_id_clause: Final[_RequestIdEquals] = {"request_id": request_id}
|
||||
call_id_clause: Final[_LitellmCallIdEquals] = {"litellm_call_id": request_id}
|
||||
return (request_id_clause, call_id_clause)
|
||||
|
||||
|
||||
async def _find_spend_log_owners(prisma_client: PrismaClient, request_id: str) -> Sequence[_SpendLogOwnerRow]:
|
||||
"""Read the distinct ``(user, team_id)`` owner pairs across every spend log row
|
||||
identified by ``request_id`` or ``litellm_call_id``.
|
||||
|
||||
``litellm_call_id`` is populated from the client-settable ``x-litellm-call-id``
|
||||
request header, so it is not guaranteed unique to one tenant: any number of rows
|
||||
can match one id. The read is uncapped because a flood of another tenant's rows
|
||||
carrying the caller's id could otherwise push the caller's own owner pair past a
|
||||
row-sample cap and lock them out of their own lookup.
|
||||
"""
|
||||
sql_query: Final = """
|
||||
SELECT DISTINCT "user", team_id
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE request_id = $1 OR litellm_call_id = $1
|
||||
"""
|
||||
owners: Final[Sequence[_SpendLogOwnerRow] | None] = await _query_raw_or_none(prisma_client, sql_query, request_id)
|
||||
return owners if owners is not None else ()
|
||||
|
||||
|
||||
async def _count_spend_logs(prisma_client: PrismaClient, where: Mapping[str, object]) -> int:
|
||||
|
|
@ -2599,10 +2627,11 @@ async def ui_view_spend_logs(
|
|||
if max_spend is not None:
|
||||
where_conditions["spend"]["lte"] = max_spend
|
||||
# A request_id lookup drops the date window, so a non-admin could otherwise
|
||||
# reach any single row by id; require they own it, mirroring the detail
|
||||
# endpoint. That ownership check fully authorizes the one row, so the
|
||||
# general scoping below is skipped for id lookups. Scoped to the UI route
|
||||
# so the public v2 contract is unchanged.
|
||||
# reach any single row by id; require they own one of the matches, mirroring
|
||||
# the detail endpoint, and keep the general scoping below so a colliding
|
||||
# foreign row is filtered out rather than served or allowed to deny the
|
||||
# caller their own row. Scoped to the UI route so the public v2 contract is
|
||||
# unchanged.
|
||||
if request_id is not None and not is_v2 and not is_admin_view:
|
||||
await _assert_user_can_view_request_id(
|
||||
prisma_client=prisma_client,
|
||||
|
|
@ -2610,10 +2639,9 @@ async def ui_view_spend_logs(
|
|||
request_id=request_id,
|
||||
)
|
||||
user_scope_applies: Final = (
|
||||
not is_request_id_lookup
|
||||
and not is_admin_view
|
||||
not is_admin_view
|
||||
and team_id is None
|
||||
and _can_user_view_spend_log(user_api_key_dict=user_api_key_dict)
|
||||
and (is_request_id_lookup or _can_user_view_spend_log(user_api_key_dict=user_api_key_dict))
|
||||
)
|
||||
permitted_team_ids: Final = (
|
||||
await _get_permitted_team_ids_for_spend_logs_or_empty(
|
||||
|
|
@ -2626,7 +2654,7 @@ async def ui_view_spend_logs(
|
|||
explicit_user_requires_caller_scope: Final = (
|
||||
user_scope_applies and not permitted_team_ids and user_id is not None
|
||||
)
|
||||
if not is_request_id_lookup and not is_admin_view:
|
||||
if not is_admin_view:
|
||||
if team_id is not None:
|
||||
can_view_team: Final = await _can_team_member_view_log(
|
||||
prisma_client=prisma_client,
|
||||
|
|
@ -2696,7 +2724,6 @@ async def ui_view_spend_logs(
|
|||
("team_id", "team_id"),
|
||||
('"user"', "user"),
|
||||
("api_key", "api_key"),
|
||||
("request_id", "request_id"),
|
||||
("model", "model"),
|
||||
("model_id", "model_id"),
|
||||
("model_group", "model_group"),
|
||||
|
|
@ -2708,6 +2735,13 @@ async def ui_view_spend_logs(
|
|||
sql_params.append(val)
|
||||
p += 1
|
||||
|
||||
request_id_filter: Final = where_conditions.get("request_id")
|
||||
exact_request_id_first: Final = f"(request_id = ${p}) DESC, " if isinstance(request_id_filter, str) else ""
|
||||
if isinstance(request_id_filter, str):
|
||||
sql_conditions.append(f"(request_id = ${p} OR litellm_call_id = ${p})")
|
||||
sql_params.append(request_id_filter)
|
||||
p += 1
|
||||
|
||||
# Multi-team OR filter: (user = $X OR team_id = ANY($Y))
|
||||
if permitted_team_ids:
|
||||
or_clause: Final = f'("user" = ${p} OR team_id = ANY(${p + 1}::text[]))'
|
||||
|
|
@ -2837,7 +2871,7 @@ async def ui_view_spend_logs(
|
|||
WHERE {joined_conditions}
|
||||
ORDER BY {_SESSION_GROUP_KEY_SQL}, call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC
|
||||
) AS session_representatives
|
||||
ORDER BY {_order_expr} {_sql_dir}{_nulls_clause}, request_id
|
||||
ORDER BY {exact_request_id_first}{_order_expr} {_sql_dir}{_nulls_clause}, request_id
|
||||
LIMIT ${p} OFFSET ${p + 1}
|
||||
"""
|
||||
if session_grouping
|
||||
|
|
@ -2846,7 +2880,7 @@ async def ui_view_spend_logs(
|
|||
{_SPEND_LOG_LIST_COLUMNS}
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE {joined_conditions}
|
||||
ORDER BY {_order_expr} {_sql_dir}{_nulls_clause}
|
||||
ORDER BY {exact_request_id_first}{_order_expr} {_sql_dir}{_nulls_clause}
|
||||
LIMIT ${p} OFFSET ${p + 1}
|
||||
"""
|
||||
)
|
||||
|
|
@ -2854,6 +2888,14 @@ async def ui_view_spend_logs(
|
|||
|
||||
data: Final = await prisma_client.db.query_raw(sql_query, *sql_params)
|
||||
|
||||
if request_id is not None and not is_v2 and not is_admin_view:
|
||||
await _assert_user_owns_fetched_spend_rows(
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
rows=data,
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
_hydrate_spend_log_metadata(data)
|
||||
|
||||
# Calculate total pages
|
||||
|
|
@ -3104,7 +3146,7 @@ def _hydrate_spend_log_metadata(rows: Sequence[Mapping[str, object]]) -> None:
|
|||
|
||||
|
||||
def _cold_storage_object_key_from_metadata(
|
||||
metadata: str | dict | None,
|
||||
metadata: str | Mapping[str, object] | None,
|
||||
) -> str | None:
|
||||
if isinstance(metadata, str):
|
||||
try:
|
||||
|
|
@ -3209,7 +3251,8 @@ async def ui_view_request_response_for_request_id(
|
|||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if not _is_admin_view_safe(user_api_key_dict=user_api_key_dict):
|
||||
caller_is_admin: Final = _is_admin_view_safe(user_api_key_dict=user_api_key_dict)
|
||||
if not caller_is_admin:
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
|
|
@ -3234,38 +3277,45 @@ async def ui_view_request_response_for_request_id(
|
|||
if end_date is not None:
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc)
|
||||
|
||||
spend_log_row: Final = (
|
||||
None
|
||||
if prisma_client is None
|
||||
else await _resolve_spend_log_payload_row(
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_id=request_id,
|
||||
caller_is_admin=caller_is_admin,
|
||||
)
|
||||
)
|
||||
stored_request_id: Final = _stored_request_id(spend_log_row, request_id)
|
||||
|
||||
for custom_logger in custom_loggers:
|
||||
payload = await custom_logger.get_request_response_payload(
|
||||
request_id=request_id,
|
||||
request_id=stored_request_id,
|
||||
start_time_utc=start_date_obj,
|
||||
end_time_utc=end_date_obj,
|
||||
)
|
||||
if payload is not None:
|
||||
if not caller_is_admin and prisma_client is not None:
|
||||
await _assert_user_owns_cold_storage_payload(
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
payload=cast(Mapping[str, object], payload), # cast-ok: custom-logger payload is untyped
|
||||
request_id=request_id,
|
||||
)
|
||||
return payload
|
||||
|
||||
if spend_log_row is None:
|
||||
return None
|
||||
|
||||
# Fallback: the list endpoint omits the heavy columns for performance, so
|
||||
# serve them here. When prompts were offloaded to cold storage the DB holds
|
||||
# only placeholders, so _resolve_request_response_payload fetches the real
|
||||
# payload from the configured cold storage backend by object key.
|
||||
if prisma_client is not None:
|
||||
from litellm.proxy.spend_tracking.cold_storage_handler import (
|
||||
ColdStorageHandler,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
|
||||
|
||||
sql_query: Final = """
|
||||
SELECT messages, response, proxy_server_request, metadata
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE request_id = $1
|
||||
LIMIT 1
|
||||
"""
|
||||
db_result: Final[Sequence[Mapping[str, object]] | None] = await _query_raw_or_none(
|
||||
prisma_client, sql_query, request_id
|
||||
)
|
||||
if db_result and len(db_result) > 0:
|
||||
resolved = await _resolve_request_response_payload(db_result[0], cold_storage_handler=ColdStorageHandler())
|
||||
return resolved._asdict()
|
||||
|
||||
return None
|
||||
resolved: Final = await _resolve_request_response_payload(spend_log_row, cold_storage_handler=ColdStorageHandler())
|
||||
return resolved._asdict()
|
||||
|
||||
|
||||
@router.get(
|
||||
|
|
@ -3391,7 +3441,7 @@ async def view_spend_logs(
|
|||
if api_key is not None and isinstance(api_key, str):
|
||||
filter_query["api_key"] = summary_api_key
|
||||
if request_id is not None and isinstance(request_id, str):
|
||||
filter_query["request_id"] = request_id
|
||||
filter_query["OR"] = _request_id_or_call_id_clause(request_id)
|
||||
if user_id is not None and isinstance(user_id, str):
|
||||
filter_query["user"] = user_id
|
||||
|
||||
|
|
@ -3439,7 +3489,7 @@ async def view_spend_logs(
|
|||
return [*summary_items, *padding]
|
||||
|
||||
else:
|
||||
scoped_filter: Final[dict[str, str]] = {}
|
||||
scoped_filter: Final[dict[str, object]] = {}
|
||||
if api_key is not None and isinstance(api_key, str):
|
||||
if api_key.startswith("sk-"):
|
||||
hashed_token = prisma_client.hash_token(token=api_key)
|
||||
|
|
@ -3447,7 +3497,7 @@ async def view_spend_logs(
|
|||
hashed_token = api_key
|
||||
scoped_filter["api_key"] = hashed_token
|
||||
if request_id is not None and isinstance(request_id, str):
|
||||
scoped_filter["request_id"] = request_id
|
||||
scoped_filter["OR"] = _request_id_or_call_id_clause(request_id)
|
||||
if user_id is not None and isinstance(user_id, str):
|
||||
scoped_filter["user"] = user_id
|
||||
|
||||
|
|
@ -4676,39 +4726,190 @@ def _can_user_view_spend_log(user_api_key_dict: UserAPIKeyAuth) -> bool:
|
|||
)
|
||||
|
||||
|
||||
async def _user_can_view_spend_log_owner(
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
owner_user: str | None,
|
||||
owner_team_id: str | None,
|
||||
) -> bool:
|
||||
if owner_user is not None and owner_user == user_api_key_dict.user_id:
|
||||
return True
|
||||
if owner_team_id:
|
||||
return await _can_team_member_view_log(
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
team_id=owner_team_id,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _spend_log_forbidden(request_id: str) -> HTTPException:
|
||||
return HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail={"error": f"Not authorized to view spend log for request_id={request_id}"},
|
||||
)
|
||||
|
||||
|
||||
async def _assert_user_can_view_request_id(
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_id: str,
|
||||
) -> None:
|
||||
"""
|
||||
Verify the requesting non-admin user is allowed to view this spend-log row.
|
||||
Allowed when the log belongs to the user directly, or to one of their
|
||||
permitted teams (admin or ``/spend/logs`` permission).
|
||||
Raises HTTP 403 if not, including when no spend-log row exists for the
|
||||
request_id (e.g. it was pruned by retention), so a missing row can't be
|
||||
used to read a payload out of cold storage via the detail endpoint.
|
||||
Verify the requesting non-admin user is allowed to view at least one spend-log
|
||||
row identified by ``request_id`` or ``litellm_call_id``. The latter is
|
||||
client-settable, so an id lookup can match rows across different tenants; the
|
||||
data queries scope a non-admin's results to rows they own directly or via a
|
||||
permitted team, so a colliding foreign row can neither be served nor deny the
|
||||
caller their own. Raises HTTP 403 when none of the matching rows is theirs to
|
||||
view, including when no row exists at all (e.g. it was pruned by retention),
|
||||
so a missing row can't be used to read a payload out of cold storage via the
|
||||
detail endpoint.
|
||||
"""
|
||||
row: Final = await _find_spend_log_row(prisma_client, request_id)
|
||||
owners: Final = await _find_spend_log_owners(prisma_client, request_id)
|
||||
for owner in owners:
|
||||
if await _user_can_view_spend_log_owner(prisma_client, user_api_key_dict, owner["user"], owner["team_id"]):
|
||||
return
|
||||
raise _spend_log_forbidden(request_id)
|
||||
|
||||
if row is not None and row.user is not None and row.user == user_api_key_dict.user_id:
|
||||
return
|
||||
|
||||
if row is not None and row.team_id:
|
||||
can_view: Final = await _can_team_member_view_log(
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _SpendLogViewer:
|
||||
user_id: str | None
|
||||
team_ids: tuple[str, ...]
|
||||
|
||||
|
||||
async def _spend_log_viewer(prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth) -> _SpendLogViewer:
|
||||
return _SpendLogViewer(
|
||||
user_id=user_api_key_dict.user_id,
|
||||
team_ids=await _get_permitted_team_ids_for_spend_logs_or_empty(
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
team_id=row.team_id,
|
||||
)
|
||||
if can_view:
|
||||
return
|
||||
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail={"error": f"Not authorized to view spend log for request_id={request_id}"},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _viewer_scope_clause(viewer: _SpendLogViewer | None) -> tuple[str, tuple[object, ...]]:
|
||||
match viewer:
|
||||
case None:
|
||||
return ("", ())
|
||||
case _SpendLogViewer(user_id=user_id, team_ids=()):
|
||||
return (' AND "user" = $2', (user_id,))
|
||||
case _SpendLogViewer(user_id=user_id, team_ids=team_ids):
|
||||
return (' AND ("user" = $2 OR team_id = ANY($3::text[]))', (user_id, team_ids))
|
||||
|
||||
|
||||
def _spend_log_payload_query(request_id: str, viewer: _SpendLogViewer | None) -> tuple[str, tuple[object, ...]]:
|
||||
"""
|
||||
Fetch the one row an id lookup resolves to, preferring the exact ``request_id``
|
||||
match over rows that merely carry the id as their client-set ``litellm_call_id``.
|
||||
A non-admin viewer only ever gets rows they own or rows of a team they may view.
|
||||
"""
|
||||
scope, scope_params = _viewer_scope_clause(viewer)
|
||||
return (
|
||||
f"""
|
||||
SELECT request_id, messages, response, proxy_server_request, metadata, "user", team_id
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE (request_id = $1 OR litellm_call_id = $1){scope}
|
||||
ORDER BY (request_id = $1) DESC
|
||||
LIMIT 1
|
||||
""",
|
||||
(request_id, *scope_params),
|
||||
)
|
||||
|
||||
|
||||
async def _resolve_spend_log_payload_row(
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_id: str,
|
||||
caller_is_admin: bool,
|
||||
) -> Mapping[str, object] | None:
|
||||
"""
|
||||
Resolve an id lookup to the caller's own spend-log row before any payload
|
||||
store is consulted. Cold storage is keyed by the provider ``request_id``, so
|
||||
asking it for the raw lookup id could hand back another tenant's payload when
|
||||
that id is only the caller's ``litellm_call_id``; the row's stored
|
||||
``request_id`` is the key that names the caller's own request.
|
||||
"""
|
||||
viewer: Final = None if caller_is_admin else await _spend_log_viewer(prisma_client, user_api_key_dict)
|
||||
sql_query, sql_params = _spend_log_payload_query(request_id, viewer)
|
||||
rows: Final[Sequence[Mapping[str, object]] | None] = await _query_raw_or_none(prisma_client, sql_query, *sql_params)
|
||||
if not rows:
|
||||
return None
|
||||
if not caller_is_admin:
|
||||
await _assert_user_owns_fetched_spend_rows(
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
rows=rows,
|
||||
request_id=request_id,
|
||||
)
|
||||
return rows[0]
|
||||
|
||||
|
||||
def _stored_request_id(row: Mapping[str, object] | None, lookup_id: str) -> str:
|
||||
stored: Final = None if row is None else row.get("request_id")
|
||||
return stored if isinstance(stored, str) else lookup_id
|
||||
|
||||
|
||||
def _fetched_row_owner(row: Mapping[str, object]) -> tuple[str | None, str | None]:
|
||||
user: Final = row.get("user")
|
||||
team_id: Final = row.get("team_id")
|
||||
return (
|
||||
user if isinstance(user, str) else None,
|
||||
team_id if isinstance(team_id, str) else None,
|
||||
)
|
||||
|
||||
|
||||
async def _assert_user_owns_fetched_spend_rows(
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
rows: Sequence[Mapping[str, object]],
|
||||
request_id: str,
|
||||
) -> None:
|
||||
"""
|
||||
Re-verify ownership on the rows an id lookup actually fetched.
|
||||
``_assert_user_can_view_request_id`` and the data query read the table at
|
||||
different moments, so a foreign row inserted between them could otherwise be
|
||||
returned even though the pre-check passed. Checking the fetched rows
|
||||
themselves means no interleaving can return another tenant's row.
|
||||
"""
|
||||
for user, team_id in frozenset(_fetched_row_owner(row) for row in rows):
|
||||
if not await _user_can_view_spend_log_owner(prisma_client, user_api_key_dict, user, team_id):
|
||||
raise _spend_log_forbidden(request_id)
|
||||
|
||||
|
||||
def _cold_storage_payload_owner(payload: Mapping[str, object]) -> tuple[str | None, str | None]:
|
||||
metadata: Final = payload.get("metadata")
|
||||
if not isinstance(metadata, Mapping):
|
||||
return (None, None)
|
||||
owner: Final = cast(Mapping[str, object], metadata) # cast-ok: cold-storage JSON is untyped
|
||||
user: Final = owner.get("user_api_key_user_id")
|
||||
team_id: Final = owner.get("user_api_key_team_id")
|
||||
return (
|
||||
user if isinstance(user, str) else None,
|
||||
team_id if isinstance(team_id, str) else None,
|
||||
)
|
||||
|
||||
|
||||
async def _assert_user_owns_cold_storage_payload(
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
payload: Mapping[str, object],
|
||||
request_id: str,
|
||||
) -> None:
|
||||
"""
|
||||
Authorize a cold-storage payload against the owner recorded inside it.
|
||||
The custom logger reads the payload straight from cold storage, written
|
||||
independently of the spend-log table and able to outlive its row, so a
|
||||
request_id lookup could otherwise hand back another tenant's stored payload
|
||||
when no row exists for the pre-check to catch. Verifying the payload's own
|
||||
owner closes that gap, and a payload that records no owner fails closed.
|
||||
"""
|
||||
owner_user, owner_team_id = _cold_storage_payload_owner(payload)
|
||||
if not await _user_can_view_spend_log_owner(prisma_client, user_api_key_dict, owner_user, owner_team_id):
|
||||
raise _spend_log_forbidden(request_id)
|
||||
|
||||
|
||||
async def _get_permitted_team_ids_for_spend_logs(
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
|
|||
|
|
@ -621,6 +621,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
|
|||
status=_get_status_for_spend_log(
|
||||
metadata=metadata,
|
||||
),
|
||||
litellm_call_id=litellm_call_id,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
|
|||
|
|
@ -268,6 +268,12 @@ class CredentialLiteLLMParams(BaseModel):
|
|||
# callers see it, breaking Azure deployments configured with
|
||||
# ``azure_ad_token`` instead of a static ``api_key`` (#30235).
|
||||
azure_ad_token: str | None = None
|
||||
tenant_id: str | None = None
|
||||
client_id: str | None = None
|
||||
client_secret: str | None = None
|
||||
azure_scope: str | None = None
|
||||
azure_username: str | None = None
|
||||
azure_password: str | None = None
|
||||
## VERTEX AI ##
|
||||
vertex_project: str | None = None
|
||||
vertex_location: str | None = None
|
||||
|
|
|
|||
|
|
@ -23083,58 +23083,206 @@
|
|||
},
|
||||
"friendliai/zai-org/GLM-5.3-Flash": {
|
||||
"litellm_provider": "friendliai",
|
||||
"supports_reasoning": true,
|
||||
"supports_function_calling": true,
|
||||
"max_input_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"max_output_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"output_cost_per_token": 5e-07,
|
||||
"cache_read_input_token_cost": 3e-08,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
"max"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"mode": "chat",
|
||||
"comment": "Native multimodal GLM model for efficient coding and long-horizon agent tasks",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models",
|
||||
"supports_vision": true,
|
||||
"supports_image_input": true,
|
||||
"supports_video_input": true
|
||||
"supports_video_input": true,
|
||||
"mode": "chat",
|
||||
"comment": "Native multimodal GLM model for efficient coding and long-horizon agent tasks",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"friendliai/zai-org/GLM-5.3": {
|
||||
"litellm_provider": "friendliai",
|
||||
"supports_reasoning": true,
|
||||
"supports_function_calling": true,
|
||||
"max_input_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"max_output_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"input_cost_per_token": 1.26e-06,
|
||||
"output_cost_per_token": 3.96e-06,
|
||||
"cache_read_input_token_cost": 2.34e-07,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [
|
||||
"low",
|
||||
"high",
|
||||
"max"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"supports_image_input": false,
|
||||
"supports_video_input": false,
|
||||
"mode": "chat",
|
||||
"comment": "Flagship GLM model for long-horizon coding, agents, and complex project delivery",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"friendliai/google/gemma-4-31B-it": {
|
||||
"litellm_provider": "friendliai",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"input_cost_per_token": 1.4e-07,
|
||||
"output_cost_per_token": 4e-07,
|
||||
"supports_prompt_caching": false,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_image_input": true,
|
||||
"supports_video_input": false,
|
||||
"mode": "chat",
|
||||
"comment": "Largest Gemma 4 instruction model for open, self-hosted chat and reasoning",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"friendliai/zai-org/GLM-5.2": {
|
||||
"litellm_provider": "friendliai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"input_cost_per_token": 1.4e-06,
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"cache_read_input_token_cost": 2.6e-07,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [
|
||||
"high",
|
||||
"max"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"supports_image_input": false
|
||||
"supports_image_input": false,
|
||||
"supports_video_input": false,
|
||||
"mode": "chat",
|
||||
"comment": "Open flagship GLM for long-horizon coding agents and million-token context work",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"friendliai/LGAI-EXAONE/K-EXAONE-2.0-750B-A37B": {
|
||||
"litellm_provider": "friendliai",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"input_cost_per_token": 6e-07,
|
||||
"output_cost_per_token": 2.4e-06,
|
||||
"cache_read_input_token_cost": 1.2e-07,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"supports_image_input": false,
|
||||
"supports_video_input": false,
|
||||
"mode": "chat",
|
||||
"comment": "Frontier-scale multilingual language model developed by LG AI Research",
|
||||
"deprecation_date": "2026-09-06",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"friendliai/deepseek-ai/DeepSeek-V3.2": {
|
||||
"litellm_provider": "friendliai",
|
||||
"max_input_tokens": 163840,
|
||||
"max_output_tokens": 163840,
|
||||
"max_tokens": 163840,
|
||||
"input_cost_per_token": 5e-07,
|
||||
"output_cost_per_token": 1.5e-06,
|
||||
"cache_read_input_token_cost": 2.5e-07,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"supports_image_input": false,
|
||||
"supports_video_input": false,
|
||||
"mode": "chat",
|
||||
"comment": "DeepSeek chat model for instruction following, coding, and analysis",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"friendliai/MiniMaxAI/MiniMax-M2.5": {
|
||||
"litellm_provider": "friendliai",
|
||||
"max_input_tokens": 196608,
|
||||
"max_output_tokens": 196608,
|
||||
"max_tokens": 196608,
|
||||
"input_cost_per_token": 3e-07,
|
||||
"output_cost_per_token": 1.2e-06,
|
||||
"cache_read_input_token_cost": 6e-08,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"supports_image_input": false,
|
||||
"supports_video_input": false,
|
||||
"mode": "chat",
|
||||
"comment": "Prior MiniMax coding model for agent workflows, office edits, and automation",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"friendliai/zai-org/GLM-5.1": {
|
||||
"litellm_provider": "friendliai",
|
||||
"max_input_tokens": 202752,
|
||||
"max_output_tokens": 202752,
|
||||
"max_tokens": 202752,
|
||||
"input_cost_per_token": 1.4e-06,
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"cache_read_input_token_cost": 2.6e-07,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"reasoning_effort_levels": [],
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false,
|
||||
"supports_image_input": false,
|
||||
"supports_video_input": false,
|
||||
"mode": "chat",
|
||||
"comment": "Strong GLM coding model for agentic engineering, terminals, and repository generation",
|
||||
"source": "https://api.friendli.ai/serverless/v1/models"
|
||||
},
|
||||
"ft:babbage-002": {
|
||||
"deprecation_date": "2026-10-23",
|
||||
|
|
|
|||
|
|
@ -659,12 +659,14 @@ model LiteLLM_SpendLogs {
|
|||
mcp_namespaced_tool_name String?
|
||||
agent_id String?
|
||||
proxy_server_request Json? @default("{}")
|
||||
litellm_call_id String?
|
||||
created_at DateTime @default(now()) @map("created_at")
|
||||
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
|
||||
@@index([startTime])
|
||||
@@index([startTime, request_id])
|
||||
@@index([end_user])
|
||||
@@index([session_id])
|
||||
@@index([litellm_call_id])
|
||||
}
|
||||
|
||||
model LiteLLM_BudgetWindowSpend {
|
||||
|
|
|
|||
|
|
@ -69,7 +69,7 @@ Each suite provides its own `client` fixture (see `llm_translation/passthrough_c
|
|||
|
||||
Request and response bodies are typed pydantic models in `models.py`; only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass
|
||||
|
||||
Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure coverage of the harness itself carries no marker and runs regardless. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache
|
||||
Mark live tests with `@pytest.mark.e2e` (on the class or the module). Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache
|
||||
|
||||
## Record and replay fixtures
|
||||
|
||||
|
|
@ -221,7 +221,7 @@ other.<area>.<case>.<assertion>
|
|||
```
|
||||
|
||||
## Hard Rules
|
||||
- no monkeypatching or mock tests, and never substitute a unit test for e2e feature coverage: a product feature is proven end to end against a live proxy, not with a unit test. if a contributor asks you to write an end to end test, do NOT stage a unit test of the feature with it; if you find a product gap, call it out in the PR description. tests that cover the harness itself are the exception and are allowed (for example `coverage_registry/test_collector.py`, which unit-tests the coverage collector): they carry no `e2e` marker, exercise harness plumbing rather than a product feature, and run whether or not a proxy is up
|
||||
- no unit tests of any kind under `tests/e2e`. a product feature is proven end to end against a live proxy, never with a unit test, and the harness itself is not unit-tested here either. no monkeypatching or mock tests. if a contributor asks you to write an end to end test, do NOT stage a unit test with it; if you find a product gap, call it out in the PR description
|
||||
|
||||
- use model management endpoints to create new models for a test. this could be in a conftest / inline for each test. ask the user what they want.
|
||||
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ verbose_logger.setLevel(logging.DEBUG)
|
|||
|
||||
ignored_keys = [
|
||||
"request_id",
|
||||
"litellm_call_id",
|
||||
"metadata.litellm_call_id",
|
||||
"session_id",
|
||||
"startTime",
|
||||
|
|
|
|||
|
|
@ -1193,7 +1193,7 @@ async def test_afile_delete_returns_managed_id_for_stored_provider_output():
|
|||
|
||||
assert response.id == unified_file_id
|
||||
assert response.object == "file"
|
||||
assert response.filename == stored_file.filename
|
||||
assert response.deleted is True
|
||||
assert stored_file.id == provider_file_id
|
||||
router.afile_delete.assert_awaited_once_with(model="model-123", file_id=provider_file_id)
|
||||
table.delete.assert_awaited_once_with(where={"unified_file_id": unified_file_id})
|
||||
|
|
@ -1730,3 +1730,100 @@ async def test_batch_retrieve_hook_does_not_claim_attribution():
|
|||
|
||||
managed_files.store_unified_object_id.assert_awaited_once()
|
||||
assert managed_files.store_unified_object_id.await_args.kwargs["persist_attribution"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_delete_passes_trusted_model_credentials_to_router():
|
||||
"""
|
||||
afile_delete must hand the deployment's credential snapshot to the router
|
||||
call, since Bedrock validates the s3:// file id against the bucket in it.
|
||||
"""
|
||||
from types import MappingProxyType
|
||||
|
||||
managed_files = _make_managed_files_instance()
|
||||
unified_file_id = "unified-file-id"
|
||||
s3_uri = "s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl"
|
||||
managed_files.get_model_file_id_mapping = AsyncMock(return_value={unified_file_id: {"model-123": s3_uri}})
|
||||
managed_files.delete_unified_file_id = AsyncMock(return_value=_make_file_object(unified_file_id))
|
||||
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_deployment_credentials_with_provider = MagicMock(
|
||||
return_value={
|
||||
"custom_llm_provider": "bedrock",
|
||||
"s3_bucket_name": "my-bucket",
|
||||
"aws_region_name": "us-west-2",
|
||||
}
|
||||
)
|
||||
mock_router.afile_delete = AsyncMock(return_value=MagicMock())
|
||||
|
||||
await managed_files.afile_delete(
|
||||
file_id=unified_file_id,
|
||||
litellm_parent_otel_span=None,
|
||||
llm_router=mock_router,
|
||||
)
|
||||
|
||||
call_kwargs = mock_router.afile_delete.call_args.kwargs
|
||||
assert call_kwargs["model"] == "model-123"
|
||||
assert call_kwargs["file_id"] == s3_uri
|
||||
trusted_credentials = call_kwargs["_litellm_internal_model_credentials"]
|
||||
assert isinstance(trusted_credentials, MappingProxyType)
|
||||
assert trusted_credentials["s3_bucket_name"] == "my-bucket"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_delete_bedrock_unified_id_end_to_end(monkeypatch):
|
||||
"""
|
||||
Proxy repro for deleting a Bedrock batch input file by unified id: the
|
||||
s3:// object must be removed via a SigV4-signed S3 DELETE using the
|
||||
deployment's s3_bucket_name (no AWS_S3_BUCKET_NAME env).
|
||||
|
||||
Regression test for "BedrockFilesConfig does not support file deletion"
|
||||
raised on this path.
|
||||
"""
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
|
||||
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "bedrock-claude",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "secret",
|
||||
"aws_region_name": "us-west-2",
|
||||
"s3_bucket_name": "my-bucket",
|
||||
},
|
||||
"model_info": {"id": "model-123"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
managed_files = _make_managed_files_instance()
|
||||
unified_file_id = "unified-file-id"
|
||||
s3_uri = "s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl"
|
||||
managed_files.get_model_file_id_mapping = AsyncMock(return_value={unified_file_id: {"model-123": s3_uri}})
|
||||
managed_files.delete_unified_file_id = AsyncMock(return_value=_make_file_object(unified_file_id))
|
||||
|
||||
expected_url = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files/job-123/input.jsonl"
|
||||
with respx.mock:
|
||||
route = respx.delete(expected_url).mock(return_value=httpx.Response(204))
|
||||
|
||||
response = await managed_files.afile_delete(
|
||||
file_id=unified_file_id,
|
||||
litellm_parent_otel_span=None,
|
||||
llm_router=router,
|
||||
)
|
||||
|
||||
assert route.called
|
||||
assert route.calls[0].request.headers["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
assert response.id == unified_file_id
|
||||
assert response.model_dump() == {"id": unified_file_id, "object": "file", "deleted": True}
|
||||
managed_files.delete_unified_file_id.assert_awaited_once_with(unified_file_id, None)
|
||||
|
|
|
|||
|
|
@ -5890,6 +5890,81 @@ def test_resolve_vertex_location_for_cost_default_region(monkeypatch):
|
|||
assert _resolve("vertex_ai", None, None, "gemini-3.5-flash") == "us-central1"
|
||||
|
||||
|
||||
def test_resolve_mantle_region_for_cost(monkeypatch):
|
||||
"""Bedrock Mantle requests resolve the served region the way dispatch does (explicit
|
||||
aws_region_name, then the api_base host, then the default); other providers get None."""
|
||||
from litellm.litellm_core_utils.litellm_logging import _resolve_mantle_region_for_cost
|
||||
|
||||
for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION_NAME", "AWS_REGION"):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
|
||||
assert _resolve_mantle_region_for_cost("bedrock", {"aws_region_name": "us-gov-west-1"}) is None
|
||||
assert _resolve_mantle_region_for_cost(None, {"aws_region_name": "us-gov-west-1"}) is None
|
||||
assert _resolve_mantle_region_for_cost("bedrock_mantle", {"aws_region_name": "us-gov-west-1"}) == "us-gov-west-1"
|
||||
assert (
|
||||
_resolve_mantle_region_for_cost(
|
||||
"bedrock_mantle",
|
||||
{"api_base": "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1/chat/completions"},
|
||||
)
|
||||
== "us-gov-west-1"
|
||||
)
|
||||
assert _resolve_mantle_region_for_cost("bedrock_mantle", None) == "us-east-1"
|
||||
|
||||
|
||||
def test_response_cost_calculator_prices_mantle_calls_on_the_served_region(monkeypatch):
|
||||
"""
|
||||
Mantle responses carry no region of their own (the OpenAI-compatible transform rebuilds the
|
||||
response, and streams never had one), so the logging layer must price them from the region
|
||||
the deployment was served in: an explicit aws_region_name or the api_base host, both of which
|
||||
must select the GovCloud row over the commercial one.
|
||||
"""
|
||||
from datetime import datetime
|
||||
|
||||
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
|
||||
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url=""))
|
||||
for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION_NAME", "AWS_REGION"):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
|
||||
def cost_with(litellm_params):
|
||||
logging_obj = LitellmLogging(
|
||||
model="xai.grok-4.3",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=False,
|
||||
call_type="completion",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="mantle-region",
|
||||
function_id="f",
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model="xai.grok-4.3",
|
||||
user="",
|
||||
optional_params={},
|
||||
litellm_params=litellm_params,
|
||||
custom_llm_provider="bedrock_mantle",
|
||||
)
|
||||
response = ModelResponse(
|
||||
id="resp-1",
|
||||
model="xai.grok-4.3",
|
||||
choices=[{"message": {"role": "assistant", "content": "hello"}, "index": 0, "finish_reason": "stop"}],
|
||||
usage={"prompt_tokens": 38, "completion_tokens": 20, "total_tokens": 58},
|
||||
)
|
||||
return logging_obj._response_cost_calculator(result=response)
|
||||
|
||||
commercial = litellm.model_cost["bedrock_mantle/xai.grok-4.3"]
|
||||
gov = litellm.model_cost["bedrock_mantle/us-gov-west-1/xai.grok-4.3"]
|
||||
expected_commercial = 38 * commercial["input_cost_per_token"] + 20 * commercial["output_cost_per_token"]
|
||||
expected_gov = 38 * gov["input_cost_per_token"] + 20 * gov["output_cost_per_token"]
|
||||
assert expected_gov != expected_commercial
|
||||
|
||||
assert cost_with({"api_base": ""}) == pytest.approx(expected_commercial)
|
||||
assert cost_with({"aws_region_name": "us-gov-west-1"}) == pytest.approx(expected_gov)
|
||||
assert cost_with(
|
||||
{"api_base": "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1/chat/completions"}
|
||||
) == pytest.approx(expected_gov)
|
||||
|
||||
|
||||
def test_response_cost_calculator_prices_proxy_vertex_calls_on_the_configured_location(monkeypatch):
|
||||
"""
|
||||
Proxy-shaped logging objects (created before the router picks a deployment) carry the
|
||||
|
|
|
|||
|
|
@ -0,0 +1,142 @@
|
|||
"""
|
||||
Regression tests for the ``/v1/messages`` async adapter dropping the socket on a
|
||||
mid-stream provider error.
|
||||
|
||||
When a non-Anthropic model (e.g. Bedrock Converse) is served through
|
||||
``/v1/messages``, the proxy hands Starlette the async SSE iterator directly. If
|
||||
the upstream provider stream raises while being pulled (Bedrock raises
|
||||
``BedrockError`` when a ConverseStream ends without a terminal ``messageStop``
|
||||
event, common on cross-region inference profiles), the exception escaped the
|
||||
request handler's try/except and tore down the connection. Clients like Claude
|
||||
Code then showed a bare "Connection closed mid-response".
|
||||
|
||||
The async SSE wrapper must instead surface the failure as a well-formed
|
||||
Anthropic ``error`` event so the stream stays valid and the client can retry.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from typing import List, Optional
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../../.."))
|
||||
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import (
|
||||
AnthropicStreamWrapper,
|
||||
_mid_stream_error_sse_event,
|
||||
)
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.types.utils import Delta, StreamingChoices
|
||||
|
||||
|
||||
def _make_chunk(delta: Delta, finish_reason: Optional[str] = None) -> MagicMock:
|
||||
chunk = MagicMock()
|
||||
chunk.choices = [
|
||||
StreamingChoices(finish_reason=finish_reason, index=0, delta=delta, logprobs=None)
|
||||
]
|
||||
chunk.usage = None
|
||||
chunk._hidden_params = {}
|
||||
return chunk
|
||||
|
||||
|
||||
class _AsyncStreamThenRaise:
|
||||
"""Yields the given chunks, then raises ``exc`` (mimics a provider stream
|
||||
that terminates mid-response)."""
|
||||
|
||||
def __init__(self, items: List[MagicMock], exc: BaseException):
|
||||
self._it = iter(items)
|
||||
self._exc = exc
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
try:
|
||||
return next(self._it)
|
||||
except StopIteration:
|
||||
raise self._exc
|
||||
|
||||
|
||||
def _parse_sse(raw: bytes) -> tuple[str, dict]:
|
||||
text = raw.decode()
|
||||
event_line, data_line = text.strip().split("\n", 1)
|
||||
return event_line.removeprefix("event: "), json.loads(data_line.removeprefix("data: "))
|
||||
|
||||
|
||||
async def _drain_sse(wrapper: AnthropicStreamWrapper) -> List[bytes]:
|
||||
return [event async for event in wrapper.async_anthropic_sse_wrapper()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_bedrock_error_becomes_anthropic_error_event():
|
||||
"""A ``BedrockError`` raised after partial content must be surfaced as a
|
||||
terminal Anthropic ``error`` event, not propagated (which drops the socket
|
||||
and yields "Connection closed mid-response")."""
|
||||
chunks = [_make_chunk(Delta(content="Creating a file"))]
|
||||
bedrock_err = BedrockError(
|
||||
status_code=500,
|
||||
message="Bedrock ConverseStream ended without a terminal 'messageStop' event",
|
||||
)
|
||||
wrapper = AnthropicStreamWrapper(
|
||||
completion_stream=_AsyncStreamThenRaise(chunks, bedrock_err),
|
||||
model="bedrock-converse-sonnet-4-6",
|
||||
)
|
||||
|
||||
events = await _drain_sse(wrapper)
|
||||
|
||||
parsed = [_parse_sse(e) for e in events]
|
||||
event_types = [name for name, _ in parsed]
|
||||
assert "message_start" in event_types
|
||||
assert event_types[-1] == "error"
|
||||
_, error_payload = parsed[-1]
|
||||
assert error_payload["type"] == "error"
|
||||
assert error_payload["error"]["type"] == "api_error"
|
||||
assert "messageStop" in error_payload["error"]["message"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_error_does_not_raise_out_of_wrapper():
|
||||
"""The async wrapper must fully drain without letting the upstream exception
|
||||
escape — escaping is exactly what tore down the connection before the fix."""
|
||||
wrapper = AnthropicStreamWrapper(
|
||||
completion_stream=_AsyncStreamThenRaise([], BedrockError(status_code=500, message="boom")),
|
||||
model="claude-x",
|
||||
)
|
||||
events = await _drain_sse(wrapper)
|
||||
assert _parse_sse(events[-1])[0] == "error"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status_code, expected_type",
|
||||
[(500, "api_error"), (529, "overloaded_error"), (429, "rate_limit_error")],
|
||||
)
|
||||
def test_error_event_maps_status_code_to_anthropic_type(status_code, expected_type):
|
||||
raw = _mid_stream_error_sse_event(BedrockError(status_code=status_code, message="upstream failed"))
|
||||
name, payload = _parse_sse(raw)
|
||||
assert name == "error"
|
||||
assert payload["error"]["type"] == expected_type
|
||||
assert payload["error"]["message"] == "upstream failed"
|
||||
|
||||
|
||||
def test_error_event_defaults_to_500_when_status_missing():
|
||||
raw = _mid_stream_error_sse_event(ValueError("no status here"))
|
||||
_, payload = _parse_sse(raw)
|
||||
assert payload["error"]["type"] == "api_error"
|
||||
assert payload["error"]["message"] == "no status here"
|
||||
|
||||
|
||||
def test_error_event_preserves_midstream_fallback_error():
|
||||
exc = MidStreamFallbackError(
|
||||
message="BedrockException - internalServerException",
|
||||
model="bedrock-converse-sonnet-4-6",
|
||||
llm_provider="bedrock",
|
||||
original_exception=BedrockError(status_code=500, message="internalServerException"),
|
||||
)
|
||||
name, payload = _parse_sse(_mid_stream_error_sse_event(exc))
|
||||
assert name == "error"
|
||||
assert payload["error"]["type"] == "api_error"
|
||||
assert "internalServerException" in payload["error"]["message"]
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -7,7 +7,7 @@ API docs: https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.ht
|
|||
|
||||
import json
|
||||
import asyncio
|
||||
from unittest.mock import patch
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
|
||||
import httpx
|
||||
|
|
@ -151,6 +151,27 @@ class TestBedrockMantleConfig:
|
|||
# /openai/v1 base per the AWS model card.
|
||||
assert api_base == "https://bedrock-mantle.us-east-2.api.aws/openai/v1"
|
||||
|
||||
def test_region_prefixed_model_routes_to_that_region(self, monkeypatch, local_cost_map):
|
||||
for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION", "AWS_REGION_NAME"):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
cfg = BedrockMantleChatConfig()
|
||||
api_base, _ = cfg._get_openai_compatible_provider_info(None, None, model="us-gov-west-1/xai.grok-4.3")
|
||||
assert api_base == "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1"
|
||||
|
||||
def test_aws_region_name_param_beats_model_region_prefix(self, monkeypatch, local_cost_map):
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION", "AWS_REGION_NAME"):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
cfg = BedrockMantleChatConfig()
|
||||
api_base, _ = cfg._get_openai_compatible_provider_info(
|
||||
None,
|
||||
None,
|
||||
litellm_params=GenericLiteLLMParams(aws_region_name="us-east-1"),
|
||||
model="us-gov-west-1/xai.grok-4.3",
|
||||
)
|
||||
assert api_base == "https://bedrock-mantle.us-east-1.api.aws/openai/v1"
|
||||
|
||||
def test_default_api_base_fallback_to_us_east_1(self, monkeypatch):
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
|
||||
|
|
@ -686,6 +707,120 @@ class TestBedrockMantleProviderResolution:
|
|||
assert model == "openai.gpt-oss-20b"
|
||||
|
||||
|
||||
def test_get_llm_provider_strips_region_prefix(self, monkeypatch, local_cost_map):
|
||||
for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION", "AWS_REGION_NAME"):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
model, provider, _, api_base = litellm.get_llm_provider("bedrock_mantle/us-gov-west-1/xai.grok-4.3")
|
||||
assert provider == "bedrock_mantle"
|
||||
assert model == "xai.grok-4.3"
|
||||
assert api_base == "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1"
|
||||
|
||||
def test_completion_region_prefixed_model_sends_bare_model_to_that_region(self, monkeypatch, local_cost_map):
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
for var in (
|
||||
"BEDROCK_MANTLE_API_KEY",
|
||||
"AWS_BEARER_TOKEN_BEDROCK",
|
||||
"BEDROCK_MANTLE_API_BASE",
|
||||
"BEDROCK_MANTLE_REGION",
|
||||
"AWS_REGION_NAME",
|
||||
"AWS_REGION",
|
||||
"AWS_PROFILE",
|
||||
):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIAEXAMPLE")
|
||||
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0")
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
status_code=200,
|
||||
json={
|
||||
"id": "chatcmpl-test",
|
||||
"object": "chat.completion",
|
||||
"created": 1733529600,
|
||||
"model": "xai.grok-4.3",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 38, "completion_tokens": 20, "total_tokens": 58},
|
||||
},
|
||||
request=request,
|
||||
)
|
||||
|
||||
handler = Mock(side_effect=respond)
|
||||
response = litellm.completion(
|
||||
model="bedrock_mantle/us-gov-west-1/xai.grok-4.3",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))),
|
||||
)
|
||||
|
||||
gov = litellm.model_cost["bedrock_mantle/us-gov-west-1/xai.grok-4.3"]
|
||||
sent = handler.call_args.args[0]
|
||||
assert str(sent.url) == "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1/chat/completions"
|
||||
assert json.loads(sent.content)["model"] == "xai.grok-4.3"
|
||||
assert "/us-gov-west-1/bedrock/aws4_request" in sent.headers["Authorization"]
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(
|
||||
38 * gov["input_cost_per_token"] + 20 * gov["output_cost_per_token"]
|
||||
)
|
||||
|
||||
def test_responses_region_prefixed_model_prices_from_that_region_over_env_region(self, monkeypatch, local_cost_map):
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
for var in (
|
||||
"BEDROCK_MANTLE_API_KEY",
|
||||
"AWS_BEARER_TOKEN_BEDROCK",
|
||||
"BEDROCK_MANTLE_API_BASE",
|
||||
"BEDROCK_MANTLE_REGION",
|
||||
"AWS_REGION",
|
||||
"AWS_PROFILE",
|
||||
):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
|
||||
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIAEXAMPLE")
|
||||
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0")
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
status_code=200,
|
||||
json={
|
||||
"id": "resp_test",
|
||||
"object": "response",
|
||||
"created_at": 1733529600,
|
||||
"status": "completed",
|
||||
"model": "xai.grok-4.3",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_test",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "ok", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"parallel_tool_calls": True,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
"top_p": 1.0,
|
||||
"usage": {"input_tokens": 38, "output_tokens": 20, "total_tokens": 58},
|
||||
},
|
||||
request=request,
|
||||
)
|
||||
|
||||
handler = Mock(side_effect=respond)
|
||||
response = litellm.responses(
|
||||
model="bedrock_mantle/us-gov-west-1/xai.grok-4.3",
|
||||
input="hello",
|
||||
client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))),
|
||||
)
|
||||
|
||||
gov = litellm.model_cost["bedrock_mantle/us-gov-west-1/xai.grok-4.3"]
|
||||
sent = handler.call_args.args[0]
|
||||
assert str(sent.url) == "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1/responses"
|
||||
assert json.loads(sent.content)["model"] == "xai.grok-4.3"
|
||||
assert "/us-gov-west-1/bedrock/aws4_request" in sent.headers["Authorization"]
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(
|
||||
38 * gov["input_cost_per_token"] + 20 * gov["output_cost_per_token"]
|
||||
)
|
||||
|
||||
|
||||
class TestBedrockMantlePricing:
|
||||
"""Tests that verify Bedrock Mantle uses correct AWS Bedrock pricing, not OpenAI pricing."""
|
||||
|
||||
|
|
|
|||
|
|
@ -6390,6 +6390,106 @@ async def test_get_team_membership_db_fetch_returns_validated_membership():
|
|||
assert result.spend == 1.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_team_membership_negative_caches_a_missing_row():
|
||||
"""
|
||||
Regression (LIT-7358): a member with no LiteLLM_TeamMembership row is the common lite-login case,
|
||||
and the session-token refresh reads this loader on every request. Before the fix a missing row
|
||||
returned None without caching, so every request re-queried the DB. The miss must be cached so the
|
||||
second request serves from cache and never touches the DB.
|
||||
"""
|
||||
from litellm.proxy.auth.auth_checks import get_team_membership
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
NO_TEAM_MEMBERSHIP_SENTINEL,
|
||||
team_membership_reservation_cache_key,
|
||||
)
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
|
||||
first = await get_team_membership(
|
||||
user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache
|
||||
)
|
||||
second = await get_team_membership(
|
||||
user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache
|
||||
)
|
||||
|
||||
assert first is None
|
||||
assert second is None
|
||||
mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once()
|
||||
cached = await cache.async_get_cache(
|
||||
key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")
|
||||
)
|
||||
assert cached == NO_TEAM_MEMBERSHIP_SENTINEL
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_team_membership_reads_sentinel_as_no_membership_not_a_model():
|
||||
"""
|
||||
The negative-cache sentinel is a plain string sharing the key a serialized membership uses.
|
||||
A pre-seeded sentinel must read back as None (no DB read), never be mistaken for a membership.
|
||||
"""
|
||||
from litellm.proxy.auth.auth_checks import get_team_membership
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
NO_TEAM_MEMBERSHIP_SENTINEL,
|
||||
team_membership_reservation_cache_key,
|
||||
)
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
await cache.async_set_cache(
|
||||
key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1"),
|
||||
value=NO_TEAM_MEMBERSHIP_SENTINEL,
|
||||
)
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
result = await get_team_membership(
|
||||
user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache
|
||||
)
|
||||
|
||||
assert result is None
|
||||
mock_prisma_client.db.litellm_teammembership.find_unique.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sentinel():
|
||||
"""
|
||||
A member who later gains a per-member budget writes a membership row and calls
|
||||
invalidate_team_member_spend_state. That must drop a cached "no membership" sentinel so the next
|
||||
request re-reads the DB and honors the new budget instead of serving the stale miss until TTL.
|
||||
"""
|
||||
from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state
|
||||
from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
membership_row = MagicMock()
|
||||
membership_row.dict = lambda: {"user_id": "u-1", "team_id": "t-1", "spend": 0.0}
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=[None, membership_row])
|
||||
|
||||
before = await get_team_membership(
|
||||
user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache
|
||||
)
|
||||
assert before is None
|
||||
|
||||
await invalidate_team_member_spend_state(user_id="u-1", team_id="t-1", user_api_key_cache=cache)
|
||||
assert (
|
||||
await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1"))
|
||||
is None
|
||||
)
|
||||
|
||||
after = await get_team_membership(
|
||||
user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache
|
||||
)
|
||||
assert after is not None
|
||||
assert after.user_id == "u-1"
|
||||
assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_access_object_db_fetch_returns_validated_access_group():
|
||||
from litellm.proxy._types import LiteLLM_AccessGroupTable
|
||||
|
|
|
|||
201
tests/test_litellm/proxy/auth/test_resolvers_grants.py
Normal file
201
tests/test_litellm/proxy/auth/test_resolvers_grants.py
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
from fastapi import HTTPException
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_TeamMembership,
|
||||
LiteLLM_TeamTableCachedObj,
|
||||
LiteLLM_UserTable,
|
||||
ProxyException,
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import TeamNotFoundError, UserNotFoundError
|
||||
from litellm.proxy.auth.resolvers.grants import (
|
||||
GrantResolver,
|
||||
LookupDegraded,
|
||||
NotAMember,
|
||||
ResolvedGrants,
|
||||
TeamGone,
|
||||
UserGone,
|
||||
UserLookup,
|
||||
raise_public,
|
||||
user_models,
|
||||
)
|
||||
|
||||
USER_ID = "user-1"
|
||||
TEAM_ID = "team-1"
|
||||
|
||||
|
||||
class _Loaders:
|
||||
"""Fake row readers standing in for the ``auth_checks`` loaders, recording every call they receive."""
|
||||
|
||||
def __init__(self, *, user=None, team=None, membership=None, user_error=None, team_error=None):
|
||||
self._user = user
|
||||
self._team = team
|
||||
self._membership = membership
|
||||
self._user_error = user_error
|
||||
self._team_error = team_error
|
||||
self.user_calls = []
|
||||
self.team_calls = []
|
||||
self.membership_calls = []
|
||||
|
||||
async def load_user(self, **kwargs):
|
||||
self.user_calls.append(kwargs)
|
||||
if self._user_error is not None:
|
||||
raise self._user_error
|
||||
return self._user
|
||||
|
||||
async def load_team(self, **kwargs):
|
||||
self.team_calls.append(kwargs)
|
||||
if self._team_error is not None:
|
||||
raise self._team_error
|
||||
return self._team
|
||||
|
||||
async def load_membership(self, **kwargs):
|
||||
self.membership_calls.append(kwargs)
|
||||
return self._membership
|
||||
|
||||
def resolver(self) -> GrantResolver:
|
||||
return GrantResolver(
|
||||
object(),
|
||||
object(),
|
||||
load_user=self.load_user,
|
||||
load_team=self.load_team,
|
||||
load_membership=self.load_membership,
|
||||
)
|
||||
|
||||
|
||||
def _user(teams=(TEAM_ID,), user_id=USER_ID) -> LiteLLM_UserTable:
|
||||
return LiteLLM_UserTable(user_id=user_id, user_role="internal_user", teams=list(teams), models=["gpt-5.5"])
|
||||
|
||||
|
||||
def _team(models=("gpt-5.5",)) -> LiteLLM_TeamTableCachedObj:
|
||||
return LiteLLM_TeamTableCachedObj(team_id=TEAM_ID, team_alias="alias", models=list(models))
|
||||
|
||||
|
||||
async def test_resolve_returns_live_rows_for_a_member():
|
||||
membership = LiteLLM_TeamMembership(user_id=USER_ID, team_id=TEAM_ID, spend=1.5)
|
||||
loaders = _Loaders(user=_user(), team=_team(models=("new-a", "new-b")), membership=membership)
|
||||
|
||||
outcome = await loaders.resolver().resolve(UserLookup(user_id=USER_ID), team_id=TEAM_ID)
|
||||
|
||||
assert outcome == ResolvedGrants(
|
||||
user_object=_user(),
|
||||
team_object=_team(models=("new-a", "new-b")),
|
||||
team_membership=membership,
|
||||
effective_user_id=USER_ID,
|
||||
)
|
||||
assert loaders.team_calls[0]["team_id"] == TEAM_ID
|
||||
assert loaders.membership_calls[0]["user_id"] == USER_ID
|
||||
assert loaders.membership_calls[0]["team_id"] == TEAM_ID
|
||||
|
||||
|
||||
async def test_resolve_denies_a_user_removed_from_the_team_without_reading_the_team():
|
||||
loaders = _Loaders(user=_user(teams=("other-team",)), team=_team())
|
||||
|
||||
outcome = await loaders.resolver().resolve(UserLookup(user_id=USER_ID), team_id=TEAM_ID)
|
||||
|
||||
assert outcome == NotAMember(user_id=USER_ID, team_id=TEAM_ID)
|
||||
assert loaders.team_calls == []
|
||||
|
||||
|
||||
async def test_resolve_reports_a_deleted_user():
|
||||
loaders = _Loaders(user_error=UserNotFoundError(user_id=USER_ID), team=_team())
|
||||
|
||||
outcome = await loaders.resolver().resolve(UserLookup(user_id=USER_ID), team_id=TEAM_ID)
|
||||
|
||||
assert outcome == UserGone(user_id=USER_ID)
|
||||
assert loaders.team_calls == []
|
||||
|
||||
|
||||
async def test_resolve_reports_a_deleted_team():
|
||||
loaders = _Loaders(user=_user(), team_error=TeamNotFoundError(team_id=TEAM_ID))
|
||||
|
||||
outcome = await loaders.resolver().resolve(UserLookup(user_id=USER_ID), team_id=TEAM_ID)
|
||||
|
||||
assert outcome == TeamGone(team_id=TEAM_ID)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"loaders",
|
||||
[
|
||||
_Loaders(user_error=Exception("No db connected")),
|
||||
_Loaders(user=_user(), team_error=HTTPException(status_code=500, detail="db timeout")),
|
||||
],
|
||||
ids=["user-read-failed", "team-read-failed"],
|
||||
)
|
||||
async def test_resolve_marks_an_unreadable_row_as_degraded_not_denied(loaders):
|
||||
outcome = await loaders.resolver().resolve(UserLookup(user_id=USER_ID), team_id=TEAM_ID)
|
||||
|
||||
assert isinstance(outcome, LookupDegraded)
|
||||
|
||||
|
||||
async def test_resolve_without_a_team_skips_team_and_membership_reads():
|
||||
loaders = _Loaders(user=_user(teams=()))
|
||||
|
||||
outcome = await loaders.resolver().resolve(UserLookup(user_id=USER_ID), team_id=None)
|
||||
|
||||
assert outcome == ResolvedGrants(
|
||||
user_object=_user(teams=()), team_object=None, team_membership=None, effective_user_id=USER_ID
|
||||
)
|
||||
assert loaders.team_calls == []
|
||||
assert loaders.membership_calls == []
|
||||
|
||||
|
||||
async def test_resolve_identity_reads_membership_under_the_matched_rows_id():
|
||||
legacy_uuid = "bb8ab11f-09aa-47ae-b063-6e80506ac3bc"
|
||||
loaders = _Loaders(user=_user(user_id=legacy_uuid))
|
||||
|
||||
user_object, _membership, effective_user_id = await loaders.resolver().resolve_identity(
|
||||
UserLookup(user_id="matt@example.com", user_email="matt@example.com", sso_user_id="matt@example.com"),
|
||||
team_id=TEAM_ID,
|
||||
)
|
||||
|
||||
assert user_object is not None and user_object.user_id == legacy_uuid
|
||||
assert effective_user_id == legacy_uuid
|
||||
assert loaders.membership_calls[0]["user_id"] == legacy_uuid
|
||||
assert loaders.user_calls[0]["user_email"] == "matt@example.com"
|
||||
|
||||
|
||||
async def test_resolve_identity_without_a_user_id_reads_nothing():
|
||||
loaders = _Loaders(user=_user())
|
||||
|
||||
outcome = await loaders.resolver().resolve_identity(UserLookup(user_id=None), team_id=TEAM_ID)
|
||||
|
||||
assert outcome == (None, None, None)
|
||||
assert loaders.user_calls == []
|
||||
assert loaders.membership_calls == []
|
||||
|
||||
|
||||
async def test_resolve_identity_lets_loader_errors_surface():
|
||||
loaders = _Loaders(user_error=UserNotFoundError(user_id=USER_ID))
|
||||
|
||||
with pytest.raises(UserNotFoundError):
|
||||
await loaders.resolver().resolve_identity(UserLookup(user_id=USER_ID), team_id=None)
|
||||
|
||||
|
||||
def test_raise_public_maps_a_deleted_user_to_401():
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
raise_public(UserGone(user_id=USER_ID))
|
||||
assert exc_info.value.code == "401"
|
||||
assert USER_ID in exc_info.value.message
|
||||
|
||||
|
||||
def test_raise_public_maps_a_removed_member_to_403():
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
raise_public(NotAMember(user_id=USER_ID, team_id=TEAM_ID))
|
||||
assert exc_info.value.status_code == 403
|
||||
assert TEAM_ID in str(exc_info.value.detail)
|
||||
|
||||
|
||||
def test_raise_public_maps_a_deleted_team_to_404():
|
||||
with pytest.raises(TeamNotFoundError) as exc_info:
|
||||
raise_public(TeamGone(team_id=TEAM_ID))
|
||||
assert exc_info.value.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("stored", "expected"),
|
||||
[(["gpt-5.5", "claude-opus-5"], ("gpt-5.5", "claude-opus-5")), ([], ()), ([{"not": "a model"}], ())],
|
||||
ids=["models", "empty", "unusable-column"],
|
||||
)
|
||||
def test_user_models_reads_the_column_as_a_tuple_of_names(stored, expected):
|
||||
assert user_models(LiteLLM_UserTable(user_id=USER_ID, models=stored)) == expected
|
||||
|
|
@ -32,7 +32,7 @@ from litellm.proxy._types import (
|
|||
JWTRoutingOverride,
|
||||
)
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
from litellm.proxy.auth.auth_checks import get_key_object, _cache_key_object
|
||||
from litellm.proxy.auth.auth_checks import TeamNotFoundError, UserNotFoundError, get_key_object, _cache_key_object
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_check_key_model_budget_with_fallback,
|
||||
|
|
@ -6266,6 +6266,192 @@ async def test_non_admin_cli_session_token_reaches_production_auth_path(monkeypa
|
|||
assert result.is_session_token is True
|
||||
|
||||
|
||||
SESSION_TEAM_ID = "team-abc"
|
||||
SESSION_USER_ID = "member-1"
|
||||
|
||||
|
||||
def _mint_session_token(
|
||||
monkeypatch,
|
||||
*,
|
||||
role=LitellmUserRoles.INTERNAL_USER,
|
||||
team_id=SESSION_TEAM_ID,
|
||||
team_models=("stale-model",),
|
||||
models=(),
|
||||
):
|
||||
"""Mint a ``lite login`` token carrying the grants as they were at login time."""
|
||||
monkeypatch.delenv("EXPERIMENTAL_UI_LOGIN", raising=False)
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-cli-test")
|
||||
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
|
||||
|
||||
user_info = LiteLLM_UserTable(
|
||||
user_id=SESSION_USER_ID, user_email="user@example.com", user_role=role.value, models=list(models)
|
||||
)
|
||||
return ExperimentalUIJWTToken.get_cli_jwt_auth_token(
|
||||
user_info, team_id=team_id, team_alias="stale-alias", team_models=list(team_models)
|
||||
)
|
||||
|
||||
|
||||
def _session_user_row(*, teams=(SESSION_TEAM_ID,), role=LitellmUserRoles.INTERNAL_USER, models=()):
|
||||
return LiteLLM_UserTable(user_id=SESSION_USER_ID, user_role=role.value, teams=list(teams), models=list(models))
|
||||
|
||||
|
||||
async def _authenticate_session_token_against_db(
|
||||
cli_token, *, user_row=None, team_row=None, membership_row=None, user_error=None, team_error=None
|
||||
):
|
||||
"""Drive the real builder for a session token with the DB row readers replaced by the given rows or
|
||||
errors. Returns the ``_return_user_api_key_auth_obj`` mock so the caller can read the token it was
|
||||
handed; a denial surfaces as the exception the builder raises."""
|
||||
import litellm.proxy.proxy_server as _proxy_server_mod
|
||||
from fastapi import Request
|
||||
from starlette.datastructures import URL
|
||||
|
||||
attrs = _proxy_attrs_for_db_lookup()
|
||||
attrs["prisma_client"].db.litellm_teammembership.find_first = AsyncMock(return_value=None)
|
||||
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
|
||||
assemble = AsyncMock(return_value=UserAPIKeyAuth(user_id=SESSION_USER_ID, is_session_token=True))
|
||||
try:
|
||||
for k, v in attrs.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
request = Request(scope={"type": "http"})
|
||||
request._url = URL(url="/chat/completions")
|
||||
with (
|
||||
patch( # test-quality-ok: the builder has no injection seam for its assembler yet
|
||||
"litellm.proxy.auth.user_api_key_auth._return_user_api_key_auth_obj", assemble
|
||||
),
|
||||
patch( # test-quality-ok: the builder reads its DB row loaders off module globals
|
||||
"litellm.proxy.auth.user_api_key_auth.get_user_object",
|
||||
AsyncMock(return_value=user_row, side_effect=user_error),
|
||||
),
|
||||
patch( # test-quality-ok: the builder reads its DB row loaders off module globals
|
||||
"litellm.proxy.auth.user_api_key_auth.get_team_object",
|
||||
AsyncMock(return_value=team_row, side_effect=team_error),
|
||||
),
|
||||
patch( # test-quality-ok: the builder reads its DB row loaders off module globals
|
||||
"litellm.proxy.auth.user_api_key_auth.get_team_membership",
|
||||
AsyncMock(return_value=membership_row),
|
||||
),
|
||||
):
|
||||
await _user_api_key_auth_builder(
|
||||
request=request,
|
||||
api_key=f"Bearer {cli_token}",
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
request_data={},
|
||||
)
|
||||
finally:
|
||||
for k, v in originals.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
return assemble
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_token_reads_team_grants_from_the_live_team_row(monkeypatch):
|
||||
"""LIT-7358: a lite login token snapshots the team's models at login, so adding a model to the team did
|
||||
nothing for that CLI until the user logged in again. The team row has to be re-read on every request."""
|
||||
from litellm.models.team import LiteLLM_ModelTable
|
||||
from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj
|
||||
|
||||
cli_token = _mint_session_token(monkeypatch, team_models=("stale-model",))
|
||||
live_team = LiteLLM_TeamTableCachedObj(
|
||||
team_id=SESSION_TEAM_ID,
|
||||
team_alias="renamed-team",
|
||||
models=["gpt-5.5", "claude-opus-5"],
|
||||
litellm_model_table=LiteLLM_ModelTable(model_aliases={"fast": "gpt-5.5"}, created_by="a", updated_by="a"),
|
||||
)
|
||||
membership = LiteLLM_TeamMembership(user_id=SESSION_USER_ID, team_id=SESSION_TEAM_ID, spend=2.5)
|
||||
|
||||
assemble = await _authenticate_session_token_against_db(
|
||||
cli_token, user_row=_session_user_row(), team_row=live_team, membership_row=membership
|
||||
)
|
||||
|
||||
token = assemble.call_args.kwargs["valid_token_dict"]
|
||||
assert token["team_models"] == ["gpt-5.5", "claude-opus-5"]
|
||||
assert token["team_alias"] == "renamed-team"
|
||||
assert token["team_model_aliases"] == {"fast": "gpt-5.5"}
|
||||
assert token["team_member_spend"] == 2.5
|
||||
assert token["is_session_token"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_token_without_a_team_reads_models_from_the_live_user_row(monkeypatch):
|
||||
cli_token = _mint_session_token(monkeypatch, team_id=None, team_models=(), models=("stale-model",))
|
||||
|
||||
assemble = await _authenticate_session_token_against_db(
|
||||
cli_token, user_row=_session_user_row(teams=(), models=("gpt-5.5",))
|
||||
)
|
||||
|
||||
assert assemble.call_args.kwargs["valid_token_dict"]["models"] == ["gpt-5.5"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_demoted_admin_session_token_loses_admin_on_the_next_request(monkeypatch):
|
||||
"""The role baked into the token used to send a former admin down the admin early return forever."""
|
||||
from litellm.proxy._types import LiteLLM_TeamTableCachedObj
|
||||
|
||||
cli_token = _mint_session_token(monkeypatch, role=LitellmUserRoles.PROXY_ADMIN)
|
||||
|
||||
assemble = await _authenticate_session_token_against_db(
|
||||
cli_token,
|
||||
user_row=_session_user_row(role=LitellmUserRoles.INTERNAL_USER),
|
||||
team_row=LiteLLM_TeamTableCachedObj(team_id=SESSION_TEAM_ID, models=["gpt-5.5"]),
|
||||
)
|
||||
|
||||
assemble.assert_awaited_once()
|
||||
assert assemble.call_args.kwargs["valid_token_dict"]["user_role"] == LitellmUserRoles.INTERNAL_USER
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_token_is_refused_once_the_user_leaves_the_team(monkeypatch):
|
||||
cli_token = _mint_session_token(monkeypatch)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _authenticate_session_token_against_db(cli_token, user_row=_session_user_row(teams=("other-team",)))
|
||||
|
||||
assert exc_info.value.code == str(status.HTTP_403_FORBIDDEN)
|
||||
assert SESSION_TEAM_ID in exc_info.value.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_token_is_refused_once_the_user_is_deleted(monkeypatch):
|
||||
cli_token = _mint_session_token(monkeypatch)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _authenticate_session_token_against_db(cli_token, user_error=UserNotFoundError(user_id=SESSION_USER_ID))
|
||||
|
||||
assert exc_info.value.code == str(status.HTTP_401_UNAUTHORIZED)
|
||||
assert exc_info.value.type == ProxyErrorTypes.auth_error
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_token_is_refused_once_the_team_is_deleted(monkeypatch):
|
||||
cli_token = _mint_session_token(monkeypatch)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _authenticate_session_token_against_db(
|
||||
cli_token, user_row=_session_user_row(), team_error=TeamNotFoundError(team_id=SESSION_TEAM_ID)
|
||||
)
|
||||
|
||||
assert exc_info.value.code == str(status.HTTP_404_NOT_FOUND)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_token_keeps_minted_grants_when_the_team_row_cannot_be_read(monkeypatch):
|
||||
"""A DB hiccup says nothing about the caller, so the grants minted at login stand for that request."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
cli_token = _mint_session_token(monkeypatch, team_models=("stale-model",))
|
||||
|
||||
assemble = await _authenticate_session_token_against_db(
|
||||
cli_token, user_row=_session_user_row(), team_error=HTTPException(status_code=500, detail="db timeout")
|
||||
)
|
||||
|
||||
token = assemble.call_args.kwargs["valid_token_dict"]
|
||||
assert token["team_models"] == ["stale-model"]
|
||||
assert token["team_alias"] == "stale-alias"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cli_session_token_authenticates_when_jwt_auth_enabled_without_license(monkeypatch):
|
||||
"""A lite login token is an encrypted (non-JWT) session blob. With
|
||||
|
|
|
|||
|
|
@ -1195,6 +1195,61 @@ async def test_health_services_endpoint_newrelic_allows_proxy_admin(admin_role):
|
|||
mock_instance.async_health_check.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"role",
|
||||
[
|
||||
None,
|
||||
LitellmUserRoles.INTERNAL_USER,
|
||||
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
|
||||
LitellmUserRoles.TEAM,
|
||||
LitellmUserRoles.CUSTOMER,
|
||||
],
|
||||
)
|
||||
async def test_health_services_endpoint_webhook_blocks_non_admin(role):
|
||||
"""
|
||||
/health/services?service=webhook fires a real budget_crossed alert for the
|
||||
caller's user_id and writes the same dedup cache entry the auth-time user
|
||||
budget alert uses, so a non-admin could suppress their own real alert for
|
||||
the cache TTL. Only proxy admins may trigger it.
|
||||
"""
|
||||
mock_proxy_logging = MagicMock()
|
||||
mock_proxy_logging.budget_alerts = AsyncMock()
|
||||
user_api_key_dict = UserAPIKeyAuth(token="non-admin-token", user_id="non-admin-user", user_role=role)
|
||||
|
||||
with patch( # test-quality-ok: endpoint reads proxy_server module globals, same pattern as sibling tests
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj",
|
||||
mock_proxy_logging,
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await health_services_endpoint(user_api_key_dict=user_api_key_dict, service="webhook")
|
||||
|
||||
assert str(exc_info.value.code) == "403"
|
||||
mock_proxy_logging.budget_alerts.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"admin_role",
|
||||
[LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY],
|
||||
)
|
||||
async def test_health_services_endpoint_webhook_allows_proxy_admin(admin_role):
|
||||
mock_proxy_logging = MagicMock()
|
||||
mock_proxy_logging.budget_alerts = AsyncMock()
|
||||
user_api_key_dict = UserAPIKeyAuth(token="admin-token", user_id="admin-user", user_role=admin_role)
|
||||
|
||||
with patch( # test-quality-ok: endpoint reads proxy_server module globals, same pattern as sibling tests
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj",
|
||||
mock_proxy_logging,
|
||||
):
|
||||
await health_services_endpoint(user_api_key_dict=user_api_key_dict, service="webhook")
|
||||
|
||||
mock_proxy_logging.budget_alerts.assert_awaited_once()
|
||||
sent = mock_proxy_logging.budget_alerts.await_args.kwargs
|
||||
assert sent["type"] == "user_budget"
|
||||
assert sent["user_info"].user_id == "admin-user"
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def proxy_client(monkeypatch):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -13941,6 +13941,52 @@ async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent(
|
|||
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 1.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_evict_created_membership_caches_drops_the_negative_sentinel():
|
||||
"""
|
||||
Regression: a membership-create path (/team/member_add, the /team/update budget backfill) must
|
||||
evict any cached "no membership" sentinel a prior session-token read left, so a per-member budget
|
||||
attached at create time is enforced on the next request instead of after the membership cache TTL.
|
||||
Uses a real cache so the assertion is that the sentinel is actually gone, not that a mock was called.
|
||||
"""
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
NO_TEAM_MEMBERSHIP_SENTINEL,
|
||||
UserApiKeyCache,
|
||||
team_membership_reservation_cache_key,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.team_endpoints import _evict_created_membership_caches
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
kept_key = team_membership_reservation_cache_key(user_id="carol", team_id="team-eviction")
|
||||
evicted_key = team_membership_reservation_cache_key(user_id="bob", team_id="team-eviction")
|
||||
await cache.async_set_cache(key=kept_key, value=NO_TEAM_MEMBERSHIP_SENTINEL)
|
||||
await cache.async_set_cache(key=evicted_key, value=NO_TEAM_MEMBERSHIP_SENTINEL)
|
||||
|
||||
await _evict_created_membership_caches(user_ids=("bob",), team_id="team-eviction", user_api_key_cache=cache)
|
||||
|
||||
assert await cache.async_get_cache(key=evicted_key) is None
|
||||
assert await cache.async_get_cache(key=kept_key) == NO_TEAM_MEMBERSHIP_SENTINEL
|
||||
|
||||
|
||||
def test_member_user_ids_keeps_only_string_user_ids():
|
||||
"""
|
||||
The /team/update backfill feeds Prisma-deserialized member dicts here; a row can be missing
|
||||
user_id or carry a non-string value. Only real string ids may reach invalidate_team_member_spend_state,
|
||||
so those get eviction and the malformed rows are dropped rather than crashing the update.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.team_endpoints import _member_user_ids
|
||||
|
||||
members = [
|
||||
{"user_id": "alice", "role": "admin"},
|
||||
{"role": "user"},
|
||||
{"user_id": None, "role": "user"},
|
||||
{"user_id": 123, "role": "user"},
|
||||
{"user_id": "bob", "role": "user"},
|
||||
]
|
||||
|
||||
assert _member_user_ids(members) == ("alice", "bob")
|
||||
|
||||
|
||||
def _team_spend_by_user_team(team_id: str, team_alias: str, member: Member, permissions: list[str]) -> MagicMock:
|
||||
team = MagicMock(spec=LiteLLM_TeamTable)
|
||||
team.team_id = team_id
|
||||
|
|
|
|||
|
|
@ -2441,6 +2441,63 @@ def test_list_files_resolves_wildcard_deployment_credentials(
|
|||
proxy_logging_obj.post_call_failure_hook.assert_not_called()
|
||||
|
||||
|
||||
def test_list_files_by_model_returns_an_openai_page_for_a_provider_listing(
|
||||
mocker: MockerFixture, monkeypatch, llm_router: Router
|
||||
):
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
from litellm.types.llms.openai import FileListPage
|
||||
|
||||
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, llm_router)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router)
|
||||
proxy_logging_obj.update_request_status = mocker.AsyncMock()
|
||||
proxy_logging_obj.post_call_success_hook = mocker.AsyncMock(return_value=None)
|
||||
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
|
||||
|
||||
listed_files = [
|
||||
OpenAIFileObject(
|
||||
id=f"file-{index}",
|
||||
bytes=index,
|
||||
created_at=index,
|
||||
filename=f"{index}.jsonl",
|
||||
object="file",
|
||||
purpose="batch",
|
||||
status="uploaded",
|
||||
)
|
||||
for index in (1, 2)
|
||||
]
|
||||
|
||||
async def _mock_afile_list(**kwargs):
|
||||
return list(listed_files)
|
||||
|
||||
monkeypatch.setattr(litellm, "afile_list", _mock_afile_list)
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_id="test-user",
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.get(
|
||||
"/v1/files?target_model_names=gpt-3.5-turbo",
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
body = response.json()
|
||||
assert body["object"] == "list"
|
||||
assert [listed["id"] for listed in body["data"]] == ["file-1", "file-2"]
|
||||
assert (body["first_id"], body["last_id"], body["has_more"]) == ("file-1", "file-2", False)
|
||||
hook_response = proxy_logging_obj.post_call_success_hook.call_args.kwargs["response"]
|
||||
assert isinstance(hook_response, FileListPage)
|
||||
assert [listed.id for listed in hook_response.data] == ["file-1", "file-2"]
|
||||
|
||||
|
||||
def test_list_files_model_routing_does_not_forward_custom_llm_provider_twice(
|
||||
mocker: MockerFixture, monkeypatch, llm_router: Router
|
||||
):
|
||||
|
|
@ -4668,6 +4725,249 @@ def test_create_file_path_traversal_filename_rejected_before_forwarding(monkeypa
|
|||
assert forwarded_calls == []
|
||||
|
||||
|
||||
def test_list_files_target_model_names_passes_trusted_bedrock_credentials(
|
||||
mocker: MockerFixture, monkeypatch
|
||||
):
|
||||
"""
|
||||
GET /v1/files?target_model_names=<bedrock model> must hand the deployment's
|
||||
immutable credential snapshot to litellm.afile_list, since Bedrock resolves
|
||||
the S3 bucket to list from that snapshot rather than from request params.
|
||||
"""
|
||||
from types import MappingProxyType
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
bedrock_router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "bedrock-claude",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "secret",
|
||||
"aws_region_name": "us-west-2",
|
||||
"s3_bucket_name": "my-bucket",
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, bedrock_router)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", bedrock_router)
|
||||
proxy_logging_obj.update_request_status = mocker.AsyncMock()
|
||||
proxy_logging_obj.post_call_success_hook = mocker.AsyncMock(return_value=[])
|
||||
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
|
||||
|
||||
captured_kwargs: dict = {}
|
||||
|
||||
async def _mock_afile_list(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(litellm, "afile_list", _mock_afile_list)
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_id="test-user",
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.get(
|
||||
"/v1/files?target_model_names=bedrock-claude&purpose=batch",
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert captured_kwargs["custom_llm_provider"] == "bedrock"
|
||||
assert captured_kwargs["purpose"] == "batch"
|
||||
trusted_credentials = captured_kwargs["_litellm_internal_model_credentials"]
|
||||
assert isinstance(trusted_credentials, MappingProxyType)
|
||||
assert trusted_credentials["s3_bucket_name"] == "my-bucket"
|
||||
proxy_logging_obj.post_call_failure_hook.assert_not_called()
|
||||
|
||||
|
||||
def test_delete_file_answers_400_for_an_id_outside_the_configured_bucket(mocker: MockerFixture, monkeypatch):
|
||||
from urllib.parse import quote
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
bedrock_router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "bedrock-claude",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "secret",
|
||||
"aws_region_name": "us-west-2",
|
||||
"s3_bucket_name": "my-bucket",
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, bedrock_router)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", bedrock_router)
|
||||
proxy_logging_obj.update_request_status = mocker.AsyncMock()
|
||||
proxy_logging_obj.post_call_success_hook = mocker.AsyncMock(return_value=[])
|
||||
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_id="test-user",
|
||||
)
|
||||
foreign_file_id: Final = quote("s3://other-bucket/litellm-bedrock-files/job-123/input.jsonl", safe="")
|
||||
|
||||
try:
|
||||
with respx.mock:
|
||||
response = client.delete(
|
||||
f"/v1/files/{foreign_file_id}?model=bedrock-claude",
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert "configured storage bucket" in response.json()["error"]["message"]
|
||||
|
||||
|
||||
def _cloud_files_router() -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "bedrock-claude",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "secret",
|
||||
"aws_region_name": "us-west-2",
|
||||
"s3_bucket_name": "my-bucket",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "vertex-gemini",
|
||||
"litellm_params": {
|
||||
"model": "vertex_ai/gemini-3.8-flash",
|
||||
"vertex_project": "my-project",
|
||||
"vertex_location": "us-central1",
|
||||
"gcs_bucket_name": "my-gcs-bucket",
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
RAW_S3_FILE_ID: Final = "s3://my-bucket/litellm-batch-outputs/job-123/abc/input.jsonl.out"
|
||||
RAW_GCS_FILE_ID: Final = "gs://my-gcs-bucket/litellm-vertex-files/publishers/google/models/gemini-3.8-flash/abc123"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route_prefix", "raw_file_id", "model_name"),
|
||||
(
|
||||
("/bedrock/v1/files", RAW_S3_FILE_ID, "bedrock-claude"),
|
||||
("/v1/files", RAW_S3_FILE_ID, "bedrock-claude"),
|
||||
("/files", RAW_S3_FILE_ID, "bedrock-claude"),
|
||||
("/vertex_ai/v1/files", RAW_GCS_FILE_ID, "vertex-gemini"),
|
||||
),
|
||||
)
|
||||
def test_delete_file_answers_403_for_a_raw_cloud_id_from_a_non_admin_key(
|
||||
mocker: MockerFixture, monkeypatch, route_prefix: str, raw_file_id: str, model_name: str
|
||||
):
|
||||
from urllib.parse import quote
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
bedrock_router = _cloud_files_router()
|
||||
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, bedrock_router)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", bedrock_router)
|
||||
proxy_logging_obj.update_request_status = mocker.AsyncMock()
|
||||
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
|
||||
afile_delete = mocker.AsyncMock()
|
||||
monkeypatch.setattr(litellm, "afile_delete", afile_delete)
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="test-user",
|
||||
models=["bedrock-claude", "vertex-gemini"],
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.delete(
|
||||
f"{route_prefix}/{quote(raw_file_id, safe='')}?model={model_name}",
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
assert "proxy admin" in response.json()["error"]["message"]
|
||||
afile_delete.assert_not_called()
|
||||
|
||||
|
||||
def test_delete_file_forwards_a_raw_cloud_id_from_a_proxy_admin_key(mocker: MockerFixture, monkeypatch):
|
||||
from urllib.parse import quote
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
bedrock_router = _cloud_files_router()
|
||||
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, bedrock_router)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", bedrock_router)
|
||||
proxy_logging_obj.update_request_status = mocker.AsyncMock()
|
||||
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
|
||||
|
||||
captured_kwargs: dict = {}
|
||||
|
||||
async def _mock_afile_delete(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return OpenAIFileObject(
|
||||
id=RAW_S3_FILE_ID,
|
||||
object="file",
|
||||
bytes=2,
|
||||
created_at=1234567890,
|
||||
filename="input.jsonl.out",
|
||||
purpose="batch_output",
|
||||
status="processed",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(litellm, "afile_delete", _mock_afile_delete)
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_id="test-user",
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.delete(
|
||||
f"/bedrock/v1/files/{quote(RAW_S3_FILE_ID, safe='')}?model=bedrock-claude",
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert captured_kwargs.get("file_id") == RAW_S3_FILE_ID
|
||||
assert captured_kwargs.get("custom_llm_provider") == "bedrock"
|
||||
proxy_logging_obj.post_call_failure_hook.assert_not_called()
|
||||
|
||||
|
||||
def _setup_managed_file_route_answering_404(
|
||||
mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -124,7 +124,10 @@ def _reconstruct_ui_where_from_sql(sql_query, params):
|
|||
sess = re.fullmatch(r"session_id LIKE \$(\d+)", cond)
|
||||
status = re.fullmatch(r"status = \$(\d+)", cond)
|
||||
api_key_not_in = re.fullmatch(r"api_key NOT IN \(\$(\d+), \$(\d+)\)", cond)
|
||||
if gte:
|
||||
req_or_call = re.fullmatch(r"\(request_id = \$(\d+) OR litellm_call_id = \$\1\)", cond)
|
||||
if req_or_call:
|
||||
where["request_id_or_call_id"] = params[int(req_or_call.group(1)) - 1]
|
||||
elif gte:
|
||||
date_bounds["gte"] = _iso(params[int(gte.group(1)) - 1])
|
||||
elif lte:
|
||||
date_bounds["lte"] = _iso(params[int(lte.group(1)) - 1])
|
||||
|
|
@ -213,6 +216,8 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No
|
|||
query_observer(sql_query, params)
|
||||
if "mcp_tool_call_count" in sql_query:
|
||||
return []
|
||||
if 'SELECT DISTINCT "user", team_id' in sql_query:
|
||||
return _emulate_spend_log_owner_lookup(mock_spend_logs, sql_query, params)
|
||||
filtered = filter_fn(_reconstruct_ui_where_from_sql(sql_query, params))
|
||||
total = len(filtered)
|
||||
if "COUNT(*)" in sql_query:
|
||||
|
|
@ -220,7 +225,13 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No
|
|||
return [{"total_count": min(total, cap_plus_one)}]
|
||||
page_size = params[-2] if len(params) >= 2 else 50
|
||||
skip = params[-1] if len(params) >= 1 else 0
|
||||
return [row for row in filtered[skip : skip + page_size]]
|
||||
exact_first = re.search(r"ORDER BY \(request_id = \$(\d+)\) DESC", sql_query)
|
||||
ordered = (
|
||||
sorted(filtered, key=lambda row: row["request_id"] == params[int(exact_first.group(1)) - 1], reverse=True)
|
||||
if exact_first
|
||||
else filtered
|
||||
)
|
||||
return [row for row in ordered[skip : skip + page_size]]
|
||||
|
||||
class MockPrismaClient:
|
||||
def __init__(self):
|
||||
|
|
@ -424,37 +435,173 @@ def test_can_user_view_spend_log_false_for_other_roles():
|
|||
assert spend_management_endpoints._can_user_view_spend_log(auth) is False
|
||||
|
||||
|
||||
def _emulate_spend_log_owner_lookup(rows, sql_query, params):
|
||||
"""Emulate the ownership lookup SQL over an in-memory spend-log corpus,
|
||||
honoring DISTINCT and any literal LIMIT the query carries so a capped or
|
||||
non-distinct query produces the truncated result it would in Postgres."""
|
||||
lookup_id = params[0]
|
||||
matches = [
|
||||
{"user": row.get("user"), "team_id": row.get("team_id")}
|
||||
for row in rows
|
||||
if lookup_id in (row.get("request_id"), row.get("litellm_call_id"))
|
||||
]
|
||||
if "DISTINCT" in sql_query:
|
||||
deduped = []
|
||||
for match in matches:
|
||||
if match not in deduped:
|
||||
deduped.append(match)
|
||||
matches = deduped
|
||||
limit = re.search(r"LIMIT\s+(\d+)", sql_query, re.IGNORECASE)
|
||||
if limit is not None:
|
||||
matches = matches[: int(limit.group(1))]
|
||||
return matches
|
||||
|
||||
|
||||
def _make_owner_lookup_prisma(rows):
|
||||
class MockDB:
|
||||
async def query_raw(self, sql_query, *params):
|
||||
return _emulate_spend_log_owner_lookup(rows, sql_query, params)
|
||||
|
||||
class MockPrisma:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
|
||||
return MockPrisma()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assert_user_can_view_request_id_rejects_both_users_none():
|
||||
"""
|
||||
API keys with user_id=None must not be treated as owning a log whose user
|
||||
field is None (avoid None == None bypass).
|
||||
"""
|
||||
|
||||
class MockRow:
|
||||
user = None
|
||||
team_id = None
|
||||
|
||||
class MockSpendLogs:
|
||||
async def find_unique(self, where, include=None):
|
||||
return MockRow()
|
||||
|
||||
class MockDB:
|
||||
def __init__(self):
|
||||
self.litellm_spendlogs = MockSpendLogs()
|
||||
|
||||
class MockPrisma:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
prisma = _make_owner_lookup_prisma(
|
||||
[{"request_id": "req-none-user", "litellm_call_id": None, "user": None, "team_id": None}]
|
||||
)
|
||||
|
||||
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id=None)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await spend_management_endpoints._assert_user_can_view_request_id(
|
||||
MockPrisma(), auth, "req-none-user"
|
||||
prisma, auth, "req-none-user"
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assert_user_can_view_request_id_rejects_when_no_match_is_owned():
|
||||
"""An id whose every matching row belongs to other tenants is refused outright,
|
||||
so the relaxed date window of an id lookup cannot reach a foreign row."""
|
||||
prisma = _make_owner_lookup_prisma(
|
||||
[
|
||||
{"request_id": "foreign-request", "litellm_call_id": "shared-id", "user": "tenant_a", "team_id": None},
|
||||
{"request_id": "shared-id", "litellm_call_id": "other-call-id", "user": "tenant_b", "team_id": None},
|
||||
]
|
||||
)
|
||||
|
||||
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="caller")
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await spend_management_endpoints._assert_user_can_view_request_id(prisma, auth, "shared-id")
|
||||
assert exc_info.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assert_user_can_view_request_id_allows_owner_despite_foreign_collision():
|
||||
"""
|
||||
litellm_call_id comes from the client-settable x-litellm-call-id header, so
|
||||
another tenant can mint a row whose call id equals the caller's request_id.
|
||||
That collision must not lock the caller out of their own row: the pre-check
|
||||
passes once one match is theirs, and the scoped data queries keep the foreign
|
||||
row out of the result. Regression for the every-match-must-be-owned rule that
|
||||
let any tenant deny another's lookup by reusing their id.
|
||||
"""
|
||||
prisma = _make_owner_lookup_prisma(
|
||||
[
|
||||
{
|
||||
"request_id": "attacker-own-request",
|
||||
"litellm_call_id": "victim-request-id",
|
||||
"user": "attacker",
|
||||
"team_id": None,
|
||||
},
|
||||
{
|
||||
"request_id": "victim-request-id",
|
||||
"litellm_call_id": "victim-call-id",
|
||||
"user": "victim",
|
||||
"team_id": None,
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="victim")
|
||||
result = await spend_management_endpoints._assert_user_can_view_request_id(prisma, auth, "victim-request-id")
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assert_user_can_view_request_id_finds_owner_past_any_row_cap():
|
||||
"""
|
||||
An attacker can mint hundreds of rows carrying the victim's request_id as
|
||||
their litellm_call_id, so a capped or sampled ownership read could exhaust
|
||||
its cap on attacker-owned rows and never see the victim's own row, locking
|
||||
the victim out of their lookup. The ownership read must consider every
|
||||
matching row's owner no matter how many rows match. Regression for the
|
||||
find_many(take=100) sample the first fix used.
|
||||
"""
|
||||
rows = [
|
||||
{
|
||||
"request_id": f"attacker-request-{i}",
|
||||
"litellm_call_id": "victim-request-id",
|
||||
"user": "attacker",
|
||||
"team_id": None,
|
||||
}
|
||||
for i in range(150)
|
||||
]
|
||||
rows.append(
|
||||
{
|
||||
"request_id": "victim-request-id",
|
||||
"litellm_call_id": "victim-call-id",
|
||||
"user": "victim",
|
||||
"team_id": None,
|
||||
}
|
||||
)
|
||||
prisma = _make_owner_lookup_prisma(rows)
|
||||
|
||||
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="victim")
|
||||
result = await spend_management_endpoints._assert_user_can_view_request_id(prisma, auth, "victim-request-id")
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assert_user_can_view_request_id_allows_when_every_match_is_owned():
|
||||
"""The same ambiguous id matching more than one row is fine when every match
|
||||
belongs to the caller (e.g. two of the caller's own requests happen to share
|
||||
a request_id/litellm_call_id pairing); only a foreign match should block it."""
|
||||
prisma = _make_owner_lookup_prisma(
|
||||
[
|
||||
{
|
||||
"request_id": "shared-request-id",
|
||||
"litellm_call_id": "caller-call-a",
|
||||
"user": "caller",
|
||||
"team_id": None,
|
||||
},
|
||||
{
|
||||
"request_id": "caller-request-b",
|
||||
"litellm_call_id": "shared-request-id",
|
||||
"user": "caller",
|
||||
"team_id": None,
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="caller")
|
||||
result = await spend_management_endpoints._assert_user_can_view_request_id(
|
||||
prisma, auth, "shared-request-id"
|
||||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assert_user_can_view_request_id_rejects_missing_row():
|
||||
"""
|
||||
|
|
@ -462,24 +609,11 @@ async def test_assert_user_can_view_request_id_rejects_missing_row():
|
|||
authorize reading the payload from cold storage; a missing row is not
|
||||
the same as an owned row.
|
||||
"""
|
||||
|
||||
class MockSpendLogs:
|
||||
async def find_unique(self, where, include=None):
|
||||
return None
|
||||
|
||||
class MockDB:
|
||||
def __init__(self):
|
||||
self.litellm_spendlogs = MockSpendLogs()
|
||||
|
||||
class MockPrisma:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
prisma = _make_owner_lookup_prisma([])
|
||||
|
||||
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1")
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await spend_management_endpoints._assert_user_can_view_request_id(
|
||||
MockPrisma(), auth, "req-missing-row"
|
||||
)
|
||||
await spend_management_endpoints._assert_user_can_view_request_id(prisma, auth, "req-missing-row")
|
||||
assert exc_info.value.status_code == 403
|
||||
|
||||
|
||||
|
|
@ -507,6 +641,7 @@ def test_ui_view_request_response_forbids_non_admin_without_db(client, monkeypat
|
|||
|
||||
ignored_keys = [
|
||||
"request_id",
|
||||
"litellm_call_id",
|
||||
"metadata.litellm_call_id",
|
||||
"session_id",
|
||||
"startTime",
|
||||
|
|
@ -2216,7 +2351,10 @@ async def test_ui_view_spend_logs_request_id_lookup_ignores_date_window(
|
|||
def filter_fn(where):
|
||||
captured["where"] = where
|
||||
rows = _filter_logs_by_date_range(mock_spend_logs, where)
|
||||
if where.get("request_id"):
|
||||
rid_either = where.get("request_id_or_call_id")
|
||||
if rid_either:
|
||||
rows = [r for r in rows if rid_either in (r["request_id"], r.get("litellm_call_id"))]
|
||||
elif where.get("request_id"):
|
||||
rows = [r for r in rows if r["request_id"] == where["request_id"]]
|
||||
return rows
|
||||
|
||||
|
|
@ -2246,9 +2384,82 @@ async def test_ui_view_spend_logs_request_id_lookup_ignores_date_window(
|
|||
data = response.json()
|
||||
assert data["total"] == 1
|
||||
assert data["data"][0]["request_id"] == "req-old"
|
||||
# Query dropped the time window and scoped solely by the primary key.
|
||||
# Query dropped the time window and scoped solely by the id lookup.
|
||||
assert "startTime" not in captured["where"]
|
||||
assert captured["where"]["request_id"] == "req-old"
|
||||
assert captured["where"]["request_id_or_call_id"] == "req-old"
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_request_id_lookup_matches_litellm_call_id(
|
||||
client, monkeypatch
|
||||
):
|
||||
"""
|
||||
LIT-6302: success rows are keyed by the upstream provider response id, so a
|
||||
lookup with the x-litellm-call-id response header value found nothing. The id
|
||||
lookup now matches request_id OR litellm_call_id, resolving the header value.
|
||||
"""
|
||||
today = datetime.datetime.now(timezone.utc)
|
||||
mock_spend_logs = [
|
||||
{
|
||||
"id": "log_provider_keyed",
|
||||
"request_id": "chatcmpl-9ZKMURhVYSi9D6r6PJ9vLcayIK0Vm",
|
||||
"litellm_call_id": "b980eea9-5cd9-4099-93cd-8291e46c76fd",
|
||||
"api_key": "sk-test-key",
|
||||
"user": "test_user_1",
|
||||
"team_id": "team1",
|
||||
"spend": 0.05,
|
||||
"startTime": today.isoformat(),
|
||||
"model": "gpt-4",
|
||||
},
|
||||
{
|
||||
"id": "log_other",
|
||||
"request_id": "chatcmpl-other",
|
||||
"litellm_call_id": "11111111-2222-3333-4444-555555555555",
|
||||
"api_key": "sk-test-key",
|
||||
"user": "test_user_1",
|
||||
"team_id": "team1",
|
||||
"spend": 0.01,
|
||||
"startTime": today.isoformat(),
|
||||
"model": "gpt-4",
|
||||
},
|
||||
]
|
||||
|
||||
def filter_fn(where):
|
||||
rid_either = where.get("request_id_or_call_id")
|
||||
if rid_either:
|
||||
return [
|
||||
r
|
||||
for r in mock_spend_logs
|
||||
if rid_either in (r["request_id"], r.get("litellm_call_id"))
|
||||
]
|
||||
if where.get("request_id"):
|
||||
return [
|
||||
r for r in mock_spend_logs if r["request_id"] == where["request_id"]
|
||||
]
|
||||
return list(mock_spend_logs)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn),
|
||||
)
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={"request_id": "b980eea9-5cd9-4099-93cd-8291e46c76fd"},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["total"] == 1
|
||||
assert (
|
||||
data["data"][0]["request_id"] == "chatcmpl-9ZKMURhVYSi9D6r6PJ9vLcayIK0Vm"
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
|
@ -2303,23 +2514,18 @@ async def test_ui_view_spend_logs_request_id_blocks_non_owner(client, monkeypatc
|
|||
"""A non-admin looking up a request_id they do not own is rejected (403), so
|
||||
the relaxed date window cannot read another tenant's log by id."""
|
||||
|
||||
class _ForeignRow:
|
||||
user = "other_user"
|
||||
team_id = None
|
||||
prisma = _make_owner_lookup_prisma(
|
||||
[
|
||||
{
|
||||
"request_id": "foreign-req",
|
||||
"litellm_call_id": None,
|
||||
"user": "other_user",
|
||||
"team_id": None,
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
class _SpendLogs:
|
||||
async def find_unique(self, where, include=None):
|
||||
return _ForeignRow()
|
||||
|
||||
class _DB:
|
||||
def __init__(self):
|
||||
self.litellm_spendlogs = _SpendLogs()
|
||||
|
||||
class _Prisma:
|
||||
def __init__(self):
|
||||
self.db = _DB()
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _Prisma())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1"
|
||||
)
|
||||
|
|
@ -2335,13 +2541,437 @@ async def test_ui_view_spend_logs_request_id_blocks_non_owner(client, monkeypatc
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only(
|
||||
async def test_ui_view_spend_logs_request_id_collision_serves_only_callers_rows(client, monkeypatch):
|
||||
"""Two tenants share one id: the attacker minted a row whose client-set
|
||||
litellm_call_id equals the victim's request_id. Each side's lookup of that id
|
||||
returns only their own row, so the collision neither leaks the other tenant's
|
||||
row nor denies the victim theirs (Veria: identifier collision could deny access)."""
|
||||
now_iso = datetime.datetime.now(timezone.utc).isoformat()
|
||||
corpus = [
|
||||
{
|
||||
"id": "log_attacker",
|
||||
"request_id": "attacker-req",
|
||||
"litellm_call_id": "victim-req",
|
||||
"api_key": "sk-attacker-key",
|
||||
"user": "attacker_user",
|
||||
"team_id": None,
|
||||
"spend": 0.05,
|
||||
"startTime": now_iso,
|
||||
"model": "gpt-4",
|
||||
},
|
||||
{
|
||||
"id": "log_victim",
|
||||
"request_id": "victim-req",
|
||||
"litellm_call_id": "victim-call-id",
|
||||
"api_key": "sk-victim-key",
|
||||
"user": "victim_user",
|
||||
"team_id": None,
|
||||
"spend": 0.07,
|
||||
"startTime": now_iso,
|
||||
"model": "gpt-4",
|
||||
},
|
||||
]
|
||||
|
||||
def filter_fn(where):
|
||||
rid_either = where.get("request_id_or_call_id")
|
||||
rows = [r for r in corpus if rid_either in (r["request_id"], r["litellm_call_id"])]
|
||||
return [r for r in rows if where.get("user") is None or r["user"] == where["user"]]
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", make_ui_spend_logs_mock_prisma(corpus, filter_fn))
|
||||
try:
|
||||
for caller, own_request_id, other in (
|
||||
("victim_user", "victim-req", "attacker_user"),
|
||||
("attacker_user", "attacker-req", "victim_user"),
|
||||
):
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda caller=caller: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER, user_id=caller
|
||||
)
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={"request_id": "victim-req"},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
data = response.json()
|
||||
assert data["total"] == 1
|
||||
assert data["data"][0]["request_id"] == own_request_id
|
||||
assert other not in response.text
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_id_lookup_scopes_every_non_admin_role(client, monkeypatch):
|
||||
"""An org admin reaches /spend/logs/ui without the internal-user row scope. An
|
||||
id lookup still fetches only rows they own, so another tenant's row carrying
|
||||
that id as its client-set litellm_call_id neither leaks nor turns the
|
||||
org admin's own lookup into a 403 (Bugbot: non-internal id lookup 403s on collision)."""
|
||||
now_iso = datetime.datetime.now(timezone.utc).isoformat()
|
||||
corpus = [
|
||||
{
|
||||
"id": "log_attacker",
|
||||
"request_id": "attacker-req",
|
||||
"litellm_call_id": "victim-req",
|
||||
"api_key": "sk-attacker-key",
|
||||
"user": "attacker_user",
|
||||
"team_id": None,
|
||||
"spend": 0.05,
|
||||
"startTime": now_iso,
|
||||
"model": "gpt-4",
|
||||
},
|
||||
{
|
||||
"id": "log_victim",
|
||||
"request_id": "victim-req",
|
||||
"litellm_call_id": "victim-call-id",
|
||||
"api_key": "sk-victim-key",
|
||||
"user": "victim_user",
|
||||
"team_id": None,
|
||||
"spend": 0.07,
|
||||
"startTime": now_iso,
|
||||
"model": "gpt-4",
|
||||
},
|
||||
]
|
||||
|
||||
def filter_fn(where):
|
||||
rid_either = where.get("request_id_or_call_id")
|
||||
rows = [r for r in corpus if rid_either in (r["request_id"], r["litellm_call_id"])]
|
||||
return [r for r in rows if where.get("user") is None or r["user"] == where["user"]]
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", make_ui_spend_logs_mock_prisma(corpus, filter_fn))
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.ORG_ADMIN, user_id="victim_user"
|
||||
)
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={"request_id": "victim-req"},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
data = response.json()
|
||||
assert data["total"] == 1
|
||||
assert [row["request_id"] for row in data["data"]] == ["victim-req"]
|
||||
assert "attacker_user" not in response.text
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_request_id_rejects_foreign_row_inserted_after_owner_check(client, monkeypatch):
|
||||
"""The SQL scope keeps foreign rows out of an id lookup; this backstop covers a
|
||||
row the scope did not filter (the mock ignores it on purpose). The rows actually
|
||||
fetched are ownership-checked again, so the lookup answers 403 instead of serving
|
||||
the other tenant's row."""
|
||||
now_iso = datetime.datetime.now(timezone.utc).isoformat()
|
||||
owned_row = {
|
||||
"id": "log_owned",
|
||||
"request_id": "attacker-req",
|
||||
"litellm_call_id": "shared-id",
|
||||
"api_key": "sk-test-key",
|
||||
"user": "user_1",
|
||||
"team_id": None,
|
||||
"spend": 0.05,
|
||||
"startTime": now_iso,
|
||||
"model": "gpt-4",
|
||||
}
|
||||
foreign_row = {
|
||||
"id": "log_foreign",
|
||||
"request_id": "shared-id",
|
||||
"litellm_call_id": None,
|
||||
"api_key": "sk-victim-key",
|
||||
"user": "victim_user",
|
||||
"team_id": None,
|
||||
"spend": 0.07,
|
||||
"startTime": now_iso,
|
||||
"model": "gpt-4",
|
||||
}
|
||||
|
||||
mock_prisma = make_ui_spend_logs_mock_prisma([owned_row], lambda where: [owned_row, foreign_row])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1"
|
||||
)
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={"request_id": "shared-id"},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 403
|
||||
assert "victim_user" not in response.text
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
def _make_payload_lookup_prisma(rows):
|
||||
"""Emulate the detail endpoint's SQL over an in-memory corpus: the owner
|
||||
pre-check, the caller scope on ``"user"`` and permitted teams, and the
|
||||
exact-request_id-first ordering with LIMIT 1."""
|
||||
|
||||
class MockDB:
|
||||
async def query_raw(self, sql_query, *params):
|
||||
if 'SELECT DISTINCT "user", team_id' in sql_query:
|
||||
return _emulate_spend_log_owner_lookup(rows, sql_query, params)
|
||||
lookup_id = params[0]
|
||||
matches = [r for r in rows if lookup_id in (r["request_id"], r["litellm_call_id"])]
|
||||
if '"user" = $2' in sql_query:
|
||||
team_ids = params[2] if "ANY($3::text[])" in sql_query else ()
|
||||
matches = [r for r in matches if r["user"] == params[1] or r["team_id"] in team_ids]
|
||||
if "ORDER BY (request_id = $1) DESC" in sql_query:
|
||||
matches = sorted(matches, key=lambda r: r["request_id"] == lookup_id, reverse=True)
|
||||
return matches[:1]
|
||||
|
||||
class MockPrisma:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
|
||||
return MockPrisma()
|
||||
|
||||
|
||||
def _payload_row(request_id, litellm_call_id, user, prompt):
|
||||
return {
|
||||
"request_id": request_id,
|
||||
"litellm_call_id": litellm_call_id,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"response": {"id": request_id},
|
||||
"proxy_server_request": None,
|
||||
"metadata": None,
|
||||
"user": user,
|
||||
"team_id": None,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_request_response_collision_serves_callers_own_row(client, monkeypatch):
|
||||
"""The attacker's row carries the victim's request_id as its client-set call id
|
||||
and was written first. Each tenant's detail lookup of that id serves only their
|
||||
own payload, and an admin's lookup resolves the exact request_id match rather
|
||||
than whichever colliding row the database happens to return first."""
|
||||
prisma = _make_payload_lookup_prisma(
|
||||
[
|
||||
_payload_row("attacker-req", "victim-req", "attacker_user", "attacker prompt"),
|
||||
_payload_row("victim-req", "victim-call-id", "victim_user", "victim prompt"),
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
|
||||
try:
|
||||
for role, user_id, own_prompt, other_prompt in (
|
||||
(LitellmUserRoles.INTERNAL_USER, "victim_user", "victim prompt", "attacker prompt"),
|
||||
(LitellmUserRoles.INTERNAL_USER, "attacker_user", "attacker prompt", "victim prompt"),
|
||||
(LitellmUserRoles.PROXY_ADMIN, "admin", "victim prompt", "attacker prompt"),
|
||||
):
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda role=role, user_id=user_id: UserAPIKeyAuth(
|
||||
user_role=role, user_id=user_id
|
||||
)
|
||||
response = client.get("/spend/logs/ui/victim-req", headers={"Authorization": "Bearer sk-test"})
|
||||
assert response.status_code == 200, response.text
|
||||
assert own_prompt in response.text
|
||||
assert other_prompt not in response.text
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_request_response_rejects_foreign_row_inserted_after_owner_check(client, monkeypatch):
|
||||
"""Backstop behind the SQL scope on the detail endpoint (the mock ignores the
|
||||
scope on purpose): the payload row fetched by id is itself ownership-checked, so
|
||||
a foreign row the scope did not filter cannot have its payload served."""
|
||||
|
||||
class MockDB:
|
||||
async def query_raw(self, sql_query, *params):
|
||||
if 'SELECT DISTINCT "user", team_id' in sql_query:
|
||||
return [{"user": "user_1", "team_id": None}]
|
||||
return [
|
||||
{
|
||||
"messages": [{"role": "user", "content": "victim prompt"}],
|
||||
"response": {"id": "resp-1"},
|
||||
"proxy_server_request": None,
|
||||
"metadata": None,
|
||||
"user": "victim_user",
|
||||
"team_id": None,
|
||||
}
|
||||
]
|
||||
|
||||
class MockPrisma:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrisma())
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1"
|
||||
)
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs/ui/shared-id",
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 403
|
||||
assert "victim prompt" not in response.text
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_request_response_custom_logger_denies_foreign_payload_owner(client, monkeypatch):
|
||||
"""The custom-logger payload comes straight from cold storage, written independently
|
||||
of the spend-log table and able to outlive its row. When an id lookup matches no row,
|
||||
the DB owner pre-check has nothing to verify, so the payload is authorized against the
|
||||
owner recorded inside it. A foreign tenant's stored payload is denied even though no
|
||||
spend-log row exists for the pre-check to catch."""
|
||||
|
||||
class MockDB:
|
||||
async def query_raw(self, sql_query, *params):
|
||||
return []
|
||||
|
||||
class MockPrisma:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
|
||||
class ColdStorageLogger:
|
||||
async def get_request_response_payload(self, request_id, start_time_utc, end_time_utc):
|
||||
return {
|
||||
"messages": [{"role": "user", "content": "victim prompt"}],
|
||||
"response": {"id": "r"},
|
||||
"metadata": {"user_api_key_user_id": "victim_user", "user_api_key_team_id": None},
|
||||
}
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrisma())
|
||||
monkeypatch.setattr(
|
||||
litellm.logging_callback_manager,
|
||||
"get_active_additional_logging_utils_from_custom_logger",
|
||||
lambda: [ColdStorageLogger()],
|
||||
)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1"
|
||||
)
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs/ui/shared-id",
|
||||
params={"start_date": "2026-01-01 00:00:00"},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 403
|
||||
assert "victim prompt" not in response.text
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_request_response_custom_logger_is_keyed_by_callers_own_request_id(client, monkeypatch):
|
||||
"""Cold storage is keyed by the provider request_id. The caller's row carries the
|
||||
lookup id only as its client-set litellm_call_id while another tenant's row owns
|
||||
that id as its request_id. The custom logger is asked for the caller's own stored
|
||||
request_id, so the caller gets their payload rather than a 403 from the foreign
|
||||
payload's owner check, and the foreign payload is never fetched."""
|
||||
prisma = _make_payload_lookup_prisma(
|
||||
[
|
||||
_payload_row("shared-id", "other-call-id", "other_user", "other tenant prompt"),
|
||||
_payload_row("caller-req", "shared-id", "caller_user", "caller prompt"),
|
||||
]
|
||||
)
|
||||
cold_storage = {
|
||||
"shared-id": {
|
||||
"messages": [{"role": "user", "content": "other tenant prompt"}],
|
||||
"response": {"id": "shared-id"},
|
||||
"metadata": {"user_api_key_user_id": "other_user", "user_api_key_team_id": None},
|
||||
},
|
||||
"caller-req": {
|
||||
"messages": [{"role": "user", "content": "caller prompt"}],
|
||||
"response": {"id": "caller-req"},
|
||||
"metadata": {"user_api_key_user_id": "caller_user", "user_api_key_team_id": None},
|
||||
},
|
||||
}
|
||||
requested_ids = []
|
||||
|
||||
class ColdStorageLogger:
|
||||
async def get_request_response_payload(self, request_id, start_time_utc, end_time_utc):
|
||||
requested_ids.append(request_id)
|
||||
return cold_storage.get(request_id)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
|
||||
monkeypatch.setattr(
|
||||
litellm.logging_callback_manager,
|
||||
"get_active_additional_logging_utils_from_custom_logger",
|
||||
lambda: [ColdStorageLogger()],
|
||||
)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER, user_id="caller_user"
|
||||
)
|
||||
try:
|
||||
response = client.get("/spend/logs/ui/shared-id", headers={"Authorization": "Bearer sk-test"})
|
||||
assert response.status_code == 200, response.text
|
||||
assert "caller prompt" in response.text
|
||||
assert "other tenant prompt" not in response.text
|
||||
assert requested_ids == ["caller-req"]
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("group_by_session", [False, True])
|
||||
async def test_ui_view_spend_logs_id_lookup_lists_exact_request_id_row_first(client, monkeypatch, group_by_session):
|
||||
"""The dashboard's deep link fetches a single row for ``?log_id=``. When a newer
|
||||
row carries that id as its client-set litellm_call_id, the row whose request_id
|
||||
is the id still comes first, so the link opens the request it names. The
|
||||
session-grouped page orders its representatives the same way."""
|
||||
today = datetime.datetime.now(timezone.utc)
|
||||
corpus = [
|
||||
{
|
||||
"id": "log_colliding",
|
||||
"request_id": "colliding-req",
|
||||
"litellm_call_id": "victim-req",
|
||||
"api_key": "sk-test-key",
|
||||
"user": "other_user",
|
||||
"team_id": None,
|
||||
"spend": 0.01,
|
||||
"startTime": today.isoformat(),
|
||||
"model": "gpt-4",
|
||||
},
|
||||
{
|
||||
"id": "log_victim",
|
||||
"request_id": "victim-req",
|
||||
"litellm_call_id": "victim-call-id",
|
||||
"api_key": "sk-test-key",
|
||||
"user": "victim_user",
|
||||
"team_id": None,
|
||||
"spend": 0.02,
|
||||
"startTime": (today - datetime.timedelta(minutes=5)).isoformat(),
|
||||
"model": "gpt-4",
|
||||
},
|
||||
]
|
||||
|
||||
def filter_fn(where):
|
||||
rows = _filter_logs_by_date_range(corpus, where)
|
||||
rid_either = where.get("request_id_or_call_id")
|
||||
if rid_either:
|
||||
return [r for r in rows if rid_either in (r["request_id"], r["litellm_call_id"])]
|
||||
return rows
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", make_ui_spend_logs_mock_prisma(corpus, filter_fn))
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"
|
||||
)
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={"request_id": "victim-req", "page_size": 1, "group_by_session": str(group_by_session).lower()},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
data = response.json()
|
||||
assert data["total"] == 2
|
||||
assert [row["request_id"] for row in data["data"]] == ["victim-req"]
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_request_id_owner_lookup_drops_window_keeps_scope(
|
||||
client, monkeypatch
|
||||
):
|
||||
"""A non-admin owner looking up their own request_id resolves across all time.
|
||||
The ownership check authorizes the single row, so the query drops both the date
|
||||
window and the general user/team scoping and filters by the primary key alone;
|
||||
without that skip an internal user would have a `user`/`OR` clause added."""
|
||||
"""A non-admin owner looking up their own request_id resolves across all time:
|
||||
the query drops the date window the dashboard sends, while the caller's own-user
|
||||
scope stays on the id lookup so a colliding foreign row can never be served."""
|
||||
today = datetime.datetime.now(timezone.utc)
|
||||
mock_spend_logs = [
|
||||
{
|
||||
|
|
@ -2361,20 +2991,14 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only(
|
|||
def filter_fn(where):
|
||||
captured["where"] = where
|
||||
rows = _filter_logs_by_date_range(mock_spend_logs, where)
|
||||
if where.get("request_id"):
|
||||
rid_either = where.get("request_id_or_call_id")
|
||||
if rid_either:
|
||||
rows = [r for r in rows if rid_either in (r["request_id"], r.get("litellm_call_id"))]
|
||||
elif where.get("request_id"):
|
||||
rows = [r for r in rows if r["request_id"] == where["request_id"]]
|
||||
return rows
|
||||
|
||||
mock_prisma = make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn)
|
||||
|
||||
class _OwnedRow:
|
||||
user = "user_1"
|
||||
team_id = "team1"
|
||||
|
||||
async def _find_unique(where, include=None):
|
||||
return _OwnedRow()
|
||||
|
||||
mock_prisma.db.find_unique = _find_unique
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
|
||||
# A 5-day window that EXCLUDES the 90-day-old log, as the dashboard sends.
|
||||
|
|
@ -2399,9 +3023,8 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only(
|
|||
assert data["total"] == 1
|
||||
assert data["data"][0]["request_id"] == "req-old"
|
||||
assert "startTime" not in captured["where"]
|
||||
assert captured["where"]["request_id"] == "req-old"
|
||||
assert "user" not in captured["where"]
|
||||
assert "OR" not in captured["where"]
|
||||
assert captured["where"]["request_id_or_call_id"] == "req-old"
|
||||
assert captured["where"]["user"] == "user_1"
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
|
@ -5095,7 +5718,7 @@ async def test_view_spend_logs_internal_user_combines_user_with_request_id(
|
|||
where = mock_client.db.captured_where
|
||||
assert where is not None
|
||||
assert where["user"] == "internal-user-2"
|
||||
assert where["request_id"] == "req-abc"
|
||||
assert where["OR"] == ({"request_id": "req-abc"}, {"litellm_call_id": "req-abc"})
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
|
@ -5122,7 +5745,7 @@ async def test_view_spend_logs_non_date_range_combines_user_with_request_id(
|
|||
where = mock_client.db.captured_where
|
||||
assert where is not None
|
||||
assert where["user"] == "internal-user-3"
|
||||
assert where["request_id"] == "req-xyz"
|
||||
assert where["OR"] == ({"request_id": "req-xyz"}, {"litellm_call_id": "req-xyz"})
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
|
@ -7052,25 +7675,19 @@ async def test_ui_view_spend_logs_search_returns_flat_rows_when_grouping_by_sess
|
|||
|
||||
|
||||
def _fake_prisma_with_owned_spend_log(owner_user_id, messages_json, response_json):
|
||||
class _Row:
|
||||
user = owner_user_id
|
||||
team_id = None
|
||||
|
||||
class _SpendLogs:
|
||||
async def find_unique(self, where, include=None):
|
||||
return _Row()
|
||||
|
||||
class _DB:
|
||||
def __init__(self):
|
||||
self.litellm_spendlogs = _SpendLogs()
|
||||
|
||||
async def query_raw(self, _sql, *_args):
|
||||
async def query_raw(self, sql, *_args):
|
||||
if 'SELECT DISTINCT "user", team_id' in sql:
|
||||
return [{"user": owner_user_id, "team_id": None}]
|
||||
return [
|
||||
{
|
||||
"request_id": "req-owned-by-user-a",
|
||||
"messages": messages_json,
|
||||
"response": response_json,
|
||||
"proxy_server_request": "{}",
|
||||
"metadata": "{}",
|
||||
"user": owner_user_id,
|
||||
"team_id": None,
|
||||
}
|
||||
]
|
||||
|
||||
|
|
@ -7127,8 +7744,8 @@ def test_ui_view_request_response_internal_user_non_owner_forbidden(client, monk
|
|||
"""
|
||||
A different internal_user requesting someone else's row is forbidden;
|
||||
guards against _assert_user_can_view_request_id being skipped in the
|
||||
detail-drawer handler. Also proves the handler stops before it ever asks
|
||||
a custom logger or the DB for the payload.
|
||||
detail-drawer handler. Also proves the handler stops at the owner lookup,
|
||||
before it ever asks a custom logger or the DB for the payload.
|
||||
"""
|
||||
messages_json = json.dumps([{"role": "user", "content": "hi"}])
|
||||
response_json = json.dumps({"choices": [{"message": {"content": "hello"}}]})
|
||||
|
|
@ -7160,7 +7777,7 @@ def test_ui_view_request_response_internal_user_non_owner_forbidden(client, monk
|
|||
)
|
||||
assert response.status_code == 403
|
||||
assert custom_logger.requested_ids == []
|
||||
assert query_raw_calls == []
|
||||
assert [args[0] for args, _kwargs in query_raw_calls if "messages" in args[0]] == []
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
|
@ -7171,24 +7788,27 @@ def test_ui_view_request_response_internal_user_missing_row_forbidden(client, mo
|
|||
request_id with no spend-log row (e.g. pruned by retention) must be
|
||||
denied before the handler ever consults a custom logger, otherwise a
|
||||
non-admin who guesses/obtains a request_id could read another tenant's
|
||||
payload out of cold storage. Fails if `if row is None: return` is
|
||||
reintroduced.
|
||||
payload out of cold storage, and an existing payload that is not theirs
|
||||
would still confirm the id exists. Even a payload recorded as the caller's
|
||||
own is never fetched once the row is gone. Fails if an empty owner lookup
|
||||
is allowed to fall through to the loggers.
|
||||
"""
|
||||
|
||||
class _SpendLogs:
|
||||
async def find_unique(self, where, include=None):
|
||||
return None
|
||||
|
||||
class _DB:
|
||||
def __init__(self):
|
||||
self.litellm_spendlogs = _SpendLogs()
|
||||
async def query_raw(self, _sql, *_args):
|
||||
return []
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
fake_prisma = SimpleNamespace(db=_DB())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", fake_prisma)
|
||||
|
||||
custom_logger = _RecordingAdditionalLoggingUtils({"messages": "should-not-be-returned"})
|
||||
custom_logger = _RecordingAdditionalLoggingUtils(
|
||||
{
|
||||
"messages": "should-not-be-returned",
|
||||
"metadata": {"user_api_key_user_id": "user_a", "user_api_key_team_id": None},
|
||||
}
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
litellm.logging_callback_manager,
|
||||
"get_active_additional_logging_utils_from_custom_logger",
|
||||
|
|
|
|||
|
|
@ -1329,6 +1329,33 @@ def test_get_logging_payload_includes_agent_id_from_kwargs():
|
|||
assert payload["agent_id"] == test_agent_id, f"Expected agent_id '{test_agent_id}', got '{payload.get('agent_id')}'"
|
||||
|
||||
|
||||
def test_get_logging_payload_populates_litellm_call_id_alongside_provider_request_id():
|
||||
"""
|
||||
LIT-6302: request_id stays the provider response id, so clients holding the
|
||||
x-litellm-call-id header value could never find their row. The payload now
|
||||
also carries litellm_call_id as its own column for lookups by either id.
|
||||
"""
|
||||
call_id = "b980eea9-5cd9-4099-93cd-8291e46c76fd"
|
||||
|
||||
payload = get_logging_payload(
|
||||
kwargs={
|
||||
"model": "gpt-4o-mini",
|
||||
"litellm_call_id": call_id,
|
||||
"litellm_params": {"metadata": {"user_api_key": "test-key"}},
|
||||
},
|
||||
response_obj=litellm.ModelResponse(
|
||||
id="chatcmpl-provider-id",
|
||||
choices=[],
|
||||
usage=litellm.Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2),
|
||||
),
|
||||
start_time=datetime.datetime.now(timezone.utc),
|
||||
end_time=datetime.datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
assert payload["request_id"] == "chatcmpl-provider-id"
|
||||
assert payload["litellm_call_id"] == call_id
|
||||
|
||||
|
||||
@patch("litellm.proxy.proxy_server.master_key", None)
|
||||
@patch("litellm.proxy.proxy_server.general_settings", {})
|
||||
def test_get_logging_payload_includes_overhead_in_spend_logs_metadata():
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ from fastapi.responses import JSONResponse, StreamingResponse
|
|||
|
||||
import litellm
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY
|
||||
from litellm.constants import MAX_LITELLM_CALL_ID_LENGTH, RETURN_RAW_MODEL_NAME_METADATA_KEY
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.opentelemetry import UserAPIKeyAuth
|
||||
from litellm.proxy.common_request_processing import (
|
||||
|
|
@ -30,6 +30,7 @@ from litellm.proxy.common_request_processing import (
|
|||
_has_attribute_error_in_chain,
|
||||
_is_azure_model_router_request,
|
||||
open_sse_before_first_byte,
|
||||
resolve_litellm_call_id,
|
||||
ttft_keepalive_interval,
|
||||
_override_openai_response_model,
|
||||
_parse_event_data_for_error,
|
||||
|
|
@ -8060,6 +8061,19 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_
|
|||
assert (records[0].exc_info is not None) is expect_traceback
|
||||
|
||||
|
||||
class TestResolveLitellmCallId:
|
||||
def test_client_call_id_within_the_bound_is_kept(self):
|
||||
assert resolve_litellm_call_id("req-abc-123") == "req-abc-123"
|
||||
at_bound: Final = "y" * MAX_LITELLM_CALL_ID_LENGTH
|
||||
assert resolve_litellm_call_id(at_bound) == at_bound
|
||||
|
||||
@pytest.mark.parametrize("client_call_id", [None, "", "x" * (MAX_LITELLM_CALL_ID_LENGTH + 1), "z" * 3000])
|
||||
def test_missing_empty_or_oversized_client_call_id_gets_a_generated_uuid(self, client_call_id):
|
||||
resolved: Final = resolve_litellm_call_id(client_call_id)
|
||||
assert resolved != client_call_id
|
||||
assert uuid.UUID(resolved).version == 4
|
||||
|
||||
|
||||
class _FailureHookRecorder:
|
||||
"""Stands in for ProxyLogging.post_call_failure_hook, recording what the detached-failure closure hands it."""
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,222 @@
|
|||
"""Unit tests for the Friendli transform in
|
||||
`.github/scripts/auto_update_price_and_context_window_file.py`."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
SCRIPT_PATH = (
|
||||
Path(__file__).resolve().parents[2]
|
||||
/ ".github"
|
||||
/ "scripts"
|
||||
/ "auto_update_price_and_context_window_file.py"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def sync_module():
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"auto_update_price_and_context_window_file", SCRIPT_PATH
|
||||
)
|
||||
assert spec and spec.loader, f"Could not load spec for {SCRIPT_PATH}"
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules["auto_update_price_and_context_window_file"] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _reasoning_model(**overrides: object) -> dict:
|
||||
model = {
|
||||
"id": "zai-org/GLM-Test",
|
||||
"base_model": "zhipuai/glm-test",
|
||||
"context_length": 1048576,
|
||||
"max_completion_tokens": 131072,
|
||||
"pricing": {"input": "0.00000015", "output": "0.0000005", "input_cache_read": "0.00000003"},
|
||||
"reasoning": True,
|
||||
"reasoning_options": [{"type": "effort", "values": ["max", "high", "low"]}],
|
||||
"functionality": {
|
||||
"tool_call": True,
|
||||
"parallel_tool_call": True,
|
||||
"structured_output": True,
|
||||
"system_messages": True,
|
||||
"tool_choice": True,
|
||||
},
|
||||
"input_modalities": ["text", "image", "video"],
|
||||
"mode": "chat",
|
||||
}
|
||||
model.update(overrides)
|
||||
return model
|
||||
|
||||
|
||||
def test_transform_emits_declared_effort_levels_in_canonical_order(sync_module):
|
||||
entry = sync_module.transform_friendli_data([_reasoning_model()], {})[
|
||||
"friendliai/zai-org/GLM-Test"
|
||||
]
|
||||
assert entry["supports_reasoning"] is True
|
||||
assert entry["reasoning_effort_levels"] == ["low", "high", "max"]
|
||||
assert not any(k.endswith("_reasoning_effort") for k in entry)
|
||||
|
||||
|
||||
def test_transform_reasoning_model_without_effort_options_declares_empty_levels(sync_module):
|
||||
model = _reasoning_model(reasoning_options=[{"type": "budget_tokens", "values": []}])
|
||||
entry = sync_module.transform_friendli_data([model], {})["friendliai/zai-org/GLM-Test"]
|
||||
assert entry["reasoning_effort_levels"] == []
|
||||
|
||||
|
||||
def test_transform_non_reasoning_model_declares_no_levels(sync_module):
|
||||
model = _reasoning_model(reasoning=False, reasoning_options=[])
|
||||
entry = sync_module.transform_friendli_data([model], {})["friendliai/zai-org/GLM-Test"]
|
||||
assert entry["supports_reasoning"] is False
|
||||
assert "reasoning_effort_levels" not in entry
|
||||
|
||||
|
||||
def test_transform_max_tokens_mirrors_output_cap_not_context(sync_module):
|
||||
entry = sync_module.transform_friendli_data([_reasoning_model()], {})[
|
||||
"friendliai/zai-org/GLM-Test"
|
||||
]
|
||||
assert entry["max_input_tokens"] == 1048576
|
||||
assert entry["max_output_tokens"] == 131072
|
||||
assert entry["max_tokens"] == entry["max_output_tokens"]
|
||||
|
||||
|
||||
def test_transform_prompt_caching_follows_cache_pricing(sync_module):
|
||||
cached = sync_module.transform_friendli_data([_reasoning_model()], {})[
|
||||
"friendliai/zai-org/GLM-Test"
|
||||
]
|
||||
assert cached["supports_prompt_caching"] is True
|
||||
assert cached["cache_read_input_token_cost"] == 3e-08
|
||||
|
||||
uncached_model = _reasoning_model(pricing={"input": "0.00000014", "output": "0.0000004"})
|
||||
uncached = sync_module.transform_friendli_data([uncached_model], {})[
|
||||
"friendliai/zai-org/GLM-Test"
|
||||
]
|
||||
assert uncached["supports_prompt_caching"] is False
|
||||
assert "cache_read_input_token_cost" not in uncached
|
||||
|
||||
|
||||
def test_transform_modalities_set_vision_image_and_video_flags(sync_module):
|
||||
entry = sync_module.transform_friendli_data([_reasoning_model()], {})[
|
||||
"friendliai/zai-org/GLM-Test"
|
||||
]
|
||||
assert entry["supports_vision"] is True
|
||||
assert entry["supports_image_input"] is True
|
||||
assert entry["supports_video_input"] is True
|
||||
|
||||
text_only = _reasoning_model(input_modalities=["text"])
|
||||
entry_text = sync_module.transform_friendli_data([text_only], {})[
|
||||
"friendliai/zai-org/GLM-Test"
|
||||
]
|
||||
assert entry_text["supports_vision"] is False
|
||||
assert entry_text["supports_image_input"] is False
|
||||
assert entry_text["supports_video_input"] is False
|
||||
|
||||
|
||||
def test_transform_skips_rows_without_valid_token_prices_so_priced_local_entries_survive(sync_module):
|
||||
local = {
|
||||
"friendliai/zai-org/GLM-Test": {
|
||||
"litellm_provider": "friendliai",
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"output_cost_per_token": 5e-07,
|
||||
}
|
||||
}
|
||||
unpriced_rows = [
|
||||
_reasoning_model(pricing={}),
|
||||
_reasoning_model(pricing=None),
|
||||
_reasoning_model(pricing={"input": "0.00000015"}),
|
||||
_reasoning_model(pricing={"output": "0.0000005"}),
|
||||
_reasoning_model(pricing={"input": "not-a-number", "output": "0.0000005"}),
|
||||
_reasoning_model(pricing={"input": "-0.00000015", "output": "0.0000005"}),
|
||||
_reasoning_model(pricing={"input": "inf", "output": "0.0000005"}),
|
||||
_reasoning_model(pricing={"input": "nan", "output": "0.0000005"}),
|
||||
]
|
||||
remote = sync_module.transform_friendli_data(unpriced_rows, local)
|
||||
assert remote == {}
|
||||
sync_module.sync_local_data_with_remote(local, remote, replace_keys=frozenset(remote))
|
||||
assert local["friendliai/zai-org/GLM-Test"]["input_cost_per_token"] == 1.5e-07
|
||||
assert local["friendliai/zai-org/GLM-Test"]["output_cost_per_token"] == 5e-07
|
||||
|
||||
|
||||
def test_transform_keeps_zero_priced_rows(sync_module):
|
||||
free_model = _reasoning_model(pricing={"input": "0", "output": "0"})
|
||||
entry = sync_module.transform_friendli_data([free_model], {})["friendliai/zai-org/GLM-Test"]
|
||||
assert entry["input_cost_per_token"] == 0.0
|
||||
assert entry["output_cost_per_token"] == 0.0
|
||||
|
||||
|
||||
def test_transforms_survive_failed_fetch(sync_module):
|
||||
assert sync_module.transform_friendli_data(None, {}) == {}
|
||||
assert sync_module.transform_friendli_data([], {}) == {}
|
||||
assert sync_module.transform_openrouter_data(None) == {}
|
||||
assert sync_module.transform_vercel_ai_gateway_data(None) == {}
|
||||
|
||||
|
||||
def test_vercel_transform_skips_rows_without_token_pricing_or_limits(sync_module):
|
||||
rows = [
|
||||
{
|
||||
"id": "wan-video",
|
||||
"pricing": {"video_duration_pricing": [{"resolution": "720p", "cost_per_second": "0.1"}]},
|
||||
},
|
||||
{
|
||||
"id": "qwen3-embedding",
|
||||
"context_window": 32768,
|
||||
"max_tokens": 32768,
|
||||
"pricing": {"input": "0.00000001"},
|
||||
},
|
||||
{
|
||||
"id": "no-limits-chat",
|
||||
"pricing": {"input": "0.000001", "output": "0.000002"},
|
||||
},
|
||||
{
|
||||
"id": "good-chat",
|
||||
"context_window": 128000,
|
||||
"max_tokens": 8192,
|
||||
"pricing": {"input": "0.000001", "output": "0.000002"},
|
||||
},
|
||||
]
|
||||
transformed = sync_module.transform_vercel_ai_gateway_data(rows)
|
||||
assert list(transformed) == ["vercel_ai_gateway/good-chat"]
|
||||
assert transformed["vercel_ai_gateway/good-chat"]["input_cost_per_token"] == 1e-06
|
||||
assert transformed["vercel_ai_gateway/good-chat"]["output_cost_per_token"] == 2e-06
|
||||
|
||||
|
||||
def test_sync_replaces_friendli_entries_so_dropped_cache_pricing_does_not_survive(sync_module):
|
||||
local = {
|
||||
"friendliai/zai-org/GLM-Test": {
|
||||
"litellm_provider": "friendliai",
|
||||
"cache_read_input_token_cost": 3e-08,
|
||||
"supports_prompt_caching": True,
|
||||
}
|
||||
}
|
||||
uncached_model = _reasoning_model(pricing={"input": "0.00000014", "output": "0.0000004"})
|
||||
remote = sync_module.transform_friendli_data([uncached_model], local)
|
||||
sync_module.sync_local_data_with_remote(local, remote, replace_keys=frozenset(remote))
|
||||
synced = local["friendliai/zai-org/GLM-Test"]
|
||||
assert "cache_read_input_token_cost" not in synced
|
||||
assert synced["supports_prompt_caching"] is False
|
||||
|
||||
|
||||
def test_sync_still_merges_entries_outside_replace_keys(sync_module):
|
||||
local = {"openrouter/some-model": {"input_cost_per_token": 1e-06, "supports_vision": True}}
|
||||
remote = {"openrouter/some-model": {"input_cost_per_token": 2e-06}}
|
||||
sync_module.sync_local_data_with_remote(local, remote)
|
||||
assert local["openrouter/some-model"] == {"input_cost_per_token": 2e-06, "supports_vision": True}
|
||||
|
||||
|
||||
def test_transform_inherits_allowlisted_keys_from_base_model_entry(sync_module):
|
||||
local = {
|
||||
"zhipuai/glm-test": {
|
||||
"supports_pdf_input": True,
|
||||
"supports_assistant_prefill": True,
|
||||
"input_cost_per_token": 9e-06,
|
||||
}
|
||||
}
|
||||
entry = sync_module.transform_friendli_data([_reasoning_model()], local)[
|
||||
"friendliai/zai-org/GLM-Test"
|
||||
]
|
||||
assert entry["supports_pdf_input"] is True
|
||||
assert entry["supports_assistant_prefill"] is True
|
||||
assert entry["input_cost_per_token"] == 1.5e-07
|
||||
|
|
@ -4313,6 +4313,59 @@ def test_select_model_name_applies_region_to_private_provider_response_model(_lo
|
|||
assert selected == "bedrock/us-east-1/anthropic.claude-v2:1"
|
||||
|
||||
|
||||
def test_completion_cost_region_name_prices_mantle_on_the_regional_row(_local_model_cost_map):
|
||||
"""completion_cost(region_name=...) must price a Bedrock Mantle call from the
|
||||
bedrock_mantle/<region>/<model> row when one exists, for the bare and the provider-prefixed
|
||||
model alike, and keep the flat row for regions without their own row."""
|
||||
|
||||
response = litellm.ModelResponse(
|
||||
id="x",
|
||||
choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
|
||||
model="xai.grok-4.3",
|
||||
usage={"prompt_tokens": 38, "completion_tokens": 20, "total_tokens": 58},
|
||||
)
|
||||
gov = litellm.model_cost["bedrock_mantle/us-gov-west-1/xai.grok-4.3"]
|
||||
flat = litellm.model_cost["bedrock_mantle/xai.grok-4.3"]
|
||||
expected_gov = 38 * gov["input_cost_per_token"] + 20 * gov["output_cost_per_token"]
|
||||
expected_flat = 38 * flat["input_cost_per_token"] + 20 * flat["output_cost_per_token"]
|
||||
assert expected_gov != expected_flat
|
||||
|
||||
for model in ("xai.grok-4.3", "bedrock_mantle/xai.grok-4.3"):
|
||||
assert litellm.completion_cost(
|
||||
completion_response=response,
|
||||
model=model,
|
||||
custom_llm_provider="bedrock_mantle",
|
||||
region_name="us-gov-west-1",
|
||||
) == pytest.approx(expected_gov)
|
||||
assert litellm.completion_cost(
|
||||
completion_response=response,
|
||||
model=model,
|
||||
custom_llm_provider="bedrock_mantle",
|
||||
region_name="eu-west-1",
|
||||
) == pytest.approx(expected_flat)
|
||||
assert litellm.completion_cost(
|
||||
completion_response=response, model="xai.grok-4.3", custom_llm_provider="bedrock_mantle"
|
||||
) == pytest.approx(expected_flat)
|
||||
|
||||
|
||||
def test_cost_per_token_region_name_applies_to_provider_prefixed_model(_local_model_cost_map):
|
||||
"""A provider-prefixed model must still find its bedrock_mantle/<region>/<model> row instead of
|
||||
composing the region key with the provider segment twice."""
|
||||
|
||||
prompt_cost, completion_cost = litellm.cost_per_token(
|
||||
model="bedrock_mantle/xai.grok-4.3",
|
||||
prompt_tokens=38,
|
||||
completion_tokens=20,
|
||||
custom_llm_provider="bedrock_mantle",
|
||||
region_name="us-gov-west-1",
|
||||
)
|
||||
gov = litellm.model_cost["bedrock_mantle/us-gov-west-1/xai.grok-4.3"]
|
||||
|
||||
assert prompt_cost + completion_cost == pytest.approx(
|
||||
38 * gov["input_cost_per_token"] + 20 * gov["output_cost_per_token"]
|
||||
)
|
||||
|
||||
|
||||
def test_select_model_name_keeps_base_model_free_of_region(_local_model_cost_map):
|
||||
"""An explicit base_model keeps pricing on that model's own key even when the request carries a
|
||||
region with different regional rates, so the private provider model never widens region pricing."""
|
||||
|
|
@ -4329,6 +4382,29 @@ def test_select_model_name_keeps_base_model_free_of_region(_local_model_cost_map
|
|||
assert selected == "bedrock/moonshotai.kimi-k2.5"
|
||||
|
||||
|
||||
def test_completion_cost_base_model_ignores_regional_row(_local_model_cost_map):
|
||||
"""A deployment with base_model set is priced from that model's own row even when the response
|
||||
carries a region whose regional row charges different rates."""
|
||||
|
||||
response = litellm.ModelResponse(
|
||||
id="x",
|
||||
choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
|
||||
model="my-bedrock-deployment",
|
||||
usage={"prompt_tokens": 1000, "completion_tokens": 0, "total_tokens": 1000},
|
||||
)
|
||||
response._hidden_params = {"custom_llm_provider": "bedrock", "region_name": "eu-central-1"}
|
||||
flat = litellm.model_cost["anthropic.claude-instant-v1"]
|
||||
regional = litellm.model_cost["bedrock/eu-central-1/anthropic.claude-instant-v1"]
|
||||
assert flat["input_cost_per_token"] != regional["input_cost_per_token"]
|
||||
|
||||
assert litellm.completion_cost(
|
||||
completion_response=response,
|
||||
model="my-bedrock-deployment",
|
||||
custom_llm_provider="bedrock",
|
||||
base_model="anthropic.claude-instant-v1",
|
||||
) == pytest.approx(1000 * flat["input_cost_per_token"])
|
||||
|
||||
|
||||
def test_completion_cost_nonzero_for_slash_alias_model_name(_local_model_cost_map):
|
||||
"""End-to-end cost through a "/"-containing alias must price above zero (#38069)."""
|
||||
|
||||
|
|
|
|||
|
|
@ -5625,6 +5625,41 @@ def test_get_deployment_credentials_with_provider_preserves_aws_auth_params():
|
|||
assert credentials.get(key) == value, key
|
||||
|
||||
|
||||
def test_get_deployment_credentials_preserves_azure_entra_id_params():
|
||||
entra_params = {
|
||||
"tenant_id": "deployment-tenant",
|
||||
"client_id": "deployment-client",
|
||||
"client_secret": "deployment-client-secret",
|
||||
"azure_scope": "https://cognitiveservices.azure.us/.default",
|
||||
"azure_username": "deployment-user",
|
||||
"azure_password": "deployment-password",
|
||||
}
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "azure-entra-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-5.4",
|
||||
"api_base": "https://example.openai.azure.com/",
|
||||
"api_version": "2024-10-21",
|
||||
**entra_params,
|
||||
},
|
||||
"model_info": {"id": "azure-entra-model-id"},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
credentials = router.get_deployment_credentials(model_id="azure-entra-model-id")
|
||||
credentials_with_provider = router.get_deployment_credentials_with_provider(model_id="azure-entra-model-id")
|
||||
|
||||
assert credentials is not None
|
||||
assert credentials_with_provider is not None
|
||||
assert "api_key" not in credentials
|
||||
for key, value in entra_params.items():
|
||||
assert credentials.get(key) == value, key
|
||||
assert credentials_with_provider.get(key) == value, key
|
||||
|
||||
|
||||
def _team_wildcard_model(api_key: str, model_id: str = "team-wildcard-id") -> dict:
|
||||
return {
|
||||
"model_name": f"model_name_team-1_{model_id}",
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ const AllModelsTab = ({
|
|||
setSelectedTeamId,
|
||||
}: AllModelsTabProps) => {
|
||||
const { data: modelCostMapData, isLoading: isLoadingModelCostMap } = useModelCostMap();
|
||||
const { accessToken, userId, userRole } = useAuthorized();
|
||||
const { accessToken, userId, userRole, isViewOnly } = useAuthorized();
|
||||
const { data: teams, isLoading: isLoadingTeams } = useTeams();
|
||||
const queryClient = useQueryClient();
|
||||
|
||||
|
|
@ -281,6 +281,7 @@ const AllModelsTab = ({
|
|||
availableModelAccessGroups={availableModelAccessGroups}
|
||||
userRole={userRole}
|
||||
userID={userId}
|
||||
isViewOnly={isViewOnly}
|
||||
onModelIdClick={setSelectedModelId}
|
||||
onTeamIdClick={setSelectedTeamId}
|
||||
onDeleteClick={handleDeleteClick}
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ const baseProps = {
|
|||
availableModelAccessGroups: ["sales-team"],
|
||||
userRole: "Admin",
|
||||
userID: "alice",
|
||||
isViewOnly: false,
|
||||
onModelIdClick: vi.fn(),
|
||||
onTeamIdClick: vi.fn(),
|
||||
onDeleteClick: vi.fn(),
|
||||
|
|
@ -254,6 +255,17 @@ describe("AllModelsTable", () => {
|
|||
expect(onTogglePauseClick).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("does not let a view-only admin toggle a model", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onTogglePauseClick = vi.fn();
|
||||
render(<AllModelsTable {...baseProps} isViewOnly onTogglePauseClick={onTogglePauseClick} />);
|
||||
|
||||
const toggle = screen.getByTestId("model-pause-toggle-model-1");
|
||||
expect(toggle).toHaveAttribute("data-disabled");
|
||||
await user.click(toggle);
|
||||
expect(onTogglePauseClick).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("does not let anyone toggle a config model", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onTogglePauseClick = vi.fn();
|
||||
|
|
@ -309,6 +321,17 @@ describe("AllModelsTable", () => {
|
|||
expect(onDeleteClick).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("blocks a view-only admin from deleting a DB model they created", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onDeleteClick = vi.fn();
|
||||
render(<AllModelsTable {...baseProps} isViewOnly onDeleteClick={onDeleteClick} />);
|
||||
|
||||
const deleteButton = screen.getByTestId("model-delete-model-1");
|
||||
expect(deleteButton).toBeDisabled();
|
||||
await user.click(deleteButton);
|
||||
expect(onDeleteClick).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("blocks deleting a config model", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onDeleteClick = vi.fn();
|
||||
|
|
|
|||
|
|
@ -73,6 +73,7 @@ interface AllModelsTableProps {
|
|||
availableModelAccessGroups: string[];
|
||||
userRole: string;
|
||||
userID: string;
|
||||
isViewOnly: boolean;
|
||||
onModelIdClick: (modelId: string) => void;
|
||||
onTeamIdClick: (teamId: string) => void;
|
||||
onDeleteClick: (modelId: string) => void;
|
||||
|
|
@ -120,6 +121,7 @@ export function AllModelsTable({
|
|||
availableModelAccessGroups,
|
||||
userRole,
|
||||
userID,
|
||||
isViewOnly,
|
||||
onModelIdClick,
|
||||
onTeamIdClick,
|
||||
onDeleteClick,
|
||||
|
|
@ -132,6 +134,7 @@ export function AllModelsTable({
|
|||
const columnDeps = {
|
||||
userRole,
|
||||
userID,
|
||||
isViewOnly,
|
||||
onModelIdClick,
|
||||
onTeamIdClick,
|
||||
onDeleteClick,
|
||||
|
|
@ -139,7 +142,7 @@ export function AllModelsTable({
|
|||
pausingModelId,
|
||||
};
|
||||
return getModelsTableColumns(columnDeps);
|
||||
}, [userRole, userID, onModelIdClick, onTeamIdClick, onDeleteClick, onTogglePauseClick, pausingModelId]);
|
||||
}, [userRole, userID, isViewOnly, onModelIdClick, onTeamIdClick, onDeleteClick, onTogglePauseClick, pausingModelId]);
|
||||
|
||||
const modelGroupOptions = useMemo(
|
||||
() => [
|
||||
|
|
|
|||
|
|
@ -247,6 +247,7 @@ interface ModelRowActionsProps {
|
|||
model: ModelData;
|
||||
userRole: string;
|
||||
userID: string;
|
||||
isViewOnly: boolean;
|
||||
isPausing: boolean;
|
||||
onDeleteClick?: (modelId: string) => void;
|
||||
onTogglePauseClick?: (modelId: string, blocked: boolean) => void | Promise<void>;
|
||||
|
|
@ -256,14 +257,15 @@ function ModelRowActions({
|
|||
model,
|
||||
userRole,
|
||||
userID,
|
||||
isViewOnly,
|
||||
isPausing,
|
||||
onDeleteClick,
|
||||
onTogglePauseClick,
|
||||
}: ModelRowActionsProps) {
|
||||
const modelId = model.model_info?.id;
|
||||
const isConfigModel = !model.model_info?.db_model;
|
||||
const isAdmin = userRole === "Admin";
|
||||
const canEditModel = isAdmin || model.model_info?.created_by === userID;
|
||||
const isAdmin = userRole === "Admin" && !isViewOnly;
|
||||
const canEditModel = !isViewOnly && (isAdmin || model.model_info?.created_by === userID);
|
||||
const isBlocked = model.model_info?.blocked === true;
|
||||
const isPauseToggleable = !isConfigModel && isAdmin && Boolean(onTogglePauseClick);
|
||||
|
||||
|
|
@ -340,6 +342,7 @@ function ModelRowActions({
|
|||
export interface ModelsTableColumnDeps {
|
||||
userRole: string;
|
||||
userID: string;
|
||||
isViewOnly: boolean;
|
||||
onModelIdClick: (modelId: string) => void;
|
||||
onTeamIdClick: (teamId: string) => void;
|
||||
onDeleteClick?: (modelId: string) => void;
|
||||
|
|
@ -350,6 +353,7 @@ export interface ModelsTableColumnDeps {
|
|||
export const getModelsTableColumns = ({
|
||||
userRole,
|
||||
userID,
|
||||
isViewOnly,
|
||||
onModelIdClick,
|
||||
onTeamIdClick,
|
||||
onDeleteClick,
|
||||
|
|
@ -479,6 +483,7 @@ export const getModelsTableColumns = ({
|
|||
model={row.original}
|
||||
userRole={userRole}
|
||||
userID={userID}
|
||||
isViewOnly={isViewOnly}
|
||||
isPausing={pausingModelId === row.original.model_info?.id}
|
||||
onDeleteClick={onDeleteClick}
|
||||
onTogglePauseClick={onTogglePauseClick}
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ vi.mock("./panels/HealthStatusPanel", () => ({ default: () => <div data-testid="
|
|||
vi.mock("./panels/ModelRetrySettingsPanel", () => ({ default: () => <div data-testid="panel-retry" /> }));
|
||||
vi.mock("./panels/ModelGroupAliasPanel", () => ({ default: () => <div data-testid="panel-alias" /> }));
|
||||
vi.mock("./panels/PriceDataPanel", () => ({ default: () => <div data-testid="panel-price" /> }));
|
||||
vi.mock("./panels/AccessGroupBudgetsPanel", () => ({ default: () => <div data-testid="panel-budgets" /> }));
|
||||
|
||||
const detailState = { modelId: null as string | null, teamId: null as string | null };
|
||||
vi.mock("./detailNavigation", () => ({
|
||||
|
|
@ -25,7 +26,11 @@ vi.mock("@/components/model_info_view", () => ({
|
|||
default: ({ modelId }: { modelId: string }) => <div data-testid="model-info">model:{modelId}</div>,
|
||||
}));
|
||||
vi.mock("@/components/team/TeamInfo", () => ({
|
||||
default: ({ teamId }: { teamId: string }) => <div data-testid="team-info">team:{teamId}</div>,
|
||||
default: ({ teamId, is_team_admin }: { teamId: string; is_team_admin: boolean }) => (
|
||||
<div data-testid="team-info" data-team-admin={String(is_team_admin)}>
|
||||
team:{teamId}
|
||||
</div>
|
||||
),
|
||||
}));
|
||||
|
||||
const mockUseAuthorized = vi.fn();
|
||||
|
|
@ -95,10 +100,19 @@ describe("ModelsAndEndpointsPage", () => {
|
|||
expect(screen.queryByRole("tab", { name: "All Models" })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("renders the team detail overlay from the ?team drill-in", () => {
|
||||
it("renders the team detail overlay from the ?team drill-in with admin edit rights", () => {
|
||||
detailState.teamId = "team-9";
|
||||
renderPage();
|
||||
expect(screen.getByTestId("team-info")).toHaveTextContent("team:team-9");
|
||||
expect(screen.getByTestId("team-info")).toHaveAttribute("data-team-admin", "true");
|
||||
});
|
||||
|
||||
it("opens the ?team drill-in without edit rights for a view-only admin", () => {
|
||||
mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN);
|
||||
detailState.teamId = "team-9";
|
||||
renderPage();
|
||||
expect(screen.getByTestId("team-info")).toHaveTextContent("team:team-9");
|
||||
expect(screen.getByTestId("team-info")).toHaveAttribute("data-team-admin", "false");
|
||||
});
|
||||
|
||||
it("hides admin-only tabs for a non-admin user", () => {
|
||||
|
|
@ -108,6 +122,35 @@ describe("ModelsAndEndpointsPage", () => {
|
|||
expect(screen.queryByRole("tab", { name: "Health Status" })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("keeps the full admin tab order for a real admin", () => {
|
||||
renderPage();
|
||||
expect(screen.getAllByRole("tab").map((tab) => tab.textContent)).toEqual([
|
||||
"All Models",
|
||||
"Add Model",
|
||||
"Auto-Routers Beta",
|
||||
"LLM Credentials",
|
||||
"Pass-Through Endpoints",
|
||||
"Health Status",
|
||||
"Model Retry Settings",
|
||||
"Model Group Alias",
|
||||
"Model Access Group Budgets Beta",
|
||||
"Price Data Reload",
|
||||
]);
|
||||
});
|
||||
|
||||
it("hides the admin write-form tabs from a view-only admin, keeping the read views", () => {
|
||||
mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN);
|
||||
renderPage();
|
||||
expect(screen.getByRole("tab", { name: "All Models" })).toBeInTheDocument();
|
||||
expect(screen.getByRole("tab", { name: "Health Status" })).toBeInTheDocument();
|
||||
expect(screen.queryByRole("tab", { name: "LLM Credentials" })).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("tab", { name: "Pass-Through Endpoints" })).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("tab", { name: "Model Retry Settings" })).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("tab", { name: "Model Group Alias" })).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("tab", { name: /Model Access Group Budgets/ })).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("tab", { name: "Price Data Reload" })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
// POST /model/new 403s a proxy_admin_viewer, so the form's tab must not render for one.
|
||||
it("hides the Add Model tab for a view-only admin session", () => {
|
||||
mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN);
|
||||
|
|
|
|||
|
|
@ -106,19 +106,16 @@ export default function ModelsAndEndpointsPage() {
|
|||
"",
|
||||
...(canCreate ? (["add"] as const) : []),
|
||||
...(isAdmin || canCreate ? (["auto-routers"] as const) : []),
|
||||
...(isAdmin
|
||||
? ([
|
||||
"llm-credentials",
|
||||
"pass-through",
|
||||
"health",
|
||||
"retry-settings",
|
||||
"model-group-alias",
|
||||
"access-group-budgets",
|
||||
"price-data",
|
||||
] as const)
|
||||
// effectiveSessionRole reports proxy_admin_viewer as "Admin", so isAdmin alone would show a
|
||||
// viewer these write-only panels; only the raw-role isViewOnly separates them. Health Status
|
||||
// stays: it is the bucket's one read view, and viewers keep read parity with admins.
|
||||
...(isAdmin && !isViewOnly ? (["llm-credentials", "pass-through"] as const) : []),
|
||||
...(isAdmin ? (["health"] as const) : []),
|
||||
...(isAdmin && !isViewOnly
|
||||
? (["retry-settings", "model-group-alias", "access-group-budgets", "price-data"] as const)
|
||||
: []),
|
||||
],
|
||||
[canCreate, isAdmin],
|
||||
[canCreate, isAdmin, isViewOnly],
|
||||
);
|
||||
|
||||
const allModelsLabel = isAdmin ? "All Models" : "Your Models";
|
||||
|
|
@ -148,7 +145,7 @@ export default function ModelsAndEndpointsPage() {
|
|||
teamId={teamId}
|
||||
onClose={close}
|
||||
accessToken={accessToken}
|
||||
is_team_admin={userRole === "Admin"}
|
||||
is_team_admin={userRole === "Admin" && !isViewOnly}
|
||||
is_proxy_admin={userRole === "Proxy Admin"}
|
||||
userModels={allModelsOnProxy}
|
||||
editTeam={false}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,97 @@
|
|||
import userEvent from "@testing-library/user-event";
|
||||
import React from "react";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
import { renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils";
|
||||
import type { LogEntry as SpendLogEntry } from "@/components/view_logs/columns";
|
||||
import { LogViewer } from "./LogViewer";
|
||||
|
||||
vi.mock("@/components/networking", async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import("@/components/networking")>();
|
||||
return { ...actual, uiSpendLogsCall: vi.fn() };
|
||||
});
|
||||
|
||||
vi.mock("@/components/view_logs/LogDetailsDrawer", () => ({
|
||||
LogDetailsDrawer: function LogDetailsDrawerMock({
|
||||
open,
|
||||
logEntry,
|
||||
}: {
|
||||
open: boolean;
|
||||
logEntry?: { request_id: string } | null;
|
||||
}) {
|
||||
return (
|
||||
<div data-testid="log-details-drawer" data-log-id={logEntry?.request_id ?? ""}>
|
||||
{open ? "open" : "closed"}
|
||||
</div>
|
||||
);
|
||||
},
|
||||
}));
|
||||
|
||||
import { uiSpendLogsCall } from "@/components/networking";
|
||||
|
||||
const spendLog = (overrides: Partial<SpendLogEntry>): SpendLogEntry => ({
|
||||
request_id: "req-1",
|
||||
api_key: "key-1",
|
||||
team_id: "team-1",
|
||||
model: "gpt-4o",
|
||||
model_id: "model-1",
|
||||
call_type: "acompletion",
|
||||
spend: 0.01,
|
||||
total_tokens: 10,
|
||||
prompt_tokens: 5,
|
||||
completion_tokens: 5,
|
||||
startTime: "2026-09-02T09:50:13Z",
|
||||
endTime: "2026-09-02T09:50:14Z",
|
||||
cache_hit: "false",
|
||||
messages: [],
|
||||
response: {},
|
||||
...overrides,
|
||||
});
|
||||
|
||||
const guardrailLog = {
|
||||
id: "provider-victim",
|
||||
timestamp: "2026-09-02 09:50:13",
|
||||
action: "passed" as const,
|
||||
input_snippet: "victim prompt",
|
||||
};
|
||||
|
||||
describe("GuardrailsMonitor LogViewer drawer", () => {
|
||||
beforeEach(() => {
|
||||
vi.mocked(uiSpendLogsCall).mockReset();
|
||||
testQueryClient.clear();
|
||||
});
|
||||
|
||||
it("opens the row whose request_id is the clicked log id even when a newer row carries that id as its call id", async () => {
|
||||
vi.mocked(uiSpendLogsCall).mockResolvedValue({
|
||||
data: [
|
||||
spendLog({ request_id: "provider-attacker", litellm_call_id: "provider-victim" }),
|
||||
spendLog({ request_id: "provider-victim", litellm_call_id: "call-victim" }),
|
||||
],
|
||||
total: 2,
|
||||
});
|
||||
|
||||
renderWithProviders(<LogViewer logs={[guardrailLog]} accessToken="sk-test" />);
|
||||
await userEvent.click(screen.getByText("victim prompt"));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByTestId("log-details-drawer")).toHaveAttribute("data-log-id", "provider-victim");
|
||||
});
|
||||
expect(vi.mocked(uiSpendLogsCall)).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ params: { request_id: "provider-victim" } }),
|
||||
);
|
||||
});
|
||||
|
||||
it("falls back to the first returned row when none carries the clicked id as its request_id", async () => {
|
||||
vi.mocked(uiSpendLogsCall).mockResolvedValue({
|
||||
data: [spendLog({ request_id: "provider-other", litellm_call_id: "provider-victim" })],
|
||||
total: 1,
|
||||
});
|
||||
|
||||
renderWithProviders(<LogViewer logs={[guardrailLog]} accessToken="sk-test" />);
|
||||
await userEvent.click(screen.getByText("victim prompt"));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByTestId("log-details-drawer")).toHaveAttribute("data-log-id", "provider-other");
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -92,7 +92,8 @@ export function LogViewer({
|
|||
enabled: Boolean(accessToken && selectedRequestId && drawerOpen),
|
||||
});
|
||||
|
||||
const selectedLog: ViewLogsLogEntry | null = fullLogResponse?.data?.[0] ?? null;
|
||||
const selectedLog: ViewLogsLogEntry | null =
|
||||
fullLogResponse?.data?.find((log) => log.request_id === selectedRequestId) ?? fullLogResponse?.data?.[0] ?? null;
|
||||
|
||||
const handleLogClick = (log: LogEntry) => {
|
||||
setSelectedRequestId(log.id);
|
||||
|
|
|
|||
|
|
@ -577,6 +577,49 @@ describe("RequestLogsPanel", () => {
|
|||
expect(byIdCall.params?.group_by_session).toBeUndefined();
|
||||
});
|
||||
|
||||
it("opens the drawer when ?log_id= is the log's litellm_call_id rather than its request_id", async () => {
|
||||
respondWith([logEntry({ request_id: "chatcmpl-provider", litellm_call_id: "call-1" })]);
|
||||
renderPanel("?log_id=call-1");
|
||||
|
||||
await waitFor(() => {
|
||||
expect(drawer()).toHaveTextContent("open");
|
||||
});
|
||||
expect(drawer()).toHaveAttribute("data-log-id", "chatcmpl-provider");
|
||||
});
|
||||
|
||||
it("fetches by litellm_call_id and opens the drawer when that log is not in the loaded page", async () => {
|
||||
vi.mocked(uiSpendLogsCall).mockImplementation(async ({ params }) =>
|
||||
params?.request_id === "call-old"
|
||||
? {
|
||||
data: [logEntry({ request_id: "chatcmpl-old", litellm_call_id: "call-old" })],
|
||||
total: 1,
|
||||
page: 1,
|
||||
page_size: 1,
|
||||
total_pages: 1,
|
||||
}
|
||||
: { data: [], total: 0, page: 1, page_size: 50, total_pages: 0 },
|
||||
);
|
||||
renderPanel("?log_id=call-old");
|
||||
|
||||
await waitFor(() => {
|
||||
expect(drawer()).toHaveTextContent("open");
|
||||
});
|
||||
expect(drawer()).toHaveAttribute("data-log-id", "chatcmpl-old");
|
||||
});
|
||||
|
||||
it("opens the exact request_id row when another log in the page carries that id as its litellm_call_id", async () => {
|
||||
respondWith([
|
||||
logEntry({ request_id: "chatcmpl-other", litellm_call_id: "victim-req" }),
|
||||
logEntry({ request_id: "victim-req", litellm_call_id: "victim-call" }),
|
||||
]);
|
||||
renderPanel("?log_id=victim-req");
|
||||
|
||||
await waitFor(() => {
|
||||
expect(drawer()).toHaveTextContent("open");
|
||||
});
|
||||
expect(drawer()).toHaveAttribute("data-log-id", "victim-req");
|
||||
});
|
||||
|
||||
it("closing the drawer removes ?log_id= from the URL and closes the drawer", async () => {
|
||||
const user = userEvent.setup();
|
||||
respondWith([logEntry({ request_id: "req-1" })]);
|
||||
|
|
|
|||
|
|
@ -28,6 +28,9 @@ import { RequestLogsTable } from "./RequestLogsTable";
|
|||
|
||||
const PAGE_SIZE = DEFAULT_PAGE_SIZE_OPTIONS[0];
|
||||
const DEFAULT_INTERVAL = { value: 24, unit: "hours" };
|
||||
const matchesLogId = (log: LogEntry, logId: string) => log.request_id === logId || log.litellm_call_id === logId;
|
||||
const findLogById = (logs: readonly LogEntry[], logId: string): LogEntry | null =>
|
||||
logs.find((log) => log.request_id === logId) ?? logs.find((log) => log.litellm_call_id === logId) ?? null;
|
||||
|
||||
interface RequestLogsPanelProps {
|
||||
accessToken: string;
|
||||
|
|
@ -141,9 +144,9 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID,
|
|||
page_size: 1,
|
||||
params: { request_id: urlLogId },
|
||||
});
|
||||
return response.data.find((log) => log.request_id === urlLogId) ?? null;
|
||||
return findLogById(response.data, urlLogId);
|
||||
},
|
||||
enabled: urlLogId !== null && selectedLog?.request_id !== urlLogId,
|
||||
enabled: urlLogId !== null && !(selectedLog !== null && matchesLogId(selectedLog, urlLogId)),
|
||||
staleTime: Infinity,
|
||||
};
|
||||
|
||||
|
|
@ -151,8 +154,8 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID,
|
|||
|
||||
const displayLog = useMemo<LogEntry | null>(() => {
|
||||
if (urlLogId === null) return null;
|
||||
if (selectedLog?.request_id === urlLogId) return selectedLog;
|
||||
return filteredLogs.data.find((log) => log.request_id === urlLogId) ?? urlLog ?? null;
|
||||
if (selectedLog !== null && matchesLogId(selectedLog, urlLogId)) return selectedLog;
|
||||
return findLogById(filteredLogs.data, urlLogId) ?? urlLog ?? null;
|
||||
}, [urlLogId, selectedLog, filteredLogs.data, urlLog]);
|
||||
|
||||
const displaySessionId = useMemo<string | null>(() => {
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ export type LogsSortField = keyof typeof LOGS_SORT_FIELD_MAP;
|
|||
|
||||
export type LogEntry = {
|
||||
request_id: string;
|
||||
litellm_call_id?: string | null;
|
||||
api_key: string;
|
||||
team_id: string;
|
||||
model: string;
|
||||
|
|
|
|||
24
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
24
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -29639,6 +29639,12 @@ export interface components {
|
|||
aws_web_identity_token?: string | null;
|
||||
/** Azure Ad Token */
|
||||
azure_ad_token?: string | null;
|
||||
/** Azure Password */
|
||||
azure_password?: string | null;
|
||||
/** Azure Scope */
|
||||
azure_scope?: string | null;
|
||||
/** Azure Username */
|
||||
azure_username?: string | null;
|
||||
/** Bedrock Tags */
|
||||
bedrock_tags?: unknown[] | null;
|
||||
/** Budget Duration */
|
||||
|
|
@ -29687,6 +29693,10 @@ export interface components {
|
|||
cache_read_input_token_cost_ultrafast?: number | null;
|
||||
/** Citation Cost Per Token */
|
||||
citation_cost_per_token?: number | null;
|
||||
/** Client Id */
|
||||
client_id?: string | null;
|
||||
/** Client Secret */
|
||||
client_secret?: string | null;
|
||||
/** Complexity Router Config */
|
||||
complexity_router_config?: {
|
||||
[key: string]: unknown;
|
||||
|
|
@ -29900,6 +29910,8 @@ export interface components {
|
|||
tag_regex?: string[] | null;
|
||||
/** Tags */
|
||||
tags?: string[] | null;
|
||||
/** Tenant Id */
|
||||
tenant_id?: string | null;
|
||||
/** Tiered Pricing */
|
||||
tiered_pricing?: {
|
||||
[key: string]: unknown;
|
||||
|
|
@ -39841,6 +39853,12 @@ export interface components {
|
|||
aws_web_identity_token?: string | null;
|
||||
/** Azure Ad Token */
|
||||
azure_ad_token?: string | null;
|
||||
/** Azure Password */
|
||||
azure_password?: string | null;
|
||||
/** Azure Scope */
|
||||
azure_scope?: string | null;
|
||||
/** Azure Username */
|
||||
azure_username?: string | null;
|
||||
/** Bedrock Tags */
|
||||
bedrock_tags?: unknown[] | null;
|
||||
/** Budget Duration */
|
||||
|
|
@ -39889,6 +39907,10 @@ export interface components {
|
|||
cache_read_input_token_cost_ultrafast?: number | null;
|
||||
/** Citation Cost Per Token */
|
||||
citation_cost_per_token?: number | null;
|
||||
/** Client Id */
|
||||
client_id?: string | null;
|
||||
/** Client Secret */
|
||||
client_secret?: string | null;
|
||||
/** Complexity Router Config */
|
||||
complexity_router_config?: {
|
||||
[key: string]: unknown;
|
||||
|
|
@ -40102,6 +40124,8 @@ export interface components {
|
|||
tag_regex?: string[] | null;
|
||||
/** Tags */
|
||||
tags?: string[] | null;
|
||||
/** Tenant Id */
|
||||
tenant_id?: string | null;
|
||||
/** Tiered Pricing */
|
||||
tiered_pricing?: {
|
||||
[key: string]: unknown;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue