mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_fix_post_call_policy_pipeline
# Conflicts: # litellm/proxy/policy_engine/pipeline_executor.py
This commit is contained in:
commit
45a6b1de23
199 changed files with 11860 additions and 831 deletions
22
.github/workflows/codspeed.yml
vendored
22
.github/workflows/codspeed.yml
vendored
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -25,6 +25,8 @@ flag_management:
|
|||
carryforward: false
|
||||
- name: proxy-db-schema-migration
|
||||
carryforward: false
|
||||
- name: circleci
|
||||
carryforward: false
|
||||
|
||||
component_management:
|
||||
individual_components:
|
||||
|
|
|
|||
|
|
@ -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": "<retrieve_batch>"}],
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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==",
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
);
|
||||
|
||||
|
|
@ -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 $$;
|
||||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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==",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
@ -1681,6 +1682,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.
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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_-]+$``
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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", "")
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,10 +3,11 @@ Tenant-isolation helpers for managed file/batch/vector-store resources.
|
|||
|
||||
Returns a Prisma filter and an ownership check that scope managed resources
|
||||
to the caller's identity: proxy admins see everything, user-keyed callers
|
||||
see records they created, and service-account keys (no user_id) fall back
|
||||
to the resource's owning team. Callers with no admin role and no
|
||||
identifying ids are denied so an empty user_id can never select an
|
||||
unscoped query.
|
||||
see records they created, service-account keys (no user_id) fall back to
|
||||
the resource's owning team, and keys with neither a user_id nor a team_id
|
||||
fall back to their own hashed token so they can still reach the resources
|
||||
they created. Callers with no admin role and no identifying ids at all
|
||||
are denied so an empty user_id can never select an unscoped query.
|
||||
"""
|
||||
|
||||
from typing import Any, Final
|
||||
|
|
@ -19,6 +20,32 @@ from litellm.proxy._types import (
|
|||
)
|
||||
|
||||
|
||||
def resolve_resource_owner_id(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> str | None:
|
||||
"""Return the identity to stamp on (and match against) a managed
|
||||
resource's ``created_by``.
|
||||
|
||||
A key with neither a user_id nor a team_id would otherwise stamp
|
||||
``created_by=None`` and be locked out of its own resources, so it owns
|
||||
them under its hashed token instead, using the ``key:`` scope prefix
|
||||
already used by ``proxy/common_utils/resource_ownership.py``. ``None``
|
||||
means the caller has no usable identity of its own and must fall back
|
||||
to team scoping, or be denied.
|
||||
"""
|
||||
if user_api_key_dict.user_id is not None:
|
||||
return user_api_key_dict.user_id
|
||||
|
||||
if user_api_key_dict.team_id is not None:
|
||||
return None
|
||||
|
||||
token: Final = user_api_key_dict.token or user_api_key_dict.api_key
|
||||
if token:
|
||||
return f"key:{token}"
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, Any]:
|
||||
"""Build the OpenAI-style paginated list response shape used by managed
|
||||
file/batch/vector-store listings. ``first_id`` and ``last_id`` are
|
||||
|
|
@ -39,7 +66,8 @@ def build_owner_filter(
|
|||
to records the caller is allowed to see.
|
||||
|
||||
- ``{}`` means no scoping (proxy admins).
|
||||
- ``{"created_by": <user_id>}`` for user-keyed callers.
|
||||
- ``{"created_by": <owner_id>}`` for user-keyed callers, and for keys
|
||||
with no user_id and no team_id (owner id is their hashed token).
|
||||
- ``{"team_id": <team_id>}`` for service-account callers
|
||||
that have a team but no user_id.
|
||||
- ``{"OR": [...]}`` when the caller has both — listing must include
|
||||
|
|
@ -62,12 +90,13 @@ def build_owner_filter(
|
|||
]
|
||||
}
|
||||
|
||||
if user_id is not None:
|
||||
return {"created_by": user_id}
|
||||
|
||||
if team_id is not None:
|
||||
return {"team_id": team_id}
|
||||
|
||||
owner_id: Final = resolve_resource_owner_id(user_api_key_dict)
|
||||
if owner_id is not None:
|
||||
return {"created_by": owner_id}
|
||||
|
||||
return None
|
||||
|
||||
|
||||
|
|
@ -86,8 +115,8 @@ def can_access_resource(
|
|||
if _user_has_admin_view(user_api_key_dict):
|
||||
return True
|
||||
|
||||
user_id: Final = user_api_key_dict.user_id
|
||||
if user_id is not None and created_by is not None and created_by == user_id:
|
||||
owner_id: Final = resolve_resource_owner_id(user_api_key_dict)
|
||||
if owner_id is not None and created_by is not None and created_by == owner_id:
|
||||
return True
|
||||
|
||||
team_id: Final = user_api_key_dict.team_id
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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),))
|
||||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -43318,6 +43318,22 @@
|
|||
"video"
|
||||
]
|
||||
},
|
||||
"vertex_ai/veo-3.1-lite-generate-001": {
|
||||
"litellm_provider": "vertex_ai-video-models",
|
||||
"max_input_tokens": 1024,
|
||||
"max_tokens": 1024,
|
||||
"mode": "video_generation",
|
||||
"output_cost_per_second": 0.05,
|
||||
"output_cost_per_second_1080p": 0.08,
|
||||
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#veo",
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"video"
|
||||
]
|
||||
},
|
||||
"voyage/rerank-2": {
|
||||
"input_cost_per_token": 5e-08,
|
||||
"litellm_provider": "voyage",
|
||||
|
|
@ -51828,6 +51844,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",
|
||||
|
|
|
|||
|
|
@ -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": [],
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -5256,6 +5606,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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -9,16 +9,18 @@ from .exceptions import UnauthorizedError
|
|||
|
||||
|
||||
class KeysManagementClient:
|
||||
def __init__(self, base_url: str, api_key: str | None = None):
|
||||
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
|
||||
"""
|
||||
Initialize the KeysManagementClient.
|
||||
|
||||
Args:
|
||||
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000")
|
||||
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
|
||||
timeout (int): Request timeout in seconds (default: 30)
|
||||
"""
|
||||
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
|
||||
self._api_key = api_key
|
||||
self._timeout = timeout
|
||||
|
||||
def _get_headers(self) -> dict[str, str]:
|
||||
"""
|
||||
|
|
@ -99,7 +101,7 @@ class KeysManagementClient:
|
|||
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
@ -174,7 +176,7 @@ class KeysManagementClient:
|
|||
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
@ -218,7 +220,7 @@ class KeysManagementClient:
|
|||
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
@ -279,7 +281,7 @@ class KeysManagementClient:
|
|||
session: Final = requests.Session()
|
||||
response_text: str | None = None
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response_text = response.text
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
|
@ -309,7 +311,7 @@ class KeysManagementClient:
|
|||
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
|
|||
|
|
@ -6,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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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.")
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -33,7 +33,11 @@ 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.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 +45,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 +97,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 +164,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({}),
|
||||
|
|
@ -610,6 +628,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 +662,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 +673,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 +699,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 +743,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 +775,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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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__"
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -86,6 +88,7 @@ class _SpendBatch(Protocol):
|
|||
litellm_organizationtable: BatchTable
|
||||
litellm_tagtable: BatchTable
|
||||
litellm_agentstable: BatchTable
|
||||
litellm_modelaccessgroupbudgettable: BatchTable
|
||||
|
||||
|
||||
class _SpendBatchManager(Protocol):
|
||||
|
|
@ -109,7 +112,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 +126,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
|
||||
|
|
@ -187,6 +236,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 +289,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),
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -320,7 +371,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 +482,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 +557,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 +874,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,
|
||||
|
|
@ -927,7 +1031,8 @@ class DBSpendUpdateWriter:
|
|||
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 +1041,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,
|
||||
|
|
@ -1433,6 +1539,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 +1569,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 +2004,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"),
|
||||
|
|
|
|||
|
|
@ -65,6 +65,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 +77,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = (
|
|||
"org_list_transactions",
|
||||
"tag_list_transactions",
|
||||
"agent_list_transactions",
|
||||
"model_access_group_list_transactions",
|
||||
)
|
||||
|
||||
_ValueT = TypeVar("_ValueT")
|
||||
|
|
@ -397,6 +399,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:
|
||||
|
|
@ -826,6 +832,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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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.<param_name>`` 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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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,6 +584,7 @@ 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(
|
||||
|
|
@ -610,6 +623,7 @@ async def _update_database_and_spend_counters(
|
|||
budget_reservation=budget_reservation,
|
||||
end_user_id=end_user_id,
|
||||
tags=request_tags,
|
||||
model_access_groups=model_access_groups,
|
||||
)
|
||||
except Exception:
|
||||
if budget_reservation is not None:
|
||||
|
|
|
|||
|
|
@ -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={
|
||||
|
|
|
|||
|
|
@ -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}",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -58,15 +58,15 @@ def _tool_call_shape(tool_call: object) -> tuple[object, object]:
|
|||
|
||||
|
||||
def _rewrote_texts(sent: Sequence[str] | None, returned: Sequence[str] | None) -> bool:
|
||||
return sent is not None and returned is not None and list(returned) != list(sent)
|
||||
return sent is not None and returned is not None and tuple(returned) != tuple(sent)
|
||||
|
||||
|
||||
def _rewrote_tool_calls(sent: Sequence[object] | None, returned: Sequence[object] | None) -> bool:
|
||||
if sent is None or returned is None:
|
||||
return False
|
||||
return [_tool_call_shape(tool_call) for tool_call in returned] != [
|
||||
return tuple(_tool_call_shape(tool_call) for tool_call in returned) != tuple(
|
||||
_tool_call_shape(tool_call) for tool_call in sent
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class _StreamRewriteObserver(CustomGuardrail):
|
||||
|
|
@ -103,7 +103,7 @@ class _StreamRewriteObserver(CustomGuardrail):
|
|||
|
||||
def _prepare_hook_input(
|
||||
step: PipelineStep,
|
||||
callback: CustomLogger,
|
||||
callback: CustomGuardrail,
|
||||
data: dict, # mutable-ok: same request-payload shape the hooks mutate
|
||||
raw_request_snapshot: dict | None, # mutable-ok: same request-payload shape as data
|
||||
) -> tuple[dict, bool]: # mutable-ok: returns that same request-payload dict
|
||||
|
|
@ -113,15 +113,17 @@ def _prepare_hook_input(
|
|||
pipeline may have already rewritten), same reason the normal sequential/parallel
|
||||
guardrail loops do this."""
|
||||
if "metadata" not in data:
|
||||
data["metadata"] = {}
|
||||
data["metadata"]["guardrails"] = [step.guardrail]
|
||||
data["metadata"] = {} # mutable-ok: request metadata bucket, hooks mutate it
|
||||
data["metadata"]["guardrails"] = [
|
||||
step.guardrail
|
||||
] # mutable-ok: guardrails list is part of the request-payload shape
|
||||
|
||||
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 else data
|
||||
)
|
||||
if hook_input is not data:
|
||||
hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail]
|
||||
hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail] # mutable-ok: request metadata shape
|
||||
return hook_input, scans_raw_request
|
||||
|
||||
|
||||
|
|
@ -342,7 +344,12 @@ class PipelineExecutor:
|
|||
if response is None or scans_raw_request:
|
||||
return ("pass", None, None, None)
|
||||
if mode == "post_call":
|
||||
return ("pass", {"response": response}, None, None)
|
||||
return (
|
||||
"pass",
|
||||
{"response": response},
|
||||
None,
|
||||
None,
|
||||
) # mutable-ok: modified-data contract is a plain dict
|
||||
return ("pass", response if isinstance(response, dict) else None, None, None)
|
||||
|
||||
except UndeliverableStreamRewrite:
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ import threading
|
|||
import time
|
||||
import traceback
|
||||
import warnings
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping, MutableMapping, Sequence
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Collection, Mapping, MutableMapping, Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType, UnionType
|
||||
from typing import (
|
||||
|
|
@ -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
|
||||
|
|
@ -2648,6 +2650,7 @@ async def increment_spend_counters(
|
|||
budget_reservation: dict | None = None,
|
||||
end_user_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
model_access_groups: Sequence[str] | None = None,
|
||||
):
|
||||
"""
|
||||
Atomically increment spend counters for budget enforcement.
|
||||
|
|
@ -2777,6 +2780,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 +2875,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,
|
||||
|
|
@ -4281,6 +4318,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 +6803,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 +7016,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 +7181,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 +12120,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 (
|
||||
|
|
@ -12792,6 +12900,7 @@ async def _fetch_db_models_for_search(
|
|||
size: int,
|
||||
sort_by: str | None,
|
||||
is_byok_outside_caller_teams: Callable[[dict[str, JsonValue]], bool],
|
||||
model_name: str | None = None,
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
"""
|
||||
Run the bounded DB query that backs `/v2/model/info?search=`. Returns
|
||||
|
|
@ -12808,7 +12917,9 @@ async def _fetch_db_models_for_search(
|
|||
filter for `team_public_model_name` instead and keep the DB cost
|
||||
bounded by `search`.
|
||||
"""
|
||||
db_where_condition: Final[dict[str, Any]] = {"model_name": {"contains": search_lower, "mode": "insensitive"}}
|
||||
db_where_condition: Final[dict[str, Any]] = {
|
||||
"model_name": {"contains": search_lower, "mode": "insensitive"} if model_name is None else model_name
|
||||
}
|
||||
if db_model_ids_in_router:
|
||||
db_where_condition["model_id"] = {"not": {"in": list(db_model_ids_in_router)}}
|
||||
|
||||
|
|
@ -12855,6 +12966,7 @@ async def _apply_search_filter_to_models(
|
|||
page: int = 1,
|
||||
size: int = 50,
|
||||
sort_by: str | None = None,
|
||||
model_name: str | None = None,
|
||||
) -> tuple[list[dict[str, Any]], int | None]:
|
||||
"""
|
||||
Apply search filter to models, querying database for additional matching models.
|
||||
|
|
@ -12875,6 +12987,11 @@ async def _apply_search_filter_to_models(
|
|||
sort_by: Sort field. When set, results must be sorted across the
|
||||
full match set, so the DB fetch is capped at
|
||||
``_SORTED_SEARCH_DB_FETCH_CAP`` instead of one page.
|
||||
model_name: Exact ``model_name`` the caller already narrowed
|
||||
``all_models`` to (``?model=``). The DB query matches it
|
||||
exactly instead of the substring, and is skipped when the
|
||||
substring cannot occur in it, otherwise rows from other model
|
||||
groups leak into the result and the count.
|
||||
|
||||
Returns:
|
||||
Tuple of (filtered_models, total_count). total_count is None if not searching.
|
||||
|
|
@ -12932,7 +13049,8 @@ async def _apply_search_filter_to_models(
|
|||
|
||||
# Query database for additional models with search term
|
||||
db_models: list[dict[str, Any]] = []
|
||||
if prisma_client is not None:
|
||||
exact_name_can_match: Final = model_name is None or search_lower in model_name.lower()
|
||||
if prisma_client is not None and exact_name_can_match:
|
||||
try:
|
||||
db_models, db_models_total_count = await _fetch_db_models_for_search(
|
||||
prisma_client=prisma_client,
|
||||
|
|
@ -12944,6 +13062,7 @@ async def _apply_search_filter_to_models(
|
|||
size=size,
|
||||
sort_by=sort_by,
|
||||
is_byok_outside_caller_teams=_is_byok_outside_caller_teams,
|
||||
model_name=model_name,
|
||||
)
|
||||
search_total_count = router_models_count + db_models_total_count
|
||||
except Exception as e:
|
||||
|
|
@ -13497,7 +13616,7 @@ async def model_info_v2(
|
|||
all_models += [user_model]
|
||||
|
||||
if model is not None:
|
||||
all_models = [m for m in all_models if m["model_name"] == model]
|
||||
all_models = [m for m in all_models if _deployment_matches_allowed_model_names(m, frozenset((model,)))]
|
||||
|
||||
# Apply search filter if provided
|
||||
all_models, search_total_count = await _apply_search_filter_to_models(
|
||||
|
|
@ -13509,6 +13628,7 @@ async def model_info_v2(
|
|||
page=page,
|
||||
size=size,
|
||||
sort_by=sortBy,
|
||||
model_name=model,
|
||||
)
|
||||
|
||||
if user_models_only:
|
||||
|
|
@ -14023,7 +14143,7 @@ async def model_metrics_exceptions(
|
|||
return {"data": response, "exception_types": list(exception_types)}
|
||||
|
||||
|
||||
def _deployment_matches_allowed_model_names(model: dict[str, JsonValue], allowed_model_names: set[str]) -> bool:
|
||||
def _deployment_matches_allowed_model_names(model: dict[str, JsonValue], allowed_model_names: Collection[str]) -> bool:
|
||||
"""Match a router deployment against allowed public model names.
|
||||
|
||||
Team-scoped rows store an internal routing key in ``model_name``; callers
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -53,6 +56,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 +162,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 +352,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 +441,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 +503,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 +542,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 +597,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 +640,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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -458,11 +458,35 @@ def _pipeline_managed_guardrail_names(
|
|||
)
|
||||
|
||||
|
||||
def _merge_pipeline_metadata_bucket(data: dict, bucket_key: str, modified_bucket_value: object) -> None:
|
||||
def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]:
|
||||
resolved: Final = tuple(
|
||||
litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
||||
cast( # cast-ok: the resolver returns None for unknown names, filtered below
|
||||
_custom_logger_compatible_callbacks_literal, callback
|
||||
)
|
||||
)
|
||||
if isinstance(callback, str)
|
||||
else callback
|
||||
for callback in litellm.callbacks
|
||||
)
|
||||
present: Final = tuple(callback for callback in resolved if callback is not None)
|
||||
guardrails: Final = tuple(callback for callback in present if isinstance(callback, CustomGuardrail))
|
||||
others: Final = cast( # cast-ok: mirrors the legacy loop, which treated every non-guardrail entry as a CustomLogger
|
||||
"tuple[CustomLogger, ...]",
|
||||
tuple(callback for callback in present if not isinstance(callback, CustomGuardrail)),
|
||||
)
|
||||
return (guardrails, others)
|
||||
|
||||
|
||||
def _merge_pipeline_metadata_bucket(
|
||||
data: dict, bucket_key: str, modified_bucket_value: object
|
||||
) -> None: # mutable-ok: request payload dict, written in place
|
||||
if not isinstance(modified_bucket_value, dict):
|
||||
return
|
||||
modified_bucket: Final = cast("dict[str, object]", modified_bucket_value) # cast-ok: metadata buckets are str-keyed
|
||||
surviving_writes: Final = {key: value for key, value in modified_bucket.items() if key != "guardrails"}
|
||||
surviving_writes: Final = {
|
||||
key: value for key, value in modified_bucket.items() if key != "guardrails"
|
||||
} # mutable-ok: merged into the live request metadata bucket in place
|
||||
existing_bucket: Final = data.get(bucket_key)
|
||||
if isinstance(existing_bucket, dict):
|
||||
cast("dict[str, object]", existing_bucket).update(surviving_writes) # cast-ok: metadata buckets are str-keyed
|
||||
|
|
@ -470,7 +494,9 @@ def _merge_pipeline_metadata_bucket(data: dict, bucket_key: str, modified_bucket
|
|||
data[bucket_key] = surviving_writes
|
||||
|
||||
|
||||
def _merge_pipeline_metadata_writes(data: dict, modified_data: Mapping[str, object]) -> None:
|
||||
def _merge_pipeline_metadata_writes(
|
||||
data: dict, modified_data: Mapping[str, object]
|
||||
) -> None: # mutable-ok: request payload dict, written in place
|
||||
"""
|
||||
Copy metadata-bucket writes from a pipeline's working copy back onto the request.
|
||||
|
||||
|
|
@ -1583,7 +1609,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
|
||||
|
|
@ -1620,7 +1646,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:
|
||||
|
|
@ -1746,7 +1772,7 @@ class ProxyLogging:
|
|||
event_hook: str,
|
||||
raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data
|
||||
response: LLMResponseTypes | None = None,
|
||||
) -> tuple[dict, LLMResponseTypes | None]:
|
||||
) -> tuple[dict, LLMResponseTypes | None]: # mutable-ok: returns the request-payload dict onward
|
||||
"""
|
||||
Execute guardrail pipelines if any are configured for this request.
|
||||
|
||||
|
|
@ -1771,7 +1797,9 @@ class ProxyLogging:
|
|||
if pipeline.mode != event_hook:
|
||||
continue
|
||||
|
||||
step_input: dict = {**data, "response": current_response} if current_response is not None else data
|
||||
step_input: dict = (
|
||||
{**data, "response": current_response} if current_response is not None else data
|
||||
) # mutable-ok: same request-payload shape as data
|
||||
|
||||
result: PipelineExecutionResult = await PipelineExecutor.execute_steps(
|
||||
steps=pipeline.steps,
|
||||
|
|
@ -1967,7 +1995,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
|
||||
|
|
@ -2129,7 +2157,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)
|
||||
|
||||
|
|
@ -2153,11 +2181,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)
|
||||
|
|
@ -2976,26 +3000,8 @@ class ProxyLogging:
|
|||
response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below
|
||||
|
||||
pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "post_call")
|
||||
guardrail_callbacks: Final[list[CustomGuardrail]] = []
|
||||
other_callbacks: Final[list[CustomLogger]] = []
|
||||
guardrail_callbacks, other_callbacks = _partition_post_call_callbacks()
|
||||
try:
|
||||
for callback in litellm.callbacks:
|
||||
_callback: CustomLogger | None = None
|
||||
if isinstance(callback, str):
|
||||
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
||||
cast(_custom_logger_compatible_callbacks_literal, callback)
|
||||
)
|
||||
else:
|
||||
_callback = callback
|
||||
|
||||
if _callback is not None:
|
||||
if isinstance(_callback, CustomGuardrail):
|
||||
guardrail_callbacks.append(_callback)
|
||||
else:
|
||||
other_callbacks.append(_callback)
|
||||
############## Handle Guardrails ########################################
|
||||
#############################################################################
|
||||
|
||||
# Merge model-level guardrails before checking which guardrails to run
|
||||
guardrail_data: Final = _check_and_merge_model_level_guardrails(data=data, llm_router=llm_router)
|
||||
|
||||
|
|
@ -3475,7 +3481,7 @@ class ProxyLogging:
|
|||
self,
|
||||
response: "AsyncGenerator[object, None]",
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_data: dict,
|
||||
request_data: dict, # mutable-ok: same request-payload shape the hooks mutate
|
||||
pipelines: "tuple[tuple[str, GuardrailPipeline], ...]",
|
||||
) -> "AsyncGenerator[Any, None]":
|
||||
"""
|
||||
|
|
@ -6348,10 +6354,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(
|
||||
|
|
@ -6397,27 +6435,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))
|
||||
|
|
|
|||
|
|
@ -144,4 +144,7 @@ class PrismaBatch(Protocol):
|
|||
@property
|
||||
def litellm_endusertable(self) -> BatchTable: ...
|
||||
|
||||
@property
|
||||
def litellm_modelaccessgroupbudgettable(self) -> BatchTable: ...
|
||||
|
||||
async def commit(self) -> None: ...
|
||||
|
|
|
|||
|
|
@ -100,6 +100,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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
(
|
||||
|
|
|
|||
72
litellm/types/llms/vertex_ai_gemini_transcription.py
Normal file
72
litellm/types/llms/vertex_ai_gemini_transcription.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
19
litellm/types/proxy/model_access_group_budget.py
Normal file
19
litellm/types/proxy/model_access_group_budget.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -43318,6 +43318,22 @@
|
|||
"video"
|
||||
]
|
||||
},
|
||||
"vertex_ai/veo-3.1-lite-generate-001": {
|
||||
"litellm_provider": "vertex_ai-video-models",
|
||||
"max_input_tokens": 1024,
|
||||
"max_tokens": 1024,
|
||||
"mode": "video_generation",
|
||||
"output_cost_per_second": 0.05,
|
||||
"output_cost_per_second_1080p": 0.08,
|
||||
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#veo",
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"video"
|
||||
]
|
||||
},
|
||||
"voyage/rerank-2": {
|
||||
"input_cost_per_token": 5e-08,
|
||||
"litellm_provider": "voyage",
|
||||
|
|
@ -51828,6 +51844,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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -177,15 +177,15 @@ quota_management.<behavior>.<variant>.<assertion>
|
|||
behavior : ratelimit | budget | spend_tracking
|
||||
variant : <ratelimit> rpm | tpm | priority_generous | priority_strict
|
||||
<budget> 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
|
||||
<spend_tracking> 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]
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
|
|
@ -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, {
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue