From b080c59e78e25d73934ad8550e832d3ee7cd882d Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 4 Nov 2026 17:39:30 -0800 Subject: [PATCH 01/39] fix(proxy): treat Postgres restart and shutdown errors as DB unavailable, not poison spend-log rows (#50132) * fix(proxy): treat Postgres restart and shutdown errors as DB unavailable, not poison spend-log rows Postgres refuses every new connection with SQLSTATE 57P03 while it starts up, shuts down or replays WAL, and prisma surfaces that as a base DataError with no code. The spend-log writer took it for a poison row, bisected the batch down to single rows and dropped each one. Spend logs now retry and requeue on 57P03 and on server-terminated sessions (57P01, 57P02, 57P05), tool and model usage rollups requeue when the write provably never reached the server, the tool registry requeues its rows, and auth answers 503 instead of 401 while the auth database restarts. Co-authored-by: Dinu Arateanu * fix(proxy): requeue tool discovery rows only when the write never reached Postgres, keep the auth fallback to 57P03 A terminated session (57P01, 57P02, 57P05) may have committed the tool upsert, and the upsert adds call_count on conflict, so replaying it counts the calls twice. The non-Redis tool discovery flush now uses the same never-sent rule as the rollups, which also stops a permanently rejected row (54000) from coming back every tick. One terminated session does not prove the database is down, so it no longer opens the allow_requests_on_db_unavailable fallback, it still answers 503 * fix(proxy): a full server (53300, P2024) also never took the write, so counters and tool rows are sent again Postgres refuses the connection before any statement runs when it is out of slots, and a prisma pool timeout never reaches the server either, so those failures are as safe to replay as 57P03 and P1001 --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Dinu Arateanu --- .../tool_discovery_queue.py | 7 + litellm/proxy/db/exception_handler.py | 70 +++++++++ litellm/proxy/db/tool_registry_writer.py | 21 ++- litellm/proxy/utils.py | 27 +++- .../test_tool_discovery_queue.py | 17 ++- tests/unit/proxy/db/test_exception_handler.py | 136 +++++++++++++++++- .../proxy/db/test_tool_registry_writer.py | 51 +++++-- .../prisma_and_spend/test_spend_functions.py | 128 ++++++++++++++++- 8 files changed, 435 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py index e9ec040023c..abbeab7f89d 100644 --- a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -55,3 +55,10 @@ class ToolDiscoveryQueue: """Return and clear all pending rows.""" items, self._pending = list(self._pending.values()), {} return items + + def requeue(self, items: Iterable[ToolDiscoveryQueueItem]) -> None: + """Put rows from a failed write back, ahead of and merged with anything queued since the flush.""" + since_flush: Final = self._pending + self._pending = {} + for item in (*items, *since_flush.values()): + self.add_update(item) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 8fbc8708f8c..dbf636e1630 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -31,6 +31,24 @@ _CONNECTION_CAPACITY_PHRASES: Final = ( "too many connections for database", "remaining connection slots are reserved", ) +_NOT_ACCEPTING_CONNECTIONS_SQLSTATE: Final = "57P03" +_NOT_ACCEPTING_CONNECTIONS_PHRASES: Final = ( + "the database system is starting up", + "the database system is shutting down", + "the database system is not yet accepting connections", + "the database system is not accepting connections", + "the database system is in recovery mode", +) +_SERVER_TERMINATED_SQLSTATES: Final = frozenset({"57P01", "57P02", "57P05"}) +_SERVER_TERMINATED_PHRASES: Final = ( + "terminating connection due to administrator command", + "terminating connection due to immediate shutdown command", + "terminating connection because of unexpected sigquit signal", + "terminating connection because of crash of another server process", + "terminating connection due to idle-session timeout", +) +_UNREACHABLE_PRISMA_CODE: Final = "P1001" +_UNREACHABLE_PHRASES: Final = ("can't reach database server", "cannot reach database server") _PRISMA_POOL_TIMEOUT_CODE: Final = "P2024" _LOCK_TIMEOUT_SQLSTATE: Final = "55P03" _LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout" @@ -139,6 +157,8 @@ class PrismaDBExceptionHandler: return True if PrismaDBExceptionHandler.is_database_capacity_error(e): return True + if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e): + return True return isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection @staticmethod @@ -289,6 +309,39 @@ class PrismaDBExceptionHandler: error_message: Final = str(e).lower() return any(phrase in error_message for phrase in _CONNECTION_CAPACITY_PHRASES) + @staticmethod + def is_database_not_accepting_connections_error(e: Exception) -> bool: + """True iff Postgres refused the connection because it is starting up, shutting + down or in crash recovery (SQLSTATE 57P03). The refusal comes before any statement + runs, so nothing the write carried was applied.""" + import prisma + + if not isinstance(e, _exception_types(prisma.errors.PrismaError)): + return False + if PrismaDBExceptionHandler.postgres_sqlstate(e) == _NOT_ACCEPTING_CONNECTIONS_SQLSTATE: + return True + error_message: Final = str(e).lower() + return any(phrase in error_message for phrase in _NOT_ACCEPTING_CONNECTIONS_PHRASES) + + @staticmethod + def is_database_lifecycle_error(e: Exception) -> bool: + """True iff Postgres failed the statement because of a server lifecycle event: + a refused connection during startup, shutdown or recovery (57P03), or a session + the server terminated (57P01, 57P02, 57P05). The server comes back on its own, + so the rows are not at fault. A terminated session may have committed before it + was cut, so only ``is_database_not_accepting_connections_error`` proves a + non-idempotent write never applied.""" + import prisma + + if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e): + return True + if not isinstance(e, _exception_types(prisma.errors.PrismaError)): + return False + if PrismaDBExceptionHandler.postgres_sqlstate(e) in _SERVER_TERMINATED_SQLSTATES: + return True + error_message: Final = str(e).lower() + return any(phrase in error_message for phrase in _SERVER_TERMINATED_PHRASES) + @staticmethod def is_lock_timeout_error(e: Exception) -> bool: """True iff Postgres cancelled a statement for exceeding ``lock_timeout`` @@ -406,6 +459,8 @@ class PrismaDBExceptionHandler: return True if PrismaDBExceptionHandler.is_database_capacity_error(e): return True + if PrismaDBExceptionHandler.is_database_lifecycle_error(e): + return True if PrismaDBExceptionHandler.is_prisma_engine_internal_error(e): return True if "cached plan must not change result type" in str(e).lower(): @@ -680,3 +735,18 @@ def db_write_failure_is_requeue_safe(e: Exception) -> bool: return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES + + +def db_write_never_sent(e: Exception) -> bool: + """Whether the database could not take the write at all, because it was unreachable (P1001), + refusing connections (57P03) or full (53300, P2024), so even a non-idempotent write can be sent again.""" + if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e): + return True + if PrismaDBExceptionHandler.is_database_capacity_error(e): + return True + if _prisma_error_code(e) == _UNREACHABLE_PRISMA_CODE: + return True + if not PrismaDBExceptionHandler.is_prisma_error(e): + return False + error_message: Final = str(e).lower() + return any(phrase in error_message for phrase in _UNREACHABLE_PHRASES) diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index f7ec88945a1..9600fb40ad3 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -23,7 +23,11 @@ from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( ToolDiscoveryQueue, aggregate_tool_discoveries, ) -from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry, db_write_failure_is_requeue_safe +from litellm.proxy.db.exception_handler import ( + call_with_db_reconnect_retry, + db_write_failure_is_requeue_safe, + db_write_never_sent, +) from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.prisma_protocols import TableActions @@ -134,8 +138,19 @@ async def flush_tool_discovery_queue(prisma_client: "PrismaClient", tool_discove return try: await batch_upsert_tools(prisma_client=prisma_client, items=items) - except Exception as e: # noqa: BLE001 # a lost tool row only delays its discovery, the scheduler job must not fail - verbose_proxy_logger.debug("flush_tool_discovery_queue error (non-blocking): %s", e) + except Exception as e: # noqa: BLE001 # the rows go back to the queue or are dropped, the scheduler job must not fail + if db_write_never_sent(e): + tool_discovery_queue.requeue(items) + verbose_proxy_logger.warning( + "Tool discovery - %d tool rows went back to the queue for the next tick. Error: %s", len(items), str(e) + ) + return + verbose_proxy_logger.error( + "Tool discovery - dropped %d tool rows: the failed statement may have applied or the database " + "refused the data, so re-sending it is not safe. Error: %s", + len(items), + str(e), + ) async def commit_tool_discoveries( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 145c8e38229..8b77e39b3e5 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -213,6 +213,7 @@ from litellm.proxy.db.db_url_settings import ( from litellm.proxy.db.exception_handler import ( PrismaDBExceptionHandler, call_with_db_reconnect_retry, + db_write_never_sent, ) from litellm.proxy.db.health_check_latest import ( LatestHealthCheckRow, @@ -8803,9 +8804,20 @@ async def update_daily_tag_spend( def _rollup_batch_never_applied(e: Exception) -> bool: """True when a drained rollup batch provably never reached its rows: Postgres had no - connection for it (53300) or cancelled it under ``lock_timeout`` before it took the row - lock (55P03). Either way re-sending the batch cannot double-count, so it is requeued.""" - return PrismaDBExceptionHandler.is_database_capacity_error(e) or PrismaDBExceptionHandler.is_lock_timeout_error(e) + connection for it (53300), was unreachable or refusing connections (P1001, 57P03), or + cancelled it under ``lock_timeout`` before it took the row lock (55P03). In each case + re-sending the batch cannot double-count, so it is requeued.""" + return ( + PrismaDBExceptionHandler.is_database_capacity_error(e) + or db_write_never_sent(e) + or PrismaDBExceptionHandler.is_lock_timeout_error(e) + ) + + +def _rollup_failure_stops_job(e: Exception) -> bool: + """True when the database cannot take any write, so the job stops instead of draining + the queues after this one into the same failure.""" + return PrismaDBExceptionHandler.is_database_capacity_error(e) or db_write_never_sent(e) async def update_spend_logs_job( @@ -8896,12 +8908,12 @@ async def _run_spend_logs_job( *prisma_client.tool_usage_transactions, ] verbose_proxy_logger.warning( - "Spend tracking - database out of connections or rollup row lock timed out during tool usage " + "Spend tracking - database unavailable or rollup row lock timed out during tool usage " "flush, requeued %s tool usage transactions for the next flush: %s", len(tool_usage_to_process), tool_tracking_err, ) - if PrismaDBExceptionHandler.is_database_capacity_error(tool_tracking_err): + if _rollup_failure_stops_job(tool_tracking_err): raise else: verbose_proxy_logger.error( @@ -8928,12 +8940,12 @@ async def _run_spend_logs_job( *prisma_client.model_usage_transactions, ] verbose_proxy_logger.warning( - "Spend tracking - database out of connections or rollup row lock timed out during model usage " + "Spend tracking - database unavailable or rollup row lock timed out during model usage " "flush, requeued %s model usage transactions for the next flush: %s", len(model_usage_to_process), model_usage_err, ) - if PrismaDBExceptionHandler.is_database_capacity_error(model_usage_err): + if _rollup_failure_stops_job(model_usage_err): raise else: verbose_proxy_logger.error( @@ -9110,6 +9122,7 @@ def _is_transient_spend_log_write_error(e: Exception) -> bool: return ( PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_prisma_engine_internal_error(e) + or PrismaDBExceptionHandler.is_database_lifecycle_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e) ) diff --git a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py index 0a015255f2e..814fdcb4eda 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py @@ -2,7 +2,6 @@ Unit tests for ToolDiscoveryQueue. """ - import pytest from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( @@ -63,3 +62,19 @@ def test_aggregate_sums_call_counts_across_workers_flushes(): {"tool_name": "get_weather", "origin": "user_defined", "call_count": 5}, {"tool_name": "search", "call_count": 1}, ) + + +def test_requeued_rows_go_ahead_of_and_merge_with_sightings_since_the_flush(queue: ToolDiscoveryQueue) -> None: + queue.add_update({"tool_name": "get_weather", "origin": "mcp_server"}) + queue.add_update({"tool_name": "search", "origin": "user_defined"}) + failed_write = queue.flush() + queue.add_update({"tool_name": "translate", "origin": "user_defined"}) + queue.add_update({"tool_name": "search", "origin": "user_defined"}) + + queue.requeue(failed_write) + + assert queue.flush() == [ + {"tool_name": "get_weather", "origin": "mcp_server", "call_count": 1}, + {"tool_name": "search", "origin": "user_defined", "call_count": 2}, + {"tool_name": "translate", "origin": "user_defined", "call_count": 1}, + ] diff --git a/tests/unit/proxy/db/test_exception_handler.py b/tests/unit/proxy/db/test_exception_handler.py index 46704bf52ef..abc0850194d 100644 --- a/tests/unit/proxy/db/test_exception_handler.py +++ b/tests/unit/proxy/db/test_exception_handler.py @@ -22,7 +22,7 @@ from prisma.errors import ( import litellm from litellm.proxy._types import ProxyErrorTypes, ProxyException -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, db_write_never_sent # Test is_database_connection_error method @@ -996,3 +996,137 @@ def test_connection_closure_code_survives_prisma_data_error_wrapping() -> None: ) assert PrismaDBExceptionHandler.is_database_transport_error(error) is True assert PrismaDBExceptionHandler.is_permanent_database_fault(error) is False + + +def _connector_error(message: str) -> DataError: + return DataError( + data={"user_facing_error": {"message": f"Error in connector: Error querying the database: FATAL: {message}"}} + ) + + +def _batched_postgres_error(sqlstate: str, message: str) -> DataError: + return DataError( + data={ + "user_facing_error": { + "message": "Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: " + f'None, kind: QueryError(PostgresError {{ code: "{sqlstate}", message: "{message}", severity: "FATAL" }}) }})' + } + } + ) + + +def _localized_error(sqlstate: str) -> DataError: + return DataError(data={"user_facing_error": {"message": "FATAL: Datenbank", "meta": {"code": sqlstate}}}) + + +_NOT_ACCEPTING_CONNECTIONS: Final = ( + _connector_error("the database system is starting up"), + _connector_error("the database system is shutting down"), + _connector_error("the database system is not yet accepting connections"), + _connector_error("the database system is not accepting connections"), + _connector_error("the database system is in recovery mode"), + _localized_error("57P03"), + _batched_postgres_error("57P03", "the database system is starting up"), + EngineRequestError( + MagicMock(status=500), + '{"is_panic":false,"message":"Error in connector: Error querying the database: FATAL: the database system is ' + 'starting up","backtrace":null}', + ), + RawQueryError( + data={ + "user_facing_error": { + "message": "Raw query failed. Code: `57P03`. Message: `db error: FATAL: the database system is starting up`", + "meta": {"code": "57P03", "message": "FATAL: the database system is starting up"}, + "error_code": "P2010", + } + } + ), +) +_SERVER_TERMINATED: Final = ( + _connector_error("terminating connection due to administrator command"), + _connector_error("terminating connection due to immediate shutdown command"), + _connector_error("terminating connection because of unexpected SIGQUIT signal"), + _connector_error("terminating connection because of crash of another server process"), + _connector_error("terminating connection due to idle-session timeout"), + _localized_error("57P01"), + _localized_error("57P02"), + _localized_error("57P05"), + _batched_postgres_error("57P01", "terminating connection due to administrator command"), +) + + +@pytest.mark.parametrize( + ("error", "connection_error"), + [ + *((error, True) for error in _NOT_ACCEPTING_CONNECTIONS), + *((error, False) for error in _SERVER_TERMINATED), + ], +) +def test_postgres_lifecycle_errors_mean_the_database_is_unavailable(error: Exception, connection_error: bool) -> None: + """A server that is restarting, shutting down or cutting sessions is a transient outage: + auth answers 503 rather than 401, and the spend-log writer must not take the rows for poison. + Only a server refusing connections opens the ``allow_requests_on_db_unavailable`` fallback, + one terminated session does not prove the database is down.""" + assert { + "lifecycle": PrismaDBExceptionHandler.is_database_lifecycle_error(error), + "service_unavailable": PrismaDBExceptionHandler.is_database_service_unavailable_error(error), + "connection_error": PrismaDBExceptionHandler.is_database_connection_error(error), + "permanent": PrismaDBExceptionHandler.is_permanent_database_fault(error), + } == { + "lifecycle": True, + "service_unavailable": True, + "connection_error": connection_error, + "permanent": False, + } + + +@pytest.mark.parametrize( + ("error", "never_sent"), + [ + *((error, True) for error in _NOT_ACCEPTING_CONNECTIONS), + *((error, False) for error in _SERVER_TERMINATED), + (DataError(data={"user_facing_error": {"message": "Can't reach database server at `db`:`5432`"}}), True), + (DataError(data={"user_facing_error": {"error_code": "P1001", "message": "unreachable"}}), True), + (_capacity_error("sorry, too many clients already"), True), + (_localized_error("53300"), True), + ( + DataError( + data={ + "user_facing_error": { + "error_code": "P2024", + "message": "Timed out fetching a new connection from the connection pool", + } + } + ), + True, + ), + (httpx.ReadTimeout("read timed out"), False), + (_localized_error("57014"), False), + ], +) +def test_db_write_never_sent_only_for_failures_before_any_statement_ran(error: Exception, never_sent: bool) -> None: + """A refused, unreachable or full server provably applied nothing, while a terminated session + or a timeout may have committed before the failure, so only the former lets a non-idempotent + write be replayed.""" + assert db_write_never_sent(error) is never_sent + + +@pytest.mark.parametrize( + "error", + [ + _localized_error("57P04"), + _localized_error("57014"), + _localized_error("23505"), + DataError(data={"user_facing_error": {"message": "invalid byte sequence for encoding UTF8: 0x00"}}), + RuntimeError("the database system is starting up"), + ], +) +def test_other_failures_are_not_postgres_lifecycle_errors(error: Exception) -> None: + assert PrismaDBExceptionHandler.is_database_lifecycle_error(error) is False + assert db_write_never_sent(error) is False + + +def test_lifecycle_error_is_seen_through_a_wrapping_exception() -> None: + wrapped: Final = RuntimeError("spend flush failed") + wrapped.__cause__ = _connector_error("the database system is starting up") + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(wrapped) is True diff --git a/tests/unit/proxy/db/test_tool_registry_writer.py b/tests/unit/proxy/db/test_tool_registry_writer.py index 76898b0eeab..a5ad897913e 100644 --- a/tests/unit/proxy/db/test_tool_registry_writer.py +++ b/tests/unit/proxy/db/test_tool_registry_writer.py @@ -69,9 +69,7 @@ def _make_prisma( prisma.db.litellm_tooltable = MagicMock() prisma.db.litellm_tooltable.upsert = AsyncMock(return_value=upsert_return) prisma.db.execute_raw = AsyncMock(return_value=1) - prisma.db.litellm_tooltable.find_many = AsyncMock( - return_value=find_many_rows if find_many_rows is not None else [] - ) + prisma.db.litellm_tooltable.find_many = AsyncMock(return_value=find_many_rows if find_many_rows is not None else []) prisma.db.litellm_tooltable.find_unique = AsyncMock(return_value=find_unique_row) prisma.db.litellm_verificationtoken = MagicMock() prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=list(key_rows)) @@ -227,9 +225,7 @@ async def test_get_tool_found(): result = await get_tool(prisma, "my_tool") assert result is not None assert result.tool_name == "my_tool" - prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with( - where={"tool_name": "my_tool"} - ) + prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with(where={"tool_name": "my_tool"}) @pytest.mark.asyncio @@ -266,9 +262,7 @@ async def test_update_tool_policy_calls_upsert_then_get_tool(): updated_by="admin", ) prisma = _make_prisma(find_unique_row=row) - result = await update_tool_policy( - prisma, "my_tool", updated_by="admin", input_policy="blocked" - ) + result = await update_tool_policy(prisma, "my_tool", updated_by="admin", input_policy="blocked") assert result is not None assert result.input_policy == "blocked" prisma.db.litellm_tooltable.upsert.assert_awaited_once() @@ -567,3 +561,42 @@ async def test_failed_tool_table_write_requeues_the_rows_only_when_they_provably left_in_redis = [json.loads(entry) for entry in redis.lists.get(REDIS_TOOL_DISCOVERY_BUFFER_KEY, [])] expected = [[{"tool_name": "get_weather", "origin": "user_defined", "call_count": 1}]] if requeued else [] assert left_in_redis == expected, f"requeued={requeued} but redis holds {left_in_redis}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "requeued"), + [ + (_postgres_rejection("57P03"), True), + (_postgres_rejection("53300"), True), + (_postgres_rejection("57P01"), False), + (_postgres_rejection("54000"), False), + (_postgres_rejection("23502"), False), + ], + ids=[ + "restarting_server_keeps_the_rows", + "full_server_keeps_the_rows", + "terminated_session_may_have_counted_so_the_rows_are_dropped", + "oversized_value_is_dropped", + "rejected_row_is_dropped", + ], +) +async def test_failed_write_without_redis_keeps_the_rows_in_the_queue_only_when_they_provably_did_not_apply( + failure: Exception, requeued: bool +): + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(side_effect=failure) + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=_InMemoryRedis()), + pod_lock_manager=PodLockManager(redis_cache=_InMemoryRedis()), + use_redis_buffer=False, + ) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + + expected = [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 2 if requeued else 1}] + assert queue.flush() == expected diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py index fc7ae531c34..293f2d49f43 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -23,9 +23,9 @@ import pytest from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY from litellm.proxy.utils import ( MAX_SPEND_LOG_DRAIN_ITERATIONS, + drain_spend_logs_queue, monitor_spend_logs_queue, raise_failed_update_spend_exception, - drain_spend_logs_queue, recover_parked_spend_logs, update_daily_tag_spend, update_spend, @@ -985,6 +985,85 @@ async def test_update_spend_logs_job_requeues_whole_batch_when_postgres_is_out_o } == {"create_many_calls": 1, "queue_after": ["r1", "r2", "r3"], "backoff_sleeps": []} +def _postgres_restarting() -> Exception: + """prisma's shape for Postgres SQLSTATE 57P03 while the server runs crash recovery after + a restart: a base ``DataError`` carrying only the connector message.""" + from prisma.errors import DataError + + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Error in connector: Error querying the database: FATAL: the database system is starting up", + "backtrace": None, + } + } + ) + + +def _postgres_unreachable() -> Exception: + from prisma.errors import DataError + + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Can't reach database server at `db.internal`:`5432`", + "error_code": "P1001", + } + } + ) + + +def _postgres_terminated_session() -> Exception: + from prisma.errors import DataError + + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Error in connector: Error querying the database: FATAL: terminating connection due to " + "administrator command", + "meta": {"code": "57P01", "message": "terminating connection due to administrator command"}, + } + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [_postgres_restarting(), _postgres_terminated_session()], ids=["57P03", "57P01"]) +async def test_update_spend_logs_job_keeps_every_row_while_postgres_restarts( + error: Exception, + mock_prisma_client: MagicMock, + make_spend_log_row: Callable[..., dict[str, object]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A restarting server refuses or cuts every statement, so the batch is not bisected + into per-row drops: it is retried whole and then goes back to the queue head.""" + monkeypatch.setattr(asyncio, "sleep", AsyncMock()) + proxy_logging: Final = MagicMock() + proxy_logging.failure_handler = AsyncMock() + mock_prisma_client.spend_log_transactions = [ + make_spend_log_row(request_id="r1"), + make_spend_log_row(request_id="r2"), + make_spend_log_row(request_id="r3"), + make_spend_log_row(request_id="r4"), + ] + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=error) + + with suppress(Exception): + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + ) + + assert { + "create_many_calls": mock_prisma_client.db.litellm_spendlogs.create_many.await_count, + "queue_after": [row["request_id"] for row in mock_prisma_client.spend_log_transactions], + } == {"create_many_calls": 4, "queue_after": ["r1", "r2", "r3", "r4"]} + + def _tool_usage_transaction(request_id: str) -> object: from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -1177,6 +1256,53 @@ async def test_model_usage_flush_requeues_and_stops_the_job_when_postgres_is_out } == {"queue_after": [first, second], "autorouter_queue_after": 1} +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [_postgres_restarting(), _postgres_unreachable()], ids=["57P03", "P1001"]) +async def test_rollups_requeue_and_stop_the_job_while_postgres_cannot_take_writes( + error: Exception, + mock_prisma_client: MagicMock, +) -> None: + """A server that is restarting or unreachable never saw the rollup batch, so the tool + usage increments are kept for the next flush, and the job stops before draining the + auto-router queue into the same outage.""" + mock_prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=error) + mock_prisma_client.spend_log_transactions = [] + first, second = _tool_usage_transaction("r1"), _tool_usage_transaction("r2") + mock_prisma_client.tool_usage_transactions = [first, second] + mock_prisma_client.autorouter_turn_transactions = [MagicMock(name="autorouter_turn")] + + with suppress(Exception): + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + + assert { + "queue_after": mock_prisma_client.tool_usage_transactions, + "autorouter_queue_after": len(mock_prisma_client.autorouter_turn_transactions), + } == {"queue_after": [first, second], "autorouter_queue_after": 1} + + +@pytest.mark.asyncio +async def test_rollups_drop_a_batch_whose_session_the_server_terminated( + mock_prisma_client: MagicMock, +) -> None: + """A terminated session (57P01) may have committed before it was cut, and the rollup + increments are not idempotent, so the batch is dropped rather than counted twice.""" + mock_prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=_postgres_terminated_session()) + mock_prisma_client.spend_log_transactions = [] + mock_prisma_client.tool_usage_transactions = [_tool_usage_transaction("r1")] + + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + + assert mock_prisma_client.tool_usage_transactions == [] + + @pytest.mark.asyncio async def test_model_usage_flush_still_drops_a_batch_that_may_have_applied( mock_prisma_client: MagicMock, From 906fa7ed280c12c905672ca39a19f3371304f7db Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Wed, 4 Nov 2026 17:49:15 -0800 Subject: [PATCH 02/39] feat(accounting): add opt-in SpendLogs ledger (#49856) * feat(accounting): add opt-in SpendLogs ledger delivery * fix(accounting): preserve missing organization daily spend parity * fix(accounting): preserve ledger receipts when replaying migration * fix(accounting): preserve callback windows and tune ledger batches * fix(accounting): preserve window charges and bound ledger statements * test(accounting): verify budget callback registration order * test(accounting): verify captured charges after repeated startup * feat(accounting): store durable payloads independently of spend logs * fix(accounting): preserve first delivery and trace ledger writes * fix(accounting): keep SDK imports independent of proxy tracing * fix(accounting): honor log retention during batch takeover * test(accounting): cover retained batch alerts and invalid settings * test(accounting): verify downgrade command against pending work * fix(accounting): drain retries without starving fresh deliveries * test(accounting): cover delivery projection, metrics and window key reads Co-Authored-By: Claude Opus 5.5 * test(accounting): read delivery metrics without the shared registry Co-Authored-By: Claude Opus 5.5 * fix(accounting): charge a leftover $0 batch row when its claim fails A failed batch claim buffers the retrieve's row and reports it charged. At flush the row conflicted with the pre-upgrade $0 SpendLogs row, so its delivery was rolled back and the charge never reached the spend tables. The flush now takes that row over in the same transaction and keeps the delivery. Co-Authored-By: Claude Opus 5.5 * fix(accounting): trace the batch takeover at flush as a SpendLogs update Co-Authored-By: Claude Opus 5.5 * fix(accounting): keep local spend logs when SPEND_LOGS_URL has no export client The export client is never created, so main always writes SpendLogs to the proxy's own database even with SPEND_LOGS_URL set. The ledger skipped local rows whenever the URL was set, which dropped those request logs. It now treats a flush as an export only when the client exists, and batch claims write local rows as before. Co-Authored-By: Claude Opus 5.5 * fix(accounting): keep every request's charge when its log row is shared, evicted or rejected - Key non-batch deliveries on a per-request delivery_id so a reused client call id is still charged once per request, while buffer replays stay deduplicated - Persist the redacted api_key instead of an opaque bearer token - Keep the accounting payload when the spend-log queue evicts a row or Postgres rejects it - Strip NUL from payloads and skip non-text tags so projection never parks a paid request Co-Authored-By: Claude Opus 5.5 * fix(accounting): isolate a raw-SQL spend log row Postgres rejects instead of requeueing the batch A rejected row on the ledger insert path raises RawQueryError, which the exact DataError check missed, so the whole batch was requeued every flush. Data and integrity SQLSTATE classes now bisect like before, while missing tables and deadlocks keep retrying Co-Authored-By: Claude Opus 5.5 * fix(accounting): keep charging a request whose tag is too long for the daily tag index A tag past Postgres's btree entry limit failed the DailyTagSpend upsert, and on the ledger path that rolled back the request's key, team and user charges with it. Such tags are now left out of the daily tag rollup only Co-Authored-By: Claude Opus 5.5 * fix(accounting): store only guardrail billing fields in the accounting payload The ledger payload copied whole guardrail records, including responses and match details, into LiteLLM_SpendLogDelivery even when spend logs were disabled. It now keeps the id, name, usage and cost fields the projection reads. NUL stripping reuses safe_json_structure instead of a recursive walk Co-Authored-By: Claude Opus 5.5 * refactor(proxy): move spend payload stamping out of update_database Co-Authored-By: Claude Opus 5.5 * test(proxy): restore the modules the import-layering guard evicts The guard popped litellm.proxy.utils from sys.modules for good, so later tests on the same xdist worker patched a fresh module while calling functions bound to the old one, and spend queue tests failed depending on scheduling Co-Authored-By: Claude Opus 5.5 * refactor(proxy): keep update_database's payload stamping on its own local Co-Authored-By: Claude Opus 5.5 * refactor(accounting): type guardrail entries and use a public key redaction helper Co-Authored-By: Claude Opus 5.5 * fix(accounting): keep charging requests whose end-user id is too long for an index An oversized end-user id failed the SpendLogs insert with SQLSTATE 54000, which was requeued forever, and failed the end-user upsert in the delivery, which parked the key and team charge with it. Class 54 rows are now isolated like other rejected rows, and such ids are left out of the end-user rollups only Co-Authored-By: Claude Opus 5.5 * refactor(proxy): type the spend log queue snapshot Co-Authored-By: Claude Opus 5.5 * fix(accounting): leave out any projected row Postgres cannot index or store Caller-supplied ids and counters, such as a long tag or end-user id or an out-of-range compression count, failed the delivery transaction and rolled back the request's key, team and user charges. One check in project() now drops only the row that cannot be stored, replacing the per-field cutoffs Co-Authored-By: Claude Opus 5.5 * fix(accounting): keep the guardrail fields compression savings are counted from Slimming the guardrail records dropped the provider and the tokens_saved figure that Headroom compression savings are read from, and dropped a bare-dict record entirely. The payload now keeps both while still leaving out guardrail content Co-Authored-By: Claude Opus 5.5 * fix(accounting): cap projected counters so daily totals cannot overflow Counters such as a caller-supplied compression count fit one by one but overflowed the bigint daily total once summed across requests, rolling back the key, team and user charges in the same transaction. Each counter is now capped at 2^40, far below any total that could overflow Co-Authored-By: Claude Opus 5.5 * fix(accounting): let a buffered batch charge take over its zero-cost receipt A batch charge queued after a failed lookup hit the $0 receipt an earlier poll left behind, and ON CONFLICT dropped it. The buffered insert now runs the same receipt takeover as a live retrieve Co-Authored-By: Claude Opus 5.5 * fix(accounting): read a rescued batch row's spend from its accounting payload Rows rescued from an overflowing queue keep only the accounting payload, so the buffered takeover never saw their positive spend and left the $0 receipt uncharged Co-Authored-By: Claude Opus 5.5 * fix(accounting): read budget windows through the delivery transaction Window seed and existing-row reads used a separate connection while the delivery transaction held one, so with a one-connection pool they waited on their own transaction until it timed out. Inside an enclosing transaction they now share its connection Co-Authored-By: Claude Opus 5.5 * fix(accounting): retry a transient failure while keeping a dropped row's charge The rescue swallowed every error, so a database blip right after a row was rejected lost the unpaid charge. Only a rejected row is still logged and skipped, anything else propagates so the flush retries Co-Authored-By: Claude Opus 5.5 * test(accounting): give delivery projection payloads their charged entities project() now reads entities only from accounting_entities, so the hand-built delivery payloads carry them like real stored payloads do. Co-Authored-By: Claude Opus 5.5 * refactor(accounting): own ledger rows by their payload and key batches on the batch id Drop the accounting_path stamp and its SpendLogs column. A buffered row that carries an accounting payload is ledger owned, one without it is legacy, and spend counters get a plain windows_from_ledger flag. Window seeding tells applied ledger rows apart by their delivery row. A ledger batch is plain transaction A keyed by the batch id, so a repeated completed retrieve conflicts and charges nothing. A $0 retrieve records no delivery, a batch a legacy retrieve already charged is not charged again, and the legacy batch path is unchanged from main apart from refusing a batch the ledger already owns. Batch takeover and its delta bookkeeping are removed. Co-Authored-By: Claude Opus 5.5 * test(accounting): drop a projection test duplicated while rebasing Co-Authored-By: Claude Opus 5.5 * refactor(accounting): drop previous_payload now that deliveries are never replaced Batch takeover was the only writer of previous_payload, so the worker applies each delivery's payload as stored and the column leaves the migration. Co-Authored-By: Claude Opus 5.5 * refactor(accounting): drop the ledger downgrade checklist script The use_spend_logs_ledger flag is the rollback path, so a separate downgrade gate is no longer needed. Co-Authored-By: Claude Opus 5.5 * fix(accounting): name the window seed reads so their Prisma span resolves Co-Authored-By: Claude Opus 5.5 * refactor(accounting): remove the unused accounting event types and protocols types.py, protocols.py and measurements.py were left from the old event scaffolding. Nothing outside the package imports them, so they go with their tests. Co-Authored-By: Claude Opus 5.5 * fix(accounting): park only deliveries whose own data fails A delivery that keeps failing on a database outage, lock timeout or deadlock stays pending and retries with capped backoff, so a recovered database applies it. Only a data error or an invalid payload parks the isolated row. Co-Authored-By: Claude Opus 5.5 * fix(accounting): isolate a spend log row the insert statement rejects The raw SpendLogs insert validates columns and JSON before Postgres sees the row, so a row that fails there is now bisected out and dropped alone while its charge is kept, instead of failing and replaying the whole flush. Co-Authored-By: Claude Opus 5.5 * fix(accounting): carry a job's model window on every completed poll in ledger mode The ledger keeps the poll whose delivery claims the job, which need not be the poll that won the limiter's in-memory charge. Every completed poll now carries the job's window, so the delivery that lands always has it, while the live model budget is still charged once. Co-Authored-By: Claude Opus 5.5 * refactor(accounting): leave model-budget windows on the existing window writer Model-budget windows are written by the model budget limiter rather than derived from a request's spend, so they stay on main's window queue like live budget counters do. This drops window capture from the accounting payload, the deferred limiter registration and the job-window capture, and the ledger keeps owning key and team budget windows. Co-Authored-By: Claude Opus 5.5 * fix(accounting): keep a delivery pending when the database is unreachable Prisma reports a can't reach database server outage as a DataError, so the worker now checks for an unavailable database before treating a data error as the delivery's own fault. Co-Authored-By: Claude Opus 5.5 * fix(accounting): isolate only rows the raw insert cannot encode Catching every ValueError also swallowed unexpected insert errors that main re-raises. The raw insert now raises RejectedSpendLogRow for an unknown column or non-JSON value, and only that is isolated as a rejected row. Co-Authored-By: Claude Opus 5.5 * fix(accounting): store zero savings when caller data makes them uncomputable A huge caller-supplied tokens_saved overflowed the savings math inside build_accounting_payload, which dropped the request's charge in ledger mode. Savings that cannot be computed are now recorded as zero and the charge lands. Co-Authored-By: Claude Opus 5.5 --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- .../migration.sql | 19 + .../litellm_proxy_extras/schema.prisma | 17 + litellm/accounting/__init__.py | 59 - litellm/accounting/application.py | 71 + litellm/accounting/delivery_projection.py | 78 + litellm/accounting/measurements.py | 13 - litellm/accounting/metrics.py | 30 + litellm/accounting/projection.py | 28 +- litellm/accounting/protocols.py | 21 - litellm/accounting/snapshot.py | 94 +- litellm/accounting/types.py | 428 ----- litellm/accounting/worker.py | 102 ++ litellm/integrations/otel/model/spans.py | 1 + litellm/proxy/_types.py | 8 + .../proxy/db/budget_window_spend_writer.py | 98 +- litellm/proxy/db/daily_spend_bulk_upsert.py | 4 +- litellm/proxy/db/db_spend_update_writer.py | 306 +++- litellm/proxy/db/exception_handler.py | 15 + litellm/proxy/db/model_usage_rollup.py | 55 +- litellm/proxy/guardrails/usage_tracking.py | 61 +- .../proxy/hooks/proxy_track_cost_callback.py | 20 +- litellm/proxy/proxy_server.py | 66 +- litellm/proxy/schema.prisma | 17 + .../spend_tracking/budget_reservation.py | 9 + .../spend_tracking/spend_tracking_utils.py | 10 +- litellm/proxy/utils.py | 96 +- .../accounting_delivery_repository.py | 193 +++ litellm/repositories/spend_log_insert.py | 88 + litellm/repositories/table_repositories.py | 115 ++ schema.prisma | 17 + .../database/test_spendlogs_delivery.py | 512 ++++++ .../database/test_spendlogs_ledger_parity.py | 1437 +++++++++++++++++ .../accounting/test_delivery_projection.py | 85 + tests/unit/accounting/test_measurements.py | 235 --- tests/unit/accounting/test_metrics.py | 28 + tests/unit/accounting/test_projection.py | 126 ++ tests/unit/accounting/test_types.py | 482 ------ .../common_utils/test_check_batch_cost.py | 50 +- .../db/test_budget_window_spend_writer.py | 33 +- .../proxy/db/test_db_spend_update_writer.py | 38 +- .../proxy/guardrails/test_usage_tracking.py | 12 + .../unit/proxy/hooks/test_proxy_hooks_init.py | 4 +- .../spend_tracking/test_budget_reservation.py | 7 + .../test_spend_tracking_utils.py | 116 +- ...test_proxy_server_endpoints_and_startup.py | 60 + tests/unit/proxy/test_update_spend.py | 122 ++ .../repositories/test_spend_log_insert.py | 31 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 + 48 files changed, 4041 insertions(+), 1482 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261103010000_spendlogs_accounting_delivery/migration.sql create mode 100644 litellm/accounting/application.py create mode 100644 litellm/accounting/delivery_projection.py delete mode 100644 litellm/accounting/measurements.py create mode 100644 litellm/accounting/metrics.py delete mode 100644 litellm/accounting/protocols.py delete mode 100644 litellm/accounting/types.py create mode 100644 litellm/accounting/worker.py create mode 100644 litellm/repositories/accounting_delivery_repository.py create mode 100644 litellm/repositories/spend_log_insert.py create mode 100644 tests/integration/database/test_spendlogs_delivery.py create mode 100644 tests/integration/database/test_spendlogs_ledger_parity.py create mode 100644 tests/unit/accounting/test_delivery_projection.py delete mode 100644 tests/unit/accounting/test_measurements.py create mode 100644 tests/unit/accounting/test_metrics.py delete mode 100644 tests/unit/accounting/test_types.py create mode 100644 tests/unit/repositories/test_spend_log_insert.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103010000_spendlogs_accounting_delivery/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103010000_spendlogs_accounting_delivery/migration.sql new file mode 100644 index 00000000000..b00cf892b68 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103010000_spendlogs_accounting_delivery/migration.sql @@ -0,0 +1,19 @@ +SET LOCAL lock_timeout = '10s'; +CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogDelivery" ( + "delivery_id" TEXT PRIMARY KEY DEFAULT gen_random_uuid()::text, + "request_id" TEXT NOT NULL, + "state" "AccountingDeliveryState" NOT NULL DEFAULT 'pending', + "attempts" INTEGER NOT NULL DEFAULT 0, + "next_attempt_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "applied_at" TIMESTAMP(3), + "last_error" TEXT, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "accounting_payload" JSONB NOT NULL, + "is_batch" BOOLEAN NOT NULL +); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_next_attempt_at_idx" + ON "LiteLLM_SpendLogDelivery" ("state", "next_attempt_at"); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_applied_at_idx" + ON "LiteLLM_SpendLogDelivery" ("state", "applied_at"); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_request_id_idx" + ON "LiteLLM_SpendLogDelivery" ("request_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index c007b047087..3aaa74d4fc2 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -2141,6 +2141,23 @@ model LiteLLM_AccountingDelivery { @@index([target, state, next_attempt_at]) } +model LiteLLM_SpendLogDelivery { + delivery_id String @id @default(dbgenerated("gen_random_uuid()::text")) + request_id String + state AccountingDeliveryState @default(pending) + accounting_payload Json + is_batch Boolean + attempts Int @default(0) + next_attempt_at DateTime @default(now()) + applied_at DateTime? + last_error String? + created_at DateTime @default(now()) + + @@index([state, next_attempt_at]) + @@index([state, applied_at]) + @@index([request_id]) +} + model LiteLLM_SubscriptionAccountTable { subscription_account_id String @id @default(uuid()) custom_llm_provider String diff --git a/litellm/accounting/__init__.py b/litellm/accounting/__init__.py index 5d59a35a987..e69de29bb2d 100644 --- a/litellm/accounting/__init__.py +++ b/litellm/accounting/__init__.py @@ -1,59 +0,0 @@ -"""Public accounting contracts. Importing these values does not enable accounting.""" - -from litellm.accounting.measurements import measurements_for -from litellm.accounting.protocols import CostAggregator -from litellm.accounting.types import ( - AdmissionDecision, - AggregateQuery, - AggregateRow, - CostEvent, - DeliveryState, - DeliveryTarget, - Dimension, - DimensionFilter, - DimensionName, - EventKind, - Granularity, - InboundUsageEvent, - Measurement, - Metric, - MetricName, - Multiplier, - NamedCost, - OutboundUsageEvent, - Outcome, - PricedQuantity, - PricingSnapshot, - Rate, - Unpriced, - validate_aggregate_rows, -) - -__all__ = [ - "AdmissionDecision", - "AggregateQuery", - "AggregateRow", - "CostAggregator", - "CostEvent", - "DeliveryState", - "DeliveryTarget", - "Dimension", - "DimensionFilter", - "DimensionName", - "EventKind", - "Granularity", - "InboundUsageEvent", - "Measurement", - "Metric", - "MetricName", - "Multiplier", - "NamedCost", - "OutboundUsageEvent", - "Outcome", - "PricedQuantity", - "PricingSnapshot", - "Rate", - "Unpriced", - "measurements_for", - "validate_aggregate_rows", -] diff --git a/litellm/accounting/application.py b/litellm/accounting/application.py new file mode 100644 index 00000000000..1e5f23e2997 --- /dev/null +++ b/litellm/accounting/application.py @@ -0,0 +1,71 @@ +from typing import TYPE_CHECKING, Final + +from litellm.accounting.delivery_projection import entity_transactions, project_deliveries +from litellm.accounting.projection import Contribution +from litellm.proxy.db.budget_window_spend_writer import commit_window_spend_updates +from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key +from litellm.proxy.db.db_span import db_span +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendTransaction +from litellm.proxy.db.model_usage_rollup import queue_model_usage_upsert +from litellm.repositories.accounting_delivery_repository import DeliveryProjection, LedgerTransaction + +if TYPE_CHECKING: + from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + from litellm.proxy.utils import PrismaClient, ProxyLogging + + +async def apply_deliveries( + transaction: LedgerTransaction, + deliveries: tuple[DeliveryProjection, ...], + *, + prisma_client: "PrismaClient", + writer: "DBSpendUpdateWriter", + proxy_logging_obj: "ProxyLogging", +) -> tuple[Contribution, ...]: + from litellm.proxy.guardrails.usage_tracking import UsageUnitIncrement, UsageUnitKey, upsert_usage_unit_row + + contributions: Final = project_deliveries(deliveries) + await writer.apply_accounting_entities( + prisma_client, proxy_logging_obj, transaction, entity_transactions(contributions) + ) + for table in sorted(DAILY_SPEND_TABLES.values(), key=lambda item: item.name): + rows = tuple( + {**dict(item.key), **dict(item.values), **dict(item.increments)} + for item in contributions + if item.table == table.name + ) + merged = merge_by_conflict_key(table, rows) + for offset in range(0, len(merged), 100): + sql, params = build_bulk_upsert(table, merged[offset : offset + 100]) + async with db_span("upsert_daily_spend", table.name): + await transaction.execute_raw(sql, *params) + async with transaction.batch_() as batcher: + for item in contributions: + if item.table == "LiteLLM_DailyModelUsage": + queue_model_usage_upsert(batcher, dict(item.key), dict(item.increments)) + for item in contributions: + if item.table == "LiteLLM_DailyGuardrailUsageUnits": + key = dict(item.key) + counters = dict(item.increments) + await upsert_usage_unit_row( + prisma_client, + UsageUnitKey(key["guardrail_id"], key["date"], key["team_id"], key["api_key"], key["usage_unit"]), + UsageUnitIncrement(int(counters["units"]), float(counters["cost"]), int(counters["untracked_units"])), + outer_transaction=transaction, + ) + windows: Final = tuple(_window(item) for item in contributions if item.table == "LiteLLM_BudgetWindowSpend") + await commit_window_spend_updates(prisma_client, windows, outer_transaction=transaction) + return contributions + + +def _window(item: Contribution) -> WindowSpendTransaction: + key: Final = dict(item.key) + values: Final = dict(item.values) + return WindowSpendTransaction( + entity_type=key["entity_type"], + entity_id=key["entity_id"], + window_duration=key["window_duration"], + window_start=values["window_start"] or "", + spend=float(dict(item.increments)["spend"]), + started_at=values["started_at"], + ) diff --git a/litellm/accounting/delivery_projection.py b/litellm/accounting/delivery_projection.py new file mode 100644 index 00000000000..da63542d5b8 --- /dev/null +++ b/litellm/accounting/delivery_projection.py @@ -0,0 +1,78 @@ +from collections.abc import Iterable +from itertools import chain, groupby +from typing import TYPE_CHECKING, Final +from urllib.parse import quote + +from litellm.accounting.projection import Contribution, project +from litellm.repositories.accounting_delivery_repository import DeliveryProjection + + +def _destination(contribution: Contribution) -> tuple[str, tuple[tuple[str, str], ...], str]: + return ( + contribution.table, + contribution.key, + (dict(contribution.values).get("window_start") or "") + if contribution.table == "LiteLLM_BudgetWindowSpend" + else "", + ) + + +def _sum(contributions: tuple[Contribution, ...]) -> Contribution: + first: Final = contributions[0] + increments: Final = tuple( + (column, sum(dict(item.increments).get(column, 0) for item in contributions)) + for column in dict.fromkeys( + column for column, _ in chain.from_iterable(item.increments for item in contributions) + ) + ) + starts: Final = tuple( + value + for column, value in chain.from_iterable(item.values for item in contributions) + if column == "started_at" and value + ) + values: Final = tuple( + (column, min(starts) if column == "started_at" and starts else value) for column, value in first.values + ) + return Contribution(first.table, first.key, increments, values) + + +def project_deliveries(deliveries: Iterable[DeliveryProjection]) -> tuple[Contribution, ...]: + ordered: Final = sorted(chain.from_iterable(project(delivery.row) for delivery in deliveries), key=_destination) + return tuple(_sum(tuple(group)) for _, group in groupby(ordered, key=_destination)) + + +def entity_transactions(contributions: tuple[Contribution, ...]) -> "DBSpendUpdateTransactions": + return { + "user_list_transactions": dict(_entity_totals(contributions, "LiteLLM_UserTable")), + "end_user_list_transactions": dict(_entity_totals(contributions, "LiteLLM_EndUserTable")), + "key_list_transactions": dict(_entity_totals(contributions, "LiteLLM_VerificationToken")), + "team_list_transactions": dict(_entity_totals(contributions, "LiteLLM_TeamTable")), + "team_member_list_transactions": dict(_entity_totals(contributions, "LiteLLM_TeamMembership")), + "org_list_transactions": dict(_entity_totals(contributions, "LiteLLM_OrganizationTable")), + "org_member_list_transactions": dict(_entity_totals(contributions, "LiteLLM_OrganizationMembership")), + "project_list_transactions": dict(_entity_totals(contributions, "LiteLLM_ProjectTable")), + "tag_list_transactions": dict(_entity_totals(contributions, "LiteLLM_TagTable")), + "agent_list_transactions": dict(_entity_totals(contributions, "LiteLLM_AgentsTable")), + "model_access_group_list_transactions": dict( + _entity_totals(contributions, "LiteLLM_ModelAccessGroupBudgetTable") + ), + } + + +def _entity_totals(contributions: tuple[Contribution, ...], table: str) -> tuple[tuple[str, float], ...]: + return tuple( + (_entity_key(item), float(dict(item.increments)["spend"])) for item in contributions if item.table == table + ) + + +def _entity_key(item: Contribution) -> str: + fields: Final = dict(item.key) + if item.table == "LiteLLM_TeamMembership": + return f"team_id::{fields['team_id']}::user_id::{fields['user_id']}" + if item.table == "LiteLLM_OrganizationMembership": + return f"org_id::{quote(fields['organization_id'], safe='')}::user_id::{quote(fields['user_id'], safe='')}" + return item.key[0][1] + + +if TYPE_CHECKING: + from litellm.proxy._types import DBSpendUpdateTransactions diff --git a/litellm/accounting/measurements.py b/litellm/accounting/measurements.py deleted file mode 100644 index 83011a280da..00000000000 --- a/litellm/accounting/measurements.py +++ /dev/null @@ -1,13 +0,0 @@ -from typing import Final - -from litellm.accounting.types import CostEvent, InboundUsageEvent, Measurement, OutboundUsageEvent - - -def measurements_for(event: InboundUsageEvent | OutboundUsageEvent | CostEvent) -> tuple[Measurement, ...]: - """Return the complete contribution set for one event, ready for one apply call.""" - metrics: Final = event.charges if isinstance(event, CostEvent) else event.metrics - occurred_at: Final = event.occurred_at if isinstance(event, CostEvent) else event.started_at - return tuple( - Measurement(source_event_id=event.event_id, occurred_at=occurred_at, dimensions=event.dimensions, metric=metric) - for metric in metrics or () - ) diff --git a/litellm/accounting/metrics.py b/litellm/accounting/metrics.py new file mode 100644 index 00000000000..ade715d5d14 --- /dev/null +++ b/litellm/accounting/metrics.py @@ -0,0 +1,30 @@ +from typing import Final + +from prometheus_client import Gauge, Histogram + +from litellm.repositories.accounting_delivery_repository import DeliveryMetrics + +_PENDING: Final = Gauge("litellm_accounting_pending", "Pending SpendLogs deliveries", multiprocess_mode="livemax") +_PARKED: Final = Gauge("litellm_accounting_parked", "Parked SpendLogs deliveries", multiprocess_mode="livemax") +_OLDEST: Final = Gauge( + "litellm_accounting_oldest_pending_seconds", + "Age of the oldest pending SpendLogs delivery", + multiprocess_mode="livemax", +) +_BATCH_TIME: Final = Histogram("litellm_accounting_batch_seconds", "SpendLogs delivery worker batch time") +_BATCH_COUNT: Final = Gauge( + "litellm_accounting_batch_deliveries", "Retained batch deliveries", multiprocess_mode="livemax" +) +_TABLE_BYTES: Final = Gauge( + "litellm_accounting_delivery_table_bytes", "Delivery table size in bytes", multiprocess_mode="livemax" +) + + +def record_deliveries(metrics: DeliveryMetrics, seconds: float) -> None: + _PENDING.set(metrics.pending_count) + _PARKED.set(metrics.parked_count) + _OLDEST.set(max(0.0, metrics.oldest_pending_seconds)) + _BATCH_TIME.observe(seconds) + + _BATCH_COUNT.set(metrics.batch_count) + _TABLE_BYTES.set(metrics.table_bytes) diff --git a/litellm/accounting/projection.py b/litellm/accounting/projection.py index fc3a89ae306..dd52b49f101 100644 --- a/litellm/accounting/projection.py +++ b/litellm/accounting/projection.py @@ -1,5 +1,5 @@ from collections.abc import Iterator, Mapping -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime, timezone from itertools import chain from types import MappingProxyType @@ -24,8 +24,11 @@ from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN _FIELDS: Final = TypeAdapter(dict[str, object]) _STRINGS: Final = TypeAdapter(tuple[str, ...]) +_VALUES: Final = TypeAdapter(tuple[object, ...]) _ROWS: Final = TypeAdapter(tuple[dict[str, object], ...]) _TASK_TYPES: Final = frozenset(load_model_insight_tasks()) +_MAX_INDEXED_TEXT_BYTES: Final = 2000 +_MAX_COUNTER: Final = 2**40 @dataclass(frozen=True, slots=True) @@ -54,7 +57,8 @@ def _tags(row: Mapping[str, object]) -> tuple[str, ...]: value: Final = row.get("request_tags") if value is None: return () - return _STRINGS.validate_json(value) if isinstance(value, str) else _STRINGS.validate_python(value) + parsed: Final = _VALUES.validate_json(value) if isinstance(value, str) else _VALUES.validate_python(value) + return tuple(tag for tag in parsed if isinstance(tag, str)) def _date(row: Mapping[str, object]) -> str | None: @@ -191,13 +195,31 @@ def _window_contributions(row: Mapping[str, object], metadata: Mapping[str, obje def project(row: Mapping[str, object]) -> tuple[Contribution, ...]: metadata: Final = _fields(row.get("metadata")) - return ( + contributions: Final = ( *_entity_contributions(row, metadata), *_daily_contributions(row, metadata), *_window_contributions(row, metadata), *_model_contributions(row, metadata), *_guardrail_contributions(row, metadata), ) + return tuple(_capped(contribution) for contribution in contributions if _indexable(contribution)) + + +def _indexable(contribution: Contribution) -> bool: + texts: Final = (value for _, value in (*contribution.key, *contribution.values) if value is not None) + return all(len(text.encode("utf-8", "surrogatepass")) <= _MAX_INDEXED_TEXT_BYTES for text in texts) + + +def _capped(contribution: Contribution) -> Contribution: + """Counters are summed into bigint columns in the charge's transaction, so one capped far below that range + can never overflow a day's total.""" + return replace( + contribution, + increments=tuple( + (name, max(-_MAX_COUNTER, min(_MAX_COUNTER, value)) if isinstance(value, int) else value) + for name, value in contribution.increments + ), + ) def _model_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: diff --git a/litellm/accounting/protocols.py b/litellm/accounting/protocols.py deleted file mode 100644 index f6ec61983a0..00000000000 --- a/litellm/accounting/protocols.py +++ /dev/null @@ -1,21 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping -from typing import Protocol - -from litellm.accounting.types import AggregateQuery, AggregateRow - - -class CostAggregator(Protocol): - async def apply(self, rows: tuple[Mapping[str, object], ...]) -> None: - """Apply stored request contributions atomically with their delivery receipts.""" - ... - - async def query(self, query: AggregateQuery) -> tuple[AggregateRow, ...]: - """Return disjoint [start, end) rows per group using existing daily projections. - - Filters combine with AND, values within a filter with IN. DAY emits daily - buckets; None emits whole-interval totals. Reject unsupported metric/scope - combinations. Different groups may overlap, including non-additive tags. - """ - ... diff --git a/litellm/accounting/snapshot.py b/litellm/accounting/snapshot.py index 51c7bbad81b..d5652f79ed0 100644 --- a/litellm/accounting/snapshot.py +++ b/litellm/accounting/snapshot.py @@ -2,12 +2,13 @@ from collections.abc import Callable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Final -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError +from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.safe_json_dumps import safe_json_structure from litellm.proxy._types import SpendLogsPayload from litellm.proxy.spend_tracking.compression_savings import extract_compression_saved_tokens -from litellm.proxy.spend_tracking.savings import compute_savings_spend, marks_gateway_injection +from litellm.proxy.spend_tracking.savings import SavingsSpend, compute_savings_spend, marks_gateway_injection if TYPE_CHECKING: from litellm.router import Router @@ -66,7 +67,47 @@ def build_accounting_payload( router: "Callable[[], Router | None]", ) -> Mapping[str, object]: metadata: Final = _FIELDS.validate_json(payload["metadata"]) - savings: Final = compute_savings_spend( + savings: Final = _savings_or_zero(payload, metadata, router) + accounting_metadata: Final = { + **{key: _accounting_metadata_field(key, value) for key, value in metadata.items() if key in _METADATA_FIELDS}, + "accounting_entities": entities, + "accounting_savings": savings._asdict(), + "accounting_model_access_groups": model_access_groups, + "accounting_windows": windows, + } + return _FIELDS.validate_python( + safe_json_structure( + { + **{ + key: value.isoformat() if isinstance(value, datetime) else value + for key, value in payload.items() + if key in _ACCOUNTING_FIELDS + }, + "metadata": accounting_metadata, + }, + non_finite_as_null=True, + ) + ) + + +def _savings_or_zero( + payload: SpendLogsPayload, metadata: Mapping[str, object], router: "Callable[[], Router | None]" +) -> SavingsSpend: + """Savings are analytics computed from caller-supplied counts, so one that cannot be computed is stored as zero + rather than losing the charge it rides with.""" + try: + return _savings(payload, metadata, router) + except (ArithmeticError, ValueError, TypeError) as error: + verbose_proxy_logger.warning( + "Recording zero savings for request_id=%s, savings could not be computed: %s", payload["request_id"], error + ) + return SavingsSpend(compression=0.0, prompt_caching=0.0) + + +def _savings( + payload: SpendLogsPayload, metadata: Mapping[str, object], router: "Callable[[], Router | None]" +) -> SavingsSpend: + return compute_savings_spend( model=payload["model"], custom_llm_provider=payload["custom_llm_provider"], compression_saved_tokens=extract_compression_saved_tokens(metadata), @@ -80,18 +121,41 @@ def build_accounting_payload( recorded_autorouter_savings_estimate=_mapping(metadata.get("autorouter_savings_estimate")), billed_at=payload["endTime"], ) - accounting_metadata: Final = { - **{key: value for key, value in metadata.items() if key in _METADATA_FIELDS}, - "accounting_entities": entities, - "accounting_savings": savings._asdict(), - "accounting_model_access_groups": model_access_groups, - "accounting_windows": windows, + + +_GUARDRAIL_ENTRIES: Final = TypeAdapter(tuple[dict[str, object], ...]) +_GUARDRAIL_METADATA_FIELDS: Final = frozenset({"guardrail_information", "standard_logging_guardrail_information"}) +_GUARDRAIL_ACCOUNTING_FIELDS: Final = frozenset( + { + "guardrail_id", + "guardrail_name", + "guardrail_provider", + "guardrail_usage", + "guardrail_cost_by_unit", + "guardrail_cost_in_spend", } +) +_GUARDRAIL_RESPONSE: Final = TypeAdapter(dict[str, object]) + + +def _accounting_guardrail_entries(value: object) -> object: + try: + entries: Final = _GUARDRAIL_ENTRIES.validate_python((value,) if isinstance(value, dict) else value) + except ValidationError: + return None + return [_accounting_guardrail_entry(entry) for entry in entries] + + +def _accounting_guardrail_entry(entry: Mapping[str, object]) -> Mapping[str, object]: + kept: Final = {key: item for key, item in entry.items() if key in _GUARDRAIL_ACCOUNTING_FIELDS} + response: Final = entry.get("guardrail_response") + if not isinstance(response, dict) or "tokens_saved" not in response: + return kept return { - **{ - key: value.isoformat() if isinstance(value, datetime) else value - for key, value in payload.items() - if key in _ACCOUNTING_FIELDS - }, - "metadata": _FIELDS.validate_python(safe_json_structure(accounting_metadata, non_finite_as_null=True)), + **kept, + "guardrail_response": {"tokens_saved": _GUARDRAIL_RESPONSE.validate_python(response)["tokens_saved"]}, } + + +def _accounting_metadata_field(key: str, value: object) -> object: + return _accounting_guardrail_entries(value) if key in _GUARDRAIL_METADATA_FIELDS else value diff --git a/litellm/accounting/types.py b/litellm/accounting/types.py deleted file mode 100644 index 2ee6f60f756..00000000000 --- a/litellm/accounting/types.py +++ /dev/null @@ -1,428 +0,0 @@ -"""Accounting values use UTC instants, decimal strings on the wire and USD charges.""" - -from collections.abc import Mapping -from datetime import datetime, timezone -from decimal import Decimal -from enum import Enum -from typing import Annotated, Final, Literal, TypeAlias - -from pydantic import AfterValidator, AwareDatetime, BaseModel, ConfigDict, Field, model_validator -from typing_extensions import Self - -from litellm.types.utils import CallTypes - -_CONFIG: Final = ConfigDict(frozen=True, extra="forbid", validate_default=True) -Identifier: TypeAlias = Annotated[str, Field(strict=True, min_length=1, pattern=r"\S")] -Quantity: TypeAlias = Annotated[Decimal, Field(strict=True, ge=0, allow_inf_nan=False)] - - -def _utc(value: datetime) -> datetime: - return value.astimezone(timezone.utc) - - -Instant: TypeAlias = Annotated[AwareDatetime, AfterValidator(_utc)] - - -class MetricName(str, Enum): - """Reasoning tokens are included in completion tokens; above-1hr writes are included in cache creation tokens.""" - - REQUESTS = "requests" - SUCCESSFUL_REQUESTS = "successful_requests" - FAILED_REQUESTS = "failed_requests" - PROMPT_TOKENS = "prompt_tokens" - COMPLETION_TOKENS = "completion_tokens" - REASONING_TOKENS = "reasoning_tokens" - TOTAL_TOKENS = "total_tokens" - CACHE_READ_INPUT_TOKENS = "cache_read_input_tokens" - CACHE_CREATION_INPUT_TOKENS = "cache_creation_input_tokens" - CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR = "cache_creation_input_tokens_above_1hr" - SPEND = "spend" - - -class DimensionName(str, Enum): - KEY = "key" - USER = "user" - TEAM = "team" - ORGANIZATION = "organization" - END_USER = "end_user" - MODEL = "model" - MODEL_GROUP = "model_group" - DEPLOYMENT = "deployment" - PROVIDER = "provider" - TOOL = "tool" - AGENT = "agent" - PROJECT = "project" - MODEL_ACCESS_GROUP = "model_access_group" - API_BASE = "api_base" - SERVICE_TIER = "service_tier" - TAG = "tag" - - -class Outcome(str, Enum): - SUCCEEDED = "succeeded" - FAILED = "failed" - CANCELLED = "cancelled" - - -class Granularity(str, Enum): - DAY = "day" - - -class EventKind(str, Enum): - INBOUND_USAGE = "inbound_usage" - OUTBOUND_USAGE = "outbound_usage" - COST = "cost" - - -class DeliveryState(str, Enum): - PENDING = "pending" - APPLIED = "applied" - PARKED = "parked" - - -class DeliveryTarget(str, Enum): - """One member per existing table the projector writes; each is applied and receipted independently.""" - - KEY_SPEND = "LiteLLM_VerificationToken" - USER_SPEND = "LiteLLM_UserTable" - TEAM_SPEND = "LiteLLM_TeamTable" - TEAM_MEMBER_SPEND = "LiteLLM_TeamMembership" - ORGANIZATION_SPEND = "LiteLLM_OrganizationTable" - ORGANIZATION_MEMBER_SPEND = "LiteLLM_OrganizationMembership" - END_USER_SPEND = "LiteLLM_EndUserTable" - AGENT_SPEND = "LiteLLM_AgentsTable" - TAG_SPEND = "LiteLLM_TagTable" - PROJECT_SPEND = "LiteLLM_ProjectTable" - MODEL_ACCESS_GROUP_SPEND = "LiteLLM_ModelAccessGroupBudgetTable" - BUDGET_SPEND = "LiteLLM_BudgetTable" - BUDGET_WINDOW_SPEND = "LiteLLM_BudgetWindowSpend" - DAILY_USER_SPEND = "LiteLLM_DailyUserSpend" - DAILY_TEAM_SPEND = "LiteLLM_DailyTeamSpend" - DAILY_ORGANIZATION_SPEND = "LiteLLM_DailyOrganizationSpend" - DAILY_END_USER_SPEND = "LiteLLM_DailyEndUserSpend" - DAILY_AGENT_SPEND = "LiteLLM_DailyAgentSpend" - DAILY_TAG_SPEND = "LiteLLM_DailyTagSpend" - DAILY_MODEL_USAGE = "LiteLLM_DailyModelUsage" - DAILY_TOOL_SPEND = "LiteLLM_DailyToolSpend" - AUTOROUTER_DAILY_SPEND = "LiteLLM_AutoRouterDailySpend" - - -_SUBSET_PAIRS: Final = ( - (MetricName.REASONING_TOKENS, MetricName.COMPLETION_TOKENS), - (MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, MetricName.CACHE_CREATION_INPUT_TOKENS), -) - - -def _check_subsets(values: Mapping[str, Decimal]) -> None: - for part, whole in _SUBSET_PAIRS: - if part.value in values and whole.value in values and values[part.value] > values[whole.value]: - raise ValueError(f"{part.value} cannot exceed {whole.value}") - - -class Metric(BaseModel): - model_config = _CONFIG - - name: MetricName - value: Quantity - unit: Literal["count", "USD"] = "count" - - @model_validator(mode="after") - def validate_unit(self) -> Self: - if (self.name is MetricName.SPEND) != (self.unit == "USD"): - raise ValueError("SPEND requires USD; usage metrics require count") - if self.unit == "count" and self.value != self.value.to_integral_value(): - raise ValueError("Count metrics require integral values") - return self - - -class Dimension(BaseModel): - model_config = _CONFIG - - name: DimensionName - value: Identifier - - -_REPEATABLE_DIMENSIONS: Final = frozenset({DimensionName.TAG, DimensionName.MODEL_ACCESS_GROUP}) - - -def _dimensions(values: tuple[Dimension, ...]) -> tuple[Dimension, ...]: - if len(frozenset(values)) != len(values): - raise ValueError("Duplicate dimension values") - single_names: Final = tuple(value.name for value in values if value.name not in _REPEATABLE_DIMENSIONS) - if len(frozenset(single_names)) != len(single_names): - raise ValueError("Only tag and model access group dimensions may have multiple values") - return values - - -Dimensions: TypeAlias = Annotated[tuple[Dimension, ...], AfterValidator(_dimensions)] -_CALLER_METRICS: Final = frozenset((MetricName.REQUESTS, MetricName.SUCCESSFUL_REQUESTS, MetricName.FAILED_REQUESTS)) - - -def _metrics(values: tuple[Metric, ...]) -> tuple[Metric, ...]: - if len(frozenset(value.name for value in values)) != len(values): - raise ValueError("Duplicate metric names") - return values - - -Metrics: TypeAlias = Annotated[tuple[Metric, ...], AfterValidator(_metrics)] - - -def _interval(start: datetime, end: datetime, *, allow_empty: bool = False) -> None: - if end < start or (end == start and not allow_empty): - raise ValueError("End must follow start") - - -class InboundUsageEvent(BaseModel): - model_config = _CONFIG - - event_id: Identifier - request_id: Identifier - call_type: CallTypes - started_at: Instant - ended_at: Instant - outcome: Outcome - dimensions: Dimensions - metrics: Metrics - - @model_validator(mode="after") - def validate_caller_activity(self) -> Self: - _interval(self.started_at, self.ended_at, allow_empty=True) - if any(metric.name not in _CALLER_METRICS for metric in self.metrics): - raise ValueError("Inbound events own only caller request metrics") - if not any(metric.name is MetricName.REQUESTS and metric.value == 1 for metric in self.metrics): - raise ValueError("Inbound events count exactly one caller request") - for metric in self.metrics: - if metric.name is MetricName.SUCCESSFUL_REQUESTS and metric.value != int(self.outcome is Outcome.SUCCEEDED): - raise ValueError("Successful request count must match outcome") - if metric.name is MetricName.FAILED_REQUESTS and metric.value != int(self.outcome is not Outcome.SUCCEEDED): - raise ValueError("Failed request count must match outcome") - return self - - -class OutboundUsageEvent(BaseModel): - model_config = _CONFIG - - event_id: Identifier - attempt_id: Identifier - request_id: Identifier | None - retry_of_attempt_id: Identifier | None - call_type: CallTypes - started_at: Instant - ended_at: Instant - outcome: Outcome - dimensions: Dimensions - metrics: Metrics | None - - @model_validator(mode="after") - def validate_attempt_usage(self) -> Self: - _interval(self.started_at, self.ended_at, allow_empty=True) - if self.retry_of_attempt_id == self.attempt_id: - raise ValueError("An attempt cannot retry itself") - if self.metrics is not None and any( - metric.name in _CALLER_METRICS or metric.name is MetricName.SPEND for metric in self.metrics - ): - raise ValueError("Outbound events own only attempt usage metrics") - if self.metrics is not None: - _check_subsets({metric.name.value: metric.value for metric in self.metrics}) - return self - - -PricingUnit: TypeAlias = Literal["token", "second", "character", "image", "request"] - - -class Rate(BaseModel): - model_config = _CONFIG - - measure: Identifier - unit: PricingUnit - usd_per_unit: Quantity - - -class PricedQuantity(BaseModel): - """A pricing input in its billing unit. Only seconds may be fractional.""" - - model_config = _CONFIG - - measure: Identifier - unit: PricingUnit - value: Quantity - - @model_validator(mode="after") - def validate_discrete_quantity(self) -> Self: - if self.unit != "second" and self.value != self.value.to_integral_value(): - raise ValueError("Discrete priced quantities require integral values") - return self - - -class Multiplier(BaseModel): - """A regional or tier factor applied on top of the priced key, e.g. data_residency=eu or vertex_location=us-east5.""" - - model_config = _CONFIG - - name: Identifier - value: Identifier - factor: Annotated[Decimal, Field(strict=True, gt=0, allow_inf_nan=False)] - - -class NamedCost(BaseModel): - """A free-form extra charge by name, e.g. azure_model_router_flat_cost.""" - - model_config = _CONFIG - - name: Identifier - amount: Quantity - - -class PricingSnapshot(BaseModel): - """Retained inputs and result of authoritative compilation, reproducible after cost-map rates change. - - Request-scoped inputs, overrides and effective rates come from the usage row's versioned payload. - Retries reuse the committed charge and snapshot. - Component names match ``CostBreakdown``: reasoning_cost is part of output_cost, not an additional charge. - Component sums and float tolerances belong to reconciliation, not this structural contract. - """ - - model_config = _CONFIG - - schema_version: Literal[1] = 1 - calculator_version: Identifier - model: Identifier - provider: Identifier | None - pricing_source: Literal["cost_map", "deployment", "per_request", "upstream_header"] - priced_model_key: Identifier - multipliers: tuple[Multiplier, ...] = () - rates: tuple[Rate, ...] = () - usage: tuple[PricedQuantity, ...] = () - input_cost: Quantity | None = None - output_cost: Quantity | None = None - cache_read_cost: Quantity | None = None - cache_creation_cost: Quantity | None = None - reasoning_cost: Quantity | None = None - tool_usage_cost: Quantity | None = None - guardrail_cost: Quantity | None = None - additional_costs: tuple[NamedCost, ...] = () - original_cost: Quantity | None = None - discount_percent: Quantity | None = None - discount_amount: Quantity | None = None - margin_percent: Quantity | None = None - margin_fixed_amount: Quantity | None = None - margin_total_amount: Quantity | None = None - upstream_reported_cost: Quantity | None = None - total_cost: Quantity - - @model_validator(mode="after") - def validate_snapshot(self) -> Self: - rate_keys: Final = frozenset((rate.measure, rate.unit) for rate in self.rates) - quantity_keys: Final = frozenset((quantity.measure, quantity.unit) for quantity in self.usage) - if len(rate_keys) != len(self.rates): - raise ValueError("Duplicate rate measures and units") - if len(quantity_keys) != len(self.usage): - raise ValueError("Duplicate priced quantity measures and units") - if not rate_keys.issubset(quantity_keys): - raise ValueError("Every rate requires one priced quantity with the same measure and unit") - if len(frozenset(cost.name for cost in self.additional_costs)) != len(self.additional_costs): - raise ValueError("Duplicate additional cost names") - if len(frozenset(multiplier.name for multiplier in self.multipliers)) != len(self.multipliers): - raise ValueError("Duplicate multiplier names") - _check_subsets({quantity.measure: quantity.value for quantity in self.usage if quantity.unit == "token"}) - return self - - -class CostEvent(BaseModel): - model_config = _CONFIG - - event_id: Identifier - usage_event_id: Identifier - occurred_at: Instant - dimensions: Dimensions - charges: Metrics - snapshot: PricingSnapshot - - @model_validator(mode="after") - def validate_charge(self) -> Self: - if len(self.charges) != 1 or self.charges[0].name is not MetricName.SPEND: - raise ValueError("A cost event requires one authoritative USD spend metric") - if self.event_id == self.usage_event_id: - raise ValueError("Cost and usage events require distinct identities") - if self.snapshot.total_cost != self.charges[0].value: - raise ValueError("Snapshot total must equal the authoritative spend") - return self - - -class Unpriced(BaseModel): - model_config = _CONFIG - - usage_event_id: Identifier - reason: Identifier - - -class Measurement(BaseModel): - model_config = _CONFIG - - source_event_id: Identifier - occurred_at: Instant - dimensions: Dimensions - metric: Metric - - -class DimensionFilter(BaseModel): - model_config = _CONFIG - - name: DimensionName - values: Annotated[tuple[Identifier, ...], Field(min_length=1)] - - -class AggregateQuery(BaseModel): - model_config = _CONFIG - - metric: MetricName - filters: tuple[DimensionFilter, ...] - group_by: tuple[DimensionName, ...] - start: Instant - end: Instant - granularity: Granularity | None - - @model_validator(mode="after") - def validate_query(self) -> Self: - _interval(self.start, self.end) - if len(frozenset(item.name for item in self.filters)) != len(self.filters): - raise ValueError("Use one IN filter per dimension; filters combine with AND") - if len(frozenset(self.group_by)) != len(self.group_by): - raise ValueError("Duplicate group-by dimensions") - if any( - (instant.hour, instant.minute, instant.second, instant.microsecond) != (0, 0, 0, 0) - for instant in (self.start, self.end) - ): - raise ValueError("Initial reporting queries require UTC day boundaries") - return self - - -class AggregateRow(BaseModel): - model_config = _CONFIG - - start: Instant - end: Instant - dimensions: Dimensions - value: Quantity - - @model_validator(mode="after") - def validate_bounds(self) -> Self: - _interval(self.start, self.end) - return self - - -def validate_aggregate_rows(rows: tuple[AggregateRow, ...]) -> tuple[AggregateRow, ...]: - """Validate [start, end) buckets per group; overlapping tag groups are not additive.""" - groups: Final = frozenset(frozenset(row.dimensions) for row in rows) - for group in groups: - ordered = sorted((row for row in rows if frozenset(row.dimensions) == group), key=lambda row: row.start) - if any(left.end > right.start for left, right in zip(ordered, ordered[1:])): - raise ValueError("Aggregate buckets overlap within a dimension group") - return rows - - -class AdmissionDecision(BaseModel): - model_config = _CONFIG - - allowed: bool - reason: Identifier | None = None diff --git a/litellm/accounting/worker.py b/litellm/accounting/worker.py new file mode 100644 index 00000000000..df71f2da31c --- /dev/null +++ b/litellm/accounting/worker.py @@ -0,0 +1,102 @@ +import os +import time +from typing import ( + TYPE_CHECKING, + Final, + Literal, + Protocol, + cast, # noqa: TID251 # generated Prisma transaction signatures cannot express the borrowed transaction protocol +) + +from litellm._logging import verbose_proxy_logger +from litellm.accounting.application import apply_deliveries +from litellm.accounting.metrics import record_deliveries +from litellm.accounting.projection import Contribution +from litellm.repositories.accounting_delivery_repository import ( + AccountingDeliveryRepository, + DeliveryProjection, + LedgerDatabase, + LedgerTransaction, +) +from litellm.types.integrations.slack_alerting import AlertType + +if TYPE_CHECKING: + from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + from litellm.proxy.utils import PrismaClient, ProxyLogging + + +class _AccountingAlerts(Protocol): + async def alerting_handler( + self, message: str, level: Literal["Low", "Medium", "High"], alert_type: AlertType + ) -> None: ... + + +_JOB: Final = "accounting_delivery" + + +async def flush_accounting_deliveries( + writer: "DBSpendUpdateWriter", prisma_client: "PrismaClient", proxy_logging_obj: "ProxyLogging" +) -> None: + acquired: Final = await writer.pod_lock_manager.acquire_lock(cronjob_id=_JOB, allow_reentrant=False) + if acquired is False: + return + applied: tuple[Contribution, ...] = () # rebind-ok: keep committed contributions for cache invalidation + + async def apply(transaction: LedgerTransaction, deliveries: tuple[DeliveryProjection, ...]) -> None: + nonlocal applied + applied = await apply_deliveries( # rebind-ok: invalidate caches only after the outer transaction commits + transaction, deliveries, prisma_client=prisma_client, writer=writer, proxy_logging_obj=proxy_logging_obj + ) + + try: + database: Final = cast( # cast-ok: PrismaClient.tx owns writer routing and yields the generated transaction + LedgerDatabase, prisma_client + ) + repository: Final = AccountingDeliveryRepository(database) + started: Final = time.monotonic() + batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "10000")) + + async def run_batch(*, retry_only: bool) -> bool: + result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only) + if result.applied: + await writer.invalidate_accounting_entities(applied, proxy_logging_obj) + if result.parked: + verbose_proxy_logger.error( + "Accounting delivery parked %d requests after isolated failures", len(result.parked) + ) + await _alert( + proxy_logging_obj, + message=f"Accounting delivery parked {len(result.parked)} requests. Included spend totals require repair.", + ) + return bool(result.applied or result.failed) + + await run_batch(retry_only=False) + retry_deadline: Final = time.monotonic() + 5 + for _ in range(batch_size): + if time.monotonic() >= retry_deadline or not await run_batch(retry_only=True): + break + await repository.cleanup_applied( + retention_days=int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")), + batch_size=batch_size * 2, + ) + metrics: Final = await repository.metrics() + record_deliveries(metrics, time.monotonic() - started) + if metrics.batch_count >= int( + os.environ.get("LITELLM_ACCOUNTING_BATCH_DELIVERY_ALERT_COUNT", "1000000") + ) or metrics.table_bytes >= int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1073741824")): + await _alert( + proxy_logging_obj, + message=f"Accounting delivery storage requires review. Retained batches={metrics.batch_count}, table bytes={metrics.table_bytes}.", + ) + if metrics.oldest_pending_seconds > 300: + await _alert( + proxy_logging_obj, + message=f"Accounting deliveries are delayed. Pending={metrics.pending_count}, oldest={metrics.oldest_pending_seconds:.0f}s.", + ) + finally: + if acquired: + await writer.pod_lock_manager.release_lock(cronjob_id=_JOB) + + +async def _alert(alerts: _AccountingAlerts, message: str) -> None: + await alerts.alerting_handler(message=message, level="High", alert_type=AlertType.failed_tracking_spend) diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index ece85b09432..c4f5c1efbcd 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -340,6 +340,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_AccessGroupTable", "LiteLLM_AccountingDelivery", "LiteLLM_AccountingEvent", + "LiteLLM_SpendLogDelivery", "LiteLLM_ClaudeCodePluginTable", "LiteLLM_ClaudeCodePluginUpload", "LiteLLM_MemoryTable", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9d84770b288..1864c3ae82e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2480,6 +2480,11 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): Documents all the fields supported by `general_settings` in config.yaml """ + use_spend_logs_ledger: bool = Field( + False, + description="Apply request accounting from stored SpendLogs deliveries. Enable only after all accounting writers are upgraded and ledger activation gates pass.", + ) + request_tag_mappings: Sequence[RequestTagMapping] | None = Field( None, description="Map W3C baggage values to request tags. Exact value/tag rules take precedence over tag_prefix fallbacks.", @@ -3800,6 +3805,9 @@ class SpendLogsMetadata(TypedDict): class SpendLogsPayload(TypedDict): + accounting_payload: NotRequired[Mapping[str, object]] # writable-ok: capture before buffering + accounting_log_disabled: NotRequired[bool] # writable-ok: capture logging mode before buffering + accounting_delivery_id: NotRequired[str] # writable-ok: assigned once before buffering request_id: str call_type: str api_key: str diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py index 99b49331932..d2b69f24397 100644 --- a/litellm/proxy/db/budget_window_spend_writer.py +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -14,11 +14,13 @@ interval of increments queued elsewhere: the same lag the SpendLogs aggregate it replaces (and every other spend column) already has. """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence +from contextlib import nullcontext from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Final, Protocol +from pydantic import TypeAdapter from typing_extensions import LiteralString, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -74,30 +76,44 @@ _OPEN_WINDOW_SPEND_SQL: Final = ( "WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')" ) +_HAS_DELIVERY: Final = 'EXISTS (SELECT 1 FROM "LiteLLM_SpendLogDelivery" delivery WHERE delivery.request_id = "LiteLLM_SpendLogs".request_id)' + +_SETTLED_SPEND_LOG_FILTER: Final = ( + ' AND NOT EXISTS (SELECT 1 FROM "LiteLLM_SpendLogDelivery" delivery ' + 'WHERE delivery.request_id = "LiteLLM_SpendLogs".request_id ' + "AND delivery.state IN ('pending', 'parked'))" +) + _SEED_FROM_SPEND_LOGS_KEY_SQL: Final = ( "SELECT COALESCE(SUM(spend), 0.0) AS total, " - "COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch " + "COALESCE(SUM(spend) FILTER (WHERE " + + _HAS_DELIVERY + + " OR \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) " + "AS before_batch " 'FROM "LiteLLM_SpendLogs" ' - "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + _SETTLED_SPEND_LOG_FILTER ) _SEED_FROM_SPEND_LOGS_TEAM_SQL: Final = ( "SELECT COALESCE(SUM(spend), 0.0) AS total, " - "COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch " + "COALESCE(SUM(spend) FILTER (WHERE " + + _HAS_DELIVERY + + " OR \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) " + "AS before_batch " 'FROM "LiteLLM_SpendLogs" ' - "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + _SETTLED_SPEND_LOG_FILTER ) _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL: Final = ( "SELECT COALESCE(SUM(spend), 0.0) AS total, COALESCE(SUM(spend), 0.0) AS before_batch " 'FROM "LiteLLM_SpendLogs" ' - "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + _SETTLED_SPEND_LOG_FILTER ) _SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL: Final = ( "SELECT COALESCE(SUM(spend), 0.0) AS total, COALESCE(SUM(spend), 0.0) AS before_batch " 'FROM "LiteLLM_SpendLogs" ' - "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + _SETTLED_SPEND_LOG_FILTER ) _UPSERT_TRANSACTION_TIMEOUT: Final = timedelta(seconds=60) @@ -114,6 +130,10 @@ class _ExistingKeyRow(TypedDict): window_duration: ReadOnly[str] +_SEED_TOTALS_ROWS: Final = TypeAdapter(tuple[_SeedTotalsRow, ...]) +_EXISTING_KEY_ROWS: Final = TypeAdapter(tuple[_ExistingKeyRow, ...]) + + class _SeedTotalsQuery(Protocol): async def __call__(self, query: LiteralString, *args: object) -> Sequence[_SeedTotalsRow]: ... @@ -135,6 +155,8 @@ class _UpsertBatchManager(Protocol): class _UpsertTransaction(Protocol): def batch_(self) -> _UpsertBatchManager: ... + async def query_raw(self, query: LiteralString, *args: object) -> Sequence[Mapping[str, object]]: ... + class _UpsertTransactionManager(Protocol): async def __aenter__(self) -> _UpsertTransaction: ... @@ -178,6 +200,7 @@ class WindowSpendLogsAggregate(Protocol): entity_id: str, window_start: datetime, batch_started_at: datetime | None, + query: _SeedTotalsQuery | None = None, ) -> WindowSeedTotals | None: ... @@ -187,6 +210,7 @@ async def spend_logs_seed_totals( entity_id: str, window_start: datetime, batch_started_at: datetime | None, + query: _SeedTotalsQuery | None = None, ) -> WindowSeedTotals | None: """LiteLLM_SpendLogs spend for one entity since window_start, both in full and up to the start of the batch being flushed, in one scan. @@ -210,11 +234,12 @@ async def spend_logs_seed_totals( bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_TEAM_SQL, _SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL else: return None + read: Final = query or _seed_totals_query(prisma_client) async with db_span("seed_window_spend", "LiteLLM_SpendLogs"): rows: Final = ( - await _seed_totals_query(prisma_client)(unbounded_sql, entity_id, window_start) + await read(unbounded_sql, entity_id, window_start) if batch_started_at is None - else await _seed_totals_query(prisma_client)( + else await read( bounded_sql, entity_id, window_start, @@ -246,9 +271,10 @@ def _primary_key(transaction: WindowSpendTransaction) -> tuple[str, str, str]: async def _existing_primary_keys( prisma_client: "PrismaClient", transactions: tuple[WindowSpendTransaction, ...], + query: _ExistingKeysQuery | None = None, ) -> frozenset[tuple[str, str, str]]: async with db_span("select_window_spend_rows", "LiteLLM_BudgetWindowSpend"): - rows: Final = await _existing_keys_query(prisma_client)( + rows: Final = await (query or _existing_keys_query(prisma_client))( _SELECT_EXISTING_ROWS_SQL, tuple(transaction["entity_type"] for transaction in transactions), tuple(transaction["entity_id"] for transaction in transactions), @@ -262,6 +288,7 @@ async def _seed_base_for_missing_row( transaction: WindowSpendTransaction, existing_primary_keys: frozenset[tuple[str, str, str]], spend_logs_aggregate: WindowSpendLogsAggregate, + query: _SeedTotalsQuery | None = None, ) -> float: """Spend already recorded for a window that has no row yet. @@ -272,12 +299,24 @@ async def _seed_base_for_missing_row( """ if _primary_key(transaction) in existing_primary_keys: return 0.0 - totals: 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), - batch_started_at=_transaction_started_at(transaction), + window_start: Final = datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc) + totals: Final = ( + await spend_logs_aggregate( + prisma_client=prisma_client, + entity_type=transaction["entity_type"], + entity_id=transaction["entity_id"], + window_start=window_start, + batch_started_at=_transaction_started_at(transaction), + ) + if query is None + else await spend_logs_aggregate( + prisma_client=prisma_client, + entity_type=transaction["entity_type"], + entity_id=transaction["entity_id"], + window_start=window_start, + batch_started_at=_transaction_started_at(transaction), + query=query, + ) ) if totals is None: return 0.0 @@ -324,10 +363,30 @@ def _upsert_params( ) +def _existing_keys_in(transaction: _UpsertTransaction) -> _ExistingKeysQuery: + async def query_raw(query: LiteralString, *args: object) -> Sequence[_ExistingKeyRow]: + return _EXISTING_KEY_ROWS.validate_python(await transaction.query_raw(query, *args)) + + return query_raw + + +def _seed_totals_in(transaction: _UpsertTransaction) -> _SeedTotalsQuery: + async def query_raw(query: LiteralString, *args: object) -> Sequence[_SeedTotalsRow]: + return _SEED_TOTALS_ROWS.validate_python(await transaction.query_raw(query, *args)) + + return query_raw + + +def _transaction_reads(transaction: _UpsertTransaction) -> tuple[_ExistingKeysQuery, _SeedTotalsQuery]: + """Reads made inside an enclosing transaction share its connection, so they cannot wait on it.""" + return _existing_keys_in(transaction), _seed_totals_in(transaction) + + async def commit_window_spend_updates( prisma_client: "PrismaClient", transactions: Sequence[WindowSpendTransaction], spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_seed_totals, + outer_transaction: _UpsertTransaction | None = None, ) -> None: """Apply aggregated window increments to LiteLLM_BudgetWindowSpend. @@ -343,9 +402,11 @@ async def commit_window_spend_updates( return ordered: Final = tuple(sorted(transactions, key=window_spend_group_key)) + reads: Final = _transaction_reads(outer_transaction) if outer_transaction is not None else (None, None) existing_primary_keys: Final = await _existing_primary_keys( prisma_client=prisma_client, transactions=ordered, + query=reads[0], ) seed_bases: Final = tuple( [ @@ -354,6 +415,7 @@ async def commit_window_spend_updates( transaction=transaction, existing_primary_keys=existing_primary_keys, spend_logs_aggregate=spend_logs_aggregate, + query=reads[1], ) for transaction in ordered ] @@ -367,7 +429,9 @@ async def commit_window_spend_updates( ) async with ( db_span("commit_window_spend_updates", "LiteLLM_BudgetWindowSpend"), - _upsert_transaction(prisma_client) as db_transaction, + ( + nullcontext(outer_transaction) if outer_transaction is not None else _upsert_transaction(prisma_client) + ) as db_transaction, db_transaction.batch_() as batcher, ): for transaction, seed_base in zip(ordered, seed_bases): diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index eb130a5196f..2d1770f4c1d 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -55,7 +55,7 @@ _ENTITY_INPUT_KEYS: Final[Mapping[DailySpendEntity, str]] = MappingProxyType( "tag": "request_tags", } ) -_TAGS: Final = TypeAdapter(tuple[str, ...]) +_TAGS: Final = TypeAdapter(tuple[object, ...]) def daily_spend_entity_ids(payload: Mapping[str, object], entity: DailySpendEntity) -> tuple[str | None, ...]: @@ -67,7 +67,7 @@ def daily_spend_entity_ids(payload: Mapping[str, object], entity: DailySpendEnti if value is None: return () tags: Final = _TAGS.validate_json(value) if isinstance(value, str) else _TAGS.validate_python(value) - return tuple(dict.fromkeys(tags)) + return tuple(dict.fromkeys(tag for tag in tags if isinstance(tag, str))) if value is None: return (None,) if entity == "user" else () if not isinstance(value, str) or (entity == "end_user" and not value): diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 35e7ae1c579..80f09df6c86 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,6 +12,7 @@ import json import os import random import time +import uuid from collections.abc import AsyncGenerator, Callable, Coroutine, Iterable, Mapping, Sequence from contextlib import asynccontextmanager from contextvars import ContextVar @@ -96,6 +97,7 @@ from litellm.types.utils import CallTypes from litellm.types.videos.utils import VIDEO_COST_REQUEST_ID_SUFFIX, VIDEO_STATUS_CALL_TYPES if TYPE_CHECKING: + from litellm.accounting.projection import Contribution from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction from litellm.proxy.db.baseline_accounting import DailyBaselineAttribution from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -160,6 +162,18 @@ def _org_member_transaction_key(org_id: str, user_id: str) -> str: return f"organization_id::{quote(org_id, safe='')}::user_id::{quote(user_id, safe='')}" +ACCOUNTING_LEDGER_KWARG: Final = "accounting_ledger" + + +def _persistable_api_key(token: str | None, logged_api_key: str) -> str | None: + from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash + from litellm.proxy.spend_tracking.spend_tracking_utils import redact_logged_api_key + + if token is None or token == logged_api_key or is_valid_sha256_hash(token): + return token + return redact_logged_api_key(token) + + async def _track_guardrail_usage_of_cost_row( payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool ) -> None: @@ -200,12 +214,20 @@ _COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "star def _cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool) -> Mapping[str, object]: """Reduce a job's cost row to what tells its polls apart when logging is off. - A proxy run with spend logs disabled still needs one row per job to charge it once, + A legacy proxy run with spend logs disabled still needs one row per job to charge it once, so the row is written either way, but it carries no request of its own: no metadata, - no requester IP, no key, model, or token counts (LIT-7048). + no requester IP, no key, model, or token counts (LIT-7048). A ledger batch is deduplicated + by its delivery row instead, so it writes no local row when logging is off. """ if disable_spend_logs is False: return payload + if "accounting_payload" in payload: + return { + "request_id": payload["request_id"], + "call_type": payload["call_type"], + "accounting_payload": payload["accounting_payload"], + "accounting_log_disabled": True, + } return MappingProxyType({field: value for field, value in payload.items() if field in _COST_CLAIM_FIELDS}) @@ -219,8 +241,17 @@ class _MemberSpendRow(TypedDict): cost: ReadOnly[float] +class SpendUpsertTable(Protocol): + def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ... + + class _SpendBatch(Protocol): litellm_usertable: BatchTable + litellm_endusertable: SpendUpsertTable + litellm_dailymodelusage: SpendUpsertTable + + def execute_raw(self, query: LiteralString, *args: object) -> None: ... + litellm_verificationtoken: BatchTable litellm_teamtable: BatchTable litellm_teammembership: BatchTable @@ -251,14 +282,14 @@ def _entity_spend_table(batcher: _SpendBatch, table_accessor: _EntitySpendTable) return _ENTITY_SPEND_TABLES[table_accessor](batcher) -class _SpendBatchManager(Protocol): +class SpendBatchManager(Protocol): async def __aenter__(self) -> _SpendBatch: ... async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... class _SpendTransaction(Protocol): - def batch_(self) -> _SpendBatchManager: ... + def batch_(self) -> SpendBatchManager: ... async def execute_raw(self, query: LiteralString, *args: object) -> int: ... @@ -413,8 +444,14 @@ _ENTITY_SPEND_MODELS: Final[Mapping[_EntitySpendTable, str]] = MappingProxyType( @asynccontextmanager async def _spend_update_tx( - prisma_client: PrismaClient, table: str, call_type: str = "commit_spend_updates" + prisma_client: PrismaClient, + table: str, + call_type: str = "commit_spend_updates", + outer_transaction: _SpendTransaction | None = None, ) -> AsyncGenerator[_SpendTransaction]: + if outer_transaction is not None: + yield outer_transaction + return tx: Final[_SpendTransactionManager] = prisma_client.db.tx(timeout=timedelta(seconds=60)) async with db_span(call_type, table), tx as transaction: await apply_rollup_lock_timeout(transaction) @@ -601,6 +638,21 @@ def debitable_model_access_groups( return tuple(group for group in ordered if group in served) +def _billable_cost(response_cost: float | None, payload: SpendLogsPayload) -> float | None: + billable_cost: Final = finite_cost(response_cost) + if response_cost is not None and billable_cost is None: + verbose_proxy_logger.warning( + "Spend tracking - recording a non-finite response_cost=%s as 0 for request_id=%s", + response_cost, + single_log_line(payload["request_id"]), + ) + return billable_cost + + +def _iso_timestamp(value: datetime | str) -> str: + return value.isoformat() if isinstance(value, datetime) else value + + class DBSpendUpdateWriter: """ Module responsible for @@ -649,6 +701,7 @@ class DBSpendUpdateWriter: so the caller leaves the key, team, and user counters alone (LIT-7048). """ from litellm.proxy.proxy_server import ( + accounting_window_keys, disable_spend_logs, general_settings, # pyright: ignore[reportUnknownVariableType] # legacy proxy configuration is narrowed below litellm_proxy_budget_name, @@ -656,6 +709,13 @@ class DBSpendUpdateWriter: ) from litellm.proxy.utils import ProxyUpdateSpend, hash_token + proxy_settings: Final = cast( # cast-ok: proxy settings use string keys + Mapping[str, object], general_settings + ) + use_ledger: Final = proxy_settings.get("use_spend_logs_ledger") is True + request_fields: Final = _SPEND_METADATA_ADAPTER.validate_python(kwargs or {}) + if kwargs is not None: + kwargs[ACCOUNTING_LEDGER_KWARG] = use_ledger # rebind-ok: spend counters leave ledger windows to the worker try: verbose_proxy_logger.debug( "Enters prisma db call, response_cost: %s, token: %s; user_id: %s; team_id: %s", @@ -666,10 +726,9 @@ class DBSpendUpdateWriter: ) if ProxyUpdateSpend.disable_spend_updates() is True: return True - if token is not None and isinstance(token, str) and token.startswith("sk-"): - hashed_token = hash_token(token=token) - else: - hashed_token = token + hashed_token: Final = ( + hash_token(token=token) if isinstance(token, str) and token.startswith("sk-") else token + ) ## CREATE SPEND LOG PAYLOAD ## from litellm.proxy.spend_tracking.spend_tracking_utils import ( @@ -685,32 +744,42 @@ class DBSpendUpdateWriter: end_time=end_time, llm_router=llm_router, ) - billable_cost: Final = finite_cost(response_cost) - if response_cost is not None and billable_cost is None: - verbose_proxy_logger.warning( - "Spend tracking - recording a non-finite response_cost=%s as 0 for request_id=%s", - response_cost, - single_log_line(payload["request_id"]), - ) + billable_cost: Final = _billable_cost(response_cost, payload) payload["spend"] = billable_cost or 0.0 - if isinstance(payload["startTime"], datetime): - payload["startTime"] = payload["startTime"].isoformat() - if isinstance(payload["endTime"], datetime): - payload["endTime"] = payload["endTime"].isoformat() - - if org_id is not None and org_id != "": + payload["startTime"] = _iso_timestamp(payload["startTime"]) + payload["endTime"] = _iso_timestamp(payload["endTime"]) + if org_id: payload["organization_id"] = org_id - - if team_id is not None and team_id != "": + if team_id: payload["team_id"] = team_id - error_log_settings: Final = cast( # cast-ok: proxy settings use string keys - Mapping[str, object], general_settings - ) + if use_ledger: + from litellm.accounting.snapshot import build_accounting_payload + + if not _is_settle_once_cost_row(payload): + payload["accounting_delivery_id"] = uuid.uuid4().hex + payload["accounting_payload"] = build_accounting_payload( + payload, + entities={ + "api_key": _persistable_api_key(hashed_token, payload.get("api_key") or ""), + "user": user_id, + "end_user": end_user_id, + "team_id": team_id, + "organization_id": org_id, + "project_id": project_id, + "agent_id": payload["agent_id"], + "proxy_budget_name": litellm_proxy_budget_name if litellm.max_budget > 0 else None, + }, + model_access_groups=debitable_model_access_groups( + get_request_model_access_groups(request_fields), payload["model_id"], llm_router + ), + windows=await accounting_window_keys(hashed_token, team_id) if billable_cost else (), + router=get_llm_router, + ) keep_spend_logs: Final = disable_spend_logs is False and not ( "status" in payload and payload["status"] == "failure" - and error_log_settings.get("disable_error_logs") is True + and proxy_settings.get("disable_error_logs") is True ) if not await self._record_spend_log( payload=payload, @@ -739,24 +808,26 @@ class DBSpendUpdateWriter: "disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur." ) - spawn_logging_task( - asyncio.get_running_loop(), - self._batch_database_updates( - response_cost=billable_cost, - user_id=user_id, - hashed_token=hashed_token, - team_id=team_id, - org_id=org_id, - project_id=project_id, - end_user_id=end_user_id, - prisma_client=prisma_client, - litellm_proxy_budget_name=litellm_proxy_budget_name, - payload=payload, - request_model_access_groups=get_request_model_access_groups(kwargs), - ), - task_name="spend rollups", - ) - + if not use_ledger: + spawn_logging_task( + asyncio.get_running_loop(), + self._batch_database_updates( + response_cost=billable_cost, + user_id=user_id, + hashed_token=hashed_token, + team_id=team_id, + org_id=org_id, + project_id=project_id, + end_user_id=end_user_id, + prisma_client=prisma_client, + litellm_proxy_budget_name=litellm_proxy_budget_name, + payload=payload, + request_model_access_groups=get_request_model_access_groups(request_fields), + ), + task_name="spend rollups", + ) + else: + self._record_request_error_safely(payload=payload, prisma_client=prisma_client) self._enqueue_tool_registry_upsert( kwargs=kwargs, completion_response=completion_response, @@ -791,7 +862,13 @@ class DBSpendUpdateWriter: return await self._claim_settle_once_cost_row( payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs ) - if disable_spend_logs is False: + if "accounting_payload" in payload: + await self._insert_spend_log_to_db( + payload={**payload, "accounting_log_disabled": disable_spend_logs}, + prisma_client=prisma_client, + llm_router=llm_router, + ) + elif disable_spend_logs is False: await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client, llm_router=llm_router) return True @@ -804,7 +881,9 @@ class DBSpendUpdateWriter: the charge and every later one finds the row and charges nothing (LIT-7048). Only a row that recorded a charge counts: a failed poll, a request whose client picked the job id as its call id, and the $0 row a poll left behind - while the job was still running all leave the charge to be made. + while the job was still running all leave the charge to be made. A ledger job + is owned by its delivery row, keyed by its cost row id, in every logging mode, and a + $0 poll records no delivery so it cannot block the poll that carries the cost. """ from litellm.repositories.table_repositories import SpendLogsRepository @@ -817,13 +896,24 @@ class DBSpendUpdateWriter: row: Final = spend_log_row_with_retained_metadata( _cost_row_to_write(payload, disable_spend_logs), configured_spend_logs_metadata_fields() ) - spend_logs: Final = SpendLogsRepository(prisma_client).table + repository: Final = SpendLogsRepository(prisma_client) try: - claimed: Final = await spend_logs.create_many( + if "accounting_payload" in payload: + if payload["spend"] <= 0: + return True + accepted: Final = await repository.insert_with_deliveries([prisma_client.jsonify_object(row)]) == 1 + if accepted: + await _track_guardrail_usage_of_cost_row(payload, prisma_client, disable_spend_logs) + return accepted + if await repository.has_delivery(request_id): + return False + claimed: Final = await repository.table.create_many( data=[prisma_client.jsonify_object(row)], skip_duplicates=True, ) - existing: Final = None if claimed == 1 else await spend_logs.find_unique(where={"request_id": request_id}) + existing: Final = ( + None if claimed == 1 else await repository.table.find_unique(where={"request_id": request_id}) + ) except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreachable DB queues the row like any other spend log verbose_proxy_logger.warning( "Could not claim spend row %s for a job's cost, queueing it: %s", request_id, e @@ -1268,10 +1358,7 @@ class DBSpendUpdateWriter: self.add_spend_log_transaction_to_daily_team_transaction(payload=payload_copy, prisma_client=prisma_client), payload=payload_copy, ) - try: - self._record_request_error(payload=payload_copy, prisma_client=prisma_client) - except Exception as e: # noqa: BLE001 # the rollup must never skip the sibling spend writes - _report_spend_update_step_failure("_record_request_error", payload=payload_copy, exc=e) + self._record_request_error_safely(payload=payload_copy, prisma_client=prisma_client) await self._run_spend_update_step( "add_spend_log_transaction_to_daily_org_transaction", self.add_spend_log_transaction_to_daily_org_transaction( @@ -1601,7 +1688,7 @@ class DBSpendUpdateWriter: async def _insert_spend_log_to_db( self, - payload: dict | SpendLogsPayload, + payload: Mapping[str, object], prisma_client: PrismaClient | None = None, spend_logs_url: str | None = os.getenv("SPEND_LOGS_URL"), llm_router: "Router | None" = None, @@ -1645,6 +1732,13 @@ class DBSpendUpdateWriter: else: - Regular flow of this method """ + from litellm.accounting.worker import flush_accounting_deliveries + + try: + await flush_accounting_deliveries(self, prisma_client, proxy_logging_obj) + except Exception: # noqa: BLE001 # pending deliveries survive outages while legacy queues must still drain + verbose_proxy_logger.exception("Accounting delivery flush failed; pending work will retry") + if RedisUpdateBuffer.should_commit_spend_updates_to_redis(): await self._commit_spend_updates_to_db_with_redis( prisma_client=prisma_client, @@ -2171,6 +2265,39 @@ class DBSpendUpdateWriter: _handle_spend_update_failure = handle_spend_update_failure + @staticmethod + async def invalidate_accounting_entities( + contributions: tuple["Contribution", ...], proxy_logging_obj: ProxyLogging + ) -> None: + await DBSpendUpdateWriter._invalidate_project_caches( + tuple(dict(item.key)["project_id"] for item in contributions if item.table == "LiteLLM_ProjectTable"), + proxy_logging_obj, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + details: Final = _SPEND_METADATA_ADAPTER.validate_python( + proxy_logging_obj.call_details # pyright: ignore[reportUnknownMemberType] # legacy logging metadata is validated at this boundary + ) + user_api_key_cache: Final = details.get("user_api_key_cache") + if isinstance(user_api_key_cache, UserApiKeyCache): + with service_target(AUTH_OBJECTS_TARGET): + for item in contributions: + if item.table == "LiteLLM_TeamMembership": + await user_api_key_cache.async_delete_cache( + key=f"team_membership:{dict(item.key)['user_id']}:{dict(item.key)['team_id']}" + ) + + async def apply_accounting_entities( + self, + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, + transaction: _SpendTransaction, + updates: DBSpendUpdateTransactions, + ) -> None: + await self._commit_spend_updates_to_db( + prisma_client, 0, proxy_logging_obj, updates, outer_transaction=transaction + ) + async def _commit_spend_updates_to_db( self, prisma_client: PrismaClient, @@ -2178,6 +2305,7 @@ class DBSpendUpdateWriter: proxy_logging_obj: ProxyLogging, db_spend_update_transactions: DBSpendUpdateTransactions, on_table_committed: Callable[[_SpendTableName], None] | None = None, + outer_transaction: _SpendTransaction | None = None, ): """ Commits all the spend `UPDATE` transactions to the Database @@ -2192,7 +2320,9 @@ class DBSpendUpdateWriter: for i in range(n_retry_times + 1): start_time = time.time() try: - async with _spend_update_tx(prisma_client, "LiteLLM_UserTable") as transaction: + async with _spend_update_tx( + prisma_client, "LiteLLM_UserTable", outer_transaction=outer_transaction + ) as transaction: async with transaction.batch_() as batcher: # Sort by ID for consistent lock ordering across pods to prevent deadlocks. # batch_() issues statements sequentially within the tx, so iteration @@ -2204,6 +2334,8 @@ class DBSpendUpdateWriter: ) break except Exception as e: + if outer_transaction is not None: + raise await self.handle_spend_update_failure( e=e, attempt=i, @@ -2223,6 +2355,7 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, end_user_list_transactions=end_user_list_transactions, + outer_transaction=outer_transaction, ) if on_table_committed is not None: on_table_committed("end_user_list_transactions") @@ -2233,7 +2366,9 @@ class DBSpendUpdateWriter: for i in range(n_retry_times + 1): start_time = time.time() try: - async with _spend_update_tx(prisma_client, "LiteLLM_VerificationToken") as transaction: + async with _spend_update_tx( + prisma_client, "LiteLLM_VerificationToken", outer_transaction=outer_transaction + ) as transaction: async with transaction.batch_() as batcher: # Sort by token for consistent lock ordering across pods to prevent deadlocks. for token, response_cost in sorted(key_list_transactions.items()): @@ -2248,6 +2383,8 @@ class DBSpendUpdateWriter: ) break except Exception as e: + if outer_transaction is not None: + raise await self.handle_spend_update_failure( e=e, attempt=i, @@ -2265,7 +2402,9 @@ class DBSpendUpdateWriter: for i in range(n_retry_times + 1): start_time = time.time() try: - async with _spend_update_tx(prisma_client, "LiteLLM_TeamTable") as transaction: + async with _spend_update_tx( + prisma_client, "LiteLLM_TeamTable", outer_transaction=outer_transaction + ) as transaction: async with transaction.batch_() as batcher: # Sort by team_id for consistent lock ordering across pods to prevent deadlocks. for team_id, response_cost in sorted(team_list_transactions.items()): @@ -2278,6 +2417,8 @@ class DBSpendUpdateWriter: ) break except Exception as e: + if outer_transaction is not None: + raise await self.handle_spend_update_failure( e=e, attempt=i, @@ -2303,11 +2444,15 @@ class DBSpendUpdateWriter: for i in range(n_retry_times + 1): start_time = time.time() try: - async with _spend_update_tx(prisma_client, "LiteLLM_TeamMembership") as transaction: + async with _spend_update_tx( + prisma_client, "LiteLLM_TeamMembership", outer_transaction=outer_transaction + ) as transaction: await _write_team_member_spend(transaction, team_member_list_transactions) # Transaction succeeded, break out of retry loop break except Exception as e: + if outer_transaction is not None: + raise await self.handle_spend_update_failure( e=e, attempt=i, @@ -2320,7 +2465,7 @@ class DBSpendUpdateWriter: # Invalidate cache for updated team memberships # This ensures budget checks read fresh spend data from the database - if team_memberships_to_invalidate and proxy_logging_obj is not None: + if outer_transaction is None and team_memberships_to_invalidate and proxy_logging_obj is not None: user_api_key_cache: Final = proxy_logging_obj.call_details.get("user_api_key_cache") if user_api_key_cache is not None: with service_target(AUTH_OBJECTS_TARGET): @@ -2340,7 +2485,9 @@ class DBSpendUpdateWriter: for i in range(n_retry_times + 1): start_time = time.time() try: - async with _spend_update_tx(prisma_client, "LiteLLM_OrganizationTable") as transaction: + async with _spend_update_tx( + prisma_client, "LiteLLM_OrganizationTable", outer_transaction=outer_transaction + ) as transaction: async with transaction.batch_() as batcher: # Sort by org_id for consistent lock ordering across pods to prevent deadlocks. for org_id, response_cost in sorted(org_list_transactions.items()): @@ -2350,6 +2497,8 @@ class DBSpendUpdateWriter: ) break except Exception as e: + if outer_transaction is not None: + raise await self.handle_spend_update_failure( e=e, attempt=i, @@ -2367,7 +2516,9 @@ class DBSpendUpdateWriter: start_time = time.time() try: async with ( - _spend_update_tx(prisma_client, "LiteLLM_OrganizationMembership") as transaction, + _spend_update_tx( + prisma_client, "LiteLLM_OrganizationMembership", outer_transaction=outer_transaction + ) as transaction, transaction.batch_() as batcher, ): for key, response_cost in sorted(org_member_list_transactions.items()): @@ -2378,6 +2529,8 @@ class DBSpendUpdateWriter: ) break except Exception as e: + if outer_transaction is not None: + raise await self.handle_spend_update_failure( e=e, attempt=i, @@ -2393,6 +2546,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_entity_spend_in_db( entity_name="Project", transactions=project_list_transactions, + outer_transaction=outer_transaction, table_accessor="litellm_projecttable", where_field="project_id", n_retry_times=n_retry_times, @@ -2401,16 +2555,18 @@ class DBSpendUpdateWriter: ) if on_table_committed is not None: on_table_committed("project_list_transactions") - await DBSpendUpdateWriter._invalidate_project_caches( - project_ids=tuple(project_list_transactions or ()), - proxy_logging_obj=proxy_logging_obj, - ) + if outer_transaction is None: + await DBSpendUpdateWriter._invalidate_project_caches( + project_ids=tuple(project_list_transactions or ()), + proxy_logging_obj=proxy_logging_obj, + ) ### UPDATE TAG TABLE ### tag_list_transactions: Final = db_spend_update_transactions["tag_list_transactions"] await DBSpendUpdateWriter._update_entity_spend_in_db( entity_name="Tag", transactions=tag_list_transactions, + outer_transaction=outer_transaction, table_accessor="litellm_tagtable", where_field="tag_name", n_retry_times=n_retry_times, @@ -2427,6 +2583,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_entity_spend_in_db( entity_name="Model access group", transactions=model_access_group_list_transactions, + outer_transaction=outer_transaction, table_accessor="litellm_modelaccessgroupbudgettable", where_field="access_group_name", n_retry_times=n_retry_times, @@ -2441,6 +2598,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_entity_spend_in_db( entity_name="Agent", transactions=agent_list_transactions, + outer_transaction=outer_transaction, table_accessor="litellm_agentstable", where_field="agent_id", n_retry_times=n_retry_times, @@ -2470,6 +2628,7 @@ class DBSpendUpdateWriter: n_retry_times: int, prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging, + outer_transaction: _SpendTransaction | None = None, ): """ Helper function to update spend for any entity type (team, org, tag, etc). @@ -2482,13 +2641,16 @@ class DBSpendUpdateWriter: n_retry_times: Number of retries on failure prisma_client: Prisma client instance proxy_logging_obj: Proxy logging object + outer_transaction: Existing ledger transaction owning the complete delivery update """ verbose_proxy_logger.debug("%s Spend transactions: %s", entity_name, transactions) if transactions is not None and len(transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: - async with _spend_update_tx(prisma_client, _ENTITY_SPEND_MODELS[table_accessor]) as transaction: + async with _spend_update_tx( + prisma_client, _ENTITY_SPEND_MODELS[table_accessor], outer_transaction=outer_transaction + ) as transaction: async with transaction.batch_() as batcher: # Sort by entity_id for consistent lock ordering across pods to prevent deadlocks. for entity_id, response_cost in sorted(transactions.items()): @@ -2505,6 +2667,8 @@ class DBSpendUpdateWriter: ) break except Exception as e: + if outer_transaction is not None: + raise await DBSpendUpdateWriter.handle_spend_update_failure( e=e, attempt=i, @@ -2838,6 +3002,12 @@ class DBSpendUpdateWriter: entity_id_field="tag", ) + def _record_request_error_safely(self, *, payload: SpendLogsPayload, prisma_client: PrismaClient | None) -> None: + try: + self._record_request_error(payload=payload, prisma_client=prisma_client) + except Exception as error: # noqa: BLE001 # error rollup failure must not skip sibling writes + _report_spend_update_step_failure("_record_request_error", payload=payload, exc=error) + @staticmethod def _record_request_error(*, payload: SpendLogsPayload, prisma_client: PrismaClient | None) -> None: if prisma_client is None: diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index dbf636e1630..c2d58512fea 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -53,6 +53,7 @@ _PRISMA_POOL_TIMEOUT_CODE: Final = "P2024" _LOCK_TIMEOUT_SQLSTATE: Final = "55P03" _LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout" _STATEMENT_TIMEOUT_SQLSTATE: Final = "57014" +_REJECTED_ROW_SQLSTATE_CLASSES: Final = frozenset({"22", "23", "54"}) _STATEMENT_TIMEOUT_PHRASE: Final = "canceling statement due to statement timeout" @@ -221,6 +222,20 @@ class PrismaDBExceptionHandler: return type(e) is prisma.errors.DataError + @staticmethod + def is_rejected_row_error(e: Exception) -> bool: + """True iff Postgres refused the rows themselves: a base ``DataError``, or a raw query failing with a data + exception (SQLSTATE class 22), integrity violation (class 23) or a value too large for an index (class 54). Systemic raw failures such as a missing + table (42P01) or a deadlock (40P01) stay False so the batch is retried instead of dropped.""" + import prisma + + if PrismaDBExceptionHandler.is_prisma_data_error(e): + return True + if not isinstance(e, _exception_types(prisma.errors.RawQueryError)): + return False + sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) + return sqlstate is not None and sqlstate[:2] in _REJECTED_ROW_SQLSTATE_CLASSES + @staticmethod def is_database_transport_error(e: Exception) -> bool: """ diff --git a/litellm/proxy/db/model_usage_rollup.py b/litellm/proxy/db/model_usage_rollup.py index e6fbc299504..f8cdabe9296 100644 --- a/litellm/proxy/db/model_usage_rollup.py +++ b/litellm/proxy/db/model_usage_rollup.py @@ -32,14 +32,15 @@ class _UpsertTable(Protocol): def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ... -class _ModelUsageBatch(Protocol): - litellm_dailymodelusage: _UpsertTable +class ModelUsageBatch(Protocol): + @property + def litellm_dailymodelusage(self) -> _UpsertTable: ... def execute_raw(self, query: LiteralString, *args: object) -> None: ... -class _ModelUsageBatchManager(Protocol): - async def __aenter__(self) -> _ModelUsageBatch: ... +class ModelUsageBatchManager(Protocol): + async def __aenter__(self) -> ModelUsageBatch: ... async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... @@ -116,8 +117,8 @@ def build_model_usage_transaction(payload: SpendLogsPayload) -> ModelUsageTransa ) -def _model_usage_batch(prisma_client: PrismaClient) -> _ModelUsageBatchManager: - batch: Final[_ModelUsageBatchManager] = prisma_client.db.batch_() +def _model_usage_batch(prisma_client: PrismaClient) -> ModelUsageBatchManager: + batch: Final[ModelUsageBatchManager] = prisma_client.db.batch_() return batch @@ -156,26 +157,16 @@ async def flush_model_usage_transactions( "custom_llm_provider": key.custom_llm_provider, "task_type": key.task_type, } - batcher.litellm_dailymodelusage.upsert( - where={"date_model_group_model_custom_llm_provider_task_type": key_fields}, - data={ - "create": { - **key_fields, - "spend": spend, - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "request_count": len(requests), - "successful_requests": successful, - "failed_requests": failed, - }, - "update": { - "spend": {"increment": spend}, - "prompt_tokens": {"increment": prompt_tokens}, - "completion_tokens": {"increment": completion_tokens}, - "request_count": {"increment": len(requests)}, - "successful_requests": {"increment": successful}, - "failed_requests": {"increment": failed}, - }, + queue_model_usage_upsert( + batcher, + key_fields, + { + "spend": spend, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "request_count": len(requests), + "successful_requests": successful, + "failed_requests": failed, }, ) return @@ -183,3 +174,15 @@ async def flush_model_usage_transactions( if attempt >= n_retry_times: raise await asyncio.sleep(2.0**attempt + random.uniform(0, 1)) + + +def queue_model_usage_upsert( + batcher: ModelUsageBatch, key_fields: Mapping[str, str], increments: Mapping[str, int | float] +) -> None: + batcher.litellm_dailymodelusage.upsert( + where={"date_model_group_model_custom_llm_provider_task_type": key_fields}, + data={ + "create": {**key_fields, **increments}, + "update": {name: {"increment": value} for name, value in increments.items()}, + }, + ) diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 1ddbc139564..55e2e129e36 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -11,7 +11,7 @@ from datetime import datetime, timezone from functools import partial from itertools import groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar +from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypeVar from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import billed_guard from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES from litellm.proxy.db.spend_log_batching import spend_log_write_batches from litellm.proxy.utils import PrismaClient +from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ( DailyGuardrailMetricsRepository, DailyGuardrailUsageUnitsRepository, @@ -40,7 +41,7 @@ _RowKey = TypeVar("_RowKey") _RowValue = TypeVar("_RowValue") -class _UsageUnitKey(NamedTuple): +class UsageUnitKey(NamedTuple): guardrail_id: str date: str team_id: str @@ -48,7 +49,7 @@ class _UsageUnitKey(NamedTuple): usage_unit: str -class _UsageUnitIncrement(NamedTuple): +class UsageUnitIncrement(NamedTuple): units: int cost: float """USD for the priced share of units.""" @@ -56,14 +57,14 @@ class _UsageUnitIncrement(NamedTuple): """Units recorded with no known price, the share cost leaves out.""" -def _usage_unit_key(entry: tuple[_UsageUnitKey, _UsageUnitIncrement]) -> _UsageUnitKey: +def _usage_unit_key(entry: tuple[UsageUnitKey, UsageUnitIncrement]) -> UsageUnitKey: return entry[0] -def _usage_unit_increment(units: int, cost: float | None) -> _UsageUnitIncrement: +def _usage_unit_increment(units: int, cost: float | None) -> UsageUnitIncrement: if cost is None: - return _UsageUnitIncrement(units=units, cost=0.0, untracked_units=units) - return _UsageUnitIncrement(units=units, cost=cost, untracked_units=0) + return UsageUnitIncrement(units=units, cost=0.0, untracked_units=units) + return UsageUnitIncrement(units=units, cost=cost, untracked_units=0) class _MetricsKey(NamedTuple): @@ -89,22 +90,22 @@ class PendingRollups: def __init__(self) -> None: self.lock: Final = asyncio.Lock() self.metrics: Mapping[_MetricsKey, Mapping[str, int]] = MappingProxyType({}) - self.units: Mapping[_UsageUnitKey, _UsageUnitIncrement] = MappingProxyType({}) + self.units: Mapping[UsageUnitKey, UsageUnitIncrement] = MappingProxyType({}) _PENDING_ROLLUPS: Final = PendingRollups() _NO_COUNTERS: Final[Mapping[str, int]] = MappingProxyType({}) -_NO_INCREMENT: Final = _UsageUnitIncrement(units=0, cost=0.0, untracked_units=0) +_NO_INCREMENT: Final = UsageUnitIncrement(units=0, cost=0.0, untracked_units=0) def _merged_keys(base: Mapping[_RowKey, object], extra: Mapping[_RowKey, object]) -> tuple[_RowKey, ...]: return (*base, *(key for key in extra if key not in base)) -def _summed_increments(increments: Iterable[_UsageUnitIncrement]) -> _UsageUnitIncrement: +def _summed_increments(increments: Iterable[UsageUnitIncrement]) -> UsageUnitIncrement: materialized: Final = tuple(increments) - return _UsageUnitIncrement( + return UsageUnitIncrement( units=sum(i.units for i in materialized), cost=sum(i.cost for i in materialized), untracked_units=sum(i.untracked_units for i in materialized), @@ -112,8 +113,8 @@ def _summed_increments(increments: Iterable[_UsageUnitIncrement]) -> _UsageUnitI def _merged_unit_rows( - base: Mapping[_UsageUnitKey, _UsageUnitIncrement], extra: Mapping[_UsageUnitKey, _UsageUnitIncrement] -) -> Mapping[_UsageUnitKey, _UsageUnitIncrement]: + base: Mapping[UsageUnitKey, UsageUnitIncrement], extra: Mapping[UsageUnitKey, UsageUnitIncrement] +) -> Mapping[UsageUnitKey, UsageUnitIncrement]: return MappingProxyType( { key: _summed_increments((base.get(key, _NO_INCREMENT), extra.get(key, _NO_INCREMENT))) @@ -250,7 +251,7 @@ def _parse_payload_start_time(payload: Mapping[str, object]) -> datetime | None: def _iter_usage_unit_increments( logs_to_process: Sequence[Mapping[str, object]], -) -> Iterator[tuple[_UsageUnitKey, _UsageUnitIncrement]]: +) -> Iterator[tuple[UsageUnitKey, UsageUnitIncrement]]: for payload in logs_to_process: start_time = _parse_payload_start_time(payload) if not payload.get("request_id") or start_time is None: @@ -266,14 +267,14 @@ def _iter_usage_unit_increments( cost_by_unit = billed_guardrail_cost_by_unit(entry) for unit_name, units in usage.items(): if isinstance(units, int) and not isinstance(units, bool) and units > 0: - key = _UsageUnitKey(guardrail_id, date_key, team_id, api_key, str(unit_name)) + key = UsageUnitKey(guardrail_id, date_key, team_id, api_key, str(unit_name)) cost = cost_by_unit.get(str(unit_name)) if cost_by_unit is not None else None yield key, _usage_unit_increment(units=units, cost=cost) def _sum_usage_unit_increments( logs_to_process: Sequence[Mapping[str, object]], -) -> Mapping[_UsageUnitKey, _UsageUnitIncrement]: +) -> Mapping[UsageUnitKey, UsageUnitIncrement]: ordered: Final = sorted(_iter_usage_unit_increments(logs_to_process), key=_usage_unit_key) return MappingProxyType( { @@ -283,8 +284,16 @@ def _sum_usage_unit_increments( ) -async def _upsert_usage_unit_row( - prisma_client: PrismaClient, key: _UsageUnitKey, increment: _UsageUnitIncrement +class GuardrailUsageDatabase(Protocol): + @property + def litellm_dailyguardrailusageunits(self) -> TableActions[object]: ... + + +async def upsert_usage_unit_row( + prisma_client: PrismaClient, + key: UsageUnitKey, + increment: UsageUnitIncrement, + outer_transaction: GuardrailUsageDatabase | None = None, ) -> None: row: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsCreateInput] = { "guardrail_id": key.guardrail_id, @@ -314,7 +323,12 @@ async def _upsert_usage_unit_row( "untracked_units": {"increment": increment.untracked_units}, }, } - await DailyGuardrailUsageUnitsRepository(prisma_client).table.upsert(where=where, data=data) + table: Final = ( + outer_transaction.litellm_dailyguardrailusageunits + if outer_transaction is not None + else DailyGuardrailUsageUnitsRepository(prisma_client).table + ) + await table.upsert(where=where, data=data) async def _upsert_metrics_row(prisma_client: PrismaClient, key: _MetricsKey, agg: Mapping[str, int]) -> None: @@ -363,7 +377,7 @@ async def process_spend_logs_guardrail_usage( ) index_rows_by_key: Final[dict[tuple[str, str], dict[str, object]]] = {} - for payload in logs_to_process: + for payload in (row for row in logs_to_process if row.get("accounting_log_disabled") is not True): request_id = payload.get("request_id") start_time = _parse_payload_start_time(payload) if not isinstance(request_id, str) or not request_id or start_time is None: @@ -416,7 +430,10 @@ async def process_spend_logs_guardrail_usage( {key: agg for key, agg in daily_guardrail.items() if int(agg["requests_evaluated"]) > 0} ) metrics_rows: Final = _merged_metric_rows(pending_metrics, evaluated_metrics) - unit_rows: Final = _merged_unit_rows(pending_units, _sum_usage_unit_increments(logs_to_process)) + unit_rows: Final = _merged_unit_rows( + pending_units, + _sum_usage_unit_increments(tuple(row for row in logs_to_process if "accounting_payload" not in row)), + ) if not metrics_rows and not index_rows and not unit_rows: return @@ -435,7 +452,7 @@ async def process_spend_logs_guardrail_usage( metrics_rows, partial(_upsert_metrics_row, prisma_client), "daily metrics", sleep ) failed_units: Final = await _upsert_rows_with_retry( - unit_rows, partial(_upsert_usage_unit_row, prisma_client), "usage unit", sleep + unit_rows, partial(upsert_usage_unit_row, prisma_client), "usage unit", sleep ) if failed_metrics or failed_units: async with pending.lock: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 90887422567..ff67f2cb7c6 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -9,6 +9,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, cast from pydantic import TypeAdapter, ValidationError +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -38,6 +39,7 @@ from litellm.proxy.auth.auth_checks import ( from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.proxy.db.db_spend_update_writer import ( + ACCOUNTING_LEDGER_KWARG, DBSpendUpdateWriter, debitable_model_access_groups, get_llm_router, @@ -657,6 +659,7 @@ class _ProxyDBLogger(CustomLogger): token: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) end_user_id: Final = None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id + accounting_rows: Final = tuple((cost, dict(row)) for cost, row in spend_rows) charged: Final = [ await self._spend_writer().update_database( token=token, @@ -664,17 +667,17 @@ class _ProxyDBLogger(CustomLogger): user_id=user_api_key_dict.user_id, end_user_id=end_user_id, team_id=user_api_key_dict.team_id, - kwargs=dict(row_kwargs), + kwargs=row_kwargs, completion_response=original_exception, start_time=start_time, end_time=datetime.now(), org_id=user_api_key_dict.org_id, project_id=user_api_key_dict.project_id, ) - for row_cost, row_kwargs in spend_rows + for row_cost, row_kwargs in accounting_rows ] counted_rows: Final = tuple( - row for row, row_charged in zip(spend_rows, charged, strict=True) if row_charged and row[0] > 0 + row for row, row_charged in zip(accounting_rows, charged, strict=True) if row_charged and row[0] > 0 ) if counted_rows and not ProxyUpdateSpend.disable_spend_updates(): await self._increment_failure_spend_counters( @@ -753,6 +756,7 @@ class _ProxyDBLogger(CustomLogger): router=get_llm_router(), ), project_id=user_api_key_dict.project_id, + **_ledger_counter_options(row_kwargs.get(ACCOUNTING_LEDGER_KWARG)), ) except Exception: # noqa: BLE001 # one row's counter failure must not keep the other rows' counters from landing verbose_proxy_logger.exception("Failed to raise spend counters for a failure spend row") @@ -1196,6 +1200,14 @@ def _get_request_tags_for_cost_tracking( return None +class _LedgerCounterOptions(TypedDict): + windows_from_ledger: ReadOnly[NotRequired[bool]] + + +def _ledger_counter_options(ledger: object) -> _LedgerCounterOptions: + return {"windows_from_ledger": True} if ledger is True else {} + + class _ChargeSpend(Protocol): async def __call__( self, @@ -1256,6 +1268,7 @@ class _IncrementSpendCounters(Protocol): request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, project_id: str | None = None, + windows_from_ledger: bool = False, ) -> None: ... @@ -1393,6 +1406,7 @@ async def _update_database_and_spend_counters_in_batch( request_started_at=start_time, model_access_groups=model_access_groups, project_id=project_id, + **_ledger_counter_options(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get(ACCOUNTING_LEDGER_KWARG)), ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b64c88521dc..6a125b9a0a4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3617,6 +3617,7 @@ async def increment_spend_counters( request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, project_id: str | None = None, + windows_from_ledger: bool = False, ): """ Atomically increment spend counters for budget enforcement. @@ -3653,6 +3654,7 @@ async def increment_spend_counters( request_started_at=request_started_at, model_access_groups=model_access_groups, project_id=project_id, + windows_from_ledger=windows_from_ledger, ) @@ -3682,6 +3684,50 @@ def _finish_pending_read(absent_key: str, pending_read: str, remember_absent: bo user_api_key_cache.budget_windows_absent_cache.delete_cache(absent_key) +async def accounting_window_keys(token: str | None, team_id: str | None) -> tuple[Mapping[str, object], ...]: + scopes: Final = tuple( + (entity_type, entity_id, cache_key) + for entity_type, entity_id, cache_key in ( + (Litellm_EntityType.KEY, token, token), + (Litellm_EntityType.TEAM, team_id, f"team_id:{team_id}"), + ) + if entity_id is not None and cache_key is not None + ) + snapshots: Final = tuple( + [ + await _accounting_scope_windows(entity_type, entity_id, cache_key) + for entity_type, entity_id, cache_key in scopes + ] + ) + from itertools import chain + + return tuple(chain.from_iterable(snapshots)) + + +async def _accounting_scope_windows( + entity_type: Litellm_EntityType, entity_id: str, cache_key: str +) -> tuple[Mapping[str, object], ...]: + from litellm.proxy.spend_tracking.budget_reservation import fixed_budget_window + + cached: Final[object] = _UNTYPED_VALUE.validate_python(await user_api_key_cache.async_get_cache(key=cache_key)) + limits: Final = ( + await _stored_budget_limits(entity_type, entity_id) if cached is None else _cached_budget_limits(cached) + ) + if not isinstance(limits, list): + return () + windows: Final = TypeAdapter(tuple[object, ...]).validate_python(limits) + return tuple( + { + "entity_type": entity_type.value, + "entity_id": entity_id, + "window_duration": fixed[0], + "window_start": fixed[1].isoformat(), + } + for window in windows + if (fixed := fixed_budget_window(window)) is not None + ) + + async def _stored_budget_limits(entity_type: Litellm_EntityType, entity_id: str) -> object: """Read only: caching the row here could race a delete and put a just-deleted key back into the auth cache.""" if prisma_client is None: @@ -3723,6 +3769,7 @@ async def _increment_spend_counters_batched( request_started_at: datetime | None, model_access_groups: Sequence[str] | None, project_id: str | None = None, + windows_from_ledger: bool = False, ): """Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET, and the reconcile adjustments go out in the same INCRBYFLOAT pipeline as the counter increments.""" @@ -3764,15 +3811,16 @@ async def _increment_spend_counters_batched( if window_counter not in reserved_counter_keys else None ) - await _enqueue_window_spend_row_update( - entity_type=entity_type, - entity_id=entity_id, - reset_at=window_reset_at, - window_duration=duration, - window_start=window_start, - increment=cost, - request_started_at=request_started_at, - ) + if not windows_from_ledger: + await _enqueue_window_spend_row_update( + entity_type=entity_type, + entity_id=entity_id, + reset_at=window_reset_at, + window_duration=duration, + window_start=window_start, + increment=cost, + request_started_at=request_started_at, + ) return pending_window async def _window_increments( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index c007b047087..3aaa74d4fc2 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -2141,6 +2141,23 @@ model LiteLLM_AccountingDelivery { @@index([target, state, next_attempt_at]) } +model LiteLLM_SpendLogDelivery { + delivery_id String @id @default(dbgenerated("gen_random_uuid()::text")) + request_id String + state AccountingDeliveryState @default(pending) + accounting_payload Json + is_batch Boolean + attempts Int @default(0) + next_attempt_at DateTime @default(now()) + applied_at DateTime? + last_error String? + created_at DateTime @default(now()) + + @@index([state, next_attempt_at]) + @@index([state, applied_at]) + @@index([request_id]) +} + model LiteLLM_SubscriptionAccountTable { subscription_account_id String @id @default(uuid()) custom_llm_provider String diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index dcfd6148f54..af7eb742aa7 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1264,6 +1264,15 @@ def _get_entry_reserved_cost(entry: dict, default_reserved_cost: float) -> float return default_reserved_cost +def fixed_budget_window(window: object) -> tuple[str, datetime] | None: + fields: Final = _coerce_window(window) + start: Final = get_budget_window_start(window) + duration: Final = fields.get("budget_duration") + if not fields.get("reset_at") or not isinstance(duration, str) or start is None: + return None + return duration, start + + def get_budget_window_start(window: object) -> datetime | None: window_dict: Final = _coerce_window(window) budget_duration: Final = window_dict.get("budget_duration") diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 85a7f870eb4..fc54d1da0b2 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -161,9 +161,7 @@ def _is_non_secret_key_value(value: str, *, key_alias: object = None) -> bool: ) -def _redact_logged_api_key( - value: str | None, *, already_redacted: bool = False, key_alias: object = None -) -> str | None: +def redact_logged_api_key(value: str | None, *, already_redacted: bool = False, key_alias: object = None) -> str | None: if not isinstance(value, str) or not value: return None stripped: Final = re.sub(r"(?i)^bearer ", "", value) @@ -327,7 +325,7 @@ def _get_spend_logs_metadata( and _is_non_secret_key_value(_trusted_hash, key_alias=_key_alias) and _trusted_hash == _raw_key ) - clean_metadata["user_api_key"] = _redact_logged_api_key( + clean_metadata["user_api_key"] = redact_logged_api_key( _raw_key, already_redacted=_already_redacted, key_alias=_key_alias ) clean_metadata["applied_guardrails"] = applied_guardrails @@ -692,14 +690,14 @@ def get_logging_payload( and _is_non_secret_key_value(_trusted_hash, key_alias=_key_alias) and _trusted_hash == api_key ) - api_key = _redact_logged_api_key(api_key, already_redacted=_key_already_redacted, key_alias=_key_alias) or "" + api_key = redact_logged_api_key(api_key, already_redacted=_key_already_redacted, key_alias=_key_alias) or "" if ( standard_logging_payload is not None ): # [TODO] migrate completely to sl payload. currently missing pass-through endpoint data api_key = ( api_key - or _redact_logged_api_key( + or redact_logged_api_key( standard_logging_payload["metadata"].get("user_api_key_hash"), already_redacted=True, key_alias=standard_logging_payload["metadata"].get("user_api_key_alias"), diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8b77e39b3e5..71ecae74058 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -278,6 +278,7 @@ from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.spend_log_insert import RejectedSpendLogRow from litellm.repositories.table_repositories import ( EndUserRepository, HealthCheckRepository, @@ -8418,14 +8419,17 @@ async def enqueue_spend_logs( """ added: Final = sum(spend_log_row_bytes(row) for row in logs) async with prisma_client._spend_log_transactions_lock: - queued: Final = ( + queued: Final[tuple[Mapping[str, object], ...]] = ( tuple(logs) + tuple(prisma_client.spend_log_transactions) if at_head else tuple(prisma_client.spend_log_transactions) + tuple(logs) ) kept, kept_bytes = spend_log_queue_within_budget(queued, PrismaClient.spend_log_queue_bytes + added, max_bytes) - prisma_client.spend_log_transactions[:] = kept - PrismaClient.spend_log_queue_bytes = kept_bytes + rescued: Final = tuple( + accounting_only_spend_log(row) for row in queued[: len(queued) - len(kept)] if _carries_a_charge(row) + ) + prisma_client.spend_log_transactions[:] = (*rescued, *kept) + PrismaClient.spend_log_queue_bytes = kept_bytes + sum(spend_log_row_bytes(row) for row in rescued) if len(kept) < len(queued): verbose_proxy_logger.error( "Spend tracking - spend log queue is at its %d byte budget; dropped the %d oldest spend logs", @@ -8480,13 +8484,18 @@ class ProxyUpdateSpend: prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging, end_user_list_transactions: dict[str, float], + outer_transaction: _EndUserSpendTransaction | None = None, ): for i in range(n_retry_times + 1): start_time = time.time() try: async with ( db_span("update_end_user_spend", "LiteLLM_EndUserTable"), - _end_user_spend_transaction(prisma_client.db, timedelta(seconds=60)) as transaction, + ( + contextlib.nullcontext(outer_transaction) + if outer_transaction is not None + else _end_user_spend_transaction(prisma_client.db, timedelta(seconds=60)) + ) as transaction, ): batcher: _EndUserSpendBatch async with transaction.batch_() as batcher: @@ -8508,6 +8517,8 @@ class ProxyUpdateSpend: break except Exception as e: + if outer_transaction is not None: + raise await DBSpendUpdateWriter.handle_spend_update_failure( e=e, attempt=i, @@ -8547,11 +8558,18 @@ class ProxyUpdateSpend: for i in range(n_retry_times + 1): try: base_url = os.getenv("SPEND_LOGS_URL", None) - if len(logs_to_process) > 0 and base_url is not None and db_writer_client is not None: + await _commit_export_accounting( + prisma_client, rows_to_write, base_url if db_writer_client is not None else None + ) + if ( + base_url is not None + and db_writer_client is not None + and any(row.get("accounting_log_disabled") is not True for row in rows_to_write) + ): if not base_url.endswith("/"): base_url += "/" verbose_proxy_logger.debug("base_url: %s", base_url) - json_data = json.dumps(rows_to_write) + json_data = json.dumps(_spend_log_export_rows(rows_to_write)) response = await db_writer_client.post( url=base_url + "spend/update", data=json_data, @@ -8629,6 +8647,28 @@ class ProxyUpdateSpend: return False +def _spend_log_export_rows(rows: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, object], ...]: + return tuple( + { + key: value + for key, value in row.items() + if key not in {"accounting_payload", "accounting_log_disabled", "accounting_delivery_id"} + } + for row in rows + if row.get("accounting_log_disabled") is not True + ) + + +async def _commit_export_accounting( + prisma_client: PrismaClient, rows: Sequence[Mapping[str, object]], base_url: str | None +) -> None: + if base_url is None: + return + await SpendLogsRepository(prisma_client).insert_with_deliveries( + tuple(row for row in rows if "accounting_payload" in row), write_spend_logs=False + ) + + async def commit_tool_discoveries(prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging) -> None: """Scheduler job writing discovered tools to LiteLLM_ToolTable, on its own lock, next to update_spend.""" writer: Final = proxy_logging_obj.db_spend_update_writer @@ -9127,6 +9167,40 @@ def _is_transient_spend_log_write_error(e: Exception) -> bool: ) +_ACCOUNTING_ONLY_KEYS: Final = frozenset({"request_id", "call_type", "accounting_payload", "accounting_delivery_id"}) + + +def _carries_a_charge(row: Mapping[str, object]) -> bool: + return "accounting_payload" in row + + +def accounting_only_spend_log(row: Mapping[str, object]) -> Mapping[str, object]: + return { + **{key: value for key, value in row.items() if key in _ACCOUNTING_ONLY_KEYS}, + "accounting_log_disabled": True, + } + + +async def _keep_charges_of_dropped_logs(repo: SpendLogsRepository, rows: Sequence[Mapping[str, object]]) -> None: + charged: Final = tuple(accounting_only_spend_log(row) for row in rows if _carries_a_charge(row)) + if not charged: + return + try: + await repo.insert_with_deliveries(charged) + except Exception as e: # noqa: BLE001 # the request log is already lost; keep the flush going for other rows + if not ( + PrismaDBExceptionHandler.is_rejected_row_error(e) + or isinstance(e, (UnicodeEncodeError, RejectedSpendLogRow)) + ): + raise + spend_log_error( + "Spend tracking - could not keep the charges of %d dropped spend log rows. error=%s", + len(charged), + str(e), + exc=e, + ) + + async def _create_spend_logs_with_poison_isolation( repo: SpendLogsRepository, rows: Sequence[Mapping[str, object]], @@ -9155,10 +9229,13 @@ async def _create_spend_logs_with_poison_isolation( persist. Returns the budget left after this subtree. """ try: - await repo.table.create_many(data=rows, skip_duplicates=True) + await repo.insert_with_deliveries(rows) return failure_budget except Exception as e: - if not (PrismaDBExceptionHandler.is_prisma_data_error(e) or isinstance(e, UnicodeEncodeError)): + if not ( + PrismaDBExceptionHandler.is_rejected_row_error(e) + or isinstance(e, (UnicodeEncodeError, RejectedSpendLogRow)) + ): raise if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): raise @@ -9173,6 +9250,7 @@ async def _create_spend_logs_with_poison_isolation( str(e), exc=e, ) + await _keep_charges_of_dropped_logs(repo, rows) return budget_left if budget_left <= 0: spend_log_error( @@ -9180,6 +9258,7 @@ async def _create_spend_logs_with_poison_isolation( "isolation failure budget exhausted for this flush", len(rows), ) + await _keep_charges_of_dropped_logs(repo, rows) return 0 mid: Final = len(rows) // 2 remaining: Final = await _create_spend_logs_with_poison_isolation(repo, rows[:mid], budget_left) @@ -9189,6 +9268,7 @@ async def _create_spend_logs_with_poison_isolation( "isolation failure budget exhausted for this flush", len(rows) - mid, ) + await _keep_charges_of_dropped_logs(repo, rows[mid:]) return 0 return await _create_spend_logs_with_poison_isolation(repo, rows[mid:], remaining) diff --git a/litellm/repositories/accounting_delivery_repository.py b/litellm/repositories/accounting_delivery_repository.py new file mode 100644 index 00000000000..7dd92aecb2d --- /dev/null +++ b/litellm/repositories/accounting_delivery_repository.py @@ -0,0 +1,193 @@ +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass +from datetime import timedelta +from itertools import chain +from types import TracebackType +from typing import TYPE_CHECKING, Final, Literal, Protocol + +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +from litellm.repositories.prisma_protocols import TableActions + +if TYPE_CHECKING: + from litellm.proxy.db.db_spend_update_writer import SpendBatchManager + + +class LedgerTransaction(Protocol): + @property + def litellm_dailyguardrailusageunits(self) -> TableActions[object]: ... + + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ... + + async def execute_raw(self, query: str, *args: object) -> int: ... + + def batch_(self) -> "SpendBatchManager": ... + + +class _TransactionManager(Protocol): + async def __aenter__(self) -> LedgerTransaction: ... + + async def __aexit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, exc_tb: TracebackType | None + ) -> None: ... + + +class LedgerDatabase(Protocol): + def tx(self, *, timeout: timedelta) -> _TransactionManager: ... + + +class _Delivery(BaseModel): + model_config = ConfigDict(frozen=True) + + delivery_id: str + attempts: int + accounting_payload: JsonValue + + +@dataclass(frozen=True, slots=True) +class DeliveryProjection: + row: Mapping[str, object] + + +@dataclass(frozen=True, slots=True) +class DeliveryBatchResult: + applied: int = 0 + failed: int = 0 + parked: tuple[str, ...] = () + error: str | None = None + + +class DeliveryMetrics(BaseModel): + model_config = ConfigDict(frozen=True) + + pending_count: int + parked_count: int + oldest_pending_seconds: float + batch_count: int + table_bytes: int + + +_DELIVERIES: Final = TypeAdapter(tuple[_Delivery, ...]) +_PAYLOAD: Final = TypeAdapter(dict[str, object]) +_TIMEOUT: Final = timedelta(seconds=60) +_CLAIM: Final = """SELECT delivery_id, attempts, accounting_payload FROM "LiteLLM_SpendLogDelivery" +WHERE state = 'pending' AND next_attempt_at <= CURRENT_TIMESTAMP + AND (attempts > 0) = $2::boolean +ORDER BY next_attempt_at, delivery_id LIMIT $1::integer FOR UPDATE SKIP LOCKED""" +_APPLIED: Final = """UPDATE "LiteLLM_SpendLogDelivery" +SET state = 'applied', applied_at = CURRENT_TIMESTAMP, last_error = NULL +WHERE delivery_id = ANY($1::text[]) AND state = 'pending' """ +_FAILED: Final = """UPDATE "LiteLLM_SpendLogDelivery" +SET attempts = attempts + 1, last_error = $3, + next_attempt_at = CURRENT_TIMESTAMP + make_interval(secs => LEAST(300, power(2, LEAST(attempts, 8)))::integer), + state = $4::"AccountingDeliveryState" +WHERE delivery_id = $1 AND state = 'pending' AND attempts = $2::integer +RETURNING delivery_id, attempts, accounting_payload""" + + +def _is_permanent_failure(error: Exception) -> bool: + """A delivery whose own data failed fails the same way on every retry, while any other error is the database's + and clears once it recovers.""" + from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler + + if PrismaDBExceptionHandler.is_database_service_unavailable_error(error): + return False + return isinstance(error, (ValueError, TypeError, KeyError)) or PrismaDBExceptionHandler.is_rejected_row_error(error) + + +class AccountingDeliveryRepository: + def __init__(self, database: LedgerDatabase) -> None: + self._database = database + + async def metrics(self) -> DeliveryMetrics: + async with self._database.tx(timeout=_TIMEOUT) as transaction: + rows: Final = await transaction.query_raw( + "SELECT COUNT(*) FILTER (WHERE state = 'pending') AS pending_count, " + "COUNT(*) FILTER (WHERE state = 'parked') AS parked_count, " + "COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(created_at) FILTER (WHERE state = 'pending')), 0) " + "AS oldest_pending_seconds, COUNT(*) FILTER (WHERE is_batch) AS batch_count, " + "pg_total_relation_size('\"LiteLLM_SpendLogDelivery\"') AS table_bytes " + 'FROM "LiteLLM_SpendLogDelivery"' + ) + return DeliveryMetrics.model_validate(rows[0]) + + async def cleanup_applied(self, *, retention_days: int = 1, batch_size: int = 1000) -> int: + if retention_days < 1 or batch_size < 1: + raise ValueError("Delivery retention and cleanup batch size must be positive") + async with self._database.tx(timeout=_TIMEOUT) as transaction: + return await transaction.execute_raw( + 'DELETE FROM "LiteLLM_SpendLogDelivery" WHERE delivery_id IN (' + 'SELECT delivery_id FROM "LiteLLM_SpendLogDelivery" ' + "WHERE state = 'applied' AND NOT is_batch AND applied_at < CURRENT_TIMESTAMP - make_interval(days => $1::integer) " + "ORDER BY applied_at, delivery_id LIMIT $2::integer FOR UPDATE SKIP LOCKED)", + retention_days, + batch_size, + ) + + async def apply_pending( + self, + apply: Callable[[LedgerTransaction, tuple[DeliveryProjection, ...]], Awaitable[None]], + *, + batch_size: int = 100, + max_attempts: int = 5, + retry_only: bool | None = None, + ) -> DeliveryBatchResult: + if batch_size < 1 or max_attempts < 1: + raise ValueError("Delivery batch size and maximum attempts must be positive") + claimed: tuple[_Delivery, ...] = () # rebind-ok: preserve claims for retry accounting after rollback + try: + async with self._database.tx(timeout=_TIMEOUT) as transaction: + retries: Final = ( + _DELIVERIES.validate_python(await transaction.query_raw(_CLAIM, 1, True)) + if retry_only is not False + else () + ) + claimed = retries or ( # rebind-ok: retain claimed identities after rollback + _DELIVERIES.validate_python(await transaction.query_raw(_CLAIM, batch_size, False)) + if retry_only is not True + else () + ) + if not claimed: + return DeliveryBatchResult() + await apply( + transaction, + tuple( + DeliveryProjection(_PAYLOAD.validate_python(delivery.accounting_payload)) + for delivery in claimed + ), + ) + await transaction.execute_raw(_APPLIED, tuple(delivery.delivery_id for delivery in claimed)) + return DeliveryBatchResult(applied=len(claimed)) + except Exception as error: # noqa: BLE001 # roll back all destinations before recording a retry + if not claimed: + raise + return await self._record_failure(claimed, error, max_attempts) + + async def _record_failure( + self, claimed: tuple[_Delivery, ...], error: Exception, max_attempts: int + ) -> DeliveryBatchResult: + parks: Final = len(claimed) == 1 and _is_permanent_failure(error) + async with self._database.tx(timeout=_TIMEOUT) as transaction: + updates: Final = tuple( + [ + await self._record_row_failure( + transaction, + delivery, + type(error).__name__, + "parked" if parks and delivery.attempts + 1 >= max_attempts else "pending", + ) + for delivery in claimed + ] + ) + return DeliveryBatchResult( + failed=len(claimed), parked=tuple(chain.from_iterable(updates)), error=type(error).__name__ + ) + + @staticmethod + async def _record_row_failure( + transaction: LedgerTransaction, delivery: _Delivery, error: str, state: Literal["pending", "parked"] + ) -> tuple[str, ...]: + updated: Final = _DELIVERIES.validate_python( + await transaction.query_raw(_FAILED, delivery.delivery_id, delivery.attempts, error, state) + ) + return tuple(row.delivery_id for row in updated) if state == "parked" else () diff --git a/litellm/repositories/spend_log_insert.py b/litellm/repositories/spend_log_insert.py new file mode 100644 index 00000000000..a958edb9f8a --- /dev/null +++ b/litellm/repositories/spend_log_insert.py @@ -0,0 +1,88 @@ +from collections.abc import Mapping, Sequence +from datetime import datetime, timezone +from typing import Final + +from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import LiteralString + +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + +_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) + +_COLUMNS: Final[tuple[LiteralString, ...]] = ( + "request_id", + "call_type", + "api_key", + "spend", + "total_tokens", + "prompt_tokens", + "completion_tokens", + "startTime", + "endTime", + "request_duration_ms", + "completionStartTime", + "model", + "model_id", + "model_group", + "custom_llm_provider", + "api_base", + "user", + "metadata", + "cache_hit", + "cache_key", + "request_tags", + "team_id", + "organization_id", + "end_user", + "requester_ip_address", + "messages", + "response", + "session_id", + "status", + "mcp_namespaced_tool_name", + "billing_agent_id", + "agent_id", + "proxy_server_request", + "litellm_call_id", + "created_at", + "updated_at", +) +_JSON_COLUMNS: Final = frozenset({"metadata", "request_tags", "messages", "response", "proxy_server_request"}) + + +class RejectedSpendLogRow(ValueError): + """A row the raw insert cannot encode, so it is isolated like a row Postgres refuses.""" + + +def _json_value(column: str, value: object) -> object: + if isinstance(value, datetime): + return value.isoformat() + if column not in _JSON_COLUMNS or not isinstance(value, str): + return value + try: + return _JSON.validate_json(value) + except ValidationError as error: + raise RejectedSpendLogRow(f"Spend log column {column} is not valid JSON") from error + + +def spend_log_insert_statement(rows: Sequence[Mapping[str, object]]) -> tuple[LiteralString, str]: + if not rows: + raise ValueError("Spend log insert requires at least one row") + fields: Final = frozenset(rows[0]) | {"created_at", "updated_at"} + if not fields.issubset(_COLUMNS) or any(frozenset(row) | {"created_at", "updated_at"} != fields for row in rows): + raise RejectedSpendLogRow("Spend log insert requires matching known columns") + columns: Final = ", ".join('"' + column + '"' for column in _COLUMNS if column in fields) + now: Final = datetime.now(timezone.utc).isoformat() + values: Final = tuple( + {"created_at": now, "updated_at": now, **{column: _json_value(column, value) for column, value in row.items()}} + for row in rows + ) + sql: Final = ( + 'INSERT INTO "LiteLLM_SpendLogs" (' + + columns + + ") SELECT " + + columns + + ' FROM jsonb_populate_recordset(NULL::"LiteLLM_SpendLogs", $1::jsonb) ' + "ON CONFLICT DO NOTHING RETURNING request_id" + ) + return sql, safe_dumps(values, non_finite_as_null=True) diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 2d3139f7a2c..10db81eca4b 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -7,8 +7,15 @@ These are thin wrappers for tables that do not (yet) need domain-specific query methods; richer repositories live in their own modules. """ +from collections.abc import Mapping, Sequence +from datetime import timedelta +from itertools import groupby from typing import TYPE_CHECKING, Any, Final, Generic +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.repositories.accounting_delivery_repository import LedgerDatabase, LedgerTransaction from litellm.repositories.prisma_protocols import RowT_co, TableActions if TYPE_CHECKING: @@ -80,9 +87,117 @@ class OrganizationMembershipRepository(PrismaTableRepository["prisma_models.Lite table_name = "litellm_organizationmembership" +_SPEND_LOG_FIELDS: Final = TypeAdapter(dict[str, object]) + + +def _accounting_payload(value: object) -> Mapping[str, object]: + return ( + _SPEND_LOG_FIELDS.validate_json(value) if isinstance(value, str) else _SPEND_LOG_FIELDS.validate_python(value) + ) + + +_ACCOUNTING_ONLY_FIELDS: Final = frozenset({"accounting_payload", "accounting_log_disabled", "accounting_delivery_id"}) + + +def _settles_once(row: Mapping[str, object]) -> bool: + """A batch or video cost row: update_database gives every other ledger row its own delivery id.""" + return "accounting_delivery_id" not in row + + +def _delivery_id(row: Mapping[str, object]) -> str: + return str(row["accounting_delivery_id"]) if not _settles_once(row) else str(row["request_id"]) + + class SpendLogsRepository(PrismaTableRepository["prisma_models.LiteLLM_SpendLogs"]): table_name = "litellm_spendlogs" + async def insert_with_deliveries( + self, rows: Sequence[Mapping[str, object]], *, write_spend_logs: bool = True + ) -> int: + legacy: Final = tuple(row for row in rows if "accounting_payload" not in row) + legacy_count: Final = await self.table.create_many(data=legacy, skip_duplicates=True) if legacy else 0 + ledger: Final = tuple( + sorted( + {_delivery_id(row): row for row in reversed(rows) if "accounting_payload" in row}.values(), + key=_delivery_id, + ) + ) + if not ledger: + return legacy_count + database: Final[LedgerDatabase] = self.prisma_client.writer_db + async with database.tx(timeout=timedelta(seconds=60)) as transaction: + accepted: Final = await self._insert_ledger(transaction, ledger, write_spend_logs=write_spend_logs) + return legacy_count + accepted + + @staticmethod + async def _insert_ledger( + transaction: LedgerTransaction, rows: Sequence[Mapping[str, object]], *, write_spend_logs: bool + ) -> int: + from litellm.proxy.db.db_span import db_span + from litellm.repositories.spend_log_insert import spend_log_insert_statement + + payloads: Final = tuple( + { + "delivery_id": _delivery_id(row), + "request_id": row["request_id"], + "accounting_payload": _accounting_payload(row["accounting_payload"]), + "is_batch": _settles_once(row), + } + for row in rows + ) + inserted_deliveries: Final = await transaction.query_raw( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (delivery_id, request_id, accounting_payload, is_batch) ' + "SELECT delivery_id, request_id, accounting_payload, is_batch FROM jsonb_to_recordset($1::jsonb) " + "AS inputs(delivery_id text, request_id text, accounting_payload jsonb, is_batch boolean) " + "ON CONFLICT (delivery_id) DO NOTHING RETURNING delivery_id", + safe_dumps(payloads, non_finite_as_null=True), + ) + accepted_ids: Final = frozenset(str(row["delivery_id"]) for row in inserted_deliveries) + charged_by_legacy: Final = await SpendLogsRepository._jobs_charged_by_legacy( + transaction, + tuple(_delivery_id(row) for row in rows if _settles_once(row) and _delivery_id(row) in accepted_ids), + ) + if charged_by_legacy: + await transaction.execute_raw( + 'DELETE FROM "LiteLLM_SpendLogDelivery" WHERE delivery_id = ANY($1::text[])', list(charged_by_legacy) + ) + accepted: Final = accepted_ids - charged_by_legacy + local_rows: Final = tuple( + {key: value for key, value in row.items() if key not in _ACCOUNTING_ONLY_FIELDS} + for row in rows + if _delivery_id(row) in accepted and row.get("accounting_log_disabled") is not True + ) + if write_spend_logs and local_rows: + groups: Final = tuple( + tuple(group) + for _, group in groupby( + sorted(local_rows, key=lambda row: tuple(sorted(row))), key=lambda row: tuple(sorted(row)) + ) + ) + async with db_span("insert_spend_logs", "LiteLLM_SpendLogs"): + for group in groups: + await transaction.query_raw(*spend_log_insert_statement(group)) + return len(accepted) + + @staticmethod + async def _jobs_charged_by_legacy(transaction: LedgerTransaction, cost_row_ids: Sequence[str]) -> frozenset[str]: + """Batches and videos a pre-ledger poll already charged through their SpendLogs cost row.""" + if not cost_row_ids: + return frozenset() + rows: Final = await transaction.query_raw( + 'SELECT DISTINCT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = ANY($1::text[]) ' + "AND status = 'success' AND spend > 0", + list(cost_row_ids), + ) + return frozenset(str(row["request_id"]) for row in rows) + + async def has_delivery(self, delivery_id: str) -> bool: + database: Final[LedgerTransaction] = self.prisma_client.writer_db + rows: Final = await database.query_raw( + 'SELECT 1 FROM "LiteLLM_SpendLogDelivery" WHERE delivery_id = $1', delivery_id + ) + return bool(rows) + class BudgetWindowSpendRepository(PrismaTableRepository["prisma_models.LiteLLM_BudgetWindowSpend"]): table_name = "litellm_budgetwindowspend" diff --git a/schema.prisma b/schema.prisma index c007b047087..3aaa74d4fc2 100644 --- a/schema.prisma +++ b/schema.prisma @@ -2141,6 +2141,23 @@ model LiteLLM_AccountingDelivery { @@index([target, state, next_attempt_at]) } +model LiteLLM_SpendLogDelivery { + delivery_id String @id @default(dbgenerated("gen_random_uuid()::text")) + request_id String + state AccountingDeliveryState @default(pending) + accounting_payload Json + is_batch Boolean + attempts Int @default(0) + next_attempt_at DateTime @default(now()) + applied_at DateTime? + last_error String? + created_at DateTime @default(now()) + + @@index([state, next_attempt_at]) + @@index([state, applied_at]) + @@index([request_id]) +} + model LiteLLM_SubscriptionAccountTable { subscription_account_id String @id @default(uuid()) custom_llm_provider String diff --git a/tests/integration/database/test_spendlogs_delivery.py b/tests/integration/database/test_spendlogs_delivery.py new file mode 100644 index 00000000000..19f01a6b902 --- /dev/null +++ b/tests/integration/database/test_spendlogs_delivery.py @@ -0,0 +1,512 @@ +import asyncio +import os +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit, urlunsplit + +import psycopg +import pytest +import pytest_asyncio +from prisma import Prisma + +from litellm.repositories.accounting_delivery_repository import ( + AccountingDeliveryRepository, + DeliveryProjection, + LedgerTransaction, +) +from litellm.repositories.spend_log_insert import spend_log_insert_statement +from tests.integration._support.database import scratch_database + +_MIGRATIONS: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras/litellm_proxy_extras/migrations" + + +@pytest.fixture +def delivery_database(request: pytest.FixtureRequest) -> Iterator[psycopg.Connection[tuple[object, ...]]]: + with scratch_database("spendlogs_delivery") as url, psycopg.connect(url, autocommit=True) as connection: + if getattr(request, "param", "public") == "ledger_tenant": + connection.execute("CREATE SCHEMA ledger_tenant") + connection.execute("SET search_path TO ledger_tenant") + connection.execute((_MIGRATIONS / "20261027000000_add_accounting_ledger_tables/migration.sql").read_text()) + connection.execute("""CREATE TABLE "LiteLLM_SpendLogs" ( + request_id text PRIMARY KEY, call_type text NOT NULL, spend float8 NOT NULL DEFAULT 0, + status text, + metadata jsonb DEFAULT '{}', "startTime" timestamp(3) NOT NULL, + created_at timestamp(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at timestamp(3) NOT NULL DEFAULT CURRENT_TIMESTAMP + )""") + connection.execute((_MIGRATIONS / "20261103010000_spendlogs_accounting_delivery/migration.sql").read_text()) + yield connection + + +def _insert( + connection: psycopg.Connection[tuple[object, ...]], rows: Sequence[Mapping[str, object]], *, ledger: bool = True +) -> None: + from psycopg.types.json import Jsonb + + statement: Final = spend_log_insert_statement(rows) + with connection.transaction(): + inserted: Final = { + row[0] for row in connection.execute(statement[0].replace("$1", "%s"), (statement[1],)).fetchall() + } + for row in rows: + if ledger and row["request_id"] in inserted: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (delivery_id, request_id, accounting_payload, is_batch) ' + "VALUES (%s, %s, %s, %s)", + (row["request_id"], row["request_id"], Jsonb(dict(row)), row.get("call_type") == "aretrieve_batch"), + ) + + +def _with_payload(row: Mapping[str, object]) -> Mapping[str, object]: + return {**row, "accounting_payload": dict(row)} + + +def _row(request_id: str) -> Mapping[str, object]: + return { + "request_id": request_id, + "call_type": "acompletion", + "spend": 0.5, + "startTime": "2026-11-02T12:00:00+00:00", + "metadata": '{"accounting_savings":{"compression":0.2}}', + } + + +@pytest.mark.parametrize("call_type", ["acompletion", "aretrieve_batch"]) +@pytest.mark.parametrize("mode", ["local", "url", "disabled"]) +@pytest.mark.parametrize("partitioned", [False, True]) +@pytest.mark.asyncio +async def test_delivery_first_acceptance_deduplicates_all_storage_modes( + delivery_database, delivery_client, mode, partitioned, call_type +): + from types import SimpleNamespace + + from litellm.repositories.table_repositories import SpendLogsRepository + + if partitioned: + delivery_database.execute('ALTER TABLE "LiteLLM_SpendLogs" RENAME TO spend_template') + delivery_database.execute( + 'CREATE TABLE "LiteLLM_SpendLogs" (LIKE spend_template INCLUDING DEFAULTS INCLUDING CONSTRAINTS, PRIMARY KEY(request_id, "startTime")) PARTITION BY RANGE ("startTime")' + ) + delivery_database.execute('CREATE TABLE "LiteLLM_SpendLogs_p20261102" PARTITION OF "LiteLLM_SpendLogs" DEFAULT') + delivery_database.execute("DROP TABLE spend_template") + repository = SpendLogsRepository(SimpleNamespace(db=delivery_client, writer_db=delivery_client)) + row = {**_with_payload({**_row("shared"), "call_type": call_type}), "accounting_log_disabled": mode == "disabled"} + results = await asyncio.gather( + repository.insert_with_deliveries((row,), write_spend_logs=mode != "url"), + repository.insert_with_deliveries( + ({**row, "startTime": "2026-11-03T12:00:00"},), write_spend_logs=mode != "url" + ), + ) + assert sorted(results) == [0, 1] + assert delivery_database.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogs"').fetchone() == (int(mode == "local"),) + assert delivery_database.execute( + "SELECT accounting_payload->>'spend', state::text FROM \"LiteLLM_SpendLogDelivery\"" + ).fetchall() == [("0.5", "pending")] + if partitioned: + delivery_database.execute('DROP TABLE "LiteLLM_SpendLogs_p20261102"') + else: + delivery_database.execute('DELETE FROM "LiteLLM_SpendLogs"') + + async def apply(transaction, deliveries): + await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(deliveries[0].row["spend"])) + + worker = AccountingDeliveryRepository(delivery_client) + assert (await worker.apply_pending(apply)).applied == 1 + assert (await worker.apply_pending(apply)).applied == 0 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + + +@pytest.mark.parametrize("mode", ["local", "url", "disabled"]) +@pytest.mark.asyncio +async def test_duplicate_rows_in_one_flush_keep_first_accounting_payload(delivery_database, delivery_client, mode): + from types import SimpleNamespace + + from litellm.repositories.table_repositories import SpendLogsRepository + + repository = SpendLogsRepository(SimpleNamespace(db=delivery_client, writer_db=delivery_client)) + first = {**_with_payload(_row("repeated")), "accounting_log_disabled": mode == "disabled"} + later = { + **_with_payload({**_row("repeated"), "spend": 9.0, "startTime": "2026-11-03T12:00:00+00:00"}), + "accounting_log_disabled": mode == "disabled", + } + assert await repository.insert_with_deliveries((first, later), write_spend_logs=mode != "url") == 1 + assert delivery_database.execute( + 'SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"' + ).fetchone() == (first["accounting_payload"],) + assert delivery_database.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == ( + [(first["spend"],)] if mode == "local" else [] + ) + + +@pytest.mark.parametrize(("legacy_spend", "accepted"), [(0.5, 0), (0.0, 1)]) +@pytest.mark.asyncio +async def test_batch_charged_before_the_ledger_is_not_charged_again( + delivery_database, delivery_client, legacy_spend, accepted +): + from types import SimpleNamespace + + from litellm.repositories.table_repositories import SpendLogsRepository + + batch: Final = {**_row("batch_legacy"), "call_type": "aretrieve_batch", "status": "success"} + _insert(delivery_database, ({**batch, "spend": legacy_spend},), ledger=False) + repository = SpendLogsRepository(SimpleNamespace(db=delivery_client, writer_db=delivery_client)) + assert await repository.insert_with_deliveries((_with_payload({**batch, "spend": 1.0}),)) == accepted + assert delivery_database.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (accepted,) + assert delivery_database.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == [(legacy_spend,)] + + +@pytest.mark.parametrize("table", ["LiteLLM_SpendLogs", "LiteLLM_SpendLogDelivery"]) +@pytest.mark.asyncio +async def test_transaction_a_rolls_back_payload_and_log_on_failure(delivery_database, delivery_client, table): + from types import SimpleNamespace + + from prisma.errors import PrismaError + from psycopg import sql + + from litellm.repositories.table_repositories import SpendLogsRepository + + delivery_database.execute( + sql.SQL("ALTER TABLE {} ADD CONSTRAINT reject_poison CHECK (request_id <> 'poison')").format( + sql.Identifier(table) + ) + ) + repository = SpendLogsRepository(SimpleNamespace(db=delivery_client, writer_db=delivery_client)) + with pytest.raises(PrismaError): + await repository.insert_with_deliveries(tuple(_with_payload(_row(name)) for name in ("good", "poison"))) + assert delivery_database.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,) + assert delivery_database.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogs"').fetchone() == (0,) + + +@pytest.mark.asyncio +async def test_worker_rolls_back_destinations_and_parks_only_isolated_failure( + delivery_database: psycopg.Connection[tuple[object, ...]], +) -> None: + url: Final = urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{delivery_database.info.dbname}")) + client: Final = Prisma(datasource={"url": url}) + delivery_database.execute("CREATE TABLE totals (spend float8 NOT NULL)") + delivery_database.execute("INSERT INTO totals VALUES (0)") + _insert(delivery_database, (_row("bad"), _row("good"))) + + async def apply(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None: + await transaction.execute_raw( + "UPDATE totals SET spend = spend + $1::float8", sum(float(str(row.row["spend"])) for row in rows) + ) + if any(row.row["request_id"] == "bad" for row in rows): + raise ValueError("Invalid projected destination") + + await client.connect() + try: + repository: Final = AccountingDeliveryRepository(client) + failed: Final = await repository.apply_pending(apply, max_attempts=2) + assert failed.failed == 2 + assert failed.parked == () + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.0,) + assert delivery_database.execute( + 'SELECT DISTINCT state::text, attempts FROM "LiteLLM_SpendLogDelivery"' + ).fetchall() == [("pending", 1)] + delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at = CURRENT_TIMESTAMP') + isolated: Final = await repository.apply_pending(apply, max_attempts=2) + assert isolated.parked == ("bad",) + applied: Final = await repository.apply_pending(apply, max_attempts=2) + assert applied.applied == 1 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + assert delivery_database.execute( + 'SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id' + ).fetchall() == [("bad", "parked"), ("good", "applied")] + assert (await repository.apply_pending(apply)).applied == 0 + finally: + await client.disconnect() + + +@pytest.mark.asyncio +async def test_a_locked_database_never_parks_a_valid_delivery(delivery_database, delivery_client): + _insert(delivery_database, (_row("valid"),)) + + async def apply(transaction, deliveries): + await transaction.execute_raw("SET LOCAL lock_timeout = '50ms'") + await transaction.execute_raw( + "UPDATE totals SET spend = spend + $1::float8", sum(float(row.row["spend"]) for row in deliveries) + ) + + repository = AccountingDeliveryRepository(delivery_client) + with psycopg.connect( + urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{delivery_database.info.dbname}")) + ) as locker: + locker.execute("LOCK TABLE totals IN ACCESS EXCLUSIVE MODE") + for _ in range(4): + result = await repository.apply_pending(apply, max_attempts=2) + assert result.failed == 1 and result.parked == () + delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at = CURRENT_TIMESTAMP') + assert delivery_database.execute('SELECT state::text, attempts FROM "LiteLLM_SpendLogDelivery"').fetchall() == [ + ("pending", 4) + ] + assert (await repository.apply_pending(apply, max_attempts=2)).applied == 1 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + + +@pytest.mark.asyncio +async def test_an_unreachable_database_reported_as_a_data_error_never_parks_a_delivery( + delivery_database, delivery_client +): + from prisma.errors import DataError + + _insert(delivery_database, (_row("valid"),)) + + async def unreachable(transaction, deliveries): + raise DataError(data={"user_facing_error": {"message": "Can't reach database server at `db`:`5432`"}}) + + repository = AccountingDeliveryRepository(delivery_client) + for _ in range(3): + assert (await repository.apply_pending(unreachable, max_attempts=2)).parked == () + delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at = CURRENT_TIMESTAMP') + assert delivery_database.execute('SELECT state::text, attempts FROM "LiteLLM_SpendLogDelivery"').fetchall() == [ + ("pending", 3) + ] + + +@pytest_asyncio.fixture(loop_scope="function") +async def delivery_client(delivery_database: psycopg.Connection[tuple[object, ...]]) -> AsyncIterator[Prisma]: + url: Final = urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{delivery_database.info.dbname}")) + client: Final = Prisma(datasource={"url": url}) + delivery_database.execute("CREATE TABLE totals (spend float8 NOT NULL)") + delivery_database.execute("INSERT INTO totals VALUES (0)") + await client.connect() + try: + yield client + finally: + await client.disconnect() + + +@pytest.mark.asyncio +async def test_two_workers_claim_disjoint_rows_and_apply_once( + delivery_database: psycopg.Connection[tuple[object, ...]], + delivery_client: Prisma, +) -> None: + _insert(delivery_database, (_row("first"), _row("second"))) + entered: Final[asyncio.Queue[str]] = asyncio.Queue() + release: Final = asyncio.Event() + + async def apply(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None: + assert len(rows) == 1 + await entered.put(str(rows[0].row["request_id"])) + await release.wait() + await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(str(rows[0].row["spend"]))) + + repository: Final = AccountingDeliveryRepository(delivery_client) + first: Final = asyncio.create_task(repository.apply_pending(apply, batch_size=1)) + first_id: Final = await asyncio.wait_for(entered.get(), timeout=10) + second: Final = asyncio.create_task(repository.apply_pending(apply, batch_size=1)) + try: + second_id: Final = await asyncio.wait_for(entered.get(), timeout=10) + assert first_id != second_id + finally: + release.set() + results: Final = await asyncio.gather(first, second) + assert sum(result.applied for result in results) == 2 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (1.0,) + assert delivery_database.execute( + "SELECT COUNT(*) FROM \"LiteLLM_SpendLogDelivery\" WHERE state = 'applied'" + ).fetchone() == (2,) + + +@pytest.mark.asyncio +async def test_cancelled_worker_rolls_back_and_restart_applies_once( + delivery_database: psycopg.Connection[tuple[object, ...]], + delivery_client: Prisma, +) -> None: + _insert(delivery_database, (_row("restart"),)) + wrote: Final = asyncio.Event() + hold: Final = asyncio.Event() + + async def interrupted(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None: + await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(str(rows[0].row["spend"]))) + wrote.set() + await hold.wait() + + repository: Final = AccountingDeliveryRepository(delivery_client) + task: Final = asyncio.create_task(repository.apply_pending(interrupted)) + await asyncio.wait_for(wrote.wait(), timeout=10) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.0,) + assert delivery_database.execute('SELECT state::text, attempts FROM "LiteLLM_SpendLogDelivery"').fetchone() == ( + "pending", + 0, + ) + + async def complete(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None: + await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(str(rows[0].row["spend"]))) + + assert (await repository.apply_pending(complete)).applied == 1 + assert (await repository.apply_pending(complete)).applied == 0 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + + +@pytest.mark.asyncio +async def test_worker_applies_accounting_after_spend_logs_are_deleted( + delivery_database: psycopg.Connection[tuple[object, ...]], delivery_client: Prisma +) -> None: + _insert(delivery_database, (_row("removed"),)) + delivery_database.execute('DELETE FROM "LiteLLM_SpendLogs"') + + async def apply(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None: + await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(str(rows[0].row["spend"]))) + + repository: Final = AccountingDeliveryRepository(delivery_client) + assert (await repository.apply_pending(apply)).applied == 1 + assert (await repository.apply_pending(apply)).applied == 0 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + + +@pytest.mark.asyncio +async def test_receipt_cleanup_keeps_pending_and_parked_rows_and_reports_lag( + delivery_database: psycopg.Connection[tuple[object, ...]], delivery_client: Prisma +) -> None: + _insert(delivery_database, tuple(_row(state) for state in ("pending", "parked", "applied"))) + delivery_database.execute( + 'UPDATE "LiteLLM_SpendLogDelivery" ' + 'SET state = request_id::"AccountingDeliveryState", ' + "created_at = CURRENT_TIMESTAMP - INTERVAL '40 days', " + "applied_at = CURRENT_TIMESTAMP - INTERVAL '40 days'" + ) + repository = AccountingDeliveryRepository(delivery_client) + assert await repository.cleanup_applied() == 1 + metrics = await repository.metrics() + assert metrics.pending_count == 1 + assert metrics.parked_count == 1 + assert metrics.oldest_pending_seconds >= 40 * 86400 + assert delivery_database.execute( + 'SELECT request_id FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id' + ).fetchall() == [("parked",), ("pending",)] + + +@pytest.mark.parametrize("delivery_database", ["public", "ledger_tenant"], indirect=True) +@pytest.mark.parametrize("state", [None, "pending", "parked", "applied"]) +def test_delivery_migration_replay_preserves_ledger_receipts(delivery_database, state): + if state is not None: + _insert(delivery_database, (_row("replay"),)) + delivery_database.execute( + 'UPDATE "LiteLLM_SpendLogDelivery" SET state=%s, attempts=3, last_error=%s', + (state, "retained failure"), + ) + before = delivery_database.execute('SELECT * FROM "LiteLLM_SpendLogDelivery"').fetchall() + migration = (_MIGRATIONS / "20261103010000_spendlogs_accounting_delivery/migration.sql").read_text() + delivery_database.execute(migration) + delivery_database.execute(migration) + assert delivery_database.execute('SELECT * FROM "LiteLLM_SpendLogDelivery"').fetchall() == before + assert delivery_database.execute('SELECT count(*) FROM "LiteLLM_SpendLogs"').fetchone() == (int(state is not None),) + + +@pytest.mark.parametrize("state", ["pending", "parked"]) +def test_delivery_migration_preserves_unused_old_receipts(state): + with scratch_database("spendlogs_old_delivery") as url, psycopg.connect(url, autocommit=True) as connection: + connection.execute((_MIGRATIONS / "20261027000000_add_accounting_ledger_tables/migration.sql").read_text()) + connection.execute('CREATE TABLE "LiteLLM_SpendLogs" (request_id text PRIMARY KEY)') + connection.execute( + 'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version, state) VALUES (%s, %s, 1, %s)', + ("old-event", "LiteLLM_UserTable", state), + ) + connection.execute((_MIGRATIONS / "20261103010000_spendlogs_accounting_delivery/migration.sql").read_text()) + assert connection.execute('SELECT event_id, state::text FROM "LiteLLM_AccountingDelivery"').fetchall() == [ + ("old-event", state) + ] + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,) + + +@pytest.mark.asyncio +async def test_malformed_payload_is_parked_without_losing_other_deliveries(delivery_database, delivery_client): + _insert(delivery_database, (_row("bad"), _row("good"))) + delivery_database.execute( + 'UPDATE "LiteLLM_SpendLogDelivery" SET accounting_payload=$$ [] $$::jsonb WHERE request_id=$$bad$$' + ) + + async def apply(transaction, deliveries): + await transaction.execute_raw( + "UPDATE totals SET spend = spend + $1::float8", sum(float(row.row["spend"]) for row in deliveries) + ) + + repository = AccountingDeliveryRepository(delivery_client) + assert (await repository.apply_pending(apply, max_attempts=2)).failed == 2 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.0,) + delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at=CURRENT_TIMESTAMP') + assert (await repository.apply_pending(apply, max_attempts=2)).parked == ("bad",) + assert (await repository.apply_pending(apply, max_attempts=2)).applied == 1 + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + + +@pytest.mark.asyncio +async def test_retention_keeps_batch_receipts_indefinitely(delivery_database, delivery_client): + _insert(delivery_database, (_row("normal"), {**_row("batch"), "call_type": "aretrieve_batch"})) + delivery_database.execute( + "UPDATE \"LiteLLM_SpendLogDelivery\" SET state='applied', applied_at=CURRENT_TIMESTAMP - INTERVAL '400 days'" + ) + repository = AccountingDeliveryRepository(delivery_client) + assert await repository.cleanup_applied() == 1 + assert await repository.cleanup_applied() == 0 + assert delivery_database.execute('SELECT request_id FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("batch",)] + metrics = await repository.metrics() + assert metrics.batch_count == 1 + assert metrics.table_bytes > 0 + + +@pytest.mark.parametrize("disabled", [False, True]) +@pytest.mark.asyncio +async def test_ledger_batch_retrieves_charge_once_in_every_logging_mode( + delivery_database, delivery_client, monkeypatch, disabled +): + from types import SimpleNamespace + + from litellm.proxy import proxy_server + from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + from litellm.proxy.utils import jsonify_object + + completed = { + **_row("batch_completed"), "call_type": "aretrieve_batch", "status": "success", "spend": 0.5, + "metadata": '{"status":"success","private":"request detail"}', + "accounting_payload": { + "request_id": "batch_completed", "call_type": "aretrieve_batch", "status": "success", "spend": 0.5, + "metadata": {"accounting_savings": {"compression": 0.2}}, + }, + } + monkeypatch.setattr(proxy_server, "general_settings", {"spend_logs_metadata_fields": {"include": ["status"]}}) + client = SimpleNamespace(db=delivery_client, writer_db=delivery_client, jsonify_object=jsonify_object) + writer = DBSpendUpdateWriter() + first = await writer._claim_settle_once_cost_row(payload=completed, prisma_client=client, disable_spend_logs=disabled) + again = await writer._claim_settle_once_cost_row( + payload={**completed, "startTime": "2026-11-05T12:00:00+00:00"}, prisma_client=client, disable_spend_logs=disabled + ) + assert (first, again) == (True, False) + assert delivery_database.execute('SELECT spend, metadata FROM "LiteLLM_SpendLogs"').fetchall() == ( + [] if disabled else [(0.5, {"status": "success"})] + ) + assert delivery_database.execute( + 'SELECT delivery_id, accounting_payload FROM "LiteLLM_SpendLogDelivery"' + ).fetchall() == [("batch_completed", completed["accounting_payload"])] + + +@pytest.mark.parametrize("setting", ["retention_days", "batch_size"]) +@pytest.mark.asyncio +async def test_invalid_cleanup_settings_preserve_receipts(delivery_database, delivery_client, setting): + _insert(delivery_database, (_row("protected"),)) + with pytest.raises(ValueError, match="positive"): + await AccountingDeliveryRepository(delivery_client).cleanup_applied(**{setting: 0}) + assert delivery_database.execute('SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [ + ("protected", "pending") + ] + + +@pytest.mark.parametrize("setting", ["max_attempts", "batch_size"]) +@pytest.mark.asyncio +async def test_invalid_worker_settings_preserve_pending_work(delivery_database, delivery_client, setting): + from unittest.mock import AsyncMock + + _insert(delivery_database, (_row("protected"),)) + apply = AsyncMock() + with pytest.raises(ValueError, match="positive"): + await AccountingDeliveryRepository(delivery_client).apply_pending(apply, **{setting: 0}) + apply.assert_not_awaited() + assert delivery_database.execute( + 'SELECT request_id, state::text, attempts FROM "LiteLLM_SpendLogDelivery"' + ).fetchall() == [("protected", "pending", 0)] diff --git a/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py new file mode 100644 index 00000000000..e1894aa66bf --- /dev/null +++ b/tests/integration/database/test_spendlogs_ledger_parity.py @@ -0,0 +1,1437 @@ +import contextlib +import asyncio +import json +import os +import uuid +import subprocess +import sys +import time +from collections.abc import Iterator +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Final + +import psycopg +from psycopg import sql +import pytest +import litellm + +from litellm.accounting.application import apply_deliveries +from litellm.caching.dual_cache import DualCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.proxy.spend_tracking.spend_tracking_utils import spend_log_row_with_retained_metadata +from litellm.proxy import proxy_server +from litellm.proxy._types import SpendLogsMetadataFields +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.db.model_usage_rollup import flush_model_usage_transactions +from litellm.proxy.guardrails.usage_tracking import PendingRollups, process_spend_logs_guardrail_usage +from litellm.proxy.utils import PrismaClient +from litellm.proxy.hooks.model_budget_shared_windows import DatabaseSharedWindows +from litellm.proxy.hooks.model_max_budget_limiter import PROXY_VirtualKeyModelMaxBudgetLimiter +from litellm.repositories.accounting_delivery_repository import AccountingDeliveryRepository +from litellm.repositories.table_repositories import SpendLogsRepository +from litellm.responses.streaming_iterator import drain_logging_tasks +from litellm.types.utils import ModelResponse, Usage +from tests.integration._support.database import scratch_database + +_ROOT: Final = Path(__file__).resolve().parents[3] +_NOW: Final = datetime(2026, 11, 2, tzinfo=timezone.utc) +_ENTITY_TABLES: Final = ( + "LiteLLM_VerificationToken", + "LiteLLM_UserTable", + "LiteLLM_TeamTable", + "LiteLLM_TeamMembership", + "LiteLLM_OrganizationTable", + "LiteLLM_OrganizationMembership", + "LiteLLM_EndUserTable", + "LiteLLM_AgentsTable", + "LiteLLM_TagTable", + "LiteLLM_ProjectTable", + "LiteLLM_ModelAccessGroupBudgetTable", +) +_DAILY_TABLES: Final = ( + "LiteLLM_DailyUserSpend", + "LiteLLM_DailyTeamSpend", + "LiteLLM_DailyOrganizationSpend", + "LiteLLM_DailyEndUserSpend", + "LiteLLM_DailyAgentSpend", + "LiteLLM_DailyTagSpend", + "LiteLLM_DailyModelUsage", + "LiteLLM_DailyGuardrailUsageUnits", + "LiteLLM_BudgetWindowSpend", +) + + +@pytest.mark.parametrize("batch_size,row_count,expected,distinct_models", [(None, 251, 251, False), ("200", 251, 200, False), (None, 10001, 10000, False), (None, 3001, 3001, True)]) +@pytest.mark.asyncio +async def test_worker_applies_the_configured_backlog_in_one_transaction( + parity_database: str, monkeypatch: pytest.MonkeyPatch, batch_size: str | None, row_count: int, expected: int, distinct_models: bool +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + if batch_size is None: + monkeypatch.delenv("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", raising=False) + else: + monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", batch_size) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + writer = DBSpendUpdateWriter() + await client.db.connect() + try: + assert await _retrieve_batch(writer, 0.25) + stored = await client.writer_db.query_raw('SELECT to_jsonb(s) AS row FROM "LiteLLM_SpendLogs" s') + delivery = (await client.writer_db.query_raw('SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"'))[0]["accounting_payload"] + assert await SpendLogsRepository(client).insert_with_deliveries(tuple( + {**stored[0]["row"], "request_id": f"backlog-{index:05}", "model": f"backlog-model-{index}" if distinct_models else stored[0]["row"]["model"], "accounting_payload": {**delivery, "request_id": f"backlog-{index:05}", "model": f"backlog-model-{index}" if distinct_models else delivery["model"]}} for index in range(row_count - 1) + )) == row_count - 1 + started = time.monotonic() + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + print(f"Ledger batch applied={expected} queued={row_count} seconds={time.monotonic() - started:.3f}") + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'applied\'').fetchone() == (expected,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'pending\'').fetchone() == (row_count - expected,) + assert connection.execute('SELECT COUNT(DISTINCT applied_at) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (1,) + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (expected * 0.25,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_DailyUserSpend"').fetchone() == (expected if distinct_models else 1,) + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize("ledger", [False, True]) +@pytest.mark.parametrize("exclude_messages", [False, True]) +@pytest.mark.asyncio +async def test_registered_proxy_callbacks_charge_model_windows_before_persisting_spend( + parity_database: str, monkeypatch: pytest.MonkeyPatch, ledger: bool, exclude_messages: bool +): + from litellm.accounting.worker import flush_accounting_deliveries + from litellm.utils import Rules, function_setup + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setenv("LITELLM_COLLECTOR_ENABLED", "false") + for name in ("callbacks", "input_callback", "success_callback", "failure_callback", "_async_success_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, name, []) + monkeypatch.setattr(litellm, "standard_logging_payload_excluded_fields", ["messages"] if exclude_messages else []) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": ledger}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + client.model_usage_transactions = [] + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server.proxy_logging_obj, "db_spend_update_writer", writer) + limiter = PROXY_VirtualKeyModelMaxBudgetLimiter( + dual_cache=DualCache(), clock=lambda: _NOW.timestamp(), + shared_windows=DatabaseSharedWindows(prisma_client_source=lambda: client), + ) + monkeypatch.setattr(proxy_server, "model_max_budget_limiter", limiter) + litellm.logging_callback_manager.add_litellm_callback(limiter) + proxy_server.cost_tracking() + logging, _ = function_setup("acompletion", Rules(), _NOW, model="ledger-test", messages=[], litellm_call_id="registered-window") + logging.update_environment_variables( + optional_params={}, custom_llm_provider="openai", + litellm_params={"metadata": {"user_api_key": "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", "user_api_key_model_max_budget": { + "ledger-test": {"max_budget": 100, "budget_duration": "1d"} + }}}, + ) + logging.model_call_details["response_cost"] = 0.25 + response = ModelResponse(id="registered-window", model="ledger-test", usage=Usage(prompt_tokens=10, completion_tokens=5)) + await client.db.connect() + try: + await logging.async_success_handler(response, _NOW, _NOW, standard_logging_object={ + "model": "ledger-test", "model_group": "ledger-test", "response_cost": 0.25, + "metadata": {"user_api_key_hash": "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683"}, + "model_map_information": {}, "hidden_params": {}, "request_tags": [], + }) + await drain_logging_tasks() + await SpendLogsRepository(client).insert_with_deliveries(tuple(client.spend_log_transactions)) + if ledger: + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + await writer._commit_spend_updates_to_db_without_redis_buffer(client, 0, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT spend FROM "LiteLLM_BudgetWindowSpend" WHERE entity_type=\'model_max_budget\'').fetchall() == [(0.25,)] + assert connection.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == [(0.25,)] + finally: + await client.db.disconnect() + + +@pytest.fixture(scope="module") +def ledger_schema() -> str: + return subprocess.run( + [ + sys.executable, + "-I", + "-m", + "prisma", + "migrate", + "diff", + "--from-empty", + "--to-schema-datamodel", + str(_ROOT / "schema.prisma"), + "--script", + ], + check=True, + capture_output=True, + text=True, + cwd=_ROOT, + ).stdout + + +@pytest.fixture +def parity_database(ledger_schema: str) -> Iterator[str]: + with scratch_database("spendlogs_parity") as url, psycopg.connect(url, autocommit=True) as connection: + connection.execute(ledger_schema) + connection.execute( + "INSERT INTO \"LiteLLM_BudgetTable\" (budget_id, created_by, updated_by) VALUES ('budget', 'test', 'test')" + ) + connection.execute("INSERT INTO \"LiteLLM_UserTable\" (user_id) VALUES ('user')") + connection.execute("INSERT INTO \"LiteLLM_TeamTable\" (team_id) VALUES ('team')") + connection.execute( + 'INSERT INTO "LiteLLM_OrganizationTable" ' + "(organization_id, organization_alias, budget_id, created_by, updated_by) " + "VALUES ('org', 'org', 'budget', 'test', 'test')" + ) + connection.execute("INSERT INTO \"LiteLLM_TeamMembership\" (team_id, user_id) VALUES ('team', 'user')") + connection.execute( + "INSERT INTO \"LiteLLM_OrganizationMembership\" (organization_id, user_id) VALUES ('org', 'user')" + ) + connection.execute("INSERT INTO \"LiteLLM_VerificationToken\" (token) VALUES ('2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683')") + connection.execute("INSERT INTO \"LiteLLM_EndUserTable\" (user_id) VALUES ('end-user')") + connection.execute( + 'INSERT INTO "LiteLLM_AgentsTable" ' + "(agent_id, agent_name, agent_card_params, created_by, updated_by) " + "VALUES ('agent', 'agent', '{}', 'test', 'test')" + ) + connection.execute( + 'INSERT INTO "LiteLLM_ProjectTable" (project_id, created_by, updated_by) ' + "VALUES ('project', 'test', 'test')" + ) + connection.execute("INSERT INTO \"LiteLLM_TagTable\" (tag_name) VALUES ('tag')") + connection.execute("INSERT INTO \"LiteLLM_ModelAccessGroupBudgetTable\" (access_group_name) VALUES ('group')") + yield url + + +def _snapshot(url: str) -> dict[str, list[object]]: + with psycopg.connect(url) as connection: + return { + table: [ + row[0] + for row in connection.execute( + sql.SQL( + "SELECT to_jsonb(t) - 'id' - 'created_at' - 'updated_at' - 'request_id' - 'last_active' FROM {} t " + "ORDER BY (to_jsonb(t) - 'id' - 'created_at' - 'updated_at' - 'request_id' - 'last_active')::text" + ).format(sql.Identifier(table)) + ).fetchall() + ] + for table in (*_ENTITY_TABLES, *_DAILY_TABLES) + } + + +def _reset(url: str) -> None: + with psycopg.connect(url, autocommit=True) as connection: + for table in _ENTITY_TABLES: + connection.execute(sql.SQL("UPDATE {} SET spend=0").format(sql.Identifier(table))) + for table in (*_DAILY_TABLES, "LiteLLM_SpendLogs", "LiteLLM_SpendLogDelivery"): + connection.execute(sql.SQL("TRUNCATE {}").format(sql.Identifier(table))) + connection.execute('UPDATE "LiteLLM_VerificationToken" SET total_spend=0') + connection.execute('UPDATE "LiteLLM_TeamMembership" SET total_spend=0') + + +async def _assert_legacy_database_parity( + parity_database: str, + monkeypatch: pytest.MonkeyPatch, + status: str, + origin: str | None, + cost: float, + retention: str, + reset_kind: str | None = None, + correction_order: str | None = None, + toggle_after_capture: bool = False, + fallback_narrowing: bool = False, + failing_table: str | None = None, + organization_id: str | None = "org", + export_url: str | None = None, + export_failure: bool = False, +): + monkeypatch.setitem( + litellm.model_cost, + "ledger-test", + { + "input_cost_per_token": 0.002, + "output_cost_per_token": 0.004, + "cache_read_input_token_cost": 0.0005, + "cache_creation_input_token_cost": 0.003, + "litellm_provider": "openai", + "mode": "chat", + "max_tokens": 4096, + }, + ) + if fallback_narrowing: + monkeypatch.setattr(proxy_server, "llm_router", litellm.Router(model_list=[{ + "model_name": "ledger-test", + "litellm_params": {"model": "openai/ledger-test", "api_key": "test"}, + "model_info": {"id": "deployment", "access_groups": ["group"]}, + }])) + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + "INSERT INTO \"LiteLLM_ModelAccessGroupBudgetTable\" (access_group_name) VALUES ('unserved-group')" + ) + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + client.model_usage_transactions = [] + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + await client.db.connect() + results = [] + try: + for path in ("legacy", "ledger"): + _reset(parity_database) + if reset_kind is not None: + _prime_reset(parity_database, reset_kind) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": path == "ledger"}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", path == "ledger" and retention == "lean") + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server.proxy_logging_obj, "db_spend_update_writer", writer) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr(proxy_server, "spend_counter_cache", DualCache()) + for cache_key in ("2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", "team_id:team"): + await proxy_server.user_api_key_cache.async_set_cache( + key=cache_key, + value={ + "budget_limits": [ + {"budget_duration": "1d", "max_budget": 100, "reset_at": "2026-11-04T00:00:00+00:00"} + ] + }, + ) + request_kwargs = { + "model": "ledger-test", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "request", + "standard_logging_object": { + "model": "ledger-test", + "model_group": "ledger-test", + "metadata": { + "user_api_key_hash": "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", + "usage_object": { + "prompt_tokens": 10, + "completion_tokens": 5, + "cache_read_input_tokens": 4, + "cache_creation_input_tokens": 1, + }, + }, + "model_map_information": {}, + "hidden_params": {}, + "request_tags": ["tag"], + "model_id": "deployment", + "response_cost": cost, + "autorouter_savings": 0.125, + "guardrail_information": [ + { + "guardrail_id": "guard", + "guardrail_name": "guard", + "guardrail_status": "success", + "guardrail_usage": {"text": 2}, + "guardrail_cost_by_unit": {"text": 0.02}, + "guardrail_cost_in_spend": True, + } + ], + }, + "agent_id": "agent", + "user": "end-user", + "litellm_params": { + "metadata": { + "user_api_key": "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", + "user_api_key_model_max_budget": {"ledger-test": {"max_budget": 100, "budget_duration": "1d"}}, + "user_api_key_user_id": "user", + "user_api_key_team_id": "team", + "user_api_key_org_id": organization_id, + "user_api_key_project_id": "project", + "user_api_key_end_user_id": "end-user", + "tags": ["tag"], + "user_api_key_matched_model_access_groups": ["group", "unserved-group"] if fallback_narrowing else ["group"], + "model_info": {"id": "deployment"}, + "litellm_gateway_injected_cache": "deployment", + "compression_savings": {"tokens_before": 20, "tokens_after": 10, "tokens_saved": 10}, + "status": status, + **({BILLED_RETRY_ATTEMPT_METADATA_KEY: True} if origin == "billed_retry" else {}), + **({INTERNAL_CALL_ORIGIN_METADATA_KEY: origin} if origin == "auto_router" else {}), + }, + "user": "end-user", + }, + } + limiter = PROXY_VirtualKeyModelMaxBudgetLimiter( + dual_cache=proxy_server.spend_counter_cache, + clock=lambda: _NOW.timestamp(), + shared_windows=DatabaseSharedWindows(prisma_client_source=lambda: client), + ) + if status == "success": + await limiter.async_log_success_event(request_kwargs, None, _NOW, _NOW) + charged = await writer.update_database( + token="2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", + user_id="user", + end_user_id="end-user", + team_id="team", + org_id=organization_id, + project_id="project", + kwargs=request_kwargs, + completion_response=ModelResponse( + id="request", model="ledger-test", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, + end_time=_NOW, + response_cost=cost, + ) + assert charged + await proxy_server.increment_spend_counters( + token="2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", + team_id="team", + user_id="user", + response_cost=cost, + request_started_at=_NOW, + **({"windows_from_ledger": True} if path == "ledger" else {}), + ) + await drain_logging_tasks() + rows = tuple(client.spend_log_transactions) + client.spend_log_transactions = [] + assert len(rows) == 1 + if path == "ledger" and retention == "lean": + assert rows[0]["accounting_log_disabled"] is True + frozen = rows[0]["accounting_payload"] + if isinstance(frozen, str): + frozen = json.loads(frozen) + assert not {"messages", "response", "proxy_server_request"}.intersection(frozen) + stored_rows = ( + tuple(spend_log_row_with_retained_metadata(row, SpendLogsMetadataFields(include=())) for row in rows) + if retention == "stripped" + else rows + ) + if toggle_after_capture: + for enabled in (True, False, True): + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": enabled}) + await SpendLogsRepository(client).insert_with_deliveries(stored_rows) + elif path == "ledger" and export_url is not None: + import httpx + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.proxy.utils import ProxyUpdateSpend + + monkeypatch.setenv("SPEND_LOGS_URL", export_url) + handler = AsyncHTTPHandler(timeout=10) + wire_rows = [ + {key: value.isoformat() if isinstance(value, datetime) else value for key, value in row.items()} + for row in stored_rows + ] + try: + if export_failure: + with pytest.raises(httpx.HTTPStatusError): + await ProxyUpdateSpend.update_spend_logs(0, client, handler, proxy_server.proxy_logging_obj, wire_rows) + else: + await ProxyUpdateSpend.update_spend_logs(0, client, handler, proxy_server.proxy_logging_obj, wire_rows) + finally: + await handler.close() + else: + await SpendLogsRepository(client).insert_with_deliveries(stored_rows) + if path == "ledger" and (retention == "lean" or export_url is not None): + assert await client.writer_db.query_raw('SELECT COUNT(*) AS count FROM "LiteLLM_SpendLogs"') == [{"count": 0}] + if reset_kind is not None: + await _apply_reset(reset_kind, client) + if correction_order == "before": + await _publish_savings_correction(client, stored_rows[0]) + if path != "ledger": + await writer._commit_spend_updates_to_db_without_redis_buffer(client, 0, proxy_server.proxy_logging_obj) + await writer.commit_daily_tag_spend_to_db(client, 0, proxy_server.proxy_logging_obj) + await flush_model_usage_transactions(prisma_client=client, transactions=client.model_usage_transactions) + client.model_usage_transactions = [] + await process_spend_logs_guardrail_usage(client, rows, pending=PendingRollups()) + else: + + async def apply(transaction, deliveries, writer: DBSpendUpdateWriter = writer): + await apply_deliveries( + transaction, + deliveries, + prisma_client=client, + writer=writer, + proxy_logging_obj=proxy_server.proxy_logging_obj, + ) + + if failing_table is not None: + before_failure = _snapshot(parity_database) + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute(sql.SQL('ALTER TABLE {} ADD CONSTRAINT reject_spend CHECK (spend = 0)').format(sql.Identifier(failing_table))) + failed = await AccountingDeliveryRepository(client.writer_db).apply_pending(apply) + assert (failed.applied, failed.failed) == (0, 1) + assert _snapshot(parity_database) == before_failure + with psycopg.connect(parity_database, autocommit=True) as connection: + assert connection.execute('SELECT state::text, attempts FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("pending", 1)] + connection.execute(sql.SQL('ALTER TABLE {} DROP CONSTRAINT reject_spend').format(sql.Identifier(failing_table))) + connection.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at=CURRENT_TIMESTAMP') + result = await AccountingDeliveryRepository(client.writer_db).apply_pending(apply) + assert result.applied == 1, result + await writer._commit_spend_updates_to_db_without_redis_buffer(client, 0, proxy_server.proxy_logging_obj) + if correction_order == "after": + await _publish_savings_correction(client, stored_rows[0]) + results.append(_snapshot(parity_database)) + differences = { + (index, table): (results[0][table], result[table]) + for index, result in enumerate(results[1:], 1) + for table in results[0] + if results[0][table] != result[table] + } + assert not differences, differences + prior_spend = 5 if reset_kind is not None and reset_kind != "scheduled" else 0 + assert results[-1]["LiteLLM_UserTable"][0]["spend"] == cost + prior_spend + daily = results[-1]["LiteLLM_DailyUserSpend"][0] + assert daily["compression_savings_spend"] > 0, results[-1]["LiteLLM_DailyUserSpend"] + assert daily["prompt_caching_savings_spend"] != 0 + assert daily["gateway_injected_caching_savings_spend"] != 0 + assert daily["autorouter_savings_spend"] == pytest.approx( + (0 if origin == "auto_router" else 0.125) + (0.2 if correction_order else 0) + ) + group_spend = {row["access_group_name"]: row["spend"] for row in results[-1]["LiteLLM_ModelAccessGroupBudgetTable"]} + assert group_spend == {"group": cost + prior_spend, **({"unserved-group": prior_spend} if fallback_narrowing else {})} + assert results[-1]["LiteLLM_DailyGuardrailUsageUnits"] + windows = results[-1]["LiteLLM_BudgetWindowSpend"] + assert len(windows) == (2 if cost else 0) + (1 if status == "success" else 0) + assert all(row["spend"] == cost for row in windows) + if reset_kind == "window_rollover": + model_window = next(row for row in windows if row["entity_type"] == "model_max_budget") + assert model_window["window_start"].startswith("2026-11-04") + if reset_kind in ("key", "global", "scheduled"): + assert results[-1]["LiteLLM_VerificationToken"][0]["spend"] == cost + if reset_kind in ("global", "scheduled"): + assert results[-1]["LiteLLM_TeamTable"][0]["spend"] == cost + if reset_kind == "team_member": + assert results[-1]["LiteLLM_TeamMembership"][0]["spend"] == cost + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize( + "status,origin,cost", + [ + ("success", None, 0.25), + ("failure", None, 0.25), + ("success", None, 0.0), + ("success", "auto_router", 0.25), + ("failure", "billed_retry", 0.25), + ], +) +@pytest.mark.parametrize("retention", ["full", "stripped", "lean"]) +@pytest.mark.asyncio +async def test_stored_ledger_matches_real_legacy_database_updates( + parity_database: str, monkeypatch: pytest.MonkeyPatch, status: str, origin: str | None, cost: float, retention: str +): + await _assert_legacy_database_parity(parity_database, monkeypatch, status, origin, cost, retention) + + +@pytest.mark.parametrize("retention", ["full", "stripped", "lean"]) +@pytest.mark.asyncio +async def test_missing_organization_matches_legacy_daily_rollups(parity_database, monkeypatch, retention): + await _assert_legacy_database_parity( + parity_database, monkeypatch, "success", None, 0.25, retention, organization_id=None + ) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT count(*) FROM "LiteLLM_DailyOrganizationSpend"').fetchone() == (0,) + + +@pytest.mark.parametrize("reset_kind", ["scheduled", "key", "team_member", "global", "window_rollover"]) +@pytest.mark.asyncio +async def test_pending_delivery_after_reset_matches_queued_legacy_spend( + parity_database: str, monkeypatch: pytest.MonkeyPatch, reset_kind: str +): + await _assert_legacy_database_parity(parity_database, monkeypatch, "success", None, 0.25, "full", reset_kind) + + +def _prime_reset(url: str, reset_kind: str) -> None: + with psycopg.connect(url, autocommit=True) as connection: + for table in _ENTITY_TABLES: + connection.execute(sql.SQL("UPDATE {} SET spend=5").format(sql.Identifier(table))) + if reset_kind == "scheduled": + for table in ("LiteLLM_VerificationToken", "LiteLLM_UserTable", "LiteLLM_TeamTable", "LiteLLM_BudgetTable"): + connection.execute( + sql.SQL("UPDATE {} SET budget_duration='1d', budget_reset_at='2026-11-01'").format( + sql.Identifier(table) + ) + ) + connection.execute("UPDATE \"LiteLLM_ModelAccessGroupBudgetTable\" SET budget_id='budget'") + + +async def _apply_reset(reset_kind: str, client: PrismaClient) -> None: + from litellm.proxy._types import ResetSpendRequest, UserAPIKeyAuth, LitellmUserRoles + from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob + from litellm.proxy.management_endpoints.key_management_endpoints import reset_key_spend_fn + from litellm.proxy.management_endpoints.team_endpoints import reset_team_member_spend_fn + from litellm.proxy.spend_tracking.spend_management_endpoints import global_spend_reset + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + if reset_kind == "window_rollover": + from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row + + windows = await client.writer_db.query_raw('SELECT * FROM "LiteLLM_BudgetWindowSpend"') + assert windows + for window in windows: + await roll_window_spend_row( + client, window["entity_type"], window["entity_id"], window["window_duration"], _NOW + timedelta(days=2) + ) + elif reset_kind == "global": + await global_spend_reset() + elif reset_kind == "key": + await reset_key_spend_fn( + key="2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", data=ResetSpendRequest(reset_to=0), user_api_key_dict=admin, litellm_changed_by=None + ) + elif reset_kind == "team_member": + await reset_team_member_spend_fn( + team_id="team", user_id="user", data=ResetSpendRequest(reset_to=0), user_api_key_dict=admin + ) + else: + await ResetBudgetJob(proxy_server.proxy_logging_obj, client).reset_budget() + + +@pytest.mark.parametrize("correction_order", ["before", "after"]) +@pytest.mark.asyncio +async def test_baseline_correction_order_preserves_frozen_request_savings( + parity_database: str, monkeypatch: pytest.MonkeyPatch, correction_order: str +): + await _assert_legacy_database_parity( + parity_database, monkeypatch, "success", None, 0.25, "full", correction_order=correction_order + ) + + +async def _publish_savings_correction(client: PrismaClient, row) -> None: + from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING + from litellm.proxy.db.baseline_accounting import ( + BaselinePublication, + DailyBaselineAttribution, + DailyBaselineTarget, + _Change, + _publish, + ) + + metadata_before = ( + await client.writer_db.query_raw( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=$1', row["request_id"] + ) + )[0]["metadata"] + if isinstance(metadata_before, str): + metadata_before = json.loads(metadata_before) + change = _Change( + request_id=row["request_id"], + publication=BaselinePublication( + comparison_id="autorouter-baseline:v3:" + "a" * 64, + comparison_started_at=_NOW.timestamp(), + status="estimated", + reason="modeled", + provenance="modeled", + actual_spend=0.25, + baseline_spend=0.575, + ), + api_key=row["api_key"], + session_id="session", + router_name="router", + baseline_model="ledger-test", + covered_delta=0, + actual_delta=0.0, + savings_delta=0.2, + daily=DailyBaselineAttribution( + date=_NOW.date().isoformat(), + api_key=row["api_key"], + model=row["model"], + custom_llm_provider=row["custom_llm_provider"], + model_group=row["model_group"], + endpoint=ROUTE_ENDPOINT_MAPPING.get(row["call_type"]), + targets=tuple( + DailyBaselineTarget(entity=entity, entity_id=identity) + for entity, identity in ( + ("user", "user"), + ("team", "team"), + ("org", "org"), + ("end_user", "end-user"), + ("agent", "agent"), + ("tag", "tag"), + ) + ), + ), + ) + async with client.writer_db.tx() as transaction: + await _publish(transaction, (change,)) + metadata_after = ( + await client.writer_db.query_raw( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=$1', row["request_id"] + ) + )[0]["metadata"] + if isinstance(metadata_after, str): + metadata_after = json.loads(metadata_after) + assert metadata_after["autorouter_savings"] == pytest.approx(0.325) + if "accounting_payload" in row: + delivery = (await client.writer_db.query_raw( + 'SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery" WHERE request_id=$1', row["request_id"] + ))[0]["accounting_payload"] + assert delivery["metadata"]["accounting_savings"]["autorouter"] == pytest.approx(0.125) + + +@pytest.mark.asyncio +async def test_flag_changes_do_not_reclassify_or_reapply_buffered_rows( + parity_database: str, + monkeypatch: pytest.MonkeyPatch, +): + await _assert_legacy_database_parity( + parity_database, + monkeypatch, + "success", + None, + 0.25, + "stripped", + toggle_after_capture=True, + ) + + +@pytest.mark.parametrize("entity_type", ["key", "team"]) +@pytest.mark.asyncio +async def test_window_seed_excludes_pending_and_parked_but_preserves_legacy_cutoff( + parity_database: str, + monkeypatch: pytest.MonkeyPatch, + entity_type: str, +): + from datetime import timedelta + from litellm.proxy.db.budget_window_spend_writer import spend_logs_seed_totals, commit_window_spend_updates + + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + await client.db.connect() + try: + for request_id, cost, offset in ( + ("old-legacy", 2.0, -1), + ("current-legacy", 5.0, 0), + ("applied", 3.0, 1), + ("pending", 7.0, 0), + ("parked", 11.0, 0), + ): + await client.writer_db.execute_raw( + 'INSERT INTO "LiteLLM_SpendLogs" ' + '(request_id,call_type,api_key,team_id,spend,"startTime","endTime") ' + "VALUES ($1,'acompletion','key','team',$2::float8,$3::timestamp,$3::timestamp)", + request_id, + cost, + _NOW + timedelta(seconds=offset), + ) + for request_id, state in (("applied", "applied"), ("pending", "pending"), ("parked", "parked")): + await client.writer_db.execute_raw( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id,state,accounting_payload,is_batch) VALUES ($1,$2::"AccountingDeliveryState",$$ {} $$::jsonb,false)', + request_id, + state, + ) + start = _NOW - timedelta(hours=1) + totals = await spend_logs_seed_totals(client, entity_type, entity_type, start, _NOW) + assert totals is not None and totals.total == 10.0 and totals.before_batch == 5.0 + unbounded = await spend_logs_seed_totals(client, entity_type, entity_type, start, None) + assert unbounded is not None and unbounded.total == unbounded.before_batch == 10.0 + await commit_window_spend_updates( + client, + ( + { + "entity_type": entity_type, + "entity_id": entity_type, + "window_duration": "1d", + "window_start": start.isoformat(), + "started_at": _NOW.isoformat(), + "spend": 5.0, + }, + ), + ) + rows = await client.writer_db.query_raw('SELECT spend FROM "LiteLLM_BudgetWindowSpend"') + assert rows[0]["spend"] == 10.0 + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize("disable_logs", [False, True]) +@pytest.mark.asyncio +async def test_batch_claim_failure_queues_an_atomic_retry_and_repeated_retrieve_does_not_charge( + parity_database: str, monkeypatch: pytest.MonkeyPatch, disable_logs: bool +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", disable_logs) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + await client.db.connect() + + + try: + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + 'ALTER TABLE "LiteLLM_SpendLogDelivery" ADD CONSTRAINT reject_batch ' + "CHECK (request_id <> 'batch_batch_cost')" + ) + assert await _retrieve_batch(writer, 0.25) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogs"').fetchone() == (0,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,) + assert len(client.spend_log_transactions) == 1 + queued = tuple(client.spend_log_transactions) + assert "accounting_payload" in queued[0] + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute('ALTER TABLE "LiteLLM_SpendLogDelivery" DROP CONSTRAINT reject_batch') + assert await SpendLogsRepository(client).insert_with_deliveries(queued) == 1 + assert await _retrieve_batch(writer, 0.25) is False + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=\'2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683\'').fetchone() == (0.25,) + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (0.25,) + assert connection.execute('SELECT state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("applied",)] + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogs"').fetchone() == (int(not disable_logs),) + finally: + await client.db.disconnect() + + +@pytest.mark.asyncio +async def test_failed_claim_over_a_legacy_zero_batch_row_still_charges_it_once( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + client.model_usage_transactions = [] + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": False}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + await client.db.connect() + try: + assert await _retrieve_batch(writer, 0) + await drain_logging_tasks() + await SpendLogsRepository(client).insert_with_deliveries(tuple(client.spend_log_transactions)) + client.spend_log_transactions = [] + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + with psycopg.connect(parity_database, autocommit=True) as connection: + assert connection.execute( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=\'batch_batch_cost\'' + ).fetchall() == [(0.0,)] + connection.execute('ALTER TABLE "LiteLLM_SpendLogDelivery" RENAME TO "LiteLLM_SpendLogDelivery_unavailable"') + assert await _retrieve_batch(writer, 0.25) + queued = tuple(client.spend_log_transactions) + assert ["accounting_payload" in row for row in queued] == [True] + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute('ALTER TABLE "LiteLLM_SpendLogDelivery_unavailable" RENAME TO "LiteLLM_SpendLogDelivery"') + assert await SpendLogsRepository(client).insert_with_deliveries(queued) == 1 + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + assert await _retrieve_batch(writer, 0.25) is False + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=\'2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683\'').fetchone() == (0.25,) + assert connection.execute( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=\'batch_batch_cost\'' + ).fetchall() == [(0.0,)] + assert connection.execute('SELECT state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("applied",)] + finally: + await client.db.disconnect() + + +@pytest.mark.asyncio +async def test_worker_exposes_invalid_accounting_as_delayed_then_parked_work( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from prometheus_client import REGISTRY + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + writer = DBSpendUpdateWriter() + await client.db.connect() + try: + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, created_at, accounting_payload, is_batch) ' + "VALUES ('invalid-accounting', CURRENT_TIMESTAMP - INTERVAL '1 hour', " + '$$ {"metadata": {"accounting_windows": [{}]}} $$::jsonb, false)' + ) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + assert REGISTRY.get_sample_value("litellm_accounting_pending") == 1 + assert REGISTRY.get_sample_value("litellm_accounting_parked") == 0 + with psycopg.connect(parity_database, autocommit=True) as connection: + assert connection.execute( + 'SELECT state::text, attempts, last_error FROM "LiteLLM_SpendLogDelivery"' + ).fetchall() == [("pending", 1, "KeyError")] + connection.execute( + 'UPDATE "LiteLLM_SpendLogDelivery" SET attempts=4, next_attempt_at=CURRENT_TIMESTAMP' + ) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + assert REGISTRY.get_sample_value("litellm_accounting_pending") == 0 + assert REGISTRY.get_sample_value("litellm_accounting_parked") == 1 + with psycopg.connect(parity_database) as connection: + assert connection.execute( + 'SELECT state::text, attempts, last_error FROM "LiteLLM_SpendLogDelivery"' + ).fetchall() == [("parked", 5, "KeyError")] + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (0,) + finally: + await client.db.disconnect() + + +@pytest.mark.asyncio +async def test_fallback_charges_only_groups_declared_by_the_served_deployment( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + await _assert_legacy_database_parity( + parity_database, monkeypatch, "success", None, 0.25, "stripped", fallback_narrowing=True + ) + + +@pytest.mark.asyncio +async def test_worker_skips_another_leader_and_releases_its_redis_lock_after_retry( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from litellm.accounting.worker import flush_accounting_deliveries + from litellm.caching.redis_cache import RedisCache + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + + monkeypatch.setenv("DATABASE_URL", parity_database) + cache = RedisCache(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), namespace=uuid.uuid4().hex) + other_leader = PodLockManager(cache) + writer = DBSpendUpdateWriter(redis_cache=cache) + writer.pod_lock_manager.redis_cache = cache + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + await client.db.connect() + try: + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute('INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, accounting_payload, is_batch) VALUES ($$invalid-row$$, $$ {"metadata": {"accounting_windows": [{}]}} $$::jsonb, false)') + assert await other_leader.acquire_lock(cronjob_id="accounting_delivery", allow_reentrant=False) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT attempts FROM "LiteLLM_SpendLogDelivery"').fetchall() == [(0,)] + await other_leader.release_lock(cronjob_id="accounting_delivery") + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT attempts FROM "LiteLLM_SpendLogDelivery"').fetchall() == [(1,)] + assert await other_leader.acquire_lock(cronjob_id="accounting_delivery", allow_reentrant=False) + finally: + await other_leader.release_lock(cronjob_id="accounting_delivery") + await client.db.disconnect() + await cache.disconnect() + + +async def _retrieve_batch(writer: DBSpendUpdateWriter, cost: float) -> bool: + return await writer.update_database( + token="2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", user_id="user", end_user_id=None, team_id="team", org_id=None, + kwargs={ + "model": "ledger-test", "call_type": "aretrieve_batch", "litellm_call_id": "batch", + "litellm_params": {"metadata": {"status": "success", "user_api_key": "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683"}}, + }, + completion_response={"id": "batch", "status": "completed"}, + start_time=_NOW, end_time=_NOW, response_cost=cost, + ) + + +@pytest.mark.parametrize("ledger_after_zero_row,delivery_unavailable", [(False, False), (True, False), (True, True)]) +@pytest.mark.asyncio +async def test_zero_cost_legacy_batch_row_is_charged_once_across_a_flag_change( + parity_database: str, monkeypatch: pytest.MonkeyPatch, ledger_after_zero_row: bool, delivery_unavailable: bool +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + client.model_usage_transactions = [] + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": False}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + await client.db.connect() + try: + assert await _retrieve_batch(writer, 0) + await drain_logging_tasks() + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": ledger_after_zero_row}) + assert await _retrieve_batch(writer, 0.25) + await drain_logging_tasks() + assert await _retrieve_batch(writer, 0.25) is False + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (0,) + assert connection.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == [ + (0.0 if ledger_after_zero_row else 0.25,) + ] + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == ( + int(ledger_after_zero_row), + ) + if delivery_unavailable: + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute('ALTER TABLE "LiteLLM_SpendLogDelivery" RENAME TO temporarily_unavailable_delivery') + try: + await writer.db_update_spend_transaction_handler(client, 0, proxy_server.proxy_logging_obj) + finally: + if delivery_unavailable: + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute('ALTER TABLE temporarily_unavailable_delivery RENAME TO "LiteLLM_SpendLogDelivery"') + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (0.25,) + assert connection.execute('SELECT spend, api_requests FROM "LiteLLM_DailyUserSpend"').fetchall() == [(0.25, 2)] + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize("failing_table", [ + "LiteLLM_UserTable", "LiteLLM_VerificationToken", "LiteLLM_TeamTable", "LiteLLM_TeamMembership", + "LiteLLM_OrganizationMembership", "LiteLLM_OrganizationTable", "LiteLLM_ProjectTable", "LiteLLM_AgentsTable", "LiteLLM_EndUserTable", +]) +@pytest.mark.asyncio +async def test_destination_failure_rolls_back_all_earlier_updates_before_retry( + parity_database: str, monkeypatch: pytest.MonkeyPatch, failing_table: str +): + await _assert_legacy_database_parity( + parity_database, monkeypatch, "success", None, 0.25, "full", failing_table=failing_table + ) + + +@pytest.mark.asyncio +async def test_spend_logs_cleanup_does_not_remove_pending_accounting(parity_database, monkeypatch): + from litellm.accounting.worker import flush_accounting_deliveries + from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup + + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + cleanup = SpendLogCleanup(clock=lambda: 0) + await client.db.connect() + try: + assert await _retrieve_batch(writer, 0.25) + delete_sql = 'DELETE FROM "LiteLLM_SpendLogs" WHERE request_id IN (SELECT request_id FROM "LiteLLM_SpendLogs" WHERE "startTime" < $1::timestamp LIMIT $2::integer)' + assert await cleanup._execute_delete_batch(client, delete_sql, _NOW + timedelta(days=1), "LiteLLM_SpendLogs", 60) == 1 + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + assert await _retrieve_batch(writer, 0.25) is False + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogs"').fetchone() == (0,) + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=$$user$$').fetchone() == (0.25,) + assert connection.execute('SELECT state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("applied",)] + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize("mode", ["local", "url", "disabled"]) +@pytest.mark.parametrize("receipt_state", ["pending", "parked", "applied", "retained"]) +@pytest.mark.asyncio +async def test_ledger_charged_batch_is_not_charged_again_after_flag_disable( + parity_database: str, monkeypatch: pytest.MonkeyPatch, receipt_state: str, mode: str +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", mode == "disabled") + if mode == "url": + monkeypatch.setenv("SPEND_LOGS_URL", "http://127.0.0.1:1") + await client.db.connect() + try: + assert await _retrieve_batch(writer, 0) + assert await _retrieve_batch(writer, 0.25) + if receipt_state in ("applied", "retained"): + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database, autocommit=True) as connection: + if receipt_state == "retained": + connection.execute("UPDATE \"LiteLLM_SpendLogDelivery\" SET applied_at = CURRENT_TIMESTAMP - INTERVAL '400 days'") + if receipt_state == "parked": + connection.execute('UPDATE "LiteLLM_SpendLogDelivery" SET state=\'parked\'') + assert await AccountingDeliveryRepository(client.writer_db).cleanup_applied() == 0 + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": False}) + assert await _retrieve_batch(writer, 0.25) is False + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == ( + [(0.25,)] if mode != "disabled" else [] + ) + expected_spend = 0 if receipt_state == "parked" else 0.25 + expected_state = "parked" if receipt_state == "parked" else "applied" + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (expected_spend,) + assert connection.execute('SELECT state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [(expected_state,)] + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize("disabled", [False, True]) +@pytest.mark.parametrize("export_status", [200, 500]) +@pytest.mark.asyncio +async def test_export_failure_does_not_change_committed_accounting(parity_database, monkeypatch, export_status, disabled): + requests = asyncio.Queue() + + async def receive(reader, writer): + header = await reader.readuntil(b"\r\n\r\n") + length = next(int(line.split(b":", 1)[1]) for line in header.split(b"\r\n") if line.lower().startswith(b"content-length:")) + body = await reader.readexactly(length) + await requests.put(json.loads(body)) + writer.write(f"HTTP/1.1 {export_status} Result\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{{}}".encode()) + await writer.drain() + writer.close() + await writer.wait_closed() + + server = await asyncio.start_server(receive, "127.0.0.1", 0) + try: + port = server.sockets[0].getsockname()[1] + await _assert_legacy_database_parity( + parity_database, monkeypatch, "success", None, 0.25, "lean" if disabled else "full", + export_url=f"http://127.0.0.1:{port}", export_failure=export_status == 500 and not disabled, + ) + assert requests.qsize() == int(not disabled) + if disabled: + return + exported = await asyncio.wait_for(requests.get(), timeout=5) + assert len(exported) == 1 + assert exported[0]["spend"] == 0.25 + assert "accounting_payload" not in exported[0] + assert "accounting_log_disabled" not in exported[0] + finally: + server.close() + await server.wait_closed() + + +@pytest.mark.asyncio +async def test_spend_logs_url_without_an_export_client_keeps_local_logs_and_charges_once( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from litellm.accounting.worker import flush_accounting_deliveries + from litellm.proxy.utils import ProxyUpdateSpend + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setenv("SPEND_LOGS_URL", "http://127.0.0.1:1") + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + client.model_usage_transactions = [] + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + await client.db.connect() + try: + assert await writer.update_database( + token="2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", user_id="user", end_user_id=None, team_id="team", org_id=None, + kwargs={ + "model": "ledger-test", "call_type": "acompletion", "litellm_call_id": "url-request", + "litellm_params": {"metadata": {"status": "success", "user_api_key": "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683"}}, + }, + completion_response=ModelResponse( + id="url-request", model="ledger-test", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, end_time=_NOW, response_cost=0.25, + ) + await drain_logging_tasks() + rows = [ + {key: value.isoformat() if isinstance(value, datetime) else value for key, value in row.items()} + for row in client.spend_log_transactions + ] + client.spend_log_transactions = [] + await ProxyUpdateSpend.update_spend_logs(0, client, None, proxy_server.proxy_logging_obj, rows) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=\'url-request\'' + ).fetchall() == [(0.25,)] + assert connection.execute('SELECT state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("applied",)] + assert connection.execute('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=\'2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683\'').fetchone() == (0.25,) + finally: + await client.db.disconnect() + + +@contextlib.asynccontextmanager +async def _ledger_client(parity_database: str, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("DATABASE_URL", parity_database) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + client.model_usage_transactions = [] + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + await client.db.connect() + try: + yield client + finally: + await client.db.disconnect() + + +async def _record_call( + writer: DBSpendUpdateWriter, call_id: str, *, token: str = "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", user: str = "user", tags: object = None +) -> None: + metadata = {"status": "success", "user_api_key": token, **({"tags": tags} if tags is not None else {})} + assert await writer.update_database( + token=token, user_id=user, end_user_id=None, team_id="team", org_id=None, + kwargs={ + "model": "ledger-test", "call_type": "aembedding", "litellm_call_id": call_id, + "litellm_params": {"metadata": metadata}, + }, + completion_response=None, start_time=_NOW, end_time=_NOW, response_cost=0.25, + ) + await drain_logging_tasks() + + +def _wire(rows) -> list[dict[str, object]]: + return [ + {key: value.isoformat() if isinstance(value, datetime) else value for key, value in row.items()} + for row in rows + ] + + +async def _flush(writer: DBSpendUpdateWriter, client: PrismaClient) -> None: + from litellm.accounting.worker import flush_accounting_deliveries + from litellm.proxy.utils import ProxyUpdateSpend + + rows = _wire(client.spend_log_transactions) + client.spend_log_transactions = [] + await ProxyUpdateSpend.update_spend_logs(0, client, None, proxy_server.proxy_logging_obj, rows) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + + +def _key_spend(parity_database: str) -> object: + with psycopg.connect(parity_database) as connection: + return connection.execute( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', ("2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683",) + ).fetchone() + + +@pytest.mark.asyncio +async def test_requests_sharing_a_client_call_id_are_each_charged_but_a_replay_is_not( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + async with _ledger_client(parity_database, monkeypatch) as client: + writer = DBSpendUpdateWriter() + await _record_call(writer, "shared-call-id") + await _record_call(writer, "shared-call-id") + buffered = _wire(client.spend_log_transactions) + await _flush(writer, client) + assert await SpendLogsRepository(client).insert_with_deliveries(buffered) == 0 + await _flush(writer, client) + assert _key_spend(parity_database) == (0.5,) + with psycopg.connect(parity_database) as connection: + assert connection.execute( + 'SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery"' + ).fetchall() == [("shared-call-id", "applied"), ("shared-call-id", "applied")] + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogs"').fetchone() == (1,) + + +@pytest.mark.asyncio +async def test_an_opaque_bearer_token_is_never_stored_in_a_delivery( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + async with _ledger_client(parity_database, monkeypatch) as client: + writer = DBSpendUpdateWriter() + await _record_call(writer, "opaque-call", token="opaque-bearer-secret") + await _flush(writer, client) + with psycopg.connect(parity_database) as connection: + stored = connection.execute('SELECT accounting_payload::text, state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() + assert [state for _, state in stored] == ["applied"] + assert "opaque-bearer-secret" not in stored[0][0] + + +@pytest.mark.asyncio +async def test_an_evicted_spend_log_still_charges_its_request(parity_database: str, monkeypatch: pytest.MonkeyPatch): + from litellm.proxy.utils import enqueue_spend_logs + + async with _ledger_client(parity_database, monkeypatch) as client: + writer = DBSpendUpdateWriter() + await _record_call(writer, "evicted-call") + await _record_call(writer, "kept-call") + queued = tuple(client.spend_log_transactions) + client.spend_log_transactions = [] + await enqueue_spend_logs(client, queued, max_bytes=1) + assert [row.get("accounting_log_disabled") for row in client.spend_log_transactions] == [True, False] + await _flush(writer, client) + assert _key_spend(parity_database) == (0.5,) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT request_id FROM "LiteLLM_SpendLogs"').fetchall() == [("kept-call",)] + + +@pytest.mark.parametrize("isolation_budget", [5, 0], ids=["isolated", "budget_exhausted"]) +@pytest.mark.asyncio +async def test_a_spend_log_postgres_rejects_still_charges_its_request_and_the_rest_of_the_batch( + parity_database: str, monkeypatch: pytest.MonkeyPatch, isolation_budget: int +): + from prisma.errors import DataError + + from litellm.accounting.worker import flush_accounting_deliveries + from litellm.proxy.utils import _create_spend_logs_with_poison_isolation + + class RejectsOneLog: + def __init__(self, repository: SpendLogsRepository) -> None: + self.repository = repository + + async def insert_with_deliveries(self, rows, *, write_spend_logs: bool = True) -> int: + if any(row["request_id"] == "poison-call" and not row.get("accounting_log_disabled") for row in rows): + raise DataError({"user_facing_error": {"message": 'invalid byte sequence for encoding "UTF8": 0x00'}}) + return await self.repository.insert_with_deliveries(rows, write_spend_logs=write_spend_logs) + + async with _ledger_client(parity_database, monkeypatch) as client: + writer = DBSpendUpdateWriter() + await _record_call(writer, "poison-call") + await _record_call(writer, "clean-call") + rows = tuple(client.jsonify_object(row) for row in client.spend_log_transactions) + client.spend_log_transactions = [] + await _create_spend_logs_with_poison_isolation(RejectsOneLog(SpendLogsRepository(client)), rows, isolation_budget) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + assert _key_spend(parity_database) == (0.5,) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT request_id FROM "LiteLLM_SpendLogs"').fetchall() == ( + [("clean-call",)] if isolation_budget else [] + ) + assert connection.execute( + 'SELECT state::text, COUNT(*) FROM "LiteLLM_SpendLogDelivery" GROUP BY 1' + ).fetchall() == [("applied", 2)] + + +@pytest.mark.asyncio +async def test_a_request_with_a_non_text_tag_is_still_charged(parity_database: str, monkeypatch: pytest.MonkeyPatch): + async with _ledger_client(parity_database, monkeypatch) as client: + writer = DBSpendUpdateWriter() + await _record_call(writer, "tagged-call", tags=[None, "real-tag"]) + await _flush(writer, client) + assert _key_spend(parity_database) == (0.25,) + with psycopg.connect(parity_database) as connection: + assert connection.execute( + 'SELECT tag, spend FROM "LiteLLM_DailyTagSpend" WHERE tag=\'real-tag\'' + ).fetchall() == [("real-tag", 0.25)] + + +@pytest.mark.parametrize("threshold", ["count", "bytes"]) +@pytest.mark.asyncio +async def test_retained_batch_storage_threshold_alerts_without_deleting_receipts( + parity_database: str, monkeypatch: pytest.MonkeyPatch, threshold: str +): + from types import SimpleNamespace + from unittest.mock import AsyncMock + + from litellm.accounting.worker import flush_accounting_deliveries + from litellm.types.integrations.slack_alerting import AlertType + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setenv("LITELLM_ACCOUNTING_BATCH_DELIVERY_ALERT_COUNT", "1" if threshold == "count" else "1000000") + monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1" if threshold == "bytes" else "1073741824") + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + alerts = SimpleNamespace(alerting_handler=AsyncMock()) + await client.db.connect() + try: + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, state, applied_at, accounting_payload, is_batch) ' + "VALUES ('old-batch', 'applied', CURRENT_TIMESTAMP - INTERVAL '400 days', '{}', true)" + ) + await flush_accounting_deliveries(DBSpendUpdateWriter(), client, alerts) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [ + ("old-batch", "applied") + ] + table_bytes = connection.execute( + 'SELECT pg_total_relation_size(\'"LiteLLM_SpendLogDelivery"\')' + ).fetchone()[0] + alerts.alerting_handler.assert_awaited_once_with( + message=f"Accounting delivery storage requires review. Retained batches=1, table bytes={table_bytes}.", + level="High", alert_type=AlertType.failed_tracking_spend, + ) + finally: + await client.db.disconnect() + + +@pytest.mark.parametrize("backlog", ["retries", "expired"]) +@pytest.mark.parametrize("batch_size", [1, 1100]) +@pytest.mark.asyncio +async def test_worker_drains_recovered_and_expired_work_at_its_configured_capacity( + parity_database, monkeypatch, backlog, batch_size +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", str(batch_size)) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + writer = DBSpendUpdateWriter() + monkeypatch.setattr(proxy_server, "prisma_client", client) + await client.db.connect() + try: + assert await _retrieve_batch(writer, 0.25) + with psycopg.connect(parity_database, autocommit=True) as connection: + if backlog == "retries": + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, accounting_payload, is_batch, attempts) ' + "SELECT 'retry-' || i, accounting_payload, false, 1 " + 'FROM "LiteLLM_SpendLogDelivery" CROSS JOIN generate_series(1, 3) i' + ) + else: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, accounting_payload, is_batch, state, applied_at) ' + "SELECT 'expired-' || i, '{}', false, 'applied', CURRENT_TIMESTAMP - INTERVAL '2 days' " + 'FROM generate_series(1, 1100) i' + ) + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == ( + 0.25 * (1 + min(batch_size, 3)) if backlog == "retries" else 0.25, + ) + assert connection.execute( + 'SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'pending\' OR request_id LIKE \'expired-%\'' + ).fetchone() == ((3 - min(batch_size, 3)) if backlog == "retries" else max(0, 1100 - batch_size * 2),) + finally: + await client.db.disconnect() + + +@pytest.mark.asyncio +async def test_analytics_overflow_across_requests_still_applies_every_charge( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": True}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + client.spend_log_transactions = [] + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + writer = DBSpendUpdateWriter() + await client.db.connect() + try: + assert await _retrieve_batch(writer, 0.25) + stored = await client.writer_db.query_raw('SELECT to_jsonb(s) AS row FROM "LiteLLM_SpendLogs" s') + delivery = (await client.writer_db.query_raw('SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"'))[0][ + "accounting_payload" + ] + oversized_savings = {**delivery["metadata"], "compression_savings": {"tokens_saved": 2**62}} + assert await SpendLogsRepository(client).insert_with_deliveries( + tuple( + { + **stored[0]["row"], + "request_id": f"overflow-{index}", + "accounting_payload": {**delivery, "request_id": f"overflow-{index}", "metadata": oversized_savings}, + } + for index in range(3) + ) + ) == 3 + await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute( + 'SELECT state::text, COUNT(*) FROM "LiteLLM_SpendLogDelivery" GROUP BY state' + ).fetchall() == [("applied", 4)] + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (1.0,) + assert connection.execute( + 'SELECT api_requests, compression_saved_tokens FROM "LiteLLM_DailyUserSpend"' + ).fetchone() == (4, 3 * 2**40) + finally: + await client.db.disconnect() diff --git a/tests/unit/accounting/test_delivery_projection.py b/tests/unit/accounting/test_delivery_projection.py new file mode 100644 index 00000000000..0b1f36d6b48 --- /dev/null +++ b/tests/unit/accounting/test_delivery_projection.py @@ -0,0 +1,85 @@ +import json +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from typing import Final +from urllib.parse import unquote + +from litellm.accounting.delivery_projection import entity_transactions, project_deliveries +from litellm.accounting.projection import Contribution +from litellm.repositories.accounting_delivery_repository import DeliveryProjection + +_NOW: Final = datetime(2026, 11, 2, tzinfo=timezone.utc) + + +def _payload(spend: float, started: datetime = _NOW, window_start: datetime = _NOW) -> Mapping[str, object]: + return { + "request_id": f"request-{spend}-{started.isoformat()}", + "startTime": started, + "api_key": "key", + "user": "user/1", + "team_id": "team", + "organization_id": "org::1", + "spend": spend, + "metadata": json.dumps( + { + "accounting_entities": { + "api_key": "key", + "user": "user/1", + "team_id": "team", + "organization_id": "org::1", + }, + "accounting_windows": [ + { + "entity_type": "key", + "entity_id": "key", + "window_duration": "1d", + "window_start": window_start.isoformat(), + } + ], + } + ), + } + + +def _spend(contributions: tuple[Contribution, ...], table: str) -> tuple[float, ...]: + return tuple(float(dict(item.increments)["spend"]) for item in contributions if item.table == table) + + +def test_deliveries_for_one_destination_sum_and_keep_the_earliest_request_start() -> None: + later: Final = _NOW + timedelta(hours=1) + + contributions: Final = project_deliveries( + (DeliveryProjection(_payload(0.5, started=later)), DeliveryProjection(_payload(0.25))) + ) + + assert _spend(contributions, "LiteLLM_VerificationToken") == (0.75,) + windows: Final = tuple(item for item in contributions if item.table == "LiteLLM_BudgetWindowSpend") + assert len(windows) == 1 + assert dict(windows[0].increments) == {"spend": 0.75} + assert dict(windows[0].values)["started_at"] == str(_NOW) + + +def test_the_same_window_in_a_new_period_is_charged_to_its_own_row() -> None: + next_day: Final = _NOW + timedelta(days=1) + + contributions: Final = project_deliveries( + (DeliveryProjection(_payload(0.25)), DeliveryProjection(_payload(0.5, window_start=next_day))) + ) + + windows: Final = sorted( + (dict(item.values)["window_start"], dict(item.increments)["spend"]) + for item in contributions + if item.table == "LiteLLM_BudgetWindowSpend" + ) + assert windows == [(_NOW.isoformat(), 0.25), (next_day.isoformat(), 0.5)] + + +def test_entity_transactions_key_memberships_the_way_the_writer_applies_them() -> None: + transactions: Final = entity_transactions(project_deliveries((DeliveryProjection(_payload(0.25)),))) + + assert transactions["key_list_transactions"] == {"key": 0.25} + assert transactions["org_list_transactions"] == {"org::1": 0.25} + assert transactions["team_member_list_transactions"] == {"team_id::team::user_id::user/1": 0.25} + ((org_member_key, org_member_spend),) = transactions["org_member_list_transactions"].items() + _, quoted_org_id, _, quoted_user_id = org_member_key.split("::") + assert (unquote(quoted_org_id), unquote(quoted_user_id), org_member_spend) == ("org::1", "user/1", 0.25) diff --git a/tests/unit/accounting/test_measurements.py b/tests/unit/accounting/test_measurements.py deleted file mode 100644 index b65c3076c5d..00000000000 --- a/tests/unit/accounting/test_measurements.py +++ /dev/null @@ -1,235 +0,0 @@ -from datetime import datetime, timezone -from decimal import Decimal -from typing import Final - -import pytest - -from litellm.accounting import ( - CostEvent, - Dimension, - DimensionName, - InboundUsageEvent, - Measurement, - Metric, - MetricName, - Outcome, - OutboundUsageEvent, - PricingSnapshot, - measurements_for, -) -from litellm.types.utils import CallTypes - -_START: Final = datetime(2026, 1, 1, 23, 59, 59, tzinfo=timezone.utc) -_END: Final = datetime(2026, 1, 2, tzinfo=timezone.utc) -_DIMENSIONS: Final = ( - Dimension(name=DimensionName.TEAM, value="team"), - Dimension(name=DimensionName.TAG, value="alpha"), - Dimension(name=DimensionName.TAG, value="beta"), -) - - -def _snapshot(amount: Decimal) -> PricingSnapshot: - return PricingSnapshot( - calculator_version="test", - model="test-model", - provider=None, - pricing_source="upstream_header", - priced_model_key="test-model", - upstream_reported_cost=amount, - total_cost=amount, - ) - - -@pytest.mark.parametrize( - ("outcome", "outcome_metric"), - [ - (Outcome.SUCCEEDED, MetricName.SUCCESSFUL_REQUESTS), - (Outcome.FAILED, MetricName.FAILED_REQUESTS), - (Outcome.CANCELLED, MetricName.FAILED_REQUESTS), - ], -) -def test_caller_measurements_keep_both_request_and_outcome_counts(outcome, outcome_metric): - event = InboundUsageEvent( - event_id="caller-event", - request_id="caller", - call_type=CallTypes.acompletion, - started_at=_START, - ended_at=_END, - outcome=outcome, - dimensions=_DIMENSIONS, - metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)), Metric(name=outcome_metric, value=Decimal(1))), - ) - assert measurements_for(event) == ( - Measurement( - source_event_id="caller-event", - occurred_at=_START, - dimensions=_DIMENSIONS, - metric=Metric(name=MetricName.REQUESTS, value=Decimal(1)), - ), - Measurement( - source_event_id="caller-event", - occurred_at=_START, - dimensions=_DIMENSIONS, - metric=Metric(name=outcome_metric, value=Decimal(1)), - ), - ) - - -def test_usage_measurements_keep_every_metric_in_the_start_time_bucket(): - event = OutboundUsageEvent( - event_id="usage-event", - attempt_id="attempt", - request_id="caller", - retry_of_attempt_id=None, - call_type=CallTypes.acompletion, - started_at=_START, - ended_at=_END, - outcome=Outcome.SUCCEEDED, - dimensions=_DIMENSIONS, - metrics=( - Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)), - Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)), - Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)), - Metric(name=MetricName.CACHE_READ_INPUT_TOKENS, value=Decimal(3)), - Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(1)), - ), - ) - assert measurements_for(event) == ( - Measurement( - source_event_id="usage-event", - occurred_at=_START, - dimensions=_DIMENSIONS, - metric=Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)), - ), - Measurement( - source_event_id="usage-event", - occurred_at=_START, - dimensions=_DIMENSIONS, - metric=Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)), - ), - Measurement( - source_event_id="usage-event", - occurred_at=_START, - dimensions=_DIMENSIONS, - metric=Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)), - ), - Measurement( - source_event_id="usage-event", - occurred_at=_START, - dimensions=_DIMENSIONS, - metric=Metric(name=MetricName.CACHE_READ_INPUT_TOKENS, value=Decimal(3)), - ), - Measurement( - source_event_id="usage-event", - occurred_at=_START, - dimensions=_DIMENSIONS, - metric=Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(1)), - ), - ) - - -@pytest.mark.parametrize("metrics", [None, ()]) -def test_unknown_or_empty_usage_does_not_fabricate_zero_measurements(metrics): - event = OutboundUsageEvent( - event_id="usage-event", - attempt_id="attempt", - request_id=None, - retry_of_attempt_id=None, - call_type=CallTypes.acompletion, - started_at=_START, - ended_at=_END, - outcome=Outcome.SUCCEEDED, - dimensions=(), - metrics=metrics, - ) - assert measurements_for(event) == () - - -@pytest.mark.parametrize("amount", [Decimal(0), Decimal("0.1234567890123456789")]) -def test_cost_measurements_preserve_authoritative_amount_identity_and_occurrence(amount): - event = CostEvent( - event_id="cost-event", - usage_event_id="usage-event", - occurred_at=_END, - dimensions=_DIMENSIONS, - charges=(Metric(name=MetricName.SPEND, value=amount, unit="USD"),), - snapshot=_snapshot(amount), - ) - assert measurements_for(event) == ( - Measurement( - source_event_id="cost-event", - occurred_at=_END, - dimensions=_DIMENSIONS, - metric=Metric(name=MetricName.SPEND, value=amount, unit="USD"), - ), - ) - - -def test_billed_retry_and_success_contribute_one_caller_request_and_both_charges(): - caller = InboundUsageEvent( - event_id="caller-event", - request_id="caller", - call_type=CallTypes.acompletion, - started_at=_START, - ended_at=_END, - outcome=Outcome.SUCCEEDED, - dimensions=_DIMENSIONS, - metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),), - ) - retry = OutboundUsageEvent( - event_id="retry-usage", - attempt_id="retry", - request_id="caller", - retry_of_attempt_id=None, - call_type=CallTypes.acompletion, - started_at=_START, - ended_at=_END, - outcome=Outcome.FAILED, - dimensions=_DIMENSIONS, - metrics=(Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(4)),), - ) - success = OutboundUsageEvent( - event_id="success-usage", - attempt_id="success", - request_id="caller", - retry_of_attempt_id="retry", - call_type=CallTypes.acompletion, - started_at=_START, - ended_at=_END, - outcome=Outcome.SUCCEEDED, - dimensions=_DIMENSIONS, - metrics=(Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)),), - ) - retry_cost = CostEvent( - event_id="retry-cost", - usage_event_id="retry-usage", - occurred_at=_START, - dimensions=_DIMENSIONS, - charges=(Metric(name=MetricName.SPEND, value=Decimal("0.025"), unit="USD"),), - snapshot=_snapshot(Decimal("0.025")), - ) - success_cost = CostEvent( - event_id="success-cost", - usage_event_id="success-usage", - occurred_at=_START, - dimensions=_DIMENSIONS, - charges=(Metric(name=MetricName.SPEND, value=Decimal("0.125"), unit="USD"),), - snapshot=_snapshot(Decimal("0.125")), - ) - measurements = ( - *measurements_for(caller), - *measurements_for(retry), - *measurements_for(success), - *measurements_for(retry_cost), - *measurements_for(success_cost), - ) - assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.REQUESTS) == Decimal(1) - assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.TOTAL_TOKENS) == Decimal(16) - assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.SPEND) == Decimal("0.150") - assert {m.source_event_id for m in measurements} == { - "caller-event", - "retry-usage", - "success-usage", - "retry-cost", - "success-cost", - } diff --git a/tests/unit/accounting/test_metrics.py b/tests/unit/accounting/test_metrics.py new file mode 100644 index 00000000000..49b46eb36c2 --- /dev/null +++ b/tests/unit/accounting/test_metrics.py @@ -0,0 +1,28 @@ +from prometheus_client.metrics import MetricWrapperBase + +from litellm.accounting import metrics +from litellm.accounting.metrics import record_deliveries +from litellm.repositories.accounting_delivery_repository import DeliveryMetrics + + +def _sample(metric: MetricWrapperBase, suffix: str = "") -> float: + family = metric.collect()[0] + return next(sample.value for sample in family.samples if sample.name == family.name + suffix) + + +def test_delivery_metrics_publish_backlog_and_clamp_clock_skew() -> None: + batches_before = _sample(metrics._BATCH_TIME, "_count") + + record_deliveries( + DeliveryMetrics(pending_count=4, parked_count=1, oldest_pending_seconds=-3.0, batch_count=7, table_bytes=8192), + 0.5, + ) + + assert ( + _sample(metrics._PENDING), + _sample(metrics._PARKED), + _sample(metrics._OLDEST), + _sample(metrics._BATCH_COUNT), + _sample(metrics._TABLE_BYTES), + ) == (4.0, 1.0, 0.0, 7.0, 8192.0) + assert _sample(metrics._BATCH_TIME, "_count") == batches_before + 1 diff --git a/tests/unit/accounting/test_projection.py b/tests/unit/accounting/test_projection.py index ec0ab44e505..17186a89922 100644 --- a/tests/unit/accounting/test_projection.py +++ b/tests/unit/accounting/test_projection.py @@ -429,3 +429,129 @@ def test_guardrail_usage_is_dated_in_utc() -> None: assert [dict(item.key)["date"] for item in projected if item.table == "LiteLLM_DailyGuardrailUsageUnits"] == [ "2026-11-03" ] + + +@pytest.mark.parametrize( + ("untrusted", "metadata"), + [ + ({"request_tags": json.dumps(["x" * 16384, "normal"])}, {}), + ({"end_user": "x" * 3072}, {}), + ({}, {"compression_savings": {"tokens_saved": 2**63}}), + ], +) +def test_untrusted_analytics_values_never_block_the_charge( + untrusted: Mapping[str, str], metadata: Mapping[str, object] +) -> None: + row: Final[Mapping[str, object]] = { + "request_id": "untrusted", + "startTime": _NOW, + "api_key": "key", + "user": "user", + "team_id": "team", + "model": "unpriced-projection-test", + "model_group": "group", + "custom_llm_provider": "openai", + "spend": 0.5, + **untrusted, + "metadata": { + **metadata, + "accounting_entities": { + "api_key": "key", + "user": "user", + "team_id": "team", + "end_user": untrusted.get("end_user"), + }, + }, + } + contributions: Final = project(row) + texts: Final = tuple(value for item in contributions for _, value in (*item.key, *item.values) if value) + numbers: Final = tuple(value for item in contributions for _, value in item.increments) + assert {"LiteLLM_VerificationToken", "LiteLLM_TeamTable", "LiteLLM_UserTable"} <= { + item.table for item in contributions + } + assert "LiteLLM_DailyUserSpend" in {item.table for item in contributions} + assert all(len(text.encode()) <= 2000 for text in texts) + assert all(abs(number) <= 2**40 for number in numbers) + + +def test_stored_payload_keeps_guardrail_billing_but_not_guardrail_content() -> None: + payload: Final = get_logging_payload( + kwargs={ + "model": "projection-test", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "guardrail-request", + "litellm_params": {"metadata": {"user_api_key": "key", "user_api_key_user_id": "user"}}, + }, + response_obj=ModelResponse( + id="guardrail-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, + end_time=_NOW, + ) + guardrail: Final = { + "guardrail_name": "pii", + "guardrail_usage": {"requests": 1}, + "guardrail_cost_by_unit": {"requests": 0.25}, + "guardrail_response": "card 4111 1111 1111 1111", + "match_details": [{"matched": "4111 1111 1111 1111"}], + } + headroom: Final = { + "guardrail_name": "compress", + "guardrail_provider": "headroom", + "guardrail_response": {"tokens_saved": 7, "compressed_prompt": "card 4111 1111 1111 1111"}, + } + metadata: Final = {**_FIELDS.validate_json(payload["metadata"]), "guardrail_information": [guardrail, headroom]} + accounting_input: Final[SpendLogsPayload] = { + **payload, + "spend": 0.5, + "metadata": json.dumps(metadata), + "request_tags": '["nul\x00tag"]', + } + accounting: Final = build_accounting_payload( + accounting_input, entities={"api_key": "key"}, model_access_groups=(), windows=(), router=lambda: None + ) + stored: Final = json.dumps(accounting, default=str) + restored: Final = _FIELDS.validate_json(stored) + usage_rows: Final = tuple(item for item in project(restored) if item.table == "LiteLLM_DailyGuardrailUsageUnits") + assert "4111" not in stored + assert "\\u0000" not in stored + assert [dict(item.increments) for item in usage_rows] == [{"units": 1, "cost": 0.25, "untracked_units": 0}] + daily_user: Final = next(item for item in project(restored) if item.table == "LiteLLM_DailyUserSpend") + assert dict(daily_user.increments)["compression_saved_tokens"] == 7 + + +def test_caller_supplied_savings_that_cannot_be_computed_never_drop_the_charge(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + + monkeypatch.setitem( + litellm.model_cost, + "savings-overflow-test", + {"input_cost_per_token": 0.002, "output_cost_per_token": 0.004, "litellm_provider": "openai", "mode": "chat"}, + ) + payload: Final[SpendLogsPayload] = { + **get_logging_payload( + kwargs={ + "model": "savings-overflow-test", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "overflow-request", + "litellm_params": {"metadata": {"user_api_key": "key"}}, + }, + response_obj=ModelResponse( + id="overflow-request", model="savings-overflow-test", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, + end_time=_NOW, + ), + "spend": 0.5, + "metadata": json.dumps({"compression_savings": {"tokens_saved": 10**400}}), + } + accounting: Final = build_accounting_payload( + payload, entities={"api_key": "key"}, model_access_groups=(), windows=(), router=lambda: None + ) + restored: Final = _FIELDS.validate_json(json.dumps(accounting, default=str)) + assert dict(restored["metadata"])["accounting_savings"]["compression"] == 0.0 + assert [dict(item.increments)["spend"] for item in project(restored) if item.table == "LiteLLM_VerificationToken"] == [ + 0.5 + ] diff --git a/tests/unit/accounting/test_types.py b/tests/unit/accounting/test_types.py deleted file mode 100644 index c1be39ab111..00000000000 --- a/tests/unit/accounting/test_types.py +++ /dev/null @@ -1,482 +0,0 @@ -import re -from collections.abc import Mapping -from datetime import datetime, timedelta, timezone -from decimal import Decimal -from enum import Enum -from pathlib import Path -from typing import Final - -import pytest -from pydantic import ValidationError - -from litellm.accounting import ( - AdmissionDecision, - AggregateQuery, - AggregateRow, - CostEvent, - DeliveryState, - DeliveryTarget, - Dimension, - DimensionFilter, - DimensionName, - Granularity, - InboundUsageEvent, - Metric, - MetricName, - Multiplier, - Outcome, - OutboundUsageEvent, - PricedQuantity, - PricingSnapshot, - Rate, - validate_aggregate_rows, -) -from litellm.accounting.types import PricingUnit -from litellm.types.utils import CallTypes - -_START: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) -_END: Final = datetime(2026, 1, 2, tzinfo=timezone.utc) -_REQUEST: Final = Metric(name=MetricName.REQUESTS, value=Decimal(1)) -_TOKEN: Final = Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)) -_SPEND: Final = Metric(name=MetricName.SPEND, value=Decimal("0.1234567890123456789"), unit="USD") -_TEAM: Final = Dimension(name=DimensionName.TEAM, value="team") - - -def _inbound(metrics=(_REQUEST,), outcome=Outcome.SUCCEEDED): - return InboundUsageEvent( - event_id="inbound:request", - request_id="request", - call_type=CallTypes.acompletion, - started_at=_START, - ended_at=_START, - outcome=outcome, - dimensions=(_TEAM,), - metrics=metrics, - ) - - -def _outbound(metrics=(_TOKEN,), **changes): - return OutboundUsageEvent.model_validate( - { - "event_id": "outbound:attempt", - "attempt_id": "attempt", - "request_id": "request", - "retry_of_attempt_id": None, - "call_type": CallTypes.acompletion, - "started_at": _START, - "ended_at": _END, - "outcome": Outcome.SUCCEEDED, - "dimensions": (_TEAM,), - "metrics": metrics, - **changes, - } - ) - - -def _query(**changes): - return AggregateQuery.model_validate( - { - "metric": MetricName.SPEND, - "filters": (DimensionFilter(name=DimensionName.TEAM, values=("one", "two")),), - "group_by": (DimensionName.MODEL,), - "start": _START, - "end": _END, - "granularity": Granularity.DAY, - **changes, - } - ) - - -def test_wire_values_preserve_decimal_precision_and_normalize_aware_times(): - offset = timezone(timedelta(hours=-8)) - event = CostEvent( - event_id="cost:attempt", - usage_event_id="outbound:attempt", - occurred_at=datetime(2025, 12, 31, 16, tzinfo=offset), - dimensions=(_TEAM,), - charges=(_SPEND,), - snapshot=_snapshot(), - ) - wire = event.model_dump(mode="json") - assert wire["occurred_at"] == "2026-01-01T00:00:00Z" - assert wire["charges"] == [{"name": "spend", "value": "0.1234567890123456789", "unit": "USD"}] - assert event.occurred_at == _START - assert CostEvent.model_validate_json(event.model_dump_json()).charges[0].value == _SPEND.value - - -@pytest.mark.parametrize( - ("name", "value", "unit", "message"), - [ - (MetricName.SPEND, Decimal(1), "count", "SPEND requires USD"), - (MetricName.REQUESTS, Decimal(1), "USD", "SPEND requires USD"), - (MetricName.PROMPT_TOKENS, Decimal("1.5"), "count", "integral"), - (MetricName.SPEND, Decimal("NaN"), "USD", "finite"), - (MetricName.SPEND, Decimal("Infinity"), "USD", "finite"), - (MetricName.SPEND, Decimal(-1), "USD", "greater than or equal"), - (MetricName.SPEND, 0.1, "USD", "instance of Decimal"), - ], -) -def test_metrics_reject_wrong_units_nonfinite_values_and_binary_float_inputs(name, value, unit, message): - with pytest.raises(ValidationError, match=message): - Metric(name=name, value=value, unit=unit) - - -@pytest.mark.parametrize("metrics", [(_TOKEN,), (_SPEND,), ()]) -def test_inbound_activity_cannot_own_usage_or_cost_or_omit_request_count(metrics): - with pytest.raises(ValidationError, match="Inbound events"): - _inbound(metrics) - - -@pytest.mark.parametrize( - ("outcome", "metric"), - [ - (Outcome.FAILED, Metric(name=MetricName.SUCCESSFUL_REQUESTS, value=Decimal(1))), - (Outcome.SUCCEEDED, Metric(name=MetricName.FAILED_REQUESTS, value=Decimal(1))), - ], -) -def test_inbound_outcome_counts_cannot_contradict_outcome(outcome, metric): - with pytest.raises(ValidationError, match="count must match outcome"): - _inbound((_REQUEST, metric), outcome) - - -@pytest.mark.parametrize("metrics", [(_REQUEST,), (_SPEND,)]) -def test_outbound_attempts_cannot_own_caller_counts_or_charges(metrics): - with pytest.raises(ValidationError, match="Outbound events own only"): - _outbound(metrics) - - -@pytest.mark.parametrize( - ("changes", "message"), - [ - ({"started_at": datetime(2026, 1, 1)}, "timezone"), - ({"ended_at": _START - timedelta(seconds=1)}, "End must follow start"), - ({"retry_of_attempt_id": "attempt"}, "cannot retry itself"), - ({"dimensions": (_TEAM, _TEAM)}, "Duplicate dimension values"), - ({"dimensions": (_TEAM, Dimension(name=DimensionName.TEAM, value="other"))}, "Only tag and model access group"), - ( - { - "metrics": ( - Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(5)), - Metric(name=MetricName.REASONING_TOKENS, value=Decimal(6)), - ) - }, - "reasoning_tokens cannot exceed completion_tokens", - ), - ( - { - "metrics": ( - Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(5)), - Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, value=Decimal(6)), - ) - }, - "cache_creation_input_tokens_above_1hr cannot exceed cache_creation_input_tokens", - ), - ({"event_id": " "}, "pattern"), - ({"metrics": (_TOKEN, _TOKEN)}, "Duplicate metric names"), - ], -) -def test_outbound_events_reject_invalid_identity_time_and_duplicate_attribution(changes, message): - with pytest.raises(ValidationError, match=message): - _outbound(**changes) - - -def test_unknown_and_known_zero_usage_are_distinct_and_values_are_immutable(): - unknown = _outbound(None) - zero = _outbound((Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(0)),)) - assert unknown.metrics is None - assert zero.metrics is not None and zero.metrics[0].value == 0 - assert _inbound().started_at == _inbound().ended_at - with pytest.raises(ValidationError, match="frozen"): - zero.attempt_id = "different" - with pytest.raises(ValidationError, match="frozen"): - zero.metrics[0].value = Decimal(10) - - -@pytest.mark.parametrize("charges", [(), (_TOKEN,), (_SPEND, _SPEND)]) -def test_cost_events_require_one_authoritative_spend(charges): - with pytest.raises(ValidationError, match=r"authoritative USD spend|Duplicate metric"): - CostEvent( - event_id="cost", - usage_event_id="usage", - occurred_at=_START, - dimensions=(), - charges=charges, - snapshot=_snapshot(), - ) - - -def test_cost_event_cannot_reuse_its_usage_identity(): - with pytest.raises(ValidationError, match="distinct identities"): - CostEvent( - event_id="usage", - usage_event_id="usage", - occurred_at=_START, - dimensions=(), - charges=(_SPEND,), - snapshot=_snapshot(), - ) - - -@pytest.mark.parametrize("granularity", [Granularity.DAY, None]) -def test_daily_and_whole_interval_queries_preserve_filter_and_group_semantics(granularity): - query = _query(granularity=granularity) - assert query.granularity is granularity - assert query.filters[0].values == ("one", "two") - assert query.group_by == (DimensionName.MODEL,) - assert query.start == _START and query.end == _END - - -@pytest.mark.parametrize( - ("changes", "message"), - [ - ({"end": _START}, "End must follow start"), - ({"start": _START + timedelta(hours=1)}, "UTC day boundaries"), - ({"end": _END + timedelta(microseconds=1)}, "UTC day boundaries"), - ({"granularity": "hour"}, "day"), - ({"group_by": (DimensionName.TEAM, DimensionName.TEAM)}, "Duplicate group-by"), - ({"filters": (DimensionFilter(name=DimensionName.TEAM, values=("one",)),) * 2}, "one IN filter"), - ({"filters": ({"name": "team", "values": ()},)}, "at least 1"), - ], -) -def test_queries_reject_unsupported_or_ambiguous_ranges_and_filters(changes, message): - with pytest.raises(ValidationError, match=message): - _query(**changes) - - -def test_rows_allow_adjacent_buckets_and_overlapping_distinct_groups(): - first = AggregateRow(start=_START, end=_END, dimensions=(_TEAM,), value=Decimal(1)) - second = AggregateRow(start=_END, end=_END + timedelta(days=1), dimensions=(_TEAM,), value=Decimal(2)) - other = AggregateRow(start=_START, end=_END, dimensions=(), value=Decimal(3)) - rows = (second, other, first) - assert validate_aggregate_rows(rows) is rows - assert validate_aggregate_rows(()) == () - - -def test_rows_reject_overlap_within_a_group_regardless_of_dimension_order(): - model = Dimension(name=DimensionName.MODEL, value="model") - first = AggregateRow(start=_START, end=_END, dimensions=(_TEAM, model), value=Decimal(1)) - overlapping = AggregateRow(start=_START, end=_END, dimensions=(model, _TEAM), value=Decimal(2)) - with pytest.raises(ValueError, match="overlap within a dimension group"): - validate_aggregate_rows((first, overlapping)) - - -def test_rows_reject_empty_intervals(): - with pytest.raises(ValidationError, match="End must follow start"): - AggregateRow(start=_START, end=_START, dimensions=(), value=Decimal(0)) - - -def test_admission_decision_preserves_a_denial_reason(): - decision = AdmissionDecision(allowed=False, reason="Team budget exhausted") - assert not decision.allowed - assert decision.reason == "Team budget exhausted" - - -def _snapshot(**changes: object) -> PricingSnapshot: - return PricingSnapshot.model_validate( - { - "calculator_version": "1.80.0", - "model": "gpt-4o", - "provider": "openai", - "pricing_source": "cost_map", - "priced_model_key": "openai/gpt-4o", - "multipliers": (Multiplier(name="data_residency", value="eu", factor=Decimal("1.1")),), - "rates": (Rate(measure="input", unit="token", usd_per_unit=Decimal("0.0000025")),), - "usage": (PricedQuantity(measure="input", unit="token", value=_TOKEN.value),), - "input_cost": Decimal("0.000025"), - "total_cost": _SPEND.value, - **changes, - } - ) - - -def test_cost_event_carries_a_snapshot_whose_total_matches_the_spend(): - event = CostEvent( - event_id="cost:attempt", - usage_event_id="outbound:attempt", - occurred_at=_START, - dimensions=(_TEAM,), - charges=(_SPEND,), - snapshot=_snapshot(), - ) - assert event.snapshot is not None - assert event.snapshot.total_cost == _SPEND.value - assert event.snapshot.multipliers[0].factor == Decimal("1.1") - with pytest.raises(ValidationError, match="Snapshot total must equal the authoritative spend"): - CostEvent( - event_id="cost:attempt", - usage_event_id="outbound:attempt", - occurred_at=_START, - dimensions=(_TEAM,), - charges=(_SPEND,), - snapshot=_snapshot(total_cost=Decimal("0.5")), - ) - - -@pytest.mark.parametrize( - ("changes", "message"), - ( - ( - { - "rates": ( - Rate(measure="input", unit="token", usd_per_unit=Decimal(1)), - Rate(measure="input", unit="token", usd_per_unit=Decimal(2)), - ) - }, - "Duplicate rate measures and units", - ), - ({"usage": (_SPEND,)}, "PricedQuantity"), - ({"usage": (_REQUEST,)}, "PricedQuantity"), - ({"pricing_source": "guess"}, "pricing_source"), - ({"multipliers": ({"name": "x", "value": "y", "factor": Decimal(0)},)}, "factor"), - ({"schema_version": 2}, "schema_version"), - ( - {"additional_costs": ({"name": "fee", "amount": Decimal(1)}, {"name": "fee", "amount": Decimal(2)})}, - "Duplicate additional cost names", - ), - ), -) -def test_snapshots_reject_ambiguous_rates_wrong_usage_and_unknown_sources(changes: Mapping[str, object], message: str): - with pytest.raises(ValidationError, match=message): - _snapshot(**changes) - - -@pytest.mark.parametrize( - ("unit", "value"), - [("token", "10"), ("second", "1.25"), ("character", "25"), ("image", "2"), ("request", "1")], -) -def test_snapshot_preserves_matching_rates_and_quantities_for_each_billing_unit(unit: PricingUnit, value: str) -> None: - quantity: Final = PricedQuantity(measure="input", unit=unit, value=Decimal(value)) - rate: Final = Rate(measure="input", unit=unit, usd_per_unit=Decimal("0.003")) - amount: Final = quantity.value * rate.usd_per_unit - snapshot: Final = _snapshot(rates=(rate,), usage=(quantity,), total_cost=amount) - restored: Final = PricingSnapshot.model_validate_json(snapshot.model_dump_json()) - assert restored == snapshot - assert restored.usage[0].value * restored.rates[0].usd_per_unit == amount - - -@pytest.mark.parametrize( - ("usage", "message"), - [ - ((), "Every rate requires"), - (({"measure": "output", "unit": "token", "value": Decimal(10)},), "Every rate requires"), - (({"measure": "input", "unit": "second", "value": Decimal(10)},), "Every rate requires"), - ( - ({"measure": "input", "unit": "token", "value": Decimal(10)},) * 2, - "Duplicate priced quantity", - ), - ], -) -def test_snapshot_rejects_missing_mismatched_or_duplicate_priced_quantities( - usage: tuple[Mapping[str, object], ...], message: str -) -> None: - with pytest.raises(ValidationError, match=message): - _snapshot(usage=usage) - - -@pytest.mark.parametrize("unit", ["token", "character", "image", "request"]) -def test_priced_counts_remain_integral(unit: PricingUnit) -> None: - with pytest.raises(ValidationError, match="integral"): - PricedQuantity(measure="input", unit=unit, value=Decimal("1.25")) - - -@pytest.mark.parametrize( - ("value", "message"), - [ - (Decimal(-1), "greater than or equal"), - (Decimal("NaN"), "finite"), - (Decimal("Infinity"), "finite"), - (1.25, "instance of Decimal"), - ], -) -def test_pricing_values_reject_invalid_amounts(value: Decimal | float, message: str) -> None: - with pytest.raises(ValidationError, match=message): - PricedQuantity.model_validate({"measure": "input", "unit": "second", "value": value}) - with pytest.raises(ValidationError, match=message): - Rate.model_validate({"measure": "input", "unit": "second", "usd_per_unit": value}) - - -@pytest.mark.parametrize("snapshot_fields", [{}, {"snapshot": None}]) -def test_cost_event_cannot_omit_its_derivation(snapshot_fields: Mapping[str, object]) -> None: - with pytest.raises(ValidationError, match="snapshot"): - CostEvent.model_validate( - { - "event_id": "cost", - "usage_event_id": "usage", - "occurred_at": _START, - "dimensions": (), - "charges": (_SPEND,), - **snapshot_fields, - } - ) - - -def test_upstream_charge_snapshot_preserves_reported_cost_without_inventing_rates() -> None: - snapshot: Final = _snapshot( - pricing_source="upstream_header", rates=(), usage=(), upstream_reported_cost=_SPEND.value - ) - restored: Final = PricingSnapshot.model_validate_json(snapshot.model_dump_json()) - assert restored.upstream_reported_cost == restored.total_cost == _SPEND.value - assert restored.rates == restored.usage == () - - -def test_subset_token_counts_round_trip_without_changing_totals() -> None: - metrics: Final = ( - Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(10)), - Metric(name=MetricName.REASONING_TOKENS, value=Decimal(3)), - Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(8)), - Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, value=Decimal(2)), - ) - event: Final = _outbound(metrics=metrics) - assert OutboundUsageEvent.model_validate_json(event.model_dump_json()).metrics == metrics - - -@pytest.mark.parametrize( - "schema_path", - [ - "schema.prisma", - "litellm/proxy/schema.prisma", - "litellm-proxy-extras/litellm_proxy_extras/schema.prisma", - ], -) -@pytest.mark.parametrize( - ("enum_name", "python_enum"), - [ - ("AccountingDeliveryState", DeliveryState), - ("AccountingDeliveryTarget", DeliveryTarget), - ], -) -def test_prisma_accepts_exactly_the_public_accounting_enum_values( - schema_path: str, enum_name: str, python_enum: type[Enum] -) -> None: - schema: Final = (Path(__file__).resolve().parents[3] / schema_path).read_text() - declaration: Final = re.search(r"enum " + enum_name + r"\s*\{([^}]+)\}", schema) - assert declaration is not None, f"Missing persisted enum {enum_name}" - assert frozenset(declaration[1].split()) == frozenset(member.value for member in python_enum) - - -def test_one_event_carries_every_matched_model_access_group(): - groups = ( - Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value="gpt-group"), - Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value="claude-group"), - ) - assert _outbound(dimensions=(_TEAM, *groups)).dimensions == (_TEAM, *groups) - with pytest.raises(ValidationError, match="Duplicate dimension values"): - _outbound(dimensions=(groups[0], groups[0])) - - -def test_snapshots_reject_duplicate_multipliers_and_contradictory_token_counts(): - multiplier = Multiplier(name="data_residency", value="eu", factor=Decimal("1.1")) - with pytest.raises(ValidationError, match="Duplicate multiplier names"): - _snapshot(multipliers=(multiplier, multiplier)) - with pytest.raises(ValidationError, match="reasoning_tokens cannot exceed completion_tokens"): - _snapshot( - usage=( - PricedQuantity(measure="input", unit="token", value=_TOKEN.value), - PricedQuantity(measure="completion_tokens", unit="token", value=Decimal(5)), - PricedQuantity(measure="reasoning_tokens", unit="token", value=Decimal(6)), - ) - ) - - -def test_agent_live_spend_is_a_delivery_target(): - assert DeliveryTarget.AGENT_SPEND.value == "LiteLLM_AgentsTable" diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 5b37cee88c5..92114b00398 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -24,12 +24,12 @@ import httpx import pytest import respx from fastapi import HTTPException +from litellm_enterprise.proxy.common_utils.check_batch_cost import GIVE_UP_POLLING_AFTER import litellm from litellm.constants import MANAGED_OBJECT_STALENESS_CUTOFF_DAYS from litellm.types.router import LiteLLM_Params from litellm.types.utils import LiteLLMBatch -from litellm_enterprise.proxy.common_utils.check_batch_cost import GIVE_UP_POLLING_AFTER if TYPE_CHECKING: from litellm.batches.batch_utils import BatchCostUsageResult @@ -254,7 +254,6 @@ class TestCheckBatchCost: @pytest.mark.asyncio 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=1) # Simulate column already known absent from a previous cycle @@ -2387,8 +2386,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: @pytest.mark.asyncio async def test_target_model_names_comes_from_input_file_not_provider_model(self): from litellm.proxy.openai_files_endpoints.common_utils import ( - is_base64_encoded_unified_file_id, get_models_from_unified_file_id, + is_base64_encoded_unified_file_id, ) output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) @@ -2422,8 +2421,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: @pytest.mark.asyncio async def test_falls_back_to_deployment_model_group_without_managed_input_file(self): from litellm.proxy.openai_files_endpoints.common_utils import ( - is_base64_encoded_unified_file_id, get_models_from_unified_file_id, + is_base64_encoded_unified_file_id, ) output_file_id = await self._run(self._job("file-raw-provider-input")) @@ -3309,7 +3308,9 @@ class TestMultiPodBatchCostClaim: 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, "total_tokens": 15}, ["gpt-4"]), + return_value=_batch_cost_result( + 0.01, {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, ["gpt-4"] + ), ), patch( "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", @@ -3329,15 +3330,15 @@ class TestMultiPodBatchCostClaim: from datetime import datetime import httpx - import litellm import respx + import litellm from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger - from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.utils import PrismaClient + from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.types.utils import SpecialEnums @@ -3383,6 +3384,14 @@ class TestMultiPodBatchCostClaim: database_url="postgresql://unused:unused@localhost:5432/unused", proxy_logging_obj=dependency ) create_spend_rows: Final = AsyncMock(return_value=1) + from litellm.proxy.db.prisma_client import PrismaWrapper + + spend_prisma.db = PrismaWrapper( + original_prisma=MagicMock( + query_raw=AsyncMock(return_value=[]), + litellm_spendlogs=MagicMock(create_many=create_spend_rows), + ) + ) input_rate: Final = 0.000002 litellm.register_model( { @@ -3434,11 +3443,6 @@ class TestMultiPodBatchCostClaim: respx.mock, patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency), patch("litellm.proxy.proxy_server.prisma_client", spend_prisma), - patch.object( - spend_prisma.writer_db._original_prisma, - "litellm_spendlogs", - MagicMock(create_many=create_spend_rows), - ), ): respx.get(f"https://api.openai.com/v1/files/{_CLAIM_OUTPUT_FILE_ID}/content").mock( return_value=httpx.Response(200, content=content) @@ -3483,9 +3487,9 @@ class TestMultiPodBatchCostClaim: from datetime import datetime import httpx - import litellm import respx + import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.types.utils import LiteLLMBatch, SpecialEnums @@ -3587,9 +3591,9 @@ class TestMultiPodBatchCostClaim: import base64 import httpx - import litellm import respx + import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import SpecialEnums @@ -4423,9 +4427,10 @@ class TestBatchLineItemsFromThePoller: @pytest.mark.asyncio async def test_a_batch_bigger_than_one_part_reaches_the_callbacks_as_numbered_parts(self, monkeypatch): - import litellm from litellm_enterprise.proxy.common_utils import check_batch_cost + import litellm + journal: list = [] _, instance = self._setup(journal) logger = self._recording_logger(journal) @@ -4933,9 +4938,10 @@ async def test_failed_retirement_keeps_batch_available() -> None: @pytest.mark.parametrize("resource_id", ("batch-123", "file-123", "gs://bucket/job/predictions.jsonl")) @pytest.mark.parametrize("other_id", ("{id}4", "prefix-{id}", "{id}/child", "{id}.backup", "{id}?generation=2")) def test_another_resource_in_a_404_is_not_permanent_loss(resource_id: str, other_id: str) -> None: - from litellm.exceptions import NotFoundError from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + from litellm.exceptions import NotFoundError + error: Final = NotFoundError( message=f"No object found with id '{other_id.format(id=resource_id)}'.", model="test", llm_provider="openai" ) @@ -4958,9 +4964,10 @@ def test_another_resource_in_a_404_is_not_permanent_loss(resource_id: str, other ), ) def test_a_404_naming_the_exact_resource_is_permanent_loss(resource_id: str, message: str) -> None: - from litellm.exceptions import NotFoundError from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + from litellm.exceptions import NotFoundError + error: Final = NotFoundError(message=message.format(id=resource_id), model="test", llm_provider="openai") assert CheckBatchCost._is_batch_gone_at_provider(error, resource_id) is True assert CheckBatchCost._is_output_file_gone_at_provider(error, resource_id) is True @@ -4970,9 +4977,10 @@ def test_a_404_naming_the_exact_resource_is_permanent_loss(resource_id: str, mes ("missing_key", "gone"), (("out/job-123/input.jsonl.out", True), ("out/job-123/other.jsonl.out", False)) ) def test_an_s3_no_such_key_for_the_output_object_is_permanent_loss(missing_key: str, gone: bool) -> None: - from litellm.llms.bedrock.common_utils import BedrockError from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + from litellm.llms.bedrock.common_utils import BedrockError + # S3 GetObject 404 body shape, checked live against a missing key on 2026-10-20 error: Final = BedrockError( status_code=404, @@ -4985,9 +4993,10 @@ def test_an_s3_no_such_key_for_the_output_object_is_permanent_loss(missing_key: @pytest.mark.parametrize("output_file_id", (None, "")) def test_a_404_without_an_output_identity_remains_retryable(output_file_id: str | None) -> None: - from litellm.exceptions import NotFoundError from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + from litellm.exceptions import NotFoundError + error: Final = NotFoundError(message="Output file not found", model="test", llm_provider="openai") assert CheckBatchCost._is_output_file_gone_at_provider(error, output_file_id) is False @@ -5003,11 +5012,12 @@ async def test_batch_cost_reaches_configured_success_callbacks( initialized: bool, claimed: bool, ) -> None: + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.types.utils import LiteLLMBatch - from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost events: Final[SimpleQueue[tuple[str, object, int, int]]] = SimpleQueue() diff --git a/tests/unit/proxy/db/test_budget_window_spend_writer.py b/tests/unit/proxy/db/test_budget_window_spend_writer.py index 27f1472c183..7315d51c5dd 100644 --- a/tests/unit/proxy/db/test_budget_window_spend_writer.py +++ b/tests/unit/proxy/db/test_budget_window_spend_writer.py @@ -407,9 +407,7 @@ async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_o @pytest.mark.asyncio -@pytest.mark.parametrize( - ("replaces_window_start", "replaced_param"), [(None, None), (WINDOW_A, datetime(2026, 8, 1))] -) +@pytest.mark.parametrize(("replaces_window_start", "replaced_param"), [(None, None), (WINDOW_A, datetime(2026, 8, 1))]) async def test_open_window_spend_row_moves_only_the_window_the_caller_saw_end(replaces_window_start, replaced_param): db = _FakeDB() @@ -425,7 +423,9 @@ async def test_open_window_spend_row_moves_only_the_window_the_caller_saw_end(re ((query, params),) = db.execute_raw_calls normalized = " ".join(query.split()) assert "ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE" in normalized - assert normalized.endswith("WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')") + assert normalized.endswith( + "WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')" + ) assert params == ("model_max_budget", "k1", "1d", datetime(2026, 8, 31), replaced_param) @@ -561,7 +561,12 @@ async def test_seed_aggregate_sql_splits_the_window_at_the_batch_start(entity_ty ((query, params),) = db.query_raw_calls normalized = " ".join(query.split()) assert expected_column in normalized - assert "FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC'))" in normalized + assert ( + 'FILTER (WHERE EXISTS (SELECT 1 FROM "LiteLLM_SpendLogDelivery" delivery ' + 'WHERE delivery.request_id = "LiteLLM_SpendLogs".request_id) ' + "OR \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC'))" + in normalized + ) assert 'FROM "LiteLLM_SpendLogs"' in normalized # startTime is TIMESTAMP(3): the bound is floored to the second so the # batch's own earliest row cannot round under it. @@ -635,3 +640,21 @@ async def test_rolling_a_window_row_renders_a_postgres_update_span( new_window_start=WINDOW_B, ) assert await postgres_span_names() == ("postgres.update LiteLLM_BudgetWindowSpend",) + + +@pytest.mark.asyncio +async def test_window_reads_inside_an_enclosing_transaction_use_its_connection(): + client_db = _FakeDB() + transaction = _FakeDB(existing_rows=[]) + + await commit_window_spend_updates( + prisma_client=_FakePrismaClient(client_db), + transactions=( + build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0, started_at=BATCH_STARTED_AT), + ), + outer_transaction=transaction, + ) + + assert client_db.query_raw_calls == [] + assert len(transaction.query_raw_calls) == 2 + assert len(transaction.batcher.calls) == 1 diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index 3a79df0370b..d90c7f6d21d 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -21,6 +21,7 @@ from redis.exceptions import DataError import litellm from litellm._logging import verbose_proxy_logger + from litellm._service_logger import ServiceTypes from litellm.constants import SPEND_ROLLUP_LOCK_TIMEOUT_MS from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendUpdateQueueItem @@ -4382,7 +4383,6 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_only_on_rows_the_bri assert prisma.spend_log_flush_requested.is_set() is expects_flush - @pytest.mark.asyncio @pytest.mark.timeout(30) async def test_a_chatgpt_responses_row_queues_without_a_default_chatgpt_login( @@ -4424,6 +4424,8 @@ def _batch_cost_payload() -> dict: def _spend_logs_prisma(inserted: int, existing: object, taken_over: int = 1) -> MagicMock: prisma = _tool_usage_prisma() prisma.jsonify_object = lambda data: dict(data) + prisma.writer_db = prisma.db + prisma.db.query_raw = AsyncMock(return_value=[]) prisma.db.litellm_spendlogs.create_many = AsyncMock(return_value=inserted) prisma.db.litellm_spendlogs.find_unique = AsyncMock(return_value=existing) prisma.db.litellm_spendlogs.update_many = AsyncMock(return_value=taken_over) @@ -4745,6 +4747,7 @@ def _spend_logs_table_prisma() -> MagicMock: prisma = _tool_usage_prisma() prisma.jsonify_object = lambda data: dict(data) prisma.db.litellm_spendlogs = _SpendLogsTable() + prisma.writer_db.query_raw = AsyncMock(return_value=[]) return prisma @@ -5839,3 +5842,36 @@ async def test_completed_spend_survives_error_log_suppression( assert len(rows) == 1 assert rows[0]["spend"] == pytest.approx(0.23) assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == (10, 5) + + +@pytest.mark.parametrize("ledger_enabled", [False, True]) +@pytest.mark.asyncio +async def test_request_error_counts_survive_the_ledger_path_switch(monkeypatch: pytest.MonkeyPatch, ledger_enabled: bool): + from litellm.proxy import proxy_server + from litellm.proxy.db import db_spend_update_writer as writer_module + from litellm.proxy.db.request_error_tracking import RequestErrorAccumulator + from litellm.types.proxy.request_errors import RequestErrorKey + + accumulator = RequestErrorAccumulator() + prisma = _tool_usage_prisma() + prisma.get_request_status.return_value = "failure" + monkeypatch.setattr(writer_module, "request_error_accumulator", accumulator) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": ledger_enabled}) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + now = datetime(2026, 11, 2, tzinfo=timezone.utc) + await DBSpendUpdateWriter().update_database( + token="key", user_id="user", team_id=None, org_id=None, end_user_id=None, + kwargs={ + "model": "ledger-test", "call_type": "acompletion", "litellm_call_id": "failed-request", + "litellm_params": {"metadata": { + "user_api_key": "key", "user_api_key_user_id": "user", "status": "failure", + "error_information": {"error_code": "503"}, + }}, + }, + completion_response=None, start_time=now, end_time=now, response_cost=0, + ) + await drain_logging_tasks() + assert accumulator.drain() == {RequestErrorKey( + date="2026-11-02", api_key=prisma.spend_log_transactions[0]["api_key"], team_id="", user_id="user", model_group="ledger-test", status_code=503, + ): 1} diff --git a/tests/unit/proxy/guardrails/test_usage_tracking.py b/tests/unit/proxy/guardrails/test_usage_tracking.py index 69ec098b840..7e92c9be4e8 100644 --- a/tests/unit/proxy/guardrails/test_usage_tracking.py +++ b/tests/unit/proxy/guardrails/test_usage_tracking.py @@ -654,3 +654,15 @@ async def test_one_failing_index_statement_does_not_drop_the_others_or_the_rollu assert prisma.db.litellm_spendlogguardrailindex.create_many.await_count == 3 assert prisma.db.litellm_dailyguardrailmetrics.upsert.await_count == len(guardrail_ids) + + +@pytest.mark.asyncio +async def test_disabled_log_deliveries_do_not_enable_guardrail_analytics(): + prisma = _prisma() + row = _payload("lean", usage={"topicPolicyUnits": 2}) + row["accounting_payload"] = {"request_id": "lean"} + row["accounting_log_disabled"] = True + await process_spend_logs_guardrail_usage(prisma, [row], pending=PendingRollups()) + prisma.db.litellm_dailyguardrailmetrics.upsert.assert_not_awaited() + prisma.db.litellm_spendlogguardrailindex.create_many.assert_not_awaited() + prisma.db.litellm_dailyguardrailusageunits.upsert.assert_not_awaited() diff --git a/tests/unit/proxy/hooks/test_proxy_hooks_init.py b/tests/unit/proxy/hooks/test_proxy_hooks_init.py index a6edd3db944..417ea5a9650 100644 --- a/tests/unit/proxy/hooks/test_proxy_hooks_init.py +++ b/tests/unit/proxy/hooks/test_proxy_hooks_init.py @@ -33,7 +33,7 @@ def test_managed_vector_stores_hook_registered(): assert hook_cls.__name__ == "_PROXY_LiteLLMManagedVectorStores" -def test_isolation_module_does_not_pull_in_proxy_utils(): +def test_isolation_module_does_not_pull_in_proxy_utils(monkeypatch): """Layering guard: litellm.llms.* must not transitively import litellm.proxy.utils, which would reintroduce the import cycle.""" import importlib @@ -44,7 +44,7 @@ def test_isolation_module_does_not_pull_in_proxy_utils(): "litellm.proxy.management_endpoints.common_utils", "litellm.llms.base_llm.managed_resources.isolation", ]: - sys.modules.pop(mod, None) + monkeypatch.delitem(sys.modules, mod, raising=False) importlib.import_module("litellm.llms.base_llm.managed_resources.isolation") assert "litellm.proxy.utils" not in sys.modules diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index a87ac7fa7c0..9267a951459 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -588,3 +588,10 @@ async def test_group_reserves_on_its_current_deployments_after_another_vendors_d by_own_name: Final = litellm.token_counter(model="prod", messages=_MULTILINGUAL_MESSAGES) assert openai_only != by_own_name assert (await _input_tokens(router, "prod"))["prod"] == by_own_name + + +@pytest.mark.parametrize("window", [{}, {"budget_duration": "1d"}, {"reset_at": "2026-11-03T00:00:00Z"}]) +def test_fixed_budget_window_requires_duration_and_reset_boundary(window: object) -> None: + from litellm.proxy.spend_tracking.budget_reservation import fixed_budget_window + + assert fixed_budget_window(window) is None diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 448cc498c7a..45ded91b61e 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -47,7 +47,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_spend_logs_metadata, _get_vector_store_request_for_spend_logs_payload, is_master_key, - _redact_logged_api_key, + redact_logged_api_key, _redact_prompt_leaks_in_error_string, sanitize_error_information_for_spend_logs, _sanitize_guardrail_information_for_spend_logs, @@ -4038,7 +4038,7 @@ def test_sanitize_error_information_redacts_pydantic_assignment_form( assert REDACTED_BY_LITELM_STRING in sanitized["error_message"] -# ── _redact_logged_api_key unit tests ────────────────────────────────────── +# ── redact_logged_api_key unit tests ────────────────────────────────────── @pytest.mark.parametrize( @@ -4067,16 +4067,16 @@ def test_sanitize_error_information_persists_no_raw_model_for_an_unknown_model_r def test_redact_logged_api_key_none_returns_none(): - assert _redact_logged_api_key(None) is None + assert redact_logged_api_key(None) is None def test_redact_logged_api_key_empty_string_returns_none(): - assert _redact_logged_api_key("") is None + assert redact_logged_api_key("") is None def test_redact_logged_api_key_sk_key_is_hashed(): raw = "sk-9876secret" - result = _redact_logged_api_key(raw) + result = redact_logged_api_key(raw) assert result == hash_token(raw) assert result is not None assert not result.startswith("sk-") @@ -4085,15 +4085,15 @@ def test_redact_logged_api_key_sk_key_is_hashed(): def test_redact_logged_api_key_bearer_sk_equals_sk_hash(): raw = "sk-9876secret" - result_plain = _redact_logged_api_key(raw) - result_bearer = _redact_logged_api_key(f"Bearer {raw}") + result_plain = redact_logged_api_key(raw) + result_bearer = redact_logged_api_key(f"Bearer {raw}") assert result_bearer == result_plain def test_redact_logged_api_key_bearer_case_insensitive(): raw = "sk-9876secret" - result_lower = _redact_logged_api_key(f"bearer {raw}") - result_upper = _redact_logged_api_key(f"BEARER {raw}") + result_lower = redact_logged_api_key(f"bearer {raw}") + result_upper = redact_logged_api_key(f"BEARER {raw}") expected = hash_token(raw) assert result_lower == expected assert result_upper == expected @@ -4101,7 +4101,7 @@ def test_redact_logged_api_key_bearer_case_insensitive(): def test_redact_logged_api_key_non_sk_raw_key_is_hashed(): raw = "anthropic-raw-key-xyz" - result = _redact_logged_api_key(raw) + result = redact_logged_api_key(raw) assert result is not None assert result != raw assert len(result) == 64 @@ -4111,7 +4111,7 @@ def test_redact_logged_api_key_non_sk_raw_key_is_hashed(): def test_redact_logged_api_key_already_valid_sha256_passes_through_with_flag(): already_hashed = hash_token("sk-some-key") assert len(already_hashed) == 64 - result = _redact_logged_api_key(already_hashed, already_redacted=True) + result = redact_logged_api_key(already_hashed, already_redacted=True) assert result == already_hashed assert hash_token(already_hashed) != result # no double-hash @@ -4119,7 +4119,7 @@ def test_redact_logged_api_key_already_valid_sha256_passes_through_with_flag(): def test_redact_logged_api_key_sha256_without_flag_is_hashed(): already_hashed = hash_token("sk-some-key") assert len(already_hashed) == 64 - result = _redact_logged_api_key(already_hashed) + result = redact_logged_api_key(already_hashed) assert result is not None assert result != already_hashed assert len(result) == 64 @@ -4129,7 +4129,7 @@ def test_redact_logged_api_key_sha256_without_flag_is_hashed(): def test_redact_logged_api_key_long_opaque_token_is_hashed(): raw = "x1" * 450 assert len(raw) == 900 - result = _redact_logged_api_key(raw) + result = redact_logged_api_key(raw) assert result is not None assert result != raw assert raw not in result @@ -4139,27 +4139,27 @@ def test_redact_logged_api_key_long_opaque_token_is_hashed(): def test_redact_logged_api_key_hashed_jwt_passes_through(): jwt_hash = "hashed-jwt-" + "a" * 64 - result = _redact_logged_api_key(jwt_hash, already_redacted=True) + result = redact_logged_api_key(jwt_hash, already_redacted=True) assert result == jwt_hash def test_redact_logged_api_key_hashed_jwt_shape_without_provenance_is_hashed(): lookalike = "hashed-jwt-" + "a" * 64 - result = _redact_logged_api_key(lookalike) + result = redact_logged_api_key(lookalike) assert result == hash_token(lookalike) assert result != lookalike def test_redact_logged_api_key_hashed_jwt_trailing_newline_is_hashed(): trailing = "hashed-jwt-" + "a" * 64 + "\n" - result = _redact_logged_api_key(trailing, already_redacted=True) + result = redact_logged_api_key(trailing, already_redacted=True) assert result == hash_token(trailing) assert result != trailing def test_redact_logged_api_key_hashed_jwt_short_suffix_is_hashed(): short_jwt = "hashed-jwt-tooshort" - result = _redact_logged_api_key(short_jwt) + result = redact_logged_api_key(short_jwt) assert result is not None assert result != short_jwt assert len(result) == 64 @@ -4169,14 +4169,14 @@ def test_redact_logged_api_key_hashed_jwt_short_suffix_is_hashed(): def test_redact_logged_api_key_master_key_alias_passes_through(): from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS - result = _redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS, already_redacted=True) + result = redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS, already_redacted=True) assert result == LITELLM_PROXY_MASTER_KEY_ALIAS def test_redact_logged_api_key_master_key_alias_without_provenance_is_hashed(): from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS - result = _redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS) + result = redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS) assert result == hash_token(LITELLM_PROXY_MASTER_KEY_ALIAS) assert result != LITELLM_PROXY_MASTER_KEY_ALIAS @@ -4282,9 +4282,9 @@ def test_get_logging_payload_reads_used_client_oauth_token_from_the_bucket_the_p def test_redact_logged_api_key_bearer_only_returns_none(): # "bearer " with nothing after stripping is equivalent to no key - assert _redact_logged_api_key("bearer ") is None - assert _redact_logged_api_key("Bearer ") is None - assert _redact_logged_api_key("BEARER ") is None + assert redact_logged_api_key("bearer ") is None + assert redact_logged_api_key("Bearer ") is None + assert redact_logged_api_key("BEARER ") is None # ── _get_spend_logs_metadata key-hash invariant tests ───────────────────── @@ -4645,7 +4645,7 @@ class TestSpendLogKeyRedaction: def test_bearer_prefixed_sk_key_is_hashed(self): raw = "Bearer sk-WLi4iRn4JmbVlTaYw12IOA" - result = _redact_logged_api_key(raw) + result = redact_logged_api_key(raw) assert result is not None assert not result.startswith("Bearer") assert not result.startswith("sk-") @@ -4653,14 +4653,14 @@ class TestSpendLogKeyRedaction: def test_bare_sk_key_is_hashed(self): raw = "sk-WLi4iRn4JmbVlTaYw12IOA" - result = _redact_logged_api_key(raw) + result = redact_logged_api_key(raw) assert result is not None assert not result.startswith("sk-") assert len(result) == 64 def test_bearer_lowercase_is_handled(self): raw = "bearer sk-WLi4iRn4JmbVlTaYw12IOA" - result = _redact_logged_api_key(raw) + result = redact_logged_api_key(raw) assert result is not None assert not result.startswith("bearer") assert not result.startswith("sk-") @@ -4668,11 +4668,11 @@ class TestSpendLogKeyRedaction: def test_already_hashed_key_unchanged(self): hashed = "bcfe8173f5447f10be0e7fb37aaa8b97829d5c9e0498232152f9d123456789ab" - assert _redact_logged_api_key(hashed, already_redacted=True) == hashed + assert redact_logged_api_key(hashed, already_redacted=True) == hashed def test_bearer_prefixed_non_sk_key_is_hashed(self): raw = "Bearer some-other-token-format" - result = _redact_logged_api_key(raw) + result = redact_logged_api_key(raw) assert result == hash_token("some-other-token-format") assert result is not None assert not result.startswith("Bearer") @@ -4680,7 +4680,7 @@ class TestSpendLogKeyRedaction: def test_bearer_and_bare_produce_same_hash(self): bare = "sk-WLi4iRn4JmbVlTaYw12IOA" bearer = "Bearer sk-WLi4iRn4JmbVlTaYw12IOA" - assert _redact_logged_api_key(bare) == _redact_logged_api_key(bearer) + assert redact_logged_api_key(bare) == redact_logged_api_key(bearer) @patch("litellm.proxy.proxy_server.master_key", None) @@ -4777,7 +4777,7 @@ def test_get_logging_payload_keeps_internal_service_account_key_readable(service def test_redact_logged_api_key_service_account_name_without_provenance_is_hashed(): - result = _redact_logged_api_key(LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME) + result = redact_logged_api_key(LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME) assert result == hash_token(LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME) @@ -5889,7 +5889,7 @@ def test_get_logging_payload_empty_key_slp_none_is_empty_string_not_none_literal assert payload["api_key"] == "", ( f"Expected empty string but got {payload['api_key']!r}; " - "dropping _redact_logged_api_key's 'or \"\"' guard would yield 'None' here" + "dropping redact_logged_api_key's 'or \"\"' guard would yield 'None' here" ) @@ -5940,7 +5940,7 @@ def test_get_logging_payload_stores_user_email_in_spend_log_metadata(user_email, def test_redact_logged_api_key_partial_sha256_is_hashed(): partial_hex = "a" * 63 - result = _redact_logged_api_key(partial_hex) + result = redact_logged_api_key(partial_hex) assert result is not None assert result != partial_hex assert len(result) == 64 @@ -5950,7 +5950,7 @@ def test_redact_logged_api_key_partial_sha256_is_hashed(): def test_redact_logged_api_key_bearer_already_hashed_passes_through_with_flag(): already_hashed = hash_token("sk-some-key") assert len(already_hashed) == 64 - result = _redact_logged_api_key(f"Bearer {already_hashed}", already_redacted=True) + result = redact_logged_api_key(f"Bearer {already_hashed}", already_redacted=True) assert result == already_hashed assert hash_token(already_hashed) != result @@ -5958,7 +5958,7 @@ def test_redact_logged_api_key_bearer_already_hashed_passes_through_with_flag(): def test_redact_logged_api_key_bearer_sha256_without_flag_is_hashed(): already_hashed = hash_token("sk-some-key") assert len(already_hashed) == 64 - result = _redact_logged_api_key(f"Bearer {already_hashed}") + result = redact_logged_api_key(f"Bearer {already_hashed}") assert result is not None assert result != already_hashed assert result == hash_token(already_hashed) @@ -6570,12 +6570,12 @@ def test_get_logging_payload_never_lands_a_raw_cli_session_token(call_type: str) def test_redact_logged_api_key_cli_session_alias_needs_alias_provenance(): - assert _redact_logged_api_key(_CLI_SESSION_ALIAS, already_redacted=True, key_alias=_CLI_SESSION_ALIAS) == ( + assert redact_logged_api_key(_CLI_SESSION_ALIAS, already_redacted=True, key_alias=_CLI_SESSION_ALIAS) == ( _CLI_SESSION_ALIAS ) - assert _redact_logged_api_key(_CLI_SESSION_ALIAS, already_redacted=True) == hash_token(_CLI_SESSION_ALIAS) - assert _redact_logged_api_key(_CLI_SESSION_ALIAS, key_alias=_CLI_SESSION_ALIAS) == hash_token(_CLI_SESSION_ALIAS) - assert _redact_logged_api_key("alice", already_redacted=True, key_alias="alice") == hash_token("alice") + assert redact_logged_api_key(_CLI_SESSION_ALIAS, already_redacted=True) == hash_token(_CLI_SESSION_ALIAS) + assert redact_logged_api_key(_CLI_SESSION_ALIAS, key_alias=_CLI_SESSION_ALIAS) == hash_token(_CLI_SESSION_ALIAS) + assert redact_logged_api_key("alice", already_redacted=True, key_alias="alice") == hash_token("alice") def test_azure_spillover_stamped_from_response_headers(): @@ -7814,3 +7814,45 @@ def test_spend_payload_preserves_unknown_completion_start(completion_start: str assert payload["completionStartTime"] == ( end if completion_start == "absent" else reported if completion_start else None ) + + +def test_spend_log_metadata_retention_drops_accounting_inputs_it_was_not_told_to_keep() -> None: + metadata: Final = { + "status": "success", + "accounting_savings": { + "compression": 0.4, + "prompt_caching": 0.1, + "gateway_injected_caching": 0.2, + "autorouter": 0.3, + }, + "accounting_windows": [ + {"entity_type": "key", "entity_id": "key", "window_duration": "1d", "window_start": "2026-11-02T00:00:00"} + ], + "accounting_model_access_groups": ["group"], + "user_api_key_project_id": "project", + "usage_object": {"prompt_tokens_details": {"cached_tokens": 2}}, + "compression_savings": {"tokens_saved": 3}, + "guardrail_information": [{"guardrail_id": "guard", "guardrail_usage": {"text": 2}}], + "internal_call_origin": "classifier", + "billed_retry_attempt": True, + "user_api_key_alias": "discard", + } + row: Final = { + "request_id": "retained-accounting", + "spend": 0.5, + "startTime": "2026-11-02T12:00:00", + "endTime": "2026-11-02T12:00:01", + "model": "model", + "model_group": "group", + "api_key": "key", + "custom_llm_provider": "openai", + "user": "user", + "team_id": "team", + "agent_id": "agent", + "request_tags": '["tag"]', + "metadata": json.dumps(metadata), + } + stored: Final = spend_log_row_with_retained_metadata(row, SpendLogsMetadataFields(include=("usage_object",))) + parsed: Final = json.loads(str(stored["metadata"])) + assert "user_api_key_alias" not in parsed + assert parsed == {"usage_object": metadata["usage_object"], "status": metadata["status"]} diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 9dd15649d3e..62d8db6f810 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -16589,3 +16589,63 @@ async def test_initialize_prints_the_banner_as_one_json_record_in_json_mode( banner_records: Final = [record for record in records if record["message"] == LITELLM_BANNER] assert len(banner_records) == 1, records assert banner_records[0]["level"] == "INFO" + + +@contextlib.contextmanager +def _accounting_window_env(cached: Mapping[str, object], client: MagicMock) -> Iterator[None]: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache: Final = UserApiKeyCache() + for key, value in cached.items(): + cache.set_cache(key=key, value=value) + originals: Final = (proxy_server_module.user_api_key_cache, proxy_server_module.prisma_client) + proxy_server_module.user_api_key_cache = cache + proxy_server_module.prisma_client = client + try: + yield + finally: + proxy_server_module.user_api_key_cache, proxy_server_module.prisma_client = originals + + +@pytest.mark.asyncio +async def test_accounting_window_keys_cover_the_stored_key_and_the_cached_team(): + key_reset: Final = datetime(2026, 11, 13, tzinfo=timezone.utc) + team_reset: Final = datetime(2026, 11, 6, tzinfo=timezone.utc) + key_row: Final = _StoredRow( + budget_limits=[{"budget_duration": "30d", "max_budget": 100.0, "reset_at": key_reset.isoformat()}] + ) + team_obj: Final = MagicMock() + team_obj.budget_limits = [{"budget_duration": "7d", "max_budget": 50.0, "reset_at": team_reset.isoformat()}] + + with _accounting_window_env({"team_id:team-1": team_obj}, _prisma_with_rows(keys={"hashed-key": key_row})): + windows = await proxy_server_module.accounting_window_keys("hashed-key", "team-1") + + assert windows == ( + { + "entity_type": "key", + "entity_id": "hashed-key", + "window_duration": "30d", + "window_start": (key_reset - timedelta(days=30)).isoformat(), + }, + { + "entity_type": "team", + "entity_id": "team-1", + "window_duration": "7d", + "window_start": (team_reset - timedelta(days=7)).isoformat(), + }, + ) + + +@pytest.mark.asyncio +async def test_accounting_window_keys_remember_a_key_without_windows_but_retry_a_failed_read(): + client: Final = _prisma_with_rows(keys={"no-windows": _StoredRow(budget_limits=[])}) + + with _accounting_window_env({}, client): + assert await proxy_server_module.accounting_window_keys("no-windows", None) == () + assert await proxy_server_module.accounting_window_keys("no-windows", None) == () + assert client.db.litellm_verificationtoken.find_unique.await_count == 1 + + client.db.litellm_verificationtoken.find_unique.side_effect = RuntimeError("database unavailable") + assert await proxy_server_module.accounting_window_keys("unreadable", None) == () + assert await proxy_server_module.accounting_window_keys("unreadable", None) == () + assert client.db.litellm_verificationtoken.find_unique.await_count == 3 diff --git a/tests/unit/proxy/test_update_spend.py b/tests/unit/proxy/test_update_spend.py index 571341e85e7..e2c027408c1 100644 --- a/tests/unit/proxy/test_update_spend.py +++ b/tests/unit/proxy/test_update_spend.py @@ -351,3 +351,125 @@ async def test_update_spend_logs_drops_row_prisma_cannot_encode_and_keeps_the_re assert sorted(written) == ["good-1", "good-2"] assert prisma_client.spend_log_transactions == [] proxy_logging_obj.failure_handler.assert_not_called() + + +def _raw_query_error(sqlstate: str): + from prisma.errors import RawQueryError + + return RawQueryError( + { + "user_facing_error": { + "error_code": "P2010", + "meta": {"code": sqlstate, "message": f"ERROR: sqlstate {sqlstate}"}, + } + } + ) + + +class _RawInsertRepository: + def __init__(self, sqlstate: str): + self.sqlstate = sqlstate + self.written: list[str] = [] + + async def insert_with_deliveries(self, rows): + if any(row["request_id"] == "rejected" for row in rows): + raise _raw_query_error(self.sqlstate) + self.written.extend(str(row["request_id"]) for row in rows) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sqlstate", ["22P02", "23502", "54000"]) +async def test_raw_insert_rejecting_one_row_drops_only_that_row(sqlstate): + from litellm.proxy.utils import _create_spend_logs_with_poison_isolation + + repo = _RawInsertRepository(sqlstate) + rows = [{"request_id": "good-1"}, {"request_id": "rejected"}, {"request_id": "good-2"}] + + await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10) + + assert sorted(repo.written) == ["good-1", "good-2"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sqlstate", ["42P01", "40P01"]) +async def test_systemic_raw_insert_failure_is_retried_not_dropped(sqlstate): + from prisma.errors import RawQueryError + + from litellm.proxy.utils import _create_spend_logs_with_poison_isolation + + repo = _RawInsertRepository(sqlstate) + rows = [{"request_id": "good-1"}, {"request_id": "rejected"}] + + with pytest.raises(RawQueryError): + await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10) + + assert repo.written == [] + + +class _RescueFailsRepository: + def __init__(self, rescue_error: Exception): + self.rescue_error = rescue_error + + async def insert_with_deliveries(self, rows): + if any(row.get("accounting_log_disabled") for row in rows): + raise self.rescue_error + if any(row["request_id"] == "rejected" for row in rows): + raise _raw_query_error("22P02") + + +@pytest.mark.asyncio +async def test_a_transient_failure_while_keeping_a_dropped_rows_charge_is_retried_not_lost(): + from litellm.proxy.utils import _create_spend_logs_with_poison_isolation + + outage = _raw_query_error("40P01") + repo = _RescueFailsRepository(outage) + rows = [{"request_id": "rejected", "accounting_payload": {"spend": 1.0}}] + + with pytest.raises(type(outage)): + await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10) + + +@pytest.mark.asyncio +async def test_a_rejected_charge_rescue_is_logged_and_the_flush_continues(): + from litellm.proxy.utils import _create_spend_logs_with_poison_isolation + + repo = _RescueFailsRepository(_raw_query_error("22P02")) + rows = [{"request_id": "rejected", "accounting_payload": {"spend": 1.0}}] + + assert await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10) == 9 + + +class _StatementValidatingRepository: + def __init__(self): + self.written = [] + self.charges_kept = [] + + async def insert_with_deliveries(self, rows): + from litellm.repositories.spend_log_insert import spend_log_insert_statement + + if all(row.get("accounting_log_disabled") for row in rows): + self.charges_kept.extend(row["request_id"] for row in rows) + return len(rows) + spend_log_insert_statement(rows) + self.written.extend(row["request_id"] for row in rows) + return len(rows) + + +@pytest.mark.asyncio +async def test_a_row_the_insert_statement_rejects_is_dropped_alone_and_keeps_its_charge(): + from litellm.proxy.utils import _create_spend_logs_with_poison_isolation + + repo = _StatementValidatingRepository() + rows = [ + {"request_id": "good", "call_type": "acompletion", "messages": "[]"}, + { + "request_id": "unparseable", + "call_type": "acompletion", + "messages": "not json", + "accounting_payload": {"spend": 1.0}, + }, + ] + + assert await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10) < 10 + assert repo.written == ["good"] + assert repo.charges_kept == ["unparseable"] diff --git a/tests/unit/repositories/test_spend_log_insert.py b/tests/unit/repositories/test_spend_log_insert.py new file mode 100644 index 00000000000..a099551086b --- /dev/null +++ b/tests/unit/repositories/test_spend_log_insert.py @@ -0,0 +1,31 @@ +from datetime import datetime, timezone +import json + +import pytest + +from litellm.repositories.spend_log_insert import spend_log_insert_statement + + +@pytest.mark.parametrize( + "rows", + [(), ({"request_id": "request", "unknown_column": 1},), ({"request_id": "first"}, {"spend": 1})], +) +def test_insert_rejects_empty_unknown_or_inconsistent_columns(rows): + with pytest.raises(ValueError, match="Spend log insert requires"): + spend_log_insert_statement(rows) + + +def test_insert_keeps_json_and_timestamps_in_bound_data(): + timestamp = datetime(2026, 11, 2, tzinfo=timezone.utc) + request_id = "request'); DROP TABLE unrelated; --" + statement, encoded = spend_log_insert_statement(({ + "request_id": request_id, "startTime": timestamp, "metadata": '{"project":"project"}', + "request_tags": '["tag"]', "messages": '[{"role":"user","content":"hello"}]', + },)) + row = json.loads(encoded)[0] + assert request_id not in statement + assert row["request_id"] == request_id + assert row["startTime"] == timestamp.isoformat() + assert row["metadata"] == {"project": "project"} + assert row["request_tags"] == ["tag"] + assert row["messages"] == [{"role": "user", "content": "hello"}] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 15d440c0ad4..e65a30f6e1d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30862,6 +30862,12 @@ export interface components { * @description decrypt keys with google kms */ use_google_kms?: boolean | null; + /** + * Use Spend Logs Ledger + * @description Apply request accounting from stored SpendLogs deliveries. Enable only after all accounting writers are upgraded and ledger activation gates pass. + * @default false + */ + use_spend_logs_ledger: boolean; /** * Use Spend Logs Partitioning * @description If True and LiteLLM_SpendLogs has been converted to a range-partitioned table (db_scripts/partition_spend_logs.sql), retention cleanup drops expired partitions instead of deleting rows, and pre-creates upcoming partitions. Default is False. From 5e91bc8e86b0481500648d9fd5d6778b614a6717 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 4 Nov 2026 18:02:18 -0800 Subject: [PATCH 03/39] fix(ui): send a team admin's member key edits on the first save (#50007) The key edit form read react-hook-form's dirtyFields only inside the submit handler. formState is a proxy that only tracks what the component reads while rendering, so on the first save of an edit session that read returned the initial empty object. A team admin editing a member's key then sent only the key id to /key/update, got a 200 and a "Key updated successfully" toast, and nothing changed: a max_budget edit was silently dropped, and a rename went through as an empty update instead of the "team admins can only change budget fields" error. Read dirtyFields during render so the form subscribes to it and the submit handler sees the fields the user actually changed. --- .../key_edit_view.integration.test.tsx | 55 +++++++++++++++++++ .../components/templates/key_edit_view.tsx | 3 +- 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx index dce14adc2e5..d590d7beef9 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx @@ -15,6 +15,7 @@ import { vectorStoreListCall, } from "../networking"; import { KeyEditView } from "./key_edit_view"; +import { teamAdminMemberKeyPayload } from "./teamAdminMemberKeyPayload"; const can = vi.fn(); vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ @@ -3234,4 +3235,58 @@ describe("KeyEditView", () => { expect(await screen.findByRole("switch", { name: /disable global guardrails/i })).toBeInTheDocument(); }); }); + + describe("team admin editing a member's key", () => { + const renderAsTeamAdmin = ( + onSubmit: (values: Record, dirtyFields: readonly string[]) => Promise, + ) => + renderWithProviders( + {}} + onSubmit={onSubmit} + accessToken="test-token" + userID="team-admin-user" + userRole="Internal User" + />, + ); + + const firstSavePayload = async (onSubmit: ReturnType) => { + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + await waitFor(() => expect(onSubmit).toHaveBeenCalledTimes(1)); + const [values, dirtyFields] = onSubmit.mock.calls[0]; + return teamAdminMemberKeyPayload({ ...values, key: values.token }, dirtyFields); + }; + + it("should block a rename on the first save instead of sending an empty update", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderAsTeamAdmin(onSubmit); + + fireEvent.change(await screen.findByRole("textbox", { name: "Key Alias" }), { target: { value: "renamed" } }); + + expect(await firstSavePayload(onSubmit)).toEqual({ kind: "blocked", fields: ["key_alias"] }); + }); + + it("should send a budget change on the first save", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderAsTeamAdmin(onSubmit); + + fireEvent.change(await screen.findByRole("spinbutton", { name: "Max Budget (USD)" }), { + target: { value: "25" }, + }); + + expect(await firstSavePayload(onSubmit)).toEqual({ + kind: "ok", + payload: { key: MOCK_KEY_DATA.token, max_budget: "25" }, + }); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 81464d0d195..940108f9721 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -116,6 +116,7 @@ function KeyEditSession({ values: toKeyEditFormValues(keyData), resetOptions: { keepDirtyValues: true }, }); + const { dirtyFields } = form.formState; const [promptsList, setPromptsList] = useState([]); const [tagsList, setTagsList] = useState>({}); const team = teams?.find((team) => team.team_id === keyData.team_id); @@ -321,7 +322,7 @@ function KeyEditSession({ ...values, ...(detachProject && enableProjectsUI && canDetachProject ? { project_id: null } : {}), }), - [...Object.keys(form.formState.dirtyFields), ...(budgetLimitsUnchanged ? [] : ["budget_limits"])], + [...Object.keys(dirtyFields), ...(budgetLimitsUnchanged ? [] : ["budget_limits"])], ); } finally { setIsKeySaving(false); From 42b4cfc5bfa2031d48bafe11cf321f4bdd98eaa5 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 4 Nov 2026 18:07:31 -0800 Subject: [PATCH 04/39] fix(nvidia_riva): send the api key exactly once by calling riva over grpc directly (#50140) * fix(nvidia_riva): send the api key once per call by talking to Riva over grpc directly Drops nvidia-riva-client from the stt-nvidia-riva extra. Depending on its version it sent the bearer token zero times on plaintext or twice on TLS, so token-checking Riva deployments rejected every request (#50059). The provider now declares the few Riva messages it uses, opens its own grpc channel and puts authorization and function-id on each call exactly once * test(nvidia_riva): run the provider against a real grpc server built from nvidia's own definitions * fix(nvidia_riva): validate endpointing settings with pydantic before building the riva message * test(nvidia_riva): check the key is sent once over tls, and drop a comment that restated the code --- .../audio_transcription/handler.py | 328 +---- .../audio_transcription/riva_asr_schema.py | 199 +++ .../audio_transcription/riva_grpc.py | 113 ++ pyproject.toml | 7 +- tests/integration/_support/riva.py | 16 +- .../providers/test_nvidia_riva_wire.py | 9 + .../audio_transcription/test_handler.py | 1253 ++++++++--------- .../test_riva_asr_schema.py | 41 + uv.lock | 10 +- 9 files changed, 985 insertions(+), 991 deletions(-) create mode 100644 litellm/llms/nvidia_riva/audio_transcription/riva_asr_schema.py create mode 100644 litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py create mode 100644 tests/unit/llms/nvidia_riva/audio_transcription/test_riva_asr_schema.py diff --git a/litellm/llms/nvidia_riva/audio_transcription/handler.py b/litellm/llms/nvidia_riva/audio_transcription/handler.py index cafeece0a4e..d2ef33b89dd 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/handler.py +++ b/litellm/llms/nvidia_riva/audio_transcription/handler.py @@ -13,24 +13,26 @@ The handler is intentionally a thin orchestration layer: 2. Build ``RecognitionConfig`` / ``StreamingRecognitionConfig`` protobufs from the structured dict produced by :class:`NvidiaRivaAudioTranscriptionConfig`. -3. Construct ``riva.client.Auth`` honoring NVCF (function-id metadata + TLS) - vs self-hosted (any host:port, optional TLS) modes. -4. Stream the audio through Riva's ``streaming_response_generator`` and - aggregate ``is_final`` results into a single transcript, or, when - ``riva_offline`` is set, send the whole buffer through the unary - ``offline_recognize`` RPC for NIMs that only ship offline engines. +3. Open a gRPC channel honoring NVCF (TLS by default) vs self-hosted (any + host:port, optional TLS) modes, and attach the ``authorization`` and NVCF + ``function-id`` metadata to each call exactly once. +4. Stream the audio through Riva's ``StreamingRecognize`` RPC and aggregate + ``is_final`` results into a single transcript, or, when ``riva_offline`` + is set, send the whole buffer through the unary ``Recognize`` RPC for NIMs + that only ship offline engines. 5. Return a normalized ``TranscriptionResponse`` with ``duration`` exposed on ``_hidden_params`` so cost calculation works. -``riva-client`` is imported lazily so ``litellm`` core remains usable -without the optional STT extras installed. +``grpcio`` and ``protobuf`` are imported lazily so ``litellm`` core remains +usable without the optional STT extras installed. """ +from __future__ import annotations + import asyncio -import inspect -from collections.abc import Callable, Iterable, Mapping, Sequence -from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Protocol +from collections.abc import Iterable, Sequence +from contextlib import closing +from typing import TYPE_CHECKING, Final import httpx @@ -64,89 +66,19 @@ if TYPE_CHECKING: Logging as LiteLLMLoggingObj, ) -# Stream audio to Riva in ~50 ms slices (1600 samples at 16 kHz). Matches -# NVIDIA's recommended chunk size for streaming ASR — small enough for -# responsive endpointing, large enough to keep per-RPC overhead low. -_DEFAULT_CHUNK_SAMPLES: Final = 1600 -_DEFAULT_CHUNK_BYTES: Final = _DEFAULT_CHUNK_SAMPLES * 2 # int16 = 2 bytes/sample - - _RIVA_INSTALL_HINT = "NVIDIA Riva client is not installed. Install with `pip install 'litellm[stt-nvidia-riva]'`." -class _RivaAuth(Protocol): - """Opaque ``riva.client.Auth`` handle.""" - - -class _RecognizeFuture(Protocol): - def result(self, timeout: float | None = None) -> object: ... - - def done(self) -> bool: ... - - def cancel(self) -> bool: ... - - -class _AsrService(Protocol): - @property - def streaming_response_generator(self) -> Callable[..., Iterable[object]]: ... - - def offline_recognize(self, audio_bytes: bytes, config: object, future: bool) -> _RecognizeFuture: ... - - -class _EndpointingConfig(Protocol): - """Opaque ``EndpointingConfig`` protobuf message.""" - - -class _EndpointingConfigField(Protocol): - CopyFrom: Callable[[_EndpointingConfig], None] - - -class _RecognitionConfig(Protocol): - @property - def endpointing_config(self) -> _EndpointingConfigField: ... - - -class _StreamingRecognitionConfig(Protocol): - """Opaque ``StreamingRecognitionConfig`` protobuf message.""" - - -class _AudioEncoding(Protocol): - @property - def LINEAR_PCM(self) -> object: ... - - -class _RivaClientModule(Protocol): - """The ``riva.client`` entry points this handler calls.""" - - @property - def Auth(self) -> Callable[..., _RivaAuth]: ... - - @property - def ASRService(self) -> Callable[[_RivaAuth], _AsrService]: ... - - -class _RivaAsrModule(Protocol): - """The protobuf constructors this handler calls, from whichever module exposes them.""" - - @property - def AudioEncoding(self) -> _AudioEncoding: ... - - @property - def RecognitionConfig(self) -> Callable[..., _RecognitionConfig]: ... - - @property - def StreamingRecognitionConfig(self) -> Callable[..., _StreamingRecognitionConfig]: ... - - @property - def EndpointingConfig(self) -> Callable[..., _EndpointingConfig]: ... - - -def _auth_factory(riva_module: _RivaClientModule) -> Callable[..., _RivaAuth]: - return riva_module.Auth - - -def _audio_encoding(riva_asr_module: _RivaAsrModule) -> _AudioEncoding: - return riva_asr_module.AudioEncoding +def _use_ssl(optional_params: dict) -> bool: + """ + NVCF endpoints are TLS-only, so ``nvcf_function_id`` defaults ``use_ssl`` + to True. Self-hosted Riva defaults to plaintext but honors an explicit + ``use_ssl``, since TLS-terminating ingresses in front of it are common. + """ + use_ssl_override: Final = optional_params.get("use_ssl") + if use_ssl_override is not None: + return bool(use_ssl_override) + return bool(optional_params.get("nvcf_function_id")) def _parse_riva_offline(value: object, drop_params: bool) -> bool: @@ -192,7 +124,7 @@ class NvidiaRivaAudioTranscription: litellm_params: dict, model_response: TranscriptionResponse, timeout: float | str | httpx.Timeout | None, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, api_key: str | None, api_base: str | None, atranscription: bool = False, @@ -237,13 +169,11 @@ class NvidiaRivaAudioTranscription: litellm_params: dict, model_response: TranscriptionResponse, timeout: float | str | httpx.Timeout | None, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, api_key: str | None, api_base: str | None, provider_config: NvidiaRivaAudioTranscriptionConfig | None = None, ) -> TranscriptionResponse: - # ``riva-client`` exposes a sync streaming generator, so we offload - # the blocking call to a worker thread to keep the event loop free. return await asyncio.to_thread( self._run_sync, model=model, @@ -267,7 +197,7 @@ class NvidiaRivaAudioTranscription: litellm_params: dict, model_response: TranscriptionResponse, timeout: float | str | httpx.Timeout | None, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, api_key: str | None, api_base: str | None, provider_config: NvidiaRivaAudioTranscriptionConfig, @@ -317,18 +247,12 @@ class NvidiaRivaAudioTranscription: ) or litellm.drop_params is True riva_offline: Final = _parse_riva_offline(optional_params.get("riva_offline"), drop_params=drop_params) - riva_module, riva_asr_module = _import_riva() - auth_obj: Final = self._construct_auth( - riva_module=riva_module, - api_base=api_base, - api_key=api_key, - optional_params=optional_params, - ) + try: + from litellm.llms.nvidia_riva.audio_transcription import riva_grpc + except ImportError as e: + raise NvidiaRivaException(status_code=500, message=_RIVA_INSTALL_HINT) from e - recognition_config: Final = self._build_recognition_config_proto( - riva_asr_module=riva_asr_module, - recognition_config_dict=recognition_config_dict, - ) + recognition_config: Final = riva_grpc.recognition_config(recognition_config_dict) logging_obj.pre_call( input=None, api_key=api_key, @@ -344,26 +268,18 @@ class NvidiaRivaAudioTranscription: }, ) + metadata: Final = riva_grpc.call_metadata(api_key, optional_params.get("nvcf_function_id")) try: - asr_service: Final[_AsrService] = riva_module.ASRService(auth_obj) - responses: Final = ( - self._offline_responses( - asr_service=asr_service, - pcm_bytes=resampled.pcm_bytes, - recognition_config=recognition_config, - timeout=None if timeout is None else _deadline_seconds(timeout, drop_params), + deadline: Final = None if timeout is None else _deadline_seconds(timeout, drop_params) + with closing(riva_grpc.open_channel(api_base, _use_ssl(optional_params))) as channel: + responses: Final = ( + riva_grpc.recognize(channel, recognition_config, resampled.pcm_bytes, metadata, deadline) + if riva_offline + else riva_grpc.streaming_recognize( + channel, recognition_config, resampled.pcm_bytes, metadata, deadline + ) ) - if riva_offline - else self._streaming_responses( - asr_service=asr_service, - riva_asr_module=riva_asr_module, - pcm_bytes=resampled.pcm_bytes, - recognition_config=recognition_config, - timeout=timeout, - drop_params=drop_params, - ) - ) - final_results: Final = self._collect_final_results(responses, require_final=not riva_offline) + final_results: Final = self._collect_final_results(responses, require_final=not riva_offline) except NvidiaRivaException: raise except Exception as e: @@ -400,142 +316,6 @@ class NvidiaRivaAudioTranscription: return final_response - def _construct_auth( - self, - riva_module: _RivaClientModule, - api_base: str, - api_key: str | None, - optional_params: dict, - ) -> _RivaAuth: - """ - Build a ``riva.client.Auth`` object. - - - When ``nvcf_function_id`` is provided we attach the NVCF - ``function-id`` and bearer ``authorization`` metadata, and default - ``use_ssl`` to True (NVCF endpoints are TLS-only). - - Otherwise (self-hosted) we default ``use_ssl`` to False but still - honor an explicit override — self-hosted Riva behind an ingress - with TLS termination is a real deployment topology. - """ - nvcf_function_id: Final = optional_params.get("nvcf_function_id") - use_ssl_override: Final = optional_params.get("use_ssl") - use_ssl: Final = bool(use_ssl_override) if use_ssl_override is not None else bool(nvcf_function_id) - - metadata: Final[list[tuple[str, str]]] = [] - if nvcf_function_id: - metadata.append(("function-id", str(nvcf_function_id))) - if api_key: - metadata.append(("authorization", f"Bearer {api_key}")) - - try: - return _auth_factory(riva_module)(uri=api_base, use_ssl=use_ssl, metadata_args=metadata) - except TypeError: - # Older riva-client signatures used positional-only args. - return _auth_factory(riva_module)(None, use_ssl, api_base, metadata) - - def _build_recognition_config_proto( - self, riva_asr_module: _RivaAsrModule, recognition_config_dict: dict[str, Any] - ) -> _RecognitionConfig: - encoding_name: Final = (recognition_config_dict.get("encoding") or "LINEAR_PCM").upper() - encoding_enum: Final[object] = getattr( - _audio_encoding(riva_asr_module), - encoding_name, - _audio_encoding(riva_asr_module).LINEAR_PCM, - ) - - config: Final[_RecognitionConfig] = riva_asr_module.RecognitionConfig( - encoding=encoding_enum, - sample_rate_hertz=int(recognition_config_dict["sample_rate_hertz"]), - language_code=recognition_config_dict["language_code"], - audio_channel_count=int(recognition_config_dict["audio_channel_count"]), - enable_automatic_punctuation=bool(recognition_config_dict.get("enable_automatic_punctuation", True)), - enable_word_time_offsets=bool(recognition_config_dict.get("enable_word_time_offsets", False)), - max_alternatives=int(recognition_config_dict.get("max_alternatives", 1)), - model=recognition_config_dict.get("model", "") or "", - verbatim_transcripts=bool(recognition_config_dict.get("verbatim_transcripts", False)), - profanity_filter=bool(recognition_config_dict.get("profanity_filter", False)), - ) - - endpointing: Final = recognition_config_dict.get("endpointing_config") - if isinstance(endpointing, dict) and endpointing: - try: - ep: Final[_EndpointingConfig] = riva_asr_module.EndpointingConfig(**endpointing) - config.endpointing_config.CopyFrom(ep) - except Exception: - # If the user supplied an unknown EndpointingConfig field - # (older Riva server), fall back to Riva's defaults rather - # than failing the whole request. - pass - - return config - - def _streaming_responses( - self, - asr_service: _AsrService, - riva_asr_module: _RivaAsrModule, - pcm_bytes: bytes, - recognition_config: _RecognitionConfig, - timeout: float | str | httpx.Timeout | None, - drop_params: bool, - ) -> Iterable[object]: - streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig( - config=recognition_config, interim_results=False - ) - return asr_service.streaming_response_generator( - audio_chunks=self._iter_audio_chunks(pcm_bytes), - streaming_config=streaming_config, - **self._timeout_kwargs(asr_service.streaming_response_generator, timeout, drop_params), - ) - - def _offline_responses( - self, - asr_service: _AsrService, - pcm_bytes: bytes, - recognition_config: _RecognitionConfig, - timeout: float | None, - ) -> Iterable[object]: - pending: Final = asr_service.offline_recognize(audio_bytes=pcm_bytes, config=recognition_config, future=True) - try: - return (pending.result(timeout=timeout),) - except Exception as e: - if pending.done(): - raise - pending.cancel() - raise NvidiaRivaException( - status_code=408, - message=f"NVIDIA Riva offline recognition did not finish within {timeout} seconds", - ) from e - - @classmethod - def _timeout_kwargs( - cls, rpc: Callable[..., object], timeout: float | str | httpx.Timeout | None, drop_params: bool - ) -> Mapping[str, float]: - # Forward the deadline so the call cannot block forever if the - # server stalls. Older riva-client versions do not accept a - # ``timeout`` kwarg, so pass it only when supported. - if timeout is None or not cls._supports_timeout_kwarg(rpc): - return MappingProxyType({}) - return MappingProxyType({"timeout": _deadline_seconds(timeout, drop_params)}) - - @staticmethod - def _supports_timeout_kwarg(callable_obj: Callable[..., object]) -> bool: - try: - sig: Final = inspect.signature(callable_obj) - except (TypeError, ValueError): - return False - params: Final = sig.parameters - if "timeout" in params: - return True - return any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) - - @staticmethod - def _iter_audio_chunks(pcm_bytes: bytes): - for offset in range(0, len(pcm_bytes), _DEFAULT_CHUNK_BYTES): - chunk = pcm_bytes[offset : offset + _DEFAULT_CHUNK_BYTES] - if not chunk: - continue - yield chunk - @staticmethod def _collect_final_results(stream: Iterable[object], require_final: bool = True) -> Sequence[RivaFinalResult]: """ @@ -570,27 +350,3 @@ class NvidiaRivaAudioTranscription: } ) return final_results - - -def _import_riva() -> tuple[_RivaClientModule, _RivaAsrModule]: - """ - Lazy import of ``riva.client`` and ``riva.client.proto.riva_asr_pb2``. - - We try the SDK first (preferred) and fall back to importing the proto - module separately when the SDK packaging changes between versions. - """ - try: - import riva.client as riva_client - except ImportError as e: - raise NvidiaRivaException(status_code=500, message=_RIVA_INSTALL_HINT) from e - - riva_asr_module = riva_client - if not hasattr(riva_asr_module, "RecognitionConfig"): - try: - from riva.client.proto import riva_asr_pb2 - - riva_asr_module = riva_asr_pb2 - except ImportError as e: - raise NvidiaRivaException(status_code=500, message=_RIVA_INSTALL_HINT) from e - - return riva_client, riva_asr_module diff --git a/litellm/llms/nvidia_riva/audio_transcription/riva_asr_schema.py b/litellm/llms/nvidia_riva/audio_transcription/riva_asr_schema.py new file mode 100644 index 00000000000..13af5820f07 --- /dev/null +++ b/litellm/llms/nvidia_riva/audio_transcription/riva_asr_schema.py @@ -0,0 +1,199 @@ +""" +The part of NVIDIA Riva's speech recognition gRPC schema that LiteLLM sends and reads. + +Names and numbers match ``riva/proto/riva_asr.proto`` and ``riva/proto/riva_audio.proto`` in +https://github.com/nvidia-riva/common. Fields LiteLLM never sets or reads are left out, and the +wire format skips them. The messages live in their own descriptor pool so they never collide +with nvidia-riva-client's copies when both are imported in one process. +""" + +from dataclasses import dataclass +from enum import Enum, IntEnum +from typing import Final + +from google.protobuf import descriptor_pb2, descriptor_pool, message_factory +from google.protobuf.message import Message + +_PACKAGE: Final = "nvidia.riva.asr" +_SERVICE: Final = f"/{_PACKAGE}.RivaSpeechRecognition" +RECOGNIZE_METHOD: Final = f"{_SERVICE}/Recognize" +STREAMING_RECOGNIZE_METHOD: Final = f"{_SERVICE}/StreamingRecognize" + + +class _Kind(Enum): + FLOAT = "TYPE_FLOAT" + INT32 = "TYPE_INT32" + BOOL = "TYPE_BOOL" + STRING = "TYPE_STRING" + MESSAGE = "TYPE_MESSAGE" + BYTES = "TYPE_BYTES" + ENUM = "TYPE_ENUM" + + +class AudioEncoding(IntEnum): + ENCODING_UNSPECIFIED = 0 + LINEAR_PCM = 1 + + +LINEAR_PCM: Final = AudioEncoding.LINEAR_PCM.value + + +@dataclass(frozen=True, slots=True) +class _Field: + name: str + number: int + kind: _Kind + type_name: str = "" + repeated: bool = False + optional: bool = False + oneof: str = "" + + +_MESSAGES: Final[tuple[tuple[str, tuple[_Field, ...]], ...]] = ( + ( + "RecognitionConfig", + ( + _Field("encoding", 1, _Kind.ENUM, "AudioEncoding"), + _Field("sample_rate_hertz", 2, _Kind.INT32), + _Field("language_code", 3, _Kind.STRING), + _Field("max_alternatives", 4, _Kind.INT32), + _Field("profanity_filter", 5, _Kind.BOOL), + _Field("audio_channel_count", 7, _Kind.INT32), + _Field("enable_word_time_offsets", 8, _Kind.BOOL), + _Field("enable_automatic_punctuation", 11, _Kind.BOOL), + _Field("model", 13, _Kind.STRING), + _Field("verbatim_transcripts", 14, _Kind.BOOL), + _Field("endpointing_config", 25, _Kind.MESSAGE, "EndpointingConfig", optional=True), + ), + ), + ( + "EndpointingConfig", + ( + _Field("start_history", 1, _Kind.INT32, optional=True), + _Field("start_threshold", 2, _Kind.FLOAT, optional=True), + _Field("stop_history", 3, _Kind.INT32, optional=True), + _Field("stop_threshold", 4, _Kind.FLOAT, optional=True), + _Field("stop_history_eou", 5, _Kind.INT32, optional=True), + _Field("stop_threshold_eou", 6, _Kind.FLOAT, optional=True), + ), + ), + ( + "StreamingRecognitionConfig", + ( + _Field("config", 1, _Kind.MESSAGE, "RecognitionConfig"), + _Field("interim_results", 2, _Kind.BOOL), + ), + ), + ( + "StreamingRecognizeRequest", + ( + _Field("streaming_config", 1, _Kind.MESSAGE, "StreamingRecognitionConfig", oneof="streaming_request"), + _Field("audio_content", 2, _Kind.BYTES, oneof="streaming_request"), + ), + ), + ( + "WordInfo", + ( + _Field("start_time", 1, _Kind.INT32), + _Field("end_time", 2, _Kind.INT32), + _Field("word", 3, _Kind.STRING), + ), + ), + ( + "SpeechRecognitionAlternative", + ( + _Field("transcript", 1, _Kind.STRING), + _Field("words", 3, _Kind.MESSAGE, "WordInfo", repeated=True), + ), + ), + ( + "StreamingRecognitionResult", + ( + _Field("alternatives", 1, _Kind.MESSAGE, "SpeechRecognitionAlternative", repeated=True), + _Field("is_final", 2, _Kind.BOOL), + ), + ), + ( + "StreamingRecognizeResponse", + (_Field("results", 1, _Kind.MESSAGE, "StreamingRecognitionResult", repeated=True),), + ), + ( + "RecognizeRequest", + ( + _Field("config", 1, _Kind.MESSAGE, "RecognitionConfig"), + _Field("audio", 2, _Kind.BYTES), + ), + ), + ( + "SpeechRecognitionResult", + (_Field("alternatives", 1, _Kind.MESSAGE, "SpeechRecognitionAlternative", repeated=True),), + ), + ( + "RecognizeResponse", + (_Field("results", 1, _Kind.MESSAGE, "SpeechRecognitionResult", repeated=True),), + ), +) + + +def _oneof_names(fields: tuple[_Field, ...]) -> tuple[str, ...]: + declared: Final = tuple(dict.fromkeys(field.oneof for field in fields if field.oneof)) + synthetic: Final = tuple(f"_{field.name}" for field in fields if field.optional) + return declared + synthetic + + +def _field_proto(field: _Field, oneofs: tuple[str, ...]) -> descriptor_pb2.FieldDescriptorProto: + oneof: Final = f"_{field.name}" if field.optional else field.oneof + return descriptor_pb2.FieldDescriptorProto( + name=field.name, + number=field.number, + type=descriptor_pb2.FieldDescriptorProto.Type.Value(field.kind.value), + label=( + descriptor_pb2.FieldDescriptorProto.LABEL_REPEATED + if field.repeated + else descriptor_pb2.FieldDescriptorProto.LABEL_OPTIONAL + ), + type_name=f".{_PACKAGE}.{field.type_name}" if field.type_name else None, + oneof_index=oneofs.index(oneof) if oneof else None, + proto3_optional=field.optional or None, + ) + + +def _message_proto(name: str, fields: tuple[_Field, ...]) -> descriptor_pb2.DescriptorProto: + oneofs: Final = _oneof_names(fields) + return descriptor_pb2.DescriptorProto( + name=name, + field=[_field_proto(field, oneofs) for field in fields], + oneof_decl=[descriptor_pb2.OneofDescriptorProto(name=oneof) for oneof in oneofs], + ) + + +_FILE: Final = descriptor_pb2.FileDescriptorProto( + name="litellm/nvidia_riva/riva_asr.proto", + package=_PACKAGE, + syntax="proto3", + enum_type=[ + descriptor_pb2.EnumDescriptorProto( + name="AudioEncoding", + value=[ + descriptor_pb2.EnumValueDescriptorProto(name=encoding.name, number=encoding.value) + for encoding in AudioEncoding + ], + ) + ], + message_type=[_message_proto(name, fields) for name, fields in _MESSAGES], +) +_POOL: Final = descriptor_pool.DescriptorPool() +_POOL.AddSerializedFile(_FILE.SerializeToString()) + + +def _message_class(name: str) -> type[Message]: + return message_factory.GetMessageClass(_POOL.FindMessageTypeByName(f"{_PACKAGE}.{name}")) + + +RecognitionConfig: Final = _message_class("RecognitionConfig") +EndpointingConfig: Final = _message_class("EndpointingConfig") +StreamingRecognitionConfig: Final = _message_class("StreamingRecognitionConfig") +StreamingRecognizeRequest: Final = _message_class("StreamingRecognizeRequest") +StreamingRecognizeResponse: Final = _message_class("StreamingRecognizeResponse") +RecognizeRequest: Final = _message_class("RecognizeRequest") +RecognizeResponse: Final = _message_class("RecognizeResponse") diff --git a/litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py b/litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py new file mode 100644 index 00000000000..8367853673f --- /dev/null +++ b/litellm/llms/nvidia_riva/audio_transcription/riva_grpc.py @@ -0,0 +1,113 @@ +""" +gRPC calls to NVIDIA Riva's speech recognition service. + +The ``authorization`` and NVCF ``function-id`` metadata ride on each call, never on the channel, so every +request carries them exactly once over both plaintext and TLS. +""" + +from collections.abc import Iterator, Mapping, Sequence +from typing import Any, Final, TypeAlias + +import grpc +from google.protobuf.message import Message +from pydantic import TypeAdapter + +from litellm.llms.nvidia_riva.audio_transcription import riva_asr_schema as schema +from litellm.llms.nvidia_riva.common_utils import NvidiaRivaException + +# Stream audio to Riva in ~50 ms slices (1600 samples at 16 kHz). Matches +# NVIDIA's recommended chunk size for streaming ASR, small enough for +# responsive endpointing, large enough to keep per-RPC overhead low. +_CHUNK_SAMPLES: Final = 1600 +_CHUNK_BYTES: Final = _CHUNK_SAMPLES * 2 + +CallMetadata: TypeAlias = tuple[tuple[str, str], ...] + +_ENDPOINTING_FIELDS: Final = TypeAdapter(Mapping[str, object]) + + +def call_metadata(api_key: str | None, nvcf_function_id: object) -> CallMetadata: + function_id: Final = (("function-id", str(nvcf_function_id)),) if nvcf_function_id else () + authorization: Final = (("authorization", f"Bearer {api_key}"),) if api_key else () + return function_id + authorization + + +def open_channel(api_base: str, use_ssl: bool) -> grpc.Channel: + if use_ssl: + return grpc.secure_channel(api_base, grpc.ssl_channel_credentials()) + return grpc.insecure_channel(api_base) + + +def recognition_config(config: Mapping[str, Any]) -> Message: + endpointing: Final = _endpointing_config(config.get("endpointing_config")) + return schema.RecognitionConfig( + encoding=schema.LINEAR_PCM, + sample_rate_hertz=int(config["sample_rate_hertz"]), + language_code=config["language_code"], + audio_channel_count=int(config["audio_channel_count"]), + enable_automatic_punctuation=bool(config.get("enable_automatic_punctuation", True)), + enable_word_time_offsets=bool(config.get("enable_word_time_offsets", False)), + max_alternatives=int(config.get("max_alternatives", 1)), + model=config.get("model", "") or "", + verbatim_transcripts=bool(config.get("verbatim_transcripts", False)), + profanity_filter=bool(config.get("profanity_filter", False)), + **({} if endpointing is None else {"endpointing_config": endpointing}), + ) + + +def _endpointing_config(endpointing: object) -> Message | None: + try: + fields: Final = _ENDPOINTING_FIELDS.validate_python(endpointing) + return schema.EndpointingConfig(**fields) if fields else None + except (ValueError, TypeError): + # A missing, unknown or mistyped field falls back to Riva's default + # endpointing rather than failing the whole request. + return None + + +def _streaming_requests(config: Message, pcm_bytes: bytes) -> Iterator[Message]: + yield schema.StreamingRecognizeRequest( + streaming_config=schema.StreamingRecognitionConfig(config=config, interim_results=False) + ) + for offset in range(0, len(pcm_bytes), _CHUNK_BYTES): + yield schema.StreamingRecognizeRequest(audio_content=pcm_bytes[offset : offset + _CHUNK_BYTES]) + + +def streaming_recognize( + channel: grpc.Channel, + config: Message, + pcm_bytes: bytes, + metadata: CallMetadata, + timeout: float | None, +) -> Iterator[Message]: + rpc: Final = channel.stream_stream( + schema.STREAMING_RECOGNIZE_METHOD, + request_serializer=schema.StreamingRecognizeRequest.SerializeToString, + response_deserializer=schema.StreamingRecognizeResponse.FromString, + ) + return rpc(_streaming_requests(config, pcm_bytes), metadata=metadata, timeout=timeout) + + +def recognize( + channel: grpc.Channel, + config: Message, + pcm_bytes: bytes, + metadata: CallMetadata, + timeout: float | None, +) -> Sequence[Message]: + rpc: Final = channel.unary_unary( + schema.RECOGNIZE_METHOD, + request_serializer=schema.RecognizeRequest.SerializeToString, + response_deserializer=schema.RecognizeResponse.FromString, + ) + pending: Final = rpc.future(schema.RecognizeRequest(config=config, audio=pcm_bytes), metadata=metadata) + try: + return (pending.result(timeout=timeout),) + except Exception as e: + if pending.done(): + raise + pending.cancel() + raise NvidiaRivaException( + status_code=408, + message=f"NVIDIA Riva offline recognition did not finish within {timeout} seconds", + ) from e diff --git a/pyproject.toml b/pyproject.toml index f2774e87b7b..2cc2d0f70a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -155,7 +155,8 @@ stt-vertex-chirp = [ stt-nvidia-riva = [ # NVIDIA Riva STT provider (gRPC). These are imported lazily inside the # provider handler so litellm core remains usable without them. - "nvidia-riva-client>=2.15.0", + "grpcio>=1.67.1,<2.0", + "protobuf>=5.26.1,<7.0", "soundfile>=0.12.1", "audioread>=3.0.1", "numpy>=1.26.0", @@ -241,6 +242,7 @@ dev = [ "pytest-timeout==2.4.0", "vcrpy==8.2.1", "pytest-recording==0.13.4", + "nvidia-riva-client==2.16.0", ] e2e-dev = [ "playwright==1.61.0", @@ -275,6 +277,9 @@ ci = [ # protobuf, Pillow is a compiled C extension). "tenacity==8.5.0", "google-generativeai==0.8.6", + # The Riva tests serve NVIDIA's own message definitions, which checks + # LiteLLM's hand-written copy of them on every run. + "nvidia-riva-client==2.16.0", "Pillow==12.3.0; python_version == '3.10'", "Pillow==13.0.0; python_version >= '3.11'", # Azure batch E2E tests still import psycopg2 directly. diff --git a/tests/integration/_support/riva.py b/tests/integration/_support/riva.py index 8552bc2197d..438d877386b 100644 --- a/tests/integration/_support/riva.py +++ b/tests/integration/_support/riva.py @@ -39,6 +39,7 @@ class RivaCall: rpc: Rpc audio: bytes cancelled: bool + authorization: tuple[str, ...] = () @dataclass(frozen=True, slots=True) @@ -68,6 +69,10 @@ def pcm_bytes(seconds: float) -> int: return int(seconds * SAMPLE_RATE_HZ) * 2 +def _authorization(context: grpc.ServicerContext) -> tuple[str, ...]: + return tuple(str(value) for key, value in context.invocation_metadata() if key == "authorization") + + def carries_clip(audio: bytes, seconds: float) -> bool: sent: Final = _clip_samples(seconds) if len(audio) != len(sent) * 2: @@ -100,9 +105,9 @@ class _Asr(riva_asr_pb2_grpc.RivaSpeechRecognitionServicer): released: Final = threading.Event() context.add_callback(released.set) released.wait(timeout=_STALL_LIMIT_SECONDS) - self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=not context.is_active())) + self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, not context.is_active(), _authorization(context))) return riva_asr_pb2.RecognizeResponse() - self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=False)) + self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, False, _authorization(context))) return riva_asr_pb2.RecognizeResponse( results=[riva_asr_pb2.SpeechRecognitionResult(alternatives=[self._alternative(request.config)])] ) @@ -111,13 +116,12 @@ class _Asr(riva_asr_pb2_grpc.RivaSpeechRecognitionServicer): self, request_iterator: Iterable[riva_asr_pb2.StreamingRecognizeRequest], context: grpc.ServicerContext ) -> Iterator[riva_asr_pb2.StreamingRecognizeResponse]: if self._offline_only: - self._calls.put(RivaCall(Rpc.STREAMING, b"", cancelled=False)) + self._calls.put(RivaCall(Rpc.STREAMING, b"", False, _authorization(context))) context.abort(grpc.StatusCode.INVALID_ARGUMENT, OFFLINE_ONLY_REFUSAL) requests: Final = tuple(request_iterator) config: Final = requests[0].streaming_config.config if requests else riva_asr_pb2.RecognitionConfig() - self._calls.put( - RivaCall(Rpc.STREAMING, b"".join(request.audio_content for request in requests), cancelled=False) - ) + audio: Final = b"".join(request.audio_content for request in requests) + self._calls.put(RivaCall(Rpc.STREAMING, audio, False, _authorization(context))) yield riva_asr_pb2.StreamingRecognizeResponse( results=[riva_asr_pb2.StreamingRecognitionResult(alternatives=[self._alternative(config)], is_final=True)] ) diff --git a/tests/integration/providers/test_nvidia_riva_wire.py b/tests/integration/providers/test_nvidia_riva_wire.py index ef36c449761..56d54a87600 100644 --- a/tests/integration/providers/test_nvidia_riva_wire.py +++ b/tests/integration/providers/test_nvidia_riva_wire.py @@ -135,3 +135,12 @@ def test_a_stalled_offline_recognize_answers_408_at_the_deployment_timeout(gatew assert timed_out.status_code == 408, timed_out.text calls: Final = eventually(server.drain, lambda observed: len(observed) > 0, seconds=30) assert all(call.rpc is Rpc.RECOGNIZE and call.cancelled for call in calls), calls + + +@pytest.mark.parametrize("riva_offline", (False, True), ids=("streaming", "offline")) +def test_the_deployment_api_key_reaches_riva_as_exactly_one_bearer_token(gateway: Gateway, riva_offline: bool) -> None: + with riva_server(_WORDS) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, api_key="riva-deployment-key", riva_offline=riva_offline) + served: Final = _transcribe(gateway, scenario.key(), model) + assert served.status_code == 200, served.text + assert [call.authorization for call in server.drain()] == [("Bearer riva-deployment-key",)] diff --git a/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py b/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py index dbb5cde6060..40ee4d61ac4 100644 --- a/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py +++ b/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py @@ -1,780 +1,641 @@ -""" -End-to-end-ish tests for NvidiaRivaAudioTranscription. - -We mock ``riva.client`` so the test does not need the real gRPC SDK or a -running Riva server. The mock also lets us assert how Auth metadata is -constructed (NVCF vs self-hosted) and how the streaming generator output -is aggregated. -""" - import asyncio import io -from types import SimpleNamespace -from unittest.mock import MagicMock +import ipaddress +import threading +import uuid +import wave +from array import array +from collections.abc import Iterable, Iterator +from concurrent import futures +from contextlib import contextmanager +from dataclasses import dataclass, field +from datetime import datetime, timezone +from queue import SimpleQueue +from typing import Final +import grpc import httpx -import numpy as np import pytest -import soundfile as sf +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ec +from cryptography.x509.oid import NameOID +from riva.client.proto import riva_asr_pb2, riva_asr_pb2_grpc, riva_audio_pb2 import litellm -from litellm.llms.nvidia_riva.audio_transcription import handler as handler_mod -from litellm.llms.nvidia_riva.audio_transcription.handler import ( - NvidiaRivaAudioTranscription, -) +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.nvidia_riva.audio_transcription.handler import NvidiaRivaAudioTranscription from litellm.llms.nvidia_riva.common_utils import NvidiaRivaException from litellm.types.utils import TranscriptionResponse +_SAMPLE_RATE_HZ: Final = 16000 +_SERVER_WAIT_SECONDS: Final = 30 +_STALL_LIMIT_SECONDS: Final = 60 +_NO_DEADLINE_SECONDS: Final = 10**9 +_STALL_TIMEOUT_SECONDS: Final = 0.5 +_API_KEY: Final = "nvapi-test" -def _make_wav_bytes(seconds: float = 1.0, sample_rate: int = 16000) -> bytes: - n = int(sample_rate * seconds) - samples = (0.05 * np.sin(np.linspace(0, 2 * np.pi * 220 * seconds, n))).astype( - np.float32 + +@dataclass(frozen=True, slots=True) +class _Script: + streaming: tuple[riva_asr_pb2.StreamingRecognizeResponse, ...] = () + offline: riva_asr_pb2.RecognizeResponse = field(default_factory=riva_asr_pb2.RecognizeResponse) + abort: tuple[grpc.StatusCode, str] | None = None + stall: bool = False + + +@dataclass(frozen=True, slots=True) +class _Call: + rpc: str + metadata: tuple[tuple[str, str], ...] + config: riva_asr_pb2.RecognitionConfig + audio: bytes + deadline_seconds: float | None + cancelled: bool + + def header(self, key: str) -> tuple[str, ...]: + return tuple(value for name, value in self.metadata if name == key) + + +@dataclass(frozen=True, slots=True) +class _Riva: + address: str + calls: SimpleQueue[_Call] + + def only_call(self) -> _Call: + call: Final = self.calls.get(timeout=_SERVER_WAIT_SECONDS) + assert self.calls.empty(), "Riva received more than one call" + return call + + +class _Asr(riva_asr_pb2_grpc.RivaSpeechRecognitionServicer): + def __init__(self, script: _Script, calls: SimpleQueue[_Call]) -> None: + self._script: Final = script + self._calls: Final = calls + + def _settle( + self, context: grpc.ServicerContext, rpc: str, config: riva_asr_pb2.RecognitionConfig, audio: bytes + ) -> None: + remaining: Final = context.time_remaining() + released: Final = threading.Event() + if self._script.stall and context.add_callback(released.set): + released.wait(timeout=_STALL_LIMIT_SECONDS) + self._calls.put( + _Call( + rpc=rpc, + metadata=tuple((item.key, item.value) for item in context.invocation_metadata()), + config=config, + audio=audio, + deadline_seconds=None if remaining is None or remaining > _NO_DEADLINE_SECONDS else remaining, + cancelled=not context.is_active(), + ) + ) + if self._script.abort is not None: + context.abort(*self._script.abort) + + def Recognize( + self, request: riva_asr_pb2.RecognizeRequest, context: grpc.ServicerContext + ) -> riva_asr_pb2.RecognizeResponse: + self._settle(context, "Recognize", request.config, request.audio) + return self._script.offline + + def StreamingRecognize( + self, request_iterator: Iterable[riva_asr_pb2.StreamingRecognizeRequest], context: grpc.ServicerContext + ) -> Iterator[riva_asr_pb2.StreamingRecognizeResponse]: + first, *rest = tuple(request_iterator) + self._settle( + context, "StreamingRecognize", first.streaming_config.config, b"".join(r.audio_content for r in rest) + ) + yield from self._script.streaming + + +@dataclass(frozen=True, slots=True) +class _Tls: + key_pem: bytes + cert_pem: bytes + + +def _self_signed_loopback_tls() -> _Tls: + key: Final = ec.generate_private_key(ec.SECP256R1()) + name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")]) + cert: Final = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(datetime(2000, 1, 1, tzinfo=timezone.utc)) + .not_valid_after(datetime(2099, 1, 1, tzinfo=timezone.utc)) + .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True) + .add_extension( + x509.SubjectAlternativeName([x509.DNSName("localhost"), x509.IPAddress(ipaddress.ip_address("127.0.0.1"))]), + critical=False, + ) + .sign(key, hashes.SHA256()) ) - buf = io.BytesIO() - sf.write(buf, samples, sample_rate, format="WAV", subtype="PCM_16") - return buf.getvalue() - - -def _fake_word(word: str, start_ms: int, end_ms: int): - return SimpleNamespace(word=word, start_time=start_ms, end_time=end_ms) - - -def _fake_alternative(transcript: str, words=None): - return SimpleNamespace(transcript=transcript, words=words or []) - - -def _fake_result(is_final: bool, alternatives): - return SimpleNamespace(is_final=is_final, alternatives=alternatives) - - -def _fake_response(results): - return SimpleNamespace(results=results) - - -@pytest.fixture -def mock_riva(monkeypatch): - """ - Stand-ins for the bits of ``riva.client`` the handler touches: - - ``Auth`` (constructor) - - ``ASRService`` with ``streaming_response_generator`` - - ``RecognitionConfig``, ``StreamingRecognitionConfig``, ``EndpointingConfig`` - - ``AudioEncoding`` namespace with ``LINEAR_PCM`` - """ - auth_calls = {} - - class FakeAuth: - def __init__(self, *args, **kwargs): - # Support both keyword and positional Auth constructors. - if kwargs: - auth_calls["uri"] = kwargs.get("uri") - auth_calls["use_ssl"] = kwargs.get("use_ssl") - auth_calls["metadata_args"] = kwargs.get("metadata_args") - else: - # positional: (None, use_ssl, uri, metadata) - auth_calls["use_ssl"] = args[1] if len(args) > 1 else None - auth_calls["uri"] = args[2] if len(args) > 2 else None - auth_calls["metadata_args"] = args[3] if len(args) > 3 else None - - class FakeRecognitionConfig: - def __init__(self, **kwargs): - self._kwargs = kwargs - self.endpointing_config = SimpleNamespace(CopyFrom=lambda _: None) - - class FakeStreamingRecognitionConfig: - def __init__(self, config, interim_results): - self.config = config - self.interim_results = interim_results - - class FakeEndpointingConfig: - def __init__(self, **kwargs): - self._kwargs = kwargs - - class FakeAudioEncoding: - LINEAR_PCM = "LINEAR_PCM" - - streaming_responses_holder = {"value": []} - offline_response_holder = {"value": None} - offline_error_holder = {"value": None} - offline_future_holder = {"value": None} - streaming_calls = [] - offline_calls = {} - - class FakeRecognizeFuture: - def result(self, timeout: float | None = None) -> object: - offline_calls["timeout"] = timeout - if offline_error_holder["value"] is not None: - raise offline_error_holder["value"] - return offline_response_holder["value"] - - def done(self) -> bool: - return True - - def cancel(self) -> bool: - return False - - class FakeASRService: - def __init__(self, auth): - self.auth = auth - - def streaming_response_generator(self, audio_chunks, streaming_config): - # Drain audio_chunks generator so we exercise the chunking path. - streaming_calls.append(list(audio_chunks)) - yield from streaming_responses_holder["value"] - - def offline_recognize(self, audio_bytes: bytes, config: object, future: bool = False) -> object: - offline_calls["audio_bytes"] = audio_bytes - offline_calls["config"] = config - pending = offline_future_holder["value"] or FakeRecognizeFuture() - return pending if future else pending.result() - - fake_riva_client = SimpleNamespace( - Auth=FakeAuth, - ASRService=FakeASRService, - RecognitionConfig=FakeRecognitionConfig, - StreamingRecognitionConfig=FakeStreamingRecognitionConfig, - EndpointingConfig=FakeEndpointingConfig, - AudioEncoding=FakeAudioEncoding, - ) - - def fake_import_riva(): - return fake_riva_client, fake_riva_client - - monkeypatch.setattr(handler_mod, "_import_riva", fake_import_riva) - - return SimpleNamespace( - auth_calls=auth_calls, - responses=streaming_responses_holder, - offline_response=offline_response_holder, - offline_error=offline_error_holder, - offline_future=offline_future_holder, - streaming_calls=streaming_calls, - offline_calls=offline_calls, - client=fake_riva_client, + return _Tls( + key_pem=key.private_bytes( + serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption() + ), + cert_pem=cert.public_bytes(serialization.Encoding.PEM), ) -@pytest.fixture -def logging_obj(): - return MagicMock() +@contextmanager +def _riva_server(script: _Script, max_receive_bytes: int | None = None, tls: _Tls | None = None) -> Iterator[_Riva]: + calls: Final[SimpleQueue[_Call]] = SimpleQueue() + options: Final = () if max_receive_bytes is None else (("grpc.max_receive_message_length", max_receive_bytes),) + server: Final = grpc.server(futures.ThreadPoolExecutor(max_workers=4), options=options) + riva_asr_pb2_grpc.add_RivaSpeechRecognitionServicer_to_server(_Asr(script, calls), server) + port: Final = ( + server.add_insecure_port("127.0.0.1:0") + if tls is None + else server.add_secure_port("127.0.0.1:0", grpc.ssl_server_credentials(((tls.key_pem, tls.cert_pem),))) + ) + server.start() + try: + yield _Riva(f"127.0.0.1:{port}", calls) + finally: + server.stop(grace=None).wait() -def test_sync_path_aggregates_only_final_results(mock_riva, logging_obj): - mock_riva.responses["value"] = [ - # Empty heartbeat chunk: ignore. - _fake_response(results=[]), - # Interim chunk (not final): ignore. - _fake_response( - results=[ - _fake_result( - is_final=False, alternatives=[_fake_alternative("partial...")] - ) - ] - ), - # Two final chunks aggregated. - _fake_response( - results=[ - _fake_result( - is_final=True, - alternatives=[ - _fake_alternative( - "Hello,", - words=[_fake_word("Hello,", 0, 320)], - ) - ], - ) - ] - ), - _fake_response( - results=[ - _fake_result( - is_final=True, - alternatives=[ - _fake_alternative( - " world.", - words=[_fake_word("world.", 480, 870)], - ) - ], - ) - ] - ), - ] +def _samples(seconds: float) -> array[int]: + return array("h", (((index * 37) % 2001 - 1000) * 16 for index in range(int(seconds * _SAMPLE_RATE_HZ)))) - impl = NvidiaRivaAudioTranscription() - response: TranscriptionResponse = impl.audio_transcriptions( - model="nvidia/parakeet-ctc-1_1b-asr", - audio_file=_make_wav_bytes(), - optional_params={ - "language_code": "en-US", - "enable_word_time_offsets": True, - "response_format": "verbose_json", - "timestamp_granularities": ["word"], - }, - litellm_params={}, + +def _wav(seconds: float = 1.0) -> bytes: + buffer: Final = io.BytesIO() + with wave.open(buffer, "wb") as writer: + writer.setnchannels(1) + writer.setsampwidth(2) + writer.setframerate(_SAMPLE_RATE_HZ) + writer.writeframes(_samples(seconds).tobytes()) + return buffer.getvalue() + + +def _carries_clip(audio: bytes, seconds: float) -> bool: + sent: Final = _samples(seconds) + if len(audio) != len(sent) * 2: + return False + return all(abs(left - right) <= 1 for left, right in zip(sent, array("h", audio), strict=True)) + + +def _alternative( + transcript: str, words: tuple[tuple[str, int, int], ...] = () +) -> riva_asr_pb2.SpeechRecognitionAlternative: + return riva_asr_pb2.SpeechRecognitionAlternative( + transcript=transcript, + words=[riva_asr_pb2.WordInfo(word=word, start_time=start, end_time=end) for word, start, end in words], + ) + + +def _streamed(transcript: str, is_final: bool = True) -> riva_asr_pb2.StreamingRecognizeResponse: + return riva_asr_pb2.StreamingRecognizeResponse( + results=[riva_asr_pb2.StreamingRecognitionResult(alternatives=[_alternative(transcript)], is_final=is_final)] + ) + + +def _offline(*transcripts: str) -> riva_asr_pb2.RecognizeResponse: + return riva_asr_pb2.RecognizeResponse( + results=[riva_asr_pb2.SpeechRecognitionResult(alternatives=[_alternative(text)]) for text in transcripts] + ) + + +def _answers(transcript: str) -> _Script: + return _Script(streaming=(_streamed(transcript),), offline=_offline(transcript)) + + +def _logging() -> Logging: + return Logging( + model="m", + messages=[], + stream=False, + call_type="transcription", + start_time=datetime(2026, 1, 1), + litellm_call_id=str(uuid.uuid4()), + function_id="riva-test", + ) + + +def _transcribe( + riva: _Riva, + optional_params: dict[str, object] | None = None, + *, + api_key: str | None = None, + timeout: float | str | httpx.Timeout | None = 60, + litellm_params: dict[str, object] | None = None, +) -> TranscriptionResponse: + return NvidiaRivaAudioTranscription().audio_transcriptions( + model="m", + audio_file=_wav(), + optional_params={"language_code": "en-US", **(optional_params or {})}, + litellm_params=litellm_params or {}, model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key="nvapi-xxx", - api_base="grpc.nvcf.nvidia.com:443", - ) - - assert response.text == "Hello, world." - # duration is propagated from the resampler. - assert response._hidden_params["audio_transcription_duration"] == pytest.approx( - 1.0, abs=0.05 - ) - # word timestamps converted from ms to seconds. - words = response["words"] - assert words[0]["start"] == pytest.approx(0.0) - assert words[1]["end"] == pytest.approx(0.87) - assert ( - logging_obj.pre_call.call_args.kwargs["additional_args"]["atranscription"] - is False + timeout=timeout, + logging_obj=_logging(), + api_key=api_key, + api_base=riva.address, ) -def test_riva_offline_uses_unary_recognize_and_keeps_results_without_is_final( - mock_riva: SimpleNamespace, logging_obj: MagicMock -) -> None: - """ - Offline-only NIMs (Parakeet TDT, Whisper) reject StreamingRecognize. With - ``riva_offline`` the handler must call ``offline_recognize`` with the whole - resampled buffer and a bare ``RecognitionConfig``, and must keep every - result even though ``RecognizeResponse`` results carry no ``is_final``. - """ - mock_riva.offline_response["value"] = _fake_response( +def test_streaming_aggregates_only_final_results_and_their_word_timings() -> None: + heartbeat: Final = riva_asr_pb2.StreamingRecognizeResponse() + hello: Final = riva_asr_pb2.StreamingRecognizeResponse( results=[ - SimpleNamespace(alternatives=[_fake_alternative("Hello,", words=[_fake_word("Hello,", 0, 320)])]), - SimpleNamespace(alternatives=[_fake_alternative(" world.", words=[_fake_word("world.", 480, 870)])]), + riva_asr_pb2.StreamingRecognitionResult( + alternatives=[_alternative("Hello,", (("Hello,", 0, 320),))], is_final=True + ) ] ) - - impl = NvidiaRivaAudioTranscription() - response: TranscriptionResponse = impl.audio_transcriptions( - model="nvidia/parakeet-0.6b-tdt", - audio_file=_make_wav_bytes(), - optional_params={ - "language_code": "multi", - "riva_offline": True, - "enable_word_time_offsets": True, - "response_format": "verbose_json", - "timestamp_granularities": ["word"], - }, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", + world: Final = riva_asr_pb2.StreamingRecognizeResponse( + results=[ + riva_asr_pb2.StreamingRecognitionResult( + alternatives=[_alternative(" world.", (("world.", 480, 870),))], is_final=True + ) + ] ) + with _riva_server(_Script(streaming=(heartbeat, _streamed("partial...", is_final=False), hello, world))) as riva: + response: Final = _transcribe( + riva, + {"enable_word_time_offsets": True, "response_format": "verbose_json", "timestamp_granularities": ["word"]}, + ) + call: Final = riva.only_call() assert response.text == "Hello, world." - assert response["words"][1]["end"] == pytest.approx(0.87) - assert response._hidden_params["audio_transcription_duration"] == pytest.approx(1.0, abs=0.05) - assert mock_riva.streaming_calls == [] - assert len(mock_riva.offline_calls["audio_bytes"]) == pytest.approx(32000, abs=1600) - assert mock_riva.offline_calls["config"]._kwargs["language_code"] == "multi" - assert logging_obj.pre_call.call_args.kwargs["additional_args"]["complete_input_dict"]["riva_offline"] is True - - -def test_default_mode_streams_and_never_calls_offline_recognize( - mock_riva: SimpleNamespace, logging_obj: MagicMock -) -> None: - mock_riva.responses["value"] = [ - _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) + assert response["words"] == [ + {"word": "Hello,", "start": 0.0, "end": 0.32}, + {"word": "world.", "start": 0.48, "end": 0.87}, ] + assert response._hidden_params["audio_transcription_duration"] == pytest.approx(1.0) + assert call.rpc == "StreamingRecognize" + assert _carries_clip(call.audio, 1.0), "streamed chunks do not add up to the resampled clip" - impl = NvidiaRivaAudioTranscription() - response = impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US"}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", + +def test_riva_offline_uses_recognize_and_keeps_results_without_is_final() -> None: + offline: Final = riva_asr_pb2.RecognizeResponse( + results=[ + riva_asr_pb2.SpeechRecognitionResult(alternatives=[_alternative("Hello,", (("Hello,", 0, 320),))]), + riva_asr_pb2.SpeechRecognitionResult(alternatives=[_alternative(" world.", (("world.", 480, 870),))]), + ] ) - assert response.text == "streamed" - assert mock_riva.offline_calls == {} - assert len(mock_riva.streaming_calls) == 1 - assert logging_obj.pre_call.call_args.kwargs["additional_args"]["complete_input_dict"]["riva_offline"] is False - - - -@pytest.mark.parametrize(("riva_offline", "expect_offline"), [("false", False), ("False", False), ("true", True)]) -def test_riva_offline_string_from_multipart_form_is_parsed( - mock_riva: SimpleNamespace, logging_obj: MagicMock, riva_offline: str, expect_offline: bool -) -> None: - mock_riva.responses["value"] = [ - _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) - ] - mock_riva.offline_response["value"] = _fake_response( - results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] - ) - - response = NvidiaRivaAudioTranscription().audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US", "riva_offline": riva_offline}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", - ) - - assert response.text == ("offline" if expect_offline else "streamed") - assert (mock_riva.offline_calls != {}) is expect_offline - - -@pytest.mark.parametrize("riva_offline", ["yes", "", 1]) -def test_riva_offline_malformed_value_is_rejected_before_any_rpc( - mock_riva: SimpleNamespace, logging_obj: MagicMock, riva_offline: object -) -> None: - with pytest.raises(NvidiaRivaException) as excinfo: - NvidiaRivaAudioTranscription().audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US", "riva_offline": riva_offline}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", + with _riva_server(_Script(offline=offline)) as riva: + response: Final = _transcribe( + riva, + { + "language_code": "multi", + "riva_offline": True, + "enable_word_time_offsets": True, + "response_format": "verbose_json", + "timestamp_granularities": ["word"], + }, ) + call: Final = riva.only_call() - assert excinfo.value.status_code == 400 - assert "riva_offline" in excinfo.value.message - assert mock_riva.offline_calls == {} - assert mock_riva.streaming_calls == [] + assert response.text == "Hello, world." + assert response["words"][1] == {"word": "world.", "start": 0.48, "end": 0.87} + assert response._hidden_params["audio_transcription_duration"] == pytest.approx(1.0) + assert call.rpc == "Recognize" + assert call.config.language_code == "multi" + assert _carries_clip(call.audio, 1.0), "offline request does not carry the whole resampled clip" -def test_riva_offline_malformed_value_is_dropped_under_drop_params( - mock_riva: SimpleNamespace, logging_obj: MagicMock -) -> None: - mock_riva.responses["value"] = [ - _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) - ] +@pytest.mark.parametrize("riva_offline", [False, True]) +def test_recognition_config_reaches_riva_as_nvidia_decodes_it(riva_offline: bool) -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe( + riva, + { + "riva_offline": riva_offline, + "enable_word_time_offsets": True, + "riva_model_name": "parakeet", + "sample_rate_hertz": 8000, + "profanity_filter": True, + }, + ) + call: Final = riva.only_call() - response = NvidiaRivaAudioTranscription().audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US", "riva_offline": "yes"}, - litellm_params={"drop_params": True}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", + assert call.config == riva_asr_pb2.RecognitionConfig( + encoding=riva_audio_pb2.AudioEncoding.LINEAR_PCM, + sample_rate_hertz=_SAMPLE_RATE_HZ, + language_code="en-US", + max_alternatives=1, + profanity_filter=True, + audio_channel_count=1, + enable_word_time_offsets=True, + enable_automatic_punctuation=True, + model="parakeet", ) + +def test_default_mode_streams_and_never_calls_recognize() -> None: + with _riva_server(_answers("streamed")) as riva: + response: Final = _transcribe(riva) + call: Final = riva.only_call() + assert response.text == "streamed" - assert mock_riva.offline_calls == {} - - -def test_deployment_riva_offline_reaches_the_handler_through_transcription(mock_riva: SimpleNamespace) -> None: - mock_riva.offline_response["value"] = _fake_response( - results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] - ) - - response = litellm.transcription( - model="nvidia_riva/nvidia/parakeet-0.6b-tdt", - file=("clip.wav", _make_wav_bytes(), "audio/wav"), - api_base="localhost:50051", - riva_offline=True, - ) - - assert response.text == "offline" - assert mock_riva.streaming_calls == [] - assert mock_riva.offline_calls != {} + assert call.rpc == "StreamingRecognize" @pytest.mark.parametrize( - ("timeout", "waited_for"), - [("7.5", 7.5), (httpx.Timeout(30.0, read=7.5), 7.5), (httpx.Timeout(5.0, read=None), 600.0), (None, None)], + ("riva_offline", "rpc"), [("false", "StreamingRecognize"), ("False", "StreamingRecognize"), ("true", "Recognize")] ) -def test_offline_recognition_waits_the_timeout_in_seconds_whatever_its_form( - mock_riva: SimpleNamespace, timeout: str | httpx.Timeout | None, waited_for: float | None -) -> None: - mock_riva.offline_response["value"] = _fake_response( - results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] - ) +def test_riva_offline_string_from_multipart_form_is_parsed(riva_offline: str, rpc: str) -> None: + with _riva_server(_Script(streaming=(_streamed("streamed"),), offline=_offline("offline"))) as riva: + response: Final = _transcribe(riva, {"riva_offline": riva_offline}) + call: Final = riva.only_call() - response = litellm.transcription( - model="nvidia_riva/nvidia/parakeet-0.6b-tdt", - file=("clip.wav", _make_wav_bytes(), "audio/wav"), - api_base="localhost:50051", - riva_offline=True, - timeout=timeout, - ) + assert call.rpc == rpc + assert response.text == ("offline" if rpc == "Recognize" else "streamed") + + +@pytest.mark.parametrize("riva_offline", ["yes", "", 1]) +def test_riva_offline_malformed_value_is_rejected_before_any_rpc(riva_offline: object) -> None: + with _riva_server(_answers("ok")) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, {"riva_offline": riva_offline}) + + assert riva.calls.empty() + + assert excinfo.value.status_code == 400 + assert "riva_offline" in excinfo.value.message + + +def test_riva_offline_malformed_value_is_dropped_under_drop_params() -> None: + with _riva_server(_Script(streaming=(_streamed("streamed"),), offline=_offline("offline"))) as riva: + response: Final = _transcribe(riva, {"riva_offline": "yes"}, litellm_params={"drop_params": True}) + call: Final = riva.only_call() + + assert response.text == "streamed" + assert call.rpc == "StreamingRecognize" + + +def test_deployment_riva_offline_reaches_the_handler_through_transcription() -> None: + with _riva_server(_Script(offline=_offline("offline"))) as riva: + response: Final = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, + riva_offline=True, + ) + call: Final = riva.only_call() assert response.text == "offline" - assert mock_riva.offline_calls["timeout"] == waited_for + assert call.rpc == "Recognize" -@pytest.mark.parametrize("timeout", ["abc", "30s"]) -def test_malformed_timeout_on_offline_call_is_a_400_before_any_rpc( - mock_riva: SimpleNamespace, timeout: str +@pytest.mark.parametrize( + ("timeout", "deadline_seconds"), + [ + (12.5, 12.5), + ("12.5", 12.5), + (httpx.Timeout(60.0, read=12.5), 12.5), + (httpx.Timeout(5.0, read=None), 600.0), + (None, None), + ], +) +def test_streaming_deadline_is_the_timeout_in_seconds_whatever_its_form( + timeout: float | str | httpx.Timeout | None, deadline_seconds: float | None ) -> None: - with pytest.raises(NvidiaRivaException) as excinfo: - litellm.transcription( - model="nvidia_riva/nvidia/parakeet-0.6b-tdt", - file=("clip.wav", _make_wav_bytes(), "audio/wav"), - api_base="localhost:50051", - riva_offline=True, + with _riva_server(_answers("streamed")) as riva: + response: Final = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-ctc-1.1b", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, timeout=timeout, ) + call: Final = riva.only_call() + + assert response.text == "streamed" + assert call.deadline_seconds == (None if deadline_seconds is None else pytest.approx(deadline_seconds, rel=0.25)) + + +@pytest.mark.parametrize("riva_offline", [False, True]) +@pytest.mark.parametrize("timeout", ["abc", "30s"]) +def test_malformed_timeout_is_a_400_before_any_rpc(riva_offline: bool, timeout: str) -> None: + with _riva_server(_answers("ok")) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, + riva_offline=riva_offline, + timeout=timeout, + ) + + assert riva.calls.empty() assert excinfo.value.status_code == 400 assert "timeout" in excinfo.value.message - assert "audio_bytes" not in mock_riva.offline_calls -def test_malformed_timeout_on_offline_call_waits_the_default_under_drop_params( - mock_riva: SimpleNamespace, -) -> None: - mock_riva.offline_response["value"] = _fake_response( - results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] - ) - - response = litellm.transcription( - model="nvidia_riva/nvidia/parakeet-0.6b-tdt", - file=("clip.wav", _make_wav_bytes(), "audio/wav"), - api_base="localhost:50051", - riva_offline=True, - timeout="abc", - drop_params=True, - ) - - assert response.text == "offline" - assert mock_riva.offline_calls["timeout"] == 600.0 - - -def test_malformed_timeout_still_streams_when_the_client_takes_no_timeout(mock_riva: SimpleNamespace) -> None: - mock_riva.responses["value"] = [ - _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) - ] - - response = litellm.transcription( - model="nvidia_riva/nvidia/parakeet-ctc-1.1b", - file=("clip.wav", _make_wav_bytes(), "audio/wav"), - api_base="localhost:50051", - timeout="abc", - ) - - assert response.text == "streamed" - assert len(mock_riva.streaming_calls) == 1 - - -def test_oversized_offline_audio_is_a_non_retryable_413(mock_riva: SimpleNamespace, logging_obj: MagicMock) -> None: - class FakeGrpcError(Exception): - def code(self) -> SimpleNamespace: - return SimpleNamespace(name="RESOURCE_EXHAUSTED") - - def details(self) -> str: - return "SERVER: Received message larger than max (4348825 vs. 4194304)" - - mock_riva.offline_error["value"] = FakeGrpcError("rpc fail") - - with pytest.raises(NvidiaRivaException) as excinfo: - NvidiaRivaAudioTranscription().audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US", "riva_offline": True}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", +def test_malformed_streaming_timeout_falls_back_to_the_default_deadline_under_drop_params() -> None: + with _riva_server(_answers("ok")) as riva: + response: Final = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-ctc-1.1b", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, + timeout="abc", + drop_params=True, ) + call: Final = riva.only_call() - assert excinfo.value.status_code == 413 + assert response.text == "ok" + assert call.rpc == "StreamingRecognize" + assert call.deadline_seconds == pytest.approx(600.0, rel=0.25) + + +def test_malformed_offline_timeout_is_ignored_under_drop_params() -> None: + with _riva_server(_answers("ok")) as riva: + response: Final = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _wav(), "audio/wav"), + api_base=riva.address, + riva_offline=True, + timeout="abc", + drop_params=True, + ) + call: Final = riva.only_call() + + assert response.text == "ok" + assert call.rpc == "Recognize" + + +@pytest.mark.parametrize( + "timeout", [_STALL_TIMEOUT_SECONDS, str(_STALL_TIMEOUT_SECONDS), httpx.Timeout(30.0, read=_STALL_TIMEOUT_SECONDS)] +) +def test_stalled_offline_recognition_is_cancelled_and_times_out_as_408(timeout: float | str | httpx.Timeout) -> None: + with _riva_server(_Script(stall=True)) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, {"riva_offline": True}, timeout=timeout) + call: Final = riva.only_call() + + assert excinfo.value.status_code == 408 + assert f"within {_STALL_TIMEOUT_SECONDS} seconds" in excinfo.value.message + assert call.cancelled is True + + +def test_stalled_streaming_recognition_hits_its_deadline_as_504() -> None: + with _riva_server(_Script(stall=True)) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, timeout=_STALL_TIMEOUT_SECONDS) + call: Final = riva.only_call() + + assert excinfo.value.status_code == 504, excinfo.value.message + assert "DEADLINE_EXCEEDED" in excinfo.value.message + assert call.cancelled is True + + +def test_oversized_offline_audio_is_a_non_retryable_413() -> None: + with _riva_server(_answers("ok"), max_receive_bytes=4096) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, {"riva_offline": True}) + + assert riva.calls.empty() + + assert excinfo.value.status_code == 413, excinfo.value.message assert litellm.utils.should_retry(excinfo.value.status_code) is False -def test_stalled_offline_recognition_is_cancelled_and_times_out_as_408( - mock_riva: SimpleNamespace, logging_obj: MagicMock +@pytest.mark.parametrize( + ("riva_offline", "code", "status_code"), + [ + (False, grpc.StatusCode.UNAUTHENTICATED, 401), + (True, grpc.StatusCode.INVALID_ARGUMENT, 400), + (False, grpc.StatusCode.RESOURCE_EXHAUSTED, 429), + ], +) +def test_grpc_error_is_wrapped_as_nvidia_riva_exception( + riva_offline: bool, code: grpc.StatusCode, status_code: int ) -> None: - class FutureTimeoutError(Exception): - pass + with _riva_server(_Script(abort=(code, "Riva said no"))) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, {"riva_offline": riva_offline}, api_key=_API_KEY) - class StalledRecognizeFuture: - def __init__(self) -> None: - self.waited_for: float | None = None - self.cancelled = False + assert excinfo.value.status_code == status_code + assert code.name in excinfo.value.message + assert "Riva said no" in excinfo.value.message - def result(self, timeout: float | None = None) -> object: - self.waited_for = timeout - raise FutureTimeoutError() - def done(self) -> bool: - return self.cancelled +@pytest.mark.parametrize("riva_offline", [False, True]) +def test_self_hosted_api_key_is_sent_as_one_bearer_token(riva_offline: bool) -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe(riva, {"riva_offline": riva_offline}, api_key=_API_KEY) + call: Final = riva.only_call() - def cancel(self) -> bool: - self.cancelled = True - return True + assert call.header("authorization") == (f"Bearer {_API_KEY}",) + assert call.header("function-id") == () - stalled = StalledRecognizeFuture() - mock_riva.offline_future["value"] = stalled +@pytest.mark.parametrize("riva_offline", [False, True]) +def test_nvcf_function_id_and_api_key_are_each_sent_once(riva_offline: bool) -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe( + riva, {"riva_offline": riva_offline, "nvcf_function_id": "abc-123", "use_ssl": False}, api_key=_API_KEY + ) + call: Final = riva.only_call() + + assert call.header("function-id") == ("abc-123",) + assert call.header("authorization") == (f"Bearer {_API_KEY}",) + + +@pytest.mark.parametrize("riva_offline", [False, True]) +def test_no_api_key_sends_no_authorization(riva_offline: bool) -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe(riva, {"riva_offline": riva_offline}) + call: Final = riva.only_call() + + assert call.header("authorization") == () + assert call.header("function-id") == () + + +@pytest.mark.parametrize("riva_offline", [False, True]) +@pytest.mark.parametrize( + ("optional_params", "function_ids"), + [({"use_ssl": True}, ()), ({"nvcf_function_id": "abc-123"}, ("abc-123",))], + ids=["use_ssl", "nvcf_default_tls"], +) +def test_tls_api_key_is_sent_as_one_bearer_token( + monkeypatch: pytest.MonkeyPatch, + optional_params: dict[str, object], + function_ids: tuple[str, ...], + riva_offline: bool, +) -> None: + tls: Final = _self_signed_loopback_tls() + system_roots: Final = grpc.ssl_channel_credentials + monkeypatch.setattr(grpc, "ssl_channel_credentials", lambda: system_roots(root_certificates=tls.cert_pem)) + with _riva_server(_answers("ok"), tls=tls) as riva: + response: Final = _transcribe(riva, {**optional_params, "riva_offline": riva_offline}, api_key=_API_KEY) + call: Final = riva.only_call() + + assert response.text == "ok" + assert call.header("authorization") == (f"Bearer {_API_KEY}",) + assert call.header("function-id") == function_ids + + +@pytest.mark.parametrize("optional_params", [{"use_ssl": True}, {"nvcf_function_id": "abc-123"}]) +def test_tls_never_falls_back_to_plaintext(optional_params: dict[str, object]) -> None: + with _riva_server(_answers("ok")) as riva: + with pytest.raises(NvidiaRivaException) as excinfo: + _transcribe(riva, optional_params, api_key=_API_KEY, timeout=_SERVER_WAIT_SECONDS) + + assert riva.calls.empty() + + assert excinfo.value.status_code == 503, excinfo.value.message + + +def test_explicit_endpointing_config_keeps_zero_thresholds() -> None: + with _riva_server(_answers("ok")) as riva: + _transcribe(riva, {"endpointing_config": {"stop_threshold": 0.0, "stop_history": 800}}) + call: Final = riva.only_call() + + assert call.config.HasField("endpointing_config") + assert call.config.endpointing_config.HasField("stop_threshold") + assert call.config.endpointing_config == riva_asr_pb2.EndpointingConfig(stop_threshold=0.0, stop_history=800) + + +def test_unknown_endpointing_field_falls_back_to_riva_defaults() -> None: + with _riva_server(_answers("ok")) as riva: + response: Final = _transcribe(riva, {"endpointing_config": {"stop_threshold": 0.5, "not_a_field": 1}}) + call: Final = riva.only_call() + + assert response.text == "ok" + assert not call.config.HasField("endpointing_config") + + +def test_missing_api_base_raises_clear_error() -> None: with pytest.raises(NvidiaRivaException) as excinfo: NvidiaRivaAudioTranscription().audio_transcriptions( model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US", "riva_offline": True}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=7.5, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", - ) - - assert excinfo.value.status_code == 408 - assert stalled.waited_for == 7.5 - assert stalled.cancelled is True - -def test_riva_offline_grpc_error_is_wrapped(mock_riva: SimpleNamespace, logging_obj: MagicMock) -> None: - class FakeGrpcError(Exception): - def code(self) -> SimpleNamespace: - return SimpleNamespace(name="INVALID_ARGUMENT") - - def details(self) -> str: - return "Unavailable model requested" - - mock_riva.offline_error["value"] = FakeGrpcError("rpc fail") - - impl = NvidiaRivaAudioTranscription() - with pytest.raises(NvidiaRivaException) as excinfo: - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "multi", "riva_offline": True}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", - ) - assert excinfo.value.status_code == 400 - assert "INVALID_ARGUMENT" in excinfo.value.message - - -def test_auth_nvcf_defaults_use_ssl_and_attaches_function_id(mock_riva, logging_obj): - mock_riva.responses["value"] = [ - _fake_response( - results=[ - _fake_result( - is_final=True, - alternatives=[_fake_alternative("ok")], - ) - ] - ) - ] - impl = NvidiaRivaAudioTranscription() - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={ - "nvcf_function_id": "abc-123", - "language_code": "en-US", - }, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key="nvapi-xxx", - api_base="grpc.nvcf.nvidia.com:443", - ) - - assert mock_riva.auth_calls["uri"] == "grpc.nvcf.nvidia.com:443" - assert mock_riva.auth_calls["use_ssl"] is True - metadata = dict(mock_riva.auth_calls["metadata_args"]) - assert metadata["function-id"] == "abc-123" - assert metadata["authorization"] == "Bearer nvapi-xxx" - - -def test_auth_self_hosted_defaults_no_ssl_and_no_function_id(mock_riva, logging_obj): - mock_riva.responses["value"] = [ - _fake_response( - results=[ - _fake_result(is_final=True, alternatives=[_fake_alternative("ok")]) - ] - ) - ] - impl = NvidiaRivaAudioTranscription() - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US"}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", - ) - - assert mock_riva.auth_calls["uri"] == "localhost:50051" - assert mock_riva.auth_calls["use_ssl"] is False - metadata = dict(mock_riva.auth_calls["metadata_args"]) - # No function-id, no authorization metadata. - assert "function-id" not in metadata - assert "authorization" not in metadata - - -def test_explicit_use_ssl_override_wins(mock_riva, logging_obj): - """ - Self-hosted Riva behind an ingress with TLS termination is a real - deployment topology. ``use_ssl=True`` must be honored even without an - NVCF function id. - """ - mock_riva.responses["value"] = [ - _fake_response( - results=[ - _fake_result(is_final=True, alternatives=[_fake_alternative("ok")]) - ] - ) - ] - impl = NvidiaRivaAudioTranscription() - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"use_ssl": True, "language_code": "en-US"}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="riva.internal.company.com:443", - ) - - assert mock_riva.auth_calls["use_ssl"] is True - - -def test_missing_api_base_raises_clear_error(mock_riva, logging_obj): - impl = NvidiaRivaAudioTranscription() - with pytest.raises(NvidiaRivaException) as excinfo: - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), + audio_file=_wav(), optional_params={}, litellm_params={}, model_response=TranscriptionResponse(), timeout=60, - logging_obj=logging_obj, + logging_obj=_logging(), api_key=None, api_base=None, ) + + assert excinfo.value.status_code == 400 assert "api_base" in excinfo.value.message -def test_async_path_uses_to_thread(mock_riva, logging_obj): - mock_riva.responses["value"] = [ - _fake_response( - results=[ - _fake_result( - is_final=True, alternatives=[_fake_alternative("async ok")] - ) - ] - ) - ] - impl = NvidiaRivaAudioTranscription() - response = asyncio.run( - impl.async_audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US"}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", - ) - ) - assert response.text == "async ok" - assert ( - logging_obj.pre_call.call_args.kwargs["additional_args"]["atranscription"] - is True - ) - - -@pytest.mark.parametrize("timeout", [12.5, "12.5", httpx.Timeout(60.0, read=12.5)]) -def test_timeout_is_forwarded_to_streaming_generator_when_supported( - mock_riva: SimpleNamespace, logging_obj: MagicMock, timeout: float | str | httpx.Timeout -) -> None: - """ - Without a deadline the gRPC stream can block forever on a stalled Riva - server. The handler must forward the call-level ``timeout`` to - ``streaming_response_generator`` whenever the installed riva-client - accepts a ``timeout`` kwarg. - """ - captured_kwargs = {} - - def streaming_with_timeout(self, audio_chunks, streaming_config, timeout=None): - captured_kwargs["timeout"] = timeout - list(audio_chunks) - yield from [ - _fake_response( - results=[ - _fake_result(is_final=True, alternatives=[_fake_alternative("ok")]) - ] +def test_async_path_transcribes() -> None: + with _riva_server(_answers("async ok")) as riva: + response: Final = asyncio.run( + NvidiaRivaAudioTranscription().async_audio_transcriptions( + model="m", + audio_file=_wav(), + optional_params={"language_code": "en-US"}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=_logging(), + api_key=_API_KEY, + api_base=riva.address, ) - ] - - mock_riva.client.ASRService.streaming_response_generator = streaming_with_timeout - - impl = NvidiaRivaAudioTranscription() - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US"}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=timeout, - logging_obj=logging_obj, - api_key=None, - api_base="localhost:50051", - ) - assert captured_kwargs["timeout"] == pytest.approx(12.5) - - -def test_grpc_error_is_wrapped_as_nvidia_riva_exception(mock_riva, logging_obj): - class FakeGrpcError(Exception): - def code(self): - return SimpleNamespace(name="UNAUTHENTICATED") - - def details(self): - return "bad token" - - def raising_streaming_response_generator(self, audio_chunks, streaming_config): - list(audio_chunks) - raise FakeGrpcError("rpc fail") - - mock_riva.client.ASRService.streaming_response_generator = ( - raising_streaming_response_generator - ) - - impl = NvidiaRivaAudioTranscription() - with pytest.raises(NvidiaRivaException) as excinfo: - impl.audio_transcriptions( - model="m", - audio_file=_make_wav_bytes(), - optional_params={"language_code": "en-US"}, - litellm_params={}, - model_response=TranscriptionResponse(), - timeout=60, - logging_obj=logging_obj, - api_key="nvapi-xxx", - api_base="grpc.nvcf.nvidia.com:443", ) + call: Final = riva.only_call() - assert excinfo.value.status_code == 401 - assert "UNAUTHENTICATED" in excinfo.value.message + assert response.text == "async ok" + assert call.header("authorization") == (f"Bearer {_API_KEY}",) diff --git a/tests/unit/llms/nvidia_riva/audio_transcription/test_riva_asr_schema.py b/tests/unit/llms/nvidia_riva/audio_transcription/test_riva_asr_schema.py new file mode 100644 index 00000000000..124e3c28291 --- /dev/null +++ b/tests/unit/llms/nvidia_riva/audio_transcription/test_riva_asr_schema.py @@ -0,0 +1,41 @@ +from collections.abc import Iterator +from typing import Final + +from google.protobuf.descriptor import FieldDescriptor +from riva.client.proto import riva_asr_pb2, riva_audio_pb2 + +from litellm.llms.nvidia_riva.audio_transcription import riva_asr_schema + + +def _declared_fields() -> Iterator[tuple[str, FieldDescriptor]]: + for message_name, _ in riva_asr_schema._MESSAGES: + message: Final = riva_asr_schema._POOL.FindMessageTypeByName(f"{riva_asr_schema._PACKAGE}.{message_name}") + yield from ((message_name, field) for field in message.fields) + + +def _wire_shape(field: FieldDescriptor | None) -> tuple[object, ...] | None: + if field is None: + return None + return ( + field.number, + field.type, + field.label, + field.has_presence, + field.message_type.full_name if field.message_type else None, + field.enum_type.name if field.enum_type else None, + field.containing_oneof.name if field.containing_oneof else None, + ) + + +def _nvidia_field(message_name: str, field_name: str) -> FieldDescriptor | None: + return riva_asr_pb2.DESCRIPTOR.message_types_by_name[message_name].fields_by_name.get(field_name) + + +def test_every_declared_field_matches_nvidias_riva_schema() -> None: + declared: Final = tuple(_declared_fields()) + + assert len(declared) == sum(len(fields) for _, fields in riva_asr_schema._MESSAGES) + assert {f"{message}.{field.name}": _wire_shape(field) for message, field in declared} == { + f"{message}.{field.name}": _wire_shape(_nvidia_field(message, field.name)) for message, field in declared + } + assert riva_asr_schema.LINEAR_PCM == riva_audio_pb2.AudioEncoding.LINEAR_PCM diff --git a/uv.lock b/uv.lock index 0c796b677e2..4f36ef20bfb 100644 --- a/uv.lock +++ b/uv.lock @@ -5134,9 +5134,10 @@ semantic-router = [ ] stt-nvidia-riva = [ { name = "audioread" }, + { name = "grpcio" }, { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "nvidia-riva-client" }, + { name = "protobuf" }, { name = "soundfile" }, ] stt-vertex-chirp = [ @@ -5176,6 +5177,7 @@ ci = [ { name = "logfire" }, { name = "lunary", version = "1.4.36", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "lunary", version = "1.4.37", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "nvidia-riva-client" }, { name = "pillow", version = "12.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "pillow", version = "13.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "psutil" }, @@ -5200,6 +5202,7 @@ dev = [ { name = "mypy" }, { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "nvidia-riva-client" }, { name = "openapi-core" }, { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp" }, @@ -5300,6 +5303,7 @@ requires-dist = [ { name = "granian", marker = "python_full_version < '3.15' and extra == 'proxy'", specifier = ">=2.7.4,<3.0" }, { name = "grpcio", marker = "extra == 'grpc'", specifier = "==1.84.1" }, { name = "grpcio", marker = "extra == 'proxy-runtime'", specifier = "==1.84.1" }, + { name = "grpcio", marker = "extra == 'stt-nvidia-riva'", specifier = ">=1.67.1,<2.0" }, { name = "gunicorn", marker = "extra == 'proxy'", specifier = ">=23.0.0,<24.0" }, { name = "hiredis", marker = "python_full_version >= '3.15' and extra == 'proxy'", specifier = ">=3.5.0,<4.0" }, { name = "hiredis", marker = "python_full_version < '3.15' and extra == 'proxy'", specifier = ">=3.0.0,<4.0" }, @@ -5323,7 +5327,6 @@ requires-dist = [ { name = "mlflow", marker = "extra == 'mlflow'", specifier = ">=3.17.2,<4.0" }, { name = "numpy", marker = "extra == 'stt-nvidia-riva'", specifier = ">=1.26.0" }, { name = "numpydoc", marker = "extra == 'utils'", specifier = ">=1.8.0,<2.0" }, - { name = "nvidia-riva-client", marker = "extra == 'stt-nvidia-riva'", specifier = ">=2.15.0" }, { name = "openai", specifier = ">=2.20.0,<4.0.0" }, { name = "opentelemetry-api", marker = "extra == 'proxy-runtime'", specifier = "==1.45.1" }, { name = "opentelemetry-exporter-otlp", marker = "extra == 'proxy-runtime'", specifier = "==1.45.1" }, @@ -5335,6 +5338,7 @@ requires-dist = [ { name = "polars", marker = "extra == 'proxy'", specifier = ">=1.38.1,<2.0" }, { name = "prisma", marker = "extra == 'extra-proxy'", specifier = ">=0.11.0,<1.0" }, { name = "prometheus-client", marker = "extra == 'proxy-runtime'", specifier = ">=0.20.0,<1.0" }, + { name = "protobuf", marker = "extra == 'stt-nvidia-riva'", specifier = ">=5.26.1,<7.0" }, { name = "psycopg", marker = "extra == 'extra-proxy'", specifier = ">=3.2,<4.0" }, { name = "psycopg-binary", marker = "extra == 'extra-proxy'", specifier = ">=3.2,<4.0" }, { name = "pydantic", marker = "python_full_version < '3.14'", specifier = ">=2.11.0,<3.0.0" }, @@ -5406,6 +5410,7 @@ ci = [ { name = "logfire", specifier = "==5.2.0" }, { name = "lunary", marker = "python_full_version == '3.10.*'", specifier = "==1.4.36" }, { name = "lunary", marker = "python_full_version >= '3.11'", specifier = "==1.4.37" }, + { name = "nvidia-riva-client", specifier = "==2.16.0" }, { name = "pillow", marker = "python_full_version == '3.10.*'", specifier = "==12.3.0" }, { name = "pillow", marker = "python_full_version >= '3.11'", specifier = "==13.0.0" }, { name = "psutil", specifier = "==7.2.2" }, @@ -5429,6 +5434,7 @@ dev = [ { name = "langfuse", specifier = ">=4.7,<5.0" }, { name = "mypy", specifier = "==1.20.1" }, { name = "numpy", specifier = ">=1.26.0,<3.0" }, + { name = "nvidia-riva-client", specifier = "==2.16.0" }, { name = "openapi-core", specifier = "==0.22.0" }, { name = "opentelemetry-api", specifier = "==1.45.1" }, { name = "opentelemetry-exporter-otlp", specifier = "==1.45.1" }, From dc741a6e16c7a15e6d467b15a12607c9c79face8 Mon Sep 17 00:00:00 2001 From: jesus-berri Date: Wed, 4 Nov 2026 18:28:35 -0800 Subject: [PATCH 05/39] fix(ui): let team admins edit models on team service account keys (#49970) --- .../components/templates/teamAdminMemberKeyPayload.test.ts | 7 +++++++ .../src/components/templates/teamAdminMemberKeyPayload.ts | 1 + 2 files changed, 8 insertions(+) diff --git a/ui/litellm-dashboard/src/components/templates/teamAdminMemberKeyPayload.test.ts b/ui/litellm-dashboard/src/components/templates/teamAdminMemberKeyPayload.test.ts index a43b41abeb0..3c0db9d2f67 100644 --- a/ui/litellm-dashboard/src/components/templates/teamAdminMemberKeyPayload.test.ts +++ b/ui/litellm-dashboard/src/components/templates/teamAdminMemberKeyPayload.test.ts @@ -22,6 +22,13 @@ describe("isTeamAdminEditingMemberKey", () => { ); }); + it("is false for a team service account key with no user", () => { + expect(isTeamAdminEditingMemberKey({ ...baseArgs, keyUserId: null, teamMembers: members("admin") })).toBe(false); + expect(isTeamAdminEditingMemberKey({ ...baseArgs, keyUserId: undefined, teamMembers: members("admin") })).toBe( + false, + ); + }); + it("is false for a personal key with no team", () => { expect(isTeamAdminEditingMemberKey({ ...baseArgs, keyTeamId: null, teamMembers: members("admin") })).toBe(false); }); diff --git a/ui/litellm-dashboard/src/components/templates/teamAdminMemberKeyPayload.ts b/ui/litellm-dashboard/src/components/templates/teamAdminMemberKeyPayload.ts index abd8653dde5..f26c8c6b550 100644 --- a/ui/litellm-dashboard/src/components/templates/teamAdminMemberKeyPayload.ts +++ b/ui/litellm-dashboard/src/components/templates/teamAdminMemberKeyPayload.ts @@ -12,6 +12,7 @@ export const isTeamAdminEditingMemberKey = (args: { }): boolean => { if (isProxyAdminRole(args.userRole)) return false; if (!args.keyTeamId) return false; + if (!args.keyUserId) return false; if (args.keyUserId === args.userId) return false; return isUserTeamAdminForSingleTeam(args.teamMembers ?? null, args.userId); }; From 429f814a043fc62656d5f02287416383a90e85d6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 4 Nov 2026 18:43:36 -0800 Subject: [PATCH 06/39] fix(cost_map): fill missing Claude capability fields from the family baseline (#49815) * fix(cost_map): fill missing Claude capability fields from the family baseline Mapped Claude entries on Anthropic, Azure Foundry, Bedrock and Vertex now get supports_parallel_function_calling and supports_system_messages from the claude-family-baseline rule when the entry leaves them out. Values set on an entry still win * test(cost_map): keep the exact fill-missing assertion for Claude thinking rules Build the expected baseline part from the loaded map so a rule that leaks an extra key into Claude entries still fails, and wrap lines over 120 characters. --- ...odel_prices_and_context_window_backup.json | 3 +- model_prices_and_context_window.json | 3 +- .../test_fallback_generalizations.py | 64 ++++++++++++++++++- 3 files changed, 67 insertions(+), 3 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7d61e731774..4d576443b30 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -64726,7 +64726,8 @@ { "name": "claude-family-baseline", "pattern": "claude-[a-z]+-\\d+(?:[-.]\\d+)?(?:-\\d{8})?", - "description": "Any Claude family-major id with an optional minor and an optional 8-digit date suffix, under any provider namespace (bare, bedrock-dotted, vertex, databricks, ...), so bare majors like claude-newfamily-5 get the same baseline as claude-newfamily-5-1. Carries the model-family facts every Claude shares; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number.", + "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], + "description": "Any Claude family-major id with an optional minor and an optional 8-digit date suffix, under any provider namespace (bare, bedrock-dotted, vertex, databricks, ...), so bare majors like claude-newfamily-5 get the same baseline as claude-newfamily-5-1. Carries the model-family facts every Claude shares; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number. On mapped Claude entries from Anthropic, Azure Foundry, Bedrock and Vertex it only fills the keys an entry leaves out, so the entry's own values always win.", "model_info": { "mode": "chat", "max_input_tokens": 200000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7d61e731774..4d576443b30 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -64726,7 +64726,8 @@ { "name": "claude-family-baseline", "pattern": "claude-[a-z]+-\\d+(?:[-.]\\d+)?(?:-\\d{8})?", - "description": "Any Claude family-major id with an optional minor and an optional 8-digit date suffix, under any provider namespace (bare, bedrock-dotted, vertex, databricks, ...), so bare majors like claude-newfamily-5 get the same baseline as claude-newfamily-5-1. Carries the model-family facts every Claude shares; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number.", + "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], + "description": "Any Claude family-major id with an optional minor and an optional 8-digit date suffix, under any provider namespace (bare, bedrock-dotted, vertex, databricks, ...), so bare majors like claude-newfamily-5 get the same baseline as claude-newfamily-5-1. Carries the model-family facts every Claude shares; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number. On mapped Claude entries from Anthropic, Azure Foundry, Bedrock and Vertex it only fills the keys an entry leaves out, so the entry's own values always win.", "model_info": { "mode": "chat", "max_input_tokens": 200000, diff --git a/tests/unit/litellm_core_utils/test_fallback_generalizations.py b/tests/unit/litellm_core_utils/test_fallback_generalizations.py index 99fa849879c..327031cfaf9 100644 --- a/tests/unit/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/unit/litellm_core_utils/test_fallback_generalizations.py @@ -9,7 +9,8 @@ resolution (get_model_info) including the shipped rules in the bundled cost map. import json import logging -from collections.abc import Callable +import re +from collections.abc import Callable, Mapping from pathlib import Path from typing import Final @@ -915,6 +916,7 @@ def test_shipped_claude_thinking_rules_backfill_only_anthropic(shipped_cost_map) assert info.get("supports_legacy_thinking") is None assert info.get("max_input_tokens") is None assert match_fill_missing_generalizations("claude-sonnet-4-6", "anthropic") == { + **_claude_baseline()["model_info"], "supports_adaptive_thinking": True, "supports_legacy_thinking": True, "supports_tool_search": True, @@ -961,3 +963,63 @@ def test_shipped_tool_search_rule_fills_mapped_claude_entries_without_flag(shipp assert match_fill_missing_generalizations("claude-opus-5", "bedrock")["supports_tool_search"] is True assert match_fill_missing_generalizations("claude-opus-5", "azure_ai")["supports_tool_search"] is True assert match_fill_missing_generalizations("claude-opus-5", "perplexity") is None + + +CLAUDE_BASELINE_FILL_PROVIDERS: Final = ( + "anthropic", + "azure_ai", + "bedrock", + "bedrock_converse", + "vertex_ai-anthropic_models", +) + + +def _claude_baseline() -> Mapping[str, object]: + return next(rule for rule in get_fallback_generalization_rules() if rule["name"] == "claude-family-baseline") + + +def _mapped_claude_chat_entries() -> tuple[tuple[str, Mapping[str, object]], ...]: + pattern: Final = re.compile(_claude_baseline()["pattern"], re.IGNORECASE) + return tuple( + (key, entry) + for key, entry in litellm.model_cost.items() + if isinstance(entry, dict) and entry.get("mode") == "chat" and pattern.search(key) is not None + ) + + +def _fields_overriding_baseline(entry: Mapping[str, object], baseline: Mapping[str, object]) -> frozenset[str]: + return frozenset(field for field in entry.keys() & baseline.keys() if entry[field] != baseline[field]) + + +@pytest.mark.parametrize("provider", CLAUDE_BASELINE_FILL_PROVIDERS) +def test_shipped_claude_baseline_fills_only_missing_keys_on_mapped_entries(shipped_cost_map, provider): + """A mapped Claude entry from a first-party host gets each baseline key it leaves out, + while every baseline key the entry sets keeps the entry's own value.""" + baseline: Final = _claude_baseline()["model_info"] + entries: Final = tuple( + (key, entry) for key, entry in _mapped_claude_chat_entries() if entry["litellm_provider"] == provider + ) + assert any(baseline.keys() - entry.keys() for _, entry in entries) + assert any(_fields_overriding_baseline(entry, baseline) for _, entry in entries) + + for key, entry in entries: + info = litellm.get_model_info(key) + assert info["key"] == key + assert {field: info[field] for field in baseline} == { + field: entry.get(field, value) for field, value in baseline.items() + }, key + + +def test_shipped_claude_baseline_leaves_other_providers_entries_unfilled(shipped_cost_map): + baseline: Final = _claude_baseline()["model_info"] + entries: Final = tuple( + (key, entry) + for key, entry in _mapped_claude_chat_entries() + if entry["litellm_provider"] not in CLAUDE_BASELINE_FILL_PROVIDERS + ) + assert any("supports_parallel_function_calling" not in entry for _, entry in entries) + + for key, entry in entries: + info = litellm.get_model_info(key) + missing = baseline.keys() - entry.keys() + assert {field: info.get(field) for field in missing} == dict.fromkeys(missing), key From c5da69cd131ee227d9a466cd068cddfa7afb0858 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 4 Nov 2026 18:50:14 -0800 Subject: [PATCH 07/39] feat(ui): show model deprecation banner on the models page (#50080) * feat(ui): show model deprecation banner on the models page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(ui): format model deprecation banner Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): refresh model deprecations when models change Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): address review feedback on model deprecation banner Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): cover deprecation role gate and agent builder refresh Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): separate deprecation banner toggle from alerts link Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): dedupe and allow dismissing the model deprecation banner Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): refresh table once and key banner entries on successor Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: quiet successor logs, count banner models, use local dates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): pin banner test date and bucket upcoming entries locally Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: cover mixed-successor log level through the deprecations endpoint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): move deprecation banner cases into integration tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): keep deprecation banner when a refetch fails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/model_deprecation.py | 46 +- litellm/proxy/proxy_server.py | 7 +- litellm/types/proxy/model_deprecation.py | 4 + .../common_utils/test_model_deprecation.py | 50 ++ .../proxy/test_model_deprecations_endpoint.py | 78 ++- .../hooks/models/invalidateModelQueries.ts | 9 + .../models/useModelDeprecations.test.tsx | 64 +++ .../hooks/models/useModelDeprecations.ts | 33 ++ .../app/(dashboard)/hooks/models/useModels.ts | 5 +- .../AllModelsTab.integration.test.tsx | 232 +++++++++ .../components/AllModelsTab.test.tsx | 9 +- .../components/AllModelsTab.tsx | 11 +- .../ModelDeprecationBanner.test.tsx | 477 ++++++++++++++++++ .../components/ModelDeprecationBanner.tsx | 98 ++++ ...riceDataManagementTab.integration.test.tsx | 83 +++ .../PriceDataManagementTab.test.tsx | 8 +- .../components/PriceDataManagementTab.tsx | 6 +- .../components/deprecationBannerText.test.ts | 140 +++++ .../components/deprecationBannerText.ts | 96 ++++ .../page.integration.test.tsx | 55 +- .../models-and-endpoints/page.test.tsx | 9 + .../(dashboard)/models-and-endpoints/page.tsx | 7 +- .../panels/AddModelPanel.integration.test.tsx | 78 ++- .../panels/AddModelPanel.tsx | 3 +- .../chat_ui/AgentBuilderView.test.tsx | 97 +++- .../components/chat_ui/AgentBuilderView.tsx | 6 + .../model_info_view.integration.test.tsx | 79 ++- .../src/components/model_info_view.tsx | 3 +- .../src/components/settings.test.tsx | 27 +- .../src/components/settings.tsx | 13 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 31 files changed, 1785 insertions(+), 53 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/models/invalidateModelQueries.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.ts diff --git a/litellm/proxy/common_utils/model_deprecation.py b/litellm/proxy/common_utils/model_deprecation.py index b86ebd830f9..832ebe175b5 100644 --- a/litellm/proxy/common_utils/model_deprecation.py +++ b/litellm/proxy/common_utils/model_deprecation.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from datetime import date, datetime, timezone @@ -166,14 +167,20 @@ def _classify(days_until: int, warn_within_days: int) -> DeprecationStatus: return "upcoming" -def _checked_successor(model_name: str, model_info: Mapping[str, object], llm_router: Router) -> str | None: +def _checked_successor( + model_name: str, + model_info: Mapping[str, object], + llm_router: Router, + successor_log_level: int, +) -> str | None: """The deployment's successor_model as set, with a warning when no model in its team scope matches it""" successor: Final = _successor_model(model_info.get("successor_model")) if successor is None: return None team_id: Final = model_info.get("team_id") if not llm_router.registry().deployments(successor, team_id=team_id if isinstance(team_id, str) else None): - verbose_logger.warning( + verbose_logger.log( + successor_log_level, "model_deprecation: %s names successor %s, which matches no model group, alias, or wildcard route here", model_name, successor, @@ -187,6 +194,7 @@ def _build_info( today: date, warn_within_days: int, catalog: Mapping[str, object], + successor_log_level: int, ) -> _DeploymentInfo | None: model_name: Final = deployment.get("model_name") if not isinstance(model_name, str) or not model_name: @@ -197,10 +205,18 @@ def _build_info( return None days_until: Final = (resolved.deprecation_date - today).days - configured: Final = _checked_successor(model_name, _mapping_field(deployment, "model_info"), llm_router) + model_info: Final = _mapping_field(deployment, "model_info") + configured: Final = _checked_successor(model_name, model_info, llm_router, successor_log_level) + team_public_model_name_value: Final = model_info.get("team_public_model_name") + team_public_model_name: Final = ( + team_public_model_name_value + if isinstance(team_public_model_name_value, str) and team_public_model_name_value.strip() + else None + ) return _DeploymentInfo( info=ModelDeprecationInfo( model_name=model_name, + team_public_model_name=team_public_model_name, litellm_model=resolved.litellm_model, deprecation_date=resolved.deprecation_date, days_until_deprecation=days_until, @@ -218,7 +234,7 @@ class _DeploymentInfo: configured_successor: bool -def _merge_group(group: Sequence[_DeploymentInfo]) -> ModelDeprecationInfo: +def _merge_group(group: Sequence[_DeploymentInfo], successor_log_level: int) -> ModelDeprecationInfo: """Keep the group's first deployment in config order, naming the first successor any of its deployments names, a successor configured in model_info ahead of one the cost map supplies""" first: Final = group[0].info @@ -234,7 +250,8 @@ def _merge_group(group: Sequence[_DeploymentInfo]) -> ModelDeprecationInfo: if not successors: return first if len(frozenset(successors)) > 1: - verbose_logger.warning( + verbose_logger.log( + successor_log_level, "model_deprecation: %s names several successors %s, reporting %s", first.model_name, sorted(frozenset(successors)), @@ -245,11 +262,12 @@ def _merge_group(group: Sequence[_DeploymentInfo]) -> ModelDeprecationInfo: def _dedupe( models: Sequence[_DeploymentInfo], + successor_log_level: int, ) -> tuple[ModelDeprecationInfo, ...]: """Report a model group carrying the same date on several deployments once""" ordered: Final = sorted(models, key=lambda model: (model.info.model_name, model.info.deprecation_date)) return tuple( - _merge_group(tuple(group)) + _merge_group(tuple(group), successor_log_level) for _, group in groupby(ordered, key=lambda model: (model.info.model_name, model.info.deprecation_date)) ) @@ -264,11 +282,18 @@ def _bucket(models: Sequence[ModelDeprecationInfo], status: DeprecationStatus) - def _collect_from_router( - llm_router: Router, today: date, warn_within_days: int, catalog: Mapping[str, object] + llm_router: Router, + today: date, + warn_within_days: int, + catalog: Mapping[str, object], + successor_log_level: int, ) -> tuple[ModelDeprecationInfo, ...]: deployments: Final = llm_router.registry().deployments() or () - infos: Final = (_build_info(deployment, llm_router, today, warn_within_days, catalog) for deployment in deployments) - deduped: Final = _dedupe(tuple(info for info in infos if info is not None)) + infos: Final = ( + _build_info(deployment, llm_router, today, warn_within_days, catalog, successor_log_level) + for deployment in deployments + ) + deduped: Final = _dedupe(tuple(info for info in infos if info is not None), successor_log_level) verbose_logger.debug( "model_deprecation: %d/%d deployments carry a deprecation date", @@ -283,12 +308,13 @@ def collect_model_deprecations( warn_within_days: int = DEFAULT_DEPRECATION_WARN_DAYS, today: date | None = None, loaded_catalog: LoadedCatalog = GetModelCostMap.loaded_model_cost_map, + successor_log_level: int = logging.WARNING, ) -> ModelDeprecationResponse: """Bucket every deployment carrying a deprecation date by how urgent it is""" snapshot_time: Final = datetime.now(timezone.utc) effective_today: Final = today or snapshot_time.date() deduped: Final = ( - _collect_from_router(llm_router, effective_today, warn_within_days, loaded_catalog()) + _collect_from_router(llm_router, effective_today, warn_within_days, loaded_catalog(), successor_log_level) if llm_router is not None else () ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6a125b9a0a4..fddecb3d6eb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6,6 +6,7 @@ import importlib import inspect import io import itertools +import logging import os import random import re @@ -18662,7 +18663,11 @@ async def model_deprecations( -H "Authorization: Bearer $LITELLM_MASTER_KEY" ``` """ - return collect_model_deprecations(llm_router=llm_router, warn_within_days=warn_within_days) + return collect_model_deprecations( + llm_router=llm_router, + warn_within_days=warn_within_days, + successor_log_level=logging.DEBUG, + ) def get_model_group_info( diff --git a/litellm/types/proxy/model_deprecation.py b/litellm/types/proxy/model_deprecation.py index 3699f01a684..7e459904d59 100644 --- a/litellm/types/proxy/model_deprecation.py +++ b/litellm/types/proxy/model_deprecation.py @@ -18,6 +18,10 @@ DeprecationStatus = Literal["upcoming", "imminent", "deprecated"] class ModelDeprecationInfo(LiteLLMBaseModel): model_name: str = Field(description="The public name of the model on the proxy (model_group).") + team_public_model_name: str | None = Field( + default=None, + description="The model name visible to the team, when configured.", + ) litellm_model: str | None = Field( default=None, description="The underlying litellm model string the deprecation date is sourced from.", diff --git a/tests/unit/proxy/common_utils/test_model_deprecation.py b/tests/unit/proxy/common_utils/test_model_deprecation.py index 7e96ad591b7..b3c135f1fad 100644 --- a/tests/unit/proxy/common_utils/test_model_deprecation.py +++ b/tests/unit/proxy/common_utils/test_model_deprecation.py @@ -467,6 +467,52 @@ class TestFormatDeprecationAlertMessage: class TestSuccessorModel: + def test_unmatched_successor_logs_warning_by_default(self, monkeypatch, caplog): + monkeypatch.setattr(litellm, "model_cost", {}) + router = Router( + model_list=[ + { + "model_name": "old", + "litellm_params": {"model": "openai/old-model"}, + "model_info": { + "id": "1", + "deprecation_date": "2026-06-10", + "successor_model": "missing-model", + }, + } + ] + ) + + with caplog.at_level(logging.DEBUG, logger=verbose_logger.name): + collect_model_deprecations(llm_router=router, today=date(2026, 6, 1)) + + successor_records: Final = tuple( + record for record in caplog.records if "names successor missing-model" in record.getMessage() + ) + assert len(successor_records) == 1 + assert successor_records[0].levelno == logging.WARNING + + def test_empty_team_public_model_name_is_normalized_to_none(self, monkeypatch): + monkeypatch.setattr(litellm, "model_cost", {}) + router = Router( + model_list=[ + { + "model_name": "team-model", + "litellm_params": {"model": "openai/team-model"}, + "model_info": { + "id": "1", + "team_id": "team-1", + "team_public_model_name": "", + "deprecation_date": "2026-06-10", + }, + } + ] + ) + + snapshot: Final = collect_model_deprecations(llm_router=router, today=date(2026, 6, 1)) + + assert snapshot.imminent[0].team_public_model_name is None + def test_should_surface_the_successor_set_on_the_deployment(self, monkeypatch): monkeypatch.setattr(litellm, "model_cost", {}) router = Router( @@ -572,6 +618,10 @@ class TestSuccessorModel: "typo-old": "gpt-nwe", "wildcard-old": "anthropic/x", } + team_deprecation: Final = next( + info for info in snapshot.imminent if info.model_name == "model_name_t1_old" + ) + assert team_deprecation.team_public_model_name == "team-old" warnings = [record.getMessage() for record in caplog.records if "names successor" in record.getMessage()] assert len(warnings) == 2 assert any("global-old names successor team-new" in warning for warning in warnings) diff --git a/tests/unit/proxy/test_model_deprecations_endpoint.py b/tests/unit/proxy/test_model_deprecations_endpoint.py index 7be90d0d1e4..705f5ec7b39 100644 --- a/tests/unit/proxy/test_model_deprecations_endpoint.py +++ b/tests/unit/proxy/test_model_deprecations_endpoint.py @@ -1,10 +1,12 @@ +import logging +from typing import Final from unittest.mock import MagicMock import pytest from fastapi.testclient import TestClient - import litellm +from litellm._logging import verbose_logger from litellm.proxy import proxy_server from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -140,3 +142,77 @@ def test_cost_map_successor_is_reported_unless_the_deployment_names_its_own(clie "dated-by-hand": "fresh-model", "orphan-alias": None, } + + +def test_unmatched_successor_logs_below_warning_through_endpoint(client_with_cost_map_successors, caplog, monkeypatch): + monkeypatch.setattr( + proxy_server.llm_router.get_model_list, + "side_effect", + lambda model_name=None, **kwargs: ( + proxy_server.llm_router.get_model_list.return_value if model_name is None else [] + ), + ) + + with caplog.at_level(logging.DEBUG, logger=verbose_logger.name): + response = client_with_cost_map_successors.get("/model/deprecations") + + assert response.status_code == 200, response.text + successor_records = tuple( + record + for record in caplog.records + if record.name == verbose_logger.name and "successor" in record.getMessage() + ) + assert successor_records + assert all(record.levelno < logging.WARNING for record in successor_records) + + +def test_mixed_successor_group_logs_below_warning_through_endpoint( + client_with_cost_map_successors: TestClient, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + router: Final = proxy_server.llm_router + models: Final[tuple[dict[str, object], ...]] = ( + { + "model_name": "mixed-successor-group", + "litellm_params": {"model": "retiring-model"}, + "model_info": {"id": "10", "successor_model": "successor-a"}, + }, + { + "model_name": "mixed-successor-group", + "litellm_params": {"model": "retiring-model"}, + "model_info": {"id": "11", "successor_model": "successor-b"}, + }, + { + "model_name": "successor-a", + "litellm_params": {"model": "fresh-model"}, + "model_info": {"id": "12"}, + }, + { + "model_name": "successor-b", + "litellm_params": {"model": "fresh-model"}, + "model_info": {"id": "13"}, + }, + ) + + def filtered_model_list( + model_name: str | None = None, **_kwargs: object + ) -> tuple[dict[str, object], ...]: + return models if model_name is None else tuple(model for model in models if model["model_name"] == model_name) + + monkeypatch.setattr(router.get_model_list, "side_effect", filtered_model_list) + + with caplog.at_level(logging.DEBUG, logger=verbose_logger.name): + response = client_with_cost_map_successors.get("/model/deprecations") + + assert response.status_code == 200, response.text + successor_records: Final = tuple( + record + for record in caplog.records + if record.name == verbose_logger.name and "successor" in record.getMessage() + ) + assert successor_records + assert all( + "names several successors" in record.getMessage() and record.levelno < logging.WARNING + for record in successor_records + ) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/invalidateModelQueries.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/invalidateModelQueries.ts new file mode 100644 index 00000000000..3382d6c9c81 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/invalidateModelQueries.ts @@ -0,0 +1,9 @@ +import type { QueryClient } from "@tanstack/react-query"; +import { modelDeprecationKeys } from "./useModelDeprecations"; + +export const invalidateModelQueries = async (queryClient: QueryClient): Promise => { + await Promise.all([ + queryClient.invalidateQueries({ queryKey: ["models", "list"] }), + queryClient.invalidateQueries({ queryKey: modelDeprecationKeys.all }), + ]); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.test.tsx new file mode 100644 index 00000000000..88efd721c26 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.test.tsx @@ -0,0 +1,64 @@ +/* @vitest-environment jsdom */ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { act, renderHook, waitFor } from "@testing-library/react"; +import type { ReactNode } from "react"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { useModelDeprecations } from "./useModelDeprecations"; + +const authState = vi.hoisted(() => ({ accessToken: "token", userRole: "internal_user" })); +const emptyDeprecationsResponse = { + checked_at: "2026-11-04T00:00:00Z", + warn_within_days: 30, + deprecated: [], + imminent: [], + upcoming: [], +}; +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ accessToken: authState.accessToken, userRole: authState.userRole }), +})); + +describe("useModelDeprecations", () => { + let fetchMock: ReturnType>; + + beforeEach(() => { + authState.accessToken = "token"; + authState.userRole = "internal_user"; + fetchMock = vi.fn().mockResolvedValue(Response.json(emptyDeprecationsResponse)); + vi.stubGlobal("fetch", fetchMock); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("does not fetch deprecations for an internal user", async () => { + const queryClient = new QueryClient(); + const wrapper = ({ children }: { children: ReactNode }) => ( + {children} + ); + + const { result } = renderHook(() => useModelDeprecations(), { wrapper }); + + await waitFor(() => expect(result.current.fetchStatus).toBe("idle")); + await act(async () => { + await new Promise((resolve) => setTimeout(resolve, 0)); + }); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("fetches deprecations for a proxy admin viewer", async () => { + authState.userRole = "proxy_admin_viewer"; + const queryClient = new QueryClient(); + const wrapper = ({ children }: { children: ReactNode }) => ( + {children} + ); + + renderHook(() => useModelDeprecations(), { wrapper }); + + await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(1)); + const request = fetchMock.mock.calls[0][0]; + expect(new URL(request instanceof Request ? request.url : String(request), "http://localhost").pathname).toBe( + "/model/deprecations", + ); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.ts new file mode 100644 index 00000000000..bb9922d8a8b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.ts @@ -0,0 +1,33 @@ +"use client"; + +import { useQuery } from "@tanstack/react-query"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { backgroundFetchClient } from "@/lib/http/api"; +import type { components } from "@/lib/http/schema"; +import { isProxyAdminTierRole } from "@/utils/roles"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +export type ModelDeprecationInfo = components["schemas"]["ModelDeprecationInfo"]; +export type ModelDeprecationResponse = components["schemas"]["ModelDeprecationResponse"]; + +export const modelDeprecationKeys = createQueryKeys("modelDeprecations"); + +const fetchModelDeprecations = async (): Promise => { + const { data } = await backgroundFetchClient.GET("/model/deprecations"); + if (data === undefined) { + throw new Error("The proxy returned no model deprecations payload"); + } + return data; +}; + +export const useModelDeprecations = () => { + const { accessToken, userRole } = useAuthorized(); + const queryOptions = { + queryKey: modelDeprecationKeys.all, + queryFn: fetchModelDeprecations, + enabled: Boolean(accessToken) && isProxyAdminTierRole(userRole), + staleTime: 5 * 60 * 1000, + retry: false, + }; + return useQuery(queryOptions); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index f9304bab0d2..7e9669e6ce5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -1,6 +1,7 @@ import { useQuery, useInfiniteQuery, useQueryClient, UseQueryResult } from "@tanstack/react-query"; import { useMemo } from "react"; import { createQueryKeys } from "../common/queryKeysFactory"; +import { invalidateModelQueries } from "./invalidateModelQueries"; import { modelInfoCall, modelHubCall, modelAvailableCall } from "@/components/networking"; import useAuthorized from "../useAuthorized"; import { EndpointType, isModeCompatibleWithEndpoint } from "@/components/chat_ui/mode_endpoint_mapping"; @@ -240,9 +241,7 @@ export const useAutoRouters = (): UseQueryResult export const useInvalidateAutoRouters = (): (() => Promise) => { const queryClient = useQueryClient(); - return async () => { - await queryClient.invalidateQueries({ queryKey: modelKeys.lists() }); - }; + return () => invalidateModelQueries(queryClient); }; export const useModelHub = () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.integration.test.tsx new file mode 100644 index 00000000000..42747bfeeca --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.integration.test.tsx @@ -0,0 +1,232 @@ +import * as useAuthorizedModule from "@/app/(dashboard)/hooks/useAuthorized"; +import { useQuery } from "@tanstack/react-query"; +import { screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { renderWithProviders, testQueryClient } from "../../../../../tests/test-utils"; +import AllModelsTab from "./AllModelsTab"; +import ModelDeprecationBanner from "./ModelDeprecationBanner"; + +const mockModelDeleteCall = vi.fn().mockResolvedValue({}); +const mockModelPatchUpdateCall = vi.fn().mockResolvedValue({}); +vi.mock("@/components/networking", () => ({ + serverRootPath: "/", + modelDeleteCall: (...args: unknown[]) => mockModelDeleteCall(...args), + modelPatchUpdateCall: (...args: unknown[]) => mockModelPatchUpdateCall(...args), +})); + +vi.mock("@/components/model_dashboard/ModelSettingsModal/ModelSettingsModal", () => ({ + default: function ModelSettingsModalMock({ isVisible }: { isVisible: boolean }) { + return isVisible ?
: null; + }, +})); + +let modelsInfoResult: Record = {}; +const mockRefetch = vi.fn(); +vi.mock("../../hooks/models/useModels", () => ({ + useModelsInfo: () => ({ ...modelsInfoResult, refetch: mockRefetch }), +})); + +vi.mock("../../hooks/models/useModelCostMap", () => ({ + useModelCostMap: () => ({ data: { "gpt-4": { litellm_provider: "openai" } }, isLoading: false, error: null }), +})); + +const mockTeams = [{ team_id: "team-1", team_alias: "Engineering" }]; +const mockCredentials = [ + { credential_name: "openai-prod", display_name: "Prod OpenAI", credential_values: {}, credential_info: {} }, +]; +const mockUseCredentials = vi.hoisted(() => vi.fn()); +vi.mock("../../hooks/credentials/useCredentials", () => ({ + useCredentials: mockUseCredentials, +})); + +vi.mock("../../hooks/teams/useTeams", () => ({ + useTeams: () => ({ data: mockTeams, isLoading: false, error: null, refetch: vi.fn() }), +})); + +const BASE_MODEL_INFO = { + id: "model-1", + db_model: true, + created_by: "user-123", + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-02T00:00:00Z", + team_id: "team-1", + access_groups: [], +}; + +const makeRow = () => ({ + model_name: "gpt-4", + litellm_params: { model: "openai/gpt-4", custom_llm_provider: "openai" }, + model_info: BASE_MODEL_INFO, +}); + +const setModelsInfo = (rows: Record[]) => { + modelsInfoResult = { + data: { data: rows, total_count: rows.length, current_page: 1, total_pages: 1, size: 50 }, + isLoading: false, + isFetching: false, + error: null, + }; +}; + +const MOCK_AUTHORIZED = { + isLoading: false, + isAuthorized: true, + token: "mock-token", + accessToken: "mock-access-token", + userId: "user-123", + userEmail: "test@example.com", + userRole: "Admin", + userRoleLabel: "Admin", + isViewOnly: false, + premiumUser: true, + disabledPersonalKeyCreation: false, + showSSOBanner: false, +}; + +const mockSetSelectedModelGroup = vi.fn(); +const mockSetSelectedModelId = vi.fn(); +const mockSetSelectedTeamId = vi.fn(); + +const defaultProps = { + selectedModelGroup: "all", + setSelectedModelGroup: mockSetSelectedModelGroup, + availableModelGroups: ["gpt-4", "gpt-3.5-turbo"], + availableModelAccessGroups: ["sales-team"], + setSelectedModelId: mockSetSelectedModelId, + setSelectedTeamId: mockSetSelectedTeamId, +}; + +const TEST_TODAY = new Date(2026, 10, 4); + +const deprecatedResponse = { + checked_at: "2026-06-01T00:00:00Z", + warn_within_days: 30, + deprecated: [ + { + model_name: "old-model", + litellm_provider: null, + deprecation_date: "2026-06-01", + days_until_deprecation: -1, + status: "deprecated", + successor_model: null, + }, + ], + imminent: [], + upcoming: [], +}; + +const currentResponse = { + checked_at: "2026-06-02T00:00:00Z", + warn_within_days: 30, + deprecated: [], + imminent: [ + { + model_name: "current-model", + litellm_provider: null, + deprecation_date: "2026-11-30", + days_until_deprecation: 26, + status: "imminent", + successor_model: null, + }, + ], + upcoming: [], +}; + +const stubDeprecationResponses = () => { + const fetchMock = vi + .fn() + .mockImplementationOnce(async (input) => { + const requestUrl = input instanceof Request ? input.url : String(input); + expect(new URL(requestUrl, "http://localhost").pathname).toBe("/model/deprecations"); + return Response.json(deprecatedResponse); + }) + .mockImplementationOnce(async (input) => { + const requestUrl = input instanceof Request ? input.url : String(input); + expect(new URL(requestUrl, "http://localhost").pathname).toBe("/model/deprecations"); + return Response.json(currentResponse); + }); + vi.stubGlobal("fetch", fetchMock); + return fetchMock; +}; + +const renderModelsWithDeprecationBanner = () => + renderWithProviders( + <> + + + , + ); + +const ModelListQueryProbe = ({ queryFn }: { queryFn: () => Promise }) => { + useQuery({ queryKey: ["models", "list"], queryFn, staleTime: Infinity }); + return null; +}; + +describe("AllModelsTab", () => { + beforeEach(() => { + vi.clearAllMocks(); + testQueryClient.clear(); + setModelsInfo([makeRow()]); + vi.spyOn(useAuthorizedModule, "default").mockReturnValue(MOCK_AUTHORIZED); + mockUseCredentials.mockImplementation(({ enabled = true }: { enabled?: boolean } = {}) => ({ + data: enabled ? { credentials: mockCredentials } : undefined, + isLoading: false, + })); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("refreshes deprecation banner data after deleting a model", async () => { + const fetchMock = stubDeprecationResponses(); + const user = userEvent.setup(); + renderModelsWithDeprecationBanner(); + + expect(await screen.findByText("old-model was deprecated on 2026-06-01")).toBeInTheDocument(); + await user.click(await screen.findByTestId("model-delete-model-1")); + await user.click(await screen.findByRole("button", { name: /^delete$/i })); + expect(await screen.findByText("current-model deprecates on 2026-11-30, in 26 days")).toBeInTheDocument(); + await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + expect(screen.queryByText("old-model was deprecated on 2026-06-01")).not.toBeInTheDocument(); + }); + + it("refreshes deprecation banner data after pausing a model", async () => { + const fetchMock = stubDeprecationResponses(); + const user = userEvent.setup(); + renderModelsWithDeprecationBanner(); + + expect(await screen.findByText("old-model was deprecated on 2026-06-01")).toBeInTheDocument(); + await user.click(await screen.findByTestId("model-pause-toggle-model-1")); + expect(await screen.findByText("current-model deprecates on 2026-11-30, in 26 days")).toBeInTheDocument(); + await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + expect(screen.queryByText("old-model was deprecated on 2026-06-01")).not.toBeInTheDocument(); + }); + + it("refreshes the table and deprecations once without refetching another models list observer", async () => { + const fetchMock = stubDeprecationResponses(); + const observerQueryFn = vi.fn<() => Promise>().mockResolvedValue("models"); + const user = userEvent.setup(); + renderWithProviders( + <> + + + + , + ); + + expect(await screen.findByText("old-model was deprecated on 2026-06-01")).toBeInTheDocument(); + await waitFor(() => expect(observerQueryFn).toHaveBeenCalledTimes(1)); + mockRefetch.mockClear(); + fetchMock.mockClear(); + + await user.click(screen.getByTestId("datatable-refresh")); + + expect(await screen.findByText("current-model deprecates on 2026-11-30, in 26 days")).toBeInTheDocument(); + await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(1)); + expect(observerQueryFn).toHaveBeenCalledTimes(1); + expect(mockRefetch).toHaveBeenCalledTimes(1); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx index 09e00db9fb0..97e7f0238bb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx @@ -4,7 +4,7 @@ import userEvent from "@testing-library/user-event"; import type { OnUrlUpdateFunction } from "nuqs/adapters/testing"; import { beforeEach, describe, expect, it, Mock, vi } from "vitest"; -import { renderWithProviders } from "../../../../../tests/test-utils"; +import { renderWithProviders, testQueryClient } from "../../../../../tests/test-utils"; import AllModelsTab from "./AllModelsTab"; import { STATUS_COLUMN_ID, toServerSortField } from "./ModelsTableColumns"; @@ -22,12 +22,6 @@ vi.mock("@/components/model_dashboard/ModelSettingsModal/ModelSettingsModal", () }, })); -const mockInvalidateQueries = vi.fn(); -vi.mock("@tanstack/react-query", async (importOriginal) => { - const actual = await importOriginal(); - return { ...actual, useQueryClient: () => ({ invalidateQueries: mockInvalidateQueries }) }; -}); - interface ModelsInfoArgs { page?: number; size?: number; @@ -157,6 +151,7 @@ const defaultProps = { describe("AllModelsTab", () => { beforeEach(() => { vi.clearAllMocks(); + testQueryClient.clear(); modelsInfoCalls.length = 0; setModelsInfo([makeRow()]); vi.spyOn(useAuthorizedModule, "default").mockReturnValue(MOCK_AUTHORIZED); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index ff027ac8b3c..cf977a6d9e3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -22,6 +22,8 @@ import { useCallback, useMemo, useState } from "react"; import { Trans, useTranslation } from "react-i18next"; import { useModelsInfo } from "../../hooks/models/useModels"; +import { modelDeprecationKeys } from "../../hooks/models/useModelDeprecations"; +import { invalidateModelQueries } from "../../hooks/models/invalidateModelQueries"; import { ROUTING_STATUS_FILTER_ID, routingStatusToBlocked } from "../utils/routingStatus"; import { transformModelData } from "../utils/modelDataTransformer"; import { @@ -266,8 +268,8 @@ const AllModelsTab = ({ try { setDeleteLoading(true); await modelDeleteCall(accessToken, deleteModalModelId); + void invalidateModelQueries(queryClient); toast.success(t("modelsAndEndpoints.allModelsTab.toasts.deleted")); - queryClient.invalidateQueries({ queryKey: ["models", "list"] }); refetchModels(); } catch (error) { console.error("Error deleting model:", error); @@ -291,7 +293,7 @@ const AllModelsTab = ({ ); // invalidateQueries already schedules a refetch for active observers // on this key — no need to also call refetchModels() (would double-fetch). - queryClient.invalidateQueries({ queryKey: ["models", "list"] }); + void invalidateModelQueries(queryClient); } catch (error) { console.error("Error toggling model pause state:", error); toast.fromError(error); @@ -303,8 +305,9 @@ const AllModelsTab = ({ ); const handleRefresh = useCallback(() => { - void refetchModels(); - }, [refetchModels]); + refetchModels(); + void queryClient.invalidateQueries({ queryKey: modelDeprecationKeys.all }); + }, [queryClient, refetchModels]); const handleDeleteClick = useCallback((modelId: string) => { setDeleteModalModelId(modelId); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.test.tsx new file mode 100644 index 00000000000..c14bf93c116 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.test.tsx @@ -0,0 +1,477 @@ +/* @vitest-environment jsdom */ +import { render, screen, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { + useModelDeprecations, + type ModelDeprecationInfo, + type ModelDeprecationResponse, +} from "@/app/(dashboard)/hooks/models/useModelDeprecations"; +import { uiHref } from "@/utils/uiHref"; +import ModelDeprecationBanner from "./ModelDeprecationBanner"; + +vi.mock("@/app/(dashboard)/hooks/models/useModelDeprecations", () => ({ + useModelDeprecations: vi.fn(), +})); + +const mockedUseModelDeprecations = vi.mocked(useModelDeprecations); + +const makeResponse = (overrides: Partial = {}): ModelDeprecationResponse => ({ + checked_at: "2026-11-04T00:00:00Z", + warn_within_days: 30, + deprecated: [], + imminent: [], + upcoming: [], + ...overrides, +}); + +const makeInfo = (overrides: Partial = {}): ModelDeprecationInfo => ({ + model_name: "gpt-4o", + deprecation_date: "2026-11-30", + days_until_deprecation: 26, + status: "imminent", + ...overrides, +}); + +const queryResult = ( + data: ModelDeprecationResponse | undefined, + isLoading = false, + isError = false, + overrides: Partial> = {}, +): ReturnType => + ({ data, isLoading, isError, ...overrides }) as ReturnType; + +const pinnedToday = (): Date => new Date(2026, 10, 4); + +const renderBanner = (today = pinnedToday()) => render(); + +describe("ModelDeprecationBanner", () => { + beforeEach(() => { + vi.clearAllMocks(); + window.localStorage.clear(); + }); + + it("renders nothing when there are no deprecated or imminent models", () => { + mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse())); + + renderBanner(); + + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + }); + + it("renders nothing when only upcoming models are present", () => { + mockedUseModelDeprecations.mockReturnValue( + queryResult( + makeResponse({ + upcoming: [makeInfo({ deprecation_date: "2026-12-05", status: "upcoming" })], + }), + ), + ); + + renderBanner(); + + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + }); + + it("includes server-upcoming models at the local warning-window boundary", () => { + const upcomingInfoOverrides: Partial = { + model_name: "due-in-window", + deprecation_date: "2026-12-04", + days_until_deprecation: 31, + status: "upcoming", + }; + const info = makeInfo(upcomingInfoOverrides); + mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ upcoming: [info] }))); + + renderBanner(); + + expect(screen.getByText("1 model deprecates within 30 days")).toBeInTheDocument(); + expect(within(screen.getByRole("list")).getByRole("listitem")).toHaveTextContent( + "due-in-window deprecates on 2026-12-04, in 30 days", + ); + }); + + it("does not show server-upcoming models one day past the local warning window", () => { + const upcomingInfoOverrides: Partial = { + model_name: "due-outside-window", + deprecation_date: "2026-12-05", + days_until_deprecation: 31, + status: "upcoming", + }; + const info = makeInfo(upcomingInfoOverrides); + mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ upcoming: [info] }))); + + renderBanner(); + + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + }); + + it("renders nothing while loading or after an error", () => { + mockedUseModelDeprecations.mockReturnValue(queryResult(undefined, true)); + const { rerender } = renderBanner(); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + + mockedUseModelDeprecations.mockReturnValue(queryResult(undefined, false, true)); + rerender(); + + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + }); + + it("keeps the banner visible when a refetch fails", () => { + const staleInfoOverrides: Partial = { + model_name: "stale-model", + deprecation_date: "2026-06-01", + days_until_deprecation: -156, + status: "deprecated", + }; + const info = makeInfo(staleInfoOverrides); + const refetchError = new Error("Refetch failed"); + const refetchErrorResult: Partial> = { + error: refetchError, + failureCount: 1, + failureReason: refetchError, + fetchStatus: "idle", + isFetching: false, + isLoadingError: false, + isRefetchError: true, + status: "error", + }; + mockedUseModelDeprecations.mockReturnValue( + queryResult(makeResponse({ deprecated: [info] }), false, true, refetchErrorResult), + ); + + renderBanner(); + + expect(screen.getByRole("alert")).toBeInTheDocument(); + expect(within(screen.getByRole("list")).getByRole("listitem")).toHaveTextContent( + "stale-model was deprecated on 2026-06-01", + ); + }); + + it("shows deprecated and imminent entries with a link to Logging & Alerts", () => { + const deprecatedInfo: ModelDeprecationInfo = { + model_name: "gpt-4", + team_public_model_name: null, + litellm_provider: "openai", + deprecation_date: "2026-06-01", + days_until_deprecation: -152, + status: "deprecated", + successor_model: "gpt-5.6-sol", + }; + const imminentInfo: ModelDeprecationInfo = { + model_name: "gpt-4o", + team_public_model_name: null, + litellm_provider: "openai", + deprecation_date: "2026-11-30", + days_until_deprecation: 26, + status: "imminent", + successor_model: "X", + }; + const response = makeResponse({ + deprecated: [deprecatedInfo], + imminent: [imminentInfo], + }); + mockedUseModelDeprecations.mockReturnValue(queryResult(response)); + + renderBanner(); + + expect(screen.getByText("1 model is deprecated, 1 deprecates within 30 days")).toBeInTheDocument(); + const entries = within(screen.getByRole("list")).getAllByRole("listitem"); + expect(entries[0]).toHaveTextContent("gpt-4 (openai) was deprecated on 2026-06-01, migrate to gpt-5.6-sol"); + expect(entries[1]).toHaveTextContent("gpt-4o (openai) deprecates on 2026-11-30, in 26 days, migrate to X"); + expect(screen.getByRole("link", { name: "Get Slack alerts for deprecations" })).toHaveAttribute( + "href", + `${uiHref("logging-and-alerts")}?tab=alerting-types`, + ); + }); + + it("treats a server-deprecated entry dated today as imminent locally", () => { + const infoOverrides: Partial = { + model_name: "due-today", + deprecation_date: "2026-11-04", + days_until_deprecation: -1, + status: "deprecated", + }; + const info = makeInfo(infoOverrides); + mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ deprecated: [info] }))); + + renderBanner(); + + expect(screen.getByText("1 model deprecates within 30 days")).toBeInTheDocument(); + expect(within(screen.getByRole("list")).getByRole("listitem")).toHaveTextContent( + "due-today deprecates today (2026-11-04)", + ); + }); + + it("treats a server-imminent entry dated yesterday as deprecated locally", () => { + const infoOverrides: Partial = { + model_name: "due-yesterday", + deprecation_date: "2026-11-03", + days_until_deprecation: 0, + status: "imminent", + }; + const info = makeInfo(infoOverrides); + mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ imminent: [info] }))); + + renderBanner(); + + expect(screen.getByText("1 model is deprecated")).toBeInTheDocument(); + expect(within(screen.getByRole("list")).getByRole("listitem")).toHaveTextContent( + "due-yesterday was deprecated on 2026-11-03", + ); + }); + + it("uses the local date difference instead of the server day count", () => { + const infoOverrides: Partial = { + model_name: "due-in-two", + deprecation_date: "2026-11-06", + days_until_deprecation: 1, + status: "imminent", + }; + const info = makeInfo(infoOverrides); + mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ imminent: [info] }))); + + renderBanner(); + + expect(within(screen.getByRole("list")).getByRole("listitem")).toHaveTextContent( + "due-in-two deprecates on 2026-11-06, in 2 days", + ); + }); + + it("deduplicates entries by their displayed name, provider, and deprecation date", () => { + const globalInfoOverrides: Partial = { + model_name: "o4-mini", + team_public_model_name: null, + litellm_provider: "openai", + deprecation_date: "2026-10-23", + status: "deprecated", + }; + const firstTeamInfoOverrides: Partial = { + model_name: "o4-mini-team-a", + team_public_model_name: "o4-mini", + litellm_provider: "openai", + deprecation_date: "2026-10-23", + status: "deprecated", + }; + const secondTeamInfoOverrides: Partial = { + model_name: "o4-mini-team-b", + team_public_model_name: "o4-mini", + litellm_provider: "openai", + deprecation_date: "2026-10-23", + status: "deprecated", + }; + mockedUseModelDeprecations.mockReturnValue( + queryResult( + makeResponse({ + deprecated: [ + makeInfo(globalInfoOverrides), + makeInfo(firstTeamInfoOverrides), + makeInfo(secondTeamInfoOverrides), + ], + }), + ), + ); + + renderBanner(); + + expect(screen.getByText("1 model is deprecated")).toBeInTheDocument(); + const entries = within(screen.getByRole("list")).getAllByRole("listitem"); + expect(entries).toHaveLength(1); + expect(entries[0]).toHaveTextContent("o4-mini (openai) was deprecated on 2026-10-23"); + }); + + it("keeps entries with different successors and renders each migration target", () => { + const firstEntryOverrides: Partial = { + model_name: "team-model-a", + team_public_model_name: "shared-model", + litellm_provider: "openai", + deprecation_date: "2026-10-23", + days_until_deprecation: -1, + status: "deprecated", + successor_model: "successor-a", + }; + const secondEntryOverrides: Partial = { + model_name: "shared-model", + litellm_provider: "openai", + deprecation_date: "2026-10-23", + days_until_deprecation: -1, + status: "deprecated", + successor_model: "successor-b", + }; + const firstEntry = makeInfo(firstEntryOverrides); + const secondEntry = makeInfo(secondEntryOverrides); + mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ deprecated: [firstEntry, secondEntry] }))); + + renderBanner(); + + const entries = within(screen.getByRole("list")).getAllByRole("listitem"); + expect(screen.getByText("1 model is deprecated")).toBeInTheDocument(); + expect(entries).toHaveLength(2); + expect(entries[0]).toHaveTextContent("shared-model (openai) was deprecated on 2026-10-23, migrate to successor-a"); + expect(entries[1]).toHaveTextContent("shared-model (openai) was deprecated on 2026-10-23, migrate to successor-b"); + }); + + it("keeps entries with different deprecation dates", () => { + const firstEntryOverrides: Partial = { + model_name: "shared-model", + litellm_provider: "openai", + deprecation_date: "2026-10-22", + days_until_deprecation: -2, + status: "deprecated", + successor_model: "successor", + }; + const secondEntryOverrides: Partial = { + model_name: "shared-model", + litellm_provider: "openai", + deprecation_date: "2026-10-23", + days_until_deprecation: -1, + status: "deprecated", + successor_model: "successor", + }; + const firstEntry = makeInfo(firstEntryOverrides); + const secondEntry = makeInfo(secondEntryOverrides); + mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ deprecated: [firstEntry, secondEntry] }))); + + renderBanner(); + + const entries = within(screen.getByRole("list")).getAllByRole("listitem"); + expect(entries).toHaveLength(2); + expect(entries[0]).toHaveTextContent("shared-model (openai) was deprecated on 2026-10-22"); + expect(entries[1]).toHaveTextContent("shared-model (openai) was deprecated on 2026-10-23"); + }); + + it("dismisses the banner and keeps it dismissed after remounting the same entries", async () => { + const info = makeInfo(); + const secondInfo = makeInfo({ model_name: "claude-sonnet-4" }); + mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ imminent: [info, secondInfo] }))); + const user = userEvent.setup(); + const { unmount } = renderBanner(); + + await user.click(screen.getByRole("button", { name: "Dismiss deprecation warning" })); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + unmount(); + + const changedSecondInfoOverrides: Partial = { + ...secondInfo, + days_until_deprecation: secondInfo.days_until_deprecation - 1, + }; + const changedInfoOverrides: Partial = { + ...info, + days_until_deprecation: info.days_until_deprecation - 1, + }; + const teamAliasInfoOverrides: Partial = { + ...info, + model_name: "team-internal-name", + team_public_model_name: info.model_name, + days_until_deprecation: info.days_until_deprecation - 1, + }; + mockedUseModelDeprecations.mockReturnValue( + queryResult( + makeResponse({ + imminent: [ + makeInfo(changedSecondInfoOverrides), + makeInfo(changedInfoOverrides), + makeInfo(teamAliasInfoOverrides), + ], + }), + ), + ); + renderBanner(); + + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + }); + + it("shows the banner again when a new entry is added after dismissal", async () => { + const firstEntryOverrides: Partial = { + model_name: "gpt-4", + deprecation_date: "2026-10-23", + days_until_deprecation: -12, + status: "deprecated", + }; + const firstEntry = makeInfo(firstEntryOverrides); + const addedEntryOverrides: Partial = { + model_name: "gpt-3.5", + deprecation_date: "2026-10-23", + days_until_deprecation: -12, + status: "deprecated", + }; + const addedEntry = makeInfo(addedEntryOverrides); + mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ deprecated: [firstEntry] }))); + const user = userEvent.setup(); + const { rerender } = renderBanner(); + + await user.click(screen.getByRole("button", { name: "Dismiss deprecation warning" })); + mockedUseModelDeprecations.mockReturnValue( + queryResult( + makeResponse({ + deprecated: [firstEntry, addedEntry], + }), + ), + ); + rerender(); + + expect(screen.getByRole("alert")).toBeInTheDocument(); + expect(screen.getByText("2 models are deprecated")).toBeInTheDocument(); + }); + + it("shows the banner again when an imminent entry becomes deprecated", async () => { + const imminentEntryOverrides: Partial = { + model_name: "gpt-4o", + deprecation_date: "2026-11-30", + days_until_deprecation: 0, + status: "imminent", + }; + const imminentEntry = makeInfo(imminentEntryOverrides); + mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ imminent: [imminentEntry] }))); + const user = userEvent.setup(); + const { rerender } = renderBanner(new Date(2026, 10, 30)); + + await user.click(screen.getByRole("button", { name: "Dismiss deprecation warning" })); + rerender(); + + expect(screen.getByRole("alert")).toBeInTheDocument(); + expect(screen.getByText("1 model is deprecated")).toBeInTheDocument(); + }); + + it("shows the banner again when an entry's successor changes after dismissal", async () => { + const entry = makeInfo({ successor_model: "successor-a" }); + mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ imminent: [entry] }))); + const user = userEvent.setup(); + const { rerender } = renderBanner(); + + await user.click(screen.getByRole("button", { name: "Dismiss deprecation warning" })); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + + mockedUseModelDeprecations.mockReturnValue( + queryResult(makeResponse({ imminent: [makeInfo({ ...entry, successor_model: "successor-b" })] })), + ); + rerender(); + + expect(screen.getByRole("alert")).toBeInTheDocument(); + expect(screen.getByText(/migrate to successor-b/)).toBeInTheDocument(); + }); + + it("expands all entries and collapses back to the first five", async () => { + const imminent = [ + ...Array.from({ length: 7 }, (_value, index) => makeInfo({ model_name: `model-${index + 1}` })), + makeInfo({ model_name: "team-model-1", team_public_model_name: "model-1" }), + ]; + mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ imminent }))); + const user = userEvent.setup(); + + renderBanner(); + + const entries = within(screen.getByRole("list")).getAllByRole("listitem"); + expect(entries).toHaveLength(5); + expect(entries[4]).toHaveTextContent("model-5 deprecates"); + expect(screen.queryByText(/model-6/)).not.toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "Show all 7" })); + expect(within(screen.getByRole("list")).getAllByRole("listitem")).toHaveLength(7); + expect(screen.getByText(/model-6 deprecates/)).toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "Show fewer" })); + expect(within(screen.getByRole("list")).getAllByRole("listitem")).toHaveLength(5); + expect(screen.queryByText(/model-6/)).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.tsx new file mode 100644 index 00000000000..3a89940f786 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.tsx @@ -0,0 +1,98 @@ +"use client"; + +import { TriangleAlert, X } from "lucide-react"; +import { useState } from "react"; +import { useModelDeprecations } from "@/app/(dashboard)/hooks/models/useModelDeprecations"; +import { Alert, AlertAction, AlertDescription, AlertTitle } from "@/components/ui/alert"; +import { Button } from "@/components/ui/button"; +import { getLocalStorageItem, setLocalStorageItem } from "@/utils/localStorageUtils"; +import { uiHref } from "@/utils/uiHref"; +import { + bucketByLocalDate, + countDistinctModels, + dedupeDeprecationEntries, + deprecationBannerFingerprint, + deprecationBannerTitle, + deprecationEntryText, +} from "./deprecationBannerText"; + +const STORAGE_KEY = "modelDeprecationBannerDismissed"; + +const ModelDeprecationBanner = ({ today = new Date() }: { today?: Date }) => { + const { data } = useModelDeprecations(); + const [isExpanded, setIsExpanded] = useState(false); + const [dismissedFingerprint, setDismissedFingerprint] = useState(() => + getLocalStorageItem(STORAGE_KEY), + ); + + if (!data) { + return null; + } + + const localBuckets = bucketByLocalDate( + [...(data.deprecated ?? []), ...(data.imminent ?? []), ...(data.upcoming ?? [])], + today, + data.warn_within_days, + ); + const deprecated = dedupeDeprecationEntries(localBuckets.deprecated); + const imminent = dedupeDeprecationEntries(localBuckets.imminent); + + if (deprecated.length === 0 && imminent.length === 0) { + return null; + } + + const entries = [...deprecated, ...imminent]; + const fingerprint = deprecationBannerFingerprint(entries); + if (dismissedFingerprint === fingerprint) { + return null; + } + + const visibleEntries = isExpanded ? entries : entries.slice(0, 5); + + return ( + 0 ? "destructive" : "default"} className="mb-4"> + + + + + + {deprecationBannerTitle(countDistinctModels(deprecated), countDistinctModels(imminent), data.warn_within_days)} + + +
    + {visibleEntries.map((info, index) => ( +
  • {deprecationEntryText(info)}
    • + ))} +
+
+ {entries.length > 5 && ( + + )} + + Get Slack alerts for deprecations + +
+ + + ); +}; + +export default ModelDeprecationBanner; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.integration.test.tsx new file mode 100644 index 00000000000..d2ddb805a0b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.integration.test.tsx @@ -0,0 +1,83 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +import ModelDeprecationBanner from "./ModelDeprecationBanner"; +import PriceDataManagementTab from "./PriceDataManagementTab"; + +const TEST_TODAY = new Date(2026, 10, 4); + +vi.mock("@/components/price_data_reload", () => ({ + default: ({ onReloadSuccess }: { onReloadSuccess: () => void }) => , +})); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ accessToken: "sk-test", userRole: "Admin" }), +})); + +const renderTabWithBanner = (queryClient: QueryClient) => + render( + + <> + + + + , + ); + +describe("PriceDataManagementTab", () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("refetches deprecations after reloading price data", async () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const deprecatedResponse = { + checked_at: "2026-06-01T00:00:00Z", + warn_within_days: 30, + deprecated: [ + { + model_name: "old-model", + litellm_provider: null, + deprecation_date: "2026-06-01", + days_until_deprecation: -1, + status: "deprecated", + successor_model: null, + }, + ], + imminent: [], + upcoming: [], + }; + const currentResponse = { + checked_at: "2026-06-02T00:00:00Z", + warn_within_days: 30, + deprecated: [], + imminent: [ + { + model_name: "current-model", + litellm_provider: null, + deprecation_date: "2026-11-30", + days_until_deprecation: 26, + status: "imminent", + successor_model: null, + }, + ], + upcoming: [], + }; + const fetchMock = vi + .fn() + .mockResolvedValueOnce(Response.json(deprecatedResponse)) + .mockResolvedValueOnce(Response.json(currentResponse)); + vi.stubGlobal("fetch", fetchMock); + renderTabWithBanner(queryClient); + + expect(await screen.findByText("old-model was deprecated on 2026-06-01")).toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "reload" })); + expect(await screen.findByText("current-model deprecates on 2026-11-30, in 26 days")).toBeInTheDocument(); + await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + expect( + fetchMock.mock.calls.map( + ([input]) => new URL(input instanceof Request ? input.url : String(input), "http://localhost").pathname, + ), + ).toEqual(["/model/deprecations", "/model/deprecations"]); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx index b2beb5beeed..445dc505ce6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx @@ -2,13 +2,16 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { fireEvent, render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; +import { modelDeprecationKeys } from "../../hooks/models/useModelDeprecations"; import { modelCostMapKeys } from "../../hooks/models/useModelCostMap"; import PriceDataManagementTab from "./PriceDataManagementTab"; vi.mock("@/components/price_data_reload", () => ({ default: ({ onReloadSuccess }: { onReloadSuccess: () => void }) => , })); -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => ({ accessToken: "sk-test" }) })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ accessToken: "sk-test", userRole: "Admin" }), +})); const renderTab = (queryClient: QueryClient) => render( @@ -27,13 +30,16 @@ describe("PriceDataManagementTab", () => { const queryClient = new QueryClient(); const liveKey = modelCostMapKeys.list({}); const catalogKey = modelCostMapKeys.list({ filters: { catalog_only: "true" } }); + const deprecationKey = modelDeprecationKeys.all; queryClient.setQueryData(liveKey, {}); queryClient.setQueryData(catalogKey, {}); + queryClient.setQueryData(deprecationKey, {}); renderTab(queryClient); fireEvent.click(screen.getByRole("button", { name: "reload" })); expect(queryClient.getQueryState(liveKey)?.isInvalidated).toBe(true); expect(queryClient.getQueryState(catalogKey)?.isInvalidated).toBe(true); + expect(queryClient.getQueryState(deprecationKey)?.isInvalidated).toBe(true); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx index 1457b5b311f..4dddc5333ea 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx @@ -3,6 +3,7 @@ import React from "react"; import { useTranslation } from "react-i18next"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useQueryClient } from "@tanstack/react-query"; +import { modelDeprecationKeys } from "../../hooks/models/useModelDeprecations"; import { modelCostMapKeys } from "../../hooks/models/useModelCostMap"; const PriceDataManagementTab = () => { @@ -20,7 +21,10 @@ const PriceDataManagementTab = () => { { - queryClient.invalidateQueries({ queryKey: modelCostMapKeys.all }); + void Promise.all([ + queryClient.invalidateQueries({ queryKey: modelCostMapKeys.all }), + queryClient.invalidateQueries({ queryKey: modelDeprecationKeys.all }), + ]); }} buttonText={t("modelsAndEndpoints.priceDataManagementTab.reloadButton")} size="middle" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.test.ts new file mode 100644 index 00000000000..66d065e6caa --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.test.ts @@ -0,0 +1,140 @@ +import { describe, expect, it } from "vitest"; +import type { ModelDeprecationInfo } from "@/app/(dashboard)/hooks/models/useModelDeprecations"; +import { + bucketByLocalDate, + countDistinctModels, + deprecationBannerTitle, + deprecationEntryText, + localDaysUntil, + parseLocalDate, +} from "./deprecationBannerText"; + +const makeInfo = (overrides: Partial = {}): ModelDeprecationInfo => ({ + model_name: "gpt-4o", + deprecation_date: "2026-11-30", + days_until_deprecation: 26, + status: "imminent", + ...overrides, +}); + +describe("local deprecation dates", () => { + it("parses date parts as a local calendar date and counts from local midnight", () => { + const parsed = parseLocalDate("2026-11-06"); + + expect([parsed.getFullYear(), parsed.getMonth(), parsed.getDate(), parsed.getHours()]).toEqual([2026, 10, 6, 0]); + expect(localDaysUntil("2026-11-06", new Date(2026, 10, 4, 18))).toBe(2); + }); + + it("buckets entries by their local date and warning window", () => { + const buckets = bucketByLocalDate( + [ + makeInfo({ model_name: "yesterday", deprecation_date: "2026-11-03", status: "imminent" }), + makeInfo({ model_name: "today", deprecation_date: "2026-11-04", status: "deprecated" }), + makeInfo({ model_name: "outside-window", deprecation_date: "2026-12-05", status: "imminent" }), + ], + new Date(2026, 10, 4), + 30, + ); + + expect( + buckets.deprecated.map(({ model_name, days_until_deprecation, status }) => [ + model_name, + days_until_deprecation, + status, + ]), + ).toEqual([["yesterday", -1, "deprecated"]]); + expect( + buckets.imminent.map(({ model_name, days_until_deprecation, status }) => [ + model_name, + days_until_deprecation, + status, + ]), + ).toEqual([["today", 0, "imminent"]]); + }); +}); + +describe("countDistinctModels", () => { + it("counts visible names and providers, not deployment aliases", () => { + expect( + countDistinctModels([ + makeInfo({ model_name: "team-a", team_public_model_name: "shared", litellm_provider: "openai" }), + makeInfo({ model_name: "team-b", team_public_model_name: "shared", litellm_provider: "openai" }), + makeInfo({ model_name: "team-c", team_public_model_name: "shared", litellm_provider: "anthropic" }), + ]), + ).toBe(2); + }); +}); + +describe("deprecationBannerTitle", () => { + it("formats both counts with singular and plural wording", () => { + expect(deprecationBannerTitle(2, 1, 30)).toBe("2 models are deprecated, 1 deprecates within 30 days"); + expect(deprecationBannerTitle(1, 2, 30)).toBe("1 model is deprecated, 2 deprecate within 30 days"); + }); + + it("omits empty counts and uses singular model wording", () => { + expect(deprecationBannerTitle(1, 0, 30)).toBe("1 model is deprecated"); + expect(deprecationBannerTitle(0, 1, 30)).toBe("1 model deprecates within 30 days"); + expect(deprecationBannerTitle(0, 3, 30)).toBe("3 models deprecate within 30 days"); + expect(deprecationBannerTitle(0, 0, 30)).toBe(""); + }); +}); + +describe("deprecationEntryText", () => { + it("formats deprecated entries with provider and successor", () => { + const deprecatedInfo: ModelDeprecationInfo = { + model_name: "gpt-4", + litellm_provider: "openai", + deprecation_date: "2026-06-01", + days_until_deprecation: -152, + status: "deprecated", + successor_model: "gpt-5.6-sol", + }; + + expect(deprecationEntryText(deprecatedInfo)).toBe( + "gpt-4 (openai) was deprecated on 2026-06-01, migrate to gpt-5.6-sol", + ); + }); + + it("omits a null provider and successor", () => { + expect( + deprecationEntryText( + makeInfo({ + litellm_provider: null, + successor_model: null, + }), + ), + ).toBe("gpt-4o deprecates on 2026-11-30, in 26 days"); + }); + + it("uses singular wording for one day", () => { + expect( + deprecationEntryText( + makeInfo({ + days_until_deprecation: 1, + }), + ), + ).toBe("gpt-4o deprecates on 2026-11-30, in 1 day"); + }); + + it("says an imminent model deprecates today", () => { + expect( + deprecationEntryText( + makeInfo({ + days_until_deprecation: 0, + deprecation_date: "2026-11-04", + }), + ), + ).toBe("gpt-4o deprecates today (2026-11-04)"); + }); + + it("uses the team-visible model name when present", () => { + expect( + deprecationEntryText( + makeInfo({ + model_name: "internal-name", + team_public_model_name: "team-name", + }), + ), + ).toBe("team-name deprecates on 2026-11-30, in 26 days"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.ts new file mode 100644 index 00000000000..0c035c5bdd6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.ts @@ -0,0 +1,96 @@ +import type { ModelDeprecationInfo } from "@/app/(dashboard)/hooks/models/useModelDeprecations"; + +const shownModelName = (info: ModelDeprecationInfo): string => info.team_public_model_name ?? info.model_name; + +export const countDistinctModels = (infos: readonly ModelDeprecationInfo[]): number => + new Set(infos.map((info) => JSON.stringify([shownModelName(info), info.litellm_provider ?? null]))).size; + +export const parseLocalDate = (value: string): Date => { + const [year, month, day] = value.split("-"); + return new Date(Number(year), Number(month) - 1, Number(day)); +}; + +const startOfLocalDay = (today: Date): Date => new Date(today.getFullYear(), today.getMonth(), today.getDate()); + +export const localDaysUntil = (deprecationDate: string, today: Date): number => + Math.round((parseLocalDate(deprecationDate).getTime() - startOfLocalDay(today).getTime()) / 86_400_000); + +export const bucketByLocalDate = ( + infos: readonly ModelDeprecationInfo[], + today: Date, + warnWithinDays: number, +): { deprecated: ModelDeprecationInfo[]; imminent: ModelDeprecationInfo[] } => { + const inWindow = infos + .map((info): ModelDeprecationInfo => { + const daysUntilDeprecation = localDaysUntil(info.deprecation_date, today); + return { + ...info, + days_until_deprecation: daysUntilDeprecation, + status: daysUntilDeprecation < 0 ? "deprecated" : "imminent", + }; + }) + .filter((info) => info.days_until_deprecation <= warnWithinDays); + + return { + deprecated: inWindow.filter((info) => info.status === "deprecated"), + imminent: inWindow.filter((info) => info.status === "imminent"), + }; +}; + +const hasSameDedupeKey = (candidate: ModelDeprecationInfo, info: ModelDeprecationInfo): boolean => { + const sameDisplayedModel = shownModelName(candidate) === shownModelName(info); + const sameProvider = (candidate.litellm_provider ?? null) === (info.litellm_provider ?? null); + const sameDeprecationDate = candidate.deprecation_date === info.deprecation_date; + const sameSuccessor = (candidate.successor_model ?? null) === (info.successor_model ?? null); + const sameVisibleIdentity = sameDisplayedModel && sameProvider; + const sameDeprecationIdentity = sameDeprecationDate && sameSuccessor; + return sameVisibleIdentity && sameDeprecationIdentity; +}; + +export const dedupeDeprecationEntries = (infos: readonly ModelDeprecationInfo[]): ModelDeprecationInfo[] => + infos.filter((info, index) => infos.findIndex((candidate) => hasSameDedupeKey(candidate, info)) === index); + +export const deprecationBannerFingerprint = (infos: readonly ModelDeprecationInfo[]): string => + JSON.stringify( + infos + .map((info) => + JSON.stringify([ + shownModelName(info), + info.litellm_provider ?? null, + info.deprecation_date, + info.status, + info.successor_model ?? null, + ]), + ) + .sort(), + ); + +export const deprecationBannerTitle = ( + deprecatedCount: number, + imminentCount: number, + warnWithinDays: number, +): string => { + const deprecatedText = + deprecatedCount === 0 ? "" : `${deprecatedCount} ${deprecatedCount === 1 ? "model is" : "models are"} deprecated`; + const imminentModelLabel = imminentCount === 1 ? "model" : "models"; + const imminentSubject = deprecatedCount > 0 ? `${imminentCount}` : `${imminentCount} ${imminentModelLabel}`; + const imminentVerb = imminentCount === 1 ? "deprecates" : "deprecate"; + const imminentText = imminentCount === 0 ? "" : `${imminentSubject} ${imminentVerb} within ${warnWithinDays} days`; + return [deprecatedText, imminentText].filter((text) => text.length > 0).join(", "); +}; + +export const deprecationEntryText = (info: ModelDeprecationInfo): string => { + const model = `${shownModelName(info)}${info.litellm_provider ? ` (${info.litellm_provider})` : ""}`; + const successor = info.successor_model ? `, migrate to ${info.successor_model}` : ""; + + if (info.status === "deprecated") { + return `${model} was deprecated on ${info.deprecation_date}${successor}`; + } + + if (info.days_until_deprecation === 0) { + return `${model} deprecates today (${info.deprecation_date})${successor}`; + } + + const dayLabel = info.days_until_deprecation === 1 ? "day" : "days"; + return `${model} deprecates on ${info.deprecation_date}, in ${info.days_until_deprecation} ${dayLabel}${successor}`; +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx index 4d18966625e..7aeff6f44f2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx @@ -1,9 +1,9 @@ /* @vitest-environment jsdom */ -import { screen } from "@testing-library/react"; +import { screen, waitForElementToBeRemoved } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import * as networking from "@/components/networking"; -import { renderWithProviders } from "../../../../tests/test-utils"; +import { renderWithProviders, testQueryClient } from "../../../../tests/test-utils"; import ModelsAndEndpointsPage from "./page"; vi.mock("./panels/AllModelsPanel", () => ({ default: () =>
})); @@ -201,6 +201,7 @@ const createMockTeamData = (overrides = {}) => ({ describe("ModelsAndEndpointsPage ?team drill-in", () => { beforeEach(() => { + testQueryClient.clear(); authState.userRole = "Admin"; authState.isViewOnly = false; can.mockReturnValue(true); @@ -221,6 +222,7 @@ describe("ModelsAndEndpointsPage ?team drill-in", () => { }); afterEach(() => { + vi.unstubAllGlobals(); vi.clearAllMocks(); }); @@ -244,4 +246,53 @@ describe("ModelsAndEndpointsPage ?team drill-in", () => { expect(screen.queryByRole("button", { name: /edit settings/i })).not.toBeInTheDocument(); expect(screen.queryByText("Disable all global guardrails")).not.toBeInTheDocument(); }); + + it("refreshes deprecation banner data when the models list is refreshed", async () => { + const deprecatedResponse = { + checked_at: "2026-06-01T00:00:00Z", + warn_within_days: 30, + deprecated: [ + { + model_name: "gpt-4", + litellm_provider: "openai", + deprecation_date: "2026-06-01", + days_until_deprecation: -1, + status: "deprecated", + successor_model: "gpt-5", + }, + ], + imminent: [], + upcoming: [], + }; + const currentResponse = { + checked_at: "2026-06-02T00:00:00Z", + warn_within_days: 30, + deprecated: [], + imminent: [], + upcoming: [], + }; + const fetchMock = vi + .fn() + .mockImplementationOnce(async (input) => { + const requestUrl = input instanceof Request ? input.url : String(input); + expect(new URL(requestUrl, "http://localhost").pathname).toBe("/model/deprecations"); + return Response.json(deprecatedResponse); + }) + .mockImplementationOnce(async (input) => { + const requestUrl = input instanceof Request ? input.url : String(input); + expect(new URL(requestUrl, "http://localhost").pathname).toBe("/model/deprecations"); + return Response.json(currentResponse); + }); + vi.stubGlobal("fetch", fetchMock); + + const user = userEvent.setup({ delay: null }); + renderWithProviders(); + + expect( + await screen.findByText("gpt-4 (openai) was deprecated on 2026-06-01, migrate to gpt-5"), + ).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Refresh models" })); + await waitForElementToBeRemoved(() => screen.queryByRole("alert")); + expect(fetchMock).toHaveBeenCalledTimes(2); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx index 3cfed7297c8..be8ee1bb8c4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx @@ -29,6 +29,9 @@ vi.mock("./detailNavigation", () => ({ vi.mock("@/components/molecules/DecisionModelsBanner", () => ({ default: () =>
, })); +vi.mock("./components/ModelDeprecationBanner", () => ({ + default: () =>
, +})); vi.mock("@/components/model_info_view", () => ({ default: ({ modelId }: { modelId: string }) =>
model:{modelId}
, })); @@ -109,6 +112,12 @@ describe("ModelsAndEndpointsPage", () => { expect(screen.queryByRole("region", { name: "Decision models" })).not.toBeInTheDocument(); }); + it("renders the model deprecation banner", () => { + renderPage(); + + expect(screen.getByRole("region", { name: "Model deprecations" })).toBeInTheDocument(); + }); + it("renders the model detail overlay from the ?model drill-in and hides the tabs", () => { detailState.modelId = "abc-123"; renderPage(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx index e0e5db8a8fd..b4082a14ffe 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx @@ -7,6 +7,7 @@ import { useQueryClient } from "@tanstack/react-query"; import type { ParseKeys } from "i18next"; import { useTranslation } from "react-i18next"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { invalidateModelQueries } from "@/app/(dashboard)/hooks/models/invalidateModelQueries"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import { useUrlTab } from "@/hooks/useUrlTab"; @@ -30,6 +31,7 @@ import AccessGroupBudgetsPanel from "@/app/(dashboard)/models-and-endpoints/pane import PriceDataPanel from "@/app/(dashboard)/models-and-endpoints/panels/PriceDataPanel"; import { Button } from "@/components/ui/button"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import ModelDeprecationBanner from "./components/ModelDeprecationBanner"; const MODEL_TAB_SLUGS = [ "add", @@ -166,10 +168,10 @@ export default function ModelsAndEndpointsPage() { const handleRefreshClick = () => { setLastRefreshed(new Date().toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" })); - queryClient.invalidateQueries({ queryKey: ["models", "list"] }); + void invalidateModelQueries(queryClient); }; - const invalidateModels = () => queryClient.invalidateQueries({ queryKey: ["models", "list"] }); + const invalidateModels = () => invalidateModelQueries(queryClient); if (teamId) { return ( @@ -204,6 +206,7 @@ export default function ModelsAndEndpointsPage() {
{!isViewOnly && } + {modelId ? ( { +const setup = async (searchParams?: string, onUrlUpdate?: OnUrlUpdateFunction, showBanner = false) => { const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); - renderWithProviders(, { searchParams, onUrlUpdate }); + renderWithProviders( + showBanner ? ( + <> + + + + ) : ( + + ), + { searchParams, onUrlUpdate }, + ); await screen.findByText("Provider"); const openAdvanced = async () => { @@ -186,11 +200,69 @@ const setup = async (searchParams?: string, onUrlUpdate?: OnUrlUpdateFunction) = describe("AddModelPanel submit payload contract", () => { beforeEach(() => { vi.clearAllMocks(); + testQueryClient.clear(); mockUseModelCostMap.mockReturnValue({ data: {} }); mockPtuEnabled.mockReturnValue(false); mockAuthorized.mockReturnValue(PROXY_ADMIN); }); + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("refreshes deprecation banner data after adding a model", async () => { + const deprecatedResponse = { + checked_at: "2026-06-01T00:00:00Z", + warn_within_days: 30, + deprecated: [ + { + model_name: "old-model", + litellm_provider: null, + deprecation_date: "2026-06-01", + days_until_deprecation: -1, + status: "deprecated", + successor_model: null, + }, + ], + imminent: [], + upcoming: [], + }; + const currentResponse = { + checked_at: "2026-06-02T00:00:00Z", + warn_within_days: 30, + deprecated: [], + imminent: [ + { + model_name: "current-model", + litellm_provider: null, + deprecation_date: "2026-11-30", + days_until_deprecation: 26, + status: "imminent", + successor_model: null, + }, + ], + upcoming: [], + }; + const fetchMock = vi + .fn() + .mockResolvedValueOnce(Response.json(deprecatedResponse)) + .mockResolvedValueOnce(Response.json(currentResponse)); + vi.stubGlobal("fetch", fetchMock); + + const { fillRequired, submit } = await setup(undefined, undefined, true); + + expect(await screen.findByText("old-model was deprecated on 2026-06-01")).toBeInTheDocument(); + await fillRequired(); + await submit(); + expect(await screen.findByText("current-model deprecates on 2026-11-30, in 26 days")).toBeInTheDocument(); + expect(fetchMock).toHaveBeenCalledTimes(2); + expect( + fetchMock.mock.calls.map( + ([input]) => new URL(input instanceof Request ? input.url : String(input), "http://localhost").pathname, + ), + ).toEqual(["/model/deprecations", "/model/deprecations"]); + }); + it("lists catalog models in the model picker, not entries registered at runtime for deployments", async () => { mockUseModelCostMap.mockImplementation((_enabled: boolean, catalogOnly: boolean) => ({ data: catalogOnly diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx index 392a3c3940a..13bec9506e4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx @@ -15,6 +15,7 @@ import { Providers, getPlaceholder, getProviderModels } from "@/components/provi import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; import { useCredentials } from "@/app/(dashboard)/hooks/credentials/useCredentials"; import { useProviderFields } from "@/app/(dashboard)/hooks/providers/useProviderFields"; +import { invalidateModelQueries } from "@/app/(dashboard)/hooks/models/invalidateModelQueries"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { buildDecisionCatalog } from "@/lib/decisionModels"; @@ -65,7 +66,7 @@ export default function AddModelPanel() { const decisionCatalog = useMemo(() => buildDecisionCatalog(modelCostMapData), [modelCostMapData]); const [showAdvancedSettings, setShowAdvancedSettings] = useState(false); - const refresh = () => queryClient.invalidateQueries({ queryKey: ["models", "list"] }); + const refresh = () => invalidateModelQueries(queryClient); const mountedValues = () => projectMountedValues(registry, form.getValues); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx index 2e27a258a84..f4ff11494da 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx @@ -1,7 +1,9 @@ import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; +import { QueryClient, QueryClientProvider, useQuery } from "@tanstack/react-query"; import userEvent from "@testing-library/user-event"; -import { useState } from "react"; +import { useState, type ReactNode } from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { modelDeprecationKeys } from "@/app/(dashboard)/hooks/models/useModelDeprecations"; import AgentBuilderView from "./AgentBuilderView"; import type { AgentModel } from "../../llm_calls/fetch_agents"; @@ -67,10 +69,41 @@ const props = { userRole: "Admin", }; +const modelListKey = ["models", "list"] as const; + +const ModelQueryProbe = ({ + modelListQuery, + modelDeprecationQuery, +}: { + modelListQuery: () => Promise; + modelDeprecationQuery: () => Promise; +}) => { + useQuery({ queryKey: modelListKey, queryFn: modelListQuery, staleTime: Infinity }); + useQuery({ queryKey: modelDeprecationKeys.all, queryFn: modelDeprecationQuery, staleTime: Infinity }); + return null; +}; + const controlUnder = (label: string): HTMLElement => within(screen.getByText(label).parentElement!).getByRole("combobox"); -const renderView = () => render(); +const renderWithQueryClient = (children: ReactNode, queryClient = new QueryClient()) => { + return render({children}); +}; + +const renderView = (queryClient?: QueryClient) => renderWithQueryClient(, queryClient); + +const renderViewWithModelQueryProbe = ( + queryClient: QueryClient, + modelListQuery: () => Promise, + modelDeprecationQuery: () => Promise, +) => + renderWithQueryClient( + <> + + + , + queryClient, + ); const waitForRoster = async () => { await screen.findByRole("button", { name: "support-agent litellm_agent" }); @@ -99,7 +132,7 @@ beforeEach(() => { describe("AgentBuilderView", () => { it("asks the visitor to sign in when there is no session", () => { - render(); + renderWithQueryClient(); expect(screen.getByText("Sign in to use Agent Builder.")).toBeInTheDocument(); expect(fetchAvailableAgentModels).not.toHaveBeenCalled(); @@ -170,6 +203,25 @@ describe("AgentBuilderView", () => { expect(payload.litellm_params.model).toBe("litellm_agent/gpt-4o"); }); + it("refetches model queries after creating an agent", async () => { + const user = userEvent.setup(); + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const modelListQuery = vi.fn().mockResolvedValue([]); + const modelDeprecationQuery = vi.fn().mockResolvedValue({}); + renderViewWithModelQueryProbe(queryClient, modelListQuery, modelDeprecationQuery); + await waitForRoster(); + await waitFor(() => expect(modelListQuery).toHaveBeenCalledTimes(1)); + await waitFor(() => expect(modelDeprecationQuery).toHaveBeenCalledTimes(1)); + + await user.click(screen.getByRole("button", { name: /New agent/i })); + fireEvent.change(screen.getByPlaceholderText("My Agent"), { target: { value: "billing-agent" } }); + await user.click(screen.getByRole("button", { name: /Save Agent/i })); + + await waitFor(() => expect(modelCreateCall).toHaveBeenCalled()); + await waitFor(() => expect(modelListQuery).toHaveBeenCalledTimes(2)); + await waitFor(() => expect(modelDeprecationQuery).toHaveBeenCalledTimes(2)); + }); + it("will not save a draft without a name", async () => { const user = userEvent.setup(); renderView(); @@ -193,6 +245,24 @@ describe("AgentBuilderView", () => { expect(modelPatchUpdateCall.mock.calls[0][2]).toBe("agent-1"); }); + it("refetches model queries after updating an agent", async () => { + const user = userEvent.setup(); + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const modelListQuery = vi.fn().mockResolvedValue([]); + const modelDeprecationQuery = vi.fn().mockResolvedValue({}); + renderViewWithModelQueryProbe(queryClient, modelListQuery, modelDeprecationQuery); + await waitForRoster(); + await screen.findByDisplayValue("support-agent"); + await waitFor(() => expect(modelListQuery).toHaveBeenCalledTimes(1)); + await waitFor(() => expect(modelDeprecationQuery).toHaveBeenCalledTimes(1)); + + await user.click(screen.getByRole("button", { name: /Update Agent/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + await waitFor(() => expect(modelListQuery).toHaveBeenCalledTimes(2)); + await waitFor(() => expect(modelDeprecationQuery).toHaveBeenCalledTimes(2)); + }); + it("deletes only after the warning is confirmed", async () => { const user = userEvent.setup(); renderView(); @@ -210,6 +280,27 @@ describe("AgentBuilderView", () => { await waitFor(() => expect(modelDeleteCall).toHaveBeenCalledWith("sk-access", "agent-1")); }); + it("refetches model queries after deleting an agent", async () => { + const user = userEvent.setup(); + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const modelListQuery = vi.fn().mockResolvedValue([]); + const modelDeprecationQuery = vi.fn().mockResolvedValue({}); + renderViewWithModelQueryProbe(queryClient, modelListQuery, modelDeprecationQuery); + await waitForRoster(); + await screen.findByDisplayValue("support-agent"); + await waitFor(() => expect(modelListQuery).toHaveBeenCalledTimes(1)); + await waitFor(() => expect(modelDeprecationQuery).toHaveBeenCalledTimes(1)); + + await user.click(screen.getAllByRole("button", { name: /Delete$/ })[0]); + await screen.findByText(/Are you sure you want to delete "support-agent"/); + const deleteButtons = screen.getAllByRole("button", { name: /Delete$/ }); + await user.click(deleteButtons[deleteButtons.length - 1]); + + await waitFor(() => expect(modelDeleteCall).toHaveBeenCalledWith("sk-access", "agent-1")); + await waitFor(() => expect(modelListQuery).toHaveBeenCalledTimes(2)); + await waitFor(() => expect(modelDeprecationQuery).toHaveBeenCalledTimes(2)); + }); + it("abandons the delete when the warning is dismissed", async () => { const user = userEvent.setup(); renderView(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx index f9909462ee9..c79c539cac9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx @@ -3,6 +3,8 @@ import { Bot, FlaskConical, Link as LinkIcon, MessageSquare, Plus, Save, Trash2 } from "lucide-react"; import React, { useCallback, useEffect, useState } from "react"; import { Trans, useTranslation } from "react-i18next"; +import { useQueryClient } from "@tanstack/react-query"; +import { invalidateModelQueries } from "@/app/(dashboard)/hooks/models/invalidateModelQueries"; import { AlertDialog, AlertDialogAction, @@ -203,6 +205,7 @@ export default function AgentBuilderView({ customProxyBaseUrl, }: AgentBuilderViewProps) { const { t } = useTranslation(); + const queryClient = useQueryClient(); const [agentModels, setAgentModels] = useState([]); const [modelGroups, setModelGroups] = useState([]); const [loadingAgents, setLoadingAgents] = useState(true); @@ -358,6 +361,7 @@ export default function AgentBuilderView({ }, model_info: {}, }); + void invalidateModelQueries(queryClient); // /model/new returns the row with `model_id` at the top level. // Prefer that id over name-matching so we land on the just-created // agent even when its public name collides with another team's. @@ -397,6 +401,7 @@ export default function AgentBuilderView({ }, selectedAgentModelId, ); + void invalidateModelQueries(queryClient); toast.success(t("playground.chatUi.agentBuilderView.updated")); const list = await loadAgents(); const stillSelected = list.find((a) => getAgentModelId(a) === selectedAgentModelId); @@ -442,6 +447,7 @@ export default function AgentBuilderView({ setDeleting(true); try { await modelDeleteCall(accessToken, selectedAgentModelId); + void invalidateModelQueries(queryClient); toast.success(t("playground.chatUi.agentBuilderView.deleted")); const list = await loadAgents(); const remaining = list.filter((a) => getAgentModelId(a) !== selectedAgentModelId); diff --git a/ui/litellm-dashboard/src/components/model_info_view.integration.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.integration.test.tsx index a42db06d251..82221826a68 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.integration.test.tsx @@ -1,9 +1,10 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import React, { ReactNode } from "react"; -import { beforeEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import ModelInfoView from "./model_info_view"; +import ModelDeprecationBanner from "@/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner"; import { toast } from "@/lib/toast"; import * as networking from "./networking"; import { useUserLookup } from "@/app/(dashboard)/hooks/users/useUsers"; @@ -45,6 +46,10 @@ vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ useModelCostMap: (...args: any[]) => mockUseModelCostMap(...args), })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ accessToken: "test-token", userRole: "Admin" }), +})); + const mockUseTeams = vi.fn(); vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ useTeams: () => mockUseTeams(), @@ -77,6 +82,8 @@ const mockModelDeleteCall = vi.mocked(networking.modelDeleteCall); const mockCredentialCreateCall = vi.mocked(networking.credentialCreateCall); const mockVectorStoreListCall = vi.mocked(networking.vectorStoreListCall); +const TEST_TODAY = new Date(2026, 10, 4); + describe("ModelInfoView", () => { let queryClient: QueryClient; @@ -195,6 +202,10 @@ describe("ModelInfoView", () => { mockCredentialCreateCall.mockResolvedValue({}); }); + afterEach(() => { + vi.unstubAllGlobals(); + }); + const wrapper = ({ children }: { children: ReactNode }) => React.createElement(QueryClientProvider, { client: queryClient }, children); @@ -205,6 +216,70 @@ describe("ModelInfoView", () => { }); }); + it("refreshes deprecation banner data after updating model credentials", async () => { + const deprecatedResponse = { + checked_at: "2026-06-01T00:00:00Z", + warn_within_days: 30, + deprecated: [ + { + model_name: "old-model", + litellm_provider: null, + deprecation_date: "2026-06-01", + days_until_deprecation: -1, + status: "deprecated", + successor_model: null, + }, + ], + imminent: [], + upcoming: [], + }; + const currentResponse = { + checked_at: "2026-06-02T00:00:00Z", + warn_within_days: 30, + deprecated: [], + imminent: [ + { + model_name: "current-model", + litellm_provider: null, + deprecation_date: "2026-11-30", + days_until_deprecation: 26, + status: "imminent", + successor_model: null, + }, + ], + upcoming: [], + }; + const fetchMock = vi + .fn() + .mockResolvedValueOnce(Response.json(deprecatedResponse)) + .mockResolvedValueOnce(Response.json(currentResponse)); + vi.stubGlobal("fetch", fetchMock); + + render( + <> + + + , + { wrapper }, + ); + + expect(await screen.findByText("old-model was deprecated on 2026-06-01")).toBeInTheDocument(); + await userEvent.click(screen.getByTestId("update-api-key-button")); + const dialog = screen.getByRole("dialog"); + fireEvent.change(within(dialog).getByPlaceholderText("Enter the new API key"), { + target: { value: "new-api-key" }, + }); + await userEvent.click(within(dialog).getByRole("button", { name: "Update API Key" })); + + expect(await screen.findByText("current-model deprecates on 2026-11-30, in 26 days")).toBeInTheDocument(); + await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + expect( + fetchMock.mock.calls.map( + ([input]) => new URL(input instanceof Request ? input.url : String(input), "http://localhost").pathname, + ), + ).toEqual(["/model/deprecations", "/model/deprecations"]); + }); + it("shows Kubernetes pods for discovery-enabled models", async () => { mockUseModelsInfo.mockReturnValue({ data: { diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index efa7caf550d..e0498fd1ea9 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -1,5 +1,6 @@ import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; import { useModelHub, useModelsInfo } from "@/app/(dashboard)/hooks/models/useModels"; +import { invalidateModelQueries } from "@/app/(dashboard)/hooks/models/invalidateModelQueries"; import { useQueryClient } from "@tanstack/react-query"; import { transformModelData } from "@/app/(dashboard)/models-and-endpoints/utils/modelDataTransformer"; import { KeyIcon, RefreshIcon, TrashIcon } from "@heroicons/react/outline"; @@ -824,7 +825,7 @@ export default function ModelInfoView({ accessToken={accessToken} modelId={modelId} onUpdated={() => { - queryClient.invalidateQueries({ queryKey: ["models", "list"] }); + void invalidateModelQueries(queryClient); }} /> )} diff --git a/ui/litellm-dashboard/src/components/settings.test.tsx b/ui/litellm-dashboard/src/components/settings.test.tsx index 280bffd8fe9..9b4145c47ad 100644 --- a/ui/litellm-dashboard/src/components/settings.test.tsx +++ b/ui/litellm-dashboard/src/components/settings.test.tsx @@ -2,6 +2,7 @@ import { act, fireEvent, render, screen, waitFor, within } from "@testing-librar import userEvent from "@testing-library/user-event"; import { FormProvider, useForm } from "react-hook-form"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../tests/test-utils"; import { alertingSettingsCall, getCallbackConfigsCall, getCallbacksCall, setCallbacksCall } from "./networking"; import Settings, { backendCallbackLogoSrc, CallbackSelector } from "./settings"; @@ -77,15 +78,21 @@ describe("Settings", () => { }); it("should render the logging callbacks tab when access token is provided", async () => { - render(); + renderWithProviders(); await waitFor(() => { expect(screen.getByText("Active Logging Callbacks")).toBeInTheDocument(); }); }); + it("opens the tab selected by the URL", async () => { + renderWithProviders(, { searchParams: { tab: "alerting-types" } }); + + expect(await screen.findByRole("tab", { name: "Alerting Types" })).toHaveAttribute("aria-selected", "true"); + }); + it("should display additional settings tabs", async () => { - render(); + renderWithProviders(); await waitFor(() => { expect(screen.getByText("CloudZero Cost Tracking")).toBeInTheDocument(); @@ -96,7 +103,7 @@ describe("Settings", () => { }); it("should load callback configs from the backend when access token is provided", async () => { - render(); + renderWithProviders(); await waitFor(() => { expect(mockGetCallbackConfigsCall).toHaveBeenCalledWith(defaultProps.accessToken); @@ -140,7 +147,7 @@ describe("Settings", () => { ]); const user = userEvent.setup(); - render(); + renderWithProviders(); await waitFor(() => { expect(screen.getByText("Active Logging Callbacks")).toBeInTheDocument(); @@ -250,7 +257,7 @@ describe("Settings", () => { const openOtelEditModal = async () => { const user = userEvent.setup(); - render(); + renderWithProviders(); await user.click(await screen.findByTestId("callback-actions-otel-success")); await user.click(await screen.findByTestId("callback-action-edit")); return user; @@ -340,7 +347,7 @@ describe("Settings", () => { const openS3EditModal = async (callbackName = "s3") => { const user = userEvent.setup(); - render(); + renderWithProviders(); await user.click(await screen.findByTestId(`callback-actions-${callbackName}-success`)); await user.click(await screen.findByTestId("callback-action-edit")); return user; @@ -492,7 +499,7 @@ describe("Settings", () => { it("should send the typed webhook url for an alert type when the alerting tab is saved", async () => { const user = userEvent.setup(); - render(); + renderWithProviders(); await user.click(await screen.findByRole("tab", { name: "Alerting Types" })); @@ -525,7 +532,7 @@ describe("Settings", () => { }), ); - render(); + renderWithProviders(); expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); @@ -540,7 +547,7 @@ describe("Settings", () => { }); it("should resolve loading without fetching when the user id is missing", async () => { - render(); + renderWithProviders(); await waitFor(() => { expect(screen.queryByTestId("skeleton-row")).not.toBeInTheDocument(); @@ -550,7 +557,7 @@ describe("Settings", () => { }); it("should display CloudZero Cost Tracking tab", async () => { - render(); + renderWithProviders(); await waitFor(() => { expect(screen.getByText("Active Logging Callbacks")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index 4153de19499..d2d65b9f526 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -24,6 +24,7 @@ import EmailSettings from "./email_settings"; import MSTeamsSettings from "./MSTeamsSettings"; import { Logo } from "@/components/molecules/logo/Logo"; import { toast } from "@/lib/toast"; +import { useUrlTab } from "@/hooks/useUrlTab"; import AlertingSettings from "./alerting/alerting_settings"; import CloudZeroCostTracking from "./CloudZeroCostTracking/CloudZeroCostTracking"; @@ -46,6 +47,15 @@ interface SettingsPageProps { } type CallbackFormValues = Record; +const SETTINGS_TABS = [ + "logging-callbacks", + "cloudzero-cost-tracking", + "alerting-types", + "alerting-settings", + "email-alerts", + "ms-teams-alerts", +] as const; +type SettingsTab = (typeof SETTINGS_TABS)[number]; const assetsLogoFolder = "/ui/assets/logos/"; @@ -313,6 +323,7 @@ const buildCallbackPayload = (formValues: Record, callbackName: str }; const Settings: React.FC = ({ accessToken, userRole, userID, premiumUser }) => { + const [activeTab, setActiveTab] = useUrlTab(SETTINGS_TABS, "logging-callbacks"); const { t } = useTranslation(); const [callbacks, setCallbacks] = useState([]); const [isLoadingCallbacks, setIsLoadingCallbacks] = useState(true); @@ -597,7 +608,7 @@ const Settings: React.FC = ({ accessToken, userRole, userID, return (
- + setActiveTab(value as SettingsTab)}> {t("loggingAndAlerts.page.tabs.loggingCallbacks")} {t("loggingAndAlerts.page.tabs.cloudZero")} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e65a30f6e1d..a2073bb00d8 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -39049,6 +39049,11 @@ export interface components { * @description The model to migrate to: model_info.successor_model on the deployment when set, else the successor_model the cost map names for the deprecated model. */ successor_model?: string | null; + /** + * Team Public Model Name + * @description The model name visible to the team, when configured. + */ + team_public_model_name?: string | null; }; /** ModelDeprecationResponse */ ModelDeprecationResponse: { From 71fc20449521b57dc4731c88a0b295883177847a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 5 Nov 2026 03:26:45 +0000 Subject: [PATCH 08/39] test(routing): expect 503 server_is_overloaded for an upstream overload (#50155) Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/routing/test_prompt_cache_pin_ttl.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/integration/routing/test_prompt_cache_pin_ttl.py b/tests/integration/routing/test_prompt_cache_pin_ttl.py index f6ee96a9206..733f7859ff0 100644 --- a/tests/integration/routing/test_prompt_cache_pin_ttl.py +++ b/tests/integration/routing/test_prompt_cache_pin_ttl.py @@ -9,7 +9,7 @@ from typing import Final import anthropic import openai import pytest -from integration._support.client import eventually, gateway_from_environment, held +from integration._support.client import eventually, gateway_from_environment, held, object_value from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.redis_process import owned_redis from integration._support.wire import Wire, wire_server @@ -56,11 +56,12 @@ from integration.routing._prompt_cache_pin_support import ( vertex_peer, wire_carries_one_hour, ) -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter from redis import Redis pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 300), pytest.mark.shares_worker_state] +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _CLAUDE: Final = "pin-claude" _SLOW: Final = "pin-slow" _MIXED: Final = "pin-mixed" @@ -539,7 +540,10 @@ def test_failed_first_turn_reports_the_error_and_writes_no_pin(rig: _Rig) -> Non "/v1/chat/completions", chat_body(_CLAUDE, _first_turn(new_marker(), f"{new_marker()} {OVERLOAD_TOKEN}", ONE_HOUR_MARK)), ) - assert response.status_code == 500, response.text + assert response.status_code == 503, response.text + error: Final = object_value(_JSON_OBJECT.validate_json(response.text)["error"]) + assert error["type"] == "service_unavailable_error", response.text + assert error["code"] == "server_is_overloaded", response.text assert "overloaded_error" in response.text, response.text healthy, _ = post_chat(rig.owned.gateway, chat_body(_CLAUDE, _first_turn(new_marker(), new_marker(), None))) spend_row_lands(healthy) From 5c60824024a70f89053f1faed9c46c0e0b5e144d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 5 Nov 2026 03:28:30 +0000 Subject: [PATCH 09/39] test(providers): expect 503 server_is_overloaded for an upstream overload (#50154) Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...nthropic_tool_search_result_matrix_wire.py | 20 ++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/tests/integration/providers/test_anthropic_tool_search_result_matrix_wire.py b/tests/integration/providers/test_anthropic_tool_search_result_matrix_wire.py index 6879d1f8acd..d4c0e2651d8 100644 --- a/tests/integration/providers/test_anthropic_tool_search_result_matrix_wire.py +++ b/tests/integration/providers/test_anthropic_tool_search_result_matrix_wire.py @@ -482,15 +482,20 @@ def test_duplicated_search_result_blocks_rebuild_the_call_once(gateway: Gateway, @pytest.mark.parametrize( - ("status", "kind", "caller_status"), + ("status", "kind", "caller_status", "caller_error_type", "caller_error_code"), ( - pytest.param(400, "invalid_request_error", 400, id="400"), - pytest.param(401, "authentication_error", 401, id="401"), - pytest.param(529, "overloaded_error", 500, id="529"), + pytest.param(400, "invalid_request_error", 400, None, None, id="400"), + pytest.param(401, "authentication_error", 401, None, None, id="401"), + pytest.param(529, "overloaded_error", 503, "service_unavailable_error", "server_is_overloaded", id="529"), ), ) def test_provider_error_on_the_replay_reaches_the_caller( - gateway: Gateway, status: int, kind: str, caller_status: int + gateway: Gateway, + status: int, + kind: str, + caller_status: int, + caller_error_type: str | None, + caller_error_code: str | None, ) -> None: marker: Final = uuid.uuid4().hex @@ -508,6 +513,11 @@ def test_provider_error_on_the_replay_reaches_the_caller( received: Final = wire.drain() assert second.status_code == caller_status, second.text assert f"scripted {kind} {marker}" in second.text, second.text + if caller_error_type is not None: + payload: Final = JSON_OBJECT.validate_json(second.content) + error: Final = JSON_OBJECT.validate_python(payload["error"]) + assert error["type"] == caller_error_type, second.text + assert error["code"] == caller_error_code, second.text assert len(received) >= 2, received assert wire_assistant(received[1]) == list(SEARCH_TURN), received[1].body assert gateway.client.get("/health/liveliness").status_code == 200 From d2b2e763745148ec2d134aca68a4402073714daa Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Thu, 5 Nov 2026 12:03:42 +0800 Subject: [PATCH 10/39] feat(proxy): add model deprecation alerts with daily requests by key (#50022) * feat(ui): show recent model deprecations in notification bell * feat(ui): simplify deprecation request trends and add model identity --- litellm/proxy/alerting_endpoints/__init__.py | 1 + .../model_deprecation_alerts.py | 109 ++++ litellm/proxy/proxy_server.py | 4 + litellm/types/proxy/model_deprecation.py | 38 +- .../test_model_deprecation_alerts.py | 180 +++++++ .../src/app/(dashboard)/alerts/page.tsx | 19 + .../hooks/models/invalidateModelQueries.ts | 9 - .../models/useModelDeprecations.test.tsx | 64 --- .../hooks/models/useModelDeprecations.ts | 33 -- .../app/(dashboard)/hooks/models/useModels.ts | 5 +- .../src/app/(dashboard)/legacyPageRoutes.ts | 1 + .../AllModelsTab.integration.test.tsx | 232 --------- .../components/AllModelsTab.test.tsx | 9 +- .../components/AllModelsTab.tsx | 11 +- .../ModelDeprecationBanner.test.tsx | 477 ------------------ .../components/ModelDeprecationBanner.tsx | 98 ---- ...riceDataManagementTab.integration.test.tsx | 83 --- .../PriceDataManagementTab.test.tsx | 8 +- .../components/PriceDataManagementTab.tsx | 6 +- .../components/deprecationBannerText.test.ts | 140 ----- .../components/deprecationBannerText.ts | 96 ---- .../page.integration.test.tsx | 55 +- .../models-and-endpoints/page.test.tsx | 9 - .../(dashboard)/models-and-endpoints/page.tsx | 7 +- .../panels/AddModelPanel.integration.test.tsx | 78 +-- .../panels/AddModelPanel.tsx | 3 +- .../chat_ui/AgentBuilderView.test.tsx | 97 +--- .../components/chat_ui/AgentBuilderView.tsx | 6 - .../components/overview/overviewData.test.ts | 24 + .../components/overview/overviewData.ts | 7 +- .../ModelDeprecationAlerts/AlertsInbox.tsx | 108 ++++ .../DeprecationAlertItem.tsx | 35 ++ .../DeprecationModel.tsx | 32 ++ .../MigrationDetail.tsx | 335 ++++++++++++ .../ModelAlertsPage.integration.test.tsx | 124 +++++ .../ModelDeprecationsPanel.tsx | 76 +++ .../ModelUsageChart.integration.test.tsx | 103 ++++ .../ModelUsageChart.tsx | 127 +++++ .../RecommendedReplacement.tsx | 7 + .../ModelDeprecationAlerts/UsageTags.test.tsx | 21 + .../ModelDeprecationAlerts/UsageTags.tsx | 40 ++ .../deprecationAlerts.test.ts | 78 +++ .../deprecationAlerts.ts | 61 +++ .../migrationData.test.ts | 43 ++ .../ModelDeprecationAlerts/migrationData.ts | 49 ++ .../ModelDeprecationAlerts/migrationRoute.ts | 14 + .../migrationState.test.ts | 37 ++ .../ModelDeprecationAlerts/migrationState.ts | 17 + .../useDeprecationAlerts.integration.test.tsx | 74 +++ .../useDeprecationAlerts.ts | 51 ++ .../NotificationsBell.test.tsx | 82 +++ .../NotificationsBell/NotificationsBell.tsx | 65 ++- .../src/components/leftnav.test.tsx | 1 + .../src/components/leftnav.tsx | 7 + .../model_info_view.integration.test.tsx | 79 +-- .../src/components/model_info_view.tsx | 3 +- .../src/components/page_metadata.ts | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 190 +++++++ 58 files changed, 2180 insertions(+), 1589 deletions(-) create mode 100644 litellm/proxy/alerting_endpoints/__init__.py create mode 100644 litellm/proxy/alerting_endpoints/model_deprecation_alerts.py create mode 100644 tests/unit/proxy/management_endpoints/test_model_deprecation_alerts.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/alerts/page.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/models/invalidateModelQueries.ts delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.test.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.ts delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.integration.test.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.test.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.integration.test.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.test.ts delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.ts create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/AlertsInbox.tsx create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/DeprecationAlertItem.tsx create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/DeprecationModel.tsx create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/MigrationDetail.tsx create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelAlertsPage.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelDeprecationsPanel.tsx create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelUsageChart.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelUsageChart.tsx create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/RecommendedReplacement.tsx create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/UsageTags.test.tsx create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/UsageTags.tsx create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/deprecationAlerts.test.ts create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/deprecationAlerts.ts create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationData.test.ts create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationData.ts create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationRoute.ts create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationState.test.ts create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationState.ts create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/useDeprecationAlerts.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/components/ModelDeprecationAlerts/useDeprecationAlerts.ts diff --git a/litellm/proxy/alerting_endpoints/__init__.py b/litellm/proxy/alerting_endpoints/__init__.py new file mode 100644 index 00000000000..6a7a204b435 --- /dev/null +++ b/litellm/proxy/alerting_endpoints/__init__.py @@ -0,0 +1 @@ +"""HTTP endpoints for gateway alerts.""" diff --git a/litellm/proxy/alerting_endpoints/model_deprecation_alerts.py b/litellm/proxy/alerting_endpoints/model_deprecation_alerts.py new file mode 100644 index 00000000000..a69dd6ee025 --- /dev/null +++ b/litellm/proxy/alerting_endpoints/model_deprecation_alerts.py @@ -0,0 +1,109 @@ +"""Shared acknowledgement state for model migration alerts.""" + +import hashlib +import json +from datetime import date, datetime, timezone +from typing import Annotated, Final + +from fastapi import APIRouter, Depends, HTTPException, Query +from pydantic import TypeAdapter + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.database import DatabaseDep +from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.daily_activity_buckets import api_key_bucket, api_key_from_bucket +from litellm.repositories.table_repositories import DailyTagSpendRepository +from litellm.types.proxy.auth.user_roles import LitellmUserRoles +from litellm.types.proxy.model_deprecation import DeprecationAlertState, DeprecationAlertUpdate, DeprecationTagUsage + +router: Final = APIRouter() + + +def alert_state_repository(database: DatabaseDep) -> ConfigRepository: + if database is None: + raise HTTPException(status_code=503, detail="Alert state requires a database connection") + return ConfigRepository(database.prisma, use_writer=True) + + +def alert_tag_repository(database: DatabaseDep) -> DailyTagSpendRepository: + if database is None: + raise HTTPException(status_code=503, detail="Model usage tags require a database connection") + return DailyTagSpendRepository(database.prisma) + + +def _state_key(model: str, deprecation_date: date) -> str: + identity: Final = json.dumps([model, deprecation_date.isoformat()]) + digest: Final = hashlib.sha256(identity.encode()).hexdigest() + return f"model_deprecation_alert:{digest}" + + +def _require_admin(user: UserAPIKeyAuth, *, write: bool) -> None: + if user.user_role == LitellmUserRoles.PROXY_ADMIN: + return + if not write and user.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY: + return + raise HTTPException(status_code=403, detail="Only gateway admins can manage model deprecation alerts") + + +@router.get("/alerting/model-deprecations/state", tags=["alerting"], response_model=DeprecationAlertState) +async def get_deprecation_alert_state( + model: Annotated[str, Query(min_length=1, max_length=1024)], + deprecation_date: date, + user: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + repository: Annotated[ConfigRepository, Depends(alert_state_repository)], +) -> DeprecationAlertState: + _require_admin(user, write=False) + record: Final = await repository.get_param(_state_key(model, deprecation_date)) + if record is None: + return DeprecationAlertState(model=model, deprecation_date=deprecation_date) + return DeprecationAlertState.model_validate(record.param_value) + + +@router.put("/alerting/model-deprecations/state", tags=["alerting"], response_model=DeprecationAlertState) +async def update_deprecation_alert_state( + data: DeprecationAlertUpdate, + user: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + repository: Annotated[ConfigRepository, Depends(alert_state_repository)], +) -> DeprecationAlertState: + _require_admin(user, write=True) + state: Final = DeprecationAlertState( + **data.model_dump(), updated_at=datetime.now(timezone.utc), updated_by=user.user_id + ) + await repository.set_param(_state_key(data.model, data.deprecation_date), state.model_dump(mode="json")) + return state + + +@router.get("/alerting/model-deprecations/tags", tags=["alerting"], response_model=tuple[DeprecationTagUsage, ...]) +async def get_deprecation_usage_tags( + model: Annotated[str, Query(min_length=1, max_length=1024)], + api_keys: Annotated[tuple[str, ...], Query(min_length=1, max_length=100)], + start_date: date, + end_date: date, + user: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + repository: Annotated[DailyTagSpendRepository, Depends(alert_tag_repository)], +) -> tuple[DeprecationTagUsage, ...]: + """Distinct tags observed with this model for a page of affected reporting key IDs.""" + _require_admin(user, write=False) + if not 0 <= (end_date - start_date).days < 30: + raise HTTPException(status_code=422, detail="Select up to 30 consecutive calendar dates") + rows: Final = await repository.table.group_by( + by=["api_key", "tag"], + where={ + "api_key": { + # bounded-ok: query validation limits api_keys to 100 IDs. + "in": [api_key_from_bucket(key) for key in api_keys] + }, + "model": model, + "date": {"gte": start_date.isoformat(), "lte": end_date.isoformat()}, + "api_requests": {"gt": 0}, + "tag": {"not": None}, + }, + ) + tags: Final = TypeAdapter(tuple[DeprecationTagUsage, ...]).validate_python(rows) + return tuple( + sorted( + (DeprecationTagUsage(api_key=api_key_bucket(usage.api_key), tag=usage.tag) for usage in tags), + key=lambda usage: (usage.api_key, usage.tag), + ) + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fddecb3d6eb..9a80761b477 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -424,6 +424,9 @@ from litellm.messages.dispatch import messages_target from litellm.proxy._lazy_features import attach_lazy_features, reserve_lazy_slot from litellm.proxy._types import * from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling +from litellm.proxy.alerting_endpoints.model_deprecation_alerts import ( + router as model_deprecation_alerts_router, +) from litellm.proxy.analytics_endpoints.analytics_endpoints import ( router as analytics_router, ) @@ -22434,6 +22437,7 @@ app.include_router(prompt_caching_requests_router) app.include_router(router_settings_router) app.include_router(fallback_management_router) app.include_router(cache_settings_router) +app.include_router(model_deprecation_alerts_router) app.include_router(coordination_redis_settings_router) app.include_router(user_agent_analytics_router) app.include_router(gateway_request_router) diff --git a/litellm/types/proxy/model_deprecation.py b/litellm/types/proxy/model_deprecation.py index 7e459904d59..af21cc770de 100644 --- a/litellm/types/proxy/model_deprecation.py +++ b/litellm/types/proxy/model_deprecation.py @@ -1,9 +1,10 @@ from __future__ import annotations from datetime import date, datetime +from enum import Enum from typing import Final, Literal -from pydantic import Field +from pydantic import ConfigDict, Field from litellm.types.llms.base import LiteLLMBaseModel @@ -63,3 +64,38 @@ class ModelDeprecationResponse(LiteLLMBaseModel): ) warn_within_days: int = Field(description="The window (in days) used to bucket 'imminent' models.") checked_at: datetime = Field(description="UTC timestamp when the deprecation snapshot was generated.") + + +class DeprecationAlertStatus(str, Enum): + OPEN = "open" + RESOLVED = "resolved" + DISMISSED = "dismissed" + + +class DeprecationAlertUpdate(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + model: str = Field(min_length=1, max_length=1024) + deprecation_date: date + status: DeprecationAlertStatus + last_seen_date: date + last_seen_requests: int = Field(ge=0) + + +class DeprecationAlertState(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + model: str + deprecation_date: date + status: DeprecationAlertStatus = DeprecationAlertStatus.OPEN + last_seen_date: date | None = None + last_seen_requests: int = 0 + updated_at: datetime | None = None + updated_by: str | None = None + + +class DeprecationTagUsage(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + api_key: str + tag: str diff --git a/tests/unit/proxy/management_endpoints/test_model_deprecation_alerts.py b/tests/unit/proxy/management_endpoints/test_model_deprecation_alerts.py new file mode 100644 index 00000000000..434dcc7a83c --- /dev/null +++ b/tests/unit/proxy/management_endpoints/test_model_deprecation_alerts.py @@ -0,0 +1,180 @@ +from collections.abc import Iterator + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.alerting_endpoints.model_deprecation_alerts import alert_state_repository, router +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.config_repository import ConfigParam, ConfigRepository +from litellm.types.proxy.auth.user_roles import LitellmUserRoles + + +class MemoryConfig(ConfigRepository): + def __init__(self) -> None: + self.rows: dict[str, object] = {} + + async def get_param(self, param_name: str) -> ConfigParam | None: + value = self.rows.get(param_name) + return None if value is None else ConfigParam(param_name, value) + + async def set_param(self, param_name: str, param_value: object) -> ConfigParam: + self.rows[param_name] = param_value + return ConfigParam(param_name, param_value) + + +@pytest.fixture() +def client() -> Iterator[tuple[TestClient, UserAPIKeyAuth]]: + app = FastAPI() + app.include_router(router) + user = UserAPIKeyAuth(user_id="admin-one", user_role=LitellmUserRoles.PROXY_ADMIN) + repository = MemoryConfig() + app.dependency_overrides[user_api_key_auth] = lambda: user + app.dependency_overrides[alert_state_repository] = lambda: repository + with TestClient(app) as http: + yield http, user + + +IDENTITY = {"model": "anthropic/claude-sonnet-4-5", "deprecation_date": "2026-11-30"} +UPDATE = {**IDENTITY, "status": "resolved", "last_seen_date": "2026-11-04", "last_seen_requests": 123} +PATH = "/alerting/model-deprecations/state" + + +def test_new_alert_is_open(client: tuple[TestClient, UserAPIKeyAuth]) -> None: + http, _ = client + response = http.get(PATH, params=IDENTITY) + assert response.status_code == 200 + assert response.json()["status"] == "open" + + +def test_resolution_is_shared_and_keeps_the_usage_snapshot(client: tuple[TestClient, UserAPIKeyAuth]) -> None: + http, user = client + response = http.put(PATH, json=UPDATE) + assert response.status_code == 200 + assert response.json()["updated_by"] == "admin-one" + user.user_id = "admin-two" + observed = http.get(PATH, params=IDENTITY).json() + assert observed["status"] == "resolved" + assert observed["last_seen_date"] == "2026-11-04" + assert observed["last_seen_requests"] == 123 + assert observed["updated_by"] == "admin-one" + + +def test_reopen_does_not_change_another_models_resolution(client: tuple[TestClient, UserAPIKeyAuth]) -> None: + http, _ = client + http.put(PATH, json=UPDATE) + other = {**UPDATE, "model": "openai/gpt-4-turbo"} + http.put(PATH, json=other) + assert http.put(PATH, json={**UPDATE, "status": "open"}).json()["status"] == "open" + assert http.get(PATH, params={**IDENTITY, "model": other["model"]}).json()["status"] == "resolved" + + +@pytest.mark.parametrize("method", ["get", "put"]) +def test_internal_users_cannot_read_or_change_global_alerts( + client: tuple[TestClient, UserAPIKeyAuth], method: str +) -> None: + http, user = client + user.user_role = LitellmUserRoles.INTERNAL_USER + response = http.get(PATH, params=IDENTITY) if method == "get" else http.put(PATH, json=UPDATE) + assert response.status_code == 403 + + +def test_viewer_can_read_but_cannot_resolve_at_the_same_path(client: tuple[TestClient, UserAPIKeyAuth]) -> None: + http, user = client + user.user_role = LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + assert http.get(PATH, params=IDENTITY).status_code == 200 + assert http.put(PATH, json=UPDATE).status_code == 403 + assert http.get(PATH, params=IDENTITY).json()["status"] == "open" + + +@pytest.mark.parametrize( + "field,value", [("status", "unknown"), ("last_seen_requests", -1), ("deprecation_date", "invalid"), ("model", "")] +) +def test_invalid_updates_are_rejected(client: tuple[TestClient, UserAPIKeyAuth], field: str, value: object) -> None: + http, _ = client + assert http.put(PATH, json={**UPDATE, field: value}).status_code == 422 + assert http.get(PATH, params=IDENTITY).json()["status"] == "open" + + +def test_tags_are_scoped_to_model_dates_and_the_visible_keys(client: tuple[TestClient, UserAPIKeyAuth]) -> None: + from types import SimpleNamespace + from unittest.mock import AsyncMock + + from litellm.proxy.alerting_endpoints.model_deprecation_alerts import alert_tag_repository + + http, _ = client + grouped = AsyncMock(return_value=[{"api_key": "reporting-one", "tag": "support"}]) + http.app.dependency_overrides[alert_tag_repository] = lambda: SimpleNamespace( + table=SimpleNamespace(group_by=grouped) + ) + response = http.get( + "/alerting/model-deprecations/tags", + params={ + "model": "openai/gpt-4-turbo", + "api_keys": ["reporting-one", "reporting-two"], + "start_date": "2026-10-06", + "end_date": "2026-11-04", + }, + ) + assert response.status_code == 200 + assert response.json() == [{"api_key": "reporting-one", "tag": "support"}] + assert grouped.call_args.kwargs["where"] == { + "api_key": {"in": ["reporting-one", "reporting-two"]}, + "model": "openai/gpt-4-turbo", + "date": {"gte": "2026-10-06", "lte": "2026-11-04"}, + "api_requests": {"gt": 0}, + "tag": {"not": None}, + } + + +@pytest.mark.parametrize( + "role,expected", [(LitellmUserRoles.INTERNAL_USER, 403), (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, 200)] +) +def test_tag_access_matches_alert_visibility( + client: tuple[TestClient, UserAPIKeyAuth], role: LitellmUserRoles, expected: int +) -> None: + from types import SimpleNamespace + from unittest.mock import AsyncMock + + from litellm.proxy.alerting_endpoints.model_deprecation_alerts import alert_tag_repository + + http, user = client + user.user_role = role + grouped = AsyncMock(return_value=[]) + http.app.dependency_overrides[alert_tag_repository] = lambda: SimpleNamespace( + table=SimpleNamespace(group_by=grouped) + ) + response = http.get( + "/alerting/model-deprecations/tags", + params={"model": "test-model", "api_keys": ["one"], "start_date": "2026-10-06", "end_date": "2026-11-04"}, + ) + assert response.status_code == expected + if expected == 403: + grouped.assert_not_called() + + +def test_tag_query_rejects_more_than_100_keys_before_database_access(client: tuple[TestClient, UserAPIKeyAuth]) -> None: + from types import SimpleNamespace + from unittest.mock import AsyncMock + + from litellm.proxy.alerting_endpoints.model_deprecation_alerts import ( + alert_tag_repository, + ) + + http, _ = client + grouped = AsyncMock(return_value=[]) + http.app.dependency_overrides[alert_tag_repository] = lambda: SimpleNamespace( + table=SimpleNamespace(group_by=grouped) + ) + response = http.get( + "/alerting/model-deprecations/tags", + params={ + "model": "test-model", + "api_keys": [f"reporting-{index}" for index in range(101)], + "start_date": "2026-10-07", + "end_date": "2026-11-05", + }, + ) + assert response.status_code == 422 + grouped.assert_not_called() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/alerts/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/alerts/page.tsx new file mode 100644 index 00000000000..ec88a69670e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/alerts/page.tsx @@ -0,0 +1,19 @@ +"use client"; + +import { Page, PageTabs, PageTabsList, PageTabsTrigger, PageTabsContent } from "@/components/shared/Page"; +import ModelDeprecationsPanel from "@/components/ModelDeprecationAlerts/ModelDeprecationsPanel"; + +export default function AlertsPage() { + return ( + + + + Model deprecations + + + + + + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/invalidateModelQueries.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/invalidateModelQueries.ts deleted file mode 100644 index 3382d6c9c81..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/invalidateModelQueries.ts +++ /dev/null @@ -1,9 +0,0 @@ -import type { QueryClient } from "@tanstack/react-query"; -import { modelDeprecationKeys } from "./useModelDeprecations"; - -export const invalidateModelQueries = async (queryClient: QueryClient): Promise => { - await Promise.all([ - queryClient.invalidateQueries({ queryKey: ["models", "list"] }), - queryClient.invalidateQueries({ queryKey: modelDeprecationKeys.all }), - ]); -}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.test.tsx deleted file mode 100644 index 88efd721c26..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.test.tsx +++ /dev/null @@ -1,64 +0,0 @@ -/* @vitest-environment jsdom */ -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { act, renderHook, waitFor } from "@testing-library/react"; -import type { ReactNode } from "react"; -import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { useModelDeprecations } from "./useModelDeprecations"; - -const authState = vi.hoisted(() => ({ accessToken: "token", userRole: "internal_user" })); -const emptyDeprecationsResponse = { - checked_at: "2026-11-04T00:00:00Z", - warn_within_days: 30, - deprecated: [], - imminent: [], - upcoming: [], -}; -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ - default: () => ({ accessToken: authState.accessToken, userRole: authState.userRole }), -})); - -describe("useModelDeprecations", () => { - let fetchMock: ReturnType>; - - beforeEach(() => { - authState.accessToken = "token"; - authState.userRole = "internal_user"; - fetchMock = vi.fn().mockResolvedValue(Response.json(emptyDeprecationsResponse)); - vi.stubGlobal("fetch", fetchMock); - }); - - afterEach(() => { - vi.unstubAllGlobals(); - }); - - it("does not fetch deprecations for an internal user", async () => { - const queryClient = new QueryClient(); - const wrapper = ({ children }: { children: ReactNode }) => ( - {children} - ); - - const { result } = renderHook(() => useModelDeprecations(), { wrapper }); - - await waitFor(() => expect(result.current.fetchStatus).toBe("idle")); - await act(async () => { - await new Promise((resolve) => setTimeout(resolve, 0)); - }); - expect(fetchMock).not.toHaveBeenCalled(); - }); - - it("fetches deprecations for a proxy admin viewer", async () => { - authState.userRole = "proxy_admin_viewer"; - const queryClient = new QueryClient(); - const wrapper = ({ children }: { children: ReactNode }) => ( - {children} - ); - - renderHook(() => useModelDeprecations(), { wrapper }); - - await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(1)); - const request = fetchMock.mock.calls[0][0]; - expect(new URL(request instanceof Request ? request.url : String(request), "http://localhost").pathname).toBe( - "/model/deprecations", - ); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.ts deleted file mode 100644 index bb9922d8a8b..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelDeprecations.ts +++ /dev/null @@ -1,33 +0,0 @@ -"use client"; - -import { useQuery } from "@tanstack/react-query"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { backgroundFetchClient } from "@/lib/http/api"; -import type { components } from "@/lib/http/schema"; -import { isProxyAdminTierRole } from "@/utils/roles"; -import { createQueryKeys } from "../common/queryKeysFactory"; - -export type ModelDeprecationInfo = components["schemas"]["ModelDeprecationInfo"]; -export type ModelDeprecationResponse = components["schemas"]["ModelDeprecationResponse"]; - -export const modelDeprecationKeys = createQueryKeys("modelDeprecations"); - -const fetchModelDeprecations = async (): Promise => { - const { data } = await backgroundFetchClient.GET("/model/deprecations"); - if (data === undefined) { - throw new Error("The proxy returned no model deprecations payload"); - } - return data; -}; - -export const useModelDeprecations = () => { - const { accessToken, userRole } = useAuthorized(); - const queryOptions = { - queryKey: modelDeprecationKeys.all, - queryFn: fetchModelDeprecations, - enabled: Boolean(accessToken) && isProxyAdminTierRole(userRole), - staleTime: 5 * 60 * 1000, - retry: false, - }; - return useQuery(queryOptions); -}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index 7e9669e6ce5..f9304bab0d2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -1,7 +1,6 @@ import { useQuery, useInfiniteQuery, useQueryClient, UseQueryResult } from "@tanstack/react-query"; import { useMemo } from "react"; import { createQueryKeys } from "../common/queryKeysFactory"; -import { invalidateModelQueries } from "./invalidateModelQueries"; import { modelInfoCall, modelHubCall, modelAvailableCall } from "@/components/networking"; import useAuthorized from "../useAuthorized"; import { EndpointType, isModeCompatibleWithEndpoint } from "@/components/chat_ui/mode_endpoint_mapping"; @@ -241,7 +240,9 @@ export const useAutoRouters = (): UseQueryResult export const useInvalidateAutoRouters = (): (() => Promise) => { const queryClient = useQueryClient(); - return () => invalidateModelQueries(queryClient); + return async () => { + await queryClient.invalidateQueries({ queryKey: modelKeys.lists() }); + }; }; export const useModelHub = () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts index fe054c7b39a..19dac805948 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts @@ -31,6 +31,7 @@ const LEGACY_PAGE_ROUTES: ReadonlyMap = new Map( "ui-theme": "ui-theme", telemetry: "telemetry", logs: "logs", + alerts: "alerts", lens: "lens", "admin-panel": "admin-panel", "logging-and-alerts": "logging-and-alerts", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.integration.test.tsx deleted file mode 100644 index 42747bfeeca..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.integration.test.tsx +++ /dev/null @@ -1,232 +0,0 @@ -import * as useAuthorizedModule from "@/app/(dashboard)/hooks/useAuthorized"; -import { useQuery } from "@tanstack/react-query"; -import { screen, waitFor } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; -import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; - -import { renderWithProviders, testQueryClient } from "../../../../../tests/test-utils"; -import AllModelsTab from "./AllModelsTab"; -import ModelDeprecationBanner from "./ModelDeprecationBanner"; - -const mockModelDeleteCall = vi.fn().mockResolvedValue({}); -const mockModelPatchUpdateCall = vi.fn().mockResolvedValue({}); -vi.mock("@/components/networking", () => ({ - serverRootPath: "/", - modelDeleteCall: (...args: unknown[]) => mockModelDeleteCall(...args), - modelPatchUpdateCall: (...args: unknown[]) => mockModelPatchUpdateCall(...args), -})); - -vi.mock("@/components/model_dashboard/ModelSettingsModal/ModelSettingsModal", () => ({ - default: function ModelSettingsModalMock({ isVisible }: { isVisible: boolean }) { - return isVisible ?
: null; - }, -})); - -let modelsInfoResult: Record = {}; -const mockRefetch = vi.fn(); -vi.mock("../../hooks/models/useModels", () => ({ - useModelsInfo: () => ({ ...modelsInfoResult, refetch: mockRefetch }), -})); - -vi.mock("../../hooks/models/useModelCostMap", () => ({ - useModelCostMap: () => ({ data: { "gpt-4": { litellm_provider: "openai" } }, isLoading: false, error: null }), -})); - -const mockTeams = [{ team_id: "team-1", team_alias: "Engineering" }]; -const mockCredentials = [ - { credential_name: "openai-prod", display_name: "Prod OpenAI", credential_values: {}, credential_info: {} }, -]; -const mockUseCredentials = vi.hoisted(() => vi.fn()); -vi.mock("../../hooks/credentials/useCredentials", () => ({ - useCredentials: mockUseCredentials, -})); - -vi.mock("../../hooks/teams/useTeams", () => ({ - useTeams: () => ({ data: mockTeams, isLoading: false, error: null, refetch: vi.fn() }), -})); - -const BASE_MODEL_INFO = { - id: "model-1", - db_model: true, - created_by: "user-123", - created_at: "2024-01-01T00:00:00Z", - updated_at: "2024-01-02T00:00:00Z", - team_id: "team-1", - access_groups: [], -}; - -const makeRow = () => ({ - model_name: "gpt-4", - litellm_params: { model: "openai/gpt-4", custom_llm_provider: "openai" }, - model_info: BASE_MODEL_INFO, -}); - -const setModelsInfo = (rows: Record[]) => { - modelsInfoResult = { - data: { data: rows, total_count: rows.length, current_page: 1, total_pages: 1, size: 50 }, - isLoading: false, - isFetching: false, - error: null, - }; -}; - -const MOCK_AUTHORIZED = { - isLoading: false, - isAuthorized: true, - token: "mock-token", - accessToken: "mock-access-token", - userId: "user-123", - userEmail: "test@example.com", - userRole: "Admin", - userRoleLabel: "Admin", - isViewOnly: false, - premiumUser: true, - disabledPersonalKeyCreation: false, - showSSOBanner: false, -}; - -const mockSetSelectedModelGroup = vi.fn(); -const mockSetSelectedModelId = vi.fn(); -const mockSetSelectedTeamId = vi.fn(); - -const defaultProps = { - selectedModelGroup: "all", - setSelectedModelGroup: mockSetSelectedModelGroup, - availableModelGroups: ["gpt-4", "gpt-3.5-turbo"], - availableModelAccessGroups: ["sales-team"], - setSelectedModelId: mockSetSelectedModelId, - setSelectedTeamId: mockSetSelectedTeamId, -}; - -const TEST_TODAY = new Date(2026, 10, 4); - -const deprecatedResponse = { - checked_at: "2026-06-01T00:00:00Z", - warn_within_days: 30, - deprecated: [ - { - model_name: "old-model", - litellm_provider: null, - deprecation_date: "2026-06-01", - days_until_deprecation: -1, - status: "deprecated", - successor_model: null, - }, - ], - imminent: [], - upcoming: [], -}; - -const currentResponse = { - checked_at: "2026-06-02T00:00:00Z", - warn_within_days: 30, - deprecated: [], - imminent: [ - { - model_name: "current-model", - litellm_provider: null, - deprecation_date: "2026-11-30", - days_until_deprecation: 26, - status: "imminent", - successor_model: null, - }, - ], - upcoming: [], -}; - -const stubDeprecationResponses = () => { - const fetchMock = vi - .fn() - .mockImplementationOnce(async (input) => { - const requestUrl = input instanceof Request ? input.url : String(input); - expect(new URL(requestUrl, "http://localhost").pathname).toBe("/model/deprecations"); - return Response.json(deprecatedResponse); - }) - .mockImplementationOnce(async (input) => { - const requestUrl = input instanceof Request ? input.url : String(input); - expect(new URL(requestUrl, "http://localhost").pathname).toBe("/model/deprecations"); - return Response.json(currentResponse); - }); - vi.stubGlobal("fetch", fetchMock); - return fetchMock; -}; - -const renderModelsWithDeprecationBanner = () => - renderWithProviders( - <> - - - , - ); - -const ModelListQueryProbe = ({ queryFn }: { queryFn: () => Promise }) => { - useQuery({ queryKey: ["models", "list"], queryFn, staleTime: Infinity }); - return null; -}; - -describe("AllModelsTab", () => { - beforeEach(() => { - vi.clearAllMocks(); - testQueryClient.clear(); - setModelsInfo([makeRow()]); - vi.spyOn(useAuthorizedModule, "default").mockReturnValue(MOCK_AUTHORIZED); - mockUseCredentials.mockImplementation(({ enabled = true }: { enabled?: boolean } = {}) => ({ - data: enabled ? { credentials: mockCredentials } : undefined, - isLoading: false, - })); - }); - - afterEach(() => { - vi.unstubAllGlobals(); - }); - - it("refreshes deprecation banner data after deleting a model", async () => { - const fetchMock = stubDeprecationResponses(); - const user = userEvent.setup(); - renderModelsWithDeprecationBanner(); - - expect(await screen.findByText("old-model was deprecated on 2026-06-01")).toBeInTheDocument(); - await user.click(await screen.findByTestId("model-delete-model-1")); - await user.click(await screen.findByRole("button", { name: /^delete$/i })); - expect(await screen.findByText("current-model deprecates on 2026-11-30, in 26 days")).toBeInTheDocument(); - await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); - expect(screen.queryByText("old-model was deprecated on 2026-06-01")).not.toBeInTheDocument(); - }); - - it("refreshes deprecation banner data after pausing a model", async () => { - const fetchMock = stubDeprecationResponses(); - const user = userEvent.setup(); - renderModelsWithDeprecationBanner(); - - expect(await screen.findByText("old-model was deprecated on 2026-06-01")).toBeInTheDocument(); - await user.click(await screen.findByTestId("model-pause-toggle-model-1")); - expect(await screen.findByText("current-model deprecates on 2026-11-30, in 26 days")).toBeInTheDocument(); - await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); - expect(screen.queryByText("old-model was deprecated on 2026-06-01")).not.toBeInTheDocument(); - }); - - it("refreshes the table and deprecations once without refetching another models list observer", async () => { - const fetchMock = stubDeprecationResponses(); - const observerQueryFn = vi.fn<() => Promise>().mockResolvedValue("models"); - const user = userEvent.setup(); - renderWithProviders( - <> - - - - , - ); - - expect(await screen.findByText("old-model was deprecated on 2026-06-01")).toBeInTheDocument(); - await waitFor(() => expect(observerQueryFn).toHaveBeenCalledTimes(1)); - mockRefetch.mockClear(); - fetchMock.mockClear(); - - await user.click(screen.getByTestId("datatable-refresh")); - - expect(await screen.findByText("current-model deprecates on 2026-11-30, in 26 days")).toBeInTheDocument(); - await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(1)); - expect(observerQueryFn).toHaveBeenCalledTimes(1); - expect(mockRefetch).toHaveBeenCalledTimes(1); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx index 97e7f0238bb..09e00db9fb0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx @@ -4,7 +4,7 @@ import userEvent from "@testing-library/user-event"; import type { OnUrlUpdateFunction } from "nuqs/adapters/testing"; import { beforeEach, describe, expect, it, Mock, vi } from "vitest"; -import { renderWithProviders, testQueryClient } from "../../../../../tests/test-utils"; +import { renderWithProviders } from "../../../../../tests/test-utils"; import AllModelsTab from "./AllModelsTab"; import { STATUS_COLUMN_ID, toServerSortField } from "./ModelsTableColumns"; @@ -22,6 +22,12 @@ vi.mock("@/components/model_dashboard/ModelSettingsModal/ModelSettingsModal", () }, })); +const mockInvalidateQueries = vi.fn(); +vi.mock("@tanstack/react-query", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, useQueryClient: () => ({ invalidateQueries: mockInvalidateQueries }) }; +}); + interface ModelsInfoArgs { page?: number; size?: number; @@ -151,7 +157,6 @@ const defaultProps = { describe("AllModelsTab", () => { beforeEach(() => { vi.clearAllMocks(); - testQueryClient.clear(); modelsInfoCalls.length = 0; setModelsInfo([makeRow()]); vi.spyOn(useAuthorizedModule, "default").mockReturnValue(MOCK_AUTHORIZED); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index cf977a6d9e3..ff027ac8b3c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -22,8 +22,6 @@ import { useCallback, useMemo, useState } from "react"; import { Trans, useTranslation } from "react-i18next"; import { useModelsInfo } from "../../hooks/models/useModels"; -import { modelDeprecationKeys } from "../../hooks/models/useModelDeprecations"; -import { invalidateModelQueries } from "../../hooks/models/invalidateModelQueries"; import { ROUTING_STATUS_FILTER_ID, routingStatusToBlocked } from "../utils/routingStatus"; import { transformModelData } from "../utils/modelDataTransformer"; import { @@ -268,8 +266,8 @@ const AllModelsTab = ({ try { setDeleteLoading(true); await modelDeleteCall(accessToken, deleteModalModelId); - void invalidateModelQueries(queryClient); toast.success(t("modelsAndEndpoints.allModelsTab.toasts.deleted")); + queryClient.invalidateQueries({ queryKey: ["models", "list"] }); refetchModels(); } catch (error) { console.error("Error deleting model:", error); @@ -293,7 +291,7 @@ const AllModelsTab = ({ ); // invalidateQueries already schedules a refetch for active observers // on this key — no need to also call refetchModels() (would double-fetch). - void invalidateModelQueries(queryClient); + queryClient.invalidateQueries({ queryKey: ["models", "list"] }); } catch (error) { console.error("Error toggling model pause state:", error); toast.fromError(error); @@ -305,9 +303,8 @@ const AllModelsTab = ({ ); const handleRefresh = useCallback(() => { - refetchModels(); - void queryClient.invalidateQueries({ queryKey: modelDeprecationKeys.all }); - }, [queryClient, refetchModels]); + void refetchModels(); + }, [refetchModels]); const handleDeleteClick = useCallback((modelId: string) => { setDeleteModalModelId(modelId); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.test.tsx deleted file mode 100644 index c14bf93c116..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.test.tsx +++ /dev/null @@ -1,477 +0,0 @@ -/* @vitest-environment jsdom */ -import { render, screen, within } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; -import { beforeEach, describe, expect, it, vi } from "vitest"; -import { - useModelDeprecations, - type ModelDeprecationInfo, - type ModelDeprecationResponse, -} from "@/app/(dashboard)/hooks/models/useModelDeprecations"; -import { uiHref } from "@/utils/uiHref"; -import ModelDeprecationBanner from "./ModelDeprecationBanner"; - -vi.mock("@/app/(dashboard)/hooks/models/useModelDeprecations", () => ({ - useModelDeprecations: vi.fn(), -})); - -const mockedUseModelDeprecations = vi.mocked(useModelDeprecations); - -const makeResponse = (overrides: Partial = {}): ModelDeprecationResponse => ({ - checked_at: "2026-11-04T00:00:00Z", - warn_within_days: 30, - deprecated: [], - imminent: [], - upcoming: [], - ...overrides, -}); - -const makeInfo = (overrides: Partial = {}): ModelDeprecationInfo => ({ - model_name: "gpt-4o", - deprecation_date: "2026-11-30", - days_until_deprecation: 26, - status: "imminent", - ...overrides, -}); - -const queryResult = ( - data: ModelDeprecationResponse | undefined, - isLoading = false, - isError = false, - overrides: Partial> = {}, -): ReturnType => - ({ data, isLoading, isError, ...overrides }) as ReturnType; - -const pinnedToday = (): Date => new Date(2026, 10, 4); - -const renderBanner = (today = pinnedToday()) => render(); - -describe("ModelDeprecationBanner", () => { - beforeEach(() => { - vi.clearAllMocks(); - window.localStorage.clear(); - }); - - it("renders nothing when there are no deprecated or imminent models", () => { - mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse())); - - renderBanner(); - - expect(screen.queryByRole("alert")).not.toBeInTheDocument(); - }); - - it("renders nothing when only upcoming models are present", () => { - mockedUseModelDeprecations.mockReturnValue( - queryResult( - makeResponse({ - upcoming: [makeInfo({ deprecation_date: "2026-12-05", status: "upcoming" })], - }), - ), - ); - - renderBanner(); - - expect(screen.queryByRole("alert")).not.toBeInTheDocument(); - }); - - it("includes server-upcoming models at the local warning-window boundary", () => { - const upcomingInfoOverrides: Partial = { - model_name: "due-in-window", - deprecation_date: "2026-12-04", - days_until_deprecation: 31, - status: "upcoming", - }; - const info = makeInfo(upcomingInfoOverrides); - mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ upcoming: [info] }))); - - renderBanner(); - - expect(screen.getByText("1 model deprecates within 30 days")).toBeInTheDocument(); - expect(within(screen.getByRole("list")).getByRole("listitem")).toHaveTextContent( - "due-in-window deprecates on 2026-12-04, in 30 days", - ); - }); - - it("does not show server-upcoming models one day past the local warning window", () => { - const upcomingInfoOverrides: Partial = { - model_name: "due-outside-window", - deprecation_date: "2026-12-05", - days_until_deprecation: 31, - status: "upcoming", - }; - const info = makeInfo(upcomingInfoOverrides); - mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ upcoming: [info] }))); - - renderBanner(); - - expect(screen.queryByRole("alert")).not.toBeInTheDocument(); - }); - - it("renders nothing while loading or after an error", () => { - mockedUseModelDeprecations.mockReturnValue(queryResult(undefined, true)); - const { rerender } = renderBanner(); - expect(screen.queryByRole("alert")).not.toBeInTheDocument(); - - mockedUseModelDeprecations.mockReturnValue(queryResult(undefined, false, true)); - rerender(); - - expect(screen.queryByRole("alert")).not.toBeInTheDocument(); - }); - - it("keeps the banner visible when a refetch fails", () => { - const staleInfoOverrides: Partial = { - model_name: "stale-model", - deprecation_date: "2026-06-01", - days_until_deprecation: -156, - status: "deprecated", - }; - const info = makeInfo(staleInfoOverrides); - const refetchError = new Error("Refetch failed"); - const refetchErrorResult: Partial> = { - error: refetchError, - failureCount: 1, - failureReason: refetchError, - fetchStatus: "idle", - isFetching: false, - isLoadingError: false, - isRefetchError: true, - status: "error", - }; - mockedUseModelDeprecations.mockReturnValue( - queryResult(makeResponse({ deprecated: [info] }), false, true, refetchErrorResult), - ); - - renderBanner(); - - expect(screen.getByRole("alert")).toBeInTheDocument(); - expect(within(screen.getByRole("list")).getByRole("listitem")).toHaveTextContent( - "stale-model was deprecated on 2026-06-01", - ); - }); - - it("shows deprecated and imminent entries with a link to Logging & Alerts", () => { - const deprecatedInfo: ModelDeprecationInfo = { - model_name: "gpt-4", - team_public_model_name: null, - litellm_provider: "openai", - deprecation_date: "2026-06-01", - days_until_deprecation: -152, - status: "deprecated", - successor_model: "gpt-5.6-sol", - }; - const imminentInfo: ModelDeprecationInfo = { - model_name: "gpt-4o", - team_public_model_name: null, - litellm_provider: "openai", - deprecation_date: "2026-11-30", - days_until_deprecation: 26, - status: "imminent", - successor_model: "X", - }; - const response = makeResponse({ - deprecated: [deprecatedInfo], - imminent: [imminentInfo], - }); - mockedUseModelDeprecations.mockReturnValue(queryResult(response)); - - renderBanner(); - - expect(screen.getByText("1 model is deprecated, 1 deprecates within 30 days")).toBeInTheDocument(); - const entries = within(screen.getByRole("list")).getAllByRole("listitem"); - expect(entries[0]).toHaveTextContent("gpt-4 (openai) was deprecated on 2026-06-01, migrate to gpt-5.6-sol"); - expect(entries[1]).toHaveTextContent("gpt-4o (openai) deprecates on 2026-11-30, in 26 days, migrate to X"); - expect(screen.getByRole("link", { name: "Get Slack alerts for deprecations" })).toHaveAttribute( - "href", - `${uiHref("logging-and-alerts")}?tab=alerting-types`, - ); - }); - - it("treats a server-deprecated entry dated today as imminent locally", () => { - const infoOverrides: Partial = { - model_name: "due-today", - deprecation_date: "2026-11-04", - days_until_deprecation: -1, - status: "deprecated", - }; - const info = makeInfo(infoOverrides); - mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ deprecated: [info] }))); - - renderBanner(); - - expect(screen.getByText("1 model deprecates within 30 days")).toBeInTheDocument(); - expect(within(screen.getByRole("list")).getByRole("listitem")).toHaveTextContent( - "due-today deprecates today (2026-11-04)", - ); - }); - - it("treats a server-imminent entry dated yesterday as deprecated locally", () => { - const infoOverrides: Partial = { - model_name: "due-yesterday", - deprecation_date: "2026-11-03", - days_until_deprecation: 0, - status: "imminent", - }; - const info = makeInfo(infoOverrides); - mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ imminent: [info] }))); - - renderBanner(); - - expect(screen.getByText("1 model is deprecated")).toBeInTheDocument(); - expect(within(screen.getByRole("list")).getByRole("listitem")).toHaveTextContent( - "due-yesterday was deprecated on 2026-11-03", - ); - }); - - it("uses the local date difference instead of the server day count", () => { - const infoOverrides: Partial = { - model_name: "due-in-two", - deprecation_date: "2026-11-06", - days_until_deprecation: 1, - status: "imminent", - }; - const info = makeInfo(infoOverrides); - mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ imminent: [info] }))); - - renderBanner(); - - expect(within(screen.getByRole("list")).getByRole("listitem")).toHaveTextContent( - "due-in-two deprecates on 2026-11-06, in 2 days", - ); - }); - - it("deduplicates entries by their displayed name, provider, and deprecation date", () => { - const globalInfoOverrides: Partial = { - model_name: "o4-mini", - team_public_model_name: null, - litellm_provider: "openai", - deprecation_date: "2026-10-23", - status: "deprecated", - }; - const firstTeamInfoOverrides: Partial = { - model_name: "o4-mini-team-a", - team_public_model_name: "o4-mini", - litellm_provider: "openai", - deprecation_date: "2026-10-23", - status: "deprecated", - }; - const secondTeamInfoOverrides: Partial = { - model_name: "o4-mini-team-b", - team_public_model_name: "o4-mini", - litellm_provider: "openai", - deprecation_date: "2026-10-23", - status: "deprecated", - }; - mockedUseModelDeprecations.mockReturnValue( - queryResult( - makeResponse({ - deprecated: [ - makeInfo(globalInfoOverrides), - makeInfo(firstTeamInfoOverrides), - makeInfo(secondTeamInfoOverrides), - ], - }), - ), - ); - - renderBanner(); - - expect(screen.getByText("1 model is deprecated")).toBeInTheDocument(); - const entries = within(screen.getByRole("list")).getAllByRole("listitem"); - expect(entries).toHaveLength(1); - expect(entries[0]).toHaveTextContent("o4-mini (openai) was deprecated on 2026-10-23"); - }); - - it("keeps entries with different successors and renders each migration target", () => { - const firstEntryOverrides: Partial = { - model_name: "team-model-a", - team_public_model_name: "shared-model", - litellm_provider: "openai", - deprecation_date: "2026-10-23", - days_until_deprecation: -1, - status: "deprecated", - successor_model: "successor-a", - }; - const secondEntryOverrides: Partial = { - model_name: "shared-model", - litellm_provider: "openai", - deprecation_date: "2026-10-23", - days_until_deprecation: -1, - status: "deprecated", - successor_model: "successor-b", - }; - const firstEntry = makeInfo(firstEntryOverrides); - const secondEntry = makeInfo(secondEntryOverrides); - mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ deprecated: [firstEntry, secondEntry] }))); - - renderBanner(); - - const entries = within(screen.getByRole("list")).getAllByRole("listitem"); - expect(screen.getByText("1 model is deprecated")).toBeInTheDocument(); - expect(entries).toHaveLength(2); - expect(entries[0]).toHaveTextContent("shared-model (openai) was deprecated on 2026-10-23, migrate to successor-a"); - expect(entries[1]).toHaveTextContent("shared-model (openai) was deprecated on 2026-10-23, migrate to successor-b"); - }); - - it("keeps entries with different deprecation dates", () => { - const firstEntryOverrides: Partial = { - model_name: "shared-model", - litellm_provider: "openai", - deprecation_date: "2026-10-22", - days_until_deprecation: -2, - status: "deprecated", - successor_model: "successor", - }; - const secondEntryOverrides: Partial = { - model_name: "shared-model", - litellm_provider: "openai", - deprecation_date: "2026-10-23", - days_until_deprecation: -1, - status: "deprecated", - successor_model: "successor", - }; - const firstEntry = makeInfo(firstEntryOverrides); - const secondEntry = makeInfo(secondEntryOverrides); - mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ deprecated: [firstEntry, secondEntry] }))); - - renderBanner(); - - const entries = within(screen.getByRole("list")).getAllByRole("listitem"); - expect(entries).toHaveLength(2); - expect(entries[0]).toHaveTextContent("shared-model (openai) was deprecated on 2026-10-22"); - expect(entries[1]).toHaveTextContent("shared-model (openai) was deprecated on 2026-10-23"); - }); - - it("dismisses the banner and keeps it dismissed after remounting the same entries", async () => { - const info = makeInfo(); - const secondInfo = makeInfo({ model_name: "claude-sonnet-4" }); - mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ imminent: [info, secondInfo] }))); - const user = userEvent.setup(); - const { unmount } = renderBanner(); - - await user.click(screen.getByRole("button", { name: "Dismiss deprecation warning" })); - expect(screen.queryByRole("alert")).not.toBeInTheDocument(); - unmount(); - - const changedSecondInfoOverrides: Partial = { - ...secondInfo, - days_until_deprecation: secondInfo.days_until_deprecation - 1, - }; - const changedInfoOverrides: Partial = { - ...info, - days_until_deprecation: info.days_until_deprecation - 1, - }; - const teamAliasInfoOverrides: Partial = { - ...info, - model_name: "team-internal-name", - team_public_model_name: info.model_name, - days_until_deprecation: info.days_until_deprecation - 1, - }; - mockedUseModelDeprecations.mockReturnValue( - queryResult( - makeResponse({ - imminent: [ - makeInfo(changedSecondInfoOverrides), - makeInfo(changedInfoOverrides), - makeInfo(teamAliasInfoOverrides), - ], - }), - ), - ); - renderBanner(); - - expect(screen.queryByRole("alert")).not.toBeInTheDocument(); - }); - - it("shows the banner again when a new entry is added after dismissal", async () => { - const firstEntryOverrides: Partial = { - model_name: "gpt-4", - deprecation_date: "2026-10-23", - days_until_deprecation: -12, - status: "deprecated", - }; - const firstEntry = makeInfo(firstEntryOverrides); - const addedEntryOverrides: Partial = { - model_name: "gpt-3.5", - deprecation_date: "2026-10-23", - days_until_deprecation: -12, - status: "deprecated", - }; - const addedEntry = makeInfo(addedEntryOverrides); - mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ deprecated: [firstEntry] }))); - const user = userEvent.setup(); - const { rerender } = renderBanner(); - - await user.click(screen.getByRole("button", { name: "Dismiss deprecation warning" })); - mockedUseModelDeprecations.mockReturnValue( - queryResult( - makeResponse({ - deprecated: [firstEntry, addedEntry], - }), - ), - ); - rerender(); - - expect(screen.getByRole("alert")).toBeInTheDocument(); - expect(screen.getByText("2 models are deprecated")).toBeInTheDocument(); - }); - - it("shows the banner again when an imminent entry becomes deprecated", async () => { - const imminentEntryOverrides: Partial = { - model_name: "gpt-4o", - deprecation_date: "2026-11-30", - days_until_deprecation: 0, - status: "imminent", - }; - const imminentEntry = makeInfo(imminentEntryOverrides); - mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ imminent: [imminentEntry] }))); - const user = userEvent.setup(); - const { rerender } = renderBanner(new Date(2026, 10, 30)); - - await user.click(screen.getByRole("button", { name: "Dismiss deprecation warning" })); - rerender(); - - expect(screen.getByRole("alert")).toBeInTheDocument(); - expect(screen.getByText("1 model is deprecated")).toBeInTheDocument(); - }); - - it("shows the banner again when an entry's successor changes after dismissal", async () => { - const entry = makeInfo({ successor_model: "successor-a" }); - mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ imminent: [entry] }))); - const user = userEvent.setup(); - const { rerender } = renderBanner(); - - await user.click(screen.getByRole("button", { name: "Dismiss deprecation warning" })); - expect(screen.queryByRole("alert")).not.toBeInTheDocument(); - - mockedUseModelDeprecations.mockReturnValue( - queryResult(makeResponse({ imminent: [makeInfo({ ...entry, successor_model: "successor-b" })] })), - ); - rerender(); - - expect(screen.getByRole("alert")).toBeInTheDocument(); - expect(screen.getByText(/migrate to successor-b/)).toBeInTheDocument(); - }); - - it("expands all entries and collapses back to the first five", async () => { - const imminent = [ - ...Array.from({ length: 7 }, (_value, index) => makeInfo({ model_name: `model-${index + 1}` })), - makeInfo({ model_name: "team-model-1", team_public_model_name: "model-1" }), - ]; - mockedUseModelDeprecations.mockReturnValue(queryResult(makeResponse({ imminent }))); - const user = userEvent.setup(); - - renderBanner(); - - const entries = within(screen.getByRole("list")).getAllByRole("listitem"); - expect(entries).toHaveLength(5); - expect(entries[4]).toHaveTextContent("model-5 deprecates"); - expect(screen.queryByText(/model-6/)).not.toBeInTheDocument(); - - await user.click(screen.getByRole("button", { name: "Show all 7" })); - expect(within(screen.getByRole("list")).getAllByRole("listitem")).toHaveLength(7); - expect(screen.getByText(/model-6 deprecates/)).toBeInTheDocument(); - - await user.click(screen.getByRole("button", { name: "Show fewer" })); - expect(within(screen.getByRole("list")).getAllByRole("listitem")).toHaveLength(5); - expect(screen.queryByText(/model-6/)).not.toBeInTheDocument(); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.tsx deleted file mode 100644 index 3a89940f786..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner.tsx +++ /dev/null @@ -1,98 +0,0 @@ -"use client"; - -import { TriangleAlert, X } from "lucide-react"; -import { useState } from "react"; -import { useModelDeprecations } from "@/app/(dashboard)/hooks/models/useModelDeprecations"; -import { Alert, AlertAction, AlertDescription, AlertTitle } from "@/components/ui/alert"; -import { Button } from "@/components/ui/button"; -import { getLocalStorageItem, setLocalStorageItem } from "@/utils/localStorageUtils"; -import { uiHref } from "@/utils/uiHref"; -import { - bucketByLocalDate, - countDistinctModels, - dedupeDeprecationEntries, - deprecationBannerFingerprint, - deprecationBannerTitle, - deprecationEntryText, -} from "./deprecationBannerText"; - -const STORAGE_KEY = "modelDeprecationBannerDismissed"; - -const ModelDeprecationBanner = ({ today = new Date() }: { today?: Date }) => { - const { data } = useModelDeprecations(); - const [isExpanded, setIsExpanded] = useState(false); - const [dismissedFingerprint, setDismissedFingerprint] = useState(() => - getLocalStorageItem(STORAGE_KEY), - ); - - if (!data) { - return null; - } - - const localBuckets = bucketByLocalDate( - [...(data.deprecated ?? []), ...(data.imminent ?? []), ...(data.upcoming ?? [])], - today, - data.warn_within_days, - ); - const deprecated = dedupeDeprecationEntries(localBuckets.deprecated); - const imminent = dedupeDeprecationEntries(localBuckets.imminent); - - if (deprecated.length === 0 && imminent.length === 0) { - return null; - } - - const entries = [...deprecated, ...imminent]; - const fingerprint = deprecationBannerFingerprint(entries); - if (dismissedFingerprint === fingerprint) { - return null; - } - - const visibleEntries = isExpanded ? entries : entries.slice(0, 5); - - return ( - 0 ? "destructive" : "default"} className="mb-4"> - - - - - - {deprecationBannerTitle(countDistinctModels(deprecated), countDistinctModels(imminent), data.warn_within_days)} - - -
    - {visibleEntries.map((info, index) => ( -
  • {deprecationEntryText(info)}
    • - ))} -
-
- {entries.length > 5 && ( - - )} - - Get Slack alerts for deprecations - -
- - - ); -}; - -export default ModelDeprecationBanner; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.integration.test.tsx deleted file mode 100644 index d2ddb805a0b..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.integration.test.tsx +++ /dev/null @@ -1,83 +0,0 @@ -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; -import { afterEach, describe, expect, it, vi } from "vitest"; - -import ModelDeprecationBanner from "./ModelDeprecationBanner"; -import PriceDataManagementTab from "./PriceDataManagementTab"; - -const TEST_TODAY = new Date(2026, 10, 4); - -vi.mock("@/components/price_data_reload", () => ({ - default: ({ onReloadSuccess }: { onReloadSuccess: () => void }) => , -})); -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ - default: () => ({ accessToken: "sk-test", userRole: "Admin" }), -})); - -const renderTabWithBanner = (queryClient: QueryClient) => - render( - - <> - - - - , - ); - -describe("PriceDataManagementTab", () => { - afterEach(() => { - vi.unstubAllGlobals(); - }); - - it("refetches deprecations after reloading price data", async () => { - const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); - const deprecatedResponse = { - checked_at: "2026-06-01T00:00:00Z", - warn_within_days: 30, - deprecated: [ - { - model_name: "old-model", - litellm_provider: null, - deprecation_date: "2026-06-01", - days_until_deprecation: -1, - status: "deprecated", - successor_model: null, - }, - ], - imminent: [], - upcoming: [], - }; - const currentResponse = { - checked_at: "2026-06-02T00:00:00Z", - warn_within_days: 30, - deprecated: [], - imminent: [ - { - model_name: "current-model", - litellm_provider: null, - deprecation_date: "2026-11-30", - days_until_deprecation: 26, - status: "imminent", - successor_model: null, - }, - ], - upcoming: [], - }; - const fetchMock = vi - .fn() - .mockResolvedValueOnce(Response.json(deprecatedResponse)) - .mockResolvedValueOnce(Response.json(currentResponse)); - vi.stubGlobal("fetch", fetchMock); - renderTabWithBanner(queryClient); - - expect(await screen.findByText("old-model was deprecated on 2026-06-01")).toBeInTheDocument(); - fireEvent.click(screen.getByRole("button", { name: "reload" })); - expect(await screen.findByText("current-model deprecates on 2026-11-30, in 26 days")).toBeInTheDocument(); - await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); - expect( - fetchMock.mock.calls.map( - ([input]) => new URL(input instanceof Request ? input.url : String(input), "http://localhost").pathname, - ), - ).toEqual(["/model/deprecations", "/model/deprecations"]); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx index 445dc505ce6..b2beb5beeed 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx @@ -2,16 +2,13 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { fireEvent, render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; -import { modelDeprecationKeys } from "../../hooks/models/useModelDeprecations"; import { modelCostMapKeys } from "../../hooks/models/useModelCostMap"; import PriceDataManagementTab from "./PriceDataManagementTab"; vi.mock("@/components/price_data_reload", () => ({ default: ({ onReloadSuccess }: { onReloadSuccess: () => void }) => , })); -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ - default: () => ({ accessToken: "sk-test", userRole: "Admin" }), -})); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => ({ accessToken: "sk-test" }) })); const renderTab = (queryClient: QueryClient) => render( @@ -30,16 +27,13 @@ describe("PriceDataManagementTab", () => { const queryClient = new QueryClient(); const liveKey = modelCostMapKeys.list({}); const catalogKey = modelCostMapKeys.list({ filters: { catalog_only: "true" } }); - const deprecationKey = modelDeprecationKeys.all; queryClient.setQueryData(liveKey, {}); queryClient.setQueryData(catalogKey, {}); - queryClient.setQueryData(deprecationKey, {}); renderTab(queryClient); fireEvent.click(screen.getByRole("button", { name: "reload" })); expect(queryClient.getQueryState(liveKey)?.isInvalidated).toBe(true); expect(queryClient.getQueryState(catalogKey)?.isInvalidated).toBe(true); - expect(queryClient.getQueryState(deprecationKey)?.isInvalidated).toBe(true); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx index 4dddc5333ea..1457b5b311f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx @@ -3,7 +3,6 @@ import React from "react"; import { useTranslation } from "react-i18next"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useQueryClient } from "@tanstack/react-query"; -import { modelDeprecationKeys } from "../../hooks/models/useModelDeprecations"; import { modelCostMapKeys } from "../../hooks/models/useModelCostMap"; const PriceDataManagementTab = () => { @@ -21,10 +20,7 @@ const PriceDataManagementTab = () => { { - void Promise.all([ - queryClient.invalidateQueries({ queryKey: modelCostMapKeys.all }), - queryClient.invalidateQueries({ queryKey: modelDeprecationKeys.all }), - ]); + queryClient.invalidateQueries({ queryKey: modelCostMapKeys.all }); }} buttonText={t("modelsAndEndpoints.priceDataManagementTab.reloadButton")} size="middle" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.test.ts deleted file mode 100644 index 66d065e6caa..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.test.ts +++ /dev/null @@ -1,140 +0,0 @@ -import { describe, expect, it } from "vitest"; -import type { ModelDeprecationInfo } from "@/app/(dashboard)/hooks/models/useModelDeprecations"; -import { - bucketByLocalDate, - countDistinctModels, - deprecationBannerTitle, - deprecationEntryText, - localDaysUntil, - parseLocalDate, -} from "./deprecationBannerText"; - -const makeInfo = (overrides: Partial = {}): ModelDeprecationInfo => ({ - model_name: "gpt-4o", - deprecation_date: "2026-11-30", - days_until_deprecation: 26, - status: "imminent", - ...overrides, -}); - -describe("local deprecation dates", () => { - it("parses date parts as a local calendar date and counts from local midnight", () => { - const parsed = parseLocalDate("2026-11-06"); - - expect([parsed.getFullYear(), parsed.getMonth(), parsed.getDate(), parsed.getHours()]).toEqual([2026, 10, 6, 0]); - expect(localDaysUntil("2026-11-06", new Date(2026, 10, 4, 18))).toBe(2); - }); - - it("buckets entries by their local date and warning window", () => { - const buckets = bucketByLocalDate( - [ - makeInfo({ model_name: "yesterday", deprecation_date: "2026-11-03", status: "imminent" }), - makeInfo({ model_name: "today", deprecation_date: "2026-11-04", status: "deprecated" }), - makeInfo({ model_name: "outside-window", deprecation_date: "2026-12-05", status: "imminent" }), - ], - new Date(2026, 10, 4), - 30, - ); - - expect( - buckets.deprecated.map(({ model_name, days_until_deprecation, status }) => [ - model_name, - days_until_deprecation, - status, - ]), - ).toEqual([["yesterday", -1, "deprecated"]]); - expect( - buckets.imminent.map(({ model_name, days_until_deprecation, status }) => [ - model_name, - days_until_deprecation, - status, - ]), - ).toEqual([["today", 0, "imminent"]]); - }); -}); - -describe("countDistinctModels", () => { - it("counts visible names and providers, not deployment aliases", () => { - expect( - countDistinctModels([ - makeInfo({ model_name: "team-a", team_public_model_name: "shared", litellm_provider: "openai" }), - makeInfo({ model_name: "team-b", team_public_model_name: "shared", litellm_provider: "openai" }), - makeInfo({ model_name: "team-c", team_public_model_name: "shared", litellm_provider: "anthropic" }), - ]), - ).toBe(2); - }); -}); - -describe("deprecationBannerTitle", () => { - it("formats both counts with singular and plural wording", () => { - expect(deprecationBannerTitle(2, 1, 30)).toBe("2 models are deprecated, 1 deprecates within 30 days"); - expect(deprecationBannerTitle(1, 2, 30)).toBe("1 model is deprecated, 2 deprecate within 30 days"); - }); - - it("omits empty counts and uses singular model wording", () => { - expect(deprecationBannerTitle(1, 0, 30)).toBe("1 model is deprecated"); - expect(deprecationBannerTitle(0, 1, 30)).toBe("1 model deprecates within 30 days"); - expect(deprecationBannerTitle(0, 3, 30)).toBe("3 models deprecate within 30 days"); - expect(deprecationBannerTitle(0, 0, 30)).toBe(""); - }); -}); - -describe("deprecationEntryText", () => { - it("formats deprecated entries with provider and successor", () => { - const deprecatedInfo: ModelDeprecationInfo = { - model_name: "gpt-4", - litellm_provider: "openai", - deprecation_date: "2026-06-01", - days_until_deprecation: -152, - status: "deprecated", - successor_model: "gpt-5.6-sol", - }; - - expect(deprecationEntryText(deprecatedInfo)).toBe( - "gpt-4 (openai) was deprecated on 2026-06-01, migrate to gpt-5.6-sol", - ); - }); - - it("omits a null provider and successor", () => { - expect( - deprecationEntryText( - makeInfo({ - litellm_provider: null, - successor_model: null, - }), - ), - ).toBe("gpt-4o deprecates on 2026-11-30, in 26 days"); - }); - - it("uses singular wording for one day", () => { - expect( - deprecationEntryText( - makeInfo({ - days_until_deprecation: 1, - }), - ), - ).toBe("gpt-4o deprecates on 2026-11-30, in 1 day"); - }); - - it("says an imminent model deprecates today", () => { - expect( - deprecationEntryText( - makeInfo({ - days_until_deprecation: 0, - deprecation_date: "2026-11-04", - }), - ), - ).toBe("gpt-4o deprecates today (2026-11-04)"); - }); - - it("uses the team-visible model name when present", () => { - expect( - deprecationEntryText( - makeInfo({ - model_name: "internal-name", - team_public_model_name: "team-name", - }), - ), - ).toBe("team-name deprecates on 2026-11-30, in 26 days"); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.ts deleted file mode 100644 index 0c035c5bdd6..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/deprecationBannerText.ts +++ /dev/null @@ -1,96 +0,0 @@ -import type { ModelDeprecationInfo } from "@/app/(dashboard)/hooks/models/useModelDeprecations"; - -const shownModelName = (info: ModelDeprecationInfo): string => info.team_public_model_name ?? info.model_name; - -export const countDistinctModels = (infos: readonly ModelDeprecationInfo[]): number => - new Set(infos.map((info) => JSON.stringify([shownModelName(info), info.litellm_provider ?? null]))).size; - -export const parseLocalDate = (value: string): Date => { - const [year, month, day] = value.split("-"); - return new Date(Number(year), Number(month) - 1, Number(day)); -}; - -const startOfLocalDay = (today: Date): Date => new Date(today.getFullYear(), today.getMonth(), today.getDate()); - -export const localDaysUntil = (deprecationDate: string, today: Date): number => - Math.round((parseLocalDate(deprecationDate).getTime() - startOfLocalDay(today).getTime()) / 86_400_000); - -export const bucketByLocalDate = ( - infos: readonly ModelDeprecationInfo[], - today: Date, - warnWithinDays: number, -): { deprecated: ModelDeprecationInfo[]; imminent: ModelDeprecationInfo[] } => { - const inWindow = infos - .map((info): ModelDeprecationInfo => { - const daysUntilDeprecation = localDaysUntil(info.deprecation_date, today); - return { - ...info, - days_until_deprecation: daysUntilDeprecation, - status: daysUntilDeprecation < 0 ? "deprecated" : "imminent", - }; - }) - .filter((info) => info.days_until_deprecation <= warnWithinDays); - - return { - deprecated: inWindow.filter((info) => info.status === "deprecated"), - imminent: inWindow.filter((info) => info.status === "imminent"), - }; -}; - -const hasSameDedupeKey = (candidate: ModelDeprecationInfo, info: ModelDeprecationInfo): boolean => { - const sameDisplayedModel = shownModelName(candidate) === shownModelName(info); - const sameProvider = (candidate.litellm_provider ?? null) === (info.litellm_provider ?? null); - const sameDeprecationDate = candidate.deprecation_date === info.deprecation_date; - const sameSuccessor = (candidate.successor_model ?? null) === (info.successor_model ?? null); - const sameVisibleIdentity = sameDisplayedModel && sameProvider; - const sameDeprecationIdentity = sameDeprecationDate && sameSuccessor; - return sameVisibleIdentity && sameDeprecationIdentity; -}; - -export const dedupeDeprecationEntries = (infos: readonly ModelDeprecationInfo[]): ModelDeprecationInfo[] => - infos.filter((info, index) => infos.findIndex((candidate) => hasSameDedupeKey(candidate, info)) === index); - -export const deprecationBannerFingerprint = (infos: readonly ModelDeprecationInfo[]): string => - JSON.stringify( - infos - .map((info) => - JSON.stringify([ - shownModelName(info), - info.litellm_provider ?? null, - info.deprecation_date, - info.status, - info.successor_model ?? null, - ]), - ) - .sort(), - ); - -export const deprecationBannerTitle = ( - deprecatedCount: number, - imminentCount: number, - warnWithinDays: number, -): string => { - const deprecatedText = - deprecatedCount === 0 ? "" : `${deprecatedCount} ${deprecatedCount === 1 ? "model is" : "models are"} deprecated`; - const imminentModelLabel = imminentCount === 1 ? "model" : "models"; - const imminentSubject = deprecatedCount > 0 ? `${imminentCount}` : `${imminentCount} ${imminentModelLabel}`; - const imminentVerb = imminentCount === 1 ? "deprecates" : "deprecate"; - const imminentText = imminentCount === 0 ? "" : `${imminentSubject} ${imminentVerb} within ${warnWithinDays} days`; - return [deprecatedText, imminentText].filter((text) => text.length > 0).join(", "); -}; - -export const deprecationEntryText = (info: ModelDeprecationInfo): string => { - const model = `${shownModelName(info)}${info.litellm_provider ? ` (${info.litellm_provider})` : ""}`; - const successor = info.successor_model ? `, migrate to ${info.successor_model}` : ""; - - if (info.status === "deprecated") { - return `${model} was deprecated on ${info.deprecation_date}${successor}`; - } - - if (info.days_until_deprecation === 0) { - return `${model} deprecates today (${info.deprecation_date})${successor}`; - } - - const dayLabel = info.days_until_deprecation === 1 ? "day" : "days"; - return `${model} deprecates on ${info.deprecation_date}, in ${info.days_until_deprecation} ${dayLabel}${successor}`; -}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx index 7aeff6f44f2..4d18966625e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx @@ -1,9 +1,9 @@ /* @vitest-environment jsdom */ -import { screen, waitForElementToBeRemoved } from "@testing-library/react"; +import { screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import * as networking from "@/components/networking"; -import { renderWithProviders, testQueryClient } from "../../../../tests/test-utils"; +import { renderWithProviders } from "../../../../tests/test-utils"; import ModelsAndEndpointsPage from "./page"; vi.mock("./panels/AllModelsPanel", () => ({ default: () =>
})); @@ -201,7 +201,6 @@ const createMockTeamData = (overrides = {}) => ({ describe("ModelsAndEndpointsPage ?team drill-in", () => { beforeEach(() => { - testQueryClient.clear(); authState.userRole = "Admin"; authState.isViewOnly = false; can.mockReturnValue(true); @@ -222,7 +221,6 @@ describe("ModelsAndEndpointsPage ?team drill-in", () => { }); afterEach(() => { - vi.unstubAllGlobals(); vi.clearAllMocks(); }); @@ -246,53 +244,4 @@ describe("ModelsAndEndpointsPage ?team drill-in", () => { expect(screen.queryByRole("button", { name: /edit settings/i })).not.toBeInTheDocument(); expect(screen.queryByText("Disable all global guardrails")).not.toBeInTheDocument(); }); - - it("refreshes deprecation banner data when the models list is refreshed", async () => { - const deprecatedResponse = { - checked_at: "2026-06-01T00:00:00Z", - warn_within_days: 30, - deprecated: [ - { - model_name: "gpt-4", - litellm_provider: "openai", - deprecation_date: "2026-06-01", - days_until_deprecation: -1, - status: "deprecated", - successor_model: "gpt-5", - }, - ], - imminent: [], - upcoming: [], - }; - const currentResponse = { - checked_at: "2026-06-02T00:00:00Z", - warn_within_days: 30, - deprecated: [], - imminent: [], - upcoming: [], - }; - const fetchMock = vi - .fn() - .mockImplementationOnce(async (input) => { - const requestUrl = input instanceof Request ? input.url : String(input); - expect(new URL(requestUrl, "http://localhost").pathname).toBe("/model/deprecations"); - return Response.json(deprecatedResponse); - }) - .mockImplementationOnce(async (input) => { - const requestUrl = input instanceof Request ? input.url : String(input); - expect(new URL(requestUrl, "http://localhost").pathname).toBe("/model/deprecations"); - return Response.json(currentResponse); - }); - vi.stubGlobal("fetch", fetchMock); - - const user = userEvent.setup({ delay: null }); - renderWithProviders(); - - expect( - await screen.findByText("gpt-4 (openai) was deprecated on 2026-06-01, migrate to gpt-5"), - ).toBeInTheDocument(); - await user.click(screen.getByRole("button", { name: "Refresh models" })); - await waitForElementToBeRemoved(() => screen.queryByRole("alert")); - expect(fetchMock).toHaveBeenCalledTimes(2); - }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx index be8ee1bb8c4..3cfed7297c8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx @@ -29,9 +29,6 @@ vi.mock("./detailNavigation", () => ({ vi.mock("@/components/molecules/DecisionModelsBanner", () => ({ default: () =>
, })); -vi.mock("./components/ModelDeprecationBanner", () => ({ - default: () =>
, -})); vi.mock("@/components/model_info_view", () => ({ default: ({ modelId }: { modelId: string }) =>
model:{modelId}
, })); @@ -112,12 +109,6 @@ describe("ModelsAndEndpointsPage", () => { expect(screen.queryByRole("region", { name: "Decision models" })).not.toBeInTheDocument(); }); - it("renders the model deprecation banner", () => { - renderPage(); - - expect(screen.getByRole("region", { name: "Model deprecations" })).toBeInTheDocument(); - }); - it("renders the model detail overlay from the ?model drill-in and hides the tabs", () => { detailState.modelId = "abc-123"; renderPage(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx index b4082a14ffe..e0e5db8a8fd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx @@ -7,7 +7,6 @@ import { useQueryClient } from "@tanstack/react-query"; import type { ParseKeys } from "i18next"; import { useTranslation } from "react-i18next"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { invalidateModelQueries } from "@/app/(dashboard)/hooks/models/invalidateModelQueries"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import { useUrlTab } from "@/hooks/useUrlTab"; @@ -31,7 +30,6 @@ import AccessGroupBudgetsPanel from "@/app/(dashboard)/models-and-endpoints/pane import PriceDataPanel from "@/app/(dashboard)/models-and-endpoints/panels/PriceDataPanel"; import { Button } from "@/components/ui/button"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; -import ModelDeprecationBanner from "./components/ModelDeprecationBanner"; const MODEL_TAB_SLUGS = [ "add", @@ -168,10 +166,10 @@ export default function ModelsAndEndpointsPage() { const handleRefreshClick = () => { setLastRefreshed(new Date().toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" })); - void invalidateModelQueries(queryClient); + queryClient.invalidateQueries({ queryKey: ["models", "list"] }); }; - const invalidateModels = () => invalidateModelQueries(queryClient); + const invalidateModels = () => queryClient.invalidateQueries({ queryKey: ["models", "list"] }); if (teamId) { return ( @@ -206,7 +204,6 @@ export default function ModelsAndEndpointsPage() {
{!isViewOnly && } - {modelId ? ( { +const setup = async (searchParams?: string, onUrlUpdate?: OnUrlUpdateFunction) => { const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); - renderWithProviders( - showBanner ? ( - <> - - - - ) : ( - - ), - { searchParams, onUrlUpdate }, - ); + renderWithProviders(, { searchParams, onUrlUpdate }); await screen.findByText("Provider"); const openAdvanced = async () => { @@ -200,69 +186,11 @@ const setup = async (searchParams?: string, onUrlUpdate?: OnUrlUpdateFunction, s describe("AddModelPanel submit payload contract", () => { beforeEach(() => { vi.clearAllMocks(); - testQueryClient.clear(); mockUseModelCostMap.mockReturnValue({ data: {} }); mockPtuEnabled.mockReturnValue(false); mockAuthorized.mockReturnValue(PROXY_ADMIN); }); - afterEach(() => { - vi.unstubAllGlobals(); - }); - - it("refreshes deprecation banner data after adding a model", async () => { - const deprecatedResponse = { - checked_at: "2026-06-01T00:00:00Z", - warn_within_days: 30, - deprecated: [ - { - model_name: "old-model", - litellm_provider: null, - deprecation_date: "2026-06-01", - days_until_deprecation: -1, - status: "deprecated", - successor_model: null, - }, - ], - imminent: [], - upcoming: [], - }; - const currentResponse = { - checked_at: "2026-06-02T00:00:00Z", - warn_within_days: 30, - deprecated: [], - imminent: [ - { - model_name: "current-model", - litellm_provider: null, - deprecation_date: "2026-11-30", - days_until_deprecation: 26, - status: "imminent", - successor_model: null, - }, - ], - upcoming: [], - }; - const fetchMock = vi - .fn() - .mockResolvedValueOnce(Response.json(deprecatedResponse)) - .mockResolvedValueOnce(Response.json(currentResponse)); - vi.stubGlobal("fetch", fetchMock); - - const { fillRequired, submit } = await setup(undefined, undefined, true); - - expect(await screen.findByText("old-model was deprecated on 2026-06-01")).toBeInTheDocument(); - await fillRequired(); - await submit(); - expect(await screen.findByText("current-model deprecates on 2026-11-30, in 26 days")).toBeInTheDocument(); - expect(fetchMock).toHaveBeenCalledTimes(2); - expect( - fetchMock.mock.calls.map( - ([input]) => new URL(input instanceof Request ? input.url : String(input), "http://localhost").pathname, - ), - ).toEqual(["/model/deprecations", "/model/deprecations"]); - }); - it("lists catalog models in the model picker, not entries registered at runtime for deployments", async () => { mockUseModelCostMap.mockImplementation((_enabled: boolean, catalogOnly: boolean) => ({ data: catalogOnly diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx index 13bec9506e4..392a3c3940a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx @@ -15,7 +15,6 @@ import { Providers, getPlaceholder, getProviderModels } from "@/components/provi import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; import { useCredentials } from "@/app/(dashboard)/hooks/credentials/useCredentials"; import { useProviderFields } from "@/app/(dashboard)/hooks/providers/useProviderFields"; -import { invalidateModelQueries } from "@/app/(dashboard)/hooks/models/invalidateModelQueries"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { buildDecisionCatalog } from "@/lib/decisionModels"; @@ -66,7 +65,7 @@ export default function AddModelPanel() { const decisionCatalog = useMemo(() => buildDecisionCatalog(modelCostMapData), [modelCostMapData]); const [showAdvancedSettings, setShowAdvancedSettings] = useState(false); - const refresh = () => invalidateModelQueries(queryClient); + const refresh = () => queryClient.invalidateQueries({ queryKey: ["models", "list"] }); const mountedValues = () => projectMountedValues(registry, form.getValues); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx index f4ff11494da..2e27a258a84 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx @@ -1,9 +1,7 @@ import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; -import { QueryClient, QueryClientProvider, useQuery } from "@tanstack/react-query"; import userEvent from "@testing-library/user-event"; -import { useState, type ReactNode } from "react"; +import { useState } from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { modelDeprecationKeys } from "@/app/(dashboard)/hooks/models/useModelDeprecations"; import AgentBuilderView from "./AgentBuilderView"; import type { AgentModel } from "../../llm_calls/fetch_agents"; @@ -69,41 +67,10 @@ const props = { userRole: "Admin", }; -const modelListKey = ["models", "list"] as const; - -const ModelQueryProbe = ({ - modelListQuery, - modelDeprecationQuery, -}: { - modelListQuery: () => Promise; - modelDeprecationQuery: () => Promise; -}) => { - useQuery({ queryKey: modelListKey, queryFn: modelListQuery, staleTime: Infinity }); - useQuery({ queryKey: modelDeprecationKeys.all, queryFn: modelDeprecationQuery, staleTime: Infinity }); - return null; -}; - const controlUnder = (label: string): HTMLElement => within(screen.getByText(label).parentElement!).getByRole("combobox"); -const renderWithQueryClient = (children: ReactNode, queryClient = new QueryClient()) => { - return render({children}); -}; - -const renderView = (queryClient?: QueryClient) => renderWithQueryClient(, queryClient); - -const renderViewWithModelQueryProbe = ( - queryClient: QueryClient, - modelListQuery: () => Promise, - modelDeprecationQuery: () => Promise, -) => - renderWithQueryClient( - <> - - - , - queryClient, - ); +const renderView = () => render(); const waitForRoster = async () => { await screen.findByRole("button", { name: "support-agent litellm_agent" }); @@ -132,7 +99,7 @@ beforeEach(() => { describe("AgentBuilderView", () => { it("asks the visitor to sign in when there is no session", () => { - renderWithQueryClient(); + render(); expect(screen.getByText("Sign in to use Agent Builder.")).toBeInTheDocument(); expect(fetchAvailableAgentModels).not.toHaveBeenCalled(); @@ -203,25 +170,6 @@ describe("AgentBuilderView", () => { expect(payload.litellm_params.model).toBe("litellm_agent/gpt-4o"); }); - it("refetches model queries after creating an agent", async () => { - const user = userEvent.setup(); - const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); - const modelListQuery = vi.fn().mockResolvedValue([]); - const modelDeprecationQuery = vi.fn().mockResolvedValue({}); - renderViewWithModelQueryProbe(queryClient, modelListQuery, modelDeprecationQuery); - await waitForRoster(); - await waitFor(() => expect(modelListQuery).toHaveBeenCalledTimes(1)); - await waitFor(() => expect(modelDeprecationQuery).toHaveBeenCalledTimes(1)); - - await user.click(screen.getByRole("button", { name: /New agent/i })); - fireEvent.change(screen.getByPlaceholderText("My Agent"), { target: { value: "billing-agent" } }); - await user.click(screen.getByRole("button", { name: /Save Agent/i })); - - await waitFor(() => expect(modelCreateCall).toHaveBeenCalled()); - await waitFor(() => expect(modelListQuery).toHaveBeenCalledTimes(2)); - await waitFor(() => expect(modelDeprecationQuery).toHaveBeenCalledTimes(2)); - }); - it("will not save a draft without a name", async () => { const user = userEvent.setup(); renderView(); @@ -245,24 +193,6 @@ describe("AgentBuilderView", () => { expect(modelPatchUpdateCall.mock.calls[0][2]).toBe("agent-1"); }); - it("refetches model queries after updating an agent", async () => { - const user = userEvent.setup(); - const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); - const modelListQuery = vi.fn().mockResolvedValue([]); - const modelDeprecationQuery = vi.fn().mockResolvedValue({}); - renderViewWithModelQueryProbe(queryClient, modelListQuery, modelDeprecationQuery); - await waitForRoster(); - await screen.findByDisplayValue("support-agent"); - await waitFor(() => expect(modelListQuery).toHaveBeenCalledTimes(1)); - await waitFor(() => expect(modelDeprecationQuery).toHaveBeenCalledTimes(1)); - - await user.click(screen.getByRole("button", { name: /Update Agent/i })); - - await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); - await waitFor(() => expect(modelListQuery).toHaveBeenCalledTimes(2)); - await waitFor(() => expect(modelDeprecationQuery).toHaveBeenCalledTimes(2)); - }); - it("deletes only after the warning is confirmed", async () => { const user = userEvent.setup(); renderView(); @@ -280,27 +210,6 @@ describe("AgentBuilderView", () => { await waitFor(() => expect(modelDeleteCall).toHaveBeenCalledWith("sk-access", "agent-1")); }); - it("refetches model queries after deleting an agent", async () => { - const user = userEvent.setup(); - const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); - const modelListQuery = vi.fn().mockResolvedValue([]); - const modelDeprecationQuery = vi.fn().mockResolvedValue({}); - renderViewWithModelQueryProbe(queryClient, modelListQuery, modelDeprecationQuery); - await waitForRoster(); - await screen.findByDisplayValue("support-agent"); - await waitFor(() => expect(modelListQuery).toHaveBeenCalledTimes(1)); - await waitFor(() => expect(modelDeprecationQuery).toHaveBeenCalledTimes(1)); - - await user.click(screen.getAllByRole("button", { name: /Delete$/ })[0]); - await screen.findByText(/Are you sure you want to delete "support-agent"/); - const deleteButtons = screen.getAllByRole("button", { name: /Delete$/ }); - await user.click(deleteButtons[deleteButtons.length - 1]); - - await waitFor(() => expect(modelDeleteCall).toHaveBeenCalledWith("sk-access", "agent-1")); - await waitFor(() => expect(modelListQuery).toHaveBeenCalledTimes(2)); - await waitFor(() => expect(modelDeprecationQuery).toHaveBeenCalledTimes(2)); - }); - it("abandons the delete when the warning is dismissed", async () => { const user = userEvent.setup(); renderView(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx index c79c539cac9..f9909462ee9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx @@ -3,8 +3,6 @@ import { Bot, FlaskConical, Link as LinkIcon, MessageSquare, Plus, Save, Trash2 } from "lucide-react"; import React, { useCallback, useEffect, useState } from "react"; import { Trans, useTranslation } from "react-i18next"; -import { useQueryClient } from "@tanstack/react-query"; -import { invalidateModelQueries } from "@/app/(dashboard)/hooks/models/invalidateModelQueries"; import { AlertDialog, AlertDialogAction, @@ -205,7 +203,6 @@ export default function AgentBuilderView({ customProxyBaseUrl, }: AgentBuilderViewProps) { const { t } = useTranslation(); - const queryClient = useQueryClient(); const [agentModels, setAgentModels] = useState([]); const [modelGroups, setModelGroups] = useState([]); const [loadingAgents, setLoadingAgents] = useState(true); @@ -361,7 +358,6 @@ export default function AgentBuilderView({ }, model_info: {}, }); - void invalidateModelQueries(queryClient); // /model/new returns the row with `model_id` at the top level. // Prefer that id over name-matching so we land on the just-created // agent even when its public name collides with another team's. @@ -401,7 +397,6 @@ export default function AgentBuilderView({ }, selectedAgentModelId, ); - void invalidateModelQueries(queryClient); toast.success(t("playground.chatUi.agentBuilderView.updated")); const list = await loadAgents(); const stillSelected = list.find((a) => getAgentModelId(a) === selectedAgentModelId); @@ -447,7 +442,6 @@ export default function AgentBuilderView({ setDeleting(true); try { await modelDeleteCall(accessToken, selectedAgentModelId); - void invalidateModelQueries(queryClient); toast.success(t("playground.chatUi.agentBuilderView.deleted")); const list = await loadAgents(); const remaining = list.filter((a) => getAgentModelId(a) !== selectedAgentModelId); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.test.ts index 1dfdfd1da78..24a0ba0d86a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.test.ts @@ -210,3 +210,27 @@ describe("model colors", () => { expect(bucketTotals(series).get("2026-10-01")).toBe(55); }); }); + +describe("daily key stacks", () => { + it("should preserve each day's total when truncated keys are grouped into Other", () => { + const first = day("2026-11-01", 10, {}); + first.metrics.api_requests = 100; + first.breakdown.api_keys = { + "key-a": { metrics: { ...metrics(3), api_requests: 30 }, metadata: { key_alias: "Support", team_id: null } }, + "key-b": { metrics: { ...metrics(2), api_requests: 20 }, metadata: { key_alias: "Research", team_id: null } }, + }; + const second = day("2026-11-02", 15, {}); + second.metrics.api_requests = 150; + second.breakdown.api_keys = { + "key-a": { metrics: { ...metrics(1), api_requests: 10 }, metadata: { key_alias: "Support", team_id: null } }, + "key-b": { metrics: { ...metrics(8), api_requests: 80 }, metadata: { key_alias: "Research", team_id: null } }, + }; + const result = seriesBy([first, second], "api_keys", "requests", 1); + expect(result.labels).toEqual(["key-b", "Other"]); + expect(result.data.map((row) => [row[result.keys[0]], row[result.keys[1]]])).toEqual([ + [20, 80], + [80, 70], + ]); + expect([...bucketTotals(result).values()]).toEqual([100, 150]); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.ts index 54f4e5dd1de..2943302dfae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/overviewData.ts @@ -21,7 +21,10 @@ interface RankedRow extends Record { const EMPTY_RANKED = { spend: 0, requests: 0, successful_requests: 0, failed_requests: 0, tokens: 0 } as const; /** Sums one breakdown dimension across every day, highest spend first. */ -export const rollUpBreakdown = (results: readonly DailyData[], dimension: BreakdownDimension): RankedRow[] => { +export const rollUpBreakdown = ( + results: readonly DailyData[], + dimension: BreakdownDimension | "api_keys", +): RankedRow[] => { const entries = results.flatMap((day) => Object.entries(day.breakdown[dimension] ?? {})); const totals = entries.reduce>((acc, [key, entry]) => { const row = acc.get(key) ?? { key, ...EMPTY_RANKED }; @@ -117,7 +120,7 @@ export const rankValue = (row: RankedRow, metric: UsageMetric): number => { */ export const seriesBy = ( results: readonly DailyData[], - dimension: BreakdownDimension, + dimension: BreakdownDimension | "api_keys", metric: UsageMetric, options: number | { top?: number; colorByKey?: ReadonlyMap } = 6, ): Series => { diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/AlertsInbox.tsx b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/AlertsInbox.tsx new file mode 100644 index 00000000000..2b3be5904e4 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/AlertsInbox.tsx @@ -0,0 +1,108 @@ +"use client"; + +import { useMemo, useState } from "react"; +import type { ColumnDef } from "@tanstack/react-table"; +import { DataTable, DataTableToolbar } from "@/components/shared/DataTable"; +import { PageTabs, PageTabsList, PageTabsTrigger, PageTabsContent } from "@/components/shared/Page"; +import type { DeprecationAlert } from "./deprecationAlerts"; +import { AlertStatus, effectiveAlertStatus } from "./migrationState"; +import RecommendedReplacement from "./RecommendedReplacement"; +import { DeprecationModel, DeprecationTiming } from "./DeprecationModel"; +import { migrationHref } from "./migrationRoute"; + +const columns: ColumnDef[] = [ + { + id: "model", + header: "Model", + cell: ({ row }) => ( +
+ + +
+ ), + }, + { + id: "deprecation", + header: "Deprecation", + cell: ({ row }) => ( +
+ + {row.original.model.deprecation_date} +
+ ), + }, + { accessorKey: "requests", header: "Requests · 30d", cell: ({ row }) => row.original.requests.toLocaleString() }, + { accessorKey: "lastUsedDate", header: "Last seen (UTC)" }, + { + id: "review", + header: "", + cell: ({ row }) => ( + + Review impact → + + ), + }, +]; + +export default function AlertsInbox({ + alerts, + pending, + status, + onStatusChange, + onRefresh, + refreshing, +}: { + alerts: readonly DeprecationAlert[]; + pending: boolean; + status: AlertStatus; + onStatusChange: (status: AlertStatus) => void; + onRefresh: () => void; + refreshing: boolean; +}) { + const [search, setSearch] = useState(""); + const counts = { open: 0, resolved: 0, dismissed: 0 }; + for (const alert of alerts) counts[effectiveAlertStatus(alert)]++; + const filtered = useMemo( + () => + alerts.filter( + (alert) => + effectiveAlertStatus(alert) === status && alert.usageModel.toLowerCase().includes(search.toLowerCase()), + ), + [alerts, status, search], + ); + return ( + + + {Object.values(AlertStatus).map((value) => ( + + {value} + {counts[value]} + + ))} + + + row.id} + fillHeight + size="compact" + isLoading={pending} + loadingMessage="Loading model deprecations…" + noDataMessage={`No ${status} model deprecations in the last 30 days.`} + toolbar={(table) => ( + + )} + /> + + + ); +} diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/DeprecationAlertItem.tsx b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/DeprecationAlertItem.tsx new file mode 100644 index 00000000000..dc2d2a848be --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/DeprecationAlertItem.tsx @@ -0,0 +1,35 @@ +"use client"; + +import { ArrowUpRight } from "lucide-react"; +import { DeprecationModel, DeprecationTiming } from "./DeprecationModel"; +import { buttonVariants } from "@/components/ui/button"; +import { cn } from "@/lib/cva.config"; +import { migrationHref } from "@/components/ModelDeprecationAlerts/migrationRoute"; +import RecommendedReplacement from "./RecommendedReplacement"; +import type { DeprecationAlert } from "./deprecationAlerts"; + +export default function DeprecationAlertItem({ alert }: { alert: DeprecationAlert }) { + const retired = alert.model.status === "deprecated"; + const title = retired + ? `${alert.usageModel} is deprecated` + : `${alert.usageModel} retires in ${alert.model.days_until_deprecation} days`; + return ( +
+
+
+ +

+ +

+

+ {alert.requests.toLocaleString()} requests in the last 30 days · Last seen {alert.lastUsedDate} UTC +

+ +
+
+ + Review migration + +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/DeprecationModel.tsx b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/DeprecationModel.tsx new file mode 100644 index 00000000000..dfb0e2b71c3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/DeprecationModel.tsx @@ -0,0 +1,32 @@ +import { ProviderLogo } from "@/components/molecules/models/ProviderLogo"; +import type { DeprecationAlert } from "./deprecationAlerts"; +import { migrationHref } from "./migrationRoute"; + +export function deprecationProvider(alert: DeprecationAlert): string { + return alert.model.litellm_provider || alert.usageModel.split("/")[0]; +} + +export function DeprecationModel({ alert }: { alert: DeprecationAlert }) { + const provider = deprecationProvider(alert); + return ( + + + {alert.usageModel} + + ); +} + +export function DeprecationTiming({ alert }: { alert: DeprecationAlert }) { + const days = alert.model.days_until_deprecation; + let label = `Deprecates in ${days} ${days === 1 ? "day" : "days"}`; + if (days === 0) label = "Deprecates today"; + if (days < 0) label = `Deprecated ${Math.abs(days)} ${days === -1 ? "day" : "days"} ago`; + return ( + + {label} + + ); +} diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/MigrationDetail.tsx b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/MigrationDetail.tsx new file mode 100644 index 00000000000..0ad842b9078 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/MigrationDetail.tsx @@ -0,0 +1,335 @@ +"use client"; + +import { useMemo } from "react"; +import { useQueryStates } from "nuqs"; +import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; +import type { ColumnDef } from "@tanstack/react-table"; +import { ArrowUpRight, Check, Copy, RotateCcw, X } from "lucide-react"; +import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { DataTable, useUrlTableState } from "@/components/shared/DataTable"; +import { Button, buttonVariants } from "@/components/ui/button"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { backgroundFetchClient } from "@/lib/http/api"; +import type { components } from "@/lib/http/schema"; +import { keyDetailHref, modelGroupHref } from "@/utils/entityLinks"; +import { copyToClipboard } from "@/utils/dataUtils"; +import { AlertStatus, effectiveAlertStatus } from "@/components/ModelDeprecationAlerts/migrationState"; +import { migrationUrlParsers } from "./migrationRoute"; +import type { DeprecationAlert, deprecationUsageRange } from "@/components/ModelDeprecationAlerts/deprecationAlerts"; +import RecommendedReplacement from "@/components/ModelDeprecationAlerts/RecommendedReplacement"; +import { DeprecationTiming } from "./DeprecationModel"; +import UsageTags from "./UsageTags"; +import ModelUsageChart from "./ModelUsageChart"; +import { affectedTeams } from "./migrationData"; + +const TABLE_STATE = { sortFields: [], defaultSort: { id: "", desc: false }, defaultPageSize: 50, filterColumns: [] }; +type KeyRow = components["schemas"]["KeySpendActivityRow"]; + +function migrationHeading(alert: DeprecationAlert, status: AlertStatus): string { + if (status === AlertStatus.Resolved) return "Migration marked resolved"; + if (status === AlertStatus.Dismissed) return "Model deprecation dismissed"; + if (alert.model.status === "deprecated") return "Migration needed: model is deprecated"; + return `Migration due ${alert.model.deprecation_date}`; +} + +function migrationFollowUp( + alert: DeprecationAlert, + context: { team: string; count: number | undefined; lastSeen: string; url: string }, +): string { + const retirement = alert.model.status === "deprecated" ? "retired on" : "retires on"; + const keys = context.count === undefined ? "the affected" : context.count; + return `${context.team}: please migrate ${keys} virtual keys using ${alert.usageModel}. This model ${retirement} ${alert.model.deprecation_date}. Last seen ${context.lastSeen} UTC. Please confirm once migration is complete. Affected keys: ${context.url}`; +} + +export default function MigrationDetail({ + selected, + scope, + range, + isViewOnly, +}: { + selected: DeprecationAlert; + scope: string; + range: ReturnType; + isViewOnly: boolean; +}) { + const [params, setParams] = useQueryStates(migrationUrlParsers); + const { data: teams = [] } = useTeams(); + const teamNames = useMemo( + () => new Map(teams.map((team) => [team.team_id, team.team_alias || team.team_id])), + [teams], + ); + const queryClient = useQueryClient(); + const table = useUrlTableState(TABLE_STATE); + const model = selected.usageModel; + const keyQuery = useQuery({ + queryKey: ["deprecationImpactKeys", scope, model, params.team, range, table.pagination], + + queryFn: async ({ signal }) => { + const query = { + ...range, + model, + offset: table.pagination.pageIndex * table.pagination.pageSize, + limit: table.pagination.pageSize, + timezone: 0, + }; + if (params.team) { + const { data } = await backgroundFetchClient.GET("/team/daily/activity/aggregated/keys", { + params: { query: { ...query, team_ids: params.team } }, + signal, + }); + if (!data) throw new Error("Affected keys are unavailable"); + return data; + } + const { data } = await backgroundFetchClient.GET("/user/daily/activity/aggregated/keys", { + params: { query }, + signal, + }); + if (!data) throw new Error("Affected keys are unavailable"); + return data; + }, + }); + const keys = useMemo(() => keyQuery.data?.api_keys.map((key) => key.api_key) ?? [], [keyQuery.data]); + const tagQueryOptions = { + queryKey: ["deprecationUsageTags", scope, model, range, keys], + enabled: keys.length > 0, + queryFn: async ({ signal }: { signal: AbortSignal }) => { + const { data } = await backgroundFetchClient.GET("/alerting/model-deprecations/tags", { + params: { query: { ...range, model, api_keys: keys } }, + signal, + }); + if (!data) throw new Error("Usage tags are unavailable"); + return data; + }, + staleTime: 5 * 60 * 1000, + }; + const tagQuery = useQuery(tagQueryOptions); + const tagsByKey = useMemo(() => { + const tags = new Map(); + for (const usage of tagQuery.data ?? []) { + const existing = tags.get(usage.api_key) ?? []; + existing.push(usage.tag); + tags.set(usage.api_key, existing); + } + return tags; + }, [tagQuery.data]); + const teamQuery = useQuery({ + queryKey: ["deprecationImpactTeams", scope, model, range], + + queryFn: async ({ signal }) => { + const query = { ...range, model, timezone: 0, api_key_limit: 1 }; + const { data } = await backgroundFetchClient.GET("/team/daily/activity/aggregated", { + params: { query }, + signal, + }); + if (!data) throw new Error("Affected teams are unavailable"); + return affectedTeams(data); + }, + }); + const statusMutation = useMutation({ + mutationFn: async (status: AlertStatus) => { + const body = { + model: selected.usageModel, + deprecation_date: selected.model.deprecation_date, + status, + last_seen_date: selected.lastUsedDate, + last_seen_requests: selected.lastDayRequests, + }; + const { data } = await backgroundFetchClient.PUT("/alerting/model-deprecations/state", { body }); + if (!data) throw new Error("Could not save alert status"); + return data; + }, + onSuccess: () => queryClient.invalidateQueries({ queryKey: ["deprecationAlerts"] }), + }); + const columns = useMemo[]>( + () => [ + { + id: "key", + header: "Virtual key", + cell: ({ row }) => ( + + {row.original.metadata.key_alias || "Unnamed key"} + + ), + }, + { + id: "team", + header: "Team", + cell: ({ row }) => + teamNames.get(row.original.metadata.team_id ?? "") || row.original.metadata.team_id || "No team", + }, + { + id: "owner", + header: "Owner", + cell: ({ row }) => row.original.metadata.user_email || row.original.metadata.user_id || "Not recorded", + }, + { + id: "tags", + header: "Usage tags", + cell: ({ row }) => ( + + ), + }, + { + id: "requests", + header: "Requests · 30d", + cell: ({ row }) => {row.original.metrics.api_requests.toLocaleString()}, + }, + { id: "spend", header: "Spend · 30d", cell: ({ row }) => `$${row.original.metrics.spend.toFixed(2)}` }, + ], + [teamNames, tagsByKey, tagQuery.isPending, tagQuery.isError], + ); + const status = effectiveAlertStatus(selected); + const heading = migrationHeading(selected, status); + const copyFollowUp = () => { + const teamName = params.team ? teamNames.get(params.team) || params.team : "Team owners"; + const count = keyQuery.data?.total_api_keys; + const lastSeen = params.team + ? teamQuery.data?.find((team) => team.id === params.team)?.lastUsedDate ?? selected.lastUsedDate + : selected.lastUsedDate; + const context = { team: teamName, count, lastSeen, url: window.location.href }; + const message = migrationFollowUp(selected, context); + void copyToClipboard(message, "Follow-up copied"); + }; + return ( + <> +
+
+ {status} +

{heading}

+

+ Last seen {selected.lastUsedDate} UTC · {selected.requests.toLocaleString()} requests in the last 30 days +

+ +
+
+ + {!isViewOnly && ( + <> + {status === AlertStatus.Open ? ( + <> + + + + ) : ( + + )} + + )} + {modelGroupHref(selected.model.model_name) && ( + + Model routing + + )} +
+
+

+ Resolve across all teams once owners confirm migration. New traffic reopens the alert. +

+ {statusMutation.isError && ( +

+ Could not save alert status. Please retry. +

+ )} + +
+
+

{params.team ? "Keys in this team" : "Affected keys"}

+

+ {keyQuery.data?.total_api_keys.toLocaleString() ?? "—"} +

+
+
+

Teams using this model

+

{teamQuery.data?.length ?? "—"}

+
+
+

Deprecation

+

+ +

+
+
+
+
+
+

+ Affected keys{" "} + + {keyQuery.data?.total_api_keys.toLocaleString() ?? "…"} + +

+

+ Keys and owners that used this model in the last 30 days. Tags show actual usage with this model. +

+
+ +
+ {keyQuery.isError ? ( +
+

Could not load affected keys.

+ +
+ ) : ( + row.api_key} + /> + )} +
+ {teamQuery.isError && ( +

Could not load team filters. The keys table still includes all matching teams.

+ )} + + ); +} diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelAlertsPage.integration.test.tsx b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelAlertsPage.integration.test.tsx new file mode 100644 index 00000000000..f8698556d49 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelAlertsPage.integration.test.tsx @@ -0,0 +1,124 @@ +import { beforeEach, afterEach, describe, expect, it, vi } from "vitest"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils"; +import ModelAlertsPage from "@/app/(dashboard)/alerts/page"; + +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useTeams: () => ({ data: [{ team_id: "support", team_alias: "Support" }] }), +})); +vi.mock("@/components/ModelDeprecationAlerts/useDeprecationAlerts", () => ({ + useDeprecationAlerts: () => ({ + scope: "test", + range: { start_date: "2026-10-06", end_date: "2026-11-04" }, + isPending: false, + isError: false, + data: [ + { + id: "retirement", + usageModel: "openai/gpt-4-turbo", + requests: 5000, + lastUsedDate: "2026-11-04", + lastDayRequests: 20, + model: { + model_name: "support-model", + deprecation_date: "2026-10-23", + days_until_deprecation: -12, + status: "deprecated", + successor_model: "replacement-model", + }, + }, + ], + }), +})); + +const calls: { path: string; query: URLSearchParams; method: string; body: unknown }[] = []; +function apiResponse(data: unknown) { + return new Response(JSON.stringify(data), { headers: { "Content-Type": "application/json" } }); +} + +describe("model migration impact", () => { + beforeEach(() => { + testQueryClient.clear(); + calls.length = 0; + vi.spyOn(globalThis, "fetch").mockImplementation(async (input) => { + const request = input instanceof Request ? input : new Request(input); + const url = new URL(request.url); + const call: (typeof calls)[number] = { + path: url.pathname, + query: url.searchParams, + method: request.method, + body: request.method === "PUT" ? await request.clone().json() : null, + }; + calls.push(call); + if (url.pathname === "/alerting/model-deprecations/tags") return apiResponse([]); + if (request.method === "PUT" && url.pathname === "/alerting/model-deprecations/state") + return apiResponse({ status: "resolved" }); + if (url.pathname.endsWith("/aggregated/keys")) { + const offset = Number(url.searchParams.get("offset") || 0); + const rows = Array.from({ length: 50 }, (_, i) => ({ + api_key: `reporting-${offset + i}`, + metadata: { key_alias: `support-key-${offset + i}`, team_id: "support", user_email: "owner@example.com" }, + metrics: { api_requests: 100, spend: 1 }, + })); + const page = { api_keys: rows, total_api_keys: 120, offset, limit: 50 }; + return apiResponse(page); + } + if (url.pathname.endsWith("/team/daily/activity/aggregated")) + return apiResponse({ results: [], metadata: { entity_total_api_keys: {} } }); + return apiResponse({}); + }); + }); + afterEach(() => vi.restoreAllMocks()); + it("should render the inbox with last seen and the recommendation supplied by the API", () => { + renderWithProviders(); + expect(screen.getByRole("tab", { name: "Model deprecations", exact: true })).toBeInTheDocument(); + expect(screen.getByText("replacement-model")).toBeInTheDocument(); + expect(screen.getByText("2026-11-04")).toBeInTheDocument(); + expect(screen.getByRole("link", { name: "Review impact →" })).toHaveAttribute( + "href", + "/ui/alerts?type=model-deprecations&model=openai%2Fgpt-4-turbo", + ); + }); + it("should page through all affected keys instead of showing a top-ten sample", async () => { + const user = userEvent.setup(); + renderWithProviders(, { searchParams: { model: "openai/gpt-4-turbo" } }); + expect(await screen.findByRole("link", { name: "support-key-0" })).toBeInTheDocument(); + expect(screen.getByText(/Showing 1-50 of 120/)).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Go to next page", exact: true })); + await screen.findByRole("link", { name: "support-key-50" }); + expect( + calls.some( + (call) => + call.path.endsWith("/aggregated/keys") && + call.query.get("offset") === "50" && + call.query.get("model") === "openai/gpt-4-turbo", + ), + ).toBe(true); + }); + it("should resolve the whole model using its global latest-day snapshot even in a team-filtered view", async () => { + const user = userEvent.setup(); + renderWithProviders(, { + searchParams: { model: "openai/gpt-4-turbo", team: "support" }, + }); + expect(await screen.findByRole("link", { name: "support-key-0" })).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Mark resolved", exact: true })); + const expectedAcknowledgement = { + model: "openai/gpt-4-turbo", + deprecation_date: "2026-10-23", + status: "resolved", + last_seen_date: "2026-11-04", + last_seen_requests: 20, + }; + await waitFor(() => + expect( + calls.find((call) => call.method === "PUT" && call.path === "/alerting/model-deprecations/state")?.body, + ).toEqual(expectedAcknowledgement), + ); + + expect( + calls.some( + (call) => call.path === "/team/daily/activity/aggregated/keys" && call.query.get("team_ids") === "support", + ), + ).toBe(true); + }); +}); diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelDeprecationsPanel.tsx b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelDeprecationsPanel.tsx new file mode 100644 index 00000000000..ef19bdea368 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelDeprecationsPanel.tsx @@ -0,0 +1,76 @@ +"use client"; + +import { useQueryStates } from "nuqs"; +import { ProviderLogo } from "@/components/molecules/models/ProviderLogo"; +import { deprecationProvider } from "./DeprecationModel"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { Page } from "@/components/shared/Page"; +import { Button, buttonVariants } from "@/components/ui/button"; +import { useDeprecationAlerts } from "@/components/ModelDeprecationAlerts/useDeprecationAlerts"; +import { hasCapability } from "@/utils/capabilities"; +import { migrationUrlParsers, migrationHref } from "./migrationRoute"; +import AlertsInbox from "./AlertsInbox"; +import MigrationDetail from "./MigrationDetail"; + +export default function ModelDeprecationsPanel() { + const { accessToken, userRole, isViewOnly } = useAuthorized(); + const allowed = hasCapability(userRole, "viewGlobalSpend"); + const [params, setParams] = useQueryStates(migrationUrlParsers); + const alerts = useDeprecationAlerts(accessToken); + const selected = alerts.data?.find((alert) => alert.usageModel === params.model); + const loadingDetail = !alerts.isError && Boolean(params.model) && alerts.isPending; + const detailReady = !alerts.isError && Boolean(params.model) && !alerts.isPending; + const missingDetail = detailReady && !selected; + if (!allowed) + return ( + +

Model deprecations

+

Gateway admin access is required to view affected keys and teams.

+ + ); + return ( + + {params.model && ( +
+
+

+ {selected && } + + {params.model} + +

+
+ + ← Model deprecations + +
+ )} + {alerts.isError ? ( +
+

Could not load alerts. No alert status has been changed.

+ +
+ ) : null} + {!alerts.isError && !params.model && ( + void alerts.refetch()} + refreshing={alerts.isFetching} + status={params.state} + onStatusChange={(state) => void setParams({ state })} + /> + )} + {loadingDetail &&

Loading alert…

} + {missingDetail &&

No matching model alert was found in the last 30 days.

} + {!alerts.isError && selected && ( + + )} + + ); +} diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelUsageChart.integration.test.tsx b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelUsageChart.integration.test.tsx new file mode 100644 index 00000000000..2ca03e64f1d --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelUsageChart.integration.test.tsx @@ -0,0 +1,103 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders, screen } from "../../../tests/test-utils"; +import ModelUsageChart from "./ModelUsageChart"; + +vi.mock("@/components/shared/charts", async (importOriginal) => ({ + ...(await importOriginal()), + StackedUsageChart: ({ data, labels }: { data: unknown; labels: string[] }) => ( +
+ {JSON.stringify({ data, labels })} +
+ ), +})); +const range = { start_date: "2026-10-07", end_date: "2026-11-05" }; +function mockActivity() { + return vi.spyOn(globalThis, "fetch").mockImplementation( + async () => + new Response( + JSON.stringify({ + results: [ + { + date: "2026-11-03", + metrics: { api_requests: 100, total_tokens: 1000, spend: 10 }, + breakdown: { + api_keys: { + "key-a": { + metrics: { api_requests: 30, spend: 3, total_tokens: 300 }, + metadata: { key_alias: "Support key" }, + }, + "key-b": { + metrics: { api_requests: 20, spend: 2, total_tokens: 200 }, + metadata: { key_alias: "Research key" }, + }, + }, + }, + }, + ], + }), + { headers: { "Content-Type": "application/json" } }, + ), + ); +} + +describe("model daily usage", () => { + afterEach(() => vi.restoreAllMocks()); + it("should render the model's daily requests stacked by key and preserve remaining usage", async () => { + const fetch = mockActivity(); + renderWithProviders( + , + ); + const chart = await screen.findByRole("img", { name: "Daily key usage chart" }); + expect(chart).toHaveTextContent('"labels":["Support key","Research key","Other"]'); + expect(chart).toHaveTextContent('"s0":30,"s1":20,"s2":50'); + const request = fetch.mock.calls[0][0] as Request; + expect(new URL(request.url).searchParams.get("model")).toBe("anthropic/retiring-model"); + }); + it("should apply the selected team and show requests without metric controls", async () => { + const fetch = mockActivity(); + renderWithProviders( + , + ); + await screen.findByRole("img", { name: "Daily key usage chart" }); + const request = fetch.mock.calls[0][0] as Request; + expect(new URL(request.url).pathname).toBe("/team/daily/activity/aggregated"); + expect(new URL(request.url).searchParams.get("team_ids")).toBe("support"); + expect(screen.queryByRole("radio")).not.toBeInTheDocument(); + expect(screen.getByRole("img", { name: "Daily key usage chart" })).toHaveTextContent('"s0":30,"s1":20,"s2":50'); + }); + it("should keep all keys in the stacks while allowing the key labels to expand", async () => { + const fetch = mockActivity(); + const keys = Object.fromEntries( + Array.from({ length: 12 }, (_, index) => [ + `key-${index}`, + { + metrics: { api_requests: 10, spend: 1, total_tokens: 100 }, + metadata: { key_alias: `Key ${index + 1}` }, + }, + ]), + ); + const response = { + results: [ + { + date: "2026-11-03", + metrics: { api_requests: 120, spend: 12, total_tokens: 1200 }, + breakdown: { api_keys: keys }, + }, + ], + }; + fetch.mockResolvedValueOnce( + new Response(JSON.stringify(response), { headers: { "Content-Type": "application/json" } }), + ); + const user = userEvent.setup(); + renderWithProviders( + , + ); + expect(await screen.findByRole("img", { name: "Daily key usage chart" })).toHaveTextContent('"s11":10'); + expect(screen.queryByText("Key 12", { exact: true })).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Show all key labels" })); + expect(screen.getByText("Key 12", { exact: true })).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Show fewer key labels" })); + expect(screen.queryByText("Key 12", { exact: true })).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelUsageChart.tsx b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelUsageChart.tsx new file mode 100644 index 00000000000..7f8262c83fd --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/ModelUsageChart.tsx @@ -0,0 +1,127 @@ +"use client"; + +import { useMemo, useState } from "react"; +import { useQuery } from "@tanstack/react-query"; +import { StackedUsageChart } from "@/components/shared/charts"; +import { Button } from "@/components/ui/button"; +import { backgroundFetchClient } from "@/lib/http/api"; +import { toDailyData } from "@/components/UsagePage/dailyActivityApi"; +import { Panel, ChartSkeleton } from "@/app/(dashboard)/usage/_components/components/overview/Primitives"; +import { seriesBy } from "@/app/(dashboard)/usage/_components/components/overview/overviewData"; +import type { deprecationUsageRange } from "./deprecationAlerts"; +import { dailyModelRequests } from "./migrationData"; + +const formatDate = (date: string) => + new Date(`${date}T00:00:00Z`).toLocaleDateString("en-US", { month: "short", day: "numeric", timeZone: "UTC" }); + +export default function ModelUsageChart({ + model, + team, + scope, + range, +}: { + model: string; + team: string | null; + scope: string; + range: ReturnType; +}) { + const [legendExpanded, setLegendExpanded] = useState(false); + const activityOptions = { + queryKey: ["deprecationDailyUsage", scope, model, team, range], + queryFn: async ({ signal }: { signal: AbortSignal }) => { + const query = { ...range, model, api_key_limit: 100, timezone: 0 }; + if (team) { + const { data } = await backgroundFetchClient.GET("/team/daily/activity/aggregated", { + params: { query: { ...query, team_ids: team } }, + signal, + }); + if (!data) throw new Error("Daily model usage is unavailable"); + return data; + } + const { data } = await backgroundFetchClient.GET("/user/daily/activity/aggregated", { + params: { query }, + signal, + }); + if (!data) throw new Error("Daily model usage is unavailable"); + return data; + }, + staleTime: 5 * 60 * 1000, + refetchInterval: 5 * 60 * 1000, + }; + const activity = useQuery(activityOptions); + const daily = useMemo(() => dailyModelRequests(activity.data?.results ?? [], range), [activity.data, range]); + const totals = useMemo(() => new Map(daily.map((day) => [day.date, day.requests])), [daily]); + const series = useMemo(() => { + const days = activity.data ? toDailyData(activity.data) : []; + const aliases = new Map( + days.flatMap((day) => + Object.entries(day.breakdown.api_keys).map( + ([id, entry]) => [id, entry.metadata.key_alias || `Key · ${id.slice(0, 6)}`] as const, + ), + ), + ); + const result = seriesBy(days, "api_keys", "requests", 100); + const byDate = new Map(result.data.map((day) => [day.date, day])); + const zeros = Object.fromEntries(result.keys.map((key) => [key, 0])); + result.data = daily.map((day) => byDate.get(day.date) ?? { date: day.date, label: formatDate(day.date), ...zeros }); + result.labels = result.labels.map((label) => aliases.get(label) || label); + return result; + }, [activity.data, daily]); + function renderContent() { + if (activity.isPending) return ; + if (activity.isError) + return ( +
+ Could not load daily requests.{" "} + +
+ ); + return ( + <> + value.toLocaleString()} + totalFor={(date) => totals.get(date)} + className="h-64" + /> +
+ {series.labels.slice(0, legendExpanded ? undefined : 8).map((label, index) => ( + + + {label} + + ))} +
+ {series.labels.length > 8 && ( + + )} + + ); + } + return ( + + {renderContent()} +

+ Last 30 days · UTC · Includes zero-usage days. Today is partial. +

+ + ); +} diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/RecommendedReplacement.tsx b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/RecommendedReplacement.tsx new file mode 100644 index 00000000000..e1c237e6ff8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/RecommendedReplacement.tsx @@ -0,0 +1,7 @@ +export default function RecommendedReplacement({ model }: { model: string | null | undefined }) { + return ( +

+ Recommended new model: {model || "Not provided"} +

+ ); +} diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/UsageTags.test.tsx b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/UsageTags.test.tsx new file mode 100644 index 00000000000..9d1767bb76b --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/UsageTags.test.tsx @@ -0,0 +1,21 @@ +import { describe, expect, it } from "vitest"; +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import UsageTags from "./UsageTags"; + +describe("model usage tags", () => { + it("should render tags used with the model", () => { + render(); + expect(screen.getByText("production")).toBeInTheDocument(); + }); + it("should show five tags and let the user expand and collapse the rest", async () => { + const user = userEvent.setup(); + render(); + expect(screen.queryByText("six")).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Show 2 more usage tags" })); + expect(screen.getByText("six")).toBeInTheDocument(); + expect(screen.getByText("seven")).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Show fewer usage tags" })); + expect(screen.queryByText("six")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/UsageTags.tsx b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/UsageTags.tsx new file mode 100644 index 00000000000..a25ef5ec5b3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/UsageTags.tsx @@ -0,0 +1,40 @@ +"use client"; + +import { useState } from "react"; +import { Badge } from "@/components/ui/badge"; + +export default function UsageTags({ + tags, + loading, + failed, +}: { + tags: readonly string[]; + loading: boolean; + failed: boolean; +}) { + const [expanded, setExpanded] = useState(false); + if (loading) return Loading tags…; + if (failed) return Tags unavailable; + if (tags.length === 0) return No recorded tags; + const visible = expanded ? tags : tags.slice(0, 5); + return ( +
+ {visible.map((tag) => ( + + {tag} + + ))} + {tags.length > 5 && ( + + )} +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/deprecationAlerts.test.ts b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/deprecationAlerts.test.ts new file mode 100644 index 00000000000..f38eee7fccf --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/deprecationAlerts.test.ts @@ -0,0 +1,78 @@ +import { describe, expect, it } from "vitest"; +import type { components } from "@/lib/http/schema"; +import { deprecationUsageRange, usedDeprecationAlerts } from "./deprecationAlerts"; + +type Model = components["schemas"]["ModelDeprecationInfo"]; +type Day = components["schemas"]["DailySpendData"]; +const model: Model = { + model_name: "support", + litellm_model: "anthropic/claude-sonnet-4-5", + deprecation_date: "2026-11-30", + days_until_deprecation: 26, + status: "imminent", +}; +// Use the API's defaults for the unrelated aggregate metrics. +const metrics: components["schemas"]["SpendMetrics"] = { + spend: 0, + api_requests: 0, + successful_requests: 0, + failed_requests: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, +}; +function day(date: string, count: number, name = model.litellm_model!): Day { + return { date, metrics, breakdown: { models: { [name]: { metrics: { ...metrics, api_requests: count } } } } }; +} + +describe("usedDeprecationAlerts", () => { + it("should aggregate requests and report the latest day of actual model usage", () => { + const result = usedDeprecationAlerts([model], [day("2026-11-04", 12), day("2026-10-20", 8)]); + expect(result).toHaveLength(1); + expect(result[0]).toMatchObject({ usageModel: model.litellm_model, requests: 20, lastUsedDate: "2026-11-04" }); + }); + it("should omit configured models without observed requests", () => { + expect(usedDeprecationAlerts([model], [])).toEqual([]); + expect(usedDeprecationAlerts([model], [day("2026-11-04", 0)])).toEqual([]); + }); + it("should avoid attributing healthy deployments in the same group to a deprecated model", () => { + const activity = day("2026-11-04", 100, "anthropic/claude-sonnet-4-6"); + activity.breakdown!.model_groups = { support: { metrics: { ...metrics, api_requests: 100 } } }; + expect(usedDeprecationAlerts([model], [activity])).toEqual([]); + }); + it("should match provider-qualified usage when the cost map names an unqualified model", () => { + const entry = { ...model, litellm_model: "claude-sonnet-4-5", litellm_provider: "anthropic" }; + const result = usedDeprecationAlerts([entry], [day("2026-11-04", 12)]); + expect(result[0]).toMatchObject({ usageModel: "anthropic/claude-sonnet-4-5", requests: 12 }); + }); + it("should keep usage for different providers separate", () => { + const entry = { ...model, litellm_provider: "anthropic" }; + expect(usedDeprecationAlerts([entry], [day("2026-11-04", 12, "bedrock/claude-sonnet-4-5")])).toEqual([]); + }); + it("should deduplicate deployments and put the earliest retirement first", () => { + const retired = { ...model, deprecation_date: "2026-10-23", status: "deprecated" as const }; + const result = usedDeprecationAlerts([model, retired, model], [day("2026-11-04", 10)]); + expect(result.map((alert) => alert.model.deprecation_date)).toEqual(["2026-10-23", "2026-11-30"]); + }); + it("should give an escalated alert a new dismissal identity", () => { + const activity = [day("2026-11-04", 10)]; + const imminent = usedDeprecationAlerts([model], activity)[0]; + const retired = usedDeprecationAlerts([{ ...model, status: "deprecated" }], activity)[0]; + expect(retired.id).not.toBe(imminent.id); + }); + it("should use the public model name when underlying model metadata is unavailable", () => { + const result = usedDeprecationAlerts([{ ...model, litellm_model: null }], [day("2026-11-04", 10, "support")]); + expect(result[0].requests).toBe(10); + }); +}); + +describe("deprecationUsageRange", () => { + it("should include exactly 30 UTC calendar dates across month boundaries", () => { + expect(deprecationUsageRange(new Date("2026-11-04T23:00:00Z"))).toEqual({ + start_date: "2026-10-06", + end_date: "2026-11-04", + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/deprecationAlerts.ts b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/deprecationAlerts.ts new file mode 100644 index 00000000000..d22d940ab16 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/deprecationAlerts.ts @@ -0,0 +1,61 @@ +import type { components } from "@/lib/http/schema"; + +type ModelDeprecation = components["schemas"]["ModelDeprecationInfo"]; +type DailySpendData = components["schemas"]["DailySpendData"]; + +export interface DeprecationAlert { + id: string; + model: ModelDeprecation; + usageModel: string; + requests: number; + lastUsedDate: string; + lastDayRequests: number; + state?: components["schemas"]["DeprecationAlertState"]; +} + +export function deprecationUsageRange(today: Date) { + const end = new Date(today); + const start = new Date(today); + start.setUTCDate(start.getUTCDate() - 29); + return { start_date: start.toISOString().slice(0, 10), end_date: end.toISOString().slice(0, 10) }; +} + +/** Match actual model usage, rather than model-group totals that may include healthy deployments. */ +export function usedDeprecationAlerts( + models: readonly ModelDeprecation[], + days: readonly DailySpendData[], +): DeprecationAlert[] { + const alerts: DeprecationAlert[] = []; + const seen = new Set(); + for (const model of models) { + const baseModel = model.litellm_model ?? model.model_name; + const candidates = new Set([baseModel]); + if (model.litellm_provider && !baseModel.startsWith(`${model.litellm_provider}/`)) { + candidates.add(`${model.litellm_provider}/${baseModel}`); + } + for (const usageModel of candidates) { + const id = JSON.stringify([usageModel, model.deprecation_date, model.status]); + if (seen.has(id)) continue; + seen.add(id); + let requests = 0; + let lastUsedDate = ""; + let lastDayRequests = 0; + for (const day of days) { + const metrics = day.breakdown?.models?.[usageModel]?.metrics; + if (!metrics || metrics.api_requests <= 0) continue; + requests += metrics.api_requests; + if (day.date > lastUsedDate) { + lastUsedDate = day.date; + lastDayRequests = metrics.api_requests; + } else if (day.date === lastUsedDate) { + lastDayRequests += metrics.api_requests; + } + } + if (requests > 0) { + const alert = { id, model, usageModel, requests, lastUsedDate, lastDayRequests }; + alerts.push(alert); + } + } + } + return alerts.sort((a, b) => a.model.deprecation_date.localeCompare(b.model.deprecation_date)); +} diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationData.test.ts b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationData.test.ts new file mode 100644 index 00000000000..17a5a47391f --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationData.test.ts @@ -0,0 +1,43 @@ +import { describe, expect, it } from "vitest"; +import type { components } from "@/lib/http/schema"; +import { dailyModelRequests } from "./migrationData"; + +type Day = components["schemas"]["DailySpendData"]; +function day(date: string, api_requests: number): Day { + return { + date, + metrics: { + api_requests, + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + successful_requests: api_requests, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + }; +} + +describe("model request trend", () => { + it("should include zero-usage days after the model stops receiving traffic", () => { + expect(dailyModelRequests([day("2026-11-03", 120)], { start_date: "2026-11-03", end_date: "2026-11-05" })).toEqual([ + { date: "2026-11-03", requests: 120 }, + { date: "2026-11-04", requests: 0 }, + { date: "2026-11-05", requests: 0 }, + ]); + }); + it("should keep dates chronological across month boundaries and preserve daily totals", () => { + expect( + dailyModelRequests([day("2026-11-01", 5), day("2026-10-31", 50)], { + start_date: "2026-10-31", + end_date: "2026-11-02", + }), + ).toEqual([ + { date: "2026-10-31", requests: 50 }, + { date: "2026-11-01", requests: 5 }, + { date: "2026-11-02", requests: 0 }, + ]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationData.ts b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationData.ts new file mode 100644 index 00000000000..99b97c0fa3a --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationData.ts @@ -0,0 +1,49 @@ +import type { components } from "@/lib/http/schema"; + +type Activity = components["schemas"]["SpendAnalyticsPaginatedResponse"]; +export interface AffectedTeam { + id: string; + requests: number; + keys: number; + lastUsedDate: string; +} + +export function affectedTeams(activity: Activity): AffectedTeam[] { + const teams = new Map(); + for (const day of activity.results ?? []) { + for (const [id, entry] of Object.entries(day.breakdown?.entities ?? {})) { + if (entry.metrics.api_requests <= 0) continue; + const team = teams.get(id) ?? { + id, + requests: 0, + keys: activity.metadata?.entity_total_api_keys?.[id] ?? 0, + lastUsedDate: "", + }; + team.requests += entry.metrics.api_requests; + if (day.date > team.lastUsedDate) team.lastUsedDate = day.date; + teams.set(id, team); + } + } + return [...teams.values()].sort((a, b) => b.requests - a.requests); +} + +export interface DailyModelRequests extends Record { + date: string; + requests: number; +} + +/** Include quiet days so a completed migration remains visible as zero traffic. */ +export function dailyModelRequests( + days: readonly components["schemas"]["DailySpendData"][], + range: { start_date: string; end_date: string }, +): DailyModelRequests[] { + const counts = new Map(); + for (const day of days) counts.set(day.date, (counts.get(day.date) ?? 0) + day.metrics.api_requests); + const result: DailyModelRequests[] = []; + const end = new Date(`${range.end_date}T00:00:00Z`); + for (const day = new Date(`${range.start_date}T00:00:00Z`); day <= end; day.setUTCDate(day.getUTCDate() + 1)) { + const date = day.toISOString().slice(0, 10); + result.push({ date, requests: counts.get(date) ?? 0 }); + } + return result; +} diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationRoute.ts b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationRoute.ts new file mode 100644 index 00000000000..4581f8042f4 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationRoute.ts @@ -0,0 +1,14 @@ +import { parseAsString, parseAsStringEnum } from "nuqs"; +import { AlertStatus } from "@/components/ModelDeprecationAlerts/migrationState"; +import { uiHref } from "@/utils/uiHref"; + +export const migrationUrlParsers = { + model: parseAsString, + team: parseAsString, + state: parseAsStringEnum(Object.values(AlertStatus)).withDefault(AlertStatus.Open), +}; + +export function migrationHref(model?: string) { + const base = `${uiHref("alerts")}?type=model-deprecations`; + return model ? `${base}&model=${encodeURIComponent(model)}` : base; +} diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationState.test.ts b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationState.test.ts new file mode 100644 index 00000000000..63581056ece --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationState.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from "vitest"; +import { AlertStatus, effectiveAlertStatus } from "./migrationState"; +import type { DeprecationAlert } from "./deprecationAlerts"; +const alert: DeprecationAlert = { + id: "model", + usageModel: "openai/gpt-4-turbo", + requests: 1000, + lastUsedDate: "2026-11-04", + lastDayRequests: 10, + model: { model_name: "support", deprecation_date: "2026-10-23", days_until_deprecation: -12, status: "deprecated" }, +}; +const state = { + model: alert.usageModel, + deprecation_date: "2026-10-23", + status: "resolved" as const, + last_seen_date: "2026-11-04", + last_seen_requests: 10, +}; +describe("effectiveAlertStatus", () => { + it("should keep a resolved migration closed when only its historic traffic remains", () => { + expect(effectiveAlertStatus({ ...alert, state })).toBe(AlertStatus.Resolved); + }); + it("should reopen on additional requests in the same UTC day", () => { + expect(effectiveAlertStatus({ ...alert, lastDayRequests: 11, state })).toBe(AlertStatus.Open); + }); + it("should reopen when the model is used on a later day", () => { + expect(effectiveAlertStatus({ ...alert, lastUsedDate: "2026-11-05", state })).toBe(AlertStatus.Open); + }); + it("should keep acknowledgement when old requests roll out of the 30-day window", () => { + expect(effectiveAlertStatus({ ...alert, requests: 900, state })).toBe(AlertStatus.Resolved); + }); + it("should allow a shared dismissed state and reopen it on fresh traffic", () => { + const dismissed = { ...state, status: "dismissed" as const }; + expect(effectiveAlertStatus({ ...alert, state: dismissed })).toBe(AlertStatus.Dismissed); + expect(effectiveAlertStatus({ ...alert, lastDayRequests: 11, state: dismissed })).toBe(AlertStatus.Open); + }); +}); diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationState.ts b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationState.ts new file mode 100644 index 00000000000..006cea4a985 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/migrationState.ts @@ -0,0 +1,17 @@ +import type { DeprecationAlert } from "./deprecationAlerts"; + +export enum AlertStatus { + Open = "open", + Resolved = "resolved", + Dismissed = "dismissed", +} + +export function effectiveAlertStatus(alert: DeprecationAlert): AlertStatus { + const state = alert.state; + if (!state || state.status === "open") return AlertStatus.Open; + if (!state.last_seen_date || alert.lastUsedDate > state.last_seen_date) return AlertStatus.Open; + if (alert.lastUsedDate === state.last_seen_date && alert.lastDayRequests > state.last_seen_requests) + return AlertStatus.Open; + if (state.status === "resolved") return AlertStatus.Resolved; + return AlertStatus.Dismissed; +} diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/useDeprecationAlerts.integration.test.tsx b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/useDeprecationAlerts.integration.test.tsx new file mode 100644 index 00000000000..897850393e2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/useDeprecationAlerts.integration.test.tsx @@ -0,0 +1,74 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders, screen, testQueryClient } from "../../../tests/test-utils"; +import { backgroundFetchClient } from "@/lib/http/api"; +import { useDeprecationAlerts } from "./useDeprecationAlerts"; +import { useCurrentUser } from "@/app/(dashboard)/hooks/currentUser/useCurrentUser"; + +vi.mock("@/lib/http/api", () => ({ backgroundFetchClient: { GET: vi.fn() } })); +vi.mock("@/app/(dashboard)/hooks/currentUser/useCurrentUser", () => ({ useCurrentUser: vi.fn() })); + +function AlertsHost() { + const { data, isError, fetchStatus } = useDeprecationAlerts("test-session"); + if (isError) return
Unavailable
; + if (fetchStatus !== "idle") return
Loading
; + return
{data?.length ?? 0} alerts
; +} + +describe("deprecation alert API wiring", () => { + beforeEach(() => { + vi.clearAllMocks(); + testQueryClient.clear(); + }); + it("should avoid querying gateway-wide usage for non-admins", () => { + vi.mocked(useCurrentUser).mockReturnValue({ + data: { user_id: "internal", user_role: "internal_user" }, + } as ReturnType); + renderWithProviders(); + expect(screen.getByText("0 alerts")).toBeInTheDocument(); + expect(backgroundFetchClient.GET).not.toHaveBeenCalled(); + }); + it("should skip the usage query when no configured models need attention", async () => { + vi.mocked(useCurrentUser).mockReturnValue({ data: { user_id: "admin", user_role: "proxy_admin" } } as ReturnType< + typeof useCurrentUser + >); + vi.mocked(backgroundFetchClient.GET).mockResolvedValue({ + data: { deprecated: [], imminent: [], upcoming: [] }, + } as never); + renderWithProviders(); + expect(await screen.findByText("0 alerts")).toBeInTheDocument(); + expect(backgroundFetchClient.GET).toHaveBeenCalledTimes(1); + expect(backgroundFetchClient.GET).toHaveBeenCalledWith( + "/model/deprecations", + expect.objectContaining({ params: { query: { warn_within_days: 30 } } }), + ); + }); + it("should load acknowledgement state from the alerting API for models with observed usage", async () => { + vi.mocked(useCurrentUser).mockReturnValue({ data: { user_id: "admin", user_role: "proxy_admin" } } as ReturnType< + typeof useCurrentUser + >); + const model = { + model_name: "support", + litellm_model: "anthropic/claude-sonnet-4-5", + deprecation_date: "2026-11-30", + status: "imminent", + }; + const schedule = { deprecated: [], imminent: [model], upcoming: [] }; + const activity = { + results: [ + { date: "2026-11-04", breakdown: { models: { [model.litellm_model]: { metrics: { api_requests: 10 } } } } }, + ], + }; + vi.mocked(backgroundFetchClient.GET) + .mockResolvedValueOnce({ data: schedule } as never) + .mockResolvedValueOnce({ data: activity } as never) + .mockResolvedValueOnce({ data: { status: "open" } } as never); + renderWithProviders(); + expect(await screen.findByText("1 alerts")).toBeInTheDocument(); + expect(backgroundFetchClient.GET).toHaveBeenLastCalledWith( + "/alerting/model-deprecations/state", + expect.objectContaining({ + params: { query: { model: model.litellm_model, deprecation_date: model.deprecation_date } }, + }), + ); + }); +}); diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/useDeprecationAlerts.ts b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/useDeprecationAlerts.ts new file mode 100644 index 00000000000..ce126d5cd4f --- /dev/null +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/useDeprecationAlerts.ts @@ -0,0 +1,51 @@ +import { useQuery } from "@tanstack/react-query"; +import { useCurrentUser } from "@/app/(dashboard)/hooks/currentUser/useCurrentUser"; +import { backgroundFetchClient } from "@/lib/http/api"; +import { getProxyBaseUrl } from "@/components/networking"; +import { hasCapability } from "@/utils/capabilities"; +import { deprecationUsageRange, usedDeprecationAlerts } from "./deprecationAlerts"; + +const POLL_INTERVAL = 5 * 60 * 1000; + +export function useDeprecationAlerts(accessToken: string | null) { + const { data: currentUser } = useCurrentUser(accessToken); + const userId = currentUser?.user_id; + const enabled = Boolean(accessToken && userId && hasCapability(currentUser?.user_role, "viewGlobalSpend")); + const scope = JSON.stringify([getProxyBaseUrl(), userId]); + const range = deprecationUsageRange(new Date()); + const options = { + queryKey: ["deprecationAlerts", scope, range], + enabled, + queryFn: async ({ signal }: { signal: AbortSignal }) => { + const { data: schedule } = await backgroundFetchClient.GET("/model/deprecations", { + params: { query: { warn_within_days: 30 } }, + signal, + }); + if (!schedule) throw new Error("Model deprecation information is unavailable"); + const models = [...(schedule.deprecated ?? []), ...(schedule.imminent ?? [])]; + if (models.length === 0) return []; + const { data: activity } = await backgroundFetchClient.GET("/user/daily/activity/aggregated", { + params: { query: { ...range, timezone: 0, api_key_limit: 1 } }, + signal, + }); + if (!activity) throw new Error("Model usage is unavailable"); + const alerts = usedDeprecationAlerts(models, activity.results ?? []); + return Promise.all( + alerts.map(async (alert) => { + const { data: state } = await backgroundFetchClient.GET("/alerting/model-deprecations/state", { + params: { query: { model: alert.usageModel, deprecation_date: alert.model.deprecation_date } }, + signal, + }); + if (!state) throw new Error("Alert acknowledgement state is unavailable"); + alert.state = state; + return alert; + }), + ); + }, + staleTime: POLL_INTERVAL, + refetchInterval: POLL_INTERVAL, + retry: false, + }; + const query = useQuery(options); + return { ...query, data: enabled ? query.data : undefined, scope, range }; +} diff --git a/ui/litellm-dashboard/src/components/Navbar/NotificationsBell/NotificationsBell.test.tsx b/ui/litellm-dashboard/src/components/Navbar/NotificationsBell/NotificationsBell.test.tsx index 93d9b6018ac..3b0eec6feec 100644 --- a/ui/litellm-dashboard/src/components/Navbar/NotificationsBell/NotificationsBell.test.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/NotificationsBell/NotificationsBell.test.tsx @@ -5,6 +5,14 @@ import React from "react"; import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +vi.mock("@/components/ModelDeprecationAlerts/useDeprecationAlerts", () => ({ + useDeprecationAlerts: () => ({ + data: [], + scope: "test", + range: { start_date: "2026-10-06", end_date: "2026-11-04" }, + }), +})); + vi.mock("@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails", () => ({ useHealthReadinessDetails: vi.fn(), })); @@ -173,3 +181,77 @@ describe("NotificationsBell", () => { expect(announcement("LiteLLM v1.103.0 is available")).toHaveTextContent("Your current version is v1.102.0"); }); }); + +const DEPRECATION = { + id: "sonnet:2026-11-30:imminent", + usageModel: "anthropic/claude-sonnet-4-5", + requests: 18432, + lastUsedDate: "2026-11-04", + lastDayRequests: 123, + model: { + model_name: "support", + litellm_model: "anthropic/claude-sonnet-4-5", + deprecation_date: "2026-11-30", + days_until_deprecation: 26, + status: "imminent" as const, + successor_model: "recommended-model", + }, +}; + +describe("model deprecation notifications", () => { + beforeEach(() => localStorage.clear()); + afterEach(() => localStorage.clear()); + it("should show last seen and link to the migration workspace and open alerts inbox", async () => { + const user = userEvent.setup(); + renderWithProviders( + , + ); + await user.click(screen.getByRole("button", { name: /1 model deprecation alert/ })); + expect(screen.getByRole("article", { name: /retires in 26 days/ })).toHaveTextContent("Last seen 2026-11-04 UTC"); + expect(screen.getByRole("link", { name: "Review migration" })).toHaveAttribute( + "href", + "/ui/alerts?type=model-deprecations&model=anthropic%2Fclaude-sonnet-4-5", + ); + expect(screen.getByRole("link", { name: "View all 1" })).toHaveAttribute( + "href", + "/ui/alerts?type=model-deprecations", + ); + expect(screen.getByText("recommended-model")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Dismiss", exact: true })).not.toBeInTheDocument(); + }); + it("should clear the warning when the alert has a shared resolved acknowledgement", () => { + const state = { + model: DEPRECATION.usageModel, + deprecation_date: "2026-11-30", + status: "resolved" as const, + last_seen_date: "2026-11-04", + last_seen_requests: 123, + }; + renderWithProviders( + , + ); + expect(bell(1)).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /model deprecation alert/ })).not.toBeInTheDocument(); + }); + it("should reopen a resolved alert when more requests arrive on the last-seen day", () => { + const state = { + model: DEPRECATION.usageModel, + deprecation_date: "2026-11-30", + status: "resolved" as const, + last_seen_date: "2026-11-04", + last_seen_requests: 122, + }; + renderWithProviders( + , + ); + expect(screen.getByRole("button", { name: /1 model deprecation alert/ })).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Navbar/NotificationsBell/NotificationsBell.tsx b/ui/litellm-dashboard/src/components/Navbar/NotificationsBell/NotificationsBell.tsx index 2c26f34b750..2756af0e60b 100644 --- a/ui/litellm-dashboard/src/components/Navbar/NotificationsBell/NotificationsBell.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/NotificationsBell/NotificationsBell.tsx @@ -17,7 +17,12 @@ import { Popover, PopoverContent, PopoverTitle, PopoverTrigger } from "@/compone import { Separator } from "@/components/ui/separator"; import { cn } from "@/lib/cva.config"; import { isNewerVersion } from "@/utils/versionUtils"; -import { ArrowUpRight, Bell } from "lucide-react"; +import { AlertTriangle, ArrowUpRight, Bell } from "lucide-react"; +import DeprecationAlertItem from "@/components/ModelDeprecationAlerts/DeprecationAlertItem"; +import { type DeprecationAlert } from "@/components/ModelDeprecationAlerts/deprecationAlerts"; +import { useDeprecationAlerts } from "@/components/ModelDeprecationAlerts/useDeprecationAlerts"; +import { AlertStatus, effectiveAlertStatus } from "@/components/ModelDeprecationAlerts/migrationState"; +import { migrationHref } from "@/components/ModelDeprecationAlerts/migrationRoute"; import React, { Fragment } from "react"; import { useTranslation } from "react-i18next"; import i18n from "@/i18n/i18n"; @@ -102,6 +107,7 @@ const AnnouncementItem: React.FC<{ announcement: Announcement }> = ({ announceme }; interface NotificationsBellViewProps { + deprecationAlerts?: readonly DeprecationAlert[]; deploymentHealthLink?: React.ReactNode; currentVersion: string | null | undefined; latestRelease: LatestReleaseInfo | null | undefined; @@ -111,25 +117,39 @@ export const NotificationsBellView: React.FC = ({ currentVersion, latestRelease, deploymentHealthLink, + deprecationAlerts, }) => { const { t } = useTranslation(); const upgrade = upgradeAnnouncement(currentVersion, latestRelease); const announcements = upgrade ? [upgrade, AUTO_ROUTER_ANNOUNCEMENT] : [AUTO_ROUTER_ANNOUNCEMENT]; const unreadCount = useUnreadAnnouncementCount(announcements.map((announcement) => announcement.readKey)); + const openAlerts = (deprecationAlerts ?? []).filter((alert) => effectiveAlertStatus(alert) === AlertStatus.Open); + const openDeprecations = openAlerts.length; + let label = + unreadCount > 0 + ? t("shell.notifications.bell.triggerUnread", { unread: unreadCount }) + : t("shell.notifications.bell.trigger"); + if (openDeprecations > 0) { + label = `Notifications, ${openDeprecations} model deprecation ${openDeprecations === 1 ? "alert needs" : "alerts need"} attention`; + } + return ( } - aria-label={ - unreadCount > 0 - ? t("shell.notifications.bell.triggerUnread", { unread: unreadCount }) - : t("shell.notifications.bell.trigger") - } + aria-label={label} > - - {unreadCount > 0 ? : null} + 0 && "text-amber-700")} aria-hidden /> + {openDeprecations > 0 ? ( + + + + ) : null} + {openDeprecations === 0 && unreadCount > 0 ? ( + + ) : null} = ({
+ {openAlerts.length > 0 && ( + <> +
+
+

Model deprecations

+ + View all {openDeprecations} + +
+

Model migrations across your keys and teams.

+
+ {openAlerts.slice(0, 3).map((alert) => ( + + + + + ))} + +

Announcements

+ + )} + {deprecationAlerts && openAlerts.length === 0 && ( + + View model deprecations + + )} + {announcements.map((announcement, index) => ( {index > 0 ? : null} @@ -162,10 +209,12 @@ export const NotificationsBell: React.FC<{ accessToken: string | null; deploymen }) => { const { data: healthData } = useHealthReadinessDetails(accessToken); const { data: latestRelease } = useLatestReleaseInfo(accessToken); + const { data: deprecationAlerts } = useDeprecationAlerts(accessToken); return ( ); diff --git a/ui/litellm-dashboard/src/components/leftnav.test.tsx b/ui/litellm-dashboard/src/components/leftnav.test.tsx index 1c44d4dae6d..bab51aab121 100644 --- a/ui/litellm-dashboard/src/components/leftnav.test.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.test.tsx @@ -693,6 +693,7 @@ describe("Sidebar (leftnav)", () => { expectHref("Playground", "/ui/playground"); expectHref("Models + Endpoints", "/ui/models-and-endpoints"); expectHref("Usage", "/ui/usage"); + expectHref("Alerts", "/ui/alerts"); expectHref("API Reference", "/ui/api-reference"); }); diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 488113d726e..33113b4ad6f 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -299,6 +299,13 @@ const menuGroups: MenuGroup[] = [ roles: [...all_admin_roles, ...internalUserRoles], label: "Usage", }, + { + key: "alerts", + page: "alerts", + label: "Alerts", + icon: , + roles: rolesWithCapability("viewGlobalSpend"), + }, { key: "model-insights", page: "model-insights", diff --git a/ui/litellm-dashboard/src/components/model_info_view.integration.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.integration.test.tsx index 82221826a68..a42db06d251 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.integration.test.tsx @@ -1,10 +1,9 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import React, { ReactNode } from "react"; -import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import ModelInfoView from "./model_info_view"; -import ModelDeprecationBanner from "@/app/(dashboard)/models-and-endpoints/components/ModelDeprecationBanner"; import { toast } from "@/lib/toast"; import * as networking from "./networking"; import { useUserLookup } from "@/app/(dashboard)/hooks/users/useUsers"; @@ -46,10 +45,6 @@ vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ useModelCostMap: (...args: any[]) => mockUseModelCostMap(...args), })); -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ - default: () => ({ accessToken: "test-token", userRole: "Admin" }), -})); - const mockUseTeams = vi.fn(); vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ useTeams: () => mockUseTeams(), @@ -82,8 +77,6 @@ const mockModelDeleteCall = vi.mocked(networking.modelDeleteCall); const mockCredentialCreateCall = vi.mocked(networking.credentialCreateCall); const mockVectorStoreListCall = vi.mocked(networking.vectorStoreListCall); -const TEST_TODAY = new Date(2026, 10, 4); - describe("ModelInfoView", () => { let queryClient: QueryClient; @@ -202,10 +195,6 @@ describe("ModelInfoView", () => { mockCredentialCreateCall.mockResolvedValue({}); }); - afterEach(() => { - vi.unstubAllGlobals(); - }); - const wrapper = ({ children }: { children: ReactNode }) => React.createElement(QueryClientProvider, { client: queryClient }, children); @@ -216,70 +205,6 @@ describe("ModelInfoView", () => { }); }); - it("refreshes deprecation banner data after updating model credentials", async () => { - const deprecatedResponse = { - checked_at: "2026-06-01T00:00:00Z", - warn_within_days: 30, - deprecated: [ - { - model_name: "old-model", - litellm_provider: null, - deprecation_date: "2026-06-01", - days_until_deprecation: -1, - status: "deprecated", - successor_model: null, - }, - ], - imminent: [], - upcoming: [], - }; - const currentResponse = { - checked_at: "2026-06-02T00:00:00Z", - warn_within_days: 30, - deprecated: [], - imminent: [ - { - model_name: "current-model", - litellm_provider: null, - deprecation_date: "2026-11-30", - days_until_deprecation: 26, - status: "imminent", - successor_model: null, - }, - ], - upcoming: [], - }; - const fetchMock = vi - .fn() - .mockResolvedValueOnce(Response.json(deprecatedResponse)) - .mockResolvedValueOnce(Response.json(currentResponse)); - vi.stubGlobal("fetch", fetchMock); - - render( - <> - - - , - { wrapper }, - ); - - expect(await screen.findByText("old-model was deprecated on 2026-06-01")).toBeInTheDocument(); - await userEvent.click(screen.getByTestId("update-api-key-button")); - const dialog = screen.getByRole("dialog"); - fireEvent.change(within(dialog).getByPlaceholderText("Enter the new API key"), { - target: { value: "new-api-key" }, - }); - await userEvent.click(within(dialog).getByRole("button", { name: "Update API Key" })); - - expect(await screen.findByText("current-model deprecates on 2026-11-30, in 26 days")).toBeInTheDocument(); - await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); - expect( - fetchMock.mock.calls.map( - ([input]) => new URL(input instanceof Request ? input.url : String(input), "http://localhost").pathname, - ), - ).toEqual(["/model/deprecations", "/model/deprecations"]); - }); - it("shows Kubernetes pods for discovery-enabled models", async () => { mockUseModelsInfo.mockReturnValue({ data: { diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index e0498fd1ea9..efa7caf550d 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -1,6 +1,5 @@ import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; import { useModelHub, useModelsInfo } from "@/app/(dashboard)/hooks/models/useModels"; -import { invalidateModelQueries } from "@/app/(dashboard)/hooks/models/invalidateModelQueries"; import { useQueryClient } from "@tanstack/react-query"; import { transformModelData } from "@/app/(dashboard)/models-and-endpoints/utils/modelDataTransformer"; import { KeyIcon, RefreshIcon, TrashIcon } from "@heroicons/react/outline"; @@ -825,7 +824,7 @@ export default function ModelInfoView({ accessToken={accessToken} modelId={modelId} onUpdated={() => { - void invalidateModelQueries(queryClient); + queryClient.invalidateQueries({ queryKey: ["models", "list"] }); }} /> )} diff --git a/ui/litellm-dashboard/src/components/page_metadata.ts b/ui/litellm-dashboard/src/components/page_metadata.ts index 97220a0ef76..80c3a9d2ed7 100644 --- a/ui/litellm-dashboard/src/components/page_metadata.ts +++ b/ui/litellm-dashboard/src/components/page_metadata.ts @@ -20,6 +20,7 @@ export const pageDescriptions: Record = { "tool-policies": "Configure tool use policies and permissions", "vector-stores": "Manage vector databases for embeddings", new_usage: "View usage analytics and metrics", + alerts: "Review and resolve model deprecation alerts across keys and teams", "cost-optimization": "Track and configure cost-saving features: prompt compression, caching, and auto routing", "model-insights": "Model Leaderboard: compare usage, spend, tokens, and task mix across this gateway", logs: "Access request and response logs", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a2073bb00d8..1f12db1c063 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -944,6 +944,44 @@ export interface paths { patch?: never; trace?: never; }; + "/alerting/model-deprecations/state": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Get Deprecation Alert State */ + get: operations["get_deprecation_alert_state_alerting_model_deprecations_state_get"]; + /** Update Deprecation Alert State */ + put: operations["update_deprecation_alert_state_alerting_model_deprecations_state_put"]; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/alerting/model-deprecations/tags": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Deprecation Usage Tags + * @description Distinct tags observed with this model for a page of affected reporting key IDs. + */ + get: operations["get_deprecation_usage_tags_alerting_model_deprecations_tags_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/alerting/settings": { parameters: { query?: never; @@ -32376,6 +32414,59 @@ export interface components { /** Warnings */ warnings: components["schemas"]["DeploymentWarning"][]; }; + /** DeprecationAlertState */ + DeprecationAlertState: { + /** + * Deprecation Date + * Format: date + */ + deprecation_date: string; + /** Last Seen Date */ + last_seen_date?: string | null; + /** + * Last Seen Requests + * @default 0 + */ + last_seen_requests: number; + /** Model */ + model: string; + /** @default open */ + status: components["schemas"]["DeprecationAlertStatus"]; + /** Updated At */ + updated_at?: string | null; + /** Updated By */ + updated_by?: string | null; + }; + /** + * DeprecationAlertStatus + * @enum {string} + */ + DeprecationAlertStatus: "open" | "resolved" | "dismissed"; + /** DeprecationAlertUpdate */ + DeprecationAlertUpdate: { + /** + * Deprecation Date + * Format: date + */ + deprecation_date: string; + /** + * Last Seen Date + * Format: date + */ + last_seen_date: string; + /** Last Seen Requests */ + last_seen_requests: number; + /** Model */ + model: string; + status: components["schemas"]["DeprecationAlertStatus"]; + }; + /** DeprecationTagUsage */ + DeprecationTagUsage: { + /** Api Key */ + api_key: string; + /** Tag */ + tag: string; + }; /** DiscoverAgentRequest */ DiscoverAgentRequest: { /** @@ -53157,6 +53248,105 @@ export interface operations { }; }; }; + get_deprecation_alert_state_alerting_model_deprecations_state_get: { + parameters: { + query: { + model: string; + deprecation_date: string; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DeprecationAlertState"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + update_deprecation_alert_state_alerting_model_deprecations_state_put: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["DeprecationAlertUpdate"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DeprecationAlertState"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_deprecation_usage_tags_alerting_model_deprecations_tags_get: { + parameters: { + query: { + model: string; + api_keys: string[]; + start_date: string; + end_date: string; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DeprecationTagUsage"][]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; alerting_settings_alerting_settings_get: { parameters: { query?: never; From 5a392b9e5f22b230be19f561d4e12e3f099781bc Mon Sep 17 00:00:00 2001 From: quietjules <349048734+quietjules@users.noreply.github.com> Date: Wed, 4 Nov 2026 20:07:11 -0800 Subject: [PATCH 11/39] fix(bedrock)!: keep Claude's thinking when response_format json_schema is set (#49696) * fix(anthropic): stop forcing the response_format json tool when Claude will think A forced tool call skips Claude's thinking. The json_schema emulation forced json_tool_call unless thinking.type was enabled, so adaptive thinking and the models that think by default (Opus 5, Sonnet 5, Haiku 5.5, and the always-on families) lost their thinking on every structured-output request. AnthropicModelInfo.response_format_tool_can_be_forced is now the one check the Anthropic, Bedrock Converse (including the native output fallback) and Databricks paths use. It leaves the tool unforced when thinking is requested in any form or the model thinks by default, and keeps forcing it when the request turns thinking off (thinking disabled or reasoning_effort none) on a model that accepts forced tools. The new thinking_on_by_default cost map key marks Opus 5, Sonnet 5 and Haiku 5.5 on the providers that reach these paths, and a miss-only fallback rule covers unmapped Claude 5 and later ids. * fix(bedrock): route native output and thinking token counts by whether Claude will think Bedrock's native output opt-in now uses the same will-think check as the json tool, so Claude 5 models that think by default with no thinking param go native. Converse also asks for /usage/output_tokens_details on those models, so reasoning tokens are reported for them. * fix(anthropic): tell Claude to answer through the unforced response_format tool Moves the integration tests that check the forced JSON tool from Claude 5 models to Opus 4.8, since Claude 5 models now keep their default thinking and take the unforced tool or native output * fix(anthropic): keep open-object schemas on the json tool when Claude will think Native output closes every object, so an open object came back empty. Also declare thinking_on_by_default in the Rust model catalog and expect the Converse thinking-token field in test_main * fix(anthropic): hold streamed prose until the JSON tool decides the answer In json mode with the unforced JSON tool, Claude can write its working before it calls the tool. The stream now holds that text, drops it when the JSON tool answers, and sends it when Claude calls a caller tool or skips the JSON tool, matching the non-streamed reply. * test(anthropic): build the stream test events without nested comprehensions * fix(anthropic): keep always-on thinking models unforced, open objects native when thinking, and the tool instruction Claude-only * test(anthropic): keep main's open-object integration cases and cover the unforced tool on a default-thinking model * fix(anthropic): treat thinking False as thinking off so the JSON tool stays forced Cover thinking False on the Anthropic and Converse forcing tests, and pin that a Converse json_schema stream with the unforced JSON tool never sends Claude's working text as content * fix(anthropic): drop the thinking False forcing case that completion never reaches litellm.completion turns thinking False into no thinking param before any route sees it, so the case was dead on the request path and False keeps meaning the model's default * test(bedrock): collect the Converse json_schema stream chunks in one assignment --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/constants.py | 3 + litellm/llms/anthropic/chat/handler.py | 40 +++- litellm/llms/anthropic/chat/transformation.py | 31 +++- litellm/llms/anthropic/common_utils.py | 28 +++ .../bedrock/chat/converse_transformation.py | 59 ++++-- .../anthropic_claude_transformation.py | 17 +- .../bedrock/chat/mantle/transformation.py | 6 +- litellm/llms/bedrock/common_utils.py | 8 +- .../llms/databricks/chat/transformation.py | 13 +- ...odel_prices_and_context_window_backup.json | 58 ++++++ litellm/types/utils.py | 1 + litellm/utils.py | 1 + model_prices_and_context_window.json | 58 ++++++ ...t_bedrock_native_structured_output_wire.py | 12 +- ...vertex_anthropic_structured_output_wire.py | 58 ++++-- .../test_reasoning_effort_levels_wire.py | 15 +- ...anthropic_json_tool_cut_stream_sync_sdk.py | 2 +- ...t_completions_response_format_anthropic.py | 22 +-- ...st_chat_completions_streaming_anthropic.py | 22 +-- .../test_fallback_generalizations.py | 17 ++ .../chat/test_anthropic_chat_handler.py | 116 +++++++++++- .../test_anthropic_chat_transformation.py | 85 +++++++++ ...mations_anthropic_claude_transformation.py | 35 ++-- .../chat/test_converse_transformation.py | 173 ++++++++++++++---- .../llms/bedrock/chat/test_invoke_handler.py | 102 +++++++++++ .../test_databricks_chat_transformation.py | 45 +++++ tests/unit/test_main.py | 1 + tests/unit/test_utils.py | 1 + 28 files changed, 887 insertions(+), 142 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 32b47bc6b46..50e090eabaf 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1673,6 +1673,9 @@ OPENAI_FINISH_REASONS: Final = [ ] HUMANLOOP_PROMPT_CACHE_TTL_SECONDS: Final = get_env_int("HUMANLOOP_PROMPT_CACHE_TTL_SECONDS", 60) # 1 minute RESPONSE_FORMAT_TOOL_NAME = "json_tool_call" # default tool name used when converting response format to tool call +RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION: Final = ( + "When you are ready to answer, call this tool with your final answer as its input. Do not reply in plain text." +) ########################### Logging Callback Constants ########################### AZURE_STORAGE_MSFT_VERSION: Final = "2019-07-07" diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 4eae06dbcdf..c273c715ba6 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -734,6 +734,7 @@ class ModelResponseIterator: self._split_answer_blocks: tuple[str, ...] = () self._withheld_answer_close: str = "" self._streamed_answer_keys: frozenset[str] | None = None + self._held_prose: str = "" # For handling partial JSON chunks from fragmentation # See: https://github.com/BerriAI/litellm/issues/17473 @@ -826,8 +827,13 @@ class ModelResponseIterator: thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] = [] if "text" in content_block["delta"]: - text = content_block["delta"]["text"] - return text, tool_use, thinking_blocks, provider_specific_fields, reasoning_content + return ( + self._hold_prose(content_block["delta"]["text"]), + tool_use, + thinking_blocks, + provider_specific_fields, + reasoning_content, + ) self.content_blocks.append(content_block) if "partial_json" in content_block["delta"]: # Only emit tool calls if we're in a tool_use or server_tool_use block @@ -1025,7 +1031,7 @@ class ModelResponseIterator: if self.current_content_block_type in _THINKING_BLOCK_TYPES: self.thinking_block_index += 1 if content_block_start["content_block"]["type"] == "text": - text = self._release_answer_close() + content_block_start["content_block"]["text"] + text = self._release_answer_close() + self._hold_prose(content_block_start["content_block"]["text"]) elif ( content_block_start["content_block"]["type"] == "tool_use" or content_block_start["content_block"]["type"] == "server_tool_use" @@ -1161,7 +1167,7 @@ class ModelResponseIterator: elif type_chunk == "message_delta": finish_reason, usage, container, stop_reason_fields = self._handle_message_delta(chunk) provider_specific_fields.update(stop_reason_fields) - text = self._release_answer_close() + self._drain_split_answer_blocks() + text = self._release_answer_close() + self._drain_split_answer_blocks() + self._release_held_prose() if container: provider_specific_fields["container"] = container # Store container_id and re-emit code_interpreter_results @@ -1211,7 +1217,9 @@ class ModelResponseIterator: status_code=500, # it looks like Anthropic API does not return a status code in the chunk error - default to 500 ) - text, tool_use = self._handle_json_mode_chunk(text=text, tool_use=tool_use) + text, tool_use = self._handle_json_mode_chunk( + text=text + self._settle_held_prose(tool_use), tool_use=tool_use + ) refusal_fields: Final = refusal_params(provider_specific_fields) returned_chunk: Final = ModelResponseStream( @@ -1309,6 +1317,24 @@ class ModelResponseIterator: self._response_format_verdict = _envelope_verdict(buffered, schema, RESPONSE_FORMAT_TOOL_NAME) return self._withhold_answer_close(buffered) if self._response_format_verdict == "pass_through" else "" + def _hold_prose(self, text: str) -> str: + if not self.json_mode or self.response_format_schema is None: + return text + self._held_prose += text + return "" + + def _settle_held_prose(self, tool_use: ChatCompletionToolCallChunk | None) -> str: + if tool_use is None or tool_use.get("id") is None or self.current_content_block_type != "tool_use": + return "" + held: Final = self._held_prose + self._held_prose = "" + return "" if tool_use.get("function", {}).get("name") == RESPONSE_FORMAT_TOOL_NAME else held + + def _release_held_prose(self) -> str: + held: Final = self._held_prose + self._held_prose = "" + return "" if self.converted_response_format_tool else held + def _withhold_answer_close(self, text: str) -> str: pending: Final = self._withheld_answer_close + text body: Final = pending.rstrip() @@ -1362,7 +1388,9 @@ class ModelResponseIterator: def cut_stream_chunk(self) -> ModelResponseStream | None: unsent: Final = self._flush_response_format_buffer() if self._response_format_buffer.getvalue() else "" - text: Final = unsent + self._release_answer_close() + self._drain_split_answer_blocks() + text: Final = ( + unsent + self._release_answer_close() + self._drain_split_answer_blocks() + self._release_held_prose() + ) if not text: return None return ModelResponseStream( diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 09d16181c3a..4f890ed922f 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -26,6 +26,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, RESPONSE_FORMAT_TOOL_NAME, + RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION, ) from litellm.litellm_core_utils.asyncify import is_event_loop_running from litellm.litellm_core_utils.core_helpers import ( @@ -1786,7 +1787,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): json_schema: Mapping[str, object] | None, request_params: Mapping[str, object], ) -> bool: - open_object_takes_forced_json_tool: Final = ( + open_object_takes_json_tool: Final = ( json_schema is not None and schema_has_open_object(json_schema) and not self.is_thinking_requested(request_params) @@ -1794,7 +1795,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): and not request_params.get("tools") and not AnthropicModelInfo.forced_tool_use_unsupported(model) ) - return self.uses_native_structured_output(model) and not open_object_takes_forced_json_tool + return self.uses_native_structured_output(model) and not open_object_takes_json_tool def native_output_format_schema(self, json_schema: Mapping[str, object]) -> Mapping[str, object] | None: return self.filter_anthropic_output_schema(dict(self._defs_unpacked(json_schema))) @@ -1809,7 +1810,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return AnthropicOutputSchema(type="json_schema", schema=dict(native_schema)) def map_response_format_to_anthropic_tool( - self, value: dict | None, optional_params: dict, is_thinking_enabled: bool + self, value: dict | None, optional_params: dict, tool_is_forced: bool ) -> AnthropicMessagesTool | None: ignore_response_format_types: Final = ["text"] if value is None or value["type"] in ignore_response_format_types: # value is a no-op @@ -1827,6 +1828,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): _tool: Final = self._create_json_tool_call_for_response_format( json_schema=json_schema, + description=None if tool_is_forced else RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION, ) return _tool @@ -1929,7 +1931,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): drop_params=drop_params, runs_hosted_mcp=self.runs_hosted_mcp, ) - is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params=request_params) + response_format_tool_is_forced: Final = ( + "response_format" in request_params + and AnthropicModelInfo.response_format_tool_can_be_forced( + model, self._resolved_provider, REQUEST_PARAMS.validate_python(request_params) + ) + ) # NB: ``map_openai_params`` deliberately does NOT sanitize tool names # here. Names are the *original* OpenAI names at this stage, and must @@ -1989,10 +1996,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if _output_format is not None: optional_params["output_format"] = _output_format else: - _tool = self.map_response_format_to_anthropic_tool(value, optional_params, is_thinking_enabled) + _tool = self.map_response_format_to_anthropic_tool( + value, optional_params, response_format_tool_is_forced + ) if _tool is None: continue - if not is_thinking_enabled and not AnthropicModelInfo.forced_tool_use_unsupported(model): + if response_format_tool_is_forced: _tool_choice = { "name": RESPONSE_FORMAT_TOOL_NAME, "type": "tool", @@ -2121,12 +2130,14 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _create_json_tool_call_for_response_format( self, json_schema: dict | None = None, + description: str | None = None, ) -> AnthropicMessagesTool: """ Handles creating a tool call for getting responses in JSON format. Args: json_schema (Optional[dict]): The JSON schema the response should be in + description (Optional[str]): The tool description, left off the tool when None Returns: AnthropicMessagesTool: The tool call to send to Anthropic API to get responses in JSON format @@ -2143,11 +2154,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): else: _input_schema.update(cast(AnthropicInputSchema, json_schema)) - _tool: Final = AnthropicMessagesTool( - name=RESPONSE_FORMAT_TOOL_NAME, - input_schema=_input_schema, + if description is None: + return AnthropicMessagesTool(name=RESPONSE_FORMAT_TOOL_NAME, input_schema=_input_schema) + return AnthropicMessagesTool( + name=RESPONSE_FORMAT_TOOL_NAME, input_schema=_input_schema, description=description ) - return _tool def should_strip_billing_metadata(self) -> bool: """ diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index bfd1341711a..bc9bd769d83 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1141,6 +1141,34 @@ class AnthropicModelInfo(BaseLLMModelInfo): """ return AnthropicModelInfo._supports_model_capability(model, "thinking_always_on", custom_llm_provider) + @staticmethod + def thinks_by_default(model: str, custom_llm_provider: str) -> bool: + return AnthropicModelInfo._is_always_on_thinking_model( + model, custom_llm_provider + ) or AnthropicModelInfo._supports_model_capability(model, "thinking_on_by_default", custom_llm_provider) + + @staticmethod + def will_think(model: str, custom_llm_provider: str, request_params: Mapping[str, object]) -> bool: + if AnthropicModelInfo._is_always_on_thinking_model(model, custom_llm_provider): + return True + match request_params: + case {"thinking": True | {"type": "enabled" | "adaptive"}}: + return True + case {"reasoning_effort": effort} if effort is not None and effort != "none": + return True + case {"thinking": {"type": "disabled"}} | {"reasoning_effort": "none"}: + return False + case _: + return AnthropicModelInfo.thinks_by_default(model, custom_llm_provider) + + @staticmethod + def response_format_tool_can_be_forced( + model: str, custom_llm_provider: str, request_params: Mapping[str, object] + ) -> bool: + return not AnthropicModelInfo.forced_tool_use_unsupported(model) and not AnthropicModelInfo.will_think( + model, custom_llm_provider, request_params + ) + @staticmethod def _supports_legacy_thinking(model: str, custom_llm_provider: str) -> bool: """Whether ``model`` is an adaptive-thinking model that still accepts legacy diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 056f34f1bf4..03fbc35a97f 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -22,6 +22,7 @@ from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_hea from litellm.constants import ( BEDROCK_MIN_THINKING_BUDGET_TOKENS, RESPONSE_FORMAT_TOOL_NAME, + RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION, ) from litellm.integrations.anthropic_cache_control_hook import ( MAX_CACHE_CONTROL_BLOCKS, @@ -950,19 +951,24 @@ class AmazonConverseConfig(BaseConfig): def _supports_native_structured_outputs( model: str, custom_llm_provider: str | None, + will_think: bool, thinking_requested: bool, model_info: Mapping[str, object] | None, ) -> bool: """Check if the Bedrock model takes native structured outputs (outputConfig.textFormat) for this request. ``supports_native_structured_output`` in ``litellm.model_cost`` turns it on for every request, and - ``supports_bedrock_native_structured_output`` (the deployment's ``model_info`` first) also when the - request asks for thinking, which the json tool path would drop. Without thinking, a model that cannot - force the json tool gets it unless that key is false. + ``supports_bedrock_native_structured_output`` (the deployment's ``model_info`` first) also when Claude + will think, which the json tool path would drop. When the request does not ask for thinking, a model + that cannot force the json tool gets it unless that key is false. Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/structured-output.html """ return bedrock_chat_uses_native_structured_output( - model, thinking_requested, custom_llm_provider or "bedrock_converse", model_info + model, + will_think=will_think, + thinking_requested=thinking_requested, + custom_llm_provider=custom_llm_provider or "bedrock_converse", + model_info=model_info, ) @staticmethod @@ -1005,11 +1011,13 @@ class AmazonConverseConfig(BaseConfig): self, model: str, json_schema: Mapping[str, object], + *, + will_think: bool, thinking_requested: bool, model_info: Mapping[str, object] | None, ) -> Mapping[str, object] | None: if not self._supports_native_structured_outputs( - model, self.custom_llm_provider, thinking_requested, model_info + model, self.custom_llm_provider, will_think, thinking_requested, model_info ): return None if BedrockModelInfo.get_base_model(model).startswith("anthropic"): @@ -1095,7 +1103,6 @@ class AmazonConverseConfig(BaseConfig): request_params: Final = tool_params_without_hosted_mcp( non_default_params, model=model, route="bedrock", drop_params=drop_params ) - is_thinking_enabled: Final = self.is_thinking_enabled(request_params) base_model: Final = BedrockModelInfo.get_base_model(model) drop_thinking_param: Final = self._is_deepseek_model(model=model, base_model=base_model) drop_reasoning_effort_param: Final = self._model_rejects_reasoning_effort_param( @@ -1109,7 +1116,6 @@ class AmazonConverseConfig(BaseConfig): model=model, optional_params=optional_params, non_default_params=request_params, - is_thinking_enabled=is_thinking_enabled, model_info=model_info, ) if param == "max_tokens" or param == "max_completion_tokens": @@ -1314,7 +1320,6 @@ class AmazonConverseConfig(BaseConfig): model: str, optional_params: dict, non_default_params: dict, - is_thinking_enabled: bool, model_info: Mapping[str, object] | None = None, ) -> dict: """ @@ -1339,9 +1344,16 @@ class AmazonConverseConfig(BaseConfig): if "type" in value and value["type"] == "text": return optional_params - thinking_requested: Final = self.is_thinking_requested(REQUEST_PARAMS.validate_python(non_default_params)) + request_params: Final = REQUEST_PARAMS.validate_python(non_default_params) + thinking_requested: Final = self.is_thinking_requested(request_params) native_schema: Final = ( - self._native_output_schema(model, json_schema, thinking_requested, model_info) + self._native_output_schema( + model, + json_schema, + will_think=AnthropicModelInfo.will_think(model, "bedrock", request_params), + thinking_requested=thinking_requested, + model_info=model_info, + ) if json_schema is not None else None ) @@ -1359,17 +1371,24 @@ class AmazonConverseConfig(BaseConfig): elif json_schema is not None: # Fallback: translate to a synthetic tool call # https://docs.anthropic.com/en/docs/build-with-claude/tool-use#json-mode + tool_is_forced: Final = litellm.utils.supports_tool_choice( + model=model, custom_llm_provider=self.custom_llm_provider + ) and AnthropicModelInfo.response_format_tool_can_be_forced(model, "bedrock", request_params) + schema_description: Final = description if isinstance(description, str) else None + instructs_unforced_tool: Final = not tool_is_forced and BedrockModelInfo.get_base_model(model).startswith( + "anthropic" + ) _tool: Final = self._create_json_tool_call_for_response_format( json_schema=json_schema, - description=description, + description=( + " ".join(filter(None, (RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION, schema_description))) + if instructs_unforced_tool + else schema_description + ), ) optional_params = self._add_tools_to_optional_params(optional_params=optional_params, tools=[_tool]) - if ( - litellm.utils.supports_tool_choice(model=model, custom_llm_provider=self.custom_llm_provider) - and not is_thinking_enabled - and not AnthropicModelInfo.forced_tool_use_unsupported(model) - ): + if tool_is_forced: optional_params["tool_choice"] = ToolChoiceValuesBlock( tool=SpecificToolChoiceBlock(name=RESPONSE_FORMAT_TOOL_NAME) ) @@ -1380,8 +1399,8 @@ class AmazonConverseConfig(BaseConfig): # returns {} instead of the requested JSON. The model already # returns JSON when the prompt asks for it. - # A request that goes native only because the model can't force the json tool keeps the json tool's faked - # stream, so it still needs only bedrock:InvokeModel. + # A request that goes native without asking for thinking keeps the json tool's faked stream, so it still + # needs only bedrock:InvokeModel, unless the model takes native output on every request. streams_natively: Final = native_schema is not None and bedrock_chat_opts_into_native_structured_output( model, thinking_requested, self.custom_llm_provider or "bedrock_converse", model_info ) @@ -2219,8 +2238,8 @@ class AmazonConverseConfig(BaseConfig): } if additional_request_params: data["additionalModelRequestFields"] = additional_request_params - if "thinking" in additional_request_params: - data["additionalModelResponseFieldPaths"] = ("/usage/output_tokens_details",) + if "thinking" in additional_request_params or AnthropicModelInfo.thinks_by_default(model, "bedrock"): + data["additionalModelResponseFieldPaths"] = ("/usage/output_tokens_details",) if system_content_blocks: data["system"] = system_content_blocks diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py index dc85028440f..62e1bee5217 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py @@ -142,14 +142,19 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): json_schema: Mapping[str, object] | None, request_params: Mapping[str, object], ) -> bool: - return self.chat_uses_native_structured_output( - model, self.is_thinking_requested(request_params), self._deployment_model_info - ) and (json_schema is None or bedrock_cleaned_output_schema(model, json_schema) is not None) + return self.chat_uses_native_structured_output(model, request_params, self._deployment_model_info) and ( + json_schema is None or bedrock_cleaned_output_schema(model, json_schema) is not None + ) def chat_uses_native_structured_output( - self, model: str, thinking_requested: bool, model_info: Mapping[str, object] | None + self, model: str, request_params: Mapping[str, object], model_info: Mapping[str, object] | None ) -> bool: - return bedrock_chat_uses_native_structured_output(model, thinking_requested, model_info=model_info) + return bedrock_chat_uses_native_structured_output( + model, + will_think=self.will_think(model, "bedrock", request_params), + thinking_requested=self.is_thinking_requested(request_params), + model_info=model_info, + ) def uses_opt_in_native_output( self, model: str, optional_params: Mapping[str, object], litellm_params: Mapping[str, object] @@ -359,7 +364,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): request_body=anthropic_request, model_info=model_info, chat_uses_native_output=lambda: self.chat_uses_native_structured_output( - model, self.is_thinking_requested(_NON_DEFAULT_PARAMS.validate_python(filtered_params)), model_info + model, _NON_DEFAULT_PARAMS.validate_python(filtered_params), model_info ), ) strip_unsupported_bedrock_invoke_output_config_keys( diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index 952ca7a2cd4..7e635a61d4d 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -55,9 +55,11 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): ) def chat_uses_native_structured_output( - self, model: str, thinking_requested: bool, model_info: Mapping[str, object] | None + self, model: str, request_params: Mapping[str, object], model_info: Mapping[str, object] | None ) -> bool: - return bedrock_chat_opts_into_native_structured_output(model, thinking_requested, model_info=model_info) + return bedrock_chat_opts_into_native_structured_output( + model, self.will_think(model, "bedrock", request_params), model_info=model_info + ) def _default_max_tokens_models(self, model: str, litellm_params: Mapping[str, object]) -> tuple[str, ...]: return (f"bedrock_mantle/{model}", *super()._default_max_tokens_models(model, litellm_params)) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index c1b3b5616b8..dfee28a7cc9 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -356,13 +356,15 @@ def bedrock_chat_opts_into_native_structured_output( def bedrock_chat_uses_native_structured_output( model: str, + *, + will_think: bool, thinking_requested: bool, custom_llm_provider: str = "bedrock", model_info: Mapping[str, object] | None = None, ) -> bool: - return bedrock_chat_opts_into_native_structured_output( - model, thinking_requested, custom_llm_provider, model_info - ) or (not thinking_requested and _bedrock_native_output_replaces_unforceable_tool(model, model_info)) + return bedrock_chat_opts_into_native_structured_output(model, will_think, custom_llm_provider, model_info) or ( + not thinking_requested and _bedrock_native_output_replaces_unforceable_tool(model, model_info) + ) def bedrock_native_output_is_deployment_opt_in( diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index d943718ed73..859211bd40d 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -31,6 +31,7 @@ from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.base_utils import ( type_to_response_format_param, # pyright: ignore[reportUnknownVariableType] # base_utils helper returns an untyped dict ) +from litellm.llms.base_llm.chat.transformation import REQUEST_PARAMS from litellm.types.llms.anthropic import AllAnthropicToolsValues from litellm.types.llms.databricks import ( AllDatabricksContentValues, @@ -541,12 +542,12 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): model: str, value: dict | None, optional_params: dict, - is_thinking_enabled: bool, + tool_is_forced: bool, ) -> DatabricksTool | None: if value is None: return None - tool: Final = self.map_response_format_to_anthropic_tool(value, optional_params, is_thinking_enabled) + tool: Final = self.map_response_format_to_anthropic_tool(value, optional_params, tool_is_forced) databricks_tool: Final = self.convert_anthropic_tool_to_databricks_tool(tool) return databricks_tool @@ -579,7 +580,6 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): if "claude" in model else non_default_params ) - is_thinking_enabled: Final = self.is_thinking_enabled(request_params) mapped_params: Final = super().map_openai_params(request_params, optional_params, model, drop_params) if "claude" in model: AnthropicConfig.translate_legacy_thinking_for_adaptive_model( @@ -596,17 +596,20 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): mapped_params.pop("max_completion_tokens", None) if "response_format" in request_params and "claude" in model: + tool_is_forced: Final = AnthropicConfig.response_format_tool_can_be_forced( + model, "databricks", REQUEST_PARAMS.validate_python(request_params) + ) _tool: Final = self.map_response_format_to_databricks_tool( model, request_params["response_format"], mapped_params, - is_thinking_enabled, + tool_is_forced, ) if _tool is not None: self._add_tools_to_optional_params(optional_params=optional_params, tools=[_tool]) optional_params["json_mode"] = True - if not is_thinking_enabled: + if tool_is_forced: _tool_choice: Final = ChatCompletionToolChoiceObjectParam( type="function", function=ChatCompletionToolChoiceFunctionParam(name=RESPONSE_FORMAT_TOOL_NAME), diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4d576443b30..a292ce0d14e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -2125,6 +2125,7 @@ "anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -2213,6 +2214,7 @@ "global.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -2301,6 +2303,7 @@ "us.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, @@ -2389,6 +2392,7 @@ "eu.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, @@ -2472,6 +2476,7 @@ "au.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, @@ -2560,6 +2565,7 @@ "jp.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, @@ -2958,6 +2964,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -3002,6 +3009,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -3046,6 +3054,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -3090,6 +3099,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -3134,6 +3144,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -3178,6 +3189,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -4063,6 +4075,7 @@ "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "input_cost_per_token": 5e-06, "output_cost_per_token": 2.5e-05, "litellm_provider": "azure_ai", @@ -4204,6 +4217,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -16507,6 +16521,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -17042,6 +17057,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -21910,6 +21926,7 @@ "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_function_calling": true, "supports_mid_conversation_system": true, @@ -22068,6 +22085,7 @@ "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_function_calling": true, "supports_mid_conversation_system": true, @@ -50492,6 +50510,7 @@ "prompt_cache_min_tokens": 1024, "source": "https://aws.amazon.com/bedrock/pricing/", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -50569,6 +50588,7 @@ "prompt_cache_min_tokens": 512, "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-west-1/index.json", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -52974,6 +52994,7 @@ "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_creation_input_token_cost_batches": 3.125e-06, @@ -53012,6 +53033,7 @@ "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_creation_input_token_cost_batches": 3.125e-06, @@ -53261,6 +53283,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -60636,6 +60659,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -64770,6 +64794,14 @@ "thinking_always_on": true } }, + { + "name": "claude-thinking-on-by-default", + "pattern": "claude-[a-z]+-[5-9](?!\\d)", + "description": "Claude at major 5 or higher, in any id shape that contains claude--. Anthropic's thinking docs list every Claude 5+ model as thinking by default (on or always on) when the request sends no thinking param, and every 4.x as off. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. The response_format json tool is then left unforced, since a forced tool call skips the model's thinking.", + "model_info": { + "thinking_on_by_default": true + } + }, { "name": "claude-mid-conversation-system", "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", @@ -69293,6 +69325,7 @@ "output_cost_per_token": 1.2e-05, "prompt_cache_min_tokens": 1024, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -69368,6 +69401,7 @@ "output_cost_per_token": 3e-05, "prompt_cache_min_tokens": 512, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -69624,6 +69658,7 @@ "output_cost_per_token": 1.2e-05, "prompt_cache_min_tokens": 1024, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -69689,6 +69724,7 @@ "output_cost_per_token": 3e-05, "prompt_cache_min_tokens": 512, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -75440,6 +75476,7 @@ "apac.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "litellm_provider": "bedrock_converse", "supports_tool_search": true, @@ -75556,6 +75593,7 @@ "max_output_tokens": 128000, "mode": "chat", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77060,6 +77098,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77112,6 +77151,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77156,6 +77196,7 @@ "thinking_always_on": false, "prompt_cache_min_tokens": 512, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -77203,6 +77244,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77239,6 +77281,7 @@ "max_output_tokens": 128000, "mode": "chat", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77295,6 +77338,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77341,6 +77385,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -77381,6 +77426,7 @@ "prompt_cache_min_tokens": 512, "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -77422,6 +77468,7 @@ "prompt_cache_min_tokens": 512, "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -77468,6 +77515,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77519,6 +77567,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77565,6 +77614,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77625,6 +77675,7 @@ "prompt_cache_min_tokens": 512, "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -77676,6 +77727,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77726,6 +77778,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -77770,6 +77823,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -79529,6 +79583,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -79743,6 +79798,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -79909,6 +79965,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -81894,6 +81951,7 @@ "prompt_cache_min_tokens": 512.0, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_anthropic_thinking_payload": true, "supports_assistant_prefill": false, "supports_function_calling": true, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4c47489e08f..ba29a67a36c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -224,6 +224,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_adaptive_thinking: bool | None supports_legacy_thinking: ReadOnly[bool | None] thinking_always_on: ReadOnly[bool | None] + thinking_on_by_default: ReadOnly[bool | None] supports_tool_search: bool | None supports_mid_conversation_system: bool | None supports_url_context: bool | None diff --git a/litellm/utils.py b/litellm/utils.py index 4add9d53cb3..7f81ece8314 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7581,6 +7581,7 @@ def get_model_info_helper( supports_adaptive_thinking=_model_info.get("supports_adaptive_thinking", None), supports_legacy_thinking=_model_info.get("supports_legacy_thinking", None), thinking_always_on=_model_info.get("thinking_always_on", None), + thinking_on_by_default=_model_info.get("thinking_on_by_default", None), supports_tool_search=_model_info.get("supports_tool_search", None), supports_mid_conversation_system=_model_info.get("supports_mid_conversation_system", None), supports_anthropic_thinking_payload=_model_info.get("supports_anthropic_thinking_payload", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4d576443b30..a292ce0d14e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -2125,6 +2125,7 @@ "anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -2213,6 +2214,7 @@ "global.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -2301,6 +2303,7 @@ "us.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, @@ -2389,6 +2392,7 @@ "eu.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, @@ -2472,6 +2476,7 @@ "au.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, @@ -2560,6 +2565,7 @@ "jp.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, @@ -2958,6 +2964,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -3002,6 +3009,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -3046,6 +3054,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -3090,6 +3099,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -3134,6 +3144,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -3178,6 +3189,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -4063,6 +4075,7 @@ "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "input_cost_per_token": 5e-06, "output_cost_per_token": 2.5e-05, "litellm_provider": "azure_ai", @@ -4204,6 +4217,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -16507,6 +16521,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -17042,6 +17057,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -21910,6 +21926,7 @@ "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_function_calling": true, "supports_mid_conversation_system": true, @@ -22068,6 +22085,7 @@ "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_function_calling": true, "supports_mid_conversation_system": true, @@ -50492,6 +50510,7 @@ "prompt_cache_min_tokens": 1024, "source": "https://aws.amazon.com/bedrock/pricing/", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -50569,6 +50588,7 @@ "prompt_cache_min_tokens": 512, "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-west-1/index.json", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -52974,6 +52994,7 @@ "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_creation_input_token_cost_batches": 3.125e-06, @@ -53012,6 +53033,7 @@ "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_creation_input_token_cost_batches": 3.125e-06, @@ -53261,6 +53283,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -60636,6 +60659,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -64770,6 +64794,14 @@ "thinking_always_on": true } }, + { + "name": "claude-thinking-on-by-default", + "pattern": "claude-[a-z]+-[5-9](?!\\d)", + "description": "Claude at major 5 or higher, in any id shape that contains claude--. Anthropic's thinking docs list every Claude 5+ model as thinking by default (on or always on) when the request sends no thinking param, and every 4.x as off. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. The response_format json tool is then left unforced, since a forced tool call skips the model's thinking.", + "model_info": { + "thinking_on_by_default": true + } + }, { "name": "claude-mid-conversation-system", "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", @@ -69293,6 +69325,7 @@ "output_cost_per_token": 1.2e-05, "prompt_cache_min_tokens": 1024, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -69368,6 +69401,7 @@ "output_cost_per_token": 3e-05, "prompt_cache_min_tokens": 512, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -69624,6 +69658,7 @@ "output_cost_per_token": 1.2e-05, "prompt_cache_min_tokens": 1024, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -69689,6 +69724,7 @@ "output_cost_per_token": 3e-05, "prompt_cache_min_tokens": 512, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -75440,6 +75476,7 @@ "apac.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "litellm_provider": "bedrock_converse", "supports_tool_search": true, @@ -75556,6 +75593,7 @@ "max_output_tokens": 128000, "mode": "chat", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77060,6 +77098,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77112,6 +77151,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77156,6 +77196,7 @@ "thinking_always_on": false, "prompt_cache_min_tokens": 512, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -77203,6 +77244,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77239,6 +77281,7 @@ "max_output_tokens": 128000, "mode": "chat", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77295,6 +77338,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77341,6 +77385,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -77381,6 +77426,7 @@ "prompt_cache_min_tokens": 512, "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -77422,6 +77468,7 @@ "prompt_cache_min_tokens": 512, "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -77468,6 +77515,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77519,6 +77567,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77565,6 +77614,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77625,6 +77675,7 @@ "prompt_cache_min_tokens": 512, "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -77676,6 +77727,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -77726,6 +77778,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -77770,6 +77823,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -79529,6 +79583,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -79743,6 +79798,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -79909,6 +79965,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -81894,6 +81951,7 @@ "prompt_cache_min_tokens": 512.0, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, + "thinking_on_by_default": true, "supports_anthropic_thinking_payload": true, "supports_assistant_prefill": false, "supports_function_calling": true, diff --git a/tests/integration/providers/test_bedrock_native_structured_output_wire.py b/tests/integration/providers/test_bedrock_native_structured_output_wire.py index 67edcf7f3e4..2502bd90bd6 100644 --- a/tests/integration/providers/test_bedrock_native_structured_output_wire.py +++ b/tests/integration/providers/test_bedrock_native_structured_output_wire.py @@ -622,10 +622,15 @@ def test_messages_thinking_on_a_bedrock_converse_native_model_gets_native_output @pytest.mark.parametrize("route", ["converse", "invoke"]) -def test_bedrock_native_model_without_thinking_keeps_the_forced_json_tool(gateway: Gateway, route: str) -> None: +@pytest.mark.parametrize( + ("backend", "thinks_by_default"), [(_OPUS_4_8, False), (_HAIKU, True)], ids=["opus_4_8", "haiku_5_5"] +) +def test_bedrock_native_model_with_no_thinking_param_gets_native_output_only_when_it_thinks_by_default( + gateway: Gateway, route: str, backend: str, thinks_by_default: bool +) -> None: marker, marker_hex = _marker() with wire_server(_bedrock) as wire, gateway.scenario() as scenario: - model: Final = _deployment(scenario, wire, route, _HAIKU) + model: Final = _deployment(scenario, wire, route, backend) response: Final = gateway.request( "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_response_format(_SCHEMA)) ) @@ -634,6 +639,9 @@ def test_bedrock_native_model_without_thinking_keeps_the_forced_json_tool(gatewa assert not message.get("tool_calls"), response.text _, body = _one_call(wire, marker_hex) invoke: Final = route == "invoke" + if thinks_by_default: + _assert_native_without_tools(body, invoke, _SCHEMA) + return assert _native_schema(body, invoke) is None, body assert _tools(body, invoke)[_JSON_TOOL]["properties"] == _SCHEMA["properties"], body diff --git a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py index db81ff70263..19928ecd6e0 100644 --- a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -34,14 +34,20 @@ from pydantic import JsonValue, TypeAdapter pytestmark: Final = pytest.mark.shares_worker_state _JSON_TOOL: Final = "json_tool_call" +_UNFORCED_JSON_TOOL_DESCRIPTION: Final = ( + "When you are ready to answer, call this tool with your final answer as its input. Do not reply in plain text." +) _USER_TOOL: Final = "lookup" _STRUCTURED_OUTPUTS_BETA: Final = "structured-outputs-2025-11-13" _TOOL_MODEL: Final = "claude-sonnet-5" _EFFORT_GATED_MODEL: Final = "claude-haiku-4-5" _NATIVE_MODEL: Final = "claude-opus-5-5" -_FORCED_TOOL_NATIVE_MODEL: Final = "claude-haiku-5-5" +_FORCED_TOOL_NATIVE_MODEL: Final = "claude-opus-4-8" +_THINKING_BY_DEFAULT_NATIVE_MODEL: Final = "claude-haiku-5-5" +_THINKING_OFF_TOOL_MODEL: Final = "claude-opus-4-8" _UNMAPPED_MODEL: Final = "claude-scripted-9" _BEDROCK_INVOKE_MODEL: Final = "anthropic.claude-sonnet-5" +_BEDROCK_TOOL_PATH_MODEL: Final = "anthropic.claude-opus-4-8" _BEDROCK_JSON_TOOL_MODEL: Final = "us.anthropic.claude-opus-4-8" _THINKING_SIGNATURE: Final = "c2NyaXB0ZWQtc2lnbmF0dXJl" _PROJECT: Final = "scripted-project" @@ -608,11 +614,15 @@ def _one_delivery(wire: Wire, marker_hex: str) -> _Delivery: return sent -def _json_tool_schema(delivery: _Delivery) -> dict[str, JsonValue]: +def _json_tool(delivery: _Delivery) -> dict[str, JsonValue]: tools: Final = tuple(object_value(tool) for tool in list_value(delivery.body["tools"])) (json_tool,) = tuple(tool for tool in tools if tool["name"] == _JSON_TOOL) - assert json_tool.keys() == {"name", "input_schema"}, json_tool - return object_value(json_tool["input_schema"]) + assert json_tool.keys() - {"description"} == {"name", "input_schema"}, json_tool + return json_tool + + +def _json_tool_schema(delivery: _Delivery) -> dict[str, JsonValue]: + return object_value(_json_tool(delivery)["input_schema"]) def _folded_format(delivery: _Delivery) -> dict[str, JsonValue]: @@ -830,10 +840,10 @@ def _family_deployment(gateway: Gateway, scenario: Scenario, api_base: str, memb case "azure_ai_unmapped": return scenario.model(model=f"azure_ai/{_UNMAPPED_MODEL}", api_base=api_base, api_key=_AZURE_KEY) case "bedrock_invoke": - return _bedrock_invoke_deployment(scenario, api_base, _BEDROCK_INVOKE_MODEL) + return _bedrock_invoke_deployment(scenario, api_base, _BEDROCK_TOOL_PATH_MODEL) case "bedrock_mantle": return scenario.model( - model=f"bedrock_mantle/{_BEDROCK_INVOKE_MODEL}", + model=f"bedrock_mantle/{_BEDROCK_TOOL_PATH_MODEL}", api_base=api_base, api_key=_MANTLE_KEY, aws_region_name="us-east-1", @@ -1448,21 +1458,24 @@ def test_tool_path_stream_cut_before_the_tool_block_closes_still_delivers_the_an _one_delivery(wire, marker_hex) -@pytest.mark.parametrize("accepts_forced_tools", [True, False], ids=["accepts_forced_tools", "rejects_forced_tools"]) -def test_vertex_chat_json_schema_forces_the_json_tool_only_when_the_model_accepts_forced_tools( - gateway: Gateway, accepts_forced_tools: bool +@pytest.mark.parametrize( + ("tool_model", "forced"), + [(_THINKING_OFF_TOOL_MODEL, True), (_TOOL_MODEL, False), (None, False)], + ids=["accepts_forced_tools", "thinks_by_default", "rejects_forced_tools"], +) +def test_vertex_chat_json_schema_forces_the_json_tool_only_when_the_model_will_not_think( + gateway: Gateway, tool_model: str | None, forced: bool ) -> None: marker, marker_hex = _marker("plain") - backend: Final = ( - _TOOL_MODEL if accepts_forced_tools else _claude_rejecting_forced_tools(gateway, "vertex_ai/", "claude") - ) + backend: Final = tool_model or _claude_rejecting_forced_tools(gateway, "vertex_ai/", "claude") with wire_server(_peer) as wire, gateway.scenario() as scenario: model: Final = _vertex_deployment(gateway, scenario, wire.url, backend) response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) assert _chat_content(response) == _payload(marker_hex), response.text sent: Final = _one_delivery(wire, marker_hex) assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body - assert sent.body.get("tool_choice") == (_FORCED_JSON_TOOL if accepts_forced_tools else None), sent.body + assert sent.body.get("tool_choice") == (_FORCED_JSON_TOOL if forced else None), sent.body + assert _json_tool(sent).get("description") == (None if forced else _UNFORCED_JSON_TOOL_DESCRIPTION), sent.body assert "output_config" not in sent.body and "output_format" not in sent.body, sent.body @@ -1566,6 +1579,25 @@ def test_native_model_json_schema_with_an_open_object_takes_the_json_tool_and_ke sent: Final = _one_delivery(wire, marker_hex) assert not _structured(sent.body), sent.body assert _json_tool_schema(sent)["properties"] == _VALUES_SCHEMA["properties"], sent.body + assert sent.body.get("tool_choice") == _FORCED_JSON_TOOL, sent.body + + +@pytest.mark.parametrize("provider", ["anthropic", "azure_ai"]) +def test_open_object_schema_takes_the_unforced_json_tool_on_a_model_that_thinks_by_default( + gateway: Gateway, provider: str +) -> None: + marker, marker_hex = _marker("values") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _native_deployment(gateway, scenario, wire.url, provider, _THINKING_BY_DEFAULT_NATIVE_MODEL) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_VALUES_RESPONSE_FORMAT) + ) + assert _chat_content(response) == _tool_input("values", marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert not _structured(sent.body), sent.body + assert _json_tool_schema(sent)["properties"] == _VALUES_SCHEMA["properties"], sent.body + assert _json_tool(sent).get("description") == _UNFORCED_JSON_TOOL_DESCRIPTION, sent.body + assert "tool_choice" not in sent.body, sent.body @pytest.mark.parametrize( diff --git a/tests/integration/routing/test_reasoning_effort_levels_wire.py b/tests/integration/routing/test_reasoning_effort_levels_wire.py index d15867ac572..efca26d38d0 100644 --- a/tests/integration/routing/test_reasoning_effort_levels_wire.py +++ b/tests/integration/routing/test_reasoning_effort_levels_wire.py @@ -1977,6 +1977,9 @@ _INVOKE_NATIVE_OUTPUT: Final[dict[str, JsonValue]] = { "format": {"type": "json_schema", "schema": _NATIVE_OUTPUT_SCHEMA}, } _JSON_TOOL: Final = "json_tool_call" +_UNFORCED_JSON_TOOL_DESCRIPTION: Final = ( + "When you are ready to answer, call this tool with your final answer as its input. Do not reply in plain text." +) _CONVERSE_TOOL_SCHEMA: Final[dict[str, JsonValue]] = { key: value for key, value in _NATIVE_OUTPUT_SCHEMA.items() if key != "additionalProperties" } @@ -1994,7 +1997,7 @@ _CONVERSE_WITH_JSON_TOOL: Final = _with_fields( { "toolSpec": { "name": _JSON_TOOL, - "description": _JSON_TOOL, + "description": _UNFORCED_JSON_TOOL_DESCRIPTION, "inputSchema": {"json": _CONVERSE_TOOL_SCHEMA}, } } @@ -2004,7 +2007,15 @@ _CONVERSE_WITH_JSON_TOOL: Final = _with_fields( ) _INVOKE_WITH_JSON_TOOL: Final = _with_fields( _without(_invoke_chat("high", _ADAPTIVE, ()), "anthropic_beta"), - {"tools": [{"name": _JSON_TOOL, "input_schema": _NATIVE_OUTPUT_SCHEMA}]}, + { + "tools": [ + { + "name": _JSON_TOOL, + "description": _UNFORCED_JSON_TOOL_DESCRIPTION, + "input_schema": _NATIVE_OUTPUT_SCHEMA, + } + ] + }, ) diff --git a/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py b/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py index 54a6df02edf..294791128ab 100644 --- a/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py +++ b/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py @@ -27,7 +27,7 @@ _RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { "json_schema": {"name": "risk", "schema": _SCHEMA}, } _ANTHROPIC_MODEL: Final = "anthropic/claude-scripted-9" -_BEDROCK_INVOKE_MODEL: Final = "bedrock/invoke/anthropic.claude-sonnet-5" +_BEDROCK_INVOKE_MODEL: Final = "bedrock/invoke/anthropic.claude-opus-4-8" def _payload(marker_hex: str) -> dict[str, JsonValue]: diff --git a/tests/integration/translation/chat_completions/response_format/test_chat_completions_response_format_anthropic.py b/tests/integration/translation/chat_completions/response_format/test_chat_completions_response_format_anthropic.py index 07e4358ff97..5cb7a7eee8e 100644 --- a/tests/integration/translation/chat_completions/response_format/test_chat_completions_response_format_anthropic.py +++ b/tests/integration/translation/chat_completions/response_format/test_chat_completions_response_format_anthropic.py @@ -6,7 +6,7 @@ import pytest from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration.translation.case import TranslationTestCase -from integration.translation.chat_completions.bases.anthropic import CLAUDE_HAIKU_5_5_TEST_CASE +from integration.translation.chat_completions.bases.anthropic import CLAUDE_OPUS_4_8_TEST_CASE from integration.translation.runner import assert_translation PLACE_SCHEMA: Final = { @@ -24,25 +24,25 @@ PLACE_TOOL_USE: Final = { "caller": {"type": "direct"}, } -CLAUDE_HAIKU_5_5_JSON_SCHEMA_TEST_CASE: Final = replace( - CLAUDE_HAIKU_5_5_TEST_CASE, +CLAUDE_OPUS_4_8_JSON_SCHEMA_TEST_CASE: Final = replace( + CLAUDE_OPUS_4_8_TEST_CASE, scenario="json_schema", litellm_request={ - "model": "anthropic/claude-haiku-5-5", + "model": "anthropic/claude-opus-4-8", "max_tokens": 128, "messages": [{"role": "user", "content": PLACE_PROMPT}], "response_format": {"type": "json_schema", "json_schema": {"name": "place", "schema": PLACE_SCHEMA}}, "cache": {"no-cache": True}, }, expected_provider_request={ - "model": "claude-haiku-5-5", + "model": "claude-opus-4-8", "messages": [{"role": "user", "content": [{"type": "text", "text": PLACE_PROMPT}]}], "max_tokens": 128, "tool_choice": {"name": "json_tool_call", "type": "tool"}, "tools": [{"name": "json_tool_call", "input_schema": PLACE_SCHEMA}], }, mock_provider_response={ - "model": "claude-haiku-5-5", + "model": "claude-opus-4-8", "id": "msg_011CgCoP2UR3vcM5ZWmEBd8D", "type": "message", "role": "assistant", @@ -68,7 +68,7 @@ CLAUDE_HAIKU_5_5_JSON_SCHEMA_TEST_CASE: Final = replace( "id": ANY, "created": ANY, "object": "chat.completion", - "model": "anthropic/claude-haiku-5-5", + "model": "anthropic/claude-opus-4-8", "choices": [ { "finish_reason": "stop", @@ -97,11 +97,11 @@ CLAUDE_HAIKU_5_5_JSON_SCHEMA_TEST_CASE: Final = replace( }, ) -CLAUDE_HAIKU_5_5_JSON_SCHEMA_TOOL_NAMED_ENVELOPE_TEST_CASE: Final = replace( - CLAUDE_HAIKU_5_5_JSON_SCHEMA_TEST_CASE, +CLAUDE_OPUS_4_8_JSON_SCHEMA_TOOL_NAMED_ENVELOPE_TEST_CASE: Final = replace( + CLAUDE_OPUS_4_8_JSON_SCHEMA_TEST_CASE, scenario="json_schema_tool_named_envelope", mock_provider_response={ - **CLAUDE_HAIKU_5_5_JSON_SCHEMA_TEST_CASE.mock_provider_response, + **CLAUDE_OPUS_4_8_JSON_SCHEMA_TEST_CASE.mock_provider_response, "content": [ { **PLACE_TOOL_USE, @@ -114,7 +114,7 @@ CLAUDE_HAIKU_5_5_JSON_SCHEMA_TOOL_NAMED_ENVELOPE_TEST_CASE: Final = replace( @pytest.mark.parametrize( "case", - [CLAUDE_HAIKU_5_5_JSON_SCHEMA_TEST_CASE, CLAUDE_HAIKU_5_5_JSON_SCHEMA_TOOL_NAMED_ENVELOPE_TEST_CASE], + [CLAUDE_OPUS_4_8_JSON_SCHEMA_TEST_CASE, CLAUDE_OPUS_4_8_JSON_SCHEMA_TOOL_NAMED_ENVELOPE_TEST_CASE], ids=lambda case: case.id, ) def test_chat_completions_response_format_anthropic( diff --git a/tests/integration/translation/chat_completions/streaming/test_chat_completions_streaming_anthropic.py b/tests/integration/translation/chat_completions/streaming/test_chat_completions_streaming_anthropic.py index fca8de69f00..a41752a96cf 100644 --- a/tests/integration/translation/chat_completions/streaming/test_chat_completions_streaming_anthropic.py +++ b/tests/integration/translation/chat_completions/streaming/test_chat_completions_streaming_anthropic.py @@ -6,7 +6,7 @@ from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration.translation.case import SSEFrame, StreamingTranslationTestCase from integration.translation.chat_completions.bases.anthropic import ( - CLAUDE_HAIKU_5_5_TEST_CASE, + CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_SONNET_5_5_TEST_CASE, ) from integration.translation.runner import assert_streaming_translation @@ -158,27 +158,27 @@ def _place_chunk(delta: dict[str, JsonValue]) -> SSEFrame: "id": ANY, "object": "chat.completion.chunk", "created": ANY, - "model": "anthropic/claude-haiku-5-5", + "model": "anthropic/claude-opus-4-8", "choices": [{"index": 0, "delta": delta}], } ) -CLAUDE_HAIKU_5_5_JSON_SCHEMA_CUT_STREAM_TEST_CASE: Final = StreamingTranslationTestCase( +CLAUDE_OPUS_4_8_JSON_SCHEMA_CUT_STREAM_TEST_CASE: Final = StreamingTranslationTestCase( scenario="json_schema_cut_stream", - litellm_endpoint=CLAUDE_HAIKU_5_5_TEST_CASE.litellm_endpoint, + litellm_endpoint=CLAUDE_OPUS_4_8_TEST_CASE.litellm_endpoint, litellm_request={ - "model": "anthropic/claude-haiku-5-5", + "model": "anthropic/claude-opus-4-8", "max_tokens": 128, "messages": [{"role": "user", "content": PLACE_PROMPT}], "response_format": {"type": "json_schema", "json_schema": {"name": "place", "schema": PLACE_SCHEMA}}, "cache": {"no-cache": True}, "stream": True, }, - expected_provider_endpoint=CLAUDE_HAIKU_5_5_TEST_CASE.expected_provider_endpoint, - expected_provider_headers=CLAUDE_HAIKU_5_5_TEST_CASE.expected_provider_headers, + expected_provider_endpoint=CLAUDE_OPUS_4_8_TEST_CASE.expected_provider_endpoint, + expected_provider_headers=CLAUDE_OPUS_4_8_TEST_CASE.expected_provider_headers, expected_provider_request={ - "model": "claude-haiku-5-5", + "model": "claude-opus-4-8", "messages": [{"role": "user", "content": [{"type": "text", "text": PLACE_PROMPT}]}], "max_tokens": 128, "tool_choice": {"name": "json_tool_call", "type": "tool"}, @@ -191,7 +191,7 @@ CLAUDE_HAIKU_5_5_JSON_SCHEMA_CUT_STREAM_TEST_CASE: Final = StreamingTranslationT data={ "type": "message_start", "message": { - "model": "claude-haiku-5-5", + "model": "claude-opus-4-8", "id": "msg_011CgCoP5RkexwE6N7PRZEca", "type": "message", "role": "assistant", @@ -263,7 +263,7 @@ CLAUDE_HAIKU_5_5_JSON_SCHEMA_CUT_STREAM_TEST_CASE: Final = StreamingTranslationT "id": ANY, "object": "chat.completion", "created": ANY, - "model": "anthropic/claude-haiku-5-5", + "model": "anthropic/claude-opus-4-8", "choices": [ { "index": 0, @@ -276,7 +276,7 @@ CLAUDE_HAIKU_5_5_JSON_SCHEMA_CUT_STREAM_TEST_CASE: Final = StreamingTranslationT @pytest.mark.parametrize( "case", - [CLAUDE_SONNET_5_5_STREAMING_TEST_CASE, CLAUDE_HAIKU_5_5_JSON_SCHEMA_CUT_STREAM_TEST_CASE], + [CLAUDE_SONNET_5_5_STREAMING_TEST_CASE, CLAUDE_OPUS_4_8_JSON_SCHEMA_CUT_STREAM_TEST_CASE], ids=lambda case: case.id, ) def test_chat_completions_streaming_anthropic( diff --git a/tests/unit/litellm_core_utils/test_fallback_generalizations.py b/tests/unit/litellm_core_utils/test_fallback_generalizations.py index 327031cfaf9..d8e56c0ae51 100644 --- a/tests/unit/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/unit/litellm_core_utils/test_fallback_generalizations.py @@ -563,6 +563,23 @@ def test_shipped_rules_flag_unmapped_fable_as_always_on_thinking(shipped_cost_ma assert other.get("thinking_always_on") is None +@pytest.mark.parametrize( + ("model", "provider", "on_by_default"), + [ + ("claude-opus-5-9", "anthropic", True), + ("claude-newfamily-6", "anthropic", True), + ("us.anthropic.claude-sonnet-5-9", "bedrock", True), + ("claude-opus-4-9", "anthropic", None), + ], +) +def test_shipped_rules_flag_unmapped_claude_5_and_later_as_thinking_by_default( + shipped_cost_map, model, provider, on_by_default +): + assert model not in litellm.model_cost + info = litellm.get_model_info(model, custom_llm_provider=provider) + assert info.get("thinking_on_by_default") is on_by_default + + @pytest.mark.parametrize( "model,provider", [ diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index 7e2786327dc..e81ee2a4e02 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -3140,6 +3140,118 @@ def test_streamed_json_mode_truncated_output_is_passed_through_untouched(): assert "".join(contents) == '{"name": "json_tool_call", "parameters": {"tier": 3' +def _block_events( + index: int, block: dict[str, object], deltas: tuple[dict[str, object], ...] +) -> tuple[dict[str, object], ...]: + return ( + {"type": "content_block_start", "index": index, "content_block": block}, + *({"type": "content_block_delta", "index": index, "delta": delta} for delta in deltas), + {"type": "content_block_stop", "index": index}, + ) + + +def _message_events(*blocks: tuple[dict[str, object], tuple[dict[str, object], ...]]) -> tuple[dict[str, object], ...]: + starts_and_deltas: Final = tuple( + chain.from_iterable(_block_events(index, block, deltas) for index, (block, deltas) in enumerate(blocks)) + ) + return ( + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-opus-5", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, + }, + *starts_and_deltas, + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None}, + "usage": {"output_tokens": 40}, + }, + ) + + +_THINKING_BLOCK: Final = ( + {"type": "thinking", "thinking": "", "signature": ""}, + ({"type": "thinking_delta", "thinking": "7 * 142 = 994"}, {"type": "signature_delta", "signature": "sig"}), +) +_PREAMBLE_BLOCK: Final = ( + {"type": "text", "text": ""}, + ({"type": "text_delta", "text": "**Working it out:** "}, {"type": "text_delta", "text": "the count is **120**"}), +) + + +def _json_tool_block(answer: dict[str, object]) -> tuple[dict[str, object], tuple[dict[str, object], ...]]: + encoded: Final = json.dumps(answer) + return ( + {"type": "tool_use", "id": "toolu_1", "name": RESPONSE_FORMAT_TOOL_NAME, "input": {}}, + tuple({"type": "input_json_delta", "partial_json": part} for part in (encoded[:5], encoded[5:])), + ) + + +def _parse_stream( + events: tuple[dict[str, object], ...], schema: dict[str, object] | None = _STREAMED_PERSON_SCHEMA +) -> tuple[ModelResponseStream, ...]: + iterator: Final = ModelResponseIterator( + streaming_response=MagicMock(), + sync_stream=True, + json_mode=True, + response_format_schema=None if schema is None else response_format_schema(schema), + ) + return tuple(iterator.chunk_parser(event) for event in events) + + +def _joined_content(chunks: tuple[ModelResponseStream, ...]) -> str: + return "".join(chunk.choices[0].delta.content or "" for chunk in chunks) + + +def test_streamed_json_mode_drops_prose_written_before_the_json_tool(): + chunks: Final = _parse_stream(_message_events(_THINKING_BLOCK, _PREAMBLE_BLOCK, _json_tool_block(_STREAMED_PERSON))) + + assert json.loads(_joined_content(chunks)) == _STREAMED_PERSON + assert "".join(getattr(chunk.choices[0].delta, "reasoning_content", None) or "" for chunk in chunks) == ( + "7 * 142 = 994" + ) + assert all(chunk.choices[0].delta.tool_calls is None for chunk in chunks) + + +def test_streamed_json_mode_returns_prose_when_claude_skips_the_json_tool(): + chunks: Final = _parse_stream(_message_events(_THINKING_BLOCK, _PREAMBLE_BLOCK)) + + assert _joined_content(chunks) == "**Working it out:** the count is **120**" + + +def test_streamed_json_mode_sends_prose_ahead_of_a_caller_tool_call(): + caller_tool: Final = ( + {"type": "tool_use", "id": "toolu_2", "name": "get_weather", "input": {}}, + ({"type": "input_json_delta", "partial_json": '{"city": "Paris"}'},), + ) + + chunks: Final = _parse_stream(_message_events(_PREAMBLE_BLOCK, caller_tool)) + tool_chunk_index: Final = next(i for i, chunk in enumerate(chunks) if chunk.choices[0].delta.tool_calls) + + assert _joined_content(chunks[: tool_chunk_index + 1]) == "**Working it out:** the count is **120**" + assert chunks[tool_chunk_index].choices[0].delta.tool_calls[0].function.name == "get_weather" + + +def test_streamed_native_output_text_is_not_held_back(): + answer: Final = json.dumps(_STREAMED_PERSON) + text_block: Final = ( + {"type": "text", "text": ""}, + tuple({"type": "text_delta", "text": part} for part in (answer[:6], answer[6:])), + ) + + chunks: Final = _parse_stream(_message_events(text_block), schema=None) + + assert tuple(chunk.choices[0].delta.content for chunk in chunks[2:4]) == (answer[:6], answer[6:]) + + def test_streamed_json_mode_tool_name_envelope_is_unwrapped_when_name_is_a_schema_property(): envelope = json.dumps({"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": _STREAMED_CONTACT}) @@ -3565,7 +3677,7 @@ def test_streamed_json_mode_partial_objects_that_repeat_a_field_keep_the_order_t assert "".join(contents) == "".join(blocks) -def test_streamed_json_mode_text_after_the_answer_follows_the_whole_answer(): +def test_streamed_json_mode_text_after_the_answer_is_dropped_like_the_non_streamed_reply(): answer = json.dumps(_STREAMED_RISK_OBJECT) message_start, *_, message_delta = _json_tool_events(()) text_block = ( @@ -3576,7 +3688,7 @@ def test_streamed_json_mode_text_after_the_answer_follows_the_whole_answer(): chunks = _json_mode_stream_chunks((message_start, *_json_tool_block_events(0, answer), *text_block, message_delta)) - assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == answer + " Done." + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == answer def test_streamed_json_mode_envelope_whose_label_is_a_property_name_is_still_unwrapped(): diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index f1c38ef4e64..80845b1f12d 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -22,6 +22,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, RESPONSE_FORMAT_TOOL_NAME, + RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION, ) from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt @@ -1424,6 +1425,33 @@ def test_dict_fields_keep_native_output_when_the_json_tool_would_not_be_forced(m assert "json_tool_call" not in [tool["name"] for tool in optional_params.get("tools", [])] +def test_dict_fields_take_the_unforced_json_tool_on_a_default_thinking_model(local_model_cost_map): + optional_params: Final = AnthropicConfig().map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "out", + "schema": { + "type": "object", + "properties": {"field": {"type": "object", "additionalProperties": {"type": "integer"}}}, + "required": ["field"], + }, + }, + } + }, + optional_params={}, + model="claude-opus-5", + drop_params=False, + ) + + assert "output_format" not in optional_params + assert "tool_choice" not in optional_params + assert [ + (tool["name"], tool["input_schema"]["properties"]["field"]) for tool in optional_params["tools"] + ] == [("json_tool_call", {"type": "object", "additionalProperties": {"type": "integer"}})] + + def test_native_structured_output_uses_bundled_capability_when_remote_map_lags( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -6956,6 +6984,63 @@ def test_response_format_tool_path_skips_forced_tool_choice_when_unsupported(loc assert "tool_choice" not in result +@pytest.mark.parametrize( + ("model", "thinking_params", "forced"), + [ + ("claude-opus-4-8", {"thinking": {"type": "adaptive"}}, False), + ("claude-opus-5", {}, False), + ("claude-sonnet-5", {}, False), + ("claude-haiku-5-5", {}, False), + ("claude-fable-5", {"reasoning_effort": "none"}, False), + ("claude-opus-4-8", {}, True), + ("claude-opus-5", {"thinking": {"type": "disabled"}}, True), + ("claude-haiku-5-5", {"reasoning_effort": "none"}, True), + ], + ids=[ + "opus-4-8-adaptive", + "opus-5-default-thinking", + "sonnet-5-default-thinking", + "haiku-5-5-default-thinking", + "fable-5-always-on-reasoning-effort-none", + "opus-4-8-default-off", + "opus-5-thinking-disabled", + "haiku-5-5-reasoning-effort-none", + ], +) +def test_response_format_tool_is_forced_only_when_the_model_will_not_think( + set_native_structured_output, model, thinking_params, forced +): + set_native_structured_output(model, None) + config: Final = AnthropicConfig() + optional_params: Final = config.map_openai_params( + non_default_params={ + **thinking_params, + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "capital", + "schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, + }, + }, + optional_params={}, + model=model, + drop_params=False, + ) + + request: Final = config.transform_request( + model=model, + messages=[{"role": "user", "content": "Name the capital of France."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert [tool["name"] for tool in request["tools"]] == ["json_tool_call"] + assert (request.get("tool_choice") == {"name": "json_tool_call", "type": "tool"}) is forced + assert request["tools"][0].get("description") == (None if forced else RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION) + + def _eager_chat_function(**extra: object) -> dict[str, object]: return { "name": "write_file", diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py index 34a4cf5225b..93f6e237fc7 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py @@ -562,19 +562,27 @@ def _invoke_params_for_schema( @pytest.mark.parametrize( - ("native_key", "thinking_params", "native"), + ("model", "native_key", "thinking_params", "native"), [ - ("legacy", {}, True), - ("bedrock", {}, False), - ("bedrock", {"thinking": {"type": "adaptive"}}, True), - ("bedrock", {"reasoning_effort": "low"}, True), + ("us.anthropic.claude-haiku-5-5", "legacy", {}, True), + ("us.anthropic.claude-haiku-5-5", "bedrock", {}, True), + ("us.anthropic.claude-opus-4-8", "bedrock", {}, False), + ("us.anthropic.claude-opus-4-8", "bedrock", {"thinking": {"type": "adaptive"}}, True), + ("us.anthropic.claude-opus-4-8", "bedrock", {"reasoning_effort": "low"}, True), + ("us.anthropic.claude-haiku-5-5", "bedrock", {"thinking": {"type": "disabled"}}, False), + ], + ids=[ + "legacy-key", + "bedrock-key-thinking-by-default", + "bedrock-key-thinking-off-by-default", + "bedrock-key-adaptive", + "bedrock-key-effort", + "bedrock-key-thinking-disabled", ], - ids=["legacy-key", "bedrock-key-no-thinking", "bedrock-key-adaptive", "bedrock-key-effort"], ) -def test_bedrock_native_output_key_applies_only_when_thinking_is_requested( - set_native_structured_output, native_key, thinking_params, native +def test_bedrock_native_output_key_applies_only_when_the_model_will_think( + set_native_structured_output, model, native_key, thinking_params, native ): - model = "us.anthropic.claude-haiku-5-5" set_native_structured_output(model, native_key) schema = {"type": "object", "properties": {"answer": {"type": "string"}}, "required": ["answer"]} @@ -914,9 +922,7 @@ def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_th assert result["output_config"] == {"effort": "high"} -def test_bedrock_chat_invoke_adaptive_thinking_keeps_the_forced_response_format_tool(set_native_structured_output): - """Anthropic accepts a forced tool_choice next to adaptive thinking, and dropping - it lets the model answer in prose instead of the requested JSON.""" +def test_bedrock_chat_invoke_adaptive_thinking_leaves_the_response_format_tool_unforced(set_native_structured_output): set_native_structured_output("us.anthropic.claude-opus-5", None) result = AmazonAnthropicClaudeConfig().map_openai_params( non_default_params={ @@ -935,7 +941,8 @@ def test_bedrock_chat_invoke_adaptive_thinking_keeps_the_forced_response_format_ drop_params=False, ) - assert result["tool_choice"] == {"type": "tool", "name": result["tools"][0]["name"]} + assert [tool["name"] for tool in result["tools"]] == ["json_tool_call"] + assert "tool_choice" not in result assert result["thinking"] == {"type": "adaptive"} @@ -2084,7 +2091,7 @@ def test_invoke_uses_opt_in_native_output_is_false_when_native_output_did_not_co (False, True, None, {"thinking": {"type": "adaptive"}}, True), (False, False, None, {}, False), (False, None, {"supports_bedrock_native_structured_output": False}, {}, False), - (True, True, None, {}, False), + (True, True, None, {"thinking": {"type": "disabled"}}, False), ], ids=[ "cannot-force-tool", diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 7c19ef5b751..b8eb5858fda 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -15,6 +15,7 @@ import respx import litellm from litellm._logging import verbose_logger from litellm import ModelResponse, Router, completion +from litellm.constants import RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION from litellm.litellm_core_utils.prompt_templates.factory import ( _bedrock_converse_messages_pt, _bedrock_tools_pt, @@ -1440,18 +1441,27 @@ def test_thinking_request_adds_output_tokens_details_response_path(): assert result["additionalModelResponseFieldPaths"] == ("/usage/output_tokens_details",) -def test_request_without_thinking_omits_response_field_paths(): - config = AmazonConverseConfig() - - result = config._transform_request( - model="bedrock/converse/us.anthropic.claude-opus-4-7", +@pytest.mark.parametrize( + ("model", "thinks_by_default"), + [ + ("bedrock/converse/us.anthropic.claude-opus-4-7", False), + ("bedrock/converse/us.anthropic.claude-opus-5", True), + ("us.anthropic.claude-haiku-5-5", True), + ], + ids=["opus-4-7-thinking-off-by-default", "opus-5-thinking-on-by-default", "haiku-5-5-thinking-on-by-default"], +) +def test_request_without_thinking_asks_for_thinking_tokens_only_when_the_model_thinks_by_default( + local_model_cost_map: None, model: str, thinks_by_default: bool +): + result = AmazonConverseConfig()._transform_request( + model=model, messages=[{"role": "user", "content": "hi"}], optional_params={"maxTokens": 256}, litellm_params={}, headers={}, ) - assert "additionalModelResponseFieldPaths" not in result + assert (result.get("additionalModelResponseFieldPaths") == ("/usage/output_tokens_details",)) is thinks_by_default def test_transform_usage_prefers_provider_reasoning_tokens(): @@ -4937,7 +4947,6 @@ def test_translate_response_format_native_output_config(monkeypatch): model="anthropic.claude-sonnet-4-5-20250929-v1:0", optional_params=optional_params, non_default_params={"response_format": response_format}, - is_thinking_enabled=False, ) # Should have outputConfig, NOT tools @@ -4965,20 +4974,20 @@ def test_translate_response_format_native_output_config(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) -CAPITAL_RESPONSE_FORMAT: Final = { - "type": "json_schema", - "json_schema": { - "name": "capital", - "schema": { - "type": "object", - "properties": {"city": {"type": "string"}, "country": {"type": "string"}}, - "required": ["city", "country"], - "additionalProperties": False, - }, +CAPITAL_JSON_SCHEMA: Final = { + "name": "capital", + "schema": { + "type": "object", + "properties": {"city": {"type": "string"}, "country": {"type": "string"}}, + "required": ["city", "country"], + "additionalProperties": False, }, } +CAPITAL_RESPONSE_FORMAT: Final = {"type": "json_schema", "json_schema": CAPITAL_JSON_SCHEMA} + + ADAPTIVE_THINKING: Final = {"type": "adaptive"} @@ -5030,24 +5039,124 @@ def test_json_schema_shape_follows_each_inference_profile_own_native_output_flag assert "fake_stream" not in native_params assert "outputConfig" not in tool_request assert [tool["toolSpec"]["name"] for tool in tool_request["toolConfig"]["tools"]] == ["json_tool_call"] - assert tool_request["toolConfig"]["toolChoice"] == {"tool": {"name": "json_tool_call"}} + assert "toolChoice" not in tool_request["toolConfig"] assert tool_params.get("fake_stream", False) is stream +_JSON_TOOL_CHOICE: Final = {"tool": {"name": "json_tool_call"}} + + @pytest.mark.parametrize( - ("native_key", "thinking_params", "native"), + ("model", "native_key", "model_info", "thinking_params", "forced"), [ - ("legacy", {}, True), - ("bedrock", {}, False), - ("bedrock", {"thinking": ADAPTIVE_THINKING}, True), - ("bedrock", {"thinking": {"type": "enabled", "budget_tokens": 1024}}, True), - ("bedrock", {"reasoning_effort": "low"}, True), - ("bedrock", {"reasoning_effort": "none"}, False), - ("bedrock", {"thinking": {"type": "disabled"}}, False), + ( + "us.anthropic.claude-opus-4-8", + "bedrock", + {"supports_bedrock_native_structured_output": False}, + {"thinking": ADAPTIVE_THINKING}, + False, + ), + ("us.anthropic.claude-opus-5", None, None, {}, False), + ("global.anthropic.claude-sonnet-5", None, None, {}, False), + ("us.anthropic.claude-haiku-5-5", None, None, {}, False), + ("us.anthropic.claude-fable-5", None, None, {}, False), + ("us.anthropic.claude-fable-5", None, None, {"reasoning_effort": "none"}, False), + ("us.anthropic.claude-opus-4-8", None, None, {}, True), + ("us.anthropic.claude-opus-5", None, None, {"thinking": {"type": "disabled"}}, True), + ("us.anthropic.claude-haiku-5-5", None, None, {"reasoning_effort": "none"}, True), + ], + ids=[ + "opus-4-8-adaptive-native-output-fallback", + "opus-5-default-thinking", + "sonnet-5-default-thinking", + "haiku-5-5-default-thinking", + "fable-5-always-on", + "fable-5-always-on-reasoning-effort-none", + "opus-4-8-default-off", + "opus-5-thinking-disabled", + "haiku-5-5-reasoning-effort-none", + ], +) +def test_json_tool_is_forced_only_when_the_model_will_not_think( + set_native_structured_output, model, native_key, model_info, thinking_params, forced +): + set_native_structured_output(model, native_key) + config: Final = AmazonConverseConfig() + + response_format: Final = { + "type": "json_schema", + "json_schema": {**CAPITAL_JSON_SCHEMA, "description": "The capital of a country"}, + } + + optional_params: Final = config.map_openai_params( + non_default_params={"response_format": response_format, **thinking_params}, + optional_params={}, + model=model, + drop_params=False, + model_info=model_info, + ) + request: Final = config._transform_request( + model=model, + messages=[{"role": "user", "content": "Name the capital of France."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert "outputConfig" not in request + assert [tool["toolSpec"]["name"] for tool in request["toolConfig"]["tools"]] == ["json_tool_call"] + assert (request["toolConfig"].get("toolChoice") == _JSON_TOOL_CHOICE) is forced + assert request["toolConfig"]["tools"][0]["toolSpec"]["description"] == ( + "The capital of a country" + if forced + else f"{RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION} The capital of a country" + ) + + +def test_unforced_json_tool_keeps_the_schema_description_on_a_non_claude_model(set_native_structured_output): + model: Final = "us.meta.llama3-3-70b-instruct-v1:0" + set_native_structured_output(model, None) + config: Final = AmazonConverseConfig() + + optional_params: Final = config.map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": {**CAPITAL_JSON_SCHEMA, "description": "The capital of a country"}, + } + }, + optional_params={}, + model=model, + drop_params=False, + ) + request: Final = config._transform_request( + model=model, + messages=[{"role": "user", "content": "Name the capital of France."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert "toolChoice" not in request["toolConfig"] + assert [tool["toolSpec"]["description"] for tool in request["toolConfig"]["tools"]] == ["The capital of a country"] + + +@pytest.mark.parametrize( + ("model", "native_key", "thinking_params", "native"), + [ + ("us.anthropic.claude-haiku-5-5", "legacy", {}, True), + ("us.anthropic.claude-haiku-5-5", "bedrock", {}, True), + ("us.anthropic.claude-opus-4-8", "bedrock", {}, False), + ("us.anthropic.claude-opus-4-8", "bedrock", {"thinking": ADAPTIVE_THINKING}, True), + ("us.anthropic.claude-haiku-5-5", "bedrock", {"thinking": {"type": "enabled", "budget_tokens": 1024}}, True), + ("us.anthropic.claude-opus-4-8", "bedrock", {"reasoning_effort": "low"}, True), + ("us.anthropic.claude-haiku-5-5", "bedrock", {"reasoning_effort": "none"}, False), + ("us.anthropic.claude-haiku-5-5", "bedrock", {"thinking": {"type": "disabled"}}, False), ], ids=[ "legacy-key", - "bedrock-key-no-thinking", + "bedrock-key-thinking-by-default", + "bedrock-key-thinking-off-by-default", "bedrock-key-adaptive", "bedrock-key-enabled", "bedrock-key-effort", @@ -5055,10 +5164,9 @@ def test_json_schema_shape_follows_each_inference_profile_own_native_output_flag "bedrock-key-thinking-disabled", ], ) -def test_bedrock_native_output_key_applies_only_when_thinking_is_requested( - set_native_structured_output, native_key, thinking_params, native +def test_bedrock_native_output_key_applies_only_when_the_model_will_think( + set_native_structured_output, model, native_key, thinking_params, native ): - model: Final = "us.anthropic.claude-haiku-5-5" set_native_structured_output(model, native_key) _, request = _converse_request_for_json_schema(model, False, thinking_params) @@ -5078,7 +5186,7 @@ _FORCED_JSON_TOOL_CHOICE: Final = {"tool": {"name": "json_tool_call"}} (False, True, None, {"thinking": {"type": "adaptive"}}, True, None), (False, False, None, None, False, None), (False, None, {"supports_bedrock_native_structured_output": False}, None, False, None), - (True, True, None, None, False, _FORCED_JSON_TOOL_CHOICE), + (True, True, None, {"thinking": {"type": "disabled"}}, False, _FORCED_JSON_TOOL_CHOICE), ], ids=[ "cannot-force-tool", @@ -5596,7 +5704,6 @@ def test_translate_response_format_fallback_tool_call(): model="anthropic.claude-3-haiku-20240307-v1:0", optional_params=optional_params, non_default_params={"response_format": response_format}, - is_thinking_enabled=False, ) # Should use tool-call approach, NOT outputConfig (model doesn't support native structured outputs) @@ -5633,7 +5740,6 @@ def test_native_structured_output_no_fake_stream(monkeypatch): model="anthropic.claude-sonnet-4-5-20250929-v1:0", optional_params=optional_params, non_default_params={"response_format": response_format, "stream": True}, - is_thinking_enabled=False, ) assert "outputConfig" in result @@ -5954,7 +6060,6 @@ def test_json_object_no_schema_skips_tool_injection(monkeypatch): model="anthropic.claude-sonnet-4-5-20250929-v1:0", optional_params=optional_params, non_default_params=non_default_params, - is_thinking_enabled=False, ) # Should NOT use native outputConfig (no schema provided) diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index 996cfec6e22..4424f7640a3 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -1149,6 +1149,108 @@ async def test_converse_streaming_completion_yields_text_and_finish_reason( _assert_converse_stream_text_and_finish_reason(sync_chunks) +_CONVERSE_PROSE: Final = "Let me work this out. " + + +async def _stream_chunks( + request: Mapping[str, object], transport: httpx.MockTransport, async_mode: bool +) -> tuple[ModelResponseStream, ...]: + if not async_mode: + return tuple(litellm.completion(**request, client=HTTPHandler(client=httpx.Client(transport=transport)))) + async_stream: Final = await litellm.acompletion(**request, client=AsyncHTTPHandler(transport=transport)) + return tuple([chunk async for chunk in async_stream]) + + +@pytest.mark.parametrize("async_mode", [False, True]) +@pytest.mark.parametrize( + ("tool_block", "stop_reason", "expected_content", "expected_tools", "expected_finish_reason"), + [ + pytest.param( + {"toolUse": {"toolUseId": "tooluse_1", "name": "json_tool_call", "input": {"answer": 120}}}, + "tool_use", + '{"answer": 120}', + (), + "stop", + id="prose_then_json_tool", + ), + pytest.param( + {"toolUse": {"toolUseId": "tooluse_1", "name": "get_weather", "input": {"city": "Paris"}}}, + "tool_use", + _CONVERSE_PROSE, + (("get_weather", '{"city": "Paris"}'),), + "tool_calls", + id="prose_then_caller_tool", + ), + pytest.param(None, "end_turn", _CONVERSE_PROSE, (), "stop", id="json_tool_skipped"), + ], +) +@pytest.mark.asyncio +async def test_converse_json_schema_stream_sends_claude_prose_only_when_claude_skips_the_json_tool( + _aws_test_credentials: None, + async_mode: bool, + tool_block: Mapping[str, object] | None, + stop_reason: str, + expected_content: str, + expected_tools: tuple[tuple[str, str], ...], + expected_finish_reason: str, +) -> None: + converse_reply: Final = { + "output": { + "message": { + "role": "assistant", + "content": [{"text": _CONVERSE_PROSE}, *((tool_block,) if tool_block is not None else ())], + } + }, + "stopReason": stop_reason, + "usage": {"inputTokens": 30, "outputTokens": 20, "totalTokens": 50}, + "metrics": {"latencyMs": 1}, + } + + def handle_request(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + assert "toolChoice" not in body["toolConfig"] + assert "json_tool_call" in [tool["toolSpec"]["name"] for tool in body["toolConfig"]["tools"]] + return httpx.Response(200, json=converse_reply) + + transport: Final = httpx.MockTransport(handle_request) + request: Final = { + "model": "bedrock/us.anthropic.claude-opus-5", + "messages": [{"role": "user", "content": "What is the LCM of 24 and 120? Reply in JSON."}], + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "answer", + "schema": {"type": "object", "properties": {"answer": {"type": "integer"}}, "required": ["answer"]}, + }, + }, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, + } + ], + "stream": True, + } + chunks: Final = await _stream_chunks(request, transport, async_mode) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == expected_content + tool_calls: Final = tuple( + itertools.chain.from_iterable(chunk.choices[0].delta.tool_calls or () for chunk in chunks) + ) + named_calls: Final = tuple( + (call.function.name, "".join(c.function.arguments or "" for c in tool_calls)) + for call in tool_calls + if call.function.name + ) + assert named_calls == expected_tools + assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == [ + expected_finish_reason + ] + + def test_mistral_invoke_stream_yields_text_and_finish_reason( _aws_test_credentials: None, ) -> None: diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index d2b78235c29..e04dd35f851 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -19,6 +19,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, + RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.databricks.common_utils import DatabricksException @@ -2591,6 +2592,50 @@ def test_databricks_claude_json_schema_preserves_refs_in_tool_parameters() -> No ] +@pytest.mark.parametrize( + ("model", "thinking_params", "forced"), + ( + ("databricks-claude-opus-5", {}, False), + ("databricks-claude-opus-4-8", {"thinking": {"type": "adaptive"}}, False), + ("databricks-claude-opus-5-5", {}, False), + ("databricks-claude-opus-4-8", {}, True), + ("databricks-claude-opus-5", {"thinking": {"type": "disabled"}}, True), + ), + ids=( + "opus-5-default-thinking", + "opus-4-8-adaptive", + "opus-5-5-without-forced-tools", + "opus-4-8-default-off", + "opus-5-thinking-disabled", + ), +) +def test_databricks_claude_json_tool_is_forced_only_when_the_model_will_not_think( + local_model_cost_map: None, model: str, thinking_params: dict[str, object], forced: bool +) -> None: + optional_params: Final = litellm.utils.get_optional_params( + model=model, + custom_llm_provider="databricks", + response_format=copy.deepcopy(JSON_SCHEMA_RESPONSE_FORMAT), + **thinking_params, + ) + request: Final[dict[str, object]] = TypeAdapter(dict[str, object]).validate_python( + DatabricksConfig().transform_request( # pyright: ignore[reportUnknownMemberType] # inherited hook is untyped + model=model, + messages=[{"role": "user", "content": "Name the capital of France."}], + optional_params=dict(optional_params), + litellm_params={}, + headers={}, + ) + ) + + assert ( + request.get("tool_choice") == {"type": "function", "function": {"name": "json_tool_call"}} + ) is forced + assert optional_params["tools"][0]["function"].get("description") == ( + None if forced else RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION + ) + + def test_databricks_pydantic_json_schema_preserves_nested_refs() -> None: class Person(BaseModel): model_config = ConfigDict(frozen=True) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 0220495032b..99c6f6c402c 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -7703,6 +7703,7 @@ async def test_bedrock_converse_parses_scripted_response( assert json.loads(route.calls[0].request.content) == { "messages": [{"role": "user", "content": [{"text": "hello"}]}], "inferenceConfig": {}, + "additionalModelResponseFieldPaths": ["/usage/output_tokens_details"], } assert response.choices[0].message.content == "Hello from Bedrock." assert response.usage.prompt_tokens == 8 diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index e293b2a5be3..e24499db071 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1244,6 +1244,7 @@ _COST_MAP_ROW_SCHEMA: Final = { "supports_anthropic_thinking_payload": {"type": "boolean"}, "supports_legacy_thinking": {"type": "boolean"}, "thinking_always_on": {"type": "boolean"}, + "thinking_on_by_default": {"type": "boolean"}, "supports_mid_conversation_system": {"type": "boolean"}, "supports_sampling_params": {"type": "boolean"}, "supports_output_config": {"type": "boolean"}, From 49f56e8b23c30ca19cc570eb4d779bf53132ac2b Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Wed, 4 Nov 2026 20:23:30 -0800 Subject: [PATCH 12/39] refactor(proxy)!: drop the unused accounting event tables (#49923) LiteLLM_AccountingEvent and LiteLLM_AccountingDelivery have had no reader since the SpendLogs projection replaced them, so the migration drops them with their enums. AccountingDeliveryState stays because LiteLLM_SpendLogDelivery uses it The schema now spells the delivery_id default the way Postgres stores it, which ends the drift Prisma reported against the ledger migration. The schema check compares base types only, so an attribute-only change is no longer reported as a type change Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- .../migration.sql | 5 ++ .../litellm_proxy_extras/schema.prisma | 66 +------------------ litellm/integrations/otel/model/spans.py | 2 - litellm/proxy/schema.prisma | 66 +------------------ schema.prisma | 66 +------------------ scripts/check_schema_compatibility.py | 15 +++-- .../test_drop_accounting_event_tables.py | 64 ++++++++++++++++++ tests/unit/proxy/test_db_schema_changes.py | 20 ++++++ 8 files changed, 100 insertions(+), 204 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261104010000_drop_accounting_event_tables/migration.sql create mode 100644 tests/integration/database/migrations/test_drop_accounting_event_tables.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261104010000_drop_accounting_event_tables/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261104010000_drop_accounting_event_tables/migration.sql new file mode 100644 index 00000000000..cba6a18368f --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261104010000_drop_accounting_event_tables/migration.sql @@ -0,0 +1,5 @@ +SET LOCAL lock_timeout = '10s'; +DROP TABLE IF EXISTS "LiteLLM_AccountingDelivery"; +DROP TABLE IF EXISTS "LiteLLM_AccountingEvent"; +DROP TYPE IF EXISTS "AccountingDeliveryTarget"; +DROP TYPE IF EXISTS "AccountingEventKind"; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 3aaa74d4fc2..78f16d9f1fd 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -2071,78 +2071,14 @@ model LiteLLM_TelemetryReport { @@index([created_at]) } -enum AccountingEventKind { - inbound_usage - outbound_usage - cost -} - enum AccountingDeliveryState { pending applied parked } -enum AccountingDeliveryTarget { - LiteLLM_VerificationToken - LiteLLM_UserTable - LiteLLM_TeamTable - LiteLLM_TeamMembership - LiteLLM_OrganizationTable - LiteLLM_OrganizationMembership - LiteLLM_EndUserTable - LiteLLM_AgentsTable - LiteLLM_TagTable - LiteLLM_ProjectTable - LiteLLM_ModelAccessGroupBudgetTable - LiteLLM_BudgetTable - LiteLLM_BudgetWindowSpend - LiteLLM_DailyUserSpend - LiteLLM_DailyTeamSpend - LiteLLM_DailyOrganizationSpend - LiteLLM_DailyEndUserSpend - LiteLLM_DailyAgentSpend - LiteLLM_DailyTagSpend - LiteLLM_DailyModelUsage - LiteLLM_DailyToolSpend - LiteLLM_AutoRouterDailySpend -} - -// Durable accounting ledger. One row per usage or cost event; payload holds the -// private event envelope versioned by schema_version. Usage rows retain request-scoped -// pricing context; cost rows include resolved rates and the finalized pricing snapshot. -model LiteLLM_AccountingEvent { - event_id String @id - kind AccountingEventKind - source_event_id String? - schema_version Int - payload_hash String - occurred_at DateTime - recorded_at DateTime @default(now()) - payload Json - - @@index([recorded_at]) - @@index([kind, source_event_id]) -} - -// One row per (event, target table, version): pending work for the projector and, -// once applied in the same transaction as the target write, the receipt that stops a retry. -model LiteLLM_AccountingDelivery { - event_id String - target AccountingDeliveryTarget - version Int - state AccountingDeliveryState @default(pending) - attempts Int @default(0) - next_attempt_at DateTime @default(now()) - applied_at DateTime? - last_error String? - - @@id([event_id, target, version]) - @@index([target, state, next_attempt_at]) -} - model LiteLLM_SpendLogDelivery { - delivery_id String @id @default(dbgenerated("gen_random_uuid()::text")) + delivery_id String @id @default(dbgenerated("(gen_random_uuid())::text")) request_id String state AccountingDeliveryState @default(pending) accounting_payload Json diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index c4f5c1efbcd..5b81fd9a331 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -338,8 +338,6 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_ToolTable", "LiteLLM_TelemetryReport", "LiteLLM_AccessGroupTable", - "LiteLLM_AccountingDelivery", - "LiteLLM_AccountingEvent", "LiteLLM_SpendLogDelivery", "LiteLLM_ClaudeCodePluginTable", "LiteLLM_ClaudeCodePluginUpload", diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 3aaa74d4fc2..78f16d9f1fd 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -2071,78 +2071,14 @@ model LiteLLM_TelemetryReport { @@index([created_at]) } -enum AccountingEventKind { - inbound_usage - outbound_usage - cost -} - enum AccountingDeliveryState { pending applied parked } -enum AccountingDeliveryTarget { - LiteLLM_VerificationToken - LiteLLM_UserTable - LiteLLM_TeamTable - LiteLLM_TeamMembership - LiteLLM_OrganizationTable - LiteLLM_OrganizationMembership - LiteLLM_EndUserTable - LiteLLM_AgentsTable - LiteLLM_TagTable - LiteLLM_ProjectTable - LiteLLM_ModelAccessGroupBudgetTable - LiteLLM_BudgetTable - LiteLLM_BudgetWindowSpend - LiteLLM_DailyUserSpend - LiteLLM_DailyTeamSpend - LiteLLM_DailyOrganizationSpend - LiteLLM_DailyEndUserSpend - LiteLLM_DailyAgentSpend - LiteLLM_DailyTagSpend - LiteLLM_DailyModelUsage - LiteLLM_DailyToolSpend - LiteLLM_AutoRouterDailySpend -} - -// Durable accounting ledger. One row per usage or cost event; payload holds the -// private event envelope versioned by schema_version. Usage rows retain request-scoped -// pricing context; cost rows include resolved rates and the finalized pricing snapshot. -model LiteLLM_AccountingEvent { - event_id String @id - kind AccountingEventKind - source_event_id String? - schema_version Int - payload_hash String - occurred_at DateTime - recorded_at DateTime @default(now()) - payload Json - - @@index([recorded_at]) - @@index([kind, source_event_id]) -} - -// One row per (event, target table, version): pending work for the projector and, -// once applied in the same transaction as the target write, the receipt that stops a retry. -model LiteLLM_AccountingDelivery { - event_id String - target AccountingDeliveryTarget - version Int - state AccountingDeliveryState @default(pending) - attempts Int @default(0) - next_attempt_at DateTime @default(now()) - applied_at DateTime? - last_error String? - - @@id([event_id, target, version]) - @@index([target, state, next_attempt_at]) -} - model LiteLLM_SpendLogDelivery { - delivery_id String @id @default(dbgenerated("gen_random_uuid()::text")) + delivery_id String @id @default(dbgenerated("(gen_random_uuid())::text")) request_id String state AccountingDeliveryState @default(pending) accounting_payload Json diff --git a/schema.prisma b/schema.prisma index 3aaa74d4fc2..78f16d9f1fd 100644 --- a/schema.prisma +++ b/schema.prisma @@ -2071,78 +2071,14 @@ model LiteLLM_TelemetryReport { @@index([created_at]) } -enum AccountingEventKind { - inbound_usage - outbound_usage - cost -} - enum AccountingDeliveryState { pending applied parked } -enum AccountingDeliveryTarget { - LiteLLM_VerificationToken - LiteLLM_UserTable - LiteLLM_TeamTable - LiteLLM_TeamMembership - LiteLLM_OrganizationTable - LiteLLM_OrganizationMembership - LiteLLM_EndUserTable - LiteLLM_AgentsTable - LiteLLM_TagTable - LiteLLM_ProjectTable - LiteLLM_ModelAccessGroupBudgetTable - LiteLLM_BudgetTable - LiteLLM_BudgetWindowSpend - LiteLLM_DailyUserSpend - LiteLLM_DailyTeamSpend - LiteLLM_DailyOrganizationSpend - LiteLLM_DailyEndUserSpend - LiteLLM_DailyAgentSpend - LiteLLM_DailyTagSpend - LiteLLM_DailyModelUsage - LiteLLM_DailyToolSpend - LiteLLM_AutoRouterDailySpend -} - -// Durable accounting ledger. One row per usage or cost event; payload holds the -// private event envelope versioned by schema_version. Usage rows retain request-scoped -// pricing context; cost rows include resolved rates and the finalized pricing snapshot. -model LiteLLM_AccountingEvent { - event_id String @id - kind AccountingEventKind - source_event_id String? - schema_version Int - payload_hash String - occurred_at DateTime - recorded_at DateTime @default(now()) - payload Json - - @@index([recorded_at]) - @@index([kind, source_event_id]) -} - -// One row per (event, target table, version): pending work for the projector and, -// once applied in the same transaction as the target write, the receipt that stops a retry. -model LiteLLM_AccountingDelivery { - event_id String - target AccountingDeliveryTarget - version Int - state AccountingDeliveryState @default(pending) - attempts Int @default(0) - next_attempt_at DateTime @default(now()) - applied_at DateTime? - last_error String? - - @@id([event_id, target, version]) - @@index([target, state, next_attempt_at]) -} - model LiteLLM_SpendLogDelivery { - delivery_id String @id @default(dbgenerated("gen_random_uuid()::text")) + delivery_id String @id @default(dbgenerated("(gen_random_uuid())::text")) request_id String state AccountingDeliveryState @default(pending) accounting_payload Json diff --git a/scripts/check_schema_compatibility.py b/scripts/check_schema_compatibility.py index b8e9ac3c07a..c56986a38c8 100644 --- a/scripts/check_schema_compatibility.py +++ b/scripts/check_schema_compatibility.py @@ -13,6 +13,7 @@ from pydantic import TypeAdapter Schema = Mapping[str, Mapping[str, str]] MODELS: Final = re.compile(r"^[ \t]*model\s+(\w+)\s*\{(.*?)^[ \t]*\}", re.MULTILINE | re.DOTALL) +DROPPED_UNUSED_MODELS: Final = frozenset({"LiteLLM_AccountingDelivery", "LiteLLM_AccountingEvent"}) def parse_model_fields(schema: str) -> Schema: @@ -27,18 +28,18 @@ def _field_changes(model: str, field: str, old_type: str, new_fields: Mapping[st if field not in new_fields: yield f"Breaking: Field {model}.{field} was removed" return - new_type: Final = new_fields[field] - if old_type == new_type: - return - if "?" in old_type and "?" not in new_type: + old_base: Final = old_type.split()[0] + new_base: Final = new_fields[field].split()[0] + if old_base.endswith("?") and not new_base.endswith("?"): yield f"Breaking: Field {model}.{field} changed from optional to required" - if not old_type.startswith(new_type.split("?")[0]): - yield f"Breaking: Field {model}.{field} changed type from {old_type} to {new_type}" + if old_base.rstrip("?") != new_base.rstrip("?"): + yield f"Breaking: Field {model}.{field} changed type from {old_base} to {new_base}" def _model_changes(model: str, old_fields: Mapping[str, str], new_schema: Schema) -> Iterator[str]: if model not in new_schema: - yield f"Breaking: Model {model} was removed" + if model not in DROPPED_UNUSED_MODELS: + yield f"Breaking: Model {model} was removed" return for field, old_type in old_fields.items(): yield from _field_changes(model, field, old_type, new_schema[model]) diff --git a/tests/integration/database/migrations/test_drop_accounting_event_tables.py b/tests/integration/database/migrations/test_drop_accounting_event_tables.py new file mode 100644 index 00000000000..583223a573b --- /dev/null +++ b/tests/integration/database/migrations/test_drop_accounting_event_tables.py @@ -0,0 +1,64 @@ +import os +from pathlib import Path +from typing import Final + +import psycopg +import pytest + +_MIGRATIONS: Final = Path(__file__).parents[4] / "litellm-proxy-extras/litellm_proxy_extras/migrations" +_OLD_LEDGER: Final = _MIGRATIONS / "20261027000000_add_accounting_ledger_tables/migration.sql" +_SPEND_LOG_DELIVERY: Final = _MIGRATIONS / "20261103010000_spendlogs_accounting_delivery/migration.sql" +_DROP: Final = _MIGRATIONS / "20261104010000_drop_accounting_event_tables/migration.sql" +_DROPPED: Final = ( + "AccountingDeliveryTarget", + "AccountingEventKind", + "LiteLLM_AccountingDelivery", + "LiteLLM_AccountingEvent", +) + + +def _apply(connection: psycopg.Connection[tuple[object, ...]], migration: Path) -> None: + with connection.transaction(): + connection.execute(migration.read_bytes()) + + +def _present(connection: psycopg.Connection[tuple[object, ...]], names: tuple[str, ...]) -> tuple[object, ...]: + return tuple( + row[0] + for row in connection.execute( + "SELECT relname FROM pg_class WHERE relname = ANY(%s) " + "UNION SELECT typname FROM pg_type WHERE typname = ANY(%s) ORDER BY 1", + (list(names), list(names)), + ).fetchall() + ) + + +@pytest.mark.parametrize("old_states", [(), ("applied", "pending", "parked")], ids=["empty", "with_rows"]) +def test_drop_removes_the_old_ledger_and_keeps_spend_log_deliveries(old_states: tuple[str, ...]) -> None: + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + _apply(connection, _OLD_LEDGER) + connection.execute('CREATE TABLE "LiteLLM_SpendLogs" (request_id text PRIMARY KEY)') + _apply(connection, _SPEND_LOG_DELIVERY) + for index, state in enumerate(old_states): + connection.execute( + 'INSERT INTO "LiteLLM_AccountingEvent" (event_id, kind, schema_version, payload_hash, occurred_at, ' + "payload) VALUES (%s, 'cost', 1, 'hash', now(), '{}')", + (f"event-{index}",), + ) + connection.execute( + 'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version, state) ' + "VALUES (%s, 'LiteLLM_UserTable', 1, %s)", + (f"event-{index}", state), + ) + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, state, accounting_payload, is_batch) ' + "VALUES ('request', 'pending', '{}', false)" + ) + + _apply(connection, _DROP) + _apply(connection, _DROP) + + assert _present(connection, _DROPPED) == () + assert connection.execute('SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [ + ("request", "pending") + ] diff --git a/tests/unit/proxy/test_db_schema_changes.py b/tests/unit/proxy/test_db_schema_changes.py index 3a5d0fb1a0a..f43e2d2a8ce 100644 --- a/tests/unit/proxy/test_db_schema_changes.py +++ b/tests/unit/proxy/test_db_schema_changes.py @@ -12,6 +12,12 @@ def test_aaaaaschema_compatibility() -> None: assert check_breaking_changes(parse_model_fields(ORIGINAL), parse_model_fields(added)) == () +def test_changing_only_a_field_attribute_is_not_breaking() -> None: + previous: Final = parse_model_fields('model Item {\n id String @id @default(dbgenerated("a()::text"))\n}\n') + current: Final = parse_model_fields('model Item {\n id String @id @default(dbgenerated("(a())::text"))\n}\n') + assert check_breaking_changes(previous, current) == () + + def test_parse_model_fields_retains_types_and_attributes() -> None: assert parse_model_fields(ORIGINAL) == {"Item": {"id": "String @id", "label": "String?"}} @@ -40,3 +46,17 @@ def test_schema_compatibility_rejects_removed_models_fields_and_narrowed_types( schema: str, expected: tuple[str, ...] ) -> None: assert check_breaking_changes(parse_model_fields(ORIGINAL), parse_model_fields(schema)) == expected + + +def test_only_the_dropped_accounting_models_may_be_removed() -> None: + accounting: Final = ( + "model LiteLLM_AccountingEvent {\n event_id String @id\n}\n" + "model LiteLLM_AccountingDelivery {\n event_id String\n note String?\n}\n" + ) + previous: Final = parse_model_fields(ORIGINAL + accounting) + + assert check_breaking_changes(previous, parse_model_fields(ORIGINAL)) == () + assert check_breaking_changes(previous, parse_model_fields("")) == ("Breaking: Model Item was removed",) + assert check_breaking_changes( + previous, parse_model_fields(ORIGINAL + "model LiteLLM_AccountingDelivery {\n event_id String\n}\n") + ) == ("Breaking: Field LiteLLM_AccountingDelivery.note was removed",) From a1ffb06e2829c4876defacdca549b99d4a0694d2 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 5 Nov 2026 04:27:27 +0000 Subject: [PATCH 13/39] chore(cost-map): add together_ai ricursive-ai/Ripple-65B-Preview from the Together models API (#50163) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 11 +++++++++++ model_prices_and_context_window.json | 11 +++++++++++ 2 files changed, 22 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a292ce0d14e..76fd85d001b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -86165,5 +86165,16 @@ "supports_tool_choice": true, "supports_vision": true, "use_openai_responses_path": true + }, + "together_ai/ricursive-ai/Ripple-65B-Preview": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://api.together.ai/v1/models", + "supports_prompt_caching": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a292ce0d14e..76fd85d001b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -86165,5 +86165,16 @@ "supports_tool_choice": true, "supports_vision": true, "use_openai_responses_path": true + }, + "together_ai/ricursive-ai/Ripple-65B-Preview": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://api.together.ai/v1/models", + "supports_prompt_caching": true } } From 2ffbdf0d7a345adbf72df8b487710da64b2f5c9f Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 4 Nov 2026 20:29:38 -0800 Subject: [PATCH 14/39] feat(ui)!: redesign public AI Hub as a catalog (#50153) --- tests/e2e/ui/tests/modelHub/modelHub.spec.ts | 170 ++++--- .../src/components/navbar.tsx | 4 + .../publicModelHub/AgentInspector.tsx | 93 ---- .../publicModelHub/CatalogDetails.tsx | 37 ++ .../publicModelHub/CatalogFilters.tsx | 44 ++ .../components/publicModelHub/FilterChip.tsx | 93 ---- .../components/publicModelHub/HubHeader.tsx | 37 +- .../publicModelHub/HubMasterDetail.tsx | 235 ---------- .../publicModelHub/InspectorHeader.tsx | 22 - .../publicModelHub/McpInspector.tsx | 36 -- .../components/publicModelHub/ModelsPane.tsx | 377 +++++++++------- .../PublicAiHub.integration.test.tsx | 422 ++++++++++++++---- .../components/publicModelHub/PublicAiHub.tsx | 131 +++--- .../publicModelHub/ResourcePanes.tsx | 371 ++++++++------- .../publicModelHub/catalog.module.css | 268 +++++++++++ .../components/publicModelHub/fonts/OFL.txt | 93 ++++ .../fonts/barlow-condensed-latin-600.ttf | Bin 0 -> 86636 bytes .../components/publicModelHub/useHubData.ts | 16 +- .../publicModelHub/useHubUrlState.ts | 34 +- ui/litellm-dashboard/src/i18n/locales/en.json | 42 ++ 20 files changed, 1422 insertions(+), 1103 deletions(-) delete mode 100644 ui/litellm-dashboard/src/components/publicModelHub/AgentInspector.tsx create mode 100644 ui/litellm-dashboard/src/components/publicModelHub/CatalogDetails.tsx create mode 100644 ui/litellm-dashboard/src/components/publicModelHub/CatalogFilters.tsx delete mode 100644 ui/litellm-dashboard/src/components/publicModelHub/FilterChip.tsx delete mode 100644 ui/litellm-dashboard/src/components/publicModelHub/HubMasterDetail.tsx delete mode 100644 ui/litellm-dashboard/src/components/publicModelHub/McpInspector.tsx create mode 100644 ui/litellm-dashboard/src/components/publicModelHub/catalog.module.css create mode 100644 ui/litellm-dashboard/src/components/publicModelHub/fonts/OFL.txt create mode 100644 ui/litellm-dashboard/src/components/publicModelHub/fonts/barlow-condensed-latin-600.ttf diff --git a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts index f8ce4faf45e..517ec64c7df 100644 --- a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts +++ b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts @@ -1,8 +1,8 @@ -import { test, expect, type APIRequestContext } from "@playwright/test"; +import { test, expect } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; -import { masterKey } from "../../helpers/traffic"; +import { masterKey, uniqueSuffix } from "../../helpers/traffic"; test.describe("AI Hub (internal admin view)", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); @@ -62,103 +62,93 @@ test.describe("AI Hub (internal admin view)", () => { }); test.describe("Public model hub (/ui/model_hub_table)", () => { - // No storageState — the public page is reached anonymously with a `key` query param. - - test("Public model_hub_table loads and renders the Models tab", async ({ page }) => { - // The page expects the proxy key as the `key` query param. Use the master - // key the e2e runner already exports — this matches what the AI Hub copy - // button hands out. - await page.goto(`/ui/model_hub_table?key=${masterKey()}`); - - // Dismiss the feedback popup before asserting on the tab, so a popup - // race can't briefly mask the tab while we're evaluating visibility. + test("shows the Model Hub and Skills Hub catalogs", async ({ page }) => { + await page.goto("/ui/model_hub_table"); await dismissFeedbackPopup(page); - await expect(page.getByRole("tab", { name: /^Models/ })).toBeVisible({ timeout: 10_000 }); + await expect( + page.getByRole("tab", { name: "Model Hub", exact: true }), + ).toBeVisible(); + await expect( + page.getByRole("textbox", { name: "Search models", exact: true }), + ).toBeVisible(); + await expect(page.getByRole("tab")).toHaveCount(2); + await page.getByRole("tab", { name: "Skills Hub", exact: true }).click(); + await expect( + page.getByRole("textbox", { name: "Search skills", exact: true }), + ).toBeVisible(); }); - test("Agents and MCP servers tabs render their public entries", async ({ page, request }) => { - const suffix = `${Date.now()}`; - const agentName = `e2e-public-agent-${suffix}`; - const mcpServerName = `e2e_public_mcp_${suffix}`; + test("finds a published skill and opens its installation instructions", async ({ + page, + request, + }) => { + const skillName = `e2e-catalog-${uniqueSuffix()}`; + const domain = `catalog-${uniqueSuffix()}`; const auth = { Authorization: `Bearer ${masterKey()}` }; - - const publicMcpServerIds = async (api: APIRequestContext): Promise => { - const res = await api.get("/public/mcp_hub"); - expect(res.ok(), `public mcp_hub read failed (${res.status()}): ${await res.text()}`).toBe(true); - const servers: { server_id: string }[] = await res.json(); - return servers.map((server) => server.server_id); - }; - - const seedPublicEntries = async ( - api: APIRequestContext, - priorMcpIds: string[], - ): Promise<{ agentId: string; serverId: string }> => { - const agentRes = await api.post("/v1/agents", { - headers: auth, - data: { - agent_name: agentName, - agent_card_params: { - name: agentName, - description: "E2E public agent", - version: "1.0.0", - url: "http://127.0.0.1:9999/", - capabilities: {}, - skills: [], - defaultInputModes: ["text"], - defaultOutputModes: ["text"], - }, - }, - }); - expect(agentRes.ok(), `agent create failed (${agentRes.status()}): ${await agentRes.text()}`).toBe(true); - const agentId = (await agentRes.json()).agent_id as string; - - const serverRes = await api.post("/v1/mcp/server", { - headers: auth, - data: { - server_name: mcpServerName, - url: "http://127.0.0.1:9999/mcp", - transport: "http", - description: "E2E public MCP server", - }, - }); - expect(serverRes.ok(), `mcp server create failed (${serverRes.status()}): ${await serverRes.text()}`).toBe(true); - const serverId = (await serverRes.json()).server_id as string; - - const agentPublicRes = await api.post(`/v1/agents/${agentId}/make_public`, { headers: auth }); - expect(agentPublicRes.ok(), `agent make_public failed: ${await agentPublicRes.text()}`).toBe(true); - const mcpPublicRes = await api.post("/v1/mcp/make_public", { - headers: auth, - data: { mcp_server_ids: [...priorMcpIds, serverId] }, - }); - expect(mcpPublicRes.ok(), `mcp make_public failed: ${await mcpPublicRes.text()}`).toBe(true); - - return { agentId, serverId }; - }; - - const priorMcpIds = await publicMcpServerIds(request); - const { agentId, serverId } = await seedPublicEntries(request, priorMcpIds); + const created = await request.post("/claude-code/plugins", { + headers: auth, + data: { + name: skillName, + source: { source: "github", repo: "example/catalog-skills" }, + description: "Catalog browser test skill", + domain, + version: "1.0.0", + }, + }); + expect( + created.ok(), + `skill registration failed (${created.status()}): ${await created.text()}`, + ).toBe(true); try { - await page.goto(`/ui/model_hub_table?key=${masterKey()}`); + await page.goto("/ui/model_hub_table?tab=skills"); await dismissFeedbackPopup(page); + await page.getByRole("checkbox", { name: domain, exact: true }).check(); + await page + .getByRole("textbox", { name: "Search skills", exact: true }) + .fill(skillName); + await page.getByRole("link").filter({ hasText: skillName }).click(); - const agentsTab = page.getByRole("tab", { name: /^Agents/ }); - await expect(agentsTab).toBeVisible({ timeout: 15_000 }); - await agentsTab.click(); - await expect(page.getByRole("textbox", { name: "Search agents" })).toBeVisible({ timeout: 10_000 }); - await expect(page.getByRole("link").filter({ hasText: agentName })).toHaveCount(1, { timeout: 10_000 }); - await expect(page.getByText("E2E public agent").first()).toBeVisible(); - - const mcpTab = page.getByRole("tab", { name: /^MCP/ }); - await expect(mcpTab).toBeVisible(); - await mcpTab.click(); - await expect(page.getByRole("textbox", { name: "Search MCP servers" })).toBeVisible({ timeout: 10_000 }); - await expect(page.getByRole("link").filter({ hasText: mcpServerName })).toHaveCount(1, { timeout: 10_000 }); - await expect(page.getByText("E2E public MCP server").first()).toBeVisible(); + const details = page.getByRole("dialog"); + await expect( + details.getByRole("heading", { name: skillName, exact: true }), + ).toBeVisible(); + await expect( + details.getByText(`/plugin install ${skillName}@litellm`, { + exact: true, + }), + ).toBeVisible(); + await expect(page).toHaveURL(new RegExp(`open=${skillName}`)); + await page.reload(); + await expect( + page + .getByRole("dialog") + .getByRole("heading", { name: skillName, exact: true }), + ).toBeVisible(); } finally { - await request.post("/v1/mcp/make_public", { headers: auth, data: { mcp_server_ids: priorMcpIds } }); - await request.delete(`/v1/agents/${agentId}`, { headers: auth }); - await request.delete(`/v1/mcp/server/${serverId}`, { headers: auth }); + const deleted = await request.delete( + `/claude-code/plugins/${skillName}`, + { headers: auth }, + ); + expect(deleted.ok(), `skill cleanup failed (${deleted.status()})`).toBe( + true, + ); + } + }); + + test("retired catalog links return to Models and clear the old selection", async ({ + page, + }) => { + for (const tab of ["agents", "mcp"]) { + await page.goto( + `/ui/model_hub_table?tab=${tab}&open=retired-entry&key=preserved`, + ); + await dismissFeedbackPopup(page); + await expect( + page.getByRole("tab", { name: "Model Hub", exact: true }), + ).toHaveAttribute("aria-selected", "true"); + await expect(page).toHaveURL(/model_hub_table\?key=preserved$/); + await expect(page.getByRole("dialog")).toHaveCount(0); } }); }); diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index 7ddbd77bb7f..331b4512bf1 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -29,6 +29,7 @@ interface NavbarProps { onToggleSidebar?: () => void; deploymentHealth?: React.ReactNode; deploymentHealthLink?: React.ReactNode; + publicNavigation?: React.ReactNode; } const NAV_LOGO_CLASS_NAME = "h-auto max-h-full w-auto max-w-full object-contain"; @@ -40,6 +41,7 @@ const Navbar: React.FC = ({ onToggleSidebar, deploymentHealth, deploymentHealthLink, + publicNavigation, }) => { const { t } = useTranslation(); const baseUrl = getProxyBaseUrl(); @@ -132,6 +134,8 @@ const Navbar: React.FC = ({
+ {publicNavigation} + {!isPublicPage && (
diff --git a/ui/litellm-dashboard/src/components/publicModelHub/AgentInspector.tsx b/ui/litellm-dashboard/src/components/publicModelHub/AgentInspector.tsx deleted file mode 100644 index 4d237877f28..00000000000 --- a/ui/litellm-dashboard/src/components/publicModelHub/AgentInspector.tsx +++ /dev/null @@ -1,93 +0,0 @@ -"use client"; - -import { ExternalLink } from "lucide-react"; -import { useTranslation } from "react-i18next"; - -import type { AgentCard } from "@/components/PublicModelHubTableColumns"; -import { Button } from "@/components/ui/button"; -import { Field, FieldDescription } from "@/components/ui/field"; -import { Item, ItemActions, ItemContent, ItemDescription, ItemGroup, ItemTitle } from "@/components/ui/item"; - -import { agentCapabilities } from "./hubFormat"; -import { Badges, CodeBlock, EndpointField, Facts, PaneSection } from "./HubParts"; -import { agentSnippets } from "./hubSnippets"; -import { InspectorHeader } from "./InspectorHeader"; - -const modes = (agent: AgentCard) => { - const input = (agent.defaultInputModes ?? []).join(", "); - const output = (agent.defaultOutputModes ?? []).join(", "); - return input || output ? `${input || "-"} / ${output || "-"}` : null; -}; - -export function AgentInspector({ agent }: { agent: AgentCard }) { - const { t } = useTranslation(); - const capabilities = agentCapabilities(agent); - return ( -
- - 0 && , - }, - ], - ]} - /> - - {agent.documentationUrl && ( - - )} - {(agent.skills ?? []).length > 0 && ( - - - {agent.skills.map((skill) => ( - - - {skill.name} - {skill.description} - - - - - - ))} - - - )} - - {agentSnippets(agent).map((snippet, index) => ( - - - {t("publicModelHub.agentInspector.step", { step: index + 1, title: snippet.title })} - - - - ))} - -
- ); -} diff --git a/ui/litellm-dashboard/src/components/publicModelHub/CatalogDetails.tsx b/ui/litellm-dashboard/src/components/publicModelHub/CatalogDetails.tsx new file mode 100644 index 00000000000..97dd6b4dbd7 --- /dev/null +++ b/ui/litellm-dashboard/src/components/publicModelHub/CatalogDetails.tsx @@ -0,0 +1,37 @@ +"use client"; + +import type { ReactNode } from "react"; +import { useTranslation } from "react-i18next"; + +import { Dialog, DialogContent, DialogTitle } from "@/components/ui/dialog"; +import { Skeleton } from "@/components/ui/skeleton"; + +export function CatalogDetails({ + name, + loading = false, + onClose, + children, +}: { + name: string | null; + loading?: boolean; + onClose: () => void; + children: ReactNode; +}) { + const { t } = useTranslation(); + return ( + !next && onClose()}> + + {t("publicModelHub.masterDetail.details")} + {loading ? ( +
+ + + +
+ ) : ( + children + )} + + + ); +} diff --git a/ui/litellm-dashboard/src/components/publicModelHub/CatalogFilters.tsx b/ui/litellm-dashboard/src/components/publicModelHub/CatalogFilters.tsx new file mode 100644 index 00000000000..7eceab1e410 --- /dev/null +++ b/ui/litellm-dashboard/src/components/publicModelHub/CatalogFilters.tsx @@ -0,0 +1,44 @@ +"use client"; + +import { ChevronDown } from "lucide-react"; +import { useState, type ReactNode } from "react"; +import { useTranslation } from "react-i18next"; + +import { Button } from "@/components/ui/button"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; + +import { useIsNarrow } from "./useIsNarrow"; + +export function CatalogFilters({ + children, + count, + onClear, +}: { + children: ReactNode; + count: number; + onClear: () => void; +}) { + const { t } = useTranslation(); + const narrow = useIsNarrow(); + const [expanded, setExpanded] = useState(false); + return ( + + ); +} diff --git a/ui/litellm-dashboard/src/components/publicModelHub/FilterChip.tsx b/ui/litellm-dashboard/src/components/publicModelHub/FilterChip.tsx deleted file mode 100644 index a5aeff6f068..00000000000 --- a/ui/litellm-dashboard/src/components/publicModelHub/FilterChip.tsx +++ /dev/null @@ -1,93 +0,0 @@ -"use client"; - -import { CirclePlus, X } from "lucide-react"; -import { useTranslation } from "react-i18next"; - -import { Button } from "@/components/ui/button"; -import { ButtonGroup } from "@/components/ui/button-group"; -import { - Combobox, - ComboboxContent, - ComboboxEmpty, - ComboboxInput, - ComboboxItem, - ComboboxList, - ComboboxTrigger, -} from "@/components/ui/combobox"; -import { cn } from "@/lib/cva.config"; - -import { firstPlus } from "./hubFormat"; -import { ProviderAvatar } from "./HubParts"; - -interface FilterOption { - value: string; - label: string; - provider?: string; -} - -export function FilterChip({ - label, - noun, - options, - value, - onChange, -}: { - label: string; - noun: string; - options: FilterOption[]; - value: string[]; - onChange: (value: string[]) => void; -}) { - const { t } = useTranslation(); - const byValue = new Map(options.map((option) => [option.value, option])); - const selected = value.map((v) => byValue.get(v)?.label ?? v); - const active = value.length > 0; - return ( - - - multiple - items={options} - value={options.filter((option) => value.includes(option.value))} - onValueChange={(next) => onChange(next.map((option) => option.value))} - itemToStringLabel={(option) => option.label} - isItemEqualToValue={(a, b) => a.value === b.value} - > - svg:last-child]:hidden", !active && "border-dashed text-muted-foreground")} - /> - } - > - {!active && } - {label} - {active && {firstPlus(selected)}} - - - - {t("publicModelHub.filterChip.empty", { noun })} - - {(option: FilterOption) => ( - - {option.provider && } - {option.label} - - )} - - - - {active && ( - - )} - - ); -} diff --git a/ui/litellm-dashboard/src/components/publicModelHub/HubHeader.tsx b/ui/litellm-dashboard/src/components/publicModelHub/HubHeader.tsx index 532df227ea2..f42e7e0ccbd 100644 --- a/ui/litellm-dashboard/src/components/publicModelHub/HubHeader.tsx +++ b/ui/litellm-dashboard/src/components/publicModelHub/HubHeader.tsx @@ -1,6 +1,6 @@ "use client"; -import { ChevronDown, ExternalLink, LayoutGrid } from "lucide-react"; +import { ChevronDown, ExternalLink } from "lucide-react"; import { useTranslation } from "react-i18next"; import { Button } from "@/components/ui/button"; @@ -18,7 +18,8 @@ import { InputGroup, InputGroupAddon, InputGroupInput, InputGroupText } from "@/ import { gatewayUrl, usefulLinkList } from "./hubFormat"; import { CopyIconButton } from "./HubParts"; -import { StatusBadge } from "./HubStatus"; +import type { HubTab } from "./useHubUrlState"; +import styles from "./catalog.module.css"; import type { HubInfo } from "./useHubData"; export type HubStatus = "loading" | "operational" | "degraded" | "unavailable"; @@ -59,7 +60,7 @@ function BaseUrlField() { const { t } = useTranslation(); const url = gatewayUrl(); return ( - + {t("publicModelHub.baseUrlField.label")} @@ -75,31 +76,33 @@ function BaseUrlField() { - + {t("publicModelHub.baseUrlField.hint")} ); } -export function HubHeader({ info, status }: { info: HubInfo | undefined; status: HubStatus }) { +export function HubHeader({ info, status, tab }: { info: HubInfo | undefined; status: HubStatus; tab: HubTab }) { const { t } = useTranslation(); return ( -
-
-
-

- - {t("publicModelHub.header.title")} -

- {status === "operational" && } -
-

- {info?.custom_docs_description || t("publicModelHub.header.description")} +

+
+

{info?.docs_title || t("publicModelHub.catalog.eyebrow")}

+

+ {t(tab === "skills" ? "publicModelHub.catalog.skillsHub" : "publicModelHub.catalog.modelHub")} +

+

+ {info?.custom_docs_description || + t( + tab === "skills" + ? "publicModelHub.catalog.skillsDescription" + : "publicModelHub.catalog.modelsDescription", + )}

{status !== "unavailable" && ( -
+
diff --git a/ui/litellm-dashboard/src/components/publicModelHub/HubMasterDetail.tsx b/ui/litellm-dashboard/src/components/publicModelHub/HubMasterDetail.tsx deleted file mode 100644 index 715e1cad07f..00000000000 --- a/ui/litellm-dashboard/src/components/publicModelHub/HubMasterDetail.tsx +++ /dev/null @@ -1,235 +0,0 @@ -"use client"; - -import * as React from "react"; - -import { Item, ItemContent, ItemGroup, ItemMedia } from "@/components/ui/item"; -import { ScrollArea } from "@/components/ui/scroll-area"; -import { Sheet, SheetContent, SheetTitle } from "@/components/ui/sheet"; -import { Skeleton } from "@/components/ui/skeleton"; -import { cn } from "@/lib/cva.config"; -import i18n from "@/i18n/i18n"; - -import { NotListed } from "./HubEmpty"; -import { InSheet } from "./InspectorHeader"; -import { useOpenHref } from "./useHubUrlState"; -import { useIsNarrow } from "./useIsNarrow"; - -function ListSkeleton() { - return ( - - {[0, 1, 2].map((i) => ( -
- - - - - - - - - -
- ))} - - ); -} - -function DetailSkeleton() { - return ( -
- -
- {[0, 1, 2, 3].map((i) => ( -
- - -
- ))} -
- -
- ); -} - -const isPlainClick = (event: React.MouseEvent) => { - const opensElsewhere = event.metaKey || event.ctrlKey || event.shiftKey; - return event.button === 0 && !opensElsewhere; -}; - -function Rows({ - rows, - activeKey, - explicit, - getKey, - onSelect, - renderItem, -}: { - rows: T[]; - activeKey: string | undefined; - explicit: boolean; - getKey: (row: T) => string; - onSelect: (key: string) => void; - renderItem: (row: T) => React.ReactNode; -}) { - const hrefFor = useOpenHref(); - const follow = (key: string) => (event: React.MouseEvent) => { - if (!isPlainClick(event)) return; - event.preventDefault(); - onSelect(key); - }; - return ( - - {rows.map((row) => { - const key = getKey(row); - const active = key === activeKey; - return ( -
- - } - className={cn(active && (explicit ? "bg-muted" : "lg:bg-muted"))} - > - {renderItem(row)} - -
- ); - })} - - ); -} - -const resolveSelection = ( - rows: T[], - selectedKey: string | null, - opened: T | undefined, - getKey: (row: T) => string, -) => { - const chosen = selectedKey === null ? undefined : opened ?? rows.find((row) => getKey(row) === selectedKey); - return { chosen, selected: selectedKey === null ? rows[0] : chosen }; -}; - -function useSheet( - chosen: T | undefined, - missingKey: string | undefined, - getKey: (row: T) => string, - renderDetail: (row: T) => React.ReactNode, -) { - const [content, setContent] = React.useState<{ row?: T; missingKey?: string }>({ row: chosen }); - if (chosen !== undefined && chosen !== content.row) setContent({ row: chosen }); - if (missingKey !== undefined && missingKey !== content.missingKey) setContent({ missingKey }); - const { row } = content; - return { - open: chosen !== undefined || missingKey !== undefined, - title: row ? getKey(row) : content.missingKey ?? i18n.t("publicModelHub.masterDetail.details"), - body: row ? renderDetail(row) : content.missingKey && , - }; -} - -export function HubMasterDetail({ - switcher, - toolbar, - rows, - loading, - selectedKey, - opened, - openPending = false, - getKey, - onSelect, - renderItem, - renderDetail, - empty, - footer, -}: { - switcher?: React.ReactNode; - toolbar: React.ReactNode; - rows: T[]; - loading: boolean; - selectedKey: string | null; - opened?: T; - openPending?: boolean; - getKey: (row: T) => string; - onSelect: (key: string | null) => void; - renderItem: (row: T) => React.ReactNode; - renderDetail: (row: T) => React.ReactNode; - empty: React.ReactNode; - footer?: React.ReactNode; -}) { - const narrow = useIsNarrow(); - const { chosen, selected } = resolveSelection(rows, selectedKey, opened, getKey); - const isEmpty = !loading && rows.length === 0; - const settled = !loading && !openPending && !isEmpty; - const missingKey = selectedKey !== null && chosen === undefined && settled ? selectedKey : undefined; - const sheet = useSheet(chosen, missingKey, getKey, renderDetail); - const listOnly = isEmpty && chosen === undefined && !openPending; - const detail = () => { - if (loading || (selected === undefined && openPending)) return ; - if (selected === undefined) return selectedKey && ; - return renderDetail(selected); - }; - const list = isEmpty ? ( - empty - ) : ( - <> - - {footer} - - ); - return ( -
-
-
- {switcher} -
{toolbar}
-
- {loading ? : list} -
- {!listOnly && {detail()}} - onSelect(null)}> - {sheet.body} - -
- ); -} - -function DetailSheet({ - open, - title, - onClose, - children, -}: { - open: boolean; - title: string; - onClose: () => void; - children: React.ReactNode; -}) { - const body = React.useRef(null); - return ( - !next && onClose()}> - - {title} -
- {children} -
- - - ); -} diff --git a/ui/litellm-dashboard/src/components/publicModelHub/InspectorHeader.tsx b/ui/litellm-dashboard/src/components/publicModelHub/InspectorHeader.tsx index bf76bd49239..9a8f5ea6bb9 100644 --- a/ui/litellm-dashboard/src/components/publicModelHub/InspectorHeader.tsx +++ b/ui/litellm-dashboard/src/components/publicModelHub/InspectorHeader.tsx @@ -1,18 +1,12 @@ "use client"; -import { X } from "lucide-react"; import * as React from "react"; -import { useTranslation } from "react-i18next"; -import { Button } from "@/components/ui/button"; -import { SheetClose } from "@/components/ui/sheet"; import { cn } from "@/lib/cva.config"; import { SECTION_TITLE } from "./hubFormat"; import { CopyIconButton } from "./HubParts"; -export const InSheet = React.createContext(false); - export function InspectorHeader({ media, title, @@ -26,8 +20,6 @@ export function InspectorHeader({ status?: React.ReactNode; description?: React.ReactNode; }) { - const { t } = useTranslation(); - const inSheet = React.useContext(InSheet); return (
@@ -36,20 +28,6 @@ export function InspectorHeader({ {status && {status}} - {inSheet && ( - - } - > - - - )}
{status && {status}} {description &&

{description}

} diff --git a/ui/litellm-dashboard/src/components/publicModelHub/McpInspector.tsx b/ui/litellm-dashboard/src/components/publicModelHub/McpInspector.tsx deleted file mode 100644 index 7ec6c5cf78b..00000000000 --- a/ui/litellm-dashboard/src/components/publicModelHub/McpInspector.tsx +++ /dev/null @@ -1,36 +0,0 @@ -"use client"; - -import { useTranslation } from "react-i18next"; - -import type { MCPServerData } from "@/components/PublicModelHubTableColumns"; - -import { mcpAuthLabel } from "./hubFormat"; -import { CodeBlock, EndpointField, Facts, PaneSection } from "./HubParts"; -import { mcpEndpoint, mcpSnippet } from "./hubSnippets"; -import { InspectorHeader } from "./InspectorHeader"; - -export function McpInspector({ server }: { server: MCPServerData }) { - const { t } = useTranslation(); - return ( -
- - - - - - -
- ); -} diff --git a/ui/litellm-dashboard/src/components/publicModelHub/ModelsPane.tsx b/ui/litellm-dashboard/src/components/publicModelHub/ModelsPane.tsx index d7fcc865036..837dd3e0b1d 100644 --- a/ui/litellm-dashboard/src/components/publicModelHub/ModelsPane.tsx +++ b/ui/litellm-dashboard/src/components/publicModelHub/ModelsPane.tsx @@ -1,24 +1,26 @@ "use client"; -import { ArrowDownWideNarrow, ArrowUpNarrowWide, Boxes, ChevronLeft, ChevronRight } from "lucide-react"; +import type { ColumnDef } from "@tanstack/react-table"; +import { ArrowDownWideNarrow, ArrowUpNarrowWide, Boxes } from "lucide-react"; +import type { MouseEvent } from "react"; import { useTranslation } from "react-i18next"; import type { ModelGroupInfo } from "@/components/PublicModelHubTableColumns"; -import { Button } from "@/components/ui/button"; +import { DataTable, DataTableSortHeader } from "@/components/shared/DataTable"; import { ButtonGroup } from "@/components/ui/button-group"; -import { ItemActions, ItemContent, ItemDescription, ItemMedia } from "@/components/ui/item"; +import { Checkbox } from "@/components/ui/checkbox"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Toggle } from "@/components/ui/toggle"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { displayMode } from "@/lib/decisionModels"; -import i18n from "@/i18n/i18n"; -import { FilterChip } from "./FilterChip"; -import { featureFacetLabel, firstPlus, formatPerMillion, formatTokens, modeLabel, providerName } from "./hubFormat"; -import { HubMasterDetail } from "./HubMasterDetail"; -import { LoadFailed, NoResults } from "./HubEmpty"; -import { Joined, ProviderAvatar, RowTitle, SearchInput } from "./HubParts"; +import { CatalogDetails } from "./CatalogDetails"; +import { CatalogFilters } from "./CatalogFilters"; +import { featureFacetLabel, formatPerMillion, formatTokens, modeLabel, providerName } from "./hubFormat"; +import { LoadFailed, NoResults, NotListed } from "./HubEmpty"; +import { ProviderAvatar, SearchInput } from "./HubParts"; import { type HubData, useOpenedModel } from "./useHubData"; +import { useOpenHref } from "./useHubUrlState"; import { ModelInspector } from "./ModelInspector"; const SORTS = [ @@ -34,42 +36,12 @@ const SORTS = [ type Models = HubData["models"]; -function ModelsPager({ models }: { models: Models }) { - const { t } = useTranslation(); - const { pagination, rowCount, rows, onPaginationChange } = models; - if (!rowCount) return null; - const start = pagination.pageIndex * pagination.pageSize + 1; - const end = start + rows.length - 1; - const go = (pageIndex: number) => onPaginationChange({ ...pagination, pageIndex }); - return ( -
- {t("publicModelHub.modelsPager.range", { start, end, total: rowCount })} - - - - -
- ); -} - function SortControl({ models }: { models: Models }) { const { t } = useTranslation(); + const options = SORTS.map((sort) => ({ + ...sort, + label: t(`publicModelHub.catalog.sort.${sort.value}`, { defaultValue: sort.label }), + })); const field = models.sorting[0]?.id ?? "model_group"; const desc = models.sorting[0]?.desc ?? false; const pick = (value: string | null) => { @@ -78,13 +50,13 @@ function SortControl({ models }: { models: Models }) { }; return ( - {t("publicModelHub.sortControl.sort")} - {SORTS.map((sort) => ( + {options.map((sort) => ( {sort.label} @@ -112,100 +84,106 @@ function SortControl({ models }: { models: Models }) { ); } -function ModelsToolbar({ hub }: { hub: HubData }) { - const { t } = useTranslation(); - const { models, facets } = hub; +interface FilterOption { + value: string; + label: string; + provider?: string; +} + +function CatalogFilter({ + label, + options, + values, + onChange, +}: { + label: string; + options: FilterOption[]; + values: string[]; + onChange: (values: string[]) => void; +}) { return ( - <> -
- - +
+ {label} +
+ {options.map((option) => ( + + ))}
-
- ({ value: p, label: providerName(p), provider: p }))} - value={models.providerValues} - onChange={models.onProvidersChange} - /> - displayMode(m) ?? m))].map((m) => ({ - value: m, - label: modeLabel(m), - }))} - value={models.modeValues} - onChange={models.onModesChange} - /> - ({ value: f, label: featureFacetLabel(f) }))} - value={models.featureValues} - onChange={models.onFeaturesChange} - /> -
- +
); } -const price = (cost: number | null | undefined, direction: "input" | "output") => { - if (cost === 0) return i18n.t("publicModelHub.modelPrices.freeDirection", { direction }); - return cost - ? i18n.t("publicModelHub.modelPrices.perMillion", { - price: formatPerMillion(cost), - direction: - direction === "input" ? i18n.t("publicModelHub.modelPrices.in") : i18n.t("publicModelHub.modelPrices.out"), - }) - : null; -}; - -const prices = (model: ModelGroupInfo) => - model.input_cost_per_token === 0 && model.output_cost_per_token === 0 - ? [i18n.t("publicModelHub.modelPrices.free"), null] - : [price(model.input_cost_per_token, "input"), price(model.output_cost_per_token, "output")]; - -function ModelRow({ model }: { model: ModelGroupInfo }) { +function ModelFilters({ hub }: { hub: HubData }) { const { t } = useTranslation(); - const [input, output] = prices(model); - const priced = Boolean(input || output); + const { models, facets } = hub; return ( - <> - - - - - {model.model_group} - - - - {priced && ( - - - - )} - - {priced && ( - - {input && {input}} - {output && {output}} - - )} - + models.onColumnFiltersChange([])}> + ({ value: provider, label: providerName(provider), provider }))} + values={models.providerValues} + onChange={models.onProvidersChange} + /> + ({ value: feature, label: featureFacetLabel(feature) }))} + values={models.featureValues} + onChange={models.onFeaturesChange} + /> + displayMode(mode) ?? mode))].map((mode) => ({ + value: mode, + label: modeLabel(mode), + }))} + values={models.modeValues} + onChange={models.onModesChange} + /> + + ); +} + +function ModelLink({ + model, + selected, + onSelect, +}: { + model: ModelGroupInfo; + selected: string | null; + onSelect: (key: string | null) => void; +}) { + const hrefFor = useOpenHref(); + const follow = (event: MouseEvent) => { + if (event.button !== 0 || event.defaultPrevented) return; + if (event.metaKey || event.ctrlKey) return; + if (event.shiftKey || event.altKey) return; + event.preventDefault(); + onSelect(model.model_group); + }; + return ( + + + + {model.model_group} + + {(model.providers ?? []).map(providerName).join(", ")} + + + ); } @@ -221,35 +199,122 @@ export function ModelsPane({ const { t } = useTranslation(); const { models } = hub; const opened = useOpenedModel(selected, models.rows, !models.isLoading); + const detailsLoading = opened.model === undefined && (models.isLoading || opened.pending); + const missingName = opened.model === undefined && !detailsLoading ? selected : null; + const countOptions = { + count: models.rowCount, + defaultValue_one: "{{count}} model", + defaultValue_other: "{{count}} models", + defaultValue: "{{count}} models", + }; + const columns: ColumnDef[] = [ + { + accessorKey: "model_group", + meta: { skeleton: "twoLine" }, + sortDescFirst: false, + header: ({ column }) => ( + + ), + cell: ({ row }) => , + }, + { + accessorKey: "max_input_tokens", + meta: { numeric: true, className: "hidden sm:table-cell" }, + size: 140, + sortDescFirst: true, + header: ({ column }) => ( + + ), + cell: ({ row }) => (row.original.max_input_tokens === 0 ? "0" : formatTokens(row.original.max_input_tokens)), + }, + { + accessorKey: "input_cost_per_token", + meta: { numeric: true }, + size: 130, + sortDescFirst: false, + header: ({ column }) => ( + + ), + cell: ({ row }) => formatPerMillion(row.original.input_cost_per_token), + }, + { + accessorKey: "output_cost_per_token", + meta: { numeric: true }, + size: 130, + sortDescFirst: false, + header: ({ column }) => ( + + ), + cell: ({ row }) => formatPerMillion(row.original.output_cost_per_token), + }, + ]; const clear = () => { models.onSearchChange(""); models.onColumnFiltersChange([]); }; return ( - model.model_group} - onSelect={onSelect} - empty={ - models.error ? ( - - ) : ( - } - noun={t("publicModelHub.modelsPane.noun")} - query={models.searchValue} - onClear={models.hasActiveQuery ? clear : undefined} +
+ +
+
+ - ) - } - toolbar={} - renderItem={(model) => } - renderDetail={(model) => } - footer={} - /> + +
+
+ {models.isFetching || models.isLoading + ? t("publicModelHub.catalog.loadingModels", { defaultValue: "Loading models…" }) + : !models.error && t("publicModelHub.catalog.modelCount", countOptions)} +
+ {models.error && models.rows.length > 0 && ( + + )} +
+ model.model_group} + sortingMode="server" + sorting={models.sorting} + onSortingChange={models.onSortingChange} + paginationMode="server" + pagination={models.pagination} + onPaginationChange={models.onPaginationChange} + rowCount={models.rowCount} + filterMode="server" + columnFilters={models.columnFilters} + onColumnFiltersChange={models.onColumnFiltersChange} + isLoading={models.isLoading} + isError={Boolean(models.error)} + errorMessage={t("publicModelHub.loadFailed.title", { noun: t("publicModelHub.modelsPane.noun") })} + onRetry={models.refetch} + noDataMessage={ + } + noun={t("publicModelHub.modelsPane.noun")} + query={models.searchValue} + onClear={models.hasActiveQuery ? clear : undefined} + /> + } + /> +
+
+ onSelect(null)}> + {opened.model && } + {missingName !== null && } + +
); } diff --git a/ui/litellm-dashboard/src/components/publicModelHub/PublicAiHub.integration.test.tsx b/ui/litellm-dashboard/src/components/publicModelHub/PublicAiHub.integration.test.tsx index bcddc1fe364..3fedeead150 100644 --- a/ui/litellm-dashboard/src/components/publicModelHub/PublicAiHub.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/publicModelHub/PublicAiHub.integration.test.tsx @@ -1,13 +1,19 @@ import { fireEvent, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; +import type { ReactNode } from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils"; +import type { Plugin } from "@/components/claude_code_plugins/types"; +import type { apiClient } from "@/components/networking"; + +import { chooseSelectOption, renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils"; import PublicAiHub from "./PublicAiHub"; const { apiGetMock, realGet, navigation, uiSettings, replaceMock, hubInfo } = vi.hoisted(() => ({ apiGetMock: vi.fn(), - realGet: { current: (..._args: unknown[]): Promise => Promise.reject(new Error("unset")) }, + realGet: { + current: (..._args: Parameters): Promise => Promise.reject(new Error("unset")), + }, navigation: { params: new URLSearchParams() }, uiSettings: { values: {} as Record }, replaceMock: vi.fn(), @@ -25,7 +31,13 @@ vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn(), replace: vi.fn(), refresh: vi.fn() }), })); -vi.mock("@/components/navbar", () => ({ default: () =>
diff --git a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/MigrationDetail.tsx b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/MigrationDetail.tsx index 0ad842b9078..62fee75071e 100644 --- a/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/MigrationDetail.tsx +++ b/ui/litellm-dashboard/src/components/ModelDeprecationAlerts/MigrationDetail.tsx @@ -203,7 +203,7 @@ export default function MigrationDetail({

Last seen {selected.lastUsedDate} UTC · {selected.requests.toLocaleString()} requests in the last 30 days

- +
+ {label} + {shortcut && ( + + {shortcut} + + )} + ); } @@ -82,28 +90,24 @@ const ChatShell: React.FC<{ children: React.ReactNode }> = ({ children }) => { const { t } = useTranslation(); const router = useRouter(); const pathname = stripTrailingSlash(usePathname() ?? ""); - const { - accessToken, - userEmail, - userId, - conversations, - activeConversationId, - deleteConversation, - renameConversation, - } = useChatShell(); + const { accessToken, conversations, activeConversationId, deleteConversation, renameConversation } = useChatShell(); const handleLogout = useLogout(accessToken); - const isDesktop = useMediaQuery("(min-width: 601px)", { initializeWithValue: false }); + const isDesktop = useMediaQuery("(min-width: 768px)", { initializeWithValue: false }); const [sidebarOpen, setSidebarOpen] = useState(false); const [sidebarCollapsed, setSidebarCollapsed] = useState(false); const [searchOpen, setSearchOpen] = useState(false); if (isDesktop && sidebarOpen) setSidebarOpen(false); const collapsed = isDesktop && sidebarCollapsed; - const { logoUrl, logoUrlDark } = useTheme(); - const lightLogo = logoUrl || logoUrlDark; - const darkLogo = logoUrlDark || logoUrl; + const [mobileNavigationHandle] = useState(() => DialogPrimitive.createHandle()); + const shortcut = useShortcutLabel(); + const openSearch = () => { + setSidebarCollapsed(false); + if (!isDesktop) setSidebarOpen(true); + setSearchOpen(true); + }; + useHotkeys("mod+k", openSearch, { enableOnFormTags: true, preventDefault: true, useKey: true }); const routes = getChatRoutes(); const isChatsRoute = pathname === routes.chats; - const accountName = userEmail?.split("@")[0] || userId || t("chat.shell.accountFallback"); const workspaceItems = [ { label: t("chat.shell.workspaceItems.credentials"), route: routes.credentials, Icon: KeyRound }, { label: t("chat.shell.workspaceItems.apiKeys"), route: routes.apiKeys, Icon: Lock }, @@ -119,43 +123,19 @@ const ChatShell: React.FC<{ children: React.ReactNode }> = ({ children }) => { const sidebar = ( <> -
- - - + setSidebarOpen(false)} + collapsed={collapsed} + onToggleCollapsed={() => (isDesktop ? setSidebarCollapsed(!sidebarCollapsed) : setSidebarOpen(false))} + /> +
+ + } label={t("chat.shell.newChat")} onClick={() => navigate(routes.chats)} />
- - -
+ } label={t("chat.shell.searchChats")} shortcut={shortcut} onClick={openSearch} /> + +
= ({ children }) => { searchOpen={searchOpen} onSearchOpenChange={setSearchOpen} hideSearchButton + enableSearchShortcut={false} />
- {lightLogo && darkLogo && ( -
- - -
- )} -
- - {accountName} - - - - -
+ + navigate(routes.credentials)} + /> + ); + const titles = [ + { route: routes.chats, label: t("shell.viewSwitcher.chat") }, + { route: routes.models, label: t("chat.shell.models") }, + { route: routes.integrations, label: t("chat.shell.integrations") }, + ...workspaceItems, + ]; + const activeConversation = conversations.find((conversation) => conversation.id === activeConversationId); + const title = + isChatsRoute && activeConversation?.messages.length + ? activeConversation.title + : titles.find((item) => item.route === pathname)?.label ?? t("shell.viewSwitcher.chat"); + return ( -
- {isDesktop ? ( - - ) : ( - -
+ + )} + + + + {t("chat.shell.sheetTitle")} + {t("chat.shell.sheetDescription")} + + {sidebar} + + +
+ } > - + - -
- - - {t("chat.shell.sheetTitle")} - {t("chat.shell.sheetDescription")} - - {sidebar} - - - )} -
{children}
+ } + /> +
{children}
+
); }; diff --git a/ui/litellm-dashboard/src/components/chat/ConversationList.tsx b/ui/litellm-dashboard/src/components/chat/ConversationList.tsx index 195c812bff4..ad759c8dc67 100644 --- a/ui/litellm-dashboard/src/components/chat/ConversationList.tsx +++ b/ui/litellm-dashboard/src/components/chat/ConversationList.tsx @@ -31,6 +31,7 @@ interface Props { searchOpen?: boolean; onSearchOpenChange?: (open: boolean) => void; hideSearchButton?: boolean; + enableSearchShortcut?: boolean; } type DateGroup = "Recents" | "Yesterday" | "Last 7 Days" | "Older"; @@ -116,7 +117,7 @@ const ConversationRow: React.FC = ({ conv, isActive, onSel return (
@@ -138,7 +139,7 @@ const ConversationRow: React.FC = ({ conv, isActive, onSel onClick={() => onSelect(conv.id)} aria-label={conv.title} aria-current={isActive ? "page" : undefined} - className={`min-w-0 flex-1 justify-start text-[13px] hover:bg-sidebar-accent ${ + className={`h-auto min-h-11 min-w-0 flex-1 justify-start text-[13px] hover:bg-sidebar-accent ${ isActive ? "font-medium" : "" }`} title={conv.title} @@ -260,7 +261,7 @@ const SearchModal: React.FC = ({ open, conversations, onSelect setQuery(e.target.value)} className="pl-9" @@ -303,6 +304,7 @@ const ConversationList: React.FC = ({ searchOpen, onSearchOpenChange, hideSearchButton = false, + enableSearchShortcut = true, }) => { const { t } = useTranslation(); const [localSearchOpen, setLocalSearchOpen] = useState(false); @@ -311,7 +313,8 @@ const ConversationList: React.FC = ({ const handleGlobalKeyDown = useCallback( (e: KeyboardEvent) => { - if (e.key === "k" && (e.metaKey || e.ctrlKey)) { + const isSearchShortcut = e.key.toLowerCase() === "k" && (e.metaKey || e.ctrlKey); + if (isSearchShortcut && !e.shiftKey) { e.preventDefault(); setSearchModalOpen(!searchModalOpen); } @@ -320,9 +323,10 @@ const ConversationList: React.FC = ({ ); useEffect(() => { + if (!enableSearchShortcut) return; document.addEventListener("keydown", handleGlobalKeyDown); return () => document.removeEventListener("keydown", handleGlobalKeyDown); - }, [handleGlobalKeyDown]); + }, [handleGlobalKeyDown, enableSearchShortcut]); const grouped = groupConversations(conversations); @@ -342,37 +346,24 @@ const ConversationList: React.FC = ({
)} - {grouped.length === 0 ? ( -
-
+ {grouped.map(({ group, items }) => ( +
+
-

- {t("chat.conversationList.empty")} -

+ {items.map((conv) => ( + + ))}
- ) : ( - grouped.map(({ group, items }) => ( -
-
-
- {items.map((conv) => ( - - ))} -
- )) - )} + ))}
diff --git a/ui/litellm-dashboard/src/components/chat/design.md b/ui/litellm-dashboard/src/components/chat/design.md index 97e0c2ebeea..98fc014b19c 100644 --- a/ui/litellm-dashboard/src/components/chat/design.md +++ b/ui/litellm-dashboard/src/components/chat/design.md @@ -2,7 +2,7 @@ Tactical guide for building chat UI features. Copy patterns exactly — don't improvise spacing, colors, or components. -The standalone `/ui/chat` workspace reuses Moyai's workspace structure and visual rules from `frontend/theme.css`, `native/web/style.css`, `brand.css`, `sidebar-layout.css`, and `home.css` at Moyai commit `1118ab0` (October 20, 2026). This is an intentional visual departure from the dashboard: neutral white/charcoal surfaces, restrained purple actions, and translucent navigation and composers. Shared controls and chat behavior remain owned by the existing LiteLLM primitives. +The standalone `/ui/chat` workspace shares the dashboard’s navigation primitives, header, account menu, typography and theme tokens. The previous Moyai glass treatment made Chat look like a separate product. Keep Chat’s conversation tools and route-specific navigation inside the shared visual language. > **Agents:** Read `AGENTS.md` alongside this file. It contains decision trees, the pre-commit checklist, and the rules that govern when to use each pattern here. This file is the _what_; `AGENTS.md` is the _when and how_. @@ -11,10 +11,10 @@ The standalone `/ui/chat` workspace reuses Moyai's workspace structure and visua ## Stack - **Next.js 16** App Router, TypeScript -- **Tailwind CSS v4** — shared primitives and utilities, with the standalone Moyai port scoped in `globals.css` (theme and rail) and `moyai/moyai-chat.css` (home, composer and transcript) +- **Tailwind CSS v4** — shared primitives and utilities, with Chat-specific content layout in `moyai/moyai-chat.css` (home, composer and transcript) - **shadcn/ui**: import from `@/components/ui/*`. Available today: `alert-dialog`, `badge`, `bubble`, `button`, `card`, `collapsible`, `dialog`, `input`, `input-group`, `label`, `popover`, `scroll-area`, `select`, `separator`, `sheet`, `skeleton`, `sonner`, `switch`, `table`, `tabs`, `textarea`, `tooltip`. Use these shared primitives when composing chat features - **lucide-react** — only icon library, no emoji in UI -- Moyai's General Sans font from its original Fontshare stylesheet, with bundled Geist fallback; Geist Mono is bundled for model names +- The dashboard font, inherited from the app root. No separate font downloads or Chat palette --- @@ -22,21 +22,17 @@ The standalone `/ui/chat` workspace reuses Moyai's workspace structure and visua Never hardcode hex values. Use these CSS variables via Tailwind classes — all verified present in `src/app/globals.css`. -`ChatShell` applies `.chat-theme`, which scopes the palette to standalone Chat. Its dialogs and mobile Sheet also carry `.chat-theme` because they render through portals. Do not apply these overrides globally or to the Playground. Light and dark modes follow LiteLLM's existing theme switch. Keep status colors semantic. +`ChatShell` and portalled dialogs carry `.chat-theme` only for Chat-specific message and editor aliases. Do not override the shared foreground, background, sidebar, primary or font tokens. Light and dark modes use the dashboard theme -### Glass surfaces +### Surfaces -- `.chat-workspace`: static neutral backdrop, with no blur on this layout ancestor. -- `.chat-glass-sidebar`: Moyai's frosted navigation rail with its original neutral sheen. -- `.chat-glass-pane`: transparent on the home view so the source glass backdrop remains visible; quiet reading fill on conversation and workspace pages, without nested backdrop filtering. -- `.moyai-chat-home-composer`: Moyai's 10px rounded glass editor, 132px minimum desktop entry height, 14px blur, original sheen/shadow and bordered lower action row. `MoyaiChatComposer` owns its layout. Model selection sits at the right beside send. Message and code surfaces stay solid, using `--chat-message` / `--chat-message-border`; editor focus uses `--chat-focus-border`. -- `.chat-mobile-sidebar`: elevated, nearly opaque Sheet surface so navigation remains legible over conversation content. +Use solid `bg-background` content and `bg-sidebar` navigation with `border-border` separators. No glass, dot/grid backgrounds, gradients or decorative shadows. The composer has a quiet border and a visible focus state. Model selection, attachments, tools and send remain in its action row -Glass progressively enhances solid fills inside `@supports`. Reduced transparency and forced colors disable filtering, sheen and decorative shadows. Reduced motion disables chrome transitions. Use the original static dot/grid backdrop and source home heading measurements. Do not add animated backgrounds or an outer rounded frame. +The welcome view contains a heading, composer and compact prompt shortcuts. Shortcuts fill the draft without submitting. Recent conversations belong in the sidebar, without an empty-count section. Show run and storage warnings when they apply, not as permanent idle helper text The standalone `ChatModelPicker` exposes the account's chat and image-generation models from the composer action row. Its portalled popover carries `.chat-theme`; keep model names, capabilities and selection states readable in both themes. Attachment and MCP tools stay beside the composer. -The **Models** sidebar destination opens `/ui/chat/models` inside the same shell. It reuses the existing Model Hub catalog, filters, pricing and details in a read-only `chat` variant. Admins see the authenticated catalog and other users see the published catalog, matching the dashboard's existing access rules. The page uses `.moyai-models-page` for scrolling and `.chat-model-hub-card` for the bordered Moyai surface. Portalled model details carry `.chat-theme`; catalog management and other hub tabs stay in the dashboard. +The **Models** sidebar destination opens `/ui/chat/models` inside the same shell. It reuses the existing Model Hub catalog, filters, pricing and details in a read-only `chat` variant. Admins see the authenticated catalog and other users see the published catalog, matching the dashboard's existing access rules. The page uses `.moyai-models-page` for scrolling and `.chat-model-hub-card` for the bordered dashboard surface. Portalled model details carry `.chat-theme`; catalog management and other hub tabs stay in the dashboard. Show available providers in a single-select dropdown beside model search. “All providers” clears the provider filter while keeping the other filters. The dropdown includes provider logos and uses the Chat theme in its portal. Provider logos reuse the existing asset mapping; white logo plates preserve the colors and contrast of the original assets in both themes. Unknown or failed logos fall back to a provider initial. The model list uses compact search and capability controls, quiet separators, and provider identities beside model names. Keep these styles scoped to Chat. @@ -58,7 +54,7 @@ Show available providers in a single-select dropdown beside model search. “All | `--sidebar-accent-foreground` | `text-sidebar-accent-foreground` | Active nav item text | | `--sidebar-border` | `border-sidebar-border` | Sidebar's own dividers/right border | -**Known gotcha:** generic `--accent`, `--secondary`, and `--muted` values may coincide in the dashboard theme. The standalone palette deliberately separates its sidebar selection from the chrome. Always use the sidebar-specific tokens rather than relying on incidental contrast from a generic token. Rules: +**Known gotcha:** generic `--accent`, `--secondary`, and `--muted` values may coincide in the dashboard theme. Always use the sidebar-specific tokens rather than relying on incidental contrast from a generic token. Rules: - Sidebar container: `bg-sidebar`, never `bg-secondary` or `bg-accent`. - Active/selected nav item: `bg-sidebar-accent text-sidebar-accent-foreground`, never `bg-accent` on its own inside the sidebar. @@ -81,7 +77,7 @@ Show available providers in a single-select dropdown beside model search. “All // Page/greeting heading

What are we working on?

-
// Moyai: 500 weight, clamp(32px, 3.4vw, 46px), 1.13 line-height +
// Section heading (panel titles: "Your API Keys", "Your Usage")

Your Usage

@@ -113,7 +109,7 @@ Show available providers in a single-select dropdown beside model search. “All | Context | Value | | ---------------------------------- | ------------------------------------------- | -| Sidebar outer padding | `8px 10px 8px` | +| Sidebar outer padding | `p-3` | | Nav item padding | `px-2.5 py-2` | | Inline icon-to-label gap | `gap-2.5` in nav items, `gap-1.5` elsewhere | | Section divider margin | `mx-2` | @@ -128,8 +124,8 @@ Show available providers in a single-select dropdown beside model search. “All ```tsx rounded - md; // 8px — buttons, inputs, nav items, bordered divs rounded - lg; // 10px — panel cards, table containers -rounded - xl; // 12px — shared controls (Moyai composer uses its source 10px rule) -rounded - lg; // 8px — Moyai user message bubbles +rounded - xl; // 12px — shared panels +rounded - lg; // 8px — user message bubbles rounded - full; // pills, avatars, status dots ``` @@ -161,198 +157,11 @@ Sizes: `default` (h-9), `sm` (h-8), `xs` (h-6), `lg` (h-10), `icon` (size-9), `i ### Sidebar nav item -No dedicated shadcn `Sidebar` block is installed (see gaps). Build nav items from `Button variant="ghost"`, full width, left-aligned, with the sidebar-specific active state: +Use `Sidebar`, `SidebarContent`, `SidebarFooter` and `SidebarMenuButton` from `components/shared/Sidebar`. Both workspaces use `SidebarBrand`, `ViewSwitcher` and `DashboardHeader`. Chat has a compact account menu with its existing identity, Settings and Sign out actions. Theme lives in the shared top bar. Gateway account preferences and admin controls stay in the gateway. Keep the shared 280px expanded and 72px collapsed widths, 56px header and 768px mobile breakpoint -```tsx - -``` +Use one New Chat action for the blank composer; do not duplicate it with a Chats destination. Chat retains Models, Integrations, Workspace and conversation history. Chat has no gateway-wide header search or worker selector. Search chats in the sidebar owns Cmd/Ctrl+K, including when the sidebar is collapsed or on mobile. Mobile navigation uses the same Sheet treatment as the dashboard -The desktop sidebar follows Moyai: 240px wide, 230px below 1250px, and 52px when collapsed. Its top row uses the shared `ViewSwitcher`, search and collapse controls above a transparent 32px New Chat button and compact destinations. The switcher exposes AI Gateway, Moyai, Chat and registered plugins, with the active workspace marked. Chat uses the same rounded speech-bubble icon in the trigger and menu, including the dashboard entry point. The collapsed rail keeps the icon-only switcher above its expand button. Workspace keeps credentials, API keys, logs, usage and preview feedback reachable in one menu. Recent conversation rows show title, model and date at 48px height. The footer shows account identity, appearance, settings and sign out. The shared dashboard Navbar is absent from this standalone surface; only its product switcher is reused. - -At 600px and below, a compact header keeps the workspace switcher beside the labelled menu button, visible even when the navigation Sheet is closed. The same navigation renders once inside the existing `Sheet`. Choosing a route or conversation closes the Sheet. Do not render two concurrent conversation lists: they each own a search shortcut listener. The collapsed desktop list remains mounted for keyboard search, but is visually hidden. The Chat layout owns an authenticated `PluginModeProvider` so returning to AI Gateway clears a previously selected plugin. - -The landing page reuses Moyai's left-aligned 720px home content, dot/grid backdrop, 46px heading, 132px editor and compact starter rows. `ChatModelPicker` sits at the right of the composer action row, matching Moyai's model selector position. Selecting a model chooses chat or image generation. Preserve LiteLLM's existing model access and request behavior; do not copy Moyai product or backend logic. - -Configured organization branding appears above the account row in the expanded sidebar and mobile navigation. Light and dark variants fall back to the available image when only one is configured. Private filesystem images use the proxy image endpoint without exposing their paths. With no configured branding, the logo row is absent. The existing collapsed rail hides it along with the account row. - -### ScrollArea - -**Known gotcha — cost real debugging time:** a `ScrollArea` won't actually scroll if its container is sized with `max-height` instead of `height`, even though the box _looks_ correctly capped. Radix's `ScrollArea` viewport is `size-full` (`height: 100%`), and CSS percentage-height resolution requires the ancestor chain to have a genuinely _definite_ height — `max-height` on a flex container doesn't count as definite for this purpose, even though the browser renders the box at the capped size. The result: the viewport silently grows to its full content height instead of 100% of the visible box, `scrollHeight === clientHeight` is reported (no overflow, so no scrollbar), and scroll events fall through to whatever's behind the popover/panel instead of scrolling the list. - -```tsx -// Wrong — looks capped, doesn't scroll (max-height isn't a definite height) -
- {items} -
- -// Right — explicit height propagates as definite through the whole chain -
- {items} -
-``` - -If you need "shrink to fit content, cap at N px" rather than "always N px," you cannot get both with pure CSS here — pick the fixed height (the common case: any list that can realistically exceed the cap, e.g. a model picker with hundreds of entries) over one that only works when short. - -### Input / Textarea - -Use `Input` and `Textarea` from `@/components/ui`. Standalone Chat uses `MoyaiChatComposer` for the source editor and toolbar arrangement, preserving submission, cancellation, attachments and IME handling. Uploads share a limit of four files per message and 10 MB per file. Images retain previews; PDF, DOCX, XLSX, text and code files use named file chips and native Responses `input_file` parts. File support depends on the selected model and provider. Attachment bytes stay in the existing user-scoped IndexedDB store; history metadata stays small. Legacy DOC/XLS files are not accepted. Other chat surfaces continue using the shared `ChatComposer` / `InputGroupTextarea`; this visual port is scoped to standalone Chat. - -```tsx -import { Input } from "@/components/ui/input"; -import { Label } from "@/components/ui/label"; - -
- - - {error &&

{error}

} -
; -``` - -Errors go **below** the field. Never a toast for form validation. - -### Card - -Use the installed Card primitives for grouped content and action reviews - -```tsx -import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; - - - - Review changes - - {children} -; -``` - -### Messages and Markdown - -`ChatMessages` owns the shared list renderer. `ChatMessageContent` renders one message without a list wrapper so a transcript scroller can own each row. Both retain user editing, reasoning, tool output, copying and response metrics - -Use the official registry `Bubble` and `BubbleContent` primitives. User messages use the muted variant and end alignment. Assistant messages use the ghost variant at full available width. Keep paragraph spacing and list markers in the shared Markdown renderer. GFM tables use the shared Table primitives, which provide horizontal scrolling. Forward Markdown cell styles so column alignment survives rendering - -LiteAdmin passes `allowImages={false}` to show image descriptions without automatically requesting external image URLs from administrative answers. Other chat consumers retain image rendering by default - -Scrolling belongs to the conversation container. LiteAdmin uses official MessageScroller items with stable chronological IDs; a pending action and its result stay in the same row. Its inline action review is an explicit product exception to the AlertDialog confirmation pattern below - -### Badge - -```tsx -import { Badge } from "@/components/ui/badge" - -connected -expiring -expired -``` - -### Dialog / AlertDialog - -```tsx -import { Dialog, DialogContent, DialogHeader, DialogTitle, DialogFooter } from "@/components/ui/dialog"; - - - - - Rotate key - -
{/* form */}
- - - - - -; -``` - -Use `AlertDialog` (not a manual confirm) for any destructive action — delete conversation, revoke credential. - -### Table - -```tsx -import { Table, TableHeader, TableBody, TableRow, TableHead, TableCell } from "@/components/ui/table"; - -
- - - - Key - - - - - sk-...abcd - - -
-
; -``` - -### Tabs - -Always pass `variant="line"` for section tabs (Apps: All/Connected). The default variant renders a boxed pill with a background/shadow meant for compact toggle groups, not section tabs, and combining it with hand-added `border-b-2` classes produces a broken double-border look. - -```tsx -import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; - - - - - All - - -; -``` - -### Loading states - -Skeleton loaders while fetching — never a bare spinner or blank screen for list/table content. `Loader2` + `animate-spin` is reserved for **inside a button** during a submit action, not as a page/panel loading state. - -```tsx -import { Skeleton } from "@/components/ui/skeleton"; - -
-
- - -
-
; -``` - -### Toast / notifications - -**Resolved:** `sonner` is installed and mounted from the root layout. Raise toasts through `toast` from `@/lib/toast`; don't introduce a second toast mechanism. - ---- - -## Page Layout Patterns - -### Chat shell (sidebar + routed content) - -```tsx -
- {/* optional full-width banner */} -
- -
{children}
-
-
-``` +The conversation title appears in the shared header. Do not repeat it in a second content header. Routed pages keep their content-specific headings and actions ### Routed panel (Integrations, Credentials, API Keys, Usage) @@ -391,14 +200,14 @@ No max-width cap on panel content — matches the rest of the dashboard's own pa Card, Textarea and sonner are installed. Earlier versions of this document listed them as missing and encouraged parallel implementations; use the shared owners above -The shadcn Sidebar primitive block (`SidebarProvider`/`SidebarMenu`/etc) is still absent. Use the "Sidebar nav item" pattern above, built from `Button variant="ghost"` and the `sidebar-*` tokens +Shared sidebar primitives are available in `components/shared/Sidebar`. Use them instead of hand-writing a separate rail ## File Structure ``` src/components/chat/ ChatShell.tsx # Desktop navigation / mobile Sheet, shared across all /chat/* routes - ConversationList.tsx # Date-grouped chat list + Cmd+K search + ConversationList.tsx # Date-grouped chat list and conversation search ChatMessages.tsx # Shared list and single-message rendering (user, assistant, tool) MCPAppsPanel.tsx # Integrations grid + detail view MCPConnectPicker.tsx # MCP server toggle popover diff --git a/ui/litellm-dashboard/src/components/chat/moyai/MoyaiChatHome.test.tsx b/ui/litellm-dashboard/src/components/chat/moyai/MoyaiChatHome.test.tsx index 9e590978bb8..c1889ad6bcc 100644 --- a/ui/litellm-dashboard/src/components/chat/moyai/MoyaiChatHome.test.tsx +++ b/ui/litellm-dashboard/src/components/chat/moyai/MoyaiChatHome.test.tsx @@ -1,8 +1,7 @@ -import { act, fireEvent, render, screen, within } from "@testing-library/react"; +import { act, fireEvent, render, screen } from "@testing-library/react"; import type { ComponentProps } from "react"; import { describe, expect, it, vi } from "vitest"; import MoyaiChatHome from "./MoyaiChatHome"; -import type { Conversation } from "../types"; function homeProps(overrides: Partial> = {}) { return { @@ -11,24 +10,16 @@ function homeProps(overrides: Partial> = {} researchMode: false, draft: "", busy: false, - conversations: [], - onSelectConversation: vi.fn(), onPrompt: vi.fn(), ...overrides, }; } -function conversation(id: string, updatedAt: number): Conversation { - return { id, title: `Chat ${id}`, model: "chat-model", messages: [], mcpServerNames: [], createdAt: 1, updatedAt }; -} - describe("MoyaiChatHome", () => { - it("should render the Moyai welcome and composer", () => { + it("should render the welcome and composer", () => { render(); expect(screen.getByRole("heading", { level: 1, name: "What are we working on?" })).toBeInTheDocument(); expect(screen.getByText("Chat composer")).toBeInTheDocument(); - expect(screen.getByRole("heading", { name: "Recent chats" })).toBeInTheDocument(); - expect(screen.getByRole("status")).toHaveTextContent("No recent chats in this browser yet."); }); it("should render the research welcome when research mode is enabled", () => { @@ -36,26 +27,6 @@ describe("MoyaiChatHome", () => { expect(screen.getByRole("heading", { level: 1, name: "What should we research?" })).toBeInTheDocument(); }); - it("should show the three most recently updated chats and open the selected one", () => { - const conversations = [ - conversation("first", 1), - conversation("second", 4), - conversation("third", 3), - conversation("fourth", 2), - ]; - const props = homeProps({ conversations }); - render(); - const recent = screen.getByRole("region", { name: "Recent chats" }); - const buttons = within(recent).getAllByRole("button"); - expect(buttons.map((button) => button.textContent)).toEqual(["Chat second", "Chat third", "Chat fourth"]); - expect(screen.queryByRole("button", { name: "Chat first" })).not.toBeInTheDocument(); - act(() => { - fireEvent.click(screen.getByRole("button", { name: "Chat third" })); - }); - expect(props.onSelectConversation).toHaveBeenCalledWith("third"); - expect(conversations.map((item) => item.id)).toEqual(["first", "second", "third", "fourth"]); - }); - it("should pass a starter prompt to the composer without submitting it", () => { const props = homeProps(); render(); @@ -63,7 +34,6 @@ describe("MoyaiChatHome", () => { fireEvent.click(screen.getByRole("button", { name: /Review code/ })); }); expect(props.onPrompt).toHaveBeenCalledWith("Review my changes for bugs and improvements."); - expect(props.onSelectConversation).not.toHaveBeenCalled(); }); it("should disable starters during a request", () => { diff --git a/ui/litellm-dashboard/src/components/chat/moyai/MoyaiChatHome.tsx b/ui/litellm-dashboard/src/components/chat/moyai/MoyaiChatHome.tsx index 45dc350acf1..818bed63b86 100644 --- a/ui/litellm-dashboard/src/components/chat/moyai/MoyaiChatHome.tsx +++ b/ui/litellm-dashboard/src/components/chat/moyai/MoyaiChatHome.tsx @@ -1,27 +1,23 @@ import type { ReactNode } from "react"; import type { ParseKeys } from "i18next"; -import { ArrowUpRight, Bug, Clock3, Code2, ListChecks, type LucideIcon } from "lucide-react"; +import { Bug, Code2, ListChecks, type LucideIcon } from "lucide-react"; import { Trans, useTranslation } from "react-i18next"; import { Button } from "@/components/ui/button"; -import type { Conversation } from "../types"; import "./moyai-chat.css"; -const starters: { label: ParseKeys; detail: ParseKeys; prompt: ParseKeys; Icon: LucideIcon }[] = [ +const starters: { label: ParseKeys; prompt: ParseKeys; Icon: LucideIcon }[] = [ { label: "chat.moyaiHome.starters.review.label", - detail: "chat.moyaiHome.starters.review.detail", prompt: "chat.moyaiHome.starters.review.prompt", Icon: Code2, }, { label: "chat.moyaiHome.starters.fixBug.label", - detail: "chat.moyaiHome.starters.fixBug.detail", prompt: "chat.moyaiHome.starters.fixBug.prompt", Icon: Bug, }, { label: "chat.moyaiHome.starters.planFeature.label", - detail: "chat.moyaiHome.starters.planFeature.detail", prompt: "chat.moyaiHome.starters.planFeature.prompt", Icon: ListChecks, }, @@ -33,8 +29,6 @@ export default function MoyaiChatHome({ researchMode, draft, busy, - conversations, - onSelectConversation, onPrompt, }: { children: ReactNode; @@ -42,15 +36,11 @@ export default function MoyaiChatHome({ researchMode: boolean; draft: string; busy: boolean; - conversations: Conversation[]; - onSelectConversation: (id: string) => void; onPrompt: (prompt: string) => void; }) { const { t } = useTranslation(); - const recent = [...conversations].sort((a, b) => b.updatedAt - a.updatedAt).slice(0, 3); return (
-