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 962d1266fd7..83969d8dedf 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 17271 + "limit": 17270 }, "reportArgumentType": { - "limit": 2539 + "limit": 2538 }, "reportAssignmentType": { "limit": 319 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 5486 + "limit": 5485 }, "reportFunctionMemberAccess": { "limit": 7 diff --git a/codecov.yaml b/codecov.yaml index bc0b3604329..4d93c18f3ac 100644 --- a/codecov.yaml +++ b/codecov.yaml @@ -25,6 +25,8 @@ flag_management: carryforward: false - name: proxy-db-schema-migration carryforward: false + - name: circleci + carryforward: false component_management: individual_components: 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/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index cf2cee9b6ef..570b306d6df 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -36,6 +36,7 @@ from litellm.llms.base_llm.managed_resources.isolation import ( build_list_page, build_owner_filter, can_access_resource, + resolve_resource_owner_id, ) from litellm.proxy._types import ( CallTypes, @@ -222,7 +223,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_object=file_object, model_mappings=model_mappings, flat_model_file_ids=list(model_mappings.values()), - created_by=user_api_key_dict.user_id, + created_by=resolve_resource_owner_id(user_api_key_dict), team_id=user_api_key_dict.team_id, updated_by=user_api_key_dict.user_id, ) @@ -238,7 +239,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "unified_file_id": file_id, "model_mappings": json.dumps(model_mappings), "flat_model_file_ids": list(model_mappings.values()), - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, } @@ -342,7 +343,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "file_object": file_object.model_dump_json(), "model_object_id": model_object_id, "file_purpose": file_purpose, - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, "status": file_object.status, diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 3653aba67ef..cac98b69793 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.61" +version = "0.1.62" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.61" +version = "0.1.62" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql new file mode 100644 index 00000000000..c3018006adb --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql @@ -0,0 +1,12 @@ +CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetWindowSpend" ( + "entity_type" TEXT NOT NULL, + "entity_id" TEXT NOT NULL, + "window_duration" TEXT NOT NULL, + "window_start" TIMESTAMP(3) NOT NULL, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_BudgetWindowSpend_pkey" PRIMARY KEY ("entity_type","entity_id","window_duration") +); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260829000000_add_model_access_group_budget_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260829000000_add_model_access_group_budget_table/migration.sql new file mode 100644 index 00000000000..62398da7f04 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260829000000_add_model_access_group_budget_table/migration.sql @@ -0,0 +1,20 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_ModelAccessGroupBudgetTable" ( + "access_group_name" TEXT NOT NULL, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "budget_id" TEXT, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT, + + CONSTRAINT "LiteLLM_ModelAccessGroupBudgetTable_pkey" PRIMARY KEY ("access_group_name") +); + +-- AddForeignKey +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_ModelAccessGroupBudgetTable_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_ModelAccessGroupBudgetTable" ADD CONSTRAINT "LiteLLM_ModelAccessGroupBudgetTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 2bb850139a2..60223265211 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -29,6 +29,7 @@ model LiteLLM_BudgetTable { keys LiteLLM_VerificationToken[] // multiple keys can have the same budget end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget + model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization } @@ -585,6 +586,20 @@ model LiteLLM_EndUserTable { blocked Boolean @default(false) } +// Budget and shared spend for a model access group. The groups themselves are not rows anywhere: +// they are free-text strings in LiteLLM_ProxyModelTable.model_info.access_groups, so a row here +// exists only once someone gives that group a budget. +model LiteLLM_ModelAccessGroupBudgetTable { + access_group_name String @id + spend Float @default(0.0) + budget_id String? + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + created_at DateTime @default(now()) @map("created_at") + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? +} + // Track tags with budgets and spend model LiteLLM_TagTable { tag_name String @id @@ -649,6 +664,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 0ef5cd1e856..d5741d479bf 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.90" +version = "0.4.91" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.90" +version = "0.4.91" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", 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 fc88086805f..cc6db6c10cc 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. @@ -288,6 +289,7 @@ REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_end_user_spend_update_buffer" REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_update_buffer" REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer" +REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) # Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth LITELLM_ASYNCIO_QUEUE_MAXSIZE: Final = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)) @@ -1681,6 +1683,7 @@ DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS: Final = 16 # Ceilings on the cached auth registries; larger tables fall back to per-row lookups # instead of holding an unbounded id set in every worker. TAG_REGISTRY_MAX_SIZE: Final = 5000 +MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE: Final = 5000 END_USER_RESTRICTED_REGISTRY_MAX_SIZE: Final = 5000 # How long a failed registry load is remembered as "unusable", so a degraded Postgres # is not re-scanned on every request on top of the per-id lookups it falls back to. 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/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index 34d5797a6d8..4d043701f40 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -25,6 +25,13 @@ from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, Inter BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) +MODEL_ACCESS_GROUP_METADATA_KEY: Final = "user_api_key_matched_model_access_groups" +"""Where auth records the model access groups that authorized the request, for the spend writer. + +The ``user_api_key`` prefix is load-bearing, not cosmetic: when a request carries both +``metadata`` and ``litellm_metadata``, ``get_litellm_metadata_from_kwargs`` returns the latter and +copies a key across only when ``user_api_key`` appears in its name.""" + _USER_API_KEY_AUTH_KEY: Final = "user_api_key_auth" FORWARDABLE_IDENTITY_METADATA_KEYS: Final = frozenset( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a8672b5c112..c27822d0479 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -64,7 +64,10 @@ from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger from litellm.litellm_core_utils.core_helpers import is_expected_client_error, reconstruct_model_name from litellm.litellm_core_utils.get_litellm_params import get_litellm_params -from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call +from litellm.litellm_core_utils.internal_call_metadata import ( + MODEL_ACCESS_GROUP_METADATA_KEY, + is_unbilled_non_inference_call, +) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( cost_breakdown_with_guardrail, guardrail_information_cost, @@ -2872,6 +2875,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 +2885,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 +2898,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, @@ -5049,6 +5054,42 @@ def is_valid_sha256_hash(value: str) -> bool: return bool(re.fullmatch(r"[a-fA-F0-9]{64}", value)) +def coerce_model_access_groups(value: object) -> tuple[str, ...]: + """Model access group names out of untrusted request metadata, deduped and order preserving.""" + if not isinstance(value, (list, tuple)): + return () + return tuple(dict.fromkeys(group for group in value if isinstance(group, str) and group)) + + +def _model_access_groups_on_auth_object(user_api_key_auth: object) -> object: + if isinstance(user_api_key_auth, Mapping): + return user_api_key_auth.get("matched_model_access_groups") + return getattr(user_api_key_auth, "matched_model_access_groups", None) + + +def _model_access_groups_from_metadata(metadata: Mapping[str, object]) -> tuple[str, ...]: + stamped: Final = coerce_model_access_groups(metadata.get(MODEL_ACCESS_GROUP_METADATA_KEY)) + if stamped: + return stamped + return coerce_model_access_groups(_model_access_groups_on_auth_object(metadata.get("user_api_key_auth"))) + + +def request_model_access_groups_from_litellm_params(litellm_params: Mapping[str, object]) -> tuple[str, ...]: + """Access groups the auth layer stamped onto this request, from whichever metadata field carries them. + + Detached internal sub-calls only inherit the identity keys, so the auth object is the + fallback there, exactly as _get_budget_reservation_from_metadata does for reservations. + """ + for metadata_variable_name in ("metadata", "litellm_metadata"): + metadata = litellm_params.get(metadata_variable_name) + if not isinstance(metadata, Mapping): + continue + model_access_groups = _model_access_groups_from_metadata(metadata) + if model_access_groups: + return model_access_groups + return () + + class StandardLoggingPayloadSetup: @staticmethod def cleanup_timestamps( @@ -5422,6 +5463,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 +5855,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, ) @@ -5896,6 +5941,7 @@ def get_standard_logging_object_payload( request_tags: Final = StandardLoggingPayloadSetup._get_request_tags( litellm_params=litellm_params, proxy_server_request=proxy_server_request ) + request_model_access_groups: Final = request_model_access_groups_from_litellm_params(litellm_params) # cleanup timestamps ( @@ -6058,6 +6104,7 @@ def get_standard_logging_object_payload( prompt_tokens=usage_dict.get("prompt_tokens", 0), completion_tokens=usage_dict.get("completion_tokens", 0), request_tags=request_tags, + request_model_access_groups=request_model_access_groups, end_user=end_user_id, api_base=StandardLoggingPayloadSetup.strip_trailing_slash(litellm_params.get("api_base", "")) or "", model_group=_model_group, @@ -6228,6 +6275,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, ) @@ -6269,6 +6318,7 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload: cache_key=None, saved_cache_cost=saved_cache_cost, request_tags=[], + request_model_access_groups=(), end_user=None, requester_ip_address="127.0.0.1", messages=messages, diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 47116a8f8fb..e1387a9068c 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1268,7 +1268,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) @staticmethod - def _cap_thinking_budget_to_max_tokens( + def cap_thinking_budget_to_max_tokens( thinking: AnthropicThinkingParam, max_tokens: int | None ) -> AnthropicThinkingParam | None: """Cap a legacy ``thinking.budget_tokens`` below ``max_tokens`` (Anthropic @@ -1530,7 +1530,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): llm_provider=self._resolved_provider, ) capped_thinking = ( - AnthropicConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens) + AnthropicConfig.cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens) if legacy_thinking is not None else 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/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index ebd514c2605..3d62b8b4784 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -40,6 +40,11 @@ DROP_UNSUPPORTED_ADAPTIVE_EFFORT_WARNING: Final = ( "minimum thinking budget." ) +DROP_UNFITTING_REASONING_EFFORT_WARNING: Final = ( + "Dropping `thinking` mapped from reasoning_effort=%s for model=%s: max_tokens=%s " + "is too small to fit the minimum thinking budget." +) + class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): @property @@ -335,11 +340,15 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): return headers, api_base @staticmethod - def _translate_reasoning_effort_to_anthropic(model: str, optional_params: dict, custom_llm_provider: str) -> None: + def _translate_reasoning_effort_to_anthropic( + model: str, optional_params: dict, max_tokens: int | None, custom_llm_provider: str + ) -> None: """Map OpenAI-style ``reasoning_effort`` to native Anthropic params. Caller-supplied ``thinking`` / ``output_config`` win over the alias. - ``effort='none'`` clears both. Invalid efforts raise a 400. + ``effort='none'`` clears both. Invalid efforts raise a 400. A mapped + thinking budget is capped below ``max_tokens`` and dropped when even + the minimum budget cannot fit. """ from litellm.exceptions import BadRequestError as _BadRequestError from litellm.llms.anthropic.chat.transformation import ( @@ -365,7 +374,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): optional_params.pop("output_config", None) return - optional_params.setdefault("thinking", mapped_thinking) + fitted_thinking: Final = AnthropicConfig.cap_thinking_budget_to_max_tokens(mapped_thinking, max_tokens) + if fitted_thinking is None: + verbose_logger.warning(DROP_UNFITTING_REASONING_EFFORT_WARNING, reasoning_effort, model, max_tokens) + return + + optional_params.setdefault("thinking", fitted_thinking) if AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider): mapped_effort: Final = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get(reasoning_effort) if mapped_effort is None: @@ -510,7 +524,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): except _BadRequestError as e: raise AnthropicError(message=str(e.message), status_code=400) capped_thinking: Final = ( - AnthropicConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens) + AnthropicConfig.cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens) if legacy_thinking is not None else None ) @@ -582,6 +596,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): self._translate_reasoning_effort_to_anthropic( model=model, optional_params=anthropic_messages_optional_request_params, + max_tokens=max_tokens, custom_llm_provider=self._resolved_provider, ) 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/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index 2a59eddf88a..cced330d873 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -12,6 +12,7 @@ from litellm.llms.base_llm.managed_resources.isolation import ( build_list_page, build_owner_filter, can_access_resource, + resolve_resource_owner_id, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import SpecialEnums @@ -157,7 +158,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): "resource_object": resource_object, "model_mappings": model_mappings, "flat_model_resource_ids": list(model_mappings.values()), - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, } @@ -179,7 +180,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): "unified_resource_id": unified_resource_id, "model_mappings": json.dumps(model_mappings), "flat_model_resource_ids": list(model_mappings.values()), - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, } diff --git a/litellm/llms/base_llm/managed_resources/isolation.py b/litellm/llms/base_llm/managed_resources/isolation.py index e1b204214d7..6a71e8e9223 100644 --- a/litellm/llms/base_llm/managed_resources/isolation.py +++ b/litellm/llms/base_llm/managed_resources/isolation.py @@ -3,10 +3,11 @@ Tenant-isolation helpers for managed file/batch/vector-store resources. Returns a Prisma filter and an ownership check that scope managed resources to the caller's identity: proxy admins see everything, user-keyed callers -see records they created, and service-account keys (no user_id) fall back -to the resource's owning team. Callers with no admin role and no -identifying ids are denied so an empty user_id can never select an -unscoped query. +see records they created, service-account keys (no user_id) fall back to +the resource's owning team, and keys with neither a user_id nor a team_id +fall back to their own hashed token so they can still reach the resources +they created. Callers with no admin role and no identifying ids at all +are denied so an empty user_id can never select an unscoped query. """ from typing import Any, Final @@ -19,6 +20,32 @@ from litellm.proxy._types import ( ) +def resolve_resource_owner_id( + user_api_key_dict: UserAPIKeyAuth, +) -> str | None: + """Return the identity to stamp on (and match against) a managed + resource's ``created_by``. + + A key with neither a user_id nor a team_id would otherwise stamp + ``created_by=None`` and be locked out of its own resources, so it owns + them under its hashed token instead, using the ``key:`` scope prefix + already used by ``proxy/common_utils/resource_ownership.py``. ``None`` + means the caller has no usable identity of its own and must fall back + to team scoping, or be denied. + """ + if user_api_key_dict.user_id is not None: + return user_api_key_dict.user_id + + if user_api_key_dict.team_id is not None: + return None + + token: Final = user_api_key_dict.token or user_api_key_dict.api_key + if token: + return f"key:{token}" + + return None + + def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, Any]: """Build the OpenAI-style paginated list response shape used by managed file/batch/vector-store listings. ``first_id`` and ``last_id`` are @@ -39,7 +66,8 @@ def build_owner_filter( to records the caller is allowed to see. - ``{}`` means no scoping (proxy admins). - - ``{"created_by": }`` for user-keyed callers. + - ``{"created_by": }`` for user-keyed callers, and for keys + with no user_id and no team_id (owner id is their hashed token). - ``{"team_id": }`` for service-account callers that have a team but no user_id. - ``{"OR": [...]}`` when the caller has both — listing must include @@ -62,12 +90,13 @@ def build_owner_filter( ] } - if user_id is not None: - return {"created_by": user_id} - if team_id is not None: return {"team_id": team_id} + owner_id: Final = resolve_resource_owner_id(user_api_key_dict) + if owner_id is not None: + return {"created_by": owner_id} + return None @@ -86,8 +115,8 @@ def can_access_resource( if _user_has_admin_view(user_api_key_dict): return True - user_id: Final = user_api_key_dict.user_id - if user_id is not None and created_by is not None and created_by == user_id: + owner_id: Final = resolve_resource_owner_id(user_api_key_dict) + if owner_id is not None and created_by is not None and created_by == owner_id: return True team_id: Final = user_api_key_dict.team_id diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index d22b225b0bd..db9c8a5cedd 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -924,7 +924,7 @@ class AmazonConverseConfig(BaseConfig): custom_llm_provider="bedrock", ) capped = ( - AnthropicConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens) + AnthropicConfig.cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens) if legacy_thinking is not None else None ) 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/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 27e4b7dc4a6..eac844a790d 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -33,6 +33,7 @@ else: _NO_TOOL_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) _MODEL_FAMILIES_REJECTING_TOP_LEVEL_SCHEMA_COMBINATORS: Final = ("gpt-4", "gpt-3.5", "chatgpt-4o", "o1", "o3", "o4") +_PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): @@ -172,7 +173,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): input = self._validate_input_param(input) tools = response_api_optional_request_params.get("tools") input, tools = self.remove_cache_control_flag_from_input_and_tools(model=model, input=input, tools=tools) - sanitized_tools: Final = self._flatten_tool_schema_combinators_for_openai(model=model, tools=tools) + sanitized_tools: Final = self._flatten_tool_schema_combinators_for_openai( + model=model, tools=tools, litellm_params=litellm_params + ) if sanitized_tools is not None: response_api_optional_request_params["tools"] = sanitized_tools final_request_params: Final = dict( @@ -217,6 +220,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): self, model: str, tools: list[ALL_RESPONSES_API_TOOL_PARAMS] | None, # mutable-ok: request tools are a JSON list + litellm_params: GenericLiteLLMParams, ) -> list[ALL_RESPONSES_API_TOOL_PARAMS] | None: # mutable-ok: request tools are a JSON list """Flatten top-level schema combinators only where OpenAI's validator rejects them. @@ -224,10 +228,14 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): Codex talks to natively) accept them, and so do GPT-5 and later models, which also call tools better with the union intact. Codex wraps MCP tools inside namespace entries, so nested ``tools`` arrays are walked too. + Azure OpenAI shares the validator but names deployments arbitrarily, so + the router's declared ``model_info.base_model`` wins over the deployment + name and an unrecognized name without one is left untouched. """ - if tools is None or self.custom_llm_provider != LlmProviders.OPENAI: + if tools is None or self.custom_llm_provider not in _PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: return tools - if not self._rejects_top_level_schema_combinators(model): + gate_model: Final = self._combinator_gate_model(model=model, litellm_params=litellm_params) + if not self._rejects_top_level_schema_combinators(gate_model): return tools flattened: Final = [ # mutable-ok: request tools are a JSON list self._flattened_tool_or_passthrough(tool) for tool in tools @@ -244,6 +252,12 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): base_model: Final = bare_model.split(":")[1] if bare_model.startswith("ft:") else bare_model return base_model.startswith(_MODEL_FAMILIES_REJECTING_TOP_LEVEL_SCHEMA_COMBINATORS) + @staticmethod + def _combinator_gate_model(model: str, litellm_params: GenericLiteLLMParams) -> str: + model_info: Final[object] = getattr(litellm_params, "model_info", None) + base_model: Final[object] = model_info.get("base_model") if isinstance(model_info, dict) else None + return base_model if isinstance(base_model, str) and base_model else model + @staticmethod def _flattened_tool_entry( entry: Mapping[str, object], @@ -714,7 +728,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): input = self._validate_input_param(input) tools = response_api_optional_request_params.get("tools") input, tools = self.remove_cache_control_flag_from_input_and_tools(model=model, input=input, tools=tools) - sanitized_tools: Final = self._flatten_tool_schema_combinators_for_openai(model=model, tools=tools) + sanitized_tools: Final = self._flatten_tool_schema_combinators_for_openai( + model=model, tools=tools, litellm_params=litellm_params + ) if sanitized_tools is not None: response_api_optional_request_params["tools"] = sanitized_tools data: Final = dict(ResponsesAPIRequestParams(model=model, input=input, **response_api_optional_request_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/audio_transcription/gemini_transcribe_transformation.py b/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py new file mode 100644 index 00000000000..f4db5eb110c --- /dev/null +++ b/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py @@ -0,0 +1,216 @@ +import base64 +from collections.abc import Mapping, Sequence +from typing import Final + +from httpx import Headers, Response + +import litellm +from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.audio_utils.utils import ( + normalize_transcription_language_to_bcp47, + process_audio_file, +) +from litellm.llms.base_llm.audio_transcription.transformation import ( + AudioTranscriptionRequestData, + BaseAudioTranscriptionConfig, +) +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.vertex_ai.audio_transcription.transformation import ( + SUPPORTED_RESPONSE_FORMATS, + validate_vertex_transcription_location, + validate_vertex_transcription_project_id, +) +from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.types.llms.openai import ( + AllMessageValues, + OpenAIAudioTranscriptionOptionalParams, +) +from litellm.types.llms.vertex_ai_gemini_transcription import ( + VertexGeminiTranscriptionAudioConfig, + VertexGeminiTranscriptionContent, + VertexGeminiTranscriptionGenerationConfig, + VertexGeminiTranscriptionInlineData, + VertexGeminiTranscriptionPart, + VertexGeminiTranscriptionRequest, + VertexGeminiTranscriptionResponse, +) +from litellm.types.utils import ( + FileTypes, + TranscriptionResponse, + TranscriptionUsageInputTokenDetailsObject, + TranscriptionUsageTokensObject, +) + +DEFAULT_GEMINI_TRANSCRIBE_LOCATION: Final = "global" +AUDIO_MODALITY: Final = "AUDIO" + + +class VertexGeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexBase): + def __init__(self) -> None: + BaseAudioTranscriptionConfig.__init__(self) + VertexBase.__init__(self) + + def get_supported_openai_params( + self, model: str + ) -> list[OpenAIAudioTranscriptionOptionalParams]: # mutable-ok: BaseAudioTranscriptionConfig signature + return ["language", "response_format"] + + def map_openai_params( + self, + non_default_params: Mapping[str, object], + optional_params: Mapping[str, object], + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: BaseAudioTranscriptionConfig signature + supported_params: Final = frozenset(self.get_supported_openai_params(model)) + mapped: Final = { + **optional_params, + **{k: v for k, v in non_default_params.items() if k in supported_params}, + } + response_format: Final = mapped.get("response_format") + if response_format is None or response_format in SUPPORTED_RESPONSE_FORMATS: + return mapped + if drop_params or litellm.drop_params: + return {k: v for k, v in mapped.items() if k != "response_format"} + raise UnsupportedParamsError( + status_code=400, + message=( + f"Vertex AI Gemini transcription does not support response_format={response_format!r}. " + f"Supported values: {', '.join(SUPPORTED_RESPONSE_FORMATS)}. " + "To drop unsupported openai params from the call, set `litellm.drop_params = True`" + ), + ) + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict | Headers, # mutable-ok: base signature and VertexAIError take dict | Headers + ) -> BaseLLMException: + return VertexAIError(status_code=status_code, message=error_message, headers=headers) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseAudioTranscriptionConfig signature + vertex_params: Final = dict(litellm_params) + access_token, project_id = self._ensure_access_token( + credentials=self.safe_get_vertex_ai_credentials(vertex_params), + project_id=self.safe_get_vertex_ai_project(vertex_params), + custom_llm_provider="vertex_ai", + ) + return { + **headers, + "Authorization": f"Bearer {access_token}", + "x-goog-user-project": project_id, + "Content-Type": "application/json", + } + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + stream: bool | None = None, + ) -> str: + vertex_params: Final = dict(litellm_params) + location: Final = validate_vertex_transcription_location( + self.safe_get_vertex_ai_location(vertex_params), default_location=DEFAULT_GEMINI_TRANSCRIBE_LOCATION + ) + project_id: Final = validate_vertex_transcription_project_id( + self.safe_get_vertex_ai_project(vertex_params) or self._resolve_project_id_from_credentials(vertex_params) + ) + base_url: Final = (api_base or get_vertex_base_url(location)).rstrip("/") + bare_model: Final = model.removeprefix("vertex_ai/") + model_path: Final = f"projects/{project_id}/locations/{location}/publishers/google/models/{bare_model}" + return f"{base_url}/v1/{model_path}:generateContent" + + def _resolve_project_id_from_credentials(self, litellm_params: Mapping[str, object]) -> str: + vertex_params: Final = dict(litellm_params) + _, project_id = self._ensure_access_token( + credentials=self.safe_get_vertex_ai_credentials(vertex_params), + project_id=None, + custom_llm_provider="vertex_ai", + ) + return project_id + + def transform_audio_transcription_request( + self, + model: str, + audio_file: FileTypes, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> AudioTranscriptionRequestData: + processed_audio: Final = process_audio_file(audio_file) + request_body: Final = VertexGeminiTranscriptionRequest( + contents=( + VertexGeminiTranscriptionContent( + role="user", + parts=( + VertexGeminiTranscriptionPart( + inlineData=VertexGeminiTranscriptionInlineData( + mimeType=processed_audio.content_type, + data=base64.b64encode(processed_audio.file_content).decode("utf-8"), + ) + ), + ), + ), + ), + generationConfig=VertexGeminiTranscriptionGenerationConfig( + audioTranscriptionConfig=_audio_transcription_config(optional_params.get("language")) + ), + ) + return AudioTranscriptionRequestData(data=dict(request_body)) + + def transform_audio_transcription_response( + self, + raw_response: Response, + ) -> TranscriptionResponse: + try: + response_json: Final = raw_response.json() + except ValueError: + raise VertexAIError( + status_code=raw_response.status_code, + message=f"Received non-JSON response from Vertex AI Gemini transcription: {raw_response.text}", + ) + parsed: Final = VertexGeminiTranscriptionResponse.model_validate(response_json) + texts: Final = tuple( + part.text + for candidate in parsed.candidates + if candidate.content is not None + for part in candidate.content.parts + if part.text + ) + response: Final = TranscriptionResponse(text=" ".join(texts)) + response["task"] = "transcribe" + usage: Final = parsed.usageMetadata + if usage is not None: + audio_tokens: Final = sum( + detail.tokenCount for detail in usage.promptTokensDetails if detail.modality == AUDIO_MODALITY + ) + response.usage = TranscriptionUsageTokensObject( + type="tokens", + input_tokens=usage.promptTokenCount, + output_tokens=usage.candidatesTokenCount, + total_tokens=usage.totalTokenCount, + input_token_details=TranscriptionUsageInputTokenDetailsObject( + audio_tokens=audio_tokens, + text_tokens=usage.promptTokenCount - audio_tokens, + ), + ) + return response + + +def _audio_transcription_config(language: object) -> VertexGeminiTranscriptionAudioConfig: + if not isinstance(language, str) or not language: + return VertexGeminiTranscriptionAudioConfig() + return VertexGeminiTranscriptionAudioConfig(languageCodes=(normalize_transcription_language_to_bcp47(language),)) diff --git a/litellm/llms/vertex_ai/audio_transcription/transformation.py b/litellm/llms/vertex_ai/audio_transcription/transformation.py index a352b2a34ca..db3504c9a6a 100644 --- a/litellm/llms/vertex_ai/audio_transcription/transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/transformation.py @@ -35,6 +35,19 @@ SUPPORTED_RESPONSE_FORMATS: Final = ("json", "text") _URL_UNSAFE_PROJECT_CHARS: Final = ("/", "?", "#", "\\", ":", " ", "\t", "\n", "\r") +def validate_vertex_transcription_location(location: str | None, default_location: str) -> str: + try: + return validate_vertex_location(location or default_location) + except ValueError as e: + raise VertexAIError(status_code=400, message=str(e)) from e + + +def validate_vertex_transcription_project_id(project_id: str) -> str: + if not project_id or ".." in project_id or any(c in project_id for c in _URL_UNSAFE_PROJECT_CHARS): + raise VertexAIError(status_code=400, message=f"Invalid vertex_project format: {project_id!r}") + return project_id + + class VertexAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexBase): def __init__(self) -> None: BaseAudioTranscriptionConfig.__init__(self) @@ -103,27 +116,16 @@ class VertexAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexBase) litellm_params: dict, stream: bool | None = None, ) -> str: - location: Final = self._validate_location(self.safe_get_vertex_ai_location(litellm_params)) - project_id: Final = self._validate_project_id( + location: Final = validate_vertex_transcription_location( + self.safe_get_vertex_ai_location(litellm_params), default_location=DEFAULT_SPEECH_TO_TEXT_LOCATION + ) + project_id: Final = validate_vertex_transcription_project_id( self.safe_get_vertex_ai_project(litellm_params) or self._resolve_project_id_from_credentials(litellm_params) ) host: Final = "speech.googleapis.com" if location == "global" else f"{location}-speech.googleapis.com" base_url: Final = (api_base or f"https://{host}").rstrip("/") return f"{base_url}/v2/projects/{project_id}/locations/{location}/recognizers/_:recognize" - @staticmethod - def _validate_location(location: str | None) -> str: - try: - return validate_vertex_location(location or DEFAULT_SPEECH_TO_TEXT_LOCATION) - except ValueError as e: - raise VertexAIError(status_code=400, message=str(e)) from e - - @staticmethod - def _validate_project_id(project_id: str) -> str: - if not project_id or ".." in project_id or any(c in project_id for c in _URL_UNSAFE_PROJECT_CHARS): - raise VertexAIError(status_code=400, message=f"Invalid vertex_project format: {project_id!r}") - return project_id - def _resolve_project_id_from_credentials(self, litellm_params: dict) -> str: _, project_id = self._ensure_access_token( credentials=self.safe_get_vertex_ai_credentials(litellm_params), 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/main.py b/litellm/main.py index cafa1e4718f..c341db08155 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8612,9 +8612,9 @@ def _joined_streamed_citations(streamed_citations: "tuple[object, ...]") -> "lis def _stream_builder_model_map_cost(response: ModelResponse) -> float | None: - model_name: Final = getattr(response, "model", None) + model_name: Final = response.model usage: Final = getattr(response, "usage", None) - if not isinstance(model_name, str) or not model_name or not isinstance(usage, Usage): + if not model_name or not isinstance(usage, Usage): return None try: prompt_cost, completion_tokens_cost = litellm.cost_per_token(model=model_name, usage_object=usage) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bebbcc32181..05c1cfd3179 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -12595,7 +12595,8 @@ "us": 1.1 }, "supports_output_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://docs.anthropic.com/en/docs/about-claude/models/overview" }, "claude-sonnet-4-6": { "deprecation_date": "2027-02-17", @@ -12997,7 +12998,8 @@ }, "supports_output_config": true, "prompt_cache_min_tokens": 512, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "source": "https://docs.anthropic.com/en/docs/about-claude/models/overview" }, "claude-opus-5": { "deprecation_date": "2027-07-24", @@ -13037,7 +13039,8 @@ }, "supports_output_config": true, "supports_speed": true, - "prompt_cache_min_tokens": 512 + "prompt_cache_min_tokens": 512, + "source": "https://docs.anthropic.com/en/docs/about-claude/models/overview" }, "claude-opus-4-8": { "deprecation_date": "2027-05-28", @@ -20562,7 +20565,8 @@ "output_cost_per_token": 2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_endpoints": [ - "/vertex_ai/live" + "/vertex_ai/live", + "/v1/realtime" ], "supported_modalities": [ "text", @@ -22015,6 +22019,49 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/nano-banana-pro-preview": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "rpm": 1000, + "tpm": 4000000, + "output_cost_per_token_batches": 6e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-3.1-flash-image": { "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, @@ -23047,7 +23094,8 @@ "supports_system_messages": true, "supports_video_input": true, "supports_vision": true, - "tpm": 800000 + "tpm": 800000, + "deprecation_date": "2026-09-30" }, "gemini/gemini-3.1-pro-preview": { "prompt_cache_min_tokens": 4096, @@ -23558,6 +23606,38 @@ "supports_tool_choice": true, "supports_vision": true }, + "gemini/gemma-4-26b-a4b-it": { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "gemini", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://ai.google.dev/gemini-api/docs/pricing" + }, + "gemini/gemma-4-31b-it": { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "gemini", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://ai.google.dev/gemini-api/docs/pricing" + }, "gemini/imagen-3.0-fast-generate-001": { "litellm_provider": "gemini", "mode": "image_generation", @@ -26097,7 +26177,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "low/1024-x-1536/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.013, @@ -26108,7 +26189,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "low/1536-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.013, @@ -26119,7 +26201,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "medium/1024-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.034, @@ -26130,7 +26213,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "medium/1024-x-1536/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.05, @@ -26141,7 +26225,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "medium/1536-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.05, @@ -26152,7 +26237,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "high/1024-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.133, @@ -26163,7 +26249,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "high/1024-x-1536/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.2, @@ -26174,7 +26261,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "high/1536-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.2, @@ -26185,7 +26273,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "standard/1024-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.009, @@ -26196,7 +26285,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "standard/1024-x-1536/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.013, @@ -26207,7 +26297,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "standard/1536-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.013, @@ -26218,7 +26309,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "1024-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.009, @@ -26229,7 +26321,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "1024-x-1536/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.013, @@ -26240,7 +26333,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "1536-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.013, @@ -26251,7 +26345,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "gpt-5": { "cache_read_input_token_cost": 1.25e-07, @@ -38583,7 +38678,6 @@ "input_cost_per_token": 1.04e-06, "litellm_provider": "together_ai", "max_input_tokens": 131072, - "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.04e-06, @@ -38705,7 +38799,6 @@ "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 131072, - "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 6e-07, @@ -38752,7 +38845,6 @@ "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_input_tokens": 200000, - "max_output_tokens": 200000, "max_tokens": 200000, "metadata": { "successor": "together_ai/zai-org/GLM-5.2" @@ -38770,7 +38862,6 @@ "input_cost_per_token": 4.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 200000, - "max_output_tokens": 200000, "max_tokens": 200000, "metadata": { "successor": "together_ai/zai-org/GLM-5.2" @@ -38788,7 +38879,6 @@ "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", "max_input_tokens": 256000, - "max_output_tokens": 256000, "max_tokens": 256000, "metadata": { "successor": "together_ai/moonshotai/Kimi-K3" @@ -38868,7 +38958,6 @@ "input_cost_per_token": 3e-07, "litellm_provider": "together_ai", "max_input_tokens": 524288, - "max_output_tokens": 524288, "max_tokens": 524288, "mode": "chat", "output_cost_per_token": 1.2e-06, @@ -38885,7 +38974,6 @@ "input_cost_per_token": 0.0, "litellm_provider": "together_ai", "max_input_tokens": 262144, - "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 0.0, @@ -38895,7 +38983,6 @@ "input_cost_per_token": 1.7e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, - "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 2.5e-07, @@ -38911,7 +38998,6 @@ "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", "max_input_tokens": 1000000, - "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 3e-06, @@ -38923,7 +39009,6 @@ "input_cost_per_token": 2.5e-06, "litellm_provider": "together_ai", "max_input_tokens": 1000000, - "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 7.5e-06, @@ -38934,21 +39019,19 @@ "input_cost_per_token": 3.2e-07, "litellm_provider": "together_ai", "max_input_tokens": 1000000, - "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 1.28e-06, "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/Qwen/Qwen3.8-2.4T-A95B": { - "cache_read_input_token_cost": 2.5e-07, - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2.5e-06, "litellm_provider": "together_ai", "max_input_tokens": 1010000, - "max_output_tokens": 1010000, "max_tokens": 1010000, "mode": "chat", - "output_cost_per_token": 6e-06, + "output_cost_per_token": 6.25e-06, "source": "https://docs.together.ai/docs/serverless-models", "supports_prompt_caching": true }, @@ -38956,7 +39039,6 @@ "input_cost_per_token": 1e-07, "litellm_provider": "together_ai", "max_input_tokens": 32768, - "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 1e-07, @@ -38967,7 +39049,6 @@ "input_cost_per_token": 1.4e-07, "litellm_provider": "together_ai", "max_input_tokens": 1048576, - "max_output_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 2.8e-07, @@ -38984,7 +39065,6 @@ "input_cost_per_token": 1.74e-06, "litellm_provider": "together_ai", "max_input_tokens": 512000, - "max_output_tokens": 512000, "max_tokens": 512000, "mode": "chat", "output_cost_per_token": 3.48e-06, @@ -39001,7 +39081,6 @@ "input_cost_per_token": 1.32e-06, "litellm_provider": "together_ai", "max_input_tokens": 1048576, - "max_output_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 3.96e-06, @@ -39017,7 +39096,6 @@ "input_cost_per_token": 6e-08, "litellm_provider": "together_ai", "max_input_tokens": 32768, - "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 1.2e-07, @@ -39027,7 +39105,6 @@ "input_cost_per_token": 3.9e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, - "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 9.7e-07, @@ -39053,7 +39130,6 @@ "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "max_input_tokens": 1048576, - "max_output_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 2e-07, @@ -39064,7 +39140,6 @@ "input_cost_per_token": 3.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 131072, - "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.5e-06, @@ -39077,7 +39152,6 @@ "input_cost_per_token": 9.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, - "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, @@ -39094,7 +39168,6 @@ "input_cost_per_token": 3e-06, "litellm_provider": "together_ai", "max_input_tokens": 1048576, - "max_output_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 1.5e-05, @@ -39118,7 +39191,6 @@ "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_input_tokens": 512288, - "max_output_tokens": 512288, "max_tokens": 512288, "mode": "chat", "output_cost_per_token": 3.6e-06, @@ -39135,7 +39207,6 @@ "input_cost_per_token": 2.8e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, - "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 8.6e-07, @@ -39146,7 +39217,6 @@ "input_cost_per_token": 1e-06, "litellm_provider": "together_ai", "max_input_tokens": 524288, - "max_output_tokens": 524288, "max_tokens": 524288, "mode": "chat", "output_cost_per_token": 4.05e-06, @@ -39162,7 +39232,6 @@ "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", "max_input_tokens": 524288, - "max_output_tokens": 524288, "max_tokens": 524288, "mode": "chat", "output_cost_per_token": 1.2e-06, @@ -39174,8 +39243,25 @@ "input_cost_per_token": 1.4e-06, "litellm_provider": "together_ai", "max_input_tokens": 1048575, - "max_output_tokens": 1048575, - "max_tokens": 1048575, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/zai-org/GLM-5.3": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048575, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 4.4e-06, "source": "https://docs.together.ai/docs/serverless-models", @@ -39191,8 +39277,8 @@ "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 1048575, - "max_output_tokens": 1048575, - "max_tokens": 1048575, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, "source": "https://docs.together.ai/docs/serverless-models", @@ -43318,6 +43404,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", @@ -44001,287 +44103,335 @@ ] }, "xai/grok-3": { - "cache_read_input_token_cost": 7.5e-07, - "input_cost_per_token": 3e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.5e-05, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-beta": { - "cache_read_input_token_cost": 7.5e-07, - "input_cost_per_token": 3e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.5e-05, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-fast-beta": { - "cache_read_input_token_cost": 1.25e-06, - "input_cost_per_token": 5e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.5e-05, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-fast-latest": { - "cache_read_input_token_cost": 1.25e-06, - "input_cost_per_token": 5e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.5e-05, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-latest": { - "cache_read_input_token_cost": 7.5e-07, - "input_cost_per_token": 3e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.5e-05, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-mini": { - "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost": 2e-07, "deprecation_date": "2026-02-28", - "input_cost_per_token": 3e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-mini-beta": { - "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost": 2e-07, "deprecation_date": "2026-02-28", - "input_cost_per_token": 3e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-mini-fast": { - "cache_read_input_token_cost": 1.5e-07, - "input_cost_per_token": 6e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 4e-06, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-02-28", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-mini-fast-beta": { - "cache_read_input_token_cost": 1.5e-07, - "input_cost_per_token": 6e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 4e-06, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-02-28", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-mini-fast-latest": { - "cache_read_input_token_cost": 1.5e-07, - "input_cost_per_token": 6e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 4e-06, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-02-28", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-mini-latest": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_token": 3e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-02-28", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4": { - "input_cost_per_token": 3e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-05, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-fast-reasoning": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 5e-07, - "output_cost_per_token_above_128k_tokens": 1e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-fast-non-reasoning": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 5e-07, - "output_cost_per_token_above_128k_tokens": 1e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-0709": { - "input_cost_per_token": 3e-06, - "input_cost_per_token_above_128k_tokens": 6e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-05, - "output_cost_per_token_above_128k_tokens": 3e-05, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-latest": { - "input_cost_per_token": 3e-06, - "input_cost_per_token_above_128k_tokens": 6e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-05, - "output_cost_per_token_above_128k_tokens": 3e-05, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-1-fast": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 5e-07, - "output_cost_per_token_above_128k_tokens": 1e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, "supports_function_calling": true, @@ -44290,19 +44440,21 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-1-fast-reasoning": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 5e-07, - "output_cost_per_token_above_128k_tokens": 1e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, "supports_function_calling": true, @@ -44312,19 +44464,20 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-1-fast-reasoning-latest": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 5e-07, - "output_cost_per_token_above_128k_tokens": 1e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, "supports_function_calling": true, @@ -44334,19 +44487,20 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-1-fast-non-reasoning": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 5e-07, - "output_cost_per_token_above_128k_tokens": 1e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-non-reasoning", "supports_audio_input": true, "supports_function_calling": true, @@ -44355,19 +44509,20 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-1-fast-non-reasoning-latest": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 5e-07, - "output_cost_per_token_above_128k_tokens": 1e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-non-reasoning", "supports_audio_input": true, "supports_function_calling": true, @@ -44376,7 +44531,10 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4.20-multi-agent-beta-0309": { "cache_read_input_token_cost": 2e-07, @@ -50390,7 +50548,7 @@ "max_tokens": 500000, "mode": "chat", "supports_function_calling": true, - "supports_prompt_caching": true, + "supports_prompt_caching": false, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true @@ -50405,7 +50563,7 @@ "max_tokens": 500000, "mode": "chat", "supports_function_calling": true, - "supports_prompt_caching": true, + "supports_prompt_caching": false, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true @@ -51828,6 +51986,43 @@ "tpm": 250000, "rpm": 10 }, + "vertex_ai/gemini-3.5-transcribe-preview": { + "input_cost_per_audio_token": 2.5e-06, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "vertex_ai", + "mode": "audio_transcription", + "output_cost_per_token": 1.2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "vertex_ai/gemini-3.5-transcribe-live-preview": { + "input_cost_per_audio_token": 3.5e-06, + "input_cost_per_token": 3.5e-06, + "litellm_provider": "vertex_ai", + "mode": "audio_transcription", + "output_cost_per_token": 2.1e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, "perplexity/pplx-embed-context-v1-0.6b": { "input_cost_per_token": 8e-09, "litellm_provider": "perplexity", @@ -54596,5 +54791,249 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true + }, + "groq/qwen/qwen3.8-27b": { + "input_cost_per_token": 8e-07, + "litellm_provider": "groq", + "max_input_tokens": 131042, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://console.groq.com/docs/model/qwen/qwen3.8-27b", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-medium-3.5": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-vibe-cli-latest": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-vibe-cli-with-tools": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-vibe-cli-fast": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-code-latest": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, + "litellm_provider": "mistral", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 9e-07, + "source": "https://docs.mistral.ai/models/model-cards/codestral-25-08", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/mistral-code-fim-latest": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, + "litellm_provider": "mistral", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 9e-07, + "source": "https://docs.mistral.ai/models/model-cards/codestral-25-08", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/mistral-code-agent-latest": { + "input_cost_per_token": 4e-07, + "litellm_provider": "mistral", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://mistral.ai/news/devstral-2-vibe-cli", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/mistral-ocr-3": { + "litellm_provider": "mistral", + "ocr_cost_per_page": 0.002, + "annotation_cost_per_page": 0.003, + "mode": "ocr", + "supported_endpoints": [ + "/v1/ocr" + ], + "source": "https://mistral.ai/pricing#api-pricing" + }, + "mistral/mistral-ocr-3-0": { + "litellm_provider": "mistral", + "ocr_cost_per_page": 0.002, + "annotation_cost_per_page": 0.003, + "mode": "ocr", + "supported_endpoints": [ + "/v1/ocr" + ], + "source": "https://mistral.ai/pricing#api-pricing" + }, + "mistral/mistral-ocr-4": { + "annotation_cost_per_page": 0.005, + "litellm_provider": "mistral", + "mode": "ocr", + "ocr_cost_per_page": 0.004, + "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", + "supported_endpoints": [ + "/v1/ocr" + ] + }, + "mistral/voxtral-mini-latest": { + "input_cost_per_second": 5e-05, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/model-cards/voxtral-mini-transcribe-26-02", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-realtime-2602": { + "input_cost_per_second": 0.0001, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/model-cards/voxtral-mini-transcribe-realtime-26-02", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-realtime-latest": { + "input_cost_per_second": 0.0001, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/model-cards/voxtral-mini-transcribe-realtime-26-02", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/labs-leanstral-1-5-1": { + "input_cost_per_token": 0.0, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://docs.mistral.ai/models/model-cards/leanstral-1-5", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "fireworks_ai/accounts/fireworks/models/glm-5p3": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/models/qwen3-embedding-8b": { + "input_cost_per_token": 1e-07, + "output_cost_per_token": 0.0, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 40960, + "max_output_tokens": 40960, + "max_tokens": 40960, + "mode": "embedding", + "source": "https://docs.fireworks.ai/serverless/pricing" } } 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..1d1b7f057cb 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -461,6 +461,145 @@ "access_groups": { "components": { "schemas": { + "AccessGroupBudget": { + "properties": { + "budget_duration": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Budget Duration" + }, + "budget_id": { + "title": "Budget Id", + "type": "string" + }, + "budget_reset_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Budget Reset At" + }, + "max_budget": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Max Budget" + }, + "soft_budget": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Soft Budget" + } + }, + "required": [ + "budget_id" + ], + "title": "AccessGroupBudget", + "type": "object" + }, + "AccessGroupBudgetRequest": { + "additionalProperties": false, + "properties": { + "budget_duration": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Budget Duration" + }, + "budget_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Budget Id" + }, + "max_budget": { + "anyOf": [ + { + "minimum": 0.0, + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Max Budget" + }, + "soft_budget": { + "anyOf": [ + { + "minimum": 0.0, + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Soft Budget" + } + }, + "title": "AccessGroupBudgetRequest", + "type": "object" + }, + "AccessGroupBudgetResponse": { + "properties": { + "access_group": { + "title": "Access Group", + "type": "string" + }, + "budget": { + "anyOf": [ + { + "$ref": "#/components/schemas/AccessGroupBudget" + }, + { + "type": "null" + } + ] + }, + "spend": { + "title": "Spend", + "type": "number" + } + }, + "required": [ + "access_group", + "spend" + ], + "title": "AccessGroupBudgetResponse", + "type": "object" + }, "AccessGroupCreateRequest": { "properties": { "access_agent_ids": { @@ -561,6 +700,16 @@ "title": "Access Group", "type": "string" }, + "budget": { + "anyOf": [ + { + "$ref": "#/components/schemas/AccessGroupBudget" + }, + { + "type": "null" + } + ] + }, "deployment_count": { "title": "Deployment Count", "type": "integer" @@ -571,6 +720,17 @@ }, "title": "Model Names", "type": "array" + }, + "spend": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Spend" } }, "required": [ @@ -782,6 +942,29 @@ "title": "AccessGroupUpdateRequest", "type": "object" }, + "DeleteAccessGroupBudgetResponse": { + "properties": { + "access_group": { + "title": "Access Group", + "type": "string" + }, + "budget_deleted": { + "title": "Budget Deleted", + "type": "boolean" + }, + "message": { + "title": "Message", + "type": "string" + } + }, + "required": [ + "access_group", + "budget_deleted", + "message" + ], + "title": "DeleteAccessGroupBudgetResponse", + "type": "object" + }, "DeleteModelGroupResponse": { "properties": { "access_group": { @@ -1072,6 +1255,156 @@ ] } }, + "/access_group/{access_group}/budget": { + "delete": { + "description": "Clear the shared budget of an access group, leaving the group itself in place.\n\nExample:\n```bash\ncurl -X DELETE 'http://localhost:4000/access_group/production-models/budget' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- DeleteAccessGroupBudgetResponse; budget_deleted is false when there was nothing to clear\n\nRaises:\n- HTTPException 404: If access group not found", + "operationId": "delete_access_group_budget_access_group__access_group__budget_delete", + "parameters": [ + { + "in": "path", + "name": "access_group", + "required": true, + "schema": { + "title": "Access Group", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/DeleteAccessGroupBudgetResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Delete Access Group Budget", + "tags": [ + "access_groups" + ] + }, + "get": { + "description": "Get the shared budget of an access group, and the spend drawn against it.\n\nExample:\n```bash\ncurl -X GET 'http://localhost:4000/access_group/production-models/budget' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- AccessGroupBudgetResponse; budget is null when the group has no budget set\n\nRaises:\n- HTTPException 404: If access group not found", + "operationId": "get_access_group_budget_access_group__access_group__budget_get", + "parameters": [ + { + "in": "path", + "name": "access_group", + "required": true, + "schema": { + "title": "Access Group", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccessGroupBudgetResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Access Group Budget", + "tags": [ + "access_groups" + ] + }, + "put": { + "description": "Set or replace the shared budget of an access group. Idempotent.\n\nEvery key that can reach a model in the group draws from this one budget.\n\nExample:\n```bash\ncurl -X PUT 'http://localhost:4000/access_group/production-models/budget' \\\n -H 'Authorization: Bearer sk-1234' \\\n -H 'Content-Type: application/json' \\\n -d '{\n \"max_budget\": 100.0,\n \"budget_duration\": \"30d\"\n }'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n- max_budget: Optional[float] - Requests fail once the group's shared spend exceeds this\n- soft_budget: Optional[float] - Fires an alert when reached; requests still succeed\n- budget_duration: Optional[str] - Frequency of resetting the group's spend (e.g. '30d')\n- budget_id: Optional[str] - Link an existing budget instead of creating one\n\nReturns:\n- AccessGroupBudgetResponse with the stored budget and current spend\n\nRaises:\n- HTTPException 400: If no budget field is given, or budget_duration cannot be parsed\n- HTTPException 404: If access group not found", + "operationId": "set_access_group_budget_access_group__access_group__budget_put", + "parameters": [ + { + "in": "path", + "name": "access_group", + "required": true, + "schema": { + "title": "Access Group", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccessGroupBudgetRequest" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccessGroupBudgetResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Set Access Group Budget", + "tags": [ + "access_groups" + ] + } + }, "/access_group/{access_group}/delete": { "delete": { "description": "Delete an access group.\n\nRemoves the access group from all deployments that have it.\n\nExample:\n```bash\ncurl -X DELETE 'http://localhost:4000/access_group/production-models/delete' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- DeleteModelGroupResponse with deletion details\n\nRaises:\n- HTTPException 404: If access group not found", @@ -1122,7 +1455,7 @@ }, "/access_group/{access_group}/info": { "get": { - "description": "Get information about a specific access group.\n\nExample:\n```bash\ncurl -X GET 'http://localhost:4000/access_group/production-models/info' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- AccessGroupInfo with the access group details\n\nRaises:\n- HTTPException 404: If access group not found", + "description": "Get information about a specific access group.\n\nExample:\n```bash\ncurl -X GET 'http://localhost:4000/access_group/production-models/info' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- AccessGroupInfo with the access group details, its shared budget and its spend\n\nRaises:\n- HTTPException 404: If access group not found", "operationId": "get_access_group_info_access_group__access_group__info_get", "parameters": [ { @@ -6459,6 +6792,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 +7090,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 +17177,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 +19769,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..cbd3571ec47 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -241,6 +241,7 @@ class Litellm_EntityType(enum.Enum): PROJECT = "project" TAG = "tag" AGENT = "agent" + MODEL_ACCESS_GROUP = "model_access_group" # global proxy level entity PROXY = "proxy" @@ -504,6 +505,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. @@ -2886,6 +2888,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob ), ) budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True) + matched_model_access_groups: list[str] | None = Field(default=None, exclude=True) budget_throttle_pct: float | None = Field(default=None, exclude=True) user: Any | None = None # Expanded user object when expand=user is used created_by_user: Any | None = None # Expanded created_by user when expand=user is used @@ -3573,6 +3576,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 @@ -4918,6 +4923,7 @@ class DBSpendUpdateTransactions(TypedDict): org_list_transactions: dict[str, float] | None tag_list_transactions: dict[str, float] | None agent_list_transactions: dict[str, float] | None + model_access_group_list_transactions: ReadOnly[dict[str, float] | None] class SpendUpdateQueueItem(TypedDict, total=False): diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index f742965ade2..7f0045c1d93 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -22,6 +22,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, create_response, + proxy_exception_from_http_exception, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.types.utils import TokenCountResponse @@ -214,6 +215,9 @@ async def anthropic_response( litellm_logging_obj=None, ) + if isinstance(e, HTTPException): + raise proxy_exception_from_http_exception(e, headers) + error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c1f9407cdad..5703c6cd5e8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -32,6 +32,7 @@ from litellm.constants import ( DEFAULT_MAX_RECURSE_DEPTH, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, END_USER_RESTRICTED_REGISTRY_MAX_SIZE, + MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE, REGISTRY_ERROR_NEGATIVE_CACHE_TTL, TAG_REGISTRY_MAX_SIZE, ) @@ -78,11 +79,15 @@ from litellm.proxy.common_utils.http_parsing_utils import ( from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import ( END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, + MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL, TAG_REGISTRY_OVERFLOW_SENTINEL, UserApiKeyCache, end_user_cache_key, end_user_restricted_registry_cache_key, get_management_object_ttl, + model_access_group_cache_key, + model_access_group_registry_cache_key, + model_access_group_spend_counter_key, object_permission_cache_key, tag_cache_key, tag_registry_cache_key, @@ -107,12 +112,14 @@ from litellm.repositories.table_repositories import ( EndUserRepository, JWTKeyMappingRepository, ManagedVectorStoresRepository, + ModelAccessGroupBudgetRepository, TagRepository, TeamMembershipRepository, ) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.router import Router +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.utils import get_utc_datetime from .auth_checks_organization import ( @@ -251,6 +258,43 @@ def _end_user_table(repo: _PrismaTableHolder[_PrismaEndUserRow]) -> _PrismaAuthT return repo.table +class _PrismaMaxBudgetRow(Protocol): + @property + def max_budget(self) -> float | None: ... + + +class _PrismaModelAccessGroupBudgetRow(Protocol): + access_group_name: str + + @property + def spend(self) -> float | None: ... + + @property + def litellm_budget_table(self) -> _PrismaMaxBudgetRow | None: ... + + +def _model_access_group_budget_table( + repo: _PrismaTableHolder[_PrismaModelAccessGroupBudgetRow], +) -> _PrismaAuthTable[_PrismaModelAccessGroupBudgetRow]: + return repo.table + + +class _MemberModelScope(Protocol): + @property + def allowed_models(self) -> Sequence[str] | None: ... + + +class _TeamMembershipModelScope(Protocol): + @property + def litellm_budget_table(self) -> _MemberModelScope | None: ... + + +def _member_allowed_models(membership: _TeamMembershipModelScope) -> Sequence[str]: + """The member's own model scope, read through a narrowed view of the membership row.""" + budget_table: Final = membership.litellm_budget_table + return () if budget_table is None else (budget_table.allowed_models or ()) + + class _RawCacheRead(Protocol): async def async_get_cache(self, *, key: str) -> object: ... @@ -807,6 +851,7 @@ async def common_checks( 1.1. If project is blocked 2. If team can call model 2.2 If project can call model + 2.3 Which model access groups authorized this request 3. If team is in budget 3.0.2. If project is in budget 3.0.3. If project is over soft budget (alert only) @@ -925,6 +970,18 @@ async def common_checks( proxy_logging_obj=proxy_logging_obj, ) + # 2.3 Which model access groups authorized this request + matched_model_access_groups: Final = await stamp_matched_model_access_groups( + model=_model, + valid_token=valid_token, + team_object=team_object, + project_object=project_object, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + # Run before apply_key_tags_pre_auth injects key metadata.tags into request_body. _reject_clientside_metadata_tags_check(general_settings, request_body, route) @@ -1004,6 +1061,13 @@ async def common_checks( proxy_logging_obj=proxy_logging_obj, valid_token=valid_token, ), + _model_access_group_max_budget_check( + matched_model_access_groups=matched_model_access_groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if matched_model_access_groups + else None, _user_max_budget_check(), _check_team_member_budget( team_object=team_object, @@ -1444,6 +1508,7 @@ _REGISTRY_NOT_CACHED: Final = _RegistryNotCached() #: One lock per registry; module-level because the stampede to collapse is worker-wide. _TAG_REGISTRY_LOAD_LOCK: Final = asyncio.Lock() _END_USER_REGISTRY_LOAD_LOCK: Final = asyncio.Lock() +_MODEL_ACCESS_GROUP_REGISTRY_LOAD_LOCK: Final = asyncio.Lock() async def _cached_registry( @@ -1836,6 +1901,105 @@ async def _load_tag_registry( ) +async def _load_model_access_group_registry( + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> frozenset[str] | None: + """The set of model access group names that have a row in ``LiteLLM_ModelAccessGroupBudgetTable``.""" + + async def fetch_ids() -> tuple[str, ...]: + registry_rows: Final = await _model_access_group_budget_table( + ModelAccessGroupBudgetRepository(prisma_client) + ).find_many(take=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE + 1) + return tuple(row.access_group_name for row in registry_rows) + + return await _load_bounded_registry( + cache_key=model_access_group_registry_cache_key(), + overflow_sentinel=MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL, + max_size=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE, + load_lock=_MODEL_ACCESS_GROUP_REGISTRY_LOAD_LOCK, + fetch_ids=fetch_ids, + user_api_key_cache=user_api_key_cache, + ) + + +async def _fetch_uncached_model_access_group_budgets( + uncached_groups: Sequence[str], + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> tuple[tuple[str, ModelAccessGroupBudget], ...]: + """Budget rows for the groups a cache probe missed. + + No registry gate here, unlike the tag path: the names only ever come from + ``matched_model_access_groups``, which :func:`collect_matched_model_access_groups` already + intersected with the registry, so a name that has no row cannot reach this. + """ + if not uncached_groups: + return () + + try: + db_rows: Final = await _model_access_group_budget_table( + ModelAccessGroupBudgetRepository(prisma_client) + ).find_many( + where={"access_group_name": {"in": list(uncached_groups)}}, + include={"litellm_budget_table": True}, + ) + fetched: Final = tuple((row.access_group_name, _model_access_group_budget(row)) for row in db_rows) + for fetched_name, fetched_obj in fetched: + await user_api_key_cache.async_set_cache( + key=model_access_group_cache_key(fetched_name), + value=fetched_obj, + model_type=ModelAccessGroupBudget, + ttl=get_management_object_ttl(user_api_key_cache), + ) + except Exception as e: # noqa: BLE001 # fail-safe: a budget fetch error must yield "no budget rows", never break auth + verbose_proxy_logger.debug("Error batch fetching model access group budgets from database: %s", e) + return () + else: + return fetched + + +def _model_access_group_budget(row: _PrismaModelAccessGroupBudgetRow) -> ModelAccessGroupBudget: + budget_table: Final = row.litellm_budget_table + return ModelAccessGroupBudget( + access_group_name=row.access_group_name, + spend=row.spend or 0.0, + max_budget=None if budget_table is None else budget_table.max_budget, + ) + + +@log_db_metrics +async def get_model_access_group_budgets_batch( + access_group_names: Sequence[str], + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> dict[str, ModelAccessGroupBudget]: + """Budget rows for the given model access groups, served from cache where possible. + + Shared by the two enforcement paths so they read one row per group per request: the + reservation counters when reservations are on, and :func:`_model_access_group_max_budget_check` + when ``disable_budget_reservation`` turns them off. + """ + if prisma_client is None or not access_group_names: + return {} + + probed: Final = [ + ( + group, + await user_api_key_cache.async_get_cache( + key=model_access_group_cache_key(group), model_type=ModelAccessGroupBudget + ), + ) + for group in access_group_names + ] + fetched: Final = await _fetch_uncached_model_access_group_budgets( + uncached_groups=tuple(group for group, budget in probed if budget is None), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return {group: budget for group, budget in (*probed, *fetched) if budget is not None} + + async def _fetch_uncached_tags( uncached_tags: Sequence[str], prisma_client: PrismaClient, @@ -3882,6 +4046,192 @@ def _resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> list[str] return models +def _model_access_groups_serving_model( + model: str | Sequence[str], + llm_router: Router, + team_id: str | None, +) -> frozenset[str]: + """Every model access group whose deployments serve the requested model(s).""" + requested: Final = (model,) if isinstance(model, str) else tuple(model) + return frozenset( + group + for requested_model in requested + for group in llm_router.get_model_access_groups(model_name=requested_model, team_id=team_id) + ) + + +async def _team_member_granted_models( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> Sequence[str]: + """The member's own ``allowed_models`` scope; empty when the member is not narrowed below the team.""" + if team_object is None or valid_token.user_id is None: + return () + + team_membership: Final = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + return () if team_membership is None else _member_allowed_models(team_membership) + + +async def _org_granted_models( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> Sequence[str]: + """The org allowlist reached through the key, or through its team when the key names no org.""" + org_id: Final = valid_token.org_id or (team_object.organization_id if team_object is not None else None) + if org_id is None: + return () + + try: + org_object: Final = await get_org_object( + org_id=org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # fail-safe: attribution degrades to "no org grant", it must never break auth + verbose_proxy_logger.debug("access group attribution: org lookup failed: %s", e) + return () + return org_object.models if org_object is not None else () + + +async def _granted_model_lists( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTableCachedObj | None, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[Sequence[str], ...]: + """One model allowlist per level that participates in authorizing the request.""" + return ( + _resolve_key_models_for_auth_check(valid_token=valid_token), + team_object.models if team_object is not None else (), + await _team_member_granted_models( + valid_token=valid_token, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + project_object.models if project_object is not None else (), + await _org_granted_models( + valid_token=valid_token, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + ) + + +async def collect_matched_model_access_groups( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTableCachedObj | None, + llm_router: Router | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[str, ...]: + """ + The budgeted model access groups that authorized this request, sorted and deduplicated. + + A group is charged only when its name appears on an allowlist the caller was granted -- key, + team, team-member scope, project or org -- *and* that group serves the requested model. Asking + for a model that merely belongs to a group attributes nothing, because nothing about the caller + named the group. + + Levels are unioned, never ranked: a team granted ``*`` whose member is scoped to one group is + still a caller gated by that group. An unrestricted allowlist (empty, ``*``) names no group and + so contributes nothing. + + The whole walk is gated on the budget registry, because collecting every match costs a full scan + of each allowlist where the plain access check stops at the first hit. An empty registry means no + group carries a budget, so there is nothing to attribute and no work worth doing. + """ + if model is None or valid_token is None or llm_router is None or prisma_client is None: + return () + + registry: Final = await _load_model_access_group_registry( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if registry is not None and not registry: + return () + + covering_groups: Final = _model_access_groups_serving_model( + model=model, + llm_router=llm_router, + team_id=valid_token.team_id, + ) + budgeted_groups: Final = covering_groups if registry is None else covering_groups & registry + if not budgeted_groups: + return () + + granted: Final = frozenset( + granted_model + for granted_models in await _granted_model_lists( + valid_token=valid_token, + team_object=team_object, + project_object=project_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + for granted_model in granted_models + ) + return tuple(sorted(budgeted_groups & granted)) + + +async def stamp_matched_model_access_groups( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTableCachedObj | None, + llm_router: Router | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[str, ...]: + """Record the groups that authorized this request on its auth object, for the post-call spend + writer and the reservation counters, and hand them back for the budget check.""" + if valid_token is None: + return () + + try: + matched: Final = await collect_matched_model_access_groups( + model=model, + valid_token=valid_token, + team_object=team_object, + project_object=project_object, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # fail-safe: attribution is spend telemetry, it must never break auth + verbose_proxy_logger.debug("model access group attribution failed: %s", e) + return () + if not matched: + return () + matched_groups: Final = list(matched) # mutable-ok: the auth field is typed list[str] | None + valid_token.matched_model_access_groups = matched_groups # rebind-ok: request-scoped carrier for the writer + return matched + + async def can_key_call_model( model: str | list[str], llm_model_list: list | None, @@ -4476,6 +4826,7 @@ async def _virtual_key_multi_budget_check( max_budget=w["max_budget"], window_entity_type="Key", window_entity_id=valid_token.token, + window_duration=str(w["budget_duration"]), window_start=get_budget_window_start(w), ) if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: @@ -4849,6 +5200,7 @@ async def _team_multi_budget_check( max_budget=w["max_budget"], window_entity_type="Team", window_entity_id=team_object.team_id, + window_duration=str(w["budget_duration"]), window_start=get_budget_window_start(w), ) if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: @@ -5256,6 +5608,61 @@ async def _tag_max_budget_check( ) +async def _model_access_group_max_budget_check( + matched_model_access_groups: Sequence[str], + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> None: + """Block the request when a model access group that authorized it is over its max budget. + + Only the groups auth already matched are charged and therefore only they are checked, so a + request that no budgeted group authorized costs nothing here. + + Like the tag check this is a plain read with no reservation, so concurrent requests can + overshoot the ceiling slightly. The reservation counters are the precise path; this one covers + the ``disable_budget_reservation`` case. + + The ceiling is exclusive, unlike the tag check it otherwise mirrors: a pool whose recorded + spend has reached ``max_budget`` has nothing left to give, so the next request is refused. + Keys and organizations already draw the line there. A non-positive budget means no budget, + matching what the reservation path treats as unbudgeted. + + Raises: + BudgetExceededError if a matched group is over its max budget. + """ + if prisma_client is None or not matched_model_access_groups: + return + + budgets: Final = await get_model_access_group_budgets_batch( + access_group_names=matched_model_access_groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + from litellm.proxy.proxy_server import get_current_spend + + for group in matched_model_access_groups: + budget = budgets.get(group) + if budget is None or budget.max_budget is None or budget.max_budget <= 0: + continue + + group_spend = await get_current_spend( + counter_key=model_access_group_spend_counter_key(group), + fallback_spend=budget.spend, + max_budget=budget.max_budget, + fallback_authoritative=True, + ) + if group_spend < budget.max_budget: + continue + raise litellm.BudgetExceededError( + current_cost=group_spend, + max_budget=budget.max_budget, + message=f"Budget has been exceeded! Model access group={group} Current cost: {group_spend}, Max budget: {budget.max_budget}", + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value, + entity_id=group, + ) + + def is_model_allowed_by_pattern(model: str, allowed_model_pattern: str) -> bool: """ Check if a model matches an allowed pattern. 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/common_request_processing.py b/litellm/proxy/common_request_processing.py index ff6c8d1b1f8..a23669af462 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -533,6 +533,21 @@ def _serialize_http_exception_detail( return str(detail), None +def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, str]) -> ProxyException: + raw_detail: Final = _getattr_object(exc, "detail", str(exc)) + message, structured_fields = _serialize_http_exception_detail(raw_detail) + existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {} + merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None) + return ProxyException( + message=message, + type=getattr(exc, "type", "None"), + param=getattr(exc, "param", "None"), + code=getattr(exc, "status_code", status.HTTP_400_BAD_REQUEST), + provider_specific_fields=merged_fields, + headers=headers, + ) + + def _collect_response_file_search_vector_store_ids(data: Mapping[str, object]) -> set[str]: vector_store_ids: Final[set[str]] = set() tools: Final = data.get("tools") @@ -1362,6 +1377,8 @@ def _classifier_cost_from_request_data(request_data: Mapping[str, object] | None routes and in `metadata` on chat-style routes, so both buckets are consulted, in the same precedence `get_or_create_metadata_bucket` writes them. """ + from litellm.proxy.spend_tracking.savings import classifier_cost_from_decision + data: Final = request_data or {} for metadata_key in ("litellm_metadata", "metadata"): metadata = data.get(metadata_key) @@ -1370,10 +1387,10 @@ def _classifier_cost_from_request_data(request_data: Mapping[str, object] | None decision = metadata.get("routing_decision") if not isinstance(decision, dict): continue - cost = decision.get("classifier_cost") - if isinstance(cost, bool) or not isinstance(cost, (int, float)): + cost = classifier_cost_from_decision(decision) + if cost is None: continue - return float(cost) + return cost return None @@ -3244,21 +3261,7 @@ class ProxyBaseLLMRequestProcessing: raise e if isinstance(e, HTTPException): - raw_detail: Final = _getattr_object(e, "detail", str(e)) - message, structured_fields = _serialize_http_exception_detail(raw_detail) - existing_fields: Final = getattr(e, "provider_specific_fields", None) or {} - if structured_fields: - merged_fields: dict | None = {**existing_fields, **structured_fields} - else: - merged_fields = existing_fields or None - raise ProxyException( - message=message, - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - provider_specific_fields=merged_fields, - headers=safe_headers, - ) + raise proxy_exception_from_http_exception(e, safe_headers) elif isinstance(e, httpx.HTTPStatusError): # Handle httpx.HTTPStatusError - extract actual error from response # This matches the original behavior before the refactor in commit 511d435f6f diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index d065b062517..1682cf12f4e 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -4,7 +4,7 @@ import math import time from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass, field -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from enum import Enum from types import MappingProxyType from typing import Final, Literal, Protocol, TypeVar, assert_never @@ -20,10 +20,12 @@ from litellm.constants import ( RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN, RESET_BUDGET_JOB_NAME, ) +from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import ( DB_RETRY_SAFE_ERROR_TYPES, LiteLLM_BudgetTableFull, LiteLLM_EndUserTable, + Litellm_EntityType, LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken, @@ -33,7 +35,12 @@ from litellm.proxy.common_utils.timezone_utils import ( compute_budget_reset_at, get_budget_reset_settings, ) -from litellm.proxy.common_utils.user_api_key_cache import tag_cache_key +from litellm.proxy.common_utils.user_api_key_cache import ( + model_access_group_cache_key, + model_access_group_spend_counter_key, + tag_cache_key, +) +from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -41,6 +48,7 @@ from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import SpendLinkedTable from litellm.repositories.table_repositories import ( EndUserRepository, + ModelAccessGroupBudgetRepository, TagRepository, TeamMembershipRepository, ) @@ -92,6 +100,11 @@ class _TagRow(_BudgetLinkedRow, Protocol): def tag_name(self) -> str: ... +class _ModelAccessGroupRow(_BudgetLinkedRow, Protocol): + @property + def access_group_name(self) -> str: ... + + class _EndUserRow(_BudgetLinkedRow, Protocol): @property def user_id(self) -> str: ... @@ -154,6 +167,14 @@ def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]: return (tag_cache_key(row.tag_name),) +def _model_access_group_counter_key(row: _ModelAccessGroupRow) -> str: + return model_access_group_spend_counter_key(row.access_group_name) + + +def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...]: + return (model_access_group_cache_key(row.access_group_name),) + + def _budget_link_where( budget_ids: Sequence[str], extra: Mapping[str, object] = MappingProxyType({}), @@ -329,6 +350,7 @@ class _WindowSource: table: str id_column: str + entity_type: Litellm_EntityType counter_prefix: str log_subject: str retry_subject: str @@ -353,6 +375,7 @@ _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = ( _WindowSource( table="LiteLLM_VerificationToken", id_column="token", + entity_type=Litellm_EntityType.KEY, counter_prefix="spend:key", log_subject="keys", retry_subject="key", @@ -361,6 +384,7 @@ _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = ( _WindowSource( table="LiteLLM_TeamTable", id_column="team_id", + entity_type=Litellm_EntityType.TEAM, counter_prefix="spend:team", log_subject="teams", retry_subject="team", @@ -610,6 +634,11 @@ class ResetBudgetJob: where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), log_subject="tags", ) + model_access_groups: Final[tuple[_ModelAccessGroupRow, ...]] = await self._fetch_linked_rows( + table=ModelAccessGroupBudgetRepository(self.prisma_client).table, + where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), + log_subject="model access groups", + ) rollover_caps: Final[Mapping[str, float]] = MappingProxyType( { # mutable-ok: MappingProxyType wraps a one-shot dict comprehension b.budget_id: cap @@ -639,6 +668,10 @@ class ResetBudgetJob: *((_key_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in keys), *((_org_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in orgs), *((_tag_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in tags), + *( + (_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps)) + for row in model_access_groups + ), ), rollover_caps=rollover_caps, cache_keys=( @@ -646,6 +679,7 @@ class ResetBudgetJob: *(key for row in keys for key in _key_cache_keys(row)), *(key for row in orgs for key in _org_cache_keys(row)), *(key for row in tags for key in _tag_cache_keys(row)), + *(key for row in model_access_groups for key in _model_access_group_cache_keys(row)), ), ) @@ -671,6 +705,7 @@ class ResetBudgetJob: _queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE) _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE) + _queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE) _queue_enduser_resets(uow.endusers, cascade) for budget_id, budget_reset_at in cascade.budget_resets: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) @@ -714,7 +749,8 @@ class ResetBudgetJob: async def reset_budget_for_litellm_budget_table(self) -> None: """ Resets the spend a budget tier gates (end users, team members, keys, - orgs, tags) and advances the tier's budget_reset_at, atomically. + orgs, tags, model access groups) and advances the tier's + budget_reset_at, atomically. Caches are invalidated only after the transaction commits, so a failed run cannot leave a zeroed counter in front of an un-reset DB row. @@ -745,8 +781,9 @@ class ResetBudgetJob: return _ChunkOutcome(fetched=len(cascade.budgets), advanced=advanced) case _BudgetCascadeFailed(cascade=cascade, error=error): verbose_proxy_logger.exception( - "Failed to reset the budget table cascade (team member, enduser, org and tag spend, plus " - "budget_reset_at); nothing was committed and the budgets stay due for the next run: %s", + "Failed to reset the budget table cascade (team member, enduser, org, tag and model access " + "group spend, plus budget_reset_at); nothing was committed and the budgets stay due for the " + "next run: %s", error, exc_info=error, ) @@ -1210,6 +1247,9 @@ class ResetBudgetJob: spend_counter_cache: DualCache, now: datetime, reset_settings: BudgetResetSettings, + prisma_client: PrismaClient, + entity_type: Litellm_EntityType, + entity_id: str, ) -> bool: """Reset a single budget window if expired. Returns True if the window was reset.""" reset_at_str: Final = window.get("reset_at") @@ -1225,11 +1265,56 @@ class ResetBudgetJob: await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=new_value) except Exception as redis_err: verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err) - window["reset_at"] = compute_budget_reset_at( - budget_duration=window["budget_duration"], settings=reset_settings - ).isoformat() + budget_duration: Final = window["budget_duration"] + next_reset_at: Final = compute_budget_reset_at(budget_duration=budget_duration, settings=reset_settings) + window["reset_at"] = next_reset_at.isoformat() + await ResetBudgetJob._roll_window_spend_row( + prisma_client=prisma_client, + entity_type=entity_type, + entity_id=entity_id, + budget_duration=budget_duration, + next_reset_at=next_reset_at, + ) return True + @staticmethod + async def _roll_window_spend_row( + prisma_client: PrismaClient, + entity_type: Litellm_EntityType, + entity_id: str, + budget_duration: str, + next_reset_at: datetime, + ) -> None: + """Move this window's LiteLLM_BudgetWindowSpend row onto the window + that just started, so the maintained total the read path uses starts + from zero alongside the counter. + + Best effort: the row is an optimization over aggregating + LiteLLM_SpendLogs, so a failure here must not stop the remaining + windows from having their counters reset. + """ + try: + window_start: Final = next_reset_at - timedelta(seconds=duration_in_seconds(budget_duration)) + except Exception as e: # noqa: BLE001 # duration_in_seconds raises bare exceptions on bad input + verbose_proxy_logger.warning("Unparseable budget_duration %s: %s", budget_duration, e) + return + try: + await roll_window_spend_row( + prisma_client=prisma_client, + entity_type=entity_type.value, + entity_id=entity_id, + window_duration=budget_duration, + new_window_start=window_start, + ) + except Exception as e: # noqa: BLE001 # the row is best effort; counter resets must still land + verbose_proxy_logger.warning( + "Failed to roll budget window spend row for %s=%s window=%s: %s", + entity_type.value, + entity_id, + budget_duration, + e, + ) + @staticmethod async def _window_carried_spend( window: Mapping[str, object], counter_key: str, spend_counter_cache: DualCache @@ -1325,6 +1410,9 @@ class ResetBudgetJob: spend_counter_cache, now, self.reset_settings, + prisma_client=self.prisma_client, + entity_type=source.entity_type, + entity_id=row_id, ): changed = True if changed: diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index b8df0105b7b..589d8fe68d1 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -185,6 +185,32 @@ def tag_registry_cache_key() -> str: return "tag_registry" +#: Cached under ``model_access_group_registry_cache_key`` when the table exceeds +#: ``MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-group lookup. +MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL: Final = "__model_access_group_registry_overflow__" + + +def model_access_group_cache_key(access_group_name: str) -> str: + """Cache key one model access group budget row is stored under; shared so auth, spend tracking and the management endpoints cannot drift.""" + return f"model_access_group:{access_group_name}" + + +def model_access_group_registry_cache_key() -> str: + """Cache key for the set of model access group names that have a budget row.""" + return "model_access_group_registry" + + +def model_access_group_spend_counter_key(access_group_name: str) -> str: + """Spend counter key for one model access group; shared so its four owners cannot drift. + + The reservation path writes it up front, the cost callback writes it after the call, auth + reads it to enforce ``max_budget``, and the reset job clears it on rollover. A copy that + drifts in any one of them silently resets or reads a counter nobody else touches, which shows + up as a budget that never trips or never resets. + """ + return f"spend:model_access_group:{access_group_name}" + + #: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds #: ``END_USER_RESTRICTED_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-id fetch. END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL: Final = "__end_user_restricted_registry_overflow__" diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index 96192b884d8..9c637a62dc1 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -185,8 +185,10 @@ def build_autorouter_turn_transaction( of a request through the router) are excluded by their internal_call_origin stamp: they are not traffic a user sent, so counting them would manufacture sessions and savings in the adoption metrics. Failed requests served nothing and are excluded. - Cache facts are derived from the payload's own usage record through the savings - owner, never handed in beside it. + The classifier's charge still lands here exactly once, via the decision's own + classifier_cost folded into this turn's spend: the excluded classifier row is how + it was billed, the decision is how it is attributed. Cache facts are derived from + the payload's own usage record through the savings owner, never handed in beside it. """ if payload.get("status") != "success": return None @@ -204,9 +206,12 @@ def build_autorouter_turn_transaction( turn_at: Final = _turn_time_utc(str(payload.get("startTime") or "")) if turn_at is None: return None + from litellm.proxy.spend_tracking.savings import classifier_cost_from_decision + usage_object_raw: Final = metadata.get("usage_object") cache: Final = turn_cache_facts(usage_object_raw if isinstance(usage_object_raw, Mapping) else None) tier_raw: Final = routing_decision.get("tier") + classifier_cost: Final = classifier_cost_from_decision(routing_decision) return AutoRouterTurnTransaction( api_key=api_key, session_id=_bounded_session_id(session_id), @@ -216,7 +221,7 @@ def build_autorouter_turn_transaction( model=model, turn_at=turn_at, total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0), - spend=float(payload.get("spend") or 0.0), + spend=float(payload.get("spend") or 0.0) + (classifier_cost or 0.0), saved_spend=saved_spend, covered=cache.covered, cache_hit=cache.read_tokens > 0, diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py new file mode 100644 index 00000000000..f9188f95cfd --- /dev/null +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -0,0 +1,313 @@ +""" +Writer for LiteLLM_BudgetWindowSpend. + +The table holds one row per configured budget window whose window_start rolls +forward in place, so budget enforcement can read a maintained running total +instead of aggregating LiteLLM_SpendLogs every time a window counter goes cold +(issue #35766). Raw SQL rather than the Prisma upsert helper because the +conditional roll cannot be expressed through the query builder. + +Seeding a row that does not exist yet reads LiteLLM_SpendLogs once, excluding +the requests whose increments are in the same batch so neither source counts +them twice. One gap survives that exclusion: without the Redis transaction +buffer every pod flushes its own increments, so a row seeded by one pod can +include spend logs whose increments are still queued on another pod, and those +increments are added again when that pod flushes. That is bounded by a single +flush interval, happens at most once per window row, and only ever over-counts: +the seed never omits spend, because every increment not yet in the row still +reaches it on its own pod's next flush. A row therefore lags real spend by at +most one flush interval of queued increments, the same lag the SpendLogs +aggregate it replaces (and every other spend column) already has. +""" + +from collections.abc import Sequence +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Final, Protocol + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendTransaction, + to_naive_utc, + window_spend_group_key, +) + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +_SELECT_EXISTING_ROWS_SQL: Final = ( + 'SELECT entity_type, entity_id, window_duration FROM "LiteLLM_BudgetWindowSpend" ' + "WHERE (entity_type, entity_id, window_duration) " + "IN (SELECT * FROM unnest($1::text[], $2::text[], $3::text[]))" +) + +_UPSERT_WINDOW_SPEND_SQL: Final = ( + 'INSERT INTO "LiteLLM_BudgetWindowSpend" ' + "(entity_type, entity_id, window_duration, window_start, spend, created_at, updated_at) " + "VALUES ($1, $2, $3, ($4::timestamptz AT TIME ZONE 'UTC'), $5, " + "($7::timestamptz AT TIME ZONE 'UTC'), ($7::timestamptz AT TIME ZONE 'UTC')) " + "ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE SET " + "spend = CASE " + 'WHEN "LiteLLM_BudgetWindowSpend".window_start >= EXCLUDED.window_start ' + 'THEN "LiteLLM_BudgetWindowSpend".spend + $6 ' + "ELSE EXCLUDED.spend " + "END, " + 'window_start = GREATEST("LiteLLM_BudgetWindowSpend".window_start, EXCLUDED.window_start), ' + "updated_at = ($7::timestamptz AT TIME ZONE 'UTC')" +) + +_ROLL_WINDOW_SPEND_SQL: Final = ( + 'UPDATE "LiteLLM_BudgetWindowSpend" SET ' + "window_start = ($4::timestamptz AT TIME ZONE 'UTC'), " + "spend = 0, " + "updated_at = ($5::timestamptz AT TIME ZONE 'UTC') " + "WHERE entity_type = $1 AND entity_id = $2 AND window_duration = $3 " + "AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')" +) + +_SEED_FROM_SPEND_LOGS_KEY_SQL: Final = ( + 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " + "AND NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" +) + +_SEED_FROM_SPEND_LOGS_TEAM_SQL: Final = ( + 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " + "AND NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" +) + +_SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL: Final = ( + 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" +) + +_SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL: Final = ( + 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" +) + +_UPSERT_TRANSACTION_TIMEOUT: Final = timedelta(seconds=60) + + +class WindowSpendLogsAggregate(Protocol): + """Sums LiteLLM_SpendLogs for one entity since window_start, ignoring the + requests whose ids are handed in. + + Injected so the flush can be exercised without a database and so the + expensive aggregate stays swappable. + """ + + async def __call__( + self, + prisma_client: "PrismaClient", + entity_type: str, + entity_id: str, + window_start: datetime, + exclude_request_ids: Sequence[str], + exclude_started_at: datetime | None, + ) -> float | None: ... + + +async def spend_logs_total_excluding( + prisma_client: "PrismaClient", + entity_type: str, + entity_id: str, + window_start: datetime, + exclude_request_ids: Sequence[str], + exclude_started_at: datetime | None, +) -> float | None: + """LiteLLM_SpendLogs spend for one entity since window_start, minus the + requests already accounted for by the increments being flushed. + + The spend log writer drains its own queue on a ~2s poll whenever anything + is queued, while window increments flush on the much slower batch tick, so + by the time a window row is seeded its batch's log rows are normally + already in the table. Counting them in the seed and again in the increment + is what made a fresh row land at twice the true spend. + + The exclusion is bounded to rows that started at or after the batch's + earliest request. request_id can be chosen by the client + (x-litellm-call-id), so an unbounded exclusion would let a replayed old id + erase a historical row from the seed while its increment still lands. + Without a known start the batch's ids are not excluded at all: that can + only over-count once, which enforcement tolerates, whereas under-counting + is a budget bypass. + """ + if entity_type == Litellm_EntityType.KEY.value: + bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_KEY_SQL, _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL + elif entity_type == Litellm_EntityType.TEAM.value: + bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_TEAM_SQL, _SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL + else: + return None + rows: Final = ( + await prisma_client.db.query_raw(unbounded_sql, entity_id, window_start) + if exclude_started_at is None or not exclude_request_ids + else await prisma_client.db.query_raw( + bounded_sql, + entity_id, + window_start, + tuple(exclude_request_ids), + _exclusion_lower_bound(exclude_started_at), + ) + ) + if not rows: + return 0.0 + return float(rows[0].get("total") or 0.0) + + +def _exclusion_lower_bound(started_at: datetime) -> datetime: + """LiteLLM_SpendLogs.startTime is TIMESTAMP(3); floor to the second so a + millisecond rounding of the batch's own earliest row cannot slip under it.""" + return to_naive_utc(started_at).replace(microsecond=0) + + +def _primary_key(transaction: WindowSpendTransaction) -> tuple[str, str, str]: + return ( + transaction["entity_type"], + transaction["entity_id"], + transaction["window_duration"], + ) + + +async def _existing_primary_keys( + prisma_client: "PrismaClient", + transactions: tuple[WindowSpendTransaction, ...], +) -> frozenset[tuple[str, str, str]]: + rows: Final = await prisma_client.db.query_raw( + _SELECT_EXISTING_ROWS_SQL, + tuple(transaction["entity_type"] for transaction in transactions), + tuple(transaction["entity_id"] for transaction in transactions), + tuple(transaction["window_duration"] for transaction in transactions), + ) + return frozenset((row["entity_type"], row["entity_id"], row["window_duration"]) for row in rows or ()) + + +async def _seed_base_for_missing_row( + prisma_client: "PrismaClient", + transaction: WindowSpendTransaction, + existing_primary_keys: frozenset[tuple[str, str, str]], + spend_logs_aggregate: WindowSpendLogsAggregate, +) -> float: + """Spend already recorded for a window that has no row yet. + + This is the LiteLLM_SpendLogs aggregate the window counter reseed runs on + every cold counter today, but here it runs once per window lifetime and off + the request path, and it excludes this batch's own requests so they are + counted by their increments alone. + """ + if _primary_key(transaction) in existing_primary_keys: + return 0.0 + base: Final = await spend_logs_aggregate( + prisma_client=prisma_client, + entity_type=transaction["entity_type"], + entity_id=transaction["entity_id"], + window_start=datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc), + exclude_request_ids=transaction["request_ids"], + exclude_started_at=_transaction_started_at(transaction), + ) + return float(base or 0.0) + + +def _transaction_started_at(transaction: WindowSpendTransaction) -> datetime | None: + started_at: Final = transaction.get("started_at") + if started_at is None: + return None + return datetime.fromisoformat(started_at).replace(tzinfo=timezone.utc) + + +def _upsert_params( + transaction: WindowSpendTransaction, + seed_base: float, + now: datetime, +) -> tuple[str, str, str, datetime, float, float, datetime]: + """$5 is what a brand new row starts at (pre-existing spend plus this + increment); $6 is the increment alone, which is all an already-current row + may add. They are equal for every row that already existed, so a row is + never seeded twice when two pods flush the same new window.""" + increment: Final = float(transaction["spend"]) + return ( + transaction["entity_type"], + transaction["entity_id"], + transaction["window_duration"], + datetime.fromisoformat(transaction["window_start"]), + seed_base + increment, + increment, + now, + ) + + +async def commit_window_spend_updates( + prisma_client: "PrismaClient", + transactions: Sequence[WindowSpendTransaction], + spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_total_excluding, +) -> None: + """Apply aggregated window increments to LiteLLM_BudgetWindowSpend. + + An increment at or behind the row's window_start adds into the row (this is + how in-flight requests that raced a reset carry into the new window); an + increment ahead of it rolls the window and starts from that increment. + + Statements are ordered by primary key so concurrent pods take row locks in + the same order, with window_start breaking ties so an older window is + applied before the roll that supersedes it. + """ + if not transactions: + return + + ordered: Final = tuple(sorted(transactions, key=window_spend_group_key)) + existing_primary_keys: Final = await _existing_primary_keys( + prisma_client=prisma_client, + transactions=ordered, + ) + seed_bases: Final = tuple( + [ + await _seed_base_for_missing_row( + prisma_client=prisma_client, + transaction=transaction, + existing_primary_keys=existing_primary_keys, + spend_logs_aggregate=spend_logs_aggregate, + ) + for transaction in ordered + ] + ) + + now: Final = to_naive_utc(datetime.now(timezone.utc)) + verbose_proxy_logger.debug( + "Spend tracking - committing %d budget window spend upserts over %d existing rows", + len(ordered), + len(existing_primary_keys), + ) + async with ( + prisma_client.db.tx(timeout=_UPSERT_TRANSACTION_TIMEOUT) as db_transaction, + db_transaction.batch_() as batcher, + ): + for transaction, seed_base in zip(ordered, seed_bases): + batcher.execute_raw( + _UPSERT_WINDOW_SPEND_SQL, + *_upsert_params(transaction=transaction, seed_base=seed_base, now=now), + ) + + +async def roll_window_spend_row( + prisma_client: "PrismaClient", + entity_type: str, + entity_id: str, + window_duration: str, + new_window_start: datetime, +) -> None: + """Move a row onto the window that just started and zero its spend. + + Conditional on the stored window_start still being behind the new one so a + pod that already rolled the row (or increments that arrived under the new + window) are not clobbered. + """ + await prisma_client.db.execute_raw( + _ROLL_WINDOW_SPEND_SQL, + entity_type, + entity_id, + window_duration, + to_naive_utc(new_window_start), + to_naive_utc(datetime.now(timezone.utc)), + ) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 3f2777ff1f3..641c07914d9 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,6 +12,7 @@ import os import random import time import traceback +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload @@ -23,6 +24,7 @@ from litellm.constants import ( DB_SPEND_UPDATE_JOB_NAME, INTERNAL_CALL_ORIGIN_METADATA_KEY, ) +from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( DB_RETRY_SAFE_ERROR_TYPES, @@ -54,6 +56,10 @@ from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdate from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( ToolDiscoveryQueue, ) +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendTransaction, + WindowSpendUpdateQueue, +) from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING from litellm.proxy.spend_tracking.compression_savings import ( extract_compression_saved_tokens, @@ -86,6 +92,7 @@ class _SpendBatch(Protocol): litellm_organizationtable: BatchTable litellm_tagtable: BatchTable litellm_agentstable: BatchTable + litellm_modelaccessgroupbudgettable: BatchTable class _SpendBatchManager(Protocol): @@ -109,7 +116,7 @@ def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager: return tx -def _get_llm_router(): +def get_llm_router(): """The proxy's router, or None outside a running proxy. Injected rather than imported where it is used, so the savings computation stays @@ -123,6 +130,52 @@ def _get_llm_router(): return None +class _DeploymentLookup(Protocol): + def get_model_info(self, id: str) -> Mapping[str, object] | None: ... + + +def _served_model_access_groups( + router: _DeploymentLookup | None, + served_model_id: str | None, +) -> frozenset[str] | None: + """Access groups declared by the deployment that actually served the request. + + None when the served deployment cannot be identified, in which case the set + attributed at auth time stands unchanged. + """ + if router is None or not served_model_id: + return None + deployment: Final = router.get_model_info(id=served_model_id) + if deployment is None: + return None + model_info: Final = deployment.get("model_info") + if not isinstance(model_info, Mapping): + return None + declared: Final = model_info.get("access_groups") + if not isinstance(declared, (list, tuple)): + return frozenset() + return frozenset(group for group in declared if isinstance(group, str)) + + +def debitable_model_access_groups( + attributed: Sequence[str] | None, + served_model_id: str | None, + router: _DeploymentLookup | None, +) -> tuple[str, ...]: + """Groups to debit: the set attributed at auth time, narrowed to those the served model belongs to. + + The router may fall back to a model outside the pool auth reserved against, so the + attributed set is the hard upper bound: a group absent from it is never debited. + """ + ordered: Final = coerce_model_access_groups(attributed) + if not ordered: + return () + served: Final = _served_model_access_groups(router=router, served_model_id=served_model_id) + if served is None: + return ordered + return tuple(group for group in ordered if group in served) + + class DBSpendUpdateWriter: """ Module responsible for @@ -146,6 +199,7 @@ class DBSpendUpdateWriter: self.daily_agent_spend_update_queue = DailySpendUpdateQueue() self.daily_org_spend_update_queue = DailySpendUpdateQueue() self.daily_tag_spend_update_queue = DailySpendUpdateQueue() + self.window_spend_update_queue = WindowSpendUpdateQueue() async def update_database( # LiteLLM management object fields @@ -161,7 +215,11 @@ class DBSpendUpdateWriter: start_time: datetime | None, end_time: datetime | None, response_cost: float | None, - ): + ) -> str | None: + """Returns the LiteLLM_SpendLogs request_id this call was recorded + under, so the caller can tell the budget-window writer which log rows + its increments already cover. None when the payload could not be built. + """ from litellm.proxy.proxy_server import ( disable_spend_logs, litellm_proxy_budget_name, @@ -178,7 +236,7 @@ class DBSpendUpdateWriter: team_id, ) if ProxyUpdateSpend.disable_spend_updates() is True: - return + return None if token is not None and isinstance(token, str) and token.startswith("sk-"): hashed_token = hash_token(token=token) else: @@ -187,6 +245,7 @@ class DBSpendUpdateWriter: ## CREATE SPEND LOG PAYLOAD ## from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_logging_payload, + get_request_model_access_groups, ) payload: Final = get_logging_payload( @@ -239,6 +298,7 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, litellm_proxy_budget_name=litellm_proxy_budget_name, payload=payload, + request_model_access_groups=get_request_model_access_groups(kwargs), ) ) @@ -250,6 +310,7 @@ class DBSpendUpdateWriter: ) verbose_proxy_logger.debug("Runs spend update on all tables") + return payload.get("request_id") except Exception: spend_log_error( "Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue " @@ -262,6 +323,7 @@ class DBSpendUpdateWriter: org_id, end_user_id, ) + return None async def _enqueue_tool_usage_transaction( self, @@ -320,7 +382,7 @@ class DBSpendUpdateWriter: routing_decision=metadata.get("routing_decision"), usage_object=usage_object_raw if isinstance(usage_object_raw, dict) else None, model_id=payload.get("model_id"), - llm_router=_get_llm_router, + llm_router=get_llm_router, cost_breakdown=metadata.get("cost_breakdown"), recorded_autorouter_savings=metadata.get("autorouter_savings"), ) @@ -431,9 +493,10 @@ class DBSpendUpdateWriter: prisma_client: PrismaClient | None, litellm_proxy_budget_name: str | None, payload: SpendLogsPayload, + request_model_access_groups: Sequence[str] = (), ): """ - Runs all 11 spend-update helpers sequentially inside a single asyncio task. + Runs all 13 spend-update helpers sequentially inside a single asyncio task. Each helper is wrapped in try/except so one failure doesn't prevent the others. @@ -505,6 +568,14 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) + await self._update_model_access_group_db( + response_cost=response_cost, + request_model_access_groups=request_model_access_groups, + served_model_id=payload_copy.get("model_id"), + prisma_client=prisma_client, + router=get_llm_router(), + ) + _agent_id_for_spend: Final = payload_copy.get("agent_id") try: await self._update_agent_db( @@ -814,6 +885,50 @@ class DBSpendUpdateWriter: ) raise e + async def _update_model_access_group_db( + self, + response_cost: float | None, + request_model_access_groups: Sequence[str] | None, + served_model_id: str | None, + prisma_client: PrismaClient | None, + router: _DeploymentLookup | None = None, + ) -> None: + """ + Update spend for every model access group this request is billed against. + + Args: + response_cost: Cost of the request, charged in full to each group + request_model_access_groups: Groups attributed at auth time, the upper bound on what may be debited + served_model_id: Deployment id actually served, used to narrow the attributed set + prisma_client: Prisma client instance + router: Deployment lookup used to re-resolve groups after a fallback + """ + try: + if prisma_client is None: + return + + for model_access_group in debitable_model_access_groups( + attributed=request_model_access_groups, + served_model_id=served_model_id, + router=router, + ): + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, + entity_id=model_access_group, + response_cost=response_cost, + ) + ) + except Exception as e: # noqa: BLE001 # isolation: a helper failure must not stop the batch + spend_log_error( + "Spend tracking - failed to enqueue model access group spend update. " + "model_access_groups=%s, response_cost=%s - %s", + request_model_access_groups, + response_cost, + str(e), + exc=e, + ) + async def _insert_spend_log_to_db( self, payload: dict | SpendLogsPayload, @@ -895,6 +1010,7 @@ class DBSpendUpdateWriter: daily_org_spend_update_queue=self.daily_org_spend_update_queue, daily_end_user_spend_update_queue=self.daily_end_user_spend_update_queue, daily_agent_spend_update_queue=self.daily_agent_spend_update_queue, + window_spend_update_queue=self.window_spend_update_queue, ) # Only commit from redis to db if this pod is the leader @@ -913,6 +1029,7 @@ class DBSpendUpdateWriter: daily_org_spend_update_transactions, daily_end_user_spend_update_transactions, daily_agent_spend_update_transactions, + window_spend_update_transactions, ) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() uncommitted = { # mutable-ok: drives which popped categories still need re-queuing @@ -922,12 +1039,14 @@ class DBSpendUpdateWriter: "daily_org_spend_update_transactions": daily_org_spend_update_transactions, "daily_end_user_spend_update_transactions": daily_end_user_spend_update_transactions, "daily_agent_spend_update_transactions": daily_agent_spend_update_transactions, + "window_spend_update_transactions": window_spend_update_transactions, } if db_spend_update_transactions is not None: verbose_proxy_logger.info( "Spend tracking - committing spend updates from Redis to DB: " - "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d, agents=%d", + "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d, agents=%d, " + "model_access_groups=%d", len(db_spend_update_transactions.get("key_list_transactions") or {}), len(db_spend_update_transactions.get("user_list_transactions") or {}), len(db_spend_update_transactions.get("team_list_transactions") or {}), @@ -936,6 +1055,7 @@ class DBSpendUpdateWriter: len(db_spend_update_transactions.get("team_member_list_transactions") or {}), len(db_spend_update_transactions.get("tag_list_transactions") or {}), len(db_spend_update_transactions.get("agent_list_transactions") or {}), + len(db_spend_update_transactions.get("model_access_group_list_transactions") or {}), ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, @@ -989,6 +1109,12 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_agent_spend_update_transactions, ) uncommitted.pop("daily_agent_spend_update_transactions", None) + if window_spend_update_transactions is not None: + await DBSpendUpdateWriter._commit_window_spend_updates( + prisma_client=prisma_client, + window_spend_transactions=window_spend_update_transactions, + ) + uncommitted.pop("window_spend_update_transactions", None) except Exception as e: spend_log_error( "Spend tracking - failed to commit spend updates from Redis to DB. " @@ -1104,6 +1230,27 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_agent_spend_update_transactions, ) + ################## Budget Window Spend Update Transactions ################## + # Aggregate all in memory budget window spend transactions and commit to db + window_spend_update_transactions: Final = ( + await self.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() + ) + + try: + await DBSpendUpdateWriter._commit_window_spend_updates( + prisma_client=prisma_client, + window_spend_transactions=window_spend_update_transactions, + ) + except Exception as e: # noqa: BLE001 # the increments go back on the queue; the rest of the flush must run + spend_log_error( + "Spend tracking - failed to commit budget window spend updates. " + "Re-queued %d window increments for retry on next tick. Error: %s", + len(window_spend_update_transactions), + str(e), + exc=e, + ) + await self.window_spend_update_queue.update_queue.put(window_spend_update_transactions) + ################## Tool Registry Upserts ################## await self._flush_tool_discovery_queue(prisma_client=prisma_client) @@ -1168,6 +1315,28 @@ class DBSpendUpdateWriter: cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME, ) + @staticmethod + async def _commit_window_spend_updates( + prisma_client: PrismaClient, + window_spend_transactions: Sequence[WindowSpendTransaction], + ) -> None: + """ + Commit per-budget-window spend increments to LiteLLM_BudgetWindowSpend. + + Raises on failure so the caller re-queues the increments: budget + enforcement trusts a current row without reconciling it against + LiteLLM_SpendLogs, so a dropped increment would let the entity spend + past its window limit after the next counter reseed. + """ + from litellm.proxy.db.budget_window_spend_writer import ( + commit_window_spend_updates, + ) + + await commit_window_spend_updates( + prisma_client=prisma_client, + transactions=window_spend_transactions, + ) + async def _drain_and_commit_daily_tag_spend_from_redis( self, prisma_client: PrismaClient, @@ -1433,6 +1602,20 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) + ### UPDATE MODEL ACCESS GROUP TABLE ### + model_access_group_list_transactions: Final = db_spend_update_transactions.get( + "model_access_group_list_transactions" + ) + await DBSpendUpdateWriter._update_entity_spend_in_db( + entity_name="Model access group", + transactions=model_access_group_list_transactions, + table_accessor="litellm_modelaccessgroupbudgettable", + where_field="access_group_name", + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + ### UPDATE AGENT TABLE ### agent_list_transactions: Final = db_spend_update_transactions["agent_list_transactions"] await DBSpendUpdateWriter._update_entity_spend_in_db( @@ -1449,7 +1632,7 @@ class DBSpendUpdateWriter: async def _update_entity_spend_in_db( entity_name: str, transactions: dict[str, float] | None, - table_accessor: Literal["litellm_tagtable", "litellm_agentstable"], + table_accessor: Literal["litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable"], where_field: str, n_retry_times: int, prisma_client: PrismaClient, @@ -1884,7 +2067,7 @@ class DBSpendUpdateWriter: gateway_injected_cache=marks_gateway_injection(_metadata, payload.get("model_id")), routing_decision=_metadata.get("routing_decision"), model_id=payload.get("model_id"), - llm_router=_get_llm_router, + llm_router=get_llm_router, usage_object=usage_obj, cost_breakdown=_metadata.get("cost_breakdown"), recorded_autorouter_savings=_metadata.get("autorouter_savings"), diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index ad92902221a..4dd23270bf8 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -23,6 +23,7 @@ from litellm.constants import ( REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY, + REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( @@ -42,6 +43,10 @@ from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, ) from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendTransaction, + WindowSpendUpdateQueue, +) from litellm.secret_managers.main import str_to_bool from litellm.types.caching import ( RedisPipelineLpopOperation, @@ -65,6 +70,7 @@ _SpendTransactionField: TypeAlias = Literal[ "org_list_transactions", "tag_list_transactions", "agent_list_transactions", + "model_access_group_list_transactions", ] _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( @@ -76,6 +82,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( "org_list_transactions", "tag_list_transactions", "agent_list_transactions", + "model_access_group_list_transactions", ) _ValueT = TypeVar("_ValueT") @@ -180,6 +187,7 @@ class RedisUpdateBuffer: daily_org_spend_update_queue: DailySpendUpdateQueue, daily_end_user_spend_update_queue: DailySpendUpdateQueue, daily_agent_spend_update_queue: DailySpendUpdateQueue, + window_spend_update_queue: WindowSpendUpdateQueue | None = None, ): """ Stores the in-memory spend updates to Redis @@ -248,6 +256,11 @@ class RedisUpdateBuffer: daily_agent_spend_update_transactions: Final = ( await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) + window_spend_update_transactions: Final = ( + await window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() + if window_spend_update_queue is not None + else () + ) verbose_proxy_logger.debug("ALL DB SPEND UPDATE TRANSACTIONS: %s", db_spend_update_transactions) verbose_proxy_logger.debug("ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions) @@ -284,6 +297,11 @@ class RedisUpdateBuffer: REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE, ), + ( + window_spend_update_transactions, + REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, + ServiceTypes.REDIS_WINDOW_SPEND_UPDATE_QUEUE, + ), ] rpush_list: Final[list[RedisPipelineRpushOperation]] = [] @@ -324,12 +342,14 @@ class RedisUpdateBuffer: daily_org_spend_update_transactions=daily_org_spend_update_transactions, daily_end_user_spend_update_transactions=daily_end_user_spend_update_transactions, daily_agent_spend_update_transactions=daily_agent_spend_update_transactions, + window_spend_update_transactions=window_spend_update_transactions, spend_update_queue=spend_update_queue, daily_spend_update_queue=daily_spend_update_queue, daily_team_spend_update_queue=daily_team_spend_update_queue, daily_org_spend_update_queue=daily_org_spend_update_queue, daily_end_user_spend_update_queue=daily_end_user_spend_update_queue, daily_agent_spend_update_queue=daily_agent_spend_update_queue, + window_spend_update_queue=window_spend_update_queue, ) return @@ -349,12 +369,14 @@ class RedisUpdateBuffer: daily_org_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None, daily_end_user_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None, daily_agent_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None, + window_spend_update_transactions: tuple[WindowSpendTransaction, ...] | None, spend_update_queue: SpendUpdateQueue, daily_spend_update_queue: DailySpendUpdateQueue, daily_team_spend_update_queue: DailySpendUpdateQueue, daily_org_spend_update_queue: DailySpendUpdateQueue, daily_end_user_spend_update_queue: DailySpendUpdateQueue, daily_agent_spend_update_queue: DailySpendUpdateQueue, + window_spend_update_queue: WindowSpendUpdateQueue | None, ) -> None: """ Put drained-but-unpushed transactions back into in-memory queues. @@ -397,6 +419,10 @@ class RedisUpdateBuffer: Litellm_EntityType.AGENT, db_spend_update_transactions.get("agent_list_transactions"), ), + ( + Litellm_EntityType.MODEL_ACCESS_GROUP, + db_spend_update_transactions.get("model_access_group_list_transactions"), + ), ] for entity_type, entities in entity_entries: if not entities: @@ -424,6 +450,9 @@ class RedisUpdateBuffer: if daily_txns: await daily_queue.update_queue.put(daily_txns) + if window_spend_update_transactions and window_spend_update_queue is not None: + await window_spend_update_queue.update_queue.put(window_spend_update_transactions) + async def restore_transactions_to_redis( self, db_spend_update_transactions: DBSpendUpdateTransactions | None = None, @@ -433,6 +462,7 @@ class RedisUpdateBuffer: daily_end_user_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, daily_agent_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, daily_tag_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, + window_spend_update_transactions: Sequence[WindowSpendTransaction] | None = None, ) -> None: """ Re-push transactions that were popped from Redis but not committed to the DB. @@ -454,6 +484,7 @@ class RedisUpdateBuffer: (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY), (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY), (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY), + (window_spend_update_transactions, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY), ) rpush_list: Final = tuple( @@ -571,20 +602,22 @@ class RedisUpdateBuffer: dict[str, DailyOrganizationSpendTransaction] | None, dict[str, DailyEndUserSpendTransaction] | None, dict[str, DailyAgentSpendTransaction] | None, + tuple[WindowSpendTransaction, ...] | None, ]: """ - Drains the main 6 Redis buffer queues in a single pipeline round-trip. + Drains the main 7 Redis buffer queues in a single pipeline round-trip. - Returns a 6-tuple of parsed results in this order: + Returns a 7-tuple of parsed results in this order: 0: DBSpendUpdateTransactions 1: daily user spend 2: daily team spend 3: daily org spend 4: daily end-user spend 5: daily agent spend + 6: budget window spend """ if self.redis_cache is None: - return None, None, None, None, None, None + return None, None, None, None, None, None, None lpop_list: Final[list[RedisPipelineLpopOperation]] = [ RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), @@ -608,12 +641,16 @@ class RedisUpdateBuffer: key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ), + RedisPipelineLpopOperation( + key=REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, + count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + ), ] raw_results: Final = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list) # Pad with None if pipeline returned fewer results than expected - while len(raw_results) < 6: + while len(raw_results) < 7: raw_results.append(None) # Slot 0: DBSpendUpdateTransactions @@ -634,6 +671,14 @@ class RedisUpdateBuffer: aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(list_of_daily) daily_results.append(aggregated) + window_spend: Final = ( + WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( + tuple(json.loads(transaction) for transaction in raw_results[6]) + ) + if raw_results[6] is not None + else None + ) + return ( db_spend, cast(dict[str, DailyUserSpendTransaction] | None, daily_results[0]), @@ -641,6 +686,7 @@ class RedisUpdateBuffer: cast(dict[str, DailyOrganizationSpendTransaction] | None, daily_results[2]), cast(dict[str, DailyEndUserSpendTransaction] | None, daily_results[3]), cast(dict[str, DailyAgentSpendTransaction] | None, daily_results[4]), + window_spend, ) async def store_in_memory_daily_tag_spend_updates_in_redis( @@ -826,6 +872,9 @@ class RedisUpdateBuffer: org_list_transactions=_merged_entity_transactions(list_of_transactions, "org_list_transactions"), tag_list_transactions=_merged_entity_transactions(list_of_transactions, "tag_list_transactions"), agent_list_transactions=_merged_entity_transactions(list_of_transactions, "agent_list_transactions"), + model_access_group_list_transactions=_merged_entity_transactions( + list_of_transactions, "model_access_group_list_transactions" + ), ) async def _emit_new_item_added_to_redis_buffer_event( diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 57cb5e73b64..8c0076b10c1 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -139,6 +139,7 @@ class SpendUpdateQueue(BaseUpdateQueue): org_list_transactions={}, tag_list_transactions={}, agent_list_transactions={}, + model_access_group_list_transactions={}, ) # Map entity types to their corresponding transaction dictionary keys @@ -151,6 +152,7 @@ class SpendUpdateQueue(BaseUpdateQueue): Litellm_EntityType.ORGANIZATION: "org_list_transactions", Litellm_EntityType.TAG: "tag_list_transactions", Litellm_EntityType.AGENT: "agent_list_transactions", + Litellm_EntityType.MODEL_ACCESS_GROUP: "model_access_group_list_transactions", } for update in updates: @@ -190,6 +192,8 @@ class SpendUpdateQueue(BaseUpdateQueue): transactions_dict = db_spend_update_transactions["tag_list_transactions"] elif dict_key == "agent_list_transactions": transactions_dict = db_spend_update_transactions["agent_list_transactions"] + elif dict_key == "model_access_group_list_transactions": + transactions_dict = db_spend_update_transactions["model_access_group_list_transactions"] else: continue diff --git a/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py new file mode 100644 index 00000000000..04dea66165e --- /dev/null +++ b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py @@ -0,0 +1,176 @@ +""" +In memory buffer for per-budget-window spend increments. + +Kept separate from SpendUpdateQueue: an increment is only meaningful together +with the window it landed in, so two increments for the same entity must not be +merged when their window_start differs. +""" + +import asyncio +import math +from collections.abc import Sequence +from datetime import datetime, timezone +from itertools import chain, groupby +from typing import Final, TypedDict + +from typing_extensions import ReadOnly + +from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE +from litellm.proxy.db.db_transaction_queue.base_update_queue import BaseUpdateQueue + + +class WindowSpendTransaction(TypedDict): + """One increment for a single (entity, budget window) pair. + + window_start is an ISO-8601 string rather than a datetime so the + transaction survives the JSON round trip through the Redis buffer. + + request_ids carries the LiteLLM_SpendLogs ids this spend came from. The + one-time seed for a window that has no row yet subtracts them from its + LiteLLM_SpendLogs aggregate, because the spend log writer flushes on its + own ~2s poll and will usually have persisted these rows before the window + queue flushes; without the exclusion the seed and the increment would each + count them. + + started_at is the earliest request start in the batch. The seed only + subtracts a request_id whose LiteLLM_SpendLogs.startTime is at or after it, + so a client that replays an old id through x-litellm-call-id cannot make the + seed drop the historical row that id already paid for. + """ + + entity_type: ReadOnly[str] + entity_id: ReadOnly[str] + window_duration: ReadOnly[str] + window_start: ReadOnly[str] + spend: ReadOnly[float] + request_ids: ReadOnly[Sequence[str]] + started_at: ReadOnly[str | None] + + +def to_naive_utc(value: datetime) -> datetime: + """LiteLLM_BudgetWindowSpend.window_start is TIMESTAMP(3), which holds naive UTC.""" + if value.tzinfo is None: + return value + return value.astimezone(timezone.utc).replace(tzinfo=None) + + +def window_spend_group_key(transaction: WindowSpendTransaction) -> tuple[str, str, str, str]: + """Identity of a window increment: the row's primary key plus the window it + belongs to. Two increments only aggregate when all four match.""" + return ( + transaction["entity_type"], + transaction["entity_id"], + transaction["window_duration"], + transaction["window_start"], + ) + + +def build_window_spend_transaction( + entity_type: str, + entity_id: str, + window_duration: str, + window_start: datetime, + spend: float, + request_id: str | None = None, + started_at: datetime | None = None, +) -> WindowSpendTransaction: + return WindowSpendTransaction( + entity_type=entity_type, + entity_id=entity_id, + window_duration=window_duration, + window_start=to_naive_utc(window_start).isoformat(timespec="microseconds"), + spend=spend, + request_ids=() if request_id is None else (request_id,), + started_at=None + if started_at is None + else to_naive_utc(started_at.astimezone(timezone.utc)).isoformat(timespec="microseconds"), + ) + + +def _merge_window_spend_transactions( + payloads: tuple[WindowSpendTransaction, ...], +) -> WindowSpendTransaction: + first: Final = payloads[0] + started_ats: Final = tuple( + started_at for payload in payloads if (started_at := payload.get("started_at")) is not None + ) + return WindowSpendTransaction( + entity_type=first["entity_type"], + entity_id=first["entity_id"], + window_duration=first["window_duration"], + window_start=first["window_start"], + spend=math.fsum(payload["spend"] for payload in payloads), + request_ids=tuple(sorted(frozenset(chain.from_iterable(payload["request_ids"] for payload in payloads)))), + started_at=min(started_ats) if started_ats else None, + ) + + +class WindowSpendUpdateQueue(BaseUpdateQueue): + """ + In memory buffer for budget-window spend increments committed to + LiteLLM_BudgetWindowSpend. + + Add an update with the payload built by build_window_spend_transaction: + window_spend_update_queue.add_update( + build_window_spend_transaction( + entity_type="key", + entity_id="", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=0.02, + ) + ) + """ + + def __init__(self) -> None: + super().__init__() + self.update_queue: asyncio.Queue[tuple[WindowSpendTransaction, ...]] = asyncio.Queue( + maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE + ) + + async def add_update(self, update: WindowSpendTransaction) -> None: + """Enqueue an update.""" + verbose_proxy_logger.debug("Adding budget window spend update to queue: %s", update) + await self.update_queue.put((update,)) + if self.update_queue.qsize() >= self.MAX_SIZE_IN_MEMORY_QUEUE: + verbose_proxy_logger.warning( + "Budget window spend update queue is full. Aggregating all entries in queue to concatenate entries." + ) + await self.aggregate_queue_updates() + + async def aggregate_queue_updates(self) -> None: + """Collapse everything currently queued into a single aggregated update.""" + updates: Final = await self.flush_all_updates_from_in_memory_queue() + await self.update_queue.put(WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(updates)) + + async def flush_and_get_aggregated_window_spend_transactions( + self, + ) -> tuple[WindowSpendTransaction, ...]: + """Drain the queue and return the increments aggregated per window.""" + updates: Final = await self.flush_all_updates_from_in_memory_queue() + if len(updates) > 0: + verbose_proxy_logger.info( + "Spend tracking - flushed %d budget window spend update batches from in-memory queue", + len(updates), + ) + return WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(updates) + + @staticmethod + def get_aggregated_window_spend_transactions( + updates: Sequence[Sequence[WindowSpendTransaction]], + ) -> tuple[WindowSpendTransaction, ...]: + """Sum spend per (entity_type, entity_id, window_duration, window_start). + + Increments belonging to different windows stay separate even when they + share a primary key, so a window boundary crossed mid-tick does not fold + the new window's spend into the previous window's total. + + The result is ordered by that same key, which is the order the flush + needs: primary key first for cross-pod lock ordering, then window_start + so an older window is applied before the roll that supersedes it. + """ + ordered: Final = tuple(sorted(chain.from_iterable(updates), key=window_spend_group_key)) + return tuple( + _merge_window_spend_transactions(tuple(group)) for _, group in groupby(ordered, key=window_spend_group_key) + ) diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 1929e7d3fc8..be515392a17 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -61,6 +61,26 @@ class _RoutedActions: return getattr(self._writer_actions, name) +class WriterPinnedClient: + """PrismaClient-shaped view whose `.db` resolves to the writer while it is available. + + Read-after-write paths (e.g. the model reconcile a /model/new triggers to + verify its own just-committed row) must not read through a lagging read + replica: the row is not replayed there yet, so the reconcile concludes the + write is missing and fails the request even though it is durable (#38556). + + While the writer is degraded (`writer_unavailable`), the pin yields to the + routed wrapper so reconcile reads keep working from the replica: a proxy + that starts during a primary outage must still load DB-backed models, and + no read-after-write hazard exists then because writes are failing anyway. + """ + + __slots__ = ("db",) + + def __init__(self, db: "PrismaWrapper | RoutingPrismaWrapper") -> None: + self.db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) and not db.writer_unavailable else db + + class RoutingPrismaWrapper: """ Routes Prisma operations between a writer and a reader Prisma client. diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index deb9cd5ae25..7b3c261036e 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -14,14 +14,18 @@ memory in long-lived deployments. import asyncio from collections import OrderedDict -from datetime import datetime +from collections.abc import Mapping +from datetime import datetime, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar, Final, Optional from litellm._logging import verbose_proxy_logger from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy._types import Litellm_EntityType from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( + BudgetWindowSpendRepository, SpendLogsRepository, TeamMembershipRepository, ) @@ -36,6 +40,25 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient +_WINDOW_SPEND_ENTITY_TYPES: Final[Mapping[str, str]] = MappingProxyType( + { + "Key": Litellm_EntityType.KEY.value, + "Team": Litellm_EntityType.TEAM.value, + } +) + +_WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + { + "Key": "api_key", + "Team": "team_id", + } +) + + +def _as_utc(value: datetime) -> datetime: + return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) + + class SpendCounterReseed: """ Reseeds spend counters from the authoritative DB and warms the cache, @@ -205,6 +228,92 @@ class SpendCounterReseed: raise return current_value + @staticmethod + async def window_from_table( + prisma_client: Optional["PrismaClient"], + entity_type: str, + entity_id: str, + window_duration: str, + expected_window_start: datetime, + ) -> float | None: + """ + Read the maintained per-window spend row by primary key. + + Returns the row's spend only when the row belongs to the window the + caller is enforcing, i.e. ``row.window_start >= expected_window_start``. + A row at or past the expected start was rolled by a pod whose reset_at + was at least as fresh as this caller's, so it is trusted; an older row + means the window boundary was crossed and nothing has rolled the row + yet, so its spend belongs to a previous window. + + Returns None for a missing, stale or unreadable row so the caller falls + back to the spend-logs aggregate. ``entity_type`` is the counter-facing + label ("Key"/"Team"); anything else has no row and returns None. + """ + if prisma_client is None: + return None + row_entity_type: Final = _WINDOW_SPEND_ENTITY_TYPES.get(entity_type) + if row_entity_type is None: + return None + + try: + row: Final = await BudgetWindowSpendRepository(prisma_client).table.find_unique( + where={ + "entity_type_entity_id_window_duration": { + "entity_type": row_entity_type, + "entity_id": entity_id, + "window_duration": window_duration, + } + } + ) + except Exception: # noqa: BLE001 # any read failure (DB, stale prisma client) must degrade to the aggregate path + verbose_proxy_logger.exception( + "SpendCounterReseed.window_from_table: failed for %s=%s window=%s", + entity_type, + entity_id, + window_duration, + ) + return None + + if row is None: + return None + if _as_utc(row.window_start) < _as_utc(expected_window_start): + return None + return float(row.spend or 0.0) + + @staticmethod + async def window_from_db( + prisma_client: Optional["PrismaClient"], + entity_type: str, + entity_id: str, + window_duration: str | None, + window_start: datetime, + ) -> float | None: + """ + Authoritative window spend: the maintained row first, falling back to + the spend-logs aggregate only when no current row exists. + + The aggregate range-scans an unindexed table, so it must stay a + transitional path (window configured before the row existed) rather + than a steady-state read. + """ + if window_duration is not None: + from_table: Final = await SpendCounterReseed.window_from_table( + prisma_client=prisma_client, + entity_type=entity_type, + entity_id=entity_id, + window_duration=window_duration, + expected_window_start=window_start, + ) + if from_table is not None: + return from_table + return await SpendCounterReseed.window_from_spend_logs( + prisma_client=prisma_client, + entity_type=entity_type, + entity_id=entity_id, + window_start=window_start, + ) + @staticmethod async def window_from_spend_logs( prisma_client: Optional["PrismaClient"], @@ -215,20 +324,13 @@ class SpendCounterReseed: if prisma_client is None: return None - if entity_type == "Key": - group_field = "api_key" - where = { - "api_key": entity_id, - "startTime": {"gte": window_start}, - } - elif entity_type == "Team": - group_field = "team_id" - where = { - "team_id": entity_id, - "startTime": {"gte": window_start}, - } - else: + group_field: Final = _WINDOW_SPEND_LOG_FIELDS.get(entity_type) + if group_field is None: return None + where: Final = { + group_field: entity_id, + "startTime": {"gte": window_start}, + } try: response: Final = await SpendLogsRepository(prisma_client).table.group_by( @@ -258,6 +360,7 @@ class SpendCounterReseed: counter_key: str, entity_type: str, entity_id: str, + window_duration: str | None, window_start: datetime, ) -> float | None: lock: Final = await SpendCounterReseed._get_lock(counter_key) @@ -276,10 +379,11 @@ class SpendCounterReseed: if val is not None: return float(val) - window_spend: Final = await SpendCounterReseed.window_from_spend_logs( + window_spend: Final = await SpendCounterReseed.window_from_db( prisma_client=prisma_client, entity_type=entity_type, entity_id=entity_id, + window_duration=window_duration, window_start=window_start, ) if window_spend is None: 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/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 7791adeb41e..bcaffa8e91c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -301,12 +301,12 @@ class LakeraAIGuardrail(CustomGuardrail): explicit sync below a hot reload that changes mode would pass validation but keep dispatching on the stale event_hook. """ - new_event_hook: Final = getattr(litellm_params, "mode", None) or self.event_hook - prospective_payload: Final = getattr(litellm_params, "payload", None) - prospective_breakdown: Final = getattr(litellm_params, "breakdown", None) + new_event_hook: Final = litellm_params.mode or self.event_hook + prospective_payload: Final = litellm_params.payload + prospective_breakdown: Final = litellm_params.breakdown self._validate_advisory_config( - on_flagged=getattr(litellm_params, "on_flagged", None) or self.on_flagged, - advisory_system_message=getattr(litellm_params, "advisory_system_message", None), + on_flagged=litellm_params.on_flagged or self.on_flagged, + advisory_system_message=litellm_params.advisory_system_message, payload=self.payload if prospective_payload is None else prospective_payload, breakdown=self.breakdown if prospective_breakdown is None else prospective_breakdown, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index daeb91eb2bd..f834426d619 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -121,7 +121,7 @@ class QualifireGuardrail(CustomGuardrail): the live instance untouched instead of raising after it's already been corrupted. Mirrors LakeraAIGuardrail's own override of this same method. """ - prospective_on_flagged: Final = getattr(litellm_params, "on_flagged", None) or self.on_flagged + prospective_on_flagged: Final = litellm_params.on_flagged or self.on_flagged self._validate_on_flagged(prospective_on_flagged) super().update_in_memory_litellm_params(litellm_params=litellm_params) diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 90d5f6f4970..dc13c09dd38 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -413,14 +413,15 @@ class GuardrailRegistry: raise Exception(f"Error getting guardrail from DB: {e}") -def _apply_configured_bool_override(instance: CustomGuardrail, litellm_params: LitellmParams, param_name: str) -> None: - """Override ``instance.`` only when ``litellm_params`` explicitly - sets it, preserving whatever default the guardrail's own constructor chose +def _apply_configured_bool_overrides(instance: CustomGuardrail, litellm_params: LitellmParams) -> None: + """Override the parallel/raw-scan flags only when ``litellm_params`` explicitly + sets them, preserving whatever default the guardrail's own constructor chose otherwise (its constructor default may be True, so blindly copying an absent/None config value would silently clobber it back to False).""" - configured: Final = getattr(litellm_params, param_name, None) - if configured is not None: - setattr(instance, param_name, bool(configured)) + if litellm_params.run_in_parallel is not None: + instance.run_in_parallel = bool(litellm_params.run_in_parallel) + if litellm_params.scan_raw_request is not None: + instance.scan_raw_request = bool(litellm_params.scan_raw_request) class InMemoryGuardrailHandler: @@ -544,8 +545,7 @@ class InMemoryGuardrailHandler: "skip_tool_message_in_guardrail are enabled together, which excludes every message from " "scanning, so no request content would ever be scanned. Remove one of the two." ) - for override_param in ("run_in_parallel", "scan_raw_request"): - _apply_configured_bool_override(custom_guardrail_callback, litellm_params, override_param) + _apply_configured_bool_overrides(custom_guardrail_callback, litellm_params) parsed_guardrail: Final = Guardrail( guardrail_id=guardrail.get("guardrail_id"), @@ -803,7 +803,6 @@ class InMemoryGuardrailHandler: previous_guardrail: Final = self.IN_MEMORY_GUARDRAILS.get(guardrail_id) previous_source: Final = self._sources.get(guardrail_id, source) - # Remove from memory if exists (also removes from callbacks) if guardrail_id in self.IN_MEMORY_GUARDRAILS: self.delete_in_memory_guardrail(guardrail_id) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index f593e94b36f..f02901f0e97 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1,5 +1,6 @@ import asyncio import traceback +from collections.abc import Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast @@ -20,6 +21,10 @@ from litellm.proxy.auth.auth_checks import ( log_db_metrics, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.db.db_spend_update_writer import ( + debitable_model_access_groups, + get_llm_router, +) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.spend_tracking.spend_log_error_logger import ( should_suppress_spend_log_tracebacks, @@ -27,6 +32,7 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import ( ) from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_error_information_for_spend_logs, + get_request_model_access_groups, ) from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( @@ -258,6 +264,11 @@ class _ProxyDBLogger(CustomLogger): sl_object=sl_object, metadata=metadata, ) + model_access_groups: Final = debitable_model_access_groups( + attributed=get_request_model_access_groups(kwargs), + served_model_id=sl_object.get("model_id") if sl_object is not None else None, + router=get_llm_router(), + ) if response_cost is not None: user_api_key: Final = metadata.get("user_api_key", None) @@ -296,6 +307,7 @@ class _ProxyDBLogger(CustomLogger): response_cost=response_cost, budget_reservation=budget_reservation, request_tags=tags, + model_access_groups=model_access_groups, ) # update cache (fire-and-forget for backward compat: @@ -572,9 +584,10 @@ async def _update_database_and_spend_counters( response_cost: float, budget_reservation: dict | None, request_tags: list[str] | None = None, + model_access_groups: Sequence[str] | None = None, ) -> None: try: - await proxy_logging_obj.db_spend_update_writer.update_database( + spend_log_request_id = await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key, response_cost=response_cost, user_id=user_id, @@ -610,6 +623,9 @@ async def _update_database_and_spend_counters( budget_reservation=budget_reservation, end_user_id=end_user_id, tags=request_tags, + request_id=spend_log_request_id, + request_started_at=start_time, + model_access_groups=model_access_groups, ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b7688790e54..ae55b7ab906 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -29,6 +29,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( _request_blocked_callback_params, iter_client_callback_metadata_dicts, ) +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import ( is_url_destination_allowed_by_host, @@ -253,6 +254,7 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = ( "_code_interpreter_interception_converted_stream", "_code_interpreter_interception_sandbox_key", "_code_interpreter_interception_session_scoped", + "_headroom_interception_converted_stream", "max_agentic_loops", # Recomputed below from the actual caller-controlled timeout sources (headers and # body fields); a client-forged value here would let a request either dodge cooldown @@ -1377,6 +1379,10 @@ class LiteLLMProxyRequestSetup: ) if user_api_key_dict.budget_reservation is not None: data[_metadata_variable_name]["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation + if user_api_key_dict.matched_model_access_groups: + data[_metadata_variable_name][MODEL_ACCESS_GROUP_METADATA_KEY] = ( + user_api_key_dict.matched_model_access_groups + ) # UserAPIKeyAuth object for MCP server access control data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict.model_copy( update={ 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/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index e1a7645e988..dabb9334b16 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -2,18 +2,33 @@ Allow proxy admin to manage model access groups Endpoints here: -- POST /model_group/new - Create a new access group with multiple model names +- POST /access_group/new - Create a new access group with multiple model names +- GET /access_group/list - List every access group +- GET /access_group/{access_group}/info - Read one access group, including its budget +- PUT /access_group/{access_group}/update - Replace an access group's deployments +- DELETE /access_group/{access_group}/delete - Delete an access group and its budget +- GET /access_group/{access_group}/budget - Read an access group's shared budget and spend +- PUT /access_group/{access_group}/budget - Set or replace an access group's shared budget +- DELETE /access_group/{access_group}/budget - Clear an access group's shared budget """ import json from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Protocol +from datetime import datetime +from typing import TYPE_CHECKING, Annotated, Any, Final, Protocol from fastapi import APIRouter, Depends, HTTPException +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + model_access_group_cache_key, + model_access_group_registry_cache_key, +) +from litellm.proxy.management_endpoints.common_utils import validate_budget_duration # Clear cache and reload models to pick up the access group changes from litellm.proxy.management_endpoints.model_management_endpoints import ( @@ -22,10 +37,16 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( model_info_as_mapping, reload_serving_verdict, ) +from litellm.proxy.management_helpers.utils import handle_budget_for_entity from litellm.proxy.utils import PrismaClient from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.table_repositories import ModelAccessGroupBudgetRepository from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudget, + AccessGroupBudgetRequest, + AccessGroupBudgetResponse, AccessGroupInfo, + DeleteAccessGroupBudgetResponse, DeleteModelGroupResponse, ListAccessGroupsResponse, NewModelGroupRequest, @@ -36,7 +57,43 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import if TYPE_CHECKING: from litellm import Router -router: Final = APIRouter() +router: Final = APIRouter(tags=["model management"]) + +_AUTH_DEPENDENCIES: Final = (Depends(user_api_key_auth),) + + +class _ErrorDetail(TypedDict): + error: ReadOnly[str] + + +class _ModelAccessGroupWhere(TypedDict): + access_group_name: ReadOnly[str] + + +class _BudgetInclude(TypedDict): + litellm_budget_table: ReadOnly[bool] + + +class _ModelAccessGroupBudgetCreate(TypedDict): + access_group_name: ReadOnly[str] + budget_id: ReadOnly[str | None] + created_by: ReadOnly[str] + updated_by: ReadOnly[str] + + +class _ModelAccessGroupBudgetUpdate(TypedDict): + budget_id: ReadOnly[str | None] + updated_by: ReadOnly[str] + + +class _ModelAccessGroupBudgetUpsert(TypedDict): + create: ReadOnly[_ModelAccessGroupBudgetCreate] + update: ReadOnly[_ModelAccessGroupBudgetUpdate] + + +def _http_error(status_code: int, message: str) -> HTTPException: + detail: Final[_ErrorDetail] = {"error": message} + return HTTPException(status_code=status_code, detail=detail) class _DeploymentRow(Protocol): @@ -58,10 +115,140 @@ class _ModelTableClient(Protocol): async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... +class _BudgetRow(Protocol): + @property + def budget_id(self) -> str: ... + + @property + def max_budget(self) -> float | None: ... + + @property + def soft_budget(self) -> float | None: ... + + @property + def budget_duration(self) -> str | None: ... + + @property + def budget_reset_at(self) -> datetime | None: ... + + +class _ModelAccessGroupBudgetRow(Protocol): + @property + def spend(self) -> float: ... + + @property + def budget_id(self) -> str | None: ... + + @property + def litellm_budget_table(self) -> _BudgetRow | None: ... + + +class _ModelAccessGroupBudgetTableClient(Protocol): + async def find_unique( + self, *, where: Mapping[str, object], include: Mapping[str, object] | None = None + ) -> _ModelAccessGroupBudgetRow | None: ... + + async def upsert( + self, + *, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _ModelAccessGroupBudgetRow: ... + + async def delete(self, *, where: Mapping[str, object]) -> _ModelAccessGroupBudgetRow | None: ... + + def _model_table(prisma_client: PrismaClient) -> _ModelTableClient: return ModelRepository(prisma_client).table +def _model_access_group_budget_table(prisma_client: PrismaClient) -> _ModelAccessGroupBudgetTableClient: + return ModelAccessGroupBudgetRepository(prisma_client).table + + +def _prisma_client_or_500() -> PrismaClient: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise _http_error(500, "Database not connected.") + return prisma_client + + +def _auth_cache() -> UserApiKeyCache: + from litellm.proxy.proxy_server import user_api_key_cache + + return user_api_key_cache + + +async def _evict_model_access_group_cache_keys(access_group: str, auth_cache: UserApiKeyCache) -> None: + """ + Every endpoint that writes an access group budget row must call this, or the budget stays + unenforced until the TTL expires: auth gates the feature on a cached registry of the groups + that have a budget row, read cache-first with no freshness check. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, + ) + + await evict_and_broadcast( + cache_keys=(model_access_group_cache_key(access_group), model_access_group_registry_cache_key()), + user_api_key_cache=auth_cache, + ) + + +async def _model_access_group_budget_row( + access_group: str, prisma_client: PrismaClient +) -> _ModelAccessGroupBudgetRow | None: + where: Final[_ModelAccessGroupWhere] = {"access_group_name": access_group} + include: Final[_BudgetInclude] = {"litellm_budget_table": True} + return await _model_access_group_budget_table(prisma_client).find_unique(where=where, include=include) + + +def _budget_or_none(row: _ModelAccessGroupBudgetRow | None) -> AccessGroupBudget | None: + budget: Final = row.litellm_budget_table if row is not None else None + if budget is None: + return None + return AccessGroupBudget( + budget_id=budget.budget_id, + max_budget=budget.max_budget, + soft_budget=budget.soft_budget, + budget_duration=budget.budget_duration, + budget_reset_at=budget.budget_reset_at, + ) + + +def _budget_response(access_group: str, row: _ModelAccessGroupBudgetRow | None) -> AccessGroupBudgetResponse: + return AccessGroupBudgetResponse( + access_group=access_group, + spend=row.spend if row is not None else 0.0, + budget=_budget_or_none(row), + ) + + +async def _delete_model_access_group_budget_row( + access_group: str, prisma_client: PrismaClient, auth_cache: UserApiKeyCache +) -> bool: + """ + Drop the group's budget row only, matching /tag/delete: the LiteLLM_BudgetTable row survives + because the link is ON DELETE SET NULL and a budget_id an admin passed in may be shared with + other entities. + + Evicts unconditionally: a group with no row of its own can still be sitting in the cached + registry, so skipping the eviction when nothing was deleted would leave that stale. + """ + where: Final[_ModelAccessGroupWhere] = {"access_group_name": access_group} + row: Final = await _model_access_group_budget_table(prisma_client).delete(where=where) + await _evict_model_access_group_cache_keys(access_group, auth_cache) + return row is not None + + +async def _raise_404_if_model_access_group_missing(access_group: str, prisma_client: PrismaClient) -> None: + access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client) + if access_group not in access_groups_map: + raise _http_error(404, f"Access group '{access_group}' not found") + + def validate_models_exist(model_names: Sequence[str], llm_router: "Router | None") -> tuple[bool, Sequence[str]]: """ Validate that all requested model names exist in the router. @@ -356,13 +543,12 @@ async def get_all_access_groups_from_db( @router.post( "/access_group/new", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=NewModelGroupResponse, ) async def create_model_group( data: NewModelGroupRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ Create a new access group containing multiple model names. @@ -503,12 +689,11 @@ async def create_model_group( @router.get( "/access_group/list", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=ListAccessGroupsResponse, ) async def list_access_groups( - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ List all access groups. @@ -553,13 +738,12 @@ async def list_access_groups( @router.get( "/access_group/{access_group}/info", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=AccessGroupInfo, ) async def get_access_group_info( access_group: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ Get information about a specific access group. @@ -574,7 +758,7 @@ async def get_access_group_info( - access_group: str - The access group name (URL path parameter) Returns: - - AccessGroupInfo with the access group details + - AccessGroupInfo with the access group details, its shared budget and its spend Raises: - HTTPException 404: If access group not found @@ -596,7 +780,15 @@ async def get_access_group_info( detail={"error": f"Access group '{access_group}' not found"}, ) - return access_groups_map[access_group] + info: Final = access_groups_map[access_group] + budget_row: Final = await _model_access_group_budget_row(access_group, prisma_client) + return AccessGroupInfo( + access_group=info.access_group, + model_names=info.model_names, + deployment_count=info.deployment_count, + spend=budget_row.spend if budget_row is not None else 0.0, + budget=_budget_or_none(budget_row), + ) except HTTPException: raise @@ -610,14 +802,13 @@ async def get_access_group_info( @router.put( "/access_group/{access_group}/update", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=NewModelGroupResponse, ) async def update_access_group( access_group: str, data: UpdateModelGroupRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ Update an access group's model names. @@ -765,13 +956,13 @@ async def update_access_group( @router.delete( "/access_group/{access_group}/delete", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=DeleteModelGroupResponse, ) async def delete_access_group( access_group: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)], ): """ Delete an access group. @@ -835,6 +1026,13 @@ async def delete_access_group( removed_pairs: Final = tuple(pair for pair in removed if pair is not None) models_updated: Final = len(removed_pairs) + # Budget last, deliberately: failing here strands a budget row for a group already on no + # deployment (clutter), where the reverse order can leave a live group enforcing nothing. + # The LiteLLM_BudgetTable row it linked is left alone, as /tag/delete leaves a tag's. + await _delete_model_access_group_budget_row( + access_group=access_group, prisma_client=prisma_client, auth_cache=auth_cache + ) + # Clear cache and reload models to pick up the access group changes live_before_reload: Final = live_model_ids_snapshot() reload_outcome: Final = await clear_cache() @@ -864,3 +1062,162 @@ async def delete_access_group( status_code=500, detail={"error": f"Failed to delete access group: {e}"}, ) + + +@router.get( + "/access_group/{access_group}/budget", + dependencies=_AUTH_DEPENDENCIES, + response_model=AccessGroupBudgetResponse, +) +async def get_access_group_budget( + access_group: str, +) -> AccessGroupBudgetResponse: + """ + Get the shared budget of an access group, and the spend drawn against it. + + Example: + ```bash + curl -X GET 'http://localhost:4000/access_group/production-models/budget' \\ + -H 'Authorization: Bearer sk-1234' + ``` + + Parameters: + - access_group: str - The access group name (URL path parameter) + + Returns: + - AccessGroupBudgetResponse; budget is null when the group has no budget set + + Raises: + - HTTPException 404: If access group not found + """ + prisma_client: Final = _prisma_client_or_500() + await _raise_404_if_model_access_group_missing(access_group=access_group, prisma_client=prisma_client) + + return _budget_response( + access_group=access_group, + row=await _model_access_group_budget_row(access_group, prisma_client), + ) + + +@router.put( + "/access_group/{access_group}/budget", + dependencies=_AUTH_DEPENDENCIES, + response_model=AccessGroupBudgetResponse, +) +async def set_access_group_budget( + access_group: str, + data: AccessGroupBudgetRequest, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)], +) -> AccessGroupBudgetResponse: + """ + Set or replace the shared budget of an access group. Idempotent. + + Every key that can reach a model in the group draws from this one budget. + + Example: + ```bash + curl -X PUT 'http://localhost:4000/access_group/production-models/budget' \\ + -H 'Authorization: Bearer sk-1234' \\ + -H 'Content-Type: application/json' \\ + -d '{ + "max_budget": 100.0, + "budget_duration": "30d" + }' + ``` + + Parameters: + - access_group: str - The access group name (URL path parameter) + - max_budget: Optional[float] - Requests fail once the group's shared spend exceeds this + - soft_budget: Optional[float] - Fires an alert when reached; requests still succeed + - budget_duration: Optional[str] - Frequency of resetting the group's spend (e.g. '30d') + - budget_id: Optional[str] - Link an existing budget instead of creating one + + Returns: + - AccessGroupBudgetResponse with the stored budget and current spend + + Raises: + - HTTPException 400: If no budget field is given, or budget_duration cannot be parsed + - HTTPException 404: If access group not found + """ + from litellm.proxy.proxy_server import litellm_proxy_admin_name + + prisma_client: Final = _prisma_client_or_500() + if not data.model_dump(exclude_none=True): + raise _http_error(400, "One of max_budget, soft_budget, budget_duration or budget_id is required") + validate_budget_duration(data.budget_duration) + await _raise_404_if_model_access_group_missing(access_group=access_group, prisma_client=prisma_client) + + existing_row: Final = await _model_access_group_budget_row(access_group, prisma_client) + budget_id: Final = await handle_budget_for_entity( + data=data, + existing_budget_id=existing_row.budget_id if existing_row is not None else None, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) + actor: Final = user_api_key_dict.user_id or litellm_proxy_admin_name + upsert_data: Final[_ModelAccessGroupBudgetUpsert] = { + "create": { + "access_group_name": access_group, + "budget_id": budget_id, + "created_by": actor, + "updated_by": actor, + }, + "update": {"budget_id": budget_id, "updated_by": actor}, + } + where: Final[_ModelAccessGroupWhere] = {"access_group_name": access_group} + include: Final[_BudgetInclude] = {"litellm_budget_table": True} + row: Final = await _model_access_group_budget_table(prisma_client).upsert( + where=where, data=upsert_data, include=include + ) + await _evict_model_access_group_cache_keys(access_group, auth_cache) + + verbose_proxy_logger.info("Set budget %s on access group '%s'", budget_id, access_group) + return _budget_response(access_group=access_group, row=row) + + +@router.delete( + "/access_group/{access_group}/budget", + dependencies=_AUTH_DEPENDENCIES, + response_model=DeleteAccessGroupBudgetResponse, +) +async def delete_access_group_budget( + access_group: str, + auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)], +) -> DeleteAccessGroupBudgetResponse: + """ + Clear the shared budget of an access group, leaving the group itself in place. + + Example: + ```bash + curl -X DELETE 'http://localhost:4000/access_group/production-models/budget' \\ + -H 'Authorization: Bearer sk-1234' + ``` + + Parameters: + - access_group: str - The access group name (URL path parameter) + + Returns: + - DeleteAccessGroupBudgetResponse; budget_deleted is false when there was nothing to clear + + Raises: + - HTTPException 404: If access group not found + """ + prisma_client: Final = _prisma_client_or_500() + await _raise_404_if_model_access_group_missing(access_group=access_group, prisma_client=prisma_client) + + budget_deleted: Final = await _delete_model_access_group_budget_row( + access_group=access_group, + prisma_client=prisma_client, + auth_cache=auth_cache, + ) + return DeleteAccessGroupBudgetResponse( + access_group=access_group, + budget_deleted=budget_deleted, + message=( + f"Budget for access group '{access_group}' deleted successfully" + if budget_deleted + else f"Access group '{access_group}' has no budget to delete" + ), + ) diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 23cfef6576c..567d8375737 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -49,6 +49,7 @@ from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.managed_resources.isolation import ( build_owner_filter, can_access_resource, + resolve_resource_owner_id, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit @@ -686,7 +687,7 @@ async def _mint_or_reuse_object( "file_object": json.dumps(body_snapshot), "model_object_id": namespaced_model_object_id, "file_purpose": file_purpose, - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, }, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 3d60f4f5f3a..09d3dedaafa 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -47,6 +47,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, ) +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -577,6 +578,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): _metadata["user_api_key"] = user_api_key_dict.api_key _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span _metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation + _metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = user_api_key_dict.matched_model_access_groups # The per-model budget counters are keyed off these. get_sanitized_user_information_from_key # returns StandardLoggingUserAPIKeyMetadata, which carries no budget field, so without this # the post-call increment finds nothing and every passthrough request goes untracked and diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index a5619821197..50cb813c6fa 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -185,7 +185,7 @@ class PipelineExecutor: # snapshot instead of `data` (which earlier pass_data steps in # this same pipeline may have already rewritten), same reason # the normal sequential/parallel guardrail loops do this. - scans_raw_request: Final = getattr(callback, "scan_raw_request", False) + scans_raw_request: Final = callback.scan_raw_request hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data independent_snapshot(raw_request_snapshot) if scans_raw_request and raw_request_snapshot is not None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3a70750528f..654fe4a3f2e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -382,6 +382,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, end_user_cache_key, get_management_object_ttl, + model_access_group_cache_key, + model_access_group_spend_counter_key, tag_cache_key, ) from litellm.proxy.config_resolvers import resolve_fields @@ -396,6 +398,9 @@ from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( SPEND_LOG_CLEANUP_BOUND_SETTINGS, SpendLogCleanup, ) +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + build_window_spend_transaction, +) from litellm.proxy.db.exception_handler import ( PrismaDBExceptionHandler, call_with_db_reconnect_retry, @@ -2431,6 +2436,7 @@ async def get_current_spend( max_budget: float | None = None, window_entity_type: str | None = None, window_entity_id: str | None = None, + window_duration: str | None = None, window_start: datetime | None = None, fallback_authoritative: bool = False, ) -> float: @@ -2455,7 +2461,8 @@ async def get_current_spend( runs and a key can leak spend past ``max_budget`` indefinitely. The authoritative source depends on the counter: primary key/team/user/org counters read the DB row; per-window counters (``window_start`` supplied) - aggregate spend logs; end-user/tag counters have no DB row, so the caller's + read the maintained window-spend row and only aggregate spend logs when + that row is missing or stale; end-user/tag counters have no DB row, so the caller's ``fallback_spend`` (loaded fresh in auth) is authoritative. The DB read is skipped for healthy primary counters (counter at or above recorded spend) and cached in-process for a few seconds, so a persistently stale counter @@ -2480,6 +2487,7 @@ async def get_current_spend( counter_key=counter_key, window_entity_type=window_entity_type, window_entity_id=window_entity_id, + window_duration=window_duration, window_start=window_start, ) if authoritative is not None: @@ -2561,6 +2569,7 @@ async def _authoritative_floor_spend( counter_key: str, window_entity_type: str | None = None, window_entity_id: str | None = None, + window_duration: str | None = None, window_start: datetime | None = None, ) -> float | None: marker_key: Final = f"spend_db_floor:{counter_key}" @@ -2575,10 +2584,11 @@ async def _authoritative_floor_spend( and window_entity_id is not None and window_start is not None ): - db_spend = await SpendCounterReseed.window_from_spend_logs( + db_spend = await SpendCounterReseed.window_from_db( prisma_client=prisma_client, entity_type=window_entity_type, entity_id=window_entity_id, + window_duration=window_duration, window_start=window_start, ) if db_spend is None: @@ -2648,6 +2658,9 @@ async def increment_spend_counters( budget_reservation: dict | None = None, end_user_id: str | None = None, tags: list[str] | None = None, + request_id: str | None = None, + request_started_at: datetime | None = None, + model_access_groups: Sequence[str] | None = None, ): """ Atomically increment spend counters for budget enforcement. @@ -2701,15 +2714,28 @@ async def increment_spend_counters( return for window in key_budget_limits: duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration + key_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at key_window_counter = f"spend:key:{hashed_token}:window:{duration}" + key_window_start = get_budget_window_start(window) if key_window_counter not in reserved_counter_keys: await _init_and_increment_window_spend_counter( counter_key=key_window_counter, entity_type="Key", entity_id=hashed_token, - window_start=get_budget_window_start(window), + window_duration=duration, + window_start=key_window_start, increment=cost, ) + await _enqueue_window_spend_row_update( + entity_type=Litellm_EntityType.KEY, + entity_id=hashed_token, + reset_at=key_window_reset_at, + window_duration=duration, + window_start=key_window_start, + increment=cost, + request_id=request_id, + request_started_at=request_started_at, + ) async def _team_scope(scope_team_id: str) -> None: team_counter_key: Final = f"spend:team:{scope_team_id}" @@ -2732,15 +2758,28 @@ async def increment_spend_counters( return for window in team_budget_limits: duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration + team_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at team_window_counter = f"spend:team:{scope_team_id}:window:{duration}" + team_window_start = get_budget_window_start(window) if team_window_counter not in reserved_counter_keys: await _init_and_increment_window_spend_counter( counter_key=team_window_counter, entity_type="Team", entity_id=scope_team_id, - window_start=get_budget_window_start(window), + window_duration=duration, + window_start=team_window_start, increment=cost, ) + await _enqueue_window_spend_row_update( + entity_type=Litellm_EntityType.TEAM, + entity_id=scope_team_id, + reset_at=team_window_reset_at, + window_duration=duration, + window_start=team_window_start, + increment=cost, + request_id=request_id, + request_started_at=request_started_at, + ) async def _team_member_scope(scope_user_id: str, scope_team_id: str) -> None: team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}" @@ -2777,6 +2816,13 @@ async def increment_spend_counters( ) if end_user_id is not None or tags is not None else None, + _increment_model_access_group_spend_counters( + model_access_groups=model_access_groups, + response_cost=cost, + reserved_counter_keys=reserved_counter_keys, + ) + if model_access_groups + else None, _increment_org_spend_counter( org_id=org_id, response_cost=cost, @@ -2865,6 +2911,33 @@ async def _increment_end_user_and_tag_spend_counters( ) +async def _increment_model_access_group_spend_counters( + model_access_groups: Sequence[object], + response_cost: float, + reserved_counter_keys: set[str], +) -> None: + """Charge the model access groups that authorized this request. + + Without this the counter auth reads is written only by the reservation path, so + ``disable_budget_reservation`` would leave ``_model_access_group_max_budget_check`` enforcing + against the DB row's spend, which lags by up to the cache TTL. + + Typed ``object`` rather than ``str`` because the names reach the cost callback out of request + metadata, which the coercion upstream filters to a list but not to strings. A non-string that + slipped through would build a counter key nothing else ever reads. + """ + unique_groups: Final = tuple( + dict.fromkeys(group for group in model_access_groups if group and isinstance(group, str)) + ) + for group in unique_groups: + await _init_and_increment_unreserved_spend_counter( + counter_key=model_access_group_spend_counter_key(group), + source_cache_key=model_access_group_cache_key(group), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + + async def _increment_org_spend_counter( org_id: str | None, response_cost: float, @@ -2925,10 +2998,62 @@ async def _init_and_increment_spend_counter( await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) +async def _enqueue_window_spend_row_update( + entity_type: Litellm_EntityType, + entity_id: str, + reset_at: datetime | str | None, + window_duration: str, + window_start: datetime | None, + increment: float, + request_id: str | None, + request_started_at: datetime | None, +) -> None: + """Queue this request's cost against the LiteLLM_BudgetWindowSpend row for + the window, so enforcement can read a maintained total instead of + aggregating LiteLLM_SpendLogs. + + request_id is the LiteLLM_SpendLogs id this cost was recorded under and + request_started_at its startTime; the flush uses them to keep the one-time + seed from counting a request that its increment already covers. + + Enqueued even when the cache increment was skipped for a reserved counter: + the reservation only pre-charged the counter, and the row still owes the + actual cost. + + Windows with no reset_at slide with wall clock, so their window_start moves + on every request and no single row can represent them. Those are left to + the read path's LiteLLM_SpendLogs fallback rather than rewritten per + request. + """ + if window_start is None or not reset_at: + return + try: + await proxy_logging_obj.db_spend_update_writer.window_spend_update_queue.add_update( + build_window_spend_transaction( + entity_type=entity_type.value, + entity_id=entity_id, + window_duration=window_duration, + window_start=window_start, + spend=increment, + request_id=request_id, + started_at=request_started_at, + ) + ) + except Exception as e: # noqa: BLE001 # spend tracking must never fail the cost callback + verbose_proxy_logger.debug( + "Unable to enqueue budget window spend update for %s=%s window=%s: %s", + entity_type.value, + entity_id, + window_duration, + e, + ) + + async def _init_and_increment_window_spend_counter( counter_key: str, entity_type: str, entity_id: str, + window_duration: str | None, window_start: datetime | None, increment: float, ): @@ -2943,6 +3068,7 @@ async def _init_and_increment_window_spend_counter( counter_key=counter_key, entity_type=entity_type, entity_id=entity_id, + window_duration=window_duration, window_start=window_start, ) if initialized is False: @@ -2988,6 +3114,7 @@ async def _ensure_window_spend_counter_initialized( counter_key: str, entity_type: str, entity_id: str, + window_duration: str | None, window_start: datetime, ) -> bool: is_warm: Final = await _is_spend_counter_cache_warm(counter_key=counter_key) @@ -3000,6 +3127,7 @@ async def _ensure_window_spend_counter_initialized( counter_key=counter_key, entity_type=entity_type, entity_id=entity_id, + window_duration=window_duration, window_start=window_start, ) if window_spend is None: @@ -4281,6 +4409,8 @@ class ProxyConfig: self.config: dict[str, Any] = {} self._last_semantic_filter_config: dict[str, object] | None = None self._last_hashicorp_vault_config: dict[str, object] | None = None + self._last_cyberark_config: dict[str, object] | None = None # mutable-ok: change-detection cache + self._cyberark_boot_env: dict[str, str | None] | None = None # mutable-ok: deployment env snapshot, set once self.worker_registry: list[WorkerRegistryEntry] = [] self.config_sync_subscriber: ConfigSyncSubscriber | None = None self.auth_cache_invalidation_subscriber: AuthCacheInvalidationSubscriber | None = None @@ -6764,9 +6894,18 @@ class ProxyConfig: - list: the rows (may be empty if no models exist) - None: signals a DB fetch *failure* — callers must not treat this as "all models deleted" and must not evict existing router deployments. + + Pinned to the writer DB: this read reconciles the router against the rows a + model write just committed, and reading it through a lagging read replica + makes the write-triggered reload report its own durable write as missing + (#38556). It also keeps a stale replica snapshot from evicting a deployment + another pod just added. While the writer is degraded the pin yields to the + replica so reader-only mode keeps loading DB-backed models. """ try: - new_models: Final[Sequence[_ProxyModelRow]] = await ModelRepository(prisma_client).table.find_many() + new_models: Final[Sequence[_ProxyModelRow]] = await ModelRepository( + WriterPinnedClient(prisma_client.db) + ).table.find_many() return new_models except Exception as e: verbose_proxy_logger.exception( @@ -6968,6 +7107,7 @@ class ProxyConfig: if self._should_load_db_object(object_type="config_overrides"): await self._init_hashicorp_vault_config_override(prisma_client=prisma_client) + await self._init_cyberark_config_override(prisma_client=prisma_client) await self._apply_safe_litellm_settings_overrides_from_db(prisma_client=prisma_client) @@ -7132,6 +7272,64 @@ class ProxyConfig: str(e), ) + async def _init_cyberark_config_override(self, prisma_client: PrismaClient) -> None: + """ + Load CyberArk Conjur config override from DB. + Decrypts sensitive fields, sets CYBERARK_* env vars, and reinitializes the secret manager. + Called periodically via _init_non_llm_objects_in_db to sync config across pods. + """ + from litellm.proxy.management_endpoints.config_override_endpoints import ( + CYBERARK_ENV_VAR_MAPPING, + _clear_cyberark_state, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module + _get_current_env_values, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module + _parse_config_value, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module + _set_env_vars, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module + _snapshot_cyberark_boot_env, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module + ) + + try: + db_record: Final[_ConfigOverridesRow | None] = cast( # cast-ok: prisma Json stub is `str`, runtime dict + "_ConfigOverridesRow | None", + await call_with_db_reconnect_retry( + prisma_client, + lambda: ConfigOverridesRepository(prisma_client).table.find_unique( + where={"config_type": "cyberark"} # mutable-ok: prisma where clause + ), + reason="init_cyberark_config_override_lookup_failure", + ), + ) + + if db_record is None or db_record.config_value is None: + if self._last_cyberark_config is not None: + _clear_cyberark_state(self) + return + + config_data: Final = _parse_config_value(db_record.config_value) + + # Skip reinit if config hasn't changed since last poll + if self._last_cyberark_config == config_data: + return + + decrypted_data: Final = self._decrypt_db_variables(config_data) + + _snapshot_cyberark_boot_env(self) + previous_env: Final = _get_current_env_values(CYBERARK_ENV_VAR_MAPPING) + _set_env_vars(decrypted_data, CYBERARK_ENV_VAR_MAPPING) + + try: + self.initialize_secret_manager(key_management_system="cyberark") + except Exception: + _set_env_vars(previous_env, CYBERARK_ENV_VAR_MAPPING) + raise + + self._last_cyberark_config = config_data.copy() + verbose_proxy_logger.debug("CyberArk config override loaded from DB") + except Exception as e: # noqa: BLE001 # any DB/decrypt/init failure must not break proxy boot + verbose_proxy_logger.exception( + "Error loading CyberArk config override from DB: %s", + str(e), + ) + async def check_periodic_reloads(self, prisma_client: PrismaClient): """ Run the admin-configured periodic model cost map reload. @@ -12013,6 +12211,7 @@ async def run_thread( # ) # async def get_available_routes(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)): from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.model_repository import ModelRepository from litellm.repositories.table_repositories import ( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 2bb850139a2..60223265211 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -29,6 +29,7 @@ model LiteLLM_BudgetTable { keys LiteLLM_VerificationToken[] // multiple keys can have the same budget end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget + model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization } @@ -585,6 +586,20 @@ model LiteLLM_EndUserTable { blocked Boolean @default(false) } +// Budget and shared spend for a model access group. The groups themselves are not rows anywhere: +// they are free-text strings in LiteLLM_ProxyModelTable.model_info.access_groups, so a row here +// exists only once someone gives that group a budget. +model LiteLLM_ModelAccessGroupBudgetTable { + access_group_name String @id + spend Float @default(0.0) + budget_id String? + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + created_at DateTime @default(now()) @map("created_at") + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? +} + // Track tags with budgets and spend model LiteLLM_TagTable { tag_name String @id @@ -649,6 +664,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 7cad3f0a022..8b2a5dd9312 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -12,7 +12,6 @@ from fastapi import HTTPException, status import litellm from litellm._logging import verbose_proxy_logger -from litellm.caching import DualCache from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate from litellm.proxy._types import ( @@ -26,12 +25,16 @@ from litellm.proxy.auth.auth_utils import get_model_from_request from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, end_user_cache_key, + model_access_group_cache_key, + model_access_group_spend_counter_key, tag_cache_key, team_membership_reservation_cache_key, ) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget @dataclass @@ -43,6 +46,7 @@ class _BudgetCounter: entity_id: str source_cache_key: str | None = None spend_log_entity_id: str | None = None + window_duration: str | None = None window_start: datetime | None = None @@ -53,6 +57,7 @@ _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = { "User": Litellm_EntityType.USER.value, "EndUser": Litellm_EntityType.END_USER.value, "Tag": Litellm_EntityType.TAG.value, + "Model access group": Litellm_EntityType.MODEL_ACCESS_GROUP.value, "Organization": Litellm_EntityType.ORGANIZATION.value, } @@ -158,7 +163,7 @@ async def reserve_budget_for_request( team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, prisma_client: PrismaClient | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, end_user_id: str | None = None, end_user_object: object = None, @@ -348,7 +353,7 @@ async def _get_budget_counters( team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, prisma_client: PrismaClient | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, end_user_id: str | None = None, end_user_object: object = None, @@ -437,6 +442,14 @@ async def _get_budget_counters( ) ) + counters.extend( + await _get_model_access_group_budget_counters( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + ) + team_member_counter: Final = await _get_team_member_budget_counter( valid_token=valid_token, team_object=team_object, @@ -491,7 +504,7 @@ async def _get_end_user_budget_counter( async def _get_tag_budget_counters( request_body: dict, prisma_client: PrismaClient | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> list[_BudgetCounter]: from litellm.proxy.auth.auth_checks import get_tag_objects_batch @@ -530,6 +543,46 @@ async def _get_tag_budget_counters( return counters +async def _get_model_access_group_budget_counters( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> list[_BudgetCounter]: + """Reservation counters for the model access groups that authorized this request. + + The names come off the auth object rather than the request body: ``common_checks`` already + resolved which granted groups serve the requested model, and re-deriving that here would both + duplicate the walk and risk disagreeing with what the spend writer attributes. + """ + from litellm.proxy.auth.auth_checks import get_model_access_group_budgets_batch + + group_names: Final = tuple(dict.fromkeys(valid_token.matched_model_access_groups or ())) + if not group_names: + return [] + + budgets: Final = await get_model_access_group_budgets_batch( + access_group_names=group_names, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + candidates: Final = (_model_access_group_counter(group, budgets.get(group)) for group in group_names) + return [counter for counter in candidates if counter is not None] + + +def _model_access_group_counter(group: str, budget: ModelAccessGroupBudget | None) -> _BudgetCounter | None: + """A counter for one group, or nothing when the group carries no budget to reserve against.""" + if budget is None or budget.max_budget is None or budget.max_budget <= 0: + return None + return _BudgetCounter( + counter_key=model_access_group_spend_counter_key(group), + source_cache_key=model_access_group_cache_key(group), + max_budget=budget.max_budget, + fallback_spend=budget.spend, + entity_type="Model access group", + entity_id=group, + ) + + def _dedupe_tags(tags: list[str]) -> list[str]: seen: Final = set() deduped_tags: Final = [] @@ -545,7 +598,7 @@ async def _get_team_member_budget_counter( valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ) -> _BudgetCounter | None: if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None: return None @@ -588,7 +641,7 @@ async def _get_team_member_budget_counter( async def _get_org_budget_counter( valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ) -> _BudgetCounter | None: org_id: str | None = None if valid_token.org_id is not None: @@ -657,6 +710,7 @@ def _get_budget_limit_counters( entity_type=entity_type, entity_id=f"{entity_id}:{budget_duration}", spend_log_entity_id=entity_id, + window_duration=str(budget_duration), window_start=window_start, ) ) @@ -700,6 +754,7 @@ async def _reserve_counter( counter_key=counter.counter_key, entity_type=counter.entity_type, entity_id=counter.spend_log_entity_id, + window_duration=counter.window_duration, window_start=counter.window_start, ) if initialized is False: diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py index 7d20aeeebac..1d0eb12da75 100644 --- a/litellm/proxy/spend_tracking/savings.py +++ b/litellm/proxy/spend_tracking/savings.py @@ -481,6 +481,20 @@ def _numeric_savings(value: object) -> float | None: return float(value) +def classifier_cost_from_decision(routing_decision: Mapping[str, object] | None) -> float | None: + """The LLM-classifier cost a routing decision recorded, or ``None`` when it holds none. + + ``None`` covers the decision-less request, the heuristic short-circuit that never + called a classifier, the unpriced classifier model, and a malformed value alike: + in every one of those cases there is no dollar figure to move, so callers treat + ``None`` as zero rather than as an error. The one owner of that reading, shared by + the savings netting, the session rollup and the response header, so the three can + never disagree about what counts as a classifier charge. + """ + decision: Final = routing_decision if isinstance(routing_decision, Mapping) else {} + return _numeric_savings(decision.get("classifier_cost")) + + def autorouter_savings_for_request( model: str | None, custom_llm_provider: str | None, @@ -490,7 +504,8 @@ def autorouter_savings_for_request( llm_router: "Callable[[], Router | None] | None" = None, cost_breakdown: Mapping[str, object] | None = None, ) -> float | None: - """Auto-router savings for one request, or ``None`` when the driver is off. + """Auto-router savings for one request, net of the classifier call that routed it, + or ``None`` when the driver is off. ``None`` and ``0.0`` are different facts: ``None`` means this request cannot carry a figure at all (no routing decision, no baseline, unusable usage), while ``0.0`` is a @@ -498,6 +513,11 @@ def autorouter_savings_for_request( Never raises: pricing failures inside degrade to zero, and the driver-off cases return ``None``, so this is safe on the logging path where a raise would fail the request's logging. + + The classifier deduction lives here, at the figure's one computation owner, rather + than in any reader: the stamped ``autorouter_savings`` is then already net, so the + session rollup, the daily tables and every logging consumer agree without each + re-deriving the deduction, and the recorded-figure-wins path cannot deduct twice. """ usage: Final = _usage_from_spend_log(usage_object) if usage is None or not model: @@ -510,7 +530,7 @@ def autorouter_savings_for_request( if not decision or not baseline_model: return None router_instance: Final = llm_router() if llm_router else None - return compute_autorouter_savings( + gross: Final = compute_autorouter_savings( baseline_model=baseline_model, selected_model=model, selected_provider=custom_llm_provider, @@ -522,6 +542,8 @@ def autorouter_savings_for_request( baseline_info=_effective_model_info(router_instance, baseline_id, baseline_model or ""), cost_breakdown=cost_breakdown, ) + classifier_cost: Final = classifier_cost_from_decision(decision) + return gross if classifier_cost is None else gross - classifier_cost def autorouter_savings_for_logging_payload( diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index a6b6375fd9c..43709e4e6ff 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -1,6 +1,7 @@ import os import re import secrets +from collections.abc import Mapping from datetime import datetime, timezone from datetime import datetime as dt from typing import Any, Final, Literal, cast @@ -23,7 +24,11 @@ from litellm.litellm_core_utils.core_helpers import ( reconstruct_model_name, ) from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call -from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash +from litellm.litellm_core_utils.litellm_logging import ( + coerce_model_access_groups, + is_valid_sha256_hash, + request_model_access_groups_from_litellm_params, +) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error @@ -92,6 +97,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 +128,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 +163,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 @@ -243,6 +254,23 @@ def _extract_usage_for_ocr_call(response_obj: Any, response_obj_dict: dict) -> d return {} +def get_request_model_access_groups(kwargs: Mapping[str, object] | None) -> tuple[str, ...]: + """Model access groups that authorized this request, as stamped onto request metadata at auth time.""" + if kwargs is None: + return () + + standard_logging_payload: Final = kwargs.get("standard_logging_object") + if isinstance(standard_logging_payload, Mapping): + from_payload: Final = coerce_model_access_groups(standard_logging_payload.get("request_model_access_groups")) + if from_payload: + return from_payload + + litellm_params: Final = kwargs.get("litellm_params") + if not isinstance(litellm_params, Mapping): + return () + return request_model_access_groups_from_litellm_params(litellm_params) + + def _sl_attribution_fallback( standard_logging_payload: StandardLoggingPayload | None, field: Literal["model_id", "model_group", "api_base", "custom_llm_provider"], @@ -360,6 +388,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/proxy/utils.py b/litellm/proxy/utils.py index 1d16fa63607..cf56fc0b1dd 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1416,7 +1416,7 @@ class ProxyLogging: mutation is discarded and a warning is logged so the misconfiguration is visible instead of silently forwarding unredacted content. """ - scans_raw_request: Final = getattr(callback, "scan_raw_request", False) + scans_raw_request: Final = callback.scan_raw_request should_use_raw_snapshot: Final = scans_raw_request and raw_request_snapshot is not None input_data: Final = ( # mutable-ok: same request-payload shape as data independent_snapshot(raw_request_snapshot) if should_use_raw_snapshot else data @@ -1453,7 +1453,7 @@ class ProxyLogging: "scan_raw_request is for block-only guardrails and this mutation is being " "discarded. Remove scan_raw_request from this guardrail's config if it needs " "to mask/rewrite content.", - getattr(callback, "guardrail_name", None) or callback.__class__.__name__, + callback.guardrail_name or callback.__class__.__name__, ) if scans_raw_request: if result is not None: @@ -1778,7 +1778,7 @@ class ProxyLogging: # guarantee must hold even under litellm.safe_memory_mode, which # otherwise makes deep copies return the original object. needs_raw_request_snapshot: Final = any( - isinstance(cb, CustomGuardrail) and getattr(cb, "scan_raw_request", False) + isinstance(cb, CustomGuardrail) and cb.scan_raw_request for cb in ProxyLogging._callback_capabilities().resolved_callbacks ) raw_request_snapshot: Final[dict | None] = ( # mutable-ok: same request-payload shape as data @@ -1938,7 +1938,7 @@ class ProxyLogging: """ def _input_for(callback: CustomGuardrail) -> dict: # mutable-ok: same request-payload shape as data - if not getattr(callback, "scan_raw_request", False) or raw_request_snapshot is None: + if not callback.scan_raw_request or raw_request_snapshot is None: return data return independent_snapshot(raw_request_snapshot) @@ -1962,11 +1962,7 @@ class ProxyLogging: # deployment-level guardrail sharing this name would see no marker # via _pre_call_hook_already_ran and re-run it a second time on # live kwargs. - if ( - getattr(callback, "scan_raw_request", False) - and not isinstance(result, BaseException) - and result is not None - ): + if callback.scan_raw_request and not isinstance(result, BaseException) and result is not None: callback.mark_pre_call_hook_ran(data) raised: Final = tuple(result for result in results if isinstance(result, BaseException)) blocking: Final = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None) @@ -6035,10 +6031,42 @@ def _should_use_smtp_ssl(smtp_port: int) -> bool: return os.getenv("SMTP_USE_SSL", "False") == "True" or smtp_port == 465 -def _create_smtp_connection(smtp_host: str, smtp_port: int) -> smtplib.SMTP: +def _create_smtp_connection(smtp_host: str, smtp_port: int, timeout: float) -> smtplib.SMTP: if _should_use_smtp_ssl(smtp_port=smtp_port): - return smtplib.SMTP_SSL(host=smtp_host, port=smtp_port, context=ssl.create_default_context()) - return smtplib.SMTP(host=smtp_host, port=smtp_port) + return smtplib.SMTP_SSL(host=smtp_host, port=smtp_port, context=ssl.create_default_context(), timeout=timeout) + return smtplib.SMTP(host=smtp_host, port=smtp_port, timeout=timeout) + + +def _send_smtp_message( + email_message: MIMEMultipart, + smtp_host: str, + smtp_port: int, + smtp_username: str | None, + smtp_password: str | None, + sender_email: str, + receiver_email: str, + timeout: float, +) -> None: + using_ssl: Final = _should_use_smtp_ssl(smtp_port=smtp_port) + with _create_smtp_connection( + smtp_host=smtp_host, + smtp_port=smtp_port, + timeout=timeout, + ) as server: + if not using_ssl and os.getenv("SMTP_TLS", "True") != "False": + server.starttls(context=ssl.create_default_context()) + + if smtp_username and smtp_password: + server.login( + user=smtp_username, + password=smtp_password, + ) + + server.send_message( + msg=email_message, + from_addr=sender_email, + to_addrs=receiver_email, + ) async def send_email( @@ -6084,27 +6112,18 @@ async def send_email( email_message.attach(MIMEText(html, "html")) try: - using_ssl: Final = _should_use_smtp_ssl(smtp_port=smtp_port) - with _create_smtp_connection( + smtp_timeout: Final = float(os.getenv("SMTP_TIMEOUT", "30")) + await asyncio.to_thread( + _send_smtp_message, + email_message=email_message, smtp_host=smtp_host, smtp_port=smtp_port, - ) as server: - if not using_ssl and os.getenv("SMTP_TLS", "True") != "False": - server.starttls(context=ssl.create_default_context()) - - # Login to your email account only if smtp_username and smtp_password are provided - if smtp_username and smtp_password: - server.login( - user=smtp_username, - password=smtp_password, - ) - - # Send the email - server.send_message( - msg=email_message, - from_addr=sender_email, - to_addrs=receiver_email, - ) + smtp_username=smtp_username, + smtp_password=smtp_password, + sender_email=sender_email, + receiver_email=receiver_email, + timeout=smtp_timeout, + ) except Exception as e: verbose_proxy_logger.exception("An error occurred while sending the email:" + str(e)) diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index 2aa1b8e0e3f..d962934dfb1 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -144,4 +144,7 @@ class PrismaBatch(Protocol): @property def litellm_endusertable(self) -> BatchTable: ... + @property + def litellm_modelaccessgroupbudgettable(self) -> BatchTable: ... + async def commit(self) -> None: ... diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index e02f652caf6..18cf884f267 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -68,6 +68,10 @@ class SpendLogsRepository(PrismaTableRepository["prisma_models.LiteLLM_SpendLogs table_name = "litellm_spendlogs" +class BudgetWindowSpendRepository(PrismaTableRepository["prisma_models.LiteLLM_BudgetWindowSpend"]): + table_name = "litellm_budgetwindowspend" + + class ClaudeCodePluginRepository(PrismaTableRepository["prisma_models.LiteLLM_ClaudeCodePluginTable"]): table_name = "litellm_claudecodeplugintable" @@ -100,6 +104,10 @@ class TagRepository(PrismaTableRepository["prisma_models.LiteLLM_TagTable"]): table_name = "litellm_tagtable" +class ModelAccessGroupBudgetRepository(PrismaTableRepository["prisma_models.LiteLLM_ModelAccessGroupBudgetTable"]): + table_name = "litellm_modelaccessgroupbudgettable" + + class InvitationLinkRepository(PrismaTableRepository["prisma_models.LiteLLM_InvitationLink"]): table_name = "litellm_invitationlink" diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index eb11ebe3b9c..a497d0580db 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -118,6 +118,7 @@ class BudgetCascadeUnitOfWork: keys: LinkedSpendResetWrites organizations: LinkedSpendResetWrites tags: LinkedSpendResetWrites + model_access_groups: LinkedSpendResetWrites endusers: LinkedSpendResetWrites budgets: BudgetWindowWrites @@ -143,6 +144,7 @@ async def budget_cascade_unit_of_work( keys=LinkedSpendResetWrites(table=batch.litellm_verificationtoken), organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable), tags=LinkedSpendResetWrites(table=batch.litellm_tagtable), + model_access_groups=LinkedSpendResetWrites(table=batch.litellm_modelaccessgroupbudgettable), endusers=LinkedSpendResetWrites(table=batch.litellm_endusertable), budgets=BudgetWindowWrites(table=batch.litellm_budgettable), ) 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/llms/vertex_ai_gemini_transcription.py b/litellm/types/llms/vertex_ai_gemini_transcription.py new file mode 100644 index 00000000000..e039bc8f2eb --- /dev/null +++ b/litellm/types/llms/vertex_ai_gemini_transcription.py @@ -0,0 +1,72 @@ +from typing import Literal + +from pydantic import BaseModel, ConfigDict +from typing_extensions import ReadOnly, TypedDict + + +class VertexGeminiTranscriptionInlineData(TypedDict): + mimeType: ReadOnly[str] + data: ReadOnly[str] + + +class VertexGeminiTranscriptionPart(TypedDict): + inlineData: ReadOnly[VertexGeminiTranscriptionInlineData] + + +class VertexGeminiTranscriptionContent(TypedDict): + role: ReadOnly[Literal["user"]] + parts: ReadOnly[tuple[VertexGeminiTranscriptionPart, ...]] + + +class VertexGeminiTranscriptionAudioConfig(TypedDict, total=False): + languageCodes: ReadOnly[tuple[str, ...]] + + +class VertexGeminiTranscriptionGenerationConfig(TypedDict): + audioTranscriptionConfig: ReadOnly[VertexGeminiTranscriptionAudioConfig] + + +class VertexGeminiTranscriptionRequest(TypedDict): + contents: ReadOnly[tuple[VertexGeminiTranscriptionContent, ...]] + generationConfig: ReadOnly[VertexGeminiTranscriptionGenerationConfig] + + +class VertexGeminiTranscriptionResponsePart(BaseModel): + model_config = ConfigDict(extra="ignore") + + text: str | None = None + + +class VertexGeminiTranscriptionResponseContent(BaseModel): + model_config = ConfigDict(extra="ignore") + + parts: tuple[VertexGeminiTranscriptionResponsePart, ...] = () + + +class VertexGeminiTranscriptionCandidate(BaseModel): + model_config = ConfigDict(extra="ignore") + + content: VertexGeminiTranscriptionResponseContent | None = None + + +class VertexGeminiTranscriptionModalityTokens(BaseModel): + model_config = ConfigDict(extra="ignore") + + modality: str | None = None + tokenCount: int = 0 + + +class VertexGeminiTranscriptionUsageMetadata(BaseModel): + model_config = ConfigDict(extra="ignore") + + promptTokenCount: int = 0 + candidatesTokenCount: int = 0 + totalTokenCount: int = 0 + promptTokensDetails: tuple[VertexGeminiTranscriptionModalityTokens, ...] = () + + +class VertexGeminiTranscriptionResponse(BaseModel): + model_config = ConfigDict(extra="ignore") + + candidates: tuple[VertexGeminiTranscriptionCandidate, ...] = () + usageMetadata: VertexGeminiTranscriptionUsageMetadata | None = None 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/proxy/management_endpoints/model_management_endpoints.py b/litellm/types/proxy/management_endpoints/model_management_endpoints.py index 6e18787a224..9d4663631fe 100644 --- a/litellm/types/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/model_management_endpoints.py @@ -1,6 +1,7 @@ +from datetime import datetime from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field from ...router import ModelGroupInfo @@ -53,10 +54,42 @@ class DeleteModelGroupResponse(BaseModel): message: str +class AccessGroupBudget(BaseModel): + budget_id: str + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + budget_reset_at: datetime | None = None + + +class AccessGroupBudgetRequest(BaseModel): + budget_id: str | None = None # Link an existing budget instead of creating one + max_budget: float | None = Field(default=None, ge=0) + soft_budget: float | None = Field(default=None, ge=0) + budget_duration: str | None = None + + # rejects tpm_limit/rpm_limit/max_parallel_requests: those are not enforced per access group + model_config = ConfigDict(extra="forbid") + + +class AccessGroupBudgetResponse(BaseModel): + access_group: str + spend: float # Shared spend accrued by every key that can reach this access group + budget: AccessGroupBudget | None = None + + +class DeleteAccessGroupBudgetResponse(BaseModel): + access_group: str + budget_deleted: bool # False when the access group had no budget to begin with + message: str + + class AccessGroupInfo(BaseModel): access_group: str model_names: list[str] # List of model names in this access group deployment_count: int # Total number of deployments with this access group + spend: float | None = None # Only populated by /access_group/{access_group}/info + budget: AccessGroupBudget | None = None class ListAccessGroupsResponse(BaseModel): diff --git a/litellm/types/proxy/model_access_group_budget.py b/litellm/types/proxy/model_access_group_budget.py new file mode 100644 index 00000000000..cccbe92b5d6 --- /dev/null +++ b/litellm/types/proxy/model_access_group_budget.py @@ -0,0 +1,19 @@ +"""The model access group budget state auth and the spend reservation path share.""" + +from __future__ import annotations + +from pydantic import BaseModel + + +class ModelAccessGroupBudget(BaseModel): + """One model access group's budget, flattened out of its joined ``LiteLLM_ModelAccessGroupBudgetTable`` row. + + Both readers want only the recorded spend and the ceiling, and this sits on the per-request hot + path behind a cache, so the linked budget row is collapsed to ``max_budget`` rather than cached + whole. ``spend`` is the DB-recorded value, which lags the live counter and is only ever a + fallback for it. + """ + + access_group_name: str + spend: float = 0.0 + max_budget: float | None = None diff --git a/litellm/types/services.py b/litellm/types/services.py index 74f908548d5..c558f6fb9d2 100644 --- a/litellm/types/services.py +++ b/litellm/types/services.py @@ -40,6 +40,8 @@ class ServiceTypes(str, enum.Enum): # spend update queue - current spend of key, user, team IN_MEMORY_SPEND_UPDATE_QUEUE = "in_memory_spend_update_queue" REDIS_SPEND_UPDATE_QUEUE = "redis_spend_update_queue" + # budget window spend queue - per-window spend of key, team + REDIS_WINDOW_SPEND_UPDATE_QUEUE = "redis_window_spend_update_queue" class ServiceConfig(TypedDict): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index f0319a7c664..4bf8289d725 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -40,7 +40,7 @@ from pydantic import ( field_serializer, field_validator, ) -from typing_extensions import ReadOnly, Required, TypedDict +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -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 @@ -3258,6 +3260,7 @@ class StandardLoggingPayload(TypedDict): cache_key: str | None saved_cache_cost: float request_tags: list + request_model_access_groups: NotRequired[ReadOnly[Sequence[str]]] end_user: str | None requester_ip_address: str | None user_agent: str | None @@ -3503,6 +3506,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..5e9e115ed54 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 @@ -8573,6 +8574,13 @@ class ProviderConfigManager: return SonioxAudioTranscriptionConfig() elif litellm.LlmProviders.VERTEX_AI == provider: + bare_vertex_model: Final = model.removeprefix("vertex_ai/") + if bare_vertex_model.startswith("gemini") and "transcribe" in bare_vertex_model: + from litellm.llms.vertex_ai.audio_transcription.gemini_transcribe_transformation import ( + VertexGeminiAudioTranscriptionConfig, + ) + + return VertexGeminiAudioTranscriptionConfig() from litellm.llms.vertex_ai.audio_transcription.transformation import ( VertexAIAudioTranscriptionConfig, ) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bebbcc32181..05c1cfd3179 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -12595,7 +12595,8 @@ "us": 1.1 }, "supports_output_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://docs.anthropic.com/en/docs/about-claude/models/overview" }, "claude-sonnet-4-6": { "deprecation_date": "2027-02-17", @@ -12997,7 +12998,8 @@ }, "supports_output_config": true, "prompt_cache_min_tokens": 512, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "source": "https://docs.anthropic.com/en/docs/about-claude/models/overview" }, "claude-opus-5": { "deprecation_date": "2027-07-24", @@ -13037,7 +13039,8 @@ }, "supports_output_config": true, "supports_speed": true, - "prompt_cache_min_tokens": 512 + "prompt_cache_min_tokens": 512, + "source": "https://docs.anthropic.com/en/docs/about-claude/models/overview" }, "claude-opus-4-8": { "deprecation_date": "2027-05-28", @@ -20562,7 +20565,8 @@ "output_cost_per_token": 2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_endpoints": [ - "/vertex_ai/live" + "/vertex_ai/live", + "/v1/realtime" ], "supported_modalities": [ "text", @@ -22015,6 +22019,49 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/nano-banana-pro-preview": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "rpm": 1000, + "tpm": 4000000, + "output_cost_per_token_batches": 6e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-3.1-flash-image": { "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, @@ -23047,7 +23094,8 @@ "supports_system_messages": true, "supports_video_input": true, "supports_vision": true, - "tpm": 800000 + "tpm": 800000, + "deprecation_date": "2026-09-30" }, "gemini/gemini-3.1-pro-preview": { "prompt_cache_min_tokens": 4096, @@ -23558,6 +23606,38 @@ "supports_tool_choice": true, "supports_vision": true }, + "gemini/gemma-4-26b-a4b-it": { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "gemini", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://ai.google.dev/gemini-api/docs/pricing" + }, + "gemini/gemma-4-31b-it": { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "gemini", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://ai.google.dev/gemini-api/docs/pricing" + }, "gemini/imagen-3.0-fast-generate-001": { "litellm_provider": "gemini", "mode": "image_generation", @@ -26097,7 +26177,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "low/1024-x-1536/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.013, @@ -26108,7 +26189,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "low/1536-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.013, @@ -26119,7 +26201,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "medium/1024-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.034, @@ -26130,7 +26213,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "medium/1024-x-1536/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.05, @@ -26141,7 +26225,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "medium/1536-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.05, @@ -26152,7 +26237,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "high/1024-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.133, @@ -26163,7 +26249,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "high/1024-x-1536/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.2, @@ -26174,7 +26261,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "high/1536-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.2, @@ -26185,7 +26273,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "standard/1024-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.009, @@ -26196,7 +26285,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "standard/1024-x-1536/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.013, @@ -26207,7 +26297,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "standard/1536-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.013, @@ -26218,7 +26309,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "1024-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.009, @@ -26229,7 +26321,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "1024-x-1536/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.013, @@ -26240,7 +26333,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "1536-x-1024/gpt-image-1.5-2025-12-16": { "input_cost_per_image": 0.013, @@ -26251,7 +26345,8 @@ "/v1/images/edits" ], "supports_vision": true, - "supports_pdf_input": true + "supports_pdf_input": true, + "deprecation_date": "2026-12-01" }, "gpt-5": { "cache_read_input_token_cost": 1.25e-07, @@ -38583,7 +38678,6 @@ "input_cost_per_token": 1.04e-06, "litellm_provider": "together_ai", "max_input_tokens": 131072, - "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.04e-06, @@ -38705,7 +38799,6 @@ "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 131072, - "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 6e-07, @@ -38752,7 +38845,6 @@ "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_input_tokens": 200000, - "max_output_tokens": 200000, "max_tokens": 200000, "metadata": { "successor": "together_ai/zai-org/GLM-5.2" @@ -38770,7 +38862,6 @@ "input_cost_per_token": 4.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 200000, - "max_output_tokens": 200000, "max_tokens": 200000, "metadata": { "successor": "together_ai/zai-org/GLM-5.2" @@ -38788,7 +38879,6 @@ "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", "max_input_tokens": 256000, - "max_output_tokens": 256000, "max_tokens": 256000, "metadata": { "successor": "together_ai/moonshotai/Kimi-K3" @@ -38868,7 +38958,6 @@ "input_cost_per_token": 3e-07, "litellm_provider": "together_ai", "max_input_tokens": 524288, - "max_output_tokens": 524288, "max_tokens": 524288, "mode": "chat", "output_cost_per_token": 1.2e-06, @@ -38885,7 +38974,6 @@ "input_cost_per_token": 0.0, "litellm_provider": "together_ai", "max_input_tokens": 262144, - "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 0.0, @@ -38895,7 +38983,6 @@ "input_cost_per_token": 1.7e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, - "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 2.5e-07, @@ -38911,7 +38998,6 @@ "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", "max_input_tokens": 1000000, - "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 3e-06, @@ -38923,7 +39009,6 @@ "input_cost_per_token": 2.5e-06, "litellm_provider": "together_ai", "max_input_tokens": 1000000, - "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 7.5e-06, @@ -38934,21 +39019,19 @@ "input_cost_per_token": 3.2e-07, "litellm_provider": "together_ai", "max_input_tokens": 1000000, - "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 1.28e-06, "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/Qwen/Qwen3.8-2.4T-A95B": { - "cache_read_input_token_cost": 2.5e-07, - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2.5e-06, "litellm_provider": "together_ai", "max_input_tokens": 1010000, - "max_output_tokens": 1010000, "max_tokens": 1010000, "mode": "chat", - "output_cost_per_token": 6e-06, + "output_cost_per_token": 6.25e-06, "source": "https://docs.together.ai/docs/serverless-models", "supports_prompt_caching": true }, @@ -38956,7 +39039,6 @@ "input_cost_per_token": 1e-07, "litellm_provider": "together_ai", "max_input_tokens": 32768, - "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 1e-07, @@ -38967,7 +39049,6 @@ "input_cost_per_token": 1.4e-07, "litellm_provider": "together_ai", "max_input_tokens": 1048576, - "max_output_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 2.8e-07, @@ -38984,7 +39065,6 @@ "input_cost_per_token": 1.74e-06, "litellm_provider": "together_ai", "max_input_tokens": 512000, - "max_output_tokens": 512000, "max_tokens": 512000, "mode": "chat", "output_cost_per_token": 3.48e-06, @@ -39001,7 +39081,6 @@ "input_cost_per_token": 1.32e-06, "litellm_provider": "together_ai", "max_input_tokens": 1048576, - "max_output_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 3.96e-06, @@ -39017,7 +39096,6 @@ "input_cost_per_token": 6e-08, "litellm_provider": "together_ai", "max_input_tokens": 32768, - "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 1.2e-07, @@ -39027,7 +39105,6 @@ "input_cost_per_token": 3.9e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, - "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 9.7e-07, @@ -39053,7 +39130,6 @@ "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "max_input_tokens": 1048576, - "max_output_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 2e-07, @@ -39064,7 +39140,6 @@ "input_cost_per_token": 3.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 131072, - "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.5e-06, @@ -39077,7 +39152,6 @@ "input_cost_per_token": 9.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, - "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, @@ -39094,7 +39168,6 @@ "input_cost_per_token": 3e-06, "litellm_provider": "together_ai", "max_input_tokens": 1048576, - "max_output_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 1.5e-05, @@ -39118,7 +39191,6 @@ "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_input_tokens": 512288, - "max_output_tokens": 512288, "max_tokens": 512288, "mode": "chat", "output_cost_per_token": 3.6e-06, @@ -39135,7 +39207,6 @@ "input_cost_per_token": 2.8e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, - "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 8.6e-07, @@ -39146,7 +39217,6 @@ "input_cost_per_token": 1e-06, "litellm_provider": "together_ai", "max_input_tokens": 524288, - "max_output_tokens": 524288, "max_tokens": 524288, "mode": "chat", "output_cost_per_token": 4.05e-06, @@ -39162,7 +39232,6 @@ "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", "max_input_tokens": 524288, - "max_output_tokens": 524288, "max_tokens": 524288, "mode": "chat", "output_cost_per_token": 1.2e-06, @@ -39174,8 +39243,25 @@ "input_cost_per_token": 1.4e-06, "litellm_provider": "together_ai", "max_input_tokens": 1048575, - "max_output_tokens": 1048575, - "max_tokens": 1048575, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/zai-org/GLM-5.3": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048575, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 4.4e-06, "source": "https://docs.together.ai/docs/serverless-models", @@ -39191,8 +39277,8 @@ "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 1048575, - "max_output_tokens": 1048575, - "max_tokens": 1048575, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 5e-07, "source": "https://docs.together.ai/docs/serverless-models", @@ -43318,6 +43404,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", @@ -44001,287 +44103,335 @@ ] }, "xai/grok-3": { - "cache_read_input_token_cost": 7.5e-07, - "input_cost_per_token": 3e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.5e-05, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-beta": { - "cache_read_input_token_cost": 7.5e-07, - "input_cost_per_token": 3e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.5e-05, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-fast-beta": { - "cache_read_input_token_cost": 1.25e-06, - "input_cost_per_token": 5e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.5e-05, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-fast-latest": { - "cache_read_input_token_cost": 1.25e-06, - "input_cost_per_token": 5e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.5e-05, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-latest": { - "cache_read_input_token_cost": 7.5e-07, - "input_cost_per_token": 3e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.5e-05, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-mini": { - "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost": 2e-07, "deprecation_date": "2026-02-28", - "input_cost_per_token": 3e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-mini-beta": { - "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost": 2e-07, "deprecation_date": "2026-02-28", - "input_cost_per_token": 3e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-mini-fast": { - "cache_read_input_token_cost": 1.5e-07, - "input_cost_per_token": 6e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 4e-06, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-02-28", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-mini-fast-beta": { - "cache_read_input_token_cost": 1.5e-07, - "input_cost_per_token": 6e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 4e-06, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-02-28", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-mini-fast-latest": { - "cache_read_input_token_cost": 1.5e-07, - "input_cost_per_token": 6e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 4e-06, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-02-28", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-3-mini-latest": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_token": 3e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-02-28", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4": { - "input_cost_per_token": 3e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-05, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-fast-reasoning": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 5e-07, - "output_cost_per_token_above_128k_tokens": 1e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-fast-non-reasoning": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 5e-07, - "output_cost_per_token_above_128k_tokens": 1e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-0709": { - "input_cost_per_token": 3e-06, - "input_cost_per_token_above_128k_tokens": 6e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-05, - "output_cost_per_token_above_128k_tokens": 3e-05, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-latest": { - "input_cost_per_token": 3e-06, - "input_cost_per_token_above_128k_tokens": 6e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-05, - "output_cost_per_token_above_128k_tokens": 3e-05, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-1-fast": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 5e-07, - "output_cost_per_token_above_128k_tokens": 1e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, "supports_function_calling": true, @@ -44290,19 +44440,21 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-1-fast-reasoning": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 5e-07, - "output_cost_per_token_above_128k_tokens": 1e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, "supports_function_calling": true, @@ -44312,19 +44464,20 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-1-fast-reasoning-latest": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 5e-07, - "output_cost_per_token_above_128k_tokens": 1e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, "supports_function_calling": true, @@ -44334,19 +44487,20 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-1-fast-non-reasoning": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 5e-07, - "output_cost_per_token_above_128k_tokens": 1e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-non-reasoning", "supports_audio_input": true, "supports_function_calling": true, @@ -44355,19 +44509,20 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4-1-fast-non-reasoning-latest": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 5e-07, - "output_cost_per_token_above_128k_tokens": 1e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-non-reasoning", "supports_audio_input": true, "supports_function_calling": true, @@ -44376,7 +44531,10 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07 }, "xai/grok-4.20-multi-agent-beta-0309": { "cache_read_input_token_cost": 2e-07, @@ -50390,7 +50548,7 @@ "max_tokens": 500000, "mode": "chat", "supports_function_calling": true, - "supports_prompt_caching": true, + "supports_prompt_caching": false, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true @@ -50405,7 +50563,7 @@ "max_tokens": 500000, "mode": "chat", "supports_function_calling": true, - "supports_prompt_caching": true, + "supports_prompt_caching": false, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true @@ -51828,6 +51986,43 @@ "tpm": 250000, "rpm": 10 }, + "vertex_ai/gemini-3.5-transcribe-preview": { + "input_cost_per_audio_token": 2.5e-06, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "vertex_ai", + "mode": "audio_transcription", + "output_cost_per_token": 1.2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "vertex_ai/gemini-3.5-transcribe-live-preview": { + "input_cost_per_audio_token": 3.5e-06, + "input_cost_per_token": 3.5e-06, + "litellm_provider": "vertex_ai", + "mode": "audio_transcription", + "output_cost_per_token": 2.1e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, "perplexity/pplx-embed-context-v1-0.6b": { "input_cost_per_token": 8e-09, "litellm_provider": "perplexity", @@ -54596,5 +54791,249 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true + }, + "groq/qwen/qwen3.8-27b": { + "input_cost_per_token": 8e-07, + "litellm_provider": "groq", + "max_input_tokens": 131042, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://console.groq.com/docs/model/qwen/qwen3.8-27b", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-medium-3.5": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-vibe-cli-latest": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-vibe-cli-with-tools": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-vibe-cli-fast": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-code-latest": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, + "litellm_provider": "mistral", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 9e-07, + "source": "https://docs.mistral.ai/models/model-cards/codestral-25-08", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/mistral-code-fim-latest": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, + "litellm_provider": "mistral", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 9e-07, + "source": "https://docs.mistral.ai/models/model-cards/codestral-25-08", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/mistral-code-agent-latest": { + "input_cost_per_token": 4e-07, + "litellm_provider": "mistral", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://mistral.ai/news/devstral-2-vibe-cli", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/mistral-ocr-3": { + "litellm_provider": "mistral", + "ocr_cost_per_page": 0.002, + "annotation_cost_per_page": 0.003, + "mode": "ocr", + "supported_endpoints": [ + "/v1/ocr" + ], + "source": "https://mistral.ai/pricing#api-pricing" + }, + "mistral/mistral-ocr-3-0": { + "litellm_provider": "mistral", + "ocr_cost_per_page": 0.002, + "annotation_cost_per_page": 0.003, + "mode": "ocr", + "supported_endpoints": [ + "/v1/ocr" + ], + "source": "https://mistral.ai/pricing#api-pricing" + }, + "mistral/mistral-ocr-4": { + "annotation_cost_per_page": 0.005, + "litellm_provider": "mistral", + "mode": "ocr", + "ocr_cost_per_page": 0.004, + "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", + "supported_endpoints": [ + "/v1/ocr" + ] + }, + "mistral/voxtral-mini-latest": { + "input_cost_per_second": 5e-05, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/model-cards/voxtral-mini-transcribe-26-02", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-realtime-2602": { + "input_cost_per_second": 0.0001, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/model-cards/voxtral-mini-transcribe-realtime-26-02", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-realtime-latest": { + "input_cost_per_second": 0.0001, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/model-cards/voxtral-mini-transcribe-realtime-26-02", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/labs-leanstral-1-5-1": { + "input_cost_per_token": 0.0, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://docs.mistral.ai/models/model-cards/leanstral-1-5", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "fireworks_ai/accounts/fireworks/models/glm-5p3": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/models/qwen3-embedding-8b": { + "input_cost_per_token": 1e-07, + "output_cost_per_token": 0.0, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 40960, + "max_output_tokens": 40960, + "max_tokens": 40960, + "mode": "embedding", + "source": "https://docs.fireworks.ai/serverless/pricing" } } diff --git a/pyproject.toml b/pyproject.toml index a0db4d49467..34c1fec1c11 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,8 +67,8 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.28.1,<2.0", - "litellm-proxy-extras==0.4.90", - "litellm-enterprise==0.1.61", + "litellm-proxy-extras==0.4.91", + "litellm-enterprise==0.1.62", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", 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/schema.prisma b/schema.prisma index 2bb850139a2..60223265211 100644 --- a/schema.prisma +++ b/schema.prisma @@ -29,6 +29,7 @@ model LiteLLM_BudgetTable { keys LiteLLM_VerificationToken[] // multiple keys can have the same budget end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget + model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization } @@ -585,6 +586,20 @@ model LiteLLM_EndUserTable { blocked Boolean @default(false) } +// Budget and shared spend for a model access group. The groups themselves are not rows anywhere: +// they are free-text strings in LiteLLM_ProxyModelTable.model_info.access_groups, so a row here +// exists only once someone gives that group a budget. +model LiteLLM_ModelAccessGroupBudgetTable { + access_group_name String @id + spend Float @default(0.0) + budget_id String? + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + created_at DateTime @default(now()) @map("created_at") + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? +} + // Track tags with budgets and spend model LiteLLM_TagTable { tag_name String @id @@ -649,6 +664,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/scripts/sync_together_ai_models.py b/scripts/sync_together_ai_models.py index 97b308fb660..12b128890f1 100644 --- a/scripts/sync_together_ai_models.py +++ b/scripts/sync_together_ai_models.py @@ -187,9 +187,22 @@ CAPABILITY_RULES: Final = ( _rule("thinkingmachines/Inkling-Small", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), _rule( "zai-org/GLM-5.2", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/glm-5-2", + "reviewed for the LIT-5968 backfill against https://www.together.ai/models/glm-5-2;" + " 128K output ceiling per https://docs.z.ai/guides/llm/glm-5.2", **_TOOLS, supports_reasoning=True, + max_output_tokens=128000, + max_tokens=128000, + ), + _rule( + "zai-org/GLM-5.3-Flash", + "reviewed for LIT-6489 against https://www.together.ai/models/glm-5-3-flash;" + " 128K output ceiling per https://docs.z.ai/guides/llm/glm-5.3", + **_TOOLS, + supports_reasoning=True, + supports_vision=True, + max_output_tokens=128000, + max_tokens=128000, ), ) @@ -288,15 +301,10 @@ def _api_fields(model: CatalogModel) -> RegistryEntry: def _new_entry(model: CatalogModel, mode: str) -> RegistryEntry: rule: Final = RULES_BY_ID.get(model.id) - length_fields: Final = ( - {} - if model.context_length is None - else {"max_input_tokens": model.context_length, "max_tokens": model.context_length} - | ({"max_output_tokens": model.context_length} if mode == "chat" else {}) - ) + legacy_ceiling: Final = {} if model.context_length is None else {"max_tokens": model.context_length} merged: Final = { **_api_fields(model), - **length_fields, + **legacy_ceiling, "litellm_provider": PROVIDER, "mode": mode, "source": SOURCE_URL, diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index d10be89b90c..052962e078e 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -126,3 +126,6 @@ POST /customer/delete # known gap: litellm_customer GET /team/{team_id}/callback # known gap: team callback resource POST /team/{team_id}/callback # known gap: team callback resource DELETE /team/{team_id}/callback/{callback_name} # known gap: team callback resource +GET /access_group/{access_group}/budget # known gap: budget attributes on litellm_access_group +PUT /access_group/{access_group}/budget # known gap: budget attributes on litellm_access_group +DELETE /access_group/{access_group}/budget # known gap: budget attributes on litellm_access_group 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/CLAUDE.md b/tests/e2e/CLAUDE.md index cf912ddaa25..14b2b4e3299 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -177,15 +177,15 @@ quota_management... behavior : ratelimit | budget | spend_tracking variant : rpm | tpm | priority_generous | priority_strict key | internal_user | end_user | organization | team | team_member | tag - | model_max | soft | key_multi_window | team_multi_window - | fallback | spend_counter + | model_access_group | model_max | soft | key_multi_window + | team_multi_window | fallback | spend_counter chat_completions | stream | messages_bridge | embeddings | cache_hit | key_rollup | concurrent_burst | tags | end_user | per_model | failure | spend_calculate | pagination assertion : blocks_over_limit | resets_after_window | headers_report_remaining | picks_under_tpm | blocks_then_resets | resets_windows_independently | alerts_without_blocking - | isolates_per_model | isolates_per_member | enforced_across_keys | routes_to_fallback - | reseed_matches_db | logs_cost | zero_cost + | isolates_per_model | isolates_per_member | isolates_per_group | enforced_across_keys + | routes_to_fallback | reseed_matches_db | reports_spend | logs_cost | zero_cost | matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows | writes_failure_row | returns_cost | keeps_total e.g. quota_management.ratelimit.rpm.blocks_over_limit exercised_on=[chat_completions, messages] diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 42a075681e0..d0afcaca848 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -20,6 +20,10 @@ - {id: quota_management.budget.organization.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: organization, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "An organization's max_budget blocks keys under its teams"} - {id: quota_management.budget.team_member.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A member's per-team budget blocks independently of the team budget"} - {id: quota_management.budget.team_member.isolates_per_member, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [isolates_per_member], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "One team member's exhausted per-team budget does not block a different member on the same team"} +- {id: quota_management.budget.model_access_group.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A model access group's shared max_budget blocks further calls to deployments in the group once the pool is spent"} +- {id: quota_management.budget.model_access_group.enforced_across_keys, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [enforced_across_keys], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "The pool is shared, so a key that spent nothing of its own is blocked once another key granted the same group drained it"} +- {id: quota_management.budget.model_access_group.isolates_per_group, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [isolates_per_group], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A request is charged only to the granted groups that serve the model it called, so an exhausted group never blocks a sibling group"} +- {id: quota_management.budget.model_access_group.reports_spend, module: quota_management, tier: P2, behavior: budget, variant: model_access_group, assertions: [reports_spend], exercised_on: [chat_completions], source: "proxy/management_endpoints/model_access_group_management_endpoints.py", rationale: "GET /access_group/{name}/budget reports the pool and the spend drawn against it, so an admin can see why calls are being refused"} - {id: quota_management.budget.tag.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: tag, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "router_strategy/budget_limiter.py", rationale: "Proxy-level tag budgets block tagged requests at the cap"} - {id: quota_management.budget.end_user_model_max.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: end_user_model_max, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "budget_management_endpoints.py", fail_before_fix: proven, rationale: "A per-model rpm_limit on an end-user budget is accepted and stored but never enforced; only key-attached budgets honour it"} - {id: quota_management.budget.model_max.isolates_per_model, module: quota_management, tier: P1, behavior: budget, variant: model_max, assertions: [isolates_per_model], exercised_on: [chat_completions], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "model_max_budget caps one model without touching a sibling's budget"} diff --git a/tests/e2e/quota_management/budgets/budget_client.py b/tests/e2e/quota_management/budgets/budget_client.py index 543d5f959e0..087dc8ca522 100644 --- a/tests/e2e/quota_management/budgets/budget_client.py +++ b/tests/e2e/quota_management/budgets/budget_client.py @@ -151,6 +151,28 @@ class TagDeleteBody(BaseModel): name: str +class AccessGroupBudgetBody(BaseModel): + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + + +class AccessGroupBudgetView(BaseModel): + budget_id: str + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + + +class AccessGroupBudgetResponse(BaseModel): + """GET/PUT /access_group/{name}/budget: the group's shared pool and the spend + every key that can reach the group has drawn against it.""" + + access_group: str + spend: float + budget: AccessGroupBudgetView | None = None + + class BudgetNewBody(BaseModel): max_budget: float | None = None soft_budget: float | None = None @@ -514,6 +536,49 @@ class BudgetClient: response_type=NoBody, ) + # ---- model access group --------------------------------------------- + + def set_access_group_budget( + self, + access_group: str, + *, + max_budget: float | None = None, + soft_budget: float | None = None, + budget_duration: str | None = None, + ) -> AccessGroupBudgetResponse: + """Give a model access group one shared budget. Every key that can reach a + deployment in the group draws from it.""" + return unwrap( + self.proxy.transport.put( + f"/access_group/{access_group}/budget", + headers=self.proxy.transport.master, + json=AccessGroupBudgetBody( + max_budget=max_budget, + soft_budget=soft_budget, + budget_duration=budget_duration, + ), + response_type=AccessGroupBudgetResponse, + ) + ) + + def access_group_budget(self, access_group: str) -> AccessGroupBudgetResponse: + return unwrap( + self.proxy.transport.get( + f"/access_group/{access_group}/budget", + headers=self.proxy.transport.master, + params=NoBody(), + response_type=AccessGroupBudgetResponse, + ) + ) + + def delete_access_group_budget(self, access_group: str) -> None: + _ = self.proxy.transport.delete( + f"/access_group/{access_group}/budget", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + # ---- budget table --------------------------------------------------- def create_budget( diff --git a/tests/e2e/quota_management/budgets/test_model_access_group_budget_e2e.py b/tests/e2e/quota_management/budgets/test_model_access_group_budget_e2e.py new file mode 100644 index 00000000000..9c927a31216 --- /dev/null +++ b/tests/e2e/quota_management/budgets/test_model_access_group_budget_e2e.py @@ -0,0 +1,161 @@ +"""Live e2e: one shared budget across every key that can reach a model access group. + +A model access group is a free-text label on a deployment (`model_info.access_groups`), +and a key is granted the group by name. The budget hangs off the group, not the key, so +the interesting behaviors are the ones a per-key budget cannot produce: a key that has +spent nothing of its own is refused once somebody else drained the pool, and draining one +group leaves a second group untouched, because a request is only charged to the groups +the caller was granted that also serve the model being called. +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import StreamingResponse, require_successful_call +from lifecycle import ResourceManager +from models import KeyGenerateBody, LiteLLMParamsBody, ModelInfoBody, ModelNewBody + +pytestmark = pytest.mark.e2e + +BACKEND: Final = "openai/gpt-5.4-nano" +TINY_BUDGET: Final = 5e-6 +MAX_TOKENS: Final = 16 +DRAIN_TIMEOUT_SECONDS: Final = 180 + + +@dataclass(frozen=True, slots=True) +class DrainedPool: + """A model access group whose shared budget has been spent to exhaustion, the + deployment inside it, the key that did the spending, and a second group holding + its own deployment that was never given a budget at all.""" + + access_group: str + model: str + spender_key: str + free_access_group: str + free_model: str + + +def _provider_key(env_var: str) -> str: + return os.environ.get(env_var) or f"os.environ/{env_var}" + + +def _grouped_model(model_name: str, access_group: str) -> ModelNewBody: + return ModelNewBody( + model_name=model_name, + litellm_params=LiteLLMParamsBody(model=BACKEND, api_key=_provider_key("OPENAI_API_KEY")), + model_info=ModelInfoBody(access_groups=[access_group]), + ) + + +def _call(client: BudgetClient, key: str, model: str) -> StreamingResponse: + return client.chat(key, model, f"hi {unique_marker()}", max_tokens=MAX_TOKENS) + + +def _drain(client: BudgetClient, key: str, model: str, access_group: str) -> None: + """Spend the group's pool until the proxy refuses the next request. The first call + lands under the cap and the block comes from the spend it recorded, so this needs at + least one round trip through the spend writer, not just one request.""" + deadline: Final = time.monotonic() + DRAIN_TIMEOUT_SECONDS + while time.monotonic() < deadline: + result = _call(client, key, model) + if is_budget_block(result): + return + require_successful_call(result) + time.sleep(1) + pytest.fail(f"budget on model access group {access_group!r} never blocked a request") + + +@pytest.fixture(scope="module") +def drained(client: BudgetClient) -> Iterator[DrainedPool]: + marker: Final = unique_marker() + pool: Final = DrainedPool( + access_group=f"e2e-mag-budget-{marker}", + model=f"e2e-mag-budgeted-{marker}", + spender_key=client.proxy.generate_key(KeyGenerateBody(models=[f"e2e-mag-budget-{marker}"])), + free_access_group=f"e2e-mag-free-{marker}", + free_model=f"e2e-mag-unbudgeted-{marker}", + ) + created: Final = ( + client.proxy.register_model(_grouped_model(pool.model, pool.access_group)), + client.proxy.register_model(_grouped_model(pool.free_model, pool.free_access_group)), + ) + try: + client.set_access_group_budget(pool.access_group, max_budget=TINY_BUDGET) + _drain(client, pool.spender_key, pool.model, pool.access_group) + yield pool + finally: + client.delete_access_group_budget(pool.access_group) + client.proxy.delete_key(pool.spender_key) + for model_id in created: + client.proxy.delete_model(model_id) + + +class TestModelAccessGroupBudget: + @pytest.mark.covers("quota_management.budget.model_access_group.blocks_over_limit") + def test_the_key_that_drained_the_pool_stays_blocked( + self, client: BudgetClient, drained: DrainedPool + ) -> None: + result = _call(client, drained.spender_key, drained.model) + assert is_budget_block(result), ( + f"an exhausted pool served {drained.model!r} again: {result.status_code} {result.body[:300]}" + ) + assert drained.access_group in result.body, ( + f"the block did not name the group that caused it: {result.body[:300]}" + ) + + @pytest.mark.covers("quota_management.budget.model_access_group.enforced_across_keys") + def test_a_key_that_spent_nothing_is_blocked_by_the_shared_pool( + self, client: BudgetClient, resources: ResourceManager, drained: DrainedPool + ) -> None: + newcomer = resources.key(models=[drained.access_group]) + + result = _call(client, newcomer, drained.model) + + assert is_budget_block(result), ( + "a freshly minted key with no spend of its own was served by an exhausted " + f"shared pool: {result.status_code} {result.body[:300]}" + ) + + @pytest.mark.covers("quota_management.budget.model_access_group.isolates_per_group") + def test_a_drained_group_does_not_block_a_different_group( + self, client: BudgetClient, resources: ResourceManager, drained: DrainedPool + ) -> None: + other = resources.key(models=[drained.free_access_group]) + + result = _call(client, other, drained.free_model) + + assert not is_budget_block(result), ( + f"{drained.free_access_group!r} has no budget of its own but was blocked by " + f"{drained.access_group!r}'s exhausted pool: {result.body[:300]}" + ) + require_successful_call(result) + + @pytest.mark.covers("quota_management.budget.model_access_group.reports_spend") + def test_the_budget_read_reports_the_spend_drawn_against_the_pool( + self, client: BudgetClient, drained: DrainedPool + ) -> None: + """Enforcement runs off a live counter while the group's row is written by the + batched spend writer, so the recorded spend an admin reads lands a beat after the + block. Poll for it: what matters is that it arrives and matches the pool.""" + deadline = time.monotonic() + client.proxy.poll_timeout + reported = client.access_group_budget(drained.access_group) + while reported.spend < TINY_BUDGET and time.monotonic() < deadline: + time.sleep(client.proxy.poll_interval) + reported = client.access_group_budget(drained.access_group) + + assert reported.budget is not None, "the group lost the budget that just blocked it" + assert reported.budget.max_budget == TINY_BUDGET + assert reported.spend >= TINY_BUDGET, ( + f"the pool blocked at {TINY_BUDGET} but only {reported.spend} was ever recorded " + f"against the group within {client.proxy.poll_timeout}s" + ) 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/llm_translation/reasoning_effort_grid/grid_spec.py b/tests/llm_translation/reasoning_effort_grid/grid_spec.py index 4fa77f38940..485f78ed8c7 100644 --- a/tests/llm_translation/reasoning_effort_grid/grid_spec.py +++ b/tests/llm_translation/reasoning_effort_grid/grid_spec.py @@ -103,7 +103,7 @@ def _bedrock_clamps_effort(model: "ModelEntry", effort: str) -> bool: return _EFFORT_RANK[effort] > _EFFORT_RANK[model.bedrock_effort_ceiling] -def expected(model: ModelEntry, effort: str) -> CellExpectation: +def expected(route_name: str, model: ModelEntry, effort: str) -> CellExpectation: if effort in ("__omit__", "none"): if model.mode == "budget": return CellExpectation( @@ -117,6 +117,15 @@ def expected(model: ModelEntry, effort: str) -> CellExpectation: if effort in ("xhigh", "max"): cap = f"supports_{effort}_reasoning_effort" if cap not in model.caps and not _bedrock_clamps_effort(model, effort): + if model.mode == "budget" and route_name == "bedrock_invoke_messages": + # the /v1/messages path caps the mapped budget below max_tokens + # (LIT-6498), so oversized tiers succeed there instead of 400ing + return CellExpectation( + status=200, + thinking_type="enabled", + thinking_budget_tokens=BUDGET_MODE_MAX_TOKENS - 1, + max_tokens=BUDGET_MODE_MAX_TOKENS, + ) return CellExpectation(status=400, thinking_type=OMIT) if model.mode == "adaptive": @@ -441,5 +450,7 @@ def all_cells() -> List[Tuple[str, ModelEntry, str, CellExpectation]]: for route in ROUTES: for model in route.models: for effort in EFFORTS: - cells.append((route.name, model, effort, expected(model, effort))) + cells.append( + (route.name, model, effort, expected(route.name, model, effort)) + ) return cells 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/enterprise/proxy/test_managed_files_access_check.py b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py index c75c8099ea1..ad46798b788 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py @@ -10,11 +10,14 @@ with deployment credentials, bypassing the managed files access-control hooks. import base64 import pytest +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch from fastapi import HTTPException -from litellm.proxy._types import UserAPIKeyAuth +from litellm.caching.dual_cache import DualCache +from litellm.proxy._types import CallTypes, UserAPIKeyAuth +from litellm.types.utils import LiteLLMBatch def _make_user_api_key_dict(user_id: str) -> UserAPIKeyAuth: @@ -161,6 +164,108 @@ async def test_service_account_blocked_from_other_team_file(): assert exc_info.value.status_code == 403 +# --- Keyless key must not be locked out of the batch it created --- + + +def _make_unified_batch_id() -> str: + raw = "litellm_proxy;model_id:my-model-id;llm_batch_id:batch_raw_123" + return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=") + + +def _make_managed_files_instance_with_object_store(): + """Managed-files hook backed by an in-memory stand-in for the managed + object table, so create and retrieve exercise the same stored row.""" + from litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) + + store = {} + + async def upsert(where, data): + store[where["unified_object_id"]] = SimpleNamespace(**data["create"]) + + async def find_first(where): + return store.get(where["unified_object_id"]) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock(side_effect=upsert) + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock( + side_effect=find_first + ) + + return ( + _PROXY_LiteLLMManagedFiles( + internal_usage_cache=DualCache(), + prisma_client=mock_prisma, + ), + store, + ) + + +async def _store_batch(managed_files, unified_batch_id: str, creator: UserAPIKeyAuth): + await managed_files.store_unified_object_id( + unified_object_id=unified_batch_id, + file_object=LiteLLMBatch( + id="batch_raw_123", + completion_window="24h", + created_at=0, + endpoint="/v1/chat/completions", + input_file_id="file-1", + object="batch", + status="validating", + ), + litellm_parent_otel_span=None, + model_object_id="batch_raw_123", + file_purpose="batch", + user_api_key_dict=creator, + ) + + +@pytest.mark.asyncio +async def test_keyless_key_can_retrieve_the_batch_it_created(): + """Regression: a key with no user_id and no team_id (what `/key/generate` + by a proxy admin and service-account keys produce) stamped + `created_by=None` and was then denied its own managed batch with + "User None does not have access".""" + unified_batch_id = _make_unified_batch_id() + managed_files, store = _make_managed_files_instance_with_object_store() + keyless = UserAPIKeyAuth(api_key="sk-keyless", parent_otel_span=None) + + await _store_batch(managed_files, unified_batch_id, keyless) + assert store[unified_batch_id].created_by == f"key:{keyless.token}" + + data = {"batch_id": unified_batch_id} + await managed_files.async_pre_call_hook( + user_api_key_dict=keyless, + cache=DualCache(), + data=data, + call_type=CallTypes.aretrieve_batch.value, + ) + assert data["batch_id"] == "batch_raw_123" + assert data["model"] == "my-model-id" + + +@pytest.mark.asyncio +async def test_other_keyless_key_still_denied_the_batch(): + unified_batch_id = _make_unified_batch_id() + managed_files, _ = _make_managed_files_instance_with_object_store() + + await _store_batch( + managed_files, + unified_batch_id, + UserAPIKeyAuth(api_key="sk-creator", parent_otel_span=None), + ) + + with pytest.raises(HTTPException) as exc_info: + await managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-other", parent_otel_span=None), + cache=DualCache(), + data={"batch_id": unified_batch_id}, + call_type=CallTypes.aretrieve_batch.value, + ) + assert exc_info.value.status_code == 403 + + # --- Option C fix test: check_batch_cost bypasses managed files hook --- diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index eddfc4fbd34..f3ad8a8592e 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -527,13 +527,33 @@ async def test_afile_list_orders_newest_first_and_breaks_ties_on_the_cursor_colu @pytest.mark.asyncio -async def test_afile_list_denies_a_caller_without_a_user_or_team(): +async def test_afile_list_scopes_a_keyless_key_to_its_own_hashed_token(): + caller = UserAPIKeyAuth(api_key="sk-test", parent_otel_span=None) + managed_files, table = _make_managed_files_over_rows( + [ + _make_managed_file_row("unified-mine", created_by=f"key:{caller.token}"), + _make_managed_file_row("unified-theirs", created_by="other-user"), + ] + ) + + response = await managed_files.afile_list( + purpose=None, + litellm_parent_otel_span=None, + user_api_key_dict=caller, + ) + + assert [file.id for file in response.data] == ["unified-mine"] + assert table.find_many_calls[0]["where"] == {"created_by": f"key:{caller.token}"} + + +@pytest.mark.asyncio +async def test_afile_list_denies_a_caller_with_no_identity_at_all(): managed_files, table = _make_managed_files_over_rows([_make_managed_file_row("unified-mine")]) response = await managed_files.afile_list( purpose=None, litellm_parent_otel_span=None, - user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", parent_otel_span=None), + user_api_key_dict=UserAPIKeyAuth(parent_otel_span=None), ) assert response.data == [] diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index 3a736e2a889..e995cbae782 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1611,6 +1611,20 @@ class TestEnableAnthropicPromptCaching: monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) assert self._points(model="anthropic.claude-3-5-sonnet-20240620-v1:0", provider="bedrock") == [] + @pytest.mark.parametrize("model", ["us.xai.grok-4.6", "global.xai.grok-4.6"]) + def test_bedrock_grok_not_injected(self, monkeypatch, local_model_cost_map, model): + """Bedrock supports only implicit prompt caching for Grok: explicit cachePoint + breakpoints make it reject the whole request ("You invoked an unsupported model + or your request did not allow prompt caching"), so supports_prompt_caching stays + false, while implicit cache hits still bill at the cache-read rate.""" + from litellm.utils import supports_prompt_caching + + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + assert supports_prompt_caching(model=model, custom_llm_provider="bedrock") is False + assert self._points(model=model, provider="bedrock") == [] + entry = litellm.model_cost[model] + assert 0 < entry["cache_read_input_token_cost"] < entry["input_cost_per_token"] + def test_stands_down_when_client_sent_cache_control(self, monkeypatch): monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) messages = [ 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..a45a6aaa2a3 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) @@ -3727,6 +3735,60 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob assert payload["litellm_call_id"] == call_id +def test_get_standard_logging_object_payload_carries_matched_access_groups(logging_obj): + """Access groups stamped at auth time reach the logging payload, so integrations see what a request billed.""" + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.now() + payload = get_standard_logging_object_payload( + kwargs={ + "model": "gpt-4o", + "messages": [], + "litellm_params": { + "metadata": { + "user_api_key_matched_model_access_groups": ["premium-pool", "shared-pool"] + }, + "proxy_server_request": {"body": {}}, + }, + }, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["request_model_access_groups"] == ("premium-pool", "shared-pool") + + +def test_get_standard_logging_object_payload_has_no_access_groups_when_unstamped( + logging_obj, +): + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.now() + payload = get_standard_logging_object_payload( + kwargs={"model": "gpt-4o", "messages": []}, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["request_model_access_groups"] == () + + def test_get_standard_logging_object_payload_preserves_absent_end_user_as_none(logging_obj): from datetime import datetime from typing import Final 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/experimental_pass_through/messages/test_anthropic_messages_effort.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py index e9d4d625421..daaa110e7b9 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py @@ -10,6 +10,10 @@ from litellm.llms.anthropic.common_utils import AnthropicError from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, ) +from litellm.llms.openai_like.json_loader import SimpleProviderConfig +from litellm.llms.openai_like.messages.transformation import ( + JSONProviderAnthropicMessagesConfig, +) def _claude_code_payload(effort="medium", max_tokens=8192, **output_config_extra): @@ -294,3 +298,39 @@ def test_non_adaptive_request_without_effort_is_untouched(): assert "thinking" not in result assert "output_config" not in result + + +def test_reasoning_effort_budget_capped_below_max_tokens(): + result = _transform("claude-haiku-4-5", {"max_tokens": 4000, "reasoning_effort": "xhigh"}) + + assert result["thinking"] == {"type": "enabled", "budget_tokens": 3999} + assert result["max_tokens"] == 4000 + + +def test_reasoning_effort_thinking_dropped_when_min_budget_cannot_fit(): + result = _transform("claude-haiku-4-5", {"max_tokens": 1024, "reasoning_effort": "xhigh"}) + + assert "thinking" not in result + assert result["max_tokens"] == 1024 + + +def test_reasoning_effort_budget_capped_for_openai_like_messages_upstream(): + provider = SimpleProviderConfig( + "meta", + { + "base_url": "https://api.meta.ai/v1", + "api_key_env": "META_API_KEY", + "supported_endpoints": ["/v1/messages"], + }, + ) + + result = JSONProviderAnthropicMessagesConfig(provider).transform_anthropic_messages_request( + model="muse-spark-1.2", + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params={"max_tokens": 4000, "reasoning_effort": "xhigh"}, + litellm_params={}, + headers={}, + ) + + assert result["thinking"] == {"type": "enabled", "budget_tokens": 3999} + assert result["max_tokens"] == 4000 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py index c9170efd18a..7e2fa356685 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py @@ -70,7 +70,7 @@ def test_reasoning_effort_none_clears_thinking_and_output_config(): def test_reasoning_effort_on_non_adaptive_model_uses_thinking_budget(): config = AnthropicMessagesConfig() - optional_params = {"max_tokens": 1024, "reasoning_effort": "high"} + optional_params = {"max_tokens": 8192, "reasoning_effort": "high"} result = config.transform_anthropic_messages_request( model="claude-opus-4-5", @@ -86,7 +86,7 @@ def test_reasoning_effort_on_non_adaptive_model_uses_thinking_budget(): assert isinstance(thinking, dict) assert thinking.get("type") == "enabled" assert isinstance(thinking.get("budget_tokens"), int) - assert thinking["budget_tokens"] >= 1024 + assert 1024 <= thinking["budget_tokens"] < result["max_tokens"] @pytest.mark.parametrize("bad_effort", ["invalid", "disabled", ""]) 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/azure/response/test_azure_transformation.py b/tests/test_litellm/llms/azure/response/test_azure_transformation.py index da44394d11d..f6bbf685f26 100644 --- a/tests/test_litellm/llms/azure/response/test_azure_transformation.py +++ b/tests/test_litellm/llms/azure/response/test_azure_transformation.py @@ -537,3 +537,79 @@ class TestAzureResponsesAPIConfig: """ supported = self.config.get_supported_openai_params(self.model) assert "context_management" not in supported + + def _anyof_tool(self): + return { + "type": "function", + "name": "automation_update", + "description": "Update an automation", + "parameters": { + "type": "object", + "anyOf": [ + { + "properties": {"id": {"type": "string"}, "enabled": {"type": "boolean"}}, + "required": ["id", "enabled"], + }, + { + "properties": {"id": {"type": "string"}, "schedule": {"type": "string"}}, + "required": ["id", "schedule"], + }, + ], + "properties": {"id": {"type": "string"}}, + "required": ["id"], + }, + } + + def test_azure_flattens_top_level_anyof_for_gpt4_family_deployment_name(self): + result = self.config.transform_responses_api_request( + model="gpt-4o", + input="hi", + response_api_optional_request_params={"tools": [self._anyof_tool()]}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + parameters = result["tools"][0]["parameters"] + assert "anyOf" not in parameters + assert parameters["type"] == "object" + assert set(parameters["properties"]) == {"id", "enabled", "schedule"} + assert parameters["required"] == ["id"] + + def test_azure_flattens_via_base_model_for_arbitrary_deployment_name(self): + result = self.config.transform_responses_api_request( + model="my-eastus-deployment", + input="hi", + response_api_optional_request_params={"tools": [self._anyof_tool()]}, + litellm_params=GenericLiteLLMParams(model_info={"base_model": "azure/gpt-4o"}), + headers={}, + ) + + assert "anyOf" not in result["tools"][0]["parameters"] + + def test_azure_keeps_combinators_for_gpt5_base_model(self): + tool = self._anyof_tool() + + result = self.config.transform_responses_api_request( + model="my-eastus-deployment", + input="hi", + response_api_optional_request_params={"tools": [tool]}, + litellm_params=GenericLiteLLMParams(model_info={"base_model": "azure/gpt-5.4-mini"}), + headers={}, + ) + + assert result["tools"][0] is tool + assert "anyOf" in result["tools"][0]["parameters"] + + def test_azure_keeps_combinators_for_unrecognized_deployment_without_base_model(self): + tool = self._anyof_tool() + + result = self.config.transform_responses_api_request( + model="my-eastus-deployment", + input="hi", + response_api_optional_request_params={"tools": [tool]}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["tools"][0] is tool + assert "anyOf" in result["tools"][0]["parameters"] diff --git a/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py b/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py index b5fcd9d8219..1746926c689 100644 --- a/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py +++ b/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py @@ -7,6 +7,7 @@ import pytest from litellm.llms.base_llm.managed_resources.isolation import ( build_owner_filter, can_access_resource, + resolve_resource_owner_id, ) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth @@ -154,3 +155,46 @@ def test_access_identity_less_caller_always_denied(created_by, resource_team_id) ) is False ) + + +# --------------------------------------------------------------------------- +# keyless keys (no user_id, no team_id) own their resources by hashed token +# --------------------------------------------------------------------------- + + +def test_owner_id_prefers_user_id_then_falls_back_to_token(): + assert resolve_resource_owner_id(UserAPIKeyAuth(user_id="alice")) == "alice" + assert resolve_resource_owner_id(UserAPIKeyAuth(team_id="team-eng")) is None + assert resolve_resource_owner_id(UserAPIKeyAuth()) is None + + keyless = UserAPIKeyAuth(api_key="sk-keyless") + assert resolve_resource_owner_id(keyless) == f"key:{keyless.token}" + + +def test_keyless_key_can_access_its_own_resource(): + """Regression for the self-lockout: a key generated by a proxy admin (or a + service-account key) has no user_id and no team_id, so it used to stamp + `created_by=None` and then be denied its own batches and files.""" + keyless = UserAPIKeyAuth(api_key="sk-keyless") + owner_id = resolve_resource_owner_id(keyless) + + assert build_owner_filter(keyless) == {"created_by": owner_id} + assert ( + can_access_resource(keyless, created_by=owner_id, resource_team_id=None) is True + ) + + +def test_keyless_key_denied_another_keyless_keys_resource(): + """The #27004 isolation invariant: two distinct keyless keys must not see + each other's resources.""" + creator = UserAPIKeyAuth(api_key="sk-creator") + other = UserAPIKeyAuth(api_key="sk-other") + + assert ( + can_access_resource( + other, + created_by=resolve_resource_owner_id(creator), + resource_team_id=None, + ) + is False + ) 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/gemini/realtime/test_gemini_realtime_transformation.py b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py index d0613403e67..3d8200bc474 100644 --- a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -1812,6 +1812,7 @@ def patch_gemini_audio_cost_map_entries(monkeypatch): "gemini-2.5-flash-native-audio", "gemini-2.5-flash-native-audio-latest", "gemini/gemini-2.5-flash-native-audio-latest", + "gemini-live-2.5-flash-native-audio", ] flash_live_models = [ "gemini-3.1-flash-live-preview", @@ -1834,6 +1835,8 @@ def patch_gemini_audio_cost_map_entries(monkeypatch): ("gemini/gemini-3.1-flash-live-preview", True), ("gemini-2.5-flash-native-audio-latest", True), ("gemini/gemini-2.5-flash-native-audio-latest", True), + ("gemini-live-2.5-flash-native-audio", True), + ("vertex_ai/gemini-live-2.5-flash-native-audio", True), ("gemini-2.0-flash", False), ("gemini-2.5-flash", False), ], @@ -1842,6 +1845,19 @@ def test_is_audio_only_live_model_uses_cost_map(model, expected, patch_gemini_au assert GeminiRealtimeConfig._is_audio_only_live_model(model) == expected +def test_gemini_live_native_audio_entry_is_vertex_only(): + import json + from pathlib import Path + from typing import Final + + catalog_path: Final = Path(__file__).parents[5] / "model_prices_and_context_window.json" + catalog: Final = json.loads(catalog_path.read_text()) + vertex_key: Final = "gemini-live-2.5-flash-native-audio" + assert catalog[vertex_key]["litellm_provider"] == "vertex_ai-language-models" + assert catalog[vertex_key].get("gemini_native_audio") is True + assert "gemini/gemini-live-2.5-flash-native-audio" not in catalog, "the Gemini API does not serve this model" + + def test_is_setup_message_and_is_content_message(): config = GeminiRealtimeConfig() assert config.is_setup_message({"setup": {}}) is True 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/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index 33e677b000e..9a6a039a470 100644 --- a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -254,7 +254,7 @@ def test_request_maps_reasoning_effort_to_thinking(config): model="claude-sonnet-4-20250514", messages=[{"role": "user", "content": "hi"}], anthropic_messages_optional_request_params={ - "max_tokens": 1024, + "max_tokens": 8192, "reasoning_effort": "medium", }, litellm_params=GenericLiteLLMParams(), @@ -264,6 +264,7 @@ def test_request_maps_reasoning_effort_to_thinking(config): assert "reasoning_effort" not in payload assert isinstance(payload.get("thinking"), dict) assert payload["thinking"].get("type") == "enabled" + assert payload["thinking"]["budget_tokens"] < payload["max_tokens"] def test_passthrough_disables_anthropic_beta_filtering(config): 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/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py new file mode 100644 index 00000000000..eadd87d9c92 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py @@ -0,0 +1,345 @@ +import base64 +import json +import os + +import httpx +import pytest + +import litellm +from litellm.llms.vertex_ai.audio_transcription.gemini_transcribe_transformation import ( + VertexGeminiAudioTranscriptionConfig, +) +from litellm.llms.vertex_ai.audio_transcription.transformation import ( + VertexAIAudioTranscriptionConfig, +) +from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.types.utils import LlmProviders, TranscriptionUsageTokensObject +from litellm.utils import ProviderConfigManager, get_optional_params_transcription + +AUDIO_BYTES = b"fake-audio-bytes" +TRANSCRIPT_TEXT = ( + "Four score and seven years ago our fathers brought forth on this continent, a new nation, " + "conceived in Liberty, and dedicated to the proposition that all men are created equal. " + "Now we are engaged in a great civil war, testing whether that nation, or any nation so " + "conceived and so dedicated, can long endure." +) +GENERATE_CONTENT_RESPONSE = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "text": TRANSCRIPT_TEXT, + "audioTranscription": {"text": TRANSCRIPT_TEXT}, + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 440, + "candidatesTokenCount": 62, + "totalTokenCount": 502, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "AUDIO", "tokenCount": 440}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 62}], + }, + "modelVersion": "gemini-3.5-transcribe-preview", + "createTime": "2026-08-29T07:25:27.591648Z", + "responseId": "Z4mSaqCOJL-O4_UP0aSh4Aw", +} + + +@pytest.fixture +def config(): + return VertexGeminiAudioTranscriptionConfig() + + +class TestProviderRouting: + @pytest.mark.parametrize( + "model", + [ + "gemini-3.5-transcribe-preview", + "gemini-3.5-transcribe-live-preview", + "vertex_ai/gemini-3.5-transcribe-preview", + ], + ) + def test_gemini_transcribe_models_use_generate_content_config(self, model): + provider_config = ProviderConfigManager.get_provider_audio_transcription_config( + model=model, + provider=LlmProviders.VERTEX_AI, + ) + assert isinstance(provider_config, VertexGeminiAudioTranscriptionConfig) + + @pytest.mark.parametrize("model", ["chirp_2", "chirp_3", "long-form", "gemini-2.5-flash"]) + def test_other_vertex_models_keep_speech_to_text_config(self, model): + provider_config = ProviderConfigManager.get_provider_audio_transcription_config( + model=model, + provider=LlmProviders.VERTEX_AI, + ) + assert isinstance(provider_config, VertexAIAudioTranscriptionConfig) + assert not isinstance(provider_config, VertexGeminiAudioTranscriptionConfig) + + +class TestGetCompleteUrl: + @pytest.fixture(autouse=True) + def _clear_ambient_vertex_location(self, monkeypatch): + monkeypatch.delenv("VERTEXAI_LOCATION", raising=False) + monkeypatch.delenv("VERTEX_LOCATION", raising=False) + + def test_defaults_to_global_location(self, config): + url = config.get_complete_url( + api_base=None, + api_key=None, + model="gemini-3.5-transcribe-preview", + optional_params={}, + litellm_params={"vertex_project": "test-project"}, + ) + assert url == ( + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global" + "/publishers/google/models/gemini-3.5-transcribe-preview:generateContent" + ) + + def test_explicit_location_is_honored(self, config): + url = config.get_complete_url( + api_base=None, + api_key=None, + model="gemini-3.5-transcribe-preview", + optional_params={}, + litellm_params={"vertex_project": "test-project", "vertex_location": "us-central1"}, + ) + assert url == ( + "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1" + "/publishers/google/models/gemini-3.5-transcribe-preview:generateContent" + ) + + def test_model_prefix_is_stripped(self, config): + url = config.get_complete_url( + api_base=None, + api_key=None, + model="vertex_ai/gemini-3.5-transcribe-preview", + optional_params={}, + litellm_params={"vertex_project": "test-project"}, + ) + assert "/models/gemini-3.5-transcribe-preview:generateContent" in url + assert "vertex_ai/" not in url + + def test_api_base_override(self, config): + url = config.get_complete_url( + api_base="http://localhost:8080/", + api_key=None, + model="gemini-3.5-transcribe-preview", + optional_params={}, + litellm_params={"vertex_project": "test-project"}, + ) + assert url == ( + "http://localhost:8080/v1/projects/test-project/locations/global" + "/publishers/google/models/gemini-3.5-transcribe-preview:generateContent" + ) + + @pytest.mark.parametrize("malicious_location", ["attacker.example/", "evil.com#", "US", "us/../.."]) + def test_malicious_location_is_rejected(self, config, malicious_location): + with pytest.raises(VertexAIError): + config.get_complete_url( + api_base=None, + api_key=None, + model="gemini-3.5-transcribe-preview", + optional_params={}, + litellm_params={"vertex_project": "test-project", "vertex_location": malicious_location}, + ) + + @pytest.mark.parametrize("malicious_project", ["proj/../../locations", "proj#frag", "proj?a=b", "proj space"]) + def test_malicious_project_is_rejected(self, config, malicious_project): + with pytest.raises(VertexAIError): + config.get_complete_url( + api_base=None, + api_key=None, + model="gemini-3.5-transcribe-preview", + optional_params={}, + litellm_params={"vertex_project": malicious_project}, + ) + + +class TestTransformRequest: + def test_request_body_shape(self, config): + request_data = config.transform_audio_transcription_request( + model="gemini-3.5-transcribe-preview", + audio_file=AUDIO_BYTES, + optional_params={}, + litellm_params={}, + ) + assert request_data.files is None + assert request_data.data == { + "contents": ( + { + "role": "user", + "parts": ( + { + "inlineData": { + "mimeType": "audio/wav", + "data": base64.b64encode(AUDIO_BYTES).decode("utf-8"), + } + }, + ), + }, + ), + "generationConfig": {"audioTranscriptionConfig": {}}, + } + + @pytest.mark.parametrize( + "language,expected_language_codes", + [ + ("en", ("en-US",)), + ("en-US", ("en-US",)), + ("fr", ("fr-FR",)), + ], + ) + def test_language_param_maps_to_language_codes(self, config, language, expected_language_codes): + request_data = config.transform_audio_transcription_request( + model="gemini-3.5-transcribe-preview", + audio_file=AUDIO_BYTES, + optional_params={"language": language}, + litellm_params={}, + ) + audio_config = request_data.data["generationConfig"]["audioTranscriptionConfig"] + assert audio_config["languageCodes"] == expected_language_codes + + def test_body_round_trips_through_json(self, config): + request_data = config.transform_audio_transcription_request( + model="gemini-3.5-transcribe-preview", + audio_file=AUDIO_BYTES, + optional_params={"language": "en"}, + litellm_params={}, + ) + round_tripped = json.loads(json.dumps(request_data.data)) + assert round_tripped["generationConfig"] == {"audioTranscriptionConfig": {"languageCodes": ["en-US"]}} + assert round_tripped["contents"][0]["role"] == "user" + + +class TestTransformResponse: + def test_generate_content_response(self, config): + raw_response = httpx.Response(status_code=200, json=GENERATE_CONTENT_RESPONSE) + response = config.transform_audio_transcription_response(raw_response) + assert response.text == TRANSCRIPT_TEXT + assert response["task"] == "transcribe" + assert isinstance(response.usage, TranscriptionUsageTokensObject) + assert response.usage.input_tokens == 440 + assert response.usage.output_tokens == 62 + assert response.usage.total_tokens == 502 + assert response.usage.input_token_details.audio_tokens == 440 + assert response.usage.input_token_details.text_tokens == 0 + + def test_multi_part_texts_are_joined(self, config): + raw_response = httpx.Response( + status_code=200, + json={ + "candidates": [ + {"content": {"role": "model", "parts": [{"text": "Hello world."}, {"text": "How are you?"}]}} + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15}, + }, + ) + response = config.transform_audio_transcription_response(raw_response) + assert response.text == "Hello world. How are you?" + + def test_empty_candidates_returns_empty_text(self, config): + raw_response = httpx.Response(status_code=200, json={}) + response = config.transform_audio_transcription_response(raw_response) + assert response.text == "" + assert response.usage is None + + def test_non_json_body_raises(self, config): + raw_response = httpx.Response(status_code=200, text="not json") + with pytest.raises(VertexAIError, match="non-JSON"): + config.transform_audio_transcription_response(raw_response) + + +class TestValidateEnvironment: + def test_sets_oauth_headers(self): + class StubbedConfig(VertexGeminiAudioTranscriptionConfig): + def _ensure_access_token(self, credentials, project_id, custom_llm_provider): + return "fake-token", "resolved-project" + + headers = StubbedConfig().validate_environment( + headers={}, + model="gemini-3.5-transcribe-preview", + messages=[], + optional_params={}, + litellm_params={"vertex_project": "resolved-project"}, + ) + assert headers["Authorization"] == "Bearer fake-token" + assert headers["x-goog-user-project"] == "resolved-project" + assert headers["Content-Type"] == "application/json" + + +class TestOptionalParams: + def test_language_and_json_response_format_pass_through(self): + optional_params = get_optional_params_transcription( + model="gemini-3.5-transcribe-preview", + custom_llm_provider="vertex_ai", + language="fr-FR", + response_format="json", + ) + assert optional_params["language"] == "fr-FR" + assert optional_params["response_format"] == "json" + + @pytest.mark.parametrize("response_format", ["verbose_json", "srt", "vtt"]) + def test_unsupported_response_format_raises(self, response_format): + with pytest.raises(litellm.utils.UnsupportedParamsError, match="response_format"): + get_optional_params_transcription( + model="gemini-3.5-transcribe-preview", + custom_llm_provider="vertex_ai", + response_format=response_format, + ) + + @pytest.mark.parametrize("response_format", ["verbose_json", "srt", "vtt"]) + def test_unsupported_response_format_dropped_with_drop_params(self, response_format): + optional_params = get_optional_params_transcription( + model="gemini-3.5-transcribe-preview", + custom_llm_provider="vertex_ai", + language="fr-FR", + response_format=response_format, + drop_params=True, + ) + assert "response_format" not in optional_params + assert optional_params["language"] == "fr-FR" + + +class TestModelCostEntry: + REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) + + @pytest.mark.parametrize( + "cost_map_path", + [ + "model_prices_and_context_window.json", + "litellm/model_prices_and_context_window_backup.json", + ], + ) + def test_transcribe_preview_pricing(self, cost_map_path): + with open(os.path.join(self.REPO_ROOT, cost_map_path)) as f: + entry = json.load(f)["vertex_ai/gemini-3.5-transcribe-preview"] + assert entry["mode"] == "audio_transcription" + assert entry["litellm_provider"] == "vertex_ai" + assert entry["input_cost_per_audio_token"] == pytest.approx(2.5e-06) + assert entry["input_cost_per_token"] == pytest.approx(2.5e-06) + assert entry["output_cost_per_token"] == pytest.approx(1.2e-05) + assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"] + + @pytest.mark.parametrize( + "cost_map_path", + [ + "model_prices_and_context_window.json", + "litellm/model_prices_and_context_window_backup.json", + ], + ) + def test_transcribe_live_preview_pricing(self, cost_map_path): + with open(os.path.join(self.REPO_ROOT, cost_map_path)) as f: + entry = json.load(f)["vertex_ai/gemini-3.5-transcribe-live-preview"] + assert entry["mode"] == "audio_transcription" + assert entry["litellm_provider"] == "vertex_ai" + assert entry["input_cost_per_audio_token"] == pytest.approx(3.5e-06) + assert entry["input_cost_per_token"] == pytest.approx(3.5e-06) + assert entry["output_cost_per_token"] == pytest.approx(2.1e-05) + assert entry["supported_endpoints"] == ["/v1/realtime"] 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/llms/xai/test_xai_cost_calculator.py b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py index 55e28dff81d..92e76fd18ab 100644 --- a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -51,8 +51,8 @@ class TestXAICostCalculator: # Expected costs for grok-3-mini: # Input: 12 tokens * $3e-7 = $0.0000036 # Output: 125 tokens * $5e-7 = $0.0000625 - expected_prompt_cost = 12 * 3e-7 - expected_completion_cost = 125 * 5e-7 + expected_prompt_cost = 12 * 1.25e-6 + expected_completion_cost = 125 * 2.5e-6 assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) @@ -77,8 +77,8 @@ class TestXAICostCalculator: # Expected costs for grok-3-mini: # Input: 12 tokens * $3e-7 = $0.0000036 # Completion: (125 + 949) tokens * $5e-7 = $0.000537 - expected_prompt_cost = 12 * 3e-7 - expected_completion_cost = (125 + 949) * 5e-7 + expected_prompt_cost = 12 * 1.25e-6 + expected_completion_cost = (125 + 949) * 2.5e-6 assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) @@ -104,8 +104,8 @@ class TestXAICostCalculator: # Input: 12 tokens * $3e-7 = $0.0000036 # Completion: (125 + 949) tokens * $5e-7 = $0.000537 # Note: text_tokens field is ignored, only completion_tokens + reasoning_tokens matters - expected_prompt_cost = 12 * 3e-7 - expected_completion_cost = (125 + 949) * 5e-7 + expected_prompt_cost = 12 * 1.25e-6 + expected_completion_cost = (125 + 949) * 2.5e-6 assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) @@ -127,11 +127,12 @@ class TestXAICostCalculator: prompt_cost, completion_cost = cost_per_token(model="grok-4", usage=usage) - # Expected costs for grok-4: - # Input: 10 tokens * $3e-6 = $0.00003 - # Completion: (200 + 150) tokens * $1.5e-5 = $0.00525 - expected_prompt_cost = 10 * 3e-6 - expected_completion_cost = (200 + 150) * 1.5e-5 + # grok-4 was retired on 2026-05-15 and now redirects to grok-4.3, so it bills + # at grok-4.3's rates: + # Input: 10 tokens * $1.25e-6 + # Completion: (200 + 150) tokens * $2.5e-6 + expected_prompt_cost = 10 * 1.25e-6 + expected_completion_cost = (200 + 150) * 2.5e-6 assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) @@ -158,8 +159,8 @@ class TestXAICostCalculator: # Expected costs for grok-3-fast-beta: # Input: 20 tokens * $5e-6 = $0.0001 # Completion: (300 + 200) tokens * $2.5e-5 = $0.0125 - expected_prompt_cost = 20 * 5e-6 - expected_completion_cost = (300 + 200) * 2.5e-5 + expected_prompt_cost = 20 * 1.25e-6 + expected_completion_cost = (300 + 200) * 2.5e-6 assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) @@ -185,46 +186,34 @@ class TestXAICostCalculator: # Expected costs: # Input: 12 tokens * $3e-7 = $0.0000036 # Completion: (50 + 100) tokens * $5e-7 = $0.000075 - expected_prompt_cost = 12 * 3e-7 - expected_completion_cost = (50 + 100) * 5e-7 + expected_prompt_cost = 12 * 1.25e-6 + expected_completion_cost = (50 + 100) * 2.5e-6 assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) - def test_tiered_pricing_above_128k_tokens(self): - """Test tiered pricing for tokens above 128k.""" - # Test with grok-4-fast-reasoning which has tiered pricing + def test_tiered_pricing_above_200k_tokens(self): usage = Usage( - prompt_tokens=150000, # Above 128k threshold - completion_tokens=100000, # Above 128k threshold - total_tokens=300000, + prompt_tokens=250000, + completion_tokens=100000, + total_tokens=400000, completion_tokens_details=CompletionTokensDetailsWrapper( accepted_prediction_tokens=0, audio_tokens=0, - reasoning_tokens=50000, # Total completion tokens = 100000 + 50000 = 150000 > 128k + reasoning_tokens=50000, rejected_prediction_tokens=0, text_tokens=None, ), ) - - prompt_cost, completion_cost = cost_per_token( - model="xai/grok-4-fast-reasoning", usage=usage - ) - - # Expected costs for grok-4-fast-reasoning with tiered pricing: - # Input: 150000 tokens * $0.4e-6 (ALL tokens at tiered rate since input > 128k) = $0.06 - # Completion: (100000 + 50000) tokens * $1e-6 (tiered rate since input > 128k) = $0.15 - expected_prompt_cost = 150000 * 0.4e-6 - expected_completion_cost = (100000 + 50000) * 1e-6 - + prompt_cost, completion_cost = cost_per_token(model="xai/grok-4.3", usage=usage) + expected_prompt_cost = 250000 * 2.5e-6 + expected_completion_cost = (100000 + 50000) * 5e-6 assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) - def test_tiered_pricing_below_128k_tokens(self): - """Test that regular pricing is used for tokens below 128k threshold.""" - # Test with grok-4-fast-reasoning which has tiered pricing + def test_tiered_pricing_below_200k_tokens(self): usage = Usage( - prompt_tokens=100000, # Below 128k threshold + prompt_tokens=100000, completion_tokens=50000, total_tokens=160000, completion_tokens_details=CompletionTokensDetailsWrapper( @@ -235,26 +224,18 @@ class TestXAICostCalculator: text_tokens=None, ), ) - - prompt_cost, completion_cost = cost_per_token( - model="xai/grok-4-fast-reasoning", usage=usage - ) - - # Expected costs for grok-4-fast-reasoning with regular pricing: - # Input: 100000 tokens * $0.2e-6 (regular rate) = $0.02 - # Completion: (50000 + 10000) tokens * $0.5e-6 (regular rate) = $0.03 - expected_prompt_cost = 100000 * 0.2e-6 - expected_completion_cost = (50000 + 10000) * 0.5e-6 - + prompt_cost, completion_cost = cost_per_token(model="xai/grok-4.3", usage=usage) + expected_prompt_cost = 100000 * 1.25e-6 + expected_completion_cost = (50000 + 10000) * 2.5e-6 assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) def test_tiered_pricing_grok_4_latest(self): """Test tiered pricing for grok-4-latest model.""" usage = Usage( - prompt_tokens=200000, # Above 128k threshold + prompt_tokens=250000, # Above the 200k threshold completion_tokens=100000, - total_tokens=350000, + total_tokens=400000, completion_tokens_details=CompletionTokensDetailsWrapper( accepted_prediction_tokens=0, audio_tokens=0, @@ -268,59 +249,45 @@ class TestXAICostCalculator: model="xai/grok-4-latest", usage=usage ) - # Expected costs for grok-4-latest with tiered pricing: - # Input: 200000 tokens * $6e-6 (ALL tokens at tiered rate since input > 128k) = $1.2 - # Completion: (100000 + 50000) tokens * $30e-6 (tiered rate since input > 128k) = $4.5 - expected_prompt_cost = 200000 * 6e-6 - expected_completion_cost = (100000 + 50000) * 30e-6 + # grok-4-latest redirects to grok-4.3, which tiers at 200k rather than 128k: + # Input: 250000 tokens * $2.5e-6 (ALL tokens at tiered rate since input > 200k) + # Completion: (100000 + 50000) tokens * $5e-6 (tiered rate since input > 200k) + expected_prompt_cost = 250000 * 2.5e-6 + expected_completion_cost = (100000 + 50000) * 5e-6 assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) - def test_tiered_pricing_output_tokens_below_128k(self): - """Test that output tokens get tiered rate when input tokens > 128k, even if output tokens < 128k.""" + def test_tiered_pricing_output_tokens_below_200k(self): usage = Usage( - prompt_tokens=150000, # Above 128k threshold - completion_tokens=50000, # Below 128k threshold - total_tokens=210000, + prompt_tokens=250000, + completion_tokens=50000, + total_tokens=310000, completion_tokens_details=CompletionTokensDetailsWrapper( accepted_prediction_tokens=0, audio_tokens=0, - reasoning_tokens=10000, # Total completion tokens = 50000 + 10000 = 60000 < 128k + reasoning_tokens=10000, rejected_prediction_tokens=0, text_tokens=None, ), ) - - prompt_cost, completion_cost = cost_per_token( - model="xai/grok-4-fast-reasoning", usage=usage - ) - - # Expected costs for grok-4-fast-reasoning: - # Input: 150000 tokens * $0.4e-6 (ALL tokens at tiered rate since input > 128k) = $0.06 - # Completion: (50000 + 10000) tokens * $1e-6 (tiered rate since input > 128k) = $0.06 - expected_prompt_cost = 150000 * 0.4e-6 - expected_completion_cost = (50000 + 10000) * 1e-6 - + prompt_cost, completion_cost = cost_per_token(model="xai/grok-4.3", usage=usage) + expected_prompt_cost = 250000 * 2.5e-6 + expected_completion_cost = (50000 + 10000) * 5e-6 assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) def test_tiered_pricing_model_without_tiered_pricing(self): - """Test that models without tiered pricing use regular pricing even above 128k.""" - usage = Usage( - prompt_tokens=150000, # Above 128k threshold - completion_tokens=50000, - total_tokens=200000, - ) - - prompt_cost, completion_cost = cost_per_token(model="grok-3-mini", usage=usage) - - # grok-3-mini doesn't have tiered pricing, so should use regular rates: - # Input: 150000 tokens * $3e-7 (regular rate) = $0.045 - # Completion: 50000 tokens * $5e-7 (regular rate) = $0.025 - expected_prompt_cost = 150000 * 3e-7 + litellm.model_cost["xai/flat-rate-fixture"] = { + "input_cost_per_token": 3e-7, + "output_cost_per_token": 5e-7, + "litellm_provider": "xai", + "mode": "chat", + } + usage = Usage(prompt_tokens=250000, completion_tokens=50000, total_tokens=300000) + prompt_cost, completion_cost = cost_per_token(model="xai/flat-rate-fixture", usage=usage) + expected_prompt_cost = 250000 * 3e-7 expected_completion_cost = 50000 * 5e-7 - assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) @@ -341,8 +308,8 @@ class TestXAICostCalculator: prompt_cost, completion_cost = cost_per_token(model="grok-3-mini", usage=usage) - expected_prompt_cost = 12 * 3e-7 - expected_completion_cost = 200 * 5e-7 + expected_prompt_cost = 12 * 1.25e-6 + expected_completion_cost = 200 * 2.5e-6 assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) diff --git a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py new file mode 100644 index 00000000000..83c3bf1ecef --- /dev/null +++ b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py @@ -0,0 +1,144 @@ +""" +xAI retired eight slugs on 2026-05-15 but kept them resolvable: chat slugs redirect to +grok-4.3 and bill at grok-4.3's rates, while the grok-code-fast slugs are aliases of +grok-build-0.1 and bill at its rates, so the registry must price them that way or spend +tracking is wrong. The grok-3-beta, grok-3-fast, grok-3-mini, and grok-4-1-fast slugs +are absent from /v1/language-models and resolve to grok-4.3 the same way (the chat +response names grok-4.3 as the served model), so they carry grok-4.3's rates too. +https://docs.x.ai/developers/migration/may-15-retirement +https://docs.x.ai/developers/models/grok-build-0.1 +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).parents[4] +PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PRICES_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" +MAP_PATHS = (PRICES_PATH, BACKUP_PRICES_PATH) + +REDIRECT_TARGET = "xai/grok-4.3" +GROK_3_MINI_SLUGS = ( + "xai/grok-3-mini", + "xai/grok-3-mini-beta", + "xai/grok-3-mini-fast", + "xai/grok-3-mini-fast-beta", + "xai/grok-3-mini-fast-latest", + "xai/grok-3-mini-latest", +) +REDIRECTED_SLUGS = ( + "xai/grok-3", + "xai/grok-3-beta", + "xai/grok-3-fast-beta", + "xai/grok-3-fast-latest", + "xai/grok-3-latest", + *GROK_3_MINI_SLUGS, + "xai/grok-4", + "xai/grok-4-0709", + "xai/grok-4-1-fast", + "xai/grok-4-1-fast-non-reasoning", + "xai/grok-4-1-fast-non-reasoning-latest", + "xai/grok-4-1-fast-reasoning", + "xai/grok-4-1-fast-reasoning-latest", + "xai/grok-4-fast-non-reasoning", + "xai/grok-4-fast-reasoning", + "xai/grok-4-latest", +) +CODE_REDIRECT_TARGET = "xai/grok-build-0.1" +CODE_SLUGS = ( + "xai/grok-code-fast", + "xai/grok-code-fast-1", + "xai/grok-code-fast-1-0825", +) +RETIREMENT_DATE = "2026-05-15" +GROK_3_MINI_RETIREMENT_DATE = "2026-02-28" + +BASE_COST_FIELDS = ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost") +TIER_COST_FIELDS = ( + "input_cost_per_token_above_200k_tokens", + "output_cost_per_token_above_200k_tokens", + "cache_read_input_token_cost_above_200k_tokens", +) +STALE_TIER_FIELDS = ( + "input_cost_per_token_above_128k_tokens", + "output_cost_per_token_above_128k_tokens", + "cache_read_input_token_cost_above_128k_tokens", +) + + +def expected_retirement_date(slug: str) -> str: + return GROK_3_MINI_RETIREMENT_DATE if slug in GROK_3_MINI_SLUGS else RETIREMENT_DATE + + +@pytest.fixture(scope="module", params=[p.name for p in MAP_PATHS]) +def cost_map(request: pytest.FixtureRequest) -> dict: + path = next(p for p in MAP_PATHS if p.name == request.param) + return json.loads(path.read_text(encoding="utf-8")) + + +@pytest.mark.parametrize("slug", REDIRECTED_SLUGS) +def test_redirected_slug_bills_at_the_target_rate(cost_map: dict, slug: str): + target = cost_map[REDIRECT_TARGET] + entry = cost_map[slug] + for field in BASE_COST_FIELDS: + assert entry[field] == target[field], field + + +@pytest.mark.parametrize("slug", CODE_SLUGS) +def test_code_slug_bills_at_grok_build_rate(cost_map: dict, slug: str): + """grok-code-fast* are aliases of grok-build-0.1, not grok-4.3 redirects.""" + target = cost_map[CODE_REDIRECT_TARGET] + entry = cost_map[slug] + for field in (*BASE_COST_FIELDS, *TIER_COST_FIELDS): + assert entry[field] == target[field], field + + +@pytest.mark.parametrize("slug", (*REDIRECTED_SLUGS, *CODE_SLUGS)) +def test_redirected_slug_keeps_its_retirement_date(cost_map: dict, slug: str): + assert cost_map[slug]["deprecation_date"] == expected_retirement_date(slug) + + +@pytest.mark.parametrize("slug", REDIRECTED_SLUGS) +def test_no_slug_keeps_the_superseded_128k_tier(cost_map: dict, slug: str): + """The 128k tier belonged to the retired model; grok-4.3 tiers at 200k.""" + for field in STALE_TIER_FIELDS: + assert field not in cost_map[slug], field + + +@pytest.mark.parametrize("slug", REDIRECTED_SLUGS) +def test_redirected_slug_carries_the_target_tier_rates(cost_map: dict, slug: str): + """The request executes as grok-4.3, so it is tiered at grok-4.3's 200k boundary.""" + target = cost_map[REDIRECT_TARGET] + entry = cost_map[slug] + for field in TIER_COST_FIELDS: + assert entry[field] == target[field], field + + +def test_a_live_xai_model_is_untouched(cost_map: dict): + """Guard against the repricing leaking onto models xAI still serves directly.""" + assert cost_map["xai/grok-4.6"]["input_cost_per_token"] != cost_map[REDIRECT_TARGET]["input_cost_per_token"] + assert "deprecation_date" not in cost_map["xai/grok-4.6"] + + +def test_both_cost_maps_agree_on_the_redirected_slugs(): + prices = json.loads(PRICES_PATH.read_text(encoding="utf-8")) + backup = json.loads(BACKUP_PRICES_PATH.read_text(encoding="utf-8")) + for slug in (*REDIRECTED_SLUGS, *CODE_SLUGS, REDIRECT_TARGET, CODE_REDIRECT_TARGET): + assert prices[slug] == backup[slug], slug + + +def test_every_retired_chat_slug_is_covered(cost_map: dict): + """The lists above must stay in step with what the registry marks retired.""" + marked = { + key + for key, entry in cost_map.items() + if isinstance(entry, dict) + and entry.get("litellm_provider") == "xai" + and "deprecation_date" in entry + and entry.get("mode") == "chat" + } + assert marked == {*REDIRECTED_SLUGS, *CODE_SLUGS} 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/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index 9a90daeccb7..c83ba142011 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -164,6 +164,49 @@ class TestProxyExceptionPassthrough: mock_logging.post_call_failure_hook.assert_awaited_once() +class TestHttpExceptionDictDetail: + @pytest.mark.asyncio + async def test_anthropic_response_serializes_dict_detail_http_exception(self): + """LIT-6466: a post_call guardrail's HTTPException(detail=) must + surface with a clean message plus provider_specific_fields, matching + /v1/chat/completions and /v1/responses, not the str() of the exception.""" + from fastapi import HTTPException + + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import ProxyException, UserAPIKeyAuth + + detail = { + "error": "Content blocked: keyword 'kumquat' detected", + "keyword": "kumquat", + "guardrail": "keyword-block", + } + exc = HTTPException(status_code=400, detail=detail) + + with ( + patch.object(ep, "_read_request_body", new=AsyncMock(return_value={})), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object( # test-quality-ok: the guardrail raise happens deep inside this call; the test targets the endpoint's except block + ep.ProxyBaseLLMRequestProcessing, + "base_process_llm_request", + new=AsyncMock(side_effect=exc), + ), + patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam + ): + mock_logging.post_call_failure_hook = AsyncMock() + with pytest.raises(ProxyException) as exc_info: + await ep.anthropic_response( + fastapi_response=MagicMock(), + request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert exc_info.value.message == "Content blocked: keyword 'kumquat' detected" + assert "{'error'" not in exc_info.value.message + assert exc_info.value.provider_specific_fields == detail + assert exc_info.value.code == "400" + mock_logging.post_call_failure_hook.assert_awaited_once() + + class TestFailureHookRequestData: @pytest.mark.asyncio async def test_failure_hook_gets_post_setup_data_with_logging_obj(self): 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/auth/test_model_access_group_budgets.py b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py new file mode 100644 index 00000000000..fb82d8708fd --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py @@ -0,0 +1,543 @@ +""" +Which model access groups a request is charged to. + +A group is attributed only when its name appears on an allowlist the caller was granted, so the +group is what authorized the call. Asking for a model that merely belongs to a group attributes +nothing, and every level that can name a group (key, team, team-member scope, project, org) is +unioned rather than ranked. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +import litellm +from litellm import Router +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + Litellm_EntityType, + LiteLLM_OrganizationTable, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import ( + _model_access_group_max_budget_check, + collect_matched_model_access_groups, + common_checks, + stamp_matched_model_access_groups, +) +from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + model_access_group_registry_cache_key, + model_access_group_spend_counter_key, + team_membership_reservation_cache_key, +) +from litellm.proxy.utils import ProxyLogging + +TEAM_ID = "team-1" +USER_ID = "user-1" +ORG_ID = "org-1" +BUDGETED_GROUPS = ("tier-a", "tier-b", "claude-tier") +MODEL_ACCESS_GROUP_COUNTER_KEY = model_access_group_spend_counter_key("tier-a") + +MODEL_LIST = [ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k"}, + "model_info": {"access_groups": ["tier-a", "tier-b"]}, + }, + { + "model_name": "claude-sonnet", + "litellm_params": {"model": "anthropic/claude-sonnet", "api_key": "k"}, + "model_info": {"access_groups": ["claude-tier"]}, + }, +] + + +class _ExplodingPrismaClient: + """Every lookup in these tests is served from the injected cache; a real DB read is a bug.""" + + def __getattr__(self, name: str) -> object: + raise AssertionError(f"unexpected database access: {name}") + + +class _CountingRouter(Router): + """Counts access-group lookups, so a test can prove the registry gate skipped them.""" + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.access_group_lookups = 0 + + def get_model_access_groups(self, *args, **kwargs): + self.access_group_lookups += 1 + return super().get_model_access_groups(*args, **kwargs) + + +async def _cache( + budgeted_groups: tuple[str, ...] = BUDGETED_GROUPS, + member_allowed_models: tuple[str, ...] = (), + org_models: tuple[str, ...] = (), +) -> UserApiKeyCache: + cache = UserApiKeyCache() + await cache.async_set_cache(key=model_access_group_registry_cache_key(), value=budgeted_groups) + if member_allowed_models: + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id=USER_ID, team_id=TEAM_ID), + value=LiteLLM_TeamMembership( + user_id=USER_ID, + team_id=TEAM_ID, + budget_id="member-budget", + litellm_budget_table=LiteLLM_BudgetTable(allowed_models=list(member_allowed_models)), + ), + model_type=LiteLLM_TeamMembership, + ) + if org_models: + await cache.async_set_cache( + key=f"org_id:{ORG_ID}", + value=LiteLLM_OrganizationTable( + organization_id=ORG_ID, + budget_id="org-budget", + models=list(org_models), + created_by=USER_ID, + updated_by=USER_ID, + ), + model_type=LiteLLM_OrganizationTable, + ) + return cache + + +async def _matched( + *, + model: str = "gpt-4o", + key_models: list[str] | None = None, + team_models: list[str] | None = None, + team_org_id: str | None = None, + project_models: list[str] | None = None, + valid_token: UserAPIKeyAuth | None = None, + cache: UserApiKeyCache | None = None, + llm_router: Router | None = None, +) -> tuple[str, ...]: + resolved_cache = cache if cache is not None else await _cache() + return await collect_matched_model_access_groups( + model=model, + valid_token=valid_token + if valid_token is not None + else UserAPIKeyAuth(api_key="hashed", models=key_models or [], team_id=TEAM_ID, user_id=USER_ID), + team_object=( + LiteLLM_TeamTable(team_id=TEAM_ID, models=team_models, organization_id=team_org_id) + if team_models is not None + else None + ), + project_object=( + LiteLLM_ProjectTableCachedObj(project_id="project-1", models=project_models) + if project_models is not None + else None + ), + llm_router=llm_router if llm_router is not None else Router(model_list=MODEL_LIST), + prisma_client=_ExplodingPrismaClient(), + user_api_key_cache=resolved_cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=resolved_cache), + ) + + +@pytest.mark.asyncio +async def test_group_named_on_the_key_is_attributed(): + assert await _matched(key_models=["tier-a"]) == ("tier-a",) + + +@pytest.mark.asyncio +async def test_model_granted_directly_on_the_key_attributes_nothing(): + assert await _matched(key_models=["gpt-4o"]) == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("key_models", [["*"], [], ["all-proxy-models"]]) +async def test_unrestricted_key_attributes_nothing(key_models: list[str]): + assert await _matched(key_models=key_models) == () + + +@pytest.mark.asyncio +async def test_group_that_does_not_serve_the_requested_model_is_not_attributed(): + assert await _matched(model="gpt-4o", key_models=["claude-tier"]) == () + + +@pytest.mark.asyncio +async def test_both_granted_groups_covering_the_model_are_attributed(): + assert await _matched(key_models=["tier-b", "tier-a"]) == ("tier-a", "tier-b") + + +@pytest.mark.asyncio +async def test_group_named_only_on_the_team_is_attributed(): + assert await _matched(key_models=[], team_models=["tier-a"]) == ("tier-a",) + + +@pytest.mark.asyncio +async def test_group_named_only_in_a_team_members_scope_is_attributed(): + assert await _matched( + model="claude-sonnet", + key_models=["*"], + team_models=["*"], + cache=await _cache(member_allowed_models=("claude-tier",)), + ) == ("claude-tier",) + + +@pytest.mark.asyncio +async def test_group_named_only_on_the_project_is_attributed(): + assert await _matched(key_models=["*"], project_models=["tier-b"]) == ("tier-b",) + + +@pytest.mark.asyncio +async def test_group_named_only_on_the_org_is_attributed(): + assert await _matched( + valid_token=UserAPIKeyAuth(api_key="hashed", models=["*"], user_id=USER_ID, org_id=ORG_ID), + cache=await _cache(org_models=("tier-a",)), + ) == ("tier-a",) + + +@pytest.mark.asyncio +async def test_group_named_on_the_teams_org_is_attributed_when_the_key_names_no_org(): + assert await _matched( + key_models=["*"], + team_models=["*"], + team_org_id=ORG_ID, + cache=await _cache(org_models=("tier-b",)), + ) == ("tier-b",) + + +@pytest.mark.asyncio +async def test_all_team_models_sentinel_on_the_key_resolves_to_the_teams_groups(): + assert await _matched( + valid_token=UserAPIKeyAuth( + api_key="hashed", + models=["all-team-models"], + team_models=["tier-a"], + team_id=TEAM_ID, + user_id=USER_ID, + ), + ) == ("tier-a",) + + +@pytest.mark.asyncio +async def test_group_without_a_budget_is_not_attributed(): + assert await _matched(key_models=["tier-a"], cache=await _cache(budgeted_groups=("tier-b",))) == () + + +@pytest.mark.asyncio +async def test_empty_registry_skips_the_access_group_matching_entirely(): + router = _CountingRouter(model_list=MODEL_LIST) + + assert await _matched(key_models=["tier-a"], cache=await _cache(budgeted_groups=()), llm_router=router) == () + assert router.access_group_lookups == 0 + + assert await _matched(key_models=["tier-a"], llm_router=router) == ("tier-a",) + assert router.access_group_lookups == 1 + + +@pytest.mark.asyncio +async def test_stamp_records_the_matched_groups_on_the_auth_object(): + cache = await _cache() + valid_token = UserAPIKeyAuth(api_key="hashed", models=["tier-a", "tier-b"], team_id=TEAM_ID, user_id=USER_ID) + + await stamp_matched_model_access_groups( + model="gpt-4o", + valid_token=valid_token, + team_object=None, + project_object=None, + llm_router=Router(model_list=MODEL_LIST), + prisma_client=_ExplodingPrismaClient(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + ) + + assert valid_token.matched_model_access_groups == ["tier-a", "tier-b"] + + +class _BrokenRouter(Router): + def get_model_access_groups(self, *args, **kwargs): + raise RuntimeError("access group store unavailable") + + +@pytest.mark.asyncio +async def test_stamp_does_not_break_auth_when_the_access_group_lookup_fails(): + cache = await _cache() + valid_token = UserAPIKeyAuth(api_key="hashed", models=["tier-a"], team_id=TEAM_ID, user_id=USER_ID) + + await stamp_matched_model_access_groups( + model="gpt-4o", + valid_token=valid_token, + team_object=None, + project_object=None, + llm_router=_BrokenRouter(model_list=MODEL_LIST), + prisma_client=_ExplodingPrismaClient(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + ) + + assert valid_token.matched_model_access_groups is None + + +@pytest.mark.asyncio +async def test_stamp_leaves_the_auth_object_untouched_when_nothing_matched(): + cache = await _cache() + valid_token = UserAPIKeyAuth(api_key="hashed", models=["gpt-4o"], team_id=TEAM_ID, user_id=USER_ID) + + await stamp_matched_model_access_groups( + model="gpt-4o", + valid_token=valid_token, + team_object=None, + project_object=None, + llm_router=Router(model_list=MODEL_LIST), + prisma_client=_ExplodingPrismaClient(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + ) + + assert valid_token.matched_model_access_groups is None + + +class _MagBudgetRow: + """One ``LiteLLM_ModelAccessGroupBudgetTable`` row as prisma hands it back.""" + + def __init__(self, access_group_name: str, spend: float = 0.0, max_budget: float | None = None) -> None: + self.access_group_name = access_group_name + self.spend = spend + self.litellm_budget_table = None if max_budget is None else SimpleNamespace(max_budget=max_budget) + + +class _RecordingPrismaClient: + """Serves budget rows and records which groups actually reached the database.""" + + def __init__(self, *rows: _MagBudgetRow) -> None: + self.rows = {row.access_group_name: row for row in rows} + self.batches: list[list[str]] = [] + self.db = SimpleNamespace( + litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many) + ) + + async def _find_many(self, **kwargs): + requested = list(kwargs["where"]["access_group_name"]["in"]) + self.batches.append(requested) + return [self.rows[group] for group in requested if group in self.rows] + + +def _spend_reader(spend_by_counter_key: dict[str, float]): + """Stand-in for proxy_server.get_current_spend, recording every counter key it is asked for.""" + seen: list[str] = [] + + async def read(counter_key, fallback_spend, max_budget=None, **kwargs): + seen.append(counter_key) + return spend_by_counter_key.get(counter_key, fallback_spend) + + return read, seen + + +async def _enforce( + matched: tuple[str, ...], + *rows: _MagBudgetRow, + spend_by_counter_key: dict[str, float] | None = None, + prisma_client: object | None = None, + cache: UserApiKeyCache | None = None, +) -> list[str]: + read, seen = _spend_reader(spend_by_counter_key or {}) + # The check takes its client and cache as arguments, injected just below. get_current_spend is the + # one collaborator it reaches by a lazy `from litellm.proxy.proxy_server import`, with no parameter. + with patch("litellm.proxy.proxy_server.get_current_spend", read): # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point + await _model_access_group_max_budget_check( + matched_model_access_groups=matched, + prisma_client=prisma_client if prisma_client is not None else _RecordingPrismaClient(*rows), + user_api_key_cache=cache if cache is not None else UserApiKeyCache(), + ) + return seen + + +@pytest.mark.asyncio +async def test_group_under_its_max_budget_passes(): + assert await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", spend=4.0, max_budget=10.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 4.0}, + ) == [MODEL_ACCESS_GROUP_COUNTER_KEY] + + +@pytest.mark.asyncio +async def test_group_exactly_at_its_max_budget_blocks_the_request(): + """A pool whose spend has reached the ceiling has nothing left, so the next request is refused. + + This is where the check departs from the tag one it otherwise mirrors, and it matches where + keys and organizations already draw the line. + """ + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", max_budget=10.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 10.0}, + ) + + assert exc_info.value.entity_id == "tier-a" + assert exc_info.value.current_cost == 10.0 + + +@pytest.mark.asyncio +async def test_group_just_under_its_max_budget_passes(): + """Asserting the counter was read is what keeps this honest: a group that got skipped entirely, + because its row never arrived or carried no budget, would also not raise.""" + assert await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", max_budget=10.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 9.99}, + ) == [MODEL_ACCESS_GROUP_COUNTER_KEY] + + +@pytest.mark.asyncio +async def test_a_non_positive_budget_means_no_budget(): + """The reservation path treats max_budget <= 0 as unbudgeted, so the read-time check must agree. + + Without this the exclusive ceiling would turn a zero into a total freeze on one path and a + no-op on the other. + """ + assert ( + await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", max_budget=0.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 5.0}, + ) + == [] + ) + + +@pytest.mark.asyncio +async def test_group_over_its_max_budget_blocks_the_request_and_names_the_group(): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", max_budget=10.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 10.5}, + ) + + assert exc_info.value.entity_id == "tier-a" + assert exc_info.value.entity_type == Litellm_EntityType.MODEL_ACCESS_GROUP.value + assert exc_info.value.current_cost == 10.5 + assert exc_info.value.max_budget == 10.0 + assert "tier-a" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_group_with_a_row_but_no_budget_never_blocks(): + """An admin can register a group without a ceiling; that must not become an implicit zero budget.""" + assert ( + await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", spend=9999.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 9999.0}, + ) + == [] + ) + + +@pytest.mark.asyncio +async def test_a_cold_counter_falls_back_to_the_spend_recorded_on_the_row(): + """After a counter expires the DB row is the only record of the spend, so it has to be read.""" + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _enforce(("tier-a",), _MagBudgetRow("tier-a", spend=12.0, max_budget=10.0)) + + assert exc_info.value.current_cost == 12.0 + + +@pytest.mark.asyncio +async def test_an_over_budget_group_blocks_even_when_another_matched_group_is_fine(): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _enforce( + ("tier-a", "tier-b"), + _MagBudgetRow("tier-a", max_budget=10.0), + _MagBudgetRow("tier-b", max_budget=1.0), + spend_by_counter_key={ + MODEL_ACCESS_GROUP_COUNTER_KEY: 1.0, + model_access_group_spend_counter_key("tier-b"): 5.0, + }, + ) + + assert exc_info.value.entity_id == "tier-b" + + +@pytest.mark.asyncio +async def test_request_that_matched_no_group_touches_neither_database_nor_counters(): + assert await _enforce((), prisma_client=_ExplodingPrismaClient()) == [] + + +@pytest.mark.asyncio +async def test_budget_check_reads_the_counter_key_the_reset_job_clears(): + """Reads and resets must agree, or a rollover clears a counter nobody reads.""" + reset_job_key = _model_access_group_counter_key(SimpleNamespace(access_group_name="tier-a")) + + assert await _enforce(("tier-a",), _MagBudgetRow("tier-a", max_budget=10.0)) == [reset_job_key] + + +@pytest.mark.asyncio +async def test_a_second_request_serves_the_budget_row_from_cache(): + cache = UserApiKeyCache() + prisma_client = _RecordingPrismaClient(_MagBudgetRow("tier-a", max_budget=10.0)) + + await _enforce(("tier-a",), prisma_client=prisma_client, cache=cache) + await _enforce(("tier-a",), prisma_client=prisma_client, cache=cache) + + assert prisma_client.batches == [["tier-a"]] + + +@pytest.mark.asyncio +async def test_a_database_error_does_not_block_the_request(): + class _FailingPrismaClient: + def __init__(self) -> None: + self.db = SimpleNamespace( + litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom) + ) + + async def _boom(self, **kwargs): + raise RuntimeError("database unavailable") + + assert await _enforce(("tier-a",), prisma_client=_FailingPrismaClient()) == [] + + +async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) -> bool: + cache = await _cache() + prisma_client = _RecordingPrismaClient(_MagBudgetRow("tier-a", max_budget=1.0)) + read, _ = _spend_reader({MODEL_ACCESS_GROUP_COUNTER_KEY: 99.0}) + + with ( + # common_checks resolves all three off the proxy_server module at call time; its signature + # has no client, cache or spend-reader parameter to pass them through instead. + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter + patch("litellm.proxy.proxy_server.get_current_spend", read), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point + ): + return await common_checks( + request_body={"model": "gpt-4o", "messages": []}, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=Router(model_list=MODEL_LIST), + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + valid_token=UserAPIKeyAuth(api_key="hashed", models=["tier-a"], user_id=USER_ID), + request=SimpleNamespace(method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions")), + skip_budget_checks=skip_budget_checks, + ) + + +@pytest.mark.asyncio +async def test_common_checks_blocks_a_request_whose_group_is_over_budget(): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _common_checks_with_over_budget_group(skip_budget_checks=False) + + assert exc_info.value.entity_id == "tier-a" + + +@pytest.mark.asyncio +async def test_free_model_routes_skip_the_model_access_group_budget_check(): + """skip_budget_checks is how free models stay free; it has to cover this budget too.""" + assert await _common_checks_with_over_budget_group(skip_budget_checks=True) is True 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/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 5d3afd95a55..03b05bd9d87 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -77,6 +77,7 @@ class MockBatcher: self.litellm_teammembership = _Table("team_membership", self) self.litellm_organizationtable = _Table("org", self) self.litellm_tagtable = _Table("tag", self) + self.litellm_modelaccessgroupbudgettable = _Table("model_access_group", self) self.litellm_endusertable = _Table("enduser", self) async def commit(self): @@ -91,6 +92,7 @@ class MockDB: self.litellm_endusertable = MockTable() self.litellm_organizationtable = MockTable() self.litellm_tagtable = MockTable() + self.litellm_modelaccessgroupbudgettable = MockTable() self.batch_calls: List[Dict[str, Any]] = [] self.batchers: List[MockBatcher] = [] @@ -521,6 +523,7 @@ _LINKED_TABLE_CASES = [ ), ("org", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ("tag", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), + ("model_access_group", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ] @@ -830,6 +833,7 @@ def _make_reset_budget_windows_job( raise AssertionError(f"Unexpected query_raw call: {query}") prisma_client.db.query_raw = AsyncMock(side_effect=fake_query_raw) + prisma_client.db.execute_raw = AsyncMock(return_value=1) prisma_client.db.litellm_verificationtoken.update = AsyncMock(return_value=None) prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None) @@ -901,6 +905,145 @@ def test_reset_budget_windows_resets_expired_key_window(monkeypatch): spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-expired:window:1d", value=0.0) +def _window_spend_rolls(prisma_client): + return [ + call.args + for call in prisma_client.db.execute_raw.await_args_list + if "LiteLLM_BudgetWindowSpend" in call.args[0] + ] + + +def test_reset_budget_windows_rolls_the_key_window_spend_row(monkeypatch): + """The maintained per-window total has to start the new window at zero + alongside the counter, or enforcement keeps reading the old window's spend.""" + now = datetime.utcnow() + expired = (now - timedelta(minutes=5)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-expired", + "budget_limits": [{"budget_duration": "1d", "reset_at": expired}], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) + + asyncio.run(job.reset_budget_windows()) + + rolls = _window_spend_rolls(prisma_client) + assert len(rolls) == 1 + query, entity_type, entity_id, window_duration, new_window_start, _updated_at = rolls[0] + assert (entity_type, entity_id, window_duration) == ("key", "sk-expired", "1d") + assert "spend = 0" in " ".join(query.split()) + + # window_start is the start of the window that just began: new reset_at minus the duration. + written_windows = json.loads( + prisma_client.db.litellm_verificationtoken.update.await_args.kwargs["data"]["budget_limits"] + ) + new_reset_at = datetime.fromisoformat(written_windows[0]["reset_at"].replace("Z", "+00:00")).replace(tzinfo=None) + assert new_window_start == pytest.approx( + new_reset_at - timedelta(days=1), + abs=timedelta(seconds=1), + ) + + +def test_reset_budget_windows_roll_is_conditional_on_an_older_stored_window(monkeypatch): + """Another pod may already have rolled the row; clobbering it would drop + spend that landed under the new window.""" + now = datetime.utcnow() + expired = (now - timedelta(minutes=5)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-expired", + "budget_limits": [{"budget_duration": "1d", "reset_at": expired}], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) + + asyncio.run(job.reset_budget_windows()) + + query = " ".join(_window_spend_rolls(prisma_client)[0][0].split()) + assert "AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')" in query + + +def test_reset_budget_windows_rolls_the_team_window_spend_row(monkeypatch): + now = datetime.utcnow() + expired = (now - timedelta(minutes=1)).isoformat() + "Z" + + team_rows = [ + { + "team_id": "team-expired", + "budget_limits": [{"budget_duration": "30d", "reset_at": expired}], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=[], team_rows=team_rows) + + asyncio.run(job.reset_budget_windows()) + + rolls = _window_spend_rolls(prisma_client) + assert len(rolls) == 1 + assert rolls[0][1:4] == ("team", "team-expired", "30d") + + +def test_reset_budget_windows_does_not_roll_an_unexpired_window(monkeypatch): + now = datetime.utcnow() + future = (now + timedelta(hours=1)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-future", + "budget_limits": [{"budget_duration": "1d", "reset_at": future}], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) + + asyncio.run(job.reset_budget_windows()) + + assert _window_spend_rolls(prisma_client) == [] + + +def test_reset_budget_windows_rolls_only_the_expired_window_of_a_key(monkeypatch): + now = datetime.utcnow() + key_rows = [ + { + "token": "sk-mixed", + "budget_limits": [ + {"budget_duration": "1d", "reset_at": (now - timedelta(minutes=5)).isoformat() + "Z"}, + {"budget_duration": "30d", "reset_at": (now + timedelta(days=2)).isoformat() + "Z"}, + ], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) + + asyncio.run(job.reset_budget_windows()) + + rolls = _window_spend_rolls(prisma_client) + assert [roll[3] for roll in rolls] == ["1d"] + + +def test_reset_budget_windows_survives_a_failed_window_spend_roll(monkeypatch): + """The row is an optimization over aggregating LiteLLM_SpendLogs; a DB + failure there must not stop the counter reset from being persisted.""" + now = datetime.utcnow() + expired = (now - timedelta(minutes=5)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-expired", + "budget_limits": [{"budget_duration": "1d", "reset_at": expired}], + } + ] + job, prisma_client, spend_counter_cache = _make_reset_budget_windows_job( + monkeypatch, key_rows=key_rows, team_rows=[] + ) + prisma_client.db.execute_raw = AsyncMock(side_effect=Exception("connection reset")) + + asyncio.run(job.reset_budget_windows()) + + prisma_client.db.litellm_verificationtoken.update.assert_awaited_once() + spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-expired:window:1d", value=0.0) + + def test_reset_budget_windows_skips_unexpired_key_window(monkeypatch): """If `reset_at` is in the future, no write should happen for that key.""" now = datetime.utcnow() @@ -1299,13 +1442,19 @@ _INVALIDATION_CASES = [ "spend:tag:tenant-42", {"tag:tenant-42"}, ), + ( + "litellm_modelaccessgroupbudgettable", + type("AccessGroup", (), {"access_group_name": "gpt-4-group"}), + "spend:model_access_group:gpt-4-group", + {"model_access_group:gpt-4-group"}, + ), ] @pytest.mark.parametrize( "table_attr, linked_row, counter_key, cache_keys", _INVALIDATION_CASES, - ids=["team_membership", "key", "org", "tag"], + ids=["team_membership", "key", "org", "tag", "model_access_group"], ) def test_budget_table_reset_invalidates_counters_and_management_cache( reset_budget_job, mock_prisma_client, monkeypatch, table_attr, linked_row, counter_key, cache_keys @@ -1359,6 +1508,102 @@ def test_budget_table_reset_commits_even_when_cache_eviction_fails(reset_budget_ assert mock_prisma_client.db.batchers[0].committed is True +# --------------------------------------------------------------------------- +# Model access group budgets ride the same cascade +# --------------------------------------------------------------------------- + + +def _model_access_group_row(name: str = "gpt-4-group", spend: float = 12.0, budget_id: str = "budget-1"): + """A LiteLLM_ModelAccessGroupBudgetTable row, shaped like prisma hands it back.""" + return type("AccessGroup", (), {"access_group_name": name, "spend": spend, "budget_id": budget_id}) + + +def test_access_group_reset_only_matches_rows_that_have_spend(reset_budget_job, mock_prisma_client, monkeypatch): + """Both the read and the write are filtered to spend > 0 on the due tiers. + + A group sitting at spend 0 has nothing to reset, and a group hanging off a + tier that is not due yet must not be swept along: both are excluded by the + filter, not by anything downstream. + """ + _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-due", budget_duration="7d")] + mock_prisma_client.db.litellm_modelaccessgroupbudgettable.set_find_many_results( + [_model_access_group_row(budget_id="budget-due")] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + expected_where = {"budget_id": {"in": ["budget-due"]}, "spend": {"gt": 0}} + assert mock_prisma_client.db.litellm_modelaccessgroupbudgettable.find_many_calls == [{"where": expected_where}] + writes = _batch_writes(mock_prisma_client, "model_access_group", op="update_many") + assert len(writes) == 1 + assert writes[0]["where"] == expected_where + assert writes[0]["data"] == {"spend": 0} + + +def test_access_groups_are_untouched_when_no_budget_is_due(reset_budget_job, mock_prisma_client, monkeypatch): + """No due tier means the group table is never read, written or evicted.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.db.litellm_modelaccessgroupbudgettable.set_find_many_results([_model_access_group_row()]) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + assert mock_prisma_client.db.litellm_modelaccessgroupbudgettable.find_many_calls == [] + assert _batch_writes(mock_prisma_client, "model_access_group") == [] + counter_cache.in_memory_cache.set_cache.assert_not_called() + counter_cache.user_api_key_cache.async_delete_cache.assert_not_awaited() + + +def test_budget_table_reset_invalidates_every_access_group_not_just_the_first( + reset_budget_job, mock_prisma_client, monkeypatch +): + """When several groups share the expiring tier, all of them are evicted.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")] + mock_prisma_client.db.litellm_modelaccessgroupbudgettable.set_find_many_results( + [_model_access_group_row(name=name) for name in ("group-a", "group-b", "group-c")] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} + assert deleted == {"model_access_group:group-a", "model_access_group:group-b", "model_access_group:group-c"} + for name in ("group-a", "group-b", "group-c"): + counter_cache.in_memory_cache.set_cache.assert_any_call(key=f"spend:model_access_group:{name}", value=0.0, ttl=60) + + +def test_budget_cascade_carries_access_group_overage_when_rollover_enabled( + rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch +): + """A group 5 over the tier cap keeps a spend of 5 in the next window, the + same way a tag or a team member does: over-cap rows are decremented by the + cap, the rest are zeroed, and the counter is seeded with the carried spend.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-roll", budget_duration="7d", max_budget=10.0)] + mock_prisma_client.db.litellm_modelaccessgroupbudgettable.set_find_many_results( + [_model_access_group_row(spend=15.0, budget_id="budget-roll")] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + writes = _batch_writes(mock_prisma_client, "model_access_group") + assert { + "table": "model_access_group", + "op": "update_many", + "where": {"budget_id": "budget-roll", "spend": {"gt": 10.0}}, + "data": {"spend": {"decrement": 10.0}}, + } in writes + assert { + "table": "model_access_group", + "op": "update_many", + "where": {"budget_id": "budget-roll", "spend": {"gt": 0, "lte": 10.0}}, + "data": {"spend": 0}, + } in writes + assert _replay_spend_writes(writes, 15.0) == 5.0 + assert _replay_spend_writes(writes, 8.0) == 0 + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:model_access_group:gpt-4-group", value=5.0, ttl=60) + + # --------------------------------------------------------------------------- # Atomicity of the budget-table cascade (LIT-5138) # --------------------------------------------------------------------------- @@ -1471,6 +1716,7 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo ("key", "update_many"), ("org", "update_many"), ("tag", "update_many"), + ("model_access_group", "update_many"), ("enduser", "update_many"), ("budget", "update_many"), } @@ -1506,7 +1752,7 @@ def test_failed_cascade_is_logged_as_a_cascade_failure(monkeypatch): assert mock_exception.call_count == 1 message = mock_exception.call_args.args[0] assert "cascade" in message - for mentioned in ("team member", "enduser", "org", "tag", "budget_reset_at"): + for mentioned in ("team member", "enduser", "org", "tag", "model access group", "budget_reset_at"): assert mentioned in message, f"failure log should mention {mentioned}: {message}" diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index fb0c994a476..504654e103a 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -1,4 +1,5 @@ import json +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -22,9 +23,7 @@ def redis_update_buffer(mock_redis_cache): @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_uses_pipeline( - redis_update_buffer, mock_redis_cache -): +async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache): """ Verify store_in_memory_spend_updates_in_redis calls async_rpush_pipeline once with the correct operations and skips empty queues. @@ -33,35 +32,29 @@ async def test_store_in_memory_spend_updates_uses_pipeline( # Create mock queues - only 3 of 6 have data spend_update_queue = AsyncMock() - spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = ( - AsyncMock(return_value={"key_list_transactions": {"key1": 1.0}}) + spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock( + return_value={"key_list_transactions": {"key1": 1.0}} ) daily_spend_queue = AsyncMock() - daily_spend_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={"user_key1": {"spend": 1.0}}) + daily_spend_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={"user_key1": {"spend": 1.0}} ) daily_team_queue = AsyncMock() - daily_team_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={"team_key1": {"spend": 2.0}}) + daily_team_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={"team_key1": {"spend": 2.0}} ) # Empty queues daily_org_queue = AsyncMock() - daily_org_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={}) - ) + daily_org_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) daily_end_user_queue = AsyncMock() - daily_end_user_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value=None) - ) + daily_end_user_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value=None) daily_agent_queue = AsyncMock() - daily_agent_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={}) - ) + daily_agent_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) await redis_update_buffer.store_in_memory_spend_updates_in_redis( spend_update_queue=spend_update_queue, @@ -82,9 +75,7 @@ async def test_store_in_memory_spend_updates_uses_pipeline( @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_restores_on_rpush_failure( - redis_update_buffer, mock_redis_cache -): +async def test_store_in_memory_spend_updates_restores_on_rpush_failure(redis_update_buffer, mock_redis_cache): """ If async_rpush_pipeline raises, the already-drained transactions must be put back into the in-memory queues so the next scheduler tick retries. @@ -98,9 +89,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure( SpendUpdateQueue, ) - mock_redis_cache.async_rpush_pipeline = AsyncMock( - side_effect=ConnectionError("redis went away") - ) + mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away")) spend_queue = SpendUpdateQueue() daily_user_queue = DailySpendUpdateQueue() @@ -145,16 +134,12 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure( # After restore, the main spend queue should hold one item per # (entity_type, entity_id) pair with the aggregated cost - restored_spend = ( - await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() - ) + restored_spend = await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() assert restored_spend["key_list_transactions"] == {"key-abc": 1.5} assert restored_spend["team_list_transactions"] == {"team-xyz": 2.5} # Daily user queue should hold the same aggregated dict - restored_daily = ( - await daily_user_queue.flush_and_get_aggregated_daily_spend_update_transactions() - ) + restored_daily = await daily_user_queue.flush_and_get_aggregated_daily_spend_update_transactions() assert restored_daily == { "user1_day_model": { "spend": 1.0, @@ -165,9 +150,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure( @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_all_empty_returns_early( - redis_update_buffer, mock_redis_cache -): +async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_update_buffer, mock_redis_cache): """ When all queues are empty, pipeline should never be called. """ @@ -175,13 +158,9 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early( # All queues return empty empty_queue = AsyncMock() - empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock( - return_value={} - ) + empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={}) empty_daily_queue = AsyncMock() - empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={}) - ) + empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) await redis_update_buffer.store_in_memory_spend_updates_in_redis( spend_update_queue=empty_queue, @@ -196,14 +175,13 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early( @pytest.mark.asyncio -async def test_get_all_transactions_from_redis_buffer_pipeline( - redis_update_buffer, mock_redis_cache -): +async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buffer, mock_redis_cache): """ Verify get_all_transactions_from_redis_buffer_pipeline correctly parses and aggregates results from async_lpop_pipeline. """ - # Simulate pipeline results: slot 0 = spend updates, slots 1-5 = daily categories + # Simulate pipeline results: slot 0 = spend updates, slots 1-5 = daily categories, + # slot 6 = budget window spend db_spend_json = json.dumps( { "key_list_transactions": {"key1": 1.0, "key2": 2.0}, @@ -217,6 +195,18 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( ) daily_user_json = json.dumps({"user_key1": {"spend": 1.0, "api_requests": 1}}) daily_team_json = json.dumps({"team_key1": {"spend": 2.0, "api_requests": 2}}) + window_spend_json = json.dumps( + [ + { + "entity_type": "key", + "entity_id": "hashed-token", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 3.0, + "request_ids": ["req-1"], + } + ] + ) mock_redis_cache.async_lpop_pipeline = AsyncMock( return_value=[ @@ -226,13 +216,30 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( None, # slot 3: daily org (empty) None, # slot 4: daily end-user (empty) None, # slot 5: daily agent (empty) + [window_spend_json, window_spend_json], # slot 6: budget window spend ] ) result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() - assert len(result) == 6 - db_spend, daily_user, daily_team, daily_org, daily_end_user, daily_agent = result + assert len(result) == 7 + ( + db_spend, + daily_user, + daily_team, + daily_org, + daily_end_user, + daily_agent, + window_spend, + ) = result + + # Budget window spend from two pods is summed per window, not overwritten, + # and both pods' request ids reach the seed exclusion. + assert window_spend is not None + assert len(window_spend) == 1 + assert window_spend[0]["spend"] == 6.0 + assert window_spend[0]["entity_id"] == "hashed-token" + assert window_spend[0]["request_ids"] == ("req-1",) # Verify db spend was parsed correctly assert db_spend is not None @@ -255,6 +262,10 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( # Verify pipeline was called once with correct keys mock_redis_cache.async_lpop_pipeline.assert_called_once() + from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY + + popped_keys = [op["key"] for op in mock_redis_cache.async_lpop_pipeline.call_args.kwargs["lpop_list"]] + assert popped_keys[6] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY @pytest.mark.asyncio @@ -262,13 +273,11 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis(): """When redis_cache is None, should return all Nones""" buffer = RedisUpdateBuffer(redis_cache=None) result = await buffer.get_all_transactions_from_redis_buffer_pipeline() - assert result == (None, None, None, None, None, None) + assert result == (None, None, None, None, None, None, None) @pytest.mark.asyncio -async def test_restore_transactions_to_redis_pushes_only_provided( - redis_update_buffer, mock_redis_cache -): +async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_buffer, mock_redis_cache): """ restore_transactions_to_redis re-pushes only the transaction sets it was given, to their matching buffer keys, so uncommitted spend can be retried. @@ -302,9 +311,42 @@ async def test_restore_transactions_to_redis_pushes_only_provided( @pytest.mark.asyncio -async def test_restore_transactions_to_redis_noop_when_empty( - redis_update_buffer, mock_redis_cache -): +async def test_restored_window_spend_transactions_drain_back_unchanged(redis_update_buffer, mock_redis_cache): + """A window commit that fails after the destructive lpop must be re-pushed + in the store path's encoding, so the next drain returns the same increments.""" + from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + build_window_spend_transaction, + ) + + window_transactions = ( + build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=3.0, + request_id="req-1", + started_at=datetime(2026, 8, 10, 12, 0, tzinfo=timezone.utc), + ), + ) + mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1]) + + await redis_update_buffer.restore_transactions_to_redis(window_spend_update_transactions=window_transactions) + + rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] + assert [op["key"] for op in rpush_list] == [REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY] + + mock_redis_cache.async_lpop_pipeline = AsyncMock( + return_value=[None, None, None, None, None, None, list(rpush_list[0]["values"])] + ) + drained = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() + + assert drained[6] == window_transactions + + +@pytest.mark.asyncio +async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer, mock_redis_cache): """Nothing to restore -> no Redis call.""" mock_redis_cache.async_rpush_pipeline = AsyncMock() await redis_update_buffer.restore_transactions_to_redis() @@ -312,15 +354,11 @@ async def test_restore_transactions_to_redis_noop_when_empty( @pytest.mark.asyncio -async def test_restore_transactions_to_redis_swallows_redis_error( - redis_update_buffer, mock_redis_cache -): +async def test_restore_transactions_to_redis_swallows_redis_error(redis_update_buffer, mock_redis_cache): """A Redis failure during restore must not propagate to the caller's finally block.""" from redis.exceptions import RedisError - mock_redis_cache.async_rpush_pipeline = AsyncMock( - side_effect=RedisError("redis down") - ) + mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=RedisError("redis down")) await redis_update_buffer.restore_transactions_to_redis( db_spend_update_transactions={"key_list_transactions": {"key1": 1.0}}, @@ -433,3 +471,108 @@ def test_get_transaction_buffer_redis_cache_parses_string_flag(monkeypatch): mock_redis_cache.assert_called_once() assert result is mock_redis_cache.return_value + + +@pytest.mark.asyncio +async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_update_buffer, mock_redis_cache): + """The budget window queue has to ride the same rpush as the daily queues, + otherwise multi-pod deployments never persist per-window spend.""" + from datetime import datetime, timezone + + from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendUpdateQueue, + build_window_spend_transaction, + ) + + mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1]) + + empty_queue = AsyncMock() + empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={}) + empty_daily_queue = AsyncMock() + empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) + + window_queue = WindowSpendUpdateQueue() + await window_queue.add_update( + build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=1.25, + request_id="req-1", + started_at=datetime(2026, 8, 10, 12, 0, 0, tzinfo=timezone.utc), + ) + ) + + await redis_update_buffer.store_in_memory_spend_updates_in_redis( + spend_update_queue=empty_queue, + daily_spend_update_queue=empty_daily_queue, + daily_team_spend_update_queue=empty_daily_queue, + daily_org_spend_update_queue=empty_daily_queue, + daily_end_user_spend_update_queue=empty_daily_queue, + daily_agent_spend_update_queue=empty_daily_queue, + window_spend_update_queue=window_queue, + ) + + rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] + assert len(rpush_list) == 1 + assert rpush_list[0]["key"] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY + pushed = json.loads(rpush_list[0]["values"][0]) + assert pushed == [ + { + "entity_type": "key", + "entity_id": "hashed-token", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 1.25, + "request_ids": ["req-1"], + "started_at": "2026-08-10T12:00:00.000000", + } + ] + + +@pytest.mark.asyncio +async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpush_failure( + redis_update_buffer, mock_redis_cache +): + """The window queue is drained before the rpush, so a Redis hiccup would + silently drop per-window spend without the restore.""" + from datetime import datetime, timezone + + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendUpdateQueue, + build_window_spend_transaction, + ) + + mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away")) + + empty_queue = AsyncMock() + empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={}) + empty_daily_queue = AsyncMock() + empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) + + window_queue = WindowSpendUpdateQueue() + await window_queue.add_update( + build_window_spend_transaction( + entity_type="team", + entity_id="team-1", + window_duration="7d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=4.0, + ) + ) + + await redis_update_buffer.store_in_memory_spend_updates_in_redis( + spend_update_queue=empty_queue, + daily_spend_update_queue=empty_daily_queue, + daily_team_spend_update_queue=empty_daily_queue, + daily_org_spend_update_queue=empty_daily_queue, + daily_end_user_spend_update_queue=empty_daily_queue, + daily_agent_spend_update_queue=empty_daily_queue, + window_spend_update_queue=window_queue, + ) + + restored = await window_queue.flush_and_get_aggregated_window_spend_transactions() + assert [payload["spend"] for payload in restored] == [4.0] + assert [payload["entity_id"] for payload in restored] == ["team-1"] diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py new file mode 100644 index 00000000000..b1ecda57afa --- /dev/null +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py @@ -0,0 +1,269 @@ +import json +from datetime import datetime, timedelta, timezone + +import pytest + +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendUpdateQueue, + build_window_spend_transaction, + to_naive_utc, +) + +WINDOW_A = datetime(2026, 8, 1, tzinfo=timezone.utc) +WINDOW_B = datetime(2026, 8, 31, tzinfo=timezone.utc) + + +def _txn( + entity_id: str, + window_start: datetime, + spend: float, + duration: str = "30d", + entity_type: str = "key", + request_id: str | None = None, + started_at: datetime | None = None, +): + return build_window_spend_transaction( + entity_type=entity_type, + entity_id=entity_id, + window_duration=duration, + window_start=window_start, + spend=spend, + request_id=request_id, + started_at=started_at, + ) + + +def test_build_window_spend_transaction_stores_naive_utc_iso(): + """window_start rides the Redis buffer as a string and lands in a naive-UTC + TIMESTAMP(3) column, so a non-UTC input must be converted, not truncated.""" + non_utc = datetime(2026, 8, 1, 20, 0, tzinfo=timezone(timedelta(hours=-4))) + + assert _txn("k1", non_utc, 1.0, request_id="req-1") == { + "entity_type": "key", + "entity_id": "k1", + "window_duration": "30d", + "window_start": "2026-08-02T00:00:00.000000", + "spend": 1.0, + "request_ids": ("req-1",), + "started_at": None, + } + + +def test_build_window_spend_transaction_stores_started_at_as_naive_utc_iso(): + """started_at is compared against LiteLLM_SpendLogs.startTime, which the + spend log writer stores after converting the request start to UTC.""" + non_utc = datetime(2026, 8, 10, 8, 30, 15, 123456, tzinfo=timezone(timedelta(hours=-4))) + + assert _txn("k1", WINDOW_A, 1.0, started_at=non_utc)["started_at"] == "2026-08-10T12:30:15.123456" + + +@pytest.mark.asyncio +async def test_aggregation_keeps_the_earliest_started_at_of_the_batch(): + """The seed bounds its request-id exclusion at the batch's earliest start, + so a later start must never win the merge.""" + queue = WindowSpendUpdateQueue() + earliest = datetime(2026, 8, 10, 12, 0, 0, tzinfo=timezone.utc) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-2", started_at=earliest + timedelta(seconds=5))) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-1", started_at=earliest)) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-3")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 1 + assert aggregated[0]["started_at"] == "2026-08-10T12:00:00.000000" + assert aggregated[0]["request_ids"] == ("req-1", "req-2", "req-3") + + +def test_to_naive_utc_leaves_naive_values_alone(): + naive = datetime(2026, 8, 1, 12, 0) + assert to_naive_utc(naive) == naive + + +@pytest.mark.asyncio +async def test_aggregation_sums_increments_within_one_window(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.5)) + await queue.add_update(_txn("k1", WINDOW_A, 2.25)) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 1 + assert aggregated[0]["spend"] == pytest.approx(3.75) + + +@pytest.mark.asyncio +async def test_aggregation_keeps_different_windows_of_same_entity_separate(): + """Merging across windows would fold spend from a window that already + rolled into the new window's total, over-counting the new window.""" + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0)) + await queue.add_update(_txn("k1", WINDOW_B, 2.0)) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 2 + assert {payload["window_start"]: payload["spend"] for payload in aggregated} == { + "2026-08-01T00:00:00.000000": 1.0, + "2026-08-31T00:00:00.000000": 2.0, + } + + +@pytest.mark.asyncio +async def test_aggregation_keeps_durations_entities_and_types_separate(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0, duration="30d")) + await queue.add_update(_txn("k1", WINDOW_A, 2.0, duration="7d")) + await queue.add_update(_txn("k2", WINDOW_A, 4.0, duration="30d")) + await queue.add_update(_txn("k1", WINDOW_A, 8.0, duration="30d", entity_type="team")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 4 + assert sorted(payload["spend"] for payload in aggregated) == [1.0, 2.0, 4.0, 8.0] + + +@pytest.mark.asyncio +async def test_aggregation_orders_by_primary_key_then_window_start(): + """The flush relies on this order: primary key first for cross-pod lock + ordering, then window_start so an older window is applied before the roll + that supersedes it.""" + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("t1", WINDOW_A, 1.0, entity_type="team")) + await queue.add_update(_txn("k2", WINDOW_B, 1.0)) + await queue.add_update(_txn("k2", WINDOW_A, 1.0)) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, duration="7d")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert [ + (payload["entity_type"], payload["entity_id"], payload["window_duration"], payload["window_start"]) + for payload in aggregated + ] == [ + ("key", "k1", "7d", "2026-08-01T00:00:00.000000"), + ("key", "k2", "30d", "2026-08-01T00:00:00.000000"), + ("key", "k2", "30d", "2026-08-31T00:00:00.000000"), + ("team", "t1", "30d", "2026-08-01T00:00:00.000000"), + ] + + +@pytest.mark.asyncio +async def test_aggregation_does_not_collide_on_entity_ids_containing_a_separator(): + """entity_id is free-form (team ids are user supplied), so grouping must not + depend on a flattened string key.""" + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("a:30d:2026-08-01T00:00:00.000000:b", WINDOW_A, 1.0)) + await queue.add_update(_txn("b", WINDOW_A, 2.0)) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 2 + + +@pytest.mark.asyncio +async def test_flush_empties_the_queue(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0)) + + assert await queue.flush_and_get_aggregated_window_spend_transactions() != () + assert await queue.flush_and_get_aggregated_window_spend_transactions() == () + + +@pytest.mark.asyncio +async def test_aggregate_queue_updates_collapses_in_place(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0)) + await queue.add_update(_txn("k1", WINDOW_A, 2.0)) + await queue.add_update(_txn("k1", WINDOW_B, 4.0)) + + await queue.aggregate_queue_updates() + + assert queue.update_queue.qsize() == 1 + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + assert sorted(payload["spend"] for payload in aggregated) == [3.0, 4.0] + + +@pytest.mark.asyncio +async def test_aggregation_does_not_mutate_the_queued_payloads(): + """The same payload can be re-aggregated after a failed Redis push, so + aggregation must not accumulate into the caller's object.""" + queue = WindowSpendUpdateQueue() + update = _txn("k1", WINDOW_A, 1.0) + await queue.add_update(update) + await queue.add_update(_txn("k1", WINDOW_A, 2.0)) + + await queue.flush_and_get_aggregated_window_spend_transactions() + + assert update["spend"] == 1.0 + + +def test_aggregation_survives_the_redis_json_round_trip(): + """The Redis buffer stores transactions as JSON, so the aggregated shape + must reload into an equivalent aggregation.""" + aggregated = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( + [(_txn("k1", WINDOW_A, 1.0),), (_txn("k1", WINDOW_B, 2.0),)] + ) + + reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions([json.loads(json.dumps(aggregated))]) + + assert reloaded == aggregated + + +@pytest.mark.asyncio +async def test_aggregation_unions_the_request_ids_of_merged_increments(): + """The seed excludes exactly the requests its batch already covers, so every + merged increment's id has to survive aggregation.""" + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-1")) + await queue.add_update(_txn("k1", WINDOW_A, 2.0, request_id="req-2")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 1 + assert aggregated[0]["request_ids"] == ("req-1", "req-2") + + +@pytest.mark.asyncio +async def test_request_ids_stay_with_their_own_window(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) + await queue.add_update(_txn("k1", WINDOW_B, 2.0, request_id="req-b")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert {payload["window_start"]: payload["request_ids"] for payload in aggregated} == { + "2026-08-01T00:00:00.000000": ("req-a",), + "2026-08-31T00:00:00.000000": ("req-b",), + } + + +@pytest.mark.asyncio +async def test_request_ids_are_deduplicated_and_ordered(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-b")) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert aggregated[0]["request_ids"] == ("req-a", "req-b") + + +@pytest.mark.asyncio +async def test_increment_without_a_request_id_carries_no_exclusion(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0)) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert aggregated[0]["request_ids"] == () + + +def test_request_ids_survive_the_redis_json_round_trip(): + aggregated = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( + [(_txn("k1", WINDOW_A, 1.0, request_id="req-1"),)] + ) + + reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions([json.loads(json.dumps(aggregated))]) + + assert reloaded[0]["request_ids"] == ("req-1",) + assert reloaded[0]["spend"] == 1.0 diff --git a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py index cb4687ef370..2ed4f843711 100644 --- a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py +++ b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py @@ -106,6 +106,28 @@ class TestBuildTransaction: transaction = _build() assert transaction is not None and transaction.tier is None + def test_a_priced_classifier_rides_the_turns_spend(self): + """The classifier row is excluded from the rollup, so its charge lands here, + folded once into the turn that paid for it (GH #38816).""" + transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "classifier_cost": 0.005})) + assert transaction is not None and transaction.spend == pytest.approx(0.015) + + @pytest.mark.parametrize( + "decision_extra", [{}, {"classifier_cost": 0.0}, {"classifier_cost": "bogus"}, {"classifier_cost": True}] + ) + def test_an_unpriced_classifier_leaves_the_spend_alone(self, decision_extra: dict): + transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, **decision_extra})) + assert transaction is not None and transaction.spend == pytest.approx(0.01) + + def test_every_turn_carries_its_own_classifier_charge(self): + first = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "classifier_cost": 0.005})) + second = _build( + payload=_payload(startTime="2026-08-01T12:01:00", spend=0.02), + metadata=_metadata(routing_decision={**ROUTING_DECISION, "classifier_cost": 0.007}), + ) + assert first is not None and first.spend == pytest.approx(0.015) + assert second is not None and second.spend == pytest.approx(0.027) + def test_router_name_falls_back_to_the_payload_model_group(self): transaction = _build(metadata=_metadata(routing_decision={"router_type": "complexity"})) assert transaction is not None and transaction.router_name == "live-auto" diff --git a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py new file mode 100644 index 00000000000..a849317c930 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py @@ -0,0 +1,596 @@ +import math +from contextlib import asynccontextmanager +from datetime import datetime, timedelta, timezone +from typing import Any + +import pytest + +from litellm.proxy.db.budget_window_spend_writer import ( + commit_window_spend_updates, + roll_window_spend_row, + spend_logs_total_excluding, +) +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + build_window_spend_transaction, +) + +WINDOW_A = datetime(2026, 8, 1, tzinfo=timezone.utc) +WINDOW_B = datetime(2026, 8, 31, tzinfo=timezone.utc) +BATCH_STARTED_AT = datetime(2026, 8, 10, 12, 0, 0, 250_000, tzinfo=timezone.utc) +BEFORE_BATCH = BATCH_STARTED_AT - timedelta(hours=1) + +ENTITY_TYPE, ENTITY_ID, WINDOW_DURATION, WINDOW_START, INSERT_SPEND, INCREMENT, NOW = range(7) + + +class _FakeBatcher: + def __init__(self) -> None: + self.calls: list[tuple[str, tuple[Any, ...]]] = [] + + def execute_raw(self, query: str, *args: Any) -> None: + self.calls.append((query, args)) + + +class _FakeDB: + """Stands in for prisma_client.db; records every statement it is handed.""" + + def __init__(self, existing_rows: list[dict[str, str]] | None = None) -> None: + self.existing_rows = existing_rows or [] + self.query_raw_calls: list[tuple[str, tuple[Any, ...]]] = [] + self.execute_raw_calls: list[tuple[str, tuple[Any, ...]]] = [] + self.batcher = _FakeBatcher() + self.committed = False + + async def query_raw(self, query: str, *args: Any) -> list[dict[str, str]]: + self.query_raw_calls.append((query, args)) + return self.existing_rows + + async def execute_raw(self, query: str, *args: Any) -> int: + self.execute_raw_calls.append((query, args)) + return 1 + + @asynccontextmanager + async def _tx(self): + yield self + + def tx(self, timeout: Any = None): + return self._tx() + + @asynccontextmanager + async def _batch(self): + yield self.batcher + self.committed = True + + def batch_(self): + return self._batch() + + +class _FakePrismaClient: + def __init__(self, db: _FakeDB) -> None: + self.db = db + + +class _RecordingAggregate: + """Stands in for the LiteLLM_SpendLogs seed aggregate.""" + + def __init__(self, value: float = 5.0) -> None: + self.value = value + self.calls: list[dict[str, Any]] = [] + + async def __call__( + self, + prisma_client: Any, + entity_type: str, + entity_id: str, + window_start: datetime, + exclude_request_ids: Any, + exclude_started_at: datetime | None, + ) -> float | None: + self.calls.append( + { + "entity_type": entity_type, + "entity_id": entity_id, + "window_start": window_start, + "exclude_request_ids": tuple(exclude_request_ids), + "exclude_started_at": exclude_started_at, + } + ) + return self.value + + +class _SpendLogsFake: + """Sums the LiteLLM_SpendLogs rows (request_id, spend, startTime) it holds, + honouring the exclusion exactly as the real aggregate's + NOT (request_id = ANY(...) AND startTime >= bound) does.""" + + def __init__(self, rows: tuple[tuple[str, float, datetime], ...]) -> None: + self.rows = rows + + async def __call__( + self, + prisma_client: Any, + entity_type: str, + entity_id: str, + window_start: datetime, + exclude_request_ids: Any, + exclude_started_at: datetime | None, + ) -> float | None: + excluded = frozenset(exclude_request_ids) if exclude_started_at is not None else frozenset() + return math.fsum( + spend + for request_id, spend, started_at in self.rows + if not (request_id in excluded and started_at >= exclude_started_at) + ) + + +def _batch(request_ids: tuple[str, ...], spend: float, started_at: datetime | None = BATCH_STARTED_AT) -> dict: + return { + "entity_type": "key", + "entity_id": "k1", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": spend, + "request_ids": request_ids, + "started_at": None + if started_at is None + else started_at.replace(tzinfo=None).isoformat(timespec="microseconds"), + } + + +def _existing(entity_type: str, entity_id: str, window_duration: str) -> dict[str, str]: + return {"entity_type": entity_type, "entity_id": entity_id, "window_duration": window_duration} + + +@pytest.mark.asyncio +async def test_no_transactions_touches_no_database(): + db = _FakeDB() + + await commit_window_spend_updates(prisma_client=_FakePrismaClient(db), transactions=()) + + assert db.query_raw_calls == [] + assert db.batcher.calls == [] + + +@pytest.mark.asyncio +async def test_missing_row_is_seeded_from_spend_logs_once(): + """A row created mid-window would undercount everything spent before it + existed, so a brand new primary key inserts the SpendLogs total plus this + increment.""" + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=5.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=aggregate, + ) + + assert len(aggregate.calls) == 1 + assert aggregate.calls[0]["entity_type"] == "key" + assert aggregate.calls[0]["entity_id"] == "k1" + assert aggregate.calls[0]["window_start"] == WINDOW_A + + ((_, params),) = db.batcher.calls + assert params[ENTITY_TYPE] == "key" + assert params[ENTITY_ID] == "k1" + assert params[WINDOW_DURATION] == "30d" + assert params[WINDOW_START] == datetime(2026, 8, 1) + assert params[INSERT_SPEND] == pytest.approx(6.0) + assert params[INCREMENT] == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_existing_row_is_never_reseeded(): + """The seed is a full LiteLLM_SpendLogs scan; running it for a row that is + already maintained would both cost a scan and double count.""" + db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")]) + aggregate = _RecordingAggregate(value=5.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=aggregate, + ) + + assert aggregate.calls == [] + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(1.0) + assert params[INCREMENT] == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_seed_runs_only_for_the_primary_keys_that_are_missing(): + db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")]) + aggregate = _RecordingAggregate(value=5.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("team", "t1", "30d", WINDOW_A, 2.0), + ), + spend_logs_aggregate=aggregate, + ) + + assert [call["entity_id"] for call in aggregate.calls] == ["t1"] + assert [call["entity_type"] for call in aggregate.calls] == ["team"] + by_entity = {params[ENTITY_ID]: params for _, params in db.batcher.calls} + assert by_entity["k1"][INSERT_SPEND] == pytest.approx(1.0) + assert by_entity["t1"][INSERT_SPEND] == pytest.approx(7.0) + + +@pytest.mark.asyncio +async def test_insert_spend_and_increment_differ_only_when_a_row_is_seeded(): + """The conflict arm adds the increment alone so two pods that both seed the + same new window cannot add the SpendLogs base twice.""" + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=9.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 0.25),), + spend_logs_aggregate=aggregate, + ) + + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(9.25) + assert params[INCREMENT] == pytest.approx(0.25) + + +@pytest.mark.asyncio +async def test_upsert_sql_adds_for_a_current_window_and_replaces_for_a_newer_one(): + """The CASE is the whole contract: an increment at or behind the stored + window_start accumulates, a newer one restarts the window.""" + db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")]) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),), + ) + + ((query, _),) = db.batcher.calls + normalized = " ".join(query.split()) + assert ( + 'spend = CASE WHEN "LiteLLM_BudgetWindowSpend".window_start >= EXCLUDED.window_start ' + 'THEN "LiteLLM_BudgetWindowSpend".spend + $6 ELSE EXCLUDED.spend END' in normalized + ) + assert 'window_start = GREATEST("LiteLLM_BudgetWindowSpend".window_start, EXCLUDED.window_start)' in normalized + assert "ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE SET" in normalized + + +@pytest.mark.asyncio +async def test_upsert_never_interpolates_values_into_the_sql(): + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "'; DROP TABLE x; --", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=aggregate, + ) + + ((query, params),) = db.batcher.calls + assert "DROP TABLE" not in query + assert params[ENTITY_ID] == "'; DROP TABLE x; --" + + +@pytest.mark.asyncio +async def test_upserts_are_ordered_by_primary_key_then_window_start(): + """Cross-pod lock ordering, plus an older window must be applied before the + roll that supersedes it or the roll would be undone.""" + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + build_window_spend_transaction("team", "t1", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("key", "k2", "30d", WINDOW_B, 1.0), + build_window_spend_transaction("key", "k2", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("key", "k1", "7d", WINDOW_A, 1.0), + ), + spend_logs_aggregate=aggregate, + ) + + ordered = [ + (params[ENTITY_TYPE], params[ENTITY_ID], params[WINDOW_DURATION], params[WINDOW_START]) + for _, params in db.batcher.calls + ] + assert ordered == [ + ("key", "k1", "7d", datetime(2026, 8, 1)), + ("key", "k2", "30d", datetime(2026, 8, 1)), + ("key", "k2", "30d", datetime(2026, 8, 31)), + ("team", "t1", "30d", datetime(2026, 8, 1)), + ] + + +@pytest.mark.asyncio +async def test_existing_row_lookup_sends_every_primary_key_as_array_params(): + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("team", "t1", "7d", WINDOW_A, 1.0), + ), + spend_logs_aggregate=aggregate, + ) + + ((query, params),) = db.query_raw_calls + assert "unnest($1::text[], $2::text[], $3::text[])" in query + assert params == (("key", "team"), ("k1", "t1"), ("30d", "7d")) + + +@pytest.mark.asyncio +async def test_all_upserts_are_committed_in_one_transaction(): + db = _FakeDB(existing_rows=[_existing("key", "k1", "30d"), _existing("key", "k2", "30d")]) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("key", "k2", "30d", WINDOW_A, 2.0), + ), + ) + + assert len(db.batcher.calls) == 2 + assert db.committed is True + + +@pytest.mark.asyncio +async def test_unknown_entity_type_contributes_no_seed(): + """Only key and team windows have a LiteLLM_SpendLogs column to aggregate; + anything else starts from its increment alone.""" + db = _FakeDB(existing_rows=[]) + + async def no_such_column( + prisma_client, entity_type, entity_id, window_start, exclude_request_ids, exclude_started_at + ): + return None + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("user", "u1", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=no_such_column, + ) + + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_unavailable_spend_logs_aggregate_seeds_zero_rather_than_failing(): + db = _FakeDB(existing_rows=[]) + + async def unavailable(prisma_client, entity_type, entity_id, window_start, exclude_request_ids, exclude_started_at): + return None + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=unavailable, + ) + + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_older(): + """Unconditional zeroing would wipe increments a pod already applied under + the new window.""" + db = _FakeDB() + + await roll_window_spend_row( + prisma_client=_FakePrismaClient(db), + entity_type="team", + entity_id="t1", + window_duration="30d", + new_window_start=WINDOW_B, + ) + + ((query, params),) = db.execute_raw_calls + normalized = " ".join(query.split()) + assert "SET window_start = ($4::timestamptz AT TIME ZONE 'UTC'), spend = 0" in normalized + assert "WHERE entity_type = $1 AND entity_id = $2 AND window_duration = $3" in normalized + assert "AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')" in normalized + assert params[:4] == ("team", "t1", "30d", datetime(2026, 8, 31)) + + +@pytest.mark.asyncio +async def test_seed_receives_the_batch_request_ids_and_earliest_start_to_exclude(): + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(_batch(("req-1", "req-2", "req-3"), 3.0),), + spend_logs_aggregate=aggregate, + ) + + assert aggregate.calls[0]["exclude_request_ids"] == ("req-1", "req-2", "req-3") + assert aggregate.calls[0]["exclude_started_at"] == BATCH_STARTED_AT + + +@pytest.mark.asyncio +async def test_seed_passes_no_start_bound_when_the_batch_has_none(): + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(_batch(("req-1",), 1.0, started_at=None),), + spend_logs_aggregate=aggregate, + ) + + assert aggregate.calls[0]["exclude_started_at"] is None + + +@pytest.mark.asyncio +async def test_new_row_is_not_double_counted_when_the_batch_logs_already_flushed(): + """The spend log writer drains on a ~2s poll while window increments flush + on the ~10s batch tick, so a new row is normally seeded from a table that + already holds this batch's rows. Counting them in both places is what made + a fresh row land at exactly twice the true spend.""" + db = _FakeDB(existing_rows=[]) + already_flushed = _SpendLogsFake( + rows=( + ("req-1", 0.000047, BATCH_STARTED_AT), + ("req-2", 0.000047, BATCH_STARTED_AT + timedelta(seconds=1)), + ("req-3", 0.000047, BATCH_STARTED_AT + timedelta(seconds=2)), + ), + ) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(_batch(("req-1", "req-2", "req-3"), 0.000141),), + spend_logs_aggregate=already_flushed, + ) + + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(0.000141) + + +@pytest.mark.asyncio +async def test_new_row_still_covers_spend_that_predates_the_batch(): + """The exclusion must not throw away the pre-existing spend the seed is for.""" + db = _FakeDB(existing_rows=[]) + spend_logs = _SpendLogsFake(rows=(("older", 0.5, BEFORE_BATCH), ("req-1", 0.000047, BATCH_STARTED_AT))) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(_batch(("req-1",), 0.000047),), + spend_logs_aggregate=spend_logs, + ) + + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(0.500047) + + +@pytest.mark.asyncio +async def test_replayed_request_id_cannot_erase_historical_spend_from_the_seed(): + """request_id can be chosen by the client via x-litellm-call-id. A request + that replays an id from before this batch writes no new LiteLLM_SpendLogs + row (the insert skips duplicates), so the seed must keep counting the + historical row that id belongs to; only its increment is new.""" + db = _FakeDB(existing_rows=[]) + spend_logs = _SpendLogsFake(rows=(("replayed", 0.5, BEFORE_BATCH),)) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(_batch(("replayed",), 0.000047),), + spend_logs_aggregate=spend_logs, + ) + + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(0.500047) + + +@pytest.mark.asyncio +async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet(): + """The other side of the race: rows absent from the aggregate are still + counted exactly once, by their increment.""" + db = _FakeDB(existing_rows=[]) + nothing_flushed = _SpendLogsFake(rows=()) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(_batch(("req-1", "req-2", "req-3"), 0.000141),), + spend_logs_aggregate=nothing_flushed, + ) + + ((_, params),) = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(0.000141) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entity_type, expected_column", + [("key", "api_key = $1"), ("team", "team_id = $1")], +) +async def test_seed_aggregate_sql_excludes_the_request_ids_only_within_the_batch_start_bound( + entity_type, expected_column +): + db = _FakeDB(existing_rows=[{"total": 1.25}]) + + total = await spend_logs_total_excluding( + prisma_client=_FakePrismaClient(db), + entity_type=entity_type, + entity_id="e1", + window_start=WINDOW_A, + exclude_request_ids=("req-1", "req-2"), + exclude_started_at=BATCH_STARTED_AT, + ) + + assert total == pytest.approx(1.25) + ((query, params),) = db.query_raw_calls + normalized = " ".join(query.split()) + assert expected_column in normalized + assert "NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" in normalized + assert 'FROM "LiteLLM_SpendLogs"' in normalized + # startTime is TIMESTAMP(3): the bound is floored to the second so the + # batch's own earliest row cannot round under it. + assert params == ("e1", WINDOW_A, ("req-1", "req-2"), datetime(2026, 8, 10, 12, 0, 0)) + # The ids are bound, never spliced into the statement. + assert "req-1" not in query + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "exclude_request_ids, exclude_started_at", + [(("req-1",), None), ((), BATCH_STARTED_AT)], +) +async def test_seed_aggregate_excludes_nothing_without_both_ids_and_a_start_bound( + exclude_request_ids, exclude_started_at +): + """Ids without a start bound would reopen the replayed-id hole, so the + seed counts everything instead; at worst that over-counts one batch.""" + db = _FakeDB(existing_rows=[{"total": 1.25}]) + + total = await spend_logs_total_excluding( + prisma_client=_FakePrismaClient(db), + entity_type="key", + entity_id="e1", + window_start=WINDOW_A, + exclude_request_ids=exclude_request_ids, + exclude_started_at=exclude_started_at, + ) + + assert total == pytest.approx(1.25) + ((query, params),) = db.query_raw_calls + assert "request_id" not in query + assert params == ("e1", WINDOW_A) + + +@pytest.mark.asyncio +async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs_column(): + db = _FakeDB(existing_rows=[]) + + total = await spend_logs_total_excluding( + prisma_client=_FakePrismaClient(db), + entity_type="user", + entity_id="u1", + window_start=WINDOW_A, + exclude_request_ids=(), + exclude_started_at=None, + ) + + assert total is None + assert db.query_raw_calls == [] + + +@pytest.mark.asyncio +async def test_seed_aggregate_treats_an_entity_with_no_rows_as_zero(): + db = _FakeDB(existing_rows=[]) + + total = await spend_logs_total_excluding( + prisma_client=_FakePrismaClient(db), + entity_type="key", + entity_id="k-unknown", + window_start=WINDOW_A, + exclude_request_ids=(), + exclude_started_at=None, + ) + + assert total == 0.0 diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index b1f647bdd3d..d28cf8c9c6a 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -4,8 +4,8 @@ import json import re - from collections.abc import Callable +from contextlib import asynccontextmanager from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, call, patch @@ -15,6 +15,9 @@ from redis.exceptions import DataError import litellm from litellm.proxy._types import Litellm_EntityType from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + build_window_spend_transaction, +) @pytest.mark.asyncio @@ -64,9 +67,7 @@ async def test_daily_spend_tracking_with_disabled_spend_logs(): assert db_writer.add_spend_log_transaction_to_daily_user_transaction.called # Verify the payload passed to add_spend_log_transaction_to_daily_user_transaction - call_args = ( - db_writer.add_spend_log_transaction_to_daily_user_transaction.call_args[1] - ) + call_args = db_writer.add_spend_log_transaction_to_daily_user_transaction.call_args[1] assert "payload" in call_args assert call_args["payload"]["spend"] == 0.1 assert call_args["payload"]["model"] == "gpt-4" @@ -406,7 +407,7 @@ async def test_update_daily_spend_sorting(): # fields, but entity_id is sufficient to test sorting. daily_spend_transactions = { f"test_key_{i}": { - "user_id": f"user{60-i}", # user60 ... user11, reverse order + "user_id": f"user{60 - i}", # user60 ... user11, reverse order "date": "2024-01-01", "api_key": "test-api-key", "model": "gpt-4", @@ -985,9 +986,9 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i transaction_dict = call[1]["update"] # Each transaction should have one key with the format tag_date_api_key_model_provider for key, transaction in transaction_dict.items(): - assert ( - transaction["request_id"] == request_id - ), f"request_id should be {request_id} but got {transaction.get('request_id')}" + assert transaction["request_id"] == request_id, ( + f"request_id should be {request_id} but got {transaction.get('request_id')}" + ) @pytest.mark.asyncio @@ -1213,21 +1214,15 @@ async def test_add_spend_log_transaction_to_daily_agent_transaction_calls_common } writer.daily_agent_spend_update_queue.add_update = AsyncMock() - original_common_helper = ( - writer._common_add_spend_log_transaction_to_daily_transaction - ) - writer._common_add_spend_log_transaction_to_daily_transaction = AsyncMock( - wraps=original_common_helper - ) + original_common_helper = writer._common_add_spend_log_transaction_to_daily_transaction + writer._common_add_spend_log_transaction_to_daily_transaction = AsyncMock(wraps=original_common_helper) await writer.add_spend_log_transaction_to_daily_agent_transaction( payload=payload, prisma_client=mock_prisma, ) - assert ( - writer._common_add_spend_log_transaction_to_daily_transaction.await_count == 1 - ) + assert writer._common_add_spend_log_transaction_to_daily_transaction.await_count == 1 @pytest.mark.asyncio @@ -1382,6 +1377,7 @@ async def test_update_daily_spend_re_raises_exception_after_logging(): Test that when batch upsert fails, the exception is properly re-raised after logging. This ensures that error handling continues to work correctly upstream. """ + def raise_connection_lost(): raise ValueError("Database connection lost") @@ -1562,9 +1558,7 @@ async def test_update_database_creates_single_task(): patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), - patch( - "litellm.proxy.db.db_spend_update_writer.asyncio.create_task" - ) as mock_create_task, + patch("litellm.proxy.db.db_spend_update_writer.asyncio.create_task") as mock_create_task, ): await db_writer.update_database( token="test-token", @@ -1663,9 +1657,7 @@ async def test_daily_agent_receives_deepcopied_payload(): db_writer._update_agent_db = AsyncMock() db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() - db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock( - side_effect=capture_agent_payload - ) + db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock(side_effect=capture_agent_payload) db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock() @@ -1727,8 +1719,8 @@ async def test_commit_spend_updates_uses_pipeline(): mock_redis_update_buffer = AsyncMock() mock_redis_update_buffer.store_in_memory_spend_updates_in_redis = AsyncMock() # Return all-None tuple (no data to commit); the pipeline yields 6 slots - mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = ( - AsyncMock(return_value=(None, None, None, None, None, None)) + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, None) ) db_writer.redis_update_buffer = mock_redis_update_buffer @@ -1782,7 +1774,7 @@ async def test_commit_with_redis_requeues_all_on_db_failure(): mock_redis_update_buffer = AsyncMock() mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( - return_value=(db_spend, daily_user, None, None, None, None) + return_value=(db_spend, daily_user, None, None, None, None, None) ) mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() db_writer.redis_update_buffer = mock_redis_update_buffer @@ -1837,7 +1829,7 @@ async def test_commit_with_redis_only_requeues_failed_category(): mock_redis_update_buffer = AsyncMock() mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( - return_value=(db_spend, daily_user, None, None, None, None) + return_value=(db_spend, daily_user, None, None, None, None, None) ) mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() db_writer.redis_update_buffer = mock_redis_update_buffer @@ -1885,7 +1877,7 @@ async def test_commit_with_redis_no_requeue_on_success(): mock_redis_update_buffer = AsyncMock() mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( - return_value=(db_spend, None, None, None, None, None) + return_value=(db_spend, None, None, None, None, None, None) ) mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() db_writer.redis_update_buffer = mock_redis_update_buffer @@ -2156,9 +2148,7 @@ async def test_update_database_does_not_deepcopy_on_request_path(): db_writer._update_org_db = AsyncMock() db_writer._update_tag_db = AsyncMock() db_writer._update_agent_db = AsyncMock() - db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock( - side_effect=capture_batch_payload - ) + db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock(side_effect=capture_batch_payload) db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock() @@ -2250,9 +2240,7 @@ async def test_spend_update_path_never_queries_user_cache_with_none_user_id(): db_writer = DBSpendUpdateWriter() strict_redis_backed_cache = MagicMock() - strict_redis_backed_cache.async_get_cache = AsyncMock( - side_effect=DataError("Invalid input of type: 'NoneType'") - ) + strict_redis_backed_cache.async_get_cache = AsyncMock(side_effect=DataError("Invalid input of type: 'NoneType'")) with ( patch.object(litellm, "max_budget", 0), @@ -2380,8 +2368,7 @@ async def test_daily_transaction_carries_compression_saved_tokens(): cache_write_cost = model_info.get("cache_creation_input_token_cost") or input_cost assert transaction["compression_savings_spend"] == pytest.approx(7600 * input_cost) assert transaction["prompt_caching_savings_spend"] == pytest.approx( - 40 * max(input_cost - cache_read_cost, 0.0) - - 15 * (cache_write_cost - input_cost) + 40 * max(input_cost - cache_read_cost, 0.0) - 15 * (cache_write_cost - input_cost) ) assert transaction["compression_savings_spend"] > 0 assert transaction["prompt_caching_savings_spend"] > 0 @@ -2421,6 +2408,304 @@ async def test_daily_transaction_compression_saved_tokens_zero_when_absent(): assert transaction["prompt_caching_savings_spend"] == 0 +# --------------------------------------------------------------------------- +# Budget window spend flush (LiteLLM_BudgetWindowSpend) +# --------------------------------------------------------------------------- + + +class _WindowSpendFakeBatcher: + def __init__(self): + self.calls = [] + + def execute_raw(self, query, *args): + self.calls.append((query, args)) + + +class _WindowSpendFakeDB: + """Minimal prisma_client.db that records the raw statements it is handed.""" + + def __init__(self, existing_rows=None): + self.existing_rows = existing_rows or [] + self.query_raw_calls = [] + self.batcher = _WindowSpendFakeBatcher() + + async def query_raw(self, query, *args): + self.query_raw_calls.append((query, args)) + if "LiteLLM_BudgetWindowSpend" in query: + return self.existing_rows + return [] + + @asynccontextmanager + async def _tx(self): + yield self + + def tx(self, timeout=None): + return self._tx() + + @asynccontextmanager + async def _batch(self): + yield self.batcher + + def batch_(self): + return self._batch() + + +class _WindowSpendFakePrisma: + def __init__(self, db): + self.db = db + + +def _window_spend_upserts(db): + return [params for query, params in db.batcher.calls if "LiteLLM_BudgetWindowSpend" in query] + + +@pytest.mark.asyncio +async def test_window_spend_queue_is_flushed_without_redis_buffer(): + """The in-memory window queue must reach the DB on the same scheduler tick + as the other spend queues when the Redis buffer is off.""" + db_writer = DBSpendUpdateWriter() + await db_writer.window_spend_update_queue.add_update( + build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=0.5, + ) + ) + db = _WindowSpendFakeDB( + existing_rows=[{"entity_type": "key", "entity_id": "hashed-token", "window_duration": "30d"}] + ) + + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + upserts = _window_spend_upserts(db) + assert len(upserts) == 1 + assert upserts[0][0] == "key" + assert upserts[0][1] == "hashed-token" + assert upserts[0][2] == "30d" + assert upserts[0][5] == pytest.approx(0.5) + assert db_writer.window_spend_update_queue.update_queue.qsize() == 0 + + +@pytest.mark.asyncio +async def test_window_spend_queue_is_handed_to_the_redis_buffer(): + """Multi-pod deployments buffer through Redis, so the window queue has to + ride the same rpush path as the daily queues.""" + db_writer = DBSpendUpdateWriter() + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, None) + ) + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=MagicMock(), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + stored = mock_redis_update_buffer.store_in_memory_spend_updates_in_redis.call_args[1] + assert stored["window_spend_update_queue"] is db_writer.window_spend_update_queue + + +@pytest.mark.asyncio +async def test_window_spend_transactions_from_redis_are_committed_by_the_lock_winner(): + db_writer = DBSpendUpdateWriter() + window_transactions = ( + build_window_spend_transaction( + entity_type="team", + entity_id="team-1", + window_duration="7d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=2.0, + ), + ) + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, window_transactions) + ) + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + db = _WindowSpendFakeDB(existing_rows=[{"entity_type": "team", "entity_id": "team-1", "window_duration": "7d"}]) + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + upserts = _window_spend_upserts(db) + assert len(upserts) == 1 + assert upserts[0][:3] == ("team", "team-1", "7d") + assert upserts[0][5] == pytest.approx(2.0) + + +@pytest.mark.asyncio +async def test_window_spend_transactions_are_not_committed_without_the_pod_lock(): + """Every pod buffers to Redis but only the lock winner may drain it.""" + db_writer = DBSpendUpdateWriter() + mock_redis_update_buffer = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=False) + db = _WindowSpendFakeDB() + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline.assert_not_called() + assert _window_spend_upserts(db) == [] + + +@pytest.mark.asyncio +async def test_failed_window_spend_commit_requeues_the_increments_and_continues_the_flush(): + """Budget enforcement trusts a current window row without reconciling it + against LiteLLM_SpendLogs, so a dropped increment would let the key spend + past its limit after the next reseed. The increments must go back on the + queue, and the tool registry flush must still run.""" + db_writer = DBSpendUpdateWriter() + transaction = build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=0.5, + request_id="req-1", + ) + await db_writer.window_spend_update_queue.add_update(transaction) + db = _WindowSpendFakeDB() + db.query_raw = AsyncMock(side_effect=Exception("connection reset")) + db_writer._flush_tool_discovery_queue = AsyncMock() + + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + db_writer._flush_tool_discovery_queue.assert_called_once() + requeued = await db_writer.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() + assert requeued == (transaction,) + + +@pytest.mark.asyncio +async def test_failed_window_spend_commit_from_redis_is_restored_to_redis(): + """The Redis drain is destructive, so a failed window commit has to push + the popped increments back exactly like the other spend categories.""" + db_writer = DBSpendUpdateWriter() + window_transactions = ( + build_window_spend_transaction( + entity_type="team", + entity_id="team-1", + window_duration="7d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=2.0, + request_id="req-1", + ), + ) + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, window_transactions) + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + db = _WindowSpendFakeDB() + db.query_raw = AsyncMock(side_effect=Exception("connection reset")) + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + assert _window_spend_upserts(db) == [] + mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once_with( + window_spend_update_transactions=window_transactions + ) + db_writer.pod_lock_manager.release_lock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_update_database_returns_the_spend_log_request_id(): + """The budget-window seed excludes the log rows its increments already + cover, so the caller needs the id this call was recorded under. It cannot + be re-derived: cache hits append time.time() to the id.""" + db_writer = DBSpendUpdateWriter() + db_writer._insert_spend_log_to_db = AsyncMock() + db_writer._enqueue_tool_usage_transaction = AsyncMock() + + with ( + patch.multiple( # test-quality-ok: update_database lazily imports these proxy_server globals; no injection seam + "litellm.proxy.proxy_server", + disable_spend_logs=False, + prisma_client=MagicMock(), + litellm_proxy_budget_name="test-budget", + ) + ): + request_id = await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id="test-team", + org_id=None, + kwargs={"model": "gpt-4", "custom_llm_provider": "openai", "litellm_call_id": "call-xyz"}, + completion_response=MagicMock(), + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.1, + ) + await asyncio.sleep(0) + + assert request_id is not None + # Same id the spend log row was queued under. + assert request_id == db_writer._insert_spend_log_to_db.call_args[1]["payload"]["request_id"] + + +@pytest.mark.asyncio +async def test_update_database_returns_none_when_the_payload_cannot_be_built(): + db_writer = DBSpendUpdateWriter() + + with ( + patch.multiple( # test-quality-ok: update_database lazily imports these proxy_server globals; no injection seam + "litellm.proxy.proxy_server", + disable_spend_logs=False, + prisma_client=MagicMock(), + litellm_proxy_budget_name="test-budget", + ), + patch( # test-quality-ok: the payload builder is called by name inside update_database; no injection seam + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + side_effect=Exception("payload boom"), + ), + ): + request_id = await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id="test-team", + org_id=None, + kwargs={"model": "gpt-4"}, + completion_response=MagicMock(), + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.1, + ) + + assert request_id is None + + @pytest.mark.asyncio async def test_commit_spend_updates_to_db_does_not_stamp_key_settings_updated_at(): """Spend flushes must leave settings_updated_at alone, or it decays into @@ -2721,9 +3006,7 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey "call_type, expects_flush", [("aresponses", True), ("responses", True), ("acompletion", False)], ) -async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls( - call_type: str, expects_flush: bool -): +async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(call_type: str, expects_flush: bool): """ A `previous_response_id` chained straight off the previous turn reads the DB, so a Responses row cannot sit in this worker's queue until the monitor's next poll. @@ -2753,9 +3036,7 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls( pytest.param("", True, id="injected-before-a-deployment-was-chosen"), ], ) -async def test_caching_savings_are_attributed_to_the_deployment_that_was_injected( - injected_deployment, attributed -): +async def test_caching_savings_are_attributed_to_the_deployment_that_was_injected(injected_deployment, attributed): """Retries, same-group failover and cross-model-group fallbacks all reuse one metadata bucket and one litellm_call_id, so a marker written by the leg that injected is visible to every sibling and nothing request-scoped can tell them apart. diff --git a/tests/test_litellm/proxy/db/test_model_access_group_spend.py b/tests/test_litellm/proxy/db/test_model_access_group_spend.py new file mode 100644 index 00000000000..d2d079bb0e4 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_model_access_group_spend.py @@ -0,0 +1,510 @@ +"""Spend accumulation for model access group budgets.""" + +import asyncio +from collections.abc import Mapping, Sequence + +import pytest + +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY +from litellm.proxy._types import DBSpendUpdateTransactions, Litellm_EntityType, SpendUpdateQueueItem +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter, debitable_model_access_groups +from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.proxy.spend_tracking.spend_tracking_utils import get_request_model_access_groups + + +class _FakeRouter: + """Deployment lookup returning the access groups each deployment declares.""" + + def __init__(self, deployments: Mapping[str, Sequence[str] | None]) -> None: + self._deployments = deployments + + def get_model_info(self, id: str) -> dict | None: + if id not in self._deployments: + return None + declared = self._deployments[id] + model_info: dict = {"id": id} + if declared is not None: + model_info["access_groups"] = list(declared) + return {"model_name": "some-model", "model_info": model_info} + + +class _FakeBatchTable: + def __init__(self) -> None: + self.calls: list[tuple[dict, dict]] = [] + + def update_many(self, where: dict, data: dict) -> None: + self.calls.append((where, data)) + + +class _FakeBatcher: + def __init__(self) -> None: + self.tables: dict[str, _FakeBatchTable] = {} + + def __getattr__(self, name: str) -> _FakeBatchTable: + return self.tables.setdefault(name, _FakeBatchTable()) + + +class _FakeBatchManager: + def __init__(self, batcher: _FakeBatcher) -> None: + self._batcher = batcher + + async def __aenter__(self) -> _FakeBatcher: + return self._batcher + + async def __aexit__(self, *exc_info: object) -> bool: + return False + + +class _FakeTransaction: + def __init__(self, batcher: _FakeBatcher) -> None: + self._batcher = batcher + + def batch_(self) -> _FakeBatchManager: + return _FakeBatchManager(self._batcher) + + async def __aenter__(self) -> "_FakeTransaction": + return self + + async def __aexit__(self, *exc_info: object) -> bool: + return False + + +class _FakeDb: + def __init__(self, batcher: _FakeBatcher) -> None: + self._batcher = batcher + + def tx(self, timeout: object = None) -> _FakeTransaction: + return _FakeTransaction(self._batcher) + + +class _FakePrismaClient: + def __init__(self) -> None: + self.batcher = _FakeBatcher() + self.db = _FakeDb(self.batcher) + + +def _empty_transactions(**overrides: dict[str, float]) -> DBSpendUpdateTransactions: + return DBSpendUpdateTransactions( + user_list_transactions=overrides.get("user_list_transactions", {}), + end_user_list_transactions=overrides.get("end_user_list_transactions", {}), + key_list_transactions=overrides.get("key_list_transactions", {}), + team_list_transactions=overrides.get("team_list_transactions", {}), + team_member_list_transactions=overrides.get("team_member_list_transactions", {}), + org_list_transactions=overrides.get("org_list_transactions", {}), + tag_list_transactions=overrides.get("tag_list_transactions", {}), + agent_list_transactions=overrides.get("agent_list_transactions", {}), + model_access_group_list_transactions=overrides.get("model_access_group_list_transactions", {}), + ) + + +async def _drain(queue: SpendUpdateQueue) -> list[SpendUpdateQueueItem]: + return await queue.flush_all_updates_from_in_memory_queue() + + +# --- enqueue --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_single_matched_group_enqueues_one_item_with_full_cost(): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=0.42, + request_model_access_groups=["premium-pool"], + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) + + updates = await _drain(writer.spend_update_queue) + assert updates == [ + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, + entity_id="premium-pool", + response_cost=0.42, + ) + ] + + +@pytest.mark.asyncio +async def test_every_matched_group_is_charged_the_full_cost_not_a_split(): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=0.30, + request_model_access_groups=["pool-a", "pool-b", "pool-c"], + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["pool-a", "pool-b", "pool-c"]}), + ) + + updates = await _drain(writer.spend_update_queue) + assert [update["entity_id"] for update in updates] == ["pool-a", "pool-b", "pool-c"] + assert [update["response_cost"] for update in updates] == [0.30, 0.30, 0.30] + assert {update["entity_type"] for update in updates} == {Litellm_EntityType.MODEL_ACCESS_GROUP} + + +@pytest.mark.parametrize("attributed", [None, [], ()]) +@pytest.mark.asyncio +async def test_no_attributed_groups_enqueues_nothing(attributed): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=1.0, + request_model_access_groups=attributed, + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) + + assert await _drain(writer.spend_update_queue) == [] + + +@pytest.mark.asyncio +async def test_no_prisma_client_enqueues_nothing(): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=1.0, + request_model_access_groups=["premium-pool"], + served_model_id="deployment-1", + prisma_client=None, + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) + + assert await _drain(writer.spend_update_queue) == [] + + +@pytest.mark.asyncio +async def test_group_outside_the_attributed_set_is_never_debited(): + """The served deployment also sits in a pool auth never attributed; that pool stays untouched.""" + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=0.10, + request_model_access_groups=["premium-pool"], + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["premium-pool", "unattributed-pool"]}), + ) + + updates = await _drain(writer.spend_update_queue) + assert [update["entity_id"] for update in updates] == ["premium-pool"] + + +# --- fallback guard -------------------------------------------------------- + + +def test_fallback_to_a_model_in_another_pool_debits_nothing(): + assert ( + debitable_model_access_groups( + attributed=["premium-pool"], + served_model_id="fallback-deployment", + router=_FakeRouter({"fallback-deployment": ["cheap-pool"]}), + ) + == () + ) + + +def test_fallback_to_a_model_in_no_pool_debits_nothing(): + assert ( + debitable_model_access_groups( + attributed=["premium-pool"], + served_model_id="fallback-deployment", + router=_FakeRouter({"fallback-deployment": None}), + ) + == () + ) + + +def test_attributed_set_stands_when_the_served_deployment_is_unknown(): + assert debitable_model_access_groups( + attributed=["premium-pool"], + served_model_id="not-in-router", + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) == ("premium-pool",) + + +def test_attributed_set_stands_without_a_router(): + assert debitable_model_access_groups( + attributed=["premium-pool", "premium-pool"], + served_model_id="deployment-1", + router=None, + ) == ("premium-pool",) + + +def test_partial_overlap_keeps_only_the_intersection(): + assert debitable_model_access_groups( + attributed=["pool-a", "pool-b"], + served_model_id="deployment-1", + router=_FakeRouter({"deployment-1": ["pool-b", "pool-c"]}), + ) == ("pool-b",) + + +def test_only_real_group_names_ever_become_entity_ids(): + """Whatever shape the attributed set arrives in, an empty or non-string name never reaches the queue.""" + assert debitable_model_access_groups( + attributed=["pool-a", "", "pool-a", None, 7], + served_model_id=None, + router=None, + ) == ("pool-a",) + + +# --- metadata extraction --------------------------------------------------- + + +def test_access_groups_read_from_request_metadata(): + kwargs = {"litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool-a", "pool-b", "pool-a"]}}} + assert get_request_model_access_groups(kwargs) == ("pool-a", "pool-b") + + +def test_access_groups_read_from_litellm_metadata(): + kwargs = {"litellm_params": {"litellm_metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool-a"]}}} + assert get_request_model_access_groups(kwargs) == ("pool-a",) + + +def test_standard_logging_payload_wins_over_metadata(): + kwargs = { + "litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["from-metadata"]}}, + "standard_logging_object": {"request_model_access_groups": ["from-payload"]}, + } + assert get_request_model_access_groups(kwargs) == ("from-payload",) + + +def test_metadata_is_used_when_the_logging_payload_carries_no_groups(): + kwargs = { + "litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["from-metadata"]}}, + "standard_logging_object": {"request_model_access_groups": []}, + } + assert get_request_model_access_groups(kwargs) == ("from-metadata",) + + +@pytest.mark.parametrize("stamped", ["pool-a", 7, {"pool-a": 1}]) +def test_non_list_access_group_metadata_is_ignored(stamped): + kwargs = {"litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: stamped}}} + assert get_request_model_access_groups(kwargs) == () + + +def test_key_absent_from_metadata_yields_no_groups(): + """The chat path only stamps the key when something matched, so absent must mean nothing to debit.""" + kwargs = {"litellm_params": {"metadata": {"user_api_key_user_id": "u-1"}}} + assert get_request_model_access_groups(kwargs) == () + + +def test_explicit_none_yields_no_groups(): + """The pass-through path stamps the key unconditionally, so it can be present and None.""" + kwargs = {"litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: None}}} + assert get_request_model_access_groups(kwargs) == () + + +@pytest.mark.parametrize( + "metadata", + [ + {"user_api_key_user_id": "u-1"}, + {MODEL_ACCESS_GROUP_METADATA_KEY: None}, + ], + ids=["key-absent", "key-present-but-none"], +) +@pytest.mark.asyncio +async def test_neither_absent_nor_none_metadata_debits_anything(metadata): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=0.5, + request_model_access_groups=get_request_model_access_groups({"litellm_params": {"metadata": metadata}}), + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) + + assert await _drain(writer.spend_update_queue) == [] + + +def test_detached_sub_call_falls_back_to_the_auth_object(): + """Sub-calls inherit only the identity keys, so the groups come off user_api_key_auth there.""" + + class _Auth: + matched_model_access_groups = ["premium-pool"] + + kwargs = {"litellm_params": {"metadata": {"user_api_key_auth": _Auth()}}} + assert get_request_model_access_groups(kwargs) == ("premium-pool",) + + +def test_stamped_metadata_wins_over_the_auth_object(): + class _Auth: + matched_model_access_groups = ["stale-pool"] + + kwargs = { + "litellm_params": { + "metadata": { + MODEL_ACCESS_GROUP_METADATA_KEY: ["fresh-pool"], + "user_api_key_auth": _Auth(), + } + } + } + assert get_request_model_access_groups(kwargs) == ("fresh-pool",) + + +def test_auth_object_without_matched_groups_yields_no_groups(): + class _Auth: + matched_model_access_groups = None + + kwargs = {"litellm_params": {"metadata": {"user_api_key_auth": _Auth()}}} + assert get_request_model_access_groups(kwargs) == () + + +def test_non_string_entries_are_dropped(): + kwargs = {"litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool-a", None, "", 3]}}} + assert get_request_model_access_groups(kwargs) == ("pool-a",) + + +def test_missing_metadata_yields_no_groups(): + assert get_request_model_access_groups(None) == () + assert get_request_model_access_groups({}) == () + assert get_request_model_access_groups({"litellm_params": {}}) == () + + +# --- queue bucketing and redis round trip ---------------------------------- + + +def test_access_group_updates_aggregate_into_their_own_bucket(): + queue = SpendUpdateQueue() + + transactions = queue.get_aggregated_db_spend_update_transactions( + [ + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, entity_id="pool-a", response_cost=0.1 + ), + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, entity_id="pool-a", response_cost=0.2 + ), + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, entity_id="pool-b", response_cost=0.5 + ), + SpendUpdateQueueItem(entity_type=Litellm_EntityType.TAG, entity_id="pool-a", response_cost=9.0), + ] + ) + + assert transactions["model_access_group_list_transactions"] == {"pool-a": pytest.approx(0.3), "pool-b": 0.5} + assert transactions["tag_list_transactions"] == {"pool-a": 9.0} + + +def test_access_group_transactions_survive_the_redis_buffer_merge(): + merged = RedisUpdateBuffer._combine_list_of_transactions( + [ + _empty_transactions(model_access_group_list_transactions={"pool-a": 0.25}), + _empty_transactions(model_access_group_list_transactions={"pool-a": 0.25, "pool-b": 1.0}), + ] + ) + + assert merged["model_access_group_list_transactions"] == {"pool-a": 0.5, "pool-b": 1.0} + + +@pytest.mark.asyncio +async def test_redis_buffer_requeues_access_group_transactions_as_queue_items(): + queue = SpendUpdateQueue() + daily_queue = DailySpendUpdateQueue() + + await RedisUpdateBuffer._restore_spend_updates_to_in_memory_queues( + db_spend_update_transactions=_empty_transactions(model_access_group_list_transactions={"pool-a": 0.75}), + daily_spend_update_transactions=None, + daily_team_spend_update_transactions=None, + daily_org_spend_update_transactions=None, + daily_end_user_spend_update_transactions=None, + daily_agent_spend_update_transactions=None, + window_spend_update_transactions=None, + spend_update_queue=queue, + daily_spend_update_queue=daily_queue, + daily_team_spend_update_queue=daily_queue, + daily_org_spend_update_queue=daily_queue, + daily_end_user_spend_update_queue=daily_queue, + daily_agent_spend_update_queue=daily_queue, + window_spend_update_queue=None, + ) + + updates = await _drain(queue) + assert updates == [ + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, + entity_id="pool-a", + response_cost=0.75, + ) + ] + + +# --- flush to postgres ----------------------------------------------------- + + +@pytest.mark.asyncio +async def test_commit_increments_spend_on_the_model_access_group_budget_table(): + prisma_client = _FakePrismaClient() + + await DBSpendUpdateWriter()._commit_spend_updates_to_db( + prisma_client=prisma_client, + n_retry_times=0, + proxy_logging_obj=None, + db_spend_update_transactions=_empty_transactions( + model_access_group_list_transactions={"pool-b": 0.5, "pool-a": 0.25} + ), + ) + + assert prisma_client.batcher.tables["litellm_modelaccessgroupbudgettable"].calls == [ + ({"access_group_name": "pool-a"}, {"spend": {"increment": 0.25}}), + ({"access_group_name": "pool-b"}, {"spend": {"increment": 0.5}}), + ] + assert "litellm_tagtable" not in prisma_client.batcher.tables + + +# --- end-to-end through the batched fan-out -------------------------------- + + +@pytest.mark.asyncio +async def test_batch_database_updates_enqueues_access_group_spend(): + writer = DBSpendUpdateWriter() + + await writer._batch_database_updates( + response_cost=0.15, + user_id=None, + hashed_token=None, + team_id=None, + org_id=None, + end_user_id=None, + prisma_client=object(), + litellm_proxy_budget_name=None, + payload={"model_id": "deployment-1", "spend": 0.15}, + request_model_access_groups=("pool-a", "pool-b"), + ) + await asyncio.sleep(0) + + access_group_updates = [ + update + for update in await _drain(writer.spend_update_queue) + if update["entity_type"] is Litellm_EntityType.MODEL_ACCESS_GROUP + ] + assert [(update["entity_id"], update["response_cost"]) for update in access_group_updates] == [ + ("pool-a", 0.15), + ("pool-b", 0.15), + ] + + +@pytest.mark.asyncio +async def test_batch_database_updates_enqueues_nothing_without_access_groups(): + writer = DBSpendUpdateWriter() + + await writer._batch_database_updates( + response_cost=0.15, + user_id=None, + hashed_token=None, + team_id=None, + org_id=None, + end_user_id=None, + prisma_client=object(), + litellm_proxy_budget_name=None, + payload={"model_id": "deployment-1", "spend": 0.15}, + ) + await asyncio.sleep(0) + + updates = await _drain(writer.spend_update_queue) + assert [update for update in updates if update["entity_type"] is Litellm_EntityType.MODEL_ACCESS_GROUP] == [] diff --git a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py index dcc0036ff04..966a638f6a4 100644 --- a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py @@ -101,6 +101,49 @@ def test_per_model_reads_route_to_reader_writes_to_writer(): assert actions.delete_many is writer_inner.litellm_usertable.delete_many +def test_writer_pinned_client_bypasses_reader_routing(): + """Regression for #38556: read-after-write reconciles must see the writer's + just-committed rows, so WriterPinnedClient must resolve reads to the writer + even when a read replica is configured.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient + + writer, writer_inner, reader, reader_inner = _make_wrappers() + writer_inner.litellm_proxymodeltable = _model_actions_mock("writer_models") + reader_inner.litellm_proxymodeltable = _model_actions_mock("reader_models") + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + + pinned = WriterPinnedClient(routing) + + assert pinned.db is writer + assert pinned.db.litellm_proxymodeltable.find_many is writer_inner.litellm_proxymodeltable.find_many + + +def test_writer_pinned_client_passes_through_single_db(): + from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient + + writer, _, _, _ = _make_wrappers() + + assert WriterPinnedClient(writer).db is writer + + +def test_writer_pinned_client_yields_to_routed_reads_when_writer_down(): + """The pin must not break reader-only degraded mode: a proxy that starts + during a primary outage still loads DB-backed models from the replica, so + while the writer is degraded the pin resolves to the routed wrapper.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient + + writer, writer_inner, reader, reader_inner = _make_wrappers() + writer_inner.litellm_proxymodeltable = _model_actions_mock("writer_models") + reader_inner.litellm_proxymodeltable = _model_actions_mock("reader_models") + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + routing._writer_unavailable = True + + pinned = WriterPinnedClient(routing) + + assert pinned.db is routing + assert pinned.db.litellm_proxymodeltable.find_many is reader_inner.litellm_proxymodeltable.find_many + + @pytest.mark.asyncio async def test_connect_invokes_both_clients(): from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py new file mode 100644 index 00000000000..816f9ae72f4 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -0,0 +1,250 @@ +"""Window-spend reads in ``SpendCounterReseed``. + +The maintained ``LiteLLM_BudgetWindowSpend`` row replaces a per-request +``LiteLLM_SpendLogs`` range scan, so these pin *when* the aggregate is still +allowed to run: only when the row is missing or belongs to an older window. +""" + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace + +import pytest + +from litellm.caching.dual_cache import DualCache +from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + +WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc) + + +class _FakeWindowSpendTable: + def __init__(self, row: SimpleNamespace | None, error: Exception | None = None) -> None: + self._row = row + self._error = error + self.where_clauses: list[dict] = [] + + async def find_unique(self, where: dict): + self.where_clauses.append(where) + if self._error is not None: + raise self._error + return self._row + + +class _FakeSpendLogsTable: + def __init__(self, total: float) -> None: + self._total = total + self.call_count = 0 + + async def group_by(self, by: list[str], where: dict, sum: dict): + self.call_count += 1 + return [{by[0]: where.get(by[0]), "_sum": {"spend": self._total}}] + + +class _FakePrismaClient: + def __init__( + self, + row: SimpleNamespace | None = None, + spend_logs_total: float = 0.0, + error: Exception | None = None, + ) -> None: + self.db = SimpleNamespace( + litellm_budgetwindowspend=_FakeWindowSpendTable(row=row, error=error), + litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total), + ) + + +def _row(window_start: datetime, spend: float) -> SimpleNamespace: + return SimpleNamespace(window_start=window_start, spend=spend) + + +@pytest.mark.asyncio +async def test_window_from_table_reads_row_by_primary_key(): + """The lookup must use the table's own entity_type values ("key"), not the + "Key"/"Team" labels the counter keys and spend-log aggregates use.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result == 4.5 + assert prisma.db.litellm_budgetwindowspend.where_clauses == [ + { + "entity_type_entity_id_window_duration": { + "entity_type": "key", + "entity_id": "tok-1", + "window_duration": "30d", + } + } + ] + + +@pytest.mark.asyncio +async def test_window_from_table_maps_team_entity_type(): + prisma = _FakePrismaClient(row=_row(WINDOW_START, 9.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Team", + entity_id="team-1", + window_duration="1d", + expected_window_start=WINDOW_START, + ) + + assert result == 9.0 + inner = prisma.db.litellm_budgetwindowspend.where_clauses[0]["entity_type_entity_id_window_duration"] + assert inner["entity_type"] == "team" + + +@pytest.mark.asyncio +async def test_window_from_table_trusts_row_newer_than_expected_window(): + """Regression: a pod holding a stale ``reset_at`` computes an expected start + behind a window another pod already rolled. Trusting only an exact match + would make it re-add the previous window's spend to the current one.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START + timedelta(days=1), 2.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result == 2.0 + + +@pytest.mark.asyncio +async def test_window_from_table_rejects_row_from_previous_window(): + prisma = _FakePrismaClient(row=_row(WINDOW_START - timedelta(seconds=1), 99.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_window_from_table_treats_naive_row_timestamp_as_utc(): + """The column is ``timestamp(3)``, so a driver that hands back a naive value + must still compare against the tz-aware expected start.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START.replace(tzinfo=None), 3.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result == 3.0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "prisma, entity_type", + [ + (_FakePrismaClient(row=None), "Key"), + (_FakePrismaClient(row=_row(WINDOW_START, 1.0)), "User"), + (_FakePrismaClient(error=RuntimeError("connection reset")), "Key"), + (None, "Key"), + ], +) +async def test_window_from_table_returns_none_without_a_usable_row(prisma, entity_type): + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type=entity_type, + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_window_from_db_prefers_the_row_over_the_spend_logs_aggregate(): + """The aggregate range-scans an unindexed table; a current row must keep it + from running at all.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=100.0) + + result = await SpendCounterReseed.window_from_db( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + window_start=WINDOW_START, + ) + + assert result == 4.5 + assert prisma.db.litellm_spendlogs.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "row", + [None, _row(WINDOW_START - timedelta(seconds=1), 99.0)], + ids=["missing_row", "previous_window_row"], +) +async def test_window_from_db_falls_back_to_spend_logs(row): + prisma = _FakePrismaClient(row=row, spend_logs_total=7.25) + + result = await SpendCounterReseed.window_from_db( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + window_start=WINDOW_START, + ) + + assert result == 7.25 + assert prisma.db.litellm_spendlogs.call_count == 1 + + +@pytest.mark.asyncio +async def test_window_from_db_without_a_duration_skips_the_row_lookup(): + """Callers that cannot name the window (no PK) keep the pre-table behavior.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=7.25) + + result = await SpendCounterReseed.window_from_db( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration=None, + window_start=WINDOW_START, + ) + + assert result == 7.25 + assert prisma.db.litellm_budgetwindowspend.where_clauses == [] + + +@pytest.mark.asyncio +async def test_coalesced_window_seeds_a_cold_counter_from_the_row(): + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=100.0) + cache = DualCache() + counter_key = "spend:key:tok-1:window:30d" + + result = await SpendCounterReseed.coalesced_window( + prisma_client=prisma, + spend_counter_cache=cache, + counter_key=counter_key, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + window_start=WINDOW_START, + ) + + assert result == 4.5 + assert cache.in_memory_cache.get_cache(key=counter_key) == 4.5 + assert prisma.db.litellm_spendlogs.call_count == 0 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/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index ca517474a5c..2a540f4f522 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,14 +1,14 @@ -import pytest - - from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch +import pytest + +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.proxy_track_cost_callback import ( - _ProxyDBLogger, _get_budget_reservation_from_metadata, + _ProxyDBLogger, _should_track_cost_callback, _update_database_and_spend_counters, ) @@ -567,9 +567,12 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re @pytest.mark.asyncio async def test_update_database_and_spend_counters_updates_counters_after_db_update(): proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( + return_value="chatcmpl-abc123" + ) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} + start_time = datetime.now() await _update_database_and_spend_counters( proxy_logging_obj=proxy_logging_obj, @@ -581,11 +584,12 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda org_id="test_org_id", kwargs={}, completion_response=None, - start_time=datetime.now(), + start_time=start_time, end_time=datetime.now(), response_cost=0.2, budget_reservation=budget_reservation, request_tags=["tag-a"], + model_access_groups=("premium",), ) proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once() @@ -598,6 +602,9 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda budget_reservation=budget_reservation, end_user_id="test_end_user_id", tags=["tag-a"], + request_id="chatcmpl-abc123", + request_started_at=start_time, + model_access_groups=("premium",), ) @@ -1875,3 +1882,168 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( 1 if expect_spend_log else 0 ) + + +@pytest.mark.asyncio +async def test_update_database_and_spend_counters_forwards_the_spend_log_request_id(): + """The budget-window flush excludes the log rows its increments already + cover. That only works if the id update_database recorded the row under is + handed to the counter update, so this seam is load-bearing.""" + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( + return_value="chatcmpl-abc123" + ) + increment_spend_counters = AsyncMock() + + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + assert increment_spend_counters.await_args.kwargs["request_id"] == "chatcmpl-abc123" + + +@pytest.mark.asyncio +async def test_update_database_and_spend_counters_forwards_a_missing_request_id_as_none(): + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(return_value=None) + increment_spend_counters = AsyncMock() + + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + assert increment_spend_counters.await_args.kwargs["request_id"] is None + + +class _FakeDeploymentLookup: + """Deployment lookup returning the access groups each deployment declares.""" + + def __init__(self, deployments): + self._deployments = deployments + + def get_model_info(self, id): + if id not in self._deployments: + return None + return {"model_name": "premium-haiku", "model_info": {"id": id, "access_groups": list(self._deployments[id])}} + + +def _model_access_group_kwargs(granted, served_model_id=None): + metadata = {"user_api_key": "hashed-key", "user_api_key_user_id": "user-1"} + if granted is not None: + metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(granted) + return { + "call_type": "acompletion", + "model": "premium-haiku", + "litellm_call_id": "test-call-id", + "litellm_params": {"metadata": metadata}, + "stream": False, + "standard_logging_object": {"response_cost": 0.25, "request_tags": None, "model_id": served_model_id}, + } + + +async def _groups_charged_by_the_callback(kwargs, deployments=None): + """The groups the callback hands the spend counters for one request. + + The callback resolves ``proxy_logging_obj`` and the router by importing them off + ``proxy_server`` inside its own body, so there is no seam to inject either through. + """ + logger = _ProxyDBLogger() + with ( + patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + patch( # test-quality-ok: the arguments to this call are the boundary under test + "litellm.proxy.hooks.proxy_track_cost_callback._update_database_and_spend_counters", + new=AsyncMock(), + ) as mock_update, + patch( # test-quality-ok: llm_router is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.llm_router", new=_FakeDeploymentLookup(deployments or {}) + ), + ): + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + return mock_update.await_args.kwargs["model_access_groups"] + + +@pytest.mark.asyncio +async def test_track_cost_callback_charges_the_model_access_groups_auth_stamped(): + """Auth stamps the matched groups onto request metadata; the callback has to carry them through. + + Without this hop nothing writes ``spend:model_access_group:*`` on the normal path, so with + reservations disabled the budget check reads a counter no one maintains. + """ + charged = await _groups_charged_by_the_callback( + kwargs=_model_access_group_kwargs(granted=["premium", "starter"]), + ) + + assert charged == ("premium", "starter") + + +@pytest.mark.asyncio +async def test_track_cost_callback_charges_no_model_access_group_when_none_were_stamped(): + """A request no budgeted group authorized must not debit anything.""" + charged = await _groups_charged_by_the_callback( + kwargs=_model_access_group_kwargs(granted=None), + ) + + assert charged == () + + +@pytest.mark.asyncio +async def test_spend_counters_only_debit_the_group_the_served_deployment_belongs_to(): + """A caller granted two pools that both cover the model group only draws down the pool that served. + + The database writer already narrows by served deployment, so passing the unnarrowed set to the + live counters let one request block a pool the persisted spend never debited. + """ + charged = await _groups_charged_by_the_callback( + kwargs=_model_access_group_kwargs(granted=["premium", "tier0"], served_model_id="deployment-premium"), + deployments={"deployment-premium": ["premium"], "deployment-tier0": ["tier0"]}, + ) + + assert charged == ("premium",) + + +@pytest.mark.asyncio +async def test_spend_counters_keep_every_granted_group_when_the_deployment_is_unknown(): + """An unidentifiable deployment leaves the auth-time set standing, so nothing silently stops billing.""" + charged = await _groups_charged_by_the_callback( + kwargs=_model_access_group_kwargs(granted=["premium", "tier0"], served_model_id="deployment-gone"), + deployments={"deployment-premium": ["premium"]}, + ) + + assert charged == ("premium", "tier0") diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py index db0557cfbf0..de5fc96c7c3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py @@ -2,6 +2,10 @@ Test access group management endpoints """ +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import datetime +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -449,6 +453,7 @@ async def test_delete_access_group_ignores_models_that_were_already_dead(): mock_prisma = MagicMock() mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[deploy_broken]) mock_prisma.db.litellm_proxymodeltable.update = AsyncMock() + mock_prisma.db.litellm_modelaccessgroupbudgettable.delete = AsyncMock(return_value=None) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( @@ -468,6 +473,7 @@ async def test_delete_access_group_ignores_models_that_were_already_dead(): response = await delete_access_group( access_group="doomed-group", user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + auth_cache=_FakeAuthCache(), ) assert response.models_updated == 1 @@ -568,3 +574,639 @@ async def test_create_access_group_model_missing_everywhere_still_400s(): assert exc_info.value.status_code == 400 assert model_name in str(exc_info.value.detail) + +@dataclass +class _FakeBudgetRow: + budget_id: str + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + budget_reset_at: datetime | None = None + + +@dataclass +class _FakeAccessGroupBudgetRow: + access_group_name: str + budget_id: str | None = None + spend: float = 0.0 + litellm_budget_table: _FakeBudgetRow | None = None + + +@dataclass +class _FakeDeployment: + model_id: str + model_name: str + model_info: dict + + +class _FakeBudgetTable: + """Stands in for litellm_budgettable so a test can see whether a budget row was created, + updated in place, or left orphaned.""" + + def __init__(self, journal: list[str]) -> None: + self.journal = journal + self.rows: dict[str, _FakeBudgetRow] = {} + self.create_calls: list[dict] = [] + self.update_calls: list[tuple[str, dict]] = [] + self.deleted_ids: list[str] = [] + self._sequence = 0 + + async def create(self, data, include=None): + self._sequence += 1 + budget_id = str(data.get("budget_id") or f"budget-{self._sequence}") + row = _FakeBudgetRow( + budget_id=budget_id, + max_budget=data.get("max_budget"), + soft_budget=data.get("soft_budget"), + budget_duration=data.get("budget_duration"), + ) + self.rows[budget_id] = row + self.create_calls.append(dict(data)) + self.journal.append(f"budget_table.create:{budget_id}") + return row + + async def update(self, where, data, include=None): + budget_id = where["budget_id"] + self.update_calls.append((budget_id, dict(data))) + self.journal.append(f"budget_table.update:{budget_id}") + row = self.rows.get(budget_id) + if row is None: + return None + for field_name in ("max_budget", "soft_budget", "budget_duration"): + if data.get(field_name) is not None: + setattr(row, field_name, data[field_name]) + return row + + async def delete(self, where, include=None): + budget_id = where["budget_id"] + self.journal.append(f"budget_table.delete:{budget_id}") + self.deleted_ids.append(budget_id) + return self.rows.pop(budget_id, None) + + +class _FakeAccessGroupBudgetTable: + """Stands in for litellm_modelaccessgroupbudgettable, resolving `include` against the fake + budget table the way prisma resolves the relation.""" + + def __init__(self, journal: list[str], budget_table: _FakeBudgetTable) -> None: + self.journal = journal + self.budget_table = budget_table + self.rows: dict[str, _FakeAccessGroupBudgetRow] = {} + self.upsert_calls: list[dict] = [] + + def _resolve(self, row, include): + if row is None: + return None + row.litellm_budget_table = ( + self.budget_table.rows.get(row.budget_id) if include and row.budget_id is not None else None + ) + return row + + async def find_unique(self, where, include=None): + return self._resolve(self.rows.get(where["access_group_name"]), include) + + async def upsert(self, where, data, include=None): + access_group_name = where["access_group_name"] + self.upsert_calls.append(dict(data)) + self.journal.append(f"access_group_budget.upsert:{access_group_name}") + existing = self.rows.get(access_group_name) + payload = data["update"] if existing is not None else data["create"] + row = existing or _FakeAccessGroupBudgetRow(access_group_name=access_group_name) + row.budget_id = payload.get("budget_id") + self.rows[access_group_name] = row + return self._resolve(row, include) + + async def delete(self, where, include=None): + access_group_name = where["access_group_name"] + self.journal.append(f"access_group_budget.delete:{access_group_name}") + return self.rows.pop(access_group_name, None) + + +class _FakeModelTable: + def __init__(self, journal: list[str], deployments) -> None: + self.journal = journal + self.deployments = list(deployments) + self.updates: list[tuple[dict, dict]] = [] + + async def find_many(self, where=None, **kwargs): + return list(self.deployments) + + async def find_unique(self, where, include=None): + return next((d for d in self.deployments if d.model_id == where["model_id"]), None) + + async def update(self, where, data, include=None): + self.journal.append(f"model_table.update:{where['model_id']}") + self.updates.append((dict(where), dict(data))) + return None + + +class _FakePrismaClient: + def __init__(self, journal: list[str], deployments=()) -> None: + self.budget_table = _FakeBudgetTable(journal) + self.access_group_budget_table = _FakeAccessGroupBudgetTable(journal, self.budget_table) + self.model_table = _FakeModelTable(journal, deployments) + self.db = SimpleNamespace( + litellm_budgettable=self.budget_table, + litellm_modelaccessgroupbudgettable=self.access_group_budget_table, + litellm_proxymodeltable=self.model_table, + ) + + def jsonify_object(self, data): + return dict(data) + + +class _FakeAuthCache: + """Spy for the auth cache the endpoints evict through. Injected into the endpoint rather than + patched over the proxy_server global, so dropping the eviction call fails a test.""" + + def __init__(self, journal: list[str] | None = None) -> None: + self.journal = journal if journal is not None else [] + self.deleted_keys: list[str] = [] + + async def async_delete_cache(self, key): + self.deleted_keys.append(key) + self.journal.append(f"auth_cache.delete:{key}") + + +def _admin(): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + return UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + +def _deployment(model_id="deploy-1", model_name="gpt-4o", access_groups=("prod-models",)): + return _FakeDeployment( + model_id=model_id, + model_name=model_name, + model_info={"access_groups": list(access_groups)}, + ) + + +def _seed_budget(prisma, access_group, spend=0.0, budget_id="budget-seed", **budget_fields): + prisma.budget_table.rows[budget_id] = _FakeBudgetRow(budget_id=budget_id, **budget_fields) + prisma.access_group_budget_table.rows[access_group] = _FakeAccessGroupBudgetRow( + access_group_name=access_group, + budget_id=budget_id, + spend=spend, + ) + + +@contextmanager +def _proxy(prisma): + with patch( # test-quality-ok: the endpoints import proxy_server.prisma_client themselves; no parameter to inject + "litellm.proxy.proxy_server.prisma_client", prisma + ): + yield + + +@contextmanager +def _proxy_with_stubbed_reload(prisma): + """delete_access_group finishes by reloading the router and judging what it serves afterwards. + Both collaborators it reaches for there are module globals it imports itself, so a fake can only + get in by patching them; auth_cache and prisma are the ones with a real seam.""" + never_served_router = MagicMock() + never_served_router.get_model_ids.return_value = [] + with ( + _proxy(prisma), + patch( # test-quality-ok: live_model_ids_snapshot() reads the llm_router global; the endpoint takes no router + "litellm.proxy.proxy_server.llm_router", never_served_router + ), + patch( # test-quality-ok: the endpoint calls its module-level clear_cache import; there is no parameter for it + "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + yield + + +def _eviction_journal(access_group): + """Both auth cache keys, in the order a write path has to evict them.""" + from litellm.proxy.common_utils.user_api_key_cache import ( + model_access_group_cache_key, + model_access_group_registry_cache_key, + ) + + return [ + f"auth_cache.delete:{model_access_group_cache_key(access_group)}", + f"auth_cache.delete:{model_access_group_registry_cache_key()}", + ] + + +def _assert_evicted_after_write(journal, access_group, write_entry): + """Exactly the two keys, in order, after the DB write. Deliberately not a tail slice: what + has to hold is that the eviction follows the write, not that nothing follows the eviction.""" + evictions = [entry for entry in journal if entry.startswith("auth_cache.delete:")] + assert evictions == _eviction_journal(access_group) + assert journal.index(write_entry) < journal.index(evictions[0]) + + +@pytest.mark.asyncio +async def test_put_access_group_budget_creates_the_row_and_its_budget(): + """First PUT has to create both halves: the budget row it links, and the access group row + that carries the link and the shared spend.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + + with _proxy(prisma): + response = await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=100.0, soft_budget=80.0, budget_duration="30d"), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert response.access_group == "prod-models" + assert response.spend == 0.0 + assert response.budget is not None + assert response.budget.max_budget == 100.0 + assert response.budget.soft_budget == 80.0 + assert response.budget.budget_duration == "30d" + assert len(prisma.budget_table.create_calls) == 1 + assert prisma.access_group_budget_table.rows["prod-models"].budget_id == response.budget.budget_id + + +@pytest.mark.asyncio +async def test_second_put_replaces_the_budget_instead_of_creating_another(): + """PUT is idempotent: a second call must update the budget already linked to the group, + not leave a second budget row (and a second group row) behind.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + + with _proxy(prisma): + first = await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=100.0), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + second = await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=250.0), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert first.budget is not None and second.budget is not None + assert second.budget.budget_id == first.budget.budget_id + assert second.budget.max_budget == 250.0 + assert len(prisma.budget_table.create_calls) == 1 + assert len(prisma.budget_table.rows) == 1 + assert len(prisma.access_group_budget_table.rows) == 1 + assert prisma.budget_table.update_calls[-1][0] == first.budget.budget_id + + +@pytest.mark.asyncio +async def test_put_access_group_budget_links_an_existing_budget_without_creating_one(): + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + prisma.budget_table.rows["shared-budget"] = _FakeBudgetRow(budget_id="shared-budget", max_budget=7.0) + cache = _FakeAuthCache() + + with _proxy(prisma): + response = await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(budget_id="shared-budget"), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert prisma.budget_table.create_calls == [] + assert response.budget is not None + assert response.budget.budget_id == "shared-budget" + assert response.budget.max_budget == 7.0 + assert prisma.access_group_budget_table.rows["prod-models"].budget_id == "shared-budget" + + +@pytest.mark.asyncio +async def test_put_access_group_budget_rejects_an_empty_body(): + """An empty PUT would register the group as budgeted while enforcing nothing.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + + with _proxy(prisma), pytest.raises(HTTPException) as exc_info: + await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert exc_info.value.status_code == 400 + assert prisma.access_group_budget_table.rows == {} + assert cache.deleted_keys == [] + + +@pytest.mark.asyncio +async def test_put_access_group_budget_rejects_an_unparseable_duration(): + """An unparseable duration can only be discovered by the reset job, long after the write.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + + with _proxy(prisma), pytest.raises(HTTPException) as exc_info: + await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=10.0, budget_duration="every other tuesday"), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert exc_info.value.status_code == 400 + assert prisma.budget_table.create_calls == [] + assert prisma.access_group_budget_table.rows == {} + + +def test_access_group_budget_request_rejects_rate_limit_fields(): + """tpm/rpm/max_parallel_requests are not enforced per access group, so accepting them would + promise rate limiting that never happens.""" + from pydantic import ValidationError + + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + for unsupported in ({"tpm_limit": 10}, {"rpm_limit": 10}, {"max_parallel_requests": 10}): + with pytest.raises(ValidationError): + AccessGroupBudgetRequest(max_budget=1.0, **unsupported) + + +@pytest.mark.asyncio +async def test_get_access_group_budget_returns_the_budget_and_the_shared_spend(): + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + get_access_group_budget, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=42.5, max_budget=100.0, budget_duration="30d") + + with _proxy(prisma): + response = await get_access_group_budget(access_group="prod-models") + + assert response.access_group == "prod-models" + assert response.spend == 42.5 + assert response.budget is not None + assert response.budget.max_budget == 100.0 + assert response.budget.budget_duration == "30d" + + +@pytest.mark.asyncio +async def test_get_access_group_budget_on_a_budgetless_group_is_200_not_404(): + """A real group that simply has no budget is not an error; only an unknown group is.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + get_access_group_budget, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + + with _proxy(prisma): + response = await get_access_group_budget(access_group="prod-models") + + assert response.spend == 0.0 + assert response.budget is None + + +@pytest.mark.asyncio +async def test_access_group_budget_routes_404_on_an_unknown_group(): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group_budget, + get_access_group_budget, + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + admin = _admin() + + calls = ( + lambda: get_access_group_budget(access_group="ghost-group"), + lambda: set_access_group_budget( + access_group="ghost-group", + data=AccessGroupBudgetRequest(max_budget=1.0), + user_api_key_dict=admin, + auth_cache=cache, + ), + lambda: delete_access_group_budget(access_group="ghost-group", auth_cache=cache), + ) + + with _proxy(prisma): + for make_call in calls: + with pytest.raises(HTTPException) as exc_info: + await make_call() + assert exc_info.value.status_code == 404 + + assert prisma.budget_table.create_calls == [] + assert prisma.access_group_budget_table.rows == {} + + +@pytest.mark.asyncio +async def test_delete_access_group_budget_drops_the_row_and_spares_the_shared_budget(): + """The group row goes; the LiteLLM_BudgetTable row it linked survives, as /tag/delete leaves a + tag's. That row can be shared, so deleting it would be data loss for whatever else points at it.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group_budget, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=12.0, max_budget=100.0) + cache = _FakeAuthCache() + + with _proxy(prisma): + response = await delete_access_group_budget(access_group="prod-models", auth_cache=cache) + + assert response.budget_deleted is True + assert prisma.access_group_budget_table.rows == {} + assert prisma.budget_table.deleted_ids == [] + assert prisma.budget_table.rows["budget-seed"].max_budget == 100.0 + + +@pytest.mark.asyncio +async def test_delete_access_group_budget_on_a_budgetless_group_still_evicts(): + """budget_deleted is False, but the group can still be sitting in the cached registry of + budgeted groups, so the eviction has to run whether or not a row was there to drop.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group_budget, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + cache = _FakeAuthCache(journal) + + with _proxy(prisma): + response = await delete_access_group_budget(access_group="prod-models", auth_cache=cache) + + assert response.budget_deleted is False + assert prisma.budget_table.deleted_ids == [] + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") + + +@pytest.mark.asyncio +async def test_deleting_the_access_group_strips_deployments_before_dropping_the_budget(): + """Ordering is the point: stripping first means a failure leaves an unreachable budget row, + while the reverse leaves a live group whose enforcement silently vanished. The shared + LiteLLM_BudgetTable row survives here too.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=3.0, max_budget=100.0) + cache = _FakeAuthCache() + + with _proxy_with_stubbed_reload(prisma): + response = await delete_access_group( + access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache + ) + + assert response.models_updated == 1 + assert prisma.access_group_budget_table.rows == {} + assert prisma.budget_table.deleted_ids == [] + assert prisma.budget_table.rows["budget-seed"].max_budget == 100.0 + assert journal.index("model_table.update:deploy-1") < journal.index("access_group_budget.delete:prod-models") + + +@pytest.mark.asyncio +async def test_access_group_info_surfaces_the_budget_and_spend(): + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + get_access_group_info, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=9.5, max_budget=100.0, soft_budget=50.0) + + with _proxy(prisma): + info = await get_access_group_info(access_group="prod-models", user_api_key_dict=_admin()) + + assert info.model_names == ["gpt-4o"] + assert info.spend == 9.5 + assert info.budget is not None + assert info.budget.max_budget == 100.0 + assert info.budget.soft_budget == 50.0 + + +@pytest.mark.asyncio +async def test_put_access_group_budget_evicts_both_auth_cache_keys(): + """Auth reads the per-group row and the registry of budgeted groups cache-first with no + freshness check, so a PUT that skips either eviction returns 200 and enforces nothing until + the TTL expires. Both keys, after the write.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + cache = _FakeAuthCache(journal) + + with _proxy(prisma): + await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=100.0), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.upsert:prod-models") + + +@pytest.mark.asyncio +async def test_delete_access_group_budget_evicts_both_auth_cache_keys(): + """Clearing a budget has the same window as setting one: until both keys are dropped, auth + keeps enforcing the budget that is already gone.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group_budget, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=12.0, max_budget=100.0) + cache = _FakeAuthCache(journal) + + with _proxy(prisma): + await delete_access_group_budget(access_group="prod-models", auth_cache=cache) + + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") + + +@pytest.mark.asyncio +async def test_deleting_the_access_group_evicts_both_auth_cache_keys(): + """The group-delete cascade drops the budget row too, so it owes the same two evictions.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=3.0, max_budget=100.0) + cache = _FakeAuthCache(journal) + + with _proxy_with_stubbed_reload(prisma): + await delete_access_group(access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache) + + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") + + +@pytest.mark.asyncio +async def test_deleting_an_access_group_that_never_had_a_budget_still_evicts(): + """The cascade's delete finds no row and reports nothing dropped, but the group can still be + sitting in the cached registry of budgeted groups, so both keys have to go regardless.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + cache = _FakeAuthCache(journal) + + with _proxy_with_stubbed_reload(prisma): + response = await delete_access_group( + access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache + ) + + assert response.models_updated == 1 + assert prisma.access_group_budget_table.rows == {} + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") 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/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 51980342a1d..fb3de990deb 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -271,6 +271,81 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch): ) +def _make_window_spend_prisma(row=None, spend_logs_total=0.0): + prisma = MagicMock() + prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=row) + prisma.db.litellm_spendlogs.group_by = AsyncMock( + return_value=[{"api_key": "tok", "_sum": {"spend": spend_logs_total}}] + ) + return prisma + + +@pytest.mark.asyncio +async def test_get_current_spend_floors_window_against_maintained_row(monkeypatch): + """The floor re-check runs every few seconds per pod, so the window branch + must read the maintained row and leave the unindexed spend-logs scan alone.""" + from datetime import timezone + from types import SimpleNamespace + + window_start = datetime(2026, 1, 1, tzinfo=timezone.utc) + fake_prisma = _make_window_spend_prisma( + row=SimpleNamespace(window_start=window_start, spend=15.0), + spend_logs_total=100.0, + ) + fake_cache = _make_spend_counter_cache(redis_get_value=2.0) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + + counter_key = "spend:key:tok:window:7d" + result = await ps.get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + max_budget=10.0, + window_entity_type="Key", + window_entity_id="tok", + window_duration="7d", + window_start=window_start, + ) + + assert result == 15.0 + fake_prisma.db.litellm_spendlogs.group_by.assert_not_awaited() + fake_cache.redis_cache.async_set_max.assert_awaited_once_with( + key=counter_key, value=15.0 + ) + + +@pytest.mark.asyncio +async def test_get_current_spend_floors_window_against_logs_when_row_stale(monkeypatch): + """A row left behind at a crossed window boundary must not be read as the + current window's spend; the aggregate stays the fallback.""" + from datetime import timedelta, timezone + from types import SimpleNamespace + + window_start = datetime(2026, 1, 8, tzinfo=timezone.utc) + fake_prisma = _make_window_spend_prisma( + row=SimpleNamespace( + window_start=window_start - timedelta(days=7), spend=999.0 + ), + spend_logs_total=15.0, + ) + fake_cache = _make_spend_counter_cache(redis_get_value=2.0) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + + result = await ps.get_current_spend( + counter_key="spend:key:tok:window:7d", + fallback_spend=0.0, + max_budget=10.0, + window_entity_type="Key", + window_entity_id="tok", + window_duration="7d", + window_start=window_start, + ) + + assert result == 15.0 + fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once() + + @pytest.mark.asyncio async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypatch): """With fail_closed_budget_enforcement on, an admit decision backed only by a @@ -895,6 +970,7 @@ async def test_init_and_increment_window_spend_counter_increments_when_initializ counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=datetime(2024, 1, 1), increment=5.0, ) @@ -922,6 +998,7 @@ async def test_init_and_increment_window_spend_counter_missing_window_start_inva counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=None, increment=5.0, ) @@ -1059,6 +1136,7 @@ async def test_ensure_window_spend_counter_initialized_warm_returns_true(monkeyp counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=datetime(2024, 1, 1), ) @@ -1091,6 +1169,7 @@ async def test_ensure_window_spend_counter_initialized_db_failure_invalid_return counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=datetime(2024, 1, 1), ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py index e8ca569763d..5fda4fb20b5 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_savings.py +++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py @@ -1151,6 +1151,61 @@ def test_logging_payload_never_stamps_internal_calls(): assert internal is None +def test_savings_are_net_of_a_priced_classifier(): + """The classifier call is part of what routing cost, so the per-request figure + deducts it; a charge big enough to outweigh the model saving goes negative, + since the figure is signed on purpose (GH #38816).""" + from litellm.proxy.spend_tracking.savings import autorouter_savings_for_request + + gross = autorouter_savings_for_request( + model="claude-haiku-4-5", + custom_llm_provider="anthropic", + routing_decision=_routed_decision(), + usage_object=_cached_usage_object(), + ) + net = autorouter_savings_for_request( + model="claude-haiku-4-5", + custom_llm_provider="anthropic", + routing_decision={**_routed_decision(), "classifier_cost": 0.005}, + usage_object=_cached_usage_object(), + ) + assert gross is not None and net == pytest.approx(gross - 0.005) + + +@pytest.mark.parametrize("classifier_cost", [0.0, "bogus", True]) +def test_an_unpriced_classifier_deducts_nothing(classifier_cost: object): + from litellm.proxy.spend_tracking.savings import autorouter_savings_for_request + + gross = autorouter_savings_for_request( + model="claude-haiku-4-5", + custom_llm_provider="anthropic", + routing_decision=_routed_decision(), + usage_object=_cached_usage_object(), + ) + with_cost_field = autorouter_savings_for_request( + model="claude-haiku-4-5", + custom_llm_provider="anthropic", + routing_decision={**_routed_decision(), "classifier_cost": classifier_cost}, + usage_object=_cached_usage_object(), + ) + assert with_cost_field == gross + + +def test_recorded_savings_are_already_net_and_not_deducted_again(): + """The deduction lives at the figure's computation owner, so a stamped figure is + net by construction; the recorded-wins path must not subtract a second time.""" + result = compute_savings_spend( + model="claude-haiku-4-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + gateway_injected_cache=False, + routing_decision={**_routed_decision(), "classifier_cost": 0.005}, + usage_object=_cached_usage_object(), + recorded_autorouter_savings=0.5, + ) + assert result.autorouter == 0.5 + + def test_caching_savings_require_a_gateway_injected_breakpoint(): """The same cached usage is attributed to the gateway only when it added a breakpoint. 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_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 38a346e7fb7..95067929ac1 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2,6 +2,7 @@ import asyncio import threading from collections.abc import Mapping from datetime import datetime, timedelta, timezone +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -10,15 +11,16 @@ from fastapi import HTTPException import litellm from litellm.caching.dual_cache import DualCache from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES -from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( - AnthropicMessagesStreamingResponse, -) from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( AgenticAnthropicStreamingIterator, ) +from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + AnthropicMessagesStreamingResponse, +) from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_EndUserTable, + Litellm_EntityType, LiteLLM_OrganizationTable, LiteLLM_TagTable, LiteLLM_TeamMembership, @@ -27,9 +29,16 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + model_access_group_cache_key, + model_access_group_spend_counter_key, +) from litellm.proxy.spend_tracking.budget_reservation import ( TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS, _approximate_input_size, + _get_model_access_group_budget_counters, estimate_request_max_cost, get_budget_window_start, invalidate_budget_reservation_counters, @@ -39,6 +48,7 @@ from litellm.proxy.spend_tracking.budget_reservation import ( ) from litellm.proxy.utils import ProxyLogging from litellm.router import Router +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget @pytest.fixture() @@ -2962,3 +2972,214 @@ async def test_small_prompt_is_tokenized_inline(spend_counter_state): assert reservation is not None assert threads == [threading.main_thread()] + + +class _ModelAccessGroupBudgetPrisma: + """Serves ``LiteLLM_ModelAccessGroupBudgetTable`` rows, recording what reached the database.""" + + def __init__(self, **max_budget_by_group) -> None: + self.rows = { + group: SimpleNamespace( + access_group_name=group, + spend=7.0, + litellm_budget_table=None if max_budget is None else SimpleNamespace(max_budget=max_budget), + ) + for group, max_budget in max_budget_by_group.items() + } + self.batches = [] + self.db = SimpleNamespace( + litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many) + ) + + async def _find_many(self, **kwargs): + requested = list(kwargs["where"]["access_group_name"]["in"]) + self.batches.append(requested) + return [self.rows[group] for group in requested if group in self.rows] + + +async def _model_access_group_counters(matched, **max_budget_by_group): + return await _get_model_access_group_budget_counters( + valid_token=UserAPIKeyAuth(api_key="hashed", matched_model_access_groups=matched), + prisma_client=_ModelAccessGroupBudgetPrisma(**max_budget_by_group), + user_api_key_cache=UserApiKeyCache(), + ) + + +@pytest.mark.asyncio +async def test_model_access_group_with_a_budget_reserves_against_the_reset_jobs_counter_key(): + counters = await _model_access_group_counters(["premium"], premium=25.0) + + assert len(counters) == 1 + counter = counters[0] + assert counter.counter_key == _model_access_group_counter_key(SimpleNamespace(access_group_name="premium")) + assert counter.source_cache_key == model_access_group_cache_key("premium") + assert counter.max_budget == 25.0 + assert counter.fallback_spend == 7.0 + assert counter.entity_type == "Model access group" + assert counter.entity_id == "premium" + + +@pytest.mark.asyncio +async def test_model_access_group_without_a_budget_reserves_nothing(): + assert await _model_access_group_counters(["premium"], premium=None) == [] + + +@pytest.mark.asyncio +async def test_model_access_group_with_a_zero_budget_reserves_nothing(): + """Zero is how a budget is cleared, not a ceiling that blocks every request.""" + assert await _model_access_group_counters(["premium"], premium=0.0) == [] + + +@pytest.mark.asyncio +async def test_model_access_group_counters_come_from_the_auth_object(): + """Auth already resolved which granted groups serve the model; re-deriving it here would drift.""" + assert await _model_access_group_counters(None, premium=25.0) == [] + + +@pytest.mark.asyncio +async def test_repeated_model_access_group_reserves_once(): + counters = await _model_access_group_counters(["premium", "premium"], premium=25.0) + + assert [counter.entity_id for counter in counters] == ["premium"] + + +@pytest.mark.asyncio +async def test_model_access_group_counter_blocks_a_request_over_the_group_budget(spend_counter_state): + """End to end through the reservation path, which is what runs when reservations are enabled.""" + counter_cache, key_cache = spend_counter_state + prisma_client = _ModelAccessGroupBudgetPrisma(premium=1.0) + valid_token = UserAPIKeyAuth(api_key="hashed", token="tok", matched_model_access_groups=["premium"]) + + with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.5, + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=prisma_client, + user_api_key_cache=key_cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=key_cache), + ) + + assert exc_info.value.entity_id == "premium" + assert exc_info.value.entity_type == Litellm_EntityType.MODEL_ACCESS_GROUP.value + + +async def _cache_model_access_group_budget(key_cache, group, spend, max_budget=None): + await key_cache.async_set_cache( + key=model_access_group_cache_key(group), + value=ModelAccessGroupBudget(access_group_name=group, spend=spend, max_budget=max_budget), + model_type=ModelAccessGroupBudget, + ) + + +async def _reserve_for_model_access_groups(key_cache, groups, estimate): + """Reserve against the given groups, whose rows are already cached, so nothing hits the DB.""" + with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=estimate, + ): + return await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=UserAPIKeyAuth( + api_key="hashed", token="tok-mag-counter", matched_model_access_groups=list(groups) + ), + team_object=None, + user_object=None, + prisma_client=_ModelAccessGroupBudgetPrisma(), + user_api_key_cache=key_cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=key_cache), + ) + + +@pytest.mark.asyncio +async def test_model_access_group_counter_accumulates_across_calls_without_a_reservation(spend_counter_state): + """With reservations disabled nothing writes the counter up front, so the cost callback must. + + Otherwise the read-time budget check enforces against the DB row's spend, which the cache + holds for the full TTL, and a caller runs past the ceiling for that whole window. + """ + counter_cache, key_cache = spend_counter_state + await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0) + + from litellm.proxy.proxy_server import increment_spend_counters + + counter_key = model_access_group_spend_counter_key("premium") + + await increment_spend_counters( + token=None, team_id=None, user_id=None, response_cost=0.25, model_access_groups=["premium"] + ) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.25) + + await increment_spend_counters( + token=None, team_id=None, user_id=None, response_cost=0.75, model_access_groups=["premium", "premium", ""] + ) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(2.0) + assert counter_cache.in_memory_cache.get_cache(key=model_access_group_spend_counter_key("")) is None + + +@pytest.mark.asyncio +async def test_reserved_model_access_group_is_not_charged_twice(spend_counter_state): + """The reservation already wrote this counter, so the post-call pass has to skip it.""" + counter_cache, key_cache = spend_counter_state + await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0) + + reservation = await _reserve_for_model_access_groups(key_cache, ["premium"], estimate=0.6) + counter_key = model_access_group_spend_counter_key("premium") + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.6) + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id=None, + user_id=None, + response_cost=0.2, + budget_reservation=reservation, + model_access_groups=["premium"], + ) + + # 1.0 recorded + the reservation reconciled down to the 0.2 actually spent. A second + # increment would land at 1.4. + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.2) + + +@pytest.mark.asyncio +async def test_unreserved_model_access_group_is_charged_alongside_a_reserved_one(spend_counter_state): + """A budgetless group reserves nothing, so only the post-call pass can charge it. + + Both groups authorized the request and both get debited, each exactly once, whether or not + the reservation path happened to hold a counter for them. + """ + counter_cache, key_cache = spend_counter_state + await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0) + await _cache_model_access_group_budget(key_cache, "starter", spend=4.0) + + reservation = await _reserve_for_model_access_groups(key_cache, ["premium", "starter"], estimate=0.6) + assert [entry["entity_id"] for entry in reservation["entries"]] == ["premium"] + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id=None, + user_id=None, + response_cost=0.2, + budget_reservation=reservation, + model_access_groups=["premium", "starter", "starter", "premium"], + ) + + assert counter_cache.in_memory_cache.get_cache( + key=model_access_group_spend_counter_key("premium") + ) == pytest.approx(1.2) + assert counter_cache.in_memory_cache.get_cache( + key=model_access_group_spend_counter_key("starter") + ) == pytest.approx(4.2) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 71d4666416d..dd5049b0f28 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1683,6 +1683,34 @@ class TestCommonRequestProcessingHelpers: assert _serialize_http_exception_detail(42) == ("42", None) + async def test_proxy_exception_from_http_exception_helper(self): + """The shared HTTPException -> ProxyException conversion keeps a clean + message, merges structured detail over existing provider_specific_fields, + and passes headers through.""" + from litellm.proxy.common_request_processing import ( + proxy_exception_from_http_exception, + ) + + exc = HTTPException( + status_code=400, + detail={"error": "Content blocked", "guardrail": "keyword-block"}, + ) + exc.provider_specific_fields = {"existing": "field", "guardrail": "stale"} + result = proxy_exception_from_http_exception(exc, {"x-litellm-call-id": "abc"}) + assert result.message == "Content blocked" + assert result.code == "400" + assert result.provider_specific_fields == { + "existing": "field", + "error": "Content blocked", + "guardrail": "keyword-block", + } + assert result.headers == {"x-litellm-call-id": "abc"} + + plain = proxy_exception_from_http_exception(HTTPException(status_code=429, detail="slow down"), {}) + assert plain.message == "slow down" + assert plain.code == "429" + assert plain.provider_specific_fields is None + async def test_create_streaming_response_first_chunk_error_string_code(self): """ Test that when the first chunk contains a string error code, a JSON error response is returned 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..ee0e2014951 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -33,6 +33,8 @@ from litellm.proxy.litellm_pre_call_utils import ( check_if_token_is_service_account, clean_headers, ) +from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, ) @@ -1073,6 +1075,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 +1110,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] @@ -7681,3 +7685,37 @@ async def test_add_litellm_data_to_request_keeps_litellm_metadata_on_litellm_met ) assert updated["litellm_metadata"]["trace_id"] == "abc" + + +def _stamp_model_access_groups(matched_model_access_groups, metadata_variable_name="metadata"): + user_api_key_dict = UserAPIKeyAuth(api_key="hashed-key") + user_api_key_dict.matched_model_access_groups = matched_model_access_groups + return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data={metadata_variable_name: {}}, + user_api_key_dict=user_api_key_dict, + _metadata_variable_name=metadata_variable_name, + )[metadata_variable_name] + + +def test_matched_model_access_groups_are_stamped_into_request_metadata(): + """The post-call spend writer reads the groups off request metadata, not off UserAPIKeyAuth.""" + stamped = _stamp_model_access_groups(["tier-a", "tier-b"]) + + assert stamped[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a", "tier-b"] + assert MODEL_ACCESS_GROUP_METADATA_KEY not in _stamp_model_access_groups(None) + + +def test_stamped_model_access_groups_survive_the_litellm_metadata_merge(): + """ + The key must keep its ``user_api_key`` prefix: when a request carries both metadata dicts, + get_litellm_metadata_from_kwargs returns litellm_metadata and copies a key over from metadata + only when that substring is in its name, so an unprefixed key is silently dropped. + """ + kwargs = { + "litellm_params": { + "metadata": _stamp_model_access_groups(["tier-a"]), + "litellm_metadata": {"trace_id": "abc"}, + } + } + + assert get_litellm_metadata_from_kwargs(kwargs)[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 949088ea3ba..1de3ed6e56d 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import importlib import json import os @@ -7816,6 +7817,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): counter_cache = DualCache() window_start = datetime.now(timezone.utc) - timedelta(hours=1) fake_prisma = MagicMock() + fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( return_value=[{"api_key": "key-window", "_sum": {"spend": 2.25}}] ) @@ -7830,6 +7832,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): counter_key="spend:key:key-window:window:1h", entity_type="Key", entity_id="key-window", + window_duration="1h", window_start=window_start, increment=0.5, ) @@ -7933,6 +7936,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() + fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( return_value=[{"api_key": "key-window-stale-local", "_sum": {"spend": 2.25}}] ) @@ -7947,6 +7951,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): counter_key=counter_key, entity_type="Key", entity_id="key-window-stale-local", + window_duration="1h", window_start=window_start, increment=0.5, ) @@ -7995,6 +8000,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() + fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( return_value=[{"api_key": "key-window-concurrent-seed", "_sum": {"spend": 2.25}}] ) @@ -8009,6 +8015,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() counter_key=counter_key, entity_type="Key", entity_id="key-window-concurrent-seed", + window_duration="1h", window_start=window_start, increment=0.5, ) @@ -8041,6 +8048,7 @@ async def test_window_spend_counter_skips_invalid_window_start(): counter_key="spend:key:key-invalid-window:window:not-a-duration", entity_type="Key", entity_id="key-invalid-window", + window_duration="not-a-duration", window_start=None, increment=0.5, ) @@ -8068,6 +8076,7 @@ async def test_window_spend_counter_does_not_seed_zero_when_db_unavailable(): counter_key=counter_key, entity_type="Key", entity_id="key-window-db-unavailable", + window_duration="1h", window_start=datetime.now(timezone.utc) - timedelta(hours=1), ) @@ -9565,6 +9574,76 @@ class TestDeleteDeploymentSync: assert result is None, f"Expected None on DB failure to signal fetch error, got {result!r}" + @pytest.mark.asyncio + async def test_get_models_from_db_reads_from_writer_not_replica(self): + """ + Regression for #38556: with DATABASE_URL_READ_REPLICA configured, the model + reconcile after /model/new used to read via the replica, so a lagging replica + made the reload miss the just-committed row and fail the request with a 500. + The reconcile read must be pinned to the writer. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.db.prisma_client import PrismaWrapper + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + from litellm.proxy.proxy_server import ProxyConfig + + writer_inner = MagicMock(name="writer_prisma") + reader_inner = MagicMock(name="reader_prisma") + committed_row = MagicMock(name="just_committed_model_row") + writer_inner.litellm_proxymodeltable.find_many = AsyncMock(return_value=[committed_row]) + reader_inner.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + mock_prisma = MagicMock() + mock_prisma.db = RoutingPrismaWrapper( + writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False), + reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False), + ) + + result = await ProxyConfig()._get_models_from_db(prisma_client=mock_prisma) + + assert result == [committed_row], f"Expected the writer's just-committed row, got {result!r}" + reader_inner.litellm_proxymodeltable.find_many.assert_not_awaited() + + @pytest.mark.asyncio + async def test_get_models_from_db_falls_back_to_replica_when_writer_down(self): + """ + The writer pin must not break reader-only degraded mode: a proxy that + starts during a primary outage (writer connect failed, replica healthy) + must still load DB-backed models through the replica instead of sending + the reconcile read to the unavailable writer. + """ + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.db.prisma_client import PrismaWrapper + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + from litellm.proxy.proxy_server import ProxyConfig + + writer_inner = MagicMock(name="writer_prisma") + reader_inner = MagicMock(name="reader_prisma") + replica_row = MagicMock(name="replica_model_row") + writer_inner.litellm_proxymodeltable = SimpleNamespace( + find_many=AsyncMock(side_effect=RuntimeError("writer unreachable")), + create=MagicMock(name="writer_create"), + ) + reader_inner.litellm_proxymodeltable = SimpleNamespace( + find_many=AsyncMock(return_value=[replica_row]), + create=MagicMock(name="reader_create"), + ) + + mock_prisma = MagicMock() + mock_prisma.db = RoutingPrismaWrapper( + writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False), + reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False), + ) + mock_prisma.db._writer_unavailable = True + + result = await ProxyConfig()._get_models_from_db(prisma_client=mock_prisma) + + assert result == [replica_row], f"Expected the replica's rows in degraded mode, got {result!r}" + writer_inner.litellm_proxymodeltable.find_many.assert_not_awaited() + def test_get_config_list_includes_cancel_on_disconnect(monkeypatch): """Follow-up to #30223: the flag must be discoverable via /config/list, @@ -11130,6 +11209,289 @@ def test_startup_is_silent_when_mock_testing_params_disabled(caplog): assert MOCK_TESTING_CONFIG_KEY not in caplog.text +# --------------------------------------------------------------------------- +# Budget window spend row enqueue (LiteLLM_BudgetWindowSpend writer) +# --------------------------------------------------------------------------- + + +@contextlib.contextmanager +def _window_spend_enqueue_env(cached_objects: dict): + """Point increment_spend_counters at throwaway caches and a real + WindowSpendUpdateQueue, and hand back the queue to inspect.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendUpdateQueue, + ) + import litellm.proxy.proxy_server as ps + + user_api_key_cache = MagicMock() + user_api_key_cache.async_get_cache = AsyncMock(side_effect=lambda key, **_: cached_objects.get(key)) + + queue = WindowSpendUpdateQueue() + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.window_spend_update_queue = queue + + originals = ( + ps.user_api_key_cache, + ps.spend_counter_cache, + ps.prisma_client, + ps.proxy_logging_obj, + ) + ps.user_api_key_cache = user_api_key_cache + ps.spend_counter_cache = DualCache() + ps.prisma_client = None + ps.proxy_logging_obj = proxy_logging_obj + try: + yield queue + finally: + ( + ps.user_api_key_cache, + ps.spend_counter_cache, + ps.prisma_client, + ps.proxy_logging_obj, + ) = originals + + +async def _drain(queue): + return list(await queue.flush_and_get_aggregated_window_spend_transactions()) + + +@pytest.mark.asyncio +async def test_key_window_spend_row_is_enqueued_with_the_actual_cost(): + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert len(enqueued) == 1 + assert enqueued[0]["entity_type"] == "key" + assert enqueued[0]["entity_id"] == "hashed-token" + assert enqueued[0]["window_duration"] == "30d" + assert enqueued[0]["spend"] == pytest.approx(0.25) + assert enqueued[0]["window_start"] == (reset_at - timedelta(days=30)).astimezone(timezone.utc).replace( + tzinfo=None + ).isoformat(timespec="microseconds") + + +@pytest.mark.asyncio +async def test_team_window_spend_row_is_enqueued(): + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=3) + team_obj = MagicMock() + team_obj.budget_limits = [ + {"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"team_id:team-1": team_obj}) as queue: + await increment_spend_counters( + token=None, team_id="team-1", user_id=None, response_cost=1.5 + ) + enqueued = await _drain(queue) + + assert len(enqueued) == 1 + assert enqueued[0]["entity_type"] == "team" + assert enqueued[0]["entity_id"] == "team-1" + assert enqueued[0]["window_duration"] == "7d" + assert enqueued[0]["spend"] == pytest.approx(1.5) + + +@pytest.mark.asyncio +async def test_window_spend_row_is_enqueued_even_when_the_counter_was_reserved(): + """A reservation only pre-charged the cache counter with an estimate; the + row still owes the actual cost, so the enqueue must not be skipped.""" + from litellm.proxy.proxy_server import increment_spend_counters + import litellm.proxy.spend_tracking.budget_reservation as br + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + reservation = { + "entries": [ + {"counter_key": "spend:key:hashed-token", "reserved": 1.0}, + {"counter_key": "spend:key:hashed-token:window:30d", "reserved": 1.0}, + ] + } + + original_reconcile = br.reconcile_budget_reservation + br.reconcile_budget_reservation = AsyncMock(return_value=None) + try: + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", + team_id=None, + user_id=None, + response_cost=0.25, + budget_reservation=reservation, + ) + enqueued = await _drain(queue) + finally: + br.reconcile_budget_reservation = original_reconcile + + assert len(enqueued) == 1 + assert enqueued[0]["spend"] == pytest.approx(0.25) + + +@pytest.mark.asyncio +async def test_sliding_window_without_reset_at_is_not_enqueued(): + """Windows with no reset_at slide with wall clock, so window_start moves on + every request and no single row can represent them; the read path keeps + using its LiteLLM_SpendLogs fallback instead.""" + from litellm.proxy.proxy_server import increment_spend_counters + + key_obj = MagicMock() + key_obj.budget_limits = [{"budget_duration": "30d", "max_budget": 100.0}] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert enqueued == [] + + +@pytest.mark.asyncio +async def test_each_configured_window_gets_its_own_row_enqueue(): + from litellm.proxy.proxy_server import increment_spend_counters + + now = datetime.now(timezone.utc) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "1d", "max_budget": 5.0, "reset_at": (now + timedelta(hours=5)).isoformat()}, + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": (now + timedelta(days=10)).isoformat()}, + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert sorted(item["window_duration"] for item in enqueued) == ["1d", "30d"] + assert all(item["spend"] == pytest.approx(0.25) for item in enqueued) + + +@pytest.mark.asyncio +async def test_no_window_spend_row_enqueued_without_budget_limits(): + from litellm.proxy.proxy_server import increment_spend_counters + + key_obj = MagicMock() + key_obj.budget_limits = None + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert enqueued == [] + + +@pytest.mark.asyncio +async def test_window_spend_row_carries_the_spend_log_request_id(): + """The flush excludes these ids from its one-time seed, so the id threaded + here has to be the same one the LiteLLM_SpendLogs row was written under.""" + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", + team_id=None, + user_id=None, + response_cost=0.25, + request_id="chatcmpl-abc123", + ) + enqueued = await _drain(queue) + + assert enqueued[0]["request_ids"] == ("chatcmpl-abc123",) + + +@pytest.mark.asyncio +async def test_window_spend_row_carries_the_request_start_time(): + """The seed only excludes a batch id whose LiteLLM_SpendLogs.startTime is at + or after this, so it must be the same start the spend log was written with.""" + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", + team_id=None, + user_id=None, + response_cost=0.25, + request_id="chatcmpl-abc123", + request_started_at=datetime(2026, 8, 10, 12, 0, 0, 500_000, tzinfo=timezone.utc), + ) + enqueued = await _drain(queue) + + assert enqueued[0]["started_at"] == "2026-08-10T12:00:00.500000" + + +@pytest.mark.asyncio +async def test_window_spend_row_without_a_request_id_excludes_nothing(): + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert enqueued[0]["request_ids"] == () + + +@pytest.mark.asyncio +async def test_team_window_spend_row_carries_the_request_id(): + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=3) + team_obj = MagicMock() + team_obj.budget_limits = [ + {"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"team_id:team-1": team_obj}) as queue: + await increment_spend_counters( + token=None, + team_id="team-1", + user_id=None, + response_cost=1.5, + request_id="chatcmpl-team", + ) + enqueued = await _drain(queue) + + assert enqueued[0]["request_ids"] == ("chatcmpl-team",) + + def _mock_startup_prisma_client(health_check_error=None, connect_error=None): client = MagicMock() client.connect = AsyncMock(side_effect=connect_error) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 6920cc0dae3..dcaad968663 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1266,13 +1266,14 @@ class TestCreateSmtpConnection: patch("smtplib.SMTP_SSL") as mock_smtp_ssl, patch("smtplib.SMTP") as mock_smtp, ): - result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=465) + result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=465, timeout=30.0) mock_smtp.assert_not_called() assert result is mock_smtp_ssl.return_value _, kwargs = mock_smtp_ssl.call_args assert kwargs["host"] == "mail.example.com" assert kwargs["port"] == 465 + assert kwargs["timeout"] == 30.0 context = kwargs["context"] assert isinstance(context, ssl.SSLContext) assert context.verify_mode == ssl.CERT_REQUIRED @@ -1286,11 +1287,11 @@ class TestCreateSmtpConnection: patch("smtplib.SMTP_SSL") as mock_smtp_ssl, patch("smtplib.SMTP") as mock_smtp, ): - result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=587) + result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=587, timeout=30.0) mock_smtp_ssl.assert_not_called() assert result is mock_smtp.return_value - mock_smtp.assert_called_once_with(host="mail.example.com", port=587) + mock_smtp.assert_called_once_with(host="mail.example.com", port=587, timeout=30.0) class TestSendEmailStartTls: diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py index 19abcb5d66d..fce51c9296c 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py @@ -14,6 +14,7 @@ from __future__ import annotations import asyncio import sys +import threading from dataclasses import dataclass, field from email.message import EmailMessage from pathlib import Path @@ -320,6 +321,7 @@ class _SentMessage: body: Optional[str] starttls_called: bool login_args: Optional[tuple] + thread_ident: int @dataclass @@ -328,6 +330,7 @@ class InMemorySMTP: sent: List[_SentMessage] = field(default_factory=list) raise_on_send: Optional[Exception] = None + connection_kwargs: List[Dict[str, Any]] = field(default_factory=list) def server_factory(self) -> Callable[..., Any]: outer = self @@ -370,10 +373,12 @@ class InMemorySMTP: body=body, starttls_called=self._starttls_called, login_args=self._login_args, + thread_ident=threading.get_ident(), ) ) def _factory(*args: Any, **kwargs: Any) -> _Conn: + outer.connection_kwargs.append(dict(kwargs)) return _Conn() return _factory diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_send_email.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_send_email.py index 739e942de52..0e8aba0a03b 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_send_email.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_send_email.py @@ -6,6 +6,7 @@ Symbols pinned here: from __future__ import annotations +import threading from typing import Any import pytest @@ -51,9 +52,7 @@ async def test_send_email_dispatches_via_smtp(in_memory_smtp: Any) -> None: @pytest.mark.asyncio -async def test_send_email_starttls_uses_ssl( - in_memory_smtp: Any, monkeypatch: pytest.MonkeyPatch -) -> None: +async def test_send_email_starttls_uses_ssl(in_memory_smtp: Any, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("SMTP_USE_SSL", "True") await send_email( receiver_email="to@invalid", @@ -82,9 +81,7 @@ async def test_send_email_error_missing_sender_email( ) -> None: monkeypatch.delenv("SMTP_SENDER_EMAIL", raising=False) with pytest.raises(ValueError, match="SMTP_SENDER_EMAIL"): - await send_email( - receiver_email="x@y", subject="s", html="

h

" - ) + await send_email(receiver_email="x@y", subject="s", html="

h

") @pytest.mark.asyncio @@ -105,6 +102,49 @@ async def test_send_email_error_missing_html() -> None: await send_email(receiver_email="x@y", subject="s", html=None) +@pytest.mark.asyncio +async def test_send_email_sets_connection_timeout(in_memory_smtp: Any) -> None: + await send_email( + receiver_email="to@invalid", + subject="Hi", + html="

x

", + ) + assert in_memory_smtp.connection_kwargs[0].get("timeout") == 30.0 + + +@pytest.mark.asyncio +async def test_send_email_timeout_env_override(in_memory_smtp: Any, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SMTP_TIMEOUT", "5") + monkeypatch.setenv("SMTP_USE_SSL", "True") + await send_email( + receiver_email="to@invalid", + subject="Hi", + html="

x

", + ) + assert in_memory_smtp.connection_kwargs[0].get("timeout") == 5.0 + + +@pytest.mark.asyncio +async def test_send_email_malformed_timeout_is_swallowed(in_memory_smtp: Any, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SMTP_TIMEOUT", "30s") + await send_email( + receiver_email="to@invalid", + subject="Hi", + html="

x

", + ) + assert in_memory_smtp.sent == [] + + +@pytest.mark.asyncio +async def test_send_email_runs_off_event_loop_thread(in_memory_smtp: Any) -> None: + await send_email( + receiver_email="to@invalid", + subject="Hi", + html="

x

", + ) + assert in_memory_smtp.sent[0].thread_ident != threading.get_ident() + + @pytest.mark.asyncio async def test_send_email_smtp_failure_is_swallowed( in_memory_smtp: Any, @@ -113,7 +153,5 @@ async def test_send_email_smtp_failure_is_swallowed( does not raise so a failing email never blocks the proxy. """ in_memory_smtp.raise_on_send = RuntimeError("smtp boom") - await send_email( - receiver_email="to@invalid", subject="Hi", html="

x

" - ) + await send_email(receiver_email="to@invalid", subject="Hi", html="

x

") assert in_memory_smtp.sent == [] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index 0971ce09d79..9d2a27ce9d3 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -9,11 +9,14 @@ import pytest from fastapi import HTTPException import litellm +from litellm.caching.caching import DualCache from litellm.exceptions import RejectedRequestError from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CallTypesLiteral def _load(module: str, name: str): @@ -473,7 +476,13 @@ class _RedactingGuardrail(CustomGuardrail): kwargs.setdefault("event_hook", GuardrailEventHooks.pre_call) super().__init__(guardrail_name="redactor", **kwargs) - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override] + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> dict | None: for msg in data.get("messages", []): if "SECRET" in msg.get("content", ""): msg["content"] = msg["content"].replace("SECRET", "[REDACTED]") @@ -488,7 +497,13 @@ class _BlockOnSecretGuardrail(CustomGuardrail): kwargs.setdefault("event_hook", GuardrailEventHooks.pre_call) super().__init__(guardrail_name="blocker", **kwargs) - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override] + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> dict | None: if any("SECRET" in msg.get("content", "") for msg in data.get("messages", [])): raise HTTPException(status_code=400, detail="blocked: SECRET detected") return None @@ -560,7 +575,13 @@ async def test_scan_raw_request_guardrail_does_not_undo_later_masking( separate marker (PII_TOKEN) that only the redactor reacts to.""" class _PiiRedactor(_RedactingGuardrail): - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override] + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> dict | None: for msg in data.get("messages", []): if "PII_TOKEN" in msg.get("content", ""): msg["content"] = msg["content"].replace("PII_TOKEN", "[REDACTED]") @@ -692,7 +713,13 @@ async def test_scan_raw_request_warns_when_guardrail_mutation_discarded( super().__init__(**kwargs) self.scan_raw_request = True - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override] + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> dict | None: for msg in data.get("messages", []): msg["content"] = msg["content"].replace("SECRET", "[REDACTED]") return data diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/test_litellm/repositories/test_unit_of_work.py index 1ebfd917e36..1a76b537e95 100644 --- a/tests/test_litellm/repositories/test_unit_of_work.py +++ b/tests/test_litellm/repositories/test_unit_of_work.py @@ -1,9 +1,11 @@ +from dataclasses import fields from datetime import datetime, timezone from typing import Any, Dict, List, Mapping, Tuple import pytest from litellm.repositories.unit_of_work import ( + LinkedSpendResetWrites, budget_cascade_unit_of_work, spend_reset_unit_of_work, ) @@ -32,6 +34,7 @@ class FakeBatch: self.litellm_teammembership = FakeBatchTable("litellm_teammembership", self.calls) self.litellm_organizationtable = FakeBatchTable("litellm_organizationtable", self.calls) self.litellm_tagtable = FakeBatchTable("litellm_tagtable", self.calls) + self.litellm_modelaccessgroupbudgettable = FakeBatchTable("litellm_modelaccessgroupbudgettable", self.calls) self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls) async def commit(self) -> None: @@ -90,6 +93,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): uow.keys.queue_spend_zero(where=linked) uow.organizations.queue_spend_zero(where=linked) uow.tags.queue_spend_zero(where=linked) + uow.model_access_groups.queue_spend_zero(where=linked) uow.endusers.queue_spend_zero(where={"user_id": {"in": ["enduser-1"]}}) uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) assert batch.commit_count == 0 @@ -100,6 +104,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): ("litellm_verificationtoken.update_many", linked, {"spend": 0}), ("litellm_organizationtable.update_many", linked, {"spend": 0}), ("litellm_tagtable.update_many", linked, {"spend": 0}), + ("litellm_modelaccessgroupbudgettable.update_many", linked, {"spend": 0}), ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-1"]}}, {"spend": 0}), ("litellm_budgettable.update_many", {"budget_id": "budget-1"}, {"budget_reset_at": reset_at}), ] @@ -117,6 +122,39 @@ async def test_budget_window_advance_tolerates_a_tier_deleted_mid_chunk(): assert [call[0] for call in batch.calls] == ["litellm_budgettable.update_many"] +async def test_every_cascade_dependent_writes_to_its_own_table_on_the_one_batch(): + """Walks the dataclass instead of naming tables, so a dependent added to + BudgetCascadeUnitOfWork later cannot go uncovered. + + The named test above only proves the tables it lists, and an unbound + dependent surfaces as an AttributeError from whichever tests happen to + open a cascade. This pins the real contract: every field writes, each to a + distinct table, all on the same batch. + """ + reset_at = datetime(2026, 8, 3, 12, 0, tzinfo=timezone.utc) + batches: List[FakeBatch] = [] + + def _new_batch() -> FakeBatch: + # Fresh per call like db.batch_(), unlike the `lambda: batch` above: a + # second transaction would otherwise alias onto the first and hide. + batches.append(FakeBatch()) + return batches[-1] + + async with budget_cascade_unit_of_work(_new_batch) as uow: + writes = [getattr(uow, field.name) for field in fields(uow)] + for write in writes: + if isinstance(write, LinkedSpendResetWrites): + write.queue_spend_zero(where={"budget_id": "budget-1"}) + else: + write.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) + + assert len(batches) == 1, "the cascade must open exactly one transaction" + batch = batches[0] + assert len(batch.calls) == len(writes), "a dependent bound to a batch of its own would not land here" + assert len({call[0] for call in batch.calls}) == len(writes), "two dependents share one table" + assert batch.commit_count == 1 + + async def test_budget_cascade_raising_inside_block_skips_commit(): """A failure part-way through must leave budget_reset_at where it was, so the tier is still due on the next tick.""" diff --git a/tests/test_litellm/test_sync_together_ai_models.py b/tests/test_litellm/test_sync_together_ai_models.py index 7c1287e94b8..b8a85bcfbdc 100644 --- a/tests/test_litellm/test_sync_together_ai_models.py +++ b/tests/test_litellm/test_sync_together_ai_models.py @@ -105,7 +105,6 @@ def test_added_chat_model_matches_reviewed_registry_shape() -> None: "input_cost_per_token": 3e-06, "litellm_provider": "together_ai", "max_input_tokens": 1048576, - "max_output_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 1.5e-05, @@ -138,7 +137,35 @@ def test_moderation_type_maps_to_chat_mode() -> None: outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) guard = outcome.cost_map["together_ai/meta-llama/Llama-Guard-4-12B"] assert guard["mode"] == "chat" - assert guard["max_output_tokens"] == 1048576 + assert "max_output_tokens" not in guard + + +def test_output_ceiling_comes_from_the_rule_never_from_context_length() -> None: + glm = next(model for model in RECORDED_CATALOG if model.id == "zai-org/GLM-5.2") + fresh = sync.compute_sync({}, [_chat_model("acme/unreviewed", ctx=1048576), glm], _doc({"x": "2026-01-01"})) + unreviewed = fresh.cost_map["together_ai/acme/unreviewed"] + assert "max_output_tokens" not in unreviewed + assert (unreviewed["max_input_tokens"], unreviewed["max_tokens"]) == (1048576, 1048576) + reviewed = fresh.cost_map["together_ai/zai-org/GLM-5.2"] + assert (reviewed["max_input_tokens"], reviewed["max_output_tokens"], reviewed["max_tokens"]) == ( + 1048575, + 128000, + 128000, + ) + inflated = { + "together_ai/zai-org/GLM-5.2": { + "input_cost_per_token": 1.4e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048575, + "max_output_tokens": 1048575, + "max_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + } + } + corrected = sync.compute_sync(inflated, [glm], _doc({"x": "2026-01-01"})) + assert corrected.cost_map["together_ai/zai-org/GLM-5.2"]["max_output_tokens"] == 128000 + assert any("max_output_tokens: 1048575 -> 128000" in line for line in corrected.updated) def test_docs_removed_but_live_model_stays_live_with_warning() -> None: diff --git a/tests/test_litellm/test_together_ai_model_metadata.py b/tests/test_litellm/test_together_ai_model_metadata.py index 45f0370386b..c9e2863d240 100644 --- a/tests/test_litellm/test_together_ai_model_metadata.py +++ b/tests/test_litellm/test_together_ai_model_metadata.py @@ -15,6 +15,7 @@ COST_MAP_ADAPTER: Final = TypeAdapter(CostMap) SERVERLESS_CHAT_MODELS: Final = ( "together_ai/moonshotai/Kimi-K3", "together_ai/zai-org/GLM-5.2", + "together_ai/zai-org/GLM-5.3", "together_ai/zai-org/GLM-5.3-Flash", "together_ai/deepseek-ai/DeepSeek-V4-Pro-0813", "together_ai/deepseek-ai/DeepSeek-V4-Flash-0731", @@ -107,6 +108,8 @@ def test_together_glm_52_pricing(cost_map: CostMap): info = cost_map["together_ai/zai-org/GLM-5.2"] assert info["input_cost_per_token"] == 1.4e-06 assert info["output_cost_per_token"] == 4.4e-06 + assert info["max_input_tokens"] == 1048575 + assert info["max_output_tokens"] == 128000 assert info["supports_function_calling"] is True assert info["supports_reasoning"] is True @@ -117,7 +120,7 @@ def test_together_glm_53_flash_pricing_and_capabilities(cost_map: CostMap): assert info["output_cost_per_token"] == 5e-07 assert info["cache_read_input_token_cost"] == 3e-08 assert info["max_input_tokens"] == 1048575 - assert info["max_output_tokens"] == 1048575 + assert info["max_output_tokens"] == 128000 assert info["supports_function_calling"] is True assert info["supports_parallel_function_calling"] is True assert info["supports_prompt_caching"] is True @@ -127,6 +130,18 @@ def test_together_glm_53_flash_pricing_and_capabilities(cost_map: CostMap): assert info["supports_reasoning"] is True +def test_together_chat_entries_never_carry_context_length_as_output_ceiling(cost_map: CostMap): + inflated = sorted( + model + for model, info in cost_map.items() + if info.get("litellm_provider") == "together_ai" + and info.get("mode") == "chat" + and "max_output_tokens" in info + and info["max_output_tokens"] == info.get("max_input_tokens") + ) + assert inflated == [] + + def test_together_multilingual_e5_embedding_entry(cost_map: CostMap): info = cost_map["together_ai/intfloat/multilingual-e5-large-instruct"] assert info["mode"] == "embedding" 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 && ( + render( + + + , + ); + describe("JsonViewer", () => { it("should render a placeholder and no tree when the log entry carries no payload", () => { - render(); + renderWithTheme("light", null); expect(screen.getByText("No data")).toBeInTheDocument(); expect(screen.queryByRole("tree")).not.toBeInTheDocument(); }); it("should render the payload as a tree exposing its keys", () => { - render(); + renderWithTheme("light", { model: "claude-opus-4-5", stream: true }); expect(screen.getByRole("tree")).toBeInTheDocument(); expect(screen.getByText(/model/)).toBeInTheDocument(); @@ -20,9 +29,26 @@ describe("JsonViewer", () => { }); it("should treat an empty payload as data rather than showing the placeholder", () => { - render(); + renderWithTheme("light", {}); expect(screen.getByRole("tree")).toBeInTheDocument(); expect(screen.queryByText("No data")).not.toBeInTheDocument(); }); + + it("should style the tree with the light palette when the dashboard theme is light", () => { + renderWithTheme("light", { model: "claude-opus-4-5" }); + + expect(screen.getByRole("tree")).toHaveClass(...defaultStyles.container.split(" ")); + }); + + it("should style the tree with the dark palette when the dashboard theme is dark", () => { + renderWithTheme("dark", { model: "claude-opus-4-5" }); + + const tree = screen.getByRole("tree"); + expect(tree).toHaveClass(...darkStyles.container.split(" ")); + defaultStyles.container + .split(" ") + .filter((className) => !darkStyles.container.split(" ").includes(className)) + .forEach((lightOnlyClassName) => expect(tree).not.toHaveClass(lightOnlyClassName)); + }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.tsx index 86b897f526f..980573a65cc 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.tsx @@ -1,4 +1,5 @@ -import { JsonView, defaultStyles } from "react-json-view-lite"; +import { useTheme } from "next-themes"; +import { JsonView, darkStyles, defaultStyles } from "react-json-view-lite"; import "react-json-view-lite/dist/index.css"; import { JSON_MAX_HEIGHT, SPACING_LARGE } from "./constants"; @@ -12,6 +13,8 @@ interface JsonViewerProps { * Uses an interactive tree component for easy navigation. */ export function JsonViewer({ data }: JsonViewerProps) { + const { resolvedTheme } = useTheme(); + if (!data) return No data; return ( @@ -24,8 +27,8 @@ export function JsonViewer({ data }: JsonViewerProps) { borderRadius: 4, }} > -
- +
+
); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx index f677a66cc54..ddd0a650c04 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx @@ -1,6 +1,5 @@ import { useEffect, useMemo, useState } from "react"; -import { Bot, Check, ChevronLeft, ChevronRight, Copy, Sparkles, Wrench } from "lucide-react"; -import { Button } from "@/components/ui/button"; +import { Bot, Check, Copy, Sparkles, Wrench } from "lucide-react"; import { Sheet, SheetContent, SheetTitle } from "@/components/ui/sheet"; import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { LogEntry } from "../columns"; @@ -9,6 +8,7 @@ import { AGENT_CALL_TYPES, MCP_CALL_TYPES } from "../constants"; import { getEventDisplayName } from "../utils"; import { ClassifyTag } from "./ClassifyTag"; import { DrawerHeader } from "./DrawerHeader"; +import { SidebarToggle } from "./SidebarToggle"; import { useKeyboardNavigation } from "./useKeyboardNavigation"; import { LogDetailContent, GuardrailJumpLink } from "./LogDetailContent"; import { sessionSpendLogsCall } from "../../networking"; @@ -313,26 +313,12 @@ export function LogDetailsDrawer({ {logEntry?.request_id ? `Request ${logEntry.request_id} details` : "Request details"}
- {!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..c762183dec4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -497,6 +497,91 @@ export interface paths { patch?: never; trace?: never; }; + "/access_group/{access_group}/budget": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Access Group Budget + * @description Get the shared budget of an access group, and the spend drawn against it. + * + * Example: + * ```bash + * curl -X GET 'http://localhost:4000/access_group/production-models/budget' \ + * -H 'Authorization: Bearer sk-1234' + * ``` + * + * Parameters: + * - access_group: str - The access group name (URL path parameter) + * + * Returns: + * - AccessGroupBudgetResponse; budget is null when the group has no budget set + * + * Raises: + * - HTTPException 404: If access group not found + */ + get: operations["get_access_group_budget_access_group__access_group__budget_get"]; + /** + * Set Access Group Budget + * @description Set or replace the shared budget of an access group. Idempotent. + * + * Every key that can reach a model in the group draws from this one budget. + * + * Example: + * ```bash + * curl -X PUT 'http://localhost:4000/access_group/production-models/budget' \ + * -H 'Authorization: Bearer sk-1234' \ + * -H 'Content-Type: application/json' \ + * -d '{ + * "max_budget": 100.0, + * "budget_duration": "30d" + * }' + * ``` + * + * Parameters: + * - access_group: str - The access group name (URL path parameter) + * - max_budget: Optional[float] - Requests fail once the group's shared spend exceeds this + * - soft_budget: Optional[float] - Fires an alert when reached; requests still succeed + * - budget_duration: Optional[str] - Frequency of resetting the group's spend (e.g. '30d') + * - budget_id: Optional[str] - Link an existing budget instead of creating one + * + * Returns: + * - AccessGroupBudgetResponse with the stored budget and current spend + * + * Raises: + * - HTTPException 400: If no budget field is given, or budget_duration cannot be parsed + * - HTTPException 404: If access group not found + */ + put: operations["set_access_group_budget_access_group__access_group__budget_put"]; + post?: never; + /** + * Delete Access Group Budget + * @description Clear the shared budget of an access group, leaving the group itself in place. + * + * Example: + * ```bash + * curl -X DELETE 'http://localhost:4000/access_group/production-models/budget' \ + * -H 'Authorization: Bearer sk-1234' + * ``` + * + * Parameters: + * - access_group: str - The access group name (URL path parameter) + * + * Returns: + * - DeleteAccessGroupBudgetResponse; budget_deleted is false when there was nothing to clear + * + * Raises: + * - HTTPException 404: If access group not found + */ + delete: operations["delete_access_group_budget_access_group__access_group__budget_delete"]; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/access_group/{access_group}/delete": { parameters: { query?: never; @@ -555,7 +640,7 @@ export interface paths { * - access_group: str - The access group name (URL path parameter) * * Returns: - * - AccessGroupInfo with the access group details + * - AccessGroupInfo with the access group details, its shared budget and its spend * * Raises: * - HTTPException 404: If access group not found @@ -2858,6 +2943,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 +7112,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; @@ -22182,6 +22344,38 @@ export interface components { */ type: "restricted_sso_group"; }; + /** AccessGroupBudget */ + AccessGroupBudget: { + /** Budget Duration */ + budget_duration?: string | null; + /** Budget Id */ + budget_id: string; + /** Budget Reset At */ + budget_reset_at?: string | null; + /** Max Budget */ + max_budget?: number | null; + /** Soft Budget */ + soft_budget?: number | null; + }; + /** AccessGroupBudgetRequest */ + AccessGroupBudgetRequest: { + /** Budget Duration */ + budget_duration?: string | null; + /** Budget Id */ + budget_id?: string | null; + /** Max Budget */ + max_budget?: number | null; + /** Soft Budget */ + soft_budget?: number | null; + }; + /** AccessGroupBudgetResponse */ + AccessGroupBudgetResponse: { + /** Access Group */ + access_group: string; + budget?: components["schemas"]["AccessGroupBudget"] | null; + /** Spend */ + spend: number; + }; /** AccessGroupCreateRequest */ AccessGroupCreateRequest: { /** Access Agent Ids */ @@ -22203,10 +22397,13 @@ export interface components { AccessGroupInfo: { /** Access Group */ access_group: string; + budget?: components["schemas"]["AccessGroupBudget"] | null; /** Deployment Count */ deployment_count: number; /** Model Names */ model_names: string[]; + /** Spend */ + spend?: number | null; }; /** AccessGroupResponse */ AccessGroupResponse: { @@ -23474,6 +23671,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 +26014,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"]; @@ -26004,6 +26252,15 @@ export interface components { [key: string]: unknown; }; }; + /** DeleteAccessGroupBudgetResponse */ + DeleteAccessGroupBudgetResponse: { + /** Access Group */ + access_group: string; + /** Budget Deleted */ + budget_deleted: boolean; + /** Message */ + message: string; + }; /** * DeleteCustomerRequest * @description Delete multiple Customers @@ -39024,6 +39281,103 @@ export interface operations { }; }; }; + get_access_group_budget_access_group__access_group__budget_get: { + parameters: { + query?: never; + header?: never; + path: { + access_group: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["AccessGroupBudgetResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + set_access_group_budget_access_group__access_group__budget_put: { + parameters: { + query?: never; + header?: never; + path: { + access_group: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["AccessGroupBudgetRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["AccessGroupBudgetResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + delete_access_group_budget_access_group__access_group__budget_delete: { + parameters: { + query?: never; + header?: never; + path: { + access_group: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DeleteAccessGroupBudgetResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; delete_access_group_access_group__access_group__delete_delete: { parameters: { query?: never; @@ -42757,6 +43111,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 +48119,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; diff --git a/uv.lock b/uv.lock index 019c5a70f9e..8ef72116466 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-08-25T23:16:47.126855Z" +exclude-newer = "2026-08-26T18:33:25.773031Z" exclude-newer-span = "P3D" [manifest] @@ -4665,12 +4665,12 @@ proxy-dev = [ [[package]] name = "litellm-enterprise" -version = "0.1.61" +version = "0.1.62" source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.90" +version = "0.4.91" source = { editable = "litellm-proxy-extras" } [[package]]