diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index a69e50b5753..7e013b7bb0b 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -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: diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index ef88ae574fb..95e69bf7f1d 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 5486 + "limit": 5485 }, "reportFunctionMemberAccess": { "limit": 7 diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index aee3295d1da..3b09dc9272e 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -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, Any, Dict, Final, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Tuple, cast from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -626,6 +627,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, ) @@ -761,16 +763,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": ""}], stream=False, call_type="aretrieve_batch", @@ -802,9 +821,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) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 2bc61aed771..3831f57a10d 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -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: diff --git a/litellm/constants.py b/litellm/constants.py index e51c0a8827d..c1f087316cc 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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. diff --git a/litellm/litellm_core_utils/audio_utils/subtitle_utils.py b/litellm/litellm_core_utils/audio_utils/subtitle_utils.py index 615873e295d..91427ba09ad 100644 --- a/litellm/litellm_core_utils/audio_utils/subtitle_utils.py +++ b/litellm/litellm_core_utils/audio_utils/subtitle_utils.py @@ -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: diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 07bed1f88ad..c8e9e2583ba 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -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 diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a8672b5c112..ff9a5d41c2d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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, ) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 681a8397f66..d8d6a7fc9f8 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1099,6 +1099,25 @@ def is_empty_thinking_block(block: object) -> bool: return not isinstance(thinking, str) or not thinking.strip() +def is_empty_unsigned_thinking_block(block: object) -> bool: + """ + True for an empty ``{"type": "thinking"}`` block carrying no signature. + + The emit-side predicate: response paths drop a thinking block only when it + holds nothing the client could need. A signature-only block is a real + provider response (Bedrock Converse under adaptive thinking emits a + reasoning block with empty text and only a signature) and the client needs + the signature to replay reasoning across tool-use turns, so it must be + emitted. Request paths keep using :func:`is_empty_thinking_block`: + Anthropic rejects empty thinking blocks in request history regardless of + signature, and the inbound strip self-heals a replayed signature-only + block. + """ + if not isinstance(block, dict) or not is_empty_thinking_block(block): + return False + return not block.get("signature") + + def normalize_anthropic_tool_use_id(raw_id: str) -> str: """ Normalize a tool_use / tool_result id for Anthropic's ``^[a-zA-Z0-9_-]+$`` diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index cefd4aa2d77..cc5879df56d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -1029,7 +1029,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): @staticmethod def _is_blank_delta(chunk: "ModelResponseStream") -> bool: - from litellm.llms.anthropic.common_utils import is_empty_thinking_block + from litellm.llms.anthropic.common_utils import is_empty_unsigned_thinking_block choice: Final = chunk.choices[0] if choice.finish_reason is not None: @@ -1041,11 +1041,14 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): return False if getattr(delta, "reasoning_content", None): return False - # thinking_blocks whose entries are all empty (even if signed) must not + # thinking_blocks whose entries are all empty AND unsigned must not # open a block: the emitted {"type": "thinking", "thinking": ""} gets - # replayed as history and Anthropic rejects it (LIT-6357). + # replayed as history and Anthropic rejects it (LIT-6357). A signed + # entry opens the block so the client receives the replay signature. thinking_blocks: Final = getattr(delta, "thinking_blocks", None) - if thinking_blocks and any(isinstance(b, dict) and not is_empty_thinking_block(b) for b in thinking_blocks): + if thinking_blocks and any( + isinstance(b, dict) and not is_empty_unsigned_thinking_block(b) for b in thinking_blocks + ): return False return True diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index a9fa00c827a..411df267442 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -90,7 +90,7 @@ from litellm.litellm_core_utils.reasoning_effort_utils import ( reasoning_effort_from_thinking_budget, ) from litellm.llms.anthropic.common_utils import ( - is_empty_thinking_block, + is_empty_unsigned_thinking_block, normalize_anthropic_tool_use_id, ) from litellm.llms.anthropic.experimental_pass_through.context_management import ( @@ -1267,7 +1267,7 @@ class LiteLLMAnthropicMessagesAdapter: if hasattr(choice.message, "thinking_blocks") and choice.message.thinking_blocks: for thinking_block in choice.message.thinking_blocks: if thinking_block.get("type") == "thinking": - if is_empty_thinking_block(thinking_block): + if is_empty_unsigned_thinking_block(thinking_block): continue thinking_value = thinking_block.get("thinking", "") signature_value = thinking_block.get("signature", "") diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 980b27cda55..2bcc830851a 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -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): diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 2c34851d275..6cb7d09cec4 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -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: diff --git a/litellm/llms/azure/image_edit/transformation.py b/litellm/llms/azure/image_edit/transformation.py index 15592968bad..e4716289a34 100644 --- a/litellm/llms/azure/image_edit/transformation.py +++ b/litellm/llms/azure/image_edit/transformation.py @@ -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) diff --git a/litellm/llms/azure/image_generation/http_utils.py b/litellm/llms/azure/image_generation/http_utils.py index 03c425eeffc..1aa5757ca95 100644 --- a/litellm/llms/azure/image_generation/http_utils.py +++ b/litellm/llms/azure/image_generation/http_utils.py @@ -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} diff --git a/litellm/llms/bedrock/embed/cohere_transformation.py b/litellm/llms/bedrock/embed/cohere_transformation.py index d1c9ceb99d1..8a17bb9d595 100644 --- a/litellm/llms/bedrock/embed/cohere_transformation.py +++ b/litellm/llms/bedrock/embed/cohere_transformation.py @@ -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 diff --git a/litellm/llms/soniox/common_utils.py b/litellm/llms/soniox/common_utils.py index cb33f0b8996..8d83b4c9218 100644 --- a/litellm/llms/soniox/common_utils.py +++ b/litellm/llms/soniox/common_utils.py @@ -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)) diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index 25db3a673a1..e6e3c2739c1 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -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, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bebbcc32181..ec175025b42 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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", diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 93b85edd88d..eff467072a8 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -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": [], diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 853a07972c1..a43e762a456 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -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) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py index 6824f96f927..0fa750a4c4a 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py @@ -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, ) diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index c435234cbbc..3f90e6c0a7a 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -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"), diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index c1e89f8aa75..08a35b41126 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -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", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f40b0632398..89f79548845 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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 diff --git a/litellm/proxy/client/chat.py b/litellm/proxy/client/chat.py index 2953ed7f683..bd4d0df3ed0 100644 --- a/litellm/proxy/client/chat.py +++ b/litellm/proxy/client/chat.py @@ -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 diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 550b11311f5..2fad9f933c1 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -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") diff --git a/litellm/proxy/client/client.py b/litellm/proxy/client/client.py index d71802e06c8..de1e45b91be 100644 --- a/litellm/proxy/client/client.py +++ b/litellm/proxy/client/client.py @@ -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) diff --git a/litellm/proxy/client/credentials.py b/litellm/proxy/client/credentials.py index 136bdf3f293..a9bff67b1c5 100644 --- a/litellm/proxy/client/credentials.py +++ b/litellm/proxy/client/credentials.py @@ -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: diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index 5b66567363d..fe100c5f676 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -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: diff --git a/litellm/proxy/client/model_groups.py b/litellm/proxy/client/model_groups.py index 9c7c38dc67c..fef307600c4 100644 --- a/litellm/proxy/client/model_groups.py +++ b/litellm/proxy/client/model_groups.py @@ -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: diff --git a/litellm/proxy/client/models.py b/litellm/proxy/client/models.py index 0f1dd2b5bab..4b16087e15b 100644 --- a/litellm/proxy/client/models.py +++ b/litellm/proxy/client/models.py @@ -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: diff --git a/litellm/proxy/client/teams.py b/litellm/proxy/client/teams.py index ef2ac53f9c4..105060e5ca9 100644 --- a/litellm/proxy/client/teams.py +++ b/litellm/proxy/client/teams.py @@ -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.") diff --git a/litellm/proxy/client/users.py b/litellm/proxy/client/users.py index df5f9aad23e..3f11fe94043 100644 --- a/litellm/proxy/client/users.py +++ b/litellm/proxy/client/users.py @@ -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() diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index b9cc90107e2..e2d2fffb2df 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py index c595952bec2..f7ea1cb632f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py @@ -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( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b7688790e54..222ba45b059 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -253,6 +253,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 diff --git a/litellm/proxy/management_endpoints/config_override_endpoints.py b/litellm/proxy/management_endpoints/config_override_endpoints.py index 7f4faddf178..84593460704 100644 --- a/litellm/proxy/management_endpoints/config_override_endpoints.py +++ b/litellm/proxy/management_endpoints/config_override_endpoints.py @@ -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}", + } diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 094877e1ef8..6ed3c86e2da 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4372,6 +4372,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 @@ -7068,6 +7070,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) @@ -7232,6 +7235,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. diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index a6b6375fd9c..ff182bd98ec 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -92,6 +92,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, @@ -121,6 +123,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, @@ -154,6 +158,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 @@ -360,6 +366,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 diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index 2cca16351af..9a714e1724e 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -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( ( diff --git a/litellm/types/proxy/management_endpoints/config_overrides.py b/litellm/types/proxy/management_endpoints/config_overrides.py index 9e1ea23ac46..f9cba6983db 100644 --- a/litellm/types/proxy/management_endpoints/config_overrides.py +++ b/litellm/types/proxy/management_endpoints/config_overrides.py @@ -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.""" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index f0319a7c664..2d06f5e0abe 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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 diff --git a/litellm/types/videos/main.py b/litellm/types/videos/main.py index 3677cec3c8f..99b08f6caf6 100644 --- a/litellm/types/videos/main.py +++ b/litellm/types/videos/main.py @@ -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 diff --git a/litellm/utils.py b/litellm/utils.py index fa2226dbf2c..bb705e4ae9f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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 diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bebbcc32181..ec175025b42 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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", diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index c60988eccc0..b4107582000 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 655 + "limit": 654 }, "ASYNC230": { "limit": 11 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1117 + "limit": 1116 }, "TRY002": { "limit": 524 diff --git a/ruff.toml b/ruff.toml index 44bdf9d8125..3ac4c1fc94d 100644 --- a/ruff.toml +++ b/ruff.toml @@ -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 diff --git a/tests/batches_tests/test_batch_custom_pricing.py b/tests/batches_tests/test_batch_custom_pricing.py index c2159b564a8..b76b865862a 100644 --- a/tests/batches_tests/test_batch_custom_pricing.py +++ b/tests/batches_tests/test_batch_custom_pricing.py @@ -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 diff --git a/tests/batches_tests/test_batch_rate_limits.py b/tests/batches_tests/test_batch_rate_limits.py index b44b8435cd9..7cbcfc1aeb1 100644 --- a/tests/batches_tests/test_batch_rate_limits.py +++ b/tests/batches_tests/test_batch_rate_limits.py @@ -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, diff --git a/tests/batches_tests/test_batches_logging_unit_tests.py b/tests/batches_tests/test_batches_logging_unit_tests.py index 5211b3ecb29..5bde40d90b0 100644 --- a/tests/batches_tests/test_batches_logging_unit_tests.py +++ b/tests/batches_tests/test_batches_logging_unit_tests.py @@ -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 diff --git a/tests/e2e/ui/tests/logs/logs.spec.ts b/tests/e2e/ui/tests/logs/logs.spec.ts index b29a72cbf81..56a3d0f0109 100644 --- a/tests/e2e/ui/tests/logs/logs.spec.ts +++ b/tests/e2e/ui/tests/logs/logs.spec.ts @@ -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, { diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json index 5789f19aa55..1838fb16e91 100644 --- a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json +++ b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json @@ -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, diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index b4fe9347581..ff5e8f89d64 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -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", diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 41b4bb8cf76..c86c7c4df03 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -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 # =========================================================================== # diff --git a/tests/test_litellm/batches/test_responses_batch_cost.py b/tests/test_litellm/batches/test_responses_batch_cost.py index 7ce026bd103..b634f5f73db 100644 --- a/tests/test_litellm/batches/test_responses_batch_cost.py +++ b/tests/test_litellm/batches/test_responses_batch_cost.py @@ -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"] ) diff --git a/tests/test_litellm/litellm_core_utils/audio_utils/test_subtitle_utils.py b/tests/test_litellm/litellm_core_utils/audio_utils/test_subtitle_utils.py index dcc4163ff10..210513f4967 100644 --- a/tests/test_litellm/litellm_core_utils/audio_utils/test_subtitle_utils.py +++ b/tests/test_litellm/litellm_core_utils/audio_utils/test_subtitle_utils.py @@ -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), diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 1193160c831..1c717228ecc 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -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) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index ea1813acb82..2d74c00071b 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1013,12 +1013,15 @@ def test_translate_openai_content_to_anthropic_thinking_and_redacted_thinking(): assert result[1]["data"] == "REDACTED" -def test_translate_openai_content_to_anthropic_drops_empty_thinking_blocks(): - """LIT-6357 non-streaming producer half: a bridged reasoning model whose - thinking_blocks entry has empty or whitespace-only text (signed or not) - must not surface as {"type": "thinking", "thinking": ""} — clients replay - it as history and Anthropic 400s with "each thinking block must contain - thinking". Non-empty thinking and redacted_thinking pass through.""" +def test_translate_openai_content_to_anthropic_drops_empty_unsigned_thinking_blocks(): + """LIT-6357 non-streaming producer half, narrowed to unsigned blocks: a + bridged reasoning model whose thinking_blocks entry has empty or + whitespace-only text and no signature must not surface as + {"type": "thinking", "thinking": ""}. A signature-only block (Bedrock + Converse adaptive thinking) must be emitted so the client keeps the + signature for tool-use replay; the inbound strip self-heals it if the + client loops it back. Non-empty thinking and redacted_thinking pass + through.""" openai_choices = [ Choices( message=Message( @@ -1037,9 +1040,11 @@ def test_translate_openai_content_to_anthropic_drops_empty_thinking_blocks(): adapter = LiteLLMAnthropicMessagesAdapter() result = adapter._translate_openai_content_to_anthropic(choices=openai_choices) - assert [b["type"] for b in result] == ["thinking", "redacted_thinking", "text"] - assert result[0]["thinking"] == "real plan" - assert result[1]["data"] == "REDACTED" + assert [b["type"] for b in result] == ["thinking", "thinking", "redacted_thinking", "text"] + assert result[0]["thinking"] == "" + assert result[0]["signature"] == "sig_abc" + assert result[1]["thinking"] == "real plan" + assert result[2]["data"] == "REDACTED" def test_translate_streaming_openai_chunk_to_anthropic_thinking_delta(): diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py index 6268cd01efe..17d42f55ae0 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py @@ -1048,19 +1048,20 @@ def _empty_thinking_then_tool_chunks(thinking: str = "", signature: str = "") -> @pytest.mark.parametrize("is_async", [False, True]) @pytest.mark.parametrize( "thinking,signature", - [("", ""), (" \n\t ", ""), ("", "sig_abc")], - ids=["empty", "whitespace-only", "empty-but-signed"], + [("", ""), (" \n\t ", "")], + ids=["empty", "whitespace-only"], ) @pytest.mark.asyncio async def test_contentless_thinking_chunk_opens_no_thinking_block(is_async: bool, thinking: str, signature: str): """LIT-6357 producer half: a reasoning model that goes straight to tool - calls streams a ``thinking_blocks`` entry with no real thinking text; the - wrapper used to open ``{"type": "thinking", "thinking": ""}`` for it and - close the block with no delta. Clients (Claude Code) replay that block as - history and Anthropic rejects the next tool-loop request with - "each thinking block must contain thinking" — empty-but-signed included. - The contentless chunk must open nothing; the tool_use block must be - unaffected.""" + calls streams a ``thinking_blocks`` entry with no real thinking text and + no signature; the wrapper used to open ``{"type": "thinking", + "thinking": ""}`` for it and close the block with no delta. Clients + (Claude Code) replay that block as history and Anthropic rejects the next + tool-loop request with "each thinking block must contain thinking". + The contentless unsigned chunk must open nothing; the tool_use block must + be unaffected. A SIGNED contentless chunk is different: see + test_signature_only_thinking_chunk_opens_signed_block.""" chunks = _empty_thinking_then_tool_chunks(thinking, signature) if is_async: wrapper = AnthropicStreamWrapper(completion_stream=_AsyncStream(chunks), model="claude-x") @@ -1138,11 +1139,38 @@ async def test_early_signature_on_blank_thinking_chunk_is_carried_to_the_opened_ @pytest.mark.parametrize("is_async", [False, True]) @pytest.mark.asyncio -async def test_early_signature_discarded_when_first_block_is_not_thinking(is_async: bool): - """An early signature from a skipped blank thinking chunk must not leak - into a text or tool_use first block, and must not resurrect an empty - thinking block on its own (an empty-but-signed block is exactly what - Anthropic rejects).""" +async def test_signature_only_thinking_chunk_opens_signed_block(is_async: bool): + """Bedrock Converse under adaptive thinking emits a reasoning delta with + empty text and only a signature. The signed chunk must open a thinking + block that carries the signature to the client (needed to replay reasoning + across tool-use turns); the tool_use block must be unaffected. Dropping it + like the unsigned case regressed the claude_code thinking e2e cells.""" + chunks = _empty_thinking_then_tool_chunks("", "sig_bedrock") + if is_async: + wrapper = AnthropicStreamWrapper(completion_stream=_AsyncStream(chunks), model="claude-x") + events = await _drain_async(wrapper) + else: + wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="claude-x") + events = _drain_sync(wrapper) + + starts = _thinking_block_starts(events) + assert len(starts) == 1 + assert starts[0].get("signature") == "sig_bedrock" or _signature_deltas(events) == ["sig_bedrock"] + assert _thinking_deltas(events) == [] + tool_starts = [ + e["content_block"] + for e in events + if e.get("type") == "content_block_start" and e["content_block"].get("type") == "tool_use" + ] + assert [b["name"] for b in tool_starts] == ["get_weather"] + _assert_deltas_match_their_block_type(events) + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_signature_only_thinking_chunk_before_text_leaks_no_signature(is_async: bool): + """The signed thinking block a signature-only chunk opens must stay its + own block: the text block that follows carries no signature.""" chunks = [ _thinking_chunk("", signature="sig_early"), _make_chunk(Delta(content="Hello")), @@ -1155,7 +1183,9 @@ async def test_early_signature_discarded_when_first_block_is_not_thinking(is_asy wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="claude-x") events = _drain_sync(wrapper) - assert _thinking_block_starts(events) == [] + starts = _thinking_block_starts(events) + assert len(starts) == 1 + assert starts[0].get("signature") == "sig_early" or _signature_deltas(events) == ["sig_early"] text_starts = [ e["content_block"] for e in events diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index a2da2cccb7c..794613942a1 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -1364,6 +1364,23 @@ class TestAnthropicThinkingSignatureSelfHeal: assert is_empty_thinking_block({"type": "text", "text": ""}) is False assert is_empty_thinking_block("not a dict") is False + def test_is_empty_unsigned_thinking_block(self): + """Emit-side predicate: a signature-only block must be kept (Bedrock + Converse adaptive thinking emits empty text with only a signature, and + the client needs it to replay reasoning in tool-use turns); only an + empty block with nothing to preserve is droppable.""" + from litellm.llms.anthropic.common_utils import is_empty_unsigned_thinking_block + + assert is_empty_unsigned_thinking_block({"type": "thinking", "thinking": ""}) is True + assert is_empty_unsigned_thinking_block({"type": "thinking", "thinking": " \n\t "}) is True + assert is_empty_unsigned_thinking_block({"type": "thinking"}) is True + assert is_empty_unsigned_thinking_block({"type": "thinking", "thinking": "", "signature": ""}) is True + assert is_empty_unsigned_thinking_block({"type": "thinking", "thinking": "", "signature": "sig_abc"}) is False + assert is_empty_unsigned_thinking_block({"type": "thinking", "thinking": " ", "signature": "sig_abc"}) is False + assert is_empty_unsigned_thinking_block({"type": "thinking", "thinking": "plan"}) is False + assert is_empty_unsigned_thinking_block({"type": "redacted_thinking", "data": "opaque"}) is False + assert is_empty_unsigned_thinking_block("not a dict") is False + def test_strip_empty_content_blocks_drops_empty_thinking_blocks(self): """LIT-6357 ingestion half: an assistant tool-loop turn carrying an empty (even signed) thinking block keeps its tool_use blocks and loses diff --git a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py b/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py index 59472d1a49d..5c7b249ae72 100644 --- a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py @@ -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"} diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py index 560fee17328..70b5eab5c37 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py @@ -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 diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 65ae719d021..08d01127eba 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -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"]), ], ) diff --git a/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py b/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py index f1226311b5e..0384fb796d9 100644 --- a/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py +++ b/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py @@ -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" ) diff --git a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py b/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py index 261efcb7b24..eadf870cb61 100644 --- a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py +++ b/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py @@ -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 diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py index 862969abbc2..6ba8706b0d8 100644 --- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -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.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 0c809940b84..3279c59acd4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -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" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 761f823076b..32a3f70c357 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -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") diff --git a/tests/test_litellm/proxy/auth/test_cli_auth.py b/tests/test_litellm/proxy/auth/test_cli_auth.py index c9b31a1d776..5cde5522376 100644 --- a/tests/test_litellm/proxy/auth/test_cli_auth.py +++ b/tests/test_litellm/proxy/auth/test_cli_auth.py @@ -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)]) diff --git a/tests/test_litellm/proxy/client/conftest.py b/tests/test_litellm/proxy/client/conftest.py new file mode 100644 index 00000000000..c8b7951e284 --- /dev/null +++ b/tests/test_litellm/proxy/client/conftest.py @@ -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) diff --git a/tests/test_litellm/proxy/client/test_chat.py b/tests/test_litellm/proxy/client/test_chat.py index b8e55c45502..67b6ee833f2 100644 --- a/tests/test_litellm/proxy/client/test_chat.py +++ b/tests/test_litellm/proxy/client/test_chat.py @@ -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 diff --git a/tests/test_litellm/proxy/client/test_client.py b/tests/test_litellm/proxy/client/test_client.py index fe3e2c52ce5..87eb3400b8c 100644 --- a/tests/test_litellm/proxy/client/test_client.py +++ b/tests/test_litellm/proxy/client/test_client.py @@ -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(): diff --git a/tests/test_litellm/proxy/client/test_credentials.py b/tests/test_litellm/proxy/client/test_credentials.py index 41886e3b292..666c5dac2b0 100644 --- a/tests/test_litellm/proxy/client/test_credentials.py +++ b/tests/test_litellm/proxy/client/test_credentials.py @@ -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 diff --git a/tests/test_litellm/proxy/client/test_keys.py b/tests/test_litellm/proxy/client/test_keys.py index 282b97b1c09..b9b07bddf1f 100644 --- a/tests/test_litellm/proxy/client/test_keys.py +++ b/tests/test_litellm/proxy/client/test_keys.py @@ -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 diff --git a/tests/test_litellm/proxy/client/test_model_groups.py b/tests/test_litellm/proxy/client/test_model_groups.py index 9ea8e94ff95..4a513a127b8 100644 --- a/tests/test_litellm/proxy/client/test_model_groups.py +++ b/tests/test_litellm/proxy/client/test_model_groups.py @@ -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 diff --git a/tests/test_litellm/proxy/client/test_models.py b/tests/test_litellm/proxy/client/test_models.py index fe053ffd683..9aa5a6cf0b3 100644 --- a/tests/test_litellm/proxy/client/test_models.py +++ b/tests/test_litellm/proxy/client/test_models.py @@ -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 diff --git a/tests/test_litellm/proxy/client/test_teams.py b/tests/test_litellm/proxy/client/test_teams.py new file mode 100644 index 00000000000..b61091ca44b --- /dev/null +++ b/tests/test_litellm/proxy/client/test_teams.py @@ -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 diff --git a/tests/test_litellm/proxy/client/test_users.py b/tests/test_litellm/proxy/client/test_users.py index 87b8392e402..5b4d89420ab 100644 --- a/tests/test_litellm/proxy/client/test_users.py +++ b/tests/test_litellm/proxy/client/test_users.py @@ -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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 7a2772ce78c..1fbc975e40a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -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. # diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py index 1b2108c837d..b140082a3bf 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_config_override_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_config_override_endpoints.py index d90d589c504..03f94fbe94c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_config_override_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_config_override_endpoints.py @@ -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 ───────────────── diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py index ac79c183ca3..1d2d7d4d5c3 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py @@ -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" diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 23eb9434585..10c3e5fecf8 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -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, diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 9fb31d2a6db..649bcbb6511 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -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] diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 1ff50bd0116..6524353aa48 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -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 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index f3f1a7defe7..7365cec4fdd 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22705 + "limit": 22704 }, "LIT002": { "limit": 26854 @@ -33,6 +33,6 @@ "limit": 5577 }, "LIT012": { - "limit": 4508 + "limit": 4506 } } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.test.tsx index 8c3ca7bd9ff..ee55c568bac 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.test.tsx @@ -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), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx index a98eb50ce77..1f35f46dcd4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx @@ -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 = ({ proxySettings }) => { label: "Hashicorp Vault", children: , }, + { + key: "cyberark", + label: "CyberArk Conjur", + children: , + }, { key: "plugins", label: "Plugins", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/cyberArkApi.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/cyberArkApi.ts new file mode 100644 index 00000000000..910fe2b17e5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/cyberArkApi.ts @@ -0,0 +1,38 @@ +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; +import { createApiClient } from "@/lib/http/client"; + +export interface CyberArkFieldSchema { + description?: string; + properties: Record; +} + +export interface CyberArkConfigResponse { + config_type: string; + values: Record; + field_schema: CyberArkFieldSchema; +} + +export interface CyberArkStatusResponse { + status: string; + message: string; +} + +const apiClient = createApiClient({ + getBaseUrl: getProxyBaseUrl, + getAuthHeaderName: getGlobalLitellmHeaderName, +}); + +export const getCyberArkConfig = async (accessToken: string): Promise => + apiClient.get("/config_overrides/cyberark", { accessToken }); + +export const updateCyberArkConfig = async ( + accessToken: string, + config: Record, +): Promise => + apiClient.post("/config_overrides/cyberark", { accessToken, body: config }); + +export const deleteCyberArkConfig = async (accessToken: string): Promise => + apiClient.delete("/config_overrides/cyberark", { accessToken }); + +export const testCyberArkConnection = async (accessToken: string): Promise => + apiClient.post("/config_overrides/cyberark/test_connection", { accessToken }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useCyberArkConfig.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useCyberArkConfig.ts new file mode 100644 index 00000000000..cfc3acdfe26 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useCyberArkConfig.ts @@ -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(queryOptions); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useDeleteCyberArkConfig.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useDeleteCyberArkConfig.ts new file mode 100644 index 00000000000..cebba3a202d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useDeleteCyberArkConfig.ts @@ -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 }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig.ts new file mode 100644 index 00000000000..f5c88e833f7 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig.ts @@ -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) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return updateCyberArkConfig(accessToken, config); + }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: cyberArkKeys.all }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.test.tsx new file mode 100644 index 00000000000..d9bd9b1ebc8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.test.tsx @@ -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 ?
Edit CyberArk Configuration
: 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(); + + expect(screen.getByRole("heading", { name: "CyberArk Conjur" })).toBeInTheDocument(); + }); + + it("should open the configuration editor from the empty state", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + 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(); + + 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(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.tsx new file mode 100644 index 00000000000..30026d0a834 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.tsx @@ -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 { + 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 ( +
+
{label}
+
{children}
+
+ ); +} + +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(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 Not configured; + if (!SENSITIVE_FIELDS.has(key)) return {value}; + + return ( +
+ {value} + +
+ ); + }; + + const fieldsToShow = Object.entries(rawValues).filter(([, value]) => value != null && value !== ""); + + const renderCard = () => { + if (isLoading) { + return ( + + + + + + + ); + } + if (isError) { + return ( + + + + Could not load CyberArk configuration + {error instanceof Error && {error.message}} + + + + ); + } + return ( + + +
+ +
+ +

CyberArk Conjur

+ + Manage secret manager configuration +
+
+ {isConfigured && ( + + + + + + )} + + + {isConfigured && ( + + + Configuration changes are hot-reloaded across all proxy instances + + + View documentation + + + + + )} + + {isConfigured ? ( + fieldsToShow.length > 0 && ( +
+ {detectAuthMethod(rawValues)} + {fieldsToShow.map(([key]) => ( + + {renderValue(key)} + + ))} +
+ ) + ) : ( + setIsEditModalVisible(true)} /> + )} + + + ); + }; + + return ( + <> + {renderCard()} + + setIsEditModalVisible(false)} + onSuccess={() => setIsEditModalVisible(false)} + /> + setIsDeleteModalOpen(false)} + onOk={handleDelete} + confirmLoading={isDeleting} + /> + setClearingField(null)} + onOk={handleClearField} + confirmLoading={isClearingField} + /> + + ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArkEmptyPlaceholder.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArkEmptyPlaceholder.tsx new file mode 100644 index 00000000000..513d34d0f65 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArkEmptyPlaceholder.tsx @@ -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 ( +
+
+ +
+

No CyberArk Configuration Found

+

+ Configure CyberArk Conjur to securely manage provider API keys and secrets for your LiteLLM deployment. +

+ +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/EditCyberArkModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/EditCyberArkModal.test.tsx new file mode 100644 index 00000000000..3f716451173 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/EditCyberArkModal.test.tsx @@ -0,0 +1,176 @@ +import { fireEvent, screen, waitFor } 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 EditCyberArkModal from "./EditCyberArkModal"; +import { useCyberArkConfig } from "@/app/(dashboard)/hooks/configOverrides/useCyberArkConfig"; +import { useUpdateCyberArkConfig } from "@/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig"; + +vi.mock("@/app/(dashboard)/hooks/configOverrides/useCyberArkConfig", () => ({ + useCyberArkConfig: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig", () => ({ + useUpdateCyberArkConfig: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ accessToken: "sk-access-token" }), +})); + +vi.mock("@/lib/toast", () => ({ + toast: { success: vi.fn(), fromError: vi.fn() }, +})); + +const ALL_FIELDS = [ + "cyberark_api_base", + "cyberark_account", + "cyberark_username", + "cyberark_api_key", + "client_cert", + "client_key", + "ssl_verify", + "refresh_interval", +] as const; + +const propertiesFor = (fields: readonly string[]) => + Object.fromEntries(fields.map((name) => [name, { description: `${name} description` }])); + +const mutate = vi.fn(); + +const setup = (options?: { values?: Record; fields?: readonly string[] }) => { + vi.mocked(useCyberArkConfig).mockReturnValue({ + data: { + field_schema: { properties: propertiesFor(options?.fields ?? ALL_FIELDS) }, + values: options?.values ?? {}, + }, + } as unknown as ReturnType); + + vi.mocked(useUpdateCyberArkConfig).mockReturnValue({ + mutate, + isPending: false, + } as unknown as ReturnType); +}; + +const renderModal = (onSuccess = vi.fn(), onCancel = vi.fn()) => + renderWithProviders(); + +const save = async (user: ReturnType) => + user.click(screen.getByRole("button", { name: "Save" })); + +describe("EditCyberArkModal", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("clears untouched non-sensitive fields and omits untouched sensitive fields", async () => { + setup({ + values: { + cyberark_api_base: "https://conjur.example.com", + cyberark_account: "myorg", + cyberark_api_key: "super-secret-key", + client_key: "super-secret-pem", + }, + }); + const user = userEvent.setup(); + renderModal(); + + await save(user); + + await waitFor(() => { + expect(mutate).toHaveBeenCalledTimes(1); + }); + const expectedPayload = { + cyberark_api_base: "https://conjur.example.com", + cyberark_account: "myorg", + cyberark_username: "", + client_cert: "", + ssl_verify: "", + refresh_interval: "", + }; + expect(mutate.mock.calls[0][0]).toEqual(expectedPayload); + }); + + it("sends a sensitive field only once it is typed into", async () => { + setup({ values: { cyberark_api_base: "https://conjur.example.com", cyberark_api_key: "super-secret-key" } }); + const user = userEvent.setup(); + renderModal(); + + fireEvent.change(screen.getByLabelText("API Key"), { target: { value: "rotated-key" } }); + await save(user); + + await waitFor(() => { + expect(mutate).toHaveBeenCalledTimes(1); + }); + expect(mutate.mock.calls[0][0]).toMatchObject({ cyberark_api_key: "rotated-key" }); + }); + + it("never seeds a stored secret into its input", () => { + setup({ values: { cyberark_api_key: "super-secret-key", client_key: "super-secret-pem" } }); + renderModal(); + + expect(screen.getByLabelText("API Key")).toHaveValue(""); + expect(screen.getByLabelText("Client Key")).toHaveValue(""); + }); + + it("renders only the fields the schema declares, and sends only those", async () => { + setup({ + fields: ["cyberark_api_base", "cyberark_api_key"], + values: { cyberark_api_base: "https://conjur.example.com" }, + }); + const user = userEvent.setup(); + renderModal(); + + expect(screen.queryByLabelText("Account")).not.toBeInTheDocument(); + expect(screen.queryByLabelText("Client Key")).not.toBeInTheDocument(); + + await save(user); + + await waitFor(() => { + expect(mutate).toHaveBeenCalledTimes(1); + }); + expect(mutate.mock.calls[0][0]).toEqual({ cyberark_api_base: "https://conjur.example.com" }); + }); + + it("blocks the submit when the server url does not start with http", async () => { + setup({ values: {} }); + const user = userEvent.setup(); + renderModal(); + + fireEvent.change(screen.getByLabelText("Conjur Server URL"), { target: { value: "conjur.example.com" } }); + await save(user); + + expect(await screen.findByText("Must start with http:// or https://")).toBeInTheDocument(); + expect(mutate).not.toHaveBeenCalled(); + }); + + it("tells the admin a stored secret is kept when the field is left blank", () => { + setup({ values: { cyberark_api_key: "super-secret-key" } }); + renderModal(); + + expect(screen.getByLabelText("API Key")).toHaveAttribute( + "placeholder", + "Leave blank to keep existing (super-secret-key)", + ); + }); + + it("falls back to the schema description when no secret is stored yet", () => { + setup({ values: {} }); + renderModal(); + + expect(screen.getByLabelText("API Key")).toHaveAttribute("placeholder", "cyberark_api_key description"); + }); + + it("closes without saving when cancelled", async () => { + setup({ values: {} }); + const onCancel = vi.fn(); + const user = userEvent.setup(); + renderModal(vi.fn(), onCancel); + + await user.click(screen.getByRole("button", { name: "Cancel" })); + + expect(onCancel).toHaveBeenCalledTimes(1); + expect(mutate).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/EditCyberArkModal.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/EditCyberArkModal.tsx new file mode 100644 index 00000000000..7a093e63f2b --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/EditCyberArkModal.tsx @@ -0,0 +1,176 @@ +"use client"; + +import { useCyberArkConfig } from "@/app/(dashboard)/hooks/configOverrides/useCyberArkConfig"; +import { useUpdateCyberArkConfig } from "@/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { toast } from "@/lib/toast"; +import React, { useMemo } from "react"; +import { z } from "zod/v4"; +import { FieldGroup } from "@/components/ui/field"; +import { FormField } from "@/components/shared/form/FormField"; +import { PasswordInput } from "@/components/shared/PasswordInput"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; +import { Separator } from "@/components/ui/separator"; +import { useZodForm } from "@/lib/forms/useZodForm"; +import { SENSITIVE_FIELDS, FIELD_LABELS } from "./constants"; +import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; + +interface CyberArkFieldGroup { + title: string; + subtitle?: string; + fields: string[]; +} + +const FIELD_GROUPS: CyberArkFieldGroup[] = [ + { + title: "Connection", + fields: ["cyberark_api_base", "cyberark_account", "cyberark_username"], + }, + { + title: "API Key Authentication", + subtitle: "Use a Conjur API key to authenticate. Only one auth method is required.", + fields: ["cyberark_api_key"], + }, + { + title: "Certificate Authentication", + subtitle: "Use a client TLS certificate and key to authenticate. Only one auth method is required.", + fields: ["client_cert", "client_key"], + }, + { + title: "Advanced", + subtitle: "Optional TLS and token caching settings.", + fields: ["ssl_verify", "refresh_interval"], + }, +]; + +type CyberArkFormValues = Record; + +const buildSchema = (fields: readonly string[]): z.ZodType => + z.object( + Object.fromEntries( + fields.map((name) => [ + name, + name === "cyberark_api_base" + ? z.string().refine((value) => value.length === 0 || /^https?:\/\/.+/.test(value), { + message: "Must start with http:// or https://", + }) + : z.string(), + ]), + ), + ) as unknown as z.ZodType; + +interface EditCyberArkModalProps { + isVisible: boolean; + onCancel: () => void; + onSuccess: () => void; +} + +const EditCyberArkModal: React.FC = ({ isVisible, onCancel, onSuccess }) => { + const { accessToken } = useAuthorized(); + const { data } = useCyberArkConfig(); + const { mutate, isPending } = useUpdateCyberArkConfig(accessToken); + + const properties: Record = useMemo( + () => data?.field_schema?.properties ?? {}, + [data], + ); + const rawValues: Record = useMemo(() => data?.values ?? {}, [data]); + + const visibleFields = useMemo( + () => FIELD_GROUPS.flatMap((group) => group.fields).filter((name) => properties[name] !== undefined), + [properties], + ); + + const seededValues = useMemo( + () => + Object.fromEntries( + visibleFields.map((name) => [name, SENSITIVE_FIELDS.has(name) ? "" : ((rawValues[name] ?? "") as string)]), + ), + [visibleFields, rawValues], + ); + + const schema = useMemo(() => buildSchema(visibleFields), [visibleFields]); + const form = useZodForm(schema, { values: seededValues }); + + const handleSubmit = (formValues: CyberArkFormValues) => { + const config: Record = Object.fromEntries( + Object.entries(formValues).flatMap(([key, value]) => { + if (value !== undefined && value !== null && value !== "") return [[key, value]]; + if (!SENSITIVE_FIELDS.has(key)) return [[key, ""]]; + return []; + }), + ); + + mutate(config, { + onSuccess: () => { + toast.success("CyberArk configuration updated successfully"); + onSuccess(); + }, + onError: (err) => { + toast.fromError(err); + }, + }); + }; + + const handleCancel = () => { + form.reset(seededValues); + onCancel(); + }; + + const renderField = (fieldName: string) => { + const fieldSchema = properties[fieldName]; + if (!fieldSchema) return null; + + const isSensitive = SENSITIVE_FIELDS.has(fieldName); + const existingValue = rawValues[fieldName]; + const hasExistingValue = isSensitive && existingValue != null && existingValue !== ""; + const placeholder = hasExistingValue ? `Leave blank to keep existing (${existingValue})` : fieldSchema?.description; + + return ( + + {({ ref, ...field }) => + isSensitive ? ( + + ) : ( + + ) + } + + ); + }; + + return ( + !open && handleCancel()}> + + + Edit CyberArk Configuration + +
+ {FIELD_GROUPS.map((group, index) => ( +
+ {index > 0 && } +
{group.title}
+ {group.subtitle &&

{group.subtitle}

} + {group.fields.map(renderField)} +
+ ))} +
+ +
+ + +
+ + + + ); +}; + +export default EditCyberArkModal; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/constants.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/constants.ts new file mode 100644 index 00000000000..5835f93d19a --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/constants.ts @@ -0,0 +1,12 @@ +export const SENSITIVE_FIELDS = new Set(["cyberark_api_key", "client_key"]); + +export const FIELD_LABELS: Record = { + cyberark_api_base: "Conjur Server URL", + cyberark_account: "Account", + cyberark_username: "Username", + cyberark_api_key: "API Key", + client_cert: "Client Certificate", + client_key: "Client Key", + ssl_verify: "SSL Verification", + refresh_interval: "Token Refresh Interval (seconds)", +}; diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index c48d15adecb..2947e29319b 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -40,6 +40,10 @@ const DEFAULT_SCORING_EXPLANATION = "The router scores each request across 7 dimensions: token count, code presence, reasoning markers, technical " + "terms, simple indicators, multi-step patterns, and question complexity. The weighted score determines the tier:"; +const CLASSIFIER_TIMEOUT_ID = "classifier-timeout-ms"; +const CLASSIFIER_CONTEXT_WINDOW_SIZE_ID = "classifier-context-window-size"; +const CLASSIFIER_CONTEXT_BUDGET_CHARS_ID = "classifier-context-budget-chars"; + const CUSTOM_PROMPT_WITH_HEURISTIC_FALLBACK = "This router classifies with your own prompt, so the tier comes from whatever rubric it states. The four tier " + "names stay fixed. The scoring below is the heuristic, which now runs only when the classifier call fails:"; @@ -204,6 +208,7 @@ const ClassificationMethodConfig: React.FC = ({ showValidationErrors = false, defaultModel, }) => { + const [draft, setDraft] = React.useState<{ id: string; raw: string } | null>(null); const hasDefaultModel = Boolean(defaultModel); const classifierType = effectiveClassifierType(value); const classifierModelMissing = @@ -261,13 +266,13 @@ const ClassificationMethodConfig: React.FC = ({ }); }; - const handleClassifierTimeoutChange = (timeoutMs: number | null) => { + const handleClassifierTimeoutChange = (timeoutMs: number) => { onChange({ ...value, classifier_llm_config: { ...value.classifier_llm_config, model: value.classifier_llm_config?.model ?? "", - timeout_ms: timeoutMs ?? DEFAULT_CLASSIFIER_TIMEOUT_MS, + timeout_ms: timeoutMs, }, }); }; @@ -300,20 +305,32 @@ const ClassificationMethodConfig: React.FC = ({ onChange({ ...value, classifier_fallback: fallback }); }; - const handleClassifierContextWindowSizeChange = (windowSize: number | null) => { + const handleClassifierContextWindowSizeChange = (windowSize: number) => { onChange({ ...value, - classifier_context_window_size: windowSize ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, + classifier_context_window_size: windowSize, }); }; - const handleClassifierContextBudgetCharsChange = (budgetChars: number | null) => { + const handleClassifierContextBudgetCharsChange = (budgetChars: number) => { onChange({ ...value, - classifier_context_budget_chars: budgetChars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS, + classifier_context_budget_chars: budgetChars, }); }; + const handleClassifierIntegerChange = ( + id: string, + raw: string, + minimum: number, + onCommit: (value: number) => void, + ) => { + setDraft({ id, raw }); + const parsed = Number(raw); + if (raw.trim() === "" || !Number.isFinite(parsed)) return; + onCommit(Math.max(minimum, Math.round(parsed))); + }; + const handleClassifierContextIncludeAssistantTurnsChange = (includeAssistantTurns: boolean) => { onChange({ ...value, @@ -366,14 +383,27 @@ const ClassificationMethodConfig: React.FC = ({ {classifierModelMissing && A classifier model is required}
- Timeout (ms) + - handleClassifierTimeoutChange(event.target.value === "" ? null : event.target.valueAsNumber) + id={CLASSIFIER_TIMEOUT_ID} + type="text" + inputMode="numeric" + value={ + draft?.id === CLASSIFIER_TIMEOUT_ID + ? draft.raw + : String(value.classifier_llm_config?.timeout_ms ?? DEFAULT_CLASSIFIER_TIMEOUT_MS) } - min={1} + onChange={(event) => + handleClassifierIntegerChange( + CLASSIFIER_TIMEOUT_ID, + event.target.value, + 1, + handleClassifierTimeoutChange, + ) + } + onBlur={() => setDraft(null)} className="w-full" /> @@ -480,14 +510,27 @@ const ClassificationMethodConfig: React.FC = ({
- Context Window Size + - handleClassifierContextWindowSizeChange(event.target.value === "" ? null : event.target.valueAsNumber) + id={CLASSIFIER_CONTEXT_WINDOW_SIZE_ID} + type="text" + inputMode="numeric" + value={ + draft?.id === CLASSIFIER_CONTEXT_WINDOW_SIZE_ID + ? draft.raw + : String(value.classifier_context_window_size ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE) } - min={0} + onChange={(event) => + handleClassifierIntegerChange( + CLASSIFIER_CONTEXT_WINDOW_SIZE_ID, + event.target.value, + 0, + handleClassifierContextWindowSizeChange, + ) + } + onBlur={() => setDraft(null)} className="w-full" /> @@ -497,14 +540,27 @@ const ClassificationMethodConfig: React.FC = ({
- Context Character Budget + - handleClassifierContextBudgetCharsChange(event.target.value === "" ? null : event.target.valueAsNumber) + id={CLASSIFIER_CONTEXT_BUDGET_CHARS_ID} + type="text" + inputMode="numeric" + value={ + draft?.id === CLASSIFIER_CONTEXT_BUDGET_CHARS_ID + ? draft.raw + : String(value.classifier_context_budget_chars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS) } - min={0} + onChange={(event) => + handleClassifierIntegerChange( + CLASSIFIER_CONTEXT_BUDGET_CHARS_ID, + event.target.value, + 0, + handleClassifierContextBudgetCharsChange, + ) + } + onBlur={() => setDraft(null)} className="w-full" /> diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 91201b51663..33ce1169c46 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -131,10 +131,8 @@ describe("ComplexityRouterConfig", () => { fireEvent.click(screen.getByText("Advanced: Classification Method")); expect(screen.getByText("Classifier Model")).toBeInTheDocument(); - expect(screen.getByText("Timeout (ms)")).toBeInTheDocument(); - expect(screen.getByDisplayValue("750")).toBeInTheDocument(); - expect(screen.getByText("Context Window Size")).toBeInTheDocument(); - expect(screen.getByDisplayValue("5")).toBeInTheDocument(); + expect(screen.getByLabelText("Timeout (ms)")).toHaveValue("750"); + expect(screen.getByLabelText("Context Window Size")).toHaveValue("5"); expect(screen.queryByText("Context Per-Turn Character Limit")).not.toBeInTheDocument(); }); @@ -148,11 +146,8 @@ describe("ComplexityRouterConfig", () => { fireEvent.click(screen.getByText("Advanced: Classification Method")); - const windowSizeSection = screen.getByText("Context Window Size").closest("div") as HTMLElement; - expect(within(windowSizeSection).getByDisplayValue("3")).toBeInTheDocument(); - - const budgetSection = screen.getByText("Context Character Budget").closest("div") as HTMLElement; - expect(within(budgetSection).getByDisplayValue("8000")).toBeInTheDocument(); + expect(screen.getByLabelText("Context Window Size")).toHaveValue("3"); + expect(screen.getByLabelText("Context Character Budget")).toHaveValue("8000"); }); it("should warn when the budget is too small to quote any turn that does not already fit", () => { @@ -247,7 +242,11 @@ describe("ComplexityRouterConfig", () => { expect(screen.queryByText("Context Per-Turn Character Limit")).not.toBeInTheDocument(); }); - it("should call onChange with the updated classifier_context_window_size when edited", () => { + it.each([ + ["Timeout (ms)", "7", { classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 7 } }], + ["Context Window Size", "0", { classifier_context_window_size: 0 }], + ["Context Character Budget", "7", { classifier_context_budget_chars: 7 }], + ])("keeps %s empty while it is being edited, then commits %s", (label, replacement, expected) => { const onChange = vi.fn(); const llmValue: ComplexityRouterConfigValue = { ...defaultValue, @@ -257,14 +256,31 @@ describe("ComplexityRouterConfig", () => { renderWithProviders(); fireEvent.click(screen.getByText("Advanced: Classification Method")); - const windowSizeSection = screen.getByText("Context Window Size").closest("div") as HTMLElement; - const input = within(windowSizeSection).getByRole("spinbutton"); - fireEvent.change(input, { target: { value: "7" } }); + const input = screen.getByLabelText(label); + fireEvent.change(input, { target: { value: "" } }); - expect(onChange).toHaveBeenCalledWith({ - ...llmValue, - classifier_context_window_size: 7, - }); + expect(input).toHaveValue(""); + expect(onChange).not.toHaveBeenCalled(); + + fireEvent.change(input, { target: { value: replacement } }); + + expect(onChange).toHaveBeenLastCalledWith({ ...llmValue, ...expected }); + }); + + it("restores the committed context window size after an empty field loses focus", () => { + const llmValue: ComplexityRouterConfigValue = { + ...defaultValue, + classifier_type: "llm", + classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 }, + }; + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + + const input = screen.getByLabelText("Context Window Size"); + fireEvent.change(input, { target: { value: "" } }); + fireEvent.blur(input); + + expect(input).toHaveValue("3"); }); it("should render the custom technical keywords field", () => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index 022dd0b7cad..a4921fcfcb5 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -295,8 +295,7 @@ describe("EditAutoRouterModal classifier context window", () => { renderLlmModal(); await user.click(await screen.findByText("Advanced: Classification Method")); - const windowSizeSection = (await screen.findByText("Context Window Size")).closest("div") as HTMLElement; - const input = within(windowSizeSection).getByRole("spinbutton"); + const input = await screen.findByLabelText("Context Window Size"); fireEvent.change(input, { target: { value: "8" } }); await user.click(screen.getByRole("button", { name: /save changes/i })); diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx index c37589f63f5..d248cf311cc 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx @@ -361,4 +361,32 @@ describe("PaginatedSearchSelect", () => { await user.click(screen.getByRole("combobox")); expect(await screen.findByTestId("paginated-search-select-loading-more")).toBeInTheDocument(); }); + + it("queries the trimmed label when a character is deleted from the end of the selection", async () => { + const user = userEvent.setup(); + const onSearchChange = vi.fn(); + renderSelect({ onSearchChange, value: "alias-alpha" }); + + const input = screen.getByRole("combobox") as HTMLInputElement; + input.focus(); + input.setSelectionRange(input.value.length, input.value.length); + await user.keyboard("{Backspace}"); + + expect(input).toHaveValue("alias-alph"); + await waitFor(() => expect(onSearchChange).toHaveBeenLastCalledWith("alias-alph")); + }); + + it("queries what is left when a character is deleted from inside the selection", async () => { + const user = userEvent.setup(); + const onSearchChange = vi.fn(); + renderSelect({ onSearchChange, value: "alias-alpha" }); + + const input = screen.getByRole("combobox") as HTMLInputElement; + input.focus(); + input.setSelectionRange(6, 6); + await user.keyboard("{Backspace}"); + + expect(input).toHaveValue("aliasalpha"); + await waitFor(() => expect(onSearchChange).toHaveBeenLastCalledWith("aliasalpha")); + }); }); diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx index 0f25260aad0..4b7ef7401c7 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx @@ -50,6 +50,11 @@ const typedInsertion = (previous: string, next: string): string => { return next.slice(start, next.length - end); }; +const editedQuery = (label: string, next: string): string => { + const inserted = typedInsertion(label, next); + return inserted === "" && next !== label ? next : inserted; +}; + export function PaginatedSearchSelect({ options, value, @@ -101,7 +106,7 @@ export function PaginatedSearchSelect({ const replacedWholeInput = wholeSelectionRef.current; wholeSelectionRef.current = false; handleInputValueChange( - typedQuery === null && !replacedWholeInput ? typedInsertion(selected?.label ?? "", next) : next, + typedQuery === null && !replacedWholeInput ? editedQuery(selected?.label ?? "", next) : next, reason, ); }; diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.test.tsx new file mode 100644 index 00000000000..a8e27504019 --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.test.tsx @@ -0,0 +1,67 @@ +import { screen, within } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; + +import { render } from "../../../../tests/test-utils"; +import type { LogEntry } from "../columns"; +import { DrawerHeader } from "./DrawerHeader"; + +const logEntry = (overrides: Partial): LogEntry => + ({ + request_id: "170d64ea-69f0-431a-be72-332f8f78c18a", + api_key: "key-1", + team_id: "team-1", + model: "gpt-4o", + model_id: "model-1", + custom_llm_provider: "openai", + call_type: "acompletion", + spend: 0.01, + total_tokens: 10, + prompt_tokens: 5, + completion_tokens: 5, + startTime: "2026-07-07T09:50:13Z", + endTime: "2026-07-07T09:50:14Z", + cache_hit: "false", + messages: [], + response: {}, + ...overrides, + }) as LogEntry; + +const renderHeader = (log: LogEntry, isSidebarCollapsed: boolean) => + render( + , + ); + +const expandToggle = () => screen.getByLabelText("Expand trace sidebar"); + +describe("DrawerHeader sidebar toggle", () => { + it("stays out of the header while the sidebar owns it", () => { + renderHeader(logEntry({}), false); + + expect(screen.queryByLabelText("Expand trace sidebar")).not.toBeInTheDocument(); + expect(screen.queryByLabelText("Collapse trace sidebar")).not.toBeInTheDocument(); + }); + + it("shares the model row once the sidebar is collapsed", () => { + renderHeader(logEntry({}), true); + + const row = expandToggle().parentElement as HTMLElement; + expect(within(row).getByText("gpt-4o")).toBeInTheDocument(); + }); + + it("falls back to the request id row when the log names no model", () => { + renderHeader(logEntry({ model: "", custom_llm_provider: "" }), true); + + const row = expandToggle().parentElement as HTMLElement; + expect(within(row).getByText("170d64ea-69f0-431a-be72-332f8f78c18a")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx index 751667b917b..65b5801602c 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx @@ -7,6 +7,7 @@ import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/comp import { LogEntry } from "../columns"; import { AutoRouterTag } from "@/components/shared/table_cells"; import { ClassifyTag } from "./ClassifyTag"; +import { SidebarToggle } from "./SidebarToggle"; import { getProviderLogoAndName } from "../../provider_info_helpers"; import { DRAWER_HEADER_PADDING, @@ -26,6 +27,8 @@ interface DrawerHeaderProps { statusLabel: string; statusColor: "error" | "success"; environment: string; + isSidebarCollapsed: boolean; + onToggleSidebar: () => void; } /** @@ -40,9 +43,13 @@ export function DrawerHeader({ statusLabel, statusColor, environment, + isSidebarCollapsed, + onToggleSidebar, }: DrawerHeaderProps) { const provider = log.custom_llm_provider || ""; const providerInfo = provider ? getProviderLogoAndName(provider) : null; + const showToggleWithProvider = isSidebarCollapsed && Boolean(providerInfo || log.model); + const showToggleWithRequestId = isSidebarCollapsed && !showToggleWithProvider; return (
{/* Row 0: Model + Provider with Logo */} - +
+ {showToggleWithProvider && } + +
{/* Row 1: Request ID + Actions */}
+ {showToggleWithRequestId && }
@@ -95,7 +112,7 @@ function ModelProviderSection({ providerName?: string; }) { return ( -
+
{providerLogo && (
- {!isSidebarCollapsed ? ( - - ) : ( - + {!isSidebarCollapsed && ( + setIsSidebarCollapsed(true)} + className="absolute top-2 left-2 z-raised" + /> )} {!isSidebarCollapsed && (
@@ -466,6 +452,8 @@ export function LogDetailsDrawer({ setIsSidebarCollapsed((collapsed) => !collapsed)} onPrevious={selectPreviousLog} onNext={selectNextLog} statusLabel={statusLabel} diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SidebarToggle.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SidebarToggle.tsx new file mode 100644 index 00000000000..a1ba7ff2491 --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SidebarToggle.tsx @@ -0,0 +1,23 @@ +import { ChevronLeft, ChevronRight } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { cn } from "@/lib/cva.config"; + +export interface SidebarToggleProps { + isCollapsed: boolean; + onToggle: () => void; + className?: string; +} + +export function SidebarToggle({ isCollapsed, onToggle, className }: SidebarToggleProps) { + return ( + + ); +} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index eaa05ddc005..6c5319f0361 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -2858,6 +2858,59 @@ export interface paths { patch?: never; trace?: never; }; + "/config_overrides/cyberark": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Cyberark Config + * @description 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. + */ + get: operations["get_cyberark_config_config_overrides_cyberark_get"]; + put?: never; + /** + * Update Cyberark Config + * @description Update CyberArk Conjur secret manager configuration. + * Sets environment variables, encrypts sensitive fields, and stores in DB. + * Reinitializes the secret manager on this pod. + */ + post: operations["update_cyberark_config_config_overrides_cyberark_post"]; + /** + * Delete Cyberark Config + * @description Delete CyberArk Conjur configuration. Idempotent. + */ + delete: operations["delete_cyberark_config_config_overrides_cyberark_delete"]; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/config_overrides/cyberark/test_connection": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Test Cyberark Connection + * @description Test the connection to the currently configured CyberArk Conjur server. + * Uses the already-initialized secret manager client. Does not modify any state. + */ + post: operations["test_cyberark_connection_config_overrides_cyberark_test_connection_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/config_overrides/hashicorp_vault": { parameters: { query?: never; @@ -6974,6 +7027,30 @@ export interface paths { patch?: never; trace?: never; }; + "/introspect": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Introspect Endpoint + * @description 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. + */ + post: operations["introspect_endpoint_introspect_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/invitation/delete": { parameters: { query?: never; @@ -23474,6 +23551,11 @@ export interface components { /** Mask[] */ "mask[]"?: string[] | null; }; + /** Body_introspect_endpoint_introspect_post */ + Body_introspect_endpoint_introspect_post: { + /** Token */ + token: string; + }; /** Body_revoke_endpoint_revoke_post */ Body_revoke_endpoint_revoke_post: { /** Client Id */ @@ -25812,6 +25894,52 @@ export interface components { /** User Id */ user_id: string; }; + /** + * CyberArkConfig + * @description Configuration for CyberArk Conjur secret manager integration. + */ + CyberArkConfig: { + /** + * Client Cert + * @description Path to the client TLS certificate for certificate-based authentication + */ + client_cert?: string | null; + /** + * Client Key + * @description Path to the client TLS private key for certificate-based authentication + */ + client_key?: string | null; + /** + * Cyberark Account + * @description The Conjur organization account name + */ + cyberark_account?: string | null; + /** + * Cyberark Api Base + * @description The address of the CyberArk Conjur server (e.g., https://conjur.example.com) + */ + cyberark_api_base?: string | null; + /** + * Cyberark Api Key + * @description API key for Conjur API-key authentication + */ + cyberark_api_key?: string | null; + /** + * Cyberark Username + * @description The Conjur username (login) to authenticate as + */ + cyberark_username?: string | null; + /** + * Refresh Interval + * @description Auth token cache TTL in seconds (default: 300) + */ + refresh_interval?: string | null; + /** + * Ssl Verify + * @description Set to false to disable SSL verification (e.g., for self-signed certificates) + */ + ssl_verify?: string | null; + }; /** DailySpendData */ DailySpendData: { breakdown?: components["schemas"]["BreakdownMetrics"]; @@ -42757,6 +42885,120 @@ export interface operations { }; }; }; + get_cyberark_config_config_overrides_cyberark_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ConfigOverrideSettingsResponse"]; + }; + }; + }; + }; + update_cyberark_config_config_overrides_cyberark_post: { + parameters: { + query?: never; + header?: { + /** @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 */ + "litellm-changed-by"?: string | null; + }; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["CyberArkConfig"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: string; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + delete_cyberark_config_config_overrides_cyberark_delete: { + parameters: { + query?: never; + header?: { + /** @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 */ + "litellm-changed-by"?: string | null; + }; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: string; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + test_cyberark_connection_config_overrides_cyberark_test_connection_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: string; + }; + }; + }; + }; + }; get_hashicorp_vault_config_config_overrides_hashicorp_vault_get: { parameters: { query?: never; @@ -47651,6 +47893,39 @@ export interface operations { }; }; }; + introspect_endpoint_introspect_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/x-www-form-urlencoded": components["schemas"]["Body_introspect_endpoint_introspect_post"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; invitation_delete_invitation_delete_post: { parameters: { query?: never;