From 3b1ab1908c5657b45bb2a0c6a6b24c5c3339ebb7 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 4 Aug 2026 16:56:37 -0700 Subject: [PATCH 001/111] feat(proxy): add LiteLLM_BudgetWindowSpend table for per-window budget spend Multi-window budgets (budget_limits on keys/teams) currently keep window spend only in cache. Every cold or expired counter recomputes the window by aggregating LiteLLM_SpendLogs, which has no usable index for that query and saturates the DB on large tables (#35766). This adds a LiteLLM_BudgetWindowSpend table holding one row per configured window, keyed (entity_type, entity_id, window_duration), with window_start identifying the period the spend belongs to. Follow-up PRs maintain these rows from the spend update writer and move window budget enforcement reads onto them. --- .../migration.sql | 13 +++++++++++++ .../litellm_proxy_extras/schema.prisma | 12 ++++++++++++ litellm/proxy/schema.prisma | 12 ++++++++++++ schema.prisma | 12 ++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 5 files changed, 50 insertions(+), 1 deletion(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql new file mode 100644 index 00000000000..45cc927a328 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql @@ -0,0 +1,13 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetWindowSpend" ( + "entity_type" TEXT NOT NULL, + "entity_id" TEXT NOT NULL, + "window_duration" TEXT NOT NULL, + "window_start" TIMESTAMP(3) NOT NULL, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_BudgetWindowSpend_pkey" PRIMARY KEY ("entity_type","entity_id","window_duration") +); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 17339541fd9..85401b82676 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -645,6 +645,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 17339541fd9..85401b82676 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -645,6 +645,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/schema.prisma b/schema.prisma index 17339541fd9..85401b82676 100644 --- a/schema.prisma +++ b/schema.prisma @@ -645,6 +645,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9133bfb5cf4..fd2155701c2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -23744,7 +23744,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; + user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; }; /** * DefaultTeamSSOParams From 187e0fab605c63dcb73fedfe756c6c80ceb396ec Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 4 Aug 2026 18:25:00 -0700 Subject: [PATCH 002/111] feat(proxy): maintain per-window budget spend rows in the spend writer Multi-window budgets (budget_limits on keys and teams) enforce off Redis counters with a 60s TTL. Every cold counter, and every authoritative floor check, aggregates LiteLLM_SpendLogs with a range scan over startTime; that table has no index on api_key or team_id, so the scan lands on the highest volume table in the schema and saturates the connection pool (#35766). Every other budget feature reads a maintained running spend value instead. This gives windows the same shape by keeping LiteLLM_BudgetWindowSpend up to date: one row per configured window, whose window_start rolls forward in place. The cost callback already iterates a key's and a team's windows with the window start computed and the actual cost in hand, so it enqueues there, onto a new WindowSpendUpdateQueue. Increments are enqueued even when the cache increment is skipped for a reserved counter, since the reservation only pre-charged an estimate and the row still owes the actual cost. Windows with no reset_at slide with wall clock and cannot be represented by a single row, so they are left to the read path's existing aggregate. The queue flushes alongside the daily spend queues, through the Redis buffer when one is configured (only the pod-lock winner commits) and directly otherwise. A flush selects the primary keys that already exist, seeds the ones that do not from LiteLLM_SpendLogs so a new row cannot undercount spend that predates it, and applies one batch of upserts ordered by primary key. An increment at or behind the stored window_start adds into the row, matching how in-flight requests carry into a window after a reset; a newer one rolls the window and starts from that increment. The conflict arm adds only the increment, never the seeded base, so two pods seeding the same new window cannot double count it. The seed excludes the requests its own batch is about to apply. Spend logs are drained by a separate monitor that fires on a ~2s poll whenever anything is queued, while window increments flush on the much slower batch tick, so by seed time the batch's log rows are normally already in the table; counting them in the aggregate and again in the increments made a fresh row land at exactly twice the true spend. Each increment therefore carries the LiteLLM_SpendLogs request_id it was recorded under, which update_database now returns rather than having the callback re-derive it (a cache hit appends time.time() to that id, so a second derivation would not match). The reset job rolls each expired window's row alongside the counter it zeroes, conditional on the stored window_start still being behind the new one so a pod that already rolled it is not clobbered. --- litellm/constants.py | 1 + .../proxy/common_utils/reset_budget_job.py | 65 ++- .../proxy/db/budget_window_spend_writer.py | 266 +++++++++ litellm/proxy/db/db_spend_update_writer.py | 65 ++- .../redis_update_buffer.py | 46 +- .../window_spend_update_queue.py | 160 ++++++ .../proxy/hooks/proxy_track_cost_callback.py | 3 +- litellm/proxy/proxy_server.py | 80 ++- litellm/types/services.py | 2 + .../common_utils/test_reset_budget_job.py | 140 +++++ .../test_redis_update_buffer.py | 147 ++++- .../test_window_spend_update_queue.py | 249 ++++++++ .../db/test_budget_window_spend_writer.py | 535 ++++++++++++++++++ .../proxy/db/test_db_spend_update_writer.py | 259 +++++++++ .../hooks/test_proxy_track_cost_callback.py | 60 +- tests/test_litellm/proxy/test_proxy_server.py | 258 +++++++++ 16 files changed, 2317 insertions(+), 19 deletions(-) create mode 100644 litellm/proxy/db/budget_window_spend_writer.py create mode 100644 litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py create mode 100644 tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py create mode 100644 tests/test_litellm/proxy/db/test_budget_window_spend_writer.py diff --git a/litellm/constants.py b/litellm/constants.py index 0c7316455d6..2c8c2918a30 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -270,6 +270,7 @@ REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_end_user_spend_update_buffer" REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_update_buffer" REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer" +REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) # Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth LITELLM_ASYNCIO_QUEUE_MAXSIZE: Final = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 450dc12a9ff..0fb82c75274 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -2,16 +2,18 @@ import asyncio import json import time from collections.abc import Callable, Mapping, Sequence -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import Final, Literal, Protocol, TypeVar import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME +from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import ( LiteLLM_BudgetTableFull, LiteLLM_EndUserTable, + Litellm_EntityType, LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken, @@ -21,6 +23,7 @@ from litellm.proxy.common_utils.timezone_utils import ( compute_budget_reset_at, get_budget_reset_settings, ) +from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( @@ -782,6 +785,9 @@ class ResetBudgetJob: spend_counter_cache: DualCache, now: datetime, reset_settings: BudgetResetSettings, + prisma_client: PrismaClient, + entity_type: Litellm_EntityType, + entity_id: str, ) -> bool: """Reset a single budget window if expired. Returns True if the window was reset.""" reset_at_str: Final = window.get("reset_at") @@ -796,11 +802,56 @@ class ResetBudgetJob: await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=0.0) except Exception as redis_err: verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err) - window["reset_at"] = compute_budget_reset_at( - budget_duration=window["budget_duration"], settings=reset_settings - ).isoformat() + budget_duration: Final = window["budget_duration"] + next_reset_at: Final = compute_budget_reset_at(budget_duration=budget_duration, settings=reset_settings) + window["reset_at"] = next_reset_at.isoformat() + await ResetBudgetJob._roll_window_spend_row( + prisma_client=prisma_client, + entity_type=entity_type, + entity_id=entity_id, + budget_duration=budget_duration, + next_reset_at=next_reset_at, + ) return True + @staticmethod + async def _roll_window_spend_row( + prisma_client: PrismaClient, + entity_type: Litellm_EntityType, + entity_id: str, + budget_duration: str, + next_reset_at: datetime, + ) -> None: + """Move this window's LiteLLM_BudgetWindowSpend row onto the window + that just started, so the maintained total the read path uses starts + from zero alongside the counter. + + Best effort: the row is an optimization over aggregating + LiteLLM_SpendLogs, so a failure here must not stop the remaining + windows from having their counters reset. + """ + try: + window_start: Final = next_reset_at - timedelta(seconds=duration_in_seconds(budget_duration)) + except Exception as e: # noqa: BLE001 # duration_in_seconds raises bare exceptions on bad input + verbose_proxy_logger.warning("Unparseable budget_duration %s: %s", budget_duration, e) + return + try: + await roll_window_spend_row( + prisma_client=prisma_client, + entity_type=entity_type.value, + entity_id=entity_id, + window_duration=budget_duration, + new_window_start=window_start, + ) + except Exception as e: # noqa: BLE001 # the row is best effort; counter resets must still land + verbose_proxy_logger.warning( + "Failed to roll budget window spend row for %s=%s window=%s: %s", + entity_type.value, + entity_id, + budget_duration, + e, + ) + async def reset_budget_windows(self) -> None: """ For keys and teams with budget_limits, reset any individual windows where @@ -836,6 +887,9 @@ class ResetBudgetJob: spend_counter_cache, now, self.reset_settings, + prisma_client=self.prisma_client, + entity_type=Litellm_EntityType.KEY, + entity_id=row["token"], ): changed = True if changed: @@ -865,6 +919,9 @@ class ResetBudgetJob: spend_counter_cache, now, self.reset_settings, + prisma_client=self.prisma_client, + entity_type=Litellm_EntityType.TEAM, + entity_id=row["team_id"], ): changed = True if changed: diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py new file mode 100644 index 00000000000..0308651f4e3 --- /dev/null +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -0,0 +1,266 @@ +""" +Writer for LiteLLM_BudgetWindowSpend. + +The table holds one row per configured budget window whose window_start rolls +forward in place, so budget enforcement can read a maintained running total +instead of aggregating LiteLLM_SpendLogs every time a window counter goes cold +(issue #35766). Raw SQL rather than the Prisma upsert helper because the +conditional roll cannot be expressed through the query builder. + +Seeding a row that does not exist yet reads LiteLLM_SpendLogs once, excluding +the requests whose increments are in the same batch so neither source counts +them twice. One gap survives that exclusion: without the Redis transaction +buffer every pod flushes its own increments, so a row seeded by one pod can +miss increments still queued on another. That is bounded by a single flush +interval, happens at most once per window row, and errs toward under-counting +for that interval only. +""" + +from collections.abc import Sequence +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Final, Protocol + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendTransaction, + to_naive_utc, + window_spend_group_key, +) + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +_SELECT_EXISTING_ROWS_SQL: Final = ( + 'SELECT entity_type, entity_id, window_duration FROM "LiteLLM_BudgetWindowSpend" ' + "WHERE (entity_type, entity_id, window_duration) " + "IN (SELECT * FROM unnest($1::text[], $2::text[], $3::text[]))" +) + +_UPSERT_WINDOW_SPEND_SQL: Final = ( + 'INSERT INTO "LiteLLM_BudgetWindowSpend" ' + "(entity_type, entity_id, window_duration, window_start, spend, created_at, updated_at) " + "VALUES ($1, $2, $3, ($4::timestamptz AT TIME ZONE 'UTC'), $5, " + "($7::timestamptz AT TIME ZONE 'UTC'), ($7::timestamptz AT TIME ZONE 'UTC')) " + "ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE SET " + "spend = CASE " + 'WHEN "LiteLLM_BudgetWindowSpend".window_start >= EXCLUDED.window_start ' + 'THEN "LiteLLM_BudgetWindowSpend".spend + $6 ' + "ELSE EXCLUDED.spend " + "END, " + 'window_start = GREATEST("LiteLLM_BudgetWindowSpend".window_start, EXCLUDED.window_start), ' + "updated_at = ($7::timestamptz AT TIME ZONE 'UTC')" +) + +_ROLL_WINDOW_SPEND_SQL: Final = ( + 'UPDATE "LiteLLM_BudgetWindowSpend" SET ' + "window_start = ($4::timestamptz AT TIME ZONE 'UTC'), " + "spend = 0, " + "updated_at = ($5::timestamptz AT TIME ZONE 'UTC') " + "WHERE entity_type = $1 AND entity_id = $2 AND window_duration = $3 " + "AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')" +) + +_SEED_FROM_SPEND_LOGS_KEY_SQL: Final = ( + 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " + "AND NOT (request_id = ANY($3::text[]))" +) + +_SEED_FROM_SPEND_LOGS_TEAM_SQL: Final = ( + 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " + "AND NOT (request_id = ANY($3::text[]))" +) + +_UPSERT_TRANSACTION_TIMEOUT: Final = timedelta(seconds=60) + + +class WindowSpendLogsAggregate(Protocol): + """Sums LiteLLM_SpendLogs for one entity since window_start, ignoring the + requests whose ids are handed in. + + Injected so the flush can be exercised without a database and so the + expensive aggregate stays swappable. + """ + + async def __call__( + self, + prisma_client: "PrismaClient", + entity_type: str, + entity_id: str, + window_start: datetime, + exclude_request_ids: Sequence[str], + ) -> float | None: ... + + +async def spend_logs_total_excluding( + prisma_client: "PrismaClient", + entity_type: str, + entity_id: str, + window_start: datetime, + exclude_request_ids: Sequence[str], +) -> float | None: + """LiteLLM_SpendLogs spend for one entity since window_start, minus the + requests already accounted for by the increments being flushed. + + The spend log writer drains its own queue on a ~2s poll whenever anything + is queued, while window increments flush on the much slower batch tick, so + by the time a window row is seeded its batch's log rows are normally + already in the table. Counting them in the seed and again in the increment + is what made a fresh row land at twice the true spend. + """ + if entity_type == Litellm_EntityType.KEY.value: + rows = await prisma_client.db.query_raw( + _SEED_FROM_SPEND_LOGS_KEY_SQL, entity_id, window_start, tuple(exclude_request_ids) + ) + elif entity_type == Litellm_EntityType.TEAM.value: + rows = await prisma_client.db.query_raw( + _SEED_FROM_SPEND_LOGS_TEAM_SQL, entity_id, window_start, tuple(exclude_request_ids) + ) + else: + return None + if not rows: + return 0.0 + return float(rows[0].get("total") or 0.0) + + +def _primary_key(transaction: WindowSpendTransaction) -> tuple[str, str, str]: + return ( + transaction["entity_type"], + transaction["entity_id"], + transaction["window_duration"], + ) + + +async def _existing_primary_keys( + prisma_client: "PrismaClient", + transactions: tuple[WindowSpendTransaction, ...], +) -> frozenset[tuple[str, str, str]]: + rows: Final = await prisma_client.db.query_raw( + _SELECT_EXISTING_ROWS_SQL, + tuple(transaction["entity_type"] for transaction in transactions), + tuple(transaction["entity_id"] for transaction in transactions), + tuple(transaction["window_duration"] for transaction in transactions), + ) + return frozenset((row["entity_type"], row["entity_id"], row["window_duration"]) for row in rows or ()) + + +async def _seed_base_for_missing_row( + prisma_client: "PrismaClient", + transaction: WindowSpendTransaction, + existing_primary_keys: frozenset[tuple[str, str, str]], + spend_logs_aggregate: WindowSpendLogsAggregate, +) -> float: + """Spend already recorded for a window that has no row yet. + + This is the LiteLLM_SpendLogs aggregate the window counter reseed runs on + every cold counter today, but here it runs once per window lifetime and off + the request path, and it excludes this batch's own requests so they are + counted by their increments alone. + """ + if _primary_key(transaction) in existing_primary_keys: + return 0.0 + base: Final = await spend_logs_aggregate( + prisma_client=prisma_client, + entity_type=transaction["entity_type"], + entity_id=transaction["entity_id"], + window_start=datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc), + exclude_request_ids=transaction["request_ids"], + ) + return float(base or 0.0) + + +def _upsert_params( + transaction: WindowSpendTransaction, + seed_base: float, + now: datetime, +) -> tuple[str, str, str, datetime, float, float, datetime]: + """$5 is what a brand new row starts at (pre-existing spend plus this + increment); $6 is the increment alone, which is all an already-current row + may add. They are equal for every row that already existed, so a row is + never seeded twice when two pods flush the same new window.""" + increment: Final = float(transaction["spend"]) + return ( + transaction["entity_type"], + transaction["entity_id"], + transaction["window_duration"], + datetime.fromisoformat(transaction["window_start"]), + seed_base + increment, + increment, + now, + ) + + +async def commit_window_spend_updates( + prisma_client: "PrismaClient", + transactions: Sequence[WindowSpendTransaction], + spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_total_excluding, +) -> None: + """Apply aggregated window increments to LiteLLM_BudgetWindowSpend. + + An increment at or behind the row's window_start adds into the row (this is + how in-flight requests that raced a reset carry into the new window); an + increment ahead of it rolls the window and starts from that increment. + + Statements are ordered by primary key so concurrent pods take row locks in + the same order, with window_start breaking ties so an older window is + applied before the roll that supersedes it. + """ + if not transactions: + return + + ordered: Final = tuple(sorted(transactions, key=window_spend_group_key)) + existing_primary_keys: Final = await _existing_primary_keys( + prisma_client=prisma_client, + transactions=ordered, + ) + seed_bases: Final = tuple( + [ + await _seed_base_for_missing_row( + prisma_client=prisma_client, + transaction=transaction, + existing_primary_keys=existing_primary_keys, + spend_logs_aggregate=spend_logs_aggregate, + ) + for transaction in ordered + ] + ) + + now: Final = to_naive_utc(datetime.now(timezone.utc)) + verbose_proxy_logger.debug( + "Spend tracking - committing %d budget window spend upserts over %d existing rows", + len(ordered), + len(existing_primary_keys), + ) + async with prisma_client.db.tx(timeout=_UPSERT_TRANSACTION_TIMEOUT) as db_transaction: + async with db_transaction.batch_() as batcher: + for transaction, seed_base in zip(ordered, seed_bases): + batcher.execute_raw( + _UPSERT_WINDOW_SPEND_SQL, + *_upsert_params(transaction=transaction, seed_base=seed_base, now=now), + ) + + +async def roll_window_spend_row( + prisma_client: "PrismaClient", + entity_type: str, + entity_id: str, + window_duration: str, + new_window_start: datetime, +) -> None: + """Move a row onto the window that just started and zero its spend. + + Conditional on the stored window_start still being behind the new one so a + pod that already rolled the row (or increments that arrived under the new + window) are not clobbered. + """ + await prisma_client.db.execute_raw( + _ROLL_WINDOW_SPEND_SQL, + entity_type, + entity_id, + window_duration, + to_naive_utc(new_window_start), + to_naive_utc(datetime.now(timezone.utc)), + ) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 4e5c5daaa32..4e0137e33c6 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,7 +12,7 @@ import os import random import time import traceback -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload @@ -50,6 +50,10 @@ from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdate from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( ToolDiscoveryQueue, ) +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendTransaction, + WindowSpendUpdateQueue, +) from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING from litellm.proxy.spend_tracking.compression_savings import ( extract_compression_saved_tokens, @@ -132,6 +136,7 @@ class DBSpendUpdateWriter: self.daily_agent_spend_update_queue = DailySpendUpdateQueue() self.daily_org_spend_update_queue = DailySpendUpdateQueue() self.daily_tag_spend_update_queue = DailySpendUpdateQueue() + self.window_spend_update_queue = WindowSpendUpdateQueue() async def update_database( # LiteLLM management object fields @@ -147,7 +152,11 @@ class DBSpendUpdateWriter: start_time: datetime | None, end_time: datetime | None, response_cost: float | None, - ): + ) -> str | None: + """Returns the LiteLLM_SpendLogs request_id this call was recorded + under, so the caller can tell the budget-window writer which log rows + its increments already cover. None when the payload could not be built. + """ from litellm.proxy.proxy_server import ( disable_spend_logs, litellm_proxy_budget_name, @@ -164,7 +173,7 @@ class DBSpendUpdateWriter: team_id, ) if ProxyUpdateSpend.disable_spend_updates() is True: - return + return None if token is not None and isinstance(token, str) and token.startswith("sk-"): hashed_token = hash_token(token=token) else: @@ -232,6 +241,7 @@ class DBSpendUpdateWriter: ) verbose_proxy_logger.debug("Runs spend update on all tables") + return payload.get("request_id") except Exception: spend_log_error( "Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue " @@ -244,6 +254,7 @@ class DBSpendUpdateWriter: org_id, end_user_id, ) + return None async def _enqueue_tool_usage_transaction( self, @@ -831,6 +842,7 @@ class DBSpendUpdateWriter: daily_org_spend_update_queue=self.daily_org_spend_update_queue, daily_end_user_spend_update_queue=self.daily_end_user_spend_update_queue, daily_agent_spend_update_queue=self.daily_agent_spend_update_queue, + window_spend_update_queue=self.window_spend_update_queue, ) # Only commit from redis to db if this pod is the leader @@ -847,6 +859,7 @@ class DBSpendUpdateWriter: daily_org_spend_update_transactions, daily_end_user_spend_update_transactions, daily_agent_spend_update_transactions, + window_spend_update_transactions, ) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() if db_spend_update_transactions is not None: @@ -906,6 +919,11 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_agent_spend_update_transactions, ) + if window_spend_update_transactions is not None: + await DBSpendUpdateWriter._commit_window_spend_updates( + prisma_client=prisma_client, + window_spend_transactions=window_spend_update_transactions, + ) except Exception as e: spend_log_error( "Spend tracking - failed to commit spend updates from Redis to DB. " @@ -1016,6 +1034,17 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_agent_spend_update_transactions, ) + ################## Budget Window Spend Update Transactions ################## + # Aggregate all in memory budget window spend transactions and commit to db + window_spend_update_transactions: Final = ( + await self.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() + ) + + await DBSpendUpdateWriter._commit_window_spend_updates( + prisma_client=prisma_client, + window_spend_transactions=window_spend_update_transactions, + ) + ################## Tool Registry Upserts ################## await self._flush_tool_discovery_queue(prisma_client=prisma_client) @@ -1086,6 +1115,36 @@ class DBSpendUpdateWriter: cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME, ) + @staticmethod + async def _commit_window_spend_updates( + prisma_client: PrismaClient, + window_spend_transactions: Sequence[WindowSpendTransaction], + ) -> None: + """ + Commit per-budget-window spend increments to LiteLLM_BudgetWindowSpend. + + Failures are logged rather than raised: these rows exist so budget + enforcement can stop aggregating LiteLLM_SpendLogs, and the read path + falls back to that aggregate, so a failed commit must not abort the + entity and daily spend commits that share this scheduler tick. + """ + from litellm.proxy.db.budget_window_spend_writer import ( + commit_window_spend_updates, + ) + + try: + await commit_window_spend_updates( + prisma_client=prisma_client, + transactions=window_spend_transactions, + ) + except Exception as e: # noqa: BLE001 # any DB failure here must stay contained to the window rows + spend_log_error( + "Spend tracking - failed to commit budget window spend updates. %d window increments lost. Error: %s", + len(window_spend_transactions), + str(e), + exc=e, + ) + async def _flush_tool_discovery_queue( self, prisma_client: PrismaClient, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index d64280efa8c..d5786aa7088 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -19,6 +19,7 @@ from litellm.constants import ( REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY, + REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( @@ -38,6 +39,10 @@ from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, ) from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendTransaction, + WindowSpendUpdateQueue, +) from litellm.secret_managers.main import str_to_bool from litellm.types.caching import ( RedisPipelineLpopOperation, @@ -130,6 +135,7 @@ class RedisUpdateBuffer: daily_org_spend_update_queue: DailySpendUpdateQueue, daily_end_user_spend_update_queue: DailySpendUpdateQueue, daily_agent_spend_update_queue: DailySpendUpdateQueue, + window_spend_update_queue: WindowSpendUpdateQueue | None = None, ): """ Stores the in-memory spend updates to Redis @@ -196,6 +202,11 @@ class RedisUpdateBuffer: daily_agent_spend_update_transactions: Final = ( await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) + window_spend_update_transactions: Final = ( + await window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() + if window_spend_update_queue is not None + else () + ) verbose_proxy_logger.debug("ALL DB SPEND UPDATE TRANSACTIONS: %s", db_spend_update_transactions) verbose_proxy_logger.debug("ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions) @@ -232,6 +243,11 @@ class RedisUpdateBuffer: REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE, ), + ( + window_spend_update_transactions, + REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, + ServiceTypes.REDIS_WINDOW_SPEND_UPDATE_QUEUE, + ), ] rpush_list: Final[list[RedisPipelineRpushOperation]] = [] @@ -272,12 +288,14 @@ class RedisUpdateBuffer: daily_org_spend_update_transactions=daily_org_spend_update_transactions, daily_end_user_spend_update_transactions=daily_end_user_spend_update_transactions, daily_agent_spend_update_transactions=daily_agent_spend_update_transactions, + window_spend_update_transactions=window_spend_update_transactions, spend_update_queue=spend_update_queue, daily_spend_update_queue=daily_spend_update_queue, daily_team_spend_update_queue=daily_team_spend_update_queue, daily_org_spend_update_queue=daily_org_spend_update_queue, daily_end_user_spend_update_queue=daily_end_user_spend_update_queue, daily_agent_spend_update_queue=daily_agent_spend_update_queue, + window_spend_update_queue=window_spend_update_queue, ) return @@ -297,12 +315,14 @@ class RedisUpdateBuffer: daily_org_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None, daily_end_user_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None, daily_agent_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None, + window_spend_update_transactions: tuple[WindowSpendTransaction, ...] | None, spend_update_queue: SpendUpdateQueue, daily_spend_update_queue: DailySpendUpdateQueue, daily_team_spend_update_queue: DailySpendUpdateQueue, daily_org_spend_update_queue: DailySpendUpdateQueue, daily_end_user_spend_update_queue: DailySpendUpdateQueue, daily_agent_spend_update_queue: DailySpendUpdateQueue, + window_spend_update_queue: WindowSpendUpdateQueue | None, ) -> None: """ Put drained-but-unpushed transactions back into in-memory queues. @@ -372,6 +392,9 @@ class RedisUpdateBuffer: if daily_txns: await daily_queue.update_queue.put(daily_txns) + if window_spend_update_transactions and window_spend_update_queue is not None: + await window_spend_update_queue.update_queue.put(window_spend_update_transactions) + @staticmethod def _number_of_transactions_to_store_in_redis( db_spend_update_transactions: DBSpendUpdateTransactions, @@ -468,20 +491,22 @@ class RedisUpdateBuffer: dict[str, DailyOrganizationSpendTransaction] | None, dict[str, DailyEndUserSpendTransaction] | None, dict[str, DailyAgentSpendTransaction] | None, + tuple[WindowSpendTransaction, ...] | None, ]: """ - Drains the main 6 Redis buffer queues in a single pipeline round-trip. + Drains the main 7 Redis buffer queues in a single pipeline round-trip. - Returns a 6-tuple of parsed results in this order: + Returns a 7-tuple of parsed results in this order: 0: DBSpendUpdateTransactions 1: daily user spend 2: daily team spend 3: daily org spend 4: daily end-user spend 5: daily agent spend + 6: budget window spend """ if self.redis_cache is None: - return None, None, None, None, None, None + return None, None, None, None, None, None, None lpop_list: Final[list[RedisPipelineLpopOperation]] = [ RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), @@ -505,12 +530,16 @@ class RedisUpdateBuffer: key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ), + RedisPipelineLpopOperation( + key=REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, + count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + ), ] raw_results: Final = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list) # Pad with None if pipeline returned fewer results than expected - while len(raw_results) < 6: + while len(raw_results) < 7: raw_results.append(None) # Slot 0: DBSpendUpdateTransactions @@ -531,6 +560,14 @@ class RedisUpdateBuffer: aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(list_of_daily) daily_results.append(aggregated) + window_spend: Final = ( + WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( + tuple(json.loads(transaction) for transaction in raw_results[6]) + ) + if raw_results[6] is not None + else None + ) + return ( db_spend, cast(dict[str, DailyUserSpendTransaction] | None, daily_results[0]), @@ -538,6 +575,7 @@ class RedisUpdateBuffer: cast(dict[str, DailyOrganizationSpendTransaction] | None, daily_results[2]), cast(dict[str, DailyEndUserSpendTransaction] | None, daily_results[3]), cast(dict[str, DailyAgentSpendTransaction] | None, daily_results[4]), + window_spend, ) async def store_in_memory_daily_tag_spend_updates_in_redis( diff --git a/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py new file mode 100644 index 00000000000..d2ad2aa61d2 --- /dev/null +++ b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py @@ -0,0 +1,160 @@ +""" +In memory buffer for per-budget-window spend increments. + +Kept separate from SpendUpdateQueue: an increment is only meaningful together +with the window it landed in, so two increments for the same entity must not be +merged when their window_start differs. +""" + +import asyncio +import math +from collections.abc import Sequence +from datetime import datetime, timezone +from itertools import chain, groupby +from typing import Final, TypedDict + +from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE +from litellm.proxy.db.db_transaction_queue.base_update_queue import BaseUpdateQueue + + +class WindowSpendTransaction(TypedDict): + """One increment for a single (entity, budget window) pair. + + window_start is an ISO-8601 string rather than a datetime so the + transaction survives the JSON round trip through the Redis buffer. + + request_ids carries the LiteLLM_SpendLogs ids this spend came from. The + one-time seed for a window that has no row yet subtracts them from its + LiteLLM_SpendLogs aggregate, because the spend log writer flushes on its + own ~2s poll and will usually have persisted these rows before the window + queue flushes; without the exclusion the seed and the increment would each + count them. + """ + + entity_type: str + entity_id: str + window_duration: str + window_start: str + spend: float + request_ids: Sequence[str] + + +def to_naive_utc(value: datetime) -> datetime: + """LiteLLM_BudgetWindowSpend.window_start is TIMESTAMP(3), which holds naive UTC.""" + if value.tzinfo is None: + return value + return value.astimezone(timezone.utc).replace(tzinfo=None) + + +def window_spend_group_key(transaction: WindowSpendTransaction) -> tuple[str, str, str, str]: + """Identity of a window increment: the row's primary key plus the window it + belongs to. Two increments only aggregate when all four match.""" + return ( + transaction["entity_type"], + transaction["entity_id"], + transaction["window_duration"], + transaction["window_start"], + ) + + +def build_window_spend_transaction( + entity_type: str, + entity_id: str, + window_duration: str, + window_start: datetime, + spend: float, + request_id: str | None = None, +) -> WindowSpendTransaction: + return WindowSpendTransaction( + entity_type=entity_type, + entity_id=entity_id, + window_duration=window_duration, + window_start=to_naive_utc(window_start).isoformat(timespec="microseconds"), + spend=spend, + request_ids=() if request_id is None else (request_id,), + ) + + +def _merge_window_spend_transactions( + payloads: tuple[WindowSpendTransaction, ...], +) -> WindowSpendTransaction: + first: Final = payloads[0] + return WindowSpendTransaction( + entity_type=first["entity_type"], + entity_id=first["entity_id"], + window_duration=first["window_duration"], + window_start=first["window_start"], + spend=math.fsum(payload["spend"] for payload in payloads), + request_ids=tuple(sorted(frozenset(chain.from_iterable(payload["request_ids"] for payload in payloads)))), + ) + + +class WindowSpendUpdateQueue(BaseUpdateQueue): + """ + In memory buffer for budget-window spend increments committed to + LiteLLM_BudgetWindowSpend. + + Add an update with the payload built by build_window_spend_transaction: + window_spend_update_queue.add_update( + build_window_spend_transaction( + entity_type="key", + entity_id="", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=0.02, + ) + ) + """ + + def __init__(self): + super().__init__() + self.update_queue: asyncio.Queue[tuple[WindowSpendTransaction, ...]] = asyncio.Queue( + maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE + ) + + async def add_update(self, update: WindowSpendTransaction) -> None: + """Enqueue an update.""" + verbose_proxy_logger.debug("Adding budget window spend update to queue: %s", update) + await self.update_queue.put((update,)) + if self.update_queue.qsize() >= self.MAX_SIZE_IN_MEMORY_QUEUE: + verbose_proxy_logger.warning( + "Budget window spend update queue is full. Aggregating all entries in queue to concatenate entries." + ) + await self.aggregate_queue_updates() + + async def aggregate_queue_updates(self) -> None: + """Collapse everything currently queued into a single aggregated update.""" + updates: Final = await self.flush_all_updates_from_in_memory_queue() + await self.update_queue.put(WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(updates)) + + async def flush_and_get_aggregated_window_spend_transactions( + self, + ) -> tuple[WindowSpendTransaction, ...]: + """Drain the queue and return the increments aggregated per window.""" + updates: Final = await self.flush_all_updates_from_in_memory_queue() + if len(updates) > 0: + verbose_proxy_logger.info( + "Spend tracking - flushed %d budget window spend update batches from in-memory queue", + len(updates), + ) + return WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(updates) + + @staticmethod + def get_aggregated_window_spend_transactions( + updates: Sequence[Sequence[WindowSpendTransaction]], + ) -> tuple[WindowSpendTransaction, ...]: + """Sum spend per (entity_type, entity_id, window_duration, window_start). + + Increments belonging to different windows stay separate even when they + share a primary key, so a window boundary crossed mid-tick does not fold + the new window's spend into the previous window's total. + + The result is ordered by that same key, which is the order the flush + needs: primary key first for cross-pod lock ordering, then window_start + so an older window is applied before the roll that supersedes it. + """ + ordered: Final = tuple(sorted(chain.from_iterable(updates), key=window_spend_group_key)) + return tuple( + _merge_window_spend_transactions(tuple(group)) for _, group in groupby(ordered, key=window_spend_group_key) + ) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3a4914d56a8..58bc9cfb301 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -498,7 +498,7 @@ async def _update_database_and_spend_counters( request_tags: list[str] | None = None, ) -> None: try: - await proxy_logging_obj.db_spend_update_writer.update_database( + spend_log_request_id = await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key, response_cost=response_cost, user_id=user_id, @@ -534,6 +534,7 @@ async def _update_database_and_spend_counters( budget_reservation=budget_reservation, end_user_id=end_user_id, tags=request_tags, + request_id=spend_log_request_id, ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fb9c4e67aad..b3b59d72a53 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -349,6 +349,9 @@ from litellm.proxy.config_resolvers.alerting import ( from litellm.proxy.container_endpoints.endpoints import router as container_router from litellm.proxy.credential_endpoints.endpoints import router as credential_router from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + build_window_spend_transaction, +) from litellm.proxy.db.exception_handler import ( PrismaDBExceptionHandler, call_with_db_reconnect_retry, @@ -2392,6 +2395,7 @@ async def increment_spend_counters( budget_reservation: dict | None = None, end_user_id: str | None = None, tags: list[str] | None = None, + request_id: str | None = None, ): """ Atomically increment spend counters for budget enforcement. @@ -2446,14 +2450,24 @@ async def increment_spend_counters( for window in key_budget_limits: duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration key_window_counter = f"spend:key:{hashed_token}:window:{duration}" + key_window_start = get_budget_window_start(window) if key_window_counter not in reserved_counter_keys: await _init_and_increment_window_spend_counter( counter_key=key_window_counter, entity_type="Key", entity_id=hashed_token, - window_start=get_budget_window_start(window), + window_start=key_window_start, increment=cost, ) + await _enqueue_window_spend_row_update( + entity_type=Litellm_EntityType.KEY, + entity_id=hashed_token, + window=window, + window_duration=duration, + window_start=key_window_start, + increment=cost, + request_id=request_id, + ) async def _team_scope(scope_team_id: str) -> None: team_counter_key: Final = f"spend:team:{scope_team_id}" @@ -2477,14 +2491,24 @@ async def increment_spend_counters( for window in team_budget_limits: duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration team_window_counter = f"spend:team:{scope_team_id}:window:{duration}" + team_window_start = get_budget_window_start(window) if team_window_counter not in reserved_counter_keys: await _init_and_increment_window_spend_counter( counter_key=team_window_counter, entity_type="Team", entity_id=scope_team_id, - window_start=get_budget_window_start(window), + window_start=team_window_start, increment=cost, ) + await _enqueue_window_spend_row_update( + entity_type=Litellm_EntityType.TEAM, + entity_id=scope_team_id, + window=window, + window_duration=duration, + window_start=team_window_start, + increment=cost, + request_id=request_id, + ) async def _team_member_scope(scope_user_id: str, scope_team_id: str) -> None: team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}" @@ -2669,6 +2693,58 @@ async def _init_and_increment_spend_counter( await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) +async def _enqueue_window_spend_row_update( + entity_type: Litellm_EntityType, + entity_id: str, + window: Any, + window_duration: str, + window_start: datetime | None, + increment: float, + request_id: str | None, +) -> None: + """Queue this request's cost against the LiteLLM_BudgetWindowSpend row for + the window, so enforcement can read a maintained total instead of + aggregating LiteLLM_SpendLogs. + + request_id is the LiteLLM_SpendLogs id this cost was recorded under; the + flush uses it to keep the one-time seed from counting a request that its + increment already covers. + + Enqueued even when the cache increment was skipped for a reserved counter: + the reservation only pre-charged the counter, and the row still owes the + actual cost. + + Windows with no reset_at slide with wall clock, so their window_start moves + on every request and no single row can represent them. Those are left to + the read path's LiteLLM_SpendLogs fallback rather than rewritten per + request. + """ + if window_start is None: + return + reset_at: Final = window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) + if not reset_at: + return + try: + await proxy_logging_obj.db_spend_update_writer.window_spend_update_queue.add_update( + build_window_spend_transaction( + entity_type=entity_type.value, + entity_id=entity_id, + window_duration=window_duration, + window_start=window_start, + spend=increment, + request_id=request_id, + ) + ) + except Exception as e: # noqa: BLE001 # spend tracking must never fail the cost callback + verbose_proxy_logger.debug( + "Unable to enqueue budget window spend update for %s=%s window=%s: %s", + entity_type.value, + entity_id, + window_duration, + e, + ) + + async def _init_and_increment_window_spend_counter( counter_key: str, entity_type: str, diff --git a/litellm/types/services.py b/litellm/types/services.py index f43494a0a09..60cfc35c0a5 100644 --- a/litellm/types/services.py +++ b/litellm/types/services.py @@ -42,6 +42,8 @@ class ServiceTypes(str, enum.Enum): # spend update queue - current spend of key, user, team IN_MEMORY_SPEND_UPDATE_QUEUE = "in_memory_spend_update_queue" REDIS_SPEND_UPDATE_QUEUE = "redis_spend_update_queue" + # budget window spend queue - per-window spend of key, team + REDIS_WINDOW_SPEND_UPDATE_QUEUE = "redis_window_spend_update_queue" class ServiceConfig(TypedDict): diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index f04d6f3cf5a..c7d31972d87 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1082,6 +1082,7 @@ def _make_reset_budget_windows_job( raise AssertionError(f"Unexpected query_raw call: {query}") prisma_client.db.query_raw = AsyncMock(side_effect=fake_query_raw) + prisma_client.db.execute_raw = AsyncMock(return_value=1) prisma_client.db.litellm_verificationtoken.update = AsyncMock(return_value=None) prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None) @@ -1153,6 +1154,145 @@ def test_reset_budget_windows_resets_expired_key_window(monkeypatch): spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-expired:window:1d", value=0.0) +def _window_spend_rolls(prisma_client): + return [ + call.args + for call in prisma_client.db.execute_raw.await_args_list + if "LiteLLM_BudgetWindowSpend" in call.args[0] + ] + + +def test_reset_budget_windows_rolls_the_key_window_spend_row(monkeypatch): + """The maintained per-window total has to start the new window at zero + alongside the counter, or enforcement keeps reading the old window's spend.""" + now = datetime.utcnow() + expired = (now - timedelta(minutes=5)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-expired", + "budget_limits": [{"budget_duration": "1d", "reset_at": expired}], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) + + asyncio.run(job.reset_budget_windows()) + + rolls = _window_spend_rolls(prisma_client) + assert len(rolls) == 1 + query, entity_type, entity_id, window_duration, new_window_start, _updated_at = rolls[0] + assert (entity_type, entity_id, window_duration) == ("key", "sk-expired", "1d") + assert "spend = 0" in " ".join(query.split()) + + # window_start is the start of the window that just began: new reset_at minus the duration. + written_windows = json.loads( + prisma_client.db.litellm_verificationtoken.update.await_args.kwargs["data"]["budget_limits"] + ) + new_reset_at = datetime.fromisoformat(written_windows[0]["reset_at"].replace("Z", "+00:00")).replace(tzinfo=None) + assert new_window_start == pytest.approx( + new_reset_at - timedelta(days=1), + abs=timedelta(seconds=1), + ) + + +def test_reset_budget_windows_roll_is_conditional_on_an_older_stored_window(monkeypatch): + """Another pod may already have rolled the row; clobbering it would drop + spend that landed under the new window.""" + now = datetime.utcnow() + expired = (now - timedelta(minutes=5)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-expired", + "budget_limits": [{"budget_duration": "1d", "reset_at": expired}], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) + + asyncio.run(job.reset_budget_windows()) + + query = " ".join(_window_spend_rolls(prisma_client)[0][0].split()) + assert "AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')" in query + + +def test_reset_budget_windows_rolls_the_team_window_spend_row(monkeypatch): + now = datetime.utcnow() + expired = (now - timedelta(minutes=1)).isoformat() + "Z" + + team_rows = [ + { + "team_id": "team-expired", + "budget_limits": [{"budget_duration": "30d", "reset_at": expired}], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=[], team_rows=team_rows) + + asyncio.run(job.reset_budget_windows()) + + rolls = _window_spend_rolls(prisma_client) + assert len(rolls) == 1 + assert rolls[0][1:4] == ("team", "team-expired", "30d") + + +def test_reset_budget_windows_does_not_roll_an_unexpired_window(monkeypatch): + now = datetime.utcnow() + future = (now + timedelta(hours=1)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-future", + "budget_limits": [{"budget_duration": "1d", "reset_at": future}], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) + + asyncio.run(job.reset_budget_windows()) + + assert _window_spend_rolls(prisma_client) == [] + + +def test_reset_budget_windows_rolls_only_the_expired_window_of_a_key(monkeypatch): + now = datetime.utcnow() + key_rows = [ + { + "token": "sk-mixed", + "budget_limits": [ + {"budget_duration": "1d", "reset_at": (now - timedelta(minutes=5)).isoformat() + "Z"}, + {"budget_duration": "30d", "reset_at": (now + timedelta(days=2)).isoformat() + "Z"}, + ], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) + + asyncio.run(job.reset_budget_windows()) + + rolls = _window_spend_rolls(prisma_client) + assert [roll[3] for roll in rolls] == ["1d"] + + +def test_reset_budget_windows_survives_a_failed_window_spend_roll(monkeypatch): + """The row is an optimization over aggregating LiteLLM_SpendLogs; a DB + failure there must not stop the counter reset from being persisted.""" + now = datetime.utcnow() + expired = (now - timedelta(minutes=5)).isoformat() + "Z" + + key_rows = [ + { + "token": "sk-expired", + "budget_limits": [{"budget_duration": "1d", "reset_at": expired}], + } + ] + job, prisma_client, spend_counter_cache = _make_reset_budget_windows_job( + monkeypatch, key_rows=key_rows, team_rows=[] + ) + prisma_client.db.execute_raw = AsyncMock(side_effect=Exception("connection reset")) + + asyncio.run(job.reset_budget_windows()) + + prisma_client.db.litellm_verificationtoken.update.assert_awaited_once() + spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-expired:window:1d", value=0.0) + + def test_reset_budget_windows_skips_unexpired_key_window(monkeypatch): """If `reset_at` is in the future, no write should happen for that key.""" now = datetime.utcnow() diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 33372e7794a..e4a653ddb51 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -208,7 +208,8 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( Verify get_all_transactions_from_redis_buffer_pipeline correctly parses and aggregates results from async_lpop_pipeline. """ - # Simulate pipeline results: slot 0 = spend updates, slots 1-5 = daily categories + # Simulate pipeline results: slot 0 = spend updates, slots 1-5 = daily categories, + # slot 6 = budget window spend db_spend_json = json.dumps( { "key_list_transactions": {"key1": 1.0, "key2": 2.0}, @@ -222,6 +223,18 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( ) daily_user_json = json.dumps({"user_key1": {"spend": 1.0, "api_requests": 1}}) daily_team_json = json.dumps({"team_key1": {"spend": 2.0, "api_requests": 2}}) + window_spend_json = json.dumps( + [ + { + "entity_type": "key", + "entity_id": "hashed-token", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 3.0, + "request_ids": ["req-1"], + } + ] + ) mock_redis_cache.async_lpop_pipeline = AsyncMock( return_value=[ @@ -231,13 +244,30 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( None, # slot 3: daily org (empty) None, # slot 4: daily end-user (empty) None, # slot 5: daily agent (empty) + [window_spend_json, window_spend_json], # slot 6: budget window spend ] ) result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() - assert len(result) == 6 - db_spend, daily_user, daily_team, daily_org, daily_end_user, daily_agent = result + assert len(result) == 7 + ( + db_spend, + daily_user, + daily_team, + daily_org, + daily_end_user, + daily_agent, + window_spend, + ) = result + + # Budget window spend from two pods is summed per window, not overwritten, + # and both pods' request ids reach the seed exclusion. + assert window_spend is not None + assert len(window_spend) == 1 + assert window_spend[0]["spend"] == 6.0 + assert window_spend[0]["entity_id"] == "hashed-token" + assert window_spend[0]["request_ids"] == ("req-1",) # Verify db spend was parsed correctly assert db_spend is not None @@ -260,6 +290,12 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( # Verify pipeline was called once with correct keys mock_redis_cache.async_lpop_pipeline.assert_called_once() + from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY + + popped_keys = [ + op["key"] for op in mock_redis_cache.async_lpop_pipeline.call_args.kwargs["lpop_list"] + ] + assert popped_keys[6] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY @pytest.mark.asyncio @@ -267,7 +303,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis(): """When redis_cache is None, should return all Nones""" buffer = RedisUpdateBuffer(redis_cache=None) result = await buffer.get_all_transactions_from_redis_buffer_pipeline() - assert result == (None, None, None, None, None, None) + assert result == (None, None, None, None, None, None, None) def test_validate_redis_transaction_buffer_raises_without_redis(): @@ -374,3 +410,106 @@ def test_get_transaction_buffer_redis_cache_parses_string_flag(monkeypatch): mock_redis_cache.assert_called_once() assert result is mock_redis_cache.return_value + + +@pytest.mark.asyncio +async def test_store_in_memory_spend_updates_pushes_budget_window_spend( + redis_update_buffer, mock_redis_cache +): + """The budget window queue has to ride the same rpush as the daily queues, + otherwise multi-pod deployments never persist per-window spend.""" + from datetime import datetime, timezone + + from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendUpdateQueue, + build_window_spend_transaction, + ) + + mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1]) + + empty_queue = AsyncMock() + empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={}) + empty_daily_queue = AsyncMock() + empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) + + window_queue = WindowSpendUpdateQueue() + await window_queue.add_update( + build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=1.25, + request_id="req-1", + ) + ) + + await redis_update_buffer.store_in_memory_spend_updates_in_redis( + spend_update_queue=empty_queue, + daily_spend_update_queue=empty_daily_queue, + daily_team_spend_update_queue=empty_daily_queue, + daily_org_spend_update_queue=empty_daily_queue, + daily_end_user_spend_update_queue=empty_daily_queue, + daily_agent_spend_update_queue=empty_daily_queue, + window_spend_update_queue=window_queue, + ) + + rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] + assert len(rpush_list) == 1 + assert rpush_list[0]["key"] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY + pushed = json.loads(rpush_list[0]["values"][0]) + assert pushed == [{ + "entity_type": "key", + "entity_id": "hashed-token", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 1.25, + "request_ids": ["req-1"], + }] + + +@pytest.mark.asyncio +async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpush_failure( + redis_update_buffer, mock_redis_cache +): + """The window queue is drained before the rpush, so a Redis hiccup would + silently drop per-window spend without the restore.""" + from datetime import datetime, timezone + + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendUpdateQueue, + build_window_spend_transaction, + ) + + mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away")) + + empty_queue = AsyncMock() + empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={}) + empty_daily_queue = AsyncMock() + empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) + + window_queue = WindowSpendUpdateQueue() + await window_queue.add_update( + build_window_spend_transaction( + entity_type="team", + entity_id="team-1", + window_duration="7d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=4.0, + ) + ) + + await redis_update_buffer.store_in_memory_spend_updates_in_redis( + spend_update_queue=empty_queue, + daily_spend_update_queue=empty_daily_queue, + daily_team_spend_update_queue=empty_daily_queue, + daily_org_spend_update_queue=empty_daily_queue, + daily_end_user_spend_update_queue=empty_daily_queue, + daily_agent_spend_update_queue=empty_daily_queue, + window_spend_update_queue=window_queue, + ) + + restored = await window_queue.flush_and_get_aggregated_window_spend_transactions() + assert [payload["spend"] for payload in restored] == [4.0] + assert [payload["entity_id"] for payload in restored] == ["team-1"] diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py new file mode 100644 index 00000000000..880cc5044f9 --- /dev/null +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py @@ -0,0 +1,249 @@ +import json +import os +import sys +from datetime import datetime, timedelta, timezone + +sys.path.insert(0, os.path.abspath("../../../../..")) + +import pytest + +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendUpdateQueue, + build_window_spend_transaction, + to_naive_utc, +) + +WINDOW_A = datetime(2026, 8, 1, tzinfo=timezone.utc) +WINDOW_B = datetime(2026, 8, 31, tzinfo=timezone.utc) + + +def _txn( + entity_id: str, + window_start: datetime, + spend: float, + duration: str = "30d", + entity_type: str = "key", + request_id: str | None = None, +): + return build_window_spend_transaction( + entity_type=entity_type, + entity_id=entity_id, + window_duration=duration, + window_start=window_start, + spend=spend, + request_id=request_id, + ) + + +def test_build_window_spend_transaction_stores_naive_utc_iso(): + """window_start rides the Redis buffer as a string and lands in a naive-UTC + TIMESTAMP(3) column, so a non-UTC input must be converted, not truncated.""" + non_utc = datetime(2026, 8, 1, 20, 0, tzinfo=timezone(timedelta(hours=-4))) + + assert _txn("k1", non_utc, 1.0, request_id="req-1") == { + "entity_type": "key", + "entity_id": "k1", + "window_duration": "30d", + "window_start": "2026-08-02T00:00:00.000000", + "spend": 1.0, + "request_ids": ("req-1",), + } + + +def test_to_naive_utc_leaves_naive_values_alone(): + naive = datetime(2026, 8, 1, 12, 0) + assert to_naive_utc(naive) == naive + + +@pytest.mark.asyncio +async def test_aggregation_sums_increments_within_one_window(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.5)) + await queue.add_update(_txn("k1", WINDOW_A, 2.25)) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 1 + assert aggregated[0]["spend"] == pytest.approx(3.75) + + +@pytest.mark.asyncio +async def test_aggregation_keeps_different_windows_of_same_entity_separate(): + """Merging across windows would fold spend from a window that already + rolled into the new window's total, over-counting the new window.""" + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0)) + await queue.add_update(_txn("k1", WINDOW_B, 2.0)) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 2 + assert {payload["window_start"]: payload["spend"] for payload in aggregated} == { + "2026-08-01T00:00:00.000000": 1.0, + "2026-08-31T00:00:00.000000": 2.0, + } + + +@pytest.mark.asyncio +async def test_aggregation_keeps_durations_entities_and_types_separate(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0, duration="30d")) + await queue.add_update(_txn("k1", WINDOW_A, 2.0, duration="7d")) + await queue.add_update(_txn("k2", WINDOW_A, 4.0, duration="30d")) + await queue.add_update(_txn("k1", WINDOW_A, 8.0, duration="30d", entity_type="team")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 4 + assert sorted(payload["spend"] for payload in aggregated) == [1.0, 2.0, 4.0, 8.0] + + +@pytest.mark.asyncio +async def test_aggregation_orders_by_primary_key_then_window_start(): + """The flush relies on this order: primary key first for cross-pod lock + ordering, then window_start so an older window is applied before the roll + that supersedes it.""" + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("t1", WINDOW_A, 1.0, entity_type="team")) + await queue.add_update(_txn("k2", WINDOW_B, 1.0)) + await queue.add_update(_txn("k2", WINDOW_A, 1.0)) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, duration="7d")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert [ + (payload["entity_type"], payload["entity_id"], payload["window_duration"], payload["window_start"]) + for payload in aggregated + ] == [ + ("key", "k1", "7d", "2026-08-01T00:00:00.000000"), + ("key", "k2", "30d", "2026-08-01T00:00:00.000000"), + ("key", "k2", "30d", "2026-08-31T00:00:00.000000"), + ("team", "t1", "30d", "2026-08-01T00:00:00.000000"), + ] + + +@pytest.mark.asyncio +async def test_aggregation_does_not_collide_on_entity_ids_containing_a_separator(): + """entity_id is free-form (team ids are user supplied), so grouping must not + depend on a flattened string key.""" + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("a:30d:2026-08-01T00:00:00.000000:b", WINDOW_A, 1.0)) + await queue.add_update(_txn("b", WINDOW_A, 2.0)) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 2 + + +@pytest.mark.asyncio +async def test_flush_empties_the_queue(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0)) + + assert await queue.flush_and_get_aggregated_window_spend_transactions() != () + assert await queue.flush_and_get_aggregated_window_spend_transactions() == () + + +@pytest.mark.asyncio +async def test_aggregate_queue_updates_collapses_in_place(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0)) + await queue.add_update(_txn("k1", WINDOW_A, 2.0)) + await queue.add_update(_txn("k1", WINDOW_B, 4.0)) + + await queue.aggregate_queue_updates() + + assert queue.update_queue.qsize() == 1 + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + assert sorted(payload["spend"] for payload in aggregated) == [3.0, 4.0] + + +@pytest.mark.asyncio +async def test_aggregation_does_not_mutate_the_queued_payloads(): + """The same payload can be re-aggregated after a failed Redis push, so + aggregation must not accumulate into the caller's object.""" + queue = WindowSpendUpdateQueue() + update = _txn("k1", WINDOW_A, 1.0) + await queue.add_update(update) + await queue.add_update(_txn("k1", WINDOW_A, 2.0)) + + await queue.flush_and_get_aggregated_window_spend_transactions() + + assert update["spend"] == 1.0 + + +def test_aggregation_survives_the_redis_json_round_trip(): + """The Redis buffer stores transactions as JSON, so the aggregated shape + must reload into an equivalent aggregation.""" + aggregated = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( + [(_txn("k1", WINDOW_A, 1.0),), (_txn("k1", WINDOW_B, 2.0),)] + ) + + reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( + [json.loads(json.dumps(aggregated))] + ) + + assert reloaded == aggregated + + +@pytest.mark.asyncio +async def test_aggregation_unions_the_request_ids_of_merged_increments(): + """The seed excludes exactly the requests its batch already covers, so every + merged increment's id has to survive aggregation.""" + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-1")) + await queue.add_update(_txn("k1", WINDOW_A, 2.0, request_id="req-2")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert len(aggregated) == 1 + assert aggregated[0]["request_ids"] == ("req-1", "req-2") + + +@pytest.mark.asyncio +async def test_request_ids_stay_with_their_own_window(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) + await queue.add_update(_txn("k1", WINDOW_B, 2.0, request_id="req-b")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert {payload["window_start"]: payload["request_ids"] for payload in aggregated} == { + "2026-08-01T00:00:00.000000": ("req-a",), + "2026-08-31T00:00:00.000000": ("req-b",), + } + + +@pytest.mark.asyncio +async def test_request_ids_are_deduplicated_and_ordered(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-b")) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert aggregated[0]["request_ids"] == ("req-a", "req-b") + + +@pytest.mark.asyncio +async def test_increment_without_a_request_id_carries_no_exclusion(): + queue = WindowSpendUpdateQueue() + await queue.add_update(_txn("k1", WINDOW_A, 1.0)) + + aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() + + assert aggregated[0]["request_ids"] == () + + +def test_request_ids_survive_the_redis_json_round_trip(): + aggregated = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( + [(_txn("k1", WINDOW_A, 1.0, request_id="req-1"),)] + ) + + reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( + [json.loads(json.dumps(aggregated))] + ) + + assert reloaded[0]["request_ids"] == ("req-1",) + assert reloaded[0]["spend"] == 1.0 diff --git a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py new file mode 100644 index 00000000000..c05e3d410f9 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py @@ -0,0 +1,535 @@ +import math +import os +import sys +from contextlib import asynccontextmanager +from datetime import datetime, timezone +from typing import Any + +sys.path.insert(0, os.path.abspath("../../../..")) + +import pytest + +from litellm.proxy.db.budget_window_spend_writer import ( + commit_window_spend_updates, + roll_window_spend_row, + spend_logs_total_excluding, +) +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + build_window_spend_transaction, +) + +WINDOW_A = datetime(2026, 8, 1, tzinfo=timezone.utc) +WINDOW_B = datetime(2026, 8, 31, tzinfo=timezone.utc) + +ENTITY_TYPE, ENTITY_ID, WINDOW_DURATION, WINDOW_START, INSERT_SPEND, INCREMENT, NOW = range(7) + + +class _FakeBatcher: + def __init__(self) -> None: + self.calls: list[tuple[str, tuple[Any, ...]]] = [] + + def execute_raw(self, query: str, *args: Any) -> None: + self.calls.append((query, args)) + + +class _FakeDB: + """Stands in for prisma_client.db; records every statement it is handed.""" + + def __init__(self, existing_rows: list[dict[str, str]] | None = None) -> None: + self.existing_rows = existing_rows or [] + self.query_raw_calls: list[tuple[str, tuple[Any, ...]]] = [] + self.execute_raw_calls: list[tuple[str, tuple[Any, ...]]] = [] + self.batcher = _FakeBatcher() + self.committed = False + + async def query_raw(self, query: str, *args: Any) -> list[dict[str, str]]: + self.query_raw_calls.append((query, args)) + return self.existing_rows + + async def execute_raw(self, query: str, *args: Any) -> int: + self.execute_raw_calls.append((query, args)) + return 1 + + @asynccontextmanager + async def _tx(self): + yield self + + def tx(self, timeout: Any = None): + return self._tx() + + @asynccontextmanager + async def _batch(self): + yield self.batcher + self.committed = True + + def batch_(self): + return self._batch() + + +class _FakePrismaClient: + def __init__(self, db: _FakeDB) -> None: + self.db = db + + +class _RecordingAggregate: + """Stands in for the LiteLLM_SpendLogs seed aggregate.""" + + def __init__(self, value: float = 5.0) -> None: + self.value = value + self.calls: list[dict[str, Any]] = [] + + async def __call__( + self, + prisma_client: Any, + entity_type: str, + entity_id: str, + window_start: datetime, + exclude_request_ids: Any, + ) -> float | None: + self.calls.append( + { + "entity_type": entity_type, + "entity_id": entity_id, + "window_start": window_start, + "exclude_request_ids": tuple(exclude_request_ids), + } + ) + return self.value + + +class _SpendLogsFake: + """Sums the LiteLLM_SpendLogs rows it holds, honouring the request-id + exclusion exactly as the real aggregate's NOT (request_id = ANY(...)) does.""" + + def __init__(self, rows: tuple[tuple[str, float], ...]) -> None: + self.rows = rows + + async def __call__( + self, + prisma_client: Any, + entity_type: str, + entity_id: str, + window_start: datetime, + exclude_request_ids: Any, + ) -> float | None: + excluded = frozenset(exclude_request_ids) + return math.fsum(spend for request_id, spend in self.rows if request_id not in excluded) + + +def _existing(entity_type: str, entity_id: str, window_duration: str) -> dict[str, str]: + return {"entity_type": entity_type, "entity_id": entity_id, "window_duration": window_duration} + + +@pytest.mark.asyncio +async def test_no_transactions_touches_no_database(): + db = _FakeDB() + + await commit_window_spend_updates(prisma_client=_FakePrismaClient(db), transactions=()) + + assert db.query_raw_calls == [] + assert db.batcher.calls == [] + + +@pytest.mark.asyncio +async def test_missing_row_is_seeded_from_spend_logs_once(): + """A row created mid-window would undercount everything spent before it + existed, so a brand new primary key inserts the SpendLogs total plus this + increment.""" + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=5.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=aggregate, + ) + + assert len(aggregate.calls) == 1 + assert aggregate.calls[0]["entity_type"] == "key" + assert aggregate.calls[0]["entity_id"] == "k1" + assert aggregate.calls[0]["window_start"] == WINDOW_A + + (_, params), = db.batcher.calls + assert params[ENTITY_TYPE] == "key" + assert params[ENTITY_ID] == "k1" + assert params[WINDOW_DURATION] == "30d" + assert params[WINDOW_START] == datetime(2026, 8, 1) + assert params[INSERT_SPEND] == pytest.approx(6.0) + assert params[INCREMENT] == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_existing_row_is_never_reseeded(): + """The seed is a full LiteLLM_SpendLogs scan; running it for a row that is + already maintained would both cost a scan and double count.""" + db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")]) + aggregate = _RecordingAggregate(value=5.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=aggregate, + ) + + assert aggregate.calls == [] + (_, params), = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(1.0) + assert params[INCREMENT] == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_seed_runs_only_for_the_primary_keys_that_are_missing(): + db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")]) + aggregate = _RecordingAggregate(value=5.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("team", "t1", "30d", WINDOW_A, 2.0), + ), + spend_logs_aggregate=aggregate, + ) + + assert [call["entity_id"] for call in aggregate.calls] == ["t1"] + assert [call["entity_type"] for call in aggregate.calls] == ["team"] + by_entity = {params[ENTITY_ID]: params for _, params in db.batcher.calls} + assert by_entity["k1"][INSERT_SPEND] == pytest.approx(1.0) + assert by_entity["t1"][INSERT_SPEND] == pytest.approx(7.0) + + +@pytest.mark.asyncio +async def test_insert_spend_and_increment_differ_only_when_a_row_is_seeded(): + """The conflict arm adds the increment alone so two pods that both seed the + same new window cannot add the SpendLogs base twice.""" + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=9.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 0.25),), + spend_logs_aggregate=aggregate, + ) + + (_, params), = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(9.25) + assert params[INCREMENT] == pytest.approx(0.25) + + +@pytest.mark.asyncio +async def test_upsert_sql_adds_for_a_current_window_and_replaces_for_a_newer_one(): + """The CASE is the whole contract: an increment at or behind the stored + window_start accumulates, a newer one restarts the window.""" + db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")]) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),), + ) + + (query, _), = db.batcher.calls + normalized = " ".join(query.split()) + assert ( + 'spend = CASE WHEN "LiteLLM_BudgetWindowSpend".window_start >= EXCLUDED.window_start ' + 'THEN "LiteLLM_BudgetWindowSpend".spend + $6 ELSE EXCLUDED.spend END' in normalized + ) + assert 'window_start = GREATEST("LiteLLM_BudgetWindowSpend".window_start, EXCLUDED.window_start)' in normalized + assert "ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE SET" in normalized + + +@pytest.mark.asyncio +async def test_upsert_never_interpolates_values_into_the_sql(): + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "'; DROP TABLE x; --", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=aggregate, + ) + + (query, params), = db.batcher.calls + assert "DROP TABLE" not in query + assert params[ENTITY_ID] == "'; DROP TABLE x; --" + + +@pytest.mark.asyncio +async def test_upserts_are_ordered_by_primary_key_then_window_start(): + """Cross-pod lock ordering, plus an older window must be applied before the + roll that supersedes it or the roll would be undone.""" + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + build_window_spend_transaction("team", "t1", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("key", "k2", "30d", WINDOW_B, 1.0), + build_window_spend_transaction("key", "k2", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("key", "k1", "7d", WINDOW_A, 1.0), + ), + spend_logs_aggregate=aggregate, + ) + + ordered = [(params[ENTITY_TYPE], params[ENTITY_ID], params[WINDOW_DURATION], params[WINDOW_START]) for _, params in db.batcher.calls] + assert ordered == [ + ("key", "k1", "7d", datetime(2026, 8, 1)), + ("key", "k2", "30d", datetime(2026, 8, 1)), + ("key", "k2", "30d", datetime(2026, 8, 31)), + ("team", "t1", "30d", datetime(2026, 8, 1)), + ] + + +@pytest.mark.asyncio +async def test_existing_row_lookup_sends_every_primary_key_as_array_params(): + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("team", "t1", "7d", WINDOW_A, 1.0), + ), + spend_logs_aggregate=aggregate, + ) + + (query, params), = db.query_raw_calls + assert "unnest($1::text[], $2::text[], $3::text[])" in query + assert params == (("key", "team"), ("k1", "t1"), ("30d", "7d")) + + +@pytest.mark.asyncio +async def test_all_upserts_are_committed_in_one_transaction(): + db = _FakeDB(existing_rows=[_existing("key", "k1", "30d"), _existing("key", "k2", "30d")]) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0), + build_window_spend_transaction("key", "k2", "30d", WINDOW_A, 2.0), + ), + ) + + assert len(db.batcher.calls) == 2 + assert db.committed is True + + +@pytest.mark.asyncio +async def test_unknown_entity_type_contributes_no_seed(): + """Only key and team windows have a LiteLLM_SpendLogs column to aggregate; + anything else starts from its increment alone.""" + db = _FakeDB(existing_rows=[]) + + async def no_such_column(prisma_client, entity_type, entity_id, window_start, exclude_request_ids): + return None + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("user", "u1", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=no_such_column, + ) + + (_, params), = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_unavailable_spend_logs_aggregate_seeds_zero_rather_than_failing(): + db = _FakeDB(existing_rows=[]) + + async def unavailable(prisma_client, entity_type, entity_id, window_start, exclude_request_ids): + return None + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),), + spend_logs_aggregate=unavailable, + ) + + (_, params), = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_older(): + """Unconditional zeroing would wipe increments a pod already applied under + the new window.""" + db = _FakeDB() + + await roll_window_spend_row( + prisma_client=_FakePrismaClient(db), + entity_type="team", + entity_id="t1", + window_duration="30d", + new_window_start=WINDOW_B, + ) + + (query, params), = db.execute_raw_calls + normalized = " ".join(query.split()) + assert "SET window_start = ($4::timestamptz AT TIME ZONE 'UTC'), spend = 0" in normalized + assert "WHERE entity_type = $1 AND entity_id = $2 AND window_duration = $3" in normalized + assert "AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')" in normalized + assert params[:4] == ("team", "t1", "30d", datetime(2026, 8, 31)) + + +@pytest.mark.asyncio +async def test_seed_receives_the_batch_request_ids_to_exclude(): + db = _FakeDB(existing_rows=[]) + aggregate = _RecordingAggregate(value=0.0) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + { + "entity_type": "key", + "entity_id": "k1", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 3.0, + "request_ids": ("req-1", "req-2", "req-3"), + }, + ), + spend_logs_aggregate=aggregate, + ) + + assert aggregate.calls[0]["exclude_request_ids"] == ("req-1", "req-2", "req-3") + + +@pytest.mark.asyncio +async def test_new_row_is_not_double_counted_when_the_batch_logs_already_flushed(): + """The spend log writer drains on a ~2s poll while window increments flush + on the ~10s batch tick, so a new row is normally seeded from a table that + already holds this batch's rows. Counting them in both places is what made + a fresh row land at exactly twice the true spend.""" + db = _FakeDB(existing_rows=[]) + already_flushed = _SpendLogsFake( + rows=(("req-1", 0.000047), ("req-2", 0.000047), ("req-3", 0.000047)), + ) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + { + "entity_type": "key", + "entity_id": "k1", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 0.000141, + "request_ids": ("req-1", "req-2", "req-3"), + }, + ), + spend_logs_aggregate=already_flushed, + ) + + (_, params), = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(0.000141) + + +@pytest.mark.asyncio +async def test_new_row_still_covers_spend_that_predates_the_batch(): + """The exclusion must not throw away the pre-existing spend the seed is for.""" + db = _FakeDB(existing_rows=[]) + spend_logs = _SpendLogsFake(rows=(("older", 0.5), ("req-1", 0.000047))) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + { + "entity_type": "key", + "entity_id": "k1", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 0.000047, + "request_ids": ("req-1",), + }, + ), + spend_logs_aggregate=spend_logs, + ) + + (_, params), = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(0.500047) + + +@pytest.mark.asyncio +async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet(): + """The other side of the race: rows absent from the aggregate are still + counted exactly once, by their increment.""" + db = _FakeDB(existing_rows=[]) + nothing_flushed = _SpendLogsFake(rows=()) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(db), + transactions=( + { + "entity_type": "key", + "entity_id": "k1", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 0.000141, + "request_ids": ("req-1", "req-2", "req-3"), + }, + ), + spend_logs_aggregate=nothing_flushed, + ) + + (_, params), = db.batcher.calls + assert params[INSERT_SPEND] == pytest.approx(0.000141) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entity_type, expected_column", + [("key", "api_key = $1"), ("team", "team_id = $1")], +) +async def test_seed_aggregate_sql_excludes_the_request_ids_by_parameter(entity_type, expected_column): + db = _FakeDB(existing_rows=[{"total": 1.25}]) + + total = await spend_logs_total_excluding( + prisma_client=_FakePrismaClient(db), + entity_type=entity_type, + entity_id="e1", + window_start=WINDOW_A, + exclude_request_ids=("req-1", "req-2"), + ) + + assert total == pytest.approx(1.25) + (query, params), = db.query_raw_calls + normalized = " ".join(query.split()) + assert expected_column in normalized + assert "NOT (request_id = ANY($3::text[]))" in normalized + assert 'FROM "LiteLLM_SpendLogs"' in normalized + assert params == ("e1", WINDOW_A, ("req-1", "req-2")) + # The ids are bound, never spliced into the statement. + assert "req-1" not in query + + +@pytest.mark.asyncio +async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs_column(): + db = _FakeDB(existing_rows=[]) + + total = await spend_logs_total_excluding( + prisma_client=_FakePrismaClient(db), + entity_type="user", + entity_id="u1", + window_start=WINDOW_A, + exclude_request_ids=(), + ) + + assert total is None + assert db.query_raw_calls == [] + + +@pytest.mark.asyncio +async def test_seed_aggregate_treats_an_entity_with_no_rows_as_zero(): + db = _FakeDB(existing_rows=[]) + + total = await spend_logs_total_excluding( + prisma_client=_FakePrismaClient(db), + entity_type="key", + entity_id="k-unknown", + window_start=WINDOW_A, + exclude_request_ids=(), + ) + + assert total == 0.0 diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 191080e3a48..bd55703408e 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -9,6 +9,7 @@ sys.path.insert( ) # Adds the parent directory to the system path +from contextlib import asynccontextmanager from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, call, patch @@ -18,6 +19,9 @@ from redis.exceptions import DataError import litellm from litellm.proxy._types import Litellm_EntityType from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + build_window_spend_transaction, +) @pytest.mark.asyncio @@ -2236,3 +2240,258 @@ async def test_daily_transaction_compression_saved_tokens_zero_when_absent(): assert transaction["compression_saved_tokens"] == 0 assert transaction["compression_savings_spend"] == 0 assert transaction["prompt_caching_savings_spend"] == 0 + + +# --------------------------------------------------------------------------- +# Budget window spend flush (LiteLLM_BudgetWindowSpend) +# --------------------------------------------------------------------------- + + +class _WindowSpendFakeBatcher: + def __init__(self): + self.calls = [] + + def execute_raw(self, query, *args): + self.calls.append((query, args)) + + +class _WindowSpendFakeDB: + """Minimal prisma_client.db that records the raw statements it is handed.""" + + def __init__(self, existing_rows=None): + self.existing_rows = existing_rows or [] + self.query_raw_calls = [] + self.batcher = _WindowSpendFakeBatcher() + + async def query_raw(self, query, *args): + self.query_raw_calls.append((query, args)) + if "LiteLLM_BudgetWindowSpend" in query: + return self.existing_rows + return [] + + @asynccontextmanager + async def _tx(self): + yield self + + def tx(self, timeout=None): + return self._tx() + + @asynccontextmanager + async def _batch(self): + yield self.batcher + + def batch_(self): + return self._batch() + + +class _WindowSpendFakePrisma: + def __init__(self, db): + self.db = db + + +def _window_spend_upserts(db): + return [ + params for query, params in db.batcher.calls if "LiteLLM_BudgetWindowSpend" in query + ] + + +@pytest.mark.asyncio +async def test_window_spend_queue_is_flushed_without_redis_buffer(): + """The in-memory window queue must reach the DB on the same scheduler tick + as the other spend queues when the Redis buffer is off.""" + db_writer = DBSpendUpdateWriter() + await db_writer.window_spend_update_queue.add_update( + build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=0.5, + ) + ) + db = _WindowSpendFakeDB( + existing_rows=[ + {"entity_type": "key", "entity_id": "hashed-token", "window_duration": "30d"} + ] + ) + + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + upserts = _window_spend_upserts(db) + assert len(upserts) == 1 + assert upserts[0][0] == "key" + assert upserts[0][1] == "hashed-token" + assert upserts[0][2] == "30d" + assert upserts[0][5] == pytest.approx(0.5) + assert db_writer.window_spend_update_queue.update_queue.qsize() == 0 + + +@pytest.mark.asyncio +async def test_window_spend_queue_is_handed_to_the_redis_buffer(): + """Multi-pod deployments buffer through Redis, so the window queue has to + ride the same rpush path as the daily queues.""" + db_writer = DBSpendUpdateWriter() + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, None) + ) + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=MagicMock(), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + stored = mock_redis_update_buffer.store_in_memory_spend_updates_in_redis.call_args[1] + assert stored["window_spend_update_queue"] is db_writer.window_spend_update_queue + + +@pytest.mark.asyncio +async def test_window_spend_transactions_from_redis_are_committed_by_the_lock_winner(): + db_writer = DBSpendUpdateWriter() + window_transactions = ( + build_window_spend_transaction( + entity_type="team", + entity_id="team-1", + window_duration="7d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=2.0, + ), + ) + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, window_transactions) + ) + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + db = _WindowSpendFakeDB( + existing_rows=[{"entity_type": "team", "entity_id": "team-1", "window_duration": "7d"}] + ) + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + upserts = _window_spend_upserts(db) + assert len(upserts) == 1 + assert upserts[0][:3] == ("team", "team-1", "7d") + assert upserts[0][5] == pytest.approx(2.0) + + +@pytest.mark.asyncio +async def test_window_spend_transactions_are_not_committed_without_the_pod_lock(): + """Every pod buffers to Redis but only the lock winner may drain it.""" + db_writer = DBSpendUpdateWriter() + mock_redis_update_buffer = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=False) + db = _WindowSpendFakeDB() + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline.assert_not_called() + assert _window_spend_upserts(db) == [] + + +@pytest.mark.asyncio +async def test_failed_window_spend_commit_does_not_abort_the_rest_of_the_flush(): + """Window rows are an optimization over aggregating LiteLLM_SpendLogs, so a + failure must not take the tool registry flush down with it.""" + db_writer = DBSpendUpdateWriter() + await db_writer.window_spend_update_queue.add_update( + build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=0.5, + ) + ) + db = _WindowSpendFakeDB() + db.query_raw = AsyncMock(side_effect=Exception("connection reset")) + db_writer._flush_tool_discovery_queue = AsyncMock() + + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + db_writer._flush_tool_discovery_queue.assert_called_once() + + +@pytest.mark.asyncio +async def test_update_database_returns_the_spend_log_request_id(): + """The budget-window seed excludes the log rows its increments already + cover, so the caller needs the id this call was recorded under. It cannot + be re-derived: cache hits append time.time() to the id.""" + db_writer = DBSpendUpdateWriter() + db_writer._insert_spend_log_to_db = AsyncMock() + db_writer._enqueue_tool_usage_transaction = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", False), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + ): + request_id = await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id="test-team", + org_id=None, + kwargs={"model": "gpt-4", "custom_llm_provider": "openai", "litellm_call_id": "call-xyz"}, + completion_response=MagicMock(), + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.1, + ) + await asyncio.sleep(0) + + assert request_id is not None + # Same id the spend log row was queued under. + assert request_id == db_writer._insert_spend_log_to_db.call_args[1]["payload"]["request_id"] + + +@pytest.mark.asyncio +async def test_update_database_returns_none_when_the_payload_cannot_be_built(): + db_writer = DBSpendUpdateWriter() + + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", False), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + patch( + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + side_effect=Exception("payload boom"), + ), + ): + request_id = await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id="test-team", + org_id=None, + kwargs={"model": "gpt-4"}, + completion_response=MagicMock(), + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.1, + ) + + assert request_id is None diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index f289148101a..0c0b4ac7889 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -440,7 +440,9 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re @pytest.mark.asyncio async def test_update_database_and_spend_counters_updates_counters_after_db_update(): proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( + return_value="chatcmpl-abc123" + ) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -471,6 +473,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda budget_reservation=budget_reservation, end_user_id="test_end_user_id", tags=["tag-a"], + request_id="chatcmpl-abc123", ) @@ -1263,3 +1266,58 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( 1 if expect_spend_log else 0 ) + + +@pytest.mark.asyncio +async def test_update_database_and_spend_counters_forwards_the_spend_log_request_id(): + """The budget-window flush excludes the log rows its increments already + cover. That only works if the id update_database recorded the row under is + handed to the counter update, so this seam is load-bearing.""" + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( + return_value="chatcmpl-abc123" + ) + increment_spend_counters = AsyncMock() + + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + assert increment_spend_counters.await_args.kwargs["request_id"] == "chatcmpl-abc123" + + +@pytest.mark.asyncio +async def test_update_database_and_spend_counters_forwards_a_missing_request_id_as_none(): + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(return_value=None) + increment_spend_counters = AsyncMock() + + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + assert increment_spend_counters.await_args.kwargs["request_id"] is None diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index c7aa376a2f7..7e7c2e01fcc 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import importlib import json import os @@ -11016,3 +11017,260 @@ def test_startup_is_silent_when_mock_testing_params_disabled(caplog): ProxyStartupEvent._warn_if_mock_testing_params_enabled(general_settings={}) assert MOCK_TESTING_CONFIG_KEY not in caplog.text + + +# --------------------------------------------------------------------------- +# Budget window spend row enqueue (LiteLLM_BudgetWindowSpend writer) +# --------------------------------------------------------------------------- + + +@contextlib.contextmanager +def _window_spend_enqueue_env(cached_objects: dict): + """Point increment_spend_counters at throwaway caches and a real + WindowSpendUpdateQueue, and hand back the queue to inspect.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendUpdateQueue, + ) + import litellm.proxy.proxy_server as ps + + user_api_key_cache = MagicMock() + user_api_key_cache.async_get_cache = AsyncMock(side_effect=lambda key, **_: cached_objects.get(key)) + + queue = WindowSpendUpdateQueue() + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.window_spend_update_queue = queue + + originals = ( + ps.user_api_key_cache, + ps.spend_counter_cache, + ps.prisma_client, + ps.proxy_logging_obj, + ) + ps.user_api_key_cache = user_api_key_cache + ps.spend_counter_cache = DualCache() + ps.prisma_client = None + ps.proxy_logging_obj = proxy_logging_obj + try: + yield queue + finally: + ( + ps.user_api_key_cache, + ps.spend_counter_cache, + ps.prisma_client, + ps.proxy_logging_obj, + ) = originals + + +async def _drain(queue): + return list(await queue.flush_and_get_aggregated_window_spend_transactions()) + + +@pytest.mark.asyncio +async def test_key_window_spend_row_is_enqueued_with_the_actual_cost(): + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert len(enqueued) == 1 + assert enqueued[0]["entity_type"] == "key" + assert enqueued[0]["entity_id"] == "hashed-token" + assert enqueued[0]["window_duration"] == "30d" + assert enqueued[0]["spend"] == pytest.approx(0.25) + assert enqueued[0]["window_start"] == (reset_at - timedelta(days=30)).astimezone(timezone.utc).replace( + tzinfo=None + ).isoformat(timespec="microseconds") + + +@pytest.mark.asyncio +async def test_team_window_spend_row_is_enqueued(): + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=3) + team_obj = MagicMock() + team_obj.budget_limits = [ + {"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"team_id:team-1": team_obj}) as queue: + await increment_spend_counters( + token=None, team_id="team-1", user_id=None, response_cost=1.5 + ) + enqueued = await _drain(queue) + + assert len(enqueued) == 1 + assert enqueued[0]["entity_type"] == "team" + assert enqueued[0]["entity_id"] == "team-1" + assert enqueued[0]["window_duration"] == "7d" + assert enqueued[0]["spend"] == pytest.approx(1.5) + + +@pytest.mark.asyncio +async def test_window_spend_row_is_enqueued_even_when_the_counter_was_reserved(): + """A reservation only pre-charged the cache counter with an estimate; the + row still owes the actual cost, so the enqueue must not be skipped.""" + from litellm.proxy.proxy_server import increment_spend_counters + import litellm.proxy.spend_tracking.budget_reservation as br + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + reservation = { + "entries": [ + {"counter_key": "spend:key:hashed-token", "reserved": 1.0}, + {"counter_key": "spend:key:hashed-token:window:30d", "reserved": 1.0}, + ] + } + + original_reconcile = br.reconcile_budget_reservation + br.reconcile_budget_reservation = AsyncMock(return_value=None) + try: + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", + team_id=None, + user_id=None, + response_cost=0.25, + budget_reservation=reservation, + ) + enqueued = await _drain(queue) + finally: + br.reconcile_budget_reservation = original_reconcile + + assert len(enqueued) == 1 + assert enqueued[0]["spend"] == pytest.approx(0.25) + + +@pytest.mark.asyncio +async def test_sliding_window_without_reset_at_is_not_enqueued(): + """Windows with no reset_at slide with wall clock, so window_start moves on + every request and no single row can represent them; the read path keeps + using its LiteLLM_SpendLogs fallback instead.""" + from litellm.proxy.proxy_server import increment_spend_counters + + key_obj = MagicMock() + key_obj.budget_limits = [{"budget_duration": "30d", "max_budget": 100.0}] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert enqueued == [] + + +@pytest.mark.asyncio +async def test_each_configured_window_gets_its_own_row_enqueue(): + from litellm.proxy.proxy_server import increment_spend_counters + + now = datetime.now(timezone.utc) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "1d", "max_budget": 5.0, "reset_at": (now + timedelta(hours=5)).isoformat()}, + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": (now + timedelta(days=10)).isoformat()}, + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert sorted(item["window_duration"] for item in enqueued) == ["1d", "30d"] + assert all(item["spend"] == pytest.approx(0.25) for item in enqueued) + + +@pytest.mark.asyncio +async def test_no_window_spend_row_enqueued_without_budget_limits(): + from litellm.proxy.proxy_server import increment_spend_counters + + key_obj = MagicMock() + key_obj.budget_limits = None + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert enqueued == [] + + +@pytest.mark.asyncio +async def test_window_spend_row_carries_the_spend_log_request_id(): + """The flush excludes these ids from its one-time seed, so the id threaded + here has to be the same one the LiteLLM_SpendLogs row was written under.""" + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", + team_id=None, + user_id=None, + response_cost=0.25, + request_id="chatcmpl-abc123", + ) + enqueued = await _drain(queue) + + assert enqueued[0]["request_ids"] == ("chatcmpl-abc123",) + + +@pytest.mark.asyncio +async def test_window_spend_row_without_a_request_id_excludes_nothing(): + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=10) + key_obj = MagicMock() + key_obj.budget_limits = [ + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: + await increment_spend_counters( + token="hashed-token", team_id=None, user_id=None, response_cost=0.25 + ) + enqueued = await _drain(queue) + + assert enqueued[0]["request_ids"] == () + + +@pytest.mark.asyncio +async def test_team_window_spend_row_carries_the_request_id(): + from litellm.proxy.proxy_server import increment_spend_counters + + reset_at = datetime.now(timezone.utc) + timedelta(days=3) + team_obj = MagicMock() + team_obj.budget_limits = [ + {"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()} + ] + + with _window_spend_enqueue_env({"team_id:team-1": team_obj}) as queue: + await increment_spend_counters( + token=None, + team_id="team-1", + user_id=None, + response_cost=1.5, + request_id="chatcmpl-team", + ) + enqueued = await _drain(queue) + + assert enqueued[0]["request_ids"] == ("chatcmpl-team",) From 124d08d59234443fb441c3baba6b23d4a5dd1756 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 4 Aug 2026 17:39:51 -0700 Subject: [PATCH 003/111] perf(proxy): read budget-window spend from the maintained window table Per-window budget enforcement aggregated LiteLLM_SpendLogs on every cold counter and on every 5s authoritative floor check. SpendLogs has no index on api_key or team_id, so each check range-scanned the highest-volume table. Reads now hit the LiteLLM_BudgetWindowSpend row by primary key and only fall back to the aggregate when no row exists for the window being enforced. A row is current when its window_start is at or past the caller's expected start, so a pod holding a stale reset_at trusts a window another pod already rolled instead of summing the previous window back in. window_duration is threaded from the budget_limits entry through to the read rather than parsed back out of the counter key. The reader never writes rows. --- litellm/proxy/auth/auth_checks.py | 2 + litellm/proxy/db/spend_counter_reseed.py | 134 ++++++++-- litellm/proxy/proxy_server.py | 15 +- .../spend_tracking/budget_reservation.py | 3 + litellm/repositories/table_repositories.py | 4 + .../proxy/db/test_spend_counter_reseed.py | 250 ++++++++++++++++++ .../proxy/proxy_server/test_spend_counters.py | 79 ++++++ tests/test_litellm/proxy/test_proxy_server.py | 8 + 8 files changed, 478 insertions(+), 17 deletions(-) create mode 100644 tests/test_litellm/proxy/db/test_spend_counter_reseed.py diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 5ef4eb471ad..3e197695710 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3710,6 +3710,7 @@ async def _virtual_key_multi_budget_check( max_budget=w["max_budget"], window_entity_type="Key", window_entity_id=valid_token.token, + window_duration=str(w["budget_duration"]), window_start=get_budget_window_start(w), ) if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: @@ -4083,6 +4084,7 @@ async def _team_multi_budget_check( max_budget=w["max_budget"], window_entity_type="Team", window_entity_id=team_object.team_id, + window_duration=str(w["budget_duration"]), window_start=get_budget_window_start(w), ) if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 72e4111e0f4..1f4fdef156c 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -14,14 +14,18 @@ memory in long-lived deployments. import asyncio from collections import OrderedDict -from datetime import datetime +from collections.abc import Mapping +from datetime import datetime, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar, Final, Optional from litellm._logging import verbose_proxy_logger from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy._types import Litellm_EntityType from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( + BudgetWindowSpendRepository, SpendLogsRepository, TeamMembershipRepository, ) @@ -36,6 +40,25 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient +_WINDOW_SPEND_ENTITY_TYPES: Final[Mapping[str, str]] = MappingProxyType( + { + "Key": Litellm_EntityType.KEY.value, + "Team": Litellm_EntityType.TEAM.value, + } +) + +_WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + { + "Key": "api_key", + "Team": "team_id", + } +) + + +def _as_utc(value: datetime) -> datetime: + return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) + + class SpendCounterReseed: """ Reseeds spend counters from the authoritative DB and warms the cache, @@ -205,6 +228,92 @@ class SpendCounterReseed: raise return current_value + @staticmethod + async def window_from_table( + prisma_client: Optional["PrismaClient"], + entity_type: str, + entity_id: str, + window_duration: str, + expected_window_start: datetime, + ) -> float | None: + """ + Read the maintained per-window spend row by primary key. + + Returns the row's spend only when the row belongs to the window the + caller is enforcing, i.e. ``row.window_start >= expected_window_start``. + A row at or past the expected start was rolled by a pod whose reset_at + was at least as fresh as this caller's, so it is trusted; an older row + means the window boundary was crossed and nothing has rolled the row + yet, so its spend belongs to a previous window. + + Returns None for a missing, stale or unreadable row so the caller falls + back to the spend-logs aggregate. ``entity_type`` is the counter-facing + label ("Key"/"Team"); anything else has no row and returns None. + """ + if prisma_client is None: + return None + row_entity_type: Final = _WINDOW_SPEND_ENTITY_TYPES.get(entity_type) + if row_entity_type is None: + return None + + try: + row: Final = await BudgetWindowSpendRepository(prisma_client).table.find_unique( + where={ + "entity_type_entity_id_window_duration": { + "entity_type": row_entity_type, + "entity_id": entity_id, + "window_duration": window_duration, + } + } + ) + except Exception: + verbose_proxy_logger.exception( + "SpendCounterReseed.window_from_table: failed for %s=%s window=%s", + entity_type, + entity_id, + window_duration, + ) + return None + + if row is None: + return None + if _as_utc(row.window_start) < _as_utc(expected_window_start): + return None + return float(row.spend or 0.0) + + @staticmethod + async def window_from_db( + prisma_client: Optional["PrismaClient"], + entity_type: str, + entity_id: str, + window_duration: str | None, + window_start: datetime, + ) -> float | None: + """ + Authoritative window spend: the maintained row first, falling back to + the spend-logs aggregate only when no current row exists. + + The aggregate range-scans an unindexed table, so it must stay a + transitional path (window configured before the row existed) rather + than a steady-state read. + """ + if window_duration is not None: + from_table: Final = await SpendCounterReseed.window_from_table( + prisma_client=prisma_client, + entity_type=entity_type, + entity_id=entity_id, + window_duration=window_duration, + expected_window_start=window_start, + ) + if from_table is not None: + return from_table + return await SpendCounterReseed.window_from_spend_logs( + prisma_client=prisma_client, + entity_type=entity_type, + entity_id=entity_id, + window_start=window_start, + ) + @staticmethod async def window_from_spend_logs( prisma_client: Optional["PrismaClient"], @@ -215,20 +324,13 @@ class SpendCounterReseed: if prisma_client is None: return None - if entity_type == "Key": - group_field = "api_key" - where = { - "api_key": entity_id, - "startTime": {"gte": window_start}, - } - elif entity_type == "Team": - group_field = "team_id" - where = { - "team_id": entity_id, - "startTime": {"gte": window_start}, - } - else: + group_field: Final = _WINDOW_SPEND_LOG_FIELDS.get(entity_type) + if group_field is None: return None + where: Final = { + group_field: entity_id, + "startTime": {"gte": window_start}, + } try: response: Final = await SpendLogsRepository(prisma_client).table.group_by( @@ -258,6 +360,7 @@ class SpendCounterReseed: counter_key: str, entity_type: str, entity_id: str, + window_duration: str | None, window_start: datetime, ) -> float | None: lock: Final = await SpendCounterReseed._get_lock(counter_key) @@ -276,10 +379,11 @@ class SpendCounterReseed: if val is not None: return float(val) - window_spend: Final = await SpendCounterReseed.window_from_spend_logs( + window_spend: Final = await SpendCounterReseed.window_from_db( prisma_client=prisma_client, entity_type=entity_type, entity_id=entity_id, + window_duration=window_duration, window_start=window_start, ) if window_spend is None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b3b59d72a53..0a017a50697 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2184,6 +2184,7 @@ async def get_current_spend( max_budget: float | None = None, window_entity_type: str | None = None, window_entity_id: str | None = None, + window_duration: str | None = None, window_start: datetime | None = None, fallback_authoritative: bool = False, ) -> float: @@ -2208,7 +2209,8 @@ async def get_current_spend( runs and a key can leak spend past ``max_budget`` indefinitely. The authoritative source depends on the counter: primary key/team/user/org counters read the DB row; per-window counters (``window_start`` supplied) - aggregate spend logs; end-user/tag counters have no DB row, so the caller's + read the maintained window-spend row and only aggregate spend logs when + that row is missing or stale; end-user/tag counters have no DB row, so the caller's ``fallback_spend`` (loaded fresh in auth) is authoritative. The DB read is skipped for healthy primary counters (counter at or above recorded spend) and cached in-process for a few seconds, so a persistently stale counter @@ -2233,6 +2235,7 @@ async def get_current_spend( counter_key=counter_key, window_entity_type=window_entity_type, window_entity_id=window_entity_id, + window_duration=window_duration, window_start=window_start, ) if authoritative is not None: @@ -2314,6 +2317,7 @@ async def _authoritative_floor_spend( counter_key: str, window_entity_type: str | None = None, window_entity_id: str | None = None, + window_duration: str | None = None, window_start: datetime | None = None, ) -> float | None: marker_key: Final = f"spend_db_floor:{counter_key}" @@ -2328,10 +2332,11 @@ async def _authoritative_floor_spend( and window_entity_id is not None and window_start is not None ): - db_spend = await SpendCounterReseed.window_from_spend_logs( + db_spend = await SpendCounterReseed.window_from_db( prisma_client=prisma_client, entity_type=window_entity_type, entity_id=window_entity_id, + window_duration=window_duration, window_start=window_start, ) if db_spend is None: @@ -2456,6 +2461,7 @@ async def increment_spend_counters( counter_key=key_window_counter, entity_type="Key", entity_id=hashed_token, + window_duration=duration, window_start=key_window_start, increment=cost, ) @@ -2497,6 +2503,7 @@ async def increment_spend_counters( counter_key=team_window_counter, entity_type="Team", entity_id=scope_team_id, + window_duration=duration, window_start=team_window_start, increment=cost, ) @@ -2749,6 +2756,7 @@ async def _init_and_increment_window_spend_counter( counter_key: str, entity_type: str, entity_id: str, + window_duration: str | None, window_start: datetime | None, increment: float, ): @@ -2763,6 +2771,7 @@ async def _init_and_increment_window_spend_counter( counter_key=counter_key, entity_type=entity_type, entity_id=entity_id, + window_duration=window_duration, window_start=window_start, ) if initialized is False: @@ -2808,6 +2817,7 @@ async def _ensure_window_spend_counter_initialized( counter_key: str, entity_type: str, entity_id: str, + window_duration: str | None, window_start: datetime, ) -> bool: is_warm: Final = await _is_spend_counter_cache_warm(counter_key=counter_key) @@ -2820,6 +2830,7 @@ async def _ensure_window_spend_counter_initialized( counter_key=counter_key, entity_type=entity_type, entity_id=entity_id, + window_duration=window_duration, window_start=window_start, ) if window_spend is None: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 907c4ac7344..de82d3cfd60 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -37,6 +37,7 @@ class _BudgetCounter: entity_id: str source_cache_key: str | None = None spend_log_entity_id: str | None = None + window_duration: str | None = None window_start: datetime | None = None @@ -633,6 +634,7 @@ def _get_budget_limit_counters( entity_type=entity_type, entity_id=f"{entity_id}:{budget_duration}", spend_log_entity_id=entity_id, + window_duration=str(budget_duration), window_start=window_start, ) ) @@ -676,6 +678,7 @@ async def _reserve_counter( counter_key=counter.counter_key, entity_type=counter.entity_type, entity_id=counter.spend_log_entity_id, + window_duration=counter.window_duration, window_start=counter.window_start, ) if initialized is False: diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index af8be986831..d3c79bb4f6f 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -62,6 +62,10 @@ class SpendLogsRepository(PrismaTableRepository): table_name = "litellm_spendlogs" +class BudgetWindowSpendRepository(PrismaTableRepository): + table_name = "litellm_budgetwindowspend" + + class ClaudeCodePluginRepository(PrismaTableRepository): table_name = "litellm_claudecodeplugintable" diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py new file mode 100644 index 00000000000..816f9ae72f4 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -0,0 +1,250 @@ +"""Window-spend reads in ``SpendCounterReseed``. + +The maintained ``LiteLLM_BudgetWindowSpend`` row replaces a per-request +``LiteLLM_SpendLogs`` range scan, so these pin *when* the aggregate is still +allowed to run: only when the row is missing or belongs to an older window. +""" + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace + +import pytest + +from litellm.caching.dual_cache import DualCache +from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + +WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc) + + +class _FakeWindowSpendTable: + def __init__(self, row: SimpleNamespace | None, error: Exception | None = None) -> None: + self._row = row + self._error = error + self.where_clauses: list[dict] = [] + + async def find_unique(self, where: dict): + self.where_clauses.append(where) + if self._error is not None: + raise self._error + return self._row + + +class _FakeSpendLogsTable: + def __init__(self, total: float) -> None: + self._total = total + self.call_count = 0 + + async def group_by(self, by: list[str], where: dict, sum: dict): + self.call_count += 1 + return [{by[0]: where.get(by[0]), "_sum": {"spend": self._total}}] + + +class _FakePrismaClient: + def __init__( + self, + row: SimpleNamespace | None = None, + spend_logs_total: float = 0.0, + error: Exception | None = None, + ) -> None: + self.db = SimpleNamespace( + litellm_budgetwindowspend=_FakeWindowSpendTable(row=row, error=error), + litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total), + ) + + +def _row(window_start: datetime, spend: float) -> SimpleNamespace: + return SimpleNamespace(window_start=window_start, spend=spend) + + +@pytest.mark.asyncio +async def test_window_from_table_reads_row_by_primary_key(): + """The lookup must use the table's own entity_type values ("key"), not the + "Key"/"Team" labels the counter keys and spend-log aggregates use.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result == 4.5 + assert prisma.db.litellm_budgetwindowspend.where_clauses == [ + { + "entity_type_entity_id_window_duration": { + "entity_type": "key", + "entity_id": "tok-1", + "window_duration": "30d", + } + } + ] + + +@pytest.mark.asyncio +async def test_window_from_table_maps_team_entity_type(): + prisma = _FakePrismaClient(row=_row(WINDOW_START, 9.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Team", + entity_id="team-1", + window_duration="1d", + expected_window_start=WINDOW_START, + ) + + assert result == 9.0 + inner = prisma.db.litellm_budgetwindowspend.where_clauses[0]["entity_type_entity_id_window_duration"] + assert inner["entity_type"] == "team" + + +@pytest.mark.asyncio +async def test_window_from_table_trusts_row_newer_than_expected_window(): + """Regression: a pod holding a stale ``reset_at`` computes an expected start + behind a window another pod already rolled. Trusting only an exact match + would make it re-add the previous window's spend to the current one.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START + timedelta(days=1), 2.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result == 2.0 + + +@pytest.mark.asyncio +async def test_window_from_table_rejects_row_from_previous_window(): + prisma = _FakePrismaClient(row=_row(WINDOW_START - timedelta(seconds=1), 99.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_window_from_table_treats_naive_row_timestamp_as_utc(): + """The column is ``timestamp(3)``, so a driver that hands back a naive value + must still compare against the tz-aware expected start.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START.replace(tzinfo=None), 3.0)) + + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result == 3.0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "prisma, entity_type", + [ + (_FakePrismaClient(row=None), "Key"), + (_FakePrismaClient(row=_row(WINDOW_START, 1.0)), "User"), + (_FakePrismaClient(error=RuntimeError("connection reset")), "Key"), + (None, "Key"), + ], +) +async def test_window_from_table_returns_none_without_a_usable_row(prisma, entity_type): + result = await SpendCounterReseed.window_from_table( + prisma_client=prisma, + entity_type=entity_type, + entity_id="tok-1", + window_duration="30d", + expected_window_start=WINDOW_START, + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_window_from_db_prefers_the_row_over_the_spend_logs_aggregate(): + """The aggregate range-scans an unindexed table; a current row must keep it + from running at all.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=100.0) + + result = await SpendCounterReseed.window_from_db( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + window_start=WINDOW_START, + ) + + assert result == 4.5 + assert prisma.db.litellm_spendlogs.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "row", + [None, _row(WINDOW_START - timedelta(seconds=1), 99.0)], + ids=["missing_row", "previous_window_row"], +) +async def test_window_from_db_falls_back_to_spend_logs(row): + prisma = _FakePrismaClient(row=row, spend_logs_total=7.25) + + result = await SpendCounterReseed.window_from_db( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + window_start=WINDOW_START, + ) + + assert result == 7.25 + assert prisma.db.litellm_spendlogs.call_count == 1 + + +@pytest.mark.asyncio +async def test_window_from_db_without_a_duration_skips_the_row_lookup(): + """Callers that cannot name the window (no PK) keep the pre-table behavior.""" + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=7.25) + + result = await SpendCounterReseed.window_from_db( + prisma_client=prisma, + entity_type="Key", + entity_id="tok-1", + window_duration=None, + window_start=WINDOW_START, + ) + + assert result == 7.25 + assert prisma.db.litellm_budgetwindowspend.where_clauses == [] + + +@pytest.mark.asyncio +async def test_coalesced_window_seeds_a_cold_counter_from_the_row(): + prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=100.0) + cache = DualCache() + counter_key = "spend:key:tok-1:window:30d" + + result = await SpendCounterReseed.coalesced_window( + prisma_client=prisma, + spend_counter_cache=cache, + counter_key=counter_key, + entity_type="Key", + entity_id="tok-1", + window_duration="30d", + window_start=WINDOW_START, + ) + + assert result == 4.5 + assert cache.in_memory_cache.get_cache(key=counter_key) == 4.5 + assert prisma.db.litellm_spendlogs.call_count == 0 diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 51980342a1d..fb3de990deb 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -271,6 +271,81 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch): ) +def _make_window_spend_prisma(row=None, spend_logs_total=0.0): + prisma = MagicMock() + prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=row) + prisma.db.litellm_spendlogs.group_by = AsyncMock( + return_value=[{"api_key": "tok", "_sum": {"spend": spend_logs_total}}] + ) + return prisma + + +@pytest.mark.asyncio +async def test_get_current_spend_floors_window_against_maintained_row(monkeypatch): + """The floor re-check runs every few seconds per pod, so the window branch + must read the maintained row and leave the unindexed spend-logs scan alone.""" + from datetime import timezone + from types import SimpleNamespace + + window_start = datetime(2026, 1, 1, tzinfo=timezone.utc) + fake_prisma = _make_window_spend_prisma( + row=SimpleNamespace(window_start=window_start, spend=15.0), + spend_logs_total=100.0, + ) + fake_cache = _make_spend_counter_cache(redis_get_value=2.0) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + + counter_key = "spend:key:tok:window:7d" + result = await ps.get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + max_budget=10.0, + window_entity_type="Key", + window_entity_id="tok", + window_duration="7d", + window_start=window_start, + ) + + assert result == 15.0 + fake_prisma.db.litellm_spendlogs.group_by.assert_not_awaited() + fake_cache.redis_cache.async_set_max.assert_awaited_once_with( + key=counter_key, value=15.0 + ) + + +@pytest.mark.asyncio +async def test_get_current_spend_floors_window_against_logs_when_row_stale(monkeypatch): + """A row left behind at a crossed window boundary must not be read as the + current window's spend; the aggregate stays the fallback.""" + from datetime import timedelta, timezone + from types import SimpleNamespace + + window_start = datetime(2026, 1, 8, tzinfo=timezone.utc) + fake_prisma = _make_window_spend_prisma( + row=SimpleNamespace( + window_start=window_start - timedelta(days=7), spend=999.0 + ), + spend_logs_total=15.0, + ) + fake_cache = _make_spend_counter_cache(redis_get_value=2.0) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + + result = await ps.get_current_spend( + counter_key="spend:key:tok:window:7d", + fallback_spend=0.0, + max_budget=10.0, + window_entity_type="Key", + window_entity_id="tok", + window_duration="7d", + window_start=window_start, + ) + + assert result == 15.0 + fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once() + + @pytest.mark.asyncio async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypatch): """With fail_closed_budget_enforcement on, an admit decision backed only by a @@ -895,6 +970,7 @@ async def test_init_and_increment_window_spend_counter_increments_when_initializ counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=datetime(2024, 1, 1), increment=5.0, ) @@ -922,6 +998,7 @@ async def test_init_and_increment_window_spend_counter_missing_window_start_inva counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=None, increment=5.0, ) @@ -1059,6 +1136,7 @@ async def test_ensure_window_spend_counter_initialized_warm_returns_true(monkeyp counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=datetime(2024, 1, 1), ) @@ -1091,6 +1169,7 @@ async def test_ensure_window_spend_counter_initialized_db_failure_invalid_return counter_key="spend:key:k:window:1d", entity_type="Key", entity_id="k", + window_duration="1d", window_start=datetime(2024, 1, 1), ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7e7c2e01fcc..bb3e3d84949 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7643,6 +7643,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): counter_cache = DualCache() window_start = datetime.now(timezone.utc) - timedelta(hours=1) fake_prisma = MagicMock() + fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( return_value=[{"api_key": "key-window", "_sum": {"spend": 2.25}}] ) @@ -7657,6 +7658,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): counter_key="spend:key:key-window:window:1h", entity_type="Key", entity_id="key-window", + window_duration="1h", window_start=window_start, increment=0.5, ) @@ -7766,6 +7768,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() + fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( return_value=[{"api_key": "key-window-stale-local", "_sum": {"spend": 2.25}}] ) @@ -7780,6 +7783,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): counter_key=counter_key, entity_type="Key", entity_id="key-window-stale-local", + window_duration="1h", window_start=window_start, increment=0.5, ) @@ -7830,6 +7834,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() + fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( return_value=[ {"api_key": "key-window-concurrent-seed", "_sum": {"spend": 2.25}} @@ -7846,6 +7851,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() counter_key=counter_key, entity_type="Key", entity_id="key-window-concurrent-seed", + window_duration="1h", window_start=window_start, increment=0.5, ) @@ -7880,6 +7886,7 @@ async def test_window_spend_counter_skips_invalid_window_start(): counter_key="spend:key:key-invalid-window:window:not-a-duration", entity_type="Key", entity_id="key-invalid-window", + window_duration="not-a-duration", window_start=None, increment=0.5, ) @@ -7912,6 +7919,7 @@ async def test_window_spend_counter_does_not_seed_zero_when_db_unavailable(): counter_key=counter_key, entity_type="Key", entity_id="key-window-db-unavailable", + window_duration="1h", window_start=datetime.now(timezone.utc) - timedelta(hours=1), ) From cb65bf08b8c937196270f14468916de3a627388b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:41:30 -0700 Subject: [PATCH 004/111] chore(typing): clear 1.2k basedpyright Any errors across 16 hotspot files Replace Any-typed seams with real types in the files carrying the highest remaining reportAny/reportExplicitAny density: Literal-keyed structural Protocols for deployment dicts in tag-based routing, typed Prisma table wrappers and row protocols in the key and internal-user management endpoints, TypedDict views for websearch interception kwargs, typed streaming state in the responses iterator and background polling, and concrete request/response types in the google_genai, vertex_ai files, runwayml, rubrik, anthropic context-management, and guardrail translation modules. Mutable annotations introduced along the way were rewritten as read-only views (Mapping/Sequence/tuple) built functionally. No casts, no type: ignore, no noqa, no suppression comments, no new Any annotations, no behavior changes. Whole-tree basedpyright: reportAny 15,496 -> 14,523, reportExplicitAny 5,356 -> 5,102, all rules 145,547 -> 143,989, with no rule increased repo-wide or per-file. Budgets ratcheted: basedpyright -1,545, ruff-strict -73, type-discipline -237. --- basedpyright-code-budget.json | 22 +- .../google_genai/adapters/transformation.py | 186 ++++++++--- litellm/integrations/rubrik.py | 228 +++++++++---- .../websearch_interception/handler.py | 157 +++++++-- .../chat/guardrail_translation/handler.py | 202 +++++++----- .../context_management/editors/compact.py | 213 ++++++++---- .../llms/runwayml/videos/transformation.py | 138 ++++---- .../llms/vertex_ai/files/transformation.py | 114 +++++-- litellm/proxy/db/tool_registry_writer.py | 197 ++++++++---- .../internal_user_endpoints.py | 302 +++++++++++++----- .../key_management_endpoints.py | 241 ++++++++++---- .../tool_management_endpoints.py | 213 ++++++++++-- litellm/proxy/prompts/prompt_endpoints.py | 161 ++++++---- .../response_polling/background_streaming.py | 117 +++++-- .../responses/file_search/emulated_handler.py | 286 +++++++++-------- litellm/responses/streaming_iterator.py | 227 +++++++++---- litellm/router_strategy/tag_based_routing.py | 174 ++++++---- ruff-strict-budget.json | 16 +- type-discipline-budget.json | 10 +- 19 files changed, 2261 insertions(+), 943 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 6e3cbdff9d0..c71ef7a0020 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 23919 + "limit": 21974 }, "reportArgumentType": { - "limit": 2580 + "limit": 2575 }, "reportAssignmentType": { "limit": 323 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 7573 + "limit": 7068 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5719 + "limit": 5697 }, "reportMissingTypeArgument": { - "limit": 15657 + "limit": 15627 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44832 + "limit": 44549 }, "reportUnknownLambdaType": { - "limit": 113 + "limit": 112 }, "reportUnknownMemberType": { - "limit": 39269 + "limit": 39156 }, "reportUnknownParameterType": { - "limit": 19988 + "limit": 19951 }, "reportUnknownVariableType": { - "limit": 30923 + "limit": 30798 }, "reportUnnecessaryCast": { "limit": 118 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 853 + "limit": 852 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 4f127f476c3..4c8e77d9feb 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -1,6 +1,9 @@ import json -from collections.abc import AsyncIterator, Iterator -from typing import Any, Final, cast +from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence +from types import MappingProxyType +from typing import Any, Final, TypeAlias, cast + +from typing_extensions import TypedDict from litellm import verbose_logger from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema @@ -9,7 +12,6 @@ from litellm.types.llms.openai import ( ChatCompletionAssistantMessage, ChatCompletionAssistantToolCall, ChatCompletionImageObject, - ChatCompletionRequest, ChatCompletionSystemMessage, ChatCompletionTextObject, ChatCompletionToolCallFunctionChunk, @@ -21,12 +23,79 @@ from litellm.types.llms.openai import ( from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( AdapterCompletionStreamWrapper, + ChatCompletionDeltaCustomToolCall, + ChatCompletionDeltaToolCall, + ChatCompletionMessageCustomToolCall, + ChatCompletionMessageToolCall, Choices, + Delta, + Function, + Message, ModelResponse, ModelResponseStream, StreamingChoices, ) +_JsonDict: TypeAlias = dict[str, object] +_JsonDictList: TypeAlias = list[_JsonDict] + + +class _ToolCallAccumulator(TypedDict): + name: str + arguments: str + + +class _GenAIFunctionCall(TypedDict): + name: str + args: Mapping[str, object] + + +class _GenAIPart(TypedDict, total=False): + text: str + functionCall: _GenAIFunctionCall + + +class _GenAIFunctionResponse(TypedDict, total=False): + name: str + response: object + + +class _GenAIRequestFunctionCall(TypedDict, total=False): + name: str + args: Mapping[str, object] + + +class _GenAIContentPart(TypedDict, total=False): + text: str + inline_data: Mapping[str, str] + functionResponse: _GenAIFunctionResponse + functionCall: _GenAIRequestFunctionCall + + +class _GenAIFunctionDeclaration(TypedDict, total=False): + name: str + description: str + parametersJsonSchema: object + + +class _GenAITool(TypedDict, total=False): + functionDeclarations: Sequence[_GenAIFunctionDeclaration] + + +class _GenAIFunctionCallingConfig(TypedDict, total=False): + mode: str + + +class _GenAIToolConfig(TypedDict, total=False): + functionCallingConfig: _GenAIFunctionCallingConfig + + +class _GenAISystemInstruction(TypedDict, total=False): + parts: Sequence[Mapping[str, str]] + + +_EMPTY_STR_MAPPING: Final[Mapping[str, str]] = MappingProxyType({}) + class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): """ @@ -35,12 +104,12 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): """ sent_first_chunk: bool = False - # State tracking for accumulating partial tool calls - accumulated_tool_calls: dict[str, dict[str, Any]] + _parse_accumulated_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads) - def __init__(self, completion_stream: Any): + def __init__(self, completion_stream: object): self.sent_first_chunk = False - self.accumulated_tool_calls = {} + # State tracking for accumulating partial tool calls + self.accumulated_tool_calls = dict[int, _ToolCallAccumulator]() self._returned_response = False super().__init__(completion_stream) @@ -85,7 +154,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): # After the stream is exhausted, check for any remaining accumulated tool calls if self.accumulated_tool_calls: try: - parts: Final = [] + parts: Final = list[_GenAIPart]() for ( tool_call_index, tool_call_data, @@ -93,8 +162,10 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): try: # For tool calls with no arguments, accumulated_args will be "", which is not valid JSON. # We default to an empty JSON object in this case. - parsed_args = json.loads(tool_call_data["arguments"] or "{}") - function_call_part = { + parsed_args: Mapping[str, object] = self._parse_accumulated_args( + tool_call_data["arguments"] or "{}" + ) + function_call_part: _GenAIPart = { "functionCall": { "name": tool_call_data["name"] or "undefined_tool_name", "args": parsed_args, @@ -172,14 +243,16 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): class GoogleGenAIAdapter: """Adapter for transforming Google GenAI generate_content requests to/from litellm.completion format""" + _parse_tool_call_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads) + def __init__(self) -> None: pass def translate_generate_content_to_completion( self, model: str, - contents: list[dict[str, Any]] | dict[str, Any], - config: dict[str, Any] | None = None, + contents: _JsonDictList | _JsonDict, + config: Mapping[str, object] | None = None, litellm_params: GenericLiteLLMParams | None = None, **kwargs, ) -> dict[str, Any]: @@ -211,7 +284,7 @@ class GoogleGenAIAdapter: messages: Final = self._transform_contents_to_messages(contents_list, system_instruction=system_instruction) # Create base request as dict (which is compatible with ChatCompletionRequest) - completion_request: Final[ChatCompletionRequest] = { + completion_request: Final[_JsonDict] = { "model": model, "messages": messages, } @@ -273,9 +346,9 @@ class GoogleGenAIAdapter: def _add_generic_litellm_params_to_request( self, - completion_request_dict: dict[str, Any], + completion_request_dict: _JsonDict, litellm_params: GenericLiteLLMParams | None = None, - ) -> dict: + ) -> _JsonDict: """Add generic litellm params to request. e.g add api_base, api_key, api_version, etc. Args: @@ -287,7 +360,7 @@ class GoogleGenAIAdapter: """ allowed_fields: Final = GenericLiteLLMParams.model_fields.keys() if litellm_params: - litellm_dict: Final = litellm_params.model_dump(exclude_none=True) + litellm_dict: Final[_JsonDict] = litellm_params.model_dump(exclude_none=True) for key, value in litellm_dict.items(): if key in allowed_fields: completion_request_dict[key] = value @@ -295,7 +368,7 @@ class GoogleGenAIAdapter: def translate_completion_output_params_streaming( self, - completion_stream: Any, + completion_stream: object, ) -> AsyncIterator[bytes] | None: """Transform streaming completion output to Google GenAI format""" google_genai_wrapper: Final = GoogleGenAIStreamWrapper(completion_stream=completion_stream) @@ -304,15 +377,15 @@ class GoogleGenAIAdapter: def _transform_google_genai_tools_to_openai( self, - tools: list[dict[str, Any]], + tools: Sequence[_GenAITool], ) -> list[ChatCompletionToolParam]: """Transform Google GenAI tools to OpenAI tools format""" - openai_tools: Final[list[dict[str, Any]]] = [] + openai_tools: Final = list[_JsonDict]() for tool in tools: if "functionDeclarations" in tool: for func_decl in tool["functionDeclarations"]: - function_chunk: dict[str, Any] = { + function_chunk: _JsonDict = { "name": func_decl.get("name", ""), } @@ -321,7 +394,7 @@ class GoogleGenAIAdapter: if "parametersJsonSchema" in func_decl: function_chunk["parameters"] = func_decl["parametersJsonSchema"] - openai_tool = {"type": "function", "function": function_chunk} + openai_tool: _JsonDict = {"type": "function", "function": function_chunk} openai_tools.append(openai_tool) # normalize the tool schemas @@ -331,7 +404,7 @@ class GoogleGenAIAdapter: def _transform_google_genai_tool_config_to_openai( self, - tool_config: dict[str, Any], + tool_config: _GenAIToolConfig, ) -> ChatCompletionToolChoiceValues | None: """Transform Google GenAI tool_config to OpenAI tool_choice""" function_calling_config: Final = tool_config.get("functionCallingConfig", {}) @@ -345,20 +418,20 @@ class GoogleGenAIAdapter: def _transform_contents_to_messages( self, contents: list[dict[str, Any]], - system_instruction: dict[str, Any] | None = None, + system_instruction: _GenAISystemInstruction | None = None, ) -> list[AllMessageValues]: """Transform Google GenAI contents to OpenAI messages format""" messages: Final[list[AllMessageValues]] = [] # Handle system instruction if system_instruction: - system_parts: Final = system_instruction.get("parts", []) + system_parts: Final[Sequence[Mapping[str, str]]] = system_instruction.get("parts", []) if system_parts and "text" in system_parts[0]: messages.append(ChatCompletionSystemMessage(role="system", content=system_parts[0]["text"])) for content in contents: role = content.get("role", "user") - parts = content.get("parts", []) + parts: Sequence[_GenAIContentPart | str | None] = content.get("parts", []) if role == "user": # Handle user messages with potential function responses @@ -461,7 +534,7 @@ class GoogleGenAIAdapter: def translate_completion_to_generate_content( self, response: ModelResponse, - ) -> dict[str, Any]: + ) -> _JsonDict: """ Transform litellm completion response to Google GenAI generate_content format @@ -484,13 +557,13 @@ class GoogleGenAIAdapter: parts = self._transform_openai_message_to_google_genai_parts(choice.message) else: # Fallback for generic choice objects - message_content = getattr(choice, "message", {}).get("content", "") or getattr(choice, "delta", {}).get( - "content", "" - ) + message_content: str = getattr(choice, "message", _EMPTY_STR_MAPPING).get("content", "") or getattr( + choice, "delta", _EMPTY_STR_MAPPING + ).get("content", "") parts = [{"text": message_content}] if message_content else [] # Create Google GenAI format response - generate_content_response: Final[dict[str, Any]] = { + generate_content_response: Final[_JsonDict] = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -524,7 +597,7 @@ class GoogleGenAIAdapter: self, response: ModelResponse | ModelResponseStream, wrapper: GoogleGenAIStreamWrapper, - ) -> dict[str, Any] | None: + ) -> Mapping[str, object] | None: """ Transform streaming litellm completion chunk to Google GenAI generate_content format @@ -548,10 +621,10 @@ class GoogleGenAIAdapter: parts = self._transform_openai_delta_to_google_genai_parts_with_accumulation(choice.delta, wrapper) else: parts = [] - finish_reason = getattr(choice, "finish_reason", None) + finish_reason: str | None = getattr(choice, "finish_reason", None) else: # Fallback for generic choice objects - message_content: Final = getattr(choice, "delta", {}).get("content", "") + message_content: Final[str] = getattr(choice, "delta", _EMPTY_STR_MAPPING).get("content", "") parts = [{"text": message_content}] if message_content else [] finish_reason = getattr(choice, "finish_reason", None) @@ -560,7 +633,7 @@ class GoogleGenAIAdapter: return None # Create Google GenAI streaming format response - streaming_chunk: Final[dict[str, Any]] = { + streaming_chunk: Final[_JsonDict] = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -596,10 +669,10 @@ class GoogleGenAIAdapter: def _transform_openai_message_to_google_genai_parts( self, - message: Any, - ) -> list[dict[str, Any]]: + message: Message, + ) -> Sequence[_GenAIPart]: """Transform OpenAI message to Google GenAI parts format""" - parts: Final[list[dict[str, Any]]] = [] + parts: Final = list[_GenAIPart]() # Add text content if present if hasattr(message, "content") and message.content: @@ -607,16 +680,22 @@ class GoogleGenAIAdapter: # Add tool calls if present if hasattr(message, "tool_calls") and message.tool_calls: - for tool_call in message.tool_calls: - if hasattr(tool_call, "function") and tool_call.function: + tool_calls: Final[Sequence[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall]] = ( + message.tool_calls + ) + for tool_call in tool_calls: + function: Function | None = getattr(tool_call, "function", None) + if function: try: - args = json.loads(tool_call.function.arguments) if tool_call.function.arguments else {} + args: Mapping[str, object] = ( + self._parse_tool_call_args(function.arguments) if function.arguments else {} + ) except json.JSONDecodeError: args = {} - function_call_part = { + function_call_part: _GenAIPart = { "functionCall": { - "name": tool_call.function.name or "undefined_tool_name", + "name": function.name or "undefined_tool_name", "args": args, } } @@ -625,28 +704,30 @@ class GoogleGenAIAdapter: return parts if parts else [{"text": ""}] def _transform_openai_delta_to_google_genai_parts_with_accumulation( - self, delta: Any, wrapper: GoogleGenAIStreamWrapper - ) -> list[dict[str, Any]]: + self, delta: Delta, wrapper: GoogleGenAIStreamWrapper + ) -> Sequence[_GenAIPart]: """Transforms OpenAI delta to Google GenAI parts, accumulating streaming tool calls.""" # 1. Initialize wrapper state if it doesn't exist if not hasattr(wrapper, "accumulated_tool_calls"): wrapper.accumulated_tool_calls = {} - parts: Final[list[dict[str, Any]]] = [] + parts: Final = list[_GenAIPart]() if hasattr(delta, "content") and delta.content: parts.append({"text": delta.content}) # 2. Ensure tool_calls is iterable - tool_calls: Final = delta.tool_calls or [] + tool_calls: Final[Sequence[ChatCompletionDeltaToolCall | ChatCompletionDeltaCustomToolCall]] = ( + delta.tool_calls or [] + ) for tool_call in tool_calls: if not hasattr(tool_call, "function"): continue # 3. Use `index` as the primary key for accumulation - tool_call_index = getattr(tool_call, "index", None) + tool_call_index: int | None = getattr(tool_call, "index", None) if tool_call_index is None: continue # Index is essential for tracking streaming tool calls @@ -658,8 +739,9 @@ class GoogleGenAIAdapter: } # Accumulate name and arguments - function_name = getattr(tool_call.function, "name", None) - args_chunk = getattr(tool_call.function, "arguments", None) + delta_function: Function | None = getattr(tool_call, "function", None) + function_name: str | None = getattr(delta_function, "name", None) + args_chunk: str | None = getattr(delta_function, "arguments", None) # Optimization: Skip chunks that have no new data if not function_name and not args_chunk: @@ -680,13 +762,13 @@ class GoogleGenAIAdapter: # 5. Attempt to parse arguments even if name hasn't arrived. try: # Attempt to parse the accumulated arguments string - parsed_args = json.loads(accumulated_args) + parsed_args: Mapping[str, object] = self._parse_tool_call_args(accumulated_args) # If parsing succeeds, but we don't have a name yet, wait. # The part will be created by a later chunk that brings the name. if accumulated_name: # If successful, create the part and clean up - function_call_part = {"functionCall": {"name": accumulated_name, "args": parsed_args}} + function_call_part: _GenAIPart = {"functionCall": {"name": accumulated_name, "args": parsed_args}} parts.append(function_call_part) # Remove the completed tool call from the accumulator @@ -714,7 +796,7 @@ class GoogleGenAIAdapter: return mapping.get(finish_reason, "STOP") - def _map_usage(self, usage: Any) -> dict[str, int]: + def _map_usage(self, usage: object) -> Mapping[str, int]: """Map OpenAI usage to Google GenAI usage format""" return { "promptTokenCount": getattr(usage, "prompt_tokens", 0) or 0, diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index 97e831f5822..c206849c86f 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -6,12 +6,14 @@ import random import time import uuid from collections import Counter -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Mapping, Sequence from dataclasses import dataclass +from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypedDict, overload import httpx +from typing_extensions import Never, Required from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -29,6 +31,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk from litellm.types.utils import ( ChatCompletionMessageToolCall, Function, @@ -48,7 +51,105 @@ _WEBHOOK_PATH_PROMPT_MODERATION: Final = "/v1/before_prompt/openai/v1" _WEBHOOK_PATH_LOGGING_BATCH: Final = "/v1/litellm/batch" _MAX_QUEUE_SIZE: Final = 10_000 _DROP_WARNING_INTERVAL_SECONDS: Final = 60.0 -_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({}) +_EMPTY_MAPPING: Final[Mapping[str, Never]] = MappingProxyType({}) + + +class _ModerationToolCall(TypedDict, total=False): + id: Required[str] + + +class _ModerationMessage(TypedDict, total=False): + content: str | None + tool_calls: Sequence[_ModerationToolCall] | None + + +class _ModerationChoice(TypedDict, total=False): + message: _ModerationMessage | None + + +class _ModerationResponse(TypedDict, total=False): + choices: Sequence[_ModerationChoice] + + +class _LogEventKwargs(TypedDict, total=False): + standard_logging_object: Required[StandardLoggingPayload] + litellm_call_id: str + + +class _HasCallId(Protocol): + def get(self, key: Literal["litellm_call_id"], /) -> str | None: ... + + +class _HasModelAttr(Protocol): + model: str | None + + +class _ResponseSource(Protocol): + def get(self, key: Literal["response"], /) -> "_HasModelAttr | None": ... + + +class _ModelSource(Protocol): + def get(self, key: Literal["model"], default: str, /) -> str: ... + + +class _FallbackSource(Protocol): + @overload + def get(self, key: Literal["start_time"], /) -> datetime | None: ... + @overload + def get(self, key: str, /) -> object | None: ... + + +class _RequestContextSource(Protocol): + @overload + def get(self, key: Literal["optional_params"], /) -> Mapping[str, object] | None: ... + @overload + def get(self, key: str, /) -> object | None: ... + def __contains__(self, key: object, /) -> bool: ... + def __getitem__(self, key: str, /) -> object: ... + + +class _ToolCallLike(Protocol): + id: str | None + type: str | None + function: Function + + +class _ModerationSourceToolCall(TypedDict, total=False): + function: Mapping[str, object] | None + + +class _ModerationSourceMessage(TypedDict, total=False): + role: str + function_call: Mapping[str, object] | None + tool_calls: Sequence[_ModerationSourceToolCall | None] | None + + +class _FlattenedModerationMessage(TypedDict): + role: str | None + content: str + + +class _CorrelatablePayload(TypedDict): + id: str + + +class _SystemPromptCarrier(TypedDict, total=False): + messages: object + + +class _BlockFailurePayload(TypedDict, total=False): + id: object + model: object + model_group: object + model_id: str + model_parameters: object + startTime: float | None + endTime: float | None + completionStartTime: float | None + messages: object + metadata: StandardLoggingUserAPIKeyMetadata + response: str + status: str class _MalformedToolBlockingResponseError(Exception): @@ -143,7 +244,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): else {"Content-Type": "application/json"} ) - self._periodic_flush_task: asyncio.Task[Any] | None = self._start_periodic_flush_task() + self._periodic_flush_task: asyncio.Task[None] | None = self._start_periodic_flush_task() @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: @@ -191,7 +292,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): params={"timeout": httpx.Timeout(5.0, connect=2.0)}, ) - def _start_periodic_flush_task(self) -> asyncio.Task[Any] | None: + def _start_periodic_flush_task(self) -> asyncio.Task[None] | None: """Start the periodic flush task only when an event loop is already running.""" try: loop: Final = asyncio.get_running_loop() @@ -212,7 +313,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): Closing them here would close the shared connection pool for every other logger instance; let LiteLLM manage their lifecycle instead. """ - task: Final = getattr(self, "_periodic_flush_task", None) + task: Final[asyncio.Task[None] | None] = getattr(self, "_periodic_flush_task", None) if task is not None: task.cancel() @@ -253,7 +354,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod async def _guarded( - coro: Any, + coro: Awaitable[GenericGuardrailAPIInputs], inputs: GenericGuardrailAPIInputs, label: str, ) -> GenericGuardrailAPIInputs: @@ -371,7 +472,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod def _stash_block_context( logging_obj: Optional["LiteLLMLoggingObj"], - request_data: dict, + request_data: dict[str, object], ) -> None: """Stash signals so the deferred success-event skips this request and ``async_post_call_failure_hook`` can build the failure payload. @@ -400,12 +501,16 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): request_data["_rubrik_logging_obj"] = logging_obj @staticmethod - def _normalize_tool_calls(tool_calls: Any) -> tuple[ChatCompletionMessageToolCall, ...]: + def _normalize_tool_calls( + tool_calls: Sequence[ChatCompletionToolCallChunk | ChatCompletionMessageToolCall | _ToolCallLike], + ) -> tuple[ChatCompletionMessageToolCall, ...]: """Convert tool_calls from inputs to ChatCompletionMessageToolCall objects.""" return tuple(RubrikLogger._normalize_tool_call(tc) for tc in tool_calls) @staticmethod - def _normalize_tool_call(tc: Any) -> ChatCompletionMessageToolCall: + def _normalize_tool_call( + tc: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall | _ToolCallLike, + ) -> ChatCompletionMessageToolCall: if isinstance(tc, ChatCompletionMessageToolCall): return tc if isinstance(tc, dict): @@ -427,7 +532,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): raise TypeError(f"Cannot normalize tool_call of type {type(tc).__name__}: {tc!r}") @staticmethod - def _join_texts(texts: Any) -> str: + def _join_texts(texts: Sequence[str] | None) -> str: """Join response text segments into the single content string the webhook evaluates. Empty when there is no assistant text.""" if not texts: @@ -439,19 +544,22 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): tool_calls: Sequence[ChatCompletionMessageToolCall], content: str, request_id: str | None, - ) -> Mapping[str, Any]: + ) -> Mapping[str, object]: """Build an OpenAI ChatCompletion-format dict (assistant text + tool calls) for the after_completion webhook. ``content`` is sent so the webhook can moderate the response text; ``None`` when the assistant produced no text (tool-call-only response). """ - message: Final[dict[str, Any]] = { + message: Final[Mapping[str, object]] = { "role": "assistant", "content": content or None, + **( + {"tool_calls": tuple(tc.model_dump(exclude_none=True) for tc in tool_calls)} + if tool_calls + else _EMPTY_MAPPING + ), } - if tool_calls: - message["tool_calls"] = tuple(tc.model_dump(exclude_none=True) for tc in tool_calls) return { "id": request_id or f"chatcmpl-{uuid.uuid4()}", "object": "chat.completion", @@ -467,7 +575,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): } @staticmethod - def _flatten_messages_for_moderation(messages: Any) -> tuple[Mapping[str, Any], ...]: + def _flatten_messages_for_moderation( + messages: Sequence[AllMessageValues | None] | None, + ) -> tuple[_FlattenedModerationMessage, ...]: """Collapse each message's content to a plain string for the webhook. litellm normalizes Anthropic ``/v1/messages`` requests to OpenAI shape, @@ -488,7 +598,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): ) @staticmethod - def _moderation_text_parts(message: Mapping[str, Any]) -> tuple[str, ...]: + def _moderation_text_parts(message: _ModerationSourceMessage) -> tuple[str, ...]: """Every attacker-controlled text segment of a message: its content plus the arguments of any tool call or deprecated function call.""" fc: Final = message.get("function_call") @@ -506,8 +616,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod def _build_prompt_moderation_payload( inputs: GenericGuardrailAPIInputs, - request_data: Mapping[str, Any], - ) -> Mapping[str, Any]: + request_data: Mapping[str, object], + ) -> Mapping[str, object]: """Build the bare OpenAI request the before_prompt webhook consumes. Unlike the after_completion envelope, this endpoint takes a raw OpenAI @@ -516,16 +626,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): ``/v1/messages`` requests too. Optional fields are sent only when present so the payload stays clean. """ - payload: Final[dict[str, Any]] = { - "model": inputs.get("model") or request_data.get("model") or "", - "messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")), - } tools: Final = inputs.get("tools") - if tools is not None: - payload["tools"] = tools user: Final = request_data.get("user") - if user: - payload["user"] = user # Fall back to litellm_call_id, the stable cross-provider join key the # response/tool path uses (see _correlation_id). LiteLLM does not # populate request_data["correlation_key"]; it carries litellm_call_id. @@ -533,15 +635,19 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # when correlation_key is empty, so without this the block fires but no # log is ever written. An explicit correlation_key still wins. correlation_key: Final = request_data.get("correlation_key") or request_data.get("litellm_call_id") - if correlation_key: - payload["correlation_key"] = correlation_key - return payload + return { + "model": inputs.get("model") or request_data.get("model") or "", + "messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")), + **({"tools": tools} if tools is not None else _EMPTY_MAPPING), + **({"user": user} if user else _EMPTY_MAPPING), + **({"correlation_key": correlation_key} if correlation_key else _EMPTY_MAPPING), + } @staticmethod def _extract_request_data( - call_details: Mapping[str, Any], - request_data: Mapping[str, Any] | None, - ) -> Mapping[str, Any]: + call_details: _RequestContextSource, + request_data: _RequestContextSource | None, + ) -> Mapping[str, object]: """Extract original request data from model_call_details for the response moderation service envelope. @@ -576,7 +682,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): } @staticmethod - def _sanitize_proxy_server_request(proxy_server_request: Any) -> Any: + def _sanitize_proxy_server_request(proxy_server_request: Mapping[str, object] | str | None) -> object: """Allowlist only routing fields (``url``, ``method``) when forwarding ``proxy_server_request`` to an external webhook, dropping inbound ``headers`` (Authorization, Cookie, x-api-key, ...) and the raw @@ -586,7 +692,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return {key: proxy_server_request[key] for key in ("url", "method") if key in proxy_server_request} @staticmethod - def _resolve_model(request_data: Mapping[str, Any], call_details: Mapping[str, Any]) -> str: + def _resolve_model(request_data: _ResponseSource, call_details: _ModelSource) -> str: """Get the model name for the ModifyResponseException.""" response: Final = request_data.get("response") if response and hasattr(response, "model"): @@ -596,7 +702,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # -- Logging hooks --------------------------------------------------------- @staticmethod - def _correlation_id(call_details: Mapping[str, Any], request_data: Mapping[str, Any] | None = None) -> str | None: + def _correlation_id( + call_details: _HasCallId | _LogEventKwargs, request_data: _HasCallId | None = None + ) -> str | None: """The id that joins a blocked request's two S3 logs by filename: the moderation (``_blocking``) log and the failure (response) log. @@ -610,7 +718,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return call_details.get("litellm_call_id") or (request_data or _EMPTY_MAPPING).get("litellm_call_id") @classmethod - def _apply_correlation_id(cls, payload: dict[str, Any], source: Mapping[str, Any]) -> None: + def _apply_correlation_id(cls, payload: _CorrelatablePayload, source: _HasCallId | _LogEventKwargs) -> None: """Pin ``payload["id"]`` to ``litellm_call_id`` in place so this log shares its S3 filename id with the moderation (``_blocking``) and failure logs for the same request -- for every provider. @@ -630,7 +738,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): payload["id"] = correlated @staticmethod - def _prepend_system_prompt(payload: dict[str, Any], source: Mapping[str, Any]) -> None: + def _prepend_system_prompt(payload: _SystemPromptCarrier, source: Mapping[str, object]) -> None: """Prepend ``source["system"]`` onto ``payload["messages"]``. Builds a NEW messages list rather than mutating ``payload["messages"]`` @@ -658,7 +766,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): exc_info=True, ) - async def _prepare_log_payload(self, kwargs: Mapping[str, Any], event_type: str) -> StandardLoggingPayload | None: + async def _prepare_log_payload(self, kwargs: _LogEventKwargs, event_type: str) -> StandardLoggingPayload | None: """Shared logic for success logging (sampled).""" if random.random() > self.sampling_rate: verbose_logger.debug("Skipping Rubrik %s logging (sampling_rate=%s)", event_type, self.sampling_rate) @@ -667,12 +775,12 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # Deep-copy so mutations don't affect other callbacks sharing this object standard_logging_payload: Final[StandardLoggingPayload] = safe_deep_copy(kwargs["standard_logging_object"]) - self._apply_correlation_id(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime + self._apply_correlation_id(standard_logging_payload, kwargs) self._prepend_system_prompt(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime return standard_logging_payload - async def _append_and_maybe_flush(self, payload) -> None: + async def _append_and_maybe_flush(self, payload: Mapping[str, object]) -> None: self._ensure_periodic_flush_task() self.log_queue.append(payload) self._enforce_max_queue_size() @@ -697,7 +805,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): self._dropped_since_warning = 0 self._last_drop_warning_time = now - async def _enqueue_log_event(self, kwargs: Mapping[str, Any], event_type: str): + async def _enqueue_log_event(self, kwargs: _LogEventKwargs, event_type: str): try: payload: Final = await self._prepare_log_payload(kwargs, event_type) if payload is None: @@ -818,7 +926,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): logging_obj: "LiteLLMLoggingObj", exception: "ModifyResponseException", user_api_key_dict: "UserAPIKeyAuth", - ) -> StandardLoggingPayload: + ) -> _BlockFailurePayload: """Build a failure-style payload using the exception text as response. Blocked-tool events are security-relevant and **bypass sampling**: @@ -860,9 +968,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): call_details: Final = logging_obj.model_call_details exception_text: Final = f"{type(exception).__name__}: {exception.message}" - base: Final = call_details.get("standard_logging_object") + base: Final[StandardLoggingPayload | None] = call_details.get("standard_logging_object") if base is not None: - payload: dict = safe_deep_copy(base) + payload: _BlockFailurePayload = self._copy_block_payload_base(base) else: verbose_logger.debug( "Rubrik: standard_logging_object not yet on model_call_details " @@ -884,6 +992,10 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return payload + @staticmethod + def _copy_block_payload_base(base: StandardLoggingPayload) -> _BlockFailurePayload: + return safe_deep_copy(base) + @staticmethod def _caller_metadata(user_api_key_dict: "UserAPIKeyAuth") -> StandardLoggingUserAPIKeyMetadata: """Identify the caller whose request was blocked. @@ -906,9 +1018,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @classmethod def _build_fallback_payload( cls, - call_details: Mapping[str, Any], + call_details: _FallbackSource, user_api_key_dict: "UserAPIKeyAuth", - ) -> dict[str, Any]: + ) -> _BlockFailurePayload: # Convert datetime to a Unix float so json.dumps can serialize it. # httpx's json= parameter uses stdlib json.dumps with no custom encoder. _raw_start: Final = call_details.get("start_time") @@ -942,7 +1054,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): response: Final = await self.async_httpx_client.post( url=self.logging_endpoint, json=data, - headers=self._headers, + headers=dict(self._headers), ) response.raise_for_status() except httpx.HTTPStatusError as e: @@ -996,7 +1108,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # -- Webhook services ------------------------------------------------------ - async def _post_json(self, endpoint: str, payload: Mapping[str, Any], service_name: str) -> Mapping[str, Any]: + async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> _ModerationResponse: """POST ``payload`` to a Rubrik webhook and return its dict response. Raises: @@ -1006,11 +1118,11 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): verbose_logger.debug("Sending request to %s: %s", service_name, endpoint) http_response: Final = await self.moderation_client.post( endpoint, - json=payload, - headers=self._headers, + json=dict(payload), + headers=dict(self._headers), ) http_response.raise_for_status() - result: Final = http_response.json() + result: Final[_ModerationResponse | None] = http_response.json() if not isinstance(result, dict): raise TypeError( f"{service_name} returned non-dict JSON " @@ -1021,9 +1133,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): async def _post_to_response_moderation_endpoint( self, - response_data: Mapping[str, Any], - request_data: Mapping[str, Any], - ) -> Mapping[str, Any]: + response_data: Mapping[str, object], + request_data: Mapping[str, object], + ) -> _ModerationResponse: """Post the ``{request, response}`` envelope to the after_completion webhook and return its (possibly rewritten) response. @@ -1039,7 +1151,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): "Response moderation service", ) - async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: + async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> _ModerationResponse: """Post a bare OpenAI request to the before_prompt webhook. Returns ``{}`` (passthrough) or a synthetic chat.completion (block). @@ -1047,7 +1159,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return await self._post_json(self.prompt_moderation_endpoint, payload, "Prompt moderation service") @staticmethod - def _extract_prompt_refusal(service_response: Mapping[str, Any]) -> str | None: + def _extract_prompt_refusal(service_response: _ModerationResponse) -> str | None: """Return the refusal text when the prompt was blocked, else None. The before_prompt webhook returns ``{}`` (passthrough) or a synthetic @@ -1063,7 +1175,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod def _extract_response_block( - service_response: Mapping[str, Any], + service_response: _ModerationResponse, all_tool_calls: Sequence[ChatCompletionMessageToolCall], sent_content: str, ) -> BlockedResponseResult | None: diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 972ae1d9856..edd3fdb8c61 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -10,7 +10,7 @@ import asyncio import math import uuid from collections.abc import AsyncIterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast import litellm from litellm._logging import verbose_logger @@ -41,7 +41,13 @@ from litellm.types.integrations.websearch_interception import ( AnthropicServerToolUseBlock, WebSearchInterceptionConfig, ) -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.anthropic import AnthropicThinkingParam +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionAudioParam, + ChatCompletionPredictionContentParam, + OpenAIWebSearchOptions, +) from litellm.types.utils import ( AgenticLoopParams, CallTypes, @@ -51,6 +57,8 @@ from litellm.types.utils import ( from litellm.utils import ProviderConfigManager if TYPE_CHECKING: + from aiohttp import ClientSession + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, @@ -72,6 +80,8 @@ WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: Final = "_websearch_interception_emit_native_b # ``web_search_tool_result`` blocks to inject into the final response. WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks" +_ResponseT = TypeVar("_ResponseT") + class _PlanMetadataView(TypedDict): websearch_native_blocks: Sequence[Mapping[str, object]] | None @@ -85,9 +95,96 @@ class _WebSearchSettingsView(TypedDict): websearch_interception_params: WebSearchInterceptionConfig +class _SearchToolLitellmParams(TypedDict, total=False): + search_provider: str | None + + class _SearchToolConfig(TypedDict, total=False): search_tool_name: str - litellm_params: Mapping[str, object] | None + litellm_params: _SearchToolLitellmParams | None + + +class _LitellmParamsProviderView(TypedDict, total=False): + custom_llm_provider: str + + +class _DeploymentCallKwargsView(TypedDict): + custom_llm_provider: str + litellm_params: _LitellmParamsProviderView + model: str + + +class _AcreateNamedParams(TypedDict, total=False): + metadata: Never + stop_sequences: Never + stream: bool | None + system: str | None + temperature: float | None + thinking: Never + tool_choice: Never + tools: Never + top_k: int | None + top_p: float | None + container: Never + + +class _AsearchNamedParams(TypedDict, total=False): + max_results: int | None + search_domain_filter: Never + max_tokens_per_page: int | None + country: str | None + api_key: str | None + api_base: str | None + timeout: float | None + extra_headers: Never + + +class _AcompletionNamedParams(TypedDict, total=False): + functions: Never + function_call: str | None + timeout: float | None + temperature: float | None + top_p: float | None + n: int | None + stream: bool | None + stream_options: Never + stop: Never + max_tokens: int | None + max_completion_tokens: int | None + modalities: Never + prediction: ChatCompletionPredictionContentParam | None + audio: ChatCompletionAudioParam | None + presence_penalty: float | None + frequency_penalty: float | None + logit_bias: Never + user: str | None + response_format: Never + seed: int | None + tools: Never + tool_choice: Never + parallel_tool_calls: bool | None + logprobs: bool | None + top_logprobs: int | None + deployment_id: str | None + reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None + verbosity: Literal["low", "medium", "high"] | None + safety_identifier: str | None + service_tier: str | None + base_url: str | None + api_version: str | None + api_key: str | None + model_list: Never + extra_headers: Never + thinking: AnthropicThinkingParam | None + web_search_options: OpenAIWebSearchOptions | None + include_server_side_tool_invocations: bool | None + shared_session: "ClientSession | None" + enable_json_schema_validation: bool | None + + +_NO_ACREATE_NAMED: Final[_AcreateNamedParams] = {} +_NO_ASEARCH_NAMED: Final[_AsearchNamedParams] = {} +_NO_ACOMPLETION_NAMED: Final[_AcompletionNamedParams] = {} class WebSearchInterceptionLogger(CustomLogger): @@ -275,12 +372,17 @@ class WebSearchInterceptionLogger(CustomLogger): """ # Check if this is for an enabled provider # Try top-level kwargs first, then nested litellm_params, then derive from model name - custom_llm_provider = kwargs.get("custom_llm_provider", "") or kwargs.get("litellm_params", {}).get( + call_kwargs_view: Final[_DeploymentCallKwargsView] = { + "custom_llm_provider": kwargs.get("custom_llm_provider", ""), + "litellm_params": kwargs.get("litellm_params", {}), + "model": kwargs.get("model", ""), + } + custom_llm_provider = call_kwargs_view["custom_llm_provider"] or call_kwargs_view["litellm_params"].get( "custom_llm_provider", "" ) if not custom_llm_provider: try: - _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs.get("model", "")) + _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=call_kwargs_view["model"]) except Exception: custom_llm_provider = "" if custom_llm_provider not in self.enabled_providers: @@ -903,17 +1005,18 @@ class WebSearchInterceptionLogger(CustomLogger): ) @staticmethod - def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any: + def _inject_native_blocks(response: _ResponseT, native_blocks: Sequence[Mapping[str, object]]) -> _ResponseT: """Prepend native blocks to response content, dict or object form.""" if not native_blocks: return response if isinstance(response, dict): - existing = response.get("content") or [] + existing: Sequence[object] = response.get("content") or [] response["content"] = list(native_blocks) + list(existing) return response existing = getattr(response, "content", None) or [] + content_attribute: Final = "content" try: - response.content = list(native_blocks) + list(existing) + setattr(response, content_attribute, list(native_blocks) + list(existing)) except (AttributeError, TypeError): # Object refused write — fall through and leave the response # untouched rather than crash the request. @@ -1169,10 +1272,10 @@ class WebSearchInterceptionLogger(CustomLogger): messages: list[dict], tool_calls: list[dict], thinking_blocks: list[dict], - anthropic_messages_optional_request_params: dict, + anthropic_messages_optional_request_params: Mapping[str, object], logging_obj: "LiteLLMLoggingObj | None", stream: bool, - kwargs: dict, + kwargs: Mapping[str, object], ) -> "AnthropicMessagesResponse | AsyncIterator[object]": """Legacy path: execute search + build patch + run follow-up call.""" request_patch, structured_results = await self._build_anthropic_request_patch( @@ -1180,9 +1283,9 @@ class WebSearchInterceptionLogger(CustomLogger): messages=messages, tool_calls=tool_calls, thinking_blocks=thinking_blocks, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + anthropic_messages_optional_request_params=dict[str, object](anthropic_messages_optional_request_params), logging_obj=logging_obj, - kwargs=kwargs, + kwargs=dict[str, object](kwargs), ) if request_patch.messages is None: raise ValueError("WebSearchInterception: missing follow-up messages") @@ -1197,12 +1300,14 @@ class WebSearchInterceptionLogger(CustomLogger): if max_tokens is None: max_tokens = cast(int, kwargs.get("max_tokens", 1024)) + patch_kwargs: Final = dict[str, object](request_patch.kwargs) response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate( max_tokens=max_tokens, messages=request_patch.messages, model=request_patch.model or model, + **_NO_ACREATE_NAMED, **optional_params, - **request_patch.kwargs, + **patch_kwargs, ) # Legacy path: the new path goes through the typed plan + core @@ -1344,12 +1449,13 @@ class WebSearchInterceptionLogger(CustomLogger): search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router) search_provider: str | None = None - search_litellm_params: dict[str, Any] = {} + search_litellm_params: Mapping[str, object] = {} search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool) if search_tool is not None: await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs) - search_litellm_params = dict(search_tool.get("litellm_params", {}) or {}) - search_provider = search_litellm_params.get("search_provider") + tool_params: Final[_SearchToolLitellmParams] = search_tool.get("litellm_params", {}) or {} + search_litellm_params = dict[str, object](tool_params) + search_provider = tool_params.get("search_provider") # Fallback to perplexity if no router or no search tools configured if not search_provider: @@ -1377,12 +1483,15 @@ class WebSearchInterceptionLogger(CustomLogger): if key != "search_provider" and value is not None } result: Final = ( - await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs) + await litellm.asearch( + query=query, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs + ) if search_metadata is None else await litellm.asearch( query=query, search_provider=search_provider, litellm_metadata=search_metadata, + **_NO_ASEARCH_NAMED, **search_kwargs, ) ) @@ -1422,7 +1531,7 @@ class WebSearchInterceptionLogger(CustomLogger): valid_token=user_api_key_auth, ) - team_id: Final = getattr(user_api_key_auth, "team_id", None) + team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None) if team_id: from litellm.proxy.proxy_server import ( prisma_client, @@ -1537,10 +1646,10 @@ class WebSearchInterceptionLogger(CustomLogger): model: str, messages: list[dict], tool_calls: list[dict], - optional_params: dict, + optional_params: Mapping[str, object], logging_obj: "LiteLLMLoggingObj | None", stream: bool, - kwargs: dict, + kwargs: Mapping[str, object], response_format: str = "openai", ) -> "ModelResponse | CustomStreamWrapper": """Legacy path: execute search + build patch + run follow-up call.""" @@ -1548,8 +1657,8 @@ class WebSearchInterceptionLogger(CustomLogger): model=model, messages=messages, tool_calls=tool_calls, - optional_params=optional_params, - kwargs=kwargs, + optional_params=dict[str, object](optional_params), + kwargs=dict[str, object](kwargs), response_format=response_format, ) if request_patch.messages is None: @@ -1557,11 +1666,13 @@ class WebSearchInterceptionLogger(CustomLogger): params: Final = dict(optional_params) params.update(request_patch.optional_params) params.pop("tool_choice", None) + patch_kwargs: Final = dict[str, object](request_patch.kwargs) return await litellm.acompletion( model=request_patch.model or model, messages=request_patch.messages, + **_NO_ACOMPLETION_NAMED, **params, - **request_patch.kwargs, + **patch_kwargs, ) async def _build_chat_completion_request_patch( diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index e4a4d23b438..47dfe8de292 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,12 +13,12 @@ Pattern Overview: """ import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from copy import deepcopy from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable -from typing_extensions import assert_never +from typing_extensions import TypedDict, assert_never from litellm._logging import verbose_proxy_logger from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -61,6 +61,7 @@ if TYPE_CHECKING: ModifyResponseException, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -95,6 +96,48 @@ InputWriteBackTarget = ( ) +class _SSEDelta(TypedDict, total=False): + type: str + text: str + stop_reason: str | None + + +class _SSEEventData(TypedDict, total=False): + delta: _SSEDelta + + +def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]: + return value + + +def _content_block_at(blocks: Sequence[object], index: int) -> object: + return blocks[index] + + +@runtime_checkable +class _ModelDumpBlock(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +@runtime_checkable +class _TextAttrBlock(Protocol): + text: str + + +class _WritableMessage(Protocol): + @overload + def get(self, key: str, /) -> object | None: ... + + @overload + def get(self, key: str, default: object, /) -> object: ... + + def __setitem__(self, key: str, value: object, /) -> None: ... + + +def _as_writable(value: _WritableMessage) -> _WritableMessage: + return value + + @dataclass(frozen=True, slots=True) class ScannedText: text: str @@ -123,7 +166,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _build_streaming_usage_response( - responses_so_far: list[Any], + responses_so_far: Sequence[object], request_data: dict | None, ) -> ModelResponse | None: chunks: Final = tuple(response for response in responses_so_far if isinstance(response, (str, bytes))) @@ -141,7 +184,7 @@ class AnthropicMessagesHandler(BaseTranslation): self, exc: "ModifyResponseException", stream_started: bool = False, - responses_so_far: list[Any] | None = None, + responses_so_far: Sequence[object] | None = None, ) -> list[bytes]: """ Build an Anthropic SSE sequence delivering the guardrail block message @@ -159,7 +202,7 @@ class AnthropicMessagesHandler(BaseTranslation): would make Anthropic clients reject the stream. """ if stream_started: - return self._block_continuation_chunks(exc, responses_so_far or []) + return list(self._block_continuation_chunks(exc, responses_so_far or [])) return self._standalone_block_chunks(exc) def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]: @@ -184,7 +227,9 @@ class AnthropicMessagesHandler(BaseTranslation): ) return list(FakeAnthropicMessagesStreamIterator(response=block_response)) - def _block_continuation_chunks(self, exc: "ModifyResponseException", responses_so_far: list[Any]) -> list[bytes]: + def _block_continuation_chunks( + self, exc: "ModifyResponseException", responses_so_far: Sequence[object] + ) -> Sequence[bytes]: """Continue an already-started message: close the open content block, append the block message as a new text block, then end the message -- without a second message_start.""" @@ -234,7 +279,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _content_block_state( - responses_so_far: list[Any], + responses_so_far: Sequence[object], ) -> tuple[int | None, int | None]: """From the SSE chunks already sent to the client, return (open content-block index or None, highest content-block index seen or None). @@ -260,7 +305,20 @@ class AnthropicMessagesHandler(BaseTranslation): return open_index, max_index @staticmethod - def _iter_sse_events(item: Any) -> list[dict]: + def _parse_sse_data_line(raw_line: str) -> tuple[Mapping[str, object], ...]: + line: Final = raw_line.strip() + if not line.startswith("data:"): + return () + try: + parsed: Final[object] = json.loads(line[len("data:") :].strip()) + except json.JSONDecodeError: + return () + if not isinstance(parsed, dict): + return () + return (_as_str_mapping(parsed),) + + @staticmethod + def _iter_sse_events(item: object) -> Sequence[Mapping[str, object]]: """Yield the event-data dicts in one stream chunk. Handles both formats this stream can carry (see @@ -268,22 +326,15 @@ class AnthropicMessagesHandler(BaseTranslation): several events separated by a blank line -- and an already-parsed event ``dict``.""" if isinstance(item, dict): - return [item] + return (_as_str_mapping(item),) if not isinstance(item, (bytes, bytearray)): - return [] - events: Final[list[dict]] = [] - for block in item.decode("utf-8", errors="replace").split("\n\n"): - for line in block.split("\n"): - line = line.strip() - if not line.startswith("data:"): - continue - try: - parsed = json.loads(line[len("data:") :].strip()) - except json.JSONDecodeError: - continue - if isinstance(parsed, dict): - events.append(parsed) - return events + return () + return tuple( + event + for block in item.decode("utf-8", errors="replace").split("\n\n") + for line in block.split("\n") + for event in AnthropicMessagesHandler._parse_sse_data_line(line) + ) def _translate_to_openai(self, data: dict) -> ChatCompletionRequest: """Translate Anthropic request to OpenAI chat completion format.""" @@ -315,8 +366,8 @@ class AnthropicMessagesHandler(BaseTranslation): self, data: dict, guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, - ) -> Any: + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> Mapping[str, object]: """ Process input messages by applying guardrails to text content. """ @@ -467,8 +518,8 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _openai_system_message_to_anthropic( - message: dict[str, Any], - ) -> dict[str, Any] | None: # mutable-ok: API message payload + message: Mapping[str, object], + ) -> dict[str, object] | None: # mutable-ok: API message payload """Convert an OpenAI system message to the client's Anthropic-shaped entry.""" content: Final = message.get("content") if isinstance(content, str): @@ -477,14 +528,14 @@ class AnthropicMessagesHandler(BaseTranslation): ) # mutable-ok: API message payload if not isinstance(content, list): return None - blocks: Final[list[dict[str, Any]]] = [] # mutable-ok: API message payload + blocks: Final[list[dict[str, object]]] = [] # mutable-ok: API message payload for block in content: if not isinstance(block, dict) or block.get("type") != "text": continue text = block.get("text") if not isinstance(text, str) or not text: continue - anthropic_block: dict[str, Any] = { # mutable-ok: API message payload + anthropic_block: dict[str, object] = { # mutable-ok: API message payload "type": "text", "text": text, } # mutable-ok: API message payload @@ -514,7 +565,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _defer_systems_inside_tool_exchanges( - structured_messages: list, # mutable-ok: API message payload + structured_messages: Sequence[Mapping[str, object]], ) -> list: """Hold a system row until the tool exchange around it completes so the call/result pair converts together.""" from litellm.litellm_core_utils.prompt_templates.factory import group_tool_exchanges @@ -602,7 +653,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _extract_midturn_system_text( - message: dict[str, Any], # mutable-ok: API message payload + message: Mapping[str, object], msg_idx: int, ) -> ExtractedInput: """Match the adapter's filtering so positional guardrail write-back stays aligned.""" @@ -636,7 +687,7 @@ class AnthropicMessagesHandler(BaseTranslation): @classmethod def _extract_input_text_and_images( cls, - message: dict[str, Any], + message: Mapping[str, object], msg_idx: int, skip_system_message: bool = False, skip_tool_message: bool = False, @@ -696,7 +747,7 @@ class AnthropicMessagesHandler(BaseTranslation): if scan_only_tool_results: return EMPTY_EXTRACTED_INPUT - text_str: Final = content_item.get("text", None) + text_str: Final[str | None] = content_item.get("text") return ExtractedInput( scanned=( () if text_str is None else (ScannedText(text_str, ContentBlockTextTarget(msg_idx, content_idx)),) @@ -707,7 +758,7 @@ class AnthropicMessagesHandler(BaseTranslation): @classmethod def _extract_tool_result( cls, - content_item: Mapping[str, Any], + content_item: Mapping[str, object], msg_idx: int, content_idx: int, ) -> ExtractedInput: @@ -736,7 +787,7 @@ class AnthropicMessagesHandler(BaseTranslation): ) @staticmethod - def _image_sources(block: Mapping[str, Any]) -> tuple[str, ...]: + def _image_sources(block: Mapping[str, object]) -> tuple[str, ...]: source: Final = block.get("source") if not isinstance(source, Mapping): return () @@ -746,7 +797,7 @@ class AnthropicMessagesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, - messages: list[dict[str, Any]], + messages: Sequence[_WritableMessage], responses: list[str], scanned: tuple[ScannedText, ...], ) -> None: @@ -788,10 +839,10 @@ class AnthropicMessagesHandler(BaseTranslation): self, response: "AnthropicMessagesResponse", guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, - user_api_key_dict: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, - ) -> Any: + ) -> "AnthropicMessagesResponse": """ Process output response by applying guardrails to text content and tool calls. @@ -869,10 +920,10 @@ class AnthropicMessagesHandler(BaseTranslation): self, responses_so_far: list[Any], guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, - user_api_key_dict: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, - ) -> list[Any]: + ) -> Sequence[object]: """ Process output streaming response by applying guardrails to text content. @@ -950,8 +1001,8 @@ class AnthropicMessagesHandler(BaseTranslation): def _prepare_request_data( self, request_data: dict | None, - response: Any, - user_api_key_dict: Any | None, + response: object, + user_api_key_dict: "UserAPIKeyAuth | None", key: str, ) -> dict: """Ensure request_data has the response/responses_so_far key and metadata.""" @@ -968,7 +1019,7 @@ class AnthropicMessagesHandler(BaseTranslation): return request_data @staticmethod - def _get_response_content(response: Any) -> list[Any]: + def _get_response_content(response: object) -> Sequence[object]: """Extract content list from a dict or object response.""" if isinstance(response, dict): return response.get("content", []) or [] @@ -978,7 +1029,7 @@ class AnthropicMessagesHandler(BaseTranslation): def _extract_from_content_blocks( self, - response_content: list[Any], + response_content: Sequence[object], texts_to_check: list[str], images_to_check: list[str], task_mappings: list[tuple[int, int | None]], @@ -986,21 +1037,10 @@ class AnthropicMessagesHandler(BaseTranslation): ) -> None: """Extract text, images, and tool calls from content blocks.""" for content_idx, content_block in enumerate(response_content): - block_dict: dict[str, Any] = {} - if isinstance(content_block, dict): - block_type = content_block.get("type") - block_dict = cast(dict[str, Any], content_block) - elif hasattr(content_block, "type"): - block_type = getattr(content_block, "type", None) - if hasattr(content_block, "model_dump"): - block_dict = content_block.model_dump() - else: - block_dict = { - "type": block_type, - "text": getattr(content_block, "text", None), - } - else: + fields = self._output_block_fields(content_block) + if fields is None: continue + block_type, block_dict = fields if block_type in ["text", "tool_use"]: self._extract_output_text_and_images( @@ -1012,12 +1052,27 @@ class AnthropicMessagesHandler(BaseTranslation): tool_calls_to_check=tool_calls_to_check, ) + @staticmethod + def _output_block_fields(content_block: object) -> "tuple[object, Mapping[str, object]] | None": + if isinstance(content_block, dict): + block_dict: Final = _as_str_mapping(content_block) + return block_dict.get("type"), block_dict + if not hasattr(content_block, "type"): + return None + block_type: Final = getattr(content_block, "type", None) + if isinstance(content_block, _ModelDumpBlock): + return block_type, content_block.model_dump() + return block_type, { + "type": block_type, + "text": getattr(content_block, "text", None), + } + @staticmethod def _build_guardrail_inputs( texts_to_check: list[str], images_to_check: list[str], tool_calls_to_check: list["ChatCompletionToolCallChunk"], - response: Any, + response: object, ) -> "GenericGuardrailAPIInputs": """Build GenericGuardrailAPIInputs with optional images, tool calls, model.""" inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) @@ -1034,7 +1089,7 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["model"] = response_model return inputs - def get_streaming_string_so_far(self, responses_so_far: list[Any]) -> str: + def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str: """ Parse streaming responses and extract accumulated text content. @@ -1105,7 +1160,7 @@ class AnthropicMessagesHandler(BaseTranslation): # Only process content_block_delta events if event_type == "content_block_delta" and data_line: try: - data = json.loads(data_line) + data: _SSEEventData = json.loads(data_line) delta = data.get("delta", {}) if delta.get("type") == "text_delta": text += delta.get("text", "") @@ -1117,7 +1172,7 @@ class AnthropicMessagesHandler(BaseTranslation): return text - def _check_streaming_has_ended(self, responses_so_far: list[Any]) -> bool: + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: """ Check if streaming response has ended by looking for non-null stop_reason. @@ -1168,7 +1223,7 @@ class AnthropicMessagesHandler(BaseTranslation): # Check for message_delta event with stop_reason if event_type == "message_delta" and data_line: try: - data = json.loads(data_line) + data: _SSEEventData = json.loads(data_line) delta = data.get("delta", {}) stop_reason = delta.get("stop_reason") if stop_reason is not None: @@ -1212,7 +1267,7 @@ class AnthropicMessagesHandler(BaseTranslation): def _extract_output_text_and_images( self, - content_block: dict[str, Any], + content_block: Mapping[str, object], content_idx: int, texts_to_check: list[str], images_to_check: list[str], @@ -1235,7 +1290,7 @@ class AnthropicMessagesHandler(BaseTranslation): task_mappings.append((content_idx, None)) # Extract tool calls - elif content_type == "tool_use": + elif content_type == "tool_use" and isinstance(content_block, dict): tool_call: Final = AnthropicConfig.convert_tool_use_to_openai_format( anthropic_tool_content=content_block, index=content_idx, @@ -1260,7 +1315,7 @@ class AnthropicMessagesHandler(BaseTranslation): content_idx = cast(int, mapping[0]) # Handle both dict and object responses - response_content: list[Any] = [] + response_content: Sequence[object] = [] if isinstance(response, dict): response_content = response.get("content", []) or [] elif hasattr(response, "content"): @@ -1276,14 +1331,15 @@ class AnthropicMessagesHandler(BaseTranslation): if content_idx >= len(response_content): continue - content_block = response_content[content_idx] + content_block = _content_block_at(response_content, content_idx) # Verify it's a text block and update the text field # Handle both dict and Pydantic object content blocks if isinstance(content_block, dict): - if content_block.get("type") == "text": - cast(dict[str, Any], content_block)["text"] = guardrail_response + block = _as_writable(content_block) + if block.get("type") == "text": + block["text"] = guardrail_response elif hasattr(content_block, "type") and getattr(content_block, "type", None) == "text": # Update Pydantic object's text attribute - if hasattr(content_block, "text"): + if isinstance(content_block, _TextAttrBlock): content_block.text = guardrail_response diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index dbeac453791..d230b438086 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -13,8 +13,10 @@ Mirrors Anthropic's native ``compact_20260112`` for non-Anthropic providers: """ import re -from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast +from collections.abc import Awaitable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, Union, cast + +from typing_extensions import NotRequired, TypedDict, Unpack import litellm from litellm._logging import verbose_logger @@ -27,11 +29,11 @@ from litellm.types.llms.anthropic import ( if TYPE_CHECKING: from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor, RateLimitResponse from litellm.router import Router from litellm.types.llms.anthropic import ( + AllAnthropicPassThroughMessageValues, AllAnthropicToolsValues, - AnthopicMessagesAssistantMessageParam, - AnthropicMessagesUserMessageParam, ) from litellm.types.llms.openai import ChatCompletionToolParam from litellm.types.utils import ModelResponse @@ -82,6 +84,69 @@ _PROPAGATED_METADATA_KEYS: Final = ( _SUMMARY_TAG_RE: Final = re.compile(r"(.*?)", re.IGNORECASE | re.DOTALL) +_MsgT: Final = TypeVar("_MsgT", bound=Mapping[str, object]) + + +def _as_object(value: object) -> object: + return value + + +def _is_tool_result_block(block: object) -> bool: + return isinstance(block, dict) and block.get("type") in ("tool_result",) + + +class _SummaryCallKwargs(TypedDict): + model: str + max_tokens: int + timeout: float + litellm_metadata: Mapping[str, object] + user: NotRequired[str] + allowed_model_region: NotRequired[str] + + +class _SummaryAcompletion(Protocol): + def __call__( + self, *, messages: Sequence[Mapping[str, object]], **kwargs: Unpack[_SummaryCallKwargs] + ) -> "Awaitable[ModelResponse | CustomStreamWrapper]": ... + + +class _CreateRateLimitDescriptors(Protocol): + def __call__( + self, + *, + user_api_key_dict: "UserAPIKeyAuth", + data: Mapping[str, str], + rpm_limit_type: object, + tpm_limit_type: object, + model_has_failures: bool, + ) -> "Sequence[RateLimitDescriptor]": ... + + +class _AddModelRateLimitDescriptor(Protocol): + def __call__( + self, + *, + user_api_key_dict: "UserAPIKeyAuth", + requested_model: str, + descriptors: "Sequence[RateLimitDescriptor]", + ) -> None: ... + + +class _CreateOrgRateLimitDescriptors(Protocol): + def __call__( + self, user_api_key_dict: "UserAPIKeyAuth", requested_model: str | None = None + ) -> "Sequence[RateLimitDescriptor]": ... + + +class _ShouldRateLimit(Protocol): + def __call__( + self, + *, + descriptors: "Sequence[RateLimitDescriptor]", + parent_otel_span: object, + read_only: bool, + ) -> "Awaitable[RateLimitResponse]": ... + def _read_summary_model_setting() -> str | None: """Look up the configured summarization model from proxy general_settings.""" @@ -157,11 +222,11 @@ async def _check_summary_model_access( return True key_models: Final = list(getattr(user_api_key_auth, "models", None) or []) - team_id: Final = getattr(user_api_key_auth, "team_id", None) + team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None) team_model_aliases: Final = getattr(user_api_key_auth, "team_model_aliases", None) team_models: Final = list(getattr(user_api_key_auth, "team_models", None) or []) - user_id: Final = getattr(user_api_key_auth, "user_id", None) - project_id: Final = getattr(user_api_key_auth, "project_id", None) + user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None) + project_id: Final[str | None] = getattr(user_api_key_auth, "project_id", None) checks: Final[tuple[tuple[Literal["key", "team"], list[str]], ...]] = ( ("key", key_models), @@ -347,7 +412,7 @@ async def _check_summary_model_budget( return False end_user_model_max_budget: Final = getattr(user_api_key_auth, "end_user_model_max_budget", None) - end_user_id: Final = getattr(user_api_key_auth, "end_user_id", None) + end_user_id: Final[str | None] = getattr(user_api_key_auth, "end_user_id", None) if isinstance(end_user_model_max_budget, dict) and end_user_model_max_budget and end_user_id is not None: try: await model_max_budget_limiter.is_end_user_within_model_budget( @@ -399,40 +464,57 @@ async def _check_summary_model_rate_limit( except Exception: return True - limiter: Final = getattr(proxy_logging_obj, "max_parallel_request_limiter", None) + limiter: Final[object] = getattr(proxy_logging_obj, "max_parallel_request_limiter", None) + should_rate_limit_check: Final[_ShouldRateLimit | None] = getattr(limiter, "should_rate_limit", None) + create_descriptors: Final[_CreateRateLimitDescriptors | None] = getattr( + limiter, "_create_rate_limit_descriptors", None + ) + add_team_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr( + limiter, "_add_team_model_rate_limit_descriptor_from_metadata", None + ) + add_project_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr( + limiter, "_add_project_model_rate_limit_descriptor_from_metadata", None + ) + create_org_descriptors: Final[_CreateOrgRateLimitDescriptors | None] = getattr( + limiter, "create_organization_rate_limit_descriptor", None + ) if ( limiter is None - or not hasattr(limiter, "should_rate_limit") - or not hasattr(limiter, "_create_rate_limit_descriptors") + or should_rate_limit_check is None + or create_descriptors is None + or add_team_descriptor is None + or add_project_descriptor is None + or create_org_descriptors is None ): return True try: - metadata: Final = getattr(user_api_key_auth, "metadata", None) or {} + metadata: Final[Mapping[str, object]] = getattr(user_api_key_auth, "metadata", None) or {} data: Final = {"model": summary_model} - descriptors: Final = limiter._create_rate_limit_descriptors( + base_descriptors: Final = create_descriptors( user_api_key_dict=user_api_key_auth, data=data, rpm_limit_type=metadata.get("rpm_limit_type"), tpm_limit_type=metadata.get("tpm_limit_type"), model_has_failures=False, ) - limiter._add_team_model_rate_limit_descriptor_from_metadata( + add_team_descriptor( user_api_key_dict=user_api_key_auth, requested_model=summary_model, - descriptors=descriptors, + descriptors=base_descriptors, ) - limiter._add_project_model_rate_limit_descriptor_from_metadata( + add_project_descriptor( user_api_key_dict=user_api_key_auth, requested_model=summary_model, - descriptors=descriptors, + descriptors=base_descriptors, ) - descriptors.extend(limiter.create_organization_rate_limit_descriptor(user_api_key_auth, summary_model)) + descriptors: Final = (*base_descriptors, *create_org_descriptors(user_api_key_auth, summary_model)) if not descriptors: return True - response: Final = await limiter.should_rate_limit( + parent_otel_span: Final[object] = getattr(user_api_key_auth, "parent_otel_span", None) + response: Final[RateLimitResponse] = await should_rate_limit_check( descriptors=descriptors, - parent_otel_span=getattr(user_api_key_auth, "parent_otel_span", None), + parent_otel_span=parent_otel_span, read_only=True, ) except Exception as e: @@ -446,7 +528,7 @@ async def _check_summary_model_rate_limit( def _find_latest_compaction_index( - messages: list[dict[str, object]], + messages: Sequence[Mapping[str, object]], ) -> tuple[int | None, int | None]: """Return (message_index, block_index) of the most recent compaction block. @@ -465,8 +547,8 @@ def _find_latest_compaction_index( def _slice_around_compaction_block( - messages: list[dict[str, Any]], -) -> tuple[list[dict[str, object]], dict[str, object] | None]: + messages: Sequence[_MsgT], +) -> tuple[Sequence[_MsgT | dict[str, object]], dict[str, object] | None]: """Apply Anthropic's "drop everything before the compaction block" rule. Returns ``(sliced_messages_with_compaction_block, compaction_block_dict)`` @@ -481,19 +563,21 @@ def _slice_around_compaction_block( original_msg: Final = messages[msg_idx] original_content: Final = original_msg["content"] - compaction_block: Final = cast(dict[str, object], original_content[blk_idx]) + if not isinstance(original_content, list): + return messages, None + original_blocks: Final = cast("Sequence[dict[str, object]]", original_content) + compaction_block: Final = original_blocks[blk_idx] # Per Anthropic's contract everything before the compaction block is # dropped, including earlier blocks within the same assistant message. - sliced_content: Final = list(original_content[blk_idx:]) + sliced_content: Final = list(original_blocks[blk_idx:]) - sliced_messages: Final[list[dict[str, object]]] = [{**original_msg, "content": sliced_content}] - sliced_messages.extend(messages[msg_idx + 1 :]) + sliced_messages: Final = [{**original_msg, "content": sliced_content}, *messages[msg_idx + 1 :]] return sliced_messages, compaction_block def _strip_compaction_blocks( - messages: list[dict[str, object]], + messages: Sequence[dict[str, object]], ) -> list[dict[str, object]]: """Drop any ``compaction`` content blocks from messages. @@ -600,7 +684,7 @@ def _propagate_metadata( def _count_effective_tokens( model: str, - effective_messages: list[dict[str, object]], + effective_messages: Sequence[dict[str, object]], compaction_block: CompactionBlock | None, tools: list[dict[str, object]] | None, system: str | list[dict[str, object]] | None = None, @@ -623,7 +707,7 @@ def _count_effective_tokens( try: openai_shape = adapter.translate_anthropic_messages_to_openai( messages=cast( - "list[AnthropicMessagesUserMessageParam | AnthopicMessagesAssistantMessageParam]", + "list[AllAnthropicPassThroughMessageValues]", messages_without_compaction, ) ) @@ -679,17 +763,18 @@ def _system_to_text( return "" if isinstance(system, str): return system - parts: Final[list[str]] = [] - for block in system: - if isinstance(block, dict) and block.get("type") == "text": - text = block.get("text") - if isinstance(text, str) and text: - parts.append(text) - return "\n".join(parts) + return "\n".join( + text + for block in system + if isinstance(block, dict) + and block.get("type") == "text" + and isinstance(text := block.get("text"), str) + and text + ) def _select_last_user_question( - messages: list[dict[str, object]], + messages: Sequence[dict[str, object]], ) -> list[dict[str, object]]: """Pick the most recent ``user`` turn that is a real question. @@ -704,16 +789,18 @@ def _select_last_user_question( turns, or contained no user turns at all). The downstream call always needs a non-empty user message. """ + blocks: Sequence[object] for msg in reversed(messages): if msg.get("role") != "user": continue content = msg.get("content") if isinstance(content, list): - filtered = [blk for blk in content if not (isinstance(blk, dict) and blk.get("type") == "tool_result")] + blocks = [*map(_as_object, content)] + filtered = [blk for blk in blocks if not _is_tool_result_block(blk)] if not filtered: # Purely tool_result — skip and look for an earlier turn. continue - if len(filtered) < len(content): + if len(filtered) < len(blocks): return [{**msg, "content": filtered}] return [msg] return [ @@ -736,7 +823,7 @@ def _extract_summary_text(raw: str | None) -> str | None: def _system_to_openai_message( system: str | list[dict[str, Any]] | None, -) -> dict[str, Any] | None: +) -> Mapping[str, object] | None: """Translate Anthropic-shaped ``system`` to an OpenAI system message. Accepts a bare string or a list of Anthropic content blocks; returns @@ -747,17 +834,19 @@ def _system_to_openai_message( if isinstance(system, str): return {"role": "system", "content": system} if system else None if isinstance(system, list): - parts = [block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text"] + parts: Final[tuple[str, ...]] = tuple( + block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text" + ) joined: Final = "\n\n".join(part for part in parts if part) return {"role": "system", "content": joined} if joined else None return None def _build_summary_messages( - effective_messages: list[dict[str, object]], + effective_messages: Sequence[dict[str, object]], prompt: str, system: str | list[dict[str, object]] | None = None, -) -> list[dict[str, object]]: +) -> Sequence[Mapping[str, object]]: """Build the OpenAI-shape message list for the summary call. The caller's ``system`` prompt is prepended (the default summarization @@ -773,7 +862,7 @@ def _build_summary_messages( try: openai_messages = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( messages=cast( - "list[AnthropicMessagesUserMessageParam | AnthopicMessagesAssistantMessageParam]", + "list[AllAnthropicPassThroughMessageValues]", stripped, ) ) @@ -785,7 +874,7 @@ def _build_summary_messages( ) openai_messages = stripped - summary_messages: Final[list[dict[str, object]]] = [] + summary_messages: Final[list[Mapping[str, object]]] = [] system_message: Final = _system_to_openai_message(system) if system_message is not None: summary_messages.append(system_message) @@ -809,7 +898,7 @@ def _is_user_message(msg: object) -> bool: return isinstance(msg, dict) and msg.get("role") == "user" -def _append_text_to_content(content: Any, extra_text: str) -> Any: +def _append_text_to_content(content: object, extra_text: str) -> object: """Append ``extra_text`` to an OpenAI-shape message ``content`` field. Handles the two common shapes: ``str`` and ``list`` of content parts. @@ -820,16 +909,17 @@ def _append_text_to_content(content: Any, extra_text: str) -> Any: if isinstance(content, str): return f"{content}\n\n{extra_text}" if isinstance(content, list): - return [*content, {"type": "text", "text": extra_text}] + appended: Final[Sequence[object]] = [*map(_as_object, content), {"type": "text", "text": extra_text}] + return appended return [content, {"type": "text", "text": extra_text}] async def _call_summary_model( *, summary_model: str, - summary_messages: list[dict[str, object]], + summary_messages: Sequence[Mapping[str, object]], metadata: Mapping[str, object], - llm_router: Any, + llm_router: object, allowed_model_region: str | None = None, max_tokens: int = COMPACT_SUMMARY_MAX_TOKENS, ) -> Union["ModelResponse", "CustomStreamWrapper"]: @@ -860,9 +950,8 @@ async def _call_summary_model( # the parent ``/v1/messages`` request. On timeout the caller catches the # exception and surfaces ``applied_edits[0].error = "summary_call_failed"``, # forwarding the request without compaction rather than hanging. - call_kwargs: Final[dict[str, Any]] = { + call_kwargs: Final[_SummaryCallKwargs] = { "model": summary_model, - "messages": summary_messages, "max_tokens": max_tokens, "timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS, "litellm_metadata": metadata, @@ -872,19 +961,23 @@ async def _call_summary_model( # than from ``litellm_metadata``, so without it the summary tokens would not # debit the caller's end-user counters. end_user_id: Final = metadata.get("user_api_key_end_user_id") - if end_user_id: + if isinstance(end_user_id, str) and end_user_id: call_kwargs["user"] = end_user_id if allowed_model_region is not None: call_kwargs["allowed_model_region"] = allowed_model_region - if llm_router is not None and hasattr(llm_router, "acompletion"): - return await llm_router.acompletion(**call_kwargs) - return await litellm.acompletion(**call_kwargs) + router_acompletion: Final[_SummaryAcompletion | None] = getattr(llm_router, "acompletion", None) + if llm_router is not None and router_acompletion is not None: + return await router_acompletion(messages=summary_messages, **call_kwargs) + return await litellm.acompletion(messages=[*summary_messages], **call_kwargs) -def _extract_response_text(response: Any) -> str | None: +def _extract_response_text(response: object) -> str | None: try: - choice: Final = response.choices[0] - message: Final = choice.message + choices: Final[Sequence[object] | None] = getattr(response, "choices", None) + if choices is None: + return None + choice: Final = choices[0] + message: Final = getattr(choice, "message", None) content: Final = getattr(message, "content", None) if isinstance(content, str): return content @@ -900,7 +993,7 @@ def _extract_response_text(response: Any) -> str | None: def _extract_usage(response: object) -> tuple[int, int]: - usage: Final = getattr(response, "usage", None) + usage: Final[object] = getattr(response, "usage", None) if usage is None: return 0, 0 return ( diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index 2e0ae30a192..6e720c058fb 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -1,8 +1,10 @@ +from collections.abc import Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict import httpx -from httpx._types import RequestFiles +from httpx._types import FileTypes, RequestFiles +from typing_extensions import NotRequired import litellm from litellm.constants import RUNWAYML_DEFAULT_API_VERSION @@ -31,6 +33,31 @@ else: LiteLLMLoggingObj = Any +class _RunwayTaskResponse(TypedDict, total=False): + id: str + status: str + createdAt: str + completedAt: str + output: Sequence[str] | str + progress: int + failureCode: str + failure: str + + +class _RunwayVideoData(TypedDict): + id: str + object: Literal["video"] + status: str + created_at: int + output_url: NotRequired[str] + completed_at: NotRequired[int] + progress: NotRequired[int] + error: NotRequired[Mapping[str, str]] + model: NotRequired[str] + size: NotRequired[str] + seconds: NotRequired[str] + + class RunwayMLVideoConfig(BaseVideoConfig): """ Configuration class for RunwayML video generation. @@ -44,6 +71,10 @@ class RunwayMLVideoConfig(BaseVideoConfig): def __init__(self): super().__init__() + @staticmethod + def _parse_task_response(raw_response: httpx.Response) -> _RunwayTaskResponse: + return raw_response.json() + def get_supported_openai_params(self, model: str) -> list: """ Get the list of supported OpenAI parameters for video generation. @@ -68,7 +99,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): video_create_optional_params: VideoCreateOptionalRequestParams, model: str, drop_params: bool, - ) -> dict: + ) -> dict[str, object]: """ Map OpenAI parameters to RunwayML format. @@ -78,37 +109,44 @@ class RunwayMLVideoConfig(BaseVideoConfig): - size -> ratio (convert "WIDTHxHEIGHT" to "WIDTH:HEIGHT") - seconds -> duration (convert to integer) """ - mapped_params: Final[dict[str, Any]] = {} + supported_openai_params: Final = self.get_supported_openai_params(model) + return { + **self._prompt_image_param(video_create_optional_params), + **self._ratio_param(video_create_optional_params), + **self._duration_param(video_create_optional_params), + # Pass through other parameters that aren't OpenAI-specific + **{key: value for key, value in video_create_optional_params.items() if key not in supported_openai_params}, + } + @staticmethod + def _prompt_image_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, object]: # Handle input_reference parameter - map to promptImage + # RunwayML supports URLs and data URIs directly if "input_reference" in video_create_optional_params: - input_reference: Final = video_create_optional_params["input_reference"] - # RunwayML supports URLs and data URIs directly - mapped_params["promptImage"] = input_reference + return {"promptImage": video_create_optional_params["input_reference"]} + return {} + @staticmethod + def _ratio_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, str]: # Handle size parameter - convert "1280x720" to "1280:720" if "size" in video_create_optional_params: size: Final = video_create_optional_params["size"] if isinstance(size, str) and "x" in size: - mapped_params["ratio"] = size.replace("x", ":") + return {"ratio": size.replace("x", ":")} + return {} + @staticmethod + def _duration_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, int]: # Handle seconds parameter - convert to integer if "seconds" in video_create_optional_params: seconds: Final = video_create_optional_params["seconds"] if seconds is not None: try: - mapped_params["duration"] = int(float(seconds)) if isinstance(seconds, str) else int(seconds) + return {"duration": int(float(seconds)) if isinstance(seconds, str) else int(seconds)} except (ValueError, TypeError): # If conversion fails, use default duration pass - - # Pass through other parameters that aren't OpenAI-specific - supported_openai_params: Final = self.get_supported_openai_params(model) - for key, value in video_create_optional_params.items(): - if key not in supported_openai_params: - mapped_params[key] = value - - return mapped_params + return {} def validate_environment( self, @@ -163,7 +201,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): model: str, prompt: str, api_base: str, - video_create_optional_request_params: dict, + video_create_optional_request_params: dict[str, object], litellm_params: GenericLiteLLMParams, headers: dict, ) -> tuple[dict, RequestFiles, str]: @@ -179,17 +217,15 @@ class RunwayMLVideoConfig(BaseVideoConfig): "duration": 5 } """ - # Build the request data - request_data: Final[dict[str, Any]] = { + # Build the request data with the mapped parameters merged in + request_data: Final = { "model": model, "promptText": prompt, + **video_create_optional_request_params, } - # Add mapped parameters - request_data.update(video_create_optional_request_params) - # RunwayML uses JSON body, no files multipart - files_list: Final[list[tuple[str, Any]]] = [] + files_list: Final[Sequence[tuple[str, FileTypes]]] = [] # Append the specific endpoint for video generation full_api_base: Final = f"{api_base}/image_to_video" @@ -216,10 +252,10 @@ class RunwayMLVideoConfig(BaseVideoConfig): We map this to OpenAI VideoObject format. """ - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) # Map RunwayML task response to VideoObject format - video_data: Final[dict[str, Any]] = { + video_data: Final[_RunwayVideoData] = { "id": response_data.get("id", ""), "object": "video", "status": self._map_runway_status(response_data.get("status", "pending")), @@ -229,9 +265,8 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Add optional fields if present if "output" in response_data and response_data["output"]: # RunwayML returns output as array of URLs when task succeeds - video_data["output_url"] = ( - response_data["output"][0] if isinstance(response_data["output"], list) else response_data["output"] - ) + output: Final = response_data["output"] + video_data["output_url"] = output if isinstance(output, str) else output[0] if "completedAt" in response_data: video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt")) @@ -254,7 +289,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): if "duration" in request_data: video_data["seconds"] = str(request_data["duration"]) - video_obj: Final = VideoObject(**video_data) + video_obj: Final = VideoObject.model_validate(video_data) if custom_llm_provider and video_obj.id: video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model) @@ -326,20 +361,18 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Get task status to retrieve video URL url: Final = f"{api_base}/tasks/{encoded_video_id}" - params: Final[dict[str, Any]] = {} + return url, dict[str, str]() - return url, params - - def _extract_video_url_from_response(self, response_data: dict[str, Any]) -> str: + def _extract_video_url_from_response(self, response_data: _RunwayTaskResponse) -> str: """ Helper method to extract video URL from RunwayML response. Shared between sync and async transforms. """ # Extract video URL from the output field video_url = None - if "output" in response_data and response_data["output"]: - output: Final = response_data["output"] - video_url = output[0] if isinstance(output, list) else output + raw_output: Final = response_data.get("output") + if raw_output: + video_url = raw_output if isinstance(raw_output, str) else raw_output[0] if not video_url: # Check if the video generation failed or is still processing @@ -373,7 +406,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): "output":["https://dnznrvs05pmza.cloudfront.net/.../video.mp4?_jwt=..."] } """ - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) video_url: Final = self._extract_video_url_from_response(response_data) # Download the video from the CloudFront URL synchronously @@ -402,7 +435,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): "output":["https://dnznrvs05pmza.cloudfront.net/.../video.mp4?_jwt=..."] } """ - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) video_url: Final = self._extract_video_url_from_response(response_data) # Download the video from the CloudFront URL asynchronously @@ -421,7 +454,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - extra_body: dict[str, Any] | None = None, + extra_body: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """ Transform the video remix request for RunwayML API. @@ -448,7 +481,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): after: str | None = None, limit: int | None = None, order: str | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """ Transform the video list request for RunwayML API. @@ -484,9 +517,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Construct the URL for task cancellation url: Final = f"{api_base}/tasks/{encoded_video_id}/cancel" - data: Final[dict[str, Any]] = {} - - return url, data + return url, dict[str, str]() def transform_video_delete_response( self, @@ -494,7 +525,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): logging_obj: LiteLLMLoggingObj, ) -> VideoObject: """Transform the RunwayML video delete/cancel response.""" - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) video_obj: Final = VideoObject( id=response_data.get("id", ""), @@ -524,9 +555,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): url: Final = f"{api_base}/tasks/{encoded_video_id}" # Empty dict for GET request (no body) - data: Final[dict[str, Any]] = {} - - return url, data + return url, dict[str, str]() def transform_video_status_retrieve_response( self, @@ -537,10 +566,10 @@ class RunwayMLVideoConfig(BaseVideoConfig): """ Transform the RunwayML video status retrieve response. """ - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) # Map RunwayML task response to VideoObject format - video_data: Final[dict[str, Any]] = { + video_data: Final[_RunwayVideoData] = { "id": response_data.get("id", ""), "object": "video", "status": self._map_runway_status(response_data.get("status", "pending")), @@ -549,9 +578,8 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Add optional fields if present if "output" in response_data and response_data["output"]: - video_data["output_url"] = ( - response_data["output"][0] if isinstance(response_data["output"], list) else response_data["output"] - ) + output: Final = response_data["output"] + video_data["output_url"] = output if isinstance(output, str) else output[0] if "completedAt" in response_data: video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt")) @@ -565,14 +593,14 @@ class RunwayMLVideoConfig(BaseVideoConfig): "message": response_data.get("failure", "Video generation failed"), } - video_obj: Final = VideoObject(**video_data) + video_obj: Final = VideoObject.model_validate(video_data) if custom_llm_provider and video_obj.id: video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) return video_obj - def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers): + def transform_video_create_character_request(self, name, video: object, api_base, litellm_params, headers): raise NotImplementedError("video create character is not supported for RunwayML") def transform_video_create_character_response(self, raw_response, logging_obj): diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 3538fc5b1a7..131ee41ea8b 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -5,12 +5,13 @@ import json import os import re import time -from collections.abc import Callable, Iterable, Iterator -from typing import Any, Final +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from typing import Final, TypedDict import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted +from typing_extensions import Required import litellm from litellm._uuid import uuid @@ -50,9 +51,10 @@ from litellm.types.llms.openai import ( HttpxBinaryResponseContent, OpenAICreateFileRequestOptionalParams, OpenAIFileObject, + OpenAIFilesPurpose, PathLike, ) -from litellm.types.llms.vertex_ai import GcsBucketResponse +from litellm.types.llms.vertex_ai import GcsBucketResponse, GenerateContentResponseBody from litellm.types.utils import LlmProviders, ModelResponse from ..common_utils import VertexAIError @@ -62,6 +64,47 @@ _GCP_LABEL_VALUE_MAX_LEN: Final = 63 _CUSTOM_ID_RAW_LABEL_PREFIX: Final = "b32_" +class _OpenAIBatchRequestBody(TypedDict, total=False): + model: str + messages: Sequence[AllMessageValues] + + +class _OpenAIBatchJsonlEntry(TypedDict, total=False): + custom_id: Required[object] + body: _OpenAIBatchRequestBody + + +class _VertexBatchOutputRequest(TypedDict, total=False): + labels: Mapping[str, str] + + +class _VertexBatchResponse(GenerateContentResponseBody, total=False): + modelVersion: str + + +class _VertexBatchOutputRow(TypedDict, total=False): + request: _VertexBatchOutputRequest + status: str + processed_time: str + response: _VertexBatchResponse + + +class _GcsObjectMetadata(TypedDict, total=False): + purpose: OpenAIFilesPurpose + + +class _GcsObjectResponse(GcsBucketResponse, total=False): + metadata: _GcsObjectMetadata + + +def _parse_gcs_object_response(raw_response: Response) -> _GcsObjectResponse: + return raw_response.json() + + +def _parse_vertex_batch_output_row(line: str) -> _VertexBatchOutputRow: + return json.loads(line) + + def _sanitize_gcp_label_value(value: str) -> str: """ Sanitize a string to meet GCP label value constraints. @@ -106,7 +149,7 @@ def _decode_gcp_label_value_chunks(values: list[str]) -> str | None: return None -def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any) -> None: +def _litellm_batch_custom_id_labels(custom_id: object) -> Mapping[str, str]: """ Store OpenAI batch custom_id for Vertex batch correlation. @@ -115,15 +158,19 @@ def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any) round-trip correlation in batch output transforms. """ custom_id_str: Final = str(custom_id) - labels["litellm_custom_id"] = _sanitize_gcp_label_value(custom_id_str) raw_label_chunks: Final = _encode_gcp_label_value_chunks(custom_id_str) - labels["litellm_custom_id_raw"] = raw_label_chunks[0] - for index, raw_label_chunk in enumerate(raw_label_chunks[1:], start=1): - labels[f"litellm_custom_id_raw_{index}"] = raw_label_chunk + return { + "litellm_custom_id": _sanitize_gcp_label_value(custom_id_str), + "litellm_custom_id_raw": raw_label_chunks[0], + **{ + f"litellm_custom_id_raw_{index}": raw_label_chunk + for index, raw_label_chunk in enumerate(raw_label_chunks[1:], start=1) + }, + } -def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str: - """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" +def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, str]) -> str: + """Prefer encoded custom_id when present (see _litellm_batch_custom_id_labels).""" raw: Final = labels.get("litellm_custom_id_raw") if raw: raw_chunks: Final = [str(raw)] @@ -141,9 +188,9 @@ def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str: def _openai_batch_jsonl_entry_to_vertex_wrapped_request( - openai_entry: dict[str, Any], - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], -) -> dict[str, Any]: + openai_entry: _OpenAIBatchJsonlEntry, + map_openai_to_vertex_params: Callable[[_OpenAIBatchRequestBody], Mapping[str, object]], +) -> Mapping[str, object]: """ Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request. @@ -151,11 +198,11 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request( Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} """ - openai_request_body: Final = openai_entry.get("body") or {} + openai_request_body: Final[_OpenAIBatchRequestBody] = openai_entry.get("body") or {} vertex_request_body: Final = _transform_request_body( - messages=openai_request_body.get("messages", []), + messages=[*openai_request_body.get("messages", [])], model=openai_request_body.get("model", ""), - optional_params=map_openai_to_vertex_params(openai_request_body), + optional_params=dict(map_openai_to_vertex_params(openai_request_body)), custom_llm_provider="vertex_ai", litellm_params={}, cached_content=None, @@ -163,9 +210,10 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request( custom_id: Final = openai_entry.get("custom_id") if custom_id is not None: - if "labels" not in vertex_request_body: - vertex_request_body["labels"] = {} - _set_litellm_batch_custom_id_labels(vertex_request_body["labels"], custom_id) + vertex_request_body["labels"] = { + **vertex_request_body.get("labels", {}), + **_litellm_batch_custom_id_labels(custom_id), + } return {"request": vertex_request_body} @@ -186,7 +234,7 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: ``str.splitlines()`` + ``line.strip()`` for ``\\n`` / ``\\r\\n`` delimited JSONL. """ - content: Any = openai_file_content + content: FileTypes | str = openai_file_content if isinstance(content, tuple): content = content[1] @@ -241,7 +289,7 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: def _iter_openai_jsonl_entries( openai_file_content: FileTypes, -) -> Iterator[dict[str, Any]]: +) -> Iterator[_OpenAIBatchJsonlEntry]: for line in _iter_openai_jsonl_lines(openai_file_content): yield json.loads(line) @@ -257,7 +305,7 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): def __init__( self, openai_file_content: FileTypes, - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], + map_openai_to_vertex_params: Callable[[_OpenAIBatchRequestBody], Mapping[str, object]], ) -> None: self._openai_file_content = openai_file_content self._map_openai_to_vertex_params = map_openai_to_vertex_params @@ -308,7 +356,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _get_gcs_object_name_from_batch_jsonl( self, - openai_jsonl_content: list[dict[str, Any]], + openai_jsonl_content: Sequence[_OpenAIBatchJsonlEntry], ) -> str: """ Gets a unique GCS object name for the VertexAI batch prediction job @@ -396,8 +444,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _map_openai_to_vertex_params( self, - openai_request_body: dict[str, Any], - ) -> dict[str, Any]: + openai_request_body: _OpenAIBatchRequestBody, + ) -> Mapping[str, object]: """ wrapper to call VertexGeminiConfig.map_openai_params """ @@ -409,7 +457,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): _model: Final = openai_request_body.get("model", "") vertex_params: Final = config.map_openai_params( model=_model, - non_default_params=openai_request_body, + non_default_params=dict(openai_request_body), optional_params={}, drop_params=False, ) @@ -463,10 +511,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Transform VertexAI File upload response into OpenAI-style FileObject """ - response_json: Final = raw_response.json() + response_json: Final = _parse_gcs_object_response(raw_response) try: - response_object: Final = GcsBucketResponse(**response_json) + response_object: Final = _GcsObjectResponse(**response_json) except Exception as e: raise VertexAIError( status_code=raw_response.status_code, @@ -523,7 +571,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> OpenAIFileObject: - response_json: Final = raw_response.json() + response_json: Final = _parse_gcs_object_response(raw_response) gcs_id = response_json.get("id", "") gcs_id = "/".join(gcs_id.split("/")[:-1]) if gcs_id else "" return OpenAIFileObject( @@ -682,7 +730,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # discriminating fields. Anything else (e.g. a binary file whose # first line is not valid UTF-8/JSON) raises and falls through to the # passthrough below, leaving the content untouched. - first_row: Final = json.loads(first_line) + first_row: Final = _parse_vertex_batch_output_row(first_line) is_vertex_batch_output: Final = ( "request" in first_row and "response" in first_row @@ -723,7 +771,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): for line in itertools.chain([first_line], lines): try: openai_output = self._transform_single_vertex_batch_output_to_openai( - vertex_output=json.loads(line), + vertex_output=_parse_vertex_batch_output_row(line), vertex_gemini_config=vertex_gemini_config, logging_obj=batch_transform_logging_obj, mock_httpx_response=mock_httpx_response, @@ -742,11 +790,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _transform_single_vertex_batch_output_to_openai( self, - vertex_output: dict[str, Any], + vertex_output: _VertexBatchOutputRow, vertex_gemini_config: VertexGeminiConfig, logging_obj: Logging, mock_httpx_response: httpx.Response, - ) -> dict[str, Any]: + ) -> Mapping[str, object]: """ Transform a single Vertex AI batch output line to OpenAI format. Uses the existing VertexGeminiConfig transformation for the response. diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index b2fa538b1cc..e6fcff548d6 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -6,8 +6,11 @@ Admins use the management endpoints to read and update input_policy / output_pol """ import uuid +from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, Protocol + +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ToolDiscoveryQueueItem @@ -23,33 +26,109 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient -def _row_to_model(row: dict | Any) -> LiteLLM_ToolTableRow: +class _ToolTableRecord(Protocol): + tool_name: str + input_policy: str | None + output_policy: str | None + + +class _TokenRelationRecord(Protocol): + token: str | None + key_alias: str | None + + +class _TeamRelationRecord(Protocol): + team_id: str | None + team_alias: str | None + + +class _ObjectPermissionRecord(Protocol): + object_permission_id: str + blocked_tools: Sequence[str] | None + verification_tokens: Sequence[_TokenRelationRecord] | None + teams: Sequence[_TeamRelationRecord] | None + + +class _ToolTable(Protocol): + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + order: Mapping[str, str] | None = None, + ) -> Sequence[_ToolTableRecord]: ... + + async def find_unique(self, *, where: Mapping[str, object]) -> _ToolTableRecord | None: ... + + async def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... + + +class _ObjectPermissionTable(Protocol): + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + include: Mapping[str, bool] | None = None, + ) -> Sequence[_ObjectPermissionRecord]: ... + + async def find_unique(self, *, where: Mapping[str, object]) -> _ObjectPermissionRecord | None: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... + + +class _ModelDumpMethod(Protocol): + def __call__(self) -> Mapping: ... + + +_ROW_DICT: Final = TypeAdapter(dict) + + +class _ToolTableHolder(Protocol): + @property + def table(self) -> _ToolTable: ... + + +class _ObjectPermissionTableHolder(Protocol): + @property + def table(self) -> _ObjectPermissionTable: ... + + +def _tool_table(repo: _ToolTableHolder) -> _ToolTable: + return repo.table + + +def _object_permission_table(repo: _ObjectPermissionTableHolder) -> _ObjectPermissionTable: + return repo.table + + +def _row_to_model(row: object) -> LiteLLM_ToolTableRow: """Convert a Prisma model instance or dict to LiteLLM_ToolTableRow.""" - model_dump: Final = getattr(row, "model_dump", None) + model_dump: Final[_ModelDumpMethod | None] = getattr(row, "model_dump", None) if callable(model_dump): row = model_dump() elif not isinstance(row, dict): - row = { - k: getattr(row, k, None) - for k in ( - "tool_id", - "tool_name", - "origin", - "input_policy", - "output_policy", - "call_count", - "assignments", - "key_hash", - "team_id", - "key_alias", - "user_agent", - "last_used_at", - "created_at", - "updated_at", - "created_by", - "updated_by", - ) - } + row = _ROW_DICT.validate_python( + { + k: getattr(row, k, None) + for k in ( + "tool_id", + "tool_name", + "origin", + "input_policy", + "output_policy", + "call_count", + "assignments", + "key_hash", + "team_id", + "key_alias", + "user_agent", + "last_used_at", + "created_at", + "updated_at", + "created_by", + "updated_by", + ) + } + ) return LiteLLM_ToolTableRow( tool_id=row.get("tool_id", ""), tool_name=row.get("tool_name", ""), @@ -87,7 +166,7 @@ async def batch_upsert_tools( if not data: return now: Final = datetime.now(timezone.utc) - table: Final = ToolRepository(prisma_client).table + table: Final = _tool_table(ToolRepository(prisma_client)) for item in data: tool_name = item.get("tool_name", "") origin = item.get("origin") or "user_defined" @@ -132,8 +211,8 @@ async def list_tools( ) -> list[LiteLLM_ToolTableRow]: """Return all tools, optionally filtered by input_policy.""" try: - where: Final = {"input_policy": input_policy} if input_policy is not None else {} - rows: Final = await ToolRepository(prisma_client).table.find_many( + where: Final[Mapping[str, str]] = {"input_policy": input_policy} if input_policy is not None else {} + rows: Final = await _tool_table(ToolRepository(prisma_client)).find_many( where=where, order={"created_at": "desc"}, ) @@ -149,7 +228,7 @@ async def get_tool( ) -> LiteLLM_ToolTableRow | None: """Return a single tool row by tool_name.""" try: - row: Final = await ToolRepository(prisma_client).table.find_unique( + row: Final = await _tool_table(ToolRepository(prisma_client)).find_unique( where={"tool_name": tool_name}, ) if row is None: @@ -172,7 +251,7 @@ async def update_tool_policy( _updated_by: Final = updated_by or "system" now: Final = datetime.now(timezone.utc) - create_data: Final[dict] = { + create_data: Final[Mapping[str, str | datetime]] = { "tool_id": str(uuid.uuid4()), "tool_name": tool_name, "input_policy": input_policy or "untrusted", @@ -182,16 +261,18 @@ async def update_tool_policy( "created_at": now, "updated_at": now, } - update_data: Final[dict] = { - "updated_by": _updated_by, - "updated_at": now, + update_data: Final[Mapping[str, str | datetime]] = { + key: value + for key, value in ( + ("updated_by", _updated_by), + ("updated_at", now), + ("input_policy", input_policy), + ("output_policy", output_policy), + ) + if value is not None } - if input_policy is not None: - update_data["input_policy"] = input_policy - if output_policy is not None: - update_data["output_policy"] = output_policy - await ToolRepository(prisma_client).table.upsert( + await _tool_table(ToolRepository(prisma_client)).upsert( where={"tool_name": tool_name}, data={ "create": create_data, @@ -214,7 +295,7 @@ async def get_tools_by_names( if not tool_names: return {} try: - rows: Final = await ToolRepository(prisma_client).table.find_many( + rows: Final = await _tool_table(ToolRepository(prisma_client)).find_many( where={"tool_name": {"in": tool_names}}, ) return { @@ -239,7 +320,7 @@ async def list_overrides_for_tool( """ out: Final[list[ToolPolicyOverrideRow]] = [] try: - perms: Final = await ObjectPermissionRepository(prisma_client).table.find_many( + perms: Final = await _object_permission_table(ObjectPermissionRepository(prisma_client)).find_many( where={"blocked_tools": {"has": tool_name}}, include={ "verification_tokens": True, @@ -248,8 +329,8 @@ async def list_overrides_for_tool( ) for perm in perms: op_id = getattr(perm, "object_permission_id", None) or "" - tokens = getattr(perm, "verification_tokens", []) or [] - teams = getattr(perm, "teams", []) or [] + tokens: Sequence[_TokenRelationRecord] = getattr(perm, "verification_tokens", []) or [] + teams: Sequence[_TeamRelationRecord] = getattr(perm, "teams", []) or [] for t in tokens: out.append( ToolPolicyOverrideRow( @@ -302,7 +383,7 @@ class ToolPolicyRegistry: try: tools: Final = await call_with_db_reconnect_retry( prisma_client, - lambda: ToolRepository(prisma_client).table.find_many(), + lambda: _tool_table(ToolRepository(prisma_client)).find_many(), reason="sync_tool_policy_from_db_tools_lookup_failure", ) self._tool_input_policies = { @@ -314,13 +395,13 @@ class ToolPolicyRegistry: perms: Final = await call_with_db_reconnect_retry( prisma_client, - lambda: ObjectPermissionRepository(prisma_client).table.find_many(), + lambda: _object_permission_table(ObjectPermissionRepository(prisma_client)).find_many(), reason="sync_tool_policy_from_db_perms_lookup_failure", ) self._blocked_tools_by_op_id = {} for row in perms: op_id = getattr(row, "object_permission_id", None) - blocked = getattr(row, "blocked_tools", None) or [] + blocked: Sequence[str] = getattr(row, "blocked_tools", None) or [] if op_id: self._blocked_tools_by_op_id[op_id] = list(blocked) @@ -352,10 +433,12 @@ class ToolPolicyRegistry: """ if not tool_names: return {} - blocked: Final[set] = set() - for op_id in (object_permission_id, team_object_permission_id): - if op_id and op_id.strip(): - blocked.update(self._blocked_tools_by_op_id.get(op_id.strip(), [])) + blocked: Final[frozenset[str]] = frozenset( + tool + for op_id in (object_permission_id, team_object_permission_id) + if op_id and op_id.strip() + for tool in self._blocked_tools_by_op_id.get(op_id.strip(), []) + ) result: Final[dict[str, str]] = {} for name in tool_names: if name in blocked: @@ -385,18 +468,17 @@ async def add_tool_to_object_permission_blocked( if not object_permission_id or not tool_name: return False try: - row: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( + row: Final = await _object_permission_table(ObjectPermissionRepository(prisma_client)).find_unique( where={"object_permission_id": object_permission_id}, ) if row is None: return False - current: Final = list(getattr(row, "blocked_tools", []) or []) + current: Final[Sequence[str]] = getattr(row, "blocked_tools", []) or [] if tool_name in current: return True - current.append(tool_name) - await ObjectPermissionRepository(prisma_client).table.update( + await _object_permission_table(ObjectPermissionRepository(prisma_client)).update( where={"object_permission_id": object_permission_id}, - data={"blocked_tools": current}, + data={"blocked_tools": [*current, tool_name]}, ) return True except Exception as e: @@ -413,18 +495,17 @@ async def remove_tool_from_object_permission_blocked( if not object_permission_id or not tool_name: return False try: - row: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( + row: Final = await _object_permission_table(ObjectPermissionRepository(prisma_client)).find_unique( where={"object_permission_id": object_permission_id}, ) if row is None: return False - current = list(getattr(row, "blocked_tools", []) or []) + current: Final[Sequence[str]] = getattr(row, "blocked_tools", []) or [] if tool_name not in current: return False - current = [t for t in current if t != tool_name] - await ObjectPermissionRepository(prisma_client).table.update( + await _object_permission_table(ObjectPermissionRepository(prisma_client)).update( where={"object_permission_id": object_permission_id}, - data={"blocked_tools": current}, + data={"blocked_tools": [t for t in current if t != tool_name]}, ) return True except Exception as e: diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index a416a197ab8..7dfc80d2b2a 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -17,7 +17,7 @@ import json import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final, Literal, cast +from typing import Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, overload import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -85,14 +85,6 @@ from litellm.types.proxy.management_endpoints.scim_v2 import ( if TYPE_CHECKING: from prisma import models as prisma_models from prisma import types as prisma_types - from prisma.actions import ( - LiteLLM_InvitationLinkActions, - LiteLLM_OrganizationMembershipActions, - LiteLLM_TeamMembershipActions, - LiteLLM_TeamTableActions, - LiteLLM_UserTableActions, - LiteLLM_VerificationTokenActions, - ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.proxy_server import PrismaClient @@ -100,55 +92,151 @@ if TYPE_CHECKING: router: Final = APIRouter() +_PrismaTableT = TypeVar("_PrismaTableT", covariant=True) + + +class _TableActions(Protocol[_PrismaTableT]): + async def find_unique( + self, + *, + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> "_PrismaTableT | None": ... + + async def find_first(self, *, where: Mapping[str, object]) -> "_PrismaTableT | None": ... + + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None, + skip: int | None = None, + take: int | None = None, + ) -> "Sequence[_PrismaTableT]": ... + + async def create(self, *, data: Mapping[str, object]) -> "_PrismaTableT": ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> "_PrismaTableT | None": ... + + async def delete_many(self, *, where: Mapping[str, object] | None = None) -> int: ... + + +class _PrismaTableHolder(Protocol[_PrismaTableT]): + @property + def table(self) -> "_TableActions[_PrismaTableT]": ... + + +def _typed_table(holder: "_PrismaTableHolder[_PrismaTableT]") -> "_TableActions[_PrismaTableT]": + return holder.table + + +class _LenientTableActions(Protocol[_PrismaTableT]): + async def find_first(self, *, where: Mapping[str, object]) -> "_PrismaTableT | None": ... + + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None, + skip: int | None = None, + take: int | None = None, + ) -> "Sequence[_PrismaTableT] | None": ... + + async def count(self, *, where: Mapping[str, object] | None = None) -> int: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + +class _LenientTableHolder(Protocol[_PrismaTableT]): + @property + def table(self) -> "_LenientTableActions[_PrismaTableT]": ... + + +def _lenient_table(holder: "_LenientTableHolder[_PrismaTableT]") -> "_LenientTableActions[_PrismaTableT]": + return holder.table + + +class _UserDeleteRow(Protocol): + user_id: str + user_email: str | None + + @property + def teams(self) -> Sequence[str]: ... + + def json(self, *, exclude_none: bool) -> str: ... + + +class _TeamCleanupRow(Protocol): + team_id: str + members_with_roles: str + + def model_dump(self) -> Mapping[str, object]: ... + def _user_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_UserTableActions[prisma_models.LiteLLM_UserTable]": - user_table: Final[LiteLLM_UserTableActions[prisma_models.LiteLLM_UserTable]] = UserRepository(prisma_client).table - return user_table +) -> "_TableActions[prisma_models.LiteLLM_UserTable]": + return _typed_table(UserRepository(prisma_client)) + + +def _user_table_lenient( + prisma_client: "PrismaClient | None", +) -> "_LenientTableActions[prisma_models.LiteLLM_UserTable]": + return _lenient_table(UserRepository(prisma_client)) + + +def _user_delete_table( + prisma_client: "PrismaClient | None", +) -> "_TableActions[_UserDeleteRow]": + return _typed_table(UserRepository(prisma_client)) def _team_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]": - team_table: Final[LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]] = TeamRepository(prisma_client).table - return team_table +) -> "_TableActions[prisma_models.LiteLLM_TeamTable]": + return _typed_table(TeamRepository(prisma_client)) + + +def _team_cleanup_table( + prisma_client: "PrismaClient | None", +) -> "_TableActions[_TeamCleanupRow]": + return _typed_table(TeamRepository(prisma_client)) def _verification_token_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]": - token_table: Final[LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]] = ( - VerificationTokenRepository(prisma_client).table - ) - return token_table +) -> "_TableActions[prisma_models.LiteLLM_VerificationToken]": + return _typed_table(VerificationTokenRepository(prisma_client)) + + +def _verification_token_table_lenient( + prisma_client: "PrismaClient | None", +) -> "_LenientTableActions[prisma_models.LiteLLM_VerificationToken]": + return _lenient_table(VerificationTokenRepository(prisma_client)) + + +def _organization_table( + prisma_client: "PrismaClient | None", +) -> "_TableActions[prisma_models.LiteLLM_OrganizationTable]": + return _typed_table(OrganizationRepository(prisma_client)) def _organization_membership_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_OrganizationMembershipActions[prisma_models.LiteLLM_OrganizationMembership]": - membership_table: Final[LiteLLM_OrganizationMembershipActions[prisma_models.LiteLLM_OrganizationMembership]] = ( - OrganizationMembershipRepository(prisma_client).table - ) - return membership_table +) -> "_TableActions[prisma_models.LiteLLM_OrganizationMembership]": + return _typed_table(OrganizationMembershipRepository(prisma_client)) def _invitation_link_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_InvitationLinkActions[prisma_models.LiteLLM_InvitationLink]": - invitation_table: LiteLLM_InvitationLinkActions[prisma_models.LiteLLM_InvitationLink] = InvitationLinkRepository( - prisma_client - ).table - return invitation_table +) -> "_TableActions[prisma_models.LiteLLM_InvitationLink]": + return _typed_table(InvitationLinkRepository(prisma_client)) def _team_membership_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership]": - team_membership_table: Final[LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership]] = ( - TeamMembershipRepository(prisma_client).table - ) - return team_membership_table +) -> "_TableActions[prisma_models.LiteLLM_TeamMembership]": + return _typed_table(TeamMembershipRepository(prisma_client)) def _hash_password_in_dict(data: dict) -> None: @@ -234,7 +322,7 @@ async def _check_duplicate_user_field( if case_insensitive: where_clause[field_name]["mode"] = "insensitive" - existing_user: Final = await UserRepository(prisma_client).table.find_first(where=where_clause) + existing_user: Final = await _user_table_lenient(prisma_client).find_first(where=where_clause) if existing_user is not None: existing_value: Final = getattr(existing_user, field_name, value) @@ -650,7 +738,7 @@ async def ui_get_available_role( def get_team_from_list( - team_list: list[LiteLLM_TeamTable] | list[TeamListResponseObject] | None, + team_list: Sequence[LiteLLM_TeamTable] | Sequence[TeamListResponseObject] | None, team_id: str, ) -> LiteLLM_TeamTable | LiteLLM_TeamMembership | None: if team_list is None: @@ -732,18 +820,59 @@ def _enforce_user_info_access(user_id: str | None, user_api_key_dict: UserAPIKey ) -async def _get_user_info_teams( - prisma_client: Any, +_TeamIdList: TypeAlias = list[str] + + +class _UserInfoDataClient(Protocol): + @overload + async def get_data(self, *, user_id: str) -> "prisma_models.LiteLLM_UserTable | None": ... + + @overload + async def get_data( + self, + *, + user_id: str | None, + table_name: Literal["key"], + query_type: Literal["find_all"], + ) -> "Sequence[LiteLLM_VerificationToken] | None": ... + + @overload + async def get_data( + self, + *, + team_id_list: _TeamIdList, + table_name: Literal["team"], + query_type: Literal["find_all"], + ) -> "Sequence[TeamListResponseObject] | None": ... + + +async def _get_user_info_row( + prisma_client: "_UserInfoDataClient", + user_id: str, +) -> "prisma_models.LiteLLM_UserTable | None": + return await prisma_client.get_data(user_id=user_id) + + +async def _get_user_info_keys( + prisma_client: "_UserInfoDataClient", user_id: str | None, - user_info: Any | None, +) -> "Sequence[LiteLLM_VerificationToken] | None": + return await prisma_client.get_data( + user_id=user_id, + table_name="key", + query_type="find_all", + ) + + +async def _get_user_info_teams( + prisma_client: "_UserInfoDataClient", + user_id: str | None, + user_info: "prisma_models.LiteLLM_UserTable", user_api_key_dict: UserAPIKeyAuth, -) -> tuple[list[Any], list[Any] | None]: +) -> tuple[Sequence[TeamListResponseObject], Sequence[TeamListResponseObject] | None]: """Fetch and merge teams from membership + user.teams field.""" from litellm.proxy.management_endpoints.team_endpoints import list_team - team_list: list[Any] = [] - team_id_list: list[str] = [] - teams_1: Final = await list_team( http_request=Request( scope={"type": "http", "path": "/user/info"}, @@ -752,11 +881,10 @@ async def _get_user_info_teams( user_api_key_dict=user_api_key_dict, ) - if teams_1 is not None and isinstance(teams_1, list): - team_list = teams_1 - team_id_list = [team.team_id for team in teams_1] + team_list: Final = teams_1 if teams_1 is not None and isinstance(teams_1, list) else list[TeamListResponseObject]() + team_id_list: Final = [team.team_id for team in team_list] - teams_2: list[Any] | None = None + teams_2: Sequence[TeamListResponseObject] | None = None target_team_ids: Final = getattr(user_info, "teams", None) if target_team_ids and isinstance(target_team_ids, list): @@ -767,7 +895,7 @@ async def _get_user_info_teams( ) elif user_api_key_dict.user_id is not None and user_id is None: caller_user_info: Final = await prisma_client.get_data(user_id=user_api_key_dict.user_id) - caller_team_ids: Final = getattr(caller_user_info, "teams", None) + caller_team_ids: Final = caller_user_info.teams if caller_user_info is not None else None if caller_team_ids: teams_2 = await prisma_client.get_data( team_id_list=caller_team_ids, @@ -804,9 +932,9 @@ def _redact_scim_enterprise_metadata( def _build_user_info_response( user_id: str | None, user_info: Any | None, - keys: list[LiteLLM_VerificationToken] | None, - team_list: list[Any], - teams_1: list[Any] | None, + keys: Sequence[LiteLLM_VerificationToken] | None, + team_list: Sequence[TeamListResponseObject], + teams_1: Sequence[TeamListResponseObject] | None, ) -> UserInfoResponse: """Create UserInfoResponse while filtering sensitive fields.""" if user_info is None and keys is not None: @@ -814,7 +942,7 @@ def _build_user_info_response( user_info = {"spend": spend} returned_keys: Final = _process_keys_for_user_info(keys=keys, all_teams=teams_1) - team_list.sort(key=lambda x: getattr(x, "team_alias", "") or "") + sorted_team_list: Final = sorted(team_list, key=lambda x: getattr(x, "team_alias", "") or "") _user_info: Final = user_info.model_dump() if isinstance(user_info, BaseModel) else user_info if isinstance(_user_info, dict): @@ -825,7 +953,7 @@ def _build_user_info_response( user_id=user_id, user_info=_user_info, keys=returned_keys, - teams=team_list, + teams=sorted_team_list, ) @@ -870,9 +998,9 @@ async def user_info( user_id = user_api_key_dict.user_id ## GET USER ROW ## - user_info = None + user_info: prisma_models.LiteLLM_UserTable | None = None if user_id is not None: - user_info = await prisma_client.get_data(user_id=user_id) + user_info = await _get_user_info_row(prisma_client, user_id) if user_info is None: raise HTTPException( @@ -888,11 +1016,7 @@ async def user_info( ) ## GET ALL KEYS ## - keys: Final = await prisma_client.get_data( - user_id=user_id, - table_name="key", - query_type="find_all", - ) + keys: Final = await _get_user_info_keys(prisma_client, user_id) response_data: Final = _build_user_info_response( user_id=user_id, @@ -1058,6 +1182,12 @@ async def user_info_v2( raise handle_exception_on_proxy(e) +async def _fetch_admin_teams_and_keys_rows( + prisma_client: "PrismaClient", sql_query: str +) -> Sequence[Mapping[str, Sequence[Mapping[str, object]] | None]]: + return await prisma_client.db.query_raw(sql_query) + + async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): """ Admin UI Endpoint - Returns All Teams and Keys when Proxy Admin is querying @@ -1081,22 +1211,25 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - results: Final = await prisma_client.db.query_raw(sql_query) + results: Final = await _fetch_admin_teams_and_keys_rows(prisma_client, sql_query) verbose_proxy_logger.debug("results_keys: %s", results) - _keys_in_db: Final[list] = results[0]["keys"] or [] + _keys_in_db: Final[Sequence[Mapping[str, object]]] = results[0]["keys"] or [] # cast all keys to LiteLLM_VerificationToken keys_in_db: Final = [] for key in _keys_in_db: - if key.get("models") is None: - key["models"] = [] - keys_in_db.append(LiteLLM_VerificationToken.model_validate(key)) + key_payload = dict[str, object](key) + if key_payload.get("models") is None: + key_payload["models"] = [] + keys_in_db.append(LiteLLM_VerificationToken.model_validate(key_payload)) # cast all teams to LiteLLM_TeamTable - _teams_in_db: list = results[0]["teams"] or [] - _teams_in_db = [LiteLLM_TeamTable.model_validate(team) for team in _teams_in_db] - _teams_in_db.sort(key=lambda x: getattr(x, "team_alias", "") or "") + _teams_rows: Final[Sequence[Mapping[str, object]]] = results[0]["teams"] or [] + _teams_in_db: Final = sorted( + (LiteLLM_TeamTable.model_validate(team) for team in _teams_rows), + key=lambda x: getattr(x, "team_alias", "") or "", + ) returned_keys: Final = _process_keys_for_user_info(keys=keys_in_db, all_teams=_teams_in_db) # Get admin's own user_id and user_info @@ -1121,8 +1254,8 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): def _process_keys_for_user_info( - keys: list[LiteLLM_VerificationToken] | None, - all_teams: list[LiteLLM_TeamTable] | list[TeamListResponseObject] | None, + keys: Sequence[LiteLLM_VerificationToken] | None, + all_teams: Sequence[LiteLLM_TeamTable] | Sequence[TeamListResponseObject] | None, ): from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy.proxy_server import general_settings, litellm_master_key_hash @@ -1212,7 +1345,7 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda async def _schedule_user_update_audit_log( - response: dict[str, Any], + response: Mapping[str, object], existing_user_row: BaseModel | None, litellm_changed_by: str | None, user_api_key_dict: UserAPIKeyAuth, @@ -1768,7 +1901,10 @@ async def bulk_user_update( # Apply update transformations (reuse existing logic) data_json: Final[dict] = data.user_updates.model_dump(exclude_unset=True) - non_default_values: Final = _update_internal_user_params(data_json=data_json, data=data.user_updates) + _raw_update_values: Final[Mapping[str, object]] = _update_internal_user_params( + data_json=data_json, data=data.user_updates + ) + non_default_values: Final = dict[str, object](_raw_update_values) # Remove user identification fields since we're updating by user_id non_default_values.pop("user_id", None) @@ -1780,7 +1916,7 @@ async def bulk_user_update( try: # Perform bulk database update - await UserRepository(prisma_client).table.update_many( + await _user_table_lenient(prisma_client).update_many( where={}, data=non_default_values, # Update all users ) @@ -1885,7 +2021,7 @@ async def get_user_key_counts( # Get count for each user_id individually for user_id in user_ids: - count = await VerificationTokenRepository(prisma_client).table.count( + count = await _verification_token_table_lenient(prisma_client).count( where={ "user_id": user_id, "OR": [ @@ -2122,7 +2258,7 @@ async def get_users( _validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None ) - users: Sequence[prisma_models.LiteLLM_UserTable] | None = await UserRepository(prisma_client).table.find_many( + users: Sequence[prisma_models.LiteLLM_UserTable] | None = await _user_table_lenient(prisma_client).find_many( where=where_conditions, skip=skip, take=page_size, @@ -2130,7 +2266,7 @@ async def get_users( ) # Get total count of user rows - total_count: Final[int] = await UserRepository(prisma_client).table.count(where=where_conditions) + total_count: Final[int] = await _user_table_lenient(prisma_client).count(where=where_conditions) # Get key count for each user if users is not None: @@ -2256,7 +2392,7 @@ async def delete_user( # check that all teams passed exist for user_id in data.user_ids: - user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) + user_row = await _user_delete_table(prisma_client).find_unique(where={"user_id": user_id}) if user_row is None: raise HTTPException( @@ -2308,8 +2444,8 @@ async def delete_user( ) ## CLEANUP MEMBERS_WITH_ROLES - fetch_all_teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_row.teams}}) - teams_to_update = [] + fetch_all_teams = await _team_cleanup_table(prisma_client).find_many(where={"team_id": {"in": user_row.teams}}) + teams_to_update = list[_TeamCleanupRow]() for team in fetch_all_teams: is_member_in_team, new_team_members = _cleanup_members_with_roles( existing_team_row=LiteLLM_TeamTable.model_validate(team.model_dump()), @@ -2327,7 +2463,7 @@ async def delete_user( ## update teams for team in teams_to_update: - await TeamRepository(prisma_client).table.update( + await _team_cleanup_table(prisma_client).update( where={"team_id": team.team_id}, data={"members_with_roles": team.members_with_roles}, ) @@ -2382,14 +2518,14 @@ async def add_internal_user_to_organization( try: # Check if organization_id exists - organization_row: Final = await OrganizationRepository(prisma_client).table.find_unique( + organization_row: Final = await _organization_table(prisma_client).find_unique( where={"organization_id": organization_id} ) if organization_row is None: raise Exception(f"Organization not found, passed organization_id={organization_id}") # Create a new organization membership entry - new_membership: Final = await OrganizationMembershipRepository(prisma_client).table.create( + new_membership: Final = await _organization_membership_table(prisma_client).create( data={ "user_id": user_id, "organization_id": organization_id, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2a385c4c42a..8fe388d2fea 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -18,9 +18,10 @@ import os import re import secrets import traceback -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from contextlib import AbstractAsyncContextManager from datetime import datetime, timedelta, timezone -from typing import Any, Final, Literal, Optional, Protocol, TypeVar, cast +from typing import Any, Final, Literal, Optional, Protocol, TypeAlias, TypeVar, cast import fastapi import yaml @@ -88,6 +89,7 @@ from litellm.proxy.management_endpoints.common_utils import ( from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, ) +from litellm.proxy.management_helpers import object_permission_utils from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.management_helpers.object_permission_utils import ( _set_object_permission, @@ -144,6 +146,7 @@ from litellm.types.utils import ( ) _PrismaRowT = TypeVar("_PrismaRowT") +_PrismaRowCoT: Final = TypeVar("_PrismaRowCoT", covariant=True) _RepositoryModelT = TypeVar("_RepositoryModelT", bound=BaseModel) @@ -169,7 +172,7 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): *, where: Mapping[str, object] | None = None, include: Mapping[str, object] | None = None, - order: Mapping[str, object] | None = None, + order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None, skip: int | None = None, take: int | None = None, ) -> list[_PrismaRowT]: ... @@ -189,17 +192,73 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): data: Mapping[str, object], ) -> _PrismaRowT | None: ... + async def upsert( + self, + *, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> _PrismaRowT: ... -class _UserRowLike(Protocol): - user_id: str | None - user_email: str | None - user_alias: str | None - def model_dump(self) -> Mapping[str, object]: ... +class _PrismaTableHolder(Protocol[_PrismaRowT]): + @property + def table(self) -> _PrismaTableActions[_PrismaRowT]: ... + +def _typed_table(holder: _PrismaTableHolder[_PrismaRowT]) -> _PrismaTableActions[_PrismaRowT]: + return holder.table + + +class _CustomKeyHooksModule(Protocol): + user_custom_key_generate: Callable[..., Awaitable[Mapping[str, object]]] | None + user_custom_key_update: Callable[..., Awaitable[Mapping[str, object]]] | None + + +def _custom_key_generate_hook( + hooks: _CustomKeyHooksModule, +) -> Callable[..., Awaitable[Mapping[str, object]]] | None: + return hooks.user_custom_key_generate + + +def _custom_key_update_hook( + hooks: _CustomKeyHooksModule, +) -> Callable[..., Awaitable[Mapping[str, object]]] | None: + return hooks.user_custom_key_update + + +class _LegacyDumpable(Protocol): def dict(self) -> Mapping[str, object]: ... +def _legacy_model_dict(row: _LegacyDumpable) -> Mapping[str, object]: + return row.dict() + + +class _PrismaTableLenient(Protocol[_PrismaRowCoT]): + async def find_unique(self, *, where: Mapping[str, object]) -> _PrismaRowCoT: ... + + async def find_many(self, *, where: Mapping[str, object] | None = None) -> Sequence[_PrismaRowCoT] | None: ... + + +class _PrismaTableLenientHolder(Protocol[_PrismaRowCoT]): + @property + def table(self) -> _PrismaTableLenient[_PrismaRowCoT]: ... + + +def _lenient_table(holder: _PrismaTableLenientHolder[_PrismaRowCoT]) -> _PrismaTableLenient[_PrismaRowCoT]: + return holder.table + + +def _prisma_table_lenient( + repository: BaseRepository[_RepositoryModelT], +) -> _PrismaTableLenient[_RepositoryModelT]: + return _lenient_table(repository) + + +def _jsonify_for_db(client: PrismaClient, data: Mapping[str, object]) -> Mapping[str, object]: + return client.jsonify_object(dict[str, object](data)) + + class _TxTables(Protocol): litellm_proxymodeltable: _PrismaTableActions[object] @@ -207,21 +266,82 @@ class _TxTables(Protocol): def _prisma_table( repository: BaseRepository[_RepositoryModelT], ) -> _PrismaTableActions[_RepositoryModelT]: - return repository.table + return _typed_table(repository) def _deleted_verification_token_table( prisma_client: PrismaClient, ) -> _PrismaTableActions[LiteLLM_DeletedVerificationToken]: - return DeletedVerificationTokenRepository(prisma_client).table + return _typed_table(DeletedVerificationTokenRepository(prisma_client)) def _credentials_table(prisma_client: PrismaClient) -> _PrismaTableActions[CredentialItem]: - return CredentialsRepository(prisma_client).table + return _typed_table(CredentialsRepository(prisma_client)) def _config_table(prisma_client: PrismaClient) -> _PrismaTableActions[ConfigParam]: - return ConfigRepository(prisma_client).table + return _typed_table(ConfigRepository(prisma_client)) + + +def _deprecated_verification_token_table(prisma_client: PrismaClient) -> _PrismaTableActions[object]: + return _typed_table(DeprecatedVerificationTokenRepository(prisma_client)) + + +_StringList: TypeAlias = list[str] + + +class _CreatedUserRow(Protocol): + models: _StringList + + +def _created_user_row(user_row: "_CreatedUserRow | None") -> "_CreatedUserRow | None": + return user_row + + +async def _query_raw_text_rows(prisma_client: PrismaClient, sql: str, *params: object) -> Sequence[Mapping[str, str]]: + return await prisma_client.db.query_raw(sql, *params) + + +def _as_object_dict(values: Mapping[str, object]) -> Mapping[str, object]: + return values + + +def _model_items(model: BaseModel) -> Iterator[tuple[str, object]]: + return iter(model) + + +class _SpendCache(Protocol): + async def async_get_cache(self, key: str) -> float | None: ... + + +def _spend_cache(cache: _SpendCache) -> _SpendCache: + return cache + + +class _ObjectPermissionUtils(Protocol): + @property + def attach_object_permission_to_dict( + self, + ) -> Callable[..., Awaitable[Mapping[str, object]]]: ... + + +def _object_permission_utils(module: _ObjectPermissionUtils) -> _ObjectPermissionUtils: + return module + + +class _EnvVarsParam(Protocol): + @property + def param_value(self) -> Mapping[str, str] | None: ... + + +def _env_vars_param_value(param: _EnvVarsParam) -> Mapping[str, str] | None: + return param.param_value + + +def _tx_tables_context( + open_tx: Callable[[], AbstractAsyncContextManager[_TxTables]], +) -> AbstractAsyncContextManager[_TxTables]: + return open_tx() async def _check_custom_key_allowed(custom_key_value: str | None) -> None: @@ -886,7 +1006,7 @@ async def _common_key_generation_helper( # check if user set default key/generate params on config.yaml if litellm.default_key_generate_params is not None: - for elem in data: + for elem in _model_items(data): key, value = elem if value is None and key in [ "max_budget", @@ -984,9 +1104,9 @@ async def _common_key_generation_helper( soft_budget=data.soft_budget, model_max_budget=data.model_max_budget or {}, ) - new_budget: Final = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) + new_budget: Final = _jsonify_for_db(prisma_client, budget_row.json(exclude_none=True)) - _budget: Final = await BudgetRepository(prisma_client).table.create( + _budget: Final = await _prisma_table(BudgetRepository(prisma_client)).create( data={ **new_budget, "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, @@ -1655,11 +1775,11 @@ async def generate_key_fn( - user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id. """ try: + from litellm.proxy import proxy_server from litellm.proxy._types import CommonProxyErrors from litellm.proxy.proxy_server import ( prisma_client, user_api_key_cache, - user_custom_key_generate, ) if prisma_client is None: @@ -1686,7 +1806,7 @@ async def generate_key_fn( ) custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = ( - user_custom_key_generate + _custom_key_generate_hook(proxy_server) ) if custom_key_generate_hook is not None: if inspect.iscoroutinefunction(custom_key_generate_hook): @@ -1855,11 +1975,11 @@ async def generate_service_account_key_fn( - user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id. """ + from litellm.proxy import proxy_server from litellm.proxy._types import CommonProxyErrors from litellm.proxy.proxy_server import ( prisma_client, user_api_key_cache, - user_custom_key_generate, ) if prisma_client is None: @@ -1887,7 +2007,9 @@ async def generate_service_account_key_fn( verbose_proxy_logger.debug("entered /key/generate") - custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_generate + custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_generate_hook( + proxy_server + ) if custom_key_generate_hook is not None: if inspect.iscoroutinefunction(custom_key_generate_hook): result: Final = await custom_key_generate_hook(data) @@ -1959,7 +2081,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ ) casted_metadata[reserved_field] = existing_value - data_json: Final[Mapping[str, object]] = data.model_dump(exclude_unset=True, exclude_none=True) + data_json: Final = _as_object_dict(data.model_dump(exclude_unset=True, exclude_none=True)) try: for k, v in data_json.items(): @@ -2737,13 +2859,13 @@ async def update_key_fn( }' ``` """ + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( llm_router, premium_user, prisma_client, proxy_logging_obj, user_api_key_cache, - user_custom_key_update, ) try: @@ -2774,7 +2896,9 @@ async def update_key_fn( ) # Custom key update hook - custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_update + custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_update_hook( + proxy_server + ) if custom_key_update_hook is not None: if inspect.iscoroutinefunction(custom_key_update_hook): result: Final = await custom_key_update_hook(data) @@ -2927,14 +3051,16 @@ async def bulk_update_keys( }' ``` """ + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( llm_router, prisma_client, proxy_logging_obj, user_api_key_cache, - user_custom_key_update, ) + custom_key_update_hook: Final = _custom_key_update_hook(proxy_server) + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=403, @@ -2980,7 +3106,7 @@ async def bulk_update_keys( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, - user_custom_key_update=user_custom_key_update, + user_custom_key_update=custom_key_update_hook, ) successful_updates.append( @@ -3058,7 +3184,7 @@ def _build_failed_team_key_update( if hasattr(existing_key_row, "model_dump"): key_info = existing_key_row.model_dump() elif hasattr(existing_key_row, "dict"): - key_info = existing_key_row.dict() + key_info = dict[str, object](_legacy_model_dict(existing_key_row)) if key_info: key_info.pop("token", None) @@ -3089,14 +3215,16 @@ async def bulk_update_team_keys( Callable by proxy admins, or by team admins with `KEY_UPDATE` permission. """ + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( llm_router, prisma_client, proxy_logging_obj, user_api_key_cache, - user_custom_key_update, ) + custom_key_update_hook: Final = _custom_key_update_hook(proxy_server) + if prisma_client is None: raise HTTPException( status_code=500, @@ -3223,7 +3351,7 @@ async def bulk_update_team_keys( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, - user_custom_key_update=user_custom_key_update, + user_custom_key_update=custom_key_update_hook, existing_key_row=existing_by_token[db_token], ) @@ -3435,7 +3563,7 @@ async def _get_model_max_budget_current_spend( virtual_key_model_spend_cache_key = ( f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{api_key_hash}:{model}:{budget_config.budget_duration}" ) - current_spend: float | None = await user_api_key_cache.async_get_cache( + current_spend: float | None = await _spend_cache(user_api_key_cache).async_get_cache( key=virtual_key_model_spend_cache_key, ) if current_spend is None: @@ -3444,7 +3572,7 @@ async def _get_model_max_budget_current_spend( f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:" f"{api_key_hash}:{model_without_prefix}:{budget_config.budget_duration}" ) - current_spend = await user_api_key_cache.async_get_cache( + current_spend = await _spend_cache(user_api_key_cache).async_get_cache( key=virtual_key_model_spend_cache_key, ) try: @@ -3635,7 +3763,7 @@ async def info_key_fn( hashed_key: str | None = key if key is not None: hashed_key = _hash_token_if_needed(token=key) - key_info = await VerificationTokenRepository(prisma_client).table.find_unique( + key_info = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique( where={"token": hashed_key}, include={"litellm_budget_table": True}, ) @@ -3945,7 +4073,7 @@ async def generate_key_helper_fn( if table_name is None or table_name == "user": # do not auto-create users for `/key/generate` ## CREATE USER (If necessary) if query_type == "insert_data": - user_row = await prisma_client.insert_data(data=user_data, table_name="user") + user_row = _created_user_row(await prisma_client.insert_data(data=user_data, table_name="user")) if user_row is None: raise Exception("Failed to create user") @@ -4232,7 +4360,7 @@ def _transform_verification_tokens_to_deleted_records( "litellm_changed_by": litellm_changed_by, } ) - record = deleted_record.model_dump() + record = dict[str, object](_as_object_dict(deleted_record.model_dump())) # Map org_id to organization_id (model uses org_id, but schema expects organization_id) org_id_value: object = record.pop("org_id", None) @@ -4355,13 +4483,12 @@ async def _rotate_master_key( should_create_model_in_db=False, ) if new_model: - _dumped = new_model.model_dump(exclude_none=True) + _dumped = dict[str, object](_as_object_dict(new_model.model_dump(exclude_none=True))) _dumped["litellm_params"] = prisma.Json(_dumped["litellm_params"]) _dumped["model_info"] = prisma.Json(_dumped["model_info"]) new_models.append(_dumped) verbose_proxy_logger.debug("Resetting proxy model table") - async with prisma_client.db.tx() as tx_ctx: - tx: Final[_TxTables] = tx_ctx + async with _tx_tables_context(prisma_client.db.tx) as tx: await tx.litellm_proxymodeltable.delete_many() verbose_proxy_logger.debug("Creating %s models", len(new_models)) await tx.litellm_proxymodeltable.create_many( @@ -4376,14 +4503,14 @@ async def _rotate_master_key( if config: """If environment_variables is found, decrypt it and encrypt it with the new master key""" - environment_variables_dict = {} + environment_variables_dict: Mapping[str, str] | None = {} for c in config: if c.param_name == "environment_variables": - environment_variables_dict = c.param_value + environment_variables_dict = _env_vars_param_value(c) if environment_variables_dict: decrypted_env_vars: Final = proxy_config._decrypt_and_set_db_env_variables( - environment_variables=environment_variables_dict + environment_variables=dict[str, str](environment_variables_dict) ) encrypted_env_vars: Final = proxy_config._encrypt_env_variables( environment_variables=decrypted_env_vars, @@ -4449,7 +4576,7 @@ async def _rotate_master_key( updated_patch=decrypted_cred, new_encryption_key=new_master_key, ) - _cred_data = encrypted_cred.model_dump(exclude_none=True) + _cred_data = dict[str, object](_as_object_dict(encrypted_cred.model_dump(exclude_none=True))) if "credential_values" in _cred_data: _cred_data["credential_values"] = prisma.Json(_cred_data["credential_values"]) if "credential_info" in _cred_data: @@ -4605,7 +4732,7 @@ async def _insert_deprecated_key( try: revoke_at: Final = datetime.now(timezone.utc) + timedelta(seconds=grace_seconds) - await DeprecatedVerificationTokenRepository(prisma_client).table.upsert( + await _deprecated_verification_token_table(prisma_client).upsert( where={"token": old_token_hash}, data={ "create": { @@ -4704,11 +4831,13 @@ async def _execute_virtual_key_regeneration( grace_period=data.grace_period if data else None, ) - updated_token: Final[Mapping[str, object] | None] = await VerificationTokenRepository(prisma_client).table.update( + updated_token: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update( where={"token": hashed_api_key}, data=with_settings_updated_at(jsonified_update_data), ) - updated_token_dict: Final[dict[str, object]] = dict(updated_token) if updated_token is not None else {} + updated_token_dict: Final = ( + dict[str, object](_as_object_dict(dict(updated_token))) if updated_token is not None else dict[str, object]() + ) updated_token_dict["key"] = new_token updated_token_dict["token_id"] = updated_token_dict.pop("token") @@ -5247,7 +5376,7 @@ async def validate_key_list_check( if key_hash: try: - key_info: Final = await VerificationTokenRepository(prisma_client).table.find_unique( + key_info: Final = await _prisma_table_lenient(VerificationTokenRepository(prisma_client)).find_unique( where={"token": key_hash}, ) except Exception: @@ -5278,7 +5407,7 @@ async def _fetch_user_team_objects( if complete_user_info is None or not complete_user_info.teams: return [] - teams: Final[list[BaseModel] | None] = await TeamRepository(prisma_client).table.find_many( + teams: Final = await _prisma_table_lenient(TeamRepository(prisma_client)).find_many( where={"team_id": {"in": complete_user_info.teams}} ) if teams is None: @@ -5653,7 +5782,7 @@ async def key_aliases( where_sql: Final = " AND ".join(where_parts) count_sql: Final = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' - count_rows: Final[Sequence[Mapping[str, int]]] = await prisma_client.db.query_raw(count_sql, *query_params) + count_rows: Final = await _query_raw_text_rows(prisma_client, count_sql, *query_params) total_count: Final = int(count_rows[0]["count"]) if count_rows else 0 aliases_params: Final = query_params + [size, (page - 1) * size] @@ -5666,7 +5795,7 @@ async def key_aliases( f" ORDER BY key_alias ASC" f" LIMIT ${limit_idx} OFFSET ${offset_idx}" ) - alias_rows: Final[Sequence[Mapping[str, str]]] = await prisma_client.db.query_raw(aliases_sql, *aliases_params) + alias_rows: Final = await _query_raw_text_rows(prisma_client, aliases_sql, *aliases_params) aliases: Final[list[str]] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] total_pages: Final = -(-total_count // size) if total_count > 0 else 0 @@ -5952,7 +6081,7 @@ async def _list_key_helper( # Fetch keys with pagination if use_deleted_table: - keys = await DeletedVerificationTokenRepository(prisma_client).table.find_many( + keys = await _deleted_verification_token_table(prisma_client).find_many( where=where, skip=skip, take=size, @@ -5966,7 +6095,7 @@ async def _list_key_helper( ), ) else: - keys = await VerificationTokenRepository(prisma_client).table.find_many( + keys = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many( where=where, skip=skip, take=size, @@ -5995,13 +6124,13 @@ async def _list_key_helper( total_pages: Final = -(-total_count // size) # Ceiling division # Fetch user information if expand includes "user" - user_map = {} + user_map: Mapping[str, LiteLLM_UserTable] = {} if expand and "user" in expand: user_ids: Final = [key.user_id for key in keys if key.user_id] created_by_ids: Final = [key.created_by for key in keys if key.created_by] all_ids: Final = list(set(user_ids + created_by_ids)) # Remove duplicates if all_ids: - users: Final[Sequence[_UserRowLike]] = await UserRepository(prisma_client).table.find_many( + users: Final = await _prisma_table(UserRepository(prisma_client)).find_many( where={"user_id": {"in": all_ids}} ) user_map = {user.user_id: user for user in users} @@ -6014,10 +6143,14 @@ async def _list_key_helper( key_dict = key.model_dump() except Exception: # Fallback for Pydantic v1 compatibility - key_dict = key.dict() + key_dict = dict[str, object](_legacy_model_dict(key)) # Attach object_permission if object_permission_id is set (only for non-deleted keys) if not use_deleted_table: - key_dict = await attach_object_permission_to_dict(key_dict, prisma_client) + key_dict = dict[str, object]( + await _object_permission_utils(object_permission_utils).attach_object_permission_to_dict( + key_dict, prisma_client + ) + ) # Include user information if expand includes "user" if expand and "user" in expand: @@ -6025,7 +6158,7 @@ async def _list_key_helper( try: key_dict["user"] = user_map[key.user_id].model_dump() except Exception: - key_dict["user"] = user_map[key.user_id].dict() + key_dict["user"] = _legacy_model_dict(user_map[key.user_id]) if key.created_by and key.created_by in user_map: created_by_user = user_map[key.created_by] key_dict["created_by_user"] = { @@ -6039,7 +6172,7 @@ async def _list_key_helper( # Use deleted key type to preserve deleted_at, deleted_by, etc. key_list.append(LiteLLM_DeletedVerificationToken.model_validate(key_dict)) else: - key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object + key_list.append(UserAPIKeyAuth.model_validate(key_dict)) # Return full key object else: _token = key_dict.get("token") key_list.append(cast(str, _token)) # Return only the token diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index 0fdedafb2bf..dfb1422308f 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -10,8 +10,9 @@ POST /v1/tool/policy - Update the input_policy / output_policy for a """ import uuid +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Annotated, Any, Final +from typing import TYPE_CHECKING, Annotated, Final, Protocol from fastapi import APIRouter, Depends, HTTPException, Query from pydantic import BaseModel, Field, TypeAdapter @@ -49,6 +50,144 @@ from litellm.types.tool_management import ( ToolUsageLogsResponse, ) + +class _DailyToolSpendRecord(Protocol): + date: str + tool_name: str + spend: float + request_count: int + + +class _SpendLogToolIndexRecord(Protocol): + request_id: str + + +class _SpendLogRecord(Protocol): + request_id: str + startTime: datetime + model: str | None + spend: float | None + total_tokens: int | None + messages: object + proxy_server_request: object + + +class _VerificationTokenRecord(Protocol): + object_permission_id: str | None + + +class _TeamRecord(Protocol): + object_permission_id: str | None + + +class _DailyToolSpendTable(Protocol): + async def group_by( + self, + *, + by: Sequence[str], + sum: Mapping[str, bool], + where: Mapping[str, object], + order: Mapping[str, object], + take: int, + ) -> Sequence[object] | None: ... + + async def find_many( + self, + *, + where: Mapping[str, object], + order: Sequence[Mapping[str, str]], + ) -> Sequence[_DailyToolSpendRecord]: ... + + +class _SpendLogToolIndexTable(Protocol): + async def count(self, *, where: Mapping[str, object]) -> int: ... + + async def find_many( + self, + *, + where: Mapping[str, object], + order: Mapping[str, str], + skip: int, + take: int, + ) -> Sequence[_SpendLogToolIndexRecord]: ... + + +class _SpendLogsTable(Protocol): + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_SpendLogRecord]: ... + + +class _VerificationTokenTable(Protocol): + async def find_unique(self, *, where: Mapping[str, object]) -> _VerificationTokenRecord | None: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + +class _TeamTable(Protocol): + async def find_unique(self, *, where: Mapping[str, object]) -> _TeamRecord | None: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + +class _ObjectPermissionTable(Protocol): + async def create(self, *, data: Mapping[str, str | Sequence[str]]) -> object: ... + + async def delete(self, *, where: Mapping[str, object]) -> object: ... + + +class _DailyToolSpendTableHolder(Protocol): + @property + def table(self) -> _DailyToolSpendTable: ... + + +class _SpendLogToolIndexTableHolder(Protocol): + @property + def table(self) -> _SpendLogToolIndexTable: ... + + +class _SpendLogsTableHolder(Protocol): + @property + def table(self) -> _SpendLogsTable: ... + + +class _VerificationTokenTableHolder(Protocol): + @property + def table(self) -> _VerificationTokenTable: ... + + +class _TeamTableHolder(Protocol): + @property + def table(self) -> _TeamTable: ... + + +class _ObjectPermissionTableHolder(Protocol): + @property + def table(self) -> _ObjectPermissionTable: ... + + +def _daily_tool_spend_table(repo: _DailyToolSpendTableHolder) -> _DailyToolSpendTable: + return repo.table + + +def _spend_log_tool_index_table(repo: _SpendLogToolIndexTableHolder) -> _SpendLogToolIndexTable: + return repo.table + + +def _spend_logs_table(repo: _SpendLogsTableHolder) -> _SpendLogsTable: + return repo.table + + +def _verification_token_table(repo: _VerificationTokenTableHolder) -> _VerificationTokenTable: + return repo.table + + +def _team_table(repo: _TeamTableHolder) -> _TeamTable: + return repo.table + + +def _object_permission_table(repo: _ObjectPermissionTableHolder) -> _ObjectPermissionTable: + return repo.table + + router: Final = APIRouter() TOOL_POLICY_OPTIONS: Final = ToolPolicyOptionsResponse( @@ -154,6 +293,7 @@ class _TopToolRow(BaseModel): _TOP_TOOL_ROWS: Final = TypeAdapter(list[_TopToolRow]) +_PARSED_JSON: Final = TypeAdapter(object) @router.get( @@ -201,7 +341,7 @@ async def get_tool_spend( end_str: Final = end_day.strftime("%Y-%m-%d") date_window: Final = {"date": {"gte": start_str, "lte": end_str}} - table: Final = DailyToolSpendRepository(prisma_client).table + table: Final = _daily_tool_spend_table(DailyToolSpendRepository(prisma_client)) top_tools: Final = _TOP_TOOL_ROWS.validate_python( await table.group_by( by=["tool_name"], @@ -222,7 +362,7 @@ async def get_tool_spend( for row in top_tools ] - daily_rows: Final = ( + daily_rows: Final[Sequence[_DailyToolSpendRecord]] = ( await table.find_many( where={**date_window, "tool_name": {"in": [row.tool_name for row in top_tools]}}, order=[{"date": "asc"}, {"spend": "desc"}], @@ -270,23 +410,23 @@ async def get_tool_detail( raise HTTPException(status_code=500, detail=str(e)) -def _input_snippet_for_tool_log(sl: Any, max_len: int = 200) -> str | None: +def _input_snippet_for_tool_log(sl: _SpendLogRecord | None, max_len: int = 200) -> str | None: """Short snippet from messages or proxy_server_request for tool usage log row.""" if sl is None: return None - messages: Final = getattr(sl, "messages", None) + messages: Final[object] = getattr(sl, "messages", None) if messages is not None: s = _snippet_str(messages, max_len) if s: return s - psr = getattr(sl, "proxy_server_request", None) + psr: object = getattr(sl, "proxy_server_request", None) if not psr: return None if isinstance(psr, str): import json try: - psr = json.loads(psr) + psr = _PARSED_JSON.validate_python(json.loads(psr)) except Exception: return _snippet_str(psr, max_len) if isinstance(psr, dict): @@ -299,7 +439,7 @@ def _input_snippet_for_tool_log(sl: Any, max_len: int = 200) -> str | None: return _snippet_str(psr, max_len) -def _snippet_str(text: Any, max_len: int = 200) -> str | None: +def _snippet_str(text: object, max_len: int = 200) -> str | None: if text is None: return None if isinstance(text, str): @@ -344,10 +484,9 @@ async def get_tool_usage_logs( raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: - where: Final[dict] = {"tool_name": tool_name} + start_time_filter: datetime | None = None + end_time_filter: datetime | None = None if start_date or end_date: - start_time_filter: datetime | None = None - end_time_filter: datetime | None = None if start_date: try: start_time_filter = datetime.strptime(start_date + "T00:00:00", "%Y-%m-%dT%H:%M:%S").replace( @@ -362,15 +501,15 @@ async def get_tool_usage_logs( ) except ValueError: pass - if start_time_filter is not None or end_time_filter is not None: - where["start_time"] = {} - if start_time_filter is not None: - where["start_time"]["gte"] = start_time_filter - if end_time_filter is not None: - where["start_time"]["lte"] = end_time_filter + start_time_range: Final[Mapping[str, datetime]] = { + key: value for key, value in (("gte", start_time_filter), ("lte", end_time_filter)) if value is not None + } + where: Final[Mapping[str, str | Mapping[str, datetime]]] = ( + {"tool_name": tool_name, "start_time": start_time_range} if start_time_range else {"tool_name": tool_name} + ) - total: Final = await SpendLogToolIndexRepository(prisma_client).table.count(where=where) - index_rows: Final = await SpendLogToolIndexRepository(prisma_client).table.find_many( + total: Final = await _spend_log_tool_index_table(SpendLogToolIndexRepository(prisma_client)).count(where=where) + index_rows: Final = await _spend_log_tool_index_table(SpendLogToolIndexRepository(prisma_client)).find_many( where=where, order={"start_time": "desc"}, skip=(page - 1) * page_size, @@ -380,7 +519,9 @@ async def get_tool_usage_logs( if not request_ids: return ToolUsageLogsResponse(logs=[], total=total, page=page, page_size=page_size) - spend_logs = await SpendLogsRepository(prisma_client).table.find_many(where={"request_id": {"in": request_ids}}) + spend_logs = await _spend_logs_table(SpendLogsRepository(prisma_client)).find_many( + where={"request_id": {"in": request_ids}} + ) log_by_id: Final = {s.request_id: s for s in spend_logs} logs_out: Final[list[ToolUsageLogEntry]] = [] @@ -449,23 +590,29 @@ async def _resolve_key_hash_to_object_permission_id( hashed: Final = key_hash if "sk-" not in (key_hash or "") else hash_token(key_hash) if not hashed: return None - row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed}) + row = await _verification_token_table(VerificationTokenRepository(prisma_client)).find_unique( + where={"token": hashed} + ) if row is None: return None - op_id: Final = getattr(row, "object_permission_id", None) + op_id: Final[str | None] = getattr(row, "object_permission_id", None) if op_id: return op_id new_id: Final = str(uuid.uuid4()) - await ObjectPermissionRepository(prisma_client).table.create( + await _object_permission_table(ObjectPermissionRepository(prisma_client)).create( data={"object_permission_id": new_id, "blocked_tools": []} ) - updated_count: Final = await VerificationTokenRepository(prisma_client).table.update_many( + updated_count: Final = await _verification_token_table(VerificationTokenRepository(prisma_client)).update_many( where={"token": hashed, "object_permission_id": None}, data={"object_permission_id": new_id}, ) if updated_count == 0: - await ObjectPermissionRepository(prisma_client).table.delete(where={"object_permission_id": new_id}) - row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed}) + await _object_permission_table(ObjectPermissionRepository(prisma_client)).delete( + where={"object_permission_id": new_id} + ) + row = await _verification_token_table(VerificationTokenRepository(prisma_client)).find_unique( + where={"token": hashed} + ) return getattr(row, "object_permission_id", None) if row else None return new_id @@ -478,23 +625,25 @@ async def _resolve_team_id_to_object_permission_id( if not team_id or not team_id.strip(): return None team_id_clean: Final = team_id.strip() - row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id_clean}) + row = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id_clean}) if row is None: return None - op_id: Final = getattr(row, "object_permission_id", None) + op_id: Final[str | None] = getattr(row, "object_permission_id", None) if op_id: return op_id new_id: Final = str(uuid.uuid4()) - await ObjectPermissionRepository(prisma_client).table.create( + await _object_permission_table(ObjectPermissionRepository(prisma_client)).create( data={"object_permission_id": new_id, "blocked_tools": []} ) - updated_count: Final = await TeamRepository(prisma_client).table.update_many( + updated_count: Final = await _team_table(TeamRepository(prisma_client)).update_many( where={"team_id": team_id_clean, "object_permission_id": None}, data={"object_permission_id": new_id}, ) if updated_count == 0: - await ObjectPermissionRepository(prisma_client).table.delete(where={"object_permission_id": new_id}) - row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id_clean}) + await _object_permission_table(ObjectPermissionRepository(prisma_client)).delete( + where={"object_permission_id": new_id} + ) + row = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id_clean}) return getattr(row, "object_permission_id", None) if row else None return new_id diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index d8e9f8dfaee..d1dc1482401 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -3,8 +3,9 @@ CRUD ENDPOINTS FOR PROMPTS """ import tempfile +from collections.abc import Mapping, Sequence from pathlib import Path -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Final, Protocol, cast from fastapi import ( APIRouter, @@ -15,7 +16,7 @@ from fastapi import ( Response, UploadFile, ) -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -38,8 +39,47 @@ from litellm.types.prompts.init_prompts import ( ) from litellm.types.proxy.prompt_endpoints import TestPromptRequest +if TYPE_CHECKING: + from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry + from litellm.proxy.utils import PrismaClient + + +class _PromptRecord(Protocol): + id: str + version: int + environment: str | None + + +class _PromptTable(Protocol): + async def find_many( + self, + *, + where: Mapping[str, object], + order: Mapping[str, str] | None = None, + take: int | None = None, + distinct: Sequence[str] | None = None, + ) -> Sequence[_PromptRecord]: ... + + async def create(self, *, data: Mapping[str, str | int | None]) -> _PromptRecord: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, str]) -> _PromptRecord: ... + + async def delete_many(self, *, where: Mapping[str, str]) -> object: ... + + +class _PromptTableHolder(Protocol): + @property + def table(self) -> _PromptTable: ... + + +def _prompt_table(repo: _PromptTableHolder) -> _PromptTable: + return repo.table + + router: Final = APIRouter() +_PARSED_VALUE: Final = TypeAdapter(object) + def get_base_prompt_id(prompt_id: str) -> str: """ @@ -132,7 +172,7 @@ def construct_versioned_prompt_id(prompt_id: str, version: int | None = None) -> return f"{base_id}.v{version}" -def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: dict[str, Any]) -> str: +def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: Mapping[str, object]) -> str: """ Find the latest version of a prompt from available prompt IDs. @@ -198,7 +238,9 @@ def get_latest_prompt_versions(prompts: list[PromptSpec]) -> list[PromptSpec]: return list(latest_prompts.values()) -async def get_next_version_for_prompt(prisma_client, prompt_id: str, environment: str = "development") -> int: +async def get_next_version_for_prompt( + prisma_client: "PrismaClient", prompt_id: str, environment: str = "development" +) -> int: """ Get the next version number for a prompt in a specific environment. @@ -210,7 +252,7 @@ async def get_next_version_for_prompt(prisma_client, prompt_id: str, environment Returns: Next version number (1 if no versions exist, max_version + 1 otherwise) """ - existing_prompts: Final = await PromptRepository(prisma_client).table.find_many( + existing_prompts: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where={"prompt_id": prompt_id, "environment": environment} ) @@ -431,10 +473,10 @@ async def get_prompt_versions( # Query DB for versions versioned_prompts: Final = [] if prisma_client is not None: - where_clause: Final[dict[str, Any]] = {"prompt_id": base_prompt_id} - if environment: - where_clause["environment"] = environment - db_prompts: Final = await PromptRepository(prisma_client).table.find_many( + where_clause: Final[Mapping[str, str]] = ( + {"prompt_id": base_prompt_id, "environment": environment} if environment else {"prompt_id": base_prompt_id} + ) + db_prompts: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where=where_clause, order={"version": "desc"}, ) @@ -590,7 +632,7 @@ async def get_prompt_info( # Query all environments this prompt exists in (lightweight: distinct on environment) all_environments: list[str] = [] if prisma_client is not None: - all_prompt_rows: Final = await PromptRepository(prisma_client).table.find_many( + all_prompt_rows: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where={"prompt_id": base_prompt_id}, distinct=["environment"], ) @@ -602,13 +644,16 @@ async def get_prompt_info( prompt_spec = None requested_version: Final = get_version_number(prompt_id=prompt_id) if prompt_id != base_prompt_id else None if environment and prisma_client is not None: - where_clause: Final[dict[str, Any]] = { - "prompt_id": base_prompt_id, - "environment": environment, + where_clause: Final[Mapping[str, str | int]] = { + key: value + for key, value in ( + ("prompt_id", base_prompt_id), + ("environment", environment), + ("version", requested_version), + ) + if value is not None } - if requested_version is not None: - where_clause["version"] = requested_version - env_prompts: Final = await PromptRepository(prisma_client).table.find_many( + env_prompts: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where=where_clause, order={"version": "desc"}, take=1, @@ -721,7 +766,7 @@ async def create_prompt( ) # Store prompt in db with version - prompt_db_entry: Final = await PromptRepository(prisma_client).table.create( + prompt_db_entry: Final = await _prompt_table(PromptRepository(prisma_client)).create( data={ "prompt_id": request.prompt_id, "version": new_version, @@ -811,7 +856,9 @@ async def update_prompt( ) # Check if any version of this prompt exists (in any environment) - existing_prompts = await PromptRepository(prisma_client).table.find_many(where={"prompt_id": base_prompt_id}) + existing_prompts = await _prompt_table(PromptRepository(prisma_client)).find_many( + where={"prompt_id": base_prompt_id} + ) if not existing_prompts: raise HTTPException( @@ -835,7 +882,7 @@ async def update_prompt( ) # Store new version in db - prompt_db_entry: Final = await PromptRepository(prisma_client).table.create( + prompt_db_entry: Final = await _prompt_table(PromptRepository(prisma_client)).create( data={ "prompt_id": base_prompt_id, "version": new_version, @@ -936,12 +983,12 @@ async def delete_prompt( base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id) # Build delete filter; scope to environment if provided - delete_where: Final[dict[str, Any]] = {"prompt_id": base_prompt_id} - if environment: - delete_where["environment"] = environment + delete_where: Final[Mapping[str, str]] = ( + {"prompt_id": base_prompt_id, "environment": environment} if environment else {"prompt_id": base_prompt_id} + ) # Delete versions from the database (scoped to environment if provided) - await PromptRepository(prisma_client).table.delete_many(where=delete_where) + await _prompt_table(PromptRepository(prisma_client)).delete_many(where=delete_where) # Remove matching prompts from memory — scope to environment if provided if environment: @@ -967,7 +1014,9 @@ async def delete_prompt( raise HTTPException(status_code=500, detail=str(e)) -def _reload_prompt_in_registry(registry: Any, versioned_id: str, updated_prompt_spec: PromptSpec) -> PromptSpec: +def _reload_prompt_in_registry( + registry: "InMemoryPromptRegistry", versioned_id: str, updated_prompt_spec: PromptSpec +) -> PromptSpec: """Remove stale entry and re-initialize the prompt in the in-memory registry.""" if versioned_id in registry.IN_MEMORY_PROMPTS: del registry.IN_MEMORY_PROMPTS[versioned_id] @@ -1033,14 +1082,13 @@ async def patch_prompt( requested_version: Final = get_version_number(prompt_id=prompt_id) if prompt_id != base_prompt_id else None # Build query to find the exact row by composite unique key - find_where: Final[dict[str, Any]] = { - "prompt_id": base_prompt_id, - "environment": env, + find_where: Final[Mapping[str, str | int]] = { + key: value + for key, value in (("prompt_id", base_prompt_id), ("environment", env), ("version", requested_version)) + if value is not None } - if requested_version is not None: - find_where["version"] = requested_version - db_rows: Final = await PromptRepository(prisma_client).table.find_many( + db_rows: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where=find_where, order={"version": "desc"}, take=1, @@ -1084,15 +1132,18 @@ async def patch_prompt( raise HTTPException(status_code=400, detail="litellm_params cannot be None") # Build update data dict - update_data: Final[dict[str, Any]] = { - "litellm_params": updated_litellm_params.model_dump_json(), - "prompt_info": updated_prompt_info.model_dump_json(), + update_data: Final[Mapping[str, str]] = { + key: value + for key, value in ( + ("litellm_params", updated_litellm_params.model_dump_json()), + ("prompt_info", updated_prompt_info.model_dump_json()), + ("created_by", user_api_key_dict.user_id), + ) + if value } - if user_api_key_dict.user_id: - update_data["created_by"] = user_api_key_dict.user_id # Update by primary key (id) to target exactly one row - updated_prompt_db_entry: Final = await PromptRepository(prisma_client).table.update( + updated_prompt_db_entry: Final = await _prompt_table(PromptRepository(prisma_client)).update( where={"id": target_row.id}, data=update_data, ) @@ -1216,23 +1267,25 @@ async def test_prompt( # Use ProxyBaseLLMRequestProcessing to go through all proxy logic base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) - result: Final = await base_llm_response_processor.base_process_llm_request( - request=fastapi_request, - fastapi_response=fastapi_response, - user_api_key_dict=user_api_key_dict, - route_type="acompletion", - proxy_logging_obj=proxy_logging_obj, - llm_router=llm_router, - general_settings=general_settings, - proxy_config=proxy_config, - select_data_generator=select_data_generator, - model=None, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - version=version, + result: Final = _PARSED_VALUE.validate_python( + await base_llm_response_processor.base_process_llm_request( + request=fastapi_request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="acompletion", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) ) if isinstance(result, BaseModel): @@ -1257,7 +1310,7 @@ async def test_prompt( async def convert_prompt_file_to_json( file: UploadFile = File(...), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -) -> dict[str, Any]: +) -> Mapping[str, object]: """ Convert a .prompt file to JSON format. diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 383ada5a1bc..a8ec8884d9b 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -10,9 +10,12 @@ https://platform.openai.com/docs/api-reference/responses-streaming import asyncio import json -from typing import Any, Final, cast +from collections.abc import Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Final, TypeAlias from fastapi import Request, Response +from fastapi.responses import StreamingResponse +from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth @@ -20,6 +23,56 @@ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessin from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler from litellm.types.llms.openai import ResponsesAPIStatus +if TYPE_CHECKING: + from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + + +_JsonDict: TypeAlias = dict[str, object] +_JsonList: TypeAlias = list[object] + + +class _OutputItem(TypedDict, total=False): + id: str + content: Sequence[object] + + +class _TerminalResponse(TypedDict, total=False): + status: ResponsesAPIStatus + error: _JsonDict + usage: _JsonDict + reasoning: _JsonDict + tool_choice: object + tools: _JsonList + model: str + instructions: str + temperature: float + top_p: float + max_output_tokens: int + previous_response_id: str + text: _JsonDict + truncation: str + parallel_tool_calls: bool + user: str + store: bool + incomplete_details: _JsonDict + output: Sequence[_OutputItem] + + +class _StreamEvent(TypedDict, total=False): + type: str + item: _OutputItem + item_id: str + content_index: int + delta: str + part: object + response: _TerminalResponse + + +class _StreamEventParser: + parse: Callable[[str], _StreamEvent] = staticmethod(json.loads) + async def background_streaming_task( polling_id: str, @@ -29,16 +82,16 @@ async def background_streaming_task( fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, general_settings: dict, - llm_router, - proxy_config, - proxy_logging_obj, + llm_router: "Router | None", + proxy_config: "ProxyConfig", + proxy_logging_obj: "ProxyLogging", select_data_generator, user_model, - user_temperature, - user_request_timeout, - user_max_tokens, - user_api_base, - version, + user_temperature: float | None, + user_request_timeout: float | None, + user_max_tokens: int | None, + user_api_base: str | None, + version: str | None, ): """ Background task to stream response and update cache @@ -69,7 +122,7 @@ async def background_streaming_task( # Make streaming request. # Pre-call checks (rate limits, guardrails, budget) were already run # before polling ID creation, so skip them here to avoid double-counting. - response: Final = await processor.base_process_llm_request( + response: Final[StreamingResponse] = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -91,8 +144,10 @@ async def background_streaming_task( # Process streaming response following OpenAI events format # https://platform.openai.com/docs/api-reference/responses-streaming - output_items: Final[dict[str, dict[str, Any]]] = {} # Track output items by ID - accumulated_text: Final = {} # Track accumulated text deltas by (item_id, content_index) + output_items: Final = dict[str, _OutputItem]() # Track output items by ID + accumulated_text: Final = dict[ + tuple[str, int], str + ]() # Track accumulated text deltas by (item_id, content_index) # ResponsesAPIResponse fields to extract from response.completed usage_data = None @@ -121,7 +176,7 @@ async def background_streaming_task( None # Will be set by response.completed/failed/incomplete/cancelled ) terminal_error = None - _event_to_status: Final = { + _event_to_status: Final[Mapping[str, ResponsesAPIStatus]] = { "response.completed": "completed", "response.failed": "failed", "response.incomplete": "incomplete", @@ -162,7 +217,7 @@ async def background_streaming_task( break try: - event = json.loads(chunk_data) + event: _StreamEvent = _StreamEventParser.parse(chunk_data) event_type = event.get("type", "") # Process different event types based on OpenAI streaming spec @@ -181,9 +236,8 @@ async def background_streaming_task( if item_id and item_id in output_items: # Update the output item with new content - if "content" not in output_items[item_id]: - output_items[item_id]["content"] = [] - output_items[item_id]["content"].append(content_part) + added_item = output_items[item_id] + added_item["content"] = (*added_item.get("content", ()), content_part) state_dirty = True elif event_type == "response.output_text.delta": @@ -201,12 +255,14 @@ async def background_streaming_task( accumulated_text[key] += delta # Update the content in output_items - if "content" in output_items[item_id]: - content_list = output_items[item_id]["content"] + delta_item = output_items[item_id] + if "content" in delta_item: + content_list = delta_item["content"] if content_index < len(content_list): # Update existing content part with accumulated text - if isinstance(content_list[content_index], dict): - content_list[content_index]["text"] = accumulated_text[key] + content_entry = content_list[content_index] + if isinstance(content_entry, dict): + content_entry["text"] = accumulated_text[key] state_dirty = True elif event_type == "response.content_part.done": @@ -217,10 +273,14 @@ async def background_streaming_task( if item_id and item_id in output_items: # Update with final content from event - if "content" in output_items[item_id]: - content_list = output_items[item_id]["content"] + done_item = output_items[item_id] + if "content" in done_item: + content_list = done_item["content"] if content_index < len(content_list): - content_list[content_index] = content_part + done_item["content"] = tuple( + content_part if part_index == content_index else existing_part + for part_index, existing_part in enumerate(content_list) + ) state_dirty = True elif event_type == "response.output_item.done": @@ -248,12 +308,9 @@ async def background_streaming_task( # Terminal event - extract all ResponsesAPIResponse fields # https://platform.openai.com/docs/api-reference/responses-streaming response_data = event.get("response", {}) - terminal_status = cast( - ResponsesAPIStatus, - response_data.get( - "status", - _event_to_status.get(event_type, "completed"), - ), + terminal_status = response_data.get( + "status", + _event_to_status.get(event_type, "completed"), ) # Extract error for failed and incomplete responses diff --git a/litellm/responses/file_search/emulated_handler.py b/litellm/responses/file_search/emulated_handler.py index 7854b17a06f..205189a6043 100644 --- a/litellm/responses/file_search/emulated_handler.py +++ b/litellm/responses/file_search/emulated_handler.py @@ -14,16 +14,19 @@ Flow: import json import time import uuid -from collections.abc import Iterable -from typing import Any, Final, cast +from collections.abc import Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, TypedDict, cast from litellm._internal_context import is_internal_call from litellm._logging import verbose_logger -from litellm.types.llms.openai import ResponseOutputItem, ResponsesAPIResponse +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.vector_stores import VectorStoreSearchResult +if TYPE_CHECKING: + from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig + # Keep ToolParam broad so we stay compatible with both dict and Pydantic forms -ToolParam = Any +ToolParam = object FILE_SEARCH_FUNCTION_NAME: Final = "litellm_file_search" @@ -35,7 +38,7 @@ FILE_SEARCH_FUNCTION_NAME: Final = "litellm_file_search" def should_use_emulated_file_search( tools: Iterable[ToolParam] | None, - provider_config: Any, # BaseResponsesAPIConfig + provider_config: "BaseResponsesAPIConfig | None", ) -> bool: """Return True when there is a file_search tool and the provider can't handle it natively.""" if not tools: @@ -51,7 +54,7 @@ def should_use_emulated_file_search( # --------------------------------------------------------------------------- -def _build_function_tool(vector_store_ids: list[str]) -> dict[str, Any]: +def _build_function_tool(vector_store_ids: Sequence[str]) -> Mapping[str, object]: """ Create a Responses API function-tool definition that describes file search. The function accepts one or more natural-language queries (like OpenAI's native @@ -94,27 +97,26 @@ def _build_function_tool(vector_store_ids: list[str]) -> dict[str, Any]: } +def _file_search_tool_vector_store_ids(tool: object) -> Sequence[str] | None: + if not (isinstance(tool, dict) and tool.get("type") == "file_search"): + return None + return tool.get("vector_store_ids") or [] + + def _replace_file_search_tools( tools: Iterable[ToolParam] | None, -) -> tuple[list[dict[str, Any]], list[str]]: +) -> tuple[Sequence[object], Sequence[str]]: """ Replace all file_search tools with a single function tool. Returns: (new_tools_list, all_vector_store_ids) """ - non_file_search: Final[list[dict[str, Any]]] = [] - vector_store_ids: Final[list[str]] = [] - - for tool in tools or []: - if isinstance(tool, dict) and tool.get("type") == "file_search": - ids = tool.get("vector_store_ids") or [] - vector_store_ids.extend(ids) - else: - non_file_search.append(tool) + ids_and_tools: Final = tuple((_file_search_tool_vector_store_ids(tool), tool) for tool in tools or ()) # Deduplicate while preserving order - unique_ids: Final[list[str]] = list(dict.fromkeys(vector_store_ids)) + unique_ids: Final = list(dict.fromkeys(vs_id for ids, _ in ids_and_tools if ids is not None for vs_id in ids)) + non_file_search: Final = [tool for ids, tool in ids_and_tools if ids is None] if unique_ids: non_file_search.append(_build_function_tool(unique_ids)) @@ -127,9 +129,9 @@ def _replace_file_search_tools( async def _run_vector_searches( - queries: list[str], - vector_store_ids: list[str], -) -> tuple[list[str], list[VectorStoreSearchResult]]: + queries: Sequence[str], + vector_store_ids: Sequence[str], +) -> tuple[Sequence[str], Sequence[VectorStoreSearchResult]]: """ Run `asearch` against all vector stores for all queries and collect results. @@ -172,7 +174,7 @@ async def _run_vector_searches( # --------------------------------------------------------------------------- -def _get_field(result: Any, key: str, default: Any = None) -> Any: +def _get_field(result: object, key: str, default: object = None) -> object: """Read a field from either a dict/TypedDict or an attribute-based object.""" if isinstance(result, dict): return result.get(key, default) @@ -180,7 +182,7 @@ def _get_field(result: Any, key: str, default: Any = None) -> Any: def _format_search_results_as_tool_output( - results: list[VectorStoreSearchResult], + results: Sequence[VectorStoreSearchResult], ) -> str: """Serialize search results into a string to pass back as the tool's output.""" if not results: @@ -191,7 +193,8 @@ def _format_search_results_as_tool_output( score = _get_field(result, "score") file_id = _get_field(result, "file_id") filename = _get_field(result, "filename") - content_items = _get_field(result, "content") or [] + raw_content = _get_field(result, "content") + content_items = raw_content if isinstance(raw_content, list) else [] text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] text = " ".join(t for t in text_chunks if t) @@ -209,9 +212,24 @@ def _format_search_results_as_tool_output( return "\n\n".join(parts) +def _format_result_for_include(result: VectorStoreSearchResult) -> Mapping[str, object]: + file_id: Final = _get_field(result, "file_id") or "" + raw_content: Final = _get_field(result, "content") + content_items: Final = raw_content if isinstance(raw_content, list) else [] + text_chunks: Final = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] + text: Final = " ".join(t for t in text_chunks if t) + return { + "file_id": file_id, + "filename": _get_field(result, "filename") or "", + "score": _get_field(result, "score"), + "text": text, + "attributes": _get_field(result, "attributes") or {}, + } + + def _build_search_results_for_include( - results: list[VectorStoreSearchResult], -) -> list[dict[str, Any]]: + results: Sequence[VectorStoreSearchResult], +) -> Sequence[Mapping[str, object]]: """ Convert VectorStoreSearchResult objects to the format expected in file_search_call.search_results (mirrors OpenAI's include= format). @@ -220,30 +238,15 @@ def _build_search_results_for_include( behaviour of OpenAI's native file_search which surfaces every relevant chunk even when multiple chunks originate from the same document. """ - formatted: Final[list[dict[str, Any]]] = [] - for result in results: - file_id = _get_field(result, "file_id") or "" - content_items = _get_field(result, "content") or [] - text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] - text = " ".join(t for t in text_chunks if t) - formatted.append( - { - "file_id": file_id, - "filename": _get_field(result, "filename") or "", - "score": _get_field(result, "score"), - "text": text, - "attributes": _get_field(result, "attributes") or {}, - } - ) - return formatted + return [_format_result_for_include(result) for result in results] def _build_file_search_call_output( call_id: str, - queries: list[str], - results: list[VectorStoreSearchResult] | None = None, + queries: Sequence[str], + results: Sequence[VectorStoreSearchResult] | None = None, include_search_results: bool = False, -) -> dict[str, Any]: +) -> Mapping[str, object]: """Build the file_search_call output item (mirrors OpenAI's format). Args: @@ -266,39 +269,34 @@ def _build_file_search_call_output( def _build_file_citation_annotations( - results: list[VectorStoreSearchResult], + results: Sequence[VectorStoreSearchResult], text: str, -) -> list[dict[str, Any]]: +) -> Sequence[Mapping[str, object]]: """ Build file_citation annotations for the text. Each result with a file_id gets a citation at the end of the text. """ - annotations: Final[list[dict[str, Any]]] = [] index: Final = len(text) # cite at end of text block - seen_file_ids: Final[set] = set() + id_filename_pairs: Final = tuple( + (_get_field(result, "file_id"), _get_field(result, "filename")) for result in results + ) + first_filename_by_id: Final = {file_id: filename for file_id, filename in reversed(id_filename_pairs) if file_id} - for result in results: - file_id = _get_field(result, "file_id") - filename = _get_field(result, "filename") - if not file_id or file_id in seen_file_ids: - continue - seen_file_ids.add(file_id) - annotations.append( - { - "type": "file_citation", - "index": index, - "file_id": file_id, - "filename": filename or "", - } - ) - - return annotations + return [ + { + "type": "file_citation", + "index": index, + "file_id": file_id, + "filename": first_filename_by_id[file_id] or "", + } + for file_id in dict.fromkeys(file_id for file_id, _ in id_filename_pairs if file_id) + ] def _build_message_output( response_text: str, - results: list[VectorStoreSearchResult], -) -> dict[str, Any]: + results: Sequence[VectorStoreSearchResult], +) -> Mapping[str, object]: """Build the message output item with optional file_citation annotations.""" annotations: Final = _build_file_citation_annotations(results, response_text) return { @@ -330,8 +328,8 @@ def _extract_text_from_responses_output(response: ResponsesAPIResponse) -> str: def _synthesize_responses_api_response( original_response: ResponsesAPIResponse, - file_search_call_output: dict[str, Any], - message_output: dict[str, Any], + file_search_call_output: Mapping[str, object], + message_output: Mapping[str, object], first_response: ResponsesAPIResponse | None = None, ) -> ResponsesAPIResponse: """ @@ -343,21 +341,20 @@ def _synthesize_responses_api_response( synthesized _hidden_params so that billing callbacks see the total cost of both provider calls that the emulated flow makes. """ - synthesized_output: Final[list[dict[str, Any]]] = [file_search_call_output, message_output] synthesized: Final = ResponsesAPIResponse( id=getattr(original_response, "id", f"resp_{uuid.uuid4().hex}"), object="response", created_at=getattr(original_response, "created_at", int(time.time())), status="completed", model=getattr(original_response, "model", ""), - output=cast(list[ResponseOutputItem | dict[str, Any]], synthesized_output), + output=[dict(file_search_call_output), dict(message_output)], usage=getattr(original_response, "usage", None), error=None, ) if hasattr(original_response, "_hidden_params"): hidden: Final = dict(getattr(original_response, "_hidden_params") or {}) if first_response is not None and hasattr(first_response, "_hidden_params"): - first_hidden: Final = getattr(first_response, "_hidden_params") or {} + first_hidden: Final[object] = getattr(first_response, "_hidden_params", None) or {} first_cost: Final = ( first_hidden.get("response_cost") if isinstance(first_hidden, dict) @@ -382,9 +379,10 @@ async def _call_aresponses(input, model, tools, **kwargs): # pragma: no cover def _prepare_emulated_file_search_call( - kwargs: dict[str, Any], -) -> tuple[bool, dict[str, Any]]: - include_items: Final[list[str]] = list(kwargs.get("include") or []) + kwargs: Mapping[str, object], +) -> tuple[bool, Mapping[str, object]]: + raw_include: Final = kwargs.get("include") + include_items: Final[Sequence[str]] = raw_include if isinstance(raw_include, list) else [] include_search_results: Final = "file_search_call.results" in include_items original_stream: Final = kwargs.get("stream") @@ -398,7 +396,7 @@ def _prepare_emulated_file_search_call( return include_search_results, updated_kwargs -def _extract_tool_call_fields(tool_call: Any, fallback_call_id: str) -> tuple[str, str]: +def _extract_tool_call_fields(tool_call: object, fallback_call_id: str) -> tuple[str, str]: """Extract (call_id, raw_arguments_string) from a dict or Pydantic tool_call item.""" if isinstance(tool_call, dict): call_id = str(tool_call.get("call_id") or tool_call.get("id") or fallback_call_id) @@ -410,7 +408,13 @@ def _extract_tool_call_fields(tool_call: Any, fallback_call_id: str) -> tuple[st return call_id, raw_args -def _resolve_queries_from_args(args: dict[str, Any], input: Any) -> list[str]: +class _FileSearchArguments(TypedDict, total=False): + queries: Sequence[str] + query: str + vector_store_id: str + + +def _resolve_queries_from_args(args: _FileSearchArguments, input: object) -> Sequence[str]: """Pull the queries list out of parsed tool-call arguments, with backward-compat fallbacks.""" queries_from_call: Final = args.get("queries") if not queries_from_call: @@ -422,76 +426,96 @@ def _resolve_queries_from_args(args: dict[str, Any], input: Any) -> list[str]: return queries_from_call -async def _execute_file_search_tool_calls( - file_search_calls: list[Any], - all_vs_ids: list[str], - input: Any, +def _parse_file_search_arguments(raw_args: str) -> _FileSearchArguments: + if not isinstance(raw_args, str): + return raw_args + try: + return json.loads(raw_args) + except json.JSONDecodeError: + return {} + + +async def _execute_single_file_search_call( + tool_call: object, + all_vs_ids: Sequence[str], + input: object, file_search_call_id: str, -) -> tuple[list[dict[str, Any]], list[str], list[VectorStoreSearchResult]]: +) -> tuple[Mapping[str, object], Sequence[str], Sequence[VectorStoreSearchResult]]: + call_id, raw_args = _extract_tool_call_fields(tool_call, fallback_call_id=file_search_call_id) + args: Final = _parse_file_search_arguments(raw_args) + queries_from_call: Final = _resolve_queries_from_args(args, input) + + vs_id_arg: Final = args.get("vector_store_id") + vs_ids_for_call: Final = [vs_id_arg] if vs_id_arg else all_vs_ids + + queries, results = await _run_vector_searches( + queries=queries_from_call, + vector_store_ids=vs_ids_for_call, + ) + + return ( + { + "type": "function_call_output", + "call_id": call_id, + "output": _format_search_results_as_tool_output(results), + }, + queries, + results, + ) + + +async def _execute_file_search_tool_calls( + file_search_calls: Sequence[object], + all_vs_ids: Sequence[str], + input: object, + file_search_call_id: str, +) -> tuple[Sequence[Mapping[str, object]], Sequence[str], Sequence[VectorStoreSearchResult]]: """Run the vector search for each file_search tool_call and collect results.""" - tool_results: Final[list[dict[str, Any]]] = [] - all_queries: Final[list[str]] = [] - all_results: Final[list[VectorStoreSearchResult]] = [] + per_call: Final = tuple( + [ + await _execute_single_file_search_call( + tool_call=tool_call, + all_vs_ids=all_vs_ids, + input=input, + file_search_call_id=file_search_call_id, + ) + for tool_call in file_search_calls + ] + ) - for tool_call in file_search_calls: - call_id, raw_args = _extract_tool_call_fields(tool_call, fallback_call_id=file_search_call_id) - - try: - args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args - except json.JSONDecodeError: - args = {} - - queries_from_call = _resolve_queries_from_args(args, input) - - vs_id_arg = args.get("vector_store_id") - vs_ids_for_call = [vs_id_arg] if vs_id_arg else all_vs_ids - - queries, results = await _run_vector_searches( - queries=queries_from_call, - vector_store_ids=vs_ids_for_call, - ) - all_queries.extend(queries) - all_results.extend(results) - - tool_results.append( - { - "type": "function_call_output", - "call_id": call_id, - "output": _format_search_results_as_tool_output(results), - } - ) - - return tool_results, all_queries, all_results + return ( + [tool_result for tool_result, _, _ in per_call], + [query for _, queries, _ in per_call for query in queries], + [result for _, _, results in per_call for result in results], + ) def _build_follow_up_input( - input: Any, + input: object, first_response: ResponsesAPIResponse, - tool_results: list[dict[str, Any]], -) -> list[Any]: + tool_results: Sequence[Mapping[str, object]], +) -> Sequence[object]: """Assemble the follow-up call input: original messages + first-response output + tool results. Including all output items (text blocks, reasoning, non-file-search calls) ensures providers like Anthropic that emit text before the tool call have complete conversation context. Serializes Pydantic model instances to plain dicts so the transformation layer can call .get(). """ - original_input_items: Final = ( - list(input) if isinstance(input, (list, tuple)) else [{"role": "user", "content": str(input)}] + original_input_items: Final[tuple[object, ...]] = ( + tuple(input) if isinstance(input, (list, tuple)) else ({"role": "user", "content": str(input)},) + ) + first_response_output_items: Final[tuple[object, ...]] = tuple( + _item + if isinstance(_item, dict) + else (_item.model_dump(exclude_none=True) if hasattr(_item, "model_dump") else _item) + for _item in first_response.output ) - first_response_output_items: Final[list[Any]] = [] - for _item in first_response.output: - if isinstance(_item, dict): - first_response_output_items.append(_item) - elif hasattr(_item, "model_dump"): - first_response_output_items.append(_item.model_dump(exclude_none=True)) - else: - first_response_output_items.append(_item) - return original_input_items + first_response_output_items + tool_results + return [*original_input_items, *first_response_output_items, *tool_results] async def aresponses_with_emulated_file_search( - input: Any, + input: object, model: str, tools: Iterable[ToolParam] | None = None, # Pass-through params — forwarded as-is to the underlying aresponses call @@ -504,7 +528,7 @@ async def aresponses_with_emulated_file_search( runs vector search, and synthesizes an OpenAI-format response. """ # Determine whether caller wants search_results populated in the output. - _include_search_results, kwargs = _prepare_emulated_file_search_call(kwargs=kwargs) + _include_search_results, call_kwargs = _prepare_emulated_file_search_call(kwargs=kwargs) # 1. Replace file_search tools with function tool transformed_tools, all_vs_ids = _replace_file_search_tools(tools) @@ -521,7 +545,7 @@ async def aresponses_with_emulated_file_search( input=input, model=model, tools=transformed_tools or None, - **kwargs, + **call_kwargs, ), ) finally: @@ -585,7 +609,7 @@ async def aresponses_with_emulated_file_search( input=follow_up_input, model=model, tools=None, # no tools needed for the answer step - **kwargs, + **call_kwargs, ), ) finally: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index d7f6ece5cd1..1385e329e93 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -5,11 +5,11 @@ import json import time import traceback import uuid -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runtime_checkable import httpx from openai._streaming import SSEDecoder @@ -42,12 +42,14 @@ from litellm.types.utils import CallTypes from litellm.utils import async_post_call_success_deployment_hook if TYPE_CHECKING: + from litellm.caching.caching_handler import LLMCachingHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.types.responses.streaming_websocket import ( PresidioGuardrailCallback, ResponsesBackendWebSocket, ResponsesClientWebSocket, ) + from litellm.types.router import LiteLLM_Params @lru_cache(maxsize=1) @@ -69,6 +71,79 @@ def _is_str_mapping(value: object) -> TypeIs[dict[str, str]]: # guard-ok: verif return _is_json_object(value) and all(isinstance(item, str) for item in value.values()) +class _MutableJsonObject(Protocol): + @overload + def get(self, key: str, /) -> object | None: ... + @overload + def get(self, key: str, default: object, /) -> object: ... + def __getitem__(self, key: str, /) -> object: ... + def __setitem__(self, key: str, value: object, /) -> None: ... + def __contains__(self, key: object, /) -> bool: ... + def items(self) -> Iterable[tuple[str, object]]: ... + + +class _LoadsJsonValue(Protocol): + def __call__(self, s: str | bytes, /) -> object: ... + + +class _LoadsJsonDict(Protocol): + def __call__(self, s: str | bytes, /) -> _MutableJsonObject: ... + + +class _GetsLitellmParams(Protocol): + def __call__(self, key: str, default: Mapping[str, object], /) -> LiteLLM_Params: ... + + +class _PopsOptionalStr(Protocol): + def __call__(self, key: str, default: None, /) -> str | None: ... + + +class _UnmasksPiiText(Protocol): + def __call__(self, text: str, pii_tokens: Mapping[str, str]) -> str: ... + + +class _ShouldStoreResultInCache(Protocol): + def __call__(self, *, original_function: Callable[..., object] | None, kwargs: Mapping[str, object]) -> bool: ... + + +class _PostStreamingDeploymentHook(Protocol): + def __call__( + self, + *, + request_data: Mapping[str, object], + response_chunk: ResponsesAPIStreamingResponse, + call_type: CallTypes | None, + ) -> Awaitable[ResponsesAPIStreamingResponse | None]: ... + + +@runtime_checkable +class _HasPostStreamingDeploymentHook(Protocol): + async_post_call_streaming_deployment_hook: _PostStreamingDeploymentHook + + +def _typed_loads_json_value(fn: _LoadsJsonValue) -> _LoadsJsonValue: + return fn + + +def _typed_loads_json_dict(fn: _LoadsJsonDict) -> _LoadsJsonDict: + return fn + + +def _typed_gets_litellm_params(fn: _GetsLitellmParams) -> _GetsLitellmParams: + return fn + + +def _typed_pops_optional_str(fn: _PopsOptionalStr) -> _PopsOptionalStr: + return fn + + +_LOADS_JSON_VALUE: Final = _typed_loads_json_value(json.loads) +_LOADS_JSON_DICT: Final = _typed_loads_json_dict(json.loads) + +_SHOULD_STORE_RESULT_IN_CACHE_ATTR: Final = "_should_store_result_in_cache" +_UNMASK_PII_TEXT_ATTR: Final = "_unmask_pii_text" + + def _model_id_from_metadata(litellm_metadata: dict[str, object] | None) -> str | None: model_info: Final = litellm_metadata.get("model_info") if litellm_metadata else None model_id: Final = model_info.get("id") if _is_json_object(model_info) else None @@ -185,7 +260,7 @@ class BaseResponsesAPIStreamingIterator: # This matches the stream wrapper in litellm/litellm_core_utils/streaming_handler.py _api_base: Final = get_api_base( model=model or "", - optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), + optional_params=_typed_gets_litellm_params(self.logging_obj.model_call_details.get)("litellm_params", {}), ) self._hidden_params: dict[str, object] = { "model_id": _model_id_from_metadata(litellm_metadata), @@ -228,7 +303,7 @@ class BaseResponsesAPIStreamingIterator: try: # Parse the JSON chunk - parsed_chunk: Final = json.loads(chunk) + parsed_chunk: Final = _LOADS_JSON_VALUE(chunk) # Format as ResponsesAPIStreamingResponse if isinstance(parsed_chunk, dict): @@ -514,7 +589,7 @@ class BaseResponsesAPIStreamingIterator: if response_obj is None: return - caching_handler: Final = getattr(self.logging_obj, "_llm_caching_handler", None) + caching_handler: Final[LLMCachingHandler | None] = getattr(self.logging_obj, "_llm_caching_handler", None) if caching_handler is None: return @@ -532,8 +607,11 @@ class BaseResponsesAPIStreamingIterator: if preset_cache_key is not None: request_kwargs["cache_key"] = preset_cache_key - if not caching_handler._should_store_result_in_cache( - original_function=caching_handler.original_function, + should_store_result_in_cache: Final[_ShouldStoreResultInCache] = getattr( + caching_handler, _SHOULD_STORE_RESULT_IN_CACHE_ATTR + ) + if not should_store_result_in_cache( + original_function=getattr(caching_handler, "original_function", None), kwargs=request_kwargs, ): return @@ -586,12 +664,15 @@ class BaseResponsesAPIStreamingIterator: typed_call_type = None request_data: Final = self.request_data or getattr(self.logging_obj, "model_call_details", {}) - callbacks: Final = getattr(litellm, "callbacks", None) or [] + callbacks: Final[Sequence[object]] = getattr(litellm, "callbacks", None) or [] hooks_ran = False for callback in callbacks: - if hasattr(callback, "async_post_call_streaming_deployment_hook"): + if isinstance(callback, _HasPostStreamingDeploymentHook): hooks_ran = True - result = await callback.async_post_call_streaming_deployment_hook( + post_streaming_hook: _PostStreamingDeploymentHook = ( + callback.async_post_call_streaming_deployment_hook + ) + result = await post_streaming_hook( request_data=request_data, response_chunk=chunk, call_type=typed_call_type, @@ -1043,8 +1124,8 @@ class _HasModelDumpJson(Protocol): def model_dump_json(self, *, exclude_none: bool = ...) -> str: ... -def _dump_response_object(obj: Any) -> dict[str, Any]: - if hasattr(obj, "model_dump"): +def _dump_response_object(obj: object) -> Mapping[str, object]: + if isinstance(obj, _HasModelDump): return obj.model_dump() if _is_json_object(obj): return obj @@ -1073,21 +1154,20 @@ def _build_content_part_done_event( item_id: str, output_index: int, content_index: int, - part_payload: dict[str, Any], + part_payload: Mapping[str, object], ) -> ResponsesAPIStreamingResponse | None: openai_types: Final = _get_openai_response_types() part_type: Final = part_payload.get("type") part: PART_UNION_TYPES if part_type == "output_text": - annotations: Final = [ - openai_types.BaseLiteLLMOpenAIResponseObject(**annotation) - for annotation in part_payload.get("annotations", []) or [] - ] - part = openai_types.ContentPartDonePartOutputText( - type="output_text", - text=str(part_payload.get("text") or ""), - annotations=annotations, - logprobs=part_payload.get("logprobs"), + raw_annotations: Final[object] = part_payload.get("annotations", []) or [] + part = openai_types.ContentPartDonePartOutputText.model_validate( + { + "type": "output_text", + "text": str(part_payload.get("text") or ""), + "annotations": raw_annotations, + "logprobs": part_payload.get("logprobs"), + } ) elif part_type == "refusal": part = openai_types.ContentPartDonePartRefusal( @@ -1117,7 +1197,7 @@ def _add_text_like_part_events( item_id: str, output_index: int, content_index: int, - part_payload: dict[str, Any], + part_payload: Mapping[str, object], chunk_size: int, ) -> None: openai_types: Final = _get_openai_response_types() @@ -1134,15 +1214,19 @@ def _add_text_like_part_events( delta=text[i : i + chunk_size], ) ) - for annotation_index, annotation in enumerate(part_payload.get("annotations", []) or []): + raw_annotation_items: Final = part_payload.get("annotations") + annotation_items: Final[Sequence[object]] = raw_annotation_items if _is_json_array(raw_annotation_items) else [] + for annotation_index, annotation in enumerate(annotation_items): events.append( - openai_types.OutputTextAnnotationAddedEvent( - type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, - item_id=item_id, - output_index=output_index, - content_index=content_index, - annotation_index=annotation_index, - annotation=annotation, + openai_types.OutputTextAnnotationAddedEvent.model_validate( + { + "type": openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, + "item_id": item_id, + "output_index": output_index, + "content_index": content_index, + "annotation_index": annotation_index, + "annotation": annotation, + } ) ) events.append( @@ -1200,7 +1284,8 @@ def _build_synthetic_response_events( ] sequence_number = 0 - for output_index, output_item in enumerate(getattr(transformed, "output", []) or []): + output_items: Final[Sequence[object]] = getattr(transformed, "output", []) or [] + for output_index, output_item in enumerate(output_items): output_item_payload = _dump_response_object(output_item) item_id = str(output_item_payload.get("id") or transformed.id) item_type = output_item_payload.get("type") @@ -1214,7 +1299,9 @@ def _build_synthetic_response_events( ) if item_type == "message": - for content_index, part in enumerate(output_item_payload.get("content", []) or []): + raw_content_parts = output_item_payload.get("content") + content_parts: Sequence[object] = raw_content_parts if _is_json_array(raw_content_parts) else [] + for content_index, part in enumerate(content_parts): part_payload = _dump_response_object(part) events.append( openai_types.ContentPartAddedEvent( @@ -1261,7 +1348,9 @@ def _build_synthetic_response_events( ) ) elif item_type == "reasoning": - for summary_index, summary in enumerate(output_item_payload.get("summary", []) or []): + raw_summary_items = output_item_payload.get("summary") + summary_items: Sequence[object] = raw_summary_items if _is_json_array(raw_summary_items) else [] + for summary_index, summary in enumerate(summary_items): summary_payload = _dump_response_object(summary) summary_text = str(summary_payload.get("text") or "") for i in range(0, len(summary_text), chunk_size): @@ -1354,7 +1443,7 @@ class ResponsesWebSocketStreaming: user_api_key_dict: UserAPIKeyAuth | None = None, request_data: dict[str, object] | None = None, first_message: str | None = None, - guardrail_callbacks: list[Any] | None = None, + guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None, output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, authorized_model: str | None = None, ): @@ -1363,16 +1452,16 @@ class ResponsesWebSocketStreaming: self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict self.request_data: dict[str, object] = request_data or {} - self.messages: list[dict[str, object]] = [] + self.messages: list[_MutableJsonObject] = [] self.input_messages: list[dict[str, object]] = [] self.first_message = first_message - self.guardrail_callbacks: list[Any] = guardrail_callbacks or [] + self.guardrail_callbacks: Sequence[PresidioGuardrailCallback] = guardrail_callbacks or [] self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or [] # Model name authorized at connection time; enforced on every # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model - def _should_store_event(self, event_obj: Mapping[str, object]) -> bool: + def _should_store_event(self, event_obj: _MutableJsonObject) -> bool: return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES def _store_event(self, event: str | bytes | dict[str, object]) -> None: @@ -1380,7 +1469,7 @@ class ResponsesWebSocketStreaming: event = event.decode("utf-8") if isinstance(event, str): try: - event_obj = json.loads(event) + event_obj = _LOADS_JSON_DICT(event) except (json.JSONDecodeError, TypeError): return else: @@ -1393,7 +1482,7 @@ class ResponsesWebSocketStreaming: """Extract user input content from response.create for logging.""" try: if isinstance(message, str): - msg_obj = json.loads(message) + msg_obj = _LOADS_JSON_DICT(message) elif _is_json_object(message): msg_obj = message else: @@ -1463,7 +1552,7 @@ class ResponsesWebSocketStreaming: # masked response.completed. if self.output_guardrail_callbacks: try: - _evt_type = json.loads(response_str).get("type") + _evt_type = _LOADS_JSON_DICT(response_str).get("type") except (json.JSONDecodeError, TypeError): _evt_type = None if _evt_type in self._DELTA_EVENT_TYPES or _evt_type in self._OUTPUT_DONE_EVENT_TYPES: @@ -1485,7 +1574,7 @@ class ResponsesWebSocketStreaming: finally: await self._log_messages() - def _enforce_authorized_model(self, msg_obj: dict[str, object]) -> bool: + def _enforce_authorized_model(self, msg_obj: _MutableJsonObject) -> bool: """ Overwrite any ``model`` field in a ``response.create`` frame with the connection-authorized model to prevent deployment-substitution attacks. @@ -1527,7 +1616,7 @@ class ResponsesWebSocketStreaming: Non-``response.create`` messages are returned unchanged. """ try: - msg_obj: Final = json.loads(message) + msg_obj: Final = _LOADS_JSON_DICT(message) except (json.JSONDecodeError, TypeError): return message @@ -1553,7 +1642,7 @@ class ResponsesWebSocketStreaming: # forwarded unmasked regardless of where the client places it. nested_candidate = msg_obj.get("response") nested_response = nested_candidate if _is_json_object(nested_candidate) else None - text_containers: list[tuple[dict[str, object], str]] = [] + text_containers: list[tuple[_MutableJsonObject, str]] = [] for container in (msg_obj, nested_response): if container is None: continue @@ -1655,11 +1744,12 @@ class ResponsesWebSocketStreaming: return response_str try: - evt_obj: Final = json.loads(response_str) + evt_obj: Final = _LOADS_JSON_DICT(response_str) except (json.JSONDecodeError, TypeError): return response_str cb: Final = self.guardrail_callbacks[0] + unmask_pii_text: Final[_UnmasksPiiText] = getattr(cb, _UNMASK_PII_TEXT_ATTR) event_type: Final = evt_obj.get("type") if event_type == "response.completed": @@ -1679,7 +1769,7 @@ class ResponsesWebSocketStreaming: continue text = content_block.get("text") if isinstance(text, str): - unmasked = cb._unmask_pii_text(text, pii_tokens) + unmasked = unmask_pii_text(text, pii_tokens) if unmasked != text: content_block["text"] = unmasked modified = True @@ -1688,7 +1778,7 @@ class ResponsesWebSocketStreaming: if event_type in self._DELTA_EVENT_TYPES: delta: Final = evt_obj.get("delta") if isinstance(delta, str): - unmasked = cb._unmask_pii_text(delta, pii_tokens) + unmasked = unmask_pii_text(delta, pii_tokens) if unmasked != delta: evt_obj["delta"] = unmasked return json.dumps(evt_obj) @@ -1711,7 +1801,7 @@ class ResponsesWebSocketStreaming: return response_str try: - evt_obj: Final[Mapping[str, object]] = json.loads(response_str) + evt_obj: Final = _LOADS_JSON_DICT(response_str) except (json.JSONDecodeError, TypeError): return response_str @@ -1859,7 +1949,7 @@ class ManagedResponsesWebSocketHandler: model: str, logging_obj: LiteLLMLoggingObj, user_api_key_dict: UserAPIKeyAuth | None = None, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, object] | None = None, api_key: str | None = None, api_base: str | None = None, timeout: float | None = None, @@ -1871,10 +1961,11 @@ class ManagedResponsesWebSocketHandler: self.model = model self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict - self.litellm_metadata: dict[str, Any] = litellm_metadata or {} - self.model_group: str | None = self.litellm_metadata.get("model_group") or self.litellm_metadata.get( + self.litellm_metadata: Mapping[str, object] = litellm_metadata or {} + raw_model_group: Final = self.litellm_metadata.get("model_group") or self.litellm_metadata.get( "deployment_model_name" ) + self.model_group: str | None = raw_model_group if isinstance(raw_model_group, str) else None self.api_key = api_key self.api_base = api_base self.timeout = timeout @@ -1894,7 +1985,7 @@ class ManagedResponsesWebSocketHandler: # ------------------------------------------------------------------ @staticmethod - def _serialize_chunk(chunk: Any) -> str | None: + def _serialize_chunk(chunk: object) -> str | None: """Serialize a streaming chunk to a JSON string for WebSocket transmission.""" try: if isinstance(chunk, _HasModelDumpJson): @@ -1937,7 +2028,7 @@ class ManagedResponsesWebSocketHandler: self._session_history[response_id] = messages @staticmethod - def _extract_response_id(completed_event: dict[str, object]) -> str | None: + def _extract_response_id(completed_event: _MutableJsonObject) -> str | None: """ Pull the raw (decoded) response ID out of a ``response.completed`` event. Returns *None* if the event doesn't contain a usable ID. @@ -1952,7 +2043,7 @@ class ManagedResponsesWebSocketHandler: @staticmethod def _extract_output_messages( - completed_event: dict[str, object], + completed_event: _MutableJsonObject, ) -> list[dict[str, object]]: """ Convert the output items in a ``response.completed`` event into @@ -2009,10 +2100,10 @@ class ManagedResponsesWebSocketHandler: # _process_response_create sub-methods # ------------------------------------------------------------------ - async def _parse_message(self, raw_message: str) -> dict[str, object] | None: + async def _parse_message(self, raw_message: str) -> _MutableJsonObject | None: """Parse raw WS text; return the message dict or None (JSON error / ignored type).""" try: - msg_obj: Final = json.loads(raw_message) + msg_obj: Final = _LOADS_JSON_DICT(raw_message) except json.JSONDecodeError: await self._send_error("Invalid JSON in response.create event", "invalid_request_error") return None @@ -2022,7 +2113,7 @@ class ManagedResponsesWebSocketHandler: return msg_obj @staticmethod - def _is_warmup_frame(msg_obj: dict[str, object]) -> bool: + def _is_warmup_frame(msg_obj: _MutableJsonObject) -> bool: """Return True for a response.create whose generate flag is false.""" nested: Final = msg_obj.get("response") source: Final = nested if _is_json_object(nested) and nested else msg_obj @@ -2038,13 +2129,13 @@ class ManagedResponsesWebSocketHandler: return str(raw_id).startswith(_WARMUP_RESPONSE_ID_PREFIX) @staticmethod - def _warmup_source_params(msg_obj: dict[str, object]) -> dict[str, object]: + def _warmup_source_params(msg_obj: _MutableJsonObject) -> dict[str, object]: nested: Final = msg_obj.get("response") if _is_json_object(nested) and nested: return nested return {k: v for k, v in msg_obj.items() if k != "type"} - def _build_warmup_response(self, msg_obj: dict[str, object]) -> dict[str, object]: + def _build_warmup_response(self, msg_obj: _MutableJsonObject) -> dict[str, object]: """Build a minimal completed Responses API object for a warmup ack.""" source: Final = self._warmup_source_params(msg_obj) wire_model: Final = source.get("model") or self.model_group or self.model @@ -2062,7 +2153,7 @@ class ManagedResponsesWebSocketHandler: }, } - async def _send_warmup_ack(self, msg_obj: dict[str, object]) -> None: + async def _send_warmup_ack(self, msg_obj: _MutableJsonObject) -> None: """ Acknowledge a generate=false prewarm without calling the provider. @@ -2085,7 +2176,7 @@ class ManagedResponsesWebSocketHandler: await self.websocket.send_text(serialized) @staticmethod - def _build_base_call_kwargs(msg_obj: dict[str, object]) -> dict[str, Any]: + def _build_base_call_kwargs(msg_obj: _MutableJsonObject) -> dict[str, Any]: """ Extract Responses API params from the event, handling both wire formats: Nested: {"type": "response.create", "response": {"input": [...], ...}} @@ -2194,7 +2285,7 @@ class ManagedResponsesWebSocketHandler: call_kwargs.setdefault("litellm_params", {}) call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request - async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> dict[str, object] | None: + async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> _MutableJsonObject | None: """ Stream ``litellm.aresponses`` and forward every chunk over the WebSocket. @@ -2202,7 +2293,7 @@ class ManagedResponsesWebSocketHandler: directly (before serialization) to avoid a redundant JSON round-trip on every chunk. Returns the completed event dict, or ``None``. """ - completed_event: dict[str, object] | None = ( + completed_event: _MutableJsonObject | None = ( None # rebind-ok: captures the completed event once the stream yields it ) stream_response: Final = await litellm.aresponses(model=model, **call_kwargs) @@ -2216,7 +2307,7 @@ class ManagedResponsesWebSocketHandler: continue if chunk_type == "response.completed" and completed_event is None: try: - completed_event = json.loads(serialized) + completed_event = _LOADS_JSON_DICT(serialized) except Exception: pass try: @@ -2228,7 +2319,7 @@ class ManagedResponsesWebSocketHandler: def _save_turn_history( self, - completed_event: dict[str, object] | None, + completed_event: _MutableJsonObject | None, prior_history: list[dict[str, object]], current_messages: list[dict[str, object]], ) -> None: @@ -2293,13 +2384,15 @@ class ManagedResponsesWebSocketHandler: # reuse the router-resolved self.model; passing the alias raw to # litellm.aresponses fails in get_llm_provider. A genuinely different # provider-prefixed per-frame model is still honored. - requested_model: Final = call_kwargs.pop("model", None) + requested_model: Final = _typed_pops_optional_str(call_kwargs.pop)("model", None) if requested_model is None or requested_model == self.model_group: model = self.model else: model = requested_model - previous_response_id: Final[str | None] = call_kwargs.pop("previous_response_id", None) + previous_response_id: Final[str | None] = _typed_pops_optional_str(call_kwargs.pop)( + "previous_response_id", None + ) current_messages: Final = self._input_to_messages(call_kwargs.get("input")) # Fetch history once; reused in both _apply_history and _save_turn_history diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index bbe97613c57..209dea87cc5 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -10,10 +10,10 @@ Use this to route requests between Teams import re from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload from litellm._logging import verbose_logger -from litellm.types.router import RouterErrors +from litellm.types.router import DeploymentTypedDict, RouterErrors if TYPE_CHECKING: from litellm.router import Router as _Router @@ -23,9 +23,68 @@ else: LitellmRouter = Any +class _TagLitellmParamsLike(Protocol): + @overload + def get(self, key: Literal["tags"], /) -> Sequence[str] | None: ... + @overload + def get(self, key: Literal["tags"], default: Sequence[str], /) -> Sequence[str]: ... + @overload + def get(self, key: Literal["tag_regex"], /) -> Sequence[str] | None: ... + + +class _ModelInfoLike(Protocol): + @overload + def get(self, key: Literal["allow_fail_open"], /) -> bool | None: ... + @overload + def get(self, key: Literal["enable_tag_filtering"], /) -> bool | None: ... + + +class _DeploymentLike(Protocol): + @overload + def get(self, key: Literal["litellm_params"], default: Mapping[str, object], /) -> _TagLitellmParamsLike: ... + @overload + def get(self, key: Literal["model_info"], /) -> _ModelInfoLike | None: ... + @overload + def get(self, key: Literal["model_name"], /) -> object: ... + + +class _MetadataLike(Protocol): + @overload + def get(self, key: Literal["tags"], /) -> Sequence[str] | None: ... + @overload + def get(self, key: Literal["tags"], default: Sequence[str], /) -> Sequence[str] | None: ... + @overload + def get(self, key: Literal["user_agent"], default: str, /) -> str: ... + @overload + def get(self, key: Literal["inherited_tags"], /) -> object: ... + def __contains__(self, key: object, /) -> bool: ... + def __setitem__(self, key: Literal["tag_routing"], value: Mapping[str, object], /) -> None: ... + + +class _NestedLitellmParamsLike(Protocol): + def get( + self, key: Literal["metadata", "litellm_metadata"], default: Mapping[str, object], / + ) -> _MetadataLike | None: ... + + +class _RequestKwargsLike(Protocol): + @overload + def get(self, key: Literal["enable_tag_filtering"], /) -> bool | None: ... + @overload + def get(self, key: Literal["metadata", "litellm_metadata"], /) -> _MetadataLike | None: ... + def __contains__(self, key: object, /) -> bool: ... + @overload + def __getitem__(self, key: Literal["metadata", "litellm_metadata"], /) -> _MetadataLike: ... + @overload + def __getitem__(self, key: Literal["litellm_params"], /) -> _NestedLitellmParamsLike: ... + + +_DeploymentPool = Sequence[_DeploymentLike] | Mapping[_DeploymentLike, object] + + def _is_valid_deployment_tag_regex( - tag_regexes: list[str], - header_strings: list[str], + tag_regexes: Sequence[str], + header_strings: Sequence[str], ) -> str | None: """ Test compiled regex patterns against "Header-Name: value" strings. @@ -46,7 +105,9 @@ def _is_valid_deployment_tag_regex( return None -def is_valid_deployment_tag(deployment_tags: list[str], request_tags: list[str], match_any: bool = True) -> bool: +def is_valid_deployment_tag( + deployment_tags: Sequence[str], request_tags: Sequence[str], match_any: bool = True +) -> bool: """ Check if a tag is valid, the matching can be either any or all based on `match_any` flag """ @@ -73,7 +134,7 @@ def is_valid_deployment_tag(deployment_tags: list[str], request_tags: list[str], def _match_deployment( - deployment: Any, + deployment: _DeploymentLike, request_tags: list[str] | None, header_strings: list[str], match_any: bool, @@ -90,8 +151,8 @@ def _match_deployment( ran and failed, so the regex cannot override strict-tag policy. """ litellm_params: Final = deployment.get("litellm_params", {}) - deployment_tags: Final[list[str] | None] = litellm_params.get("tags") - deployment_tag_regex: Final[list[str] | None] = litellm_params.get("tag_regex") + deployment_tags: Final[Sequence[str] | None] = litellm_params.get("tags") + deployment_tag_regex: Final[Sequence[str] | None] = litellm_params.get("tag_regex") # 1. Exact tag match (existing behaviour). if deployment_tags and request_tags: @@ -162,38 +223,38 @@ def _split_tags(tags: Sequence[str]) -> tuple[tuple[str, ...], list[str], tuple[ def _exclude_deployments( - deployments: Sequence[Any] | Mapping[Any, Any], + deployments: _DeploymentPool, excluded_set: frozenset[str], -) -> list[Any]: +) -> Sequence[_DeploymentLike]: if not excluded_set: return list(deployments) return [d for d in deployments if not excluded_set.intersection(d.get("litellm_params", {}).get("tags") or [])] def _require_all_tags( - deployments: Sequence[Any] | Mapping[Any, Any], + deployments: _DeploymentPool, required_set: frozenset[str], -) -> tuple[Any, ...]: +) -> tuple[_DeploymentLike, ...]: if not required_set: return tuple(deployments) return tuple(d for d in deployments if required_set.issubset(d.get("litellm_params", {}).get("tags") or [])) def _default_tagged_pool( - deployments: Sequence[Any] | Mapping[Any, Any], -) -> tuple[Any, ...]: + deployments: _DeploymentPool, +) -> tuple[_DeploymentLike, ...]: defaults: Final = tuple(d for d in deployments if "default" in (d.get("litellm_params", {}).get("tags") or [])) return defaults if defaults else tuple(deployments) -def _known_tag_values(deployments: Sequence[Any] | Mapping[Any, Any]) -> frozenset[str]: +def _known_tag_values(deployments: _DeploymentPool) -> frozenset[str]: return frozenset( tag for d in deployments for tag in (d.get("litellm_params", MappingProxyType({})).get("tags") or ()) ) def _unknown_required_tag_hides_an_answer( - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], routing_confirmed: frozenset[str], @@ -217,7 +278,7 @@ def _unknown_required_tag_hides_an_answer( def _chain_allows_fail_open( - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], routing_confirmed: frozenset[str], @@ -228,12 +289,12 @@ def _chain_allows_fail_open( def _trusted_only_pool( - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, inherited_required_set: frozenset[str] | None, -) -> tuple[Any, ...]: +) -> tuple[_DeploymentLike, ...]: # inherited_*_set is None only when this request carries no origin information # at all (e.g. direct SDK Router usage, bypassing the proxy layer that # populates metadata.inherited_tags) -- treat every constraint as @@ -260,8 +321,8 @@ def _trusted_only_pool( def _resolve_or_fail_open( - pool: Sequence[Any], - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + pool: Sequence[_DeploymentLike], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, @@ -269,7 +330,7 @@ def _resolve_or_fail_open( routing_confirmed: frozenset[str], model: str, request_tags: object, -) -> tuple[Any, ...]: +) -> tuple[_DeploymentLike, ...]: if pool: return tuple(pool) if _chain_allows_fail_open(healthy_deployments, excluded_set, required_set, routing_confirmed): @@ -289,7 +350,7 @@ def _resolve_or_fail_open( def _resolve_constraint_only_pool( - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, @@ -297,7 +358,7 @@ def _resolve_constraint_only_pool( routing_confirmed: frozenset[str], model: str, request_tags: object, -) -> tuple[Any, ...]: +) -> tuple[_DeploymentLike, ...]: pool: Final = ( _require_all_tags(_exclude_deployments(healthy_deployments, excluded_set), required_set) if required_set @@ -319,8 +380,8 @@ def _resolve_constraint_only_pool( def _all_deployments_or_fallback( llm_router_instance: LitellmRouter, model: str, - fallback: Sequence[Any] | Mapping[Any, Any], -) -> Sequence[Any] | Mapping[Any, Any]: + fallback: _DeploymentPool, +) -> Sequence[_DeploymentLike | DeploymentTypedDict] | Mapping[_DeploymentLike, object]: try: return llm_router_instance._get_all_deployments(model_name=model) except Exception: # noqa: BLE001 # fail safe toward today's healthy-only behavior on lookup errors @@ -330,7 +391,7 @@ def _all_deployments_or_fallback( def _chain_tag_filtering_override( llm_router_instance: LitellmRouter, model: str, - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: _DeploymentPool, ) -> bool | None: # Resolved from every deployment configured for this model group, not just the # ones that survived cooldown/health filtering (async_get_healthy_deployments @@ -392,10 +453,10 @@ def _tag_known_to_group( async def get_deployments_for_tag( llm_router_instance: LitellmRouter, model: str, # used to raise the correct error - healthy_deployments: list[Any] | dict[Any, Any], - request_kwargs: dict[Any, Any] | None = None, + healthy_deployments: _DeploymentPool, + request_kwargs: _RequestKwargsLike | None = None, metadata_variable_name: Literal["metadata", "litellm_metadata"] = "metadata", -): +) -> _DeploymentPool: """ Returns a list of deployments that match the requested model and tags in the request. @@ -473,25 +534,25 @@ async def get_deployments_for_tag( request_tags, ) - new_healthy_deployments: Final[list[Any]] = [] - default_deployments: Final[list[Any]] = [] - if has_positive_filter: verbose_logger.debug( "get_deployments_for_tag routing: request_tags=%s user_agent=%s", request_tags, user_agent, ) - for deployment in candidates: - deployment_tags = deployment.get("litellm_params", {}).get("tags") - - match_result = _match_deployment( - deployment=deployment, - request_tags=positive_tags, - header_strings=header_strings, - match_any=match_any, + deployment_matches: Final = tuple( + ( + deployment, + _match_deployment( + deployment=deployment, + request_tags=positive_tags, + header_strings=header_strings, + match_any=match_any, + ), ) - + for deployment in candidates + ) + for deployment, match_result in deployment_matches: if match_result is not None: verbose_logger.debug( "tag routing match: deployment=%s matched_via=%s matched_value=%s", @@ -507,10 +568,10 @@ async def get_deployments_for_tag( "request_tags": request_tags or [], "user_agent": user_agent, } - new_healthy_deployments.append(deployment) - - if deployment_tags and "default" in deployment_tags: - default_deployments.append(deployment) + new_healthy_deployments: Final = [d for d, result in deployment_matches if result is not None] + default_deployments: Final = [ + d for d, _ in deployment_matches if "default" in (d.get("litellm_params", {}).get("tags") or ()) + ] if len(new_healthy_deployments) == 0 and len(default_deployments) == 0: return _resolve_or_fail_open( @@ -545,10 +606,11 @@ async def get_deployments_for_tag( return new_healthy_deployments if len(new_healthy_deployments) > 0 else default_deployments # for Untagged requests use default deployments if set - _default_deployments_with_tags: Final = [] - for deployment in healthy_deployments: - if "default" in deployment.get("litellm_params", {}).get("tags", []): - _default_deployments_with_tags.append(deployment) + _default_deployments_with_tags: Final = [ + deployment + for deployment in healthy_deployments + if "default" in deployment.get("litellm_params", {}).get("tags", []) + ] if len(_default_deployments_with_tags) > 0: return _default_deployments_with_tags @@ -562,7 +624,7 @@ async def get_deployments_for_tag( def _get_tags_from_request_kwargs( - request_kwargs: dict[Any, Any] | None = None, + request_kwargs: _RequestKwargsLike | None = None, metadata_variable_name: Literal["metadata", "litellm_metadata"] = "metadata", ) -> list[str]: """ @@ -577,12 +639,12 @@ def _get_tags_from_request_kwargs( if request_kwargs is None: return [] if metadata_variable_name in request_kwargs: - metadata: Final = request_kwargs[metadata_variable_name] or {} + metadata: Final[_MetadataLike] = request_kwargs[metadata_variable_name] or {} tags = metadata.get("tags", []) - return tags if tags is not None else [] + return list(tags) if tags is not None else [] elif "litellm_params" in request_kwargs: - litellm_params: Final = request_kwargs["litellm_params"] or {} - _metadata: Final = litellm_params.get(metadata_variable_name, {}) or {} + litellm_params: Final[_NestedLitellmParamsLike] = request_kwargs["litellm_params"] or {} + _metadata: Final[_MetadataLike] = litellm_params.get(metadata_variable_name, {}) or {} tags = _metadata.get("tags", []) - return tags if tags is not None else [] + return list(tags) if tags is not None else [] return [] diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index dff010bfd30..90033af024b 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3058 + "limit": 3036 }, "ANN002": { "limit": 71 @@ -9,7 +9,7 @@ "limit": 827 }, "ANN201": { - "limit": 2022 + "limit": 2020 }, "ANN202": { "limit": 855 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1384 + "limit": 1286 }, "ASYNC230": { "limit": 11 @@ -39,7 +39,7 @@ "limit": 505 }, "B009": { - "limit": 64 + "limit": 63 }, "B010": { "limit": 190 @@ -123,7 +123,7 @@ "limit": 12 }, "PERF403": { - "limit": 34 + "limit": 33 }, "PIE804": { "limit": 18 @@ -180,7 +180,7 @@ "limit": 8 }, "RUF019": { - "limit": 38 + "limit": 36 }, "RUF046": { "limit": 4 @@ -201,7 +201,7 @@ "limit": 58 }, "SIM102": { - "limit": 321 + "limit": 318 }, "SIM103": { "limit": 119 @@ -234,7 +234,7 @@ "limit": 5 }, "TID251": { - "limit": 1224 + "limit": 1214 }, "TRY002": { "limit": 524 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index fdacf375844..83def5fe2e3 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23003 + "limit": 22780 }, "LIT002": { - "limit": 27146 + "limit": 27144 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1077 + "limit": 1069 }, "LIT007": { "limit": 0 @@ -27,9 +27,9 @@ "limit": 0 }, "LIT010": { - "limit": 16731 + "limit": 16725 }, "LIT011": { - "limit": 5596 + "limit": 5590 } } From 141ada1118bc64c011b67ef9fd7ee2f15854865c Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 17 Aug 2026 11:25:55 -0400 Subject: [PATCH 005/111] feat(batches): aggregate reasoning tokens and per-line pass/fail counts Batch retrieval already computed cost/usage on completion, but silently dropped reasoning tokens and never counted per-line success/failure. Adds BatchCostUsageResult (replacing bare cost/usage/models tuples) with successful_requests/failed_requests, and threads reasoning_tokens through the aggregated Usage. Both surface on SpendLogs the same way batch_models already does. --- .../proxy/common_utils/check_batch_cost.py | 22 +- litellm/batches/batch_utils.py | 92 +++++-- litellm/litellm_core_utils/litellm_logging.py | 24 +- litellm/proxy/_types.py | 2 + .../spend_tracking/spend_tracking_utils.py | 16 ++ litellm/types/utils.py | 4 +- .../test_batch_custom_pricing.py | 16 +- tests/batches_tests/test_batch_rate_limits.py | 14 +- .../test_batches_logging_unit_tests.py | 58 ++-- .../proxy_unit_tests/test_check_batch_cost.py | 32 ++- .../test_litellm/batches/test_batch_utils.py | 259 ++++++++++++------ .../test_vertex_ai_batch_passthrough.py | 78 +++--- 12 files changed, 406 insertions(+), 211 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 25b00597355..85b9bd77bc3 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -651,16 +651,14 @@ class CheckBatchCost: # Pass deployment model_info so custom batch pricing # (input_cost_per_token_batches etc.) is used for cost calc deployment_model_info = deployment_info.model_info.model_dump() if deployment_info.model_info else {} - 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, # type: ignore[arg-type] - ) + batch_result = 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, # type: ignore[arg-type] ) logging_obj = LiteLLMLogging( - model=batch_models[0], + model=batch_result.models[0], messages=[{"role": "user", "content": ""}], stream=False, call_type="aretrieve_batch", @@ -684,9 +682,11 @@ class CheckBatchCost: 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, ) # Record batch duration (completed_at - created_at) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 9681d64f656..20aa3c755bd 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -12,12 +12,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 +43,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, @@ -48,7 +58,7 @@ async def _handle_completed_batch( custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], model_name: str | None = None, litellm_params: dict | 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. @@ -66,10 +76,7 @@ async def _handle_completed_batch( 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 calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name) return _aggregate_batch_cost_usage_models( entries=_iter_batch_input_entries(file_content), @@ -86,19 +93,24 @@ 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 | None]: + """Classify every output line in a single pass: yields stats for a + successful line, ``None`` for a failed one (per ``_batch_response_was_successful``). + Counting failures this way avoids a second pass over a potentially huge output file.""" from litellm.cost_calculator import batch_cost_calculator for entry in entries: if not _batch_response_was_successful(entry, custom_llm_provider): + yield None continue response_body = _get_response_from_batch_job_output_file(entry, custom_llm_provider) usage = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider) @@ -123,6 +135,7 @@ def _iter_successful_output_line_stats( custom_llm_provider=custom_llm_provider, call_type=CallTypes.aretrieve_batch.value, ) + reasoning_tokens = usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None yield _BatchOutputLineStats( cost=line_cost, prompt_tokens=usage.prompt_tokens, @@ -130,6 +143,7 @@ def _iter_successful_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=reasoning_tokens or 0, model=response_model, ) @@ -139,10 +153,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(stats for stats in all_results if stats is not None) + successful_requests: Final = len(line_stats) + failed_requests: Final = len(all_results) - successful_requests cache_token_params: Final = { key: tokens @@ -156,18 +174,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. @@ -178,6 +210,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 @@ -185,12 +221,16 @@ def calculate_vertex_ai_batch_cost_and_usage( total_tokens = 0 prompt_tokens = 0 completion_tokens = 0 + successful_requests = 0 + failed_requests = 0 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 @@ -218,17 +258,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, ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a72d46e3fe8..9d035afbde2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2574,6 +2574,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 = kwargs.get("batch_successful_requests", None) + batch_failed_requests = 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 = ( @@ -2582,22 +2584,22 @@ 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 + result._hidden_params["batch_failed_requests"] = batch_failed_requests result.usage = batch_usage elif should_compute_batch_data: - ( - response_cost, - batch_usage, - batch_models, - ) = await _handle_completed_batch( + batch_result = await _handle_completed_batch( batch=result, custom_llm_provider=self.custom_llm_provider, litellm_params=self.litellm_params, ) - 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 + result._hidden_params["batch_failed_requests"] = batch_result.failed_requests + result.usage = batch_result.usage start_time, end_time, result = self._success_handler_helper_fn( start_time=start_time, @@ -5062,6 +5064,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, ) @@ -5451,6 +5455,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, ) @@ -5819,6 +5825,8 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload: additional_headers=None, litellm_overhead_time_ms=None, batch_models=None, + batch_successful_requests=None, + batch_failed_requests=None, litellm_model_name=None, usage_object=None, ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a566d491597..297fc1bd201 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3462,6 +3462,8 @@ class SpendLogsMetadata(TypedDict): status: StandardLoggingPayloadStatus proxy_server_request: str | None batch_models: list[str] | None + batch_successful_requests: int | None + batch_failed_requests: int | None error_information: StandardLoggingPayloadErrorInformation | None usage_object: dict | None model_map_information: StandardLoggingModelInformation | None diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 3146d8bccfb..f08236adab2 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -74,6 +74,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, @@ -102,6 +104,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, @@ -128,6 +132,8 @@ def _get_spend_logs_metadata( clean_metadata["user_api_key"] = _hash_api_key_for_spend_log(raw_user_api_key) 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 @@ -310,6 +316,16 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs if standard_logging_payload is not None else None ), + batch_successful_requests=( + standard_logging_payload.get("hidden_params", {}).get("batch_successful_requests", None) + if standard_logging_payload is not None + else None + ), + batch_failed_requests=( + standard_logging_payload.get("hidden_params", {}).get("batch_failed_requests", None) + if standard_logging_payload is not None + else None + ), mcp_tool_call_metadata=( standard_logging_payload["metadata"].get("mcp_tool_call_metadata", None) if standard_logging_payload is not None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 272fbabf807..dc8adf9b88e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -39,7 +39,7 @@ from pydantic import ( field_serializer, field_validator, ) -from typing_extensions import Required, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -2880,6 +2880,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 diff --git a/tests/batches_tests/test_batch_custom_pricing.py b/tests/batches_tests/test_batch_custom_pricing.py index c2159b564a8..b76b865862a 100644 --- a/tests/batches_tests/test_batch_custom_pricing.py +++ b/tests/batches_tests/test_batch_custom_pricing.py @@ -116,16 +116,16 @@ def test_aggregate_batch_cost_uses_custom_model_info(): """_aggregate_batch_cost_usage_models should thread model_info to batch_cost_calculator.""" file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)] - cost, _, _ = _aggregate_batch_cost_usage_models( + result = _aggregate_batch_cost_usage_models( entries=file_content, custom_llm_provider="openai", model_info=CUSTOM_MODEL_INFO, ) expected = (10 * 0.00125) + (5 * 0.005) - assert cost == pytest.approx( + assert result.cost == pytest.approx( expected - ), f"Expected total cost {expected}, got {cost}" + ), f"Expected total cost {expected}, got {result.cost}" @pytest.mark.parametrize("data_residency", ["eu", "us"]) @@ -164,15 +164,15 @@ async def test_calculate_batch_cost_and_usage_uses_custom_model_info(): """calculate_batch_cost_and_usage should thread model_info.""" file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)] - batch_cost, batch_usage, batch_models = await calculate_batch_cost_and_usage( + result = await calculate_batch_cost_and_usage( file_content_dictionary=file_content, custom_llm_provider="openai", model_info=CUSTOM_MODEL_INFO, ) expected = (10 * 0.00125) + (5 * 0.005) - assert batch_cost == pytest.approx( + assert result.cost == pytest.approx( expected - ), f"Expected total cost {expected}, got {batch_cost}" - assert batch_usage.prompt_tokens == 10 - assert batch_usage.completion_tokens == 5 + ), f"Expected total cost {expected}, got {result.cost}" + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 5 diff --git a/tests/batches_tests/test_batch_rate_limits.py b/tests/batches_tests/test_batch_rate_limits.py index 2c804d21ace..0baff6c17be 100644 --- a/tests/batches_tests/test_batch_rate_limits.py +++ b/tests/batches_tests/test_batch_rate_limits.py @@ -1027,7 +1027,7 @@ async def test_batch_logging_azure_credentials_regression(): with patch( "litellm.files.main.afile_content", side_effect=mock_afile_content_tracker ): - cost, usage, models = await _handle_completed_batch( + result = await _handle_completed_batch( batch=mock_batch, custom_llm_provider="azure", litellm_params=azure_credentials, @@ -1039,13 +1039,13 @@ async def test_batch_logging_azure_credentials_regression(): ], "REGRESSION: Credentials not passed through _handle_completed_batch" # Verify cost and usage were calculated - assert cost > 0, "Cost should be calculated" - assert usage.total_tokens == 40, "Usage should be calculated correctly" + assert result.cost > 0, "Cost should be calculated" + assert result.usage.total_tokens == 40, "Usage should be calculated correctly" print(" ✓ Credentials passed through full flow") - print(f" ✓ Cost: {cost}") - print(f" ✓ Usage: {usage.total_tokens} tokens") - print(f" ✓ Models: {models}") + print(f" ✓ Cost: {result.cost}") + print(f" ✓ Usage: {result.usage.total_tokens} tokens") + print(f" ✓ Models: {result.models}") # Test 4: Verify error prevention print("\n4. Testing 'Missing credentials' error prevention...") @@ -1064,7 +1064,7 @@ async def test_batch_logging_azure_credentials_regression(): "litellm.files.main.afile_content", side_effect=mock_afile_content_tracker ): try: - cost, usage, models = await _handle_completed_batch( + result = await _handle_completed_batch( batch=mock_batch, custom_llm_provider="azure", litellm_params=azure_credentials, diff --git a/tests/batches_tests/test_batches_logging_unit_tests.py b/tests/batches_tests/test_batches_logging_unit_tests.py index 62b6f5b08e4..9c26514f872 100644 --- a/tests/batches_tests/test_batches_logging_unit_tests.py +++ b/tests/batches_tests/test_batches_logging_unit_tests.py @@ -138,12 +138,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 @@ -156,11 +156,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): @@ -226,6 +226,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, @@ -236,7 +238,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( @@ -251,6 +261,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 @@ -284,7 +296,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" ) @@ -294,16 +306,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 @@ -542,9 +556,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( @@ -560,4 +584,6 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data(): # Verify computed cost data was used (not partial explicit data) assert mock_batch._hidden_params["response_cost"] == expected_cost assert mock_batch._hidden_params["batch_models"] == expected_models + assert mock_batch._hidden_params["batch_successful_requests"] == 8 + assert mock_batch._hidden_params["batch_failed_requests"] == 0 assert mock_batch.usage == expected_usage diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 72f8b87dd16..79739669ff4 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -13,6 +13,20 @@ import pytest _IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id" +def _batch_cost_result(cost, usage, models, successful_requests=1, failed_requests=0): + """Build the BatchCostUsageResult calculate_batch_cost_and_usage now returns, + for mocking it in tests that only care about cost/usage/models.""" + from litellm.batches.batch_utils import BatchCostUsageResult + + return BatchCostUsageResult( + cost=cost, + usage=usage, + models=models, + successful_requests=successful_requests, + failed_requests=failed_requests, + ) + + def _unmanaged_vertex_file_object( input_file_id="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash/abc.jsonl", status="validating", @@ -321,7 +335,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -426,7 +440,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), + return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), ), patch( "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", @@ -526,7 +540,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -656,7 +670,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -1142,7 +1156,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -1271,7 +1285,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -1529,7 +1543,7 @@ class TestUnmanagedVertexRouting: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gemini-2.5-flash"], @@ -1759,7 +1773,7 @@ class TestUnmanagedBedrockRouting: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.02, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-sonnet-4"], @@ -1951,7 +1965,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=(0.01, {"prompt_tokens": 10}, ["gpt-5.5"]), + return_value=_batch_cost_result(0.01, {"prompt_tokens": 10}, ["gpt-5.5"]), ), patch("litellm.litellm_core_utils.litellm_logging.Logging") as logging_cls, ): diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index d2074853f2b..46969dfc033 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -210,10 +210,10 @@ def test_estimate_tokens_never_zero_for_short_rows(): def test_output_models_uses_model_name_override(monkeypatch): monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) - _, _, models = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=[_success_row(model="ignored")], custom_llm_provider="openai", model_name="forced-model" ) - assert models == ["forced-model"] + assert result.models == ["forced-model"] def test_output_models_collects_from_successful_only(monkeypatch): @@ -223,15 +223,15 @@ def test_output_models_collects_from_successful_only(monkeypatch): _failed_row(model="should-be-skipped"), _success_row(model="claude-3"), ] - _, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert models == ["gpt-4o", "claude-3"] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert result.models == ["gpt-4o", "claude-3"] def test_output_models_skips_successful_without_model(monkeypatch): monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) rows = [{"response": {"status_code": 200, "body": {}}}] - _, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert models == [] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert result.models == [] # =========================================================================== # @@ -398,8 +398,8 @@ def test_total_usage_sums_successful_only(monkeypatch): _failed_row(), # excluded _success_row(usage=_usage(20, 10)), # 30 ] - _, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( 30, 15, 45, @@ -417,7 +417,7 @@ def test_total_usage_and_cost_normalize_mixed_responses_and_chat(): ) chat_row = _success_row(usage=_usage(10, 5)) - cost, usage, _ = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=[responses_row, chat_row], custom_llm_provider="openai", model_info={ @@ -426,22 +426,79 @@ def test_total_usage_and_cost_normalize_mixed_responses_and_chat(): }, ) - assert usage.prompt_tokens == 30 - assert usage.completion_tokens == 12 - assert usage.total_tokens == 42 - assert usage.cache_read_input_tokens == 3 - assert cost == pytest.approx((30 * 0.00125) + (12 * 0.005)) + assert result.usage.prompt_tokens == 30 + assert result.usage.completion_tokens == 12 + assert result.usage.total_tokens == 42 + assert result.usage.cache_read_input_tokens == 3 + assert result.cost == pytest.approx((30 * 0.00125) + (12 * 0.005)) def test_total_usage_empty_is_zero(): - cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=[], custom_llm_provider="openai") - assert cost == 0.0 - assert models == [] - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + result = bu._aggregate_batch_cost_usage_models(entries=[], custom_llm_provider="openai") + assert result.cost == 0.0 + assert result.models == [] + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( 0, 0, 0, ) + assert result.successful_requests == 0 + assert result.failed_requests == 0 + + +def test_total_usage_includes_reasoning_tokens(monkeypatch): + monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) + rows = [ + _success_row( + usage={ + "prompt_tokens": 10, + "completion_tokens": 50, + "total_tokens": 60, + "completion_tokens_details": {"reasoning_tokens": 30}, + } + ), + _success_row( + usage={ + "prompt_tokens": 5, + "completion_tokens": 20, + "total_tokens": 25, + "completion_tokens_details": {"reasoning_tokens": 8}, + } + ), + _failed_row(), # excluded, must not contribute reasoning tokens either + ] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert result.usage.completion_tokens_details is not None + assert result.usage.completion_tokens_details.reasoning_tokens == 38 + + +def test_aggregate_counts_successful_and_failed_requests(monkeypatch): + monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) + rows = [ + _success_row(usage=_usage(10, 5)), + _failed_row(), + _success_row(usage=_usage(20, 10)), + _failed_row(), + _failed_row(), + ] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert result.successful_requests == 2 + assert result.failed_requests == 3 + assert result.successful_requests + result.failed_requests == len(rows) + + +def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch): + monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 1.0) + result = bu._aggregate_batch_cost_usage_models( + entries=[_success_row(usage=_usage(10, 5))], custom_llm_provider="openai" + ) + assert isinstance(result, bu.BatchCostUsageResult) + assert (result.cost, result.models, result.successful_requests, result.failed_requests) == ( + 1.0, + ["gpt-4o"], + 1, + 0, + ) # =========================================================================== # @@ -464,10 +521,12 @@ def test_cost_from_content_completion_cost_path(monkeypatch): _success_row(usage=_usage(20, 10)), ] - total, _, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert total == 1.0 # 2 successful * 0.5 + assert result.cost == 1.0 # 2 successful * 0.5 assert len(calls) == 2 # failed row not costed + assert result.successful_requests == 2 + assert result.failed_requests == 1 def test_cost_from_content_model_info_path(monkeypatch): @@ -480,13 +539,13 @@ def test_cost_from_content_model_info_path(monkeypatch): _success_row(usage=_usage(20, 10)), ] - total, _, _ = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=rows, custom_llm_provider="openai", model_info={"input_cost_per_token": 0.0}, # type: ignore[arg-type] # truthy -> model_info path ) - assert total == pytest.approx(0.6) # 2 * (0.1 + 0.2) + assert result.cost == pytest.approx(0.6) # 2 * (0.1 + 0.2) def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch): @@ -496,11 +555,13 @@ def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch): monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.5) one_shot = (row for row in [_success_row(usage=_usage(10, 5)), _failed_row(), _success_row(usage=_usage(20, 10))]) - cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai") + result = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai") - assert cost == 1.0 - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45) - assert models == ["gpt-4o", "gpt-4o"] + assert result.cost == 1.0 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45) + assert result.models == ["gpt-4o", "gpt-4o"] + assert result.successful_requests == 2 + assert result.failed_requests == 1 # =========================================================================== # @@ -514,7 +575,13 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch): monkeypatch.setattr( bu, "calculate_vertex_ai_batch_cost_and_usage", - lambda content, model: (9.9, Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3)), + lambda content, model: bu.BatchCostUsageResult( + cost=9.9, + usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), + models=["gemini-2.0-flash-001"], + successful_requests=1, + failed_requests=0, + ), ) # generic path must NOT be taken monkeypatch.setattr( @@ -523,12 +590,12 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch): lambda **kw: pytest.fail("generic path should not run"), ) - cost, usage, models = await bu.calculate_batch_cost_and_usage( + result = await bu.calculate_batch_cost_and_usage( file_content_dictionary=[], custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001" ) - assert cost == 9.9 - assert usage.total_tokens == 3 - assert models == ["gemini-2.0-flash-001"] + assert result.cost == 9.9 + assert result.usage.total_tokens == 3 + assert result.models == ["gemini-2.0-flash-001"] @pytest.mark.asyncio @@ -542,12 +609,12 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch): lambda content, model: pytest.fail("raw vertex path should not run"), ) - cost, usage, models = await bu.calculate_batch_cost_and_usage( + result = await bu.calculate_batch_cost_and_usage( file_content_dictionary=[], custom_llm_provider="vertex_ai" ) - assert cost == 0.0 - assert usage.total_tokens == 0 - assert models == [] + assert result.cost == 0.0 + assert result.usage.total_tokens == 0 + assert result.models == [] # =========================================================================== # @@ -580,14 +647,16 @@ def test_vertex_cost_and_usage_aggregation(monkeypatch): }, ] - cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") - assert cost == pytest.approx(0.6) # 2 * (0.1 + 0.2) - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + assert result.cost == pytest.approx(0.6) # 2 * (0.1 + 0.2) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( 30, 15, 45, ) + assert result.successful_requests == 2 + assert result.failed_requests == 0 def test_vertex_cost_skips_none_response_body(monkeypatch): @@ -607,10 +676,12 @@ def test_vertex_cost_skips_none_response_body(monkeypatch): }, ] - cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") - assert cost == pytest.approx(1.0) # only one line costed - assert usage.total_tokens == 10 + assert result.cost == pytest.approx(1.0) # only one line costed + assert result.usage.total_tokens == 10 + assert result.successful_requests == 1 + assert result.failed_requests == 1 def test_vertex_usage_total_token_fallback(monkeypatch): @@ -620,8 +691,8 @@ def test_vertex_usage_total_token_fallback(monkeypatch): monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0)) responses = [{"response": {"usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 4}}}] - _, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") - assert usage.total_tokens == 12 + result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + assert result.usage.total_tokens == 12 def test_vertex_cost_error_in_line_is_swallowed(monkeypatch): @@ -644,9 +715,9 @@ def test_vertex_cost_error_in_line_is_swallowed(monkeypatch): } ] - cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") - assert cost == 0.0 - assert usage.total_tokens == 10 + result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + assert result.cost == 0.0 + assert result.usage.total_tokens == 10 # =========================================================================== # @@ -659,13 +730,11 @@ async def test_calculate_batch_cost_and_usage_orchestration(monkeypatch): rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 2.5) - cost, usage, models = await bu.calculate_batch_cost_and_usage( - file_content_dictionary=rows, custom_llm_provider="openai" - ) + result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=rows, custom_llm_provider="openai") - assert cost == 2.5 - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15) - assert models == ["gpt-4o"] + assert result.cost == 2.5 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15) + assert result.models == ["gpt-4o"] # =========================================================================== # @@ -883,16 +952,18 @@ async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monk monkeypatch.setattr(files_main, "afile_content", fake_afile_content) - cost, usage, models = await bu._handle_completed_batch( + result = await bu._handle_completed_batch( _batch("gs://litellm-bucket/output/predictions.jsonl"), custom_llm_provider="vertex_ai", litellm_params={"vertex_project": "proj-1", "vertex_location": "us-central1"}, ) - assert cost > 0 - assert cost == pytest.approx(30 * 7.5e-07 + 15 * 3.75e-06) - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45) - assert models == ["gemini-3.6-flash", "gemini-3.6-flash"] + assert result.cost > 0 + assert result.cost == pytest.approx(30 * 7.5e-07 + 15 * 3.75e-06) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45) + assert result.models == ["gemini-3.6-flash", "gemini-3.6-flash"] + assert result.successful_requests == 2 + assert result.failed_requests == 0 @pytest.mark.asyncio @@ -970,11 +1041,11 @@ async def test_handle_completed_batch_orchestration(monkeypatch): monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 3.3) - cost, usage, models = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai") + result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai") - assert cost == 3.3 - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15) - assert models == ["gpt-4o"] + assert result.cost == 3.3 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15) + assert result.models == ["gpt-4o"] @pytest.mark.asyncio @@ -991,19 +1062,25 @@ async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch) def fake_vertex_calc(content, model): seen["content"] = content seen["model"] = model - return 7.7, Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3) + return bu.BatchCostUsageResult( + cost=7.7, + usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), + models=["gemini-x"], + successful_requests=1, + failed_requests=0, + ) monkeypatch.setattr(bu, "calculate_vertex_ai_batch_cost_and_usage", fake_vertex_calc) - cost, usage, models = await bu._handle_completed_batch( + result = await bu._handle_completed_batch( _batch("gs://litellm-bucket/output/predictions.jsonl"), custom_llm_provider="vertex_ai", model_name="gemini-x", ) - assert cost == 7.7 - assert usage.total_tokens == 3 - assert models == ["gemini-x"] + assert result.cost == 7.7 + assert result.usage.total_tokens == 3 + assert result.models == ["gemini-x"] assert seen["content"] == raw_rows assert seen["model"] == "gemini-x" @@ -1105,14 +1182,14 @@ def test_bedrock_cost_uses_deployment_model_name(): "recordId": "1", "modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}}, } - cost, _, models = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=[row], custom_llm_provider="bedrock", model_name="us.anthropic.claude-sonnet-4-6", model_info={}, ) - assert cost > 0 - assert models == ["us.anthropic.claude-sonnet-4-6"] + assert result.cost > 0 + assert result.models == ["us.anthropic.claude-sonnet-4-6"] def test_anthropic_total_usage_sums_succeeded_only(monkeypatch): @@ -1124,8 +1201,10 @@ def test_anthropic_total_usage_sums_succeeded_only(monkeypatch): _anthropic_errored_row(), _anthropic_succeeded_row(usage=_anthropic_usage(20, 10, cache_read=100)), ] - _, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (130, 15, 145) + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (130, 15, 145) + assert result.successful_requests == 2 + assert result.failed_requests == 1 def test_anthropic_total_usage_aggregates_cache_token_details(monkeypatch): @@ -1137,11 +1216,11 @@ def test_anthropic_total_usage_aggregates_cache_token_details(monkeypatch): _anthropic_errored_row(), _anthropic_succeeded_row(usage=_anthropic_usage(50, 20, cache_creation=300, cache_read=700)), ] - _, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") - assert usage.prompt_tokens_details.cached_tokens == 8700 - assert usage.prompt_tokens_details.cache_creation_tokens == 2300 - assert usage.cache_read_input_tokens == 8700 - assert usage.cache_creation_input_tokens == 2300 + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") + assert result.usage.prompt_tokens_details.cached_tokens == 8700 + assert result.usage.prompt_tokens_details.cache_creation_tokens == 2300 + assert result.usage.cache_read_input_tokens == 8700 + assert result.usage.cache_creation_input_tokens == 2300 def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch): @@ -1152,9 +1231,9 @@ def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch): "response": {"status_code": 200, "body": {"model": "gpt-5.2", "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}}, } ] - _, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15) - assert usage.prompt_tokens_details is None + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15) + assert result.usage.prompt_tokens_details is None def test_anthropic_cost_applies_batch_discount_and_cache_pricing(): @@ -1165,14 +1244,14 @@ def test_anthropic_cost_applies_batch_discount_and_cache_pricing(): _anthropic_errored_row(), ] - total, _, _ = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=rows, custom_llm_provider="anthropic", model_info=_ANTHROPIC_MODEL_INFO, # type: ignore[arg-type] ) expected_half_price = (1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3.75e-6 + 200 * 15e-6) / 2 - assert total == pytest.approx(expected_half_price) + assert result.cost == pytest.approx(expected_half_price) def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatch): @@ -1191,11 +1270,9 @@ def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatc lambda **kw: pytest.fail("anthropic rows must not go through completion_cost"), ) - total, _, _ = bu._aggregate_batch_cost_usage_models( - entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic" - ) + result = bu._aggregate_batch_cost_usage_models(entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic") - assert total == pytest.approx(0.3) + assert result.cost == pytest.approx(0.3) assert seen[0]["model"] == "claude-sonnet-4-5-20250929" assert seen[0]["custom_llm_provider"] == "anthropic" assert seen[0]["usage"].prompt_tokens == 10 @@ -1209,8 +1286,8 @@ def test_anthropic_batch_models_collected_from_succeeded_rows(monkeypatch): _anthropic_succeeded_row(model="claude-sonnet-4-5-20250929"), _anthropic_errored_row(), ] - _, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") - assert models == ["claude-sonnet-4-5-20250929"] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") + assert result.models == ["claude-sonnet-4-5-20250929"] @pytest.mark.asyncio @@ -1220,16 +1297,16 @@ async def test_calculate_batch_cost_and_usage_anthropic_end_to_end(): _anthropic_errored_row(), ] - cost, usage, models = await bu.calculate_batch_cost_and_usage( + result = await bu.calculate_batch_cost_and_usage( file_content_dictionary=rows, custom_llm_provider="anthropic", model_name="claude-sonnet-4-5", model_info=_ANTHROPIC_MODEL_INFO, # type: ignore[arg-type] ) - assert cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2) - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (11000, 200, 11200) - assert models == ["claude-sonnet-4-5"] + assert result.cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (11000, 200, 11200) + assert result.models == ["claude-sonnet-4-5"] def test_extract_credentials_forwards_the_trusted_model_credential_snapshot(): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py index ac79c183ca3..1d2d7d4d5c3 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py @@ -478,14 +478,14 @@ class TestVertexAIBatchPassthroughHandler: } ] - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( vertex_ai_batch_responses, model_name="gemini-2.0-flash-001" ) - assert usage.total_tokens == 15 - assert usage.prompt_tokens == 10 - assert usage.completion_tokens == 5 - assert total_cost > 0, "batch_cost_calculator should return a non-zero cost" + assert result.usage.total_tokens == 15 + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 5 + assert result.cost > 0, "batch_cost_calculator should return a non-zero cost" def test_batch_response_transformation(self): """Test transformation of Vertex AI batch responses to OpenAI format""" @@ -664,14 +664,14 @@ class TestVertexAIBatchCostCalculation: }, ] - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( responses, model_name="gemini-2.0-flash-001" ) - assert usage.prompt_tokens == 18 - assert usage.completion_tokens == 8 - assert usage.total_tokens == 26 - assert total_cost > 0, "batch_cost_calculator should return a non-zero cost" + assert result.usage.prompt_tokens == 18 + assert result.usage.completion_tokens == 8 + assert result.usage.total_tokens == 26 + assert result.cost > 0, "batch_cost_calculator should return a non-zero cost" def test_should_skip_responses_with_null_response_body(self): """Failed lines (response: None) are skipped without error.""" @@ -699,27 +699,29 @@ class TestVertexAIBatchCostCalculation: }, ] - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( responses, model_name="gemini-2.0-flash-001" ) - assert usage.prompt_tokens == 18 - assert usage.completion_tokens == 8 - assert usage.total_tokens == 26 - assert total_cost > 0 + assert result.usage.prompt_tokens == 18 + assert result.usage.completion_tokens == 8 + assert result.usage.total_tokens == 26 + assert result.cost > 0 + assert result.successful_requests == 2 + assert result.failed_requests == 1 def test_should_return_zeros_for_empty_response_list(self): """Empty input → zero cost and zero usage.""" from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( [], model_name="gemini-2.0-flash-001" ) - assert total_cost == 0.0 - assert usage.total_tokens == 0 - assert usage.prompt_tokens == 0 - assert usage.completion_tokens == 0 + assert result.cost == 0.0 + assert result.usage.total_tokens == 0 + assert result.usage.prompt_tokens == 0 + assert result.usage.completion_tokens == 0 def test_should_handle_missing_usage_metadata_gracefully(self): """Response without usageMetadata → 0 tokens, 0 cost for that line.""" @@ -729,13 +731,13 @@ class TestVertexAIBatchCostCalculation: {"response": {"candidates": [{"content": {"parts": [{"text": "hi"}]}}]}}, ] - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( responses, model_name="gemini-2.0-flash-001" ) - assert usage.prompt_tokens == 0 - assert usage.completion_tokens == 0 - assert usage.total_tokens == 0 + assert result.usage.prompt_tokens == 0 + assert result.usage.completion_tokens == 0 + assert result.usage.total_tokens == 0 @pytest.mark.asyncio async def test_openai_shaped_output_records_nonzero_cost_and_usage(self): @@ -813,7 +815,7 @@ class TestVertexAIBatchCostCalculation: try: litellm.disable_vertex_batch_output_transformation = False - cost, usage, _ = await calculate_batch_cost_and_usage( + result = await calculate_batch_cost_and_usage( file_content_dictionary=openai_shaped_responses, custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001", @@ -822,17 +824,17 @@ class TestVertexAIBatchCostCalculation: litellm.disable_vertex_batch_output_transformation = original_flag assert ( - usage.prompt_tokens == 18 - ), f"expected 18 prompt tokens, got {usage.prompt_tokens}" + result.usage.prompt_tokens == 18 + ), f"expected 18 prompt tokens, got {result.usage.prompt_tokens}" assert ( - usage.completion_tokens == 8 - ), f"expected 8 completion tokens, got {usage.completion_tokens}" + result.usage.completion_tokens == 8 + ), f"expected 8 completion tokens, got {result.usage.completion_tokens}" assert ( - usage.total_tokens == 26 - ), f"expected 26 total tokens, got {usage.total_tokens}" + result.usage.total_tokens == 26 + ), f"expected 26 total tokens, got {result.usage.total_tokens}" assert ( - cost > 0 - ), f"expected non-zero cost for completed Vertex batch, got {cost}" + result.cost > 0 + ), f"expected non-zero cost for completed Vertex batch, got {result.cost}" @pytest.mark.asyncio async def test_raw_vertex_output_still_works_when_transformation_disabled(self): @@ -865,7 +867,7 @@ class TestVertexAIBatchCostCalculation: try: litellm.disable_vertex_batch_output_transformation = True - cost, usage, _ = await calculate_batch_cost_and_usage( + result = await calculate_batch_cost_and_usage( file_content_dictionary=raw_vertex_responses, custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001", @@ -873,7 +875,7 @@ class TestVertexAIBatchCostCalculation: finally: litellm.disable_vertex_batch_output_transformation = original_flag - assert usage.prompt_tokens == 10 - assert usage.completion_tokens == 5 - assert usage.total_tokens == 15 - assert cost > 0, "raw Vertex shape should also produce non-zero cost" + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 5 + assert result.usage.total_tokens == 15 + assert result.cost > 0, "raw Vertex shape should also produce non-zero cost" From 2bfa1613b4837db3bc01e48297588a123c0c08e9 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 17 Aug 2026 11:48:04 -0400 Subject: [PATCH 006/111] fix(batches): satisfy LIT010/LIT011/reportPrivateUsage budgets for new fields Suppress the mutation/private-access lints the new batch_successful_requests/batch_failed_requests plumbing triggers, matching the existing suppressed pattern already used for response_cost/batch_models on the same lines. --- litellm/batches/batch_utils.py | 4 ++-- litellm/litellm_core_utils/litellm_logging.py | 14 +++++++------- litellm/proxy/_types.py | 6 +++--- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 20aa3c755bd..6f6de18b04e 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -221,8 +221,8 @@ def calculate_vertex_ai_batch_cost_and_usage( total_tokens = 0 prompt_tokens = 0 completion_tokens = 0 - successful_requests = 0 - failed_requests = 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: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9d035afbde2..c795b181064 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2574,8 +2574,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 = kwargs.get("batch_successful_requests", None) - batch_failed_requests = kwargs.get("batch_failed_requests", 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 = ( @@ -2584,12 +2584,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 - result._hidden_params["batch_failed_requests"] = batch_failed_requests + 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: - batch_result = await _handle_completed_batch( + batch_result: Final = await _handle_completed_batch( batch=result, custom_llm_provider=self.custom_llm_provider, litellm_params=self.litellm_params, @@ -2597,8 +2597,8 @@ class Logging(LiteLLMLoggingBaseClass): 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 - result._hidden_params["batch_failed_requests"] = batch_result.failed_requests + 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( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 297fc1bd201..cbc8c003868 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -15,7 +15,7 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import NotRequired, Required, TypedDict +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict from litellm._uuid import uuid from litellm.constants import DEFAULT_STAGGER_WINDOW_SECONDS, MCP_STDIO_ALLOWED_COMMANDS @@ -3462,8 +3462,8 @@ class SpendLogsMetadata(TypedDict): status: StandardLoggingPayloadStatus proxy_server_request: str | None batch_models: list[str] | None - batch_successful_requests: int | None - batch_failed_requests: int | None + batch_successful_requests: ReadOnly[int | None] + batch_failed_requests: ReadOnly[int | None] error_information: StandardLoggingPayloadErrorInformation | None usage_object: dict | None model_map_information: StandardLoggingModelInformation | None From 02cf319b483d8c6961b5464f0a72875c0ee6ae4a Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 17 Aug 2026 14:58:14 -0400 Subject: [PATCH 007/111] fix(batches): count failures reported only in the batch error file Live verification against a real OpenAI batch showed per-request failures (e.g. a rejected param) land in error_file_id, never in the output file, so failed_requests silently undercounted them (0 instead of the real 1). _handle_completed_batch now also fetches error_file_id when present and folds its line count into failed_requests. --- litellm/batches/batch_utils.py | 124 +++++++++++++----- .../test_litellm/batches/test_batch_utils.py | 63 +++++++++ 2 files changed, 152 insertions(+), 35 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 6f6de18b04e..4889be3c0f6 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -1,6 +1,7 @@ import json from collections.abc import Iterable, Iterator from dataclasses import dataclass +from dataclasses import replace as dataclasses_replace from typing import Any, Final, Literal import litellm @@ -70,18 +71,28 @@ async def _handle_completed_batch( litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.) """ file_content = await _fetch_batch_output_file_content(batch, custom_llm_provider, litellm_params=litellm_params) + error_file_failed_requests: Final = await _count_error_file_failed_requests( + batch, custom_llm_provider=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) - ): - return calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name) + 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_input_entries(file_content), + custom_llm_provider=custom_llm_provider, + model_name=model_name, + ) + ) - return _aggregate_batch_cost_usage_models( - entries=_iter_batch_input_entries(file_content), - custom_llm_provider=custom_llm_provider, - model_name=model_name, + 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 ) @@ -280,6 +291,50 @@ def calculate_vertex_ai_batch_cost_and_usage( ) +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 + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + ) + + resolved_file_id = file_id + is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(file_id) + if is_base64_unified_file_id: + try: + resolved_file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] + verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", resolved_file_id) + except (IndexError, AttributeError) as e: + verbose_logger.error( + "Failed to extract LLM output file ID from unified file ID: %s, error: %s", file_id, e + ) + + # Build kwargs for afile_content with credentials from litellm_params + file_content_kwargs: Final = { + "file_id": resolved_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", @@ -294,37 +349,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 - from litellm.proxy.openai_files_endpoints.common_utils import ( - _is_base64_encoded_unified_file_id, - ) - if batch.output_file_id is None: raise ValueError("Output file id is None cannot retrieve file content") - file_id = batch.output_file_id - is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(file_id) - if is_base64_unified_file_id: - try: - file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] - verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", file_id) - except (IndexError, AttributeError) as e: - verbose_logger.error( - "Failed to extract LLM output file ID from unified file ID: %s, error: %s", batch.output_file_id, e - ) + return await _fetch_batch_managed_file_content( + batch.output_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params + ) - # Build kwargs for afile_content with credentials from litellm_params - file_content_kwargs: Final = { - "file_id": file_id, - "custom_llm_provider": custom_llm_provider, - } - # Extract and add credentials for file access - credentials: Final = _extract_file_access_credentials(litellm_params) - file_content_kwargs.update(credentials) +async def _count_error_file_failed_requests( + batch: Batch, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], + litellm_params: dict | None, +) -> int: + """Count failed requests reported only in the batch's separate error file. - _file_content: Final = await afile_content(**file_content_kwargs) - return _file_content.content + OpenAI-shaped batch providers write successful lines to ``output_file_id`` + and per-request failures (e.g. a rejected param) to a distinct + ``error_file_id`` - they never appear in the output file at all, so + counting failures from the output file alone silently undercounts them. + """ + if batch.error_file_id is None: + return 0 + try: + error_file_content = await _fetch_batch_managed_file_content( + batch.error_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params + ) + except Exception as e: # noqa: BLE001 # a failed/missing error file must not abort cost tracking for the batch + verbose_logger.debug("Failed to fetch batch error file %s: %s", batch.error_file_id, e) + return 0 + return sum(1 for _ in _iter_batch_input_lines(error_file_content)) def _extract_file_access_credentials(litellm_params: dict | None) -> dict: diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 46969dfc033..06f42ce7b51 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1048,6 +1048,69 @@ async def test_handle_completed_batch_orchestration(monkeypatch): assert result.models == ["gpt-4o"] +@pytest.mark.asyncio +async def test_handle_completed_batch_counts_error_file_failures(monkeypatch): + """Regression test: OpenAI writes per-request failures (e.g. a rejected param) + to a separate error_file_id, never into the output file - so failed_requests + must include them or it silently undercounts real batch failures.""" + from litellm.types.llms.openai import Batch + + rows = [_success_row(model="gpt-5-mini", usage=_usage(24, 107))] + error_rows = [ + { + "id": "batch_req_err1", + "custom_id": "req-2-bad", + "response": {"status_code": 400, "body": {"error": {"message": "Invalid 'temperature'"}}}, + "error": None, + } + ] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return _vertex_jsonl(rows) + + async def fake_afile_content(**kw): + return type("R", (), {"content": _vertex_jsonl(error_rows)})() + + import litellm.files.main as files_main + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) + + batch = Batch( + id="b", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + output_file_id="of", + error_file_id="ef", + ) + + result = await bu._handle_completed_batch(batch, custom_llm_provider="openai") + + assert result.successful_requests == 1 + assert result.failed_requests == 1 + + +@pytest.mark.asyncio +async def test_handle_completed_batch_no_error_file_id_reports_zero_error_failures(monkeypatch): + rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return _vertex_jsonl(rows) + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) + monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) + + result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai") + + assert result.successful_requests == 1 + assert result.failed_requests == 0 + + @pytest.mark.asyncio async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch): raw_rows = [{"response": {"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 2}}}] From 682032dd1b23ba49aab8ba486b44716805619ec1 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 19:29:25 -0400 Subject: [PATCH 008/111] fix(batches): fix CI failures from lint format and merged upstream guard - ruff format litellm/batches/batch_utils.py - reconcile the upstream output_file_id=None guard (merged in from litellm_internal_staging) with BatchCostUsageResult, and count that batch's error_file_id failures instead of always reporting 0 - fix test_handle_completed_batch_no_output_file_is_zero's tuple unpacking, which predated the BatchCostUsageResult refactor - commit the batch_successful_requests/batch_failed_requests fixture fix to test_spend_management_endpoints.py that was left uncommitted --- litellm/batches/batch_utils.py | 14 ++++++++++---- tests/test_litellm/batches/test_batch_utils.py | 10 ++++++---- .../test_spend_management_endpoints.py | 6 +++--- 3 files changed, 19 insertions(+), 11 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 8a42efd6dc5..4eb9c7a5dfa 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -79,7 +79,15 @@ 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=[], + 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) error_file_failed_requests: Final = await _count_error_file_failed_requests( @@ -328,9 +336,7 @@ async def _fetch_batch_managed_file_content( resolved_file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", resolved_file_id) except (IndexError, AttributeError) as e: - verbose_logger.error( - "Failed to extract LLM output file ID from unified file ID: %s, error: %s", file_id, e - ) + verbose_logger.error("Failed to extract LLM output file ID from unified file ID: %s, error: %s", file_id, e) # Build kwargs for afile_content with credentials from litellm_params file_content_kwargs: Final = { diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index daf3b91a26b..3a69911dfe6 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1126,11 +1126,13 @@ async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch): monkeypatch.setattr(bu, "_fetch_batch_output_file_content", _must_not_fetch) - cost, usage, models = await bu._handle_completed_batch(_batch(None), custom_llm_provider="openai") + result = await bu._handle_completed_batch(_batch(None), custom_llm_provider="openai") - assert cost == 0.0 - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (0, 0, 0) - assert models == [] + assert result.cost == 0.0 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (0, 0, 0) + assert result.models == [] + assert result.successful_requests == 0 + assert result.failed_requests == 0 @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 7052e050806..b8ed2b04b9d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -2633,7 +2633,7 @@ class TestSpendLogsPayload: "model": "gpt-4o", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', + "metadata": '{"applied_guardrails": [], "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, @@ -2729,7 +2729,7 @@ class TestSpendLogsPayload: "model": "claude-4-sonnet-20250514", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "metadata": '{"applied_guardrails": [], "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', "cache_key": "Cache OFF", "spend": 0.01383, "total_tokens": 2598, @@ -2823,7 +2823,7 @@ class TestSpendLogsPayload: "model": "claude-4-sonnet-20250514", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "metadata": '{"applied_guardrails": [], "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', "cache_key": "Cache OFF", "spend": 0.01383, "total_tokens": 2598, From ff833a5872cc167a7166a58fdf78aeb2d175b570 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 19:37:07 -0400 Subject: [PATCH 009/111] fix(batches): address Greptile type-discipline feedback Add Final to the reasoning_tokens local var and type the _batch_cost_result test helper's parameters, per review feedback on PR #37208. --- litellm/batches/batch_utils.py | 4 +- .../proxy_unit_tests/test_check_batch_cost.py | 519 ++++++------------ 2 files changed, 167 insertions(+), 356 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 4eb9c7a5dfa..057c9978879 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -165,7 +165,9 @@ def _classify_output_line_stats( custom_llm_provider=custom_llm_provider, call_type=CallTypes.aretrieve_batch.value, ) - reasoning_tokens = usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None + reasoning_tokens: Final = ( + usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None + ) yield _BatchOutputLineStats( cost=line_cost, prompt_tokens=usage.prompt_tokens, diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 139bf583e6a..ff2dce498f0 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -6,14 +6,24 @@ Vertex (raw gs:// input_file_id) and Bedrock (raw s3:// input_file_id, ARN unified_object_id) batches with no managed unified id. """ +from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch import pytest +if TYPE_CHECKING: + from litellm.batches.batch_utils import BatchCostUsageResult + _IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id" -def _batch_cost_result(cost, usage, models, successful_requests=1, failed_requests=0): +def _batch_cost_result( + cost: float, + usage: dict, + models: list[str], + successful_requests: int = 1, + failed_requests: int = 0, +) -> "BatchCostUsageResult": """Build the BatchCostUsageResult calculate_batch_cost_and_usage now returns, for mocking it in tests that only care about cost/usage/models.""" from litellm.batches.batch_utils import BatchCostUsageResult @@ -90,9 +100,7 @@ class TestCheckBatchCost: return MagicMock() @pytest.fixture - def check_batch_cost_instance( - self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router - ): + def check_batch_cost_instance(self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router): from litellm_enterprise.proxy.common_utils.check_batch_cost import ( CheckBatchCost, ) @@ -104,23 +112,15 @@ class TestCheckBatchCost: ) @pytest.mark.asyncio - async def test_cleanup_scoped_to_batch_file_purpose( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_cleanup_scoped_to_batch_file_purpose(self, check_batch_cost_instance, mock_prisma_client): """_cleanup_stale_managed_objects scopes its update to file_purpose='batch' only.""" - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) # Return empty so the main poll loop exits immediately - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) await check_batch_cost_instance.check_batch_cost() - calls = ( - mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list - ) + calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list stale_call = calls[0] assert stale_call[1]["data"] == {"status": "stale_expired"} where = stale_call[1]["where"] @@ -129,9 +129,7 @@ class TestCheckBatchCost: assert "created_at" in where @pytest.mark.asyncio - async def test_startup_probe_confirms_batch_processed_support( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_startup_probe_confirms_batch_processed_support(self, check_batch_cost_instance, mock_prisma_client): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) await check_batch_cost_instance.confirm_batch_processed_support() @@ -142,9 +140,7 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True @pytest.mark.asyncio - async def test_startup_probe_marks_column_absent( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_startup_probe_marks_column_absent(self, check_batch_cost_instance, mock_prisma_client): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( side_effect=Exception("column batch_processed does not exist") ) @@ -168,18 +164,12 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True @pytest.mark.asyncio - async def test_find_many_uses_pagination_and_excludes_stale( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_find_many_uses_pagination_and_excludes_stale(self, check_batch_cost_instance, mock_prisma_client): """find_many is called with take, order, and all terminal statuses excluded.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) await check_batch_cost_instance.check_batch_cost() @@ -205,9 +195,7 @@ class TestCheckBatchCost: """Falls back to query without batch_processed when primary query raises.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) # First find_many (primary query) raises with a schema error; second (fallback) returns empty mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( side_effect=[Exception("column batch_processed does not exist"), []] @@ -215,9 +203,7 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - calls = ( - mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list - ) + calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list assert len(calls) == 2 fallback_where = calls[1][1]["where"] assert "batch_processed" not in fallback_where @@ -228,32 +214,20 @@ class TestCheckBatchCost: assert check_batch_cost_instance.batch_processed_support_confirmed is False @pytest.mark.asyncio - async def test_column_absence_cached_across_cycles( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_column_absence_cached_across_cycles(self, check_batch_cost_instance, mock_prisma_client): """After column absence is discovered, subsequent cycles skip the primary query entirely.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) # Simulate column already known absent from a previous cycle check_batch_cost_instance._has_batch_processed_column = False - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) await check_batch_cost_instance.check_batch_cost() # Only one find_many call — the fallback directly, no primary query attempt - assert ( - mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 - ) - fallback_where = ( - mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1][ - "where" - ] - ) + assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 + fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"] assert "batch_processed" not in fallback_where @pytest.mark.asyncio @@ -267,13 +241,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-fallback-1" @@ -282,22 +252,16 @@ class TestCheckBatchCost: # Simulate column already known absent (e.g. discovered on a previous cycle) check_batch_cost_instance._has_batch_processed_column = False - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) # Build a fake batch response whose status triggers the completion branch mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -345,9 +309,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -356,15 +318,11 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() # The update must have been called — this is the core assertion. - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "Expected update() to be called exactly once for the completed job" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] - assert ( - "batch_processed" not in update_data - ), "update() must NOT include batch_processed when column is absent" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "Expected update() to be called exactly once for the completed job" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert "batch_processed" not in update_data, "update() must NOT include batch_processed when column is absent" assert update_data["status"] == "complete" @pytest.mark.asyncio @@ -440,7 +398,9 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), + return_value=_batch_cost_result( + 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"] + ), ), patch( "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", @@ -458,9 +418,9 @@ class TestCheckBatchCost: passed_kwargs = mock_afile_content.await_args[1] snapshot = passed_kwargs.get("_litellm_internal_model_credentials") assert snapshot is not None, "cost poller must pass the trusted credential snapshot" - assert isinstance( - snapshot, MappingProxyType - ), "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" + assert isinstance(snapshot, MappingProxyType), ( + "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" + ) assert snapshot["s3_bucket_name"] == "configured-batch-bucket" @pytest.mark.asyncio @@ -474,13 +434,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-primary-1" @@ -488,21 +444,15 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -550,9 +500,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -560,15 +508,13 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "Expected update() to be called exactly once for the completed job" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] - assert ( - update_data["batch_processed"] is True - ), "update() must include batch_processed=True when column is present" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "Expected update() to be called exactly once for the completed job" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert update_data["batch_processed"] is True, ( + "update() must include batch_processed=True when column is present" + ) assert update_data["status"] == "complete" @pytest.mark.asyncio @@ -712,22 +658,16 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-anthropic-1" mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = "completed" @@ -761,9 +701,9 @@ class TestCheckBatchCost: ): await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 - ), "a failed cost tracking attempt must not mark the job processed" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0, ( + "a failed cost tracking attempt must not mark the job processed" + ) @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["failed", "expired", "cancelled"]) @@ -780,13 +720,9 @@ class TestCheckBatchCost: """ import base64 - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-terminal-1" @@ -796,31 +732,25 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = terminal_status mock_response.output_file_id = None - mock_response.model_dump_json.return_value = ( - f'{{"id":"batch-1","status":"{terminal_status}"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), f"Expected update() to be called exactly once for a {terminal_status} job" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + f"Expected update() to be called exactly once for a {terminal_status} job" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["status"] == terminal_status - assert ( - update_data["batch_processed"] is True - ), "terminal-status update() must set batch_processed=True so polling stops" + assert update_data["batch_processed"] is True, ( + "terminal-status update() must set batch_processed=True so polling stops" + ) @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["failed", "cancelled"]) @@ -855,13 +785,9 @@ class TestCheckBatchCost: f"litellm_proxy:application/octet-stream;unified_id,u-2;llm_output_file_id,{raw_error_file_id}".encode() ).decode() - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) input_file_row = MagicMock() input_file_row.unified_file_id = unified_input_file_id @@ -871,9 +797,7 @@ class TestCheckBatchCost: return input_file_row return None - mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( - side_effect=find_managed_file - ) + mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(side_effect=find_managed_file) mock_job = MagicMock() mock_job.id = "job-terminal-mint-1" @@ -882,9 +806,7 @@ class TestCheckBatchCost: mock_job.team_id = "team-1" check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) response = LiteLLMBatch( id="batch-456", @@ -902,9 +824,7 @@ class TestCheckBatchCost: mock_hook = MagicMock() mock_hook.get_unified_output_file_id.side_effect = [unified_error_file_id] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( - mock_hook - ) + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook await check_batch_cost_instance.check_batch_cost() @@ -955,13 +875,9 @@ class TestCheckBatchCost: import base64 from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-completed-no-output-1" @@ -971,24 +887,18 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = completed_status mock_response.output_file_id = None mock_response.error_file_id = "file-error-123" - mock_response.model_dump_json.return_value = ( - f'{{"id":"batch-1","status":"{completed_status}"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{completed_status}"}}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) # Billing reads credentials off the router; if it is touched we billed a batch # that has no output, which is the behaviour this test guards against. - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) with patch( "litellm.files.main.afile_content", @@ -996,22 +906,18 @@ class TestCheckBatchCost: ) as mock_afile_content: await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "a completed batch with no output file must be marked processed exactly once" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "a completed batch with no output file must be marked processed exactly once" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["status"] == completed_status - assert ( - update_data["batch_processed"] is True - ), "completed-without-output update() must set batch_processed=True so polling stops" - assert ( - mock_afile_content.await_count == 0 - ), "a batch with no output file must not be billed" - assert ( - mock_llm_router.get_deployment_credentials_with_provider.call_count == 0 - ), "a batch with no output file must not enter the cost-tracking path" + assert update_data["batch_processed"] is True, ( + "completed-without-output update() must set batch_processed=True so polling stops" + ) + assert mock_afile_content.await_count == 0, "a batch with no output file must not be billed" + assert mock_llm_router.get_deployment_credentials_with_provider.call_count == 0, ( + "a batch with no output file must not enter the cost-tracking path" + ) @pytest.mark.asyncio async def test_non_terminal_status_left_unprocessed( @@ -1022,9 +928,7 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_job = MagicMock() @@ -1032,9 +936,7 @@ class TestCheckBatchCost: mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = "in_progress" @@ -1060,9 +962,9 @@ class TestCheckBatchCost: ): await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 - ), "a non-terminal batch must not be written back (would stop polling prematurely)" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0, ( + "a non-terminal batch must not be written back (would stop polling prematurely)" + ) @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["expired", "cancelled", "failed"]) @@ -1079,13 +981,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-terminal-with-output-1" @@ -1093,21 +991,15 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = terminal_status mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = ( - f'{{"id":"batch-1","status":"{terminal_status}"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -1155,9 +1047,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1165,20 +1055,16 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert ( - mock_afile_content.await_count == 1 - ), f"{terminal_status} batch with an output file must fetch results and be billed" - mock_logging_obj.async_success_handler.assert_awaited_once() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + assert mock_afile_content.await_count == 1, ( + f"{terminal_status} batch with an output file must fetch results and be billed" ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + mock_logging_obj.async_success_handler.assert_awaited_once() + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["batch_processed"] is True - assert ( - update_data["status"] == terminal_status - ), f"billed {terminal_status} batch must keep its real terminal status in the DB" + assert update_data["status"] == terminal_status, ( + f"billed {terminal_status} batch must keep its real terminal status in the DB" + ) @pytest.mark.asyncio async def test_terminal_batch_with_missing_output_file_is_retired_unbilled( @@ -1195,13 +1081,9 @@ class TestCheckBatchCost: from litellm.exceptions import NotFoundError - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-output-gone-1" @@ -1211,23 +1093,17 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) missing_output_file_id = "gs://batch-out/job-1/predictions.jsonl" mock_response = MagicMock() mock_response.status = "failed" mock_response.output_file_id = missing_output_file_id mock_response.error_file_id = None - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"failed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"failed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) with ( patch( @@ -1248,12 +1124,10 @@ class TestCheckBatchCost: assert mock_afile_content.await_count == 1 mock_calculate.assert_not_awaited() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "a terminal batch with a 404ing output file must be retired, not retried forever" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "a terminal batch with a 404ing output file must be retired, not retried forever" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["status"] == "failed" assert update_data["batch_processed"] is True @@ -1266,13 +1140,9 @@ class TestCheckBatchCost: Without this, GET /batches/{id} returns a raw file ID that cannot be routed through the proxy, causing API_KEY errors when clients call GET /files/{id}/content. """ - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-raw-file-1" @@ -1281,9 +1151,7 @@ class TestCheckBatchCost: mock_job.team_id = None check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) raw_output_file_id = "file-batch-output-abc123" raw_error_file_id = "file-batch-error-xyz456" @@ -1294,14 +1162,10 @@ class TestCheckBatchCost: mock_response.status = "completed" mock_response.output_file_id = raw_output_file_id mock_response.error_file_id = raw_error_file_id - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "azure" @@ -1316,9 +1180,7 @@ class TestCheckBatchCost: fake_managed_error_id, ] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( - mock_hook - ) + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook mock_file_content = MagicMock() mock_file_content.content = b'{"id":"req-1"}' @@ -1361,9 +1223,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-5-mini", "azure", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1414,9 +1274,7 @@ class TestUnmanagedVertexRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = "8823717160934178816" - job.file_object = ( - file_object if file_object is not None else _unmanaged_vertex_file_object() - ) + job.file_object = file_object if file_object is not None else _unmanaged_vertex_file_object() return job def test_flag_off_skips_unmanaged_id_unchanged(self): @@ -1454,9 +1312,7 @@ class TestUnmanagedVertexRouting: assert result == ("deploy-1", "8823717160934178816") # bare model name (trailing GCS segment), not the full publishers/.. path - router.resolve_model_name_from_model_id.assert_called_once_with( - "gemini-2.5-flash" - ) + router.resolve_model_name_from_model_id.assert_called_once_with("gemini-2.5-flash") router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): @@ -1476,9 +1332,7 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_uses_later_vertex_deployment_with_matching_suffix(self): router = MagicMock() @@ -1526,9 +1380,7 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_non_gcs_input_is_not_unmanaged_vertex(self): """Flag on, but input_file_id is not a gs:// publishers path: treat as unroutable, @@ -1536,9 +1388,7 @@ class TestUnmanagedVertexRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job( - file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123") - ) + job = self._job(file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123")) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1562,9 +1412,7 @@ class TestUnmanagedVertexRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = ( - '{"id":"8823717160934178816","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}' router.aretrieve_batch = AsyncMock(return_value=mock_response) router.get_deployment_credentials_with_provider = MagicMock( return_value={"vertex_project": "p", "vertex_location": "us-central1"} @@ -1586,9 +1434,7 @@ class TestUnmanagedVertexRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[self._job()] - ) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1619,9 +1465,7 @@ class TestUnmanagedVertexRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gemini-2.5-flash", "vertex_ai", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1662,9 +1506,7 @@ class TestUnmanagedBedrockRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = self._ARN - job.file_object = ( - file_object if file_object is not None else _unmanaged_bedrock_file_object() - ) + job.file_object = file_object if file_object is not None else _unmanaged_bedrock_file_object() return job def _bedrock_deployment(self): @@ -1719,9 +1561,7 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_matches_deployment_despite_colon_dash_mismatch(self): """The S3 object key has ':' replaced with '-' (e.g. 'v1-0'), but the configured @@ -1759,9 +1599,7 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_non_s3_input_is_not_unmanaged_bedrock(self): """Flag on, but input_file_id is not a litellm-bedrock-files- s3:// key: treat as @@ -1769,9 +1607,7 @@ class TestUnmanagedBedrockRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job( - file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123") - ) + job = self._job(file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123")) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1794,13 +1630,9 @@ class TestUnmanagedBedrockRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = ( - f'{{"id":"{self._ARN}","status":"completed"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}' router.aretrieve_batch = AsyncMock(return_value=mock_response) - router.get_deployment_credentials_with_provider = MagicMock( - return_value={"aws_region_name": "us-east-1"} - ) + router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"}) deployment = self._bedrock_deployment() deployment.model_name = "claude-sonnet-4" @@ -1816,9 +1648,7 @@ class TestUnmanagedBedrockRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[self._job()] - ) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1849,9 +1679,7 @@ class TestUnmanagedBedrockRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("claude-sonnet-4", "bedrock", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1975,9 +1803,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: ) router = MagicMock() - router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) deployment = MagicMock() deployment.litellm_params.custom_llm_provider = "azure" deployment.litellm_params.model = "azure/gpt-5.5" @@ -1986,8 +1812,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: router.get_deployment = MagicMock(return_value=deployment) hook = MagicMock() - hook.get_unified_output_file_id = ( - lambda output_file_id, model_id, model_name: _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( + hook.get_unified_output_file_id = lambda output_file_id, model_id, model_name: ( + _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( None, output_file_id=output_file_id, model_id=model_id, model_name=model_name ) ) @@ -2056,9 +1882,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: get_models_from_unified_file_id, ) - output_file_id = await self._run( - self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) - ) + output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] @@ -2072,9 +1896,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: _extract_models_from_managed_resource_id, ) - output_file_id = await self._run( - self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) - ) + output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) models = _extract_models_from_managed_resource_id(output_file_id, "file_id", None) assert models == [self._PUBLIC_MODEL_GROUP] @@ -2082,9 +1904,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: await can_key_call_model( model=models[0], llm_model_list=None, - valid_token=UserAPIKeyAuth( - api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP] - ), + valid_token=UserAPIKeyAuth(api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP]), llm_router=None, ) is True @@ -2101,6 +1921,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] + + class TestBatchCostAttribution: """CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so the batch-cost log is attributed like a non-batch request.""" @@ -2196,9 +2018,7 @@ class TestBatchCostAttribution: """An alias lookup failure must not lose the spend row; the key hash and team still attribute it.""" instance = self._instance() - instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( - side_effect=Exception("db down") - ) + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down")) metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") @@ -2294,9 +2114,7 @@ class TestPollPageStarvation: async def test_unified_id_without_model_id_is_retired(self): """A unified id that decodes but carries no model_id is unroutable no matter what the config says, so it must leave the poll page instead of being retried forever.""" - prisma = self._prisma( - [self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] - ) + prisma = self._prisma([self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) llm_router = MagicMock() llm_router.aretrieve_batch = AsyncMock() @@ -2334,9 +2152,7 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() prisma.db.litellm_managedobjecttable.update.assert_awaited_once() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { - "batch_processed": True - } + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"batch_processed": True} @pytest.mark.asyncio async def test_provider_404_with_deployment_gone_keeps_job(self): @@ -2389,17 +2205,13 @@ class TestPollPageStarvation: async def test_retirement_falls_back_to_status_without_batch_processed_column(self): """Older schemas have no batch_processed column, so the only way to stop selecting the row is the status filter the poll query already applies.""" - prisma = self._prisma( - [self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] - ) + prisma = self._prisma([self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) instance = self._instance(prisma, MagicMock()) instance._has_batch_processed_column = False await instance.check_batch_cost() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { - "status": "stale_expired" - } + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"status": "stale_expired"} @pytest.mark.asyncio async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self): @@ -2454,14 +2266,11 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() - retired = [ - call[1]["where"]["id"] - for call in prisma.db.litellm_managedobjecttable.update.call_args_list - ] + retired = [call[1]["where"]["id"] for call in prisma.db.litellm_managedobjecttable.update.call_args_list] assert retired == ["job-no-model", "job-gone"] - assert ( - llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live" - ), "the newer healthy batch must still be polled in the same cycle" + assert llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live", ( + "the newer healthy batch must still be polled in the same cycle" + ) @pytest.mark.asyncio async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self): From 3bfeaa78ddd1bb7cae9ab6576e21a699a466b642 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 19:39:28 -0400 Subject: [PATCH 010/111] revert: undo accidental whole-file ruff-format of test_check_batch_cost.py The prior commit ran ruff format on the whole file to type the _batch_cost_result helper, reflowing hundreds of unrelated pre-existing lines that were never ruff-format-clean to begin with (confirmed at the PR's base commit, before any of these changes). CI's ruff-format gate only checks litellm/**/*.py, not tests/, so this reformatting served no CI purpose and only bloated the diff. Restores everything except the intended TYPE_CHECKING import and _batch_cost_result annotations. --- .../proxy_unit_tests/test_check_batch_cost.py | 507 ++++++++++++------ 1 file changed, 354 insertions(+), 153 deletions(-) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index ff2dce498f0..cd7f28007af 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -100,7 +100,9 @@ class TestCheckBatchCost: return MagicMock() @pytest.fixture - def check_batch_cost_instance(self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router): + def check_batch_cost_instance( + self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router + ): from litellm_enterprise.proxy.common_utils.check_batch_cost import ( CheckBatchCost, ) @@ -112,15 +114,23 @@ class TestCheckBatchCost: ) @pytest.mark.asyncio - async def test_cleanup_scoped_to_batch_file_purpose(self, check_batch_cost_instance, mock_prisma_client): + async def test_cleanup_scoped_to_batch_file_purpose( + self, check_batch_cost_instance, mock_prisma_client + ): """_cleanup_stale_managed_objects scopes its update to file_purpose='batch' only.""" - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) # Return empty so the main poll loop exits immediately - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[] + ) await check_batch_cost_instance.check_batch_cost() - calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list + calls = ( + mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list + ) stale_call = calls[0] assert stale_call[1]["data"] == {"status": "stale_expired"} where = stale_call[1]["where"] @@ -129,7 +139,9 @@ class TestCheckBatchCost: assert "created_at" in where @pytest.mark.asyncio - async def test_startup_probe_confirms_batch_processed_support(self, check_batch_cost_instance, mock_prisma_client): + async def test_startup_probe_confirms_batch_processed_support( + self, check_batch_cost_instance, mock_prisma_client + ): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) await check_batch_cost_instance.confirm_batch_processed_support() @@ -140,7 +152,9 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True @pytest.mark.asyncio - async def test_startup_probe_marks_column_absent(self, check_batch_cost_instance, mock_prisma_client): + async def test_startup_probe_marks_column_absent( + self, check_batch_cost_instance, mock_prisma_client + ): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( side_effect=Exception("column batch_processed does not exist") ) @@ -164,12 +178,18 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True @pytest.mark.asyncio - async def test_find_many_uses_pagination_and_excludes_stale(self, check_batch_cost_instance, mock_prisma_client): + async def test_find_many_uses_pagination_and_excludes_stale( + self, check_batch_cost_instance, mock_prisma_client + ): """find_many is called with take, order, and all terminal statuses excluded.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[] + ) await check_batch_cost_instance.check_batch_cost() @@ -195,7 +215,9 @@ class TestCheckBatchCost: """Falls back to query without batch_processed when primary query raises.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) # First find_many (primary query) raises with a schema error; second (fallback) returns empty mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( side_effect=[Exception("column batch_processed does not exist"), []] @@ -203,7 +225,9 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list + calls = ( + mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list + ) assert len(calls) == 2 fallback_where = calls[1][1]["where"] assert "batch_processed" not in fallback_where @@ -214,20 +238,32 @@ class TestCheckBatchCost: assert check_batch_cost_instance.batch_processed_support_confirmed is False @pytest.mark.asyncio - async def test_column_absence_cached_across_cycles(self, check_batch_cost_instance, mock_prisma_client): + async def test_column_absence_cached_across_cycles( + self, check_batch_cost_instance, mock_prisma_client + ): """After column absence is discovered, subsequent cycles skip the primary query entirely.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) # Simulate column already known absent from a previous cycle check_batch_cost_instance._has_batch_processed_column = False - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[] + ) await check_batch_cost_instance.check_batch_cost() # Only one find_many call — the fallback directly, no primary query attempt - assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 - fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"] + assert ( + mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 + ) + fallback_where = ( + mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1][ + "where" + ] + ) assert "batch_processed" not in fallback_where @pytest.mark.asyncio @@ -241,9 +277,13 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-fallback-1" @@ -252,16 +292,22 @@ class TestCheckBatchCost: # Simulate column already known absent (e.g. discovered on a previous cycle) check_batch_cost_instance._has_batch_processed_column = False - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) # Build a fake batch response whose status triggers the completion branch mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"completed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -309,7 +355,9 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -318,11 +366,15 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() # The update must have been called — this is the core assertion. - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - "Expected update() to be called exactly once for the completed job" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] - assert "batch_processed" not in update_data, "update() must NOT include batch_processed when column is absent" + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "Expected update() to be called exactly once for the completed job" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] + assert ( + "batch_processed" not in update_data + ), "update() must NOT include batch_processed when column is absent" assert update_data["status"] == "complete" @pytest.mark.asyncio @@ -398,9 +450,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=_batch_cost_result( - 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"] - ), + return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), ), patch( "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", @@ -418,9 +468,9 @@ class TestCheckBatchCost: passed_kwargs = mock_afile_content.await_args[1] snapshot = passed_kwargs.get("_litellm_internal_model_credentials") assert snapshot is not None, "cost poller must pass the trusted credential snapshot" - assert isinstance(snapshot, MappingProxyType), ( - "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" - ) + assert isinstance( + snapshot, MappingProxyType + ), "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" assert snapshot["s3_bucket_name"] == "configured-batch-bucket" @pytest.mark.asyncio @@ -434,9 +484,13 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-primary-1" @@ -444,15 +498,21 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"completed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -500,7 +560,9 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -508,13 +570,15 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - "Expected update() to be called exactly once for the completed job" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] - assert update_data["batch_processed"] is True, ( - "update() must include batch_processed=True when column is present" - ) + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "Expected update() to be called exactly once for the completed job" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] + assert ( + update_data["batch_processed"] is True + ), "update() must include batch_processed=True when column is present" assert update_data["status"] == "complete" @pytest.mark.asyncio @@ -658,16 +722,22 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-anthropic-1" mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = "completed" @@ -701,9 +771,9 @@ class TestCheckBatchCost: ): await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0, ( - "a failed cost tracking attempt must not mark the job processed" - ) + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 + ), "a failed cost tracking attempt must not mark the job processed" @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["failed", "expired", "cancelled"]) @@ -720,9 +790,13 @@ class TestCheckBatchCost: """ import base64 - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-terminal-1" @@ -732,25 +806,31 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = terminal_status mock_response.output_file_id = None - mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' + mock_response.model_dump_json.return_value = ( + f'{{"id":"batch-1","status":"{terminal_status}"}}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - f"Expected update() to be called exactly once for a {terminal_status} job" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), f"Expected update() to be called exactly once for a {terminal_status} job" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] assert update_data["status"] == terminal_status - assert update_data["batch_processed"] is True, ( - "terminal-status update() must set batch_processed=True so polling stops" - ) + assert ( + update_data["batch_processed"] is True + ), "terminal-status update() must set batch_processed=True so polling stops" @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["failed", "cancelled"]) @@ -785,9 +865,13 @@ class TestCheckBatchCost: f"litellm_proxy:application/octet-stream;unified_id,u-2;llm_output_file_id,{raw_error_file_id}".encode() ).decode() - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) input_file_row = MagicMock() input_file_row.unified_file_id = unified_input_file_id @@ -797,7 +881,9 @@ class TestCheckBatchCost: return input_file_row return None - mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(side_effect=find_managed_file) + mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( + side_effect=find_managed_file + ) mock_job = MagicMock() mock_job.id = "job-terminal-mint-1" @@ -806,7 +892,9 @@ class TestCheckBatchCost: mock_job.team_id = "team-1" check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) response = LiteLLMBatch( id="batch-456", @@ -824,7 +912,9 @@ class TestCheckBatchCost: mock_hook = MagicMock() mock_hook.get_unified_output_file_id.side_effect = [unified_error_file_id] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( + mock_hook + ) await check_batch_cost_instance.check_batch_cost() @@ -875,9 +965,13 @@ class TestCheckBatchCost: import base64 from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-completed-no-output-1" @@ -887,18 +981,24 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = completed_status mock_response.output_file_id = None mock_response.error_file_id = "file-error-123" - mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{completed_status}"}}' + mock_response.model_dump_json.return_value = ( + f'{{"id":"batch-1","status":"{completed_status}"}}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) # Billing reads credentials off the router; if it is touched we billed a batch # that has no output, which is the behaviour this test guards against. - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) with patch( "litellm.files.main.afile_content", @@ -906,18 +1006,22 @@ class TestCheckBatchCost: ) as mock_afile_content: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - "a completed batch with no output file must be marked processed exactly once" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "a completed batch with no output file must be marked processed exactly once" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] assert update_data["status"] == completed_status - assert update_data["batch_processed"] is True, ( - "completed-without-output update() must set batch_processed=True so polling stops" - ) - assert mock_afile_content.await_count == 0, "a batch with no output file must not be billed" - assert mock_llm_router.get_deployment_credentials_with_provider.call_count == 0, ( - "a batch with no output file must not enter the cost-tracking path" - ) + assert ( + update_data["batch_processed"] is True + ), "completed-without-output update() must set batch_processed=True so polling stops" + assert ( + mock_afile_content.await_count == 0 + ), "a batch with no output file must not be billed" + assert ( + mock_llm_router.get_deployment_credentials_with_provider.call_count == 0 + ), "a batch with no output file must not enter the cost-tracking path" @pytest.mark.asyncio async def test_non_terminal_status_left_unprocessed( @@ -928,7 +1032,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_job = MagicMock() @@ -936,7 +1042,9 @@ class TestCheckBatchCost: mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = "in_progress" @@ -962,9 +1070,9 @@ class TestCheckBatchCost: ): await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0, ( - "a non-terminal batch must not be written back (would stop polling prematurely)" - ) + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 + ), "a non-terminal batch must not be written back (would stop polling prematurely)" @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["expired", "cancelled", "failed"]) @@ -981,9 +1089,13 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-terminal-with-output-1" @@ -991,15 +1103,21 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = terminal_status mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' + mock_response.model_dump_json.return_value = ( + f'{{"id":"batch-1","status":"{terminal_status}"}}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -1047,7 +1165,9 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1055,16 +1175,20 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_afile_content.await_count == 1, ( - f"{terminal_status} batch with an output file must fetch results and be billed" - ) + assert ( + mock_afile_content.await_count == 1 + ), f"{terminal_status} batch with an output file must fetch results and be billed" mock_logging_obj.async_success_handler.assert_awaited_once() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] - assert update_data["batch_processed"] is True - assert update_data["status"] == terminal_status, ( - f"billed {terminal_status} batch must keep its real terminal status in the DB" + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] + assert update_data["batch_processed"] is True + assert ( + update_data["status"] == terminal_status + ), f"billed {terminal_status} batch must keep its real terminal status in the DB" @pytest.mark.asyncio async def test_terminal_batch_with_missing_output_file_is_retired_unbilled( @@ -1081,9 +1205,13 @@ class TestCheckBatchCost: from litellm.exceptions import NotFoundError - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-output-gone-1" @@ -1093,17 +1221,23 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) missing_output_file_id = "gs://batch-out/job-1/predictions.jsonl" mock_response = MagicMock() mock_response.status = "failed" mock_response.output_file_id = missing_output_file_id mock_response.error_file_id = None - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"failed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"failed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) with ( patch( @@ -1124,10 +1258,12 @@ class TestCheckBatchCost: assert mock_afile_content.await_count == 1 mock_calculate.assert_not_awaited() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - "a terminal batch with a 404ing output file must be retired, not retried forever" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "a terminal batch with a 404ing output file must be retired, not retried forever" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] assert update_data["status"] == "failed" assert update_data["batch_processed"] is True @@ -1140,9 +1276,13 @@ class TestCheckBatchCost: Without this, GET /batches/{id} returns a raw file ID that cannot be routed through the proxy, causing API_KEY errors when clients call GET /files/{id}/content. """ - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-raw-file-1" @@ -1151,7 +1291,9 @@ class TestCheckBatchCost: mock_job.team_id = None check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) raw_output_file_id = "file-batch-output-abc123" raw_error_file_id = "file-batch-error-xyz456" @@ -1162,10 +1304,14 @@ class TestCheckBatchCost: mock_response.status = "completed" mock_response.output_file_id = raw_output_file_id mock_response.error_file_id = raw_error_file_id - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"completed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "azure" @@ -1180,7 +1326,9 @@ class TestCheckBatchCost: fake_managed_error_id, ] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( + mock_hook + ) mock_file_content = MagicMock() mock_file_content.content = b'{"id":"req-1"}' @@ -1223,7 +1371,9 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-5-mini", "azure", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1274,7 +1424,9 @@ class TestUnmanagedVertexRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = "8823717160934178816" - job.file_object = file_object if file_object is not None else _unmanaged_vertex_file_object() + job.file_object = ( + file_object if file_object is not None else _unmanaged_vertex_file_object() + ) return job def test_flag_off_skips_unmanaged_id_unchanged(self): @@ -1312,7 +1464,9 @@ class TestUnmanagedVertexRouting: assert result == ("deploy-1", "8823717160934178816") # bare model name (trailing GCS segment), not the full publishers/.. path - router.resolve_model_name_from_model_id.assert_called_once_with("gemini-2.5-flash") + router.resolve_model_name_from_model_id.assert_called_once_with( + "gemini-2.5-flash" + ) router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): @@ -1332,7 +1486,9 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_uses_later_vertex_deployment_with_matching_suffix(self): router = MagicMock() @@ -1380,7 +1536,9 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_non_gcs_input_is_not_unmanaged_vertex(self): """Flag on, but input_file_id is not a gs:// publishers path: treat as unroutable, @@ -1388,7 +1546,9 @@ class TestUnmanagedVertexRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job(file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123")) + job = self._job( + file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123") + ) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1412,7 +1572,9 @@ class TestUnmanagedVertexRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"8823717160934178816","status":"completed"}' + ) router.aretrieve_batch = AsyncMock(return_value=mock_response) router.get_deployment_credentials_with_provider = MagicMock( return_value={"vertex_project": "p", "vertex_location": "us-central1"} @@ -1434,7 +1596,9 @@ class TestUnmanagedVertexRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[self._job()] + ) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1465,7 +1629,9 @@ class TestUnmanagedVertexRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gemini-2.5-flash", "vertex_ai", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1506,7 +1672,9 @@ class TestUnmanagedBedrockRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = self._ARN - job.file_object = file_object if file_object is not None else _unmanaged_bedrock_file_object() + job.file_object = ( + file_object if file_object is not None else _unmanaged_bedrock_file_object() + ) return job def _bedrock_deployment(self): @@ -1561,7 +1729,9 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_matches_deployment_despite_colon_dash_mismatch(self): """The S3 object key has ':' replaced with '-' (e.g. 'v1-0'), but the configured @@ -1599,7 +1769,9 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_non_s3_input_is_not_unmanaged_bedrock(self): """Flag on, but input_file_id is not a litellm-bedrock-files- s3:// key: treat as @@ -1607,7 +1779,9 @@ class TestUnmanagedBedrockRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job(file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123")) + job = self._job( + file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123") + ) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1630,9 +1804,13 @@ class TestUnmanagedBedrockRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}' + mock_response.model_dump_json.return_value = ( + f'{{"id":"{self._ARN}","status":"completed"}}' + ) router.aretrieve_batch = AsyncMock(return_value=mock_response) - router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"}) + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"aws_region_name": "us-east-1"} + ) deployment = self._bedrock_deployment() deployment.model_name = "claude-sonnet-4" @@ -1648,7 +1826,9 @@ class TestUnmanagedBedrockRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[self._job()] + ) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1679,7 +1859,9 @@ class TestUnmanagedBedrockRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("claude-sonnet-4", "bedrock", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1803,7 +1985,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: ) router = MagicMock() - router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) deployment = MagicMock() deployment.litellm_params.custom_llm_provider = "azure" deployment.litellm_params.model = "azure/gpt-5.5" @@ -1812,8 +1996,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: router.get_deployment = MagicMock(return_value=deployment) hook = MagicMock() - hook.get_unified_output_file_id = lambda output_file_id, model_id, model_name: ( - _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( + hook.get_unified_output_file_id = ( + lambda output_file_id, model_id, model_name: _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( None, output_file_id=output_file_id, model_id=model_id, model_name=model_name ) ) @@ -1882,7 +2066,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: get_models_from_unified_file_id, ) - output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) + output_file_id = await self._run( + self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) + ) decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] @@ -1896,7 +2082,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: _extract_models_from_managed_resource_id, ) - output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) + output_file_id = await self._run( + self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) + ) models = _extract_models_from_managed_resource_id(output_file_id, "file_id", None) assert models == [self._PUBLIC_MODEL_GROUP] @@ -1904,7 +2092,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: await can_key_call_model( model=models[0], llm_model_list=None, - valid_token=UserAPIKeyAuth(api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP]), + valid_token=UserAPIKeyAuth( + api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP] + ), llm_router=None, ) is True @@ -1921,8 +2111,6 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] - - class TestBatchCostAttribution: """CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so the batch-cost log is attributed like a non-batch request.""" @@ -2018,7 +2206,9 @@ class TestBatchCostAttribution: """An alias lookup failure must not lose the spend row; the key hash and team still attribute it.""" instance = self._instance() - instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down")) + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + side_effect=Exception("db down") + ) metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") @@ -2114,7 +2304,9 @@ class TestPollPageStarvation: async def test_unified_id_without_model_id_is_retired(self): """A unified id that decodes but carries no model_id is unroutable no matter what the config says, so it must leave the poll page instead of being retried forever.""" - prisma = self._prisma([self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) + prisma = self._prisma( + [self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] + ) llm_router = MagicMock() llm_router.aretrieve_batch = AsyncMock() @@ -2152,7 +2344,9 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() prisma.db.litellm_managedobjecttable.update.assert_awaited_once() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"batch_processed": True} + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { + "batch_processed": True + } @pytest.mark.asyncio async def test_provider_404_with_deployment_gone_keeps_job(self): @@ -2205,13 +2399,17 @@ class TestPollPageStarvation: async def test_retirement_falls_back_to_status_without_batch_processed_column(self): """Older schemas have no batch_processed column, so the only way to stop selecting the row is the status filter the poll query already applies.""" - prisma = self._prisma([self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) + prisma = self._prisma( + [self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] + ) instance = self._instance(prisma, MagicMock()) instance._has_batch_processed_column = False await instance.check_batch_cost() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"status": "stale_expired"} + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { + "status": "stale_expired" + } @pytest.mark.asyncio async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self): @@ -2266,11 +2464,14 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() - retired = [call[1]["where"]["id"] for call in prisma.db.litellm_managedobjecttable.update.call_args_list] + retired = [ + call[1]["where"]["id"] + for call in prisma.db.litellm_managedobjecttable.update.call_args_list + ] assert retired == ["job-no-model", "job-gone"] - assert llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live", ( - "the newer healthy batch must still be polled in the same cycle" - ) + assert ( + llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live" + ), "the newer healthy batch must still be polled in the same cycle" @pytest.mark.asyncio async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self): From 98447e00436443f26ac5538bb4ad646cc0b32f75 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 19:58:48 -0400 Subject: [PATCH 011/111] fix(batches): reconcile merge with upstream batch pricing changes - Resolved merge conflicts in _handle_completed_batch and CheckBatchCost._track_completed_batch_cost, keeping BatchCostUsageResult while adopting upstream's model_info threading and improved deployment-pricing lookup - Fixed two upstream tests and one mock that still expected the old tuple(cost, usage, models) return shape - Suppressed the one new LIT002 violation from the empty-output-file BatchCostUsageResult literal --- litellm/batches/batch_utils.py | 2 +- .../test_litellm/batches/test_batch_utils.py | 20 +++++++++---------- .../test_litellm_logging.py | 12 +++++++++-- 3 files changed, 21 insertions(+), 13 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index bae8e198223..d0874b5aff5 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -86,7 +86,7 @@ async def _handle_completed_batch( return BatchCostUsageResult( cost=0.0, usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), - models=[], + 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 diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 0adebf2506b..64b3df4180b 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1501,24 +1501,24 @@ async def test_handle_completed_bedrock_batch_prices_from_deployment_model(monke monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) - cost, usage, _ = await bu._handle_completed_batch( + result = await bu._handle_completed_batch( _batch("of"), custom_llm_provider="bedrock", model_name="bedrock/global.anthropic.claude-sonnet-4-6", ) - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (1800, 1000, 2800) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (1800, 1000, 2800) # 3e-06 / 1.5e-05 on-demand, halved for batch. - assert cost == pytest.approx(1800 * 3e-06 / 2 + 1000 * 1.5e-05 / 2) + assert result.cost == pytest.approx(1800 * 3e-06 / 2 + 1000 * 1.5e-05 / 2) # The response model alone cannot price a bedrock batch: this is the $0 bug. - zero_cost, zero_usage, _ = await bu._handle_completed_batch( + zero_result = await bu._handle_completed_batch( _batch("of"), custom_llm_provider="bedrock", model_name=None, ) - assert zero_cost == 0.0 - assert zero_usage.total_tokens == 2800 + assert zero_result.cost == 0.0 + assert zero_result.usage.total_tokens == 2800 @pytest.mark.asyncio @@ -1531,7 +1531,7 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) -> monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) - free_cost, _, _ = await bu._handle_completed_batch( + free_result = await bu._handle_completed_batch( _batch("of"), custom_llm_provider="vertex_ai", model_name="vertex_ai/gemini-2.5-flash", @@ -1542,15 +1542,15 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) -> "output_cost_per_token_batches": 0.0, }, ) - assert free_cost == 0.0 + assert free_result.cost == 0.0 - billed_cost, _, _ = await bu._handle_completed_batch( + billed_result = await bu._handle_completed_batch( _batch("of"), custom_llm_provider="vertex_ai", model_name="vertex_ai/gemini-2.5-flash", model_info=None, ) - assert billed_cost > 0.0 + assert billed_result.cost > 0.0 # =========================================================================== # diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 0d54680fa81..e46057fe11a 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -586,9 +586,17 @@ class TestRetrieveBatchCostPassesModelIdentity: captured: dict[str, object] = {} - async def fake_handle_completed_batch(**kwargs: object) -> tuple[float, Usage, list[str]]: + from litellm.batches.batch_utils import BatchCostUsageResult + + async def fake_handle_completed_batch(**kwargs: object) -> BatchCostUsageResult: captured.update(kwargs) - return 1.25, Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800), ["m"] + return BatchCostUsageResult( + cost=1.25, + usage=Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800), + models=["m"], + successful_requests=1, + failed_requests=0, + ) monkeypatch.setattr(logging_module, "_handle_completed_batch", fake_handle_completed_batch) From d989f172e9d78ff63c5867f45850d5065b1ad7c8 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 20:32:11 -0400 Subject: [PATCH 012/111] fix(batches): revert invalid Final on a loop-scoped local basedpyright rejects a Final variable assigned inside a loop body (reportGeneralTypeIssues); the type-discipline checker doesn't flag this line without Final either, so the annotation only bought a basedpyright budget regression. --- litellm/batches/batch_utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index d0874b5aff5..4e62c4a0130 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -170,9 +170,7 @@ def _classify_output_line_stats( custom_llm_provider=custom_llm_provider, call_type=CallTypes.aretrieve_batch.value, ) - reasoning_tokens: Final = ( - usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None - ) + reasoning_tokens = usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None yield _BatchOutputLineStats( cost=line_cost, prompt_tokens=usage.prompt_tokens, From 64e993773da57f5bd39dc60b1d45b2eaafc5acd6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 05:25:22 +0000 Subject: [PATCH 013/111] chore(typing): clear 1.5k basedpyright Any errors across 54 files Retypes the 54 highest-density reportAny/reportExplicitAny sources with real types instead of shuffling the ceilings around: typed prisma table Protocols so the untyped client surface stops at the query, local TypedDicts for JSON and dict payloads, concrete chunk and logging types on the streaming and callback surfaces, and 3-argument getattr with a Callable annotation where an SDK object is genuinely duck-typed No cast(), no type: ignore, no noqa, no suppression comments, and no new Any annotations. Whole-tree basedpyright drops 1,941 errors with no rule rising anywhere, and all three budget files are ratcheted so the cleared headroom cannot silently grow back Adds a GDC regression test pinning the named AttributeError that the typed credential accessor now raises when with_gdch_audience is missing --- basedpyright-code-budget.json | 26 +-- .../proxy/audit_logging_endpoints.py | 108 +++++++---- litellm/_lazy_imports.py | 79 ++++---- .../transformation.py | 46 ++--- litellm/caching/valkey_semantic_cache.py | 71 ++++--- .../bitbucket/bitbucket_client.py | 49 +++-- .../compression_interception/handler.py | 84 +++++---- litellm/integrations/custom_logger.py | 28 +-- .../opik_payload_builder/payload_builders.py | 12 +- .../prometheus_helpers/prometheus_api.py | 54 +++++- .../websearch_interception/handler.py | 16 +- .../llm_response_utils/response_metadata.py | 31 ++- .../model_response_utils.py | 66 +++++-- .../litellm_core_utils/streaming_handler.py | 8 +- litellm/litellm_core_utils/url_utils.py | 25 +-- .../llms/anthropic/batches/transformation.py | 59 ++++-- .../context_management/dispatcher.py | 81 +++++--- .../responses_adapters/handler.py | 40 ++-- .../llms/anthropic/skills/transformation.py | 30 ++- litellm/llms/azure/files/handler.py | 13 +- litellm/llms/bedrock/realtime/handler.py | 104 ++++++++-- litellm/llms/chatgpt/chat/streaming_utils.py | 41 +++- .../llms/compactifai/chat/transformation.py | 33 +++- .../llms/custom_httpx/container_handler.py | 178 ++++++++++++------ litellm/llms/gdc/chat/transformation.py | 45 ++++- .../llms/infinity/rerank/transformation.py | 29 ++- .../litellm_proxy/skills/code_execution.py | 76 ++++++-- litellm/llms/oci/chat/cohere.py | 38 +++- litellm/llms/ollama/completion/handler.py | 36 +++- .../llms/openai/containers/transformation.py | 72 +++++-- .../vector_stores/rag_api/transformation.py | 127 ++++++++++--- litellm/models/base.py | 2 +- .../mcp_server/oauth2_flow_backfill.py | 63 ++++++- .../_experimental/mcp_server/toolset_db.py | 93 +++++++-- .../proxy/agent_endpoints/agent_registry.py | 81 +++++--- .../proxy/client/cli/commands/credentials.py | 51 ++++- litellm/proxy/client/cli/commands/teams.py | 52 +++-- litellm/proxy/common_utils/callback_utils.py | 14 +- litellm/proxy/common_utils/get_routes.py | 65 ++++--- .../proxy/common_utils/user_api_key_cache.py | 48 ++--- litellm/proxy/db/routing_prisma_wrapper.py | 36 ++-- .../guardrail_hooks/custom_code/sandbox.py | 14 +- .../hiddenlayer/hiddenlayer.py | 33 ++-- .../llm_as_a_judge/__init__.py | 130 ++++++++++--- .../guardrail_hooks/noma/noma_v2.py | 26 +-- .../promptguard/promptguard.py | 20 +- .../jwt_key_mapping_endpoints.py | 76 ++++++-- .../response_polling/background_streaming.py | 47 ++++- .../search_endpoints/search_tool_registry.py | 57 ++++-- litellm/rag/rag_query.py | 65 +++++-- litellm/repositories/base_repository.py | 51 +++-- .../repositories/credentials_repository.py | 58 ++++-- litellm/repositories/team_repository.py | 76 ++++++-- litellm/rust_bridge/responses_websocket.py | 23 ++- .../secret_managers/secret_manager_handler.py | 60 +++++- ruff-strict-budget.json | 18 +- .../gdc/chat/test_gdc_chat_transformation.py | 19 ++ type-discipline-budget.json | 12 +- 58 files changed, 2172 insertions(+), 823 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index b4c324a2c4c..58b8ae01cb0 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 19955 + "limit": 18773 }, "reportArgumentType": { "limit": 2566 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 6049 + "limit": 5774 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5663 + "limit": 5640 }, "reportMissingTypeArgument": { - "limit": 15555 + "limit": 15498 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,31 +99,31 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44655 + "limit": 44589 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 39017 + "limit": 38882 }, "reportUnknownParameterType": { - "limit": 19885 + "limit": 19806 }, "reportUnknownVariableType": { - "limit": 30572 + "limit": 30456 }, "reportUnnecessaryCast": { - "limit": 117 + "limit": 116 }, "reportUnnecessaryComparison": { - "limit": 699 + "limit": 698 }, "reportUnnecessaryContains": { "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 836 + "limit": 835 }, "reportUntypedBaseClass": { "limit": 0 @@ -138,9 +138,9 @@ "limit": 139 }, "reportUnusedImport": { - "limit": 545 + "limit": 544 }, "reportUnusedVariable": { - "limit": 146 + "limit": 142 } } diff --git a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py index 18ac29b9781..21bde06e86a 100644 --- a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py @@ -7,7 +7,9 @@ GET - /audit/{id} - Get audit log by id GET - /audit - Get all audit logs """ -from typing import Any, Dict, List, Optional +from collections.abc import Mapping, Sequence +from datetime import datetime +from typing import Final, Protocol #### AUDIT LOGGING #### from fastapi import APIRouter, Depends, HTTPException, Query @@ -15,6 +17,7 @@ from litellm_enterprise.types.proxy.audit_logging_endpoints import ( AuditLogResponse, PaginatedAuditLogResponse, ) +from typing_extensions import ReadOnly, TypedDict from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -22,7 +25,44 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth router = APIRouter() -def _build_json_field_or_condition(json_key: str, value: str) -> Dict[str, Any]: +class _AuditLogFields(TypedDict): + """Columns of the `LiteLLM_AuditLog` table, as returned by `model_dump()`.""" + + id: ReadOnly[str] + updated_at: ReadOnly[datetime] + changed_by: ReadOnly[str] + changed_by_api_key: ReadOnly[str] + action: ReadOnly[str] + table_name: ReadOnly[str] + object_id: ReadOnly[str] + before_value: ReadOnly[dict[str, object] | None] + updated_values: ReadOnly[dict[str, object] | None] + + +class _AuditLogRecord(Protocol): + """Row of the `LiteLLM_AuditLog` table as materialised by the Prisma client.""" + + def model_dump(self) -> _AuditLogFields: ... + + +class _AuditLogTable(Protocol): + """The `litellm_auditlog` accessor of the Prisma client.""" + + async def find_many( + self, + *, + where: Mapping[str, object], + order: Mapping[str, str], + skip: int, + take: int, + ) -> Sequence[_AuditLogRecord]: ... + + async def count(self, *, where: Mapping[str, object]) -> int: ... + + async def find_unique(self, *, where: Mapping[str, str]) -> _AuditLogRecord | None: ... + + +def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, object]: """ Build an OR condition that matches a value inside a JSON column at the given key, checking both before_value and updated_values. @@ -53,33 +93,33 @@ async def get_audit_logs( page: int = Query(1, ge=1), page_size: int = Query(10, ge=1, le=100), # Filter parameters - changed_by: Optional[str] = Query( + changed_by: str | None = Query( None, description="Filter by user or system that performed the action" ), - changed_by_api_key: Optional[str] = Query( + changed_by_api_key: str | None = Query( None, description="Filter by API key hash that performed the action" ), - action: Optional[str] = Query( + action: str | None = Query( None, description="Filter by action type (create, update, delete)" ), - table_name: Optional[str] = Query( + table_name: str | None = Query( None, description="Filter by table name that was modified" ), - object_id: Optional[str] = Query( + object_id: str | None = Query( None, description="Filter by ID of the object that was modified" ), - start_date: Optional[str] = Query(None, description="Filter logs after this date"), - end_date: Optional[str] = Query(None, description="Filter logs before this date"), - object_team_id: Optional[str] = Query( + start_date: str | None = Query(None, description="Filter logs after this date"), + end_date: str | None = Query(None, description="Filter logs before this date"), + object_team_id: str | None = Query( None, description="Filter by team_id present in before_value or updated_values JSON (PostgreSQL only)", ), - object_key_hash: Optional[str] = Query( + object_key_hash: str | None = Query( None, description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)", ), # Sorting parameters - sort_by: Optional[str] = Query( + sort_by: str | None = Query( None, description="Column to sort by (e.g. 'updated_at', 'action', 'table_name')", ), @@ -102,7 +142,7 @@ async def get_audit_logs( ) # Build filter conditions - where_conditions: Dict[str, Any] = {} + where_conditions: Final[dict[str, object]] = {} if changed_by: where_conditions["changed_by"] = changed_by if changed_by_api_key: @@ -114,33 +154,31 @@ async def get_audit_logs( if object_id: where_conditions["object_id"] = object_id if start_date or end_date: - date_filter: Dict[str, Any] = {} - if start_date: - date_filter["gte"] = start_date - if end_date: - date_filter["lte"] = end_date + date_filter: Final[Mapping[str, str]] = { + bound: bound_value for bound, bound_value in (("gte", start_date), ("lte", end_date)) if bound_value + } where_conditions["updated_at"] = date_filter # JSON field filters (PostgreSQL only) — each filter is AND'd with the # others, but checks both before_value and updated_values internally (OR). - if object_team_id: - where_conditions["AND"] = where_conditions.get("AND", []) + [ - _build_json_field_or_condition("team_id", object_team_id) - ] - if object_key_hash: - where_conditions["AND"] = where_conditions.get("AND", []) + [ - _build_json_field_or_condition("token", object_key_hash) + if object_team_id or object_key_hash: + where_conditions["AND"] = [ + _build_json_field_or_condition(json_key, json_value) + for json_key, json_value in ( + ("team_id", object_team_id), + ("token", object_key_hash), + ) + if json_value ] # Build sort conditions - order_by: Dict[str, Any] = {} - if sort_by and isinstance(sort_by, str): - order_by[sort_by] = sort_order - else: - order_by["updated_at"] = sort_order # Default sort by updated_at + sort_column: Final[str] = sort_by if sort_by and isinstance(sort_by, str) else "updated_at" + order_by: Final[Mapping[str, str]] = {sort_column: sort_order} + + audit_log_table: Final[_AuditLogTable] = prisma_client.db.litellm_auditlog # Get paginated results - audit_logs = await prisma_client.db.litellm_auditlog.find_many( + audit_logs: Final = await audit_log_table.find_many( where=where_conditions, order=order_by, skip=(page - 1) * page_size, @@ -148,8 +186,8 @@ async def get_audit_logs( ) # Get total count for pagination - total_count = await prisma_client.db.litellm_auditlog.count(where=where_conditions) - total_pages = -(-total_count // page_size) # Ceiling division + total_count: Final = await audit_log_table.count(where=where_conditions) + total_pages: Final = -(-total_count // page_size) # Ceiling division # Return paginated response return PaginatedAuditLogResponse( @@ -198,8 +236,10 @@ async def get_audit_log_by_id( detail={"message": CommonProxyErrors.db_not_connected_error.value}, ) + audit_log_table: Final[_AuditLogTable] = prisma_client.db.litellm_auditlog + # Get the audit log by ID - audit_log = await prisma_client.db.litellm_auditlog.find_unique(where={"id": id}) + audit_log: Final = await audit_log_table.find_unique(where={"id": id}) if audit_log is None: raise HTTPException( diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 933464d3f23..dff3b11e353 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -17,8 +17,14 @@ until they're actually needed. import importlib import sys -from collections.abc import Callable -from typing import Any, Final, cast +from collections.abc import Callable, Mapping +from typing import TYPE_CHECKING, Any, Final, cast + +if TYPE_CHECKING: + import httpx + import tiktoken + + from .caching.llm_caching_handler import LLMClientCache as LLMClientCacheType # Import all the data structures that define what can be lazy-loaded # These are just lists of names and maps of where to find them @@ -54,7 +60,7 @@ from ._lazy_imports_registry import ( ) -def get_litellm_globals() -> dict: +def get_litellm_globals() -> dict[str, object]: """ Get the globals dictionary of the litellm module. @@ -64,7 +70,7 @@ def get_litellm_globals() -> dict: return sys.modules["litellm"].__dict__ -def _get_utils_globals() -> dict: +def _get_utils_globals() -> dict[str, object]: """ Get the globals dictionary of the utils module. @@ -74,14 +80,19 @@ def _get_utils_globals() -> dict: return sys.modules["litellm.utils"].__dict__ +def _get_module_level_client_timeout(litellm_globals: Mapping[str, Any]) -> "float | httpx.Timeout | None": + """Read the configured `litellm.request_timeout` used for the module level http clients.""" + return litellm_globals.get("request_timeout") + + # These are special lazy loaders for things that are used internally # They're separate from the main lazy import system because they have specific use cases # Lazy loader for default encoding - avoids importing heavy tiktoken library at startup -_default_encoding: Any | None = None +_default_encoding: "tiktoken.Encoding | None" = None -def _get_default_encoding() -> Any: +def _get_default_encoding() -> "tiktoken.Encoding": """ Lazily load and cache the default OpenAI encoding. @@ -100,10 +111,10 @@ def _get_default_encoding() -> Any: # Lazy loader for get_modified_max_tokens to avoid importing token_counter at module import time -_get_modified_max_tokens_func: Any | None = None +_get_modified_max_tokens_func: Callable[..., int | None] | None = None -def _get_modified_max_tokens() -> Any: +def _get_modified_max_tokens() -> Callable[..., int | None]: """ Lazily load and cache the get_modified_max_tokens function. @@ -124,10 +135,10 @@ def _get_modified_max_tokens() -> Any: # Lazy loader for token_counter to avoid importing token_counter module at module import time -_token_counter_new_func: Any | None = None +_token_counter_new_func: Callable[..., int] | None = None -def _get_token_counter_new() -> Any: +def _get_token_counter_new() -> Callable[..., int]: """ Lazily load and cache the token_counter function (aliased as token_counter_new). @@ -154,10 +165,10 @@ def _get_token_counter_new() -> Any: # This registry maps attribute names (like "ModelResponse") to handler functions # It's built once the first time someone accesses a lazy-loaded attribute # Example: {"ModelResponse": _lazy_import_utils, "Cache": _lazy_import_caching, ...} -_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], Any]] | None = None +_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], object]] | None = None -def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]: +def _get_lazy_import_registry() -> dict[str, Callable[[str], object]]: """ Build the registry that maps attribute names to their handler functions. @@ -206,7 +217,7 @@ def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]: return _LAZY_IMPORT_REGISTRY -def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> Any: +def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> object: """ Generic function that handles lazy importing for most attributes. @@ -255,7 +266,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate # Step 6: Get the actual attribute from the module # Example: getattr(utils_module, "ModelResponse") returns the ModelResponse class - value: Final = getattr(module, attr_name) + value: Final[object] = getattr(module, attr_name) # Step 7: Cache it so we don't have to import again next time _globals[name] = value @@ -272,62 +283,62 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate # The registry (above) maps attribute names to these handler functions. -def _lazy_import_utils(name: str) -> Any: +def _lazy_import_utils(name: str) -> object: """Handler for utils module attributes (ModelResponse, token_counter, etc.)""" return _generic_lazy_import(name, _UTILS_IMPORT_MAP, "Utils") -def _lazy_import_cost_calculator(name: str) -> Any: +def _lazy_import_cost_calculator(name: str) -> object: """Handler for cost calculator functions (completion_cost, cost_per_token, etc.)""" return _generic_lazy_import(name, _COST_CALCULATOR_IMPORT_MAP, "Cost calculator") -def _lazy_import_token_counter(name: str) -> Any: +def _lazy_import_token_counter(name: str) -> object: """Handler for token counter utilities""" return _generic_lazy_import(name, _TOKEN_COUNTER_IMPORT_MAP, "Token counter") -def _lazy_import_bedrock_types(name: str) -> Any: +def _lazy_import_bedrock_types(name: str) -> object: """Handler for Bedrock type aliases""" return _generic_lazy_import(name, _BEDROCK_TYPES_IMPORT_MAP, "Bedrock types") -def _lazy_import_types_utils(name: str) -> Any: +def _lazy_import_types_utils(name: str) -> object: """Handler for types from litellm.types.utils (BudgetConfig, ImageObject, etc.)""" return _generic_lazy_import(name, _TYPES_UTILS_IMPORT_MAP, "Types utils") -def _lazy_import_caching(name: str) -> Any: +def _lazy_import_caching(name: str) -> object: """Handler for caching classes (Cache, DualCache, RedisCache, etc.)""" return _generic_lazy_import(name, _CACHING_IMPORT_MAP, "Caching") -def _lazy_import_dotprompt(name: str) -> Any: +def _lazy_import_dotprompt(name: str) -> object: """Handler for dotprompt integration globals""" return _generic_lazy_import(name, _DOTPROMPT_IMPORT_MAP, "Dotprompt") -def _lazy_import_types(name: str) -> Any: +def _lazy_import_types(name: str) -> object: """Handler for type classes (GuardrailItem, etc.)""" return _generic_lazy_import(name, _TYPES_IMPORT_MAP, "Types") -def _lazy_import_llm_configs(name: str) -> Any: +def _lazy_import_llm_configs(name: str) -> object: """Handler for LLM config classes (AnthropicConfig, OpenAILikeChatConfig, etc.)""" return _generic_lazy_import(name, _LLM_CONFIGS_IMPORT_MAP, "LLM config") -def _lazy_import_litellm_logging(name: str) -> Any: +def _lazy_import_litellm_logging(name: str) -> object: """Handler for litellm_logging module (Logging, modify_integration)""" return _generic_lazy_import(name, _LITELLM_LOGGING_IMPORT_MAP, "Litellm logging") -def _lazy_import_llm_provider_logic(name: str) -> Any: +def _lazy_import_llm_provider_logic(name: str) -> object: """Handler for LLM provider logic functions (get_llm_provider, etc.)""" return _generic_lazy_import(name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic") -def _lazy_import_utils_module(name: str) -> Any: +def _lazy_import_utils_module(name: str) -> object: """ Handler for utils module lazy imports. @@ -355,7 +366,7 @@ def _lazy_import_utils_module(name: str) -> Any: module = importlib.import_module(module_path) # Get the actual attribute from the module - value: Final = getattr(module, attr_name) + value: Final[object] = getattr(module, attr_name) # Cache it so we don't have to import again next time _globals[name] = value @@ -370,7 +381,7 @@ def _lazy_import_utils_module(name: str) -> Any: # These handlers have custom logic that doesn't fit the generic pattern -def _lazy_import_llm_client_cache(name: str) -> Any: +def _lazy_import_llm_client_cache(name: str) -> object: """ Handler for LLM client cache - has special logic for singleton instance. @@ -387,7 +398,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any: # Import the class module: Final = importlib.import_module("litellm.caching.llm_caching_handler") - LLMClientCache: Final = getattr(module, "LLMClientCache") + LLMClientCache: Final[type[LLMClientCacheType]] = getattr(module, "LLMClientCache") # If they want the class itself, return it if name == "LLMClientCache": @@ -403,7 +414,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any: raise AttributeError(f"LLM client cache lazy import: unknown attribute {name!r}") -def _lazy_import_http_handlers(name: str) -> Any: +def _lazy_import_http_handlers(name: str) -> object: """ Handler for HTTP clients - has special logic for creating client instances. @@ -419,8 +430,8 @@ def _lazy_import_http_handlers(name: str) -> Any: from litellm.llms.custom_httpx.http_handler import get_async_httpx_client # Get timeout from module config (if set) - timeout = _globals.get("request_timeout") - params: Final = {"timeout": timeout, "client_alias": "module level aclient"} + async_timeout: Final = _get_module_level_client_timeout(_globals) + params: Final = {"timeout": async_timeout, "client_alias": "module level aclient"} # Create the client instance provider_id: Final = cast(Any, "litellm_module_level_client") @@ -437,8 +448,8 @@ def _lazy_import_http_handlers(name: str) -> Any: # Create a sync HTTP client from litellm.llms.custom_httpx.http_handler import HTTPHandler - timeout = _globals.get("request_timeout") - sync_client: Final = HTTPHandler(timeout=timeout) + sync_timeout: Final = _get_module_level_client_timeout(_globals) + sync_client: Final = HTTPHandler(timeout=sync_timeout) # Cache it _globals["module_level_client"] = sync_client diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index 15cf77708f9..90afa5adf9e 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -17,12 +17,16 @@ A2A Streaming Events: - Artifact update (kind: "artifact-update") - Content/artifact delivery """ +from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final from uuid import uuid4 from litellm._logging import verbose_logger +if TYPE_CHECKING: + from litellm.types.utils import Choices + class A2AStreamingContext: """ @@ -30,7 +34,7 @@ class A2AStreamingContext: Tracks task_id, context_id, and message accumulation. """ - def __init__(self, request_id: str, input_message: dict[str, Any]): + def __init__(self, request_id: str, input_message: Mapping[str, object]): self.request_id = request_id self.task_id = str(uuid4()) self.context_id = str(uuid4()) @@ -46,7 +50,7 @@ class A2ACompletionBridgeTransformation: """ @staticmethod - def _extract_text_from_a2a_parts(parts: list[dict[str, Any]]) -> str: + def _extract_text_from_a2a_parts(parts: Sequence[Mapping[str, object]]) -> str: """Extract text from A2A parts (with or without explicit ``kind``).""" content_parts: Final[list[str]] = [] for part in parts: @@ -62,16 +66,16 @@ class A2ACompletionBridgeTransformation: @staticmethod def get_forward_metadata( - a2a_message: dict[str, Any], + a2a_message: Mapping[str, object], params: dict[str, Any] | None = None, - ) -> dict[str, Any] | None: + ) -> dict[str, object] | None: """ Merge A2A metadata from MessageSendParams and the message for downstream providers. Forwarded once on the LangGraph run payload (``metadata``), not duplicated on each input message — see ``apply_forward_metadata_to_completion_params``. """ - merged: Final[dict[str, Any]] = {} + merged: Final[dict[str, object]] = {} if params and isinstance(params.get("metadata"), dict): merged.update(params["metadata"]) message_metadata: Final = a2a_message.get("metadata") @@ -81,8 +85,8 @@ class A2ACompletionBridgeTransformation: @staticmethod def apply_forward_metadata_to_completion_params( - completion_params: dict[str, Any], - a2a_message: dict[str, Any], + completion_params: dict[str, object], + a2a_message: Mapping[str, object], params: dict[str, Any] | None = None, ) -> None: """ @@ -104,8 +108,8 @@ class A2ACompletionBridgeTransformation: # ``extra_body.metadata`` so the configured keys remain authoritative # and an A2A caller cannot overwrite server-set run metadata. existing_metadata: Final = extra_body.get("metadata") - existing_dict: Final[dict[str, Any]] = existing_metadata if isinstance(existing_metadata, dict) else {} - merged_metadata: Final[dict[str, Any]] = {**forward_metadata, **existing_dict} + existing_dict: Final[dict[str, object]] = existing_metadata if isinstance(existing_metadata, dict) else {} + merged_metadata: Final[dict[str, object]] = {**forward_metadata, **existing_dict} extra_body = {**extra_body, "metadata": merged_metadata} completion_params["extra_body"] = extra_body @@ -114,7 +118,7 @@ class A2ACompletionBridgeTransformation: @staticmethod def a2a_message_to_openai_messages( a2a_message: dict[str, Any], - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """ Transform an A2A message to OpenAI message format. @@ -124,8 +128,8 @@ class A2ACompletionBridgeTransformation: Returns: List of OpenAI-format messages """ - role: Final = a2a_message.get("role", "user") - parts = a2a_message.get("parts", []) + role: Final[object] = a2a_message.get("role", "user") + parts: Sequence[Mapping[str, object]] = a2a_message.get("parts", []) # Map A2A roles to OpenAI roles openai_role = role @@ -143,7 +147,7 @@ class A2ACompletionBridgeTransformation: # Do not attach A2A message.metadata here — the completion bridge forwards it # once at run level via extra_body.metadata (LangGraph POST /runs/wait shape). - openai_message: Final[dict[str, Any]] = {"role": openai_role, "content": content} + openai_message: Final[dict[str, object]] = {"role": openai_role, "content": content} verbose_logger.debug( "A2A -> OpenAI transform: role=%s -> %s, content_length=%s", role, openai_role, len(content) @@ -155,7 +159,7 @@ class A2ACompletionBridgeTransformation: def openai_response_to_a2a_response( response: Any, request_id: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform a LiteLLM ModelResponse to A2A SendMessageResponse format. @@ -169,7 +173,7 @@ class A2ACompletionBridgeTransformation: # Extract content from response content = "" if hasattr(response, "choices") and response.choices: - choice: Final = response.choices[0] + choice: Final[Choices] = response.choices[0] if hasattr(choice, "message") and choice.message: content = choice.message.content or "" @@ -182,7 +186,7 @@ class A2ACompletionBridgeTransformation: } # Build A2A response - a2a_response: Final = { + a2a_response: Final[dict[str, object]] = { "jsonrpc": "2.0", "id": request_id, "result": a2a_message, @@ -200,7 +204,7 @@ class A2ACompletionBridgeTransformation: @staticmethod def create_task_event( ctx: A2AStreamingContext, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Create the initial task event with status 'submitted'. @@ -235,7 +239,7 @@ class A2ACompletionBridgeTransformation: state: str, final: bool = False, message_text: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Create a status update event. @@ -245,7 +249,7 @@ class A2ACompletionBridgeTransformation: final: Whether this is the final event message_text: Optional message text for 'working' status """ - status: Final[dict[str, Any]] = { + status: Final[dict[str, object]] = { "state": state, "timestamp": A2ACompletionBridgeTransformation._get_timestamp(), } @@ -277,7 +281,7 @@ class A2ACompletionBridgeTransformation: def create_artifact_update_event( ctx: A2AStreamingContext, text: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Create an artifact update event with content. diff --git a/litellm/caching/valkey_semantic_cache.py b/litellm/caching/valkey_semantic_cache.py index 737d212a89d..ac4d4033546 100644 --- a/litellm/caching/valkey_semantic_cache.py +++ b/litellm/caching/valkey_semantic_cache.py @@ -17,6 +17,7 @@ RedisSemanticCache since those are backend agnostic. import asyncio import hashlib import os +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Any, Final @@ -62,7 +63,7 @@ class ValkeySemanticCache(RedisSemanticCache): sync_client: Redis | None = None, async_client: AsyncRedis | None = None, embedding_max_input_tokens: int | None = None, - **kwargs: Any, + **kwargs: object, ): if similarity_threshold is None: raise ValueError("similarity_threshold must be provided, passed None") @@ -115,7 +116,7 @@ class ValkeySemanticCache(RedisSemanticCache): return hashlib.sha256(str(key).encode("utf-8")).hexdigest() @staticmethod - def _embedding_to_bytes(embedding: list[float]) -> bytes: + def _embedding_to_bytes(embedding: Sequence[float]) -> bytes: return pack_vector(embedding) def _index_schema(self, dim: int) -> tuple[TagField, VectorField]: @@ -189,7 +190,9 @@ class ValkeySemanticCache(RedisSemanticCache): def _doc_key(self, key: str) -> str: return f"{self.key_prefix}{self._scope_tag(key)}:{uuid.uuid4()}" - def _doc_mapping(self, key: str, prompt: str, value_str: str, embedding: list[float]) -> dict: + def _doc_mapping( + self, key: str, prompt: str, value_str: str, embedding: Sequence[float] + ) -> Mapping[str | bytes, str | bytes]: return { self.CACHE_KEY_FIELD_NAME: self._scope_tag(key), self.PROMPT_FIELD_NAME: prompt, @@ -205,30 +208,49 @@ class ValkeySemanticCache(RedisSemanticCache): ) return Query(query_string).return_fields(self.RESPONSE_FIELD_NAME, self.DISTANCE_FIELD_NAME).dialect(2) + async def _async_search(self, key: str, embedding: Sequence[float]) -> object: + """Run the KNN query on the async client, stopping the untyped search surface here.""" + return await self.async_client.ft(self.index_name).search( + self._knn_query(key), + query_params={"vec": self._embedding_to_bytes(embedding)}, + ) + @classmethod - def _first_hit(cls, search_result: Any) -> _ValkeyCacheHit | None: - docs: Final = getattr(search_result, "docs", []) + def _first_hit(cls, search_result: object) -> _ValkeyCacheHit | None: + docs: Final[Sequence[object]] = getattr(search_result, "docs", []) if not docs: return None doc: Final = docs[0] + response_field: Final[object] = getattr(doc, cls.RESPONSE_FIELD_NAME) + distance_field: Final[str | bytes | float] = getattr(doc, cls.DISTANCE_FIELD_NAME) return _ValkeyCacheHit( - response=str(getattr(doc, cls.RESPONSE_FIELD_NAME)), - distance=float(getattr(doc, cls.DISTANCE_FIELD_NAME)), + response=str(response_field), + distance=float(distance_field), ) - def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: Any) -> Any: + @staticmethod + def _record_similarity(kwargs: dict[str, Any], similarity: float) -> None: + """Stamp the semantic-similarity score onto the request metadata carried in ``kwargs``.""" + kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity + + @staticmethod + def _embedding_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None: + """The request metadata forwarded to the embedding call.""" + return kwargs.get("metadata") + + def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: object) -> object: if hit is None: - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) return None similarity: Final = 1 - hit.distance - kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity + self._record_similarity(kwargs, similarity) if similarity < self.similarity_threshold: return None return self._get_cache_logic(cached_response=hit.response) - def set_cache(self, key: str, value: Any, **kwargs: Any) -> None: + def set_cache(self, key: str, value: object, **kwargs: object) -> None: print_verbose(f"Valkey semantic-cache set_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) @@ -247,12 +269,12 @@ class ValkeySemanticCache(RedisSemanticCache): except Exception as e: print_verbose(f"Error in Valkey semantic-cache set_cache: {e}") - def get_cache(self, key: str, **kwargs: Any) -> Any: + def get_cache(self, key: str, **kwargs: object) -> object: print_verbose(f"Valkey semantic-cache get_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) if prompt is None: - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) return None embedding: Final = self._get_embedding(prompt) @@ -265,9 +287,9 @@ class ValkeySemanticCache(RedisSemanticCache): return self._resolve_hit(self._first_hit(search_result), key, **kwargs) except Exception as e: print_verbose(f"Error in Valkey semantic-cache get_cache: {e}") - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) - async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> None: + async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: print_verbose(f"Async Valkey semantic-cache set_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) @@ -275,7 +297,7 @@ class ValkeySemanticCache(RedisSemanticCache): print_verbose("No prompt provided for semantic caching") return - embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) + embedding: Final = await self._get_async_embedding(prompt, metadata=self._embedding_metadata(kwargs)) await self._ensure_index_async(len(embedding)) doc_key: Final = self._doc_key(key) @@ -286,31 +308,28 @@ class ValkeySemanticCache(RedisSemanticCache): except Exception as e: print_verbose(f"Error in async Valkey semantic-cache set_cache: {e}") - async def async_get_cache(self, key: str, **kwargs: Any) -> Any: + async def async_get_cache(self, key: str, **kwargs: object) -> object: print_verbose(f"Async Valkey semantic-cache get_cache, kwargs: {kwargs}") try: prompt: Final = self._get_prompt_from_kwargs(**kwargs) if prompt is None: - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) return None - embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) + embedding: Final = await self._get_async_embedding(prompt, metadata=self._embedding_metadata(kwargs)) await self._ensure_index_async(len(embedding)) - search_result: Final = await self.async_client.ft(self.index_name).search( - self._knn_query(key), - query_params={"vec": self._embedding_to_bytes(embedding)}, - ) + search_result: Final[object] = await self._async_search(key, embedding) return self._resolve_hit(self._first_hit(search_result), key, **kwargs) except Exception as e: print_verbose(f"Error in async Valkey semantic-cache get_cache: {e}") - kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + self._record_similarity(kwargs, 0.0) - async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: Any) -> None: + async def async_set_cache_pipeline(self, cache_list: list[tuple[str, object]], **kwargs: object) -> None: try: await asyncio.gather(*[self.async_set_cache(key, value, **kwargs) for key, value in cache_list]) except Exception as e: print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {e}") - async def _index_info(self) -> dict: + async def _index_info(self) -> Mapping[str, object]: return await self.async_client.ft(self.index_name).info() diff --git a/litellm/integrations/bitbucket/bitbucket_client.py b/litellm/integrations/bitbucket/bitbucket_client.py index e06e5ab358f..c6967ef3340 100644 --- a/litellm/integrations/bitbucket/bitbucket_client.py +++ b/litellm/integrations/bitbucket/bitbucket_client.py @@ -4,11 +4,33 @@ BitBucket API client for fetching .prompt files from BitBucket repositories. import base64 import urllib.parse -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final + +from typing_extensions import ReadOnly, TypedDict from litellm.llms.custom_httpx.http_handler import HTTPHandler +class BitBucketSrcEntry(TypedDict, total=False): + """One entry of a BitBucket ``src`` directory listing.""" + + type: ReadOnly[str] + path: ReadOnly[str] + + +class BitBucketSrcListing(TypedDict, total=False): + """A page of a BitBucket ``src`` directory listing.""" + + values: ReadOnly[Sequence[BitBucketSrcEntry]] + + +class BitBucketBranchListing(TypedDict, total=False): + """A page of a BitBucket ``refs/branches`` listing.""" + + values: ReadOnly[Sequence[Mapping[str, object]]] + + def _sanitize_file_path(file_path: str) -> str: """Reject path traversal and URL-encode each path segment.""" if "#" in file_path or "?" in file_path: @@ -31,7 +53,7 @@ class BitBucketClient: - Branch-specific file fetching """ - def __init__(self, config: dict[str, Any]): + def __init__(self, config: Mapping[str, object]): """ Initialize the BitBucket client. @@ -135,16 +157,13 @@ class BitBucketClient: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data: Final = response.json() - files: Final = [] + data: Final[BitBucketSrcListing] = response.json() - for item in data.get("values", []): - if item.get("type") == "commit_file": - file_path = item.get("path", "") - if file_path.endswith(file_extension): - files.append(file_path) - - return files + return [ + file_path + for item in data.get("values", []) + if item.get("type") == "commit_file" and (file_path := item.get("path", "")).endswith(file_extension) + ] except Exception as e: # Check if it's an HTTP error @@ -162,7 +181,7 @@ class BitBucketClient: else: raise Exception(f"Error listing files in '{directory_path}': {e}") - def get_repository_info(self) -> dict[str, Any]: + def get_repository_info(self) -> Mapping[str, object]: """ Get information about the repository. @@ -191,7 +210,7 @@ class BitBucketClient: except Exception: return False - def get_branches(self) -> list[dict[str, Any]]: + def get_branches(self) -> Sequence[Mapping[str, object]]: """ Get list of branches in the repository. @@ -204,12 +223,12 @@ class BitBucketClient: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data: Final = response.json() + data: Final[BitBucketBranchListing] = response.json() return data.get("values", []) except Exception as e: raise Exception(f"Failed to get branches: {e}") - def get_file_metadata(self, file_path: str) -> dict[str, Any] | None: + def get_file_metadata(self, file_path: str) -> Mapping[str, object] | None: """ Get metadata about a file (size, last modified, etc.). diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index 7ea60053e6f..219718f6771 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -7,7 +7,10 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan. import time import uuid -from typing import Any, Final, cast +from collections.abc import Mapping, Sequence +from typing import Any, Final, Protocol, cast + +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.compression import compress @@ -26,6 +29,19 @@ LITELLM_CONTENT_RETRIEVE_TOOL_NAME: Final = "litellm_content_retrieve" _CACHE_TTL_SECONDS: Final = 15 * 60 +class _AgenticLoopParams(TypedDict, total=False): + """The ``agentic_loop_params`` entry the agentic loop driver records on the logging object.""" + + model: ReadOnly[str] + + +class _AgenticLoopLoggingObj(Protocol): + """Logging object view exposing the untyped call details this handler reads.""" + + @property + def model_call_details(self) -> Mapping[str, _AgenticLoopParams]: ... + + def _compression_savings_from_counts( original_tokens: object, compressed_tokens: object ) -> CompressionSavingsMetadata | None: @@ -78,7 +94,7 @@ class CompressionInterceptionLogger(CustomLogger): compression_trigger: int = 200_000, compression_target: int | None = None, embedding_model: str | None = None, - embedding_model_params: dict[str, Any] | None = None, + embedding_model_params: dict[str, object] | None = None, ): super().__init__() self.enabled = enabled @@ -101,7 +117,7 @@ class CompressionInterceptionLogger(CustomLogger): @staticmethod def initialize_from_proxy_config( litellm_settings: dict[str, Any], - callback_specific_params: dict[str, Any], + callback_specific_params: Mapping[str, object], ) -> "CompressionInterceptionLogger": compression_params: CompressionInterceptionConfig = {} if "compression_interception_params" in litellm_settings: @@ -115,7 +131,9 @@ class CompressionInterceptionLogger(CustomLogger): ) return CompressionInterceptionLogger.from_config_yaml(compression_params) - async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None: + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, Any], call_type: CallTypes | None + ) -> dict[str, object] | None: if not self.enabled: return None if call_type is not None and call_type != CallTypes.anthropic_messages: @@ -145,7 +163,7 @@ class CompressionInterceptionLogger(CustomLogger): cache: Final = cast(dict[str, str], compressed.get("cache", {})) skip_reason: Final = cast(str | None, compressed.get("compression_skipped_reason")) - compressed_tools: Final = cast(list[dict[str, Any]], compressed.get("tools", [])) + compressed_tools: Final = cast(list[dict[str, object]], compressed.get("tools", [])) # Only mutate kwargs when compression actually produced a result. # If compression was a no-op (below trigger, invalid tool sequence, etc.), @@ -156,7 +174,7 @@ class CompressionInterceptionLogger(CustomLogger): kwargs["messages"] = compressed["messages"] if compressed_tools: kwargs["tools"] = self._merge_tools( - existing_tools=cast(list[dict[str, Any]] | None, kwargs.get("tools")), + existing_tools=cast(list[dict[str, object]] | None, kwargs.get("tools")), compressed_tools=compressed_tools, ) call_id = cast(str | None, kwargs.get("litellm_call_id")) @@ -189,14 +207,14 @@ class CompressionInterceptionLogger(CustomLogger): async def async_should_run_agentic_loop( self, - response: Any, + response: object, model: str, - messages: list[dict], - tools: list[dict] | None, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, stream: bool, custom_llm_provider: str, - kwargs: dict, - ) -> tuple[bool, dict]: + kwargs: Mapping[str, object], + ) -> tuple[bool, dict[str, object]]: if not self.enabled: return False, {} if not self._has_retrieval_tool(tools): @@ -214,19 +232,19 @@ class CompressionInterceptionLogger(CustomLogger): async def async_build_agentic_loop_plan( self, - tools: dict, + tools: Mapping[str, object], model: str, - messages: list[dict], - response: Any, - anthropic_messages_provider_config: Any, - anthropic_messages_optional_request_params: dict, - logging_obj: Any, + messages: list[dict[str, object]], + response: object, + anthropic_messages_provider_config: object, + anthropic_messages_optional_request_params: Mapping[str, object], + logging_obj: _AgenticLoopLoggingObj | None, stream: bool, - kwargs: dict, + kwargs: Mapping[str, object], ) -> AgenticLoopPlan: self._prune_expired_cache() - tool_calls: Final = cast(list[dict[str, Any]], tools.get("tool_calls", [])) - thinking_blocks: Final = cast(list[dict[str, Any]], tools.get("thinking_blocks", [])) + tool_calls: Final = cast(list[dict[str, object]], tools.get("tool_calls", [])) + thinking_blocks: Final = cast(list[dict[str, object]], tools.get("thinking_blocks", [])) call_id: Final = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs) cache: Final = self._get_cache(call_id=call_id) @@ -269,7 +287,7 @@ class CompressionInterceptionLogger(CustomLogger): full_model_name = model if logging_obj is not None: agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {}) - full_model_name = cast(str, agentic_params.get("model", model)) + full_model_name = agentic_params.get("model", model) request_patch: Final = AgenticLoopRequestPatch( model=full_model_name, @@ -304,15 +322,15 @@ class CompressionInterceptionLogger(CustomLogger): return {} return cache_entry[0] - def _resolve_call_id(self, logging_obj: Any, kwargs: dict[str, Any]) -> str | None: + def _resolve_call_id(self, logging_obj: _AgenticLoopLoggingObj | None, kwargs: Mapping[str, object]) -> str | None: if logging_obj is not None: logging_call_id: Final = getattr(logging_obj, "litellm_call_id", None) if isinstance(logging_call_id, str) and logging_call_id: return logging_call_id kwargs_call_id: Final = kwargs.get("litellm_call_id") - return cast(str | None, kwargs_call_id if isinstance(kwargs_call_id, str) else None) + return kwargs_call_id if isinstance(kwargs_call_id, str) else None - def _resolve_retrieval_content(self, tool_call: dict[str, Any], cache: dict[str, str]) -> str: + def _resolve_retrieval_content(self, tool_call: Mapping[str, object], cache: Mapping[str, str]) -> str: raw_input: Final = tool_call.get("input", {}) key = "" if isinstance(raw_input, dict): @@ -323,7 +341,9 @@ class CompressionInterceptionLogger(CustomLogger): return cache[key] return f"[compressed content key '{key}' not found]" - def _extract_retrieval_tool_calls(self, response: Any) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + def _extract_retrieval_tool_calls( + self, response: object + ) -> tuple[list[dict[str, object]], list[dict[str, object]]]: if isinstance(response, dict): content = response.get("content", []) else: @@ -332,8 +352,8 @@ class CompressionInterceptionLogger(CustomLogger): if not isinstance(content, list): return [], [] - tool_calls: Final[list[dict[str, Any]]] = [] - thinking_blocks: Final[list[dict[str, Any]]] = [] + tool_calls: Final[list[dict[str, object]]] = [] + thinking_blocks: Final[list[dict[str, object]]] = [] for block in content: if isinstance(block, dict): @@ -380,13 +400,13 @@ class CompressionInterceptionLogger(CustomLogger): return tool_calls, thinking_blocks - def _prepare_followup_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: + def _prepare_followup_kwargs(self, kwargs: Mapping[str, object]) -> dict[str, object]: internal_keys: Final = {"litellm_logging_obj"} return { k: v for k, v in kwargs.items() if not k.startswith("_compression_interception") and k not in internal_keys } - def _has_retrieval_tool(self, tools: Any) -> bool: + def _has_retrieval_tool(self, tools: object) -> bool: if not isinstance(tools, list): return False for tool in tools: @@ -402,9 +422,9 @@ class CompressionInterceptionLogger(CustomLogger): def _merge_tools( self, - existing_tools: list[dict[str, Any]] | None, - compressed_tools: list[dict[str, Any]], - ) -> list[dict[str, Any]]: + existing_tools: Sequence[Mapping[str, object]] | None, + compressed_tools: Sequence[Mapping[str, object]], + ) -> list[Mapping[str, object]]: merged: Final = list(existing_tools or []) if self._has_retrieval_tool(merged): return merged diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index a0c78674ac8..417edc77e9c 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -2,7 +2,7 @@ # On success, logs events to Promptlayer import re import traceback -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Sequence from typing import TYPE_CHECKING, Any, Final, Optional from pydantic import BaseModel @@ -103,11 +103,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return [] callbacks: Final = AllCallbacks() - callback_info: Final = getattr(callbacks, lookup_name, None) + callback_info: Final[object] = getattr(callbacks, lookup_name, None) if callback_info is None: return [] - params: Final = getattr(callback_info, "litellm_callback_params", None) + params: Final[Sequence[str] | None] = getattr(callback_info, "litellm_callback_params", None) if not params: return [] @@ -783,7 +783,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac - Converting to string and then truncating the logged content catches this 2. We want to avoid modifying the original `messages`, `response`, and `error_str` in the logging payload since these are in kwargs and could be returned to the user """ - field_value: Final = standard_logging_object.get(field_name) + field_value: Final[object] = standard_logging_object.get(field_name) if field_value: str_value: Final = str(field_value) if len(str_value) > max_length: @@ -937,8 +937,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac • Keep untyped or text content. • Recursively redact inline base64 blobs in *any* string field, at any depth. """ - raw_messages: Final[Any] = payload.get("messages", []) - messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else [] + raw_messages: Final[object] = payload.get("messages", []) + messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else [] verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages)) if messages: @@ -969,8 +969,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac • Keep untyped or text content. • Recursively redact inline base64 blobs in *any* string field, at any depth. """ - raw_messages: Final[Any] = payload.get("messages", []) - messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else [] + raw_messages: Final[object] = payload.get("messages", []) + messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else [] verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages)) if messages: @@ -991,7 +991,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac value: Any, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, - ) -> Any: + ) -> object: """Recursively redact inline base64 from any nested structure with a max recursion depth limit.""" if depth > max_depth: verbose_logger.warning("[CustomLogger] Max recursion depth %s reached while redacting base64", max_depth) @@ -1022,16 +1022,16 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac def _process_messages( self, - messages: list[Any], + messages: list[object], max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, - ) -> list[dict[str, Any]]: - filtered_messages: Final[list[dict[str, Any]]] = [] + ) -> list[dict[str, object]]: + filtered_messages: Final[list[dict[str, object]]] = [] for msg in messages: if not isinstance(msg, dict): continue - contents: Any = msg.get("content") + contents: object = msg.get("content") if isinstance(contents, list): - cleaned: list[Any] = [] + cleaned: list[object] = [] for c in contents: if self._should_keep_content(content=c): cleaned.append(self._redact_base64(value=c, max_depth=max_depth)) diff --git a/litellm/integrations/opik/opik_payload_builder/payload_builders.py b/litellm/integrations/opik/opik_payload_builder/payload_builders.py index e40d72ea542..855b84ba4c8 100644 --- a/litellm/integrations/opik/opik_payload_builder/payload_builders.py +++ b/litellm/integrations/opik/opik_payload_builder/payload_builders.py @@ -17,12 +17,12 @@ def build_trace_payload( end_time: datetime, input_data: Any, output_data: Any, - metadata: dict[str, Any], + metadata: dict[str, object], tags: list[str], thread_id: str | None, ) -> types.TracePayload: """Build a complete trace payload.""" - trace_name: Final = response_obj.get("object", "unknown type") + trace_name: Final[str] = response_obj.get("object", "unknown type") return types.TracePayload( project_name=project_name, @@ -47,7 +47,7 @@ def build_span_payload( end_time: datetime, input_data: Any, output_data: Any, - metadata: dict[str, Any], + metadata: dict[str, object], tags: list[str], usage: dict[str, int], provider: str | None = None, @@ -56,9 +56,9 @@ def build_span_payload( """Build a complete span payload.""" span_id: Final = utils.create_uuid7() - model: Final = response_obj.get("model", "unknown-model") - obj_type: Final = response_obj.get("object", "unknown-object") - created: Final = response_obj.get("created", 0) + model: Final[str] = response_obj.get("model", "unknown-model") + obj_type: Final[str] = response_obj.get("object", "unknown-object") + created: Final[int] = response_obj.get("created", 0) span_name: Final = f"{model}_{obj_type}_{created}" _logging.verbose_logger.debug("OpikLogger creating span with id %s for trace %s", span_id, trace_id) diff --git a/litellm/integrations/prometheus_helpers/prometheus_api.py b/litellm/integrations/prometheus_helpers/prometheus_api.py index 9f77f87a670..f677db17648 100644 --- a/litellm/integrations/prometheus_helpers/prometheus_api.py +++ b/litellm/integrations/prometheus_helpers/prometheus_api.py @@ -4,9 +4,13 @@ Helper functions to query prometheus API import json import time +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta from typing import Final +from httpx import Response +from typing_extensions import ReadOnly, TypedDict + from litellm import get_secret from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( @@ -19,9 +23,45 @@ PROMETHEUS_SELECTED_INSTANCE: Final[str | None] = get_secret("PROMETHEUS_SELECTE async_http_handler: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) +class _PrometheusSeries(TypedDict): + """One time series in a Prometheus query result, with its labels and its ``[timestamp, value]`` samples.""" + + metric: ReadOnly[Mapping[str, str]] + values: ReadOnly[Sequence[Sequence[float | str]]] + + +class _PrometheusResultData(TypedDict): + """The ``data`` envelope of a Prometheus query response.""" + + result: ReadOnly[Sequence[_PrometheusSeries]] + + +class _PrometheusQueryResponse(TypedDict): + """The JSON body returned by the Prometheus ``/api/v1/query`` and ``/api/v1/query_range`` endpoints.""" + + data: ReadOnly[_PrometheusResultData] + + +class _DailySpend(TypedDict): + """One day of spend, in the shape ``get_daily_spend_from_prometheus`` returns.""" + + date: ReadOnly[str] + spend: ReadOnly[float] + + +def _query_body(response: Response) -> _PrometheusQueryResponse: + """Read the untyped JSON body of a Prometheus query response.""" + return response.json() + + +def _query_series(response: Response) -> Sequence[_PrometheusSeries]: + """Read the time series list out of the untyped JSON body of a Prometheus query response.""" + return response.json()["data"]["result"] + + async def get_metric_from_prometheus( metric_name: str, -): +) -> Sequence[_PrometheusSeries]: # Get the start of the current day in Unix timestamp if PROMETHEUS_URL is None: raise ValueError("PROMETHEUS_URL not set please set 'PROMETHEUS_URL=<>' in .env") @@ -31,13 +71,13 @@ async def get_metric_from_prometheus( response: Final = await async_http_handler.get( f"{PROMETHEUS_URL}/api/v1/query", params={"query": query, "time": now} ) # End of the day - _json_response: Final = response.json() + _json_response: Final = _query_body(response) verbose_logger.debug("json response from prometheus /query api %s", _json_response) - results: Final = response.json()["data"]["result"] + results: Final = _query_series(response) return results -async def get_fallback_metric_from_prometheus(): +async def get_fallback_metric_from_prometheus() -> str: """ Gets fallback metrics from prometheus for the last 24 hours """ @@ -96,7 +136,7 @@ def _quote_promql_string_literal(value: str) -> str: return json.dumps(value, ensure_ascii=False) -async def get_daily_spend_from_prometheus(api_key: str | None): +async def get_daily_spend_from_prometheus(api_key: str | None) -> Sequence[_DailySpend]: """ Expected Response Format: [ @@ -133,9 +173,9 @@ async def get_daily_spend_from_prometheus(api_key: str | None): } response: Final = await async_http_handler.get(url, params=params) - _json_response: Final = response.json() + _json_response: Final = _query_body(response) verbose_logger.debug("json response from prometheus /query api %s", _json_response) - results: Final = response.json()["data"]["result"] + results: Final = _query_series(response) formatted_results: Final = [] for result in results: diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index e59ef0449d0..e1ec3360351 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -10,7 +10,7 @@ import asyncio import math import uuid from collections.abc import AsyncIterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, TypedDict, TypeVar, cast from typing_extensions import ReadOnly @@ -74,6 +74,10 @@ WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: Final = "_websearch_interception_emit_native_b # ``web_search_tool_result`` blocks to inject into the final response. WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks" +_RESPONSE_CONTENT_FIELD: Final = "content" + +_ResponseT: Final = TypeVar("_ResponseT") + class _PlanMetadataView(TypedDict): websearch_native_blocks: Sequence[Mapping[str, object]] | None @@ -926,17 +930,17 @@ class WebSearchInterceptionLogger(CustomLogger): ) @staticmethod - def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any: + def _inject_native_blocks(response: _ResponseT, native_blocks: Sequence[Mapping[str, object]]) -> _ResponseT: """Prepend native blocks to response content, dict or object form.""" if not native_blocks: return response if isinstance(response, dict): - existing = response.get("content") or [] - response["content"] = list(native_blocks) + list(existing) + existing = response.get(_RESPONSE_CONTENT_FIELD) or [] + response[_RESPONSE_CONTENT_FIELD] = list(native_blocks) + list(existing) return response - existing = getattr(response, "content", None) or [] + existing = getattr(response, _RESPONSE_CONTENT_FIELD, None) or [] try: - response.content = list(native_blocks) + list(existing) + setattr(response, _RESPONSE_CONTENT_FIELD, list(native_blocks) + list(existing)) except (AttributeError, TypeError): # Object refused write — fall through and leave the response # untouched rather than crash the request. diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index 44fed944d2a..1737e2b8cb0 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -1,6 +1,9 @@ import datetime +from collections.abc import Mapping from typing import Any, Final +import httpx + from litellm.constants import LITELLM_DETAILED_TIMING from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base @@ -25,11 +28,7 @@ class ResponseMetadata: @property def supports_response_time(self) -> bool: """Check if response type supports timing metrics""" - return ( - isinstance(self.result, ModelResponse) - or isinstance(self.result, EmbeddingResponse) - or isinstance(self.result, TranscriptionResponse) - ) + return isinstance(self.result, (ModelResponse, EmbeddingResponse, TranscriptionResponse)) def set_hidden_params(self, logging_obj: LiteLLMLoggingObject, model: str | None, kwargs: dict) -> None: """Set hidden parameters on the response""" @@ -45,14 +44,14 @@ class ResponseMetadata: result=self.result, litellm_model_name=model, router_model_id=model_id ), "additional_headers": process_response_headers( - self._get_value_from_hidden_params("additional_headers") or {}, + self._get_additional_headers_from_hidden_params() or {}, preserve_litellm_internal_headers=True, ), "litellm_model_name": model, } self._update_hidden_params(new_params) - def _update_hidden_params(self, new_params: dict) -> None: + def _update_hidden_params(self, new_params: Mapping[str, object]) -> None: """ Update hidden params - handles when self._hidden_params is a dict or HiddenParams object """ @@ -64,12 +63,12 @@ class ResponseMetadata: for key, value in new_params.items(): setattr(self._hidden_params, key, value) - def _get_value_from_hidden_params(self, key: str) -> Any | None: - """Get value from hidden params - handles when self._hidden_params is a dict or HiddenParams object""" + def _get_additional_headers_from_hidden_params(self) -> httpx.Headers | dict[str, str] | None: + """Get `additional_headers` from hidden params - handles when self._hidden_params is a dict or HiddenParams object""" if isinstance(self._hidden_params, dict): - return self._hidden_params.get(key, None) + return self._hidden_params.get("additional_headers", None) elif isinstance(self._hidden_params, HiddenParams): - return getattr(self._hidden_params, key, None) + return getattr(self._hidden_params, "additional_headers", None) def set_timing_metrics( self, @@ -96,7 +95,7 @@ class ResponseMetadata: ######################################################### # 2. Add LiteLLM overhead duration ######################################################### - llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms") + llm_api_duration_ms: Final[float | None] = logging_obj.model_call_details.get("llm_api_duration_ms") if llm_api_duration_ms is not None: overhead_ms = round(total_response_time_ms - llm_api_duration_ms, 4) self._update_hidden_params( @@ -108,7 +107,7 @@ class ResponseMetadata: ######################################################### # 3. Add callback processing duration ######################################################### - callback_duration_ms: Final = getattr(logging_obj, "callback_duration_ms", None) + callback_duration_ms: Final[float | None] = getattr(logging_obj, "callback_duration_ms", None) if callback_duration_ms is not None: self._update_hidden_params( { @@ -136,17 +135,17 @@ class ResponseMetadata: # 5. Detailed per-phase timing (opt-in via env var) ######################################################### if LITELLM_DETAILED_TIMING and llm_api_duration_ms is not None: - detailed: Final[dict] = { + detailed: Final[dict[str, float]] = { "timing_llm_api_ms": round(llm_api_duration_ms, 4), } # message copy time from Logging.__init__() - msg_copy_ms: Final = getattr(logging_obj, "message_copy_duration_ms", None) + msg_copy_ms: Final[float | None] = getattr(logging_obj, "message_copy_duration_ms", None) if msg_copy_ms is not None: detailed["timing_message_copy_ms"] = round(msg_copy_ms, 4) # pre-processing = time from request start to LLM API call start - api_call_start: Final = logging_obj.model_call_details.get("api_call_start_time") + api_call_start: Final[datetime.datetime | None] = logging_obj.model_call_details.get("api_call_start_time") if api_call_start is not None and start_time is not None: pre_ms: Final = (api_call_start - start_time).total_seconds() * 1000 detailed["timing_pre_processing_ms"] = round(pre_ms, 4) diff --git a/litellm/litellm_core_utils/model_response_utils.py b/litellm/litellm_core_utils/model_response_utils.py index ea4be1c856f..7d412229a96 100644 --- a/litellm/litellm_core_utils/model_response_utils.py +++ b/litellm/litellm_core_utils/model_response_utils.py @@ -2,10 +2,54 @@ Utility functions for ModelResponse and ModelResponseStream objects. """ -from typing import Any, Final +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final, Protocol from litellm.types.utils import Delta, ModelResponseBase, ModelResponseStream +_NO_EXTRA_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +class _HasModelExtra(Protocol): + """A Pydantic model, seen through the extra fields it collected.""" + + @property + def model_extra(self) -> Mapping[str, object] | None: ... + + +class _StreamingChoice(_HasModelExtra, Protocol): + """The streaming-choice fields this emptiness check reads.""" + + @property + def finish_reason(self) -> object: ... + + @property + def logprobs(self) -> object: ... + + @property + def enhancements(self) -> object: ... + + @property + def delta(self) -> Delta | None: ... + + +def _extra_fields(model: _HasModelExtra) -> Mapping[str, object]: + """The dynamically added fields Pydantic stored on ``model``.""" + return model.model_extra or _NO_EXTRA_FIELDS + + +def _attribute(obj: object, name: str) -> object: + """The named attribute of ``obj``, or ``None`` when it is absent.""" + attribute: Final[object] = getattr(obj, name, None) + return attribute + + +def _has_callable_attribute(obj: object, name: str) -> bool: + """Whether the named attribute of ``obj`` is callable.""" + attribute: Final[object] = getattr(obj, name) + return callable(attribute) + def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: """ @@ -41,7 +85,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: # Check model_extra for dynamically added fields (this is where Pydantic stores them) if hasattr(model_response, "model_extra") and model_response.model_extra: - for extra_field_name, extra_field_value in model_response.model_extra.items(): + for extra_field_name, extra_field_value in _extra_fields(model_response).items(): if _has_meaningful_content(extra_field_value): return False @@ -57,7 +101,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: continue # Check if any other field has meaningful content - model_response_value = getattr(model_response, model_response_field, None) + model_response_value = _attribute(model_response, model_response_field) if _has_meaningful_content(model_response_value): return False @@ -71,7 +115,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: return True -def _has_meaningful_content(value: Any) -> bool: +def _has_meaningful_content(value: object) -> bool: """ Check if a value contains meaningful content. @@ -102,7 +146,7 @@ def _has_meaningful_content(value: Any) -> bool: return True -def _is_choice_non_empty(choice: Any) -> bool: +def _is_choice_non_empty(choice: _StreamingChoice) -> bool: """ Deep check if a choice contains any meaningful content. @@ -131,7 +175,7 @@ def _is_choice_non_empty(choice: Any) -> bool: # Check model_extra for dynamically added fields on the choice if hasattr(choice, "model_extra") and choice.model_extra: - for extra_field_name, extra_field_value in choice.model_extra.items(): + for extra_field_name, extra_field_value in _extra_fields(choice).items(): # Skip certain structural fields that are just default/None placeholders if extra_field_name == "index" and extra_field_value == 0: continue @@ -147,7 +191,7 @@ def _is_choice_non_empty(choice: Any) -> bool: # Skip private attributes, methods, and known empty fields if ( attr_name.startswith("_") - or callable(getattr(choice, attr_name)) + or _has_callable_attribute(choice, attr_name) or attr_name.startswith("model_") or attr_name in { @@ -160,7 +204,7 @@ def _is_choice_non_empty(choice: Any) -> bool: ): continue - attr_value = getattr(choice, attr_name, None) + attr_value = _attribute(choice, attr_name) if _has_meaningful_content(attr_value): return True @@ -179,7 +223,7 @@ def _is_delta_non_empty(delta: Delta) -> bool: """ # Check model_extra for dynamically added fields (this is where Pydantic stores them) if hasattr(delta, "model_extra") and delta.model_extra: - for extra_field_name, extra_field_value in delta.model_extra.items(): + for extra_field_name, extra_field_value in _extra_fields(delta).items(): # Even structural fields are meaningful if they have actual content if _has_meaningful_content(extra_field_value): return True @@ -187,10 +231,10 @@ def _is_delta_non_empty(delta: Delta) -> bool: # Check all regular attributes of the delta object for attr_name in dir(delta): # Skip private attributes, methods, and Pydantic-specific fields - if attr_name.startswith("_") or callable(getattr(delta, attr_name)) or attr_name.startswith("model_"): + if attr_name.startswith("_") or _has_callable_attribute(delta, attr_name) or attr_name.startswith("model_"): continue - attr_value = getattr(delta, attr_name, None) + attr_value = _attribute(delta, attr_name) if _has_meaningful_content(attr_value): return True diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 485091bccd0..300ba427cda 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -92,7 +92,7 @@ def print_verbose(print_statement: object): @dataclass(frozen=True, slots=True) class _ProviderChunkParsed: - response_obj: dict[str, Any] + response_obj: dict[str, object] @dataclass(frozen=True, slots=True) @@ -1242,7 +1242,7 @@ class CustomStreamWrapper: for key, value in anthropic_response_obj["provider_specific_fields"].items(): setattr(model_response, key, value) - response_obj = cast(dict[str, Any], anthropic_response_obj) + response_obj = cast(dict[str, object], anthropic_response_obj) elif self.model == "replicate" or self.custom_llm_provider == "replicate": response_obj = self.handle_replicate_chunk(chunk) completion_obj["content"] = response_obj["text"] @@ -1398,7 +1398,7 @@ class CustomStreamWrapper: if not isinstance(chunk, str): raise ValueError(f"chunk is not a string: {chunk}") response_obj = cast( - dict[str, Any], + dict[str, object], litellm.CodestralTextCompletionConfig()._chunk_parser(chunk), ) completion_obj["content"] = response_obj["text"] @@ -2462,7 +2462,7 @@ def calculate_total_usage(chunks: list[ModelResponse]) -> Usage: prompt_tokens: int = 0 completion_tokens: int = 0 - latest_usage_chunk = None + latest_usage_chunk: Usage | Mapping[str, int] | None = None prompt_tokens_details: PromptTokensDetailsWrapper | None = None completion_tokens_details: CompletionTokensDetailsWrapper | None = None cache_creation_token_details: CacheCreationTokenDetails | None = None diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 0a59eaa75d3..53ae25b928c 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -20,6 +20,7 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config): """ import socket +from collections.abc import Sequence from ipaddress import ip_address, ip_network from typing import Any, Final from urllib.parse import quote, urlparse, urlunparse @@ -44,7 +45,7 @@ class SSRFError(ValueError): """Raised when a URL targets a blocked network.""" -def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") -> str: +def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str: """Percent-encode one user-controlled URL path segment. ``urllib.parse.quote(..., safe="")`` intentionally leaves RFC 3986 @@ -64,7 +65,7 @@ def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") - return quote(value_str, safe="") -def encode_url_path_segments(value: Any, *, field_name: str = "path") -> str: +def encode_url_path_segments(value: object, *, field_name: str = "path") -> str: """Percent-encode a user-controlled URL path made of multiple segments. Empty segments are rejected, so leading, trailing, or consecutive slashes @@ -77,9 +78,9 @@ def encode_url_path_segments(value: Any, *, field_name: str = "path") -> str: if value_str == "": raise ValueError(f"{field_name} is required") - encoded_segments: Final = [] - for segment in value_str.split("/"): - encoded_segments.append(encode_url_path_segment(segment, field_name=field_name)) + encoded_segments: Final = tuple( + encode_url_path_segment(segment, field_name=field_name) for segment in value_str.split("/") + ) return "/".join(encoded_segments) @@ -202,7 +203,7 @@ def _format_host_header(hostname: str, port: int, default_port: int) -> str: return f"{bracketed}:{port}" -def _sockaddr_host(sockaddr: Any) -> str: +def _sockaddr_host(sockaddr: Sequence[object]) -> str: """Return the host element of a ``getaddrinfo`` sockaddr as ``str``. ``getaddrinfo`` with ``IPPROTO_TCP`` returns AF_INET / AF_INET6 sockaddrs @@ -285,7 +286,7 @@ def validate_url(url: str) -> tuple[str, str]: raise SSRFError(f"No addresses found for '{hostname}'") if not is_allowlisted: - for family, type_, proto, canonname, sockaddr in addrinfo: + for _family, _type, _proto, _canonname, sockaddr in addrinfo: resolved_ip = _sockaddr_host(sockaddr) if _is_blocked_ip(resolved_ip): raise SSRFError( @@ -363,7 +364,7 @@ def assert_same_origin(candidate_url: str, expected_url: str) -> None: _MAX_REDIRECTS: Final = 10 -def _extract_redirect_url(response: Any, request_url: str) -> str: +def _extract_redirect_url(response: httpx.Response, request_url: str) -> str: """Extract and resolve the redirect target from a response's Location header.""" location: Final = response.headers.get("location") if not isinstance(location, str) or not location: @@ -372,7 +373,7 @@ def _extract_redirect_url(response: Any, request_url: str) -> str: return str(httpx.URL(request_url).join(location)) -def safe_get(client: Any, url: str, **kwargs: Any) -> Any: +def safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: """ Fetch a user-supplied URL with SSRF protection on every redirect hop. @@ -398,7 +399,7 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any: caller_headers: Final = kwargs.pop("headers", {}) for _ in range(_MAX_REDIRECTS): validated_url, original_host = validate_url(url) - response = client.get( + response: httpx.Response = client.get( validated_url, headers={**caller_headers, "Host": original_host}, follow_redirects=False, @@ -412,7 +413,7 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any: raise SSRFError("Too many redirects") -async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any: +async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: """Async version of safe_get.""" if not getattr(litellm, "user_url_validation", True): kwargs.setdefault("follow_redirects", True) @@ -421,7 +422,7 @@ async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any: caller_headers: Final = kwargs.pop("headers", {}) for _ in range(_MAX_REDIRECTS): validated_url, original_host = validate_url(url) - response = await client.get( + response: httpx.Response = await client.get( validated_url, headers={**caller_headers, "Host": original_host}, follow_redirects=False, diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index 3f8fd2c27f4..3071c3ddd58 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -1,9 +1,11 @@ import json import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal, cast import httpx from httpx import Headers, Response +from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig @@ -19,6 +21,29 @@ else: LoggingClass = Any +class AnthropicBatchRequestCounts(TypedDict, total=False): + """The ``request_counts`` object of an Anthropic Message Batch.""" + + processing: ReadOnly[int] + succeeded: ReadOnly[int] + errored: ReadOnly[int] + canceled: ReadOnly[int] + expired: ReadOnly[int] + + +class AnthropicMessageBatch(TypedDict, total=False): + """The fields of an Anthropic Message Batch that map onto an OpenAI Batch.""" + + id: ReadOnly[str] + processing_status: ReadOnly[str] + created_at: ReadOnly[str | None] + ended_at: ReadOnly[str | None] + expires_at: ReadOnly[str | None] + cancel_initiated_at: ReadOnly[str | None] + archived_at: ReadOnly[str | None] + request_counts: ReadOnly[AnthropicBatchRequestCounts] + + class AnthropicBatchesConfig(BaseBatchesConfig): def __init__(self): from ..chat.transformation import AnthropicConfig @@ -83,7 +108,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): create_batch_data: CreateBatchRequest, optional_params: dict, litellm_params: dict, - ) -> bytes | str | dict[str, Any]: + ) -> bytes | str | dict[str, object]: """ Transform the batch creation request to Anthropic format. @@ -133,7 +158,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): batch_id: str, optional_params: dict, litellm_params: dict, - ) -> bytes | str | dict[str, Any]: + ) -> bytes | str | dict[str, object]: """ Transform batch retrieval request for Anthropic. @@ -152,7 +177,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): ) -> LiteLLMBatch: """Transform Anthropic MessageBatch retrieval response to LiteLLM format.""" try: - response_data: Final = raw_response.json() + response_data: Final[AnthropicMessageBatch] = raw_response.json() except Exception as e: raise ValueError(f"Failed to parse Anthropic batch response: {e}") @@ -161,18 +186,20 @@ class AnthropicBatchesConfig(BaseBatchesConfig): processing_status: Final = response_data.get("processing_status", "in_progress") # Map Anthropic processing_status to OpenAI status - status_mapping: dict[ - str, - Literal[ - "validating", - "failed", - "in_progress", - "finalizing", - "completed", - "expired", - "cancelling", - "cancelled", - ], + status_mapping: Final[ + Mapping[ + str, + Literal[ + "validating", + "failed", + "in_progress", + "finalizing", + "completed", + "expired", + "cancelling", + "cancelled", + ], + ] ] = { "in_progress": "in_progress", "canceling": "cancelling", @@ -279,7 +306,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): if not line: continue try: - response_json = json.loads(line) + response_json: Mapping[str, Mapping[str, dict[str, object]]] = json.loads(line) # Update model_response with the parsed JSON completion_response = response_json["result"]["message"] transformed_response = self.anthropic_chat_config.transform_parsed_response( diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py index 41795fa0f32..8e06b73a92c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py @@ -1,8 +1,10 @@ """Dispatch ``context_management`` edits to registered polyfill editors.""" import inspect -from collections.abc import Awaitable, Callable -from typing import Any, Final, cast +from collections.abc import Awaitable, Callable, Mapping +from typing import TYPE_CHECKING, Final, TypeAlias, TypedDict, cast + +from typing_extensions import ReadOnly from litellm._logging import verbose_logger from litellm.types.llms.anthropic import AppliedEdit @@ -11,31 +13,54 @@ from .constants import CLEAR_TOOL_USES_EDIT_TYPE, COMPACT_EDIT_TYPE from .editors import apply_clear_tool_uses_20250919, apply_compact_20260112 from .result import PolyfillResult -EditorFn = Callable[..., Any] +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.router import Router -_EDITOR_REGISTRY: Final[dict[str, EditorFn]] = { +AnthropicMessages: TypeAlias = list[dict[str, object]] +AnthropicSystem: TypeAlias = str | list[dict[str, object]] | None +AnthropicTools: TypeAlias = list[dict[str, object]] | None +EditSpec: TypeAlias = dict[str, object] +ContextManagementSpec: TypeAlias = EditSpec | list[EditSpec] | None +SyncEditorReturn: TypeAlias = tuple[AnthropicMessages, AppliedEdit | None] +EditorFn: TypeAlias = Callable[..., object] + + +class EditorKwargs(TypedDict): + """The keyword payload every registered editor accepts.""" + + model: ReadOnly[str] + messages: ReadOnly[AnthropicMessages] + tools: ReadOnly[AnthropicTools] + system: ReadOnly[AnthropicSystem] + edit_spec: ReadOnly[EditSpec] + + +_EDITOR_REGISTRY: Final[Mapping[str, EditorFn]] = { CLEAR_TOOL_USES_EDIT_TYPE: apply_clear_tool_uses_20250919, COMPACT_EDIT_TYPE: apply_compact_20260112, } -def _normalize_spec( - spec: dict[str, Any] | list[dict[str, Any]] | None, -) -> list[dict[str, Any]] | None: +def _map_openai_spec(spec: list[EditSpec]) -> EditSpec | None: + """Translate the OpenAI list form into the Anthropic-native dict form.""" + # Local import to avoid an import cycle at module load. + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + return AnthropicConfig.map_openai_context_management_to_anthropic(spec) + + +def _normalize_spec(spec: ContextManagementSpec) -> list[EditSpec] | None: """Accept Anthropic-native dict form or OpenAI list form; return edits list.""" - if isinstance(spec, list): - # Local import to avoid an import cycle at module load. - from litellm.llms.anthropic.chat.transformation import AnthropicConfig + normalized: Final = _map_openai_spec(spec) if isinstance(spec, list) else spec - spec = AnthropicConfig.map_openai_context_management_to_anthropic(spec) - - edits: Final = spec.get("edits") if isinstance(spec, dict) else None + edits: Final = normalized.get("edits") if isinstance(normalized, dict) else None if not edits or not isinstance(edits, list): return None return [edit for edit in edits if isinstance(edit, dict)] -def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult: +def _wrap_editor_return(raw: object, *, fallback_system: AnthropicSystem) -> PolyfillResult: """Coerce an editor's native return shape into a ``PolyfillResult``. v0 sync editors (e.g. ``clear_tool_uses_20250919``) return a 2-tuple @@ -46,7 +71,7 @@ def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult: return raw # Legacy 2-tuple return — sync editors don't mutate ``system``, so # carry the caller's value forward. - messages, applied = cast(tuple[list[dict[str, Any]], Any], raw) + messages, applied = cast(SyncEditorReturn, raw) return PolyfillResult( messages=messages, system=fallback_system, @@ -57,13 +82,13 @@ def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult: async def apply_context_management( *, model: str, - messages: list[dict[str, Any]], - tools: list[dict[str, Any]] | None, - system: Any, - context_management_spec: dict[str, Any] | list[dict[str, Any]] | None, - litellm_metadata: dict[str, Any] | None = None, - llm_router: Any = None, - user_api_key_auth: Any = None, + messages: AnthropicMessages, + tools: AnthropicTools, + system: AnthropicSystem, + context_management_spec: ContextManagementSpec, + litellm_metadata: Mapping[str, object] | None = None, + llm_router: "Router | None" = None, + user_api_key_auth: "UserAPIKeyAuth | None" = None, ) -> PolyfillResult: """Run edits in order; return a single ``PolyfillResult``. @@ -92,7 +117,7 @@ async def apply_context_management( ) continue - kwargs: dict[str, Any] = { + kwargs: EditorKwargs = { "model": model, "messages": current_messages, "tools": tools, @@ -102,10 +127,12 @@ async def apply_context_management( # Only async editors accept these — passing them to sync v0 editors # would break their signature. if inspect.iscoroutinefunction(editor): - kwargs["litellm_metadata"] = litellm_metadata - kwargs["llm_router"] = llm_router - kwargs["user_api_key_auth"] = user_api_key_auth - raw_result = await cast(Callable[..., Awaitable[Any]], editor)(**kwargs) + raw_result = await cast(Callable[..., Awaitable[PolyfillResult]], editor)( + **kwargs, + litellm_metadata=litellm_metadata, + llm_router=llm_router, + user_api_key_auth=user_api_key_auth, + ) else: raw_result = editor(**kwargs) diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index e6d8686b466..d99c7f556fb 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -33,18 +33,18 @@ def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str, def _build_responses_kwargs( *, max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - context_management: dict | None = None, - metadata: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[AllAnthropicToolsValues | dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: AnthropicOutputSchema | None = None, @@ -134,22 +134,22 @@ class LiteLLMMessagesToResponsesAPIHandler: @staticmethod async def async_anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - context_management: dict | None = None, - metadata: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[AllAnthropicToolsValues | dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: AnthropicOutputSchema | None = None, - **kwargs, + **kwargs: object, ) -> AnthropicMessagesResponse | AsyncIterator[bytes]: responses_kwargs: Final = _build_responses_kwargs( max_tokens=max_tokens, @@ -185,23 +185,23 @@ class LiteLLMMessagesToResponsesAPIHandler: @staticmethod def anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: list[dict[str, object]], model: str, - context_management: dict | None = None, - metadata: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[AllAnthropicToolsValues | dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: AnthropicOutputSchema | None = None, _is_async: bool = False, - **kwargs, + **kwargs: object, ) -> ( AnthropicMessagesResponse | AsyncIterator[bytes] diff --git a/litellm/llms/anthropic/skills/transformation.py b/litellm/llms/anthropic/skills/transformation.py index 566322bbdd6..20f3a3d74f0 100644 --- a/litellm/llms/anthropic/skills/transformation.py +++ b/litellm/llms/anthropic/skills/transformation.py @@ -2,9 +2,10 @@ Anthropic Skills API configuration and transformations """ -from typing import Any, Final +from typing import Final import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -23,6 +24,25 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +class _SkillPayload(TypedDict): + """The JSON body Anthropic returns for a single skill, before ``Skill`` validates it.""" + + id: ReadOnly[str] + created_at: ReadOnly[str] + source: ReadOnly[str] + updated_at: ReadOnly[str] + display_title: NotRequired[ReadOnly[str | None]] + latest_version: NotRequired[ReadOnly[str | None]] + type: NotRequired[ReadOnly[str]] + + +class _DeleteSkillPayload(TypedDict): + """The JSON body Anthropic returns for a skill deletion, before ``DeleteSkillResponse`` validates it.""" + + id: ReadOnly[str] + type: NotRequired[ReadOnly[str]] + + class AnthropicSkillsConfig(BaseSkillsAPIConfig): """Anthropic-specific Skills API configuration""" @@ -104,7 +124,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> Skill: """Transform Anthropic response to Skill object""" - response_json: Final = raw_response.json() + response_json: Final[_SkillPayload] = raw_response.json() verbose_logger.debug("Transforming create skill response: %s", response_json) return Skill(**response_json) @@ -122,7 +142,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): url: Final = self.get_complete_url(api_base=api_base, endpoint="skills") # Build query parameters - query_params: Final[dict[str, Any]] = {} + query_params: Final[dict[str, int | str]] = {} if "limit" in list_params and list_params["limit"]: query_params["limit"] = list_params["limit"] if "page" in list_params and list_params["page"]: @@ -168,7 +188,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> Skill: """Transform Anthropic response to Skill object""" - response_json: Final = raw_response.json() + response_json: Final[_SkillPayload] = raw_response.json() verbose_logger.debug("Transforming get skill response: %s", response_json) return Skill(**response_json) @@ -193,7 +213,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> DeleteSkillResponse: """Transform Anthropic response to DeleteSkillResponse""" - response_json: Final = raw_response.json() + response_json: Final[_DeleteSkillPayload] = raw_response.json() verbose_logger.debug("Transforming delete skill response: %s", response_json) return DeleteSkillResponse(**response_json) diff --git a/litellm/llms/azure/files/handler.py b/litellm/llms/azure/files/handler.py index 4f93896699f..f3ed2e6d6ed 100644 --- a/litellm/llms/azure/files/handler.py +++ b/litellm/llms/azure/files/handler.py @@ -40,6 +40,11 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): data.pop("expires_after", None) return data + @staticmethod + def _to_openai_file_object(response: FileObject) -> OpenAIFileObject: + """Re-wrap the SDK's file object as litellm's, carrying every field across.""" + return OpenAIFileObject(**response.model_dump()) + async def acreate_file( self, create_file_data: CreateFileRequest, @@ -48,7 +53,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): verbose_logger.debug("create_file_data=%s", create_file_data) response = await openai_client.files.create(**self._prepare_create_file_data(create_file_data)) verbose_logger.debug("create_file_response=%s", response) - return OpenAIFileObject(**response.model_dump()) + return self._to_openai_file_object(response) def create_file( self, @@ -61,7 +66,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): max_retries: int | None, client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, litellm_params: dict | None = None, - ) -> OpenAIFileObject | Coroutine[Any, Any, OpenAIFileObject]: + ) -> OpenAIFileObject | Coroutine[None, None, OpenAIFileObject]: openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( litellm_params=litellm_params or {}, api_key=api_key, @@ -84,7 +89,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): response: Final = cast(AzureOpenAI | OpenAI, openai_client).files.create( **self._prepare_create_file_data(create_file_data) ) - return OpenAIFileObject(**response.model_dump()) + return self._to_openai_file_object(response) async def afile_content( self, @@ -105,7 +110,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): api_version: str | None = None, client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, litellm_params: dict | None = None, - ) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]: + ) -> HttpxBinaryResponseContent | Coroutine[None, None, HttpxBinaryResponseContent]: openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( litellm_params=litellm_params or {}, api_key=api_key, diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 3eeb3cb9fc6..e795e907eb4 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -7,9 +7,10 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic. import asyncio import contextlib import json -from typing import Any, Final +from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict from litellm._logging import _redact_string, verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -18,9 +19,74 @@ from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig +if TYPE_CHECKING: + from litellm.types.realtime import RealtimeResponseTransformInput + _CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) +class _ClientWebSocket(Protocol): + """The client-facing websocket surface used by the Bedrock realtime bridge.""" + + async def send_text(self, data: str) -> None: ... + + async def receive_text(self) -> str: ... + + async def close(self, code: int = ..., reason: str | None = ...) -> None: ... + + +class _BedrockInputStream(Protocol): + """The write half of a Bedrock bidirectional stream.""" + + async def send(self, chunk: object) -> None: ... + + async def close(self) -> None: ... + + +class _BedrockPayloadPart(Protocol): + """A single Bedrock bidirectional output payload.""" + + bytes_: bytes | None + + +class _BedrockOutputChunk(Protocol): + """A chunk read off the read half of a Bedrock bidirectional stream.""" + + value: _BedrockPayloadPart | None + + +class _BedrockOutputStream(Protocol): + """The read half of a Bedrock bidirectional stream.""" + + async def receive(self) -> _BedrockOutputChunk | None: ... + + +class _BedrockBidirectionalStream(Protocol): + """The bidirectional stream returned by ``invoke_model_with_bidirectional_stream``.""" + + input_stream: _BedrockInputStream + + async def await_output(self) -> tuple[object, _BedrockOutputStream]: ... + + +class _ClientSessionPayload(TypedDict, total=False): + """The ``session`` body of a client ``session.update`` frame.""" + + modalities: ReadOnly[object] + + +class _ClientRealtimeFrame(TypedDict, total=False): + """The fields read off a client realtime frame.""" + + type: ReadOnly[str] + session: ReadOnly[_ClientSessionPayload] + + +def _decode_client_frame(payload: str) -> _ClientRealtimeFrame: + """Decode a client realtime frame into the fields this bridge reads.""" + return json.loads(payload) + + class BedrockRealtime(BaseAWSLLM): """Handler for Bedrock Nova Sonic realtime speech-to-speech API.""" @@ -30,7 +96,7 @@ class BedrockRealtime(BaseAWSLLM): async def async_realtime( self, model: str, - websocket: Any, + websocket: _ClientWebSocket, logging_obj: LiteLLMLogging, api_base: str | None = None, api_key: str | None = None, @@ -46,7 +112,7 @@ class BedrockRealtime(BaseAWSLLM): aws_sts_endpoint: str | None = None, aws_bedrock_runtime_endpoint: str | None = None, aws_external_id: str | None = None, - **kwargs, + **kwargs: object, ): """ Establish bidirectional streaming connection with Bedrock Nova Sonic. @@ -118,13 +184,16 @@ class BedrockRealtime(BaseAWSLLM): ) bedrock_client: Final = BedrockRuntimeClient(config=config) + async def open_bidirectional_stream() -> _BedrockBidirectionalStream: + return await bedrock_client.invoke_model_with_bidirectional_stream( + InvokeModelWithBidirectionalStreamOperationInput(model_id=model) + ) + transformation_config: Final = BedrockRealtimeConfig() try: # Initialize the bidirectional stream - bedrock_stream: Final = await bedrock_client.invoke_model_with_bidirectional_stream( - InvokeModelWithBidirectionalStreamOperationInput(model_id=model) - ) + bedrock_stream: Final = await open_bidirectional_stream() verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") @@ -132,7 +201,7 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect") # Track state for transformation - session_state: Final = { + session_state: Final[RealtimeResponseTransformInput] = { "current_output_item_id": None, "current_response_id": None, "current_conversation_id": None, @@ -182,11 +251,11 @@ class BedrockRealtime(BaseAWSLLM): async def _forward_client_to_bedrock( self, - client_ws: Any, - bedrock_stream: Any, + client_ws: _ClientWebSocket, + bedrock_stream: _BedrockBidirectionalStream, transformation_config: BedrockRealtimeConfig, model: str, - session_state: dict, + session_state: "RealtimeResponseTransformInput", logging_obj: LiteLLMLogging | None = None, ): """Forward messages from client WebSocket to Bedrock stream.""" @@ -195,10 +264,11 @@ class BedrockRealtime(BaseAWSLLM): InvokeModelWithBidirectionalStreamInputChunk, ) + def build_input_chunk(payload: bytes) -> object: + return InvokeModelWithBidirectionalStreamInputChunk(value=BidirectionalInputPayloadPart(bytes_=payload)) + async def send_to_bedrock(bedrock_message: str) -> None: - event: Final = InvokeModelWithBidirectionalStreamInputChunk( - value=BidirectionalInputPayloadPart(bytes_=bedrock_message.encode("utf-8")) - ) + event: Final = build_input_chunk(bedrock_message.encode("utf-8")) await bedrock_stream.input_stream.send(event) verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200]) @@ -223,7 +293,7 @@ class BedrockRealtime(BaseAWSLLM): client_message_type: str | None = None requested_modalities: list[str] | None = None with contextlib.suppress(Exception): - parsed_client_message = json.loads(message) + parsed_client_message = _decode_client_frame(message) client_message_type = parsed_client_message.get("type") if client_message_type == "session.update": requested_modalities = _CLIENT_MODALITIES_ADAPTER.validate_python( @@ -246,12 +316,12 @@ class BedrockRealtime(BaseAWSLLM): async def _forward_bedrock_to_client( self, - bedrock_stream: Any, - client_ws: Any, + bedrock_stream: _BedrockBidirectionalStream, + client_ws: _ClientWebSocket, transformation_config: BedrockRealtimeConfig, model: str, logging_obj: LiteLLMLogging, - session_state: dict, + session_state: "RealtimeResponseTransformInput", ): """Forward messages from Bedrock stream to client WebSocket.""" try: diff --git a/litellm/llms/chatgpt/chat/streaming_utils.py b/litellm/llms/chatgpt/chat/streaming_utils.py index 57f679947f6..f8c54c0a1d8 100644 --- a/litellm/llms/chatgpt/chat/streaming_utils.py +++ b/litellm/llms/chatgpt/chat/streaming_utils.py @@ -4,7 +4,32 @@ Streaming utilities for ChatGPT provider. Normalizes non-spec-compliant tool_call chunks from the ChatGPT backend API. """ -from typing import Any, Final +from collections.abc import Sequence +from typing import Final, Protocol + +from litellm.types.utils import Delta + + +class ChatGPTStreamChoice(Protocol): + """Streaming choice as read by :class:`ChatGPTToolCallNormalizer`.""" + + @property + def delta(self) -> Delta | None: ... + + +class ChatGPTStreamChunk(Protocol): + """Streaming chunk as read by :class:`ChatGPTToolCallNormalizer`.""" + + @property + def choices(self) -> Sequence[ChatGPTStreamChoice]: ... + + +class ChatGPTChunkStream(Protocol): + """Sync/async chunk source wrapped by :class:`ChatGPTToolCallNormalizer`.""" + + def __next__(self) -> ChatGPTStreamChunk: ... + + async def __anext__(self) -> ChatGPTStreamChunk: ... class ChatGPTToolCallNormalizer: @@ -20,36 +45,36 @@ class ChatGPTToolCallNormalizer: chunks to the consumer. """ - def __init__(self, stream: Any): + def __init__(self, stream: ChatGPTChunkStream): self._stream = stream self._seen_ids: dict[str, int] = {} # tool_call_id -> assigned_index self._next_index: int = 0 self._last_id: str | None = None # tracks which tool call the next delta belongs to - def __getattr__(self, name: str) -> Any: + def __getattr__(self, name: str) -> object: return getattr(self._stream, name) - def __iter__(self): + def __iter__(self) -> "ChatGPTToolCallNormalizer": return self - def __aiter__(self): + def __aiter__(self) -> "ChatGPTToolCallNormalizer": return self - def __next__(self): + def __next__(self) -> ChatGPTStreamChunk: while True: chunk = next(self._stream) result = self._normalize(chunk) if result is not None: return result - async def __anext__(self): + async def __anext__(self) -> ChatGPTStreamChunk: while True: chunk = await self._stream.__anext__() result = self._normalize(chunk) if result is not None: return result - def _normalize(self, chunk: Any) -> Any: + def _normalize(self, chunk: ChatGPTStreamChunk) -> ChatGPTStreamChunk | None: """Fix tool_calls in the chunk. Returns None to skip duplicate chunks.""" if not chunk.choices: return chunk diff --git a/litellm/llms/compactifai/chat/transformation.py b/litellm/llms/compactifai/chat/transformation.py index 44e1ab15801..3d189911e9e 100644 --- a/litellm/llms/compactifai/chat/transformation.py +++ b/litellm/llms/compactifai/chat/transformation.py @@ -2,13 +2,16 @@ CompactifAI chat completion transformation """ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.common_utils import OpenAIError from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse from ...openai.chat.gpt_transformation import OpenAIGPTConfig @@ -21,6 +24,18 @@ else: LiteLLMLoggingObj = Any +class CompactifAIResponseFields(TypedDict, total=False): + """The chat completion fields of a CompactifAI response body.""" + + id: ReadOnly[str] + choices: ReadOnly[Sequence[Mapping[str, object]]] + created: ReadOnly[int] + model: ReadOnly[str] + system_fingerprint: ReadOnly[str | None] + usage: ReadOnly[Mapping[str, object]] + object: ReadOnly[str] + + class CompactifAIChatConfig(OpenAIGPTConfig): """ Configuration class for CompactifAI chat completions. @@ -45,11 +60,11 @@ class CompactifAIChatConfig(OpenAIGPTConfig): raw_response: httpx.Response, model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, - request_data: dict, - messages: list, - optional_params: dict, - litellm_params: dict, - encoding: Any, + request_data: Mapping[str, object], + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + encoding: object, api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: @@ -79,14 +94,18 @@ class CompactifAIChatConfig(OpenAIGPTConfig): message["content"] = tool_calls[0]["function"].get("arguments", "") message["tool_calls"] = None - returned_response: Final = ModelResponse(**response_json) + response_fields: Final[CompactifAIResponseFields] = response_json + + returned_response: Final = ModelResponse(**response_fields) # Set model name with provider prefix returned_response.model = f"compactifai/{model}" return returned_response - def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, str] | httpx.Headers + ) -> BaseLLMException: """ Get the appropriate error class for CompactifAI errors. Since CompactifAI is OpenAI-compatible, we use OpenAI error handling. diff --git a/litellm/llms/custom_httpx/container_handler.py b/litellm/llms/custom_httpx/container_handler.py index 7690351e3b2..343cf043187 100644 --- a/litellm/llms/custom_httpx/container_handler.py +++ b/litellm/llms/custom_httpx/container_handler.py @@ -6,11 +6,13 @@ endpoint defined in endpoints.json, eliminating the need for individual handler """ import json -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping, Sequence from pathlib import Path from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import BaseModel +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -32,22 +34,54 @@ if TYPE_CHECKING: from litellm.llms.base_llm.containers.transformation import BaseContainerConfig +class EndpointConfig(TypedDict): + """One endpoint entry of ``litellm/containers/endpoints.json``.""" + + name: ReadOnly[str] + async_name: ReadOnly[str] + path: ReadOnly[str] + method: ReadOnly[str] + path_params: ReadOnly[Sequence[str]] + query_params: ReadOnly[Sequence[str]] + response_type: ReadOnly[str] + is_multipart: NotRequired[ReadOnly[bool]] + returns_binary: NotRequired[ReadOnly[bool]] + + +class EndpointsConfig(TypedDict): + """The parsed ``litellm/containers/endpoints.json`` document.""" + + endpoints: ReadOnly[Sequence[EndpointConfig]] + + +class ContainerErrorDetail(TypedDict, total=False): + """The ``error`` object of a container API error body.""" + + message: ReadOnly[str] + + +class ContainerResponseBody(TypedDict, total=False): + """The fields this handler reads off a container API JSON body.""" + + error: ReadOnly[ContainerErrorDetail] + + # Response type mapping -RESPONSE_TYPES: Final[dict[str, type]] = { +RESPONSE_TYPES: Final[Mapping[str, type[BaseModel]]] = { "ContainerFileListResponse": ContainerFileListResponse, "ContainerFileObject": ContainerFileObject, "DeleteContainerFileResponse": DeleteContainerFileResponse, } -def _load_endpoints_config() -> dict: +def _load_endpoints_config() -> EndpointsConfig: """Load the endpoints configuration from JSON file.""" config_path: Final = Path(__file__).parent.parent.parent / "containers" / "endpoints.json" with open(config_path) as f: return json.load(f) -def _get_endpoint_config(endpoint_name: str) -> dict | None: +def _get_endpoint_config(endpoint_name: str) -> EndpointConfig | None: """Get config for a specific endpoint by name.""" config: Final = _load_endpoints_config() for endpoint in config["endpoints"]: @@ -56,10 +90,15 @@ def _get_endpoint_config(endpoint_name: str) -> dict | None: return None +def _response_model(response_type_name: str) -> type[BaseModel] | None: + """The pydantic model a container endpoint's ``response_type`` names.""" + return RESPONSE_TYPES.get(response_type_name) + + def _build_url( api_base: str, path_template: str, - path_params: dict[str, str], + path_params: Mapping[str, object], ) -> str: """Build the full URL by substituting path parameters. @@ -89,22 +128,18 @@ def _build_url( def _build_query_params( - query_param_names: list, - kwargs: dict[str, Any], -) -> dict[str, str]: + query_param_names: Sequence[str], + kwargs: Mapping[str, object], +) -> dict[str, object]: """Build query parameters from kwargs.""" - params: Final = {} - for param_name in query_param_names: - value = kwargs.get(param_name) - if value is not None: - params[param_name] = str(value) if not isinstance(value, str) else value - return params + supplied: Final = ((param_name, kwargs.get(param_name)) for param_name in query_param_names) + return {name: value if isinstance(value, str) else str(value) for name, value in supplied if value is not None} def _prepare_multipart_file_upload( file: Any, - headers: dict[str, Any], -) -> tuple: + headers: dict[str, object], +) -> tuple[dict[str, tuple[str, bytes, str]], dict[str, object]]: """ Prepare file and headers for multipart upload. @@ -129,6 +164,52 @@ def _prepare_multipart_file_upload( return files, headers_copy +def _request_headers( + container_provider_config: "BaseContainerConfig", + extra_headers: dict[str, object] | None, + litellm_params: GenericLiteLLMParams, +) -> dict[str, object]: + """The provider auth headers for a container request.""" + return container_provider_config.validate_environment( + headers=extra_headers or {}, + api_key=litellm_params.get("api_key", None), + ) + + +def _request_api_base( + container_provider_config: "BaseContainerConfig", + litellm_params: GenericLiteLLMParams, +) -> str: + """The provider base URL for a container request.""" + return container_provider_config.get_complete_url( + api_base=litellm_params.get("api_base", None), + litellm_params=dict(litellm_params), + ) + + +def _sync_http_client( + client: HTTPHandler | AsyncHTTPHandler | None, + litellm_params: GenericLiteLLMParams, +) -> HTTPHandler: + """The sync HTTP client for a container request, reusing the caller's when usable.""" + if client is None or not isinstance(client, HTTPHandler): + return _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) + return client + + +def _async_http_client( + client: HTTPHandler | AsyncHTTPHandler | None, + litellm_params: GenericLiteLLMParams, +) -> AsyncHTTPHandler: + """The async HTTP client for a container request, reusing the caller's when usable.""" + if client is None or not isinstance(client, AsyncHTTPHandler): + return get_async_httpx_client( + llm_provider=litellm.LlmProviders.OPENAI, + params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + ) + return client + + class GenericContainerHandler: """ Generic handler for container file API endpoints. @@ -143,13 +224,13 @@ class GenericContainerHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout = 600, _is_async: bool = False, client: HTTPHandler | AsyncHTTPHandler | None = None, - **kwargs, - ) -> Any | Coroutine[Any, Any, Any]: + **kwargs: object, + ) -> Any | Coroutine[object, object, Any]: """ Generic handler for any container file endpoint. @@ -196,11 +277,11 @@ class GenericContainerHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout = 600, client: HTTPHandler | AsyncHTTPHandler | None = None, - **kwargs, + **kwargs: object, ) -> Any: """Synchronous request handler.""" endpoint_config: Final = _get_endpoint_config(endpoint_name) @@ -208,23 +289,14 @@ class GenericContainerHandler: raise ValueError(f"Unknown endpoint: {endpoint_name}") # Get HTTP client - if client is None or not isinstance(client, HTTPHandler): - http_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) - else: - http_client = client + http_client: Final = _sync_http_client(client, litellm_params) # Build request - headers = container_provider_config.validate_environment( - headers=extra_headers or {}, - api_key=litellm_params.get("api_key", None), - ) + headers = _request_headers(container_provider_config, extra_headers, litellm_params) if extra_headers: headers.update(extra_headers) - api_base: Final = container_provider_config.get_complete_url( - api_base=litellm_params.get("api_base", None), - litellm_params=dict(litellm_params), - ) + api_base: Final = _request_api_base(container_provider_config, litellm_params) # Build URL with path params path_params: Final = {p: kwargs.get(p, "") for p in endpoint_config.get("path_params", [])} @@ -275,11 +347,11 @@ class GenericContainerHandler: return response.content # Check for error response - response_json: Final = response.json() + response_json: Final[ContainerResponseBody] = response.json() if "error" in response_json: from litellm.llms.base_llm.chat.transformation import BaseLLMException - error_msg: Final = response_json.get("error", {}).get("message", str(response_json)) + error_msg: Final = response_json["error"].get("message", str(response_json)) raise BaseLLMException( status_code=response.status_code, message=error_msg, @@ -287,7 +359,7 @@ class GenericContainerHandler: ) # Parse response - response_type: Final = RESPONSE_TYPES.get(endpoint_config["response_type"]) + response_type: Final = _response_model(endpoint_config["response_type"]) if response_type: return response_type(**response_json) return response_json @@ -301,11 +373,11 @@ class GenericContainerHandler: container_provider_config: "BaseContainerConfig", litellm_params: GenericLiteLLMParams, logging_obj: "LiteLLMLoggingObj", - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout = 600, client: HTTPHandler | AsyncHTTPHandler | None = None, - **kwargs, + **kwargs: object, ) -> Any: """Asynchronous request handler.""" endpoint_config: Final = _get_endpoint_config(endpoint_name) @@ -313,26 +385,14 @@ class GenericContainerHandler: raise ValueError(f"Unknown endpoint: {endpoint_name}") # Get HTTP client - if client is None or not isinstance(client, AsyncHTTPHandler): - http_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders.OPENAI, - params={"ssl_verify": litellm_params.get("ssl_verify", None)}, - ) - else: - http_client = client + http_client: Final = _async_http_client(client, litellm_params) # Build request - headers = container_provider_config.validate_environment( - headers=extra_headers or {}, - api_key=litellm_params.get("api_key", None), - ) + headers = _request_headers(container_provider_config, extra_headers, litellm_params) if extra_headers: headers.update(extra_headers) - api_base: Final = container_provider_config.get_complete_url( - api_base=litellm_params.get("api_base", None), - litellm_params=dict(litellm_params), - ) + api_base: Final = _request_api_base(container_provider_config, litellm_params) # Build URL with path params path_params: Final = {p: kwargs.get(p, "") for p in endpoint_config.get("path_params", [])} @@ -383,11 +443,11 @@ class GenericContainerHandler: return response.content # Check for error response - response_json: Final = response.json() + response_json: Final[ContainerResponseBody] = response.json() if "error" in response_json: from litellm.llms.base_llm.chat.transformation import BaseLLMException - error_msg: Final = response_json.get("error", {}).get("message", str(response_json)) + error_msg: Final = response_json["error"].get("message", str(response_json)) raise BaseLLMException( status_code=response.status_code, message=error_msg, @@ -395,7 +455,7 @@ class GenericContainerHandler: ) # Parse response - response_type: Final = RESPONSE_TYPES.get(endpoint_config["response_type"]) + response_type: Final = _response_model(endpoint_config["response_type"]) if response_type: return response_type(**response_json) return response_json diff --git a/litellm/llms/gdc/chat/transformation.py b/litellm/llms/gdc/chat/transformation.py index 2d0322bf10f..03037512551 100644 --- a/litellm/llms/gdc/chat/transformation.py +++ b/litellm/llms/gdc/chat/transformation.py @@ -6,11 +6,25 @@ import json import os import re import threading -from typing import Any, Final +from collections.abc import Callable +from typing import Any, Final, Protocol from urllib.parse import urlsplit import litellm from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig +from litellm.types.llms.openai import AllMessageValues + + +class _GDCHAudienceCredentials(Protocol): + """A GDCH service account credential already bound to an audience, ready to mint a bearer token.""" + + @property + def valid(self) -> bool: ... + + @property + def token(self) -> str: ... + + def refresh(self, request: object) -> None: ... class GDCGeminiConfig(OpenAILikeChatConfig): @@ -21,7 +35,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self._creds_lock = threading.Lock() - self._gdch_creds_cache: dict = {} + self._gdch_creds_cache: dict[tuple[str, str], _GDCHAudienceCredentials] = {} def get_supported_openai_params(self, model: str) -> list: return [ @@ -110,7 +124,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): return f"{api_base}/v1/projects/{project}/locations/{location}/chat/completions" - def _read_env_bool(self, val: Any, env_var: str, default: bool = True) -> bool | str: + def _read_env_bool(self, val: bool | str | None, env_var: str, default: bool = True) -> bool | str: def _parse(s: str) -> bool | str: cleaned: Final = s.strip().lower() if cleaned in ("false", "0", "no", "off"): @@ -129,7 +143,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): return default return _parse(_env_val) - def _fetch_auth(self, gdch_creds: Any, ssl_verify: bool | str) -> None: + def _fetch_auth(self, gdch_creds: _GDCHAudienceCredentials, ssl_verify: bool | str) -> None: import requests from google.auth.transport import requests as auth_requests @@ -138,13 +152,24 @@ class GDCGeminiConfig(OpenAILikeChatConfig): auth_request: Final = auth_requests.Request(session=auth_session) gdch_creds.refresh(auth_request) - def _cached_fetch_token(self, creds: Any, audience: str, ssl_verify: bool | str, api_key: str | None = None) -> str: + def _with_gdch_audience(self, creds: object, audience: str) -> _GDCHAudienceCredentials: + """The credential rebound to ``audience``, which GDCH requires before a token refresh.""" + bind_audience: Final[Callable[[str], _GDCHAudienceCredentials] | None] = getattr( + creds, "with_gdch_audience", None + ) + if bind_audience is None: + raise AttributeError("GDC credentials must expose with_gdch_audience to be bound to a request audience") + return bind_audience(audience) + + def _cached_fetch_token( + self, creds: object, audience: str, ssl_verify: bool | str, api_key: str | None = None + ) -> str: # Key cache by both audience and credential identity to prevent cross-caller contamination cache_key: Final = (audience.rstrip("/"), api_key or str(id(creds))) with self._creds_lock: if cache_key not in self._gdch_creds_cache: - self._gdch_creds_cache[cache_key] = creds.with_gdch_audience(audience.rstrip("/")) + self._gdch_creds_cache[cache_key] = self._with_gdch_audience(creds, audience.rstrip("/")) gdch_creds: Final = self._gdch_creds_cache[cache_key] @@ -155,7 +180,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): return token - def _load_creds_from_key(self, api_key: str) -> tuple[Any, bool]: + def _load_creds_from_key(self, api_key: str) -> tuple[object | None, bool]: import google.auth try: @@ -175,7 +200,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): self, headers: dict, model: str, - messages: list[Any], + messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, api_key: str | None = None, @@ -230,7 +255,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): if self._read_env_bool(litellm_params.get("gdc_token_caching"), "GDC_TOKEN_CACHING", default=False): token = self._cached_fetch_token(creds, audience, ssl_verify, api_key) else: - gdch_creds: Final = creds.with_gdch_audience(audience) + gdch_creds: Final = self._with_gdch_audience(creds, audience) self._fetch_auth(gdch_creds, ssl_verify) token = gdch_creds.token headers["Authorization"] = f"Bearer {token}" @@ -252,7 +277,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): def transform_request( self, model: str, - messages: list[Any], + messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, headers: dict, diff --git a/litellm/llms/infinity/rerank/transformation.py b/litellm/llms/infinity/rerank/transformation.py index ebcbf1b5a07..df4adf0c9a2 100644 --- a/litellm/llms/infinity/rerank/transformation.py +++ b/litellm/llms/infinity/rerank/transformation.py @@ -4,9 +4,11 @@ Transformation logic from Cohere's /v1/rerank format to Infinity's `/v1/rerank` Why separate file? Make it easy to see how transformation works """ +from collections.abc import Sequence from typing import Final import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._uuid import uuid @@ -25,6 +27,31 @@ from litellm.types.rerank import ( from ..common_utils import InfinityError +class _InfinityRerankUsage(TypedDict, extra_items=ReadOnly[int]): + """The token counters Infinity reports in the ``usage`` block of a rerank response.""" + + +class _InfinityRerankResult(TypedDict): + """One scored document in an Infinity ``/v1/rerank`` response.""" + + index: ReadOnly[int] + relevance_score: ReadOnly[float] + document: ReadOnly[str] + + +class _InfinityRerankResponse(TypedDict): + """The JSON body returned by Infinity's ``/v1/rerank`` endpoint.""" + + id: ReadOnly[NotRequired[str]] + usage: ReadOnly[NotRequired[_InfinityRerankUsage]] + results: ReadOnly[Sequence[_InfinityRerankResult]] + + +def _parse_rerank_response(raw_response: httpx.Response) -> _InfinityRerankResponse: + """Read the untyped JSON body of an Infinity rerank response.""" + return raw_response.json() + + class InfinityRerankConfig(CohereRerankConfig): def get_complete_url( self, @@ -80,7 +107,7 @@ class InfinityRerankConfig(CohereRerankConfig): No transformation required, Infinity follows Cohere API response format """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _parse_rerank_response(raw_response) except Exception: raise InfinityError(message=raw_response.text, status_code=raw_response.status_code) diff --git a/litellm/llms/litellm_proxy/skills/code_execution.py b/litellm/llms/litellm_proxy/skills/code_execution.py index d435994ce20..f51213ca1c3 100644 --- a/litellm/llms/litellm_proxy/skills/code_execution.py +++ b/litellm/llms/litellm_proxy/skills/code_execution.py @@ -13,12 +13,57 @@ Generated files are returned directly in the response - no separate storage need import base64 import json +from collections.abc import Sequence from enum import Enum -from typing import Any, Final +from typing import Any, Final, Protocol + +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger +class _ToolCallFunction(Protocol): + """Function payload of an assistant tool call.""" + + name: str | None + arguments: str + + +class _ToolCall(Protocol): + """Tool call requested by the assistant on a chat completion choice.""" + + id: str + function: _ToolCallFunction + + +class _AssistantMessage(Protocol): + """Assistant message carried by a chat completion choice.""" + + content: str | None + tool_calls: Sequence[_ToolCall] | None + + +class _CompletionChoice(Protocol): + """Single choice of a chat completion response.""" + + finish_reason: str + message: _AssistantMessage + + +class _SandboxFile(TypedDict): + """File generated inside the sandbox during a code execution run.""" + + name: ReadOnly[str] + mime_type: ReadOnly[str] + content_base64: ReadOnly[str] + + +class _CodeExecutionArguments(TypedDict): + """Arguments the model passes to the `litellm_code_execution` tool.""" + + code: NotRequired[ReadOnly[str]] + + class LiteLLMInternalTools(str, Enum): """ Enum for internal LiteLLM tools that are injected into requests. @@ -30,7 +75,7 @@ class LiteLLMInternalTools(str, Enum): CODE_EXECUTION = "litellm_code_execution" -def get_litellm_code_execution_tool() -> dict[str, Any]: +def get_litellm_code_execution_tool() -> dict[str, object]: """ Returns the litellm_code_execution tool definition in OpenAI format. @@ -51,7 +96,7 @@ def get_litellm_code_execution_tool() -> dict[str, Any]: } -def get_litellm_code_execution_tool_anthropic() -> dict[str, Any]: +def get_litellm_code_execution_tool_anthropic() -> dict[str, object]: """ Returns the litellm_code_execution tool definition in Anthropic/messages API format. @@ -103,7 +148,7 @@ class CodeExecutionHandler: skill_files: dict[str, bytes], skill_id: str | None = None, **kwargs, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Execute an LLM call with automatic code execution handling. @@ -134,8 +179,8 @@ class CodeExecutionHandler: ) current_messages: Final = list(messages) - generated_files: Final[list[dict[str, Any]]] = [] # Files returned directly - execution_results: Final[list[dict]] = [] + generated_files: Final[list[dict[str, object]]] = [] # Files returned directly + execution_results: Final[list[dict[str, object]]] = [] executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) response: Any = None # Initialize to avoid possibly unbound error @@ -151,11 +196,12 @@ class CodeExecutionHandler: **kwargs, ) - assistant_message = response.choices[0].message - stop_reason = response.choices[0].finish_reason + choice: _CompletionChoice = response.choices[0] + assistant_message = choice.message + stop_reason: str = choice.finish_reason # Build assistant message for conversation history - assistant_msg_dict: dict[str, Any] = { + assistant_msg_dict: dict[str, object] = { "role": "assistant", "content": assistant_message.content, } @@ -190,8 +236,8 @@ class CodeExecutionHandler: if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value: # Execute code in sandbox try: - args = json.loads(tool_call.function.arguments) - code = args.get("code", "") + args: _CodeExecutionArguments = json.loads(tool_call.function.arguments) + code: str = args.get("code", "") verbose_logger.debug("CodeExecutionHandler: Executing code (%s chars)", len(code)) @@ -202,13 +248,15 @@ class CodeExecutionHandler: verbose_logger.debug("CodeExecutionHandler: Execution result: %s", exec_result) + sandbox_files: Sequence[_SandboxFile] = exec_result["files"] + execution_results.append( { "iteration": iteration, "success": exec_result["success"], "output": exec_result["output"], "error": exec_result["error"], - "files": [f["name"] for f in exec_result["files"]], + "files": [f["name"] for f in sandbox_files], } ) @@ -216,9 +264,9 @@ class CodeExecutionHandler: tool_result = exec_result["output"] or "" # Collect generated files (returned directly, no storage) - if exec_result["files"]: + if sandbox_files: tool_result += "\n\nGenerated files:" - for f in exec_result["files"]: + for f in sandbox_files: file_content = base64.b64decode(f["content_base64"]) # Add to generated files list (returned in response) generated_files.append( diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 7ae438fd4cd..18f75f0e45c 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -8,10 +8,12 @@ response parsing, and streaming chunk parsing for models served with import datetime import json -from typing import Any, Final +from collections.abc import Sequence +from typing import Any, Final, TypedDict import httpx from pydantic import ValidationError +from typing_extensions import ReadOnly from litellm.llms.oci.chat.generic import ( _normalize_oci_finish_reason, @@ -46,6 +48,20 @@ from litellm.types.utils import ( ) +class _OpenAIToolCallFunction(TypedDict, total=False): + """The ``function`` block of an OpenAI-format assistant tool call.""" + + name: ReadOnly[str | None] + arguments: ReadOnly[str | dict[str, object]] + + +class _OpenAIToolCall(TypedDict, total=False): + """An entry of an OpenAI-format assistant message's ``tool_calls``.""" + + id: ReadOnly[str | None] + function: ReadOnly[_OpenAIToolCallFunction] + + def _extract_text_content(content: Any) -> str: """Return the plain-text representation of a message content value.""" if content is None: @@ -78,10 +94,10 @@ def adapt_messages_to_cohere_standard( """ # First pass: build tool_call_id → CohereToolCall so tool-result messages can # reference the originating call by name and parameters. - tool_call_lookup: Final[dict[str, CohereToolCall]] = {} + tool_call_lookup: Final[dict[str | None, CohereToolCall]] = {} for msg in messages: if msg.get("role") == "assistant": - tool_calls_raw: Any = msg.get("tool_calls") or [] + tool_calls_raw: Sequence[_OpenAIToolCall] = msg.get("tool_calls") or [] for tc in tool_calls_raw: tc_id = tc.get("id", "") raw_args = tc.get("function", {}).get("arguments", "{}") @@ -150,8 +166,22 @@ def adapt_messages_to_cohere_standard( return chat_history +class _OpenAIToolDefinitionFunction(TypedDict, total=False): + """The ``function`` block of an OpenAI-format tool definition.""" + + name: ReadOnly[str] + description: ReadOnly[str] + parameters: ReadOnly[dict[str, object]] + + +class _OpenAIToolDefinition(TypedDict, total=False): + """An entry of an OpenAI-format ``tools`` array.""" + + function: ReadOnly[_OpenAIToolDefinitionFunction] + + def adapt_tool_definitions_to_cohere_standard( - tools: list[dict[str, Any]], + tools: Sequence[_OpenAIToolDefinition], ) -> list[CohereTool]: """Adapt OpenAI-format tool definitions to the OCI Cohere format. diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 6e490f3ff15..449952217b5 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -4,16 +4,32 @@ Ollama /chat/completion calls handled in llm_http_handler.py [TODO]: migrate embeddings to a base handler as well. """ -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Any, Final, Protocol, TypedDict + +from typing_extensions import NotRequired, ReadOnly import litellm from litellm.types.utils import EmbeddingResponse +class TokenEncoder(Protocol): + """The tokenizer surface used to estimate prompt tokens.""" + + def encode(self, text: str, /) -> Sequence[int]: ... + + +class OllamaEmbeddingResponse(TypedDict): + """Body of an Ollama ``/api/embed`` response.""" + + embeddings: ReadOnly[list[list[float]]] + prompt_eval_count: ReadOnly[NotRequired[int]] + + def _prepare_ollama_embedding_payload( - model: str, prompts: list[str], optional_params: dict[str, Any] -) -> dict[str, Any]: - data: Final[dict[str, Any]] = {"model": model, "input": prompts} + model: str, prompts: list[str], optional_params: Mapping[str, object] +) -> dict[str, object]: + data: Final[dict[str, object]] = {"model": model, "input": prompts} special_optional_params: Final = ["truncate", "options", "keep_alive", "dimensions"] for k, v in optional_params.items(): @@ -27,12 +43,12 @@ def _prepare_ollama_embedding_payload( def _process_ollama_embedding_response( - response_json: dict, + response_json: OllamaEmbeddingResponse, prompts: list[str], model: str, model_response: EmbeddingResponse, logging_obj: Any, - encoding: Any, + encoding: TokenEncoder | None, ) -> EmbeddingResponse: output_data: Final = [] embeddings: Final[list[list[float]]] = response_json["embeddings"] @@ -72,7 +88,7 @@ async def ollama_aembeddings( model_response: EmbeddingResponse, optional_params: dict, logging_obj: Any, - encoding: Any, + encoding: TokenEncoder | None, ): if not api_base.endswith("/api/embed"): api_base += "/api/embed" @@ -80,7 +96,7 @@ async def ollama_aembeddings( data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) response: Final = await litellm.module_level_aclient.post(url=api_base, json=data) - response_json: Final = response.json() + response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( response_json=response_json, @@ -99,7 +115,7 @@ def ollama_embeddings( optional_params: dict, model_response: EmbeddingResponse, logging_obj: Any, - encoding: Any = None, + encoding: TokenEncoder | None = None, ): if not api_base.endswith("/api/embed"): api_base += "/api/embed" @@ -107,7 +123,7 @@ def ollama_embeddings( data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) response: Final = litellm.module_level_client.post(url=api_base, json=data) - response_json: Final = response.json() + response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( response_json=response_json, diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index 6fc50458aa3..44bd401c115 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -1,6 +1,8 @@ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, Literal import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( @@ -11,9 +13,11 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.containers.main import ( ContainerCreateOptionalRequestParams, ContainerFileListResponse, + ContainerFileObject, ContainerListResponse, ContainerObject, DeleteContainerResult, + ExpiresAfter, ) from litellm.types.router import GenericLiteLLMParams @@ -32,6 +36,46 @@ else: BaseLLMException = Any +class OpenAIContainerPayload(TypedDict): + """The JSON body OpenAI returns for a single container.""" + + id: ReadOnly[str] + object: ReadOnly[Literal["container"]] + created_at: ReadOnly[int] + status: ReadOnly[str] + expires_after: ReadOnly[ExpiresAfter | None] + last_active_at: ReadOnly[int | None] + name: ReadOnly[str | None] + + +class OpenAIContainerListPayload(TypedDict): + """The JSON body OpenAI returns for a page of containers.""" + + object: ReadOnly[Literal["list"]] + data: ReadOnly[list[ContainerObject]] + first_id: ReadOnly[str | None] + last_id: ReadOnly[str | None] + has_more: ReadOnly[bool] + + +class OpenAIContainerDeletedPayload(TypedDict): + """The JSON body OpenAI returns for a deleted container.""" + + id: ReadOnly[str] + object: ReadOnly[Literal["container.deleted"]] + deleted: ReadOnly[bool] + + +class OpenAIContainerFileListPayload(TypedDict): + """The JSON body OpenAI returns for a page of container files.""" + + object: ReadOnly[Literal["list"]] + data: ReadOnly[list[ContainerFileObject]] + first_id: ReadOnly[str | None] + last_id: ReadOnly[str | None] + has_more: ReadOnly[bool] + + class OpenAIContainerConfig(BaseContainerConfig): """Configuration class for OpenAI container API.""" @@ -87,7 +131,7 @@ class OpenAIContainerConfig(BaseContainerConfig): def transform_container_create_request( self, name: str, - container_create_optional_request_params: dict, + container_create_optional_request_params: Mapping[str, object], litellm_params: GenericLiteLLMParams, headers: dict, ) -> dict: @@ -111,7 +155,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerObject: """Transform the OpenAI container creation response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerPayload] = raw_response.json() # Transform the response data container_obj: Final = ContainerObject(**response_data) @@ -140,7 +184,7 @@ class OpenAIContainerConfig(BaseContainerConfig): after: str | None = None, limit: int | None = None, order: str | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """Transform the container list request for OpenAI API. @@ -151,7 +195,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = api_base # Prepare query parameters - params: Final = {} + params: Final[dict[str, object]] = {} if after is not None: params["after"] = after if limit is not None: @@ -171,7 +215,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerListResponse: """Transform the OpenAI container list response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerListPayload] = raw_response.json() # Transform the response data container_list: Final = ContainerListResponse(**response_data) @@ -191,7 +235,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No additional data needed for GET request - data: Final[dict[str, Any]] = {} + data: Final[dict[str, object]] = {} return url, data @@ -201,7 +245,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerObject: """Transform the OpenAI container retrieve response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerPayload] = raw_response.json() # Transform the response data container_obj: Final = ContainerObject(**response_data) @@ -224,7 +268,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No data needed for DELETE request - data: Final[dict[str, Any]] = {} + data: Final[dict[str, object]] = {} return url, data @@ -234,7 +278,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> DeleteContainerResult: """Transform the OpenAI container delete response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerDeletedPayload] = raw_response.json() # Transform the response data delete_result: Final = DeleteContainerResult(**response_data) @@ -250,7 +294,7 @@ class OpenAIContainerConfig(BaseContainerConfig): after: str | None = None, limit: int | None = None, order: str | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """Transform the container file list request for OpenAI API. @@ -262,7 +306,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}/files") # Prepare query parameters - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} if after is not None: params["after"] = after if limit is not None: @@ -282,7 +326,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerFileListResponse: """Transform the OpenAI container file list response.""" - response_data: Final = raw_response.json() + response_data: Final[OpenAIContainerFileListPayload] = raw_response.json() # Transform the response data file_list: Final = ContainerFileListResponse(**response_data) @@ -308,7 +352,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}/files/{encoded_file_id}/content") # No query parameters needed - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} return url, params diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index 962dfe52c0a..deedf237a50 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.vertex_ai.common_utils import get_vertex_base_url @@ -25,6 +27,68 @@ else: LiteLLMLoggingObj = Any +class _AuthHeadersView(TypedDict): + """Auth headers read out of an untyped ``BaseVectorStoreAuthCredentials``.""" + + headers: ReadOnly[Mapping[str, str]] + + +class _RagContext(TypedDict, total=False): + """One context entry of a Vertex RAG ``:retrieveContexts`` response.""" + + text: ReadOnly[str] + sourceUri: ReadOnly[str] + sourceDisplayName: ReadOnly[str] + pageSpan: ReadOnly[Mapping[str, object]] + score: ReadOnly[float] + + +class _RagContexts(TypedDict, total=False): + """The ``contexts`` envelope wrapping the context list.""" + + contexts: ReadOnly[Sequence[_RagContext]] + + +class _RetrieveContextsBody(TypedDict, total=False): + """Body of a Vertex RAG ``:retrieveContexts`` response.""" + + contexts: ReadOnly[_RagContexts] + + +class _RetrieveContextsView(TypedDict): + """Typed view over the untyped ``:retrieveContexts`` JSON payload.""" + + body: ReadOnly[_RetrieveContextsBody] + + +class _RagCorpusBody(TypedDict, total=False): + """Body of a Vertex RAG ``ragCorpora`` create response.""" + + name: ReadOnly[str] + display_name: ReadOnly[str] + createTime: ReadOnly[str | int | float] + labels: ReadOnly[Mapping[str, str]] + + +class _RagCorpusView(TypedDict): + """Typed view over the untyped ``ragCorpora`` create JSON payload.""" + + body: ReadOnly[_RagCorpusBody] + + +class _RagSearchQuery(TypedDict, total=False): + """The ``query`` block of a Vertex RAG ``:retrieveContexts`` request.""" + + text: ReadOnly[str] + rag_retrieval_config: ReadOnly[Mapping[str, object]] + + +class _LoggedQueryView(TypedDict): + """The search query recovered from the logging object's call details.""" + + search_query: ReadOnly[str] + + class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): """ Configuration for Vertex AI Vector Store RAG API @@ -35,7 +99,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): def __init__(self): super().__init__() - def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: + def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials: # Get credentials and project info vertex_credentials: Final = self.get_vertex_ai_credentials(dict(litellm_params)) vertex_project: Final = self.get_vertex_ai_project(dict(litellm_params)) @@ -60,20 +124,23 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): "write": [("POST", "/ragCorpora")], } - def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict: + def validate_environment( + self, headers: dict[str, str], litellm_params: GenericLiteLLMParams | None + ) -> dict[str, str]: """ Validate and set up authentication for Vertex AI RAG API """ litellm_params = litellm_params or GenericLiteLLMParams() auth_headers: Final = self.get_auth_credentials(litellm_params.model_dump()) - headers.update(auth_headers.get("headers", {})) + auth_view: Final[_AuthHeadersView] = {"headers": auth_headers.get("headers", {})} + headers.update(auth_view["headers"]) return headers def get_complete_url( self, api_base: str | None, - litellm_params: dict, + litellm_params: dict[str, object], ) -> str: """ Get the Base endpoint for Vertex AI RAG API @@ -95,9 +162,9 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, api_base: str, litellm_logging_obj: LiteLLMLoggingObj, - litellm_params: dict, - extra_body: dict[str, Any] | None = None, - ) -> tuple[str, dict[str, Any]]: + litellm_params: dict[str, object], + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: """ Transform search request for Vertex AI RAG API """ @@ -120,35 +187,34 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): # Just the corpus ID, construct full path full_rag_corpus = f"projects/{vertex_project}/locations/{vertex_location}/ragCorpora/{vector_store_id}" - # Build the request body for Vertex AI RAG API - request_body: Final[dict[str, Any]] = { - "vertex_rag_store": {"rag_resources": [{"rag_corpus": full_rag_corpus}]}, - "query": {"text": query}, - } - ######################################################### # Update logging object with details of the request ######################################################### litellm_logging_obj.model_call_details["query"] = query # Add optional parameters + rag_retrieval_config: Final[dict[str, object]] = {} max_num_results: Final = vector_store_search_optional_params.get("max_num_results") if max_num_results is not None: - request_body["query"]["rag_retrieval_config"] = {"top_k": max_num_results} + rag_retrieval_config["top_k"] = max_num_results # Add filters if provided - filters: Final = vector_store_search_optional_params.get("filters") + filters: Final[object] = vector_store_search_optional_params.get("filters") if filters is not None: - if "rag_retrieval_config" not in request_body["query"]: - request_body["query"]["rag_retrieval_config"] = {} - request_body["query"]["rag_retrieval_config"]["filter"] = filters + rag_retrieval_config["filter"] = filters # Add ranking options if provided - ranking_options: Final = vector_store_search_optional_params.get("ranking_options") + ranking_options: Final[object] = vector_store_search_optional_params.get("ranking_options") if ranking_options is not None: - if "rag_retrieval_config" not in request_body["query"]: - request_body["query"]["rag_retrieval_config"] = {} - request_body["query"]["rag_retrieval_config"]["ranking"] = ranking_options + rag_retrieval_config["ranking"] = ranking_options + + query_body: Final[_RagSearchQuery] = ( + {"text": query, "rag_retrieval_config": rag_retrieval_config} if rag_retrieval_config else {"text": query} + ) + request_body: Final[dict[str, object]] = { + "vertex_rag_store": {"rag_resources": [{"rag_corpus": full_rag_corpus}]}, + "query": query_body, + } return url, request_body @@ -159,12 +225,13 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): Transform Vertex AI RAG API response to standard vector store search response """ try: - response_json: Final = response.json() + response_view: Final[_RetrieveContextsView] = {"body": response.json()} + response_json: Final = response_view["body"] # Extract contexts from Vertex AI response - handle nested structure contexts: Final = response_json.get("contexts", {}).get("contexts", []) # Transform contexts to standard format - search_results: Final = [] + search_results: Final[list[VectorStoreSearchResult]] = [] for context in contexts: content = [ VectorStoreResultContent( @@ -202,9 +269,12 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): ) search_results.append(result) + query_view: Final[_LoggedQueryView] = { + "search_query": litellm_logging_obj.model_call_details.get("query", "") + } return VectorStoreSearchResponse( object="vector_store.search_results.page", - search_query=litellm_logging_obj.model_call_details.get("query", ""), + search_query=query_view["search_query"], data=search_results, ) @@ -219,14 +289,14 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): self, vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, api_base: str, - ) -> tuple[str, dict[str, Any]]: + ) -> tuple[str, dict[str, object]]: """ Transform create request for Vertex AI RAG Corpus """ url: Final = f"{api_base}/ragCorpora" # Base URL for creating RAG corpus # Build the request body for Vertex AI RAG Corpus creation - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "display_name": vector_store_create_optional_params.get("name", "litellm-vector-store"), "description": "Vector store created via LiteLLM", } @@ -243,7 +313,8 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): Transform Vertex AI RAG Corpus creation response to standard vector store response """ try: - response_json: Final = response.json() + response_view: Final[_RagCorpusView] = {"body": response.json()} + response_json: Final = response_view["body"] # Extract the corpus ID from the response name corpus_name: Final = response_json.get("name", "") diff --git a/litellm/models/base.py b/litellm/models/base.py index 7eedf10212e..8125bfd0205 100644 --- a/litellm/models/base.py +++ b/litellm/models/base.py @@ -33,6 +33,6 @@ class DomainModel(BaseModel): return cls(**record.dict()) return cls(**dict(record)) - def to_db_dict(self, exclude_unset: bool = False) -> dict[str, Any]: + def to_db_dict(self, exclude_unset: bool = False) -> dict[str, object]: """Convert domain model to a dictionary for database operations.""" return self.model_dump(exclude_none=True, exclude_unset=exclude_unset) diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py b/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py index c09106273e1..150900e7ff2 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py @@ -36,7 +36,10 @@ a healed fleet has no null rows and the backfill exits after one query. import json from collections import Counter -from typing import Any, Final, Literal +from collections.abc import Mapping, Sequence +from typing import Final, Literal, Protocol + +from pydantic import JsonValue from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.db import _decode_oauth_payload, decrypt_credentials @@ -55,9 +58,59 @@ BackfillRule = Literal[ _BACKFILL_AUDIT_ACTOR: Final = "oauth2_flow_backfill" -def _decrypted_credentials(raw_credentials: Any) -> MCPCredentials | None: +class _MCPServerRow(Protocol): + """The ``LiteLLM_MCPServerTable`` columns this backfill reads.""" + + @property + def server_id(self) -> str: ... + + @property + def authorization_url(self) -> str | None: ... + + @property + def registration_url(self) -> str | None: ... + + @property + def token_url(self) -> str | None: ... + + @property + def credentials(self) -> str | Mapping[str, JsonValue] | None: ... + + +class _MCPUserCredentialRow(Protocol): + """The ``LiteLLM_MCPUserCredentials`` columns this backfill reads.""" + + @property + def server_id(self) -> str: ... + + @property + def credential_b64(self) -> str: ... + + +class _MCPServerTable(Protocol): + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_MCPServerRow]: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, str]) -> object: ... + + +class _MCPUserCredentialsTable(Protocol): + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_MCPUserCredentialRow]: ... + + +def _mcp_server_table(prisma_client: PrismaClient) -> _MCPServerTable: + """The MCP server table, typed so the untyped prisma client surface stops here.""" + return prisma_client.db.litellm_mcpservertable + + +def _mcp_user_credentials_table(prisma_client: PrismaClient) -> _MCPUserCredentialsTable: + """The per-user MCP credential table, typed so the untyped prisma client surface stops here.""" + return prisma_client.db.litellm_mcpusercredentials + + +def _decrypted_credentials(raw_credentials: str | Mapping[str, JsonValue] | None) -> MCPCredentials | None: if raw_credentials is None: return None + parsed: JsonValue | Mapping[str, JsonValue] if isinstance(raw_credentials, str): try: parsed = json.loads(raw_credentials) @@ -92,14 +145,14 @@ def classify_null_flow_row( async def backfill_null_oauth2_flows(prisma_client: PrismaClient) -> dict[BackfillRule, int]: """Classify every ``auth_type=oauth2`` row whose ``oauth2_flow`` is null; stamp the provable ones, warn on the ambiguous ones, and return counts per rule.""" - null_rows: Final[list[Any]] = await prisma_client.db.litellm_mcpservertable.find_many( + null_rows: Final[Sequence[_MCPServerRow]] = await _mcp_server_table(prisma_client).find_many( where={"auth_type": "oauth2", "oauth2_flow": None}, ) if not null_rows: return {} server_ids: Final = [row.server_id for row in null_rows] - token_rows: Final[list[Any]] = await prisma_client.db.litellm_mcpusercredentials.find_many( + token_rows: Final[Sequence[_MCPUserCredentialRow]] = await _mcp_user_credentials_table(prisma_client).find_many( where={"server_id": {"in": server_ids}}, ) server_ids_with_oauth_tokens: Final[set[str]] = { @@ -141,7 +194,7 @@ async def backfill_null_oauth2_flows(prisma_client: PrismaClient) -> dict[Backfi stamped_flows: Final = {flow for _, (flow, _) in classified if flow is not None} for stamped_flow in stamped_flows: server_ids_for_flow = [row.server_id for row, (row_flow, _) in classified if row_flow == stamped_flow] - await prisma_client.db.litellm_mcpservertable.update_many( + await _mcp_server_table(prisma_client).update_many( where={"server_id": {"in": server_ids_for_flow}, "oauth2_flow": None}, data={"oauth2_flow": stamped_flow, "updated_by": _BACKFILL_AUDIT_ACTOR}, ) diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index 9672383a572..ecaaf35e817 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -1,5 +1,9 @@ import json -from typing import Final +from collections.abc import Mapping, Sequence +from datetime import datetime +from typing import Final, Protocol + +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -7,18 +11,73 @@ from litellm.proxy.utils import PrismaClient from litellm.repositories.table_repositories import MCPToolsetRepository from litellm.types.mcp_server.mcp_toolset import ( MCPToolset, + MCPToolsetTool, NewMCPToolsetRequest, UpdateMCPToolsetRequest, ) -def _toolset_from_row(row) -> MCPToolset: +class MCPToolsetFields(TypedDict): + """The ``MCPToolset`` constructor keywords a toolset row expands into.""" + + toolset_id: ReadOnly[str] + toolset_name: ReadOnly[str] + description: NotRequired[ReadOnly[str | None]] + tools: NotRequired[ReadOnly[list[MCPToolsetTool]]] + created_at: NotRequired[ReadOnly[datetime | None]] + created_by: NotRequired[ReadOnly[str | None]] + updated_at: NotRequired[ReadOnly[datetime | None]] + updated_by: NotRequired[ReadOnly[str | None]] + + +class MCPToolsetRowData(TypedDict): + """A toolset table row, whose ``tools`` column is stored as JSON.""" + + toolset_id: ReadOnly[str] + toolset_name: ReadOnly[str] + description: NotRequired[ReadOnly[str | None]] + tools: NotRequired[ReadOnly[str | list[MCPToolsetTool]]] + created_at: NotRequired[ReadOnly[datetime | None]] + created_by: NotRequired[ReadOnly[str | None]] + updated_at: NotRequired[ReadOnly[datetime | None]] + updated_by: NotRequired[ReadOnly[str | None]] + + +class MCPToolsetRow(Protocol): + """A row of the toolset table, as the prisma client returns it.""" + + def model_dump(self) -> MCPToolsetRowData: ... + + +class MCPToolsetTable(Protocol): + """The prisma table actions this module runs against the toolset table.""" + + async def create(self, data: Mapping[str, object]) -> MCPToolsetRow: ... + + async def find_unique(self, where: Mapping[str, object]) -> MCPToolsetRow | None: ... + + async def find_first(self, where: Mapping[str, object]) -> MCPToolsetRow | None: ... + + async def find_many(self, where: Mapping[str, object]) -> Sequence[MCPToolsetRow]: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> MCPToolsetRow: ... + + async def delete(self, where: Mapping[str, object]) -> MCPToolsetRow: ... + + +def _toolset_table(prisma_client: PrismaClient) -> MCPToolsetTable: + """The toolset table actions of the prisma client.""" + return MCPToolsetRepository(prisma_client).table + + +def _toolset_from_row(row: MCPToolsetRow) -> MCPToolset: data: Final = row.model_dump() - tools = data.get("tools") or [] - if isinstance(tools, str): - tools = json.loads(tools) - data["tools"] = tools - return MCPToolset(**data) + tools: Final = data.get("tools") or [] + resolved: Final[MCPToolsetFields] = { + **data, + "tools": json.loads(tools) if isinstance(tools, str) else tools, + } + return MCPToolset(**resolved) async def create_mcp_toolset( @@ -31,7 +90,7 @@ async def create_mcp_toolset( data_dict["tools"] = json.dumps(data_dict.get("tools", [])) data_dict["created_by"] = touched_by data_dict["updated_by"] = touched_by - row: Final = await MCPToolsetRepository(prisma_client).table.create(data=data_dict) + row: Final = await _toolset_table(prisma_client).create(data=data_dict) return _toolset_from_row(row) @@ -39,7 +98,7 @@ async def get_mcp_toolset( prisma_client: PrismaClient, toolset_id: str, ) -> MCPToolset | None: - row: Final = await MCPToolsetRepository(prisma_client).table.find_unique(where={"toolset_id": toolset_id}) + row: Final = await _toolset_table(prisma_client).find_unique(where={"toolset_id": toolset_id}) if row is None: return None return _toolset_from_row(row) @@ -47,13 +106,11 @@ async def get_mcp_toolset( async def list_mcp_toolsets( prisma_client: PrismaClient, - toolset_ids: list[str] | None = None, -) -> list[MCPToolset]: + toolset_ids: Sequence[str] | None = None, +) -> Sequence[MCPToolset]: try: - where = {} - if toolset_ids is not None: - where = {"toolset_id": {"in": toolset_ids}} - rows: Final = await MCPToolsetRepository(prisma_client).table.find_many(where=where) + where: Final[Mapping[str, object]] = {} if toolset_ids is None else {"toolset_id": {"in": toolset_ids}} + rows: Final = await _toolset_table(prisma_client).find_many(where=where) return [_toolset_from_row(r) for r in rows] except Exception as e: verbose_proxy_logger.warning("litellm.proxy._experimental.mcp_server.toolset_db::list_mcp_toolsets - %s", e) @@ -64,7 +121,7 @@ async def get_mcp_toolset_by_name( prisma_client: PrismaClient, toolset_name: str, ) -> MCPToolset | None: - row: Final = await MCPToolsetRepository(prisma_client).table.find_first(where={"toolset_name": toolset_name}) + row: Final = await _toolset_table(prisma_client).find_first(where={"toolset_name": toolset_name}) if row is None: return None return _toolset_from_row(row) @@ -80,7 +137,7 @@ async def update_mcp_toolset( data_dict["tools"] = json.dumps(data_dict["tools"]) data_dict["updated_by"] = touched_by try: - row: Final = await MCPToolsetRepository(prisma_client).table.update( + row: Final = await _toolset_table(prisma_client).update( where={"toolset_id": data.toolset_id}, data=data_dict, ) @@ -98,7 +155,7 @@ async def delete_mcp_toolset( toolset_id: str, ) -> MCPToolset | None: try: - row: Final = await MCPToolsetRepository(prisma_client).table.delete(where={"toolset_id": toolset_id}) + row: Final = await _toolset_table(prisma_client).delete(where={"toolset_id": toolset_id}) except Exception as e: from prisma.errors import RecordNotFoundError diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 64de6827679..edb945fc674 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -12,6 +12,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.table_repositories import AgentsRepository, ObjectPermissionRepository from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -83,9 +84,17 @@ class AgentTableClient(Protocol): async def delete(self, where: Mapping[str, object]) -> AgentRecord: ... +class _AgentsRepositoryView(Protocol): + @property + def table(self) -> AgentTableClient: ... + + +def _agents_table_of(repository: _AgentsRepositoryView) -> AgentTableClient: + return repository.table + + def agents_table(prisma_client: PrismaClient) -> AgentTableClient: - table: Final[AgentTableClient] = AgentsRepository(prisma_client).table - return table + return _agents_table_of(AgentsRepository(prisma_client)) class ObjectPermissionGrantRecord(Protocol): @@ -99,9 +108,17 @@ class ObjectPermissionTableClient(Protocol): async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... +class _ObjectPermissionRepositoryView(Protocol): + @property + def table(self) -> ObjectPermissionTableClient: ... + + +def _object_permission_table_of(repository: _ObjectPermissionRepositoryView) -> ObjectPermissionTableClient: + return repository.table + + def object_permission_table(prisma_client: PrismaClient) -> ObjectPermissionTableClient: - table: Final[ObjectPermissionTableClient] = ObjectPermissionRepository(prisma_client).table - return table + return _object_permission_table_of(ObjectPermissionRepository(prisma_client)) class GrantMigrationResult(NamedTuple): @@ -283,19 +300,21 @@ class AgentRegistry: agent_name: Final = agent.get("agent_name") # Serialize litellm_params - litellm_params_obj: Final[Any] = agent.get("litellm_params", {}) - if hasattr(litellm_params_obj, "model_dump"): - litellm_params_dict = litellm_params_obj.model_dump() - else: - litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {} + litellm_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("litellm_params", {}) + litellm_params_dict: Final[Mapping[str, object]] = ( + litellm_params_obj.model_dump() + if isinstance(litellm_params_obj, SupportsModelDump) + else (dict(litellm_params_obj) if litellm_params_obj else {}) + ) litellm_params: Final[str] = safe_dumps(litellm_params_dict) # Serialize agent_card_params - agent_card_params_obj: Final[Any] = agent.get("agent_card_params", {}) - if hasattr(agent_card_params_obj, "model_dump"): - agent_card_params_dict = agent_card_params_obj.model_dump() - else: - agent_card_params_dict = dict(agent_card_params_obj) if agent_card_params_obj else {} + agent_card_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("agent_card_params", {}) + agent_card_params_dict: Final[Mapping[str, object]] = ( + agent_card_params_obj.model_dump() + if isinstance(agent_card_params_obj, SupportsModelDump) + else (dict(agent_card_params_obj) if agent_card_params_obj else {}) + ) agent_card_params: Final[str] = safe_dumps(agent_card_params_dict) # Handle object_permission (MCP tool access for agent) @@ -386,15 +405,13 @@ class AgentRegistry: The patched agent """ try: - existing_agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) - if existing_agent is not None: - existing_agent = dict(existing_agent) - - if existing_agent is None: + existing_record: Final = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id}) + if existing_record is None: raise Exception(f"Agent with ID {agent_id} not found") + existing_agent: Final[Mapping[str, object]] = dict(existing_record) augment_agent: Final = {**existing_agent, **agent} - update_data: Final[dict[str, Any]] = {} + update_data: Final[dict[str, object]] = {} if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") if augment_agent.get("litellm_params"): @@ -418,7 +435,7 @@ class AgentRegistry: update_data["extra_headers"] = extra_headers_value if extra_headers_value is not None else [] if agent.get("object_permission") is not None: agent_copy: Final = dict(augment_agent) - existing_object_permission_id: Final = existing_agent.get("object_permission_id") + existing_object_permission_id: Final = existing_record.object_permission_id object_permission_id: Final = await handle_update_object_permission_common( agent_copy, existing_object_permission_id, @@ -460,19 +477,21 @@ class AgentRegistry: agent_name: Final = agent.get("agent_name") # Serialize litellm_params - litellm_params_obj: Final[Any] = agent.get("litellm_params", {}) - if hasattr(litellm_params_obj, "model_dump"): - litellm_params_dict = litellm_params_obj.model_dump() - else: - litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {} + litellm_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("litellm_params", {}) + litellm_params_dict: Final[Mapping[str, object]] = ( + litellm_params_obj.model_dump() + if isinstance(litellm_params_obj, SupportsModelDump) + else (dict(litellm_params_obj) if litellm_params_obj else {}) + ) litellm_params: Final[str] = safe_dumps(litellm_params_dict) # Serialize agent_card_params - agent_card_params_obj: Final[Any] = agent.get("agent_card_params", {}) - if hasattr(agent_card_params_obj, "model_dump"): - agent_card_params_dict = agent_card_params_obj.model_dump() - else: - agent_card_params_dict = dict(agent_card_params_obj) if agent_card_params_obj else {} + agent_card_params_obj: Final[Mapping[str, object] | SupportsModelDump] = agent.get("agent_card_params", {}) + agent_card_params_dict: Final[Mapping[str, object]] = ( + agent_card_params_obj.model_dump() + if isinstance(agent_card_params_obj, SupportsModelDump) + else (dict(agent_card_params_obj) if agent_card_params_obj else {}) + ) agent_card_params: Final[str] = safe_dumps(agent_card_params_dict) # Serialize static_headers for update diff --git a/litellm/proxy/client/cli/commands/credentials.py b/litellm/proxy/client/cli/commands/credentials.py index c550b39d33f..2beff05b375 100644 --- a/litellm/proxy/client/cli/commands/credentials.py +++ b/litellm/proxy/client/cli/commands/credentials.py @@ -1,14 +1,42 @@ import json +from collections.abc import Mapping, Sequence from typing import Final, Literal import click import requests import rich from rich.table import Table +from typing_extensions import NotRequired, ReadOnly, TypedDict from ...credentials import CredentialsManagementClient +class _CliContext(TypedDict): + """Values the top-level CLI group stores on the click context.""" + + base_url: ReadOnly[str] + api_key: ReadOnly[str | None] + + +class _CliContextView(TypedDict): + obj: ReadOnly[_CliContext] + + +class _CredentialRow(TypedDict): + """Single credential entry as returned by ``GET /credentials``.""" + + credential_name: ReadOnly[NotRequired[str]] + credential_info: ReadOnly[NotRequired[Mapping[str, object]]] + + +class _CredentialRowsView(TypedDict): + rows: ReadOnly[Sequence[_CredentialRow]] + + +class _JsonBodyView(TypedDict): + body: ReadOnly[object] + + @click.group() def credentials(): """Manage credentials for the LiteLLM proxy server""" @@ -25,7 +53,8 @@ def credentials(): @click.pass_context def list(ctx: click.Context, output_format: Literal["table", "json"]): """List all credentials""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final[_CliContextView] = {"obj": ctx.obj} + client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) response: Final = client.list() assert isinstance(response, dict) @@ -39,7 +68,8 @@ def list(ctx: click.Context, output_format: Literal["table", "json"]): table.add_column("Custom LLM Provider", style="green") # Add rows - for cred in response.get("credentials", []): + credential_rows: Final[_CredentialRowsView] = {"rows": response.get("credentials", [])} + for cred in credential_rows["rows"]: info = cred.get("credential_info", {}) table.add_row( str(cred.get("credential_name", "")), @@ -66,7 +96,8 @@ def list(ctx: click.Context, output_format: Literal["table", "json"]): @click.pass_context def create(ctx: click.Context, credential_name: str, info: str, values: str): """Create a new credential""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final[_CliContextView] = {"obj": ctx.obj} + client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) try: credential_info: Final = json.loads(info) credential_values: Final = json.loads(values) @@ -79,8 +110,8 @@ def create(ctx: click.Context, credential_name: str, info: str, values: str): except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) try: - error_body: Final = e.response.json() - rich.print_json(data=error_body) + error_body: Final[_JsonBodyView] = {"body": e.response.json()} + rich.print_json(data=error_body["body"]) except json.JSONDecodeError: click.echo(e.response.text, err=True) raise click.Abort() @@ -91,15 +122,16 @@ def create(ctx: click.Context, credential_name: str, info: str, values: str): @click.pass_context def delete(ctx: click.Context, credential_name: str): """Delete a credential by name""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final[_CliContextView] = {"obj": ctx.obj} + client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) try: response: Final = client.delete(credential_name) rich.print_json(data=response) except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) try: - error_body: Final = e.response.json() - rich.print_json(data=error_body) + error_body: Final[_JsonBodyView] = {"body": e.response.json()} + rich.print_json(data=error_body["body"]) except json.JSONDecodeError: click.echo(e.response.text, err=True) raise click.Abort() @@ -110,6 +142,7 @@ def delete(ctx: click.Context, credential_name: str): @click.pass_context def get(ctx: click.Context, credential_name: str): """Get a credential by name""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final[_CliContextView] = {"obj": ctx.obj} + client: Final = CredentialsManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) response: Final = client.get(credential_name) rich.print_json(data=response) diff --git a/litellm/proxy/client/cli/commands/teams.py b/litellm/proxy/client/cli/commands/teams.py index e814ac84ebb..1212e194bde 100644 --- a/litellm/proxy/client/cli/commands/teams.py +++ b/litellm/proxy/client/cli/commands/teams.py @@ -1,21 +1,48 @@ """Team management commands for LiteLLM CLI.""" +from collections.abc import Mapping, Sequence from typing import Any, Final import click import requests from rich.console import Console from rich.table import Table +from typing_extensions import ReadOnly, TypedDict from litellm.proxy.client import Client +class _CliContext(TypedDict): + """The proxy connection settings the CLI group stores on the click context.""" + + base_url: ReadOnly[str] + api_key: ReadOnly[str | None] + + +def _cli_context(ctx: click.Context) -> _CliContext: + """The proxy connection settings the CLI group stored on the click context.""" + ctx_obj: Final[_CliContext] = ctx.obj + return ctx_obj + + +def _proxy_client(ctx: click.Context) -> Client: + """A proxy client for the base URL and API key on the click context.""" + ctx_obj: Final = _cli_context(ctx) + return Client(ctx_obj["base_url"], ctx_obj["api_key"]) + + +def _http_error_detail(error: requests.exceptions.HTTPError) -> object: + """The ``detail`` the proxy reported for a failed request.""" + error_body: Final[Mapping[str, object]] = error.response.json() + return error_body.get("detail", "Unknown error") + + @click.group() def teams(): """Manage teams and team assignments""" -def display_teams_table(teams: list[dict[str, Any]]) -> None: +def display_teams_table(teams: Sequence[dict[str, Any]]) -> None: """Display teams in a formatted table""" console: Final = Console() @@ -33,8 +60,8 @@ def display_teams_table(teams: list[dict[str, Any]]) -> None: for i, team in enumerate(teams): team_alias = team.get("team_alias") or "N/A" - team_id = team.get("team_id", "N/A") - models = team.get("models", []) + team_id: str = team.get("team_id", "N/A") + models: Sequence[str] = team.get("models", []) max_budget = team.get("max_budget") # Format models list @@ -64,7 +91,7 @@ def display_teams_table(teams: list[dict[str, Any]]) -> None: @click.pass_context def list(ctx: click.Context): """List teams that you belong to""" - client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"]) + client: Final = _proxy_client(ctx) try: # Use list() for simpler response structure (returns array directly) @@ -72,8 +99,7 @@ def list(ctx: click.Context): display_teams_table(teams) except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) - error_body: Final = e.response.json() - click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) + click.echo(f"Details: {_http_error_detail(e)}", err=True) raise click.Abort() except Exception as e: click.echo(f"Error: {e}", err=True) @@ -84,7 +110,7 @@ def list(ctx: click.Context): @click.pass_context def available(ctx: click.Context): """List teams that are available to join""" - client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"]) + client: Final = _proxy_client(ctx) try: teams: Final = client.teams.get_available() @@ -96,8 +122,7 @@ def available(ctx: click.Context): click.echo("No available teams to join.") except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) - error_body: Final = e.response.json() - click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) + click.echo(f"Details: {_http_error_detail(e)}", err=True) except Exception as e: click.echo(f"Error: {e}", err=True) raise click.Abort() @@ -108,8 +133,8 @@ def available(ctx: click.Context): @click.pass_context def assign_key(ctx: click.Context, team_id: str | None): """Assign your current CLI key to a team""" - client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"]) - api_key: Final = ctx.obj["api_key"] + client: Final = _proxy_client(ctx) + api_key: Final = _cli_context(ctx)["api_key"] if not api_key: click.echo("No API key found. Please login first using 'litellm login'") @@ -145,7 +170,7 @@ def assign_key(ctx: click.Context, team_id: str | None): teams = client.teams.list() for team in teams: if team.get("team_id") == team_id: - models = team.get("models", []) + models: Sequence[str] = team.get("models", []) if models: click.echo(f"You can now access models: {', '.join(models)}") else: @@ -154,8 +179,7 @@ def assign_key(ctx: click.Context, team_id: str | None): except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) - error_body: Final = e.response.json() - click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) + click.echo(f"Details: {_http_error_detail(e)}", err=True) raise click.Abort() except Exception as e: click.echo(f"Error: {e}", err=True) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 9379a8577a3..a28d03eebf0 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,6 +1,6 @@ import copy import os -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias @@ -525,8 +525,8 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset( def sanitize_openai_provider_metadata( - metadata: dict[str, Any] | None, -) -> dict[str, str] | None: + metadata: Mapping[str, object] | None, +) -> Mapping[str, object] | None: """ Keep only provider-safe OpenAI metadata entries (string keys -> string values). @@ -644,7 +644,7 @@ def process_callback(_callback: str, callback_type: str, environment_variables: return {"name": _callback, "variables": env_vars_dict, "type": callback_type} -def normalize_callback_names(callbacks: Iterable[Any]) -> list[Any]: +def normalize_callback_names(callbacks: Iterable[object] | None) -> list[object]: if callbacks is None: return [] return [c.lower() if isinstance(c, str) else c for c in callbacks] @@ -674,7 +674,7 @@ def decrypt_callback_vars(metadata: Any) -> Any: return _transform_callback_vars(metadata, _decrypt_or_passthrough) -def _transform_callback_vars(metadata: Any, transform: Callable[[str, Any], Any]) -> Any: +def _transform_callback_vars(metadata: object, transform: Callable[[str, Any], Any]) -> object: if not isinstance(metadata, dict): return metadata out: Final = copy.deepcopy(metadata) @@ -704,7 +704,7 @@ def is_sensitive_callback_key( return _CALLBACK_VAR_MASKER.is_sensitive_key(key) -def _encrypt_if_plaintext(key: str, value: Any) -> Any: +def _encrypt_if_plaintext(key: str, value: object) -> object: if not isinstance(value, str) or not value: return value if not is_sensitive_callback_key(key): @@ -725,7 +725,7 @@ def _encrypt_if_plaintext(key: str, value: Any) -> Any: return value -def _decrypt_or_passthrough(key: str, value: Any) -> Any: +def _decrypt_or_passthrough(key: str, value: object) -> object: if not isinstance(value, str) or not value: return value if not value.startswith(_CALLBACK_VAR_ENCRYPTED_PREFIX): diff --git a/litellm/proxy/common_utils/get_routes.py b/litellm/proxy/common_utils/get_routes.py index 2118a6610b4..5aeb755071c 100644 --- a/litellm/proxy/common_utils/get_routes.py +++ b/litellm/proxy/common_utils/get_routes.py @@ -2,63 +2,78 @@ Utility class for getting routes from a FastAPI app. """ +from collections.abc import Sequence from typing import Any, Final from starlette.routing import BaseRoute +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger +class RouteInfo(TypedDict, total=False): + """One entry of the app's route listing.""" + + path: ReadOnly[object] + methods: ReadOnly[object] + name: ReadOnly[object] + endpoint: ReadOnly[str | None] + mounted_app: ReadOnly[bool] + + class GetRoutes: @staticmethod def get_app_routes( route: BaseRoute, endpoint_route: Any, - ) -> list[dict[str, Any]]: + ) -> Sequence[RouteInfo]: """ Get routes for a regular route. """ - routes: Final[list[dict[str, Any]]] = [] - route_info: Final = { + route_info: Final[RouteInfo] = { "path": getattr(route, "path", None), "methods": getattr(route, "methods", None), "name": getattr(route, "name", None), "endpoint": (endpoint_route.__name__ if getattr(route, "endpoint", None) else None), } - routes.append(route_info) - return routes + return [route_info] @staticmethod def get_routes_for_mounted_app( route: BaseRoute, - ) -> list[dict[str, Any]]: + ) -> Sequence[RouteInfo]: """ Get routes for a mounted sub-application. """ - routes: Final[list[dict[str, Any]]] = [] + routes: Final[list[RouteInfo]] = [] mount_path: Final = getattr(route, "path", "") - sub_app: Final = getattr(route, "app", None) - if sub_app and hasattr(sub_app, "routes"): - for sub_route in sub_app.routes: - # Get endpoint - either from endpoint attribute or app attribute - endpoint_func = getattr(sub_route, "endpoint", None) or getattr(sub_route, "app", None) + for sub_route in GetRoutes._mounted_app_routes(route): + endpoint_func: object = getattr(sub_route, "endpoint", None) or getattr(sub_route, "app", None) - if endpoint_func is not None: - sub_route_path = getattr(sub_route, "path", "") - full_path = mount_path.rstrip("/") + sub_route_path + if endpoint_func is not None: + sub_route_path = getattr(sub_route, "path", "") + full_path = mount_path.rstrip("/") + sub_route_path - route_info = { - "path": full_path, - "methods": getattr(sub_route, "methods", ["GET", "POST"]), - "name": getattr(sub_route, "name", None), - "endpoint": GetRoutes._safe_get_endpoint_name(endpoint_func), - "mounted_app": True, - } - routes.append(route_info) + route_info: RouteInfo = { + "path": full_path, + "methods": getattr(sub_route, "methods", ["GET", "POST"]), + "name": getattr(sub_route, "name", None), + "endpoint": GetRoutes._safe_get_endpoint_name(endpoint_func), + "mounted_app": True, + } + routes.append(route_info) return routes @staticmethod - def _safe_get_endpoint_name(endpoint_function: Any) -> str | None: + def _mounted_app_routes(route: BaseRoute) -> Sequence[BaseRoute]: + """The routes of the sub-application mounted at ``route``, if it mounts one.""" + sub_app: Final[object] = getattr(route, "app", None) + if sub_app and hasattr(sub_app, "routes"): + return getattr(sub_app, "routes") + return () + + @staticmethod + def _safe_get_endpoint_name(endpoint_function: object) -> str | None: """ Safely get the name of the endpoint function. """ @@ -66,7 +81,7 @@ class GetRoutes: if hasattr(endpoint_function, "__name__"): return getattr(endpoint_function, "__name__") elif hasattr(endpoint_function, "__class__") and hasattr(endpoint_function.__class__, "__name__"): - return getattr(endpoint_function.__class__, "__name__") + return endpoint_function.__class__.__name__ else: return None except Exception: diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 93d51bdd461..a8a71134034 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -40,30 +40,31 @@ class UserApiKeyCache(DualCache): @overload def get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: object, + parent_otel_span: object = None, local_only: bool = False, *, model_type: type[T], - **kwargs: Any, + **kwargs: object, ) -> T | None: ... @overload def get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: object, + parent_otel_span: object = None, local_only: bool = False, - **kwargs: Any, + model_type: None = None, + **kwargs: object, ) -> Any: ... def get_cache( self, - key, - parent_otel_span=None, + key: object, + parent_otel_span: object = None, local_only: bool = False, model_type: type[BaseModel] | None = None, - **kwargs, + **kwargs: object, ) -> Any | BaseModel | None: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) @@ -85,30 +86,31 @@ class UserApiKeyCache(DualCache): @overload async def async_get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: object, + parent_otel_span: object = None, local_only: bool = False, *, model_type: type[T], - **kwargs: Any, + **kwargs: object, ) -> T | None: ... @overload async def async_get_cache( self, - key: Any, - parent_otel_span: Any = None, + key: object, + parent_otel_span: object = None, local_only: bool = False, - **kwargs: Any, + model_type: None = None, + **kwargs: object, ) -> Any: ... async def async_get_cache( self, - key, - parent_otel_span=None, + key: object, + parent_otel_span: object = None, local_only: bool = False, model_type: type[BaseModel] | None = None, - **kwargs, + **kwargs: object, ) -> Any | BaseModel | None: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) @@ -129,17 +131,17 @@ class UserApiKeyCache(DualCache): return None return decoded - def set_cache(self, key, value, local_only: bool = False, **kwargs): + def set_cache(self, key: object, value: object, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) - payload: Final = CacheCodec.serialize(value, model_type=model_type) + payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs) - async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): + async def async_set_cache(self, key: object, value: object, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) - payload: Final = CacheCodec.serialize(value, model_type=model_type) + payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) - async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs) -> None: + async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs: object) -> None: """ Batch writes with the same Codec boundary as ``async_set_cache`` without ``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged. diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 5aeb52be535..9cc4369fea6 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -6,11 +6,14 @@ otherwise PrismaClient uses the writer-only PrismaWrapper directly. import os from collections.abc import Callable -from typing import Any, Final +from datetime import timedelta +from typing import Any, Final, TypeAlias from litellm._logging import verbose_proxy_logger from litellm.proxy.db.prisma_client import PrismaWrapper +ConnectTimeout: TypeAlias = "int | timedelta | None" + # Per-model action methods that read from the database. These are routed to # the read replica when one is configured. _MODEL_READ_METHODS: Final = frozenset( @@ -31,6 +34,11 @@ _MODEL_READ_METHODS: Final = frozenset( _TOP_LEVEL_READ_METHODS: Final = frozenset({"query_first", "query_raw"}) +def _dynamic_attr(target: object, name: str) -> object: + """Fetch `name` off an unstubbed Prisma object as an opaque value.""" + return getattr(target, name) + + class _RoutedActions: """Per-model accessor that sends reads to the reader and writes to the writer. @@ -43,18 +51,18 @@ class _RoutedActions: def __init__( self, - writer_actions: Any, - reader_actions: Any, + writer_actions: object, + reader_actions: object, should_use_reader: Callable[[], bool], ): self._writer_actions = writer_actions self._reader_actions = reader_actions self._should_use_reader = should_use_reader - def __getattr__(self, name: str) -> Any: + def __getattr__(self, name: str) -> object: if name in _MODEL_READ_METHODS and self._should_use_reader(): - return getattr(self._reader_actions, name) - return getattr(self._writer_actions, name) + return _dynamic_attr(self._reader_actions, name) + return _dynamic_attr(self._writer_actions, name) class RoutingPrismaWrapper: @@ -135,7 +143,7 @@ class RoutingPrismaWrapper: return not self._reader_unavailable @staticmethod - async def _try_connect(client: PrismaWrapper, *args: Any, **kwargs: Any) -> Exception | None: + async def _try_connect(client: PrismaWrapper, *args: ConnectTimeout, **kwargs: ConnectTimeout) -> Exception | None: if client.is_connected() is True: return None try: @@ -144,7 +152,7 @@ class RoutingPrismaWrapper: except Exception as e: return e - async def connect(self, *args: Any, **kwargs: Any) -> None: + async def connect(self, *args: ConnectTimeout, **kwargs: ConnectTimeout) -> None: writer_error: Final = await self._try_connect(self._writer, *args, **kwargs) if writer_error is None: self._writer_unavailable = False @@ -176,7 +184,7 @@ class RoutingPrismaWrapper: writer_error, ) - async def disconnect(self, *args: Any, **kwargs: Any) -> None: + async def disconnect(self, *args: object, **kwargs: object) -> None: first_error: BaseException | None = None for client in (self._writer, self._reader): try: @@ -206,7 +214,7 @@ class RoutingPrismaWrapper: async def recreate_prisma_client( self, new_db_url: str, - http_client: Any | None = None, + http_client: object | None = None, *, expected_generation: int | None = None, ) -> bool: @@ -245,7 +253,7 @@ class RoutingPrismaWrapper: ) return True - async def _recreate_reader(self, http_client: Any | None = None) -> None: + async def _recreate_reader(self, http_client: object | None = None) -> None: """Resolve the reader URL and recreate its Prisma client. IAM-enabled readers regenerate their token (host/port/user came from @@ -265,14 +273,14 @@ class RoutingPrismaWrapper: def __getattr__(self, name: str) -> Any: if name in _TOP_LEVEL_READ_METHODS: - return getattr(self.read_target, name) - writer_attr: Final = getattr(self._writer, name) + return _dynamic_attr(self.read_target, name) + writer_attr: Final = _dynamic_attr(self._writer, name) # Per-model action accessors are non-callable instances that expose # both `find_many` and `create`. Methods like execute_raw / batch_ / # tx are callables and stay on the writer untouched. if not callable(writer_attr) and hasattr(writer_attr, "find_many") and hasattr(writer_attr, "create"): try: - reader_attr: Final = getattr(self._reader, name) + reader_attr: Final = _dynamic_attr(self._reader, name) except AttributeError: return writer_attr return _RoutedActions(writer_attr, reader_attr, self._should_use_reader) diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py index da12222f233..707118bf016 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py @@ -15,6 +15,8 @@ restriction intact. """ import operator +from collections.abc import Callable +from types import CodeType from typing import Any, Final from RestrictedPython import ( @@ -58,7 +60,7 @@ class AsyncAwareTransformer(RestrictingNodeTransformer): return self.node_contents_visit(node) -_INPLACE_OPS: Final[dict[str, Any]] = { +_INPLACE_OPS: Final[dict[str, Callable[[object, object], object]]] = { "+=": operator.iadd, "-=": operator.isub, "*=": operator.imul, @@ -75,7 +77,7 @@ _INPLACE_OPS: Final[dict[str, Any]] = { } -def _inplacevar_(op: str, x: Any, y: Any) -> Any: +def _inplacevar_(op: str, x: object, y: object) -> object: # RestrictedPython rewrites ``x += 1`` on a simple name into # ``x = _inplacevar_("+=", x, 1)``. The package deliberately ships no # default, so we dispatch through ``operator``'s in-place helpers, which @@ -86,7 +88,7 @@ def _inplacevar_(op: str, x: Any, y: Any) -> Any: return fn(x, y) -def _build_sandbox_builtins() -> dict[str, Any]: +def _build_sandbox_builtins() -> dict[str, object]: # ``limited_builtins`` overrides ``list``/``tuple``/``range`` from # ``safe_builtins`` with bounds-checking variants (e.g. ``limited_range`` # rejects ``range(10**18)``). ``utility_builtins`` adds ``set``, @@ -98,14 +100,14 @@ def _build_sandbox_builtins() -> dict[str, Any]: } -def build_sandbox_globals() -> dict[str, Any]: +def build_sandbox_globals() -> dict[str, object]: """Assemble the globals dict for executing guardrail code. Includes the LiteLLM-provided primitives (``regex_match``, ``http_get``, ``allow``/``block``/``modify``, etc.) plus the RestrictedPython guards that the compiled bytecode expects to find by name. """ - sandbox: Final[dict[str, Any]] = get_custom_code_primitives().copy() + sandbox: Final[dict[str, object]] = get_custom_code_primitives().copy() sandbox["__builtins__"] = _build_sandbox_builtins() sandbox["_getattr_"] = safer_getattr sandbox["_getitem_"] = default_guarded_getitem @@ -116,7 +118,7 @@ def build_sandbox_globals() -> dict[str, Any]: return sandbox -def compile_sandboxed(source: str, filename: str = "") -> Any: +def compile_sandboxed(source: str, filename: str = "") -> CodeType: """Compile guardrail source with RestrictedPython's AST transformer. Raises ``SyntaxError`` on either a Python syntax error or a restricted diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py index 507dd645953..0ae048c040d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py @@ -2,7 +2,7 @@ from __future__ import annotations import os from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import urlparse from uuid import uuid4 @@ -11,7 +11,7 @@ import requests from fastapi import HTTPException from httpx import HTTPStatusError from requests.auth import HTTPBasicAuth -from typing_extensions import ReadOnly +from typing_extensions import ReadOnly, TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -36,24 +36,31 @@ if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + """Base-class constructor options carried by this guardrail's forwarded keyword arguments.""" + + guardrail_name: ReadOnly[str | None] + supported_event_hooks: list[GuardrailEventHooks] | None + + class _HiddenlayerEvaluation(TypedDict, total=False): - action: str - threat_level: str + action: ReadOnly[str] + threat_level: ReadOnly[str] class _HiddenlayerAnalysisEntry(TypedDict, total=False): - name: str - detected: bool + name: ReadOnly[str] + detected: ReadOnly[bool] class _HiddenlayerModifiedSide(TypedDict): - messages: Any + messages: ReadOnly[Any] class _HiddenlayerResponse(TypedDict, total=False): - evaluation: _HiddenlayerEvaluation - analysis: Sequence[_HiddenlayerAnalysisEntry] - modified_data: Mapping[str, _HiddenlayerModifiedSide] + evaluation: ReadOnly[_HiddenlayerEvaluation] + analysis: ReadOnly[Sequence[_HiddenlayerAnalysisEntry]] + modified_data: ReadOnly[Mapping[str, _HiddenlayerModifiedSide]] class _LoggedCallMetadata(TypedDict, total=False): @@ -151,7 +158,7 @@ class HiddenlayerGuardrail(CustomGuardrail): api_key: str | None = None, api_base: str | None = None, auth_url: str | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID") @@ -356,7 +363,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): api_key: str | None = None, api_base: str | None = None, auth_url: str | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID") self.hiddenlayer_client_secret = api_key or os.getenv("HIDDENLAYER_CLIENT_SECRET") @@ -486,7 +493,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): self, payload: Any, input_type: Literal["request", "response"], - hl_headers: dict[str, str], + hl_headers: Mapping[str, str], ) -> httpx.Response: if input_type == "request": path = "detection/v2/request-evaluations" diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py index e3f67f0024b..c5c36ae51c4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py @@ -1,10 +1,11 @@ """LLM-as-a-Judge guardrail: uses an LLM to score responses against weighted criteria.""" -from collections.abc import Callable +from collections.abc import Callable, Mapping, MutableMapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional from fastapi import HTTPException +from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack import litellm from litellm._logging import verbose_logger @@ -41,28 +42,87 @@ _parse_judge_verdict: Final = parse_json_verdict _extract_text_from_content: Final = extract_text_from_content +class _JudgeMessage(TypedDict): + """Chat message, as far as the judge prompt builder reads it.""" + + role: ReadOnly[NotRequired[str]] + content: ReadOnly[NotRequired[object]] + + +class _GuardrailOptions(TypedDict, total=False): + """Base :class:`CustomGuardrail` options forwarded untouched.""" + + mask_request_content: ReadOnly[bool] + mask_response_content: ReadOnly[bool] + violation_message_template: ReadOnly[str | None] + end_session_after_n_fails: ReadOnly[int | None] + on_violation: ReadOnly[str | None] + realtime_violation_message: ReadOnly[str | None] + on_sensitive_data: ReadOnly[str | None] + sensitive_data_route_to_model: ReadOnly[str | None] + sticky_session_routing: ReadOnly[bool] + run_in_parallel: ReadOnly[bool] + only_scan_new_messages: ReadOnly[bool] + + +class _RequestMessagesView(TypedDict): + messages: ReadOnly[Sequence[_JudgeMessage]] + + +class _RequestMetadataView(TypedDict): + metadata: ReadOnly[MutableMapping[str, object]] + + +class _OverallScoreView(TypedDict): + overall_score: ReadOnly[str | float] + + +class _JudgeModelView(TypedDict): + judge_model: ReadOnly[str] + + +class _CriteriaView(TypedDict): + criteria: ReadOnly[Sequence[Mapping[str, str | float]]] + + +class _OnFailureView(TypedDict): + on_failure: ReadOnly[Literal["block", "log"]] + + +class _ThresholdView(TypedDict): + overall_threshold: ReadOnly[str | float] + + +class _ModeView(TypedDict): + mode: ReadOnly[object] + + +class _DefaultOnView(TypedDict): + default_on: ReadOnly[object] + + def _get_litellm_param( litellm_params: "LitellmParams", guardrail: "Guardrail", key: str, - default: Any = None, + default: str | float | bool | None = None, ) -> Any: - val: Final = getattr(litellm_params, key, None) + val: Final[object] = getattr(litellm_params, key, None) if val is not None: return val raw: Final = guardrail.get("litellm_params") if isinstance(raw, dict) and key in raw: return raw[key] if raw is not None and not isinstance(raw, dict): - attr: Final = getattr(raw, key, None) + attr: Final[object] = getattr(raw, key, None) if attr is not None: return attr return default def _build_judge_prompt( - criteria: list[dict[str, Any]], - messages: list[dict[str, Any]], + criteria: Sequence[Mapping[str, object]], + messages: Sequence[_JudgeMessage], response_text: str, ) -> str: criteria_block: Final = "\n".join( @@ -87,13 +147,13 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): self, guardrail_name: str, judge_model: str, - criteria: list[dict[str, Any]], + criteria: Sequence[Mapping[str, object]], overall_threshold: float = 80.0, on_failure: Literal["block", "log"] = "block", event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None, default_on: bool = False, router_provider: "Callable[[], Router | None] | None" = None, - **kwargs: Any, + **kwargs: Unpack[_GuardrailOptions], ) -> None: _event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None if event_hook is not None: @@ -121,9 +181,9 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): async def _run_judge( self, - messages: list[dict[str, Any]], + messages: Sequence[_JudgeMessage], response_text: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: judge_messages: Final = [ {"role": "system", "content": JUDGE_SYSTEM_PROMPT}, { @@ -162,10 +222,10 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): judge_result: dict[str, Any] = {} try: - messages: Final[list[dict[str, Any]]] = request_data.get("messages") or [] + request_messages: Final[_RequestMessagesView] = {"messages": request_data.get("messages") or []} try: - judge_result = await self._run_judge(messages, response_text) + judge_result = await self._run_judge(request_messages["messages"], response_text) except Exception as judge_err: verbose_logger.warning( "llm_as_a_judge guardrail: judge call failed, failing open. Error: %s", judge_err @@ -174,7 +234,8 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): return inputs try: - overall_score: Final = max(0.0, min(100.0, float(judge_result.get("overall_score", 100)))) + raw_score: Final[_OverallScoreView] = {"overall_score": judge_result.get("overall_score", 100)} + overall_score: Final = max(0.0, min(100.0, float(raw_score["overall_score"]))) except (TypeError, ValueError): verbose_logger.warning("llm_as_a_judge: invalid overall_score from judge, failing open") return inputs @@ -189,7 +250,8 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): "threshold": self.overall_threshold, "verdicts": judge_result.get("verdicts", []), } - _metadata: Final = request_data.setdefault("metadata", {}) + request_metadata: Final[_RequestMetadataView] = {"metadata": request_data.setdefault("metadata", {})} + _metadata: Final = request_metadata["metadata"] existing: Final = _metadata.get("eval_information") if isinstance(existing, list): existing.append(eval_info) @@ -238,37 +300,45 @@ def initialize_guardrail( if not guardrail_name: raise ValueError("llm_as_a_judge guardrail requires a guardrail_name") - judge_model: Final = _get_litellm_param(litellm_params, guardrail, "judge_model") - if not judge_model: + judge_model: Final[_JudgeModelView] = {"judge_model": _get_litellm_param(litellm_params, guardrail, "judge_model")} + if not judge_model["judge_model"]: raise ValueError("llm_as_a_judge guardrail requires judge_model in litellm_params") - criteria: Final = _get_litellm_param(litellm_params, guardrail, "criteria") or [] - if not criteria: + criteria: Final[_CriteriaView] = {"criteria": _get_litellm_param(litellm_params, guardrail, "criteria") or []} + if not criteria["criteria"]: raise ValueError("llm_as_a_judge guardrail requires at least one criterion") - weight_total: Final = sum(float(c.get("weight", 0)) for c in criteria) + weight_total: Final = sum(float(c.get("weight", 0)) for c in criteria["criteria"]) if abs(weight_total - 100) > 0.5: raise ValueError(f"llm_as_a_judge criterion weights must sum to 100 (got {weight_total})") - on_failure: Final = _get_litellm_param(litellm_params, guardrail, "on_failure", "block") - if on_failure not in _VALID_ON_FAILURE: - raise ValueError(f"llm_as_a_judge on_failure must be 'block' or 'log', got '{on_failure}'") + on_failure: Final[_OnFailureView] = { + "on_failure": _get_litellm_param(litellm_params, guardrail, "on_failure", "block") + } + if on_failure["on_failure"] not in _VALID_ON_FAILURE: + raise ValueError(f"llm_as_a_judge on_failure must be 'block' or 'log', got '{on_failure['on_failure']}'") - overall_threshold: Final = float(_get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0)) + threshold: Final[_ThresholdView] = { + "overall_threshold": _get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0) + } + overall_threshold: Final = float(threshold["overall_threshold"]) - mode: Final = _get_litellm_param(litellm_params, guardrail, "mode") + mode: Final[_ModeView] = {"mode": _get_litellm_param(litellm_params, guardrail, "mode")} event_hook: GuardrailEventHooks | None = None - if isinstance(mode, str) and mode in {e.value for e in GuardrailEventHooks}: - event_hook = GuardrailEventHooks(mode) + if isinstance(mode["mode"], str) and mode["mode"] in {e.value for e in GuardrailEventHooks}: + event_hook = GuardrailEventHooks(mode["mode"]) + default_on: Final[_DefaultOnView] = { + "default_on": _get_litellm_param(litellm_params, guardrail, "default_on", False) + } instance: Final = LLMAsAJudgeGuardrail( guardrail_name=guardrail_name, - judge_model=judge_model, - criteria=criteria, + judge_model=judge_model["judge_model"], + criteria=criteria["criteria"], overall_threshold=overall_threshold, - on_failure=on_failure, + on_failure=on_failure["on_failure"], event_hook=event_hook, - default_on=bool(_get_litellm_param(litellm_params, guardrail, "default_on", False)), + default_on=bool(default_on["default_on"]), ) litellm.logging_callback_manager.add_litellm_callback(instance) return instance diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index e9cd6addef8..704e2564ef5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -7,8 +7,9 @@ import enum import json import os +from collections.abc import Callable, Mapping from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, cast from urllib.parse import urlparse from litellm._logging import verbose_proxy_logger @@ -36,6 +37,8 @@ _AIDR_SCAN_ENDPOINT: Final = "/litellm/guardrail" _INTERVENED_INPUT_FIELDS: Final = ("texts", "images", "tools", "tool_calls") _DEFAULT_API_BASE_HOSTNAME: Final = urlparse(_DEFAULT_API_BASE).hostname +_GuardrailJsonResponse: TypeAlias = Exception | str | dict[str, object] + _KEYS_DUPLICATING_SCAN_INPUTS: Final = ("messages", "input") _LOGGING_KEYS_DUPLICATING_SCAN_INPUTS: Final = _KEYS_DUPLICATING_SCAN_INPUTS + ( "additional_args", @@ -119,7 +122,7 @@ class NomaV2Guardrail(CustomGuardrail): def _resolve_action_from_response( self, - response_json: dict, + response_json: Mapping[str, object], ) -> _Action: action: Final = response_json.get("action") if isinstance(action, str): @@ -165,10 +168,11 @@ class NomaV2Guardrail(CustomGuardrail): @staticmethod def _sanitize_payload_for_transport(payload: dict) -> dict: - def _default(obj: Any) -> Any: - if hasattr(obj, "model_dump"): + def _default(obj: object) -> object: + model_dump: Final[Callable[[], Mapping[str, object]] | None] = getattr(obj, "model_dump", None) + if model_dump is not None: try: - return obj.model_dump() + return model_dump() except Exception: pass return str(obj) @@ -178,7 +182,7 @@ class NomaV2Guardrail(CustomGuardrail): except (ValueError, TypeError): json_str = safe_dumps(payload) - safe_payload: Final = safe_json_loads(json_str, default={}) + safe_payload: Final[object] = safe_json_loads(json_str, default={}) if safe_payload == {} and payload: verbose_proxy_logger.warning( "Noma v2 guardrail: payload serialization failed, falling back to empty payload" @@ -196,7 +200,7 @@ class NomaV2Guardrail(CustomGuardrail): async def _call_noma_scan( self, payload: dict, - ) -> dict: + ) -> dict[str, object]: headers: Final[dict[str, str]] = {"Content-Type": "application/json"} authorization_header: Final = self._get_authorization_header() if authorization_header: @@ -215,7 +219,7 @@ class NomaV2Guardrail(CustomGuardrail): response.text, ) response.raise_for_status() - response_json: Final = response.json() + response_json: Final[dict[str, object]] = response.json() verbose_proxy_logger.debug( "Noma v2 AIDR response parsed: %s", json.dumps(response_json, default=str), @@ -227,7 +231,7 @@ class NomaV2Guardrail(CustomGuardrail): request_data: dict, start_time: datetime, guardrail_status: GuardrailStatus, - guardrail_json_response: Any, + guardrail_json_response: _GuardrailJsonResponse, ) -> None: end_time: Final = datetime.now() duration: Final = (end_time - start_time).total_seconds() @@ -270,11 +274,11 @@ class NomaV2Guardrail(CustomGuardrail): ) -> GenericGuardrailAPIInputs: start_time: Final = datetime.now() guardrail_status: GuardrailStatus = "success" - guardrail_json_response: Any = {} + guardrail_json_response: _GuardrailJsonResponse = {} dynamic_params = self.get_guardrail_dynamic_request_body_params(request_data) if not isinstance(dynamic_params, dict): dynamic_params = {} - response_json: dict | None = None + response_json: dict[str, object] | None = None # Per-request dynamic params can override configured application context. application_id = self._get_non_empty_str(dynamic_params.get("application_id")) diff --git a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py index 4775a8b3caa..155db816e94 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py @@ -7,8 +7,11 @@ before and after LLM calls. """ import os +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing_extensions import ReadOnly, TypedDict + from litellm._logging import verbose_proxy_logger from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_guardrail import ( @@ -20,6 +23,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -34,6 +38,16 @@ _DEFAULT_API_BASE: Final = "https://api.promptguard.co" _GUARD_ENDPOINT: Final = "/api/v1/guard" +class PromptGuardResult(TypedDict, total=False): + """The fields this guardrail reads off a PromptGuard Guard API response.""" + + decision: ReadOnly[str] + threat_type: ReadOnly[str] + event_id: ReadOnly[str] + confidence: ReadOnly[float] + redacted_messages: ReadOnly[list[AllMessageValues]] + + class PromptGuardMissingCredentials(Exception): pass @@ -96,7 +110,7 @@ class PromptGuardGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: Mapping[str, object], input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: @@ -114,7 +128,7 @@ class PromptGuardGuardrail(CustomGuardrail): direction: Final = "input" if input_type == "request" else "output" - payload: Final[dict[str, Any]] = { + payload: Final[dict[str, object]] = { "messages": messages, "direction": direction, } @@ -144,7 +158,7 @@ class PromptGuardGuardrail(CustomGuardrail): timeout=10.0, ) response.raise_for_status() - result: Final = response.json() + result: Final[PromptGuardResult] = response.json() except Exception as exc: verbose_proxy_logger.error("PromptGuard API error: %s", str(exc)) if self.block_on_error: diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 41e52f05c01..7b45afa0ad2 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -1,4 +1,6 @@ -from typing import Final +from collections.abc import Mapping, Sequence +from datetime import datetime +from typing import Final, Protocol from fastapi import APIRouter, Depends, HTTPException, Query @@ -18,7 +20,59 @@ from litellm.repositories.table_repositories import JWTKeyMappingRepository router: Final = APIRouter() -def _to_response(mapping) -> JWTKeyMappingResponse: +class _JWTKeyMappingRecord(Protocol): + """A ``LiteLLM_JWTKeyMapping`` row, viewed through the columns these endpoints read.""" + + @property + def id(self) -> str: ... + + @property + def jwt_claim_name(self) -> str: ... + + @property + def jwt_claim_value(self) -> str: ... + + @property + def description(self) -> str | None: ... + + @property + def is_active(self) -> bool: ... + + @property + def created_at(self) -> datetime: ... + + @property + def updated_at(self) -> datetime: ... + + @property + def created_by(self) -> str | None: ... + + @property + def updated_by(self) -> str | None: ... + + +class _JWTKeyMappingTable(Protocol): + """The Prisma table actions these endpoints issue against the JWT key mapping table.""" + + async def create(self, *, data: Mapping[str, object]) -> _JWTKeyMappingRecord: ... + + async def find_unique(self, *, where: Mapping[str, object]) -> _JWTKeyMappingRecord | None: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> _JWTKeyMappingRecord: ... + + async def delete(self, *, where: Mapping[str, object]) -> _JWTKeyMappingRecord | None: ... + + async def find_many(self, *, skip: int, take: int, order: Mapping[str, str]) -> Sequence[_JWTKeyMappingRecord]: ... + + async def count(self) -> int: ... + + +def _mapping_table(prisma_client: object) -> _JWTKeyMappingTable: + """View the JWT key mapping repository's untyped Prisma table through the actions used here.""" + return JWTKeyMappingRepository(prisma_client).table + + +def _to_response(mapping: _JWTKeyMappingRecord) -> JWTKeyMappingResponse: """Convert a Prisma mapping object to a safe response (no hashed token).""" return JWTKeyMappingResponse( id=mapping.id, @@ -62,7 +116,7 @@ async def create_jwt_key_mapping( if data.description is not None: create_data["description"] = data.description - new_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.create(data=create_data) + new_mapping: Final = await _mapping_table(prisma_client).create(data=create_data) # Invalidate cache cache_key: Final = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}" @@ -110,7 +164,7 @@ async def update_jwt_key_mapping( try: # Get old mapping for cache invalidation - old_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": data.id}) + old_mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": data.id}) if old_mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") @@ -118,9 +172,7 @@ async def update_jwt_key_mapping( cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" await user_api_key_cache.async_delete_cache(cache_key) - updated_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.update( - where={"id": data.id}, data=update_data - ) + updated_mapping: Final = await _mapping_table(prisma_client).update(where={"id": data.id}, data=update_data) # Invalidate new cache key if claim fields changed cache_key = f"jwt_key_mapping:{updated_mapping.jwt_claim_name}:{updated_mapping.jwt_claim_value}" @@ -159,7 +211,7 @@ async def delete_jwt_key_mapping( try: # Get old mapping for cache invalidation - old_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": data.id}) + old_mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": data.id}) if old_mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") @@ -167,7 +219,7 @@ async def delete_jwt_key_mapping( cache_key: Final = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" await user_api_key_cache.async_delete_cache(cache_key) - await JWTKeyMappingRepository(prisma_client).table.delete(where={"id": data.id}) + await _mapping_table(prisma_client).delete(where={"id": data.id}) return {"status": "success"} except HTTPException: raise @@ -195,12 +247,12 @@ async def list_jwt_key_mappings( try: skip: Final = (page - 1) * size - mappings: Final = await JWTKeyMappingRepository(prisma_client).table.find_many( + mappings: Final = await _mapping_table(prisma_client).find_many( skip=skip, take=size, order={"created_at": "desc"}, ) - total_count: Final = await JWTKeyMappingRepository(prisma_client).table.count() + total_count: Final = await _mapping_table(prisma_client).count() return { "mappings": [_to_response(m) for m in mappings], "total_count": total_count, @@ -232,7 +284,7 @@ async def info_jwt_key_mapping( raise HTTPException(status_code=500, detail="Database not connected") try: - mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": id}) + mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": id}) if mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") return _to_response(mapping) diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 020698dabd9..3c5d889d23f 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -10,8 +10,8 @@ https://platform.openai.com/docs/api-reference/responses-streaming import asyncio import json -from collections.abc import Sequence -from typing import TYPE_CHECKING, Final, TypedDict, cast +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final, TypedDict from fastapi import Request, Response from fastapi.responses import StreamingResponse @@ -38,6 +38,36 @@ class _StreamOutputItem(TypedDict, total=False): content: ReadOnly[Sequence[_StreamContentPart | None]] +class _StreamTerminalResponse(TypedDict, total=False): + """Fields of the ``response`` payload carried by a terminal streaming event.""" + + status: ReadOnly[ResponsesAPIStatus] + tool_choice: ReadOnly[object] + model: ReadOnly[str] + instructions: ReadOnly[str] + temperature: ReadOnly[float] + top_p: ReadOnly[float] + max_output_tokens: ReadOnly[int] + previous_response_id: ReadOnly[str] + truncation: ReadOnly[str] + parallel_tool_calls: ReadOnly[bool] + user: ReadOnly[str] + store: ReadOnly[bool] + output: ReadOnly[Sequence[_StreamOutputItem]] + + +class _StreamEvent(TypedDict, total=False): + """One decoded ``data:`` frame of an OpenAI Responses streaming body.""" + + type: ReadOnly[str] + item: ReadOnly[_StreamOutputItem] + item_id: ReadOnly[str] + part: ReadOnly[_StreamContentPart] + content_index: ReadOnly[int] + delta: ReadOnly[str] + response: ReadOnly[_StreamTerminalResponse] + + async def background_streaming_task( polling_id: str, data, @@ -139,7 +169,7 @@ async def background_streaming_task( None # Will be set by response.completed/failed/incomplete/cancelled ) terminal_error = None - _event_to_status: Final = { + _event_to_status: Final[Mapping[str, ResponsesAPIStatus]] = { "response.completed": "completed", "response.failed": "failed", "response.incomplete": "incomplete", @@ -180,7 +210,7 @@ async def background_streaming_task( break try: - event = json.loads(chunk_data) + event: _StreamEvent = json.loads(chunk_data) event_type = event.get("type", "") # Process different event types based on OpenAI streaming spec @@ -288,12 +318,9 @@ async def background_streaming_task( # Terminal event - extract all ResponsesAPIResponse fields # https://platform.openai.com/docs/api-reference/responses-streaming response_data = event.get("response", {}) - terminal_status = cast( - ResponsesAPIStatus, - response_data.get( - "status", - _event_to_status.get(event_type, "completed"), - ), + terminal_status = response_data.get( + "status", + _event_to_status.get(event_type, "completed"), ) # Extract error for failed and incomplete responses diff --git a/litellm/proxy/search_endpoints/search_tool_registry.py b/litellm/proxy/search_endpoints/search_tool_registry.py index fe4794f3ba1..b25263e4c64 100644 --- a/litellm/proxy/search_endpoints/search_tool_registry.py +++ b/litellm/proxy/search_endpoints/search_tool_registry.py @@ -2,8 +2,9 @@ Search Tool Registry for managing search tool configurations. """ +from collections.abc import Iterator, Mapping, Sequence from datetime import datetime, timezone -from typing import Final +from typing import Final, Protocol from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -13,6 +14,40 @@ from litellm.repositories.table_repositories import SearchToolsRepository from litellm.types.search import SearchTool +class SearchToolRecord(Protocol): + search_tool_id: str + search_tool_name: str + created_at: datetime + updated_at: datetime + + def __iter__(self) -> Iterator[tuple[str, object]]: ... + + +class SearchToolTableClient(Protocol): + async def create(self, data: Mapping[str, object]) -> SearchToolRecord: ... + + async def find_unique(self, where: Mapping[str, object]) -> SearchToolRecord | None: ... + + async def find_many(self, order: Mapping[str, str] | None = None) -> Sequence[SearchToolRecord]: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> SearchToolRecord: ... + + async def delete(self, where: Mapping[str, object]) -> SearchToolRecord: ... + + +class _SearchToolsRepositoryView(Protocol): + @property + def table(self) -> SearchToolTableClient: ... + + +def _search_tools_table_of(repository: _SearchToolsRepositoryView) -> SearchToolTableClient: + return repository.table + + +def _search_tools_table(prisma_client: PrismaClient) -> SearchToolTableClient: + return _search_tools_table_of(SearchToolsRepository(prisma_client)) + + class SearchToolRegistry: """ Handles adding, removing, and getting search tools in DB + in memory. @@ -22,7 +57,7 @@ class SearchToolRegistry: pass @staticmethod - def _convert_prisma_to_dict(prisma_obj) -> dict: + def _convert_prisma_to_dict(prisma_obj: SearchToolRecord) -> dict: """ Convert Prisma result to dict with datetime objects as ISO format strings. @@ -35,9 +70,9 @@ class SearchToolRegistry: result: Final = dict(prisma_obj) # Convert datetime objects to ISO format strings if "created_at" in result and result["created_at"]: - result["created_at"] = result["created_at"].isoformat() + result["created_at"] = prisma_obj.created_at.isoformat() if "updated_at" in result and result["updated_at"]: - result["updated_at"] = result["updated_at"].isoformat() + result["updated_at"] = prisma_obj.updated_at.isoformat() return result ########################################################### @@ -61,7 +96,7 @@ class SearchToolRegistry: search_tool_info: Final[str] = safe_dumps(search_tool.get("search_tool_info", {})) # Create search tool in DB - created_search_tool: Final = await SearchToolsRepository(prisma_client).table.create( + created_search_tool: Final = await _search_tools_table(prisma_client).create( data={ "search_tool_name": search_tool_name, "litellm_params": litellm_params, @@ -95,7 +130,7 @@ class SearchToolRegistry: """ try: # Get search tool before deletion for response - existing_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique( + existing_tool: Final = await _search_tools_table(prisma_client).find_unique( where={"search_tool_id": search_tool_id} ) @@ -103,7 +138,7 @@ class SearchToolRegistry: raise Exception(f"Search tool with ID {search_tool_id} not found") # Delete from DB - await SearchToolsRepository(prisma_client).table.delete(where={"search_tool_id": search_tool_id}) + await _search_tools_table(prisma_client).delete(where={"search_tool_id": search_tool_id}) return { "message": f"Search tool {search_tool_id} deleted successfully", @@ -131,7 +166,7 @@ class SearchToolRegistry: search_tool_info: Final[str] = safe_dumps(search_tool.get("search_tool_info", {})) # Update in DB - updated_search_tool: Final = await SearchToolsRepository(prisma_client).table.update( + updated_search_tool: Final = await _search_tools_table(prisma_client).update( where={"search_tool_id": search_tool_id}, data={ "search_tool_name": search_tool_name, @@ -163,7 +198,7 @@ class SearchToolRegistry: try: search_tools_from_db: Final = await call_with_db_reconnect_retry( prisma_client, - lambda: SearchToolsRepository(prisma_client).table.find_many( + lambda: _search_tools_table(prisma_client).find_many( order={"created_at": "desc"}, ), reason="get_all_search_tools_from_db_lookup_failure", @@ -194,7 +229,7 @@ class SearchToolRegistry: Search tool configuration or None if not found """ try: - search_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique( + search_tool: Final = await _search_tools_table(prisma_client).find_unique( where={"search_tool_id": search_tool_id} ) @@ -222,7 +257,7 @@ class SearchToolRegistry: Search tool configuration or None if not found """ try: - search_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique( + search_tool: Final = await _search_tools_table(prisma_client).find_unique( where={"search_tool_name": search_tool_name} ) diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py index 16b8f82815c..255faf94402 100644 --- a/litellm/rag/rag_query.py +++ b/litellm/rag/rag_query.py @@ -1,11 +1,45 @@ +from collections.abc import Sequence from typing import Any, Final +from typing_extensions import NotRequired, ReadOnly, TypedDict + from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage from litellm.types.utils import ModelResponse -from litellm.types.vector_stores import ( - VectorStoreResultContent, - VectorStoreSearchResponse, -) +from litellm.types.vector_stores import VectorStoreSearchResponse + + +class _ResultContentView(TypedDict): + """Content entry carried by a vector store search result.""" + + type: ReadOnly[NotRequired[str]] + text: ReadOnly[str] + + +class _SearchResultView(TypedDict): + """Vector store search result, as far as :class:`RAGQuery` reads it.""" + + content: ReadOnly[NotRequired[Sequence[_ResultContentView]]] + text: ReadOnly[NotRequired[str]] + + +class _SearchDataView(TypedDict): + results: ReadOnly[Sequence[_SearchResultView]] + + +class _ContextChunksView(TypedDict): + chunks: ReadOnly[Sequence[_SearchResultView | str | None]] + + +class _RerankResultView(TypedDict): + index: ReadOnly[NotRequired[int]] + + +class _RerankResultsView(TypedDict): + results: ReadOnly[Sequence[_RerankResultView]] + + +class _MessageView(TypedDict): + message: ReadOnly[object] class RAGQuery: @@ -42,9 +76,10 @@ class RAGQuery: """ context_content = RAGQuery.CONTENT_PREFIX_STRING - for chunk in context_chunks: + chunks: Final[_ContextChunksView] = {"chunks": context_chunks} + for chunk in chunks["chunks"]: if isinstance(chunk, dict): - result_content: list[VectorStoreResultContent] | None = chunk.get("content") + result_content: Sequence[_ResultContentView] | None = chunk.get("content") if result_content: for content_item in result_content: content_text: str | None = content_item.get("text") @@ -64,14 +99,15 @@ class RAGQuery: def add_search_results_to_response( response: ModelResponse, search_results: VectorStoreSearchResponse, - rerank_results: Any | None = None, + rerank_results: object = None, ) -> ModelResponse: """ Add search results to the response choices. """ if hasattr(response, "choices") and response.choices: for choice in response.choices: - message = getattr(choice, "message", None) + message_view: _MessageView = {"message": getattr(choice, "message", None)} + message = message_view["message"] if message is not None: # Get existing provider_specific_fields or create new dict provider_fields = getattr(message, "provider_specific_fields", None) or {} @@ -91,7 +127,8 @@ class RAGQuery: ) -> list[str | dict[str, Any]]: """Extract text documents from vector store search response.""" documents: Final[list[str | dict[str, Any]]] = [] - for result in search_response.get("data", []): + search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])} + for result in search_data["results"]: content_list = result.get("content", []) for content in content_list: if content.get("type") == "text" and content.get("text"): @@ -99,11 +136,13 @@ class RAGQuery: return documents @staticmethod - def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> list[Any]: + def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> list[_SearchResultView]: """Get the original search results corresponding to the top reranked results.""" - top_chunks: Final = [] - original_results: Final = search_response.get("data", []) - for result in rerank_response.get("results", []): + top_chunks: Final[list[_SearchResultView]] = [] + search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])} + original_results: Final = search_data["results"] + reranked: Final[_RerankResultsView] = {"results": rerank_response.get("results", [])} + for result in reranked["results"]: index = result.get("index") if index is not None and index < len(original_results): top_chunks.append(original_results[index]) diff --git a/litellm/repositories/base_repository.py b/litellm/repositories/base_repository.py index 7008099fe8c..8dfe3ed962e 100644 --- a/litellm/repositories/base_repository.py +++ b/litellm/repositories/base_repository.py @@ -3,7 +3,7 @@ Base repository class with common functionality. """ from abc import ABC, abstractmethod -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from typing import Any, Final, Generic, Protocol, TypeVar, runtime_checkable from pydantic import BaseModel @@ -24,6 +24,22 @@ class SupportsDict(Protocol): DbRecord = Mapping[str, object] | SupportsModelDump | SupportsDict | Sequence[tuple[str, object]] +class PrismaCrudActions(Protocol): + """The Prisma table actions reached by the generic repository CRUD helpers.""" + + async def find_unique(self, *, where: Mapping[str, object]) -> DbRecord | None: ... + + find_many: Callable[..., Awaitable[Sequence[DbRecord]]] + + async def create(self, *, data: Mapping[str, object]) -> DbRecord: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> DbRecord | None: ... + + async def delete(self, *, where: Mapping[str, object]) -> DbRecord | None: ... + + async def count(self, *, where: Mapping[str, object] | None = None) -> int: ... + + def record_to_dict(record: DbRecord) -> Mapping[str, object]: """Project a database record into a mapping of column name to value.""" if isinstance(record, SupportsModelDump): @@ -38,7 +54,7 @@ def record_to_dict(record: DbRecord) -> Mapping[str, object]: class BaseRepository(ABC, Generic[T]): """Abstract base class for all repositories.""" - def __init__(self, prisma_client: Any): # any-ok: PrismaClient is an untyped runtime wrapper + def __init__(self, prisma_client: object): self._prisma_client = prisma_client @property @@ -53,6 +69,11 @@ class BaseRepository(ABC, Generic[T]): """Return the Prisma table for this repository.""" ... + @property + def _crud_actions(self) -> PrismaCrudActions: + """View ``table`` through the action surface the CRUD helpers below use.""" + return self.table + @property @abstractmethod def model_class(self) -> type[T]: @@ -71,18 +92,18 @@ class BaseRepository(ABC, Generic[T]): async def find_by_id(self, id_value: str, id_field: str = "id") -> T | None: """Find a record by its primary key.""" - record: Final = await self.table.find_unique(where={id_field: id_value}) + record: Final = await self._crud_actions.find_unique(where={id_field: id_value}) return self._to_model(record) async def find_many( self, - where: dict[str, Any] | None = None, + where: Mapping[str, object] | None = None, skip: int | None = None, take: int | None = None, - order: dict[str, str] | None = None, + order: Mapping[str, str] | None = None, ) -> list[T]: """Find multiple records matching the criteria.""" - kwargs: Final[dict[str, Any]] = {} + kwargs: Final[dict[str, object]] = {} if where: kwargs["where"] = where if skip is not None: @@ -92,31 +113,31 @@ class BaseRepository(ABC, Generic[T]): if order: kwargs["order"] = order - records: Final = await self.table.find_many(**kwargs) + records: Final = await self._crud_actions.find_many(**kwargs) return self._to_model_list(records) - async def create(self, data: dict[str, Any]) -> T: + async def create(self, data: Mapping[str, object]) -> T: """Create a new record.""" - record: Final = await self.table.create(data=data) + record: Final = await self._crud_actions.create(data=data) model: Final = self._to_model(record) assert model is not None return model - async def update(self, id_value: str, data: dict[str, Any], id_field: str = "id") -> T | None: + async def update(self, id_value: str, data: Mapping[str, object], id_field: str = "id") -> T | None: """Update an existing record.""" - record: Final = await self.table.update(where={id_field: id_value}, data=data) + record: Final = await self._crud_actions.update(where={id_field: id_value}, data=data) return self._to_model(record) async def delete(self, id_value: str, id_field: str = "id") -> T | None: """Delete a record by its primary key.""" - record: Final = await self.table.delete(where={id_field: id_value}) + record: Final = await self._crud_actions.delete(where={id_field: id_value}) return self._to_model(record) - async def count(self, where: dict[str, Any] | None = None) -> int: + async def count(self, where: Mapping[str, object] | None = None) -> int: """Count records matching the criteria.""" - return await self.table.count(where=where) + return await self._crud_actions.count(where=where) async def exists(self, id_value: str, id_field: str = "id") -> bool: """Check if a record exists.""" - record: Final = await self.table.find_unique(where={id_field: id_value}) + record: Final = await self._crud_actions.find_unique(where={id_field: id_value}) return record is not None diff --git a/litellm/repositories/credentials_repository.py b/litellm/repositories/credentials_repository.py index 9fdb6e4aca7..b22f24e0de5 100644 --- a/litellm/repositories/credentials_repository.py +++ b/litellm/repositories/credentials_repository.py @@ -6,11 +6,41 @@ credential values is the caller's responsibility (see ``CredentialHelperUtils``) so reads return the stored values verbatim. """ -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Protocol from litellm.models.credentials import CredentialItem from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync +if TYPE_CHECKING: + from prisma.models import LiteLLM_CredentialsTable + + +class _CredentialsDb(Protocol): + @property + def litellm_credentialstable(self) -> object: ... + + +class _PrismaClientView(Protocol): + @property + def db(self) -> _CredentialsDb: ... + + +class _CredentialsActions(Protocol): + """Prisma table actions used by :class:`CredentialsRepository`.""" + + async def find_many(self) -> "Sequence[LiteLLM_CredentialsTable]": ... + + async def create(self, *, data: Mapping[str, object]) -> "LiteLLM_CredentialsTable": ... + + async def find_unique(self, *, where: Mapping[str, object]) -> "LiteLLM_CredentialsTable | None": ... + + async def update( + self, *, where: Mapping[str, object], data: Mapping[str, object] + ) -> "LiteLLM_CredentialsTable | None": ... + + async def delete(self, *, where: Mapping[str, object]) -> "LiteLLM_CredentialsTable | None": ... + class CredentialsRepository: """Repository for credentials database operations, keyed by credential name.""" @@ -19,7 +49,7 @@ class CredentialsRepository: self._prisma_client = prisma_client @property - def prisma_client(self) -> Any: + def prisma_client(self) -> _PrismaClientView: if self._prisma_client is None: raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") return self._prisma_client @@ -31,6 +61,10 @@ class CredentialsRepository: table_name="litellm_credentialstable", ) + @property + def _credentials_table(self) -> _CredentialsActions: + return self.table + @staticmethod def _to_model(record: Any) -> CredentialItem | None: if record is None: @@ -42,18 +76,20 @@ class CredentialsRepository: credential_info=data.get("credential_info") or {}, ) - async def find_all(self) -> Any: - return await self.table.find_many() + async def find_all(self) -> "Sequence[LiteLLM_CredentialsTable]": + return await self._credentials_table.find_many() - async def create(self, data: dict[str, Any]) -> Any: - return await self.table.create(data=data) + async def create(self, data: Mapping[str, object]) -> "LiteLLM_CredentialsTable": + return await self._credentials_table.create(data=data) async def find_by_name(self, credential_name: str) -> CredentialItem | None: - record: Final = await self.table.find_unique(where={"credential_name": credential_name}) + record: Final = await self._credentials_table.find_unique(where={"credential_name": credential_name}) return self._to_model(record) - async def update_by_name(self, credential_name: str, data: dict[str, Any]) -> Any: - return await self.table.update(where={"credential_name": credential_name}, data=data) + async def update_by_name( + self, credential_name: str, data: Mapping[str, object] + ) -> "LiteLLM_CredentialsTable | None": + return await self._credentials_table.update(where={"credential_name": credential_name}, data=data) - async def delete_by_name(self, credential_name: str) -> Any: - return await self.table.delete(where={"credential_name": credential_name}) + async def delete_by_name(self, credential_name: str) -> "LiteLLM_CredentialsTable | None": + return await self._credentials_table.delete(where={"credential_name": credential_name}) diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index 7efd32288e4..68490797348 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -3,9 +3,10 @@ Team repository for database operations on LiteLLM_TeamTable. """ import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from contextlib import AbstractAsyncContextManager from datetime import datetime -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from pydantic import TypeAdapter @@ -13,12 +14,45 @@ from litellm.models.team import LiteLLM_TeamTable, Member from litellm.repositories.base_repository import ( BaseRepository, DbRecord, + PrismaCrudActions, record_to_dict, ) if TYPE_CHECKING: from prisma import Prisma + +class _TeamTables(Protocol): + """The Prisma tables this repository reaches, on the client or inside a transaction.""" + + litellm_teamtable: PrismaCrudActions + litellm_deletedteamtable: PrismaCrudActions + + +class _TeamArrays(Protocol): + """The string array columns of a team row, which the domain model leaves untyped.""" + + @property + def members(self) -> Sequence[str]: ... + + @property + def admins(self) -> Sequence[str]: ... + + @property + def models(self) -> Sequence[str]: ... + + +def _team_arrays(team: LiteLLM_TeamTable) -> _TeamArrays: + """View a team's untyped list columns as sequences of ids.""" + return team + + +class _TeamDatabase(_TeamTables, Protocol): + """The Prisma client surface used for team reads, writes, and archival transactions.""" + + def tx(self) -> AbstractAsyncContextManager[_TeamTables]: ... + + _MEMBERS_WITH_ROLES_ADAPTER: Final = TypeAdapter(list[Member]) _JSON_ENCODED_TEAM_FIELDS: Final = ( "metadata", @@ -34,12 +68,16 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): """Repository for team database operations.""" @property - def table(self) -> Any: # any-ok: PrismaClient.db is an untyped runtime wrapper - return self.prisma_client.db.litellm_teamtable + def _db(self) -> _TeamDatabase: + return self.prisma_client.db @property - def deleted_table(self) -> Any: # any-ok: PrismaClient.db is an untyped runtime wrapper - return self.prisma_client.db.litellm_deletedteamtable + def table(self) -> Any: # any-ok: callers reach model-specific actions this repository does not use + return self._db.litellm_teamtable + + @property + def deleted_table(self) -> PrismaCrudActions: + return self._db.litellm_deletedteamtable @property def model_class(self) -> type[LiteLLM_TeamTable]: @@ -75,8 +113,8 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): ) if not rows: return None - raw_value: Final = rows[0]["members_with_roles"] - parsed: Final = json.loads(raw_value) if isinstance(raw_value, str) else raw_value + raw_value: Final[object] = rows[0]["members_with_roles"] + parsed: Final[object] = json.loads(raw_value) if isinstance(raw_value, str) else raw_value if not parsed: return [] return _MEMBERS_WITH_ROLES_ADAPTER.validate_python(parsed) @@ -86,24 +124,24 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): async def find_by_alias(self, team_alias: str) -> LiteLLM_TeamTable | None: """Find a team by alias.""" - records: Final = await self.table.find_many(where={"team_alias": team_alias}) + records: Final = await self._crud_actions.find_many(where={"team_alias": team_alias}) if records: return self._to_model(records[0]) return None async def find_by_organization_id(self, organization_id: str) -> list[LiteLLM_TeamTable]: """Find all teams belonging to an organization.""" - records: Final = await self.table.find_many(where={"organization_id": organization_id}) + records: Final = await self._crud_actions.find_many(where={"organization_id": organization_id}) return self._to_model_list(records) async def find_by_member(self, user_id: str) -> list[LiteLLM_TeamTable]: """Find all teams where user is a member.""" - records: Final = await self.table.find_many(where={"members": {"has": user_id}}) + records: Final = await self._crud_actions.find_many(where={"members": {"has": user_id}}) return self._to_model_list(records) async def find_by_admin(self, user_id: str) -> list[LiteLLM_TeamTable]: """Find all teams where user is an admin.""" - records: Final = await self.table.find_many(where={"admins": {"has": user_id}}) + records: Final = await self._crud_actions.find_many(where={"admins": {"has": user_id}}) return self._to_model_list(records) async def create_team( @@ -232,7 +270,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): archive_data["litellm_changed_by"] = litellm_changed_by archive_data["deleted_at"] = datetime.utcnow() - async with self.prisma_client.db.tx() as tx: + async with self._db.tx() as tx: await tx.litellm_deletedteamtable.create(data=archive_data) await tx.litellm_teamtable.delete(where={"team_id": team_id}) @@ -293,7 +331,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if not await self.exists(team_id, id_field="team_id"): return None - record: Final = await self.table.update( + record: Final = await self._crud_actions.update( where={"team_id": team_id}, data={"members": {"push": user_id}}, ) @@ -310,7 +348,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if team is None: return None - members: Final = [m for m in team.members if m != user_id] + members: Final = [m for m in _team_arrays(team).members if m != user_id] return await self.update(team_id, {"members": members}, id_field="team_id") async def add_admin(self, team_id: str, user_id: str) -> LiteLLM_TeamTable | None: @@ -318,7 +356,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if not await self.exists(team_id, id_field="team_id"): return None - record: Final = await self.table.update( + record: Final = await self._crud_actions.update( where={"team_id": team_id}, data={"admins": {"push": user_id}}, ) @@ -335,7 +373,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if team is None: return None - admins: Final = [a for a in team.admins if a != user_id] + admins: Final = [a for a in _team_arrays(team).admins if a != user_id] return await self.update(team_id, {"admins": admins}, id_field="team_id") async def add_models(self, team_id: str, models: list[str]) -> LiteLLM_TeamTable | None: @@ -343,7 +381,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if not await self.exists(team_id, id_field="team_id"): return None - record: Final = await self.table.update( + record: Final = await self._crud_actions.update( where={"team_id": team_id}, data={"models": {"push": models}}, ) @@ -360,5 +398,5 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): if team is None: return None - current_models: Final = [m for m in team.models if m not in models] + current_models: Final = [m for m in _team_arrays(team).models if m not in models] return await self.update(team_id, {"models": current_models}, id_field="team_id") diff --git a/litellm/rust_bridge/responses_websocket.py b/litellm/rust_bridge/responses_websocket.py index 20fc2634a8a..c57a2b14dcb 100644 --- a/litellm/rust_bridge/responses_websocket.py +++ b/litellm/rust_bridge/responses_websocket.py @@ -2,8 +2,9 @@ from __future__ import annotations +from collections.abc import Awaitable from dataclasses import dataclass -from typing import Any, Final, Protocol +from typing import Final, Protocol import httpx from websockets.exceptions import ConnectionClosedOK @@ -12,6 +13,16 @@ from litellm.rust_bridge.loader import get_native_bridge from litellm.rust_bridge.timeouts import timeout_to_seconds +class RustResponsesWebSocketSocket(Protocol): + """Open socket handle handed back by the native bridge.""" + + def send_text(self, text: str) -> Awaitable[None]: ... + + def recv_text(self) -> Awaitable[str | None]: ... + + def close(self) -> Awaitable[None]: ... + + class RustResponsesWebSocketConnection(Protocol): @classmethod def connect( @@ -19,7 +30,7 @@ class RustResponsesWebSocketConnection(Protocol): url: str, headers: dict[str, str], timeout_seconds: float | None, - ) -> Any: + ) -> Awaitable[RustResponsesWebSocketSocket]: raise NotImplementedError @@ -32,7 +43,7 @@ _UNSET: Final[_Unset] = _Unset() @dataclass(slots=True) class _RustResponsesWebSocketState: - connection: Any = None + connection: type[RustResponsesWebSocketConnection] | None = None _STATE: Final[_RustResponsesWebSocketState] = _RustResponsesWebSocketState() @@ -40,13 +51,13 @@ _STATE: Final[_RustResponsesWebSocketState] = _RustResponsesWebSocketState() def set_rust_responses_websocket( *, - connection: Any = _UNSET, + connection: type[RustResponsesWebSocketConnection] | None | _Unset = _UNSET, ) -> None: if not isinstance(connection, _Unset): _STATE.connection = connection -def load_rust_responses_websocket() -> Any: +def load_rust_responses_websocket() -> type[RustResponsesWebSocketConnection] | None: if _STATE.connection is not None: return _STATE.connection native_bridge: Final = get_native_bridge() @@ -59,7 +70,7 @@ def load_rust_responses_websocket() -> Any: class _ConnectionAdapter: - def __init__(self, connection: Any): + def __init__(self, connection: RustResponsesWebSocketSocket): self._connection = connection async def send(self, text: str) -> None: diff --git a/litellm/secret_managers/secret_manager_handler.py b/litellm/secret_managers/secret_manager_handler.py index c77d2505d0e..a85382f863e 100644 --- a/litellm/secret_managers/secret_manager_handler.py +++ b/litellm/secret_managers/secret_manager_handler.py @@ -6,14 +6,62 @@ Handles retrieving secrets from different secret management systems. import base64 import os -from typing import Any, Final +from collections.abc import Mapping +from typing import Final, Protocol, overload import litellm from litellm._logging import print_verbose -from litellm.types.secret_managers.main import KeyManagementSystem +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem -def _is_base64(s): +class _VaultSecret(Protocol): + """The secret object returned by the Azure Key Vault and Infisical clients.""" + + @property + def value(self) -> str | None: ... + + @property + def secret_value(self) -> str | None: ... + + +class _KmsPlaintext(Protocol): + """The decrypted payload the AWS KMS client exposes under the ``Plaintext`` response key.""" + + def decode(self, encoding: str) -> str | None: ... + + +class _KmsDecryptResponse(Protocol): + """The decrypt response returned by the Google KMS and AWS KMS clients.""" + + @property + def plaintext(self) -> bytes: ... + + def __getitem__(self, key: str) -> _KmsPlaintext: ... + + +class _SecretManagerClient(Protocol): + """The untyped secret manager client surface reached by ``get_secret_from_manager``.""" + + def get_secret(self, secret_name: str, /) -> _VaultSecret: ... + + @overload + def decrypt(self, *, request: Mapping[str, object]) -> _KmsDecryptResponse: ... + + @overload + def decrypt(self, *, CiphertextBlob: bytes) -> _KmsDecryptResponse: ... + + def sync_read_secret( + self, + *, + secret_name: str, + primary_secret_name: str | None = None, + optional_params: Mapping[str, object] | None = None, + ) -> str | None: ... + + def get_secret_from_google_secret_manager(self, secret_name: str, /) -> str | None: ... + + +def _is_base64(s: str | bytes) -> bool: """Check if a string is valid base64.""" import binascii @@ -24,10 +72,10 @@ def _is_base64(s): def get_secret_from_manager( - client: Any, + client: _SecretManagerClient, key_manager: str, secret_name: str, - key_management_settings: Any | None = None, + key_management_settings: KeyManagementSettings | None = None, ) -> str | None: """ Get a secret from the configured secret manager. @@ -56,7 +104,7 @@ def get_secret_from_manager( elif ( key_manager == KeyManagementSystem.GOOGLE_KMS.value or client.__class__.__name__ == "KeyManagementServiceClient" ): - encrypted_secret: Any = os.getenv(secret_name) + encrypted_secret: str | bytes | None = os.getenv(secret_name) if encrypted_secret is None: raise ValueError("Google KMS requires the encrypted secret to be in the environment!") b64_flag: Final = _is_base64(encrypted_secret) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 5c312dcf1c8..ceaf1ec504f 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,21 +1,21 @@ { "ANN001": { - "limit": 3020 + "limit": 3012 }, "ANN002": { "limit": 71 }, "ANN003": { - "limit": 827 + "limit": 818 }, "ANN201": { - "limit": 2017 + "limit": 2014 }, "ANN202": { - "limit": 852 + "limit": 851 }, "ANN204": { - "limit": 711 + "limit": 707 }, "ANN205": { "limit": 112 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1188 + "limit": 1061 }, "ASYNC230": { "limit": 11 @@ -198,7 +198,7 @@ "limit": 22 }, "SIM101": { - "limit": 58 + "limit": 57 }, "SIM102": { "limit": 317 @@ -234,7 +234,7 @@ "limit": 5 }, "TID251": { - "limit": 1212 + "limit": 1202 }, "TRY002": { "limit": 524 @@ -249,7 +249,7 @@ "limit": 113 }, "TRY300": { - "limit": 859 + "limit": 858 }, "UP028": { "limit": 2 diff --git a/tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py b/tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py index d106cf7ea21..a768f8ccd8c 100644 --- a/tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py +++ b/tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py @@ -313,6 +313,25 @@ class TestGDCGeminiConfig: api_base=TEST_API_BASE, ) + def test_validate_environment_credentials_missing_audience_binding_are_named(self): + config = GDCGeminiConfig() + creds_without_audience_binding = MagicMock(spec=[]) + + with patch( + "google.auth.load_credentials_from_dict", + return_value=(creds_without_audience_binding, None), + ): + with pytest.raises(AttributeError, match="must expose with_gdch_audience"): + config.validate_environment( + headers={}, + model=TEST_MODEL, + messages=[], + optional_params={}, + litellm_params={"vertex_project": TEST_PROJECT}, + api_key=TEST_API_KEY, + api_base=TEST_API_BASE, + ) + def test_validate_environment_string_false_disables_token_caching(self): config = GDCGeminiConfig() mock_creds = MagicMock() diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 20cd1165577..890184fad78 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22805 + "limit": 22728 }, "LIT002": { - "limit": 26878 + "limit": 26860 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1069 + "limit": 1066 }, "LIT007": { "limit": 0 @@ -27,12 +27,12 @@ "limit": 0 }, "LIT010": { - "limit": 16695 + "limit": 16673 }, "LIT011": { - "limit": 5588 + "limit": 5586 }, "LIT012": { - "limit": 4519 + "limit": 4512 } } From 5a7edc9c77836f8f47634d8d47719bb44059a1fc Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 24 Aug 2026 19:26:52 -0400 Subject: [PATCH 014/111] fix(spend-tracking): make the new batch count keys writable SpendLogsMetadata is built by assigning each key in turn, so ReadOnly on the two new ones breached the basedpyright reportTypedDictNotRequiredAccess ceiling. Every sibling key in this TypedDict is writable for the same reason. --- litellm/proxy/_types.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bf0dbc97482..377eb7dce0e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3528,8 +3528,8 @@ class SpendLogsMetadata(TypedDict): status: StandardLoggingPayloadStatus proxy_server_request: str | None batch_models: list[str] | None - batch_successful_requests: ReadOnly[int | None] - batch_failed_requests: ReadOnly[int | 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 From 4c2f0f3632c39ae8015f136b02eb002ee24ad248 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 24 Aug 2026 19:34:16 -0400 Subject: [PATCH 015/111] test(batches): cover decoding a model-encoded error file id The error file now resolves through _provider_output_file_id like the output file does. Sending the encoded id straight to the provider 404s, and the swallowed fetch failure would silently report zero failures. --- .../test_litellm/batches/test_batch_utils.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 1540a587349..c86c7c4df03 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1164,6 +1164,53 @@ async def test_handle_completed_batch_counts_error_file_failures(monkeypatch): assert result.failed_requests == 1 +@pytest.mark.asyncio +async def test_handle_completed_batch_decodes_model_encoded_error_file_id(monkeypatch): + """A model-encoded error file id must be decoded to the raw provider id before + the fetch, exactly like the output file id. Sending the encoded id straight to + the provider 404s, and the swallowed fetch failure silently reports 0 failures.""" + import base64 + + from litellm.types.llms.openai import Batch + + provider_error_file_id = "file-real-error-id" + encoded_error_file_id = "file-" + base64.urlsafe_b64encode( + f"litellm:{provider_error_file_id};model,model-abc".encode() + ).decode().rstrip("=") + + requested_file_ids = [] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return _vertex_jsonl([_success_row(model="gpt-4o", usage=_usage(10, 5))]) + + async def fake_afile_content(**kw): + requested_file_ids.append(kw["file_id"]) + return type("R", (), {"content": _vertex_jsonl([{"custom_id": "bad-1"}])})() + + import litellm.files.main as files_main + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) + + batch = Batch( + id="b", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + output_file_id="of", + error_file_id=encoded_error_file_id, + ) + + result = await bu._handle_completed_batch(batch, custom_llm_provider="openai") + + assert requested_file_ids == [provider_error_file_id] + assert result.failed_requests == 1 + + @pytest.mark.asyncio async def test_handle_completed_batch_no_error_file_id_reports_zero_error_failures(monkeypatch): rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] From 3ebf09464a9fd81c4aa0f0fa8edde0016ed54104 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 25 Aug 2026 10:12:33 -0700 Subject: [PATCH 016/111] fix(proxy): give every `requests` call a timeout so a silent server cannot hang the caller `requests` has no default timeout, so a host that accepts the connection and never answers blocks the calling thread forever. The one on the request path is the HiddenLayer guardrail's `_get_jwt`. It runs synchronously inside `_call_hiddenlayer` whenever the hour-long JWT expires and the API answers 401, so a stalled auth host parked the worker's whole event loop, not just the guarded request. The other eight are the teams and users CLI clients, which pin the operator's terminal instead. `TeamsManagementClient` and `UsersManagementClient` now take the same `timeout: int = 30` their `HTTPClient` sibling already had, and `Client` threads its own timeout down to teams. `_poll_for_ready_data` already passed a timeout through a TypedDict that ruff could not see into; passing the argument directly retires both the TypedDict and the suppression it would have needed. Graduate S113 into ruff.toml so the next `requests` call without a timeout fails the lint step. --- litellm/proxy/client/cli/commands/auth.py | 10 +--- litellm/proxy/client/client.py | 2 +- litellm/proxy/client/teams.py | 10 ++-- litellm/proxy/client/users.py | 15 +++--- .../hiddenlayer/hiddenlayer.py | 7 ++- ruff.toml | 3 +- tests/test_litellm/proxy/client/conftest.py | 38 +++++++++++++++ tests/test_litellm/proxy/client/test_teams.py | 20 ++++++++ tests/test_litellm/proxy/client/test_users.py | 16 +++++++ .../guardrail_hooks/test_hiddenlayer.py | 48 +++++++++++++++++++ type-discipline-budget.json | 4 +- 11 files changed, 148 insertions(+), 25 deletions(-) create mode 100644 tests/test_litellm/proxy/client/conftest.py create mode 100644 tests/test_litellm/proxy/client/test_teams.py diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 550b11311f5..2fad9f933c1 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -99,11 +99,6 @@ class CliPollData(TypedDict, total=False): team_id: str -class CliPollRequestKwargs(TypedDict, total=False): - timeout: int - headers: dict[str, str] - - class CliSsoStartData(TypedDict): login_id: str poll_secret: str @@ -518,10 +513,7 @@ def _poll_for_ready_data( ) -> CliPollData | None: for attempt in range(total_timeout // poll_interval): try: - request_kwargs: CliPollRequestKwargs = {"timeout": request_timeout} - if headers is not None: - request_kwargs["headers"] = headers - response = requests.get(url, **request_kwargs) + response = requests.get(url, headers=headers, timeout=request_timeout) if response.status_code == 200: data: CliPollData = response.json() status = data.get("status") diff --git a/litellm/proxy/client/client.py b/litellm/proxy/client/client.py index d71802e06c8..560523db189 100644 --- a/litellm/proxy/client/client.py +++ b/litellm/proxy/client/client.py @@ -38,4 +38,4 @@ class Client: 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.teams = TeamsManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout) diff --git a/litellm/proxy/client/teams.py b/litellm/proxy/client/teams.py index ef2ac53f9c4..105060e5ca9 100644 --- a/litellm/proxy/client/teams.py +++ b/litellm/proxy/client/teams.py @@ -11,16 +11,18 @@ from .exceptions import UnauthorizedError class TeamsManagementClient: """Client for managing teams in LiteLLM proxy.""" - def __init__(self, base_url: str, api_key: str | None = None): + def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): """ Initialize the TeamsManagementClient. Args: base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:4000") api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token. + timeout (int): Request timeout in seconds (default: 30) """ self._base_url = base_url.rstrip("/") # Remove trailing slash if present self._api_key = api_key + self._timeout = timeout def _get_headers(self) -> dict[str, str]: """ @@ -60,7 +62,7 @@ class TeamsManagementClient: if organization_id: params["organization_id"] = organization_id - response: Final = requests.get(url, headers=self._get_headers(), params=params) + response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self._timeout) if response.status_code == 401: raise UnauthorizedError("Authentication failed. Check your API key.") @@ -117,7 +119,7 @@ class TeamsManagementClient: if sort_by: params["sort_by"] = sort_by - response: Final = requests.get(url, headers=self._get_headers(), params=params) + response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self._timeout) if response.status_code == 401: raise UnauthorizedError("Authentication failed. Check your API key.") @@ -138,7 +140,7 @@ class TeamsManagementClient: """ url: Final = f"{self._base_url}/team/available" - response: Final = requests.get(url, headers=self._get_headers()) + response: Final = requests.get(url, headers=self._get_headers(), timeout=self._timeout) if response.status_code == 401: raise UnauthorizedError("Authentication failed. Check your API key.") diff --git a/litellm/proxy/client/users.py b/litellm/proxy/client/users.py index df5f9aad23e..3f11fe94043 100644 --- a/litellm/proxy/client/users.py +++ b/litellm/proxy/client/users.py @@ -6,9 +6,10 @@ from .exceptions import NotFoundError, UnauthorizedError class UsersManagementClient: - def __init__(self, base_url: str, api_key: str | None = None): + def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): self.base_url = base_url.rstrip("/") self.api_key = api_key + self.timeout = timeout def _get_headers(self) -> dict[str, str]: headers: Final = {"Content-Type": "application/json"} @@ -19,7 +20,7 @@ class UsersManagementClient: def list_users(self, params: dict[str, Any] | None = None) -> list[dict[str, Any]]: """List users (GET /user/list)""" url: Final = f"{self.base_url}/user/list" - response: Final = requests.get(url, headers=self._get_headers(), params=params) + response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self.timeout) if response.status_code == 401: raise UnauthorizedError(response.text) response.raise_for_status() @@ -29,7 +30,7 @@ class UsersManagementClient: """Get user info (GET /user/info)""" url: Final = f"{self.base_url}/user/info" params: Final = {"user_id": user_id} if user_id else {} - response: Final = requests.get(url, headers=self._get_headers(), params=params) + response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self.timeout) if response.status_code == 401: raise UnauthorizedError(response.text) if response.status_code == 404: @@ -41,7 +42,7 @@ class UsersManagementClient: """Get user info v2 - lightweight, returns only user object (GET /v2/user/info)""" url: Final = f"{self.base_url}/v2/user/info" params: Final = {"user_id": user_id} if user_id else {} - response: Final = requests.get(url, headers=self._get_headers(), params=params) + response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self.timeout) if response.status_code == 401: raise UnauthorizedError(response.text) if response.status_code == 404: @@ -52,7 +53,7 @@ class UsersManagementClient: def create_user(self, user_data: dict[str, Any]) -> dict[str, Any]: """Create a new user (POST /user/new)""" url: Final = f"{self.base_url}/user/new" - response: Final = requests.post(url, headers=self._get_headers(), json=user_data) + response: Final = requests.post(url, headers=self._get_headers(), json=user_data, timeout=self.timeout) if response.status_code == 401: raise UnauthorizedError(response.text) response.raise_for_status() @@ -61,7 +62,9 @@ class UsersManagementClient: def delete_user(self, user_ids: list[str]) -> dict[str, Any]: """Delete users (POST /user/delete)""" url: Final = f"{self.base_url}/user/delete" - response: Final = requests.post(url, headers=self._get_headers(), json={"user_ids": user_ids}) + response: Final = requests.post( + url, headers=self._get_headers(), json={"user_ids": user_ids}, timeout=self.timeout + ) if response.status_code == 401: raise UnauthorizedError(response.text) response.raise_for_status() diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py index 507dd645953..f15b8ec1e74 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py @@ -36,6 +36,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 @@ -117,10 +120,10 @@ def is_saas(host: str) -> bool: return False -def _get_jwt(auth_url, api_id, api_key) -> str: +def _get_jwt(auth_url, api_id, api_key, timeout: float = _AUTH_TIMEOUT_SECONDS) -> str: token_url: Final = f"{auth_url}/oauth2/token?grant_type=client_credentials" - resp: Final = requests.post(token_url, auth=HTTPBasicAuth(api_id, api_key)) + resp: Final = requests.post(token_url, auth=HTTPBasicAuth(api_id, api_key), timeout=timeout) if not resp.ok: raise RuntimeError( diff --git a/ruff.toml b/ruff.toml index 44bdf9d8125..3ac4c1fc94d 100644 --- a/ruff.toml +++ b/ruff.toml @@ -6,7 +6,8 @@ lint.extend-select = [ "T20", "PGH004", "RUF008", "RUF009", "RUF100", "B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790", "PIE800", "PLC0208", "PLR0402", "PLR1711", "PLR1730", "PLR2044", "PLW0133", "PYI030", "PYI041", "PYI064", "RET501", - "RUF010", "RUF022", "RUF023", "RUF051", "SIM114", "SIM118", "TC005", "UP006", "UP007", "UP008", + "RUF010", "RUF022", "RUF023", "RUF051", "S113", "SIM114", "SIM118", "TC005", "UP006", "UP007", + "UP008", "UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP037", "UP045", ] # RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip diff --git a/tests/test_litellm/proxy/client/conftest.py b/tests/test_litellm/proxy/client/conftest.py new file mode 100644 index 00000000000..c8b7951e284 --- /dev/null +++ b/tests/test_litellm/proxy/client/conftest.py @@ -0,0 +1,38 @@ +import threading + +import pytest + + +@pytest.fixture +def hanging_server(): + """A server that accepts the connection and never answers, so only a timeout ends the call.""" + from http.server import BaseHTTPRequestHandler, HTTPServer + from socketserver import ThreadingMixIn + + stop: threading.Event = threading.Event() + + class SilentRequestHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def _hang(self): + stop.wait(timeout=30) + + do_GET = _hang + do_POST = _hang + + def log_message(self, format, *args): + pass + + class ThreadedServer(ThreadingMixIn, HTTPServer): + daemon_threads = True + + server = ThreadedServer(("127.0.0.1", 0), SilentRequestHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}" + finally: + stop.set() + server.shutdown() + server.server_close() + thread.join(timeout=5) diff --git a/tests/test_litellm/proxy/client/test_teams.py b/tests/test_litellm/proxy/client/test_teams.py new file mode 100644 index 00000000000..b61091ca44b --- /dev/null +++ b/tests/test_litellm/proxy/client/test_teams.py @@ -0,0 +1,20 @@ +import time + +import pytest +import requests + +from litellm.proxy.client.teams import TeamsManagementClient + + +def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server): + """ + A proxy that accepts the connection but never answers used to pin the caller's + process forever, since the request carried no timeout at all. + """ + client = TeamsManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1) + + started = time.monotonic() + with pytest.raises(requests.exceptions.Timeout): + client.list() + + assert time.monotonic() - started < 10 diff --git a/tests/test_litellm/proxy/client/test_users.py b/tests/test_litellm/proxy/client/test_users.py index 87b8392e402..5b4d89420ab 100644 --- a/tests/test_litellm/proxy/client/test_users.py +++ b/tests/test_litellm/proxy/client/test_users.py @@ -1,6 +1,8 @@ +import time from unittest.mock import MagicMock, patch import pytest +import requests @@ -82,3 +84,17 @@ def test_delete_user_unauthorized(mock_post, client): mock_post.return_value.text = "unauthorized" with pytest.raises(UnauthorizedError): client.delete_user(["u1"]) + + +def test_delete_user_gives_up_at_the_timeout_instead_of_hanging(hanging_server): + """ + A proxy that accepts the connection but never answers used to pin the caller's + process forever, since the request carried no timeout at all. + """ + client = UsersManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1) + + started = time.monotonic() + with pytest.raises(requests.exceptions.Timeout): + client.delete_user(["u1"]) + + assert time.monotonic() - started < 10 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py index 1b2108c837d..b140082a3bf 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py @@ -1,4 +1,6 @@ import os +import threading +import time import uuid from typing import List, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -6,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException from httpx import Request, Response +import requests import litellm @@ -14,6 +17,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.proxy.guardrails.guardrail_hooks.hiddenlayer.hiddenlayer import ( HiddenlayerGuardrail, HiddenlayerGuardrailV2, + _get_jwt, ) from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.types.utils import ( @@ -1088,3 +1092,47 @@ class TestHiddenlayerGuardrailV2: config_model = HiddenlayerGuardrailV2.get_config_model() assert config_model is not None assert config_model.__name__ == "HiddenlayerGuardrailConfigModel" + + +@pytest.fixture +def hanging_auth_server(): + """A server that accepts the connection and never answers, so only a timeout ends the call.""" + from http.server import BaseHTTPRequestHandler, HTTPServer + from socketserver import ThreadingMixIn + + stop: threading.Event = threading.Event() + + class SilentRequestHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self): + stop.wait(timeout=30) + + def log_message(self, format, *args): + pass + + class ThreadedServer(ThreadingMixIn, HTTPServer): + daemon_threads = True + + server = ThreadedServer(("127.0.0.1", 0), SilentRequestHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}" + finally: + stop.set() + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +def test_get_jwt_gives_up_at_the_timeout_instead_of_blocking_the_event_loop(hanging_auth_server): + """ + `_get_jwt` runs synchronously inside `_call_hiddenlayer`, so an auth host that + accepts and never answers used to park the whole worker's event loop. + """ + started = time.monotonic() + with pytest.raises(requests.exceptions.Timeout): + _get_jwt(auth_url=hanging_auth_server, api_id="id", api_key="secret", timeout=1) + + assert time.monotonic() - started < 10 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 627811a7f1d..6dcabe076c9 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22805 + "limit": 22804 }, "LIT002": { "limit": 26873 @@ -33,6 +33,6 @@ "limit": 5588 }, "LIT012": { - "limit": 4510 + "limit": 4508 } } From cbc931da5465e1ea08fa9d5cf702c2cc38f1ab35 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 25 Aug 2026 11:00:29 -0700 Subject: [PATCH 017/111] test: assert the poll call shape after the timeout refactor --- tests/test_litellm/proxy/auth/test_cli_auth.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_cli_auth.py b/tests/test_litellm/proxy/auth/test_cli_auth.py index c9b31a1d776..5cde5522376 100644 --- a/tests/test_litellm/proxy/auth/test_cli_auth.py +++ b/tests/test_litellm/proxy/auth/test_cli_auth.py @@ -82,7 +82,7 @@ async def test_poll_for_ready_404(sleep_mock, request_mock): _poll_for_ready_data( "https://litellm.com", poll_interval=1, total_timeout=1, request_timeout=42 ) - request_mock.assert_called_once_with("https://litellm.com", timeout=42) + request_mock.assert_called_once_with("https://litellm.com", headers=None, timeout=42) @pytest.mark.asyncio @@ -103,7 +103,7 @@ async def test_poll_for_ready_200_ready(sleep_mock, click_mock, request_mock): ) assert actual == {"status": "ready", "json": "data"} click_mock.assert_not_called() - request_mock.assert_called_once_with("https://litellm.com", timeout=42) + request_mock.assert_called_once_with("https://litellm.com", headers=None, timeout=42) sleep_mock.assert_not_called() @@ -131,8 +131,8 @@ async def test_poll_for_ready_single_pending(sleep_mock, click_mock, request_moc click_mock.assert_not_called() request_mock.assert_has_calls( [ - call("https://litellm.com", timeout=42), - call("https://litellm.com", timeout=42), + call("https://litellm.com", headers=None, timeout=42), + call("https://litellm.com", headers=None, timeout=42), ] ) sleep_mock.assert_called_once_with(1) @@ -168,8 +168,8 @@ async def test_poll_for_ready_pending(sleep_mock, click_mock, request_mock): click_mock.assert_has_calls([call("Pending message"), call("Pending message")]) request_mock.assert_has_calls( [ - call("https://litellm.com", timeout=42), - call("https://litellm.com", timeout=42), + call("https://litellm.com", headers=None, timeout=42), + call("https://litellm.com", headers=None, timeout=42), ] ) sleep_mock.assert_has_calls([call(1), call(1)]) @@ -194,7 +194,7 @@ async def test_poll_for_ready_connection_failure(sleep_mock, click_mock, request click_mock.assert_called_once_with("Connection error (will retry): ERROR") request_mock.assert_has_calls( [ - call("https://litellm.com", timeout=42), + call("https://litellm.com", headers=None, timeout=42), ] ) sleep_mock.assert_has_calls([call(1), call(1)]) From 6f24d85490f03c8af4f6b94e93aff0ffa2240c1b Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 25 Aug 2026 18:26:57 +0000 Subject: [PATCH 018/111] fix(proxy): grant users with empty models list direct access in model listing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 4 +- .../test_team_model_name_translation.py | 55 +++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0abcdeaf3f6..3d3169eeab9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12382,8 +12382,10 @@ def get_direct_access_models( The 'all-proxy-models' sentinel grants direct access to every non-team deployment, mirroring how get_key_models expands it for the key/team path. + An empty models list means unrestricted access at call time (see + can_user_call_model), so it resolves the same way. """ - if SpecialModelNames.all_proxy_models.value in user_db_object.models: + if not user_db_object.models or SpecialModelNames.all_proxy_models.value in user_db_object.models: return llm_router.get_model_ids(exclude_team_models=True) return [ diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 2f4018b55ab..cb1521f10ae 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -1507,6 +1507,61 @@ def test_get_direct_access_models_resolves_explicit_model_names(): router.get_model_list.assert_called_once_with(model_name="gpt-4o") +def test_get_direct_access_models_empty_models_grants_all_non_team_models(): + """An empty user.models list means unrestricted access at call time + (can_user_call_model), so the listing must resolve it like 'all-proxy-models' + instead of returning nothing. Regression for a user with models=[] and no + teams seeing an empty Models+Endpoints page.""" + router = MagicMock() + router.get_model_ids.return_value = ["global-id-1", "global-id-2"] + + user = LiteLLM_UserTable(user_id="u", models=[], teams=[]) + + result = ps.get_direct_access_models(user_db_object=user, llm_router=router) + + assert result == ["global-id-1", "global-id-2"] + router.get_model_ids.assert_called_once_with(exclude_team_models=True) + router.get_model_list.assert_not_called() + + +@pytest.mark.asyncio +async def test_populate_team_access_grants_empty_models_user_direct_access(monkeypatch): + """An internal user with models=[] and no teams can call every non-team model, + so the Models+Endpoints page must list them instead of rendering empty.""" + global_row = { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + + router = MagicMock() + router.get_model_ids.return_value = ["global-id-1"] + + user_row = LiteLLM_UserTable( + user_id="u", + user_role=LitellmUserRoles.INTERNAL_USER.value, + models=[], + teams=[], + ) + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) + + monkeypatch.setattr(ps, "get_all_team_models", AsyncMock(return_value={})) + + caller = UserAPIKeyAuth(user_id="u", user_role=LitellmUserRoles.INTERNAL_USER, team_models=[]) + + populated = await ps._populate_team_access_on_models( + user_api_key_dict=caller, + prisma_client=prisma_client, + llm_router=router, + all_models=[global_row], + ) + visible = ps._filter_models_to_user_accessible(populated) + + assert [m["model_info"]["id"] for m in visible] == ["global-id-1"] + assert visible[0]["model_info"]["direct_access"] is True + + @pytest.mark.asyncio async def test_populate_team_access_grants_all_proxy_models_user_direct_access( monkeypatch, From bd2b93c6889614f04986fcd9d94067e061a20567 Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 25 Aug 2026 20:19:47 +0000 Subject: [PATCH 019/111] chore(proxy): drop redundant docstring note in get_direct_access_models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3d3169eeab9..f536f521d57 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12382,8 +12382,6 @@ def get_direct_access_models( The 'all-proxy-models' sentinel grants direct access to every non-team deployment, mirroring how get_key_models expands it for the key/team path. - An empty models list means unrestricted access at call time (see - can_user_call_model), so it resolves the same way. """ if not user_db_object.models or SpecialModelNames.all_proxy_models.value in user_db_object.models: return llm_router.get_model_ids(exclude_team_models=True) From f81408535447db7cf6be81caa48c1447965f2288 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 27 Aug 2026 01:34:33 +0000 Subject: [PATCH 020/111] feat(mcp): bulk-import Anthropic MCP connectors via API and admin UI Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 1147 ++++++++++++++++- .../mcp_connector_import.py | 168 +++ .../mcp_management_endpoints.py | 117 ++ .../test_mcp_connector_import.py | 157 +++ .../test_mcp_management_endpoints.py | 140 ++ .../_components/ImportMCPServers.tsx | 134 ++ .../_components/importConnectorConfig.test.ts | 74 ++ .../_components/importConnectorConfig.ts | 64 + .../mcp-servers/_components/mcp_servers.tsx | 19 +- .../src/components/networking.tsx | 9 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 486 ++++++- 11 files changed, 2506 insertions(+), 9 deletions(-) create mode 100644 litellm/proxy/management_endpoints/mcp_connector_import.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ImportMCPServers.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/importConnectorConfig.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/importConnectorConfig.ts diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 0e3cbbfd560..5e0b5ba3faf 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -14727,6 +14727,17 @@ "title": "Args", "type": "array" }, + "audience": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Audience" + }, "auth_type": { "anyOf": [ { @@ -14738,7 +14749,11 @@ "authorization", "oauth2", "aws_sigv4", - "token" + "token", + "oauth2_token_exchange", + "oauth2_id_jag", + "true_passthrough", + "oauth_delegate" ], "type": "string" }, @@ -14793,6 +14808,17 @@ ], "title": "Command" }, + "connected_app_reachable": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Connected App Reachable" + }, "created_at": { "anyOf": [ { @@ -14826,6 +14852,22 @@ } ] }, + "dcr_bridge": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Dcr Bridge" + }, + "delegate_auth_to_upstream": { + "default": false, + "title": "Delegate Auth To Upstream", + "type": "boolean" + }, "description": { "anyOf": [ { @@ -14844,6 +14886,20 @@ "title": "Env", "type": "object" }, + "env_vars": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/MCPEnvVar" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Env Vars" + }, "extra_headers": { "items": { "type": "string" @@ -14889,6 +14945,17 @@ "title": "Is Byok", "type": "boolean" }, + "issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Issuer" + }, "last_health_check": { "anyOf": [ { @@ -14901,6 +14968,17 @@ ], "title": "Last Health Check" }, + "max_concurrent_requests": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Max Concurrent Requests" + }, "mcp_access_groups": { "items": { "type": "string" @@ -14920,6 +14998,26 @@ ], "title": "Mcp Info" }, + "oauth2_flow": { + "anyOf": [ + { + "enum": [ + "client_credentials", + "authorization_code" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Oauth2 Flow" + }, + "oauth_passthrough": { + "default": false, + "title": "Oauth Passthrough", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -15023,6 +15121,17 @@ "description": "Health status: 'healthy', 'unhealthy', 'unknown'", "title": "Status" }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" + }, "submitted_at": { "anyOf": [ { @@ -15063,6 +15172,39 @@ "title": "Teams", "type": "array" }, + "timeout": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Timeout" + }, + "token_exchange_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Endpoint" + }, + "token_exchange_profile": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Profile" + }, "token_url": { "anyOf": [ { @@ -15153,8 +15295,236 @@ "title": "LiteLLM_MCPServerTable", "type": "object" }, + "MCPConnectorEntry": { + "properties": { + "args": { + "items": { + "type": "string" + }, + "title": "Args", + "type": "array" + }, + "authorization_token": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Authorization Token" + }, + "command": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Command" + }, + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "env": { + "additionalProperties": { + "type": "string" + }, + "title": "Env", + "type": "object" + }, + "headers": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Headers" + }, + "name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Name" + }, + "type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Type" + }, + "url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Url" + } + }, + "title": "MCPConnectorEntry", + "type": "object" + }, + "MCPConnectorImportFailure": { + "properties": { + "error": { + "title": "Error", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "name", + "error" + ], + "title": "MCPConnectorImportFailure", + "type": "object" + }, + "MCPConnectorImportRequest": { + "properties": { + "mcp_servers": { + "anyOf": [ + { + "additionalProperties": { + "$ref": "#/components/schemas/MCPConnectorEntry" + }, + "type": "object" + }, + { + "items": { + "$ref": "#/components/schemas/MCPConnectorEntry" + }, + "type": "array" + } + ], + "title": "Mcp Servers" + } + }, + "required": [ + "mcp_servers" + ], + "title": "MCPConnectorImportRequest", + "type": "object" + }, + "MCPConnectorImportResponse": { + "properties": { + "errors": { + "items": { + "$ref": "#/components/schemas/MCPConnectorImportFailure" + }, + "title": "Errors", + "type": "array" + }, + "imported": { + "items": { + "$ref": "#/components/schemas/MCPConnectorImportResult" + }, + "title": "Imported", + "type": "array" + }, + "skipped": { + "items": { + "$ref": "#/components/schemas/MCPConnectorImportSkipped" + }, + "title": "Skipped", + "type": "array" + } + }, + "required": [ + "imported", + "skipped", + "errors" + ], + "title": "MCPConnectorImportResponse", + "type": "object" + }, + "MCPConnectorImportResult": { + "properties": { + "alias": { + "title": "Alias", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "server_id": { + "title": "Server Id", + "type": "string" + } + }, + "required": [ + "name", + "server_id", + "alias" + ], + "title": "MCPConnectorImportResult", + "type": "object" + }, + "MCPConnectorImportSkipped": { + "properties": { + "name": { + "title": "Name", + "type": "string" + }, + "reason": { + "title": "Reason", + "type": "string" + } + }, + "required": [ + "name", + "reason" + ], + "title": "MCPConnectorImportSkipped", + "type": "object" + }, "MCPCredentials": { "properties": { + "audience": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Audience" + }, "auth_value": { "anyOf": [ { @@ -15243,6 +15613,17 @@ ], "title": "Aws Session Token" }, + "client_assertion_signing_alg": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Assertion Signing Alg" + }, "client_id": { "anyOf": [ { @@ -15254,6 +15635,28 @@ ], "title": "Client Id" }, + "client_private_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Private Key" + }, + "client_private_key_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Private Key Id" + }, "client_secret": { "anyOf": [ { @@ -15265,6 +15668,42 @@ ], "title": "Client Secret" }, + "id_jag_resource": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Id Jag Resource" + }, + "id_jag_resource_token_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Id Jag Resource Token Endpoint" + }, + "redirect_uris": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Redirect Uris" + }, "scopes": { "anyOf": [ { @@ -15278,11 +15717,113 @@ } ], "title": "Scopes" + }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" + }, + "token_endpoint_auth_method": { + "anyOf": [ + { + "enum": [ + "client_secret_basic", + "client_secret_post" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Endpoint Auth Method" + }, + "token_exchange_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Endpoint" + }, + "token_exchange_profile": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Profile" + }, + "upstream_resource": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Upstream Resource" } }, "title": "MCPCredentials", "type": "object" }, + "MCPEnvVar": { + "description": "One environment variable for an MCP server.\n\nVariables can be interpolated into ``static_headers`` using ``${NAME}``\nsyntax. ``scope=global`` values are stored on the server. ``scope=user``\nvalues are stored per-user in ``LiteLLM_MCPUserEnvVars`` and supplied by\neach user.", + "properties": { + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "name": { + "title": "Name", + "type": "string" + }, + "scope": { + "$ref": "#/components/schemas/MCPEnvVarScope", + "default": "global" + }, + "value": { + "default": "", + "title": "Value", + "type": "string" + } + }, + "required": [ + "name" + ], + "title": "MCPEnvVar", + "type": "object" + }, + "MCPEnvVarScope": { + "description": "Scope for an MCP server environment variable.\n\n- ``global``: value is provided by the admin and used for all users.\n- ``user``: each user must provide their own value via the per-user\n env-var endpoint. The admin-supplied ``value`` is treated as a\n placeholder/hint and is not used at request time.", + "enum": [ + "global", + "user" + ], + "title": "MCPEnvVarScope", + "type": "string" + }, "MCPOAuthUserCredentialRequest": { "description": "Stores a user's OAuth2 token for an OpenAPI MCP server.", "properties": { @@ -15537,6 +16078,112 @@ "title": "MCPUserCredentialResponse", "type": "object" }, + "MCPUserEnvVarSpec": { + "description": "Describes one per-user env var slot for the calling user.\n\nStored values are write-only: the status only reports whether a value\n``is_set`` and never echoes the decrypted secret back to the client.", + "properties": { + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "is_set": { + "default": false, + "title": "Is Set", + "type": "boolean" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "name" + ], + "title": "MCPUserEnvVarSpec", + "type": "object" + }, + "MCPUserEnvVarsRequest": { + "description": "Payload for storing the calling user's per-user env var values.", + "properties": { + "values": { + "additionalProperties": { + "type": "string" + }, + "title": "Values", + "type": "object" + } + }, + "required": [ + "values" + ], + "title": "MCPUserEnvVarsRequest", + "type": "object" + }, + "MCPUserEnvVarsStatus": { + "description": "Per-user env var status for a single MCP server.", + "properties": { + "alias": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Alias" + }, + "missing_count": { + "default": 0, + "title": "Missing Count", + "type": "integer" + }, + "required": { + "items": { + "$ref": "#/components/schemas/MCPUserEnvVarSpec" + }, + "title": "Required", + "type": "array" + }, + "server_id": { + "title": "Server Id", + "type": "string" + }, + "server_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Server Name" + }, + "setup_url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Setup Url" + } + }, + "required": [ + "server_id" + ], + "title": "MCPUserEnvVarsStatus", + "type": "object" + }, "MakeMCPServersPublicRequest": { "properties": { "mcp_server_ids": { @@ -15604,6 +16251,17 @@ "title": "Args", "type": "array" }, + "audience": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Audience" + }, "auth_type": { "anyOf": [ { @@ -15615,7 +16273,11 @@ "authorization", "oauth2", "aws_sigv4", - "token" + "token", + "oauth2_token_exchange", + "oauth2_id_jag", + "true_passthrough", + "oauth_delegate" ], "type": "string" }, @@ -15680,6 +16342,22 @@ } ] }, + "dcr_bridge": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Dcr Bridge" + }, + "delegate_auth_to_upstream": { + "default": false, + "title": "Delegate Auth To Upstream", + "type": "boolean" + }, "description": { "anyOf": [ { @@ -15698,6 +16376,20 @@ "title": "Env", "type": "object" }, + "env_vars": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/MCPEnvVar" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Env Vars" + }, "extra_headers": { "anyOf": [ { @@ -15728,6 +16420,28 @@ "title": "Is Byok", "type": "boolean" }, + "issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Issuer" + }, + "max_concurrent_requests": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Max Concurrent Requests" + }, "mcp_access_groups": { "items": { "type": "string" @@ -15762,6 +16476,11 @@ ], "title": "Oauth2 Flow" }, + "oauth_passthrough": { + "default": false, + "title": "Oauth Passthrough", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -15831,6 +16550,17 @@ ], "title": "Static Headers" }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" + }, "submitted_at": { "anyOf": [ { @@ -15856,6 +16586,39 @@ "description": "Server-managed: set by the endpoint; caller values are overridden.", "title": "Submitted By" }, + "timeout": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Timeout" + }, + "token_exchange_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Endpoint" + }, + "token_exchange_profile": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Profile" + }, "token_url": { "anyOf": [ { @@ -16008,6 +16771,17 @@ "title": "Args", "type": "array" }, + "audience": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Audience" + }, "auth_type": { "anyOf": [ { @@ -16019,7 +16793,11 @@ "authorization", "oauth2", "aws_sigv4", - "token" + "token", + "oauth2_token_exchange", + "oauth2_id_jag", + "true_passthrough", + "oauth_delegate" ], "type": "string" }, @@ -16084,6 +16862,22 @@ } ] }, + "dcr_bridge": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Dcr Bridge" + }, + "delegate_auth_to_upstream": { + "default": false, + "title": "Delegate Auth To Upstream", + "type": "boolean" + }, "description": { "anyOf": [ { @@ -16102,6 +16896,20 @@ "title": "Env", "type": "object" }, + "env_vars": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/MCPEnvVar" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Env Vars" + }, "extra_headers": { "anyOf": [ { @@ -16132,6 +16940,28 @@ "title": "Is Byok", "type": "boolean" }, + "issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Issuer" + }, + "max_concurrent_requests": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Max Concurrent Requests" + }, "mcp_access_groups": { "items": { "type": "string" @@ -16151,6 +16981,26 @@ ], "title": "Mcp Info" }, + "oauth2_flow": { + "anyOf": [ + { + "enum": [ + "client_credentials", + "authorization_code" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Oauth2 Flow" + }, + "oauth_passthrough": { + "default": false, + "title": "Oauth Passthrough", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -16213,6 +17063,50 @@ ], "title": "Static Headers" }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" + }, + "timeout": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Timeout" + }, + "token_exchange_endpoint": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Endpoint" + }, + "token_exchange_profile": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token Exchange Profile" + }, "token_url": { "anyOf": [ { @@ -16331,6 +17225,13 @@ }, "ValidationError": { "properties": { + "ctx": { + "title": "Context", + "type": "object" + }, + "input": { + "title": "Input" + }, "loc": { "items": { "anyOf": [ @@ -16600,6 +17501,18 @@ "description": "Filter MCP servers by team scope. When provided, returns only servers the team has access to plus globally available (allow_all_keys) servers. Used by the Create Key UI to show team-scoped MCP servers.", "title": "Team Id" } + }, + { + "description": "Annotate each returned server with connected_app_reachable: whether a connected app authorized by the calling user (a gateway OAuth session) is served this server on the aggregate MCP endpoint.", + "in": "query", + "name": "connected_app_view", + "required": false, + "schema": { + "default": false, + "description": "Annotate each returned server with connected_app_reachable: whether a connected app authorized by the calling user (a gateway OAuth session) is served this server on the aggregate MCP endpoint.", + "title": "Connected App View", + "type": "boolean" + } } ], "responses": { @@ -16827,6 +17740,53 @@ ] } }, + "/v1/mcp/server/import": { + "post": { + "description": "Bulk-import MCP connectors from Anthropic mcpServers or mcp_servers JSON", + "operationId": "import_mcp_servers_v1_mcp_server_import_post", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPConnectorImportRequest" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPConnectorImportResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Import Mcp Servers", + "tags": [ + "mcp_management" + ] + } + }, "/v1/mcp/server/oauth/session": { "post": { "description": "Temporarily cache an MCP server in memory without writing to the database", @@ -17438,6 +18398,156 @@ ] } }, + "/v1/mcp/server/{server_id}/user-env-vars": { + "delete": { + "description": "Clear the calling user's per-user MCP env var values for this server.", + "operationId": "clear_mcp_user_env_vars_v1_mcp_server__server_id__user_env_vars_delete", + "parameters": [ + { + "in": "path", + "name": "server_id", + "required": true, + "schema": { + "title": "Server Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPUserEnvVarsStatus" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Clear Mcp User Env Vars", + "tags": [ + "mcp_management" + ] + }, + "get": { + "description": "Return the calling user's per-user MCP env var status for this server.", + "operationId": "get_mcp_user_env_vars_v1_mcp_server__server_id__user_env_vars_get", + "parameters": [ + { + "in": "path", + "name": "server_id", + "required": true, + "schema": { + "title": "Server Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPUserEnvVarsStatus" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Mcp User Env Vars", + "tags": [ + "mcp_management" + ] + }, + "post": { + "description": "Store the calling user's per-user MCP env var values for this server. Submitted values are merged over any previously stored values, so you only send the fields you want to set or change; a variable omitted (or sent empty) keeps its stored value. Use DELETE to clear all stored values.", + "operationId": "store_mcp_user_env_vars_v1_mcp_server__server_id__user_env_vars_post", + "parameters": [ + { + "in": "path", + "name": "server_id", + "required": true, + "schema": { + "title": "Server Id", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPUserEnvVarsRequest" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPUserEnvVarsStatus" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Store Mcp User Env Vars", + "tags": [ + "mcp_management" + ] + } + }, "/v1/mcp/tools": { "get": { "description": "Get all MCP tools available for the current key, including those from access groups", @@ -17746,6 +18856,37 @@ "mcp_management" ] } + }, + "/v1/mcp/user-env-vars/status": { + "get": { + "description": "Per-user MCP env var status across every server the user can access. Used by the dashboard to highlight servers with missing per-user vars.", + "operationId": "list_mcp_user_env_var_status_v1_mcp_user_env_vars_status_get", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "items": { + "$ref": "#/components/schemas/MCPUserEnvVarsStatus" + }, + "title": "Response List Mcp User Env Var Status V1 Mcp User Env Vars Status Get", + "type": "array" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "List Mcp User Env Var Status", + "tags": [ + "mcp_management" + ] + } } } }, diff --git a/litellm/proxy/management_endpoints/mcp_connector_import.py b/litellm/proxy/management_endpoints/mcp_connector_import.py new file mode 100644 index 00000000000..d576131ae1b --- /dev/null +++ b/litellm/proxy/management_endpoints/mcp_connector_import.py @@ -0,0 +1,168 @@ +""" +Convert Anthropic MCP connector definitions into LiteLLM MCP server create requests. + +Two interchange shapes are accepted: +- the ``mcpServers`` mapping used by Claude Desktop / Claude Code config files +- the ``mcp_servers`` array used by the Anthropic Messages API MCP connector +""" + +import re +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +from pydantic import AliasChoices, BaseModel, ConfigDict, Field, ValidationError + +from litellm.proxy._types import MCPApprovalStatus, NewMCPServerRequest +from litellm.types.mcp import MCPAuth, MCPCredentials, MCPTransport + + +class MCPConnectorEntry(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + name: str | None = None + type: str | None = None + url: str | None = None + authorization_token: str | None = Field( + default=None, validation_alias=AliasChoices("authorization_token", "authorizationToken") + ) + headers: Mapping[str, str] | None = None + command: str | None = None + args: tuple[str, ...] = Field(default_factory=tuple) + env: Mapping[str, str] = Field(default_factory=dict) + description: str | None = None + + +class MCPConnectorImportRequest(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + mcp_servers: Mapping[str, MCPConnectorEntry] | tuple[MCPConnectorEntry, ...] = Field( + validation_alias=AliasChoices("mcp_servers", "mcpServers") + ) + + +@dataclass(frozen=True, slots=True) +class ConvertedConnector: + name: str + request: NewMCPServerRequest + + +@dataclass(frozen=True, slots=True) +class ConnectorConversionError: + name: str + error: str + + +class MCPConnectorImportResult(BaseModel): + name: str + server_id: str + alias: str + + +class MCPConnectorImportSkipped(BaseModel): + name: str + reason: str + + +class MCPConnectorImportFailure(BaseModel): + name: str + error: str + + +class MCPConnectorImportResponse(BaseModel): + imported: tuple[MCPConnectorImportResult, ...] + skipped: tuple[MCPConnectorImportSkipped, ...] + errors: tuple[MCPConnectorImportFailure, ...] + + +_INVALID_SERVER_NAME_CHARS: Final = re.compile(r"[^A-Za-z0-9_]") + + +def sanitize_connector_name(name: str) -> str: + sanitized: Final = re.sub(r"_+", "_", _INVALID_SERVER_NAME_CHARS.sub("_", name.strip())).strip("_") + return sanitized + + +_SSE_TYPES: Final = frozenset({"sse"}) +_URL_TYPES: Final = frozenset({"url", "http", "streamable_http", "streamable-http", "sse", ""}) + + +def _convert_entry(name: str, entry: MCPConnectorEntry) -> ConvertedConnector | ConnectorConversionError: + sanitized_name: Final = sanitize_connector_name(name) + if not sanitized_name: + return ConnectorConversionError(name=name, error="Connector name is empty after sanitization.") + + if entry.url and entry.command: + return ConnectorConversionError(name=name, error="Connector cannot have both a url and a command.") + + if entry.command: + try: + stdio_request: Final = NewMCPServerRequest( + server_name=sanitized_name, + alias=sanitized_name, + description=entry.description, + approval_status=MCPApprovalStatus.active, + transport=MCPTransport.stdio, + command=entry.command, + args=entry.args, + env=entry.env, + ) + except ValidationError as e: + return ConnectorConversionError(name=name, error=_first_validation_message(e)) + return ConvertedConnector(name=name, request=stdio_request) + + if not entry.url: + return ConnectorConversionError(name=name, error="Connector must have either a url or a command.") + + entry_type: Final = (entry.type or "").lower() + if entry_type not in _URL_TYPES: + return ConnectorConversionError(name=name, error=f"Unsupported connector type '{entry.type}'.") + + transport: Final = MCPTransport.sse if entry_type in _SSE_TYPES else MCPTransport.http + credentials: Final = _bearer_credentials(entry.authorization_token) + try: + remote_request: Final = NewMCPServerRequest( + server_name=sanitized_name, + alias=sanitized_name, + description=entry.description, + approval_status=MCPApprovalStatus.active, + transport=transport, + url=entry.url, + auth_type=MCPAuth.bearer_token if entry.authorization_token else MCPAuth.none, + credentials=credentials, + static_headers=entry.headers, + ) + except ValidationError as e: + return ConnectorConversionError(name=name, error=_first_validation_message(e)) + return ConvertedConnector(name=name, request=remote_request) + + +def _bearer_credentials(token: str | None) -> MCPCredentials | None: + if not token: + return None + credentials: Final[MCPCredentials] = {"auth_value": token} + return credentials + + +def _first_validation_message(error: ValidationError) -> str: + messages: Final = tuple(str(detail.get("msg", "")) for detail in error.errors()) + return messages[0] if messages else str(error) + + +def convert_connector_entries( + payload: MCPConnectorImportRequest, +) -> tuple[ConvertedConnector | ConnectorConversionError, ...]: + servers: Final = payload.mcp_servers + if isinstance(servers, Mapping): + return tuple(_convert_entry(name, entry) for name, entry in servers.items()) + return tuple( + _convert_entry(entry.name or "", entry) if entry.name else _named_entry_error(index, entry) + for index, entry in enumerate(servers) + ) + + +def _named_entry_error(index: int, entry: MCPConnectorEntry) -> ConnectorConversionError: + return ConnectorConversionError( + name=entry.url or f"entry {index}", + error="Connector entries in list form must have a name.", + ) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 54a591a5e1a..6583d5405f8 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -133,6 +133,7 @@ if MCP_AVAILABLE: delete_mcp_server, delete_user_credential, delete_user_env_vars, + get_all_mcp_servers, get_all_mcp_servers_for_user, get_draft_mcp_server, get_mcp_server, @@ -199,6 +200,16 @@ if MCP_AVAILABLE: populate_request_with_path_params, ) from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view + from litellm.proxy.management_endpoints.mcp_connector_import import ( + ConnectorConversionError, + ConvertedConnector, + MCPConnectorImportFailure, + MCPConnectorImportRequest, + MCPConnectorImportResponse, + MCPConnectorImportResult, + MCPConnectorImportSkipped, + convert_connector_entries, + ) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.types.mcp import ( MCP_ADMIN_CONFIG_CREDENTIAL_KEYS, @@ -1607,6 +1618,112 @@ if MCP_AVAILABLE: return _redact_mcp_credentials(new_mcp_server) + @router.post( + "/server/import", + description="Bulk-import MCP connectors from Anthropic mcpServers or mcp_servers JSON", + dependencies=(Depends(user_api_key_auth),), + response_model=MCPConnectorImportResponse, + status_code=status.HTTP_200_OK, + ) + @management_endpoint_wrapper + async def import_mcp_servers( + payload: MCPConnectorImportRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection + ): + """ + Bulk-import MCP connectors. Accepts the Claude Desktop / Claude Code + ``mcpServers`` mapping or the Anthropic Messages API ``mcp_servers`` + array, creates each entry as a LiteLLM MCP server, and returns + per-entry results so partial imports are visible to the caller. + """ + prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") + + if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict + "error": "User does not have permission to import mcp servers. You can only import mcp servers if you are a PROXY_ADMIN." + }, + ) + + conversions: Final = convert_connector_entries(payload) + existing_servers: Final = await get_all_mcp_servers(prisma_client) + existing_names: Final = frozenset( + name for server in existing_servers for name in (server.alias, server.server_name) if name + ) + + def _classify( + index: int, conversion: ConvertedConnector | ConnectorConversionError + ) -> ConvertedConnector | ConnectorConversionError | MCPConnectorImportSkipped: + if isinstance(conversion, ConnectorConversionError): + return conversion + alias: Final = conversion.request.alias or "" + if alias in existing_names: + return MCPConnectorImportSkipped( + name=conversion.name, reason=f"An MCP server named '{alias}' already exists." + ) + earlier_aliases: Final = frozenset( + earlier.request.alias or "" + for earlier in conversions[:index] + if isinstance(earlier, ConvertedConnector) + ) + if alias in earlier_aliases: + return MCPConnectorImportSkipped( + name=conversion.name, reason=f"Duplicate connector name '{alias}' in the import payload." + ) + return conversion + + async def _create( + conversion: ConvertedConnector, + ) -> MCPConnectorImportResult | MCPConnectorImportFailure: + try: + validate_and_normalize_mcp_server_payload(conversion.request) + except HTTPException as e: + error_text: Final = ( + str(e.detail.get("error", e.detail)) if isinstance(e.detail, dict) else str(e.detail) + ) + return MCPConnectorImportFailure(name=conversion.name, error=error_text) + try: + created: Final = await create_mcp_server( + prisma_client, + conversion.request, + touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, + ) + except Exception as e: # noqa: BLE001 # any create failure must become a per-entry error, not a 500 + verbose_proxy_logger.exception("Error importing mcp server %s: %s", conversion.name, e) + return MCPConnectorImportFailure(name=conversion.name, error=str(e)) + return MCPConnectorImportResult( + name=conversion.name, server_id=created.server_id, alias=created.alias or "" + ) + + classified: Final = tuple(_classify(index, conversion) for index, conversion in enumerate(conversions)) + outcomes: Final = tuple( + [ + await _create(entry) if isinstance(entry, ConvertedConnector) else entry for entry in classified + ] # mutable-ok: await is illegal in a generator expression here + ) + + imported: Final = tuple(entry for entry in outcomes if isinstance(entry, MCPConnectorImportResult)) + if imported: + # Best-effort registry refresh, mirroring add_mcp_server: rows are + # already committed, so a refresh failure must not surface as a 500. + try: + await global_mcp_server_manager.reload_servers_from_database() + except Exception as e: # noqa: BLE001 # rows are committed; a refresh failure must not surface as a 500 + verbose_proxy_logger.exception("MCP connector import committed but registry refresh failed: %s", e) + + return MCPConnectorImportResponse( + imported=imported, + skipped=tuple(entry for entry in outcomes if isinstance(entry, MCPConnectorImportSkipped)), + errors=tuple( + MCPConnectorImportFailure(name=entry.name, error=entry.error) + if isinstance(entry, ConnectorConversionError) + else entry + for entry in outcomes + if isinstance(entry, (ConnectorConversionError, MCPConnectorImportFailure)) + ), + ) + @router.post( "/server/oauth/session", description="Temporarily cache an MCP server in memory without writing to the database", diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py new file mode 100644 index 00000000000..9f2874bcbf3 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py @@ -0,0 +1,157 @@ +import pytest + +from litellm.proxy.management_endpoints.mcp_connector_import import ( + ConnectorConversionError, + ConvertedConnector, + MCPConnectorImportRequest, + convert_connector_entries, + sanitize_connector_name, +) +from litellm.types.mcp import MCPAuth, MCPTransport + + +def _single(payload: dict) -> ConvertedConnector | ConnectorConversionError: + results = convert_connector_entries(MCPConnectorImportRequest.model_validate(payload)) + assert len(results) == 1 + return results[0] + + +class TestSanitizeConnectorName: + @pytest.mark.parametrize( + "raw,expected", + [ + ("my-server", "my_server"), + (" spaced name ", "spaced_name"), + ("already_ok", "already_ok"), + ("a.b.c", "a_b_c"), + ("---", ""), + ], + ) + def test_sanitizes_to_mcp_safe_names(self, raw, expected): + assert sanitize_connector_name(raw) == expected + + +class TestConvertMcpServersMapping: + def test_url_connector_with_authorization_token(self): + result = _single( + { + "mcpServers": { + "github-mcp": { + "url": "https://api.example.com/mcp", + "authorization_token": "secret-token", + "headers": {"X-Env": "prod"}, + "description": "GitHub connector", + } + } + } + ) + assert isinstance(result, ConvertedConnector) + assert result.request.server_name == "github_mcp" + assert result.request.alias == "github_mcp" + assert result.request.transport == MCPTransport.http + assert result.request.url == "https://api.example.com/mcp" + assert result.request.auth_type == MCPAuth.bearer_token + assert result.request.credentials == {"auth_value": "secret-token"} + assert result.request.static_headers == {"X-Env": "prod"} + assert result.request.description == "GitHub connector" + + def test_camel_case_authorization_token_alias(self): + result = _single( + {"mcpServers": {"srv": {"url": "https://x.example/mcp", "authorizationToken": "tok"}}} + ) + assert isinstance(result, ConvertedConnector) + assert result.request.credentials == {"auth_value": "tok"} + + def test_url_connector_without_token_uses_no_auth(self): + result = _single({"mcpServers": {"open": {"url": "https://open.example/mcp"}}}) + assert isinstance(result, ConvertedConnector) + assert result.request.auth_type == MCPAuth.none + assert result.request.credentials is None + + def test_sse_type_maps_to_sse_transport(self): + result = _single({"mcpServers": {"legacy": {"type": "sse", "url": "https://sse.example/mcp"}}}) + assert isinstance(result, ConvertedConnector) + assert result.request.transport == MCPTransport.sse + + def test_stdio_connector(self): + result = _single( + { + "mcpServers": { + "local": { + "command": "npx", + "args": ["-y", "@example/mcp-server"], + "env": {"API_KEY": "value"}, + } + } + } + ) + assert isinstance(result, ConvertedConnector) + assert result.request.transport == MCPTransport.stdio + assert result.request.command == "npx" + assert result.request.args == ["-y", "@example/mcp-server"] + assert result.request.env == {"API_KEY": "value"} + + def test_disallowed_stdio_command_returns_error(self): + result = _single({"mcpServers": {"evil": {"command": "rm", "args": ["-rf", "/"]}}}) + assert isinstance(result, ConnectorConversionError) + assert "not in the allowed commands list" in result.error + + def test_unsupported_type_returns_error(self): + result = _single({"mcpServers": {"ws": {"type": "websocket", "url": "wss://x.example"}}}) + assert isinstance(result, ConnectorConversionError) + assert "Unsupported connector type" in result.error + + def test_missing_url_and_command_returns_error(self): + result = _single({"mcpServers": {"empty": {}}}) + assert isinstance(result, ConnectorConversionError) + assert "either a url or a command" in result.error + + def test_url_and_command_together_returns_error(self): + result = _single({"mcpServers": {"both": {"url": "https://x.example/mcp", "command": "npx"}}}) + assert isinstance(result, ConnectorConversionError) + assert "both a url and a command" in result.error + + def test_name_empty_after_sanitization_returns_error(self): + result = _single({"mcpServers": {"---": {"url": "https://x.example/mcp"}}}) + assert isinstance(result, ConnectorConversionError) + assert "empty after sanitization" in result.error + + +class TestConvertMcpServersList: + def test_anthropic_messages_api_list_shape(self): + result = _single( + { + "mcp_servers": [ + { + "type": "url", + "url": "https://mcp.example.com/sse", + "name": "deepwiki", + "authorization_token": "tok", + } + ] + } + ) + assert isinstance(result, ConvertedConnector) + assert result.request.server_name == "deepwiki" + assert result.request.transport == MCPTransport.http + assert result.request.credentials == {"auth_value": "tok"} + + def test_list_entry_without_name_returns_error(self): + result = _single({"mcp_servers": [{"type": "url", "url": "https://x.example/mcp"}]}) + assert isinstance(result, ConnectorConversionError) + assert "must have a name" in result.error + + def test_partial_conversion_preserves_per_entry_results(self): + results = convert_connector_entries( + MCPConnectorImportRequest.model_validate( + { + "mcpServers": { + "good": {"url": "https://good.example/mcp"}, + "bad": {"type": "websocket", "url": "wss://bad.example"}, + } + } + ) + ) + assert len(results) == 2 + assert isinstance(results[0], ConvertedConnector) + assert isinstance(results[1], ConnectorConversionError) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 0d639e1cb6a..7b3580f166a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -6766,3 +6766,143 @@ class TestConnectedAppViewAnnotation: assert all(server.connected_app_reachable is None for server in result) reload_mock.assert_not_awaited() + + +class TestImportMCPServers: + """Bulk connector import must be admin-only and report per-entry outcomes.""" + + @staticmethod + def _import_patches(existing_servers, create_mock, mock_manager): + return ( + patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers", + AsyncMock(return_value=existing_servers), + ), + patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", + create_mock, + ), + patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ) + + @pytest.mark.asyncio + async def test_non_admin_is_rejected(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + MCPConnectorImportRequest, + import_mcp_servers, + ) + + payload = MCPConnectorImportRequest.model_validate( + {"mcpServers": {"srv": {"url": "https://x.example/mcp"}}} + ) + caller = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.INTERNAL_USER) + + with patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ): + with pytest.raises(HTTPException) as exc_info: + await import_mcp_servers(payload=payload, user_api_key_dict=caller) + + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_import_reports_imported_skipped_and_errors(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + MCPConnectorImportRequest, + import_mcp_servers, + ) + + payload = MCPConnectorImportRequest.model_validate( + { + "mcpServers": { + "new-server": {"url": "https://new.example/mcp", "authorization_token": "tok"}, + "existing": {"url": "https://existing.example/mcp"}, + "broken": {"type": "websocket", "url": "wss://x.example"}, + } + } + ) + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") + existing = generate_mock_mcp_server_db_record(server_id="existing-1", alias="existing") + created = generate_mock_mcp_server_db_record(server_id="created-1", alias="new_server") + create_mock = AsyncMock(return_value=created) + mock_manager = MagicMock() + mock_manager.reload_servers_from_database = AsyncMock() + + with ExitStack() as stack: + for p in self._import_patches([existing], create_mock, mock_manager): + stack.enter_context(p) + result = await import_mcp_servers(payload=payload, user_api_key_dict=admin) + + assert [entry.name for entry in result.imported] == ["new-server"] + assert result.imported[0].server_id == "created-1" + assert [entry.name for entry in result.skipped] == ["existing"] + assert "already exists" in result.skipped[0].reason + assert [entry.name for entry in result.errors] == ["broken"] + create_mock.assert_awaited_once() + sent_request = create_mock.await_args[0][1] + assert sent_request.credentials == {"auth_value": "tok"} + mock_manager.reload_servers_from_database.assert_awaited_once() + + @pytest.mark.asyncio + async def test_duplicate_names_within_payload_are_skipped(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + MCPConnectorImportRequest, + import_mcp_servers, + ) + + payload = MCPConnectorImportRequest.model_validate( + { + "mcp_servers": [ + {"type": "url", "url": "https://a.example/mcp", "name": "dup srv"}, + {"type": "url", "url": "https://b.example/mcp", "name": "dup-srv"}, + ] + } + ) + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") + created = generate_mock_mcp_server_db_record(server_id="created-1", alias="dup_srv") + create_mock = AsyncMock(return_value=created) + mock_manager = MagicMock() + mock_manager.reload_servers_from_database = AsyncMock() + + with ExitStack() as stack: + for p in self._import_patches([], create_mock, mock_manager): + stack.enter_context(p) + result = await import_mcp_servers(payload=payload, user_api_key_dict=admin) + + assert len(result.imported) == 1 + assert len(result.skipped) == 1 + assert "Duplicate connector name" in result.skipped[0].reason + create_mock.assert_awaited_once() + + @pytest.mark.asyncio + async def test_no_imports_skips_registry_refresh(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + MCPConnectorImportRequest, + import_mcp_servers, + ) + + payload = MCPConnectorImportRequest.model_validate( + {"mcpServers": {"existing": {"url": "https://existing.example/mcp"}}} + ) + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") + existing = generate_mock_mcp_server_db_record(server_id="existing-1", alias="existing") + create_mock = AsyncMock() + mock_manager = MagicMock() + mock_manager.reload_servers_from_database = AsyncMock() + + with ExitStack() as stack: + for p in self._import_patches([existing], create_mock, mock_manager): + stack.enter_context(p) + result = await import_mcp_servers(payload=payload, user_api_key_dict=admin) + + assert result.imported == () + create_mock.assert_not_awaited() + mock_manager.reload_servers_from_database.assert_not_awaited() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ImportMCPServers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ImportMCPServers.tsx new file mode 100644 index 00000000000..d04160bca3e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ImportMCPServers.tsx @@ -0,0 +1,134 @@ +import React, { useState } from "react"; +import { Button } from "@/components/ui/button"; +import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { Textarea } from "@/components/ui/textarea"; +import { Alert, AlertTitle } from "@/components/shared/Alert"; +import { importMCPServers } from "@/components/networking"; +import { toast } from "@/lib/toast"; +import { MCPConnectorImportResponse, parseConnectorConfig } from "./importConnectorConfig"; + +interface ImportMCPServersProps { + accessToken: string; + open: boolean; + onClose: () => void; + onImported: () => void; +} + +const PLACEHOLDER = `{ + "mcpServers": { + "my_server": { + "url": "https://example.com/mcp", + "authorization_token": "..." + } + } +}`; + +const ImportMCPServers: React.FC = ({ accessToken, open, onClose, onImported }) => { + const [configText, setConfigText] = useState(""); + const [parseError, setParseError] = useState(null); + const [isImporting, setIsImporting] = useState(false); + const [result, setResult] = useState(null); + + const handleClose = () => { + setConfigText(""); + setParseError(null); + setResult(null); + onClose(); + }; + + const handleImport = async () => { + const parsed = parseConnectorConfig(configText); + if (!parsed.ok) { + setParseError(parsed.error); + return; + } + setParseError(null); + setIsImporting(true); + try { + const response = (await importMCPServers(accessToken, parsed.payload)) as MCPConnectorImportResponse; + setResult(response); + if (response.imported.length > 0) { + toast.success(`Imported ${response.imported.length} MCP server${response.imported.length === 1 ? "" : "s"}`); + onImported(); + } + } catch (error) { + console.error("Failed to import MCP servers:", error); + setParseError("Import request failed. Check the proxy logs for details."); + } finally { + setIsImporting(false); + } + }; + + return ( + !isOpen && handleClose()}> + + + Import MCP Connectors + +
+

+ Paste an Anthropic connector configuration: the mcpServers mapping from a Claude Desktop / + Claude Code config file, or the mcp_servers array from the Anthropic Messages API. +

+