scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_decrease_anys_opus5_r2

Browse source 
# Conflicts:
#	basedpyright-code-budget.json
#	enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
#	ruff-strict-budget.json
#	type-discipline-budget.json
This commit is contained in:
mateo-berri 2026-08-29 15:11:45 -07:00
commit 048e82ec65
111 changed files with 5702 additions and 560 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
.github/workflows/codspeed.yml vendored
View file

@ -12,6 +12,7 @@ on:
- "uv.lock"
- ".github/workflows/codspeed.yml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/actions/cache-cargo-build/**"
pull_request:
branches:
- main
@ -23,6 +24,7 @@ on:
- "uv.lock"
- ".github/workflows/codspeed.yml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/actions/cache-cargo-build/**"
# Allow CodSpeed to trigger backtest performance analysis
# in order to generate initial data
workflow_dispatch:
@ -55,6 +57,26 @@ jobs:
with:
version: "0.10.9"
- name: Cache the Rust build
uses: ./.github/actions/cache-cargo-build
# Build the wheel and resolve every dependency outside the CodSpeed
# runner: the same maturin build took 42 minutes inside `codspeed run`
# versus under 3 minutes as a plain step (LIT-6183)
- name: Build environment
run: >
env PYTEST_DISABLE_PLUGIN_AUTOLOAD=1
uv run --frozen --no-default-groups
--with pytest==8.3.5
--with pytest-codspeed==4.3.0
--with "mcp>=1.26.0,<2.0"
--with "a2a-sdk>=1.1.0,<2.0"
pytest
-p pytest_codspeed.plugin
tests/benchmarks/
--codspeed
--collect-only -q
- name: Run benchmarks
uses: CodSpeedHQ/action@1c8ae4843586d3ba879736b7f6b7b0c990757fab # v4.12.1
with:

2
basedpyright-code-budget.json
View file

@ -24,7 +24,7 @@
"limit": 19
},
"reportExplicitAny": {
"limit": 5063
"limit": 5062
},
"reportFunctionMemberAccess": {
"limit": 7

43
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -2,9 +2,10 @@
Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
"""
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, List, Optional, Tuple
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Tuple, cast
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -632,6 +633,7 @@ class CheckBatchCost:
later poll.
"""
from litellm.batches.batch_utils import (
count_error_file_failed_requests,
_get_file_content_as_dictionary,
calculate_batch_cost_and_usage,
)
@ -767,16 +769,33 @@ class CheckBatchCost:
model_id=model_id,
deployment_model=litellm_model_name,
)
batch_cost, batch_usage, batch_models = (
await calculate_batch_cost_and_usage(
file_content_dictionary=file_content_as_dict,
custom_llm_provider=llm_provider, # type: ignore
model_name=model_name,
model_info=deployment_model_info,
batch_file_provider: Final = cast(
Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], llm_provider
)
output_file_result: Final = await calculate_batch_cost_and_usage(
file_content_dictionary=file_content_as_dict,
custom_llm_provider=batch_file_provider,
model_name=model_name,
model_info=deployment_model_info,
)
error_file_failed_requests: Final = await count_error_file_failed_requests(
response,
custom_llm_provider=batch_file_provider,
litellm_params={
**credentials,
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
},
)
batch_result: Final = (
output_file_result
if not error_file_failed_requests
else dataclasses_replace(
output_file_result,
failed_requests=output_file_result.failed_requests + error_file_failed_requests,
)
)
logging_obj = LiteLLMLogging(
model=batch_models[0],
model=batch_result.models[0],
messages=[{"role": "user", "content": "<retrieve_batch>"}],
stream=False,
call_type="aretrieve_batch",
@ -808,9 +827,11 @@ class CheckBatchCost:
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_cost,
batch_usage=batch_usage,
batch_models=batch_models,
batch_cost=batch_result.cost,
batch_usage=batch_result.usage,
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
)
except Exception:
await self._release_job_claim(job)

7
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -36,6 +36,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
build_list_page,
build_owner_filter,
can_access_resource,
resolve_resource_owner_id,
)
from litellm.proxy._types import (
CallTypes,
@ -226,7 +227,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
file_object=file_object,
model_mappings=model_mappings,
flat_model_file_ids=list(model_mappings.values()),
created_by=user_api_key_dict.user_id,
created_by=resolve_resource_owner_id(user_api_key_dict),
team_id=user_api_key_dict.team_id,
updated_by=user_api_key_dict.user_id,
)
@ -242,7 +243,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"unified_file_id": file_id,
"model_mappings": json.dumps(model_mappings),
"flat_model_file_ids": list(model_mappings.values()),
"created_by": user_api_key_dict.user_id,
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
"updated_by": user_api_key_dict.user_id,
}
@ -346,7 +347,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"file_object": file_object.model_dump_json(),
"model_object_id": model_object_id,
"file_purpose": file_purpose,
"created_by": user_api_key_dict.user_id,
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
"updated_by": user_api_key_dict.user_id,
"status": file_object.status,

227
litellm/batches/batch_utils.py
View file

@ -1,6 +1,8 @@
import json
from collections.abc import Iterable, Iterator, Mapping
from dataclasses import dataclass
from dataclasses import replace as dataclasses_replace
from enum import Enum
from typing import Any, Final, Literal
import litellm
@ -12,12 +14,23 @@ from litellm.types.utils import CallTypes, ModelInfo, Usage
from litellm.utils import token_counter
@dataclass(frozen=True, slots=True)
class BatchCostUsageResult:
"""Aggregate cost, usage, and per-line pass/fail counts for a completed batch."""
cost: float
usage: Usage
models: list[str]
successful_requests: int
failed_requests: int
async def calculate_batch_cost_and_usage(
file_content_dictionary: list[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
model_name: str | None = None,
model_info: ModelInfo | None = None,
) -> tuple[float, Usage, list[str]]:
) -> BatchCostUsageResult:
"""
Calculate the cost and usage of a batch.
@ -32,8 +45,7 @@ async def calculate_batch_cost_and_usage(
and model_name
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
):
batch_cost, batch_usage = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
return batch_cost, batch_usage, [model_name]
return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
return _aggregate_batch_cost_usage_models(
entries=file_content_dictionary,
@ -49,7 +61,7 @@ async def _handle_completed_batch(
model_name: str | None = None,
litellm_params: dict | None = None,
model_info: ModelInfo | None = None,
) -> tuple[float, Usage, list[str]]:
) -> BatchCostUsageResult:
"""Fetch a completed batch's output file and aggregate its cost, usage, and
models in a single pass over the JSONL lines, so the parsed file content is
never materialized in memory.
@ -72,27 +84,49 @@ async def _handle_completed_batch(
# The generic retrieval helper keeps raising for callers that explicitly ask
# for a missing output file.
if batch.output_file_id is None:
return 0.0, Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), []
return BatchCostUsageResult(
cost=0.0,
usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
models=[], # mutable-ok: no output file means no model was ever priced; BatchCostUsageResult.models requires list[str]
successful_requests=0,
failed_requests=await count_error_file_failed_requests(
batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
),
)
file_content = await _fetch_batch_output_file_content(batch, custom_llm_provider, litellm_params=litellm_params)
if (
custom_llm_provider == "vertex_ai"
and model_name
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
):
batch_cost, batch_usage = calculate_vertex_ai_batch_cost_and_usage(
_get_file_content_as_dictionary(file_content), model_name
)
return batch_cost, batch_usage, [model_name]
return _aggregate_batch_cost_usage_models(
entries=_iter_batch_output_entries(file_content),
custom_llm_provider=custom_llm_provider,
model_name=model_name,
model_info=model_info,
error_file_failed_requests: Final = await count_error_file_failed_requests(
batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
)
output_file_result: Final = (
calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name)
if (
custom_llm_provider == "vertex_ai"
and model_name
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
)
else _aggregate_batch_cost_usage_models(
entries=_iter_batch_output_entries(file_content),
custom_llm_provider=custom_llm_provider,
model_name=model_name,
model_info=model_info,
)
)
if not error_file_failed_requests:
return output_file_result
return dataclasses_replace(
output_file_result, failed_requests=output_file_result.failed_requests + error_file_failed_requests
)
class _LineOutcome(Enum):
"""A batch output line that yielded no billable stats."""
PROVIDER_FAILED = "provider_failed"
UNCOSTABLE = "uncostable"
@dataclass(frozen=True, slots=True)
class _BatchOutputLineStats:
@ -102,19 +136,27 @@ class _BatchOutputLineStats:
total_tokens: int
cache_read_tokens: int
cache_creation_tokens: int
reasoning_tokens: int
model: str | None
def _iter_successful_output_line_stats(
def _classify_output_line_stats(
entries: Iterable[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
model_name: str | None,
model_info: ModelInfo | None,
) -> Iterator[_BatchOutputLineStats]:
) -> Iterator[_BatchOutputLineStats | _LineOutcome]:
"""Classify every output line in a single pass, so counting failures never needs
a second read of a potentially huge output file. A line the provider reported as
failed yields ``PROVIDER_FAILED``; a successful line litellm could not price
yields ``UNCOSTABLE`` and still counts as a successful request billed at $0, so
the counts stay reconcilable with the provider's own ``request_counts``."""
for entry in entries:
if not _batch_response_was_successful(entry, custom_llm_provider):
yield _LineOutcome.PROVIDER_FAILED
continue
stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info)
if stats is not None:
yield stats
yield stats if stats is not None else _LineOutcome.UNCOSTABLE
def _safe_output_line_stats(
@ -123,13 +165,11 @@ def _safe_output_line_stats(
model_name: str | None,
model_info: ModelInfo | None,
) -> _BatchOutputLineStats | None:
"""Return the stats for one batch output line, or None for a line that is
unsuccessful or cannot be costed, so a single bad line never aborts the
whole batch's cost accounting."""
"""Return the stats for one provider-successful batch output line, or None when
it cannot be costed, so a single bad line never aborts the whole batch's cost
accounting."""
custom_id: Final = entry.get("custom_id") if isinstance(entry, dict) else None
try:
if not _batch_response_was_successful(entry, custom_llm_provider):
return None
return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info)
except Exception as e: # noqa: BLE001 # any single line's costing failure must not abort the whole batch
verbose_logger.warning(
@ -152,6 +192,7 @@ def _compute_output_line_stats(
prompt_details: Final = parse_prompt_tokens_details(usage)
raw_model: Final = response_body.get("model")
response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None
completion_details: Final = usage.completion_tokens_details
return _BatchOutputLineStats(
cost=_output_line_cost(
response_body=response_body,
@ -166,6 +207,7 @@ def _compute_output_line_stats(
total_tokens=usage.total_tokens,
cache_read_tokens=prompt_details["cache_hit_tokens"],
cache_creation_tokens=prompt_details["cache_creation_tokens"],
reasoning_tokens=(completion_details.reasoning_tokens if completion_details else None) or 0,
model=response_model,
)
@ -203,10 +245,14 @@ def _aggregate_batch_cost_usage_models(
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
model_name: str | None = None,
model_info: ModelInfo | None = None,
) -> tuple[float, Usage, list[str]]:
"""Aggregate cost, usage, and models from batch output entries in a single
pass, holding one small stats record per line instead of the parsed file."""
line_stats: Final = tuple(_iter_successful_output_line_stats(entries, custom_llm_provider, model_name, model_info))
) -> BatchCostUsageResult:
"""Aggregate cost, usage, models, and pass/fail counts from batch output
entries in a single pass, holding one small stats record per line instead
of the parsed file."""
all_results: Final = tuple(_classify_output_line_stats(entries, custom_llm_provider, model_name, model_info))
line_stats: Final = tuple(result for result in all_results if isinstance(result, _BatchOutputLineStats))
failed_requests: Final = sum(1 for result in all_results if result is _LineOutcome.PROVIDER_FAILED)
successful_requests: Final = len(all_results) - failed_requests
cache_token_params: Final = {
key: tokens
@ -220,18 +266,32 @@ def _aggregate_batch_cost_usage_models(
total_tokens=sum(stats.total_tokens for stats in line_stats),
prompt_tokens=sum(stats.prompt_tokens for stats in line_stats),
completion_tokens=sum(stats.completion_tokens for stats in line_stats),
reasoning_tokens=sum(stats.reasoning_tokens for stats in line_stats),
**cache_token_params,
)
batch_models: Final = [model_name] if model_name else [stats.model for stats in line_stats if stats.model]
total_cost: Final = sum((stats.cost for stats in line_stats), 0.0)
verbose_logger.debug("batch output aggregate: cost=%s usage=%s models=%s", total_cost, batch_usage, batch_models)
return total_cost, batch_usage, batch_models
verbose_logger.debug(
"batch output aggregate: cost=%s usage=%s models=%s successful=%d failed=%d",
total_cost,
batch_usage,
batch_models,
successful_requests,
failed_requests,
)
return BatchCostUsageResult(
cost=total_cost,
usage=batch_usage,
models=batch_models,
successful_requests=successful_requests,
failed_requests=failed_requests,
)
def calculate_vertex_ai_batch_cost_and_usage(
vertex_ai_batch_responses: list[dict],
model_name: str | None = None,
) -> tuple[float, Usage]:
) -> BatchCostUsageResult:
"""
Calculate both cost and usage from raw Vertex AI batch responses.
@ -242,6 +302,10 @@ def calculate_vertex_ai_batch_cost_and_usage(
{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}}}
usageMetadata contains promptTokenCount, candidatesTokenCount, totalTokenCount.
A row with no ``response`` is counted as failed - the same signal already
used to skip it from cost/usage aggregation, since Vertex batch prediction
output doesn't establish a distinct error shape in this (non-default) path.
"""
from litellm.cost_calculator import batch_cost_calculator
@ -249,12 +313,16 @@ def calculate_vertex_ai_batch_cost_and_usage(
total_tokens = 0
prompt_tokens = 0
completion_tokens = 0
successful_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above
failed_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above
actual_model_name: Final = model_name or "gemini-2.0-flash-001"
for response in vertex_ai_batch_responses:
response_body = response.get("response")
if response_body is None:
failed_requests += 1
continue
successful_requests += 1
usage_metadata = response_body.get("usageMetadata", {})
_prompt = usage_metadata.get("promptTokenCount", 0) or 0
@ -282,17 +350,25 @@ def calculate_vertex_ai_batch_cost_and_usage(
total_tokens += _total
verbose_logger.info(
"vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d",
"vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d",
total_cost,
prompt_tokens,
completion_tokens,
total_tokens,
successful_requests,
failed_requests,
)
return total_cost, Usage(
total_tokens=total_tokens,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
return BatchCostUsageResult(
cost=total_cost,
usage=Usage(
total_tokens=total_tokens,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
),
models=[actual_model_name],
successful_requests=successful_requests,
failed_requests=failed_requests,
)
@ -322,6 +398,36 @@ def _provider_output_file_id(output_file_id: str) -> str:
return extracted
async def _fetch_batch_managed_file_content(
file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
litellm_params: dict | None = None,
) -> bytes:
"""
Fetch a batch's output or error file and return its raw JSONL bytes.
Args:
file_id: The provider or unified (litellm-managed) file id to fetch
custom_llm_provider: The LLM provider
litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.)
Required for Azure and other providers that need authentication
"""
from litellm.files.main import afile_content
# Build kwargs for afile_content with credentials from litellm_params
file_content_kwargs: Final = {
"file_id": _provider_output_file_id(file_id),
"custom_llm_provider": custom_llm_provider,
}
# Extract and add credentials for file access
credentials: Final = _extract_file_access_credentials(litellm_params)
file_content_kwargs.update(credentials)
_file_content: Final = await afile_content(**file_content_kwargs)
return _file_content.content
async def _fetch_batch_output_file_content(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
@ -336,25 +442,36 @@ async def _fetch_batch_output_file_content(
litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.)
Required for Azure and other providers that need authentication
"""
from litellm.files.main import afile_content
if batch.output_file_id is None:
raise ValueError("Output file id is None cannot retrieve file content")
file_id: Final = _provider_output_file_id(batch.output_file_id)
return await _fetch_batch_managed_file_content(
batch.output_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
)
# Build kwargs for afile_content with credentials from litellm_params
file_content_kwargs: Final = {
"file_id": file_id,
"custom_llm_provider": custom_llm_provider,
}
# Extract and add credentials for file access
credentials: Final = _extract_file_access_credentials(litellm_params)
file_content_kwargs.update(credentials)
async def count_error_file_failed_requests(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
litellm_params: dict | None,
) -> int:
"""Count failed requests reported only in the batch's separate error file.
_file_content: Final = await afile_content(**file_content_kwargs)
return _file_content.content
OpenAI-shaped batch providers write successful lines to ``output_file_id``
and per-request failures (e.g. a rejected param) to a distinct
``error_file_id`` - they never appear in the output file at all, so
counting failures from the output file alone silently undercounts them.
"""
if batch.error_file_id is None:
return 0
try:
error_file_content = await _fetch_batch_managed_file_content(
batch.error_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
)
except Exception as e: # noqa: BLE001 # a failed/missing error file must not abort cost tracking for the batch
verbose_logger.debug("Failed to fetch batch error file %s: %s", batch.error_file_id, e)
return 0
return sum(1 for _ in _iter_batch_input_lines(error_file_content))
def _extract_file_access_credentials(litellm_params: dict | None) -> dict:

1
litellm/constants.py
View file

@ -35,6 +35,7 @@ DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECO
DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5))
DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1))
DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250))
HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0
# Maximum wall-clock seconds a streaming response is allowed to run.
# Streams exceeding this duration are terminated with a Timeout error.

190
litellm/litellm_core_utils/audio_utils/subtitle_utils.py
View file

@ -1,19 +1,36 @@
"""Provider-agnostic SRT/WebVTT subtitle synthesis from timestamped transcription tokens."""
import unicodedata
from collections.abc import Sequence
from dataclasses import dataclass
from itertools import accumulate, chain
from itertools import accumulate, groupby
from typing import Final
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
CUE_MAX_TOKENS: Final = 15
CUE_MAX_DURATION_MS: Final = 5000
CUE_MAX_CHARS: Final = 84
CUE_MAX_DURATION_MS: Final = 7000
CUE_GAP_MS: Final = 700
SRT_RESPONSE_FORMAT: Final = "srt"
VTT_RESPONSE_FORMAT: Final = "vtt"
SUBTITLE_RESPONSE_FORMATS: Final = frozenset((SRT_RESPONSE_FORMAT, VTT_RESPONSE_FORMAT))
_SENTENCE_END_CHARS: Final = (".", "!", "?", "。", "！", "？", "؟", "۔", "।", "॥", "։", "።")
_CJK_RANGES: Final = (
(0x3400, 0x4DBF),
(0x4E00, 0x9FFF),
(0xF900, 0xFAFF),
(0x3040, 0x309F),
(0x30A0, 0x30FF),
(0x31F0, 0x31FF),
)
_CJK_NO_BREAK_BEFORE: Final = "、。，．！？：；・ー…」』）〉》】〕"
_CJK_NO_BREAK_AFTER: Final = "「『（〈《【〔"
@dataclass(frozen=True, slots=True)
class SubtitleToken:
@ -31,69 +48,138 @@ class SubtitleCue:
@dataclass(frozen=True, slots=True)
class _CueAccumulator:
texts: tuple[str, ...] = ()
start_ms: int | None = None
end_ms: int | None = None
speaker: str | int | None = None
class _Word:
text: str
start_ms: int | None
end_ms: int | None
speaker: str | int | None
def _completed_cue(accumulator: _CueAccumulator) -> tuple[SubtitleCue, ...]:
if not accumulator.texts or accumulator.start_ms is None:
return ()
text: Final = "".join(accumulator.texts).strip()
if not text:
return ()
end_ms: Final = accumulator.end_ms if accumulator.end_ms is not None else accumulator.start_ms
return (SubtitleCue(start_ms=accumulator.start_ms, end_ms=end_ms, text=text),)
def _is_cjk(ch: str) -> bool:
cp: Final = ord(ch)
return any(lo <= cp <= hi for lo, hi in _CJK_RANGES)
def _cue_break_reached(accumulator: _CueAccumulator, token: SubtitleToken) -> bool:
if len(accumulator.texts) >= CUE_MAX_TOKENS:
return True
def _is_cjk_word_boundary(prev_ch: str, next_ch: str) -> bool:
if not (_is_cjk(prev_ch) or _is_cjk(next_ch)):
return False
return next_ch not in _CJK_NO_BREAK_BEFORE and prev_ch not in _CJK_NO_BREAK_AFTER
def _text_width(text: str) -> int:
return sum(2 if unicodedata.east_asian_width(ch) in ("W", "F") else 1 for ch in text)
def _starts_new_word(prev: SubtitleToken, token: SubtitleToken) -> bool:
prev_last: Final = prev.text[-1:]
first: Final = token.text[0]
return (
accumulator.start_ms is not None
and token.start_ms is not None
and token.start_ms - accumulator.start_ms >= CUE_MAX_DURATION_MS
first.isspace()
or prev_last.isspace()
or token.speaker != prev.speaker
or _is_cjk_word_boundary(prev_last, first)
)
_AbsorbStep = tuple[tuple[SubtitleCue, ...], _CueAccumulator]
def _absorb_token(accumulator: _CueAccumulator, token: SubtitleToken) -> _AbsorbStep:
if token.start_ms is None and accumulator.start_ms is None:
return (), accumulator
if token.speaker is not None and token.speaker != accumulator.speaker:
return _completed_cue(accumulator), _CueAccumulator(
texts=(token.text,),
start_ms=token.start_ms,
end_ms=token.end_ms,
speaker=token.speaker,
)
if _cue_break_reached(accumulator, token):
return _completed_cue(accumulator), _CueAccumulator(
texts=(token.text,),
start_ms=token.start_ms,
end_ms=token.end_ms,
speaker=accumulator.speaker,
)
return (), _CueAccumulator(
texts=(*accumulator.texts, token.text),
start_ms=accumulator.start_ms if accumulator.start_ms is not None else token.start_ms,
end_ms=token.end_ms if token.end_ms is not None else accumulator.end_ms,
speaker=accumulator.speaker,
def _build_word(group: Sequence[SubtitleToken]) -> _Word:
return _Word(
text="".join(t.text for t in group),
start_ms=next((t.start_ms for t in group if t.start_ms is not None), None),
end_ms=next((t.end_ms for t in reversed(group) if t.end_ms is not None), None),
speaker=group[0].speaker,
)
def _absorb_step(carry: _AbsorbStep, token: SubtitleToken) -> _AbsorbStep:
return _absorb_token(carry[1], token)
def _merge_tokens_into_words(tokens: Sequence[SubtitleToken]) -> tuple[_Word, ...]:
"""
Merge subword tokens (e.g. ``"Hel"``, ``"lo"``) into whole words.
A token starts a new word when its text begins with whitespace, when the
previous token's text ends with whitespace, when the speaker changes, or
at a CJK character boundary (CJK scripts carry no spaces, so without this
an entire utterance would fuse into a single unbreakable "word"; CJK
punctuation stays attached to the preceding character per kinsoku rules).
Each word carries the first/last available timestamps of its tokens.
"""
kept: Final = tuple(t for t in tokens if t.text != "")
starts: Final = tuple(i for i, t in enumerate(kept) if i == 0 or _starts_new_word(kept[i - 1], t))
return tuple(_build_word(kept[begin:end]) for begin, end in zip(starts, (*starts[1:], len(kept))))
def _cue_start(ws: Sequence[_Word]) -> int | None:
return next((w.start_ms for w in ws if w.start_ms is not None), None)
def _cue_end(ws: Sequence[_Word]) -> int | None:
return next((w.end_ms for w in reversed(ws) if w.end_ms is not None), _cue_start(ws))
def _cue_text(ws: Sequence[_Word]) -> str:
return "".join(w.text for w in ws).strip()
def _should_break(cue: Sequence[_Word], word: _Word) -> bool:
speaker_changed: Final = word.speaker is not None and any(
w.speaker is not None and w.speaker != word.speaker for w in cue
)
cue_start: Final = _cue_start(cue)
cue_end: Final = _cue_end(cue)
gap_exceeded: Final = word.start_ms is not None and cue_end is not None and (word.start_ms - cue_end) >= CUE_GAP_MS
chars_exceeded: Final = _text_width(_cue_text(cue)) + _text_width(word.text) > CUE_MAX_CHARS
word_end: Final = word.end_ms if word.end_ms is not None else word.start_ms
duration_exceeded: Final = (
word_end is not None and cue_start is not None and (word_end - cue_start) > CUE_MAX_DURATION_MS
)
return speaker_changed or gap_exceeded or chars_exceeded or duration_exceeded
def _cue_start_indices(words: Sequence[_Word]) -> tuple[int, ...]:
def next_start(start: int, index: int) -> int:
if words[index - 1].text.rstrip().endswith(_SENTENCE_END_CHARS):
return index
if _should_break(words[start:index], words[index]):
return index
return start
if not words:
return ()
return tuple(start for start, _ in groupby(accumulate(range(1, len(words)), next_start, initial=0)))
def _build_cue(ws: Sequence[_Word]) -> SubtitleCue | None:
text: Final = _cue_text(ws)
start: Final = _cue_start(ws)
if not text or start is None:
return None
end: Final = _cue_end(ws)
return SubtitleCue(start_ms=start, end_ms=end if end is not None else start, text=text)
def group_subtitle_tokens_into_cues(tokens: Sequence[SubtitleToken]) -> tuple[SubtitleCue, ...]:
steps: Final = tuple(accumulate(tokens, _absorb_step, initial=((), _CueAccumulator())))
completed: Final = chain.from_iterable(emitted for emitted, _ in steps)
return (*completed, *_completed_cue(steps[-1][1]))
"""
Group transcription tokens into subtitle cues aligned to the actual speech.
Cues only ever break at word boundaries (tokens may be subwords, so they
are first merged into words). A new cue starts when:
- the speaker changes (if diarization is on),
- a silence gap of at least CUE_GAP_MS separates two words, so
subtitles never bridge pauses in speech,
- adding the next word would exceed CUE_MAX_CHARS of display width
(~two subtitle lines; East-Asian wide characters count double), or
- adding the next word would make the cue span more than
CUE_MAX_DURATION_MS.
A cue also ends after sentence-final punctuation, which keeps cue breaks
at natural seams. Cue timestamps come straight from token timestamps;
words without timestamps stay attached to the surrounding cue, and a cue
whose words carry no timestamps at all is dropped.
"""
words: Final = _merge_tokens_into_words(tokens)
starts: Final = _cue_start_indices(words)
return tuple(
cue
for begin, end in zip(starts, (*starts[1:], len(words)))
if (cue := _build_cue(words[begin:end])) is not None
)
def _format_timestamp(total_ms: int, millis_separator: str) -> str:

50
litellm/litellm_core_utils/chat_completion_agentic_loop.py
View file

@ -1,6 +1,7 @@
# this is a patch to allow for agentic loops covering llm_http_handler.py and openai sdk based calling flows for the .completion() api
import json
from collections.abc import Mapping
from typing import Final, cast
from litellm._logging import verbose_logger
@ -9,8 +10,11 @@ from litellm.litellm_core_utils.agentic_loop_settings import (
DEFAULT_MAX_AGENTIC_LOOPS,
validated_max_agentic_loops,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
from litellm.types.integrations.custom_logger import (
CHAT_COMPLETION_AGENTIC_SURFACE,
HEADROOM_CONVERTED_STREAM_KEY,
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
AgenticLoopPlan,
AgenticLoopRequestPatch,
@ -50,6 +54,12 @@ def _post_hook_overridden(callback: CustomLogger) -> bool:
return getattr(func, "__func__", func) is not getattr(base, "__func__", base)
def _converted_stream_requested(kwargs: Mapping[str, object]) -> bool:
return bool(
kwargs.get("_code_interpreter_interception_converted_stream") or kwargs.get(HEADROOM_CONVERTED_STREAM_KEY)
)
def _coerce_int(value: object, default: int) -> int:
return int(value) if isinstance(value, (int, str)) else default
@ -87,16 +97,24 @@ def _check_agentic_loop_safety(
return fingerprint
def _wrap_response_as_fake_stream(response: object) -> object:
if getattr(response, "object", None) == "chat.completion.chunk":
def _wrap_response_as_fake_stream(
response: object,
*,
model: str,
custom_llm_provider: str,
logging_obj: object,
) -> object:
if isinstance(response, CustomStreamWrapper):
return response
if not hasattr(response, "choices"):
if not isinstance(response, ModelResponse) or not isinstance(logging_obj, LiteLLMLoggingObject):
return response
from litellm.llms.base_llm.base_model_iterator import (
convert_model_response_to_streaming,
)
return convert_model_response_to_streaming(cast(ModelResponse, response))
return CustomStreamWrapper(
completion_stream=MockResponseIterator(model_response=response),
model=model,
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
)
def _add_agentic_loop_metadata(kwargs_for_followup: dict[str, object]) -> None:
@ -177,8 +195,13 @@ async def _execute_chat_completion_agentic_plan(
model,
str(e),
)
if kwargs.get("_code_interpreter_interception_converted_stream") and not depth:
return _wrap_response_as_fake_stream(response_followup)
if _converted_stream_requested(kwargs) and not depth:
return _wrap_response_as_fake_stream(
response_followup,
model=model,
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
)
return response_followup
finally:
try:
@ -302,9 +325,14 @@ async def maybe_run_chat_completion_agentic_loop(
str(e),
)
if kwargs.get("_code_interpreter_interception_converted_stream") and not depth and hasattr(response, "choices"):
if _converted_stream_requested(kwargs) and not depth:
return cast(
"ModelResponse | CustomStreamWrapper",
_wrap_response_as_fake_stream(response),
_wrap_response_as_fake_stream(
response,
model=model,
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
),
)
return None

24
litellm/litellm_core_utils/litellm_logging.py
View file

@ -2872,6 +2872,8 @@ class Logging(LiteLLMLoggingBaseClass):
batch_cost: Final = kwargs.get("batch_cost", None)
batch_usage = kwargs.get("batch_usage", None)
batch_models = kwargs.get("batch_models", None)
batch_successful_requests: Final = kwargs.get("batch_successful_requests", None)
batch_failed_requests: Final = kwargs.get("batch_failed_requests", None)
has_explicit_batch_data: Final = all(x is not None for x in (batch_cost, batch_usage, batch_models))
should_compute_batch_data: Final = (
@ -2880,14 +2882,12 @@ class Logging(LiteLLMLoggingBaseClass):
if has_explicit_batch_data:
result._hidden_params["response_cost"] = batch_cost
result._hidden_params["batch_models"] = batch_models
result._hidden_params["batch_successful_requests"] = batch_successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same result._hidden_params pattern as response_cost/batch_models above
result._hidden_params["batch_failed_requests"] = batch_failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above
result.usage = batch_usage
elif should_compute_batch_data:
(
response_cost,
batch_usage,
batch_models,
) = await _handle_completed_batch(
batch_result: Final = await _handle_completed_batch(
batch=result,
custom_llm_provider=self.custom_llm_provider,
model_name=self.get_deployment_model_for_cost(),
@ -2895,9 +2895,11 @@ class Logging(LiteLLMLoggingBaseClass):
model_info=self.get_router_deployment_model_info(),
)
result._hidden_params["response_cost"] = response_cost
result._hidden_params["batch_models"] = batch_models
result.usage = batch_usage
result._hidden_params["response_cost"] = batch_result.cost
result._hidden_params["batch_models"] = batch_result.models
result._hidden_params["batch_successful_requests"] = batch_result.successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above
result._hidden_params["batch_failed_requests"] = batch_result.failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above
result.usage = batch_result.usage
start_time, end_time, result = self._success_handler_helper_fn(
start_time=start_time,
@ -5422,6 +5424,8 @@ class StandardLoggingPayloadSetup:
additional_headers=None,
litellm_overhead_time_ms=None,
batch_models=None,
batch_successful_requests=None,
batch_failed_requests=None,
litellm_model_name=None,
usage_object=None,
)
@ -5812,6 +5816,8 @@ def _extract_response_obj_and_hidden_params(
response_cost=None,
litellm_overhead_time_ms=None,
batch_models=None,
batch_successful_requests=None,
batch_failed_requests=None,
litellm_model_name=None,
usage_object=None,
)
@ -6228,6 +6234,8 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload:
additional_headers=None,
litellm_overhead_time_ms=None,
batch_models=None,
batch_successful_requests=None,
batch_failed_requests=None,
litellm_model_name=None,
usage_object=None,
)

23
litellm/llms/azure/azure.py
View file

@ -846,6 +846,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
api_key: str,
data: dict,
headers: dict,
deployment_name: str | None = None,
) -> httpx.Response:
"""
Implemented for azure dall-e-2 image gen calls
@ -957,7 +958,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
content=json.dumps(result).encode("utf-8"),
request=httpx.Request(method="POST", url="https://api.openai.com/v1"),
)
request_json: Final = azure_deployment_image_generation_json_body(api_base, data)
request_json: Final = azure_deployment_image_generation_json_body(api_base, data, deployment_name)
return await async_handler.post(
url=api_base,
json=request_json,
@ -973,6 +974,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
api_key: str,
data: dict,
headers: dict,
deployment_name: str | None = None,
) -> httpx.Response:
"""
Implemented for azure dall-e-2 image gen calls
@ -1073,7 +1075,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
content=json.dumps(result).encode("utf-8"),
request=httpx.Request(method="POST", url="https://api.openai.com/v1"),
)
request_json: Final = azure_deployment_image_generation_json_body(api_base, data)
request_json: Final = azure_deployment_image_generation_json_body(api_base, data, deployment_name)
return sync_handler.post(
url=api_base,
json=request_json,
@ -1091,9 +1093,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
AzureFoundryMAIImageGenerationConfig,
)
api_base: str = azure_client_params.get("azure_endpoint", "") # "https://example-endpoint.openai.azure.com"
if api_base.endswith("/"):
api_base = api_base.rstrip("/")
# deployment-scoped endpoints are moved to "base_url" by select_azure_base_url_or_endpoint
api_base: str = (azure_client_params.get("azure_endpoint") or azure_client_params.get("base_url") or "").rstrip(
"/"
)
api_version: Final[str] = azure_client_params.get("api_version", "")
if model is None:
model = ""
@ -1113,6 +1116,14 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
api_version=api_version,
)
v1_url: Final = BaseAzureLLM.get_azure_v1_image_url(
api_base=api_base,
api_version=api_version,
route="/openai/images/generations",
)
if v1_url is not None:
return v1_url
if "/openai/deployments/" in api_base:
base_url_with_deployment = api_base
else:
@ -1167,6 +1178,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
api_key=api_key,
data=data,
headers=headers,
deployment_name=model,
)
provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2"))
@ -1302,6 +1314,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
api_key=api_key or "",
data=data,
headers=headers,
deployment_name=model,
)
provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2"))
if isinstance(provider_config, AzureFoundryMAIImageGenerationConfig):

27
litellm/llms/azure/common_utils.py
View file

@ -4,6 +4,7 @@ import json
import os
from collections.abc import Callable, Mapping
from functools import lru_cache
from types import MappingProxyType
from typing import Any, Final, Literal, NamedTuple, cast
import httpx
@ -789,6 +790,32 @@ class BaseAzureLLM(BaseOpenAILLM):
return str(final_url)
@staticmethod
def get_azure_v1_image_url(api_base: str, api_version: str | None, route: str) -> str | None:
"""
Azure's v1 surface serves images at ``/openai/v1/images/{generations,edits}`` and routes by
``model`` in the request body, so any deployment path and stale ``api-version`` in
``api_base`` have to be dropped.
Returns None when ``api_version`` is a dated one, which still uses the deployment route.
"""
if not BaseAzureLLM._is_azure_v1_api_version(api_version):
return None
base_url: Final = httpx.URL(api_base)
openai_path_start: Final = base_url.path.find("/openai")
resource_base: Final = str(
base_url.copy_with(
path=base_url.path if openai_path_start == -1 else base_url.path[:openai_path_start],
params=httpx.QueryParams(tuple((k, v) for k, v in base_url.params.multi_items() if k != "api-version")),
)
)
return BaseAzureLLM._get_base_azure_url(
api_base=resource_base,
litellm_params=MappingProxyType({"api_version": api_version}),
route=route,
)
@staticmethod
def _is_azure_v1_api_version(api_version: str | None) -> bool:
if api_version is None:

12
litellm/llms/azure/image_edit/transformation.py
View file

@ -93,8 +93,6 @@ class AzureImageEditConfig(OpenAIImageEditConfig):
raise ValueError(
f"api_base is required for Azure AI Studio. Please set the api_base parameter. Passed `api_base={api_base}`"
)
original_url: Final = httpx.URL(api_base)
# Resolve api_version: litellm_params > litellm.api_version > AZURE_API_VERSION env > default.
# Mirrors the fallback chain used by the Azure chat path in common_utils.py,
# so callers that set a global / env api_version don't get an unversioned URL.
@ -105,6 +103,16 @@ class AzureImageEditConfig(OpenAIImageEditConfig):
or litellm.AZURE_DEFAULT_API_VERSION
)
v1_url: Final = BaseAzureLLM.get_azure_v1_image_url(
api_base=api_base,
api_version=api_version,
route="/openai/images/edits",
)
if v1_url is not None:
return v1_url
original_url: Final = httpx.URL(api_base)
# Create a new dictionary with existing params
query_params: Final = dict(original_url.params)

21
litellm/llms/azure/image_generation/http_utils.py
View file

@ -1,7 +1,9 @@
"""HTTP helpers for Azure OpenAI image generation (REST, not SDK)."""
from typing import Final
def azure_deployment_image_generation_json_body(api_base: str, data: dict) -> dict:
def azure_deployment_image_generation_json_body(api_base: str, data: dict, deployment_name: str | None = None) -> dict:
"""
Build the JSON body for Azure OpenAI image generation POSTs.
@ -9,9 +11,20 @@ def azure_deployment_image_generation_json_body(api_base: str, data: dict) -> di
deployment in the URL only; sending ``model`` in the body (especially the deployment
name) breaks some models (e.g. gpt-image-2). See LiteLLM #26316.
For the v1 surface (``.../openai/v1/images/...``), Azure routes by the deployment
name in the body ``model`` field, so the deployment name must replace any base
model name there or Azure answers 404 DeploymentNotFound.
Provider-style URLs (e.g. ``/providers/...`` for FLUX on Azure AI) keep all keys
so non–OpenAI-deployment payloads still work.
"""
if "images/generations" in api_base and "/openai/deployments/" in api_base:
return {k: v for k, v in data.items() if k != "model"}
return data
drop_model: Final = "images/generations" in api_base and "/openai/deployments/" in api_base
v1_route: Final = "/openai/v1/images/" in api_base and bool(deployment_name)
if not drop_model and not v1_route:
return data
entries: Final = (
tuple((k, v) for k, v in data.items() if k != "model")
if drop_model
else (*data.items(), ("model", deployment_name))
)
return {k: v for k, v in entries}

5
litellm/llms/base_llm/managed_resources/base_managed_resource.py
View file

@ -13,6 +13,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
build_list_page,
build_owner_filter,
can_access_resource,
resolve_resource_owner_id,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import SpecialEnums
@ -185,7 +186,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]):
"resource_object": resource_object,
"model_mappings": model_mappings,
"flat_model_resource_ids": list(model_mappings.values()),
"created_by": user_api_key_dict.user_id,
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
"updated_by": user_api_key_dict.user_id,
}
@ -207,7 +208,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]):
"unified_resource_id": unified_resource_id,
"model_mappings": json.dumps(model_mappings),
"flat_model_resource_ids": list(model_mappings.values()),
"created_by": user_api_key_dict.user_id,
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
"updated_by": user_api_key_dict.user_id,
}

49
litellm/llms/base_llm/managed_resources/isolation.py
View file

@ -3,10 +3,11 @@ Tenant-isolation helpers for managed file/batch/vector-store resources.
Returns a Prisma filter and an ownership check that scope managed resources
to the caller's identity: proxy admins see everything, user-keyed callers
see records they created, and service-account keys (no user_id) fall back
to the resource's owning team. Callers with no admin role and no
identifying ids are denied so an empty user_id can never select an
unscoped query.
see records they created, service-account keys (no user_id) fall back to
the resource's owning team, and keys with neither a user_id nor a team_id
fall back to their own hashed token so they can still reach the resources
they created. Callers with no admin role and no identifying ids at all
are denied so an empty user_id can never select an unscoped query.
"""
from typing import Any, Final
@ -19,6 +20,32 @@ from litellm.proxy._types import (
)
def resolve_resource_owner_id(
user_api_key_dict: UserAPIKeyAuth,
) -> str | None:
"""Return the identity to stamp on (and match against) a managed
resource's ``created_by``.
A key with neither a user_id nor a team_id would otherwise stamp
``created_by=None`` and be locked out of its own resources, so it owns
them under its hashed token instead, using the ``key:`` scope prefix
already used by ``proxy/common_utils/resource_ownership.py``. ``None``
means the caller has no usable identity of its own and must fall back
to team scoping, or be denied.
"""
if user_api_key_dict.user_id is not None:
return user_api_key_dict.user_id
if user_api_key_dict.team_id is not None:
return None
token: Final = user_api_key_dict.token or user_api_key_dict.api_key
if token:
return f"key:{token}"
return None
def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, Any]:
"""Build the OpenAI-style paginated list response shape used by managed
file/batch/vector-store listings. ``first_id`` and ``last_id`` are
@ -39,7 +66,8 @@ def build_owner_filter(
to records the caller is allowed to see.
- ``{}`` means no scoping (proxy admins).
- ``{"created_by": <user_id>}`` for user-keyed callers.
- ``{"created_by": <owner_id>}`` for user-keyed callers, and for keys
with no user_id and no team_id (owner id is their hashed token).
- ``{"team_id": <team_id>}`` for service-account callers
that have a team but no user_id.
- ``{"OR": [...]}`` when the caller has both — listing must include
@ -62,12 +90,13 @@ def build_owner_filter(
]
}
if user_id is not None:
return {"created_by": user_id}
if team_id is not None:
return {"team_id": team_id}
owner_id: Final = resolve_resource_owner_id(user_api_key_dict)
if owner_id is not None:
return {"created_by": owner_id}
return None
@ -86,8 +115,8 @@ def can_access_resource(
if _user_has_admin_view(user_api_key_dict):
return True
user_id: Final = user_api_key_dict.user_id
if user_id is not None and created_by is not None and created_by == user_id:
owner_id: Final = resolve_resource_owner_id(user_api_key_dict)
if owner_id is not None and created_by is not None and created_by == owner_id:
return True
team_id: Final = user_api_key_dict.team_id

4
litellm/llms/bedrock/embed/cohere_transformation.py
View file

@ -20,7 +20,9 @@ class BedrockCohereEmbeddingConfig:
def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict:
for k, v in non_default_params.items():
if k == "encoding_format":
optional_params["embedding_types"] = v if isinstance(v, list) else [v]
optional_params["embedding_types"] = [
"float" if fmt == "base64" else fmt for fmt in (tuple(v) if isinstance(v, list) else (v,))
]
elif k == "dimensions":
optional_params["output_dimension"] = v
return optional_params

16
litellm/llms/soniox/common_utils.py
View file

@ -138,13 +138,25 @@ def _soniox_token_to_subtitle_token(token: SonioxToken) -> SubtitleToken:
)
def _subtitle_tokens(tokens: Sequence[SonioxToken]) -> tuple[SubtitleToken, ...]:
"""
Convert Soniox tokens for subtitle rendering, excluding translation tokens
(``translation_status == "translation"``): Soniox does not timestamp them,
so they cannot be aligned to the audio and would otherwise mix translated
text into original-language cues.
"""
return tuple(
_soniox_token_to_subtitle_token(token) for token in tokens if token.get("translation_status") != "translation"
)
def render_soniox_tokens_as_srt(tokens: Sequence[SonioxToken]) -> str:
"""
Render Soniox tokens as SRT (SubRip) subtitle format.
Returns an empty string if no tokens have timestamp data.
"""
return render_subtitle_tokens_as_srt(tuple(_soniox_token_to_subtitle_token(token) for token in tokens))
return render_subtitle_tokens_as_srt(_subtitle_tokens(tokens))
def render_soniox_tokens_as_vtt(tokens: Sequence[SonioxToken]) -> str:
@ -153,4 +165,4 @@ def render_soniox_tokens_as_vtt(tokens: Sequence[SonioxToken]) -> str:
Returns the VTT header even if no cues are present.
"""
return render_subtitle_tokens_as_vtt(tuple(_soniox_token_to_subtitle_token(token) for token in tokens))
return render_subtitle_tokens_as_vtt(_subtitle_tokens(tokens))

52
litellm/llms/vertex_ai/videos/transformation.py
View file

@ -8,12 +8,14 @@ Based on: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-refer
import base64
import time
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, TypedDict, cast
import httpx
from httpx._types import FileContent, RequestFiles
from typing_extensions import ReadOnly
import litellm
from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS
from litellm.images.utils import ImageEditRequestUtils
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
@ -119,6 +121,23 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
3. Extract video data (base64) from response
"""
_OPENAI_VIDEO_SIZE_TO_ASPECT_RATIO: ClassVar[Mapping[str, str]] = MappingProxyType(
{
"1280x720": "16:9",
"1920x1080": "16:9",
"720x1280": "9:16",
"1080x1920": "9:16",
}
)
_OPENAI_VIDEO_SIZE_TO_RESOLUTION: ClassVar[Mapping[str, str]] = MappingProxyType(
{
"1280x720": "720p",
"1920x1080": "1080p",
"720x1280": "720p",
"1080x1920": "1080p",
}
)
def __init__(self):
BaseVideoConfig.__init__(self)
VertexBase.__init__(self)
@ -161,6 +180,9 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
- prompt → prompt (in instances)
- input_reference → image (in instances)
- size → aspectRatio (e.g., "1280x720" → "16:9")
- size → resolution for models with resolution-tier pricing when inferable
("1280x720"/"720x1280" → "720p", "1920x1080"/"1080x1920" → "1080p");
skipped if ``resolution`` is already set
- seconds → durationSeconds (defaults to 4 seconds if not provided)
"""
mapped_params: Final[dict[str, object]] = {}
@ -175,6 +197,9 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
if "parameters" in video_create_optional_params:
mapped_params["parameters"] = video_create_optional_params["parameters"]
if "resolution" in video_create_optional_params:
mapped_params["resolution"] = video_create_optional_params["resolution"]
# Map size to aspectRatio
if "size" in video_create_optional_params:
size: Final = video_create_optional_params["size"]
@ -182,6 +207,15 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
aspect_ratio: Final = self._convert_size_to_aspect_ratio(size)
if aspect_ratio:
mapped_params["aspectRatio"] = aspect_ratio
nested_params: Final = video_create_optional_params.get("parameters")
has_resolution = "resolution" in mapped_params or (
isinstance(nested_params, dict) and nested_params.get("resolution") is not None
)
supports_resolution = self._supports_resolution_inference(model)
if supports_resolution and not has_resolution:
inferred_resolution = self._convert_size_to_resolution(size)
if inferred_resolution is not None:
mapped_params["resolution"] = inferred_resolution
# Map seconds to durationSeconds, default to 4 seconds (matching OpenAI)
if "seconds" in video_create_optional_params:
@ -205,14 +239,16 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
if not size:
return None
aspect_ratio_map: Final = {
"1280x720": "16:9",
"1920x1080": "16:9",
"720x1280": "9:16",
"1080x1920": "9:16",
}
return self._OPENAI_VIDEO_SIZE_TO_ASPECT_RATIO.get(size, "16:9")
return aspect_ratio_map.get(size, "16:9")
def _convert_size_to_resolution(self, size: str) -> str | None:
return self._OPENAI_VIDEO_SIZE_TO_RESOLUTION.get(size)
@staticmethod
def _supports_resolution_inference(model: str) -> bool:
model_key: Final = model if model.startswith("vertex_ai/") else f"vertex_ai/{model}"
model_info: Final = litellm.model_cost.get(model_key)
return model_info is not None and model_info.get("output_cost_per_second_1080p") is not None
def validate_environment(
self,

16
litellm/model_prices_and_context_window_backup.json
View file

@ -43318,6 +43318,22 @@
"video"
]
},
"vertex_ai/veo-3.1-lite-generate-001": {
"litellm_provider": "vertex_ai-video-models",
"max_input_tokens": 1024,
"max_tokens": 1024,
"mode": "video_generation",
"output_cost_per_second": 0.05,
"output_cost_per_second_1080p": 0.08,
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#veo",
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"video"
]
},
"voyage/rerank-2": {
"input_cost_per_token": 5e-08,
"litellm_provider": "voyage",

25
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Optional
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
import httpx
from fastapi import APIRouter, Form, HTTPException, Request
from fastapi import APIRouter, Depends, Form, HTTPException, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError
@ -46,6 +46,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
aggregate_authorize,
aggregate_token,
complete_connect_flow,
introspect_gateway_token,
is_gateway_dcr_client_id,
is_proxy_api_resource,
native_client_auth_contract,
@ -67,6 +68,7 @@ from litellm.proxy._experimental.mcp_server.proxy_api_credentials import (
mint_proxy_credential,
)
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
@ -1951,6 +1953,26 @@ async def revoke_endpoint(request: Request, token: str = Form(...), client_id: s
return await revoke_refresh_token(token=token, client_id=client_id, master_key=master_key, cache=user_api_key_cache)
@router.post("/introspect", dependencies=[Depends(user_api_key_auth)])
async def introspect_endpoint(token: str = Form(...)) -> Response:
"""RFC 7662 introspection for gateway-issued session tokens (``llm_session_`` /
``llm_srefresh_``), so an external gateway can validate them without the signing
secret. The caller authenticates with a LiteLLM virtual key (section 2.1, enforced by
the route dependency); any token the gateway cannot vouch for answers
``{"active": false}`` with no further detail."""
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # circular import at module load
master_key,
user_api_key_cache,
)
return await introspect_gateway_token(
token=token,
master_key=master_key,
reload_user=_reload_active_user_by_id,
cache=user_api_key_cache,
)
@router.get("/.well-known/litellm-cli-auth")
async def native_client_auth_discovery(request: Request) -> JSONResponse:
"""The versioned contract a native client (``lite login --pkce``, or a CLI in any other
@ -2456,6 +2478,7 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict:
"issuer": f"{request_base_url}/mcp",
"authorization_endpoint": f"{request_base_url}/authorize",
"token_endpoint": f"{request_base_url}/token",
"introspection_endpoint": f"{request_base_url}/introspect",
"registration_endpoint": f"{request_base_url}/register",
"response_types_supported": ["code"],
"scopes_supported": [],

103
litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
View file

@ -70,13 +70,19 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent
open_session_refresh_bearer,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
SESSION_ISSUER,
SESSION_REFRESH_TTL_SECONDS,
MintedSessionToken,
OpenedSessionToken,
SessionAudience,
SessionPrincipal,
SessionSigningKeys,
is_session_refresh_token,
is_session_token,
mint_session_refresh_token,
mint_session_token,
open_session_refresh_token,
open_session_token,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
@ -885,6 +891,23 @@ class _SingleUseGuard:
count = await self._cache.async_increment_cache(key, 1, ttl=ttl_seconds, local_only=True)
return "first" if count == 1 else "replayed"
async def peek(self, key: str) -> Literal["unclaimed", "claimed", "unavailable"]:
"""Read-only view of a single-use marker, resolved against the same shared authority as
:meth:`claim` so introspection observes exactly the record redemption and revocation wrote.
A backend fault is ``"unavailable"`` (fail closed) rather than a guess either way."""
from litellm.proxy.proxy_server import redis_usage_cache # noqa: PLC0415 # circular import at module load
redis_cache: Final = redis_usage_cache or getattr(self._cache, "redis_cache", None)
if redis_cache is not None:
try:
value = await redis_cache.async_get_cache(key)
except Exception as e: # noqa: BLE001 # ANY Redis fault fails the read closed
verbose_logger.warning("mcp gateway single-use peek: shared cache backend unavailable: %s", e)
return "unavailable"
return "unclaimed" if value is None else "claimed"
local: Final = await self._cache.async_get_cache(key, local_only=True)
return "unclaimed" if local is None else "claimed"
def _session_token_pair(principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime) -> Response:
access: Final = mint_session_token(principal, keys, now)
@ -1199,3 +1222,83 @@ async def revoke_refresh_token(token: str, client_id: str, master_key: str | Non
if burned == "unavailable":
return _oauth_error(503, "temporarily_unavailable", _CLAIM_UNAVAILABLE_DESCRIPTION)
return Response(content="{}", media_type="application/json", headers=TOKEN_NO_CACHE_HEADERS)
def _inactive_introspection_response() -> Response:
"""RFC 7662 section 2.2: any token the gateway cannot vouch for, whatever the reason
(wrong family, bad signature, expired, revoked, or a deactivated user), answers 200
with ``active: false`` and nothing else, so introspection is not a token oracle."""
return JSONResponse(status_code=200, content={"active": False}, headers=TOKEN_NO_CACHE_HEADERS)
def _active_introspection_response(opened: OpenedSessionToken) -> Response:
principal: Final = opened.principal
optional_claims: Final = {
key: value
for key, value in (
("token_type", "Bearer" if opened.kind == "session" else None),
("team_id", principal.team_id),
("resource_server_id", principal.resource_server_id),
("audience", principal.audience),
)
if value is not None
}
return JSONResponse(
status_code=200,
content={
"active": True,
"iss": SESSION_ISSUER,
"sub": principal.user_id,
"client_id": principal.client_id,
"jti": opened.jti,
"iat": opened.iat,
"exp": opened.exp,
"kind": opened.kind,
**optional_claims,
},
headers=TOKEN_NO_CACHE_HEADERS,
)
async def introspect_gateway_token(
token: str,
master_key: str | None,
reload_user: ReloadUser,
cache: DualCache,
) -> Response:
"""RFC 7662 introspection for the gateway's session tokens, so an external gateway
(Kong, an API management layer) can validate a LiteLLM-issued MCP session credential
without holding the signing secret. The caller is already authenticated by the route
(section 2.1). Active means everything admission itself would require: valid signature
under the configured session signing keys, unexpired, not a revoked or rotated refresh
token, and a litellm user that is still live, so a deactivated user's outstanding
tokens introspect as inactive immediately. A shared-backend or DB outage answers 503
rather than guessing in either direction."""
if master_key is None:
verbose_logger.error("mcp_gateway_dcr introspect rejected: no master_key configured")
return _oauth_error(500, "server_error", "the gateway has no master key configured")
keys: Final = active_session_signing_keys(master_key)
if isinstance(keys, SessionSigningConfigError):
verbose_logger.error("mcp_gateway_dcr introspect rejected: %s", keys.detail)
return _oauth_error(500, "server_error", keys.detail)
now: Final = datetime.now(timezone.utc)
if is_session_token(token):
opened = open_session_token(token, keys, now)
elif is_session_refresh_token(token):
opened = open_session_refresh_token(token, keys, now)
else:
return _inactive_introspection_response()
if not isinstance(opened, OpenedSessionToken):
return _inactive_introspection_response()
if opened.kind == "session_refresh":
peeked: Final = await _SingleUseGuard(cache).peek(f"{_USED_REFRESH_CACHE_PREFIX}{opened.jti}")
if peeked == "unavailable":
return _oauth_error(503, "temporarily_unavailable", _CLAIM_UNAVAILABLE_DESCRIPTION)
if peeked == "claimed":
return _inactive_introspection_response()
failure: Final = await reload_user(opened.principal.user_id)
if failure == "unavailable":
return _oauth_error(503, "temporarily_unavailable", "the gateway database is unavailable; retry")
if failure is not None:
return _inactive_introspection_response()
return _active_introspection_response(opened)

12
litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py
View file

@ -221,12 +221,17 @@ class MintedSessionToken(BaseModel):
class OpenedSessionToken(BaseModel):
"""A validated session token of either kind: the principal it was minted for, plus the
``jti`` so the token endpoint can enforce single-use rotation on a refresh token."""
"""A validated session token of either kind: the principal it was minted for, the
``jti`` so the token endpoint can enforce single-use rotation on a refresh token, and
the signed ``kind``/``iat``/``exp`` so an introspection response can report the
token's metadata without re-decoding."""
model_config = ConfigDict(frozen=True)
principal: SessionPrincipal
jti: str
kind: SessionTokenKind
iat: int
exp: int
class SessionTokenTooLarge(BaseModel):
@ -458,6 +463,9 @@ def _open(
team_id=claims.team_id,
),
jti=claims.jti,
kind=claims.kind,
iat=claims.iat,
exp=claims.exp,
)

1
litellm/proxy/_lazy_features.py
View file

@ -161,6 +161,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
"/callback",
"/register",
"/revoke",
"/introspect",
),
# Catches the /{mcp_server_name}/authorize|token|register variants.
path_suffixes=("/authorize", "/token", "/register"),

347
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -6459,6 +6459,109 @@
"title": "ConfigOverrideSettingsResponse",
"type": "object"
},
"CyberArkConfig": {
"description": "Configuration for CyberArk Conjur secret manager integration.",
"properties": {
"client_cert": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"description": "Path to the client TLS certificate for certificate-based authentication",
"title": "Client Cert"
},
"client_key": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"description": "Path to the client TLS private key for certificate-based authentication",
"title": "Client Key"
},
"cyberark_account": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"description": "The Conjur organization account name",
"title": "Cyberark Account"
},
"cyberark_api_base": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"description": "The address of the CyberArk Conjur server (e.g., https://conjur.example.com)",
"title": "Cyberark Api Base"
},
"cyberark_api_key": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"description": "API key for Conjur API-key authentication",
"title": "Cyberark Api Key"
},
"cyberark_username": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"description": "The Conjur username (login) to authenticate as",
"title": "Cyberark Username"
},
"refresh_interval": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"description": "Auth token cache TTL in seconds (default: 300)",
"title": "Refresh Interval"
},
"ssl_verify": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"description": "Set to false to disable SSL verification (e.g., for self-signed certificates)",
"title": "Ssl Verify"
}
},
"title": "CyberArkConfig",
"type": "object"
},
"HTTPValidationError": {
"properties": {
"detail": {
@ -6654,6 +6757,192 @@
}
},
"paths": {
"/config_overrides/cyberark": {
"delete": {
"description": "Delete CyberArk Conjur configuration. Idempotent.",
"operationId": "delete_cyberark_config_config_overrides_cyberark_delete",
"parameters": [
{
"description": "The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
"in": "header",
"name": "litellm-changed-by",
"required": false,
"schema": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"description": "The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
"title": "Litellm-Changed-By"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"additionalProperties": {
"type": "string"
},
"title": "Response Delete Cyberark Config Config Overrides Cyberark Delete",
"type": "object"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Delete Cyberark Config",
"tags": [
"config_overrides"
]
},
"get": {
"description": "Get current CyberArk Conjur configuration.\nReturns decrypted values from DB, or falls back to current env vars.\nSensitive fields are masked before leaving the server.",
"operationId": "get_cyberark_config_config_overrides_cyberark_get",
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/ConfigOverrideSettingsResponse"
}
}
},
"description": "Successful Response"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Get Cyberark Config",
"tags": [
"config_overrides"
]
},
"post": {
"description": "Update CyberArk Conjur secret manager configuration.\nSets environment variables, encrypts sensitive fields, and stores in DB.\nReinitializes the secret manager on this pod.",
"operationId": "update_cyberark_config_config_overrides_cyberark_post",
"parameters": [
{
"description": "The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
"in": "header",
"name": "litellm-changed-by",
"required": false,
"schema": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"description": "The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
"title": "Litellm-Changed-By"
}
}
],
"requestBody": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/CyberArkConfig"
}
}
},
"required": true
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"additionalProperties": {
"type": "string"
},
"title": "Response Update Cyberark Config Config Overrides Cyberark Post",
"type": "object"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Update Cyberark Config",
"tags": [
"config_overrides"
]
}
},
"/config_overrides/cyberark/test_connection": {
"post": {
"description": "Test the connection to the currently configured CyberArk Conjur server.\nUses the already-initialized secret manager client. Does not modify any state.",
"operationId": "test_cyberark_connection_config_overrides_cyberark_test_connection_post",
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"additionalProperties": {
"type": "string"
},
"title": "Response Test Cyberark Connection Config Overrides Cyberark Test Connection Post",
"type": "object"
}
}
},
"description": "Successful Response"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Test Cyberark Connection",
"tags": [
"config_overrides"
]
}
},
"/config_overrides/hashicorp_vault": {
"delete": {
"description": "Delete Hashicorp Vault configuration. Idempotent.",
@ -16555,6 +16844,19 @@
"title": "Body_authorize_complete_authorize_complete_post",
"type": "object"
},
"Body_introspect_endpoint_introspect_post": {
"properties": {
"token": {
"title": "Token",
"type": "string"
}
},
"required": [
"token"
],
"title": "Body_introspect_endpoint_introspect_post",
"type": "object"
},
"Body_revoke_endpoint_revoke_post": {
"properties": {
"client_id": {
@ -19134,6 +19436,51 @@
]
}
},
"/introspect": {
"post": {
"description": "RFC 7662 introspection for gateway-issued session tokens (``llm_session_`` /\n``llm_srefresh_``), so an external gateway can validate them without the signing\nsecret. The caller authenticates with a LiteLLM virtual key (section 2.1, enforced by\nthe route dependency); any token the gateway cannot vouch for answers\n``{\"active\": false}`` with no further detail.",
"operationId": "introspect_endpoint_introspect_post",
"requestBody": {
"content": {
"application/x-www-form-urlencoded": {
"schema": {
"$ref": "#/components/schemas/Body_introspect_endpoint_introspect_post"
}
}
},
"required": true
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Introspect Endpoint",
"tags": [
"mcp_discoverable"
]
}
},
"/register": {
"post": {
"operationId": "register_client_register_post",

3
litellm/proxy/_types.py
View file

@ -504,6 +504,7 @@ class LiteLLMRoutes(enum.Enum):
"/mcp-rest/tools/list",
"/mcp-rest/tools/call",
"/v1/mcp/tools",
"/introspect",
]
# MCP server CRUD routes — control-plane. Gated by DISABLE_ADMIN_ENDPOINTS.
@ -3573,6 +3574,8 @@ class SpendLogsMetadata(TypedDict):
status: StandardLoggingPayloadStatus
proxy_server_request: str | None
batch_models: list[str] | None
batch_successful_requests: int | None # writable-ok: built by assignment like every sibling key in this TypedDict
batch_failed_requests: int | None # writable-ok: built by assignment like every sibling key in this TypedDict
error_information: StandardLoggingPayloadErrorInformation | None
usage_object: dict | None
model_map_information: StandardLoggingModelInformation | None

11
litellm/proxy/client/chat.py
View file

@ -8,16 +8,19 @@ from .exceptions import UnauthorizedError
class ChatClient:
def __init__(self, base_url: str, api_key: str | None = None):
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 600):
"""
Initialize the ChatClient.
Args:
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000")
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
timeout (int): Request timeout in seconds (default: 600, the OpenAI SDK default, since a completion
can legitimately take minutes)
"""
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
self._api_key = api_key
self._timeout = timeout
def _get_headers(self) -> dict[str, str]:
"""
@ -96,7 +99,7 @@ class ChatClient:
# Prepare and send the request
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
except requests.exceptions.HTTPError as e:
@ -161,7 +164,9 @@ class ChatClient:
# Make streaming request
session: Final = requests.Session()
try:
response: Final = session.post(url, headers=self._get_headers(), json=data, stream=True)
response: Final = session.post(
url, headers=self._get_headers(), json=data, stream=True, timeout=self._timeout
)
response.raise_for_status()
# Parse SSE stream

10
litellm/proxy/client/cli/commands/auth.py
View file

@ -99,11 +99,6 @@ class CliPollData(TypedDict, total=False):
team_id: str
class CliPollRequestKwargs(TypedDict, total=False):
timeout: int
headers: dict[str, str]
class CliSsoStartData(TypedDict):
login_id: str
poll_secret: str
@ -518,10 +513,7 @@ def _poll_for_ready_data(
) -> CliPollData | None:
for attempt in range(total_timeout // poll_interval):
try:
request_kwargs: CliPollRequestKwargs = {"timeout": request_timeout}
if headers is not None:
request_kwargs["headers"] = headers
response = requests.get(url, **request_kwargs)
response = requests.get(url, headers=headers, timeout=request_timeout)
if response.status_code == 200:
data: CliPollData = response.json()
status = data.get("status")

13
litellm/proxy/client/client.py
View file

@ -24,7 +24,8 @@ class Client:
Args:
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:4000")
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
timeout: Request timeout in seconds (default: 30)
timeout: Request timeout in seconds for management calls (default: 30). Chat completions keep
ChatClient's own 600 second default, since a completion can legitimately take minutes
"""
self._base_url = base_url.rstrip("/")
# Only use the stored CLI key when it was issued for this server.
@ -33,9 +34,9 @@ class Client:
# Initialize resource clients
self.http = HTTPClient(base_url=base_url, api_key=self._api_key, timeout=timeout)
self.models = ModelsManagementClient(base_url=self._base_url, api_key=self._api_key)
self.model_groups = ModelGroupsManagementClient(base_url=self._base_url, api_key=self._api_key)
self.models = ModelsManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout)
self.model_groups = ModelGroupsManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout)
self.chat = ChatClient(base_url=self._base_url, api_key=self._api_key)
self.keys = KeysManagementClient(base_url=self._base_url, api_key=self._api_key)
self.credentials = CredentialsManagementClient(base_url=self._base_url, api_key=self._api_key)
self.teams = TeamsManagementClient(base_url=self._base_url, api_key=self._api_key)
self.keys = KeysManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout)
self.credentials = CredentialsManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout)
self.teams = TeamsManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout)

12
litellm/proxy/client/credentials.py
View file

@ -6,16 +6,18 @@ from .exceptions import UnauthorizedError
class CredentialsManagementClient:
def __init__(self, base_url: str, api_key: str | None = None):
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
"""
Initialize the CredentialsManagementClient.
Args:
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000")
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
timeout (int): Request timeout in seconds (default: 30)
"""
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
self._api_key = api_key
self._timeout = timeout
def _get_headers(self) -> dict[str, str]:
"""
@ -56,7 +58,7 @@ class CredentialsManagementClient:
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
except requests.exceptions.HTTPError as e:
@ -103,7 +105,7 @@ class CredentialsManagementClient:
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
except requests.exceptions.HTTPError as e:
@ -140,7 +142,7 @@ class CredentialsManagementClient:
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
except requests.exceptions.HTTPError as e:
@ -177,7 +179,7 @@ class CredentialsManagementClient:
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
except requests.exceptions.HTTPError as e:

14
litellm/proxy/client/keys.py
View file

@ -9,16 +9,18 @@ from .exceptions import UnauthorizedError
class KeysManagementClient:
def __init__(self, base_url: str, api_key: str | None = None):
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
"""
Initialize the KeysManagementClient.
Args:
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000")
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
timeout (int): Request timeout in seconds (default: 30)
"""
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
self._api_key = api_key
self._timeout = timeout
def _get_headers(self) -> dict[str, str]:
"""
@ -99,7 +101,7 @@ class KeysManagementClient:
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
except requests.exceptions.HTTPError as e:
@ -174,7 +176,7 @@ class KeysManagementClient:
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
except requests.exceptions.HTTPError as e:
@ -218,7 +220,7 @@ class KeysManagementClient:
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
except requests.exceptions.HTTPError as e:
@ -279,7 +281,7 @@ class KeysManagementClient:
session: Final = requests.Session()
response_text: str | None = None
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response_text = response.text
response.raise_for_status()
return response.json()
@ -309,7 +311,7 @@ class KeysManagementClient:
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
except requests.exceptions.HTTPError as e:

6
litellm/proxy/client/model_groups.py
View file

@ -6,16 +6,18 @@ from .exceptions import UnauthorizedError
class ModelGroupsManagementClient:
def __init__(self, base_url: str, api_key: str | None = None):
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
"""
Initialize the ModelGroupsManagementClient.
Args:
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000")
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
timeout (int): Request timeout in seconds (default: 30)
"""
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
self._api_key = api_key
self._timeout = timeout
def _get_headers(self) -> dict[str, str]:
"""
@ -53,7 +55,7 @@ class ModelGroupsManagementClient:
# Prepare and send the request
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()["data"]
except requests.exceptions.HTTPError as e:

14
litellm/proxy/client/models.py
View file

@ -7,16 +7,18 @@ from .exceptions import NotFoundError, UnauthorizedError
class ModelsManagementClient:
def __init__(self, base_url: str, api_key: str | None = None):
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
"""
Initialize the ModelsManagementClient.
Args:
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000")
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
timeout (int): Request timeout in seconds (default: 30)
"""
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
self._api_key = api_key
self._timeout = timeout
def _get_headers(self) -> dict[str, str]:
"""
@ -55,7 +57,7 @@ class ModelsManagementClient:
# Prepare and send the request
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()["data"]
except requests.exceptions.HTTPError as e:
@ -104,7 +106,7 @@ class ModelsManagementClient:
# Prepare and send the request
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
except requests.exceptions.HTTPError as e:
@ -140,7 +142,7 @@ class ModelsManagementClient:
# Prepare and send the request
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
except requests.exceptions.HTTPError as e:
@ -232,7 +234,7 @@ class ModelsManagementClient:
# Prepare and send the request
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()["data"]
except requests.exceptions.HTTPError as e:
@ -282,7 +284,7 @@ class ModelsManagementClient:
# Prepare and send the request
session: Final = requests.Session()
try:
response: Final = session.send(request.prepare())
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
except requests.exceptions.HTTPError as e:

10
litellm/proxy/client/teams.py
View file

@ -11,16 +11,18 @@ from .exceptions import UnauthorizedError
class TeamsManagementClient:
"""Client for managing teams in LiteLLM proxy."""
def __init__(self, base_url: str, api_key: str | None = None):
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
"""
Initialize the TeamsManagementClient.
Args:
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:4000")
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
timeout (int): Request timeout in seconds (default: 30)
"""
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
self._api_key = api_key
self._timeout = timeout
def _get_headers(self) -> dict[str, str]:
"""
@ -60,7 +62,7 @@ class TeamsManagementClient:
if organization_id:
params["organization_id"] = organization_id
response: Final = requests.get(url, headers=self._get_headers(), params=params)
response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self._timeout)
if response.status_code == 401:
raise UnauthorizedError("Authentication failed. Check your API key.")
@ -117,7 +119,7 @@ class TeamsManagementClient:
if sort_by:
params["sort_by"] = sort_by
response: Final = requests.get(url, headers=self._get_headers(), params=params)
response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self._timeout)
if response.status_code == 401:
raise UnauthorizedError("Authentication failed. Check your API key.")
@ -138,7 +140,7 @@ class TeamsManagementClient:
"""
url: Final = f"{self._base_url}/team/available"
response: Final = requests.get(url, headers=self._get_headers())
response: Final = requests.get(url, headers=self._get_headers(), timeout=self._timeout)
if response.status_code == 401:
raise UnauthorizedError("Authentication failed. Check your API key.")

15
litellm/proxy/client/users.py
View file

@ -6,9 +6,10 @@ from .exceptions import NotFoundError, UnauthorizedError
class UsersManagementClient:
def __init__(self, base_url: str, api_key: str | None = None):
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
self.base_url = base_url.rstrip("/")
self.api_key = api_key
self.timeout = timeout
def _get_headers(self) -> dict[str, str]:
headers: Final = {"Content-Type": "application/json"}
@ -19,7 +20,7 @@ class UsersManagementClient:
def list_users(self, params: dict[str, Any] | None = None) -> list[dict[str, Any]]:
"""List users (GET /user/list)"""
url: Final = f"{self.base_url}/user/list"
response: Final = requests.get(url, headers=self._get_headers(), params=params)
response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self.timeout)
if response.status_code == 401:
raise UnauthorizedError(response.text)
response.raise_for_status()
@ -29,7 +30,7 @@ class UsersManagementClient:
"""Get user info (GET /user/info)"""
url: Final = f"{self.base_url}/user/info"
params: Final = {"user_id": user_id} if user_id else {}
response: Final = requests.get(url, headers=self._get_headers(), params=params)
response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self.timeout)
if response.status_code == 401:
raise UnauthorizedError(response.text)
if response.status_code == 404:
@ -41,7 +42,7 @@ class UsersManagementClient:
"""Get user info v2 - lightweight, returns only user object (GET /v2/user/info)"""
url: Final = f"{self.base_url}/v2/user/info"
params: Final = {"user_id": user_id} if user_id else {}
response: Final = requests.get(url, headers=self._get_headers(), params=params)
response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self.timeout)
if response.status_code == 401:
raise UnauthorizedError(response.text)
if response.status_code == 404:
@ -52,7 +53,7 @@ class UsersManagementClient:
def create_user(self, user_data: dict[str, Any]) -> dict[str, Any]:
"""Create a new user (POST /user/new)"""
url: Final = f"{self.base_url}/user/new"
response: Final = requests.post(url, headers=self._get_headers(), json=user_data)
response: Final = requests.post(url, headers=self._get_headers(), json=user_data, timeout=self.timeout)
if response.status_code == 401:
raise UnauthorizedError(response.text)
response.raise_for_status()
@ -61,7 +62,9 @@ class UsersManagementClient:
def delete_user(self, user_ids: list[str]) -> dict[str, Any]:
"""Delete users (POST /user/delete)"""
url: Final = f"{self.base_url}/user/delete"
response: Final = requests.post(url, headers=self._get_headers(), json={"user_ids": user_ids})
response: Final = requests.post(
url, headers=self._get_headers(), json={"user_ids": user_ids}, timeout=self.timeout
)
if response.status_code == 401:
raise UnauthorizedError(response.text)
response.raise_for_status()

20
litellm/proxy/db/routing_prisma_wrapper.py
View file

@ -61,6 +61,26 @@ class _RoutedActions:
return getattr(self._writer_actions, name)
class WriterPinnedClient:
"""PrismaClient-shaped view whose `.db` resolves to the writer while it is available.
Read-after-write paths (e.g. the model reconcile a /model/new triggers to
verify its own just-committed row) must not read through a lagging read
replica: the row is not replayed there yet, so the reconcile concludes the
write is missing and fails the request even though it is durable (#38556).
While the writer is degraded (`writer_unavailable`), the pin yields to the
routed wrapper so reconcile reads keep working from the replica: a proxy
that starts during a primary outage must still load DB-backed models, and
no read-after-write hazard exists then because writes are failing anyway.
"""
__slots__ = ("db",)
def __init__(self, db: "PrismaWrapper | RoutingPrismaWrapper") -> None:
self.db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) and not db.writer_unavailable else db
class RoutingPrismaWrapper:
"""
Routes Prisma operations between a writer and a reader Prisma client.

27
litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
View file

@ -37,8 +37,12 @@ from litellm.proxy.guardrails.guardrail_hooks.content_text import (
from litellm.proxy.spend_tracking.compression_savings import HEADROOM_GUARDRAIL_PROVIDER
from litellm.secret_managers.main import get_secret_str
from litellm.types.guardrails import GuardrailEventHooks, Mode
from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch
from litellm.types.utils import GenericGuardrailAPIInputs
from litellm.types.integrations.custom_logger import (
HEADROOM_CONVERTED_STREAM_KEY,
AgenticLoopPlan,
AgenticLoopRequestPatch,
)
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -713,6 +717,25 @@ class HeadroomGuardrail(CustomGuardrail):
return {**inputs, "structured_messages": compressed, "tools": merged_tools} # pyright: ignore[reportReturnType]
async def async_pre_call_deployment_hook(
self,
kwargs: dict[str, Any],
call_type: CallTypes | None,
) -> dict[str, Any] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict
base_result: Final = await super().async_pre_call_deployment_hook(kwargs, call_type)
effective: Final = base_result if base_result is not None else kwargs
if call_type not in (CallTypes.completion, CallTypes.acompletion):
return base_result
if not effective.get("stream"):
return base_result
if not has_headroom_retrieve_tool(effective.get("tools")):
return base_result
return { # mutable-ok: the hook contract is a plain dict the router merges into the request kwargs
**effective,
"stream": False,
HEADROOM_CONVERTED_STREAM_KEY: True,
}
async def async_should_run_agentic_loop(
self,
response: Any,

7
litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py
View file

@ -37,6 +37,9 @@ if TYPE_CHECKING:
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
_AUTH_TIMEOUT_SECONDS: Final[float] = 30.0
class _HiddenlayerEvaluation(TypedDict, total=False):
action: str
threat_level: str
@ -157,10 +160,10 @@ def is_saas(host: str) -> bool:
return False
def _get_jwt(auth_url, api_id, api_key) -> str:
def _get_jwt(auth_url, api_id, api_key, timeout: float = _AUTH_TIMEOUT_SECONDS) -> str:
token_url: Final = f"{auth_url}/oauth2/token?grant_type=client_credentials"
resp: Final = requests.post(token_url, auth=HTTPBasicAuth(api_id, api_key))
resp: Final = requests.post(token_url, auth=HTTPBasicAuth(api_id, api_key), timeout=timeout)
if not resp.ok:
raise RuntimeError(

1
litellm/proxy/litellm_pre_call_utils.py
View file

@ -254,6 +254,7 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = (
"_code_interpreter_interception_converted_stream",
"_code_interpreter_interception_sandbox_key",
"_code_interpreter_interception_session_scoped",
"_headroom_interception_converted_stream",
"max_agentic_loops",
# Recomputed below from the actual caller-controlled timeout sources (headers and
# body fields); a client-forged value here would let a request either dodge cooldown

408
litellm/proxy/management_endpoints/config_override_endpoints.py
View file

@ -3,7 +3,7 @@ import json
import os
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Final, Protocol
from typing import TYPE_CHECKING, Final, Protocol
from fastapi import APIRouter, Depends, Header, HTTPException
from pydantic import BaseModel, TypeAdapter
@ -36,10 +36,12 @@ from litellm.repositories.table_repositories import ConfigOverridesRepository
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.proxy.management_endpoints.config_overrides import (
ConfigOverrideSettingsResponse,
CyberArkConfig,
HashicorpVaultConfig,
)
if TYPE_CHECKING:
from litellm.proxy.proxy_server import ProxyConfig
from litellm.proxy.utils import PrismaClient
router: Final = APIRouter()
@ -83,18 +85,19 @@ def _log_audit_task_exception(task: "asyncio.Task[None]") -> None:
return
exc: Final = task.exception()
if exc is not None:
verbose_proxy_logger.warning("Failed to write hashicorp-vault config audit log: %s", exc)
verbose_proxy_logger.warning("Failed to write config override audit log: %s", exc)
async def _emit_hashicorp_vault_audit_log(
async def _emit_config_override_audit_log(
*,
object_id: str,
action: AUDIT_ACTIONS,
before_config: Mapping[str, object] | None,
after_config: Mapping[str, object] | None,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None,
) -> None:
"""Emit an audit-log row for a /config_overrides/hashicorp_vault mutation.
"""Emit an audit-log row for a /config_overrides/{object_id} mutation.
Mirrors the ``store_audit_logs``-gated pattern from
``team_callback_endpoints.py``. Captured under
@ -118,7 +121,7 @@ async def _emit_hashicorp_vault_audit_log(
changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name,
changed_by_api_key=user_api_key_dict.api_key,
table_name=LitellmTableNames.CONFIG_OVERRIDES_TABLE_NAME,
object_id="hashicorp_vault",
object_id=object_id,
action=action,
updated_values=json.dumps({"config": _redact_config(after_config)}, default=str),
before_value=json.dumps({"config": _redact_config(before_config)}, default=str),
@ -150,6 +153,24 @@ HASHICORP_SENSITIVE_FIELDS: Final[set[str]] = {
"client_key",
}
# --- CyberArk Conjur constants ---
CYBERARK_ENV_VAR_MAPPING: Final[dict[str, str]] = { # mutable-ok: module-level env mapping
"cyberark_api_base": "CYBERARK_API_BASE",
"cyberark_account": "CYBERARK_ACCOUNT",
"cyberark_username": "CYBERARK_USERNAME",
"cyberark_api_key": "CYBERARK_API_KEY",
"client_cert": "CYBERARK_CLIENT_CERT",
"client_key": "CYBERARK_CLIENT_KEY",
"ssl_verify": "CYBERARK_SSL_VERIFY",
"refresh_interval": "CYBERARK_REFRESH_INTERVAL",
}
CYBERARK_SENSITIVE_FIELDS: Final[set[str]] = { # mutable-ok: module-level constant, mirrors HASHICORP_SENSITIVE_FIELDS
"cyberark_api_key",
"client_key",
}
_sensitive_masker: Final = SensitiveDataMasker()
@ -215,9 +236,12 @@ def _parse_config_value(raw: str | Mapping[str, object]) -> dict[str, object]:
return dict(raw)
def _set_env_vars(config_data: Mapping[str, object]) -> None:
"""Set HCP_VAULT_* env vars from config data. Unsets vars for missing/None/empty fields."""
for field_name, env_var_name in HASHICORP_ENV_VAR_MAPPING.items():
def _set_env_vars(
config_data: Mapping[str, object],
env_var_mapping: Mapping[str, str] = HASHICORP_ENV_VAR_MAPPING,
) -> None:
"""Set mapped env vars from config data. Unsets vars for missing/None/empty fields."""
for field_name, env_var_name in env_var_mapping.items():
value = config_data.get(field_name)
if value is not None and value != "":
os.environ[env_var_name] = str(value)
@ -225,13 +249,74 @@ def _set_env_vars(config_data: Mapping[str, object]) -> None:
os.environ.pop(env_var_name, None)
def _clear_hashicorp_vault_state(proxy_config: Any) -> None:
def _clear_hashicorp_vault_state(proxy_config: "ProxyConfig") -> None:
"""Clear all Hashicorp Vault state: env vars, secret manager, and change-detection cache."""
_set_env_vars({})
if litellm._key_management_system == KeyManagementSystem.HASHICORP_VAULT:
litellm.secret_manager_client = None
litellm._key_management_system = None
proxy_config._last_hashicorp_vault_config = None
proxy_config._last_hashicorp_vault_config = None # pyright: ignore[reportPrivateUsage] # proxy-internal change-detection cache
def _snapshot_cyberark_boot_env(proxy_config: "ProxyConfig") -> None:
"""Capture deployment-provided CYBERARK_* env vars once, before the first DB-driven overwrite."""
if proxy_config._cyberark_boot_env is None: # pyright: ignore[reportPrivateUsage] # proxy-internal boot snapshot
proxy_config._cyberark_boot_env = _get_current_env_values(CYBERARK_ENV_VAR_MAPPING) # pyright: ignore[reportPrivateUsage] # proxy-internal boot snapshot
def _restore_cyberark_runtime(proxy_config: "ProxyConfig", env_values: Mapping[str, str | None]) -> None:
"""Restore CYBERARK_* env vars and reinitialize (or drop) the secret manager to match them."""
_set_env_vars(env_values, CYBERARK_ENV_VAR_MAPPING)
if env_values.get("cyberark_api_base"):
try:
proxy_config.initialize_secret_manager(key_management_system="cyberark")
except Exception: # noqa: BLE001 # restore is best-effort; fall through to dropping the manager
verbose_proxy_logger.exception("Failed to restore previous CyberArk configuration")
else:
return
if litellm._key_management_system != KeyManagementSystem.CYBERARK: # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage
return
litellm.secret_manager_client = None
litellm._key_management_system = None # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage
# Force the vault reload to re-init from its own row so no manager is stranded inactive
proxy_config._last_hashicorp_vault_config = None # pyright: ignore[reportPrivateUsage] # proxy-internal change-detection cache
if os.environ.get("HCP_VAULT_ADDR"):
try:
proxy_config.initialize_secret_manager(key_management_system="hashicorp_vault")
except Exception: # noqa: BLE001 # restore is best-effort; the vault reload loop retries from its own row
verbose_proxy_logger.exception("Failed to reinitialize Hashicorp Vault after CyberArk rollback")
def _clear_cyberark_state(proxy_config: "ProxyConfig") -> None:
"""Drop DB-driven CyberArk state, restoring deployment-provided env vars if any."""
boot_env: Final[Mapping[str, str | None]] = (
proxy_config._cyberark_boot_env or {} # pyright: ignore[reportPrivateUsage] # proxy-internal boot snapshot
)
_restore_cyberark_runtime(proxy_config, boot_env)
proxy_config._last_cyberark_config = None # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage
async def _persist_cyberark_config(
prisma_client: "PrismaClient",
proxy_config: "ProxyConfig",
config_data: Mapping[str, object],
) -> dict[str, object]:
"""Encrypt and upsert the CyberArk config row; returns the stored (encrypted) payload."""
encrypted_data: Final = proxy_config._encrypt_env_variables(dict(config_data)) # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage
config_value: Final = safe_dumps(encrypted_data)
await _config_overrides_table(prisma_client).upsert(
where={"config_type": "cyberark"}, # mutable-ok: prisma upsert payload
data={ # mutable-ok: prisma upsert payload
"create": { # mutable-ok: prisma upsert payload
"config_type": "cyberark",
"config_value": config_value,
},
"update": { # mutable-ok: prisma upsert payload
"config_value": config_value,
},
},
)
return safe_json_loads(config_value)
# --- Hashicorp Vault endpoints ---
@ -358,7 +443,8 @@ async def update_hashicorp_vault_config(
# row was absent or its ``config_value`` was NULL.
before_config: Final = existing_decrypted if existing_decrypted is not None else env_values
action: Final[AUDIT_ACTIONS] = "updated" if existing_record is not None else "created"
await _emit_hashicorp_vault_audit_log(
await _emit_config_override_audit_log(
object_id="hashicorp_vault",
action=action,
before_config=before_config,
after_config=config_data,
@ -484,7 +570,8 @@ async def delete_hashicorp_vault_config(
# Only emit audit log if a row was actually removed; an idempotent
# delete on a non-existent row produces no security-relevant change.
if deleted:
await _emit_hashicorp_vault_audit_log(
await _emit_config_override_audit_log(
object_id="hashicorp_vault",
action="deleted",
before_config=before_config,
after_config=None,
@ -529,7 +616,7 @@ async def test_hashicorp_vault_connection(
# Step 1: Authenticate (exercises AppRole login, TLS cert login, or direct token)
try:
headers: Final[dict[str, str]] = await asyncio.to_thread(client._get_request_headers)
headers: Final[Mapping[str, str]] = await asyncio.to_thread(client._get_request_headers)
except Exception as e:
raise HTTPException(
status_code=502,
@ -554,3 +641,298 @@ async def test_hashicorp_vault_connection(
"status": "success",
"message": f"Successfully connected to Vault at {client.vault_addr}",
}
# --- CyberArk Conjur endpoints ---
@router.post(
"/config_overrides/cyberark",
tags=["Config Overrides"], # mutable-ok: FastAPI route decorator metadata
dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI route decorator metadata
)
async def update_cyberark_config(
config: CyberArkConfig,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection
litellm_changed_by: str | None = Header(
None,
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
),
) -> dict[str, str]:
"""
Update CyberArk Conjur secret manager configuration.
Sets environment variables, encrypts sensitive fields, and stores in DB.
Reinitializes the secret manager on this pod.
"""
from litellm.proxy.proxy_server import prisma_client, proxy_config
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Only admin users can update config overrides",
)
if prisma_client is None:
raise HTTPException(
status_code=500,
detail=CommonProxyErrors.db_not_connected_error.value,
)
config_data: dict[str, object] = config.model_dump(exclude_none=True) # mutable-ok: merged # rebind-ok: stripped
# Merge ALL fields the user didn't send: try DB first, fall back to env vars.
# Omitted field = keep existing; empty string = clear/remove the field.
existing_record: Final = await _config_overrides_table(prisma_client).find_unique(
where={"config_type": "cyberark"} # mutable-ok: prisma where clause
)
existing_decrypted: dict[str, object] | None = None # mutable-ok: DB payload # rebind-ok: set when record exists
env_values: dict[str, str | None] = {} # mutable-ok: env snapshot # rebind-ok: populated when no DB record exists
if existing_record is not None and existing_record.config_value is not None:
existing_data: Final = _parse_config_value(existing_record.config_value)
existing_decrypted = proxy_config._decrypt_db_variables(existing_data) # pyright: ignore[reportPrivateUsage] # rebind-ok: populated when a prior record decrypts
for field in CYBERARK_ENV_VAR_MAPPING:
if field not in config_data and existing_decrypted.get(field):
config_data[field] = existing_decrypted[field]
else:
env_values = _get_current_env_values(CYBERARK_ENV_VAR_MAPPING) # rebind-ok: populated when no DB record exists
for field in CYBERARK_ENV_VAR_MAPPING:
if field not in config_data and env_values.get(field):
config_data[field] = env_values[field]
config_data = {k: v for k, v in config_data.items() if v != ""} # mutable-ok: dict # rebind-ok: "" means clear
has_api_base: Final = bool(config_data.get("cyberark_api_base"))
has_api_key_auth: Final = bool(config_data.get("cyberark_api_key"))
has_tls_cert_auth: Final = bool(config_data.get("client_cert") and config_data.get("client_key"))
if not has_api_base:
raise HTTPException(
status_code=400,
detail="CyberArk API Base is required",
)
if not has_api_key_auth and not has_tls_cert_auth:
raise HTTPException(
status_code=400,
detail="At least one authentication method is required: "
"provide an API Key, or both Client Certificate and Client Key",
)
_snapshot_cyberark_boot_env(proxy_config)
previous_env: Final = _get_current_env_values(CYBERARK_ENV_VAR_MAPPING)
_set_env_vars(config_data, CYBERARK_ENV_VAR_MAPPING)
try:
proxy_config.initialize_secret_manager(key_management_system="cyberark")
except Exception as e: # noqa: BLE001 # any init failure must roll back env vars
_set_env_vars(previous_env, CYBERARK_ENV_VAR_MAPPING)
verbose_proxy_logger.exception("Error reinitializing CyberArk secret manager: %s", str(e))
raise HTTPException(
status_code=500,
detail=f"Failed to initialize secret manager: {e}",
)
try:
proxy_config._last_cyberark_config = await _persist_cyberark_config( # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage
prisma_client, proxy_config, config_data
)
except Exception as e: # noqa: BLE001 # persistence failure must roll back the runtime state set above
_restore_cyberark_runtime(proxy_config, previous_env)
verbose_proxy_logger.exception("Error persisting CyberArk configuration: %s", str(e))
raise HTTPException(
status_code=500,
detail=f"Failed to persist CyberArk configuration: {e}",
)
before_config: Final = existing_decrypted if existing_decrypted is not None else env_values
action: Final[AUDIT_ACTIONS] = "updated" if existing_record is not None else "created"
await _emit_config_override_audit_log(
object_id="cyberark",
action=action,
before_config=before_config,
after_config=config_data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
return { # mutable-ok: JSON response payload
"message": "CyberArk configuration updated successfully",
"status": "success",
}
@router.get(
"/config_overrides/cyberark",
tags=["Config Overrides"], # mutable-ok: FastAPI route decorator metadata
dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI route decorator metadata
response_model=ConfigOverrideSettingsResponse,
)
async def get_cyberark_config(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection
) -> ConfigOverrideSettingsResponse:
"""
Get current CyberArk Conjur configuration.
Returns decrypted values from DB, or falls back to current env vars.
Sensitive fields are masked before leaving the server.
"""
from litellm.proxy.management_endpoints.common_utils import (
_user_has_admin_view, # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage
)
from litellm.proxy.proxy_server import prisma_client, proxy_config
if not _user_has_admin_view(user_api_key_dict):
raise HTTPException(
status_code=403,
detail="Only admin users can view config overrides",
)
if prisma_client is None:
raise HTTPException(
status_code=500,
detail=CommonProxyErrors.db_not_connected_error.value,
)
field_schema: Final = _build_field_schema(CyberArkConfig)
db_record: Final = await _config_overrides_table(prisma_client).find_unique(
where={"config_type": "cyberark"}
) # mutable-ok: prisma where clause
if db_record is not None and db_record.config_value is not None:
config_data: Final = _parse_config_value(db_record.config_value)
decrypted_data: Final[Mapping[str, object]] = proxy_config._decrypt_db_variables(config_data) # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage
masked_data: Final = _mask_sensitive_fields(decrypted_data, CYBERARK_SENSITIVE_FIELDS)
return ConfigOverrideSettingsResponse(
config_type="cyberark",
values=masked_data,
field_schema=field_schema,
)
env_values: Final = _get_current_env_values(CYBERARK_ENV_VAR_MAPPING)
masked_env_values: Final = _mask_sensitive_fields(env_values, CYBERARK_SENSITIVE_FIELDS)
return ConfigOverrideSettingsResponse(
config_type="cyberark",
values=masked_env_values,
field_schema=field_schema,
)
@router.delete(
"/config_overrides/cyberark",
tags=["Config Overrides"], # mutable-ok: FastAPI route decorator metadata
dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI route decorator metadata
)
async def delete_cyberark_config(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection
litellm_changed_by: str | None = Header(
None,
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
),
) -> dict[str, str]:
"""Delete CyberArk Conjur configuration. Idempotent."""
from litellm.proxy.proxy_server import prisma_client, proxy_config
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Only admin users can delete config overrides",
)
if prisma_client is None:
raise HTTPException(
status_code=500,
detail=CommonProxyErrors.db_not_connected_error.value,
)
existing_record: Final = await _config_overrides_table(prisma_client).find_unique(
where={"config_type": "cyberark"} # mutable-ok: prisma where clause
)
before_config: dict[str, object] | None = None # mutable-ok: audit snapshot # rebind-ok: set when decrypts
if existing_record is not None and existing_record.config_value is not None:
try:
before_config = proxy_config._decrypt_db_variables(_parse_config_value(existing_record.config_value)) # pyright: ignore[reportPrivateUsage] # rebind-ok: populated when the prior record decrypts
except Exception: # noqa: BLE001 # undecryptable prior config must not block deletion
before_config = None # rebind-ok: reset when decryption fails
deleted = False # rebind-ok: set true once the DB row is removed
try:
await _config_overrides_table(prisma_client).delete(
where={"config_type": "cyberark"}
) # mutable-ok: prisma where clause
deleted = True # rebind-ok: set true once the DB row is removed
except RecordNotFoundError:
verbose_proxy_logger.debug("No existing CyberArk config record to delete")
_clear_cyberark_state(proxy_config)
if deleted:
await _emit_config_override_audit_log(
object_id="cyberark",
action="deleted",
before_config=before_config,
after_config=None,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
return { # mutable-ok: JSON response payload
"message": "CyberArk configuration deleted successfully",
"status": "success",
}
@router.post(
"/config_overrides/cyberark/test_connection",
tags=["Config Overrides"], # mutable-ok: FastAPI route decorator metadata
dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI route decorator metadata
)
async def test_cyberark_connection(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection
) -> dict[str, str]:
"""
Test the connection to the currently configured CyberArk Conjur server.
Uses the already-initialized secret manager client. Does not modify any state.
"""
from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Only admin users can test CyberArk connection",
)
client: Final = litellm.secret_manager_client
if not isinstance(client, CyberArkSecretManager):
raise HTTPException(
status_code=400,
detail="CyberArk is not configured. Save a configuration first.",
)
try:
headers: Final[Mapping[str, str]] = await asyncio.to_thread(client._get_request_headers) # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage
except Exception as e: # noqa: BLE001 # surface any auth failure as a 502 with detail
raise HTTPException(
status_code=502,
detail=f"CyberArk authentication failed: {e}",
)
try:
async_client: Final = get_async_httpx_client(
llm_provider=httpxSpecialProvider.SecretManager,
params={"ssl_verify": client.ssl_verify}, # mutable-ok: httpx client params
)
whoami_url: Final = f"{client.conjur_addr}/whoami"
response: Final = await async_client.get(whoami_url, headers=headers)
response.raise_for_status()
except Exception as e: # noqa: BLE001 # surface any connectivity/TLS failure as a 502 with detail
raise HTTPException(
status_code=502,
detail=f"CyberArk token validation failed: {e}",
)
return { # mutable-ok: JSON response payload
"status": "success",
"message": f"Successfully connected to CyberArk Conjur at {client.conjur_addr}",
}

3
litellm/proxy/pass_through_endpoints/managed_id_rewriter.py
View file

@ -49,6 +49,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.managed_resources.isolation import (
build_owner_filter,
can_access_resource,
resolve_resource_owner_id,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
@ -686,7 +687,7 @@ async def _mint_or_reuse_object(
"file_object": json.dumps(body_snapshot),
"model_object_id": namespaced_model_object_id,
"file_purpose": file_purpose,
"created_by": user_api_key_dict.user_id,
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
"updated_by": user_api_key_dict.user_id,
},

73
litellm/proxy/proxy_server.py
View file

@ -4281,6 +4281,8 @@ class ProxyConfig:
self.config: dict[str, Any] = {}
self._last_semantic_filter_config: dict[str, object] | None = None
self._last_hashicorp_vault_config: dict[str, object] | None = None
self._last_cyberark_config: dict[str, object] | None = None # mutable-ok: change-detection cache
self._cyberark_boot_env: dict[str, str | None] | None = None # mutable-ok: deployment env snapshot, set once
self.worker_registry: list[WorkerRegistryEntry] = []
self.config_sync_subscriber: ConfigSyncSubscriber | None = None
self.auth_cache_invalidation_subscriber: AuthCacheInvalidationSubscriber | None = None
@ -6764,9 +6766,18 @@ class ProxyConfig:
- list: the rows (may be empty if no models exist)
- None: signals a DB fetch *failure* — callers must not treat this
as "all models deleted" and must not evict existing router deployments.
Pinned to the writer DB: this read reconciles the router against the rows a
model write just committed, and reading it through a lagging read replica
makes the write-triggered reload report its own durable write as missing
(#38556). It also keeps a stale replica snapshot from evicting a deployment
another pod just added. While the writer is degraded the pin yields to the
replica so reader-only mode keeps loading DB-backed models.
"""
try:
new_models: Final[Sequence[_ProxyModelRow]] = await ModelRepository(prisma_client).table.find_many()
new_models: Final[Sequence[_ProxyModelRow]] = await ModelRepository(
WriterPinnedClient(prisma_client.db)
).table.find_many()
return new_models
except Exception as e:
verbose_proxy_logger.exception(
@ -6968,6 +6979,7 @@ class ProxyConfig:
if self._should_load_db_object(object_type="config_overrides"):
await self._init_hashicorp_vault_config_override(prisma_client=prisma_client)
await self._init_cyberark_config_override(prisma_client=prisma_client)
await self._apply_safe_litellm_settings_overrides_from_db(prisma_client=prisma_client)
@ -7132,6 +7144,64 @@ class ProxyConfig:
str(e),
)
async def _init_cyberark_config_override(self, prisma_client: PrismaClient) -> None:
"""
Load CyberArk Conjur config override from DB.
Decrypts sensitive fields, sets CYBERARK_* env vars, and reinitializes the secret manager.
Called periodically via _init_non_llm_objects_in_db to sync config across pods.
"""
from litellm.proxy.management_endpoints.config_override_endpoints import (
CYBERARK_ENV_VAR_MAPPING,
_clear_cyberark_state, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module
_get_current_env_values, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module
_parse_config_value, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module
_set_env_vars, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module
_snapshot_cyberark_boot_env, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module
)
try:
db_record: Final[_ConfigOverridesRow | None] = cast( # cast-ok: prisma Json stub is `str`, runtime dict
"_ConfigOverridesRow | None",
await call_with_db_reconnect_retry(
prisma_client,
lambda: ConfigOverridesRepository(prisma_client).table.find_unique(
where={"config_type": "cyberark"} # mutable-ok: prisma where clause
),
reason="init_cyberark_config_override_lookup_failure",
),
)
if db_record is None or db_record.config_value is None:
if self._last_cyberark_config is not None:
_clear_cyberark_state(self)
return
config_data: Final = _parse_config_value(db_record.config_value)
# Skip reinit if config hasn't changed since last poll
if self._last_cyberark_config == config_data:
return
decrypted_data: Final = self._decrypt_db_variables(config_data)
_snapshot_cyberark_boot_env(self)
previous_env: Final = _get_current_env_values(CYBERARK_ENV_VAR_MAPPING)
_set_env_vars(decrypted_data, CYBERARK_ENV_VAR_MAPPING)
try:
self.initialize_secret_manager(key_management_system="cyberark")
except Exception:
_set_env_vars(previous_env, CYBERARK_ENV_VAR_MAPPING)
raise
self._last_cyberark_config = config_data.copy()
verbose_proxy_logger.debug("CyberArk config override loaded from DB")
except Exception as e: # noqa: BLE001 # any DB/decrypt/init failure must not break proxy boot
verbose_proxy_logger.exception(
"Error loading CyberArk config override from DB: %s",
str(e),
)
async def check_periodic_reloads(self, prisma_client: PrismaClient):
"""
Run the admin-configured periodic model cost map reload.
@ -12013,6 +12083,7 @@ async def run_thread(
# )
# async def get_available_routes(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)):
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient
from litellm.repositories.config_repository import ConfigRepository
from litellm.repositories.model_repository import ModelRepository
from litellm.repositories.table_repositories import (

16
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -93,6 +93,8 @@ def _get_spend_logs_metadata(
metadata: dict | None,
applied_guardrails: list[str] | None = None,
batch_models: list[str] | None = None,
batch_successful_requests: int | None = None,
batch_failed_requests: int | None = None,
mcp_tool_call_metadata: StandardLoggingMCPToolCall | None = None,
vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None = None,
guardrail_information: list[StandardLoggingGuardrailInformation] | None = None,
@ -122,6 +124,8 @@ def _get_spend_logs_metadata(
error_information=None,
proxy_server_request=None,
batch_models=None,
batch_successful_requests=None,
batch_failed_requests=None,
mcp_tool_call_metadata=None,
vector_store_request_metadata=None,
model_map_information=None,
@ -155,6 +159,8 @@ def _get_spend_logs_metadata(
clean_metadata["user_api_key"] = _redact_logged_api_key(_raw_key, already_redacted=_already_redacted)
clean_metadata["applied_guardrails"] = applied_guardrails
clean_metadata["batch_models"] = batch_models
clean_metadata["batch_successful_requests"] = batch_successful_requests
clean_metadata["batch_failed_requests"] = batch_failed_requests
clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata
clean_metadata["vector_store_request_metadata"] = _get_vector_store_request_for_spend_logs_payload(
vector_store_request_metadata
@ -380,6 +386,16 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
if standard_logging_payload is not None
else None
),
batch_successful_requests=(
standard_logging_payload.get("hidden_params", {}).get("batch_successful_requests", None)
if standard_logging_payload is not None
else None
),
batch_failed_requests=(
standard_logging_payload.get("hidden_params", {}).get("batch_failed_requests", None)
if standard_logging_payload is not None
else None
),
mcp_tool_call_metadata=(
standard_logging_payload["metadata"].get("mcp_tool_call_metadata", None)
if standard_logging_payload is not None

8
litellm/types/integrations/custom_logger.py
View file

@ -5,8 +5,14 @@ from pydantic import BaseModel, Field
CHAT_COMPLETION_AGENTIC_SURFACE: Final = "chat_completions"
RESPONSES_AGENTIC_SURFACE: Final = "responses"
CODE_INTERPRETER_INTERCEPTION_PREFIX: Final = "_code_interpreter_interception"
HEADROOM_INTERCEPTION_PREFIX: Final = "_headroom_interception"
HEADROOM_CONVERTED_STREAM_KEY: Final = f"{HEADROOM_INTERCEPTION_PREFIX}_converted_stream"
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES: Final = frozenset(
("_websearch_interception", "_compression_interception")
(
"_websearch_interception",
"_compression_interception",
HEADROOM_INTERCEPTION_PREFIX,
)
)
INTERCEPTION_INTERNAL_PREFIXES: Final = frozenset(
(

37
litellm/types/proxy/management_endpoints/config_overrides.py
View file

@ -52,6 +52,43 @@ class HashicorpVaultConfig(BaseModel):
)
class CyberArkConfig(BaseModel):
"""Configuration for CyberArk Conjur secret manager integration."""
cyberark_api_base: str | None = Field(
default=None,
description="The address of the CyberArk Conjur server (e.g., https://conjur.example.com)",
)
cyberark_account: str | None = Field(
default=None,
description="The Conjur organization account name",
)
cyberark_username: str | None = Field(
default=None,
description="The Conjur username (login) to authenticate as",
)
cyberark_api_key: str | None = Field(
default=None,
description="API key for Conjur API-key authentication",
)
client_cert: str | None = Field(
default=None,
description="Path to the client TLS certificate for certificate-based authentication",
)
client_key: str | None = Field(
default=None,
description="Path to the client TLS private key for certificate-based authentication",
)
ssl_verify: str | None = Field(
default=None,
description="Set to false to disable SSL verification (e.g., for self-signed certificates)",
)
refresh_interval: str | None = Field(
default=None,
description="Auth token cache TTL in seconds (default: 300)",
)
class ConfigOverrideSettingsResponse(BaseModel):
"""Response model for config override settings GET endpoints."""

3
litellm/types/utils.py
View file

@ -2957,6 +2957,8 @@ class StandardLoggingHiddenParams(TypedDict):
litellm_overhead_time_ms: float | None
additional_headers: StandardLoggingAdditionalHeaders | None
batch_models: list[str] | None
batch_successful_requests: ReadOnly[int | None]
batch_failed_requests: ReadOnly[int | None]
litellm_model_name: str | None # the model name sent to the provider by litellm
usage_object: dict | None
@ -3503,6 +3505,7 @@ agentic_loop_internal_litellm_params: Final = [
"_code_interpreter_interception_converted_stream",
"_websearch_interception_emit_native_blocks",
"_websearch_interception_converted_stream",
"_headroom_interception_converted_stream",
]
# Proxy-owned callback credentials, stamped from admin-configured team/key callback

3
litellm/types/videos/main.py
View file

@ -2,7 +2,7 @@ from typing import Any, Literal
from openai.types.audio.transcription_create_params import FileTypes
from pydantic import BaseModel
from typing_extensions import TypedDict
from typing_extensions import ReadOnly, TypedDict
class VideoObject(BaseModel):
@ -76,6 +76,7 @@ class VideoCreateOptionalRequestParams(TypedDict, total=False):
image: Any | None # Image for image-to-video; dict with gcsUri/bytesBase64Encoded, or file-like object
parameters: dict[str, Any] | None # Provider-specific parameters block passed directly to the API
model: str | None
resolution: ReadOnly[str | None]
seconds: str | None
size: str | None
characters: list[dict[str, str]] | None

7
litellm/utils.py
View file

@ -69,6 +69,7 @@ from litellm.constants import (
DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
DEFAULT_TRIM_RATIO,
FUNCTION_DEFINITION_TOKEN_COUNT,
HF_CONFIG_FETCH_TIMEOUT_SECONDS,
INITIAL_RETRY_DELAY,
JITTER,
MAX_RETRY_DELAY,
@ -3581,7 +3582,7 @@ def get_optional_params_embeddings(
object = litellm.AmazonTitanMultimodalEmbeddingG1Config()
elif "amazon.titan-embed-text-v2:0" in model:
object = litellm.AmazonTitanV2Config()
elif "cohere.embed-multilingual-v3" in model or "cohere.embed-v4" in model:
elif "cohere.embed" in model:
object = litellm.BedrockCohereEmbeddingConfig()
elif "twelvelabs" in model or "marengo" in model:
object = litellm.TwelveLabsMarengoEmbeddingConfig()
@ -5168,7 +5169,7 @@ def get_max_tokens(model: str) -> int | None:
config_url: Final = f"https://huggingface.co/{model_name}/raw/main/config.json"
try:
# Make the HTTP request to get the raw JSON file
response: Final = litellm.module_level_client.get(config_url)
response: Final = litellm.module_level_client.get(config_url, timeout=HF_CONFIG_FETCH_TIMEOUT_SECONDS)
response.raise_for_status() # Raise an exception for bad responses (4xx or 5xx)
# Parse the JSON response
@ -5522,7 +5523,7 @@ def _get_max_position_embeddings(model_name: str) -> int | None:
try:
# Make the HTTP request to get the raw JSON file
response: Final = litellm.module_level_client.get(config_url)
response: Final = litellm.module_level_client.get(config_url, timeout=HF_CONFIG_FETCH_TIMEOUT_SECONDS)
response.raise_for_status() # Raise an exception for bad responses (4xx or 5xx)
# Parse the JSON response

16
model_prices_and_context_window.json
View file

@ -43318,6 +43318,22 @@
"video"
]
},
"vertex_ai/veo-3.1-lite-generate-001": {
"litellm_provider": "vertex_ai-video-models",
"max_input_tokens": 1024,
"max_tokens": 1024,
"mode": "video_generation",
"output_cost_per_second": 0.05,
"output_cost_per_second_1080p": 0.08,
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#veo",
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"video"
]
},
"voyage/rerank-2": {
"input_cost_per_token": 5e-08,
"litellm_provider": "voyage",

4
ruff-strict-budget.json
View file

@ -24,7 +24,7 @@
"limit": 133
},
"ANN401": {
"limit": 517
"limit": 516
},
"ASYNC230": {
"limit": 11
@ -231,7 +231,7 @@
"limit": 5
},
"TID251": {
"limit": 1105
"limit": 1104
},
"TRY002": {
"limit": 524

3
ruff.toml
View file

@ -6,7 +6,8 @@ lint.extend-select = [
"T20", "PGH004", "RUF008", "RUF009", "RUF100",
"B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790", "PIE800", "PLC0208",
"PLR0402", "PLR1711", "PLR1730", "PLR2044", "PLW0133", "PYI030", "PYI041", "PYI064", "RET501",
"RUF010", "RUF022", "RUF023", "RUF051", "SIM114", "SIM118", "TC005", "UP006", "UP007", "UP008",
"RUF010", "RUF022", "RUF023", "RUF051", "S113", "SIM114", "SIM118", "TC005", "UP006", "UP007",
"UP008",
"UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP037", "UP045",
]
# RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip

16
tests/batches_tests/test_batch_custom_pricing.py
View file

@ -116,16 +116,16 @@ def test_aggregate_batch_cost_uses_custom_model_info():
"""_aggregate_batch_cost_usage_models should thread model_info to batch_cost_calculator."""
file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)]
cost, _, _ = _aggregate_batch_cost_usage_models(
result = _aggregate_batch_cost_usage_models(
entries=file_content,
custom_llm_provider="openai",
model_info=CUSTOM_MODEL_INFO,
)
expected = (10 * 0.00125) + (5 * 0.005)
assert cost == pytest.approx(
assert result.cost == pytest.approx(
expected
), f"Expected total cost {expected}, got {cost}"
), f"Expected total cost {expected}, got {result.cost}"
@pytest.mark.parametrize("data_residency", ["eu", "us"])
@ -164,15 +164,15 @@ async def test_calculate_batch_cost_and_usage_uses_custom_model_info():
"""calculate_batch_cost_and_usage should thread model_info."""
file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)]
batch_cost, batch_usage, batch_models = await calculate_batch_cost_and_usage(
result = await calculate_batch_cost_and_usage(
file_content_dictionary=file_content,
custom_llm_provider="openai",
model_info=CUSTOM_MODEL_INFO,
)
expected = (10 * 0.00125) + (5 * 0.005)
assert batch_cost == pytest.approx(
assert result.cost == pytest.approx(
expected
), f"Expected total cost {expected}, got {batch_cost}"
assert batch_usage.prompt_tokens == 10
assert batch_usage.completion_tokens == 5
), f"Expected total cost {expected}, got {result.cost}"
assert result.usage.prompt_tokens == 10
assert result.usage.completion_tokens == 5

14
tests/batches_tests/test_batch_rate_limits.py
View file

@ -1027,7 +1027,7 @@ async def test_batch_logging_azure_credentials_regression():
with patch(
"litellm.files.main.afile_content", side_effect=mock_afile_content_tracker
):
cost, usage, models = await _handle_completed_batch(
result = await _handle_completed_batch(
batch=mock_batch,
custom_llm_provider="azure",
litellm_params=azure_credentials,
@ -1039,13 +1039,13 @@ async def test_batch_logging_azure_credentials_regression():
], "REGRESSION: Credentials not passed through _handle_completed_batch"
# Verify cost and usage were calculated
assert cost > 0, "Cost should be calculated"
assert usage.total_tokens == 40, "Usage should be calculated correctly"
assert result.cost > 0, "Cost should be calculated"
assert result.usage.total_tokens == 40, "Usage should be calculated correctly"
print(" ✓ Credentials passed through full flow")
print(f" ✓ Cost: {cost}")
print(f" ✓ Usage: {usage.total_tokens} tokens")
print(f" ✓ Models: {models}")
print(f" ✓ Cost: {result.cost}")
print(f" ✓ Usage: {result.usage.total_tokens} tokens")
print(f" ✓ Models: {result.models}")
# Test 4: Verify error prevention
print("\n4. Testing 'Missing credentials' error prevention...")
@ -1064,7 +1064,7 @@ async def test_batch_logging_azure_credentials_regression():
"litellm.files.main.afile_content", side_effect=mock_afile_content_tracker
):
try:
cost, usage, models = await _handle_completed_batch(
result = await _handle_completed_batch(
batch=mock_batch,
custom_llm_provider="azure",
litellm_params=azure_credentials,

58
tests/batches_tests/test_batches_logging_unit_tests.py
View file

@ -133,12 +133,12 @@ def test_get_file_content_as_dictionary(sample_file_content):
def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict):
with patch("litellm.completion_cost", return_value=0.0):
_, usage, _ = _aggregate_batch_cost_usage_models(
result = _aggregate_batch_cost_usage_models(
entries=sample_file_content_dict, custom_llm_provider="openai"
)
assert usage.total_tokens == 62 # 30 + 32
assert usage.prompt_tokens == 42 # 20 + 22
assert usage.completion_tokens == 20 # 10 + 10
assert result.usage.total_tokens == 62 # 30 + 32
assert result.usage.prompt_tokens == 42 # 20 + 22
assert result.usage.completion_tokens == 20 # 10 + 10
@pytest.mark.asyncio
@ -151,11 +151,11 @@ async def test_batch_cost_calculator(sample_file_content_dict):
so we expect the cost to be 0.5 * 2 = 1.0
"""
with patch("litellm.completion_cost", return_value=0.5):
cost, _, _ = _aggregate_batch_cost_usage_models(
result = _aggregate_batch_cost_usage_models(
entries=sample_file_content_dict,
custom_llm_provider="openai",
)
assert cost == 1.0 # 0.5 * 2 successful responses
assert result.cost == 1.0 # 0.5 * 2 successful responses
def test_get_response_from_batch_job_output_file(sample_file_content_dict):
@ -221,6 +221,8 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos
logging_obj.custom_llm_provider = "openai"
# Mock _handle_completed_batch to return cost data
from litellm.batches.batch_utils import BatchCostUsageResult
expected_cost = 0.05
expected_usage = litellm.Usage(
prompt_tokens=100,
@ -231,7 +233,15 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos
with patch(
"litellm.litellm_core_utils.litellm_logging._handle_completed_batch",
new=AsyncMock(return_value=(expected_cost, expected_usage, expected_models)),
new=AsyncMock(
return_value=BatchCostUsageResult(
cost=expected_cost,
usage=expected_usage,
models=expected_models,
successful_requests=10,
failed_requests=0,
)
),
) as mock_handle_batch:
# Call async_success_handler
await logging_obj.async_success_handler(
@ -246,6 +256,8 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos
# Verify cost and usage were set on the batch result
assert mock_batch._hidden_params["response_cost"] == expected_cost
assert mock_batch._hidden_params["batch_models"] == expected_models
assert mock_batch._hidden_params["batch_successful_requests"] == 10
assert mock_batch._hidden_params["batch_failed_requests"] == 0
assert mock_batch.usage == expected_usage
@ -279,7 +291,7 @@ async def test_handle_completed_batch_computes_real_cost_from_output_file(
"litellm.batches.batch_utils._fetch_batch_output_file_content",
new=AsyncMock(return_value=sample_file_content_bytes),
):
cost, usage, models = await _handle_completed_batch(
result = await _handle_completed_batch(
batch=batch, custom_llm_provider="openai"
)
@ -289,16 +301,18 @@ async def test_handle_completed_batch_computes_real_cost_from_output_file(
+ 20 * pricing["output_cost_per_token_batches"]
)
assert cost == pytest.approx(expected_cost)
assert cost > 0
assert result.cost == pytest.approx(expected_cost)
assert result.cost > 0
assert (
cost
result.cost
< 42 * pricing["input_cost_per_token"] + 20 * pricing["output_cost_per_token"]
)
assert usage.prompt_tokens == 42
assert usage.completion_tokens == 20
assert usage.total_tokens == 62
assert models == ["gpt-4o-mini-2024-07-18", "gpt-4o-mini-2024-07-18"]
assert result.usage.prompt_tokens == 42
assert result.usage.completion_tokens == 20
assert result.usage.total_tokens == 62
assert result.models == ["gpt-4o-mini-2024-07-18", "gpt-4o-mini-2024-07-18"]
assert result.successful_requests == 2
assert result.failed_requests == 0
@pytest.mark.asyncio
@ -537,9 +551,19 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data():
)
expected_models = ["gpt-5-mini"]
from litellm.batches.batch_utils import BatchCostUsageResult
with patch(
"litellm.litellm_core_utils.litellm_logging._handle_completed_batch",
new=AsyncMock(return_value=(expected_cost, expected_usage, expected_models)),
new=AsyncMock(
return_value=BatchCostUsageResult(
cost=expected_cost,
usage=expected_usage,
models=expected_models,
successful_requests=8,
failed_requests=0,
)
),
) as mock_handle_batch:
# Call async_success_handler with partial explicit data
await logging_obj.async_success_handler(
@ -555,4 +579,6 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data():
# Verify computed cost data was used (not partial explicit data)
assert mock_batch._hidden_params["response_cost"] == expected_cost
assert mock_batch._hidden_params["batch_models"] == expected_models
assert mock_batch._hidden_params["batch_successful_requests"] == 8
assert mock_batch._hidden_params["batch_failed_requests"] == 0
assert mock_batch.usage == expected_usage

26
tests/e2e/ui/tests/logs/logs.spec.ts
View file

@ -148,6 +148,32 @@ test.describe("Logs page", () => {
});
});
test("the trace sidebar collapses and expands again", async ({ page, request }) => {
const prompt = `logs-sidebar-prompt-${uniqueSuffix()}`;
const requestId = await sendChatCompletion(request, {
model: CHAT_MODEL_A,
prompt,
});
await waitForSpendLog(request, requestId);
const row = await openLogsForRequest(page, requestId);
await row.click();
const drawer = page.getByRole("dialog").first();
await expect(drawer.getByText("Request & Response")).toBeVisible({ timeout: 20_000 });
const toggle = drawer.getByLabel("Collapse trace sidebar");
await expect(toggle).toBeVisible({ timeout: 10_000 });
await toggle.click();
const expandToggle = drawer.getByLabel("Expand trace sidebar");
await expect(expandToggle).toBeVisible({ timeout: 10_000 });
await expandToggle.click({ timeout: 10_000 });
await expect(drawer.getByLabel("Collapse trace sidebar")).toBeVisible({ timeout: 10_000 });
await expect(drawer.getByText(prompt, { exact: false })).toBeVisible({ timeout: 10_000 });
});
test("the JSON view exposes Request and Response tabs", async ({ page, request }) => {
const prompt = `logs-json-prompt-${uniqueSuffix()}`;
const requestId = await sendChatCompletion(request, {

2
tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json
View file

@ -11,7 +11,7 @@
"user": "",
"team_id": "",
"organization_id": "",
"metadata": "{\"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}",
"metadata": "{\"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}",
"cache_key": "Cache OFF",
"spend": 0.00022500000000000002,
"total_tokens": 30,

156
tests/proxy_unit_tests/test_check_batch_cost.py
View file

@ -9,16 +9,40 @@ ARN unified_object_id) batches with no managed unified id.
import asyncio
import json
from contextlib import contextmanager
from typing import TYPE_CHECKING
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
if TYPE_CHECKING:
from litellm.batches.batch_utils import BatchCostUsageResult
_IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id"
_CLAIM_UNIFIED_BATCH_ID = "dW5pZmllZF9iYXRjaF9pZA=="
_CLAIM_OUTPUT_FILE_ID = "file-output-123"
def _batch_cost_result(
cost: float,
usage: dict,
models: list[str],
successful_requests: int = 1,
failed_requests: int = 0,
) -> "BatchCostUsageResult":
"""Build the BatchCostUsageResult calculate_batch_cost_and_usage now returns,
for mocking it in tests that only care about cost/usage/models."""
from litellm.batches.batch_utils import BatchCostUsageResult
return BatchCostUsageResult(
cost=cost,
usage=usage,
models=models,
successful_requests=successful_requests,
failed_requests=failed_requests,
)
def _unmanaged_vertex_file_object(
input_file_id="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash/abc.jsonl",
status="validating",
@ -327,7 +351,7 @@ class TestCheckBatchCost:
patch(
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
new_callable=AsyncMock,
return_value=(
return_value=_batch_cost_result(
0.01,
{"prompt_tokens": 10, "completion_tokens": 5},
["gpt-4"],
@ -432,7 +456,7 @@ class TestCheckBatchCost:
patch(
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
new_callable=AsyncMock,
return_value=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]),
return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]),
),
patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
@ -535,7 +559,9 @@ class TestCheckBatchCost:
patch(
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
new_callable=AsyncMock,
return_value=(0.0052, {"prompt_tokens": 1400, "completion_tokens": 600}, ["claude-haiku-4-5"]),
return_value=_batch_cost_result(
0.0052, {"prompt_tokens": 1400, "completion_tokens": 600}, ["claude-haiku-4-5"]
),
) as mock_calculate,
patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
@ -634,7 +660,7 @@ class TestCheckBatchCost:
patch(
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
new_callable=AsyncMock,
return_value=(
return_value=_batch_cost_result(
0.01,
{"prompt_tokens": 10, "completion_tokens": 5},
["gpt-4"],
@ -764,7 +790,7 @@ class TestCheckBatchCost:
patch(
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
new_callable=AsyncMock,
return_value=(
return_value=_batch_cost_result(
0.01,
{"prompt_tokens": 10, "completion_tokens": 5},
["gpt-4"],
@ -1312,7 +1338,7 @@ class TestCheckBatchCost:
patch(
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
new_callable=AsyncMock,
return_value=(
return_value=_batch_cost_result(
0.01,
{"prompt_tokens": 10, "completion_tokens": 5},
["gpt-4"],
@ -1347,6 +1373,114 @@ class TestCheckBatchCost:
update_data["status"] == terminal_status
), f"billed {terminal_status} batch must keep its real terminal status in the DB"
@pytest.mark.asyncio
async def test_error_file_failures_add_to_failed_request_count(
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
):
"""OpenAI-shaped providers report per-request failures only in a separate
error file. The poller prices from the output file, so without also counting
the error file's lines, batch_failed_requests on the spend log undercounts:
regression test for the poller path merging error-file failures.
"""
import base64
from unittest.mock import patch
import httpx
import respx
from litellm.litellm_core_utils.litellm_logging import Logging
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_job = MagicMock()
mock_job.id = "job-error-file-1"
mock_job.unified_object_id = base64.urlsafe_b64encode(
b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456"
).decode()
mock_job.created_by = "user-1"
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job])
mock_response = MagicMock()
mock_response.status = "completed"
mock_response.output_file_id = "file-output-123"
mock_response.error_file_id = "file-error-456"
mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}'
mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"})
mock_deployment = MagicMock()
mock_deployment.litellm_params.custom_llm_provider = "openai"
mock_deployment.litellm_params.model = "gpt-4"
mock_deployment.model_info.model_dump.return_value = {}
mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment)
succeeded_line = json.dumps(
{
"custom_id": "req-1",
"response": {
"status_code": 200,
"body": {
"id": "chatcmpl-1",
"object": "chat.completion",
"model": "gpt-4",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 5,
"total_tokens": 15,
},
},
},
"error": None,
}
)
rejected_line = json.dumps(
{
"custom_id": "req-2",
"response": {
"status_code": 400,
"body": {"error": {"message": "bad request"}},
},
"error": None,
}
)
error_file_lines = "\n".join(
json.dumps({"custom_id": custom_id, "error": {"message": "rejected"}}) for custom_id in ("req-3", "req-4")
)
with (
respx.mock(assert_all_called=True) as provider,
patch.object( # test-quality-ok: the poller builds Logging inline, the only seam to its handler kwargs
Logging, "async_success_handler", new_callable=AsyncMock
) as success_handler,
):
provider.get("https://api.openai.com/v1/files/file-output-123/content").mock(
return_value=httpx.Response(200, content=f"{succeeded_line}\n{rejected_line}\n".encode())
)
provider.get("https://api.openai.com/v1/files/file-error-456/content").mock(
return_value=httpx.Response(200, content=f"{error_file_lines}\n\n".encode())
)
await check_batch_cost_instance.check_batch_cost()
spend_log_calls = [call.kwargs for call in success_handler.await_args_list if "batch_cost" in call.kwargs]
assert len(spend_log_calls) == 1
handler_kwargs = spend_log_calls[0]
assert handler_kwargs["batch_successful_requests"] == 1
assert handler_kwargs["batch_failed_requests"] == 3, (
"2 error-file lines must add to the output file's 1 rejected request"
)
assert handler_kwargs["batch_models"] == ["gpt-4"]
assert handler_kwargs["batch_usage"].total_tokens == 15
assert handler_kwargs["batch_cost"] > 0
@pytest.mark.asyncio
async def test_terminal_batch_with_missing_output_file_is_retired_unbilled(
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
@ -1518,7 +1652,7 @@ class TestCheckBatchCost:
patch(
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
new_callable=AsyncMock,
return_value=(
return_value=_batch_cost_result(
0.01,
{"prompt_tokens": 10, "completion_tokens": 5},
["gpt-4"],
@ -1776,7 +1910,7 @@ class TestUnmanagedVertexRouting:
patch(
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
new_callable=AsyncMock,
return_value=(
return_value=_batch_cost_result(
0.01,
{"prompt_tokens": 10, "completion_tokens": 5},
["gemini-2.5-flash"],
@ -2006,7 +2140,7 @@ class TestUnmanagedBedrockRouting:
patch(
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
new_callable=AsyncMock,
return_value=(
return_value=_batch_cost_result(
0.02,
{"prompt_tokens": 10, "completion_tokens": 5},
["claude-sonnet-4"],
@ -2198,7 +2332,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
patch(
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
new_callable=AsyncMock,
return_value=(0.01, {"prompt_tokens": 10}, ["gpt-5.5"]),
return_value=_batch_cost_result(0.01, {"prompt_tokens": 10}, ["gpt-5.5"]),
),
patch("litellm.litellm_core_utils.litellm_logging.Logging") as logging_cls,
):
@ -2826,7 +2960,7 @@ class TestMultiPodBatchCostClaim:
patch(
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
new_callable=AsyncMock,
return_value=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"]),
return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"]),
),
patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",

415
tests/test_litellm/batches/test_batch_utils.py
View file

@ -211,10 +211,10 @@ def test_estimate_tokens_never_zero_for_short_rows():
def test_output_models_uses_model_name_override(monkeypatch):
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
_, _, models = bu._aggregate_batch_cost_usage_models(
result = bu._aggregate_batch_cost_usage_models(
entries=[_success_row(model="ignored")], custom_llm_provider="openai", model_name="forced-model"
)
assert models == ["forced-model"]
assert result.models == ["forced-model"]
def test_output_models_collects_from_successful_only(monkeypatch):
@ -224,15 +224,15 @@ def test_output_models_collects_from_successful_only(monkeypatch):
_failed_row(model="should-be-skipped"),
_success_row(model="claude-3"),
]
_, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert models == ["gpt-4o", "claude-3"]
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert result.models == ["gpt-4o", "claude-3"]
def test_output_models_skips_successful_without_model(monkeypatch):
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
rows = [{"response": {"status_code": 200, "body": {}}}]
_, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert models == []
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert result.models == []
# =========================================================================== #
@ -399,8 +399,8 @@ def test_total_usage_sums_successful_only(monkeypatch):
_failed_row(), # excluded
_success_row(usage=_usage(20, 10)), # 30
]
_, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
30,
15,
45,
@ -418,7 +418,7 @@ def test_total_usage_and_cost_normalize_mixed_responses_and_chat():
)
chat_row = _success_row(usage=_usage(10, 5))
cost, usage, _ = bu._aggregate_batch_cost_usage_models(
result = bu._aggregate_batch_cost_usage_models(
entries=[responses_row, chat_row],
custom_llm_provider="openai",
model_info={
@ -427,22 +427,79 @@ def test_total_usage_and_cost_normalize_mixed_responses_and_chat():
},
)
assert usage.prompt_tokens == 30
assert usage.completion_tokens == 12
assert usage.total_tokens == 42
assert usage.cache_read_input_tokens == 3
assert cost == pytest.approx((30 * 0.00125) + (12 * 0.005))
assert result.usage.prompt_tokens == 30
assert result.usage.completion_tokens == 12
assert result.usage.total_tokens == 42
assert result.usage.cache_read_input_tokens == 3
assert result.cost == pytest.approx((30 * 0.00125) + (12 * 0.005))
def test_total_usage_empty_is_zero():
cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=[], custom_llm_provider="openai")
assert cost == 0.0
assert models == []
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
result = bu._aggregate_batch_cost_usage_models(entries=[], custom_llm_provider="openai")
assert result.cost == 0.0
assert result.models == []
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
0,
0,
0,
)
assert result.successful_requests == 0
assert result.failed_requests == 0
def test_total_usage_includes_reasoning_tokens(monkeypatch):
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
rows = [
_success_row(
usage={
"prompt_tokens": 10,
"completion_tokens": 50,
"total_tokens": 60,
"completion_tokens_details": {"reasoning_tokens": 30},
}
),
_success_row(
usage={
"prompt_tokens": 5,
"completion_tokens": 20,
"total_tokens": 25,
"completion_tokens_details": {"reasoning_tokens": 8},
}
),
_failed_row(), # excluded, must not contribute reasoning tokens either
]
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert result.usage.completion_tokens_details is not None
assert result.usage.completion_tokens_details.reasoning_tokens == 38
def test_aggregate_counts_successful_and_failed_requests(monkeypatch):
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
rows = [
_success_row(usage=_usage(10, 5)),
_failed_row(),
_success_row(usage=_usage(20, 10)),
_failed_row(),
_failed_row(),
]
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert result.successful_requests == 2
assert result.failed_requests == 3
assert result.successful_requests + result.failed_requests == len(rows)
def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch):
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 1.0)
result = bu._aggregate_batch_cost_usage_models(
entries=[_success_row(usage=_usage(10, 5))], custom_llm_provider="openai"
)
assert isinstance(result, bu.BatchCostUsageResult)
assert (result.cost, result.models, result.successful_requests, result.failed_requests) == (
1.0,
["gpt-4o"],
1,
0,
)
# =========================================================================== #
@ -465,15 +522,22 @@ def test_cost_from_content_completion_cost_path(monkeypatch):
_success_row(usage=_usage(20, 10)),
]
total, _, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert total == 1.0 # 2 successful * 0.5
assert result.cost == 1.0 # 2 successful * 0.5
assert len(calls) == 2 # failed row not costed
assert result.successful_requests == 2
assert result.failed_requests == 1
def test_empty_body_line_does_not_zero_whole_batch():
"""A status-200 row with an empty body makes litellm.completion_cost raise;
that line must be skipped instead of zeroing the whole batch."""
that line must be skipped from pricing instead of zeroing the whole batch.
The provider still reported it as a success, so it stays in
successful_requests and out of failed_requests - otherwise the counts stop
reconciling with the provider's own request_counts over a litellm-side
pricing gap the customer never caused."""
rows = [
_success_row(usage=_usage(10, 5)),
{
@ -483,11 +547,12 @@ def test_empty_body_line_does_not_zero_whole_batch():
_success_row(usage=_usage(20, 10)),
]
cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert cost > 0.0
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45)
assert models == ["gpt-4o", "gpt-4o"]
assert result.cost > 0.0
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45)
assert result.models == ["gpt-4o", "gpt-4o"]
assert (result.successful_requests, result.failed_requests) == (3, 0)
def test_cost_from_content_model_info_path(monkeypatch):
@ -500,13 +565,13 @@ def test_cost_from_content_model_info_path(monkeypatch):
_success_row(usage=_usage(20, 10)),
]
total, _, _ = bu._aggregate_batch_cost_usage_models(
result = bu._aggregate_batch_cost_usage_models(
entries=rows,
custom_llm_provider="openai",
model_info={"input_cost_per_token": 0.0}, # type: ignore[arg-type] # truthy -> model_info path
)
assert total == pytest.approx(0.6) # 2 * (0.1 + 0.2)
assert result.cost == pytest.approx(0.6) # 2 * (0.1 + 0.2)
def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch):
@ -516,11 +581,13 @@ def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch):
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.5)
one_shot = (row for row in [_success_row(usage=_usage(10, 5)), _failed_row(), _success_row(usage=_usage(20, 10))])
cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai")
result = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai")
assert cost == 1.0
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45)
assert models == ["gpt-4o", "gpt-4o"]
assert result.cost == 1.0
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45)
assert result.models == ["gpt-4o", "gpt-4o"]
assert result.successful_requests == 2
assert result.failed_requests == 1
# =========================================================================== #
@ -534,7 +601,13 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch):
monkeypatch.setattr(
bu,
"calculate_vertex_ai_batch_cost_and_usage",
lambda content, model: (9.9, Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3)),
lambda content, model: bu.BatchCostUsageResult(
cost=9.9,
usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3),
models=["gemini-2.0-flash-001"],
successful_requests=1,
failed_requests=0,
),
)
# generic path must NOT be taken
monkeypatch.setattr(
@ -543,12 +616,12 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch):
lambda **kw: pytest.fail("generic path should not run"),
)
cost, usage, models = await bu.calculate_batch_cost_and_usage(
result = await bu.calculate_batch_cost_and_usage(
file_content_dictionary=[], custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001"
)
assert cost == 9.9
assert usage.total_tokens == 3
assert models == ["gemini-2.0-flash-001"]
assert result.cost == 9.9
assert result.usage.total_tokens == 3
assert result.models == ["gemini-2.0-flash-001"]
@pytest.mark.asyncio
@ -562,12 +635,12 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch):
lambda content, model: pytest.fail("raw vertex path should not run"),
)
cost, usage, models = await bu.calculate_batch_cost_and_usage(
result = await bu.calculate_batch_cost_and_usage(
file_content_dictionary=[], custom_llm_provider="vertex_ai"
)
assert cost == 0.0
assert usage.total_tokens == 0
assert models == []
assert result.cost == 0.0
assert result.usage.total_tokens == 0
assert result.models == []
# =========================================================================== #
@ -600,14 +673,16 @@ def test_vertex_cost_and_usage_aggregation(monkeypatch):
},
]
cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
assert cost == pytest.approx(0.6) # 2 * (0.1 + 0.2)
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
assert result.cost == pytest.approx(0.6) # 2 * (0.1 + 0.2)
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
30,
15,
45,
)
assert result.successful_requests == 2
assert result.failed_requests == 0
def test_vertex_cost_skips_none_response_body(monkeypatch):
@ -627,10 +702,12 @@ def test_vertex_cost_skips_none_response_body(monkeypatch):
},
]
cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
assert cost == pytest.approx(1.0) # only one line costed
assert usage.total_tokens == 10
assert result.cost == pytest.approx(1.0) # only one line costed
assert result.usage.total_tokens == 10
assert result.successful_requests == 1
assert result.failed_requests == 1
def test_vertex_usage_total_token_fallback(monkeypatch):
@ -640,8 +717,8 @@ def test_vertex_usage_total_token_fallback(monkeypatch):
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0))
responses = [{"response": {"usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 4}}}]
_, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
assert usage.total_tokens == 12
result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
assert result.usage.total_tokens == 12
def test_vertex_cost_error_in_line_is_swallowed(monkeypatch):
@ -664,9 +741,9 @@ def test_vertex_cost_error_in_line_is_swallowed(monkeypatch):
}
]
cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
assert cost == 0.0
assert usage.total_tokens == 10
result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
assert result.cost == 0.0
assert result.usage.total_tokens == 10
# =========================================================================== #
@ -679,13 +756,11 @@ async def test_calculate_batch_cost_and_usage_orchestration(monkeypatch):
rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 2.5)
cost, usage, models = await bu.calculate_batch_cost_and_usage(
file_content_dictionary=rows, custom_llm_provider="openai"
)
result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=rows, custom_llm_provider="openai")
assert cost == 2.5
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
assert models == ["gpt-4o"]
assert result.cost == 2.5
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15)
assert result.models == ["gpt-4o"]
# =========================================================================== #
@ -940,7 +1015,7 @@ async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monk
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
cost, usage, models = await bu._handle_completed_batch(
result = await bu._handle_completed_batch(
_batch("gs://litellm-bucket/output/predictions.jsonl"),
custom_llm_provider="vertex_ai",
litellm_params={"vertex_project": "proj-1", "vertex_location": "us-central1"},
@ -952,10 +1027,12 @@ async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monk
assert batch_input < pricing["input_cost_per_token"]
assert batch_output < pricing["output_cost_per_token"]
assert cost > 0
assert cost == pytest.approx(30 * batch_input + 15 * batch_output)
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45)
assert models == ["gemini-3.6-flash", "gemini-3.6-flash"]
assert result.cost > 0
assert result.cost == pytest.approx(30 * batch_input + 15 * batch_output)
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45)
assert result.models == ["gemini-3.6-flash", "gemini-3.6-flash"]
assert result.successful_requests == 2
assert result.failed_requests == 0
@pytest.mark.asyncio
@ -1033,11 +1110,121 @@ async def test_handle_completed_batch_orchestration(monkeypatch):
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 3.3)
cost, usage, models = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai")
result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai")
assert cost == 3.3
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
assert models == ["gpt-4o"]
assert result.cost == 3.3
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15)
assert result.models == ["gpt-4o"]
@pytest.mark.asyncio
async def test_handle_completed_batch_counts_error_file_failures(monkeypatch):
"""Regression test: OpenAI writes per-request failures (e.g. a rejected param)
to a separate error_file_id, never into the output file - so failed_requests
must include them or it silently undercounts real batch failures."""
from litellm.types.llms.openai import Batch
rows = [_success_row(model="gpt-5-mini", usage=_usage(24, 107))]
error_rows = [
{
"id": "batch_req_err1",
"custom_id": "req-2-bad",
"response": {"status_code": 400, "body": {"error": {"message": "Invalid 'temperature'"}}},
"error": None,
}
]
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
return _vertex_jsonl(rows)
async def fake_afile_content(**kw):
return type("R", (), {"content": _vertex_jsonl(error_rows)})()
import litellm.files.main as files_main
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
batch = Batch(
id="b",
completion_window="24h",
created_at=1,
endpoint="/v1/chat/completions",
input_file_id="f",
object="batch",
status="completed",
output_file_id="of",
error_file_id="ef",
)
result = await bu._handle_completed_batch(batch, custom_llm_provider="openai")
assert result.successful_requests == 1
assert result.failed_requests == 1
@pytest.mark.asyncio
async def test_handle_completed_batch_decodes_model_encoded_error_file_id(monkeypatch):
"""A model-encoded error file id must be decoded to the raw provider id before
the fetch, exactly like the output file id. Sending the encoded id straight to
the provider 404s, and the swallowed fetch failure silently reports 0 failures."""
import base64
from litellm.types.llms.openai import Batch
provider_error_file_id = "file-real-error-id"
encoded_error_file_id = "file-" + base64.urlsafe_b64encode(
f"litellm:{provider_error_file_id};model,model-abc".encode()
).decode().rstrip("=")
requested_file_ids = []
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
return _vertex_jsonl([_success_row(model="gpt-4o", usage=_usage(10, 5))])
async def fake_afile_content(**kw):
requested_file_ids.append(kw["file_id"])
return type("R", (), {"content": _vertex_jsonl([{"custom_id": "bad-1"}])})()
import litellm.files.main as files_main
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
batch = Batch(
id="b",
completion_window="24h",
created_at=1,
endpoint="/v1/chat/completions",
input_file_id="f",
object="batch",
status="completed",
output_file_id="of",
error_file_id=encoded_error_file_id,
)
result = await bu._handle_completed_batch(batch, custom_llm_provider="openai")
assert requested_file_ids == [provider_error_file_id]
assert result.failed_requests == 1
@pytest.mark.asyncio
async def test_handle_completed_batch_no_error_file_id_reports_zero_error_failures(monkeypatch):
rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
return _vertex_jsonl(rows)
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai")
assert result.successful_requests == 1
assert result.failed_requests == 0
@pytest.mark.asyncio
@ -1054,11 +1241,13 @@ async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch):
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", _must_not_fetch)
cost, usage, models = await bu._handle_completed_batch(_batch(None), custom_llm_provider="openai")
result = await bu._handle_completed_batch(_batch(None), custom_llm_provider="openai")
assert cost == 0.0
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (0, 0, 0)
assert models == []
assert result.cost == 0.0
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (0, 0, 0)
assert result.models == []
assert result.successful_requests == 0
assert result.failed_requests == 0
@pytest.mark.asyncio
@ -1075,19 +1264,25 @@ async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch)
def fake_vertex_calc(content, model):
seen["content"] = content
seen["model"] = model
return 7.7, Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3)
return bu.BatchCostUsageResult(
cost=7.7,
usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3),
models=["gemini-x"],
successful_requests=1,
failed_requests=0,
)
monkeypatch.setattr(bu, "calculate_vertex_ai_batch_cost_and_usage", fake_vertex_calc)
cost, usage, models = await bu._handle_completed_batch(
result = await bu._handle_completed_batch(
_batch("gs://litellm-bucket/output/predictions.jsonl"),
custom_llm_provider="vertex_ai",
model_name="gemini-x",
)
assert cost == 7.7
assert usage.total_tokens == 3
assert models == ["gemini-x"]
assert result.cost == 7.7
assert result.usage.total_tokens == 3
assert result.models == ["gemini-x"]
assert seen["content"] == raw_rows
assert seen["model"] == "gemini-x"
@ -1189,14 +1384,14 @@ def test_bedrock_cost_uses_deployment_model_name():
"recordId": "1",
"modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}},
}
cost, _, models = bu._aggregate_batch_cost_usage_models(
result = bu._aggregate_batch_cost_usage_models(
entries=[row],
custom_llm_provider="bedrock",
model_name="us.anthropic.claude-sonnet-4-6",
model_info={},
)
assert cost > 0
assert models == ["us.anthropic.claude-sonnet-4-6"]
assert result.cost > 0
assert result.models == ["us.anthropic.claude-sonnet-4-6"]
def test_anthropic_total_usage_sums_succeeded_only(monkeypatch):
@ -1208,8 +1403,10 @@ def test_anthropic_total_usage_sums_succeeded_only(monkeypatch):
_anthropic_errored_row(),
_anthropic_succeeded_row(usage=_anthropic_usage(20, 10, cache_read=100)),
]
_, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (130, 15, 145)
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (130, 15, 145)
assert result.successful_requests == 2
assert result.failed_requests == 1
def test_anthropic_total_usage_aggregates_cache_token_details(monkeypatch):
@ -1221,11 +1418,11 @@ def test_anthropic_total_usage_aggregates_cache_token_details(monkeypatch):
_anthropic_errored_row(),
_anthropic_succeeded_row(usage=_anthropic_usage(50, 20, cache_creation=300, cache_read=700)),
]
_, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
assert usage.prompt_tokens_details.cached_tokens == 8700
assert usage.prompt_tokens_details.cache_creation_tokens == 2300
assert usage.cache_read_input_tokens == 8700
assert usage.cache_creation_input_tokens == 2300
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
assert result.usage.prompt_tokens_details.cached_tokens == 8700
assert result.usage.prompt_tokens_details.cache_creation_tokens == 2300
assert result.usage.cache_read_input_tokens == 8700
assert result.usage.cache_creation_input_tokens == 2300
def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch):
@ -1236,9 +1433,9 @@ def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch):
"response": {"status_code": 200, "body": {"model": "gpt-5.2", "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}},
}
]
_, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
assert usage.prompt_tokens_details is None
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15)
assert result.usage.prompt_tokens_details is None
def test_anthropic_cost_applies_batch_discount_and_cache_pricing():
@ -1249,14 +1446,14 @@ def test_anthropic_cost_applies_batch_discount_and_cache_pricing():
_anthropic_errored_row(),
]
total, _, _ = bu._aggregate_batch_cost_usage_models(
result = bu._aggregate_batch_cost_usage_models(
entries=rows,
custom_llm_provider="anthropic",
model_info=_ANTHROPIC_MODEL_INFO, # type: ignore[arg-type]
)
expected_half_price = (1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3.75e-6 + 200 * 15e-6) / 2
assert total == pytest.approx(expected_half_price)
assert result.cost == pytest.approx(expected_half_price)
def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatch):
@ -1275,11 +1472,9 @@ def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatc
lambda **kw: pytest.fail("anthropic rows must not go through completion_cost"),
)
total, _, _ = bu._aggregate_batch_cost_usage_models(
entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic"
)
result = bu._aggregate_batch_cost_usage_models(entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic")
assert total == pytest.approx(0.3)
assert result.cost == pytest.approx(0.3)
assert seen[0]["model"] == "claude-sonnet-4-5-20250929"
assert seen[0]["custom_llm_provider"] == "anthropic"
assert seen[0]["usage"].prompt_tokens == 10
@ -1293,8 +1488,8 @@ def test_anthropic_batch_models_collected_from_succeeded_rows(monkeypatch):
_anthropic_succeeded_row(model="claude-sonnet-4-5-20250929"),
_anthropic_errored_row(),
]
_, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
assert models == ["claude-sonnet-4-5-20250929"]
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
assert result.models == ["claude-sonnet-4-5-20250929"]
@pytest.mark.asyncio
@ -1304,16 +1499,16 @@ async def test_calculate_batch_cost_and_usage_anthropic_end_to_end():
_anthropic_errored_row(),
]
cost, usage, models = await bu.calculate_batch_cost_and_usage(
result = await bu.calculate_batch_cost_and_usage(
file_content_dictionary=rows,
custom_llm_provider="anthropic",
model_name="claude-sonnet-4-5",
model_info=_ANTHROPIC_MODEL_INFO, # type: ignore[arg-type]
)
assert cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2)
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (11000, 200, 11200)
assert models == ["claude-sonnet-4-5"]
assert result.cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2)
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (11000, 200, 11200)
assert result.models == ["claude-sonnet-4-5"]
def test_extract_credentials_forwards_the_trusted_model_credential_snapshot():
@ -1421,24 +1616,24 @@ async def test_handle_completed_bedrock_batch_prices_from_deployment_model(monke
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
cost, usage, _ = await bu._handle_completed_batch(
result = await bu._handle_completed_batch(
_batch("of"),
custom_llm_provider="bedrock",
model_name="bedrock/global.anthropic.claude-sonnet-4-6",
)
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (1800, 1000, 2800)
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (1800, 1000, 2800)
# 3e-06 / 1.5e-05 on-demand, halved for batch.
assert cost == pytest.approx(1800 * 3e-06 / 2 + 1000 * 1.5e-05 / 2)
assert result.cost == pytest.approx(1800 * 3e-06 / 2 + 1000 * 1.5e-05 / 2)
# The response model alone cannot price a bedrock batch: this is the $0 bug.
zero_cost, zero_usage, _ = await bu._handle_completed_batch(
zero_result = await bu._handle_completed_batch(
_batch("of"),
custom_llm_provider="bedrock",
model_name=None,
)
assert zero_cost == 0.0
assert zero_usage.total_tokens == 2800
assert zero_result.cost == 0.0
assert zero_result.usage.total_tokens == 2800
@pytest.mark.asyncio
@ -1451,7 +1646,7 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) ->
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
free_cost, _, _ = await bu._handle_completed_batch(
free_result = await bu._handle_completed_batch(
_batch("of"),
custom_llm_provider="vertex_ai",
model_name="vertex_ai/gemini-2.5-flash",
@ -1462,15 +1657,15 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) ->
"output_cost_per_token_batches": 0.0,
},
)
assert free_cost == 0.0
assert free_result.cost == 0.0
billed_cost, _, _ = await bu._handle_completed_batch(
billed_result = await bu._handle_completed_batch(
_batch("of"),
custom_llm_provider="vertex_ai",
model_name="vertex_ai/gemini-2.5-flash",
model_info=None,
)
assert billed_cost > 0.0
assert billed_result.cost > 0.0
# =========================================================================== #

16
tests/test_litellm/batches/test_responses_batch_cost.py
View file

@ -71,24 +71,24 @@ async def test_responses_batch_reconciles_to_real_tokens_and_spend(local_model_c
input_tokens = 33
output_tokens = 57
cost, usage, models = await bu.calculate_batch_cost_and_usage(
result = await bu.calculate_batch_cost_and_usage(
file_content_dictionary=[_responses_line(input_tokens, output_tokens)],
custom_llm_provider="openai",
model_name=MODEL,
model_info=model_info,
)
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
input_tokens,
output_tokens,
input_tokens + output_tokens,
)
assert models == [MODEL]
assert cost == pytest.approx(
assert result.models == [MODEL]
assert result.cost == pytest.approx(
input_tokens * model_info["input_cost_per_token_batches"]
+ output_tokens * model_info["output_cost_per_token_batches"]
)
assert cost > 0.0
assert result.cost > 0.0
async def test_mixed_shape_batch_output_sums_across_both_line_shapes(local_model_cost_map):
@ -96,15 +96,15 @@ async def test_mixed_shape_batch_output_sums_across_both_line_shapes(local_model
batch's declared endpoint rather than each line's shape would miss this."""
model_info = litellm.get_model_info(model=MODEL, custom_llm_provider="openai")
cost, usage, _ = await bu.calculate_batch_cost_and_usage(
result = await bu.calculate_batch_cost_and_usage(
file_content_dictionary=[_responses_line(100, 50), _chat_line(33, 57)],
custom_llm_provider="openai",
model_name=MODEL,
model_info=model_info,
)
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (133, 107, 240)
assert cost == pytest.approx(
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (133, 107, 240)
assert result.cost == pytest.approx(
133 * model_info["input_cost_per_token_batches"] + 107 * model_info["output_cost_per_token_batches"]
)

107
tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py
View file

@ -10,11 +10,14 @@ with deployment credentials, bypassing the managed files access-control hooks.
import base64
import pytest
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
from fastapi import HTTPException
from litellm.proxy._types import UserAPIKeyAuth
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
from litellm.types.utils import LiteLLMBatch
def _make_user_api_key_dict(user_id: str) -> UserAPIKeyAuth:
@ -161,6 +164,108 @@ async def test_service_account_blocked_from_other_team_file():
assert exc_info.value.status_code == 403
# --- Keyless key must not be locked out of the batch it created ---
def _make_unified_batch_id() -> str:
raw = "litellm_proxy;model_id:my-model-id;llm_batch_id:batch_raw_123"
return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=")
def _make_managed_files_instance_with_object_store():
"""Managed-files hook backed by an in-memory stand-in for the managed
object table, so create and retrieve exercise the same stored row."""
from litellm_enterprise.proxy.hooks.managed_files import (
_PROXY_LiteLLMManagedFiles,
)
store = {}
async def upsert(where, data):
store[where["unified_object_id"]] = SimpleNamespace(**data["create"])
async def find_first(where):
return store.get(where["unified_object_id"])
mock_prisma = MagicMock()
mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock(side_effect=upsert)
mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(
side_effect=find_first
)
return (
_PROXY_LiteLLMManagedFiles(
internal_usage_cache=DualCache(),
prisma_client=mock_prisma,
),
store,
)
async def _store_batch(managed_files, unified_batch_id: str, creator: UserAPIKeyAuth):
await managed_files.store_unified_object_id(
unified_object_id=unified_batch_id,
file_object=LiteLLMBatch(
id="batch_raw_123",
completion_window="24h",
created_at=0,
endpoint="/v1/chat/completions",
input_file_id="file-1",
object="batch",
status="validating",
),
litellm_parent_otel_span=None,
model_object_id="batch_raw_123",
file_purpose="batch",
user_api_key_dict=creator,
)
@pytest.mark.asyncio
async def test_keyless_key_can_retrieve_the_batch_it_created():
"""Regression: a key with no user_id and no team_id (what `/key/generate`
by a proxy admin and service-account keys produce) stamped
`created_by=None` and was then denied its own managed batch with
"User None does not have access"."""
unified_batch_id = _make_unified_batch_id()
managed_files, store = _make_managed_files_instance_with_object_store()
keyless = UserAPIKeyAuth(api_key="sk-keyless", parent_otel_span=None)
await _store_batch(managed_files, unified_batch_id, keyless)
assert store[unified_batch_id].created_by == f"key:{keyless.token}"
data = {"batch_id": unified_batch_id}
await managed_files.async_pre_call_hook(
user_api_key_dict=keyless,
cache=DualCache(),
data=data,
call_type=CallTypes.aretrieve_batch.value,
)
assert data["batch_id"] == "batch_raw_123"
assert data["model"] == "my-model-id"
@pytest.mark.asyncio
async def test_other_keyless_key_still_denied_the_batch():
unified_batch_id = _make_unified_batch_id()
managed_files, _ = _make_managed_files_instance_with_object_store()
await _store_batch(
managed_files,
unified_batch_id,
UserAPIKeyAuth(api_key="sk-creator", parent_otel_span=None),
)
with pytest.raises(HTTPException) as exc_info:
await managed_files.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-other", parent_otel_span=None),
cache=DualCache(),
data={"batch_id": unified_batch_id},
call_type=CallTypes.aretrieve_batch.value,
)
assert exc_info.value.status_code == 403
# --- Option C fix test: check_batch_cost bypasses managed files hook ---

24
tests/test_litellm/enterprise/proxy/test_managed_files_hook.py
View file

@ -527,13 +527,33 @@ async def test_afile_list_orders_newest_first_and_breaks_ties_on_the_cursor_colu
@pytest.mark.asyncio
async def test_afile_list_denies_a_caller_without_a_user_or_team():
async def test_afile_list_scopes_a_keyless_key_to_its_own_hashed_token():
caller = UserAPIKeyAuth(api_key="sk-test", parent_otel_span=None)
managed_files, table = _make_managed_files_over_rows(
[
_make_managed_file_row("unified-mine", created_by=f"key:{caller.token}"),
_make_managed_file_row("unified-theirs", created_by="other-user"),
]
)
response = await managed_files.afile_list(
purpose=None,
litellm_parent_otel_span=None,
user_api_key_dict=caller,
)
assert [file.id for file in response.data] == ["unified-mine"]
assert table.find_many_calls[0]["where"] == {"created_by": f"key:{caller.token}"}
@pytest.mark.asyncio
async def test_afile_list_denies_a_caller_with_no_identity_at_all():
managed_files, table = _make_managed_files_over_rows([_make_managed_file_row("unified-mine")])
response = await managed_files.afile_list(
purpose=None,
litellm_parent_otel_span=None,
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", parent_otel_span=None),
user_api_key_dict=UserAPIKeyAuth(parent_otel_span=None),
)
assert response.data == []

53
tests/test_litellm/litellm_core_utils/audio_utils/test_subtitle_utils.py
View file

@ -1,5 +1,6 @@
from litellm.litellm_core_utils.audio_utils.subtitle_utils import (
SubtitleToken,
_merge_tokens_into_words,
render_subtitle_tokens_as_srt,
render_subtitle_tokens_as_vtt,
synthesize_subtitle_document,
@ -23,25 +24,59 @@ class TestRenderSubtitleTokensAsSrt:
"1\n00:00:00,000 --> 00:00:01,000\nHi.\n\n2\n00:00:01,500 --> 00:00:02,500\nHey.\n"
)
def test_token_cap_starts_a_new_cue_after_15_tokens(self):
tokens = tuple(
SubtitleToken(text=f"{index} ", start_ms=index * 100, end_ms=index * 100 + 100) for index in range(16)
def test_width_budget_starts_a_new_cue_at_word_boundaries(self):
tokens = tuple(SubtitleToken(text="abcdefghi ", start_ms=i * 100, end_ms=i * 100 + 90) for i in range(20))
result = render_subtitle_tokens_as_srt(tokens)
texts = [cue.split("\n", 2)[2] for cue in result.strip().split("\n\n")]
assert len(texts) == 3
assert all(len(text) <= 84 for text in texts)
assert all(set(text.split()) == {"abcdefghi"} for text in texts)
def test_duration_cap_starts_a_new_cue_before_word_crossing_7000ms(self):
tokens = (
SubtitleToken(text="Alpha ", start_ms=0, end_ms=3400),
SubtitleToken(text="beta ", start_ms=3400, end_ms=6800),
SubtitleToken(text="gamma", start_ms=6800, end_ms=7400),
)
assert render_subtitle_tokens_as_srt(tokens) == (
"1\n00:00:00,000 --> 00:00:01,500\n0 1 2 3 4 5 6 7 8 9 10 11 12 13 14\n"
"\n2\n00:00:01,500 --> 00:00:01,600\n15\n"
"1\n00:00:00,000 --> 00:00:06,800\nAlpha beta\n\n2\n00:00:06,800 --> 00:00:07,400\ngamma\n"
)
def test_duration_cap_starts_a_new_cue_at_5000ms(self):
def test_silence_gap_starts_a_new_cue(self):
tokens = (
SubtitleToken(text="Alpha ", start_ms=0, end_ms=400),
SubtitleToken(text="beta ", start_ms=2000, end_ms=2400),
SubtitleToken(text="gamma.", start_ms=5000, end_ms=5400),
SubtitleToken(text="beta", start_ms=2000, end_ms=2400),
)
assert render_subtitle_tokens_as_srt(tokens) == (
"1\n00:00:00,000 --> 00:00:02,400\nAlpha beta\n\n2\n00:00:05,000 --> 00:00:05,400\ngamma.\n"
"1\n00:00:00,000 --> 00:00:00,400\nAlpha\n\n2\n00:00:02,000 --> 00:00:02,400\nbeta\n"
)
def test_sentence_final_punctuation_starts_a_new_cue(self):
tokens = (
SubtitleToken(text="Done. ", start_ms=0, end_ms=400),
SubtitleToken(text="Next", start_ms=500, end_ms=800),
)
assert render_subtitle_tokens_as_srt(tokens) == (
"1\n00:00:00,000 --> 00:00:00,400\nDone.\n\n2\n00:00:00,500 --> 00:00:00,800\nNext\n"
)
def test_subword_tokens_merge_into_words_before_grouping(self):
tokens = (
SubtitleToken(text=" hel", start_ms=0, end_ms=150),
SubtitleToken(text="lo", start_ms=150, end_ms=300),
SubtitleToken(text=" world.", start_ms=350, end_ms=600),
)
assert render_subtitle_tokens_as_srt(tokens) == "1\n00:00:00,000 --> 00:00:00,600\nhello world.\n"
def test_cjk_tokens_merge_and_keep_punctuation_attached(self):
tokens = (
SubtitleToken(text="編", start_ms=0, end_ms=100),
SubtitleToken(text="集", start_ms=100, end_ms=200),
SubtitleToken(text="、", start_ms=200, end_ms=250),
SubtitleToken(text="保存", start_ms=250, end_ms=400),
)
assert [word.text for word in _merge_tokens_into_words(tokens)] == ["編", "集、", "保存"]
def test_timestampless_token_joins_the_current_cue(self):
tokens = (
SubtitleToken(text="Hello ", start_ms=0, end_ms=500),

12
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -580,9 +580,17 @@ class TestRetrieveBatchCostPassesModelIdentity:
captured: dict[str, object] = {}
async def fake_handle_completed_batch(**kwargs: object) -> tuple[float, Usage, list[str]]:
from litellm.batches.batch_utils import BatchCostUsageResult
async def fake_handle_completed_batch(**kwargs: object) -> BatchCostUsageResult:
captured.update(kwargs)
return 1.25, Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800), ["m"]
return BatchCostUsageResult(
cost=1.25,
usage=Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800),
models=["m"],
successful_requests=1,
failed_requests=0,
)
monkeypatch.setattr(logging_module, "_handle_completed_batch", fake_handle_completed_batch)

59
tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py
View file

@ -233,3 +233,62 @@ def test_api_version_in_api_base_query_is_preserved(monkeypatch):
)
assert _query_params(url) == {"api-version": "2024-05-01-preview"}
def test_v1_api_version_uses_v1_route_and_keeps_model(monkeypatch):
monkeypatch.setattr(litellm, "api_version", None, raising=False)
monkeypatch.delenv("AZURE_API_VERSION", raising=False)
config = AzureImageEditConfig()
for api_version in ("v1", "preview", "latest"):
url = config.get_complete_url(
model=_FALLBACK_MODEL,
api_base=_FALLBACK_API_BASE,
litellm_params={"api_version": api_version},
)
assert urllib.parse.urlparse(url).path == "/openai/v1/images/edits"
assert _query_params(url) == {"api-version": api_version}
assert config.finalize_image_edit_request_data({"model": _FALLBACK_MODEL, "prompt": "x"}, url) == {
"model": _FALLBACK_MODEL,
"prompt": "x",
}
def test_v1_api_version_from_global_uses_v1_route(monkeypatch):
monkeypatch.setattr(litellm, "api_version", "preview", raising=False)
monkeypatch.delenv("AZURE_API_VERSION", raising=False)
url = AzureImageEditConfig().get_complete_url(
model=_FALLBACK_MODEL,
api_base=_FALLBACK_API_BASE,
litellm_params={},
)
assert urllib.parse.urlparse(url).path == "/openai/v1/images/edits"
def test_dated_api_version_still_uses_deployment_route(monkeypatch):
monkeypatch.setattr(litellm, "api_version", None, raising=False)
monkeypatch.delenv("AZURE_API_VERSION", raising=False)
url = AzureImageEditConfig().get_complete_url(
model=_FALLBACK_MODEL,
api_base=_FALLBACK_API_BASE,
litellm_params={"api_version": "2024-10-21"},
)
assert urllib.parse.urlparse(url).path == f"/openai/deployments/{_FALLBACK_MODEL}/images/edits"
def test_v1_api_version_replaces_deployment_scoped_api_base(monkeypatch):
monkeypatch.setattr(litellm, "api_version", None, raising=False)
monkeypatch.delenv("AZURE_API_VERSION", raising=False)
url = AzureImageEditConfig().get_complete_url(
model=_FALLBACK_MODEL,
api_base=f"{_FALLBACK_API_BASE}/openai/deployments/{_FALLBACK_MODEL}/images/edits?api-version=2024-10-21",
litellm_params={"api_version": "preview"},
)
assert urllib.parse.urlparse(url).path == "/openai/v1/images/edits"
assert _query_params(url) == {"api-version": "preview"}

154
tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py
View file

@ -3,9 +3,12 @@ import traceback
from typing import Callable, Optional
from unittest.mock import AsyncMock, MagicMock, Mock, patch
import httpx
import pytest
import respx
import litellm
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.llms.azure.azure import AzureChatCompletion
from litellm.llms.azure.image_generation.http_utils import (
azure_deployment_image_generation_json_body,
@ -433,3 +436,154 @@ async def test_azure_aimage_generation_base_model_vs_deployment_name():
wire_json = post_kwargs.get("json") or {}
assert "model" not in wire_json
assert data.get("model") == base_model
@pytest.mark.parametrize("api_version", ["v1", "preview", "latest"])
def test_azure_image_generation_v1_api_version_uses_v1_route(api_version):
"""The v1 Azure surface exposes /openai/v1/images/generations and routes by body ``model``."""
url = AzureChatCompletion().create_azure_base_url(
azure_client_params={
"azure_endpoint": "https://my-resource.openai.azure.com",
"api_version": api_version,
},
model="gpt-image-1",
base_model=None,
)
assert url == f"https://my-resource.openai.azure.com/openai/v1/images/generations?api-version={api_version}"
data = {"model": "gpt-image-1", "prompt": "x"}
assert azure_deployment_image_generation_json_body(url, data) == data
def test_azure_image_generation_dated_api_version_uses_deployment_route():
url = AzureChatCompletion().create_azure_base_url(
azure_client_params={
"azure_endpoint": "https://my-resource.openai.azure.com",
"api_version": "2024-10-21",
},
model="gpt-image-1",
base_model=None,
)
assert (
url
== "https://my-resource.openai.azure.com/openai/deployments/gpt-image-1/images/generations?api-version=2024-10-21"
)
assert "model" not in azure_deployment_image_generation_json_body(url, {"model": "gpt-image-1", "prompt": "x"})
def test_azure_image_generation_v1_api_version_replaces_deployment_scoped_api_base():
url = AzureChatCompletion().create_azure_base_url(
azure_client_params={
"azure_endpoint": "https://my-resource.openai.azure.com/openai/deployments/gpt-image-1/images/generations",
"api_version": "preview",
},
model="gpt-image-1",
base_model=None,
)
assert url == "https://my-resource.openai.azure.com/openai/v1/images/generations?api-version=preview"
def test_azure_image_generation_v1_api_version_uses_base_url_client_param():
url = AzureChatCompletion().create_azure_base_url(
azure_client_params={
"base_url": "https://my-resource.openai.azure.com/openai/deployments/gpt-image-1?api-version=2024-10-21",
"api_version": "preview",
},
model="gpt-image-1",
base_model=None,
)
assert url == "https://my-resource.openai.azure.com/openai/v1/images/generations?api-version=preview"
def test_azure_v1_image_generation_json_body_sends_deployment_name():
"""The v1 route ignores the URL and routes by body ``model``, which must be the deployment name."""
url = "https://my-resource.openai.azure.com/openai/v1/images/generations?api-version=preview"
data = {"model": "gpt-image-2", "prompt": "x", "n": 1}
out = azure_deployment_image_generation_json_body(url, data, deployment_name="img-dep")
assert out["model"] == "img-dep"
assert out["prompt"] == "x"
assert data["model"] == "gpt-image-2"
assert azure_deployment_image_generation_json_body(url, data) == data
@pytest.mark.asyncio
async def test_azure_aimage_generation_v1_route_sends_deployment_name_in_body(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
azure_chat_completion = AzureChatCompletion()
model = "img-dep"
base_model = "gpt-image-2"
data = {"model": base_model, "prompt": "A beautiful image of a cat", "n": 1}
azure_client_params = {
"azure_endpoint": "https://my-resource.openai.azure.com",
"api_version": "preview",
}
route = respx_mock.post("https://my-resource.openai.azure.com/openai/v1/images/generations").mock(
return_value=httpx.Response(200, json={"created": 1234567890, "data": [{"b64_json": "aaaa"}]})
)
logging_obj = MagicMock()
logging_obj.pre_call = MagicMock()
logging_obj.post_call = MagicMock()
await azure_chat_completion.aimage_generation(
data=data,
model_response=None,
azure_client_params=azure_client_params,
api_key="test-api-key",
input=[],
logging_obj=logging_obj,
headers={},
model=model,
timeout=60.0,
)
request = route.calls.last.request
assert str(request.url) == ("https://my-resource.openai.azure.com/openai/v1/images/generations?api-version=preview")
sent_body = json.loads(request.content)
assert sent_body["model"] == model
assert sent_body["prompt"] == data["prompt"]
def test_azure_image_generation_v1_route_base_model_vs_deployment_name(respx_mock: respx.MockRouter):
"""On the v1 surface the body ``model`` must be the deployment name, never base_model."""
azure_chat_completion = AzureChatCompletion()
prompt = "A beautiful image of a cat"
model = "img-dep"
base_model = "gpt-image-2"
api_base = "https://my-resource.openai.azure.com"
api_version = "v1"
litellm_params = {
"base_model": base_model,
"api_base": api_base,
"api_version": api_version,
}
route = respx_mock.post(f"{api_base}/openai/v1/images/generations").mock(
return_value=httpx.Response(200, json={"created": 1234567890, "data": [{"b64_json": "aaaa"}]})
)
logging_obj = MagicMock()
logging_obj.pre_call = MagicMock()
logging_obj.post_call = MagicMock()
azure_chat_completion.image_generation(
prompt=prompt,
timeout=60.0,
optional_params={"n": 1, "size": "1024x1024"},
logging_obj=logging_obj,
headers={},
model=model,
api_key="test-api-key",
api_base=api_base,
api_version=api_version,
litellm_params=litellm_params,
)
request = route.calls.last.request
assert str(request.url) == f"{api_base}/openai/v1/images/generations?api-version={api_version}"
sent_body = json.loads(request.content)
assert sent_body["model"] == model
assert sent_body["prompt"] == prompt

44
tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py
View file

@ -7,6 +7,7 @@ import pytest
from litellm.llms.base_llm.managed_resources.isolation import (
build_owner_filter,
can_access_resource,
resolve_resource_owner_id,
)
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
@ -154,3 +155,46 @@ def test_access_identity_less_caller_always_denied(created_by, resource_team_id)
)
is False
)
# ---------------------------------------------------------------------------
# keyless keys (no user_id, no team_id) own their resources by hashed token
# ---------------------------------------------------------------------------
def test_owner_id_prefers_user_id_then_falls_back_to_token():
assert resolve_resource_owner_id(UserAPIKeyAuth(user_id="alice")) == "alice"
assert resolve_resource_owner_id(UserAPIKeyAuth(team_id="team-eng")) is None
assert resolve_resource_owner_id(UserAPIKeyAuth()) is None
keyless = UserAPIKeyAuth(api_key="sk-keyless")
assert resolve_resource_owner_id(keyless) == f"key:{keyless.token}"
def test_keyless_key_can_access_its_own_resource():
"""Regression for the self-lockout: a key generated by a proxy admin (or a
service-account key) has no user_id and no team_id, so it used to stamp
`created_by=None` and then be denied its own batches and files."""
keyless = UserAPIKeyAuth(api_key="sk-keyless")
owner_id = resolve_resource_owner_id(keyless)
assert build_owner_filter(keyless) == {"created_by": owner_id}
assert (
can_access_resource(keyless, created_by=owner_id, resource_team_id=None) is True
)
def test_keyless_key_denied_another_keyless_keys_resource():
"""The #27004 isolation invariant: two distinct keyless keys must not see
each other's resources."""
creator = UserAPIKeyAuth(api_key="sk-creator")
other = UserAPIKeyAuth(api_key="sk-other")
assert (
can_access_resource(
other,
created_by=resolve_resource_owner_id(creator),
resource_team_id=None,
)
is False
)

2
tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py
View file

@ -945,7 +945,7 @@ def test_titan_image_embedding_cost_uses_per_image_rate():
"encoding_format,expected_embedding_types",
[
("float", ["float"]),
("base64", ["base64"]),
("base64", ["float"]),
(["float", "int8"], ["float", "int8"]),
],
)

13
tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py
View file

@ -4,6 +4,7 @@ from unittest.mock import patch, MagicMock, AsyncMock
import litellm
import pytest
import respx
MOCK_EMBEDDING_RESPONSE = [[0.1, 0.2, 0.3, 0.4, 0.5]]
@ -21,6 +22,16 @@ def mock_embedding_http_handler():
yield mock_post
@pytest.fixture
def mock_hf_config_fetch():
"""Serve the Hugging Face config.json fetched during cost calculation, so no test leaves the process"""
with respx.mock(assert_all_called=False) as respx_mock:
respx_mock.get(url__regex=r"https://huggingface\.co/.*/config\.json").respond(
json={"max_position_embeddings": 512}
)
yield respx_mock
@pytest.fixture
def mock_embedding_async_http_handler():
"""Fixture to mock the async HTTP handler for embedding tests"""
@ -39,7 +50,7 @@ def mock_embedding_async_http_handler():
class TestHuggingFaceEmbedding:
@pytest.fixture(autouse=True)
def setup(self, mock_embedding_http_handler, mock_embedding_async_http_handler):
def setup(self, mock_embedding_http_handler, mock_embedding_async_http_handler, mock_hf_config_fetch):
self.mock_get_task_patcher = patch(
"litellm.llms.huggingface.embedding.handler.get_hf_task_embedding_for_model"
)

185
tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py
View file

@ -369,6 +369,191 @@ class TestRenderSonioxTokensAsSrt:
assert "01:01:01,000" in result
def _subword_tokens(words, start_ms=0, subword_ms=150, inter_word_gap_ms=50):
tokens = []
t = start_ms
for word in words:
halves = [word[: len(word) // 2], word[len(word) // 2 :]] if len(word) > 3 else [word]
for i, piece in enumerate(halves):
text = (" " + piece) if i == 0 else piece
tokens.append({"text": text, "start_ms": t, "end_ms": t + subword_ms})
t += subword_ms
t += inter_word_gap_ms
return tokens, t
class TestCueGroupingAlignment:
def test_should_split_cue_on_silence_gap_with_exact_timestamps(self):
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
before, t = _subword_tokens(["hello", "there"])
after, _ = _subword_tokens(["welcome", "back"], start_ms=t + 5000)
result = render_soniox_tokens_as_srt(before + after)
cues = result.strip().split("\n\n")
assert len(cues) == 2
assert "00:00:00,000 --> 00:00:00,650" in cues[0]
assert "hello there" in cues[0]
assert "00:00:05,700 --> 00:00:06,350" in cues[1]
assert "welcome back" in cues[1]
def test_should_not_bridge_pause_shorter_than_old_duration_cap(self):
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
before, t = _subword_tokens(["first", "part"])
after, _ = _subword_tokens(["second", "part"], start_ms=t + 3000)
result = render_soniox_tokens_as_srt(before + after)
cues = result.strip().split("\n\n")
assert len(cues) == 2
assert "first part" in cues[0]
assert "second part" in cues[1]
def test_should_never_split_mid_word(self):
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
tokens, _ = _subword_tokens(["hello"] * 20)
result = render_soniox_tokens_as_srt(tokens)
text_lines = [
line for line in result.split("\n") if line and "-->" not in line and not line.isdigit()
]
assert len(text_lines) >= 2
for line in text_lines:
assert set(line.split()) == {"hello"}
def test_should_split_after_sentence_final_punctuation(self):
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
tokens, _ = _subword_tokens(["That", "is", "done.", "Next", "topic"])
result = render_soniox_tokens_as_srt(tokens)
cues = result.strip().split("\n\n")
assert len(cues) == 2
assert cues[0].endswith("That is done.")
assert cues[1].endswith("Next topic")
def test_should_split_on_char_budget_at_word_boundary(self):
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
tokens, _ = _subword_tokens(["wonderful"] * 12)
result = render_soniox_tokens_as_srt(tokens)
text_lines = [
line for line in result.split("\n") if line and "-->" not in line and not line.isdigit()
]
assert len(text_lines) >= 2
for line in text_lines:
assert len(line) <= 84
assert set(line.split()) == {"wonderful"}
def test_should_exclude_untimestamped_translation_tokens_from_cues(self):
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
tokens = [
{"text": " Good", "start_ms": 0, "end_ms": 200, "translation_status": "original", "language": "en"},
{"text": " Guten", "translation_status": "translation", "language": "de", "source_language": "en"},
{"text": " morning.", "start_ms": 250, "end_ms": 600, "translation_status": "original", "language": "en"},
]
result = render_soniox_tokens_as_srt(tokens)
assert "Good morning." in result
assert "Guten" not in result
assert "00:00:00,000 --> 00:00:00,600" in result
def test_should_split_before_word_whose_end_crosses_duration_cap(self):
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
tokens = [{"text": " hm", "start_ms": i * 650, "end_ms": i * 650 + 600} for i in range(10)] + [
{"text": " boom", "start_ms": 6900, "end_ms": 7600}
]
result = render_soniox_tokens_as_srt(tokens)
cues = result.strip().split("\n\n")
assert len(cues) == 2
assert "00:00:00,000 --> 00:00:06,450" in cues[0]
assert "00:00:06,900 --> 00:00:07,600" in cues[1]
assert cues[1].endswith("boom")
def test_should_keep_untimestamped_word_in_cue(self):
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
tokens = [
{"text": " uh", "start_ms": None, "end_ms": None},
{"text": " hello", "start_ms": 100, "end_ms": 500},
]
result = render_soniox_tokens_as_srt(tokens)
assert "uh hello" in result
assert "00:00:00,100 --> 00:00:00,500" in result
def _cue_texts(srt: str) -> list:
return [cue.split("\n", 2)[2] for cue in srt.strip().split("\n\n")]
class TestMultilingualCueGrouping:
def test_should_split_spaceless_chinese_on_width_budget(self):
from litellm.litellm_core_utils.audio_utils.subtitle_utils import _text_width
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
tokens = [{"text": "你好", "start_ms": i * 100, "end_ms": i * 100 + 90} for i in range(60)]
result = render_soniox_tokens_as_srt(tokens)
texts = _cue_texts(result)
assert len(texts) >= 3
for text in texts:
assert _text_width(text) <= 84
assert set(text) <= {"你", "好"}
def test_should_split_japanese_after_sentence_end_and_keep_punctuation_attached(self):
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
tokens = [
{"text": "今日は", "start_ms": 0, "end_ms": 300},
{"text": "いい", "start_ms": 300, "end_ms": 500},
{"text": "天気です", "start_ms": 500, "end_ms": 900},
{"text": "。", "start_ms": 900, "end_ms": 950},
{"text": "明日も", "start_ms": 1000, "end_ms": 1300},
{"text": "晴れ", "start_ms": 1300, "end_ms": 1500},
]
texts = _cue_texts(render_soniox_tokens_as_srt(tokens))
assert texts == ["今日はいい天気です。", "明日も晴れ"]
def test_should_split_arabic_after_arabic_question_mark(self):
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
tokens = [
{"text": " كيف", "start_ms": 0, "end_ms": 300},
{"text": " حالك؟", "start_ms": 300, "end_ms": 700},
{"text": " أنا", "start_ms": 800, "end_ms": 1000},
{"text": " بخير", "start_ms": 1000, "end_ms": 1300},
]
texts = _cue_texts(render_soniox_tokens_as_srt(tokens))
assert texts == ["كيف حالك؟", "أنا بخير"]
def test_should_split_after_devanagari_and_urdu_terminators(self):
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
tokens = [
{"text": " नमस्ते।", "start_ms": 0, "end_ms": 400},
{"text": " آپ", "start_ms": 500, "end_ms": 700},
{"text": " ٹھیک۔", "start_ms": 700, "end_ms": 1100},
{"text": " शुभ", "start_ms": 1200, "end_ms": 1400},
]
texts = _cue_texts(render_soniox_tokens_as_srt(tokens))
assert texts == ["नमस्ते।", "آپ ٹھیک۔", "शुभ"]
def test_should_split_russian_after_sentence_end(self):
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
tokens = [
{"text": " Как", "start_ms": 0, "end_ms": 200},
{"text": " дела?", "start_ms": 200, "end_ms": 600},
{"text": " Хорошо.", "start_ms": 700, "end_ms": 1200},
]
texts = _cue_texts(render_soniox_tokens_as_srt(tokens))
assert texts == ["Как дела?", "Хорошо."]
def test_should_not_split_latin_text_within_width_budget(self):
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt
tokens = [{"text": f" word{i}", "start_ms": i * 100, "end_ms": i * 100 + 90} for i in range(12)]
texts = _cue_texts(render_soniox_tokens_as_srt(tokens))
assert len(texts) == 1
class TestRenderSonioxTokensAsVtt:
def test_should_render_basic_vtt_with_header(self):
from litellm.llms.soniox.common_utils import render_soniox_tokens_as_vtt

162
tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py
View file

@ -4,13 +4,17 @@ Tests for Vertex AI (Veo) video generation transformation.
import base64
import json
import os
from unittest.mock import MagicMock, Mock, patch
from collections.abc import Mapping
from pathlib import Path
from typing import cast
from unittest.mock import Mock, patch
import httpx
import pytest
import litellm
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.llms.openai.cost_calculation import video_generation_cost
from litellm.llms.vertex_ai.videos.transformation import (
VertexAIVideoConfig,
_convert_image_to_vertex_format,
@ -18,6 +22,21 @@ from litellm.llms.vertex_ai.videos.transformation import (
from litellm.types.router import GenericLiteLLMParams
from litellm.types.videos.main import VideoObject
VEO_31_LITE_VERTEX_MODEL = "vertex_ai/veo-3.1-lite-generate-001"
ROOT_MODEL_COST_PATH = (
Path(__file__).parents[5] / "model_prices_and_context_window.json"
)
BACKUP_MODEL_COST_PATH = (
Path(__file__).parents[5]
/ "litellm"
/ "model_prices_and_context_window_backup.json"
)
ModelCostMap = Mapping[str, Mapping[str, object]]
def _load_model_cost_map(path: Path) -> ModelCostMap:
return cast(ModelCostMap, json.loads(path.read_text()))
class TestVertexAIVideoConfig:
"""Test VertexAIVideoConfig transformation class."""
@ -117,6 +136,56 @@ class TestVertexAIVideoConfig:
# Should NOT include endpoint
assert not url.endswith(":predictLongRunning")
def test_veo_31_lite_model_cost_entries_match_pricing(self):
for path in (ROOT_MODEL_COST_PATH, BACKUP_MODEL_COST_PATH):
model_cost = _load_model_cost_map(path)
info = model_cost.get(VEO_31_LITE_VERTEX_MODEL)
assert info is not None, f"{VEO_31_LITE_VERTEX_MODEL} missing from {path}"
assert info["litellm_provider"] == "vertex_ai-video-models"
assert info["mode"] == "video_generation"
assert info["max_input_tokens"] == 1024
assert info["output_cost_per_second"] == 0.05
assert info["output_cost_per_second_1080p"] == 0.08
assert info["supported_modalities"] == ["text", "image"]
def test_veo_31_lite_provider_routing_from_local_model_map(
self, monkeypatch: pytest.MonkeyPatch
):
model_cost = _load_model_cost_map(BACKUP_MODEL_COST_PATH)
vertex_video_models = {
model_name.removeprefix("vertex_ai/")
for model_name, info in model_cost.items()
if info.get("litellm_provider") == "vertex_ai-video-models"
}
monkeypatch.setattr(litellm, "vertex_ai_video_models", vertex_video_models)
model, custom_llm_provider, _, _ = get_llm_provider(
model="veo-3.1-lite-generate-001"
)
assert model == "veo-3.1-lite-generate-001"
assert custom_llm_provider == "vertex_ai"
def test_veo_31_lite_cost_uses_resolution_tiers(self):
model_cost = _load_model_cost_map(BACKUP_MODEL_COST_PATH)
model_info = model_cost[VEO_31_LITE_VERTEX_MODEL]
assert video_generation_cost(
model=VEO_31_LITE_VERTEX_MODEL,
duration_seconds=10.0,
custom_llm_provider="vertex_ai",
model_info=dict(model_info),
video_resolution="720p",
) == pytest.approx(0.5)
assert video_generation_cost(
model=VEO_31_LITE_VERTEX_MODEL,
duration_seconds=10.0,
custom_llm_provider="vertex_ai",
model_info=dict(model_info),
video_resolution="1080p",
) == pytest.approx(0.8)
def test_transform_video_create_request(self):
"""Test transformation of video creation request."""
prompt = "A cat playing with a ball of yarn"
@ -210,6 +279,95 @@ class TestVertexAIVideoConfig:
assert mapped["durationSeconds"] == 8
assert mapped["aspectRatio"] == "16:9"
assert "resolution" not in mapped
@pytest.mark.parametrize(
("model", "size", "expected_resolution"),
(
(VEO_31_LITE_VERTEX_MODEL, "1280x720", "720p"),
(
VEO_31_LITE_VERTEX_MODEL.removeprefix("vertex_ai/"),
"1920x1080",
"1080p",
),
),
)
def test_map_openai_size_to_resolution_for_resolution_tier_model(
self,
model: str,
size: str,
expected_resolution: str,
monkeypatch: pytest.MonkeyPatch,
):
model_cost = _load_model_cost_map(BACKUP_MODEL_COST_PATH)
monkeypatch.setitem(
litellm.model_cost,
VEO_31_LITE_VERTEX_MODEL,
dict(model_cost[VEO_31_LITE_VERTEX_MODEL]),
)
mapped = self.config.map_openai_params(
video_create_optional_params={"size": size},
model=model,
drop_params=False,
)
assert mapped["aspectRatio"] == "16:9"
assert mapped["resolution"] == expected_resolution
def test_map_openai_size_does_not_infer_resolution_for_veo_2(self):
mapped = self.config.map_openai_params(
video_create_optional_params={"size": "1920x1080"},
model="vertex_ai/veo-2.0-generate-001",
drop_params=False,
)
assert mapped["aspectRatio"] == "16:9"
assert "resolution" not in mapped
def test_map_openai_size_does_not_infer_resolution_for_existing_veo_3(
self, monkeypatch: pytest.MonkeyPatch
):
model = "veo-3.1-generate-001"
model_key = f"vertex_ai/{model}"
model_cost = _load_model_cost_map(BACKUP_MODEL_COST_PATH)
monkeypatch.setitem(litellm.model_cost, model_key, dict(model_cost[model_key]))
mapped = self.config.map_openai_params(
video_create_optional_params={"size": "1920x1080"},
model=model,
drop_params=False,
)
assert mapped["aspectRatio"] == "16:9"
assert "resolution" not in mapped
def test_map_openai_size_does_not_override_provider_resolution(self):
mapped = self.config.map_openai_params(
video_create_optional_params={
"size": "1920x1080",
"parameters": {"resolution": "720p"},
},
model=VEO_31_LITE_VERTEX_MODEL,
drop_params=False,
)
assert mapped["aspectRatio"] == "16:9"
assert "resolution" not in mapped
assert mapped["parameters"] == {"resolution": "720p"}
def test_map_openai_size_does_not_override_direct_resolution(self):
mapped = self.config.map_openai_params(
video_create_optional_params={
"size": "1920x1080",
"resolution": "720p",
},
model=VEO_31_LITE_VERTEX_MODEL,
drop_params=False,
)
assert mapped["aspectRatio"] == "16:9"
assert mapped["resolution"] == "720p"
def test_map_openai_params_default_duration(self):
"""Test that durationSeconds is omitted when not provided."""

79
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -10290,3 +10290,82 @@ def test_native_client_authorize_without_the_proxy_resource_keeps_the_mcp_flow(m
assert 'name="decision"' not in response.text
assert "team-b" not in response.text
assert minted == []
def test_introspect_route_requires_virtual_key_auth_and_is_advertised():
"""RFC 7662 section 2.1: introspection must not be anonymous. Pins the route-level
user_api_key_auth dependency (structure, so removing it fails here without a proxy),
and that the aggregate AS metadata advertises the endpoint for discovery."""
from fastapi import FastAPI
from fastapi.routing import APIRoute
from fastapi.testclient import TestClient
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
route = next(r for r in router.routes if isinstance(r, APIRoute) and r.path == "/introspect")
assert route.methods == {"POST"}
assert any(dependency.call is user_api_key_auth for dependency in route.dependant.dependencies)
from litellm.proxy._types import LiteLLMRoutes
assert "/introspect" in LiteLLMRoutes.mcp_routes.value
from litellm.proxy._lazy_features import LAZY_FEATURES
discoverable = next(feature for feature in LAZY_FEATURES if feature.name == "mcp_discoverable")
assert "/introspect" in discoverable.path_prefixes
app = FastAPI()
app.include_router(router)
client = TestClient(app)
asm = client.get("/.well-known/oauth-authorization-server/mcp")
assert asm.json()["introspection_endpoint"] == "http://testserver/introspect"
def test_introspect_route_answers_for_authenticated_caller(monkeypatch):
"""End-to-end over the real route with the auth dependency satisfied: a garbage token
is active false, a freshly minted session access token is active true with its claims."""
from datetime import datetime, timezone
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
session_keys_from_master_key,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
SessionPrincipal,
mint_session_token,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
introspect_master_key = "sk-introspect-route-test"
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", introspect_master_key, raising=False)
async def fake_reload(user_id: str):
return None
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_user_by_id", fake_reload
)
app = FastAPI()
app.include_router(router)
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth()
client = TestClient(app)
garbage = client.post("/introspect", data={"token": "llm_session_garbage"})
assert garbage.status_code == 200
assert garbage.json() == {"active": False}
minted = mint_session_token(
SessionPrincipal(user_id="u1", client_id="llm_dcrc_client"),
session_keys_from_master_key(introspect_master_key),
datetime.now(timezone.utc),
)
active = client.post("/introspect", data={"token": minted.token.get_secret_value()})
assert active.status_code == 200
assert active.json()["active"] is True
assert active.json()["sub"] == "u1"

141
tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
View file

@ -26,6 +26,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
aggregate_authorize,
aggregate_token,
complete_connect_flow,
introspect_gateway_token,
is_gateway_dcr_client_id,
is_proxy_api_resource,
native_client_auth_contract,
@ -41,7 +42,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent
resolve_session_bearer,
session_keys_from_master_key,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import SESSION_REFRESH_PREFIX
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
SESSION_ISSUER,
SESSION_REFRESH_PREFIX,
SessionPrincipal,
mint_session_refresh_token,
mint_session_token,
)
MASTER_KEY = "sk-gateway-dcr-flow-tests"
REDIRECT_URI = "https://claude.ai/api/mcp/auth_callback"
@ -1598,17 +1605,24 @@ async def test_refresh_answers_503_without_burning_the_token_while_redis_is_down
)
redis_down = await _refresh_native(
payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(side_effect=ConnectionError("redis down")))
payload["refresh_token"],
client_id,
_Minter(),
_redis_that(AsyncMock(side_effect=ConnectionError("redis down"))),
)
assert redis_down.status_code == 503
assert json.loads(redis_down.body)["error"] == "temporarily_unavailable"
assert "refresh_token" not in json.loads(redis_down.body)
redis_back = await _refresh_native(payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=1)))
redis_back = await _refresh_native(
payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=1))
)
assert redis_back.status_code == 200
assert json.loads(redis_back.body)["refresh_token"] != payload["refresh_token"]
replayed = await _refresh_native(payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=2)))
replayed = await _refresh_native(
payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=2))
)
assert replayed.status_code == 400
assert json.loads(replayed.body)["error"] == "invalid_grant"
@ -1674,3 +1688,122 @@ def test_native_client_auth_contract_points_every_endpoint_at_this_proxy():
)
def test_is_proxy_api_resource_matches_only_this_proxy(resource, expected):
assert is_proxy_api_resource(_request(), resource) is expected
def _introspection_fixtures():
keys = session_keys_from_master_key(MASTER_KEY)
now = datetime.now(timezone.utc)
principal = SessionPrincipal(user_id="u1", client_id="llm_dcrc_client", team_id="t1")
return keys, now, principal
async def _introspect(token, cache=None, reload_user=_reload_user_active, master_key=MASTER_KEY):
response = await introspect_gateway_token(
token=token, master_key=master_key, reload_user=reload_user, cache=cache or DualCache()
)
return response.status_code, json.loads(response.body)
@pytest.mark.asyncio
async def test_introspect_active_access_token_reports_rfc7662_claims():
keys, now, principal = _introspection_fixtures()
minted = mint_session_token(principal, keys, now)
status, body = await _introspect(minted.token.get_secret_value())
assert status == 200
assert body["active"] is True
assert body["token_type"] == "Bearer"
assert body["iss"] == SESSION_ISSUER
assert body["sub"] == "u1"
assert body["client_id"] == "llm_dcrc_client"
assert body["kind"] == "session"
assert body["team_id"] == "t1"
assert body["exp"] - body["iat"] == 3600
assert body["jti"]
@pytest.mark.asyncio
async def test_introspect_invalid_tokens_answer_active_false():
keys, now, principal = _introspection_fixtures()
wrong_key = mint_session_token(principal, session_keys_from_master_key("sk-a-rotated-master-key"), now)
expired = mint_session_token(principal, keys, now - timedelta(seconds=7200))
for candidate in (
"sk-not-a-session-token",
"llm_session_malformed",
wrong_key.token.get_secret_value(),
expired.token.get_secret_value(),
):
status, body = await _introspect(candidate)
assert (status, body) == (200, {"active": False})
@pytest.mark.asyncio
async def test_introspect_refresh_token_goes_inactive_once_rotated():
keys, now, _ = _introspection_fixtures()
client_id = (await _register([REDIRECT_URI]))["client_id"]
minted = mint_session_refresh_token(SessionPrincipal(user_id="u1", client_id=client_id), keys, now)
cache = DualCache()
status, body = await _introspect(minted.token.get_secret_value(), cache=cache)
assert (status, body["active"], body["kind"]) == (200, True, "session_refresh")
assert "token_type" not in body
revoked = await revoke_refresh_token(
token=minted.token.get_secret_value(), client_id=client_id, master_key=MASTER_KEY, cache=cache
)
assert revoked.status_code == 200
status, body = await _introspect(minted.token.get_secret_value(), cache=cache)
assert (status, body) == (200, {"active": False})
@pytest.mark.asyncio
async def test_introspect_accepts_rs256_signed_tokens_under_configured_signing(monkeypatch):
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from pydantic import SecretStr
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import AsymmetricSessionKeys
private_pem = (
rsa.generate_private_key(public_exponent=65537, key_size=2048)
.private_bytes(
serialization.Encoding.PEM,
serialization.PrivateFormat.PKCS8,
serialization.NoEncryption(),
)
.decode()
)
monkeypatch.setitem(
proxy_server.general_settings,
"mcp_session_token_signing",
{"algorithm": "RS256", "kid": "k1", "private_key": private_pem},
)
_, now, principal = _introspection_fixtures()
rs_keys = AsymmetricSessionKeys(private_key_pem=SecretStr(private_pem), kid="k1")
minted = mint_session_token(principal, rs_keys, now)
status, body = await _introspect(minted.token.get_secret_value())
assert (status, body["active"], body["kind"]) == (200, True, "session")
hs_signed = mint_session_token(principal, session_keys_from_master_key(MASTER_KEY), now)
status, body = await _introspect(hs_signed.token.get_secret_value())
assert (status, body) == (200, {"active": False})
@pytest.mark.asyncio
async def test_introspect_fails_closed_on_dead_user_and_503s_on_outage():
keys, now, principal = _introspection_fixtures()
minted = mint_session_token(principal, keys, now)
async def _reload_user_gone(user_id: str):
return "unresolvable"
status, body = await _introspect(minted.token.get_secret_value(), reload_user=_reload_user_gone)
assert (status, body) == (200, {"active": False})
async def _reload_user_outage(user_id: str):
return "unavailable"
status, body = await _introspect(minted.token.get_secret_value(), reload_user=_reload_user_outage)
assert (status, body["error"]) == (503, "temporarily_unavailable")
status, body = await _introspect(minted.token.get_secret_value(), master_key=None)
assert (status, body["error"]) == (500, "server_error")

14
tests/test_litellm/proxy/auth/test_cli_auth.py
View file

@ -82,7 +82,7 @@ async def test_poll_for_ready_404(sleep_mock, request_mock):
_poll_for_ready_data(
"https://litellm.com", poll_interval=1, total_timeout=1, request_timeout=42
)
request_mock.assert_called_once_with("https://litellm.com", timeout=42)
request_mock.assert_called_once_with("https://litellm.com", headers=None, timeout=42)
@pytest.mark.asyncio
@ -103,7 +103,7 @@ async def test_poll_for_ready_200_ready(sleep_mock, click_mock, request_mock):
)
assert actual == {"status": "ready", "json": "data"}
click_mock.assert_not_called()
request_mock.assert_called_once_with("https://litellm.com", timeout=42)
request_mock.assert_called_once_with("https://litellm.com", headers=None, timeout=42)
sleep_mock.assert_not_called()
@ -131,8 +131,8 @@ async def test_poll_for_ready_single_pending(sleep_mock, click_mock, request_moc
click_mock.assert_not_called()
request_mock.assert_has_calls(
[
call("https://litellm.com", timeout=42),
call("https://litellm.com", timeout=42),
call("https://litellm.com", headers=None, timeout=42),
call("https://litellm.com", headers=None, timeout=42),
]
)
sleep_mock.assert_called_once_with(1)
@ -168,8 +168,8 @@ async def test_poll_for_ready_pending(sleep_mock, click_mock, request_mock):
click_mock.assert_has_calls([call("Pending message"), call("Pending message")])
request_mock.assert_has_calls(
[
call("https://litellm.com", timeout=42),
call("https://litellm.com", timeout=42),
call("https://litellm.com", headers=None, timeout=42),
call("https://litellm.com", headers=None, timeout=42),
]
)
sleep_mock.assert_has_calls([call(1), call(1)])
@ -194,7 +194,7 @@ async def test_poll_for_ready_connection_failure(sleep_mock, click_mock, request
click_mock.assert_called_once_with("Connection error (will retry): ERROR")
request_mock.assert_has_calls(
[
call("https://litellm.com", timeout=42),
call("https://litellm.com", headers=None, timeout=42),
]
)
sleep_mock.assert_has_calls([call(1), call(1)])

38
tests/test_litellm/proxy/client/conftest.py Normal file
View file

@ -0,0 +1,38 @@
import threading
import pytest
@pytest.fixture
def hanging_server():
"""A server that accepts the connection and never answers, so only a timeout ends the call."""
from http.server import BaseHTTPRequestHandler, HTTPServer
from socketserver import ThreadingMixIn
stop: threading.Event = threading.Event()
class SilentRequestHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def _hang(self):
stop.wait(timeout=30)
do_GET = _hang
do_POST = _hang
def log_message(self, format, *args):
pass
class ThreadedServer(ThreadingMixIn, HTTPServer):
daemon_threads = True
server = ThreadedServer(("127.0.0.1", 0), SilentRequestHandler)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
yield f"http://127.0.0.1:{server.server_port}"
finally:
stop.set()
server.shutdown()
server.server_close()
thread.join(timeout=5)

29
tests/test_litellm/proxy/client/test_chat.py
View file

@ -1,6 +1,7 @@
import importlib
import importlib.util
from importlib.machinery import PathFinder
import time
import site
import sys
@ -227,3 +228,31 @@ def test_completions_other_errors(client, sample_messages):
with pytest.raises(requests.exceptions.HTTPError) as exc_info:
client.completions(model="gpt-4", messages=sample_messages)
assert exc_info.value.response.status_code == 500
def test_completions_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
"""
A proxy that accepts the connection but never answers used to pin the caller's
process forever, since the request carried no timeout at all.
"""
client = ChatClient(base_url=hanging_server, api_key="sk-test", timeout=1)
started = time.monotonic()
with pytest.raises(requests.exceptions.Timeout):
client.completions(model="gpt-5.4", messages=[{"role": "user", "content": "hi"}])
assert time.monotonic() - started < 10
def test_completions_stream_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
"""
The streaming call opens the response before reading chunks, so a proxy that never
sends its headers used to hang here forever too.
"""
client = ChatClient(base_url=hanging_server, api_key="sk-test", timeout=1)
started = time.monotonic()
with pytest.raises(requests.exceptions.Timeout):
next(client.completions_stream(model="gpt-5.4", messages=[{"role": "user", "content": "hi"}]))
assert time.monotonic() - started < 10

8
tests/test_litellm/proxy/client/test_client.py
View file

@ -82,6 +82,12 @@ def test_client_initialization():
assert client.http._base_url == "http://localhost:4000"
assert client.http._api_key == "test-key"
assert client.http._timeout == 60
assert client.teams._timeout == 60
assert client.keys._timeout == 60
assert client.credentials._timeout == 60
assert client.models._timeout == 60
assert client.model_groups._timeout == 60
assert client.chat._timeout == 600
def test_client_default_timeout():
@ -92,6 +98,8 @@ def test_client_default_timeout():
)
assert client.http._timeout == 30
assert client.keys._timeout == 30
assert client.chat._timeout == 600
def test_client_without_api_key():

15
tests/test_litellm/proxy/client/test_credentials.py
View file

@ -1,4 +1,5 @@
import time
import pytest
import requests
@ -276,3 +277,17 @@ def test_encrypt_credential_values_does_not_mutate_original(monkeypatch):
assert encrypted.credential_values["api_key"] != "sk-123"
assert credential.credential_values["api_key"] == "sk-123"
assert encrypted.credential_name == credential.credential_name
def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
"""
A proxy that accepts the connection but never answers used to pin the caller's
process forever, since the request carried no timeout at all.
"""
client = CredentialsManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1)
started = time.monotonic()
with pytest.raises(requests.exceptions.Timeout):
client.list()
assert time.monotonic() - started < 10

15
tests/test_litellm/proxy/client/test_keys.py
View file

@ -1,3 +1,4 @@
import time
import traceback
import pytest
@ -509,3 +510,17 @@ def test_not_found_error_redacts_wrapped_key():
assert "REDACTED" in str(wrapped)
assert LEAKY_KEY not in str(wrapped.orig_exception)
assert wrapped.orig_exception.response.status_code == 404
def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
"""
A proxy that accepts the connection but never answers used to pin the caller's
process forever, since the request carried no timeout at all.
"""
client = KeysManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1)
started = time.monotonic()
with pytest.raises(requests.exceptions.Timeout):
client.list()
assert time.monotonic() - started < 10

15
tests/test_litellm/proxy/client/test_model_groups.py
View file

@ -1,4 +1,5 @@
import time
import pytest
import requests
@ -172,3 +173,17 @@ def test_client_initialization_without_api_key(base_url):
assert client._api_key is None
assert client.model_groups._api_key is None
def test_info_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
"""
A proxy that accepts the connection but never answers used to pin the caller's
process forever, since the request carried no timeout at all.
"""
client = ModelGroupsManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1)
started = time.monotonic()
with pytest.raises(requests.exceptions.Timeout):
client.info()
assert time.monotonic() - started < 10

15
tests/test_litellm/proxy/client/test_models.py
View file

@ -1,4 +1,5 @@
import time
import pytest
import requests
@ -732,3 +733,17 @@ def test_update_other_errors(client):
with pytest.raises(requests.exceptions.HTTPError) as exc_info:
client.update(model_id=model_id, model_params=model_params)
assert exc_info.value.response.status_code == 500
def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
"""
A proxy that accepts the connection but never answers used to pin the caller's
process forever, since the request carried no timeout at all.
"""
client = ModelsManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1)
started = time.monotonic()
with pytest.raises(requests.exceptions.Timeout):
client.list()
assert time.monotonic() - started < 10

20
tests/test_litellm/proxy/client/test_teams.py Normal file
View file

@ -0,0 +1,20 @@
import time
import pytest
import requests
from litellm.proxy.client.teams import TeamsManagementClient
def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
"""
A proxy that accepts the connection but never answers used to pin the caller's
process forever, since the request carried no timeout at all.
"""
client = TeamsManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1)
started = time.monotonic()
with pytest.raises(requests.exceptions.Timeout):
client.list()
assert time.monotonic() - started < 10

16
tests/test_litellm/proxy/client/test_users.py
View file

@ -1,6 +1,8 @@
import time
from unittest.mock import MagicMock, patch
import pytest
import requests
@ -82,3 +84,17 @@ def test_delete_user_unauthorized(mock_post, client):
mock_post.return_value.text = "unauthorized"
with pytest.raises(UnauthorizedError):
client.delete_user(["u1"])
def test_delete_user_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
"""
A proxy that accepts the connection but never answers used to pin the caller's
process forever, since the request carried no timeout at all.
"""
client = UsersManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1)
started = time.monotonic()
with pytest.raises(requests.exceptions.Timeout):
client.delete_user(["u1"])
assert time.monotonic() - started < 10

43
tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py
View file

@ -101,6 +101,49 @@ def test_per_model_reads_route_to_reader_writes_to_writer():
assert actions.delete_many is writer_inner.litellm_usertable.delete_many
def test_writer_pinned_client_bypasses_reader_routing():
"""Regression for #38556: read-after-write reconciles must see the writer's
just-committed rows, so WriterPinnedClient must resolve reads to the writer
even when a read replica is configured."""
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient
writer, writer_inner, reader, reader_inner = _make_wrappers()
writer_inner.litellm_proxymodeltable = _model_actions_mock("writer_models")
reader_inner.litellm_proxymodeltable = _model_actions_mock("reader_models")
routing = RoutingPrismaWrapper(writer=writer, reader=reader)
pinned = WriterPinnedClient(routing)
assert pinned.db is writer
assert pinned.db.litellm_proxymodeltable.find_many is writer_inner.litellm_proxymodeltable.find_many
def test_writer_pinned_client_passes_through_single_db():
from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient
writer, _, _, _ = _make_wrappers()
assert WriterPinnedClient(writer).db is writer
def test_writer_pinned_client_yields_to_routed_reads_when_writer_down():
"""The pin must not break reader-only degraded mode: a proxy that starts
during a primary outage still loads DB-backed models from the replica, so
while the writer is degraded the pin resolves to the routed wrapper."""
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient
writer, writer_inner, reader, reader_inner = _make_wrappers()
writer_inner.litellm_proxymodeltable = _model_actions_mock("writer_models")
reader_inner.litellm_proxymodeltable = _model_actions_mock("reader_models")
routing = RoutingPrismaWrapper(writer=writer, reader=reader)
routing._writer_unavailable = True
pinned = WriterPinnedClient(routing)
assert pinned.db is routing
assert pinned.db.litellm_proxymodeltable.find_many is reader_inner.litellm_proxymodeltable.find_many
@pytest.mark.asyncio
async def test_connect_invokes_both_clients():
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper

203
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
View file

@ -17,14 +17,18 @@ Tests cover:
- CCR: headroom_retrieve tool injected when compressed messages contain hashes
- CCR: async_should_run_agentic_loop returns True when response has headroom_retrieve tool calls
- CCR: async_build_agentic_loop_plan calls retrieve endpoint and builds follow-up messages
- CCR: streaming /chat/completions is converted to a non-streaming call so the agentic
loop resolves the retrieve tool call, then fake-streamed back to the client
"""
import json
import time
from typing import Optional
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
from fastapi import HTTPException
import litellm
@ -38,7 +42,11 @@ from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import (
from litellm.proxy.spend_tracking.compression_savings import (
extract_compression_saved_tokens,
)
from litellm.types.utils import GenericGuardrailAPIInputs
from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY
from litellm.types.utils import (
CallTypes,
GenericGuardrailAPIInputs,
)
FAKE_API_BASE = "https://headroom.example.com"
FAKE_API_KEY = "test-key"
@ -1893,6 +1901,199 @@ async def test_fail_open_returns_original_parts_shapes():
assert [m["content"] for m in messages] == [m["content"] for m in PARTS_MESSAGES]
CCR_HASH = "b573993006976af767214fac"
def _retrieve_tool_definition() -> dict:
return {
"type": "function",
"function": {
"name": HEADROOM_RETRIEVE_TOOL_NAME,
"description": "retrieve compressed content",
"parameters": {"type": "object", "properties": {"hash": {"type": "string"}}},
},
}
def _openai_completion_payload(message: dict, finish_reason: str) -> dict:
return {
"id": "chatcmpl-ccr",
"object": "chat.completion",
"created": 1700000000,
"model": "gpt-4o",
"choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
}
def _openai_tool_call_payload() -> dict:
return _openai_completion_payload(
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_ccr",
"type": "function",
"function": {
"name": HEADROOM_RETRIEVE_TOOL_NAME,
"arguments": json.dumps({"hash": CCR_HASH}),
},
}
],
},
"tool_calls",
)
def _openai_text_payload(content: str) -> dict:
return _openai_completion_payload({"role": "assistant", "content": content}, "stop")
@pytest.mark.parametrize(
"call_type, stream, tools, expect_conversion",
[
(CallTypes.acompletion, True, [_retrieve_tool_definition()], True),
(CallTypes.completion, True, [_retrieve_tool_definition()], True),
(CallTypes.acompletion, False, [_retrieve_tool_definition()], False),
(CallTypes.acompletion, True, [{"type": "function", "function": {"name": "get_weather"}}], False),
(CallTypes.acompletion, True, None, False),
(CallTypes.aresponses, True, [_retrieve_tool_definition()], False),
(CallTypes.anthropic_messages, True, [_retrieve_tool_definition()], False),
],
)
@pytest.mark.asyncio
async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions(
guardrail: HeadroomGuardrail,
call_type: CallTypes,
stream: bool,
tools: Optional[list],
expect_conversion: bool,
):
kwargs = {"model": "gpt-4o", "stream": stream, "tools": tools}
result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=call_type)
if not expect_conversion:
assert result is kwargs
assert HEADROOM_CONVERTED_STREAM_KEY not in kwargs
assert kwargs["stream"] is stream
return
assert result is not None
assert result["stream"] is False
assert result[HEADROOM_CONVERTED_STREAM_KEY] is True
assert kwargs["stream"] is True
@pytest.mark.asyncio
async def test_pre_call_deployment_hook_still_compresses_for_deployment_level_configs(
guardrail: HeadroomGuardrail,
):
"""Regression for the stream-conversion override swallowing the parent hook:
when the guardrail is attached at the deployment level and proxy pre_call never
ran, the deployment hook is the only place compression executes, so the
override must delegate to CustomGuardrail.async_pre_call_deployment_hook."""
kwargs = {
"model": "gpt-4o",
"messages": [dict(m) for m in ORIGINAL_MESSAGES],
"stream": False,
"guardrails": ["headroom"],
"metadata": {},
}
with patch.object(
guardrail.async_handler,
"post",
new_callable=AsyncMock,
return_value=_make_compress_response(COMPRESSED_MESSAGES),
):
result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acompletion)
assert result is not None
assert result["messages"] == EXPECTED_MESSAGES
@pytest.mark.asyncio
async def test_pre_call_deployment_hook_converts_stream_after_deployment_level_compression(
guardrail: HeadroomGuardrail,
):
kwargs = {
"model": "gpt-4o",
"messages": [dict(m) for m in ORIGINAL_MESSAGES],
"stream": True,
"guardrails": ["headroom"],
"metadata": {},
}
with patch.object(
guardrail.async_handler,
"post",
new_callable=AsyncMock,
return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH),
):
result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acompletion)
assert result is not None
assert has_headroom_retrieve_tool(result["tools"])
assert result["stream"] is False
assert result[HEADROOM_CONVERTED_STREAM_KEY] is True
@pytest.mark.asyncio
async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end(
guardrail: HeadroomGuardrail,
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
):
"""Regression test for streaming /chat/completions: the retrieve tool call the
model emits must be resolved by the agentic loop instead of being streamed back
to a client that never declared the tool."""
original_content = "the full uncompressed document"
final_answer = "the document says hello"
guardrail._issued_hashes_by_call_id["ccr-call-id"] = (
frozenset({CCR_HASH}),
time.monotonic() + 999,
)
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
monkeypatch.setattr(litellm, "callbacks", [guardrail])
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
upstream = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
side_effect=[
httpx.Response(200, json=_openai_tool_call_payload()),
httpx.Response(200, json=_openai_text_payload(final_answer)),
]
)
with patch.object(
guardrail.async_handler,
"get",
new_callable=AsyncMock,
return_value=_make_retrieve_response(original_content),
) as mock_get:
response = await litellm.acompletion(
model="openai/gpt-4o",
messages=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}],
tools=[_retrieve_tool_definition()],
stream=True,
litellm_call_id="ccr-call-id",
)
chunks = [chunk async for chunk in response]
streamed_text = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices)
assert streamed_text == final_answer
assert not any(chunk.choices and chunk.choices[0].delta.tool_calls for chunk in chunks)
mock_get.assert_called_once()
assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0])
assert len(upstream.calls) == 2
followup_body = json.loads(upstream.calls[1].request.content)
assert not followup_body.get("stream")
assert original_content in json.dumps(followup_body["messages"])
assert not any(key.startswith("_headroom_interception") for key in followup_body)
# ---------------------------------------------------------------------------
# LIT-5018: the turn the model is being asked to act on is never compressed.
#

48
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py
View file

@ -1,4 +1,6 @@
import os
import threading
import time
import uuid
from typing import List, cast
from unittest.mock import AsyncMock, MagicMock, patch
@ -6,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
from httpx import Request, Response
import requests
import litellm
@ -14,6 +17,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.proxy.guardrails.guardrail_hooks.hiddenlayer.hiddenlayer import (
HiddenlayerGuardrail,
HiddenlayerGuardrailV2,
_get_jwt,
)
from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
from litellm.types.utils import (
@ -1088,3 +1092,47 @@ class TestHiddenlayerGuardrailV2:
config_model = HiddenlayerGuardrailV2.get_config_model()
assert config_model is not None
assert config_model.__name__ == "HiddenlayerGuardrailConfigModel"
@pytest.fixture
def hanging_auth_server():
"""A server that accepts the connection and never answers, so only a timeout ends the call."""
from http.server import BaseHTTPRequestHandler, HTTPServer
from socketserver import ThreadingMixIn
stop: threading.Event = threading.Event()
class SilentRequestHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def do_POST(self):
stop.wait(timeout=30)
def log_message(self, format, *args):
pass
class ThreadedServer(ThreadingMixIn, HTTPServer):
daemon_threads = True
server = ThreadedServer(("127.0.0.1", 0), SilentRequestHandler)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
yield f"http://127.0.0.1:{server.server_port}"
finally:
stop.set()
server.shutdown()
server.server_close()
thread.join(timeout=5)
def test_get_jwt_gives_up_at_the_timeout_instead_of_blocking_the_event_loop(hanging_auth_server):
"""
`_get_jwt` runs synchronously inside `_call_hiddenlayer`, so an auth host that
accepts and never answers used to park the whole worker's event loop.
"""
started = time.monotonic()
with pytest.raises(requests.exceptions.Timeout):
_get_jwt(auth_url=hanging_auth_server, api_id="id", api_key="secret", timeout=1)
assert time.monotonic() - started < 10

391
tests/test_litellm/proxy/management_endpoints/test_config_override_endpoints.py
View file

@ -11,16 +11,19 @@ import litellm
import litellm.proxy.proxy_server as ps
from litellm.proxy._types import KeyManagementSystem, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.config_override_endpoints import (
CYBERARK_ENV_VAR_MAPPING,
HASHICORP_ENV_VAR_MAPPING,
_build_field_schema,
_set_env_vars,
)
from litellm.proxy.proxy_server import app
from litellm.types.proxy.management_endpoints.config_overrides import (
CyberArkConfig,
HashicorpVaultConfig,
)
VAULT_URL = "/config_overrides/hashicorp_vault"
CYBERARK_URL = "/config_overrides/cyberark"
@pytest.fixture
@ -42,6 +45,7 @@ def _make_mock_proxy_config():
cfg = MagicMock()
cfg.initialize_secret_manager = MagicMock()
cfg._last_hashicorp_vault_config = None
cfg._cyberark_boot_env = None
cfg._encrypt_env_variables = MagicMock(
side_effect=lambda d: {k: f"enc_{v}" for k, v in d.items()}
)
@ -67,6 +71,8 @@ def _cleanup():
app.dependency_overrides.pop(ps.user_api_key_auth, None)
for env_var in HASHICORP_ENV_VAR_MAPPING.values():
os.environ.pop(env_var, None)
for env_var in CYBERARK_ENV_VAR_MAPPING.values():
os.environ.pop(env_var, None)
def _set_admin():
@ -275,6 +281,391 @@ async def test_hashicorp_vault_validation_errors_and_access_control(
_cleanup()
@pytest.mark.asyncio
async def test_cyberark_crud_lifecycle(client, monkeypatch):
"""Create → read (masked) → partial update (merge from DB) → clear field →
delete → idempotent delete → env fallback → merge from env → schema."""
mock_prisma, mock_db = _make_mock_db()
mock_cfg = _make_mock_proxy_config()
mock_cfg._last_cyberark_config = None
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
monkeypatch.setattr(ps, "proxy_config", mock_cfg)
old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system
_set_admin()
try:
# 1. POST: create with API-key auth
r = client.post(
CYBERARK_URL,
json={
"cyberark_api_base": "https://conjur.example.com",
"cyberark_account": "myorg",
"cyberark_username": "litellm-user",
"cyberark_api_key": "my-secret-api-key",
},
)
assert r.status_code == 200
assert os.environ["CYBERARK_API_BASE"] == "https://conjur.example.com"
assert os.environ["CYBERARK_API_KEY"] == "my-secret-api-key"
data = _upserted_data(mock_db)
assert data["cyberark_api_key"] == "enc_my-secret-api-key"
mock_cfg.initialize_secret_manager.assert_called_with(
key_management_system="cyberark"
)
assert mock_cfg._last_cyberark_config is not None
# 2. GET: sensitive fields masked
mock_db.find_unique = AsyncMock(return_value=_db_record(data))
r = client.get(CYBERARK_URL)
assert r.status_code == 200
vals = r.json()["values"]
assert vals["cyberark_api_base"] == "https://conjur.example.com"
assert "*" in vals["cyberark_api_key"]
assert "properties" in r.json()["field_schema"]
# 3. POST partial: omitted fields merge from DB
r = client.post(CYBERARK_URL, json={"cyberark_api_base": "https://conjur.new.com"})
assert r.status_code == 200
data = _upserted_data(mock_db)
assert data["cyberark_api_base"] == "enc_https://conjur.new.com"
assert data["cyberark_api_key"] == "enc_my-secret-api-key"
assert data["cyberark_account"] == "enc_myorg"
# 4. POST empty string: clears field, switches to cert auth
step3 = {
**data,
"client_cert": "enc_/certs/client.pem",
"client_key": "enc_/certs/client.key",
}
mock_db.find_unique = AsyncMock(return_value=_db_record(step3))
mock_db.upsert = AsyncMock(return_value=None)
r = client.post(CYBERARK_URL, json={"cyberark_api_key": ""})
assert r.status_code == 200
data = _upserted_data(mock_db)
assert "cyberark_api_key" not in data
assert data["client_cert"] == "enc_/certs/client.pem"
# 5. DELETE: clears everything
litellm.secret_manager_client = MagicMock() # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
litellm._key_management_system = KeyManagementSystem.CYBERARK # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
r = client.delete(CYBERARK_URL)
assert r.status_code == 200
assert os.environ.get("CYBERARK_API_BASE") is None
assert litellm.secret_manager_client is None
assert mock_cfg._last_cyberark_config is None
# 6. DELETE idempotent
mock_db.delete = AsyncMock(
side_effect=RecordNotFoundError(
data={"clientVersion": "0.0.0"}, message="Not found"
)
)
assert client.delete(CYBERARK_URL).status_code == 200
# 7. GET: env var fallback with masking
mock_db.find_unique = AsyncMock(return_value=None)
monkeypatch.setenv("CYBERARK_API_BASE", "https://conjur.env.com")
monkeypatch.setenv("CYBERARK_API_KEY", "env-api-key")
r = client.get(CYBERARK_URL)
vals = r.json()["values"]
assert vals["cyberark_api_base"] == "https://conjur.env.com"
assert "*" in vals["cyberark_api_key"]
# 8. POST: merge from env vars
mock_cfg.initialize_secret_manager = MagicMock()
mock_db.upsert = AsyncMock(return_value=None)
r = client.post(CYBERARK_URL, json={"cyberark_api_base": "https://conjur.merged.com"})
assert r.status_code == 200
data = _upserted_data(mock_db)
assert data["cyberark_api_key"] == "enc_env-api-key"
# 9. _build_field_schema
schema = _build_field_schema(CyberArkConfig)
assert "cyberark_api_base" in schema["properties"]
assert len(schema["properties"]["cyberark_api_base"]["description"]) > 0
finally:
litellm.secret_manager_client = old_client # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
litellm._key_management_system = old_kms # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
_cleanup()
@pytest.mark.asyncio
async def test_cyberark_validation_errors_and_access_control(client, monkeypatch):
"""Validation (missing api base, missing auth, init failure rollback),
DELETE preserves non-CyberArk secret managers, non-admin 403."""
mock_prisma, mock_db = _make_mock_db()
mock_cfg = MagicMock()
mock_cfg._last_cyberark_config = {"cyberark_api_base": "old"}
mock_cfg._cyberark_boot_env = None
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
monkeypatch.setattr(ps, "proxy_config", mock_cfg)
old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system
_set_admin()
try:
# 1. Missing cyberark_api_base → 400
r = client.post(CYBERARK_URL, json={"cyberark_api_key": "key"})
assert r.status_code == 400
assert "API Base" in r.json()["detail"]
# 2. Missing auth → 400 (cert without key is not valid auth)
r = client.post(
CYBERARK_URL,
json={"cyberark_api_base": "https://c.com", "client_cert": "/c.pem"},
)
assert r.status_code == 400
assert "authentication" in r.json()["detail"].lower()
# 3. Init failure → 500, env vars restored, nothing persisted
mock_cfg.initialize_secret_manager = MagicMock(side_effect=Exception("fail"))
monkeypatch.setenv("CYBERARK_API_BASE", "https://conjur.old.com")
monkeypatch.setenv("CYBERARK_API_KEY", "old-key")
r = client.post(
CYBERARK_URL,
json={"cyberark_api_base": "https://bad.com", "cyberark_api_key": "bad"},
)
assert r.status_code == 500
assert os.environ["CYBERARK_API_BASE"] == "https://conjur.old.com"
mock_db.upsert.assert_not_awaited()
# 4. DELETE preserves non-CyberArk secret manager
aws = MagicMock()
litellm.secret_manager_client = aws # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
litellm._key_management_system = KeyManagementSystem.AWS_SECRET_MANAGER # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
assert client.delete(CYBERARK_URL).status_code == 200
assert litellm.secret_manager_client is aws
# 5. Non-admin → 403
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, user_id="user"
)
assert client.get(CYBERARK_URL).status_code == 403
assert (
client.post(
CYBERARK_URL, json={"cyberark_api_base": "https://c.com"}
).status_code
== 403
)
assert client.delete(CYBERARK_URL).status_code == 403
assert client.post(CYBERARK_URL + "/test_connection").status_code == 403
finally:
litellm.secret_manager_client = old_client # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
litellm._key_management_system = old_kms # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
_cleanup()
@pytest.mark.asyncio
async def test_cyberark_delete_restores_deployment_env_config(client, monkeypatch):
"""Deleting the DB override must restore env vars the deployment started with,
and reinitialize the manager from them, instead of wiping CyberArk entirely."""
mock_prisma, mock_db = _make_mock_db()
mock_cfg = _make_mock_proxy_config()
mock_cfg._last_cyberark_config = None
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
monkeypatch.setattr(ps, "proxy_config", mock_cfg)
old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system
_set_admin()
try:
monkeypatch.setenv("CYBERARK_API_BASE", "https://conjur.boot.com")
monkeypatch.setenv("CYBERARK_API_KEY", "boot-key")
r = client.post(
CYBERARK_URL,
json={"cyberark_api_base": "https://conjur.db.com", "cyberark_api_key": "db-key"},
)
assert r.status_code == 200
assert os.environ["CYBERARK_API_BASE"] == "https://conjur.db.com"
mock_cfg.initialize_secret_manager.reset_mock()
r = client.delete(CYBERARK_URL)
assert r.status_code == 200
assert os.environ["CYBERARK_API_BASE"] == "https://conjur.boot.com"
assert os.environ["CYBERARK_API_KEY"] == "boot-key"
mock_cfg.initialize_secret_manager.assert_called_with(key_management_system="cyberark")
assert mock_cfg._last_cyberark_config is None
finally:
litellm.secret_manager_client = old_client # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
litellm._key_management_system = old_kms # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
_cleanup()
@pytest.mark.asyncio
async def test_cyberark_persist_failure_rolls_back_runtime_state(client, monkeypatch):
"""If the DB upsert fails after the manager was reinitialized, the endpoint
must restore the previous env vars and reinitialize from them, so this pod
does not keep serving credentials that were never committed to the DB."""
mock_prisma, mock_db = _make_mock_db()
mock_cfg = _make_mock_proxy_config()
mock_cfg._last_cyberark_config = None
mock_db.upsert = AsyncMock(side_effect=Exception("db write failed"))
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
monkeypatch.setattr(ps, "proxy_config", mock_cfg)
old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system
_set_admin()
try:
monkeypatch.setenv("CYBERARK_API_BASE", "https://conjur.prev.com")
monkeypatch.setenv("CYBERARK_API_KEY", "prev-key")
r = client.post(
CYBERARK_URL,
json={"cyberark_api_base": "https://conjur.new.com", "cyberark_api_key": "new-key"},
)
assert r.status_code == 500
assert "persist" in r.json()["detail"].lower()
assert os.environ["CYBERARK_API_BASE"] == "https://conjur.prev.com"
assert os.environ["CYBERARK_API_KEY"] == "prev-key"
# last call must be the rollback reinit against the restored env
assert (
mock_cfg.initialize_secret_manager.call_args_list[-1].kwargs["key_management_system"] == "cyberark"
)
assert os.environ.get("CYBERARK_API_BASE") != "https://conjur.new.com"
finally:
litellm.secret_manager_client = old_client # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
litellm._key_management_system = old_kms # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
_cleanup()
@pytest.mark.asyncio
async def test_cyberark_persist_failure_restores_hashicorp_manager(client, monkeypatch):
"""If CyberArk init displaced an env-configured Hashicorp manager and the DB
upsert then fails, rollback must bring the Hashicorp manager back."""
mock_prisma, mock_db = _make_mock_db()
mock_cfg = _make_mock_proxy_config()
mock_cfg._last_cyberark_config = None
mock_db.upsert = AsyncMock(side_effect=Exception("db write failed"))
def _fake_init(key_management_system):
litellm._key_management_system = ( # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
KeyManagementSystem.CYBERARK
if key_management_system == "cyberark"
else KeyManagementSystem.HASHICORP_VAULT
)
mock_cfg.initialize_secret_manager = MagicMock(side_effect=_fake_init)
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
monkeypatch.setattr(ps, "proxy_config", mock_cfg)
old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system
_set_admin()
try:
monkeypatch.setenv("HCP_VAULT_ADDR", "https://vault.example.com")
litellm._key_management_system = KeyManagementSystem.HASHICORP_VAULT # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
r = client.post(
CYBERARK_URL,
json={"cyberark_api_base": "https://conjur.new.com", "cyberark_api_key": "new-key"},
)
assert r.status_code == 500
assert litellm._key_management_system == KeyManagementSystem.HASHICORP_VAULT
assert (
mock_cfg.initialize_secret_manager.call_args_list[-1].kwargs["key_management_system"] == "hashicorp_vault"
)
finally:
litellm.secret_manager_client = old_client # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
litellm._key_management_system = old_kms # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
os.environ.pop("HCP_VAULT_ADDR", None)
_cleanup()
@pytest.mark.asyncio
async def test_cyberark_audit_log_redacts_values(client, monkeypatch):
monkeypatch.setattr(litellm, "store_audit_logs", True)
mock_prisma, mock_db = _make_mock_db()
mock_cfg = _make_mock_proxy_config()
mock_cfg._last_cyberark_config = None
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
monkeypatch.setattr(ps, "proxy_config", mock_cfg)
_set_admin()
audit_calls = []
async def capture(request_data):
audit_calls.append(request_data)
try:
with patch( # test-quality-ok: patching proxy-internal collaborator to isolate the endpoint
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
new=capture,
):
r = client.post(
CYBERARK_URL,
json={
"cyberark_api_base": "https://conjur.example.com",
"cyberark_api_key": "my-very-secret-key",
},
)
assert r.status_code == 200
for _ in range(3):
await asyncio.sleep(0)
assert len(audit_calls) == 1
log = audit_calls[0]
assert log.action == "created"
assert log.object_id == "cyberark"
assert "my-very-secret-key" not in log.updated_values
assert "conjur.example.com" not in log.updated_values
after = json.loads(log.updated_values)
assert "cyberark_api_key" in after["config"]
assert "cyberark_api_base" in after["config"]
finally:
_cleanup()
@pytest.mark.asyncio
async def test_cyberark_test_connection(client, monkeypatch):
"""400 when not configured; success path authenticates and hits /whoami."""
from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager
old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system
_set_admin()
try:
# Not configured → 400
litellm.secret_manager_client = None # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
r = client.post(CYBERARK_URL + "/test_connection")
assert r.status_code == 400
assert "not configured" in r.json()["detail"].lower()
# Configured → authenticates and calls /whoami
mock_manager = MagicMock(spec=CyberArkSecretManager)
mock_manager.conjur_addr = "https://conjur.example.com"
mock_manager.ssl_verify = True
mock_manager._get_request_headers = MagicMock(
return_value={"Authorization": "Token abc"}
)
litellm.secret_manager_client = mock_manager # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
mock_response = MagicMock()
mock_response.raise_for_status = MagicMock()
mock_http = MagicMock()
mock_http.get = AsyncMock(return_value=mock_response)
with patch( # test-quality-ok: patching proxy-internal collaborator to isolate the endpoint
"litellm.proxy.management_endpoints.config_override_endpoints.get_async_httpx_client",
return_value=mock_http,
):
r = client.post(CYBERARK_URL + "/test_connection")
assert r.status_code == 200
assert "conjur.example.com" in r.json()["message"]
called_url = mock_http.get.call_args.args[0]
assert called_url == "https://conjur.example.com/whoami"
# Auth failure → 502
mock_manager._get_request_headers = MagicMock(
side_effect=Exception("bad credentials")
)
r = client.post(CYBERARK_URL + "/test_connection")
assert r.status_code == 502
assert "authentication failed" in r.json()["detail"].lower()
finally:
litellm.secret_manager_client = old_client # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
litellm._key_management_system = old_kms # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them
_cleanup()
# ── Audit-log emission for /config_overrides/hashicorp_vault ─────────────────

78
tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py
View file

@ -478,14 +478,14 @@ class TestVertexAIBatchPassthroughHandler:
}
]
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
result = calculate_vertex_ai_batch_cost_and_usage(
vertex_ai_batch_responses, model_name="gemini-2.0-flash-001"
)
assert usage.total_tokens == 15
assert usage.prompt_tokens == 10
assert usage.completion_tokens == 5
assert total_cost > 0, "batch_cost_calculator should return a non-zero cost"
assert result.usage.total_tokens == 15
assert result.usage.prompt_tokens == 10
assert result.usage.completion_tokens == 5
assert result.cost > 0, "batch_cost_calculator should return a non-zero cost"
def test_batch_response_transformation(self):
"""Test transformation of Vertex AI batch responses to OpenAI format"""
@ -664,14 +664,14 @@ class TestVertexAIBatchCostCalculation:
},
]
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
result = calculate_vertex_ai_batch_cost_and_usage(
responses, model_name="gemini-2.0-flash-001"
)
assert usage.prompt_tokens == 18
assert usage.completion_tokens == 8
assert usage.total_tokens == 26
assert total_cost > 0, "batch_cost_calculator should return a non-zero cost"
assert result.usage.prompt_tokens == 18
assert result.usage.completion_tokens == 8
assert result.usage.total_tokens == 26
assert result.cost > 0, "batch_cost_calculator should return a non-zero cost"
def test_should_skip_responses_with_null_response_body(self):
"""Failed lines (response: None) are skipped without error."""
@ -699,27 +699,29 @@ class TestVertexAIBatchCostCalculation:
},
]
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
result = calculate_vertex_ai_batch_cost_and_usage(
responses, model_name="gemini-2.0-flash-001"
)
assert usage.prompt_tokens == 18
assert usage.completion_tokens == 8
assert usage.total_tokens == 26
assert total_cost > 0
assert result.usage.prompt_tokens == 18
assert result.usage.completion_tokens == 8
assert result.usage.total_tokens == 26
assert result.cost > 0
assert result.successful_requests == 2
assert result.failed_requests == 1
def test_should_return_zeros_for_empty_response_list(self):
"""Empty input → zero cost and zero usage."""
from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
result = calculate_vertex_ai_batch_cost_and_usage(
[], model_name="gemini-2.0-flash-001"
)
assert total_cost == 0.0
assert usage.total_tokens == 0
assert usage.prompt_tokens == 0
assert usage.completion_tokens == 0
assert result.cost == 0.0
assert result.usage.total_tokens == 0
assert result.usage.prompt_tokens == 0
assert result.usage.completion_tokens == 0
def test_should_handle_missing_usage_metadata_gracefully(self):
"""Response without usageMetadata → 0 tokens, 0 cost for that line."""
@ -729,13 +731,13 @@ class TestVertexAIBatchCostCalculation:
{"response": {"candidates": [{"content": {"parts": [{"text": "hi"}]}}]}},
]
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
result = calculate_vertex_ai_batch_cost_and_usage(
responses, model_name="gemini-2.0-flash-001"
)
assert usage.prompt_tokens == 0
assert usage.completion_tokens == 0
assert usage.total_tokens == 0
assert result.usage.prompt_tokens == 0
assert result.usage.completion_tokens == 0
assert result.usage.total_tokens == 0
@pytest.mark.asyncio
async def test_openai_shaped_output_records_nonzero_cost_and_usage(self):
@ -813,7 +815,7 @@ class TestVertexAIBatchCostCalculation:
try:
litellm.disable_vertex_batch_output_transformation = False
cost, usage, _ = await calculate_batch_cost_and_usage(
result = await calculate_batch_cost_and_usage(
file_content_dictionary=openai_shaped_responses,
custom_llm_provider="vertex_ai",
model_name="gemini-2.0-flash-001",
@ -822,17 +824,17 @@ class TestVertexAIBatchCostCalculation:
litellm.disable_vertex_batch_output_transformation = original_flag
assert (
usage.prompt_tokens == 18
), f"expected 18 prompt tokens, got {usage.prompt_tokens}"
result.usage.prompt_tokens == 18
), f"expected 18 prompt tokens, got {result.usage.prompt_tokens}"
assert (
usage.completion_tokens == 8
), f"expected 8 completion tokens, got {usage.completion_tokens}"
result.usage.completion_tokens == 8
), f"expected 8 completion tokens, got {result.usage.completion_tokens}"
assert (
usage.total_tokens == 26
), f"expected 26 total tokens, got {usage.total_tokens}"
result.usage.total_tokens == 26
), f"expected 26 total tokens, got {result.usage.total_tokens}"
assert (
cost > 0
), f"expected non-zero cost for completed Vertex batch, got {cost}"
result.cost > 0
), f"expected non-zero cost for completed Vertex batch, got {result.cost}"
@pytest.mark.asyncio
async def test_raw_vertex_output_still_works_when_transformation_disabled(self):
@ -865,7 +867,7 @@ class TestVertexAIBatchCostCalculation:
try:
litellm.disable_vertex_batch_output_transformation = True
cost, usage, _ = await calculate_batch_cost_and_usage(
result = await calculate_batch_cost_and_usage(
file_content_dictionary=raw_vertex_responses,
custom_llm_provider="vertex_ai",
model_name="gemini-2.0-flash-001",
@ -873,7 +875,7 @@ class TestVertexAIBatchCostCalculation:
finally:
litellm.disable_vertex_batch_output_transformation = original_flag
assert usage.prompt_tokens == 10
assert usage.completion_tokens == 5
assert usage.total_tokens == 15
assert cost > 0, "raw Vertex shape should also produce non-zero cost"
assert result.usage.prompt_tokens == 10
assert result.usage.completion_tokens == 5
assert result.usage.total_tokens == 15
assert result.cost > 0, "raw Vertex shape should also produce non-zero cost"

6
tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
View file

@ -2865,7 +2865,7 @@ class TestSpendLogsPayload:
"model": "gpt-4o",
"user": "",
"team_id": "",
"metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}',
"metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}',
"cache_key": "Cache OFF",
"spend": 0.00022500000000000002,
"total_tokens": 30,
@ -2961,7 +2961,7 @@ class TestSpendLogsPayload:
"model": "claude-4-sonnet-20250514",
"user": "",
"team_id": "",
"metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}',
"metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}',
"cache_key": "Cache OFF",
"spend": 0.01383,
"total_tokens": 2598,
@ -3055,7 +3055,7 @@ class TestSpendLogsPayload:
"model": "claude-4-sonnet-20250514",
"user": "",
"team_id": "",
"metadata": '{"applied_guardrails": [], "attempted_fallbacks": 0, "original_model_group": "my-anthropic-model-group", "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}',
"metadata": '{"applied_guardrails": [], "attempted_fallbacks": 0, "original_model_group": "my-anthropic-model-group", "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}',
"cache_key": "Cache OFF",
"spend": 0.01383,
"total_tokens": 2598,

2
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -1073,6 +1073,7 @@ async def test_key_metadata_enable_prompt_caching_promoted_to_request_root(key_v
"_code_interpreter_interception_active",
"_code_interpreter_interception_converted_stream",
"_code_interpreter_interception_sandbox_key",
"_headroom_interception_converted_stream",
"max_agentic_loops",
],
)
@ -1107,6 +1108,7 @@ async def test_add_litellm_data_to_request_strips_callback_control_fields(
"_code_interpreter_interception_active": True,
"_code_interpreter_interception_converted_stream": True,
"_code_interpreter_interception_sandbox_key": "forged-key",
"_headroom_interception_converted_stream": True,
"max_agentic_loops": 9999,
}
sample_value = sample_values[control_field]

70
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -9565,6 +9565,76 @@ class TestDeleteDeploymentSync:
assert result is None, f"Expected None on DB failure to signal fetch error, got {result!r}"
@pytest.mark.asyncio
async def test_get_models_from_db_reads_from_writer_not_replica(self):
"""
Regression for #38556: with DATABASE_URL_READ_REPLICA configured, the model
reconcile after /model/new used to read via the replica, so a lagging replica
made the reload miss the just-committed row and fail the request with a 500.
The reconcile read must be pinned to the writer.
"""
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
from litellm.proxy.proxy_server import ProxyConfig
writer_inner = MagicMock(name="writer_prisma")
reader_inner = MagicMock(name="reader_prisma")
committed_row = MagicMock(name="just_committed_model_row")
writer_inner.litellm_proxymodeltable.find_many = AsyncMock(return_value=[committed_row])
reader_inner.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
mock_prisma = MagicMock()
mock_prisma.db = RoutingPrismaWrapper(
writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False),
reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False),
)
result = await ProxyConfig()._get_models_from_db(prisma_client=mock_prisma)
assert result == [committed_row], f"Expected the writer's just-committed row, got {result!r}"
reader_inner.litellm_proxymodeltable.find_many.assert_not_awaited()
@pytest.mark.asyncio
async def test_get_models_from_db_falls_back_to_replica_when_writer_down(self):
"""
The writer pin must not break reader-only degraded mode: a proxy that
starts during a primary outage (writer connect failed, replica healthy)
must still load DB-backed models through the replica instead of sending
the reconcile read to the unavailable writer.
"""
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
from litellm.proxy.proxy_server import ProxyConfig
writer_inner = MagicMock(name="writer_prisma")
reader_inner = MagicMock(name="reader_prisma")
replica_row = MagicMock(name="replica_model_row")
writer_inner.litellm_proxymodeltable = SimpleNamespace(
find_many=AsyncMock(side_effect=RuntimeError("writer unreachable")),
create=MagicMock(name="writer_create"),
)
reader_inner.litellm_proxymodeltable = SimpleNamespace(
find_many=AsyncMock(return_value=[replica_row]),
create=MagicMock(name="reader_create"),
)
mock_prisma = MagicMock()
mock_prisma.db = RoutingPrismaWrapper(
writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False),
reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False),
)
mock_prisma.db._writer_unavailable = True
result = await ProxyConfig()._get_models_from_db(prisma_client=mock_prisma)
assert result == [replica_row], f"Expected the replica's rows in degraded mode, got {result!r}"
writer_inner.litellm_proxymodeltable.find_many.assert_not_awaited()
def test_get_config_list_includes_cancel_on_disconnect(monkeypatch):
"""Follow-up to #30223: the flag must be discoverable via /config/list,

76
tests/test_litellm/test_utils.py
View file

@ -6,6 +6,7 @@ from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import respx
from jsonschema import validate
@ -4432,6 +4433,53 @@ class TestVertexEmbeddingEncodingFormat:
assert optional_params.get("outputDimensionality") == 256
class TestBedrockCohereEmbeddingDispatch:
"""All bedrock cohere.embed models must route to BedrockCohereEmbeddingConfig,
not just multilingual-v3/v4: english-v3 was falling into the unmapped
else-branch and rejecting encoding_format. Issue #38659."""
@pytest.mark.parametrize(
"model",
[
"cohere.embed-english-v3",
"cohere.embed-multilingual-v3",
"cohere.embed-v4:0",
],
)
def test_cohere_embed_models_accept_encoding_format(self, model):
optional_params = litellm.utils.get_optional_params_embeddings(
model=model,
encoding_format="float",
custom_llm_provider="bedrock",
)
assert optional_params.get("embedding_types") == ["float"]
@pytest.mark.parametrize(
"model",
[
"cohere.embed-english-v3",
"cohere.embed-multilingual-v3",
"cohere.embed-v4:0",
],
)
def test_cohere_embed_models_map_base64_to_float(self, model):
optional_params = litellm.utils.get_optional_params_embeddings(
model=model,
encoding_format="base64",
custom_llm_provider="bedrock",
)
assert optional_params.get("embedding_types") == ["float"]
def test_cohere_embed_english_v3_maps_dimensions(self):
optional_params = litellm.utils.get_optional_params_embeddings(
model="cohere.embed-english-v3",
encoding_format="float",
dimensions=512,
custom_llm_provider="bedrock",
)
assert optional_params.get("output_dimension") == 512
@pytest.mark.parametrize(
"model",
[
@ -5689,3 +5737,31 @@ class TestDefaultReasoningEffortHydration:
model_info = dict(_get_model_info_helper(model="gpt-5.6-terra", custom_llm_provider="openai"))
assert model_info.get("default_reasoning_effort") is None
class TestHuggingFaceConfigFetch:
"""The Hugging Face config.json fetch runs on background logging threads during cost
calculation, so an unbounded request can hang a whole test job; the timeout is the fix."""
@pytest.fixture
def hf_config_route(self):
with respx.mock(assert_all_called=True) as respx_mock:
yield respx_mock.get(url__regex=r"https://huggingface\.co/.*/config\.json").respond(
json={"max_position_embeddings": 512}
)
def test_get_max_tokens_reads_hf_config_with_a_bounded_timeout(self, hf_config_route):
from litellm.constants import HF_CONFIG_FETCH_TIMEOUT_SECONDS
from litellm.utils import get_max_tokens
assert get_max_tokens("huggingface/some-org/some-model") == 512
request_timeout = hf_config_route.calls.last.request.extensions["timeout"]
assert request_timeout["read"] == HF_CONFIG_FETCH_TIMEOUT_SECONDS
def test_get_max_position_embeddings_reads_hf_config_with_a_bounded_timeout(self, hf_config_route):
from litellm.constants import HF_CONFIG_FETCH_TIMEOUT_SECONDS
from litellm.utils import _get_max_position_embeddings
assert _get_max_position_embeddings("some-org/some-model") == 512
request_timeout = hf_config_route.calls.last.request.extensions["timeout"]
assert request_timeout["read"] == HF_CONFIG_FETCH_TIMEOUT_SECONDS

4
type-discipline-budget.json
View file

@ -1,6 +1,6 @@
{
"LIT001": {
"limit": 22655
"limit": 22654
},
"LIT002": {
"limit": 26830
@ -33,6 +33,6 @@
"limit": 5558
},
"LIT012": {
"limit": 4508
"limit": 4506
}
}

2
ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.test.tsx
View file

@ -9,6 +9,8 @@ const mockAddAllowedIP = vi.fn();
const mockDeleteAllowedIP = vi.fn();
vi.mock("@/components/networking", () => ({
getProxyBaseUrl: () => "http://localhost:4000",
getGlobalLitellmHeaderName: () => "Authorization",
getSSOSettings: (...args: unknown[]) => mockGetSSOSettings(...args),
getAllowedIPs: (...args: unknown[]) => mockGetAllowedIPs(...args),
addAllowedIP: (...args: unknown[]) => mockAddAllowedIP(...args),

6
ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx
View file

@ -18,6 +18,7 @@ import LoggingSettings from "@/components/Settings/AdminSettings/LoggingSettings
import SSOSettings from "@/components/Settings/AdminSettings/SSOSettings/SSOSettings";
import UISettings from "@/components/Settings/AdminSettings/UISettings/UISettings";
import UserBannerSettings from "@/components/Settings/AdminSettings/UserBannerSettings/UserBannerSettings";
import CyberArk from "@/components/Settings/AdminSettings/CyberArk/CyberArk";
import HashicorpVault from "@/components/Settings/AdminSettings/HashicorpVault/HashicorpVault";
import PluginSettings from "@/components/Settings/AdminSettings/PluginSettings/PluginSettings";
import SSOModals from "@/components/SSOModals";
@ -395,6 +396,11 @@ const AdminPanel: React.FC<AdminPanelProps> = ({ proxySettings }) => {
label: "Hashicorp Vault",
children: <HashicorpVault />,
},
{
key: "cyberark",
label: "CyberArk Conjur",
children: <CyberArk />,
},
{
key: "plugins",
label: "Plugins",

38
ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/cyberArkApi.ts Normal file
View file

@ -0,0 +1,38 @@
import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking";
import { createApiClient } from "@/lib/http/client";
export interface CyberArkFieldSchema {
description?: string;
properties: Record<string, { description?: string; type?: string }>;
}
export interface CyberArkConfigResponse {
config_type: string;
values: Record<string, string | null>;
field_schema: CyberArkFieldSchema;
}
export interface CyberArkStatusResponse {
status: string;
message: string;
}
const apiClient = createApiClient({
getBaseUrl: getProxyBaseUrl,
getAuthHeaderName: getGlobalLitellmHeaderName,
});
export const getCyberArkConfig = async (accessToken: string): Promise<CyberArkConfigResponse> =>
apiClient.get<CyberArkConfigResponse>("/config_overrides/cyberark", { accessToken });
export const updateCyberArkConfig = async (
accessToken: string,
config: Record<string, string>,
): Promise<CyberArkStatusResponse> =>
apiClient.post<CyberArkStatusResponse>("/config_overrides/cyberark", { accessToken, body: config });
export const deleteCyberArkConfig = async (accessToken: string): Promise<CyberArkStatusResponse> =>
apiClient.delete<CyberArkStatusResponse>("/config_overrides/cyberark", { accessToken });
export const testCyberArkConnection = async (accessToken: string): Promise<CyberArkStatusResponse> =>
apiClient.post<CyberArkStatusResponse>("/config_overrides/cyberark/test_connection", { accessToken });

24
ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useCyberArkConfig.ts Normal file
View file

@ -0,0 +1,24 @@
import { getCyberArkConfig, type CyberArkConfigResponse } from "./cyberArkApi";
import { useQuery } from "@tanstack/react-query";
import useAuthorized from "../useAuthorized";
import { createQueryKeys } from "../common/queryKeysFactory";
export const cyberArkKeys = createQueryKeys("cyberArkConfig");
export const useCyberArkConfig = () => {
const { accessToken } = useAuthorized();
const queryOptions = {
queryKey: cyberArkKeys.list({}),
queryFn: async () => {
if (!accessToken) {
throw new Error("Access token is required");
}
return getCyberArkConfig(accessToken);
},
enabled: !!accessToken,
staleTime: 60 * 60 * 1000,
gcTime: 60 * 60 * 1000,
};
return useQuery<CyberArkConfigResponse>(queryOptions);
};

19
ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useDeleteCyberArkConfig.ts Normal file
View file

@ -0,0 +1,19 @@
import { deleteCyberArkConfig } from "./cyberArkApi";
import { useMutation, useQueryClient } from "@tanstack/react-query";
import { cyberArkKeys } from "./useCyberArkConfig";
export const useDeleteCyberArkConfig = (accessToken: string | null) => {
const queryClient = useQueryClient();
return useMutation({
mutationFn: async () => {
if (!accessToken) {
throw new Error("Access token is required");
}
return deleteCyberArkConfig(accessToken);
},
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: cyberArkKeys.all });
},
});
};

19
ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig.ts Normal file
View file

@ -0,0 +1,19 @@
import { updateCyberArkConfig } from "./cyberArkApi";
import { useMutation, useQueryClient } from "@tanstack/react-query";
import { cyberArkKeys } from "./useCyberArkConfig";
export const useUpdateCyberArkConfig = (accessToken: string | null) => {
const queryClient = useQueryClient();
return useMutation({
mutationFn: async (config: Record<string, string>) => {
if (!accessToken) {
throw new Error("Access token is required");
}
return updateCyberArkConfig(accessToken, config);
},
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: cyberArkKeys.all });
},
});
};

78
ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.test.tsx Normal file
View file

@ -0,0 +1,78 @@
import { screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { renderWithProviders } from "../../../../../tests/test-utils";
import CyberArk from "./CyberArk";
const mockUseAuthorized = vi.hoisted(() => vi.fn());
const mockUseCyberArkConfig = vi.hoisted(() => vi.fn());
vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
default: mockUseAuthorized,
}));
vi.mock("@/app/(dashboard)/hooks/configOverrides/useCyberArkConfig", () => ({
useCyberArkConfig: mockUseCyberArkConfig,
}));
vi.mock("@/app/(dashboard)/hooks/configOverrides/useDeleteCyberArkConfig", () => ({
useDeleteCyberArkConfig: () => ({ mutate: vi.fn(), isPending: false }),
}));
vi.mock("@/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig", () => ({
useUpdateCyberArkConfig: () => ({ mutate: vi.fn(), isPending: false }),
}));
vi.mock("./EditCyberArkModal", () => ({
default: ({ isVisible }: { isVisible: boolean }) => (isVisible ? <div>Edit CyberArk Configuration</div> : null),
}));
vi.mock("@/components/common_components/DeleteResourceModal", () => ({
default: () => null,
}));
describe("CyberArk", () => {
beforeEach(() => {
vi.clearAllMocks();
mockUseAuthorized.mockReturnValue({ accessToken: "test-token" });
const emptyConfigResult = {
data: { values: {} },
isLoading: false,
isError: false,
error: null,
};
mockUseCyberArkConfig.mockReturnValue(emptyConfigResult);
});
it("should render", () => {
renderWithProviders(<CyberArk />);
expect(screen.getByRole("heading", { name: "CyberArk Conjur" })).toBeInTheDocument();
});
it("should open the configuration editor from the empty state", async () => {
const user = userEvent.setup();
renderWithProviders(<CyberArk />);
await user.click(screen.getByRole("button", { name: /configure cyberark/i }));
expect(screen.getByText("Edit CyberArk Configuration")).toBeInTheDocument();
});
it("should display configured values and management actions", () => {
const configuredResult = {
data: { values: { cyberark_api_base: "https://conjur.example.com", cyberark_api_key: "secret" } },
isLoading: false,
isError: false,
error: null,
};
mockUseCyberArkConfig.mockReturnValue(configuredResult);
renderWithProviders(<CyberArk />);
expect(screen.getByText("https://conjur.example.com")).toBeInTheDocument();
expect(screen.getByText("Auth Method")).toBeInTheDocument();
expect(screen.getAllByText("API Key")).toHaveLength(2);
expect(screen.getByRole("button", { name: /test connection/i })).toBeInTheDocument();
});
});

232
ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.tsx Normal file
View file

@ -0,0 +1,232 @@
"use client";
import { Edit, ExternalLink, Info, KeyRound, PlugZap, Trash2 } from "lucide-react";
import { useState } from "react";
import { testCyberArkConnection } from "@/app/(dashboard)/hooks/configOverrides/cyberArkApi";
import { useCyberArkConfig } from "@/app/(dashboard)/hooks/configOverrides/useCyberArkConfig";
import { useDeleteCyberArkConfig } from "@/app/(dashboard)/hooks/configOverrides/useDeleteCyberArkConfig";
import { useUpdateCyberArkConfig } from "@/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import DeleteResourceModal from "@/components/common_components/DeleteResourceModal";
import { toast } from "@/lib/toast";
import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert";
import { Button } from "@/components/ui/button";
import { Card, CardAction, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card";
import { Skeleton } from "@/components/ui/skeleton";
import CyberArkEmptyPlaceholder from "./CyberArkEmptyPlaceholder";
import EditCyberArkModal from "./EditCyberArkModal";
import { FIELD_LABELS, SENSITIVE_FIELDS } from "./constants";
function detectAuthMethod(values: Record<string, unknown>): string {
if (values.cyberark_api_key) return "API Key";
if (values.client_cert && values.client_key) return "TLS Certificate";
return "None";
}
function DetailRow({ children, label }: { children: React.ReactNode; label: string }) {
return (
<div className="grid grid-cols-1 sm:grid-cols-3">
<dt className="bg-muted/50 px-4 py-3 text-sm font-medium text-foreground">{label}</dt>
<dd className="px-4 py-3 text-sm text-foreground sm:col-span-2">{children}</dd>
</div>
);
}
export default function CyberArk() {
const { accessToken } = useAuthorized();
const { data, isLoading, isError, error } = useCyberArkConfig();
const { mutate: deleteConfig, isPending: isDeleting } = useDeleteCyberArkConfig(accessToken);
const { mutate: updateConfig, isPending: isClearingField } = useUpdateCyberArkConfig(accessToken);
const [isEditModalVisible, setIsEditModalVisible] = useState(false);
const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
const [clearingField, setClearingField] = useState<string | null>(null);
const [isTesting, setIsTesting] = useState(false);
const rawValues = data?.values ?? {};
const isConfigured = Boolean(rawValues.cyberark_api_base);
const handleTestConnection = async () => {
if (!accessToken) return;
setIsTesting(true);
try {
const result = await testCyberArkConnection(accessToken);
toast.success(result.message || "Connection to CyberArk Conjur successful!");
} catch (err) {
toast.fromError(err);
} finally {
setIsTesting(false);
}
};
const handleDelete = () => {
deleteConfig(undefined, {
onSuccess: () => {
toast.success("CyberArk configuration deleted");
setIsDeleteModalOpen(false);
},
onError: (err) => toast.fromError(err),
});
};
const handleClearField = () => {
if (!clearingField) return;
updateConfig(
{ [clearingField]: "" },
{
onSuccess: () => {
toast.success(`${FIELD_LABELS[clearingField] ?? clearingField} cleared`);
setClearingField(null);
},
onError: (err) => toast.fromError(err),
},
);
};
const renderValue = (key: string) => {
const value = rawValues[key];
if (!value) return <span className="text-muted-foreground italic">Not configured</span>;
if (!SENSITIVE_FIELDS.has(key)) return <span className="font-mono text-muted-foreground">{value}</span>;
return (
<div className="flex items-center justify-between gap-2">
<span className="font-mono text-muted-foreground">{value}</span>
<Button
type="button"
variant="ghost"
size="icon-sm"
aria-label={`Clear ${FIELD_LABELS[key] ?? key}`}
onClick={() => setClearingField(key)}
>
<Trash2 className="size-3.5" />
</Button>
</div>
);
};
const fieldsToShow = Object.entries(rawValues).filter(([, value]) => value != null && value !== "");
const renderCard = () => {
if (isLoading) {
return (
<Card role="status" aria-label="Loading CyberArk configuration">
<CardContent className="space-y-3">
<Skeleton className="h-8 w-64" />
<Skeleton className="h-40 w-full" />
</CardContent>
</Card>
);
}
if (isError) {
return (
<Card>
<CardContent>
<Alert variant="error">
<AlertTitle>Could not load CyberArk configuration</AlertTitle>
{error instanceof Error && <AlertDescription>{error.message}</AlertDescription>}
</Alert>
</CardContent>
</Card>
);
}
return (
<Card>
<CardHeader>
<div className="flex items-center gap-3">
<KeyRound className="size-6 text-muted-foreground" />
<div>
<CardTitle>
<h3>CyberArk Conjur</h3>
</CardTitle>
<CardDescription>Manage secret manager configuration</CardDescription>
</div>
</div>
{isConfigured && (
<CardAction className="flex flex-wrap gap-2">
<Button type="button" variant="outline" disabled={isTesting} onClick={handleTestConnection}>
<PlugZap />
{isTesting ? "Testing..." : "Test Connection"}
</Button>
<Button type="button" variant="outline" onClick={() => setIsEditModalVisible(true)}>
<Edit />
Edit Configuration
</Button>
<Button type="button" variant="destructive" onClick={() => setIsDeleteModalOpen(true)}>
<Trash2 />
Delete Configuration
</Button>
</CardAction>
)}
</CardHeader>
<CardContent className="space-y-6">
{isConfigured && (
<Alert variant="info">
<Info />
<AlertTitle>Configuration changes are hot-reloaded across all proxy instances</AlertTitle>
<AlertDescription>
<a
href="https://docs.litellm.ai/docs/secret_managers/cyberark"
target="_blank"
rel="noreferrer"
className="inline-flex items-center gap-1"
>
View documentation
<ExternalLink className="size-3" />
</a>
</AlertDescription>
</Alert>
)}
{isConfigured ? (
fieldsToShow.length > 0 && (
<dl className="divide-y divide-border overflow-hidden rounded-md border border-border">
<DetailRow label="Auth Method">{detectAuthMethod(rawValues)}</DetailRow>
{fieldsToShow.map(([key]) => (
<DetailRow key={key} label={FIELD_LABELS[key] ?? key}>
{renderValue(key)}
</DetailRow>
))}
</dl>
)
) : (
<CyberArkEmptyPlaceholder onAdd={() => setIsEditModalVisible(true)} />
)}
</CardContent>
</Card>
);
};
return (
<>
{renderCard()}
<EditCyberArkModal
isVisible={isEditModalVisible}
onCancel={() => setIsEditModalVisible(false)}
onSuccess={() => setIsEditModalVisible(false)}
/>
<DeleteResourceModal
isOpen={isDeleteModalOpen}
title="Delete CyberArk Configuration?"
message="Models using CyberArk secrets will lose access to their API keys until a new configuration is saved."
resourceInformationTitle="CyberArk Configuration"
resourceInformation={[{ label: "Conjur Server URL", value: rawValues.cyberark_api_base }]}
onCancel={() => setIsDeleteModalOpen(false)}
onOk={handleDelete}
confirmLoading={isDeleting}
/>
<DeleteResourceModal
isOpen={clearingField !== null}
title={`Clear ${clearingField ? FIELD_LABELS[clearingField] ?? clearingField : ""}?`}
message="This will remove the stored value."
resourceInformationTitle="Field"
resourceInformation={[
{ label: "Field", value: clearingField ? FIELD_LABELS[clearingField] ?? clearingField : "" },
]}
onCancel={() => setClearingField(null)}
onOk={handleClearField}
confirmLoading={isClearingField}
/>
</>
);
}

24
ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArkEmptyPlaceholder.tsx Normal file
View file

@ -0,0 +1,24 @@
import { KeyRound } from "lucide-react";
import { Button } from "@/components/ui/button";
interface CyberArkEmptyPlaceholderProps {
onAdd: () => void;
}
export default function CyberArkEmptyPlaceholder({ onAdd }: CyberArkEmptyPlaceholderProps) {
return (
<div className="flex w-full flex-col items-center rounded-lg border border-dashed border-border bg-card p-12 text-center">
<div className="mb-4 flex size-12 items-center justify-center rounded-full bg-muted">
<KeyRound className="size-6 text-muted-foreground" />
</div>
<h4 className="text-base font-semibold text-foreground">No CyberArk Configuration Found</h4>
<p className="mx-auto mt-2 max-w-md text-sm text-muted-foreground">
Configure CyberArk Conjur to securely manage provider API keys and secrets for your LiteLLM deployment.
</p>
<Button size="lg" onClick={onAdd} className="mt-4">
Configure CyberArk
</Button>
</div>
);
}

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