From bde00952b6aa68372e739e1f377cc9c32a9063f1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 23:19:21 +0000 Subject: [PATCH 001/576] fix(proxy): requeue Redis spend buffer transactions when DB commit fails The Redis transaction buffer leader drains the spend buffers with a destructive lpop before committing to the database. When the DB commit failed after exhausting retries, the popped transactions were only logged and then lost, permanently undercounting key/user/team/org/end-user/ team-member/tag/agent and daily spend after a database outage. Track each popped category and re-push the ones that were not committed back to their Redis buffers so a later scheduler tick retries them. Categories that already committed are not re-queued, so their spend is not double-counted. The daily tag spend path gets the same treatment. --- litellm/proxy/db/db_spend_update_writer.py | 45 +++++- .../redis_update_buffer.py | 55 +++++++ .../test_redis_update_buffer.py | 46 ++++++ .../proxy/db/test_db_spend_update_writer.py | 151 +++++++++++++++++- 4 files changed, 289 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 54a4c2dad91..cc266019ff7 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -797,6 +797,12 @@ class DBSpendUpdateWriter: ): verbose_proxy_logger.debug("acquired lock for spend updates") + # Track everything popped from Redis. Each category is removed once it + # has been committed to the DB, so whatever is left after a failure can + # be re-queued for the next tick instead of being lost. Committed + # categories are never re-queued, so their spend is not double-counted. + uncommitted: dict[str, Any] = {} # mutable-ok: drives which popped categories still need re-queuing + try: ( db_spend_update_transactions, @@ -807,6 +813,15 @@ class DBSpendUpdateWriter: daily_agent_spend_update_transactions, ) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() + uncommitted = { # mutable-ok: drives which popped categories still need re-queuing + "db_spend_update_transactions": db_spend_update_transactions, + "daily_spend_update_transactions": daily_spend_update_transactions, + "daily_team_spend_update_transactions": daily_team_spend_update_transactions, + "daily_org_spend_update_transactions": daily_org_spend_update_transactions, + "daily_end_user_spend_update_transactions": daily_end_user_spend_update_transactions, + "daily_agent_spend_update_transactions": daily_agent_spend_update_transactions, + } + if db_spend_update_transactions is not None: verbose_proxy_logger.info( "Spend tracking - committing spend updates from Redis to DB: " @@ -826,6 +841,7 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, db_spend_update_transactions=db_spend_update_transactions, ) + uncommitted.pop("db_spend_update_transactions", None) if daily_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_user_spend( @@ -834,6 +850,8 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_spend_update_transactions, ) + uncommitted.pop("daily_spend_update_transactions", None) + if daily_team_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_team_spend( n_retry_times=n_retry_times, @@ -841,6 +859,7 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_team_spend_update_transactions, ) + uncommitted.pop("daily_team_spend_update_transactions", None) if daily_org_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_org_spend( @@ -849,6 +868,7 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_org_spend_update_transactions, ) + uncommitted.pop("daily_org_spend_update_transactions", None) if daily_end_user_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_end_user_spend( @@ -857,6 +877,8 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_end_user_spend_update_transactions, ) + uncommitted.pop("daily_end_user_spend_update_transactions", None) + if daily_agent_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_agent_spend( n_retry_times=n_retry_times, @@ -864,14 +886,20 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_agent_spend_update_transactions, ) + uncommitted.pop("daily_agent_spend_update_transactions", None) except Exception as e: spend_log_error( "Spend tracking - failed to commit spend updates from Redis to DB. " - "Data already popped from Redis may be lost. Error: %s", + "Re-queuing uncommitted transactions to Redis for retry on next tick. Error: %s", str(e), exc=e, ) finally: + to_restore = { # mutable-ok: transient kwargs payload consumed immediately below + name: txns for name, txns in uncommitted.items() if txns is not None + } + if to_restore: + await self.redis_update_buffer.restore_transactions_to_redis(**to_restore) await self.pod_lock_manager.release_lock( cronjob_id=DB_SPEND_UPDATE_JOB_NAME, ) @@ -1020,11 +1048,11 @@ class DBSpendUpdateWriter: cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME, ): verbose_proxy_logger.debug("acquired lock for daily tag spend updates") + daily_tag_spend_update_transactions = ( + await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() + ) + committed = False try: - daily_tag_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() - ) - if daily_tag_spend_update_transactions: await DBSpendUpdateWriter.update_daily_tag_spend( n_retry_times=n_retry_times, @@ -1032,14 +1060,19 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_tag_spend_update_transactions, ) + committed = True except Exception as e: spend_log_error( "Spend tracking - failed to commit daily tag spend updates from Redis to DB. " - "Data already popped from Redis may be lost. Error: %s", + "Re-queuing to Redis for retry on next tick. Error: %s", str(e), exc=e, ) finally: + if not committed and daily_tag_spend_update_transactions: + await self.redis_update_buffer.restore_transactions_to_redis( + daily_tag_spend_update_transactions=daily_tag_spend_update_transactions, + ) await self.pod_lock_manager.release_lock( cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME, ) diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index c924448669d..b30fadd86ab 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -8,6 +8,8 @@ import asyncio import json from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast +from redis.exceptions import RedisError + from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache from litellm.constants import ( @@ -374,6 +376,59 @@ class RedisUpdateBuffer: if daily_txns: await daily_queue.update_queue.put(daily_txns) + async def restore_transactions_to_redis( + self, + db_spend_update_transactions: DBSpendUpdateTransactions | None = None, + daily_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, + daily_team_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, + daily_org_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, + daily_end_user_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, + daily_agent_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, + daily_tag_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, + ) -> None: + """ + Re-push transactions that were popped from Redis but not committed to the DB. + + The leader drains the buffers with a destructive ``lpop`` before committing to + the database. When a commit fails after its retries are exhausted, the popped + transactions must be pushed back so a later scheduler tick can retry them; + otherwise the aggregated spend is lost permanently. The re-pushed payloads use + the same JSON encoding as the store path, so the next drain parses them normally. + """ + if self.redis_cache is None: + return + + _configs = ( + (db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY), + (daily_spend_update_transactions, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY), + (daily_team_spend_update_transactions, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY), + (daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY), + (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY), + (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY), + (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY), + ) + + rpush_list: list[RedisPipelineRpushOperation] = [ # mutable-ok: async_rpush_pipeline requires a list arg + RedisPipelineRpushOperation(key=redis_key, values=[safe_dumps(transactions)]) + for transactions, redis_key in _configs + if transactions + ] + if len(rpush_list) == 0: + return + + try: + await self.redis_cache.async_rpush_pipeline(rpush_list=rpush_list) + verbose_proxy_logger.info( + "Spend tracking - restored %d uncommitted transaction set(s) to Redis for retry on next tick.", + len(rpush_list), + ) + except RedisError as e: + verbose_proxy_logger.error( + "Spend tracking - failed to restore uncommitted transactions to Redis. " + "These spend updates are lost. Error: %s", + str(e), + ) + @staticmethod def _number_of_transactions_to_store_in_redis( db_spend_update_transactions: DBSpendUpdateTransactions, diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 33372e7794a..79909561683 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -270,6 +270,52 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis(): assert result == (None, None, None, None, None, None) +@pytest.mark.asyncio +async def test_restore_transactions_to_redis_pushes_only_provided( + redis_update_buffer, mock_redis_cache +): + """ + restore_transactions_to_redis re-pushes only the transaction sets it was + given, to their matching buffer keys, so uncommitted spend can be retried. + """ + from litellm.constants import ( + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + REDIS_UPDATE_BUFFER_KEY, + ) + + mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1, 1]) + + db_spend = {"key_list_transactions": {"key1": 1.0}} + daily_user = {"user_key1": {"spend": 1.0}} + + await redis_update_buffer.restore_transactions_to_redis( + db_spend_update_transactions=db_spend, + daily_spend_update_transactions=daily_user, + ) + + mock_redis_cache.async_rpush_pipeline.assert_called_once() + rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] + pushed_keys = {op["key"] for op in rpush_list} + assert pushed_keys == { + REDIS_UPDATE_BUFFER_KEY, + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + } + # Payloads round-trip through the same JSON encoding used on the store path + payloads = {op["key"]: json.loads(op["values"][0]) for op in rpush_list} + assert payloads[REDIS_UPDATE_BUFFER_KEY] == db_spend + assert payloads[REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY] == daily_user + + +@pytest.mark.asyncio +async def test_restore_transactions_to_redis_noop_when_empty( + redis_update_buffer, mock_redis_cache +): + """Nothing to restore -> no Redis call.""" + mock_redis_cache.async_rpush_pipeline = AsyncMock() + await redis_update_buffer.restore_transactions_to_redis() + mock_redis_cache.async_rpush_pipeline.assert_not_called() + + def test_validate_redis_transaction_buffer_raises_without_redis(): """ When use_redis_transaction_buffer=true but no Redis cache is configured, diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 4c17c5d3482..10544e82453 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1532,9 +1532,9 @@ async def test_commit_spend_updates_uses_pipeline(): mock_redis_update_buffer = AsyncMock() mock_redis_update_buffer.store_in_memory_spend_updates_in_redis = AsyncMock() - # Return all-None tuple (no data to commit) + # Return all-None tuple (no data to commit); the pipeline yields 6 slots mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = ( - AsyncMock(return_value=(None, None, None, None, None, None, None)) + AsyncMock(return_value=(None, None, None, None, None, None)) ) db_writer.redis_update_buffer = mock_redis_update_buffer @@ -1565,6 +1565,153 @@ async def test_commit_spend_updates_uses_pipeline(): mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer.assert_not_called() +@pytest.mark.asyncio +async def test_commit_with_redis_requeues_all_on_db_failure(): + """ + Regression for #33872: if the DB commit fails after the leader has already + popped transactions from Redis, the popped transactions must be re-queued to + Redis so a later tick can retry them, instead of being silently lost. + """ + db_writer = DBSpendUpdateWriter() + + db_spend = { + "user_list_transactions": {"user1": 1.5}, + "end_user_list_transactions": {}, + "key_list_transactions": {"key1": 1.5}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + daily_user = {"user_key1": {"spend": 1.5, "api_requests": 1}} + + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(db_spend, daily_user, None, None, None, None) + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + + mock_pod_lock_manager = AsyncMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + db_writer.pod_lock_manager = mock_pod_lock_manager + + # Every DB write raises -> simulates a full database outage + db_writer._commit_spend_updates_to_db = AsyncMock(side_effect=Exception("db down")) + + with patch.object( + DBSpendUpdateWriter, + "update_daily_user_spend", + new=AsyncMock(side_effect=Exception("db down")), + ): + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=MagicMock(), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + # Both failed categories must be re-queued to Redis, nothing lost + mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once() + _, kwargs = mock_redis_update_buffer.restore_transactions_to_redis.call_args + assert kwargs["db_spend_update_transactions"] == db_spend + assert kwargs["daily_spend_update_transactions"] == daily_user + # The lock must still be released + mock_pod_lock_manager.release_lock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_commit_with_redis_only_requeues_failed_category(): + """ + A partial DB failure must not re-queue categories that already committed, + otherwise their spend would be double-counted on the next tick. + """ + db_writer = DBSpendUpdateWriter() + + db_spend = { + "user_list_transactions": {"user1": 1.5}, + "end_user_list_transactions": {}, + "key_list_transactions": {}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + daily_user = {"user_key1": {"spend": 1.5, "api_requests": 1}} + + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(db_spend, daily_user, None, None, None, None) + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + + mock_pod_lock_manager = AsyncMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + db_writer.pod_lock_manager = mock_pod_lock_manager + + # db_spend commits fine; only the daily user commit fails + db_writer._commit_spend_updates_to_db = AsyncMock() + + with patch.object( + DBSpendUpdateWriter, + "update_daily_user_spend", + new=AsyncMock(side_effect=Exception("db down")), + ): + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=MagicMock(), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once() + _, kwargs = mock_redis_update_buffer.restore_transactions_to_redis.call_args + # Only the failed daily category is requeued; the committed db_spend is not + assert kwargs == {"daily_spend_update_transactions": daily_user} + + +@pytest.mark.asyncio +async def test_commit_with_redis_no_requeue_on_success(): + """When all commits succeed, nothing should be re-queued to Redis.""" + db_writer = DBSpendUpdateWriter() + + db_spend = { + "user_list_transactions": {"user1": 1.5}, + "end_user_list_transactions": {}, + "key_list_transactions": {}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(db_spend, None, None, None, None, None) + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + + mock_pod_lock_manager = AsyncMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + db_writer.pod_lock_manager = mock_pod_lock_manager + + db_writer._commit_spend_updates_to_db = AsyncMock() + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=MagicMock(), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + mock_redis_update_buffer.restore_transactions_to_redis.assert_not_awaited() + + @pytest.mark.parametrize( "bucket_name,input_dict,table_attr,method_name,where_key,expected_order", [ From 118b47a8a39a79922f09c851364e43e3c73d0ce1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 23:45:23 +0000 Subject: [PATCH 002/576] fix(proxy): keep tag drain inside try and cover requeue paths with tests Move the destructive daily-tag Redis drain back inside the try so a Redis read failure still releases the pod lock via the finally block, and use a covariant Mapping for the restore signature. Add regression tests for the daily-tag requeue-on-failure/no-requeue-on-success paths and the RedisError swallow branch in restore_transactions_to_redis. --- litellm/proxy/db/db_spend_update_writer.py | 13 ++-- .../redis_update_buffer.py | 13 ++-- .../test_redis_update_buffer.py | 18 +++++ .../proxy/db/test_db_spend_update_writer.py | 72 +++++++++++++++++++ 4 files changed, 102 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index cc266019ff7..f13bf2e2105 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -797,11 +797,7 @@ class DBSpendUpdateWriter: ): verbose_proxy_logger.debug("acquired lock for spend updates") - # Track everything popped from Redis. Each category is removed once it - # has been committed to the DB, so whatever is left after a failure can - # be re-queued for the next tick instead of being lost. Committed - # categories are never re-queued, so their spend is not double-counted. - uncommitted: dict[str, Any] = {} # mutable-ok: drives which popped categories still need re-queuing + uncommitted: dict[str, Any] = {} # mutable-ok: tracks popped categories still needing commit try: ( @@ -1048,11 +1044,12 @@ class DBSpendUpdateWriter: cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME, ): verbose_proxy_logger.debug("acquired lock for daily tag spend updates") - daily_tag_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() - ) + daily_tag_spend_update_transactions: dict[str, DailyTagSpendTransaction] | None = None committed = False try: + daily_tag_spend_update_transactions = ( + await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() + ) if daily_tag_spend_update_transactions: await DBSpendUpdateWriter.update_daily_tag_spend( n_retry_times=n_retry_times, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index b30fadd86ab..660fd514d99 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -6,6 +6,7 @@ This is to prevent deadlocks and improve reliability import asyncio import json +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from redis.exceptions import RedisError @@ -379,12 +380,12 @@ class RedisUpdateBuffer: async def restore_transactions_to_redis( self, db_spend_update_transactions: DBSpendUpdateTransactions | None = None, - daily_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, - daily_team_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, - daily_org_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, - daily_end_user_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, - daily_agent_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, - daily_tag_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, + daily_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, + daily_team_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, + daily_org_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, + daily_end_user_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, + daily_agent_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, + daily_tag_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, ) -> None: """ Re-push transactions that were popped from Redis but not committed to the DB. diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 79909561683..3325893c5f6 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -316,6 +316,24 @@ async def test_restore_transactions_to_redis_noop_when_empty( mock_redis_cache.async_rpush_pipeline.assert_not_called() +@pytest.mark.asyncio +async def test_restore_transactions_to_redis_swallows_redis_error( + redis_update_buffer, mock_redis_cache +): + """A Redis failure during restore must not propagate to the caller's finally block.""" + from redis.exceptions import RedisError + + mock_redis_cache.async_rpush_pipeline = AsyncMock( + side_effect=RedisError("redis down") + ) + + await redis_update_buffer.restore_transactions_to_redis( + db_spend_update_transactions={"key_list_transactions": {"key1": 1.0}}, + ) + + mock_redis_cache.async_rpush_pipeline.assert_called_once() + + def test_validate_redis_transaction_buffer_raises_without_redis(): """ When use_redis_transaction_buffer=true but no Redis cache is configured, diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 10544e82453..06d06b50234 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1712,6 +1712,78 @@ async def test_commit_with_redis_no_requeue_on_success(): mock_redis_update_buffer.restore_transactions_to_redis.assert_not_awaited() +@pytest.mark.asyncio +async def test_commit_daily_tag_spend_requeues_on_db_failure(): + """A failed daily tag commit must re-queue the popped tag transactions and release the lock.""" + db_writer = DBSpendUpdateWriter() + + daily_tag = {"tag_key1": {"spend": 1.5, "api_requests": 1}} + + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis = AsyncMock() + mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer = AsyncMock( + return_value=daily_tag + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + + mock_pod_lock_manager = AsyncMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + db_writer.pod_lock_manager = mock_pod_lock_manager + + with patch.object( + DBSpendUpdateWriter, + "update_daily_tag_spend", + new=AsyncMock(side_effect=Exception("db down")), + ): + await db_writer._commit_daily_tag_spend_to_db_with_redis( + prisma_client=MagicMock(), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once_with( + daily_tag_spend_update_transactions=daily_tag, + ) + mock_pod_lock_manager.release_lock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_commit_daily_tag_spend_no_requeue_on_success(): + """A successful daily tag commit must not re-queue anything.""" + db_writer = DBSpendUpdateWriter() + + daily_tag = {"tag_key1": {"spend": 1.5, "api_requests": 1}} + + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis = AsyncMock() + mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer = AsyncMock( + return_value=daily_tag + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + + mock_pod_lock_manager = AsyncMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + db_writer.pod_lock_manager = mock_pod_lock_manager + + with patch.object( + DBSpendUpdateWriter, + "update_daily_tag_spend", + new=AsyncMock(), + ): + await db_writer._commit_daily_tag_spend_to_db_with_redis( + prisma_client=MagicMock(), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + mock_redis_update_buffer.restore_transactions_to_redis.assert_not_awaited() + mock_pod_lock_manager.release_lock.assert_awaited_once() + + @pytest.mark.parametrize( "bucket_name,input_dict,table_attr,method_name,where_key,expected_order", [ From ab997e04eb4f0f50bc2c6ae738231455cdd97329 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 25 Jul 2026 00:09:21 +0000 Subject: [PATCH 003/576] fix(caching): cache anthropic /v1/messages responses, including streaming anthropic_messages was missing from the cache's supported call types, so every /v1/messages request went to the provider. Adding it alone is not enough: the cache key is built from the OpenAI-ish param set, which has no system, top_k or stop_sequences, so two requests differing only by system prompt shared an entry and the second got the first one's answer. The Anthropic Messages request shape now feeds the key set as well. Streaming responses return to the caller before async_set_cache runs, so they are teed on the way out and the SSE events are stored verbatim once the stream reaches message_stop without a provider error. A hit replays those bytes and logs the request as a cache hit with zero cost. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/caching.py | 45 +---- litellm/caching/caching_handler.py | 41 +++- .../litellm_core_utils/model_param_helper.py | 19 +- .../messages/response_cache.py | 163 ++++++++++++++++ .../anthropic_passthrough_logging_handler.py | 16 +- .../streaming_handler.py | 2 +- litellm/types/caching.py | 19 ++ litellm/utils.py | 5 +- tests/test_litellm/caching/test_caching.py | 21 ++ .../messages/test_response_cache.py | 179 ++++++++++++++++++ 10 files changed, 457 insertions(+), 53 deletions(-) create mode 100644 litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 34badaa3e8a..88a5e08604e 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -67,20 +67,7 @@ class Cache: default_in_memory_ttl: Optional[float] = None, default_in_redis_ttl: Optional[float] = None, similarity_threshold: Optional[float] = None, - supported_call_types: Optional[List[CachingSupportedCallTypes]] = [ - "completion", - "acompletion", - "embedding", - "aembedding", - "atranscription", - "transcription", - "atext_completion", - "text_completion", - "arerank", - "rerank", - "responses", - "aresponses", - ], + supported_call_types: list[CachingSupportedCallTypes] | None = list(DEFAULT_CACHING_SUPPORTED_CALL_TYPES), # s3 Bucket, boto3 configuration azure_account_url: Optional[str] = None, azure_blob_container: Optional[str] = None, @@ -930,20 +917,7 @@ def enable_cache( host: Optional[str] = None, port: Optional[str] = None, password: Optional[str] = None, - supported_call_types: Optional[List[CachingSupportedCallTypes]] = [ - "completion", - "acompletion", - "embedding", - "aembedding", - "atranscription", - "transcription", - "atext_completion", - "text_completion", - "arerank", - "rerank", - "responses", - "aresponses", - ], + supported_call_types: list[CachingSupportedCallTypes] | None = list(DEFAULT_CACHING_SUPPORTED_CALL_TYPES), **kwargs, ): """ @@ -990,20 +964,7 @@ def update_cache( host: Optional[str] = None, port: Optional[str] = None, password: Optional[str] = None, - supported_call_types: Optional[List[CachingSupportedCallTypes]] = [ - "completion", - "acompletion", - "embedding", - "aembedding", - "atranscription", - "transcription", - "atext_completion", - "text_completion", - "arerank", - "rerank", - "responses", - "aresponses", - ], + supported_call_types: list[CachingSupportedCallTypes] | None = list(DEFAULT_CACHING_SUPPORTED_CALL_TYPES), **kwargs, ): """ diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index b17e055c7ea..8b2d033f24a 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -116,7 +116,8 @@ def _should_defer_streaming_cache_hit_callbacks(*, kwargs: Dict[str, Any]) -> bo When stream=True, do not run success callbacks at cache-hit time. Cached chat/text completion replay uses CustomStreamWrapper; cached Responses - replay uses CachedResponsesAPIStreamingIterator. Both invoke logging success + replay uses CachedResponsesAPIStreamingIterator; cached Anthropic Messages + replay uses CachedAnthropicMessagesStreamIterator. All invoke logging success handlers when the stream finishes; firing them here too would double-count spend and callback records. """ @@ -848,6 +849,18 @@ class LLMCachingHandler: response_type="audio_transcription", hidden_params=hidden_params, ) + elif ( + call_type == CallTypes.anthropic_messages.value or call_type == CallTypes.aanthropic_messages.value + ) and isinstance(cached_result, dict): + from litellm.llms.anthropic.experimental_pass_through.messages.response_cache import ( + convert_cached_anthropic_messages_result, + ) + + cached_result = convert_cached_anthropic_messages_result( + cached_result=cached_result, + logging_obj=logging_obj, + kwargs=kwargs, + ) elif (call_type == "aresponses" or call_type == "responses") and isinstance(cached_result, dict): use_chat_completion_cache = _is_chat_completion_cached_dict(cached_result) if use_chat_completion_cache: @@ -1044,6 +1057,32 @@ class LLMCachingHandler: and (kwargs.get("cache", {}).get("no-store", False) is not True) ) + def wrap_streaming_result_for_cache(self, result: Any, call_type: str) -> Any: + """ + Tee a streaming result so it still reaches the cache. + + Streaming responses are returned to the caller before ``async_set_cache`` + runs. Chat/text completion streams are teed inside ``CustomStreamWrapper`` + and Responses API streams inside their own iterator; Anthropic Messages + streams have no such hook, so they are wrapped here. + """ + if call_type not in ( + CallTypes.anthropic_messages.value, + CallTypes.aanthropic_messages.value, + ): + return result + if litellm.cache is None or not self._should_store_result_in_cache( + original_function=self.original_function, kwargs=self.request_kwargs + ): + return result + if not hasattr(result, "__anext__"): + return result + from litellm.llms.anthropic.experimental_pass_through.messages.response_cache import ( + AnthropicMessagesStreamCacheWriter, + ) + + return AnthropicMessagesStreamCacheWriter(stream=result, caching_handler=self) + def _is_call_type_supported_by_cache( self, original_function: Callable, diff --git a/litellm/litellm_core_utils/model_param_helper.py b/litellm/litellm_core_utils/model_param_helper.py index 39b3f0d5376..cf4eba933b8 100644 --- a/litellm/litellm_core_utils/model_param_helper.py +++ b/litellm/litellm_core_utils/model_param_helper.py @@ -18,6 +18,7 @@ from openai.types.responses.response_create_params import ( ) from litellm._logging import verbose_logger +from litellm.types.llms.anthropic import AnthropicMessagesRequest from litellm.types.rerank import RerankRequest @@ -40,7 +41,7 @@ class ModelParamHelper: @staticmethod def get_exclude_params_for_model_parameters() -> Set[str]: - return set(["messages", "prompt", "input"]) + return set(["messages", "prompt", "input", "system"]) @staticmethod def _get_relevant_args_to_use_for_logging() -> Set[str]: @@ -73,6 +74,7 @@ class ModelParamHelper: transcription_kwargs = ModelParamHelper._get_litellm_supported_transcription_kwargs() rerank_kwargs = ModelParamHelper._get_litellm_supported_rerank_kwargs() responses_api_kwargs = ModelParamHelper._get_litellm_supported_responses_api_kwargs() + anthropic_messages_kwargs = ModelParamHelper._get_litellm_supported_anthropic_messages_kwargs() exclude_kwargs = ModelParamHelper._get_exclude_kwargs() combined_kwargs = chat_completion_kwargs.union( @@ -81,6 +83,7 @@ class ModelParamHelper: transcription_kwargs, rerank_kwargs, responses_api_kwargs, + anthropic_messages_kwargs, ) combined_kwargs = combined_kwargs.difference(exclude_kwargs) return combined_kwargs @@ -167,12 +170,24 @@ class ModelParamHelper: streaming_params: Set[str] = set(getattr(ResponseCreateParamsStreaming, "__annotations__", {}).keys()) return non_streaming_params.union(streaming_params) + @staticmethod + def _get_litellm_supported_anthropic_messages_kwargs() -> set[str]: + """ + Get the litellm supported Anthropic /v1/messages kwargs + + This follows the Anthropic Messages API spec. `system`, `top_k` and + `stop_sequences` have no OpenAI equivalent, so without them the cache key + for a /v1/messages request ignores them and collides across requests that + differ only by system prompt. + """ + return set(getattr(AnthropicMessagesRequest, "__annotations__", {}).keys()) + @staticmethod def _get_exclude_kwargs() -> Set[str]: """ Get the kwargs to exclude from the cache key """ - return set(["metadata"]) + return set(["metadata", "litellm_metadata"]) ModelParamHelper._relevant_logging_args = frozenset(ModelParamHelper._get_relevant_args_to_use_for_logging()) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py new file mode 100644 index 00000000000..e94d8f6bbaf --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py @@ -0,0 +1,163 @@ +""" +Response caching for Anthropic Messages (`/v1/messages`) requests. + +Non-streaming responses are plain dicts and are stored by the generic caching +handler. Streaming responses are returned to the caller before +``LLMCachingHandler.async_set_cache`` runs, so they are teed here instead: the +SSE events are buffered while they are forwarded and persisted verbatim once the +stream completes, and a hit replays exactly what the provider sent. +""" + +from collections.abc import AsyncIterator +from typing import TYPE_CHECKING, Any, cast + +import litellm +from litellm._logging import verbose_logger +from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + BaseAnthropicMessagesStreamingIterator, + _is_message_stop_chunk, + _is_provider_error_chunk, + aclose_if_supported, +) +from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, +) + +if TYPE_CHECKING: + from litellm.caching.caching_handler import LLMCachingHandler + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +else: + LLMCachingHandler = Any + LiteLLMLoggingObj = Any + +CACHED_STREAM_EVENTS_KEY = "litellm_cached_anthropic_sse_events" + + +def _decode(chunk: bytes | str) -> str: + return chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk + + +class AnthropicMessagesStreamCacheWriter: + """ + Forwards a `/v1/messages` SSE stream unchanged while buffering it, then + writes the collected events to the response cache on normal completion. + + Only a stream that ran to a ``message_stop`` without a provider ``error`` + event is written, so partial or failed responses cannot be replayed. + """ + + def __init__( + self, + stream: AsyncIterator[bytes | str], + caching_handler: "LLMCachingHandler", + ) -> None: + self.stream = stream + self.caching_handler = caching_handler + self.collected_events: list[str] = [] + self.saw_message_stop = False + self.saw_provider_error = False + self.persisted = False + self._hidden_params: dict[str, Any] = getattr(stream, "_hidden_params", {}) or {} + + def __aiter__(self) -> "AnthropicMessagesStreamCacheWriter": + return self + + async def __anext__(self) -> bytes | str: + try: + chunk = await self.stream.__anext__() + except StopAsyncIteration: + await self._persist() + raise + chunk_bytes = chunk.encode("utf-8") if isinstance(chunk, str) else chunk + self.saw_message_stop = self.saw_message_stop or _is_message_stop_chunk(chunk_bytes) + self.saw_provider_error = self.saw_provider_error or _is_provider_error_chunk(chunk_bytes) + self.collected_events.append(_decode(chunk)) + return chunk + + async def aclose(self) -> None: + await aclose_if_supported(self.stream) + + async def _persist(self) -> None: + if self.persisted or litellm.cache is None: + return + if not self.saw_message_stop or self.saw_provider_error: + return + self.persisted = True + + request_kwargs = dict(self.caching_handler.request_kwargs) + if not self.caching_handler._should_store_result_in_cache( + original_function=self.caching_handler.original_function, + kwargs=request_kwargs, + ): + return + preset_cache_key = self.caching_handler.preset_cache_key + if preset_cache_key is not None: + request_kwargs["cache_key"] = preset_cache_key + + try: + await litellm.cache.async_add_cache( + {CACHED_STREAM_EVENTS_KEY: self.collected_events}, + dynamic_cache_object=self.caching_handler.dual_cache, + **request_kwargs, + ) + except Exception as e: + verbose_logger.exception("Anthropic Messages stream cache write failed: %s", e) + + +class CachedAnthropicMessagesStreamIterator(BaseAnthropicMessagesStreamingIterator): + """ + Replays cached `/v1/messages` SSE events and logs the request as a cache hit + once the replay finishes, mirroring what the live stream logs at end of stream. + """ + + def __init__( + self, + events: list[str], + litellm_logging_obj: LiteLLMLoggingObj, + request_body: dict[str, Any], + ) -> None: + super().__init__(litellm_logging_obj=litellm_logging_obj, request_body=request_body) + self.chunks: list[bytes] = [event.encode("utf-8") for event in events] + self.current_index = 0 + self._hidden_params: dict[str, Any] = {"cache_hit": True} + litellm_logging_obj.model_call_details["cache_hit"] = True + + def __aiter__(self) -> "CachedAnthropicMessagesStreamIterator": + return self + + async def __anext__(self) -> bytes: + if self.current_index >= len(self.chunks): + await self._handle_streaming_logging(self.chunks) + raise StopAsyncIteration + chunk = self.chunks[self.current_index] + self.current_index += 1 + return chunk + + +def get_cached_stream_events(cached_result: dict[str, Any]) -> list[str] | None: + events = cached_result.get(CACHED_STREAM_EVENTS_KEY) + if isinstance(events, list): + return [_decode(event) for event in events] + return None + + +def convert_cached_anthropic_messages_result( + cached_result: dict[str, Any], + logging_obj: LiteLLMLoggingObj, + kwargs: dict[str, Any], +) -> AnthropicMessagesResponse | CachedAnthropicMessagesStreamIterator: + """ + Turn a cached `/v1/messages` entry back into what the caller expects: an + SSE replay iterator for a streamed entry, otherwise the response itself + (``AnthropicMessagesResponse`` is a TypedDict, i.e. a dict at runtime). + """ + events = get_cached_stream_events(cached_result) + if events is not None: + return CachedAnthropicMessagesStreamIterator( + events=events, + litellm_logging_obj=logging_obj, + request_body=kwargs, + ) + return cast( # cast-ok: AnthropicMessagesResponse is a TypedDict; validating would drop provider fields we must replay verbatim + AnthropicMessagesResponse, cached_result + ) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 50e90699194..51813983876 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -255,12 +255,16 @@ class AnthropicPassthroughLoggingHandler: litellm_params=(logging_obj.litellm_params if hasattr(logging_obj, "litellm_params") else None) ) - response_cost = litellm.completion_cost( - completion_response=litellm_model_response, - model=model_for_cost, - custom_llm_provider=custom_llm_provider, - custom_pricing=custom_pricing, - router_model_id=router_model_id, + response_cost = ( + 0.0 + if logging_obj.model_call_details.get("cache_hit") is True + else litellm.completion_cost( + completion_response=litellm_model_response, + model=model_for_cost, + custom_llm_provider=custom_llm_provider, + custom_pricing=custom_pricing, + router_model_id=router_model_id, + ) ) kwargs["response_cost"] = response_cost diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 4dc1e0e70dd..24e5f1d16d5 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -161,7 +161,7 @@ class PassThroughStreamingHandler: result=standard_logging_response_object, start_time=start_time, end_time=end_time, - cache_hit=False, + cache_hit=litellm_logging_obj.model_call_details.get("cache_hit") is True, prefer_async_handlers=True, **kwargs, ) diff --git a/litellm/types/caching.py b/litellm/types/caching.py index eaa80c2f525..4255a8bd7fc 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -30,8 +30,27 @@ CachingSupportedCallTypes = Literal[ "rerank", "responses", "aresponses", + "anthropic_messages", + "aanthropic_messages", ] +DEFAULT_CACHING_SUPPORTED_CALL_TYPES: tuple[CachingSupportedCallTypes, ...] = ( + "completion", + "acompletion", + "embedding", + "aembedding", + "atranscription", + "transcription", + "atext_completion", + "text_completion", + "arerank", + "rerank", + "responses", + "aresponses", + "anthropic_messages", + "aanthropic_messages", +) + class RedisPipelineIncrementOperation(TypedDict): """ diff --git a/litellm/utils.py b/litellm/utils.py index a11c5500503..f5c8330c284 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1708,7 +1708,10 @@ def client(original_function): start_time=start_time, end_time=end_time, ) - return result + return _llm_caching_handler.wrap_streaming_result_for_cache( + result=result, + call_type=call_type, + ) elif call_type == CallTypes.arealtime.value: return result ### POST-CALL RULES ### diff --git a/tests/test_litellm/caching/test_caching.py b/tests/test_litellm/caching/test_caching.py index eaee54bac5a..b65e8773c85 100644 --- a/tests/test_litellm/caching/test_caching.py +++ b/tests/test_litellm/caching/test_caching.py @@ -1,6 +1,8 @@ import logging import re +import pytest + from litellm.caching.caching import Cache from litellm.types.caching import LiteLLMCacheType from litellm.types.utils import Embedding, EmbeddingResponse, Usage @@ -146,3 +148,22 @@ def test_exact_cache_key_still_includes_prompt(): model="gpt-4o-mini", messages=[{"role": "user", "content": "b"}] ) assert key_a != key_b + + +@pytest.mark.parametrize( + "anthropic_param", + [ + {"system": "answer ALPHA"}, + {"top_k": 5}, + {"stop_sequences": ["STOP"]}, + ], +) +def test_exact_cache_key_includes_anthropic_messages_params(anthropic_param): + """Anthropic /v1/messages params with no OpenAI equivalent must still key the + cache; without them two requests that differ only by system prompt collide.""" + cache = Cache(type=LiteLLMCacheType.LOCAL) + messages = [{"role": "user", "content": "which greek letter?"}] + baseline = cache.get_cache_key(model="claude-sonnet-4-5", messages=messages) + assert baseline != cache.get_cache_key( + model="claude-sonnet-4-5", messages=messages, **anthropic_param + ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py new file mode 100644 index 00000000000..344152cd828 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py @@ -0,0 +1,179 @@ +import asyncio +import os +import sys +from typing import Any, AsyncIterator, Dict, List + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +import litellm +from litellm.caching.caching import Cache, LiteLLMCacheType +from litellm.llms.anthropic.experimental_pass_through.messages import handler + +STREAM_EVENTS: List[bytes] = [ + b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_stream_1", "type": "message", ' + b'"role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": null, ' + b'"usage": {"input_tokens": 10, "output_tokens": 0}}}\n\n', + b'event: content_block_start\ndata: {"type": "content_block_start", "index": 0, ' + b'"content_block": {"type": "text", "text": ""}}\n\n', + b'event: content_block_delta\ndata: {"type": "content_block_delta", "index": 0, ' + b'"delta": {"type": "text_delta", "text": "ALPHA"}}\n\n', + b'event: content_block_stop\ndata: {"type": "content_block_stop", "index": 0}\n\n', + b'event: message_delta\ndata: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, ' + b'"usage": {"output_tokens": 3}}\n\n', + b'event: message_stop\ndata: {"type": "message_stop"}\n\n', +] + + +def _anthropic_response(message_id: str, text: str) -> Dict[str, Any]: + return { + "id": message_id, + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": text}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 3}, + } + + +class _CountingHandler: + """Stands in for the provider dispatch so cache hits are observable as skipped calls.""" + + def __init__(self, results: List[Any]) -> None: + self.results = results + self.calls: List[Dict[str, Any]] = [] + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + self.calls.append(kwargs) + return self.results[min(len(self.calls) - 1, len(self.results) - 1)] + + +async def _byte_stream(chunks: List[bytes]) -> AsyncIterator[bytes]: + for chunk in chunks: + yield chunk + + +async def _collect(stream: AsyncIterator[bytes]) -> List[bytes]: + return [chunk async for chunk in stream] + + +@pytest.fixture +def local_cache(): + previous_cache = litellm.cache + litellm.cache = Cache(type=LiteLLMCacheType.LOCAL) + yield litellm.cache + litellm.cache = previous_cache + + +@pytest.fixture +def request_kwargs() -> Dict[str, Any]: + return { + "model": "anthropic/claude-sonnet-4-5", + "custom_llm_provider": "anthropic", + "api_key": "fake-key", + "max_tokens": 64, + "messages": [{"role": "user", "content": "which greek letter?"}], + } + + +@pytest.mark.asyncio +async def test_non_streaming_request_is_served_from_cache(local_cache, request_kwargs, monkeypatch): + fake_handler = _CountingHandler([_anthropic_response("msg_1", "ALPHA"), _anthropic_response("msg_2", "BETA")]) + monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) + + first = await litellm.anthropic_messages(**request_kwargs) + await asyncio.sleep(0) + second = await litellm.anthropic_messages(**request_kwargs) + + assert len(fake_handler.calls) == 1 + assert first == second + assert second["content"][0]["text"] == "ALPHA" + + +@pytest.mark.asyncio +async def test_cache_key_separates_different_system_prompts(local_cache, request_kwargs, monkeypatch): + """`system` has no OpenAI equivalent; if it is dropped from the cache key the + second request is answered with the first system prompt's response.""" + fake_handler = _CountingHandler([_anthropic_response("msg_1", "ALPHA"), _anthropic_response("msg_2", "BETA")]) + monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) + + first = await litellm.anthropic_messages(**request_kwargs, system="Always answer ALPHA") + await asyncio.sleep(0) + second = await litellm.anthropic_messages(**request_kwargs, system="Always answer BETA") + + assert len(fake_handler.calls) == 2 + assert first["content"][0]["text"] == "ALPHA" + assert second["content"][0]["text"] == "BETA" + + +@pytest.mark.parametrize("anthropic_param", [{"top_k": 5}, {"stop_sequences": ["STOP"]}]) +@pytest.mark.asyncio +async def test_cache_key_separates_anthropic_native_params(local_cache, request_kwargs, monkeypatch, anthropic_param): + fake_handler = _CountingHandler([_anthropic_response("msg_1", "ALPHA"), _anthropic_response("msg_2", "BETA")]) + monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) + + await litellm.anthropic_messages(**request_kwargs) + await asyncio.sleep(0) + await litellm.anthropic_messages(**request_kwargs, **anthropic_param) + + assert len(fake_handler.calls) == 2 + + +@pytest.mark.asyncio +async def test_streaming_request_is_replayed_from_cache(local_cache, request_kwargs, monkeypatch): + fake_handler = _CountingHandler([_byte_stream(STREAM_EVENTS), _byte_stream([b"event: never_used\n\n"])]) + monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) + + first = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + second_stream = await litellm.anthropic_messages(**request_kwargs, stream=True) + second = await _collect(second_stream) + + assert len(fake_handler.calls) == 1 + assert first == STREAM_EVENTS + assert second == STREAM_EVENTS + assert second_stream._hidden_params["cache_hit"] is True + + +@pytest.mark.asyncio +async def test_streaming_cache_is_not_shared_with_non_streaming(local_cache, request_kwargs, monkeypatch): + fake_handler = _CountingHandler([_byte_stream(STREAM_EVENTS), _anthropic_response("msg_2", "ALPHA")]) + monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) + + await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + non_streaming = await litellm.anthropic_messages(**request_kwargs) + + assert len(fake_handler.calls) == 2 + assert non_streaming["content"][0]["text"] == "ALPHA" + + +@pytest.mark.asyncio +async def test_failed_stream_is_not_cached(local_cache, request_kwargs, monkeypatch): + error_events = STREAM_EVENTS[:3] + [ + b'event: error\ndata: {"type": "error", "error": {"type": "overloaded_error", "message": "overloaded"}}\n\n' + ] + fake_handler = _CountingHandler([_byte_stream(error_events), _byte_stream(STREAM_EVENTS)]) + monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) + + failed = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + replayed = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + + assert failed == error_events + assert len(fake_handler.calls) == 2 + assert replayed == STREAM_EVENTS + + +@pytest.mark.asyncio +async def test_abandoned_stream_is_not_cached(local_cache, request_kwargs, monkeypatch): + fake_handler = _CountingHandler([_byte_stream(STREAM_EVENTS), _byte_stream(STREAM_EVENTS)]) + monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) + + partial_stream = await litellm.anthropic_messages(**request_kwargs, stream=True) + await partial_stream.__anext__() + await partial_stream.aclose() + + replayed = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + + assert len(fake_handler.calls) == 2 + assert replayed == STREAM_EVENTS From d2a5de2e04d10042060c1e68c157857bdacfb165 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 25 Jul 2026 00:15:55 +0000 Subject: [PATCH 004/576] refactor(caching): tighten anthropic messages cache types and drop comments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/caching_handler.py | 11 ++------ .../litellm_core_utils/model_param_helper.py | 7 +---- .../messages/response_cache.py | 28 ------------------- 3 files changed, 3 insertions(+), 43 deletions(-) diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 8b2d033f24a..70a2e3fd1b2 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -18,6 +18,7 @@ import asyncio import datetime import inspect import time +from collections.abc import AsyncIterator from typing import ( TYPE_CHECKING, Any, @@ -1058,14 +1059,6 @@ class LLMCachingHandler: ) def wrap_streaming_result_for_cache(self, result: Any, call_type: str) -> Any: - """ - Tee a streaming result so it still reaches the cache. - - Streaming responses are returned to the caller before ``async_set_cache`` - runs. Chat/text completion streams are teed inside ``CustomStreamWrapper`` - and Responses API streams inside their own iterator; Anthropic Messages - streams have no such hook, so they are wrapped here. - """ if call_type not in ( CallTypes.anthropic_messages.value, CallTypes.aanthropic_messages.value, @@ -1075,7 +1068,7 @@ class LLMCachingHandler: original_function=self.original_function, kwargs=self.request_kwargs ): return result - if not hasattr(result, "__anext__"): + if not isinstance(result, AsyncIterator): return result from litellm.llms.anthropic.experimental_pass_through.messages.response_cache import ( AnthropicMessagesStreamCacheWriter, diff --git a/litellm/litellm_core_utils/model_param_helper.py b/litellm/litellm_core_utils/model_param_helper.py index cf4eba933b8..7e99e5fc5b2 100644 --- a/litellm/litellm_core_utils/model_param_helper.py +++ b/litellm/litellm_core_utils/model_param_helper.py @@ -174,13 +174,8 @@ class ModelParamHelper: def _get_litellm_supported_anthropic_messages_kwargs() -> set[str]: """ Get the litellm supported Anthropic /v1/messages kwargs - - This follows the Anthropic Messages API spec. `system`, `top_k` and - `stop_sequences` have no OpenAI equivalent, so without them the cache key - for a /v1/messages request ignores them and collides across requests that - differ only by system prompt. """ - return set(getattr(AnthropicMessagesRequest, "__annotations__", {}).keys()) + return set(AnthropicMessagesRequest.__annotations__.keys()) @staticmethod def _get_exclude_kwargs() -> Set[str]: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py index e94d8f6bbaf..4873b1fdb96 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py @@ -1,13 +1,3 @@ -""" -Response caching for Anthropic Messages (`/v1/messages`) requests. - -Non-streaming responses are plain dicts and are stored by the generic caching -handler. Streaming responses are returned to the caller before -``LLMCachingHandler.async_set_cache`` runs, so they are teed here instead: the -SSE events are buffered while they are forwarded and persisted verbatim once the -stream completes, and a hit replays exactly what the provider sent. -""" - from collections.abc import AsyncIterator from typing import TYPE_CHECKING, Any, cast @@ -38,14 +28,6 @@ def _decode(chunk: bytes | str) -> str: class AnthropicMessagesStreamCacheWriter: - """ - Forwards a `/v1/messages` SSE stream unchanged while buffering it, then - writes the collected events to the response cache on normal completion. - - Only a stream that ran to a ``message_stop`` without a provider ``error`` - event is written, so partial or failed responses cannot be replayed. - """ - def __init__( self, stream: AsyncIterator[bytes | str], @@ -105,11 +87,6 @@ class AnthropicMessagesStreamCacheWriter: class CachedAnthropicMessagesStreamIterator(BaseAnthropicMessagesStreamingIterator): - """ - Replays cached `/v1/messages` SSE events and logs the request as a cache hit - once the replay finishes, mirroring what the live stream logs at end of stream. - """ - def __init__( self, events: list[str], @@ -146,11 +123,6 @@ def convert_cached_anthropic_messages_result( logging_obj: LiteLLMLoggingObj, kwargs: dict[str, Any], ) -> AnthropicMessagesResponse | CachedAnthropicMessagesStreamIterator: - """ - Turn a cached `/v1/messages` entry back into what the caller expects: an - SSE replay iterator for a streamed entry, otherwise the response itself - (``AnthropicMessagesResponse`` is a TypedDict, i.e. a dict at runtime). - """ events = get_cached_stream_events(cached_result) if events is not None: return CachedAnthropicMessagesStreamIterator( From b6cf4066f4e907c03f11065f52f4da149e649128 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 25 Jul 2026 01:17:43 +0000 Subject: [PATCH 005/576] fix(caching): log cached anthropic stream replay only once Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../messages/response_cache.py | 5 ++- .../messages/test_response_cache.py | 32 +++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py index 4873b1fdb96..d5b68c99130 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py @@ -96,6 +96,7 @@ class CachedAnthropicMessagesStreamIterator(BaseAnthropicMessagesStreamingIterat super().__init__(litellm_logging_obj=litellm_logging_obj, request_body=request_body) self.chunks: list[bytes] = [event.encode("utf-8") for event in events] self.current_index = 0 + self.logged = False self._hidden_params: dict[str, Any] = {"cache_hit": True} litellm_logging_obj.model_call_details["cache_hit"] = True @@ -104,7 +105,9 @@ class CachedAnthropicMessagesStreamIterator(BaseAnthropicMessagesStreamingIterat async def __anext__(self) -> bytes: if self.current_index >= len(self.chunks): - await self._handle_streaming_logging(self.chunks) + if not self.logged: + self.logged = True + await self._handle_streaming_logging(self.chunks) raise StopAsyncIteration chunk = self.chunks[self.current_index] self.current_index += 1 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py index 344152cd828..071580347a6 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py @@ -177,3 +177,35 @@ async def test_abandoned_stream_is_not_cached(local_cache, request_kwargs, monke assert len(fake_handler.calls) == 2 assert replayed == STREAM_EVENTS + +@pytest.mark.asyncio +async def test_cached_stream_replay_logs_once_when_polled_after_exhaustion(): + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.llms.anthropic.experimental_pass_through.messages.response_cache import ( + CachedAnthropicMessagesStreamIterator, + ) + from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, + ) + + logging_obj = MagicMock() + logging_obj.model_call_details = {} + iterator = CachedAnthropicMessagesStreamIterator( + events=[event.decode("utf-8") for event in STREAM_EVENTS], + litellm_logging_obj=logging_obj, + request_body={"model": "claude-sonnet-4-5"}, + ) + + with patch.object( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + new=AsyncMock(), + ) as mock_route: + assert await _collect(iterator) == STREAM_EVENTS + for _ in range(2): + with pytest.raises(StopAsyncIteration): + await iterator.__anext__() + await asyncio.sleep(0) + + mock_route.assert_called_once() From e4a047526334dc97a93ef364878b14760e85408d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 25 Jul 2026 10:57:19 -0700 Subject: [PATCH 006/576] test: remove four mirror test files that exercise none of their module A second mutation batch scored the previously unmapped mirror files on current staging. These four generate mutants for the module they are named after, yet no test in the file executes any of them; their test-context coverage lands on generic shared machinery or, for the guardrail translation handler remainder, on no litellm line at all. Eight sibling findings that do exercise a different real module are kept for retargeting instead of removal. --- .../datadog/test_datadog_llm_observability.py | 1195 ----------------- .../guardrail_translation/test_handler.py | 37 - .../test_reasoning_content_transformation.py | 296 ---- tests/test_litellm/test_azure_video_router.py | 53 - 4 files changed, 1581 deletions(-) delete mode 100644 tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py delete mode 100644 tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py delete mode 100644 tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py delete mode 100644 tests/test_litellm/test_azure_video_router.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py deleted file mode 100644 index 1cc3591392b..00000000000 --- a/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py +++ /dev/null @@ -1,1195 +0,0 @@ -import asyncio -import os -import sys -from datetime import datetime, timedelta, timezone -from typing import Optional -from unittest.mock import MagicMock, Mock, patch - -import pytest - -# Adds the grandparent directory to sys.path to allow importing project modules -sys.path.insert(0, os.path.abspath("../..")) -import litellm -from litellm.integrations.custom_logger import CustomLogger -from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger -from litellm.types.integrations.datadog_llm_obs import ( - DatadogLLMObsInitParams, -) -from litellm.types.utils import ( - StandardLoggingGuardrailInformation, - StandardLoggingHiddenParams, - StandardLoggingMetadata, - StandardLoggingModelInformation, - StandardLoggingPayload, - StandardLoggingPayloadErrorInformation, -) - - -def create_standard_logging_payload_with_cache() -> StandardLoggingPayload: - """Create a real StandardLoggingPayload object for testing""" - return StandardLoggingPayload( - id="test-request-id-456", - call_type="completion", - response_cost=0.05, - response_cost_failure_debug_info=None, - status="success", - total_tokens=30, - prompt_tokens=10, - completion_tokens=20, - startTime=1234567890.0, - endTime=1234567891.0, - completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-4", model_map_value=None - ), - model="gpt-4", - model_id="model-123", - model_group="openai-gpt", - api_base="https://api.openai.com", - metadata=StandardLoggingMetadata( - user_api_key_hash="test_hash", - user_api_key_org_id=None, - user_api_key_alias="test_alias", - user_api_key_team_id="test_team", - user_api_key_user_id="test_user", - user_api_key_team_alias="test_team_alias", - spend_logs_metadata=None, - requester_ip_address="127.0.0.1", - requester_metadata=None, - ), - cache_hit=True, - cache_key="test-cache-key-789", - saved_cache_cost=0.02, - request_tags=[], - end_user=None, - requester_ip_address="127.0.0.1", - messages=[{"role": "user", "content": "Hello, world!"}], - response={"choices": [{"message": {"content": "Hi there!"}}]}, - error_str=None, - model_parameters={"stream": True}, - hidden_params=StandardLoggingHiddenParams( - model_id="model-123", - cache_key="test-cache-key-789", - api_base="https://api.openai.com", - response_cost="0.05", - additional_headers=None, - ), - trace_id="test-trace-id-123", - custom_llm_provider="openai", - ) - - -def create_standard_logging_payload_with_failure() -> StandardLoggingPayload: - """Create a StandardLoggingPayload object for failure testing""" - return StandardLoggingPayload( - id="test-request-id-failure-789", - call_type="completion", - response_cost=0.0, - response_cost_failure_debug_info=None, - status="failure", - total_tokens=0, - prompt_tokens=10, - completion_tokens=0, - startTime=1234567890.0, - endTime=1234567891.0, - completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-4", model_map_value=None - ), - model="gpt-4", - model_id="model-123", - model_group="openai-gpt", - api_base="https://api.openai.com", - metadata=StandardLoggingMetadata( - user_api_key_hash="test_hash", - user_api_key_org_id=None, - user_api_key_alias="test_alias", - user_api_key_team_id="test_team", - user_api_key_user_id="test_user", - user_api_key_team_alias="test_team_alias", - spend_logs_metadata=None, - requester_ip_address="127.0.0.1", - requester_metadata=None, - ), - cache_hit=False, - cache_key=None, - saved_cache_cost=0.0, - request_tags=[], - end_user=None, - requester_ip_address="127.0.0.1", - messages=[{"role": "user", "content": "Hello, world!"}], - response=None, - error_str="RateLimitError: You exceeded your current quota", - error_information=StandardLoggingPayloadErrorInformation( - error_code="rate_limit_exceeded", - error_class="RateLimitError", - llm_provider="openai", - traceback="Traceback (most recent call last):\n File test.py, line 1\n RateLimitError: You exceeded your current quota", - error_message="RateLimitError: You exceeded your current quota", - ), - model_parameters={"stream": False}, - hidden_params=StandardLoggingHiddenParams( - model_id="model-123", - cache_key=None, - api_base="https://api.openai.com", - response_cost="0.0", - additional_headers=None, - ), - trace_id="test-trace-id-failure-456", - custom_llm_provider="openai", - ) - - -class TestDataDogLLMObsLogger: - """Test suite for DataDog LLM Observability Logger""" - - @pytest.fixture - def mock_env_vars(self): - """Mock environment variables for DataDog""" - with patch.dict( - os.environ, {"DD_API_KEY": "test_api_key", "DD_SITE": "us5.datadoghq.com"} - ): - yield - - @pytest.fixture - def mock_response_obj(self): - """Create a mock response object""" - mock_response = Mock() - mock_response.__getitem__ = Mock( - return_value={ - "choices": [ - { - "message": Mock( - json=Mock( - return_value={"role": "assistant", "content": "Hello!"} - ) - ) - } - ] - } - ) - return mock_response - - def test_cost_and_trace_id_integration(self, mock_env_vars, mock_response_obj): - """Test that total_cost is passed and trace_id from standard payload is used""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - standard_payload = create_standard_logging_payload_with_cache() - - kwargs = { - "standard_logging_object": standard_payload, - "litellm_params": { - "metadata": {"trace_id": "old-trace-id-should-be-ignored"} - }, - } - - start_time = datetime.now() - end_time = datetime.now() - - payload = logger.create_llm_obs_payload(kwargs, start_time, end_time) - - # Test 1: Verify total_cost is correctly extracted from response_cost - assert payload["metrics"].get("total_cost") == 0.05 - - # Test 2: Verify trace_id comes from standard_logging_payload, not metadata - assert payload["trace_id"] == "test-trace-id-123" - - # Test 3: Verify saved_cache_cost is in metadata - metadata = payload["meta"]["metadata"] - assert metadata["saved_cache_cost"] == 0.02 - assert metadata["cache_hit"] is True - assert metadata["cache_key"] == "test-cache-key-789" - - # Test 4: Verify is_streamed_request is in metadata - assert metadata["is_streamed_request"] is True - - def test_cache_metadata_fields(self, mock_env_vars, mock_response_obj): - """Test that cache-related metadata fields are correctly tracked""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - standard_payload = create_standard_logging_payload_with_cache() - - # Test the _get_dd_llm_obs_payload_metadata method directly - metadata = logger._get_dd_llm_obs_payload_metadata(standard_payload) - - # Verify all cache-related fields are present - assert metadata["cache_hit"] is True - assert metadata["cache_key"] == "test-cache-key-789" - assert metadata["saved_cache_cost"] == 0.02 - assert metadata["id"] == "test-request-id-456" - assert metadata["trace_id"] == "test-trace-id-123" - assert metadata["model_name"] == "gpt-4" - assert metadata["model_provider"] == "openai" - - def test_get_time_to_first_token_seconds(self, mock_env_vars): - """Test the _get_time_to_first_token_seconds method for streaming calls""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - # Test streaming case (completion_start_time available) - streaming_payload = create_standard_logging_payload_with_cache() - # Modify times for testing: start=1000, completion_start=1002, end=1005 - streaming_payload["startTime"] = 1000.0 - streaming_payload["completionStartTime"] = 1002.0 - streaming_payload["endTime"] = 1005.0 - - # Test streaming case: should use completion_start_time - start_time - time_to_first_token = logger._get_time_to_first_token_seconds( - streaming_payload - ) - assert time_to_first_token == 2.0 # 1002.0 - 1000.0 = 2.0 seconds - - def test_datadog_span_kind_mapping(self, mock_env_vars): - """Test that call_type values are correctly mapped to DataDog span kinds""" - from litellm.types.utils import CallTypes - - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - # Test embedding operations - assert ( - logger._get_datadog_span_kind(CallTypes.embedding.value, "123") - == "embedding" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.aembedding.value, "123") - == "embedding" - ) - - # Test LLM completion operations - assert logger._get_datadog_span_kind(CallTypes.completion.value, None) == "llm" - assert logger._get_datadog_span_kind(CallTypes.acompletion.value, None) == "llm" - assert ( - logger._get_datadog_span_kind(CallTypes.text_completion.value, None) - == "llm" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.generate_content.value, None) - == "llm" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.anthropic_messages.value, None) - == "llm" - ) - assert logger._get_datadog_span_kind(CallTypes.responses.value, None) == "llm" - assert logger._get_datadog_span_kind(CallTypes.aresponses.value, None) == "llm" - - # Test tool operations - assert ( - logger._get_datadog_span_kind(CallTypes.call_mcp_tool.value, "123") - == "tool" - ) - - # Test retrieval operations - assert ( - logger._get_datadog_span_kind(CallTypes.get_assistants.value, "123") - == "retrieval" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.file_retrieve.value, "123") - == "retrieval" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.retrieve_batch.value, "123") - == "retrieval" - ) - - # Test task operations - assert ( - logger._get_datadog_span_kind(CallTypes.create_batch.value, "123") == "task" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.image_generation.value, "123") - == "task" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.moderation.value, "123") == "task" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.transcription.value, "123") - == "task" - ) - - # Test default fallback - assert logger._get_datadog_span_kind("unknown_call_type", None) == "llm" - assert logger._get_datadog_span_kind(None, None) == "llm" - - def test_datadog_span_kind_defaults_without_parent(self, mock_env_vars): - """Test that non-llm kinds fallback to llm when no parent span is provided""" - from litellm.types.utils import CallTypes - - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - # Tool/task/retrieval span kinds should fallback to llm when parent_id missing - assert ( - logger._get_datadog_span_kind(CallTypes.call_mcp_tool.value, None) == "llm" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.create_batch.value, None) == "llm" - ) - assert ( - logger._get_datadog_span_kind(CallTypes.get_assistants.value, None) == "llm" - ) - - @pytest.mark.asyncio - async def test_async_log_failure_event(self, mock_env_vars): - """Test that async_log_failure_event correctly processes failure payloads according to DD LLM Obs API spec""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - # Ensure log_queue starts empty - logger.log_queue = [] - - standard_failure_payload = create_standard_logging_payload_with_failure() - - kwargs = { - "standard_logging_object": standard_failure_payload, - "model": "gpt-4", - "litellm_params": {"metadata": {}}, - } - - start_time = datetime.now() - end_time = datetime.now() + timedelta(seconds=2) - - # Mock async_send_batch to prevent actual network calls - with patch.object(logger, "async_send_batch") as mock_send_batch: - # Call the method under test - await logger.async_log_failure_event(kwargs, None, start_time, end_time) - - # Verify payload was added to queue - assert len(logger.log_queue) == 1 - - # Verify the payload has correct failure characteristics according to DD LLM Obs API spec - payload = logger.log_queue[0] - assert payload["trace_id"] == "test-trace-id-failure-456" - assert ( - payload["meta"]["metadata"]["id"] == "test-request-id-failure-789" - ) - assert payload["status"] == "error" - - # Verify error information follows DD LLM Obs API spec - assert ( - payload["meta"]["error"]["message"] - == "RateLimitError: You exceeded your current quota" - ) - assert payload["meta"]["error"]["type"] == "RateLimitError" - assert ( - payload["meta"]["error"]["stack"] - == "Traceback (most recent call last):\n File test.py, line 1\n RateLimitError: You exceeded your current quota" - ) - - assert payload["metrics"]["total_cost"] == 0.0 - assert payload["metrics"]["total_tokens"] == 0 - assert payload["metrics"]["output_tokens"] == 0 - - # Verify batch sending not triggered (queue size < batch_size) - mock_send_batch.assert_not_called() - - -class TestDataDogLLMObsLoggerForRedaction(DataDogLLMObsLogger): - """Test suite for DataDog LLM Observability Logger""" - - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - self.logged_standard_logging_payload = kwargs.get("standard_logging_object") - - -class TestS3Logger(CustomLogger): - """Test suite for S3 Logger""" - - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - self.logged_standard_logging_payload = kwargs.get("standard_logging_object") - - -@pytest.mark.asyncio -async def test_dd_llms_obs_redaction(mock_env_vars): - # init DD with turn_off_message_logging=True - litellm._turn_on_debug() - from litellm.types.utils import LiteLLMCommonStrings - - litellm.datadog_llm_observability_params = DatadogLLMObsInitParams( - turn_off_message_logging=True - ) - dd_llms_obs_logger = TestDataDogLLMObsLoggerForRedaction() - test_s3_logger = TestS3Logger() - litellm.callbacks = [dd_llms_obs_logger, test_s3_logger] - - # call litellm - await litellm.acompletion( - model="gpt-4o", - mock_response="Hi there!", - messages=[{"role": "user", "content": "Hello, world!"}], - ) - - # sleep 1 second for logging to complete - await asyncio.sleep(1) - - ################# - # test validation - # 1. both loggers logged a standard_logging_payload - # 2. DD LLM Obs standard_logging_payload has messages and response redacted - # 3. S3 standard_logging_payload does not have messages and response redacted - - assert dd_llms_obs_logger.logged_standard_logging_payload is not None - assert test_s3_logger.logged_standard_logging_payload is not None - - assert ( - dd_llms_obs_logger.logged_standard_logging_payload["messages"][0]["content"] - == "redacted-by-litellm" - ) - assert ( - dd_llms_obs_logger.logged_standard_logging_payload["response"]["choices"][0][ - "message" - ]["content"] - == "redacted-by-litellm" - ) - - assert test_s3_logger.logged_standard_logging_payload["messages"] == [ - {"role": "user", "content": "Hello, world!"} - ] - assert ( - test_s3_logger.logged_standard_logging_payload["response"]["choices"][0][ - "message" - ]["content"] - == "Hi there!" - ) - - -@pytest.fixture -def mock_env_vars(): - """Mock environment variables for DataDog""" - with patch.dict( - os.environ, {"DD_API_KEY": "test_api_key", "DD_SITE": "us5.datadoghq.com"} - ): - yield - - -@pytest.mark.asyncio -async def test_create_llm_obs_payload(mock_env_vars): - datadog_llm_obs_logger = DataDogLLMObsLogger() - standard_logging_payload = create_standard_logging_payload_with_cache() - payload = datadog_llm_obs_logger.create_llm_obs_payload( - kwargs={ - "model": "gpt-4", - "messages": [{"role": "user", "content": "Hello"}], - "standard_logging_object": standard_logging_payload, - }, - start_time=datetime.now(), - end_time=datetime.now() + timedelta(seconds=1), - ) - - assert payload["name"] == "litellm_llm_call" - assert payload["meta"]["kind"] == "llm" - assert payload["meta"]["input"]["messages"] == [ - {"role": "user", "content": "Hello, world!"} - ] - assert payload["meta"]["output"]["messages"][0]["content"] == "Hi there!" - assert payload["metrics"]["input_tokens"] == 10 - assert payload["metrics"]["output_tokens"] == 20 - assert payload["metrics"]["total_tokens"] == 30 - - -def create_standard_logging_payload_with_latency_metrics() -> StandardLoggingPayload: - """Create a StandardLoggingPayload object with latency metrics for testing""" - guardrail_info = StandardLoggingGuardrailInformation( - guardrail_name="test_guardrail", - guardrail_status="success", - start_time=1234567890.0, - end_time=1234567890.5, - duration=0.5, # 500ms - guardrail_request={"input": "test input message", "user_id": "test_user"}, - guardrail_response={ - "output": "filtered output", - "flagged": False, - "score": 0.1, - }, - ) - - hidden_params = StandardLoggingHiddenParams( - model_id="model-123", - cache_key="test-cache-key", - api_base="https://api.openai.com", - response_cost="0.05", - litellm_overhead_time_ms=150.0, # 150ms - additional_headers=None, - ) - - return StandardLoggingPayload( - id="test-request-id-latency", - call_type="completion", - response_cost=0.05, - response_cost_failure_debug_info=None, - status="success", - total_tokens=30, - prompt_tokens=10, - completion_tokens=20, - startTime=1234567890.0, - endTime=1234567892.0, - completionStartTime=1234567890.8, # 800ms after start - response_time=2.0, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-4", model_map_value=None - ), - model="gpt-4", - model_id="model-123", - model_group="openai-gpt", - api_base="https://api.openai.com", - metadata=StandardLoggingMetadata( - user_api_key_hash="test_hash", - user_api_key_org_id=None, - user_api_key_alias="test_alias", - user_api_key_team_id="test_team", - user_api_key_user_id="test_user", - user_api_key_team_alias="test_team_alias", - spend_logs_metadata=None, - requester_ip_address="127.0.0.1", - requester_metadata=None, - ), - cache_hit=False, - cache_key=None, - saved_cache_cost=0.0, - request_tags=[], - end_user=None, - requester_ip_address="127.0.0.1", - messages=[{"role": "user", "content": "Hello, world!"}], - response={"choices": [{"message": {"content": "Hi there!"}}]}, - error_str=None, - error_information=None, - model_parameters={"stream": True}, - hidden_params=hidden_params, - guardrail_information=[guardrail_info], - trace_id="test-trace-id-latency", - custom_llm_provider="openai", - ) - - -def test_latency_metrics_in_metadata(mock_env_vars): - """Test that time to first token, litellm overhead, and guardrail overhead are included in metadata""" - with ( - patch("litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client"), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - standard_payload = create_standard_logging_payload_with_latency_metrics() - - kwargs = { - "standard_logging_object": standard_payload, - "litellm_params": {"metadata": {}}, - } - - start_time = datetime.now() - end_time = datetime.now() - - # Test the metadata generation directly - metadata = logger._get_dd_llm_obs_payload_metadata(standard_payload) - latency_metadata = metadata.get("latency_metrics", {}) - - # Verify time to first token is included (800ms) - assert "time_to_first_token_ms" in latency_metadata - assert ( - abs(latency_metadata["time_to_first_token_ms"] - 800.0) < 0.001 - ) # 0.8 seconds * 1000 with tolerance for floating-point precision - - # Verify litellm overhead is included (150ms) - assert "litellm_overhead_time_ms" in latency_metadata - assert latency_metadata["litellm_overhead_time_ms"] == 150.0 - - # Verify guardrail overhead is included (500ms) - assert "guardrail_overhead_time_ms" in latency_metadata - assert ( - latency_metadata["guardrail_overhead_time_ms"] == 500.0 - ) # 0.5 seconds * 1000 - - # Verify these metrics are also included in the full payload - payload = logger.create_llm_obs_payload(kwargs, start_time, end_time) - payload_metadata_latency = payload["meta"]["metadata"]["latency_metrics"] - - assert abs(payload_metadata_latency["time_to_first_token_ms"] - 800.0) < 0.001 - assert payload_metadata_latency["litellm_overhead_time_ms"] == 150.0 - assert payload_metadata_latency["guardrail_overhead_time_ms"] == 500.0 - - -def test_latency_metrics_edge_cases(mock_env_vars): - """Test latency metrics with edge cases (missing fields, zero values, etc.)""" - with ( - patch("litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client"), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - # Test case 1: No latency metrics present - standard_payload = create_standard_logging_payload_with_cache() - metadata = logger._get_dd_llm_obs_payload_metadata(standard_payload) - - # Should not have latency fields if data is missing/zero - assert "time_to_first_token_ms" not in metadata # Will be 0, so not included - assert ( - "litellm_overhead_time_ms" not in metadata - ) # Not present in hidden_params - assert "guardrail_overhead_time_ms" not in metadata # No guardrail_information - - # Test case 2: Zero time to first token should not be included - standard_payload = create_standard_logging_payload_with_cache() - standard_payload["startTime"] = 1000.0 - standard_payload["completionStartTime"] = 1000.0 # Same time = 0 difference - metadata = logger._get_dd_llm_obs_payload_metadata(standard_payload) - assert "time_to_first_token_ms" not in metadata - - # Test case 3: Missing guardrail duration should not crash - standard_payload = create_standard_logging_payload_with_cache() - standard_payload["guardrail_information"] = [ - StandardLoggingGuardrailInformation( - guardrail_name="test", - guardrail_status="success", - # duration is missing - ) - ] - metadata = logger._get_dd_llm_obs_payload_metadata(standard_payload) - assert "guardrail_overhead_time_ms" not in metadata - - -def test_guardrail_information_in_metadata(mock_env_vars): - """Test that guardrail_information is included in metadata with input/output fields""" - with ( - patch("litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client"), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - # Create a standard payload with guardrail information - standard_payload = create_standard_logging_payload_with_latency_metrics() - - kwargs = { - "standard_logging_object": standard_payload, - "litellm_params": {"metadata": {}}, - } - - start_time = datetime.now() - end_time = datetime.now() - - # Create the payload and verify guardrail_information is in metadata - payload = logger.create_llm_obs_payload(kwargs, start_time, end_time) - metadata = payload["meta"]["metadata"] - - # Verify guardrail_information is present in metadata - assert "guardrail_information" in metadata - assert metadata["guardrail_information"] is not None - - # Verify the guardrail information structure - guardrail_info = metadata["guardrail_information"] - assert guardrail_info[0]["guardrail_name"] == "test_guardrail" - assert guardrail_info[0]["guardrail_status"] == "success" - assert guardrail_info[0]["duration"] == 0.5 - - # Verify input/output fields are present - assert "guardrail_request" in guardrail_info[0] - assert "guardrail_response" in guardrail_info[0] - - # Validate the input/output content - assert guardrail_info[0]["guardrail_request"]["input"] == "test input message" - assert guardrail_info[0]["guardrail_request"]["user_id"] == "test_user" - assert guardrail_info[0]["guardrail_response"]["output"] == "filtered output" - assert guardrail_info[0]["guardrail_response"]["flagged"] is False - assert guardrail_info[0]["guardrail_response"]["score"] == 0.1 - - -def create_standard_logging_payload_with_tool_calls() -> StandardLoggingPayload: - """Create a StandardLoggingPayload object with tool calls for testing""" - return { - "id": "test-request-id-tool-calls", - "trace_id": "test-trace-id-tool-calls", - "call_type": "completion", - "stream": None, - "response_cost": 0.05, - "response_cost_failure_debug_info": None, - "status": "success", - "custom_llm_provider": "openai", - "total_tokens": 50, - "prompt_tokens": 20, - "completion_tokens": 30, - "startTime": 1234567890.0, - "endTime": 1234567891.0, - "completionStartTime": 1234567890.5, - "response_time": 1.0, - "model_map_information": {"model_map_key": "gpt-4", "model_map_value": None}, - "model": "gpt-4", - "model_id": "model-123", - "model_group": "openai-gpt", - "api_base": "https://api.openai.com", - "metadata": { - "user_api_key_hash": "test_hash", - "user_api_key_org_id": None, - "user_api_key_alias": "test_alias", - "user_api_key_team_id": "test_team", - "user_api_key_user_id": "test_user", - "user_api_key_team_alias": "test_team_alias", - "user_api_key_user_email": None, - "user_api_key_end_user_id": None, - "user_api_key_request_route": None, - "spend_logs_metadata": None, - "requester_ip_address": "127.0.0.1", - "requester_metadata": None, - "requester_custom_headers": None, - "prompt_management_metadata": None, - "mcp_tool_call_metadata": None, - "vector_store_request_metadata": None, - "applied_guardrails": None, - "usage_object": None, - "cold_storage_object_key": None, - }, - "cache_hit": False, - "cache_key": None, - "saved_cache_cost": 0.0, - "request_tags": [], - "end_user": None, - "requester_ip_address": "127.0.0.1", - "messages": [ - {"role": "user", "content": "What's the weather?"}, - { - "role": "assistant", - "content": "I'll check the weather for you.", - "tool_calls": [ - { - "id": "call_123", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"location": "NYC"}', - }, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_123", - "content": '{"temperature": 72, "condition": "sunny"}', - }, - ], - "response": { - "choices": [ - { - "message": { - "role": "assistant", - "content": "It's 72°F and sunny in NYC!", - "tool_calls": [ - { - "id": "call_456", - "type": "function", - "function": { - "name": "format_response", - "arguments": '{"temp": 72, "condition": "sunny"}', - }, - } - ], - } - } - ] - }, - "error_str": None, - "error_information": None, - "model_parameters": {"temperature": 0.7}, - "hidden_params": { - "model_id": "model-123", - "cache_key": None, - "api_base": "https://api.openai.com", - "response_cost": "0.05", - "litellm_overhead_time_ms": None, - "additional_headers": None, - "batch_models": None, - "litellm_model_name": None, - "usage_object": None, - }, - "guardrail_information": None, - "standard_built_in_tools_params": None, - } # type: ignore - - -class TestDataDogLLMObsLoggerToolCalls: - """Simple test suite for DataDog LLM Observability Logger tool call handling""" - - @pytest.fixture - def mock_env_vars(self): - """Mock environment variables for DataDog""" - with patch.dict( - os.environ, {"DD_API_KEY": "test_api_key", "DD_SITE": "us5.datadoghq.com"} - ): - yield - - def test_tool_call_span_kind_mapping(self, mock_env_vars): - """Test that tool call operations are correctly mapped to 'tool' span kind""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - # Test MCP tool call mapping - from litellm.types.utils import CallTypes - - assert ( - logger._get_datadog_span_kind(CallTypes.call_mcp_tool.value, "123") - == "tool" - ) - - def test_tool_call_payload_creation(self, mock_env_vars): - """Test that tool call payloads are created correctly""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - standard_payload = create_standard_logging_payload_with_tool_calls() - - kwargs = { - "standard_logging_object": standard_payload, - "litellm_params": {"metadata": {}}, - } - - start_time = datetime.now() - end_time = datetime.now() - - payload = logger.create_llm_obs_payload(kwargs, start_time, end_time) - - # Verify basic payload structure - assert payload.get("name") == "litellm_llm_call" - assert payload.get("status") == "ok" - assert ( - payload.get("meta", {}).get("kind") == "llm" - ) # Regular completion, not tool call - - # Verify metrics - metrics = payload.get("metrics", {}) - assert metrics.get("input_tokens") == 20 - assert metrics.get("output_tokens") == 30 - assert metrics.get("total_tokens") == 50 - - def test_tool_call_messages_preserved(self, mock_env_vars): - """Test that tool call messages are preserved in the payload""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - standard_payload = create_standard_logging_payload_with_tool_calls() - - kwargs = { - "standard_logging_object": standard_payload, - "litellm_params": {"metadata": {}}, - } - - start_time = datetime.now() - end_time = datetime.now() - - payload = logger.create_llm_obs_payload(kwargs, start_time, end_time) - - # Verify input messages include tool calls - meta = payload.get("meta", {}) - input_meta = meta.get("input", {}) - input_messages = input_meta.get("messages", []) - assert len(input_messages) == 3 - - # Check assistant message has tool calls - assistant_msg = input_messages[1] - assert assistant_msg.get("role") == "assistant" - assert "tool_calls" in assistant_msg - tool_calls = assistant_msg.get("tool_calls", []) - assert len(tool_calls) == 1 - tool_call = tool_calls[0] - function_info = tool_call.get("function", {}) - assert function_info.get("name") == "get_weather" - - # Check tool message - tool_msg = input_messages[2] - assert tool_msg.get("role") == "tool" - assert tool_msg.get("tool_call_id") == "call_123" - - def test_tool_call_response_handling(self, mock_env_vars): - """Test that tool calls in response are handled correctly""" - with ( - patch( - "litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client" - ), - patch("asyncio.create_task"), - ): - logger = DataDogLLMObsLogger() - - standard_payload = create_standard_logging_payload_with_tool_calls() - - kwargs = { - "standard_logging_object": standard_payload, - "litellm_params": {"metadata": {}}, - } - - start_time = datetime.now() - end_time = datetime.now() - - payload = logger.create_llm_obs_payload(kwargs, start_time, end_time) - - # Verify output messages include tool calls - meta = payload.get("meta", {}) - output_meta = meta.get("output", {}) - output_messages = output_meta.get("messages", []) - assert len(output_messages) == 1 - - output_msg = output_messages[0] - assert output_msg.get("role") == "assistant" - assert "tool_calls" in output_msg - output_tool_calls = output_msg.get("tool_calls", []) - assert len(output_tool_calls) == 1 - output_function_info = output_tool_calls[0].get("function", {}) - assert output_function_info.get("name") == "format_response" - - -def create_standard_logging_payload_with_spend_metrics() -> StandardLoggingPayload: - """Create a StandardLoggingPayload object with spend metrics for testing""" - from datetime import datetime, timezone - - # Create a budget reset time 10 days from now (using "10d" format) - budget_reset_at = datetime.now(timezone.utc) + timedelta(days=10) - - return { - "id": "test-request-id-spend", - "trace_id": "test-trace-id-spend", - "call_type": "completion", - "stream": None, - "response_cost": 0.15, - "response_cost_failure_debug_info": None, - "status": "success", - "custom_llm_provider": "openai", - "total_tokens": 30, - "prompt_tokens": 10, - "completion_tokens": 20, - "startTime": 1234567890.0, - "endTime": 1234567891.0, - "completionStartTime": 1234567890.5, - "response_time": 1.0, - "model_map_information": {"model_map_key": "gpt-4", "model_map_value": None}, - "model": "gpt-4", - "model_id": "model-123", - "model_group": "openai-gpt", - "api_base": "https://api.openai.com", - "metadata": { - "user_api_key_hash": "test_hash", - "user_api_key_org_id": None, - "user_api_key_alias": "test_alias", - "user_api_key_team_id": "test_team", - "user_api_key_user_id": "test_user", - "user_api_key_team_alias": "test_team_alias", - "user_api_key_user_email": None, - "user_api_key_end_user_id": None, - "user_api_key_request_route": None, - "user_api_key_spend": 0.67, - "user_api_key_max_budget": 10.0, # $10 max budget - "user_api_key_budget_reset_at": budget_reset_at.isoformat(), # ISO format: 2025-09-26T... - "spend_logs_metadata": None, - "requester_ip_address": "127.0.0.1", - "requester_metadata": None, - "requester_custom_headers": None, - "prompt_management_metadata": None, - "mcp_tool_call_metadata": None, - "vector_store_request_metadata": None, - "applied_guardrails": None, - "usage_object": None, - "cold_storage_object_key": None, - }, - "cache_hit": False, - "cache_key": None, - "saved_cache_cost": 0.0, - "request_tags": [], - "end_user": None, - "requester_ip_address": "127.0.0.1", - "messages": [{"role": "user", "content": "Hello, world!"}], - "response": {"choices": [{"message": {"content": "Hi there!"}}]}, - "error_str": None, - "error_information": None, - "model_parameters": {"stream": False}, - "hidden_params": { - "model_id": "model-123", - "cache_key": None, - "api_base": "https://api.openai.com", - "response_cost": "0.15", - "litellm_overhead_time_ms": None, - "additional_headers": None, - "batch_models": None, - "litellm_model_name": None, - "usage_object": None, - }, - "guardrail_information": None, - "standard_built_in_tools_params": None, - } # type: ignore - - -@pytest.mark.asyncio -async def test_datadog_llm_obs_spend_metrics(mock_env_vars): - """Test that budget metrics are properly extracted and logged""" - datadog_llm_obs_logger = DataDogLLMObsLogger() - - # Create a standard logging payload with spend metrics - payload = create_standard_logging_payload_with_spend_metrics() - - # Show the budget reset time in ISO format - budget_reset_iso = payload["metadata"]["user_api_key_budget_reset_at"] - print(f"Budget reset time (ISO format): {budget_reset_iso}") - from datetime import datetime, timezone - - print(f"Current time: {datetime.now(timezone.utc).isoformat()}") - - # Test the _get_spend_metrics method - spend_metrics = datadog_llm_obs_logger._get_spend_metrics(payload) - - # Verify budget metrics are present - assert "user_api_key_max_budget" in spend_metrics - assert spend_metrics["user_api_key_max_budget"] == 10.0 - - assert "user_api_key_budget_reset_at" in spend_metrics - # The budget reset should be a datetime string in ISO format - budget_reset = spend_metrics["user_api_key_budget_reset_at"] - assert isinstance(budget_reset, str) - print(f"Budget reset datetime: {budget_reset}") - # Should be close to 10 days from now - budget_reset_dt = datetime.fromisoformat(budget_reset.replace("Z", "+00:00")) - now = datetime.now(timezone.utc) - time_diff = (budget_reset_dt - now).total_seconds() / 86400 # days - assert 9.5 <= time_diff <= 10.5 # Should be close to 10 days - - print(f"Spend metrics: {spend_metrics}") - - -@pytest.mark.asyncio -async def test_datadog_llm_obs_spend_metrics_no_budget(mock_env_vars): - """Test that spend metrics work when no budget is set""" - datadog_llm_obs_logger = DataDogLLMObsLogger() - - # Create a standard logging payload without budget metadata - payload = create_standard_logging_payload_with_spend_metrics() - - # Remove budget-related metadata to test no-budget scenario - payload["metadata"].pop("user_api_key_max_budget", None) - payload["metadata"].pop("user_api_key_budget_reset_at", None) - - # Test the _get_spend_metrics method - spend_metrics = datadog_llm_obs_logger._get_spend_metrics(payload) - - # Verify only response cost is present - assert "response_cost" in spend_metrics - assert spend_metrics["response_cost"] == 0.15 - - # Budget metrics should not be present - assert "user_api_key_max_budget" not in spend_metrics - assert "user_api_key_budget_reset_at" not in spend_metrics - - print(f"Spend metrics (no budget): {spend_metrics}") - - -@pytest.mark.asyncio -async def test_spend_metrics_in_datadog_payload(mock_env_vars): - """Test that spend metrics are correctly included in DataDog LLM Observability payloads""" - from datetime import datetime - - datadog_llm_obs_logger = DataDogLLMObsLogger() - - standard_payload = create_standard_logging_payload_with_spend_metrics() - - kwargs = { - "standard_logging_object": standard_payload, - "litellm_params": {"metadata": {}}, - } - - start_time = datetime.now() - end_time = datetime.now() - - payload = datadog_llm_obs_logger.create_llm_obs_payload( - kwargs, start_time, end_time - ) - - # Verify basic payload structure - assert payload.get("name") == "litellm_llm_call" - assert payload.get("status") == "ok" - - # Verify spend metrics are included in metadata - meta = payload.get("meta", {}) - assert meta is not None, "Meta section should exist in payload" - - metadata = meta.get("metadata", {}) - assert metadata is not None, "Metadata section should exist in meta" - - spend_metrics = metadata.get("spend_metrics", {}) - assert spend_metrics, "Spend metrics should exist in metadata" - - # Check that all metrics are present - assert "response_cost" in spend_metrics - assert "user_api_key_spend" in spend_metrics - assert "user_api_key_max_budget" in spend_metrics - assert "user_api_key_budget_reset_at" in spend_metrics - - # Verify the values are correct - assert spend_metrics["response_cost"] == 0.15 # response_cost - assert spend_metrics["user_api_key_spend"] == 0.67 # lol - assert spend_metrics["user_api_key_max_budget"] == 10.0 # max budget - - # Verify budget reset is a datetime string in ISO format - budget_reset = spend_metrics["user_api_key_budget_reset_at"] - assert isinstance(budget_reset, str) - print( - f"Budget reset in payload: {budget_reset}" - ) # In StandardLoggingUserAPIKeyMetadata - user_api_key_budget_reset_at: Optional[str] = None - - # In DDLLMObsSpendMetrics - user_api_key_budget_reset_at: str - # Should be close to 10 days from now - from datetime import datetime, timezone - - budget_reset_dt = datetime.fromisoformat(budget_reset.replace("Z", "+00:00")) - now = datetime.now(timezone.utc) - time_diff = (budget_reset_dt - now).total_seconds() / 86400 # days - assert 9.5 <= time_diff <= 10.5 # Should be close to 10 days diff --git a/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py b/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py deleted file mode 100644 index 1043c26c6ec..00000000000 --- a/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py +++ /dev/null @@ -1,37 +0,0 @@ -""" -Tests for the guardrail_translation_mappings registry. - -Validates: -- allm_passthrough_route is registered in the mappings (regression: this was the bug) -""" - -from litellm.llms.pass_through.guardrail_translation import ( - guardrail_translation_mappings, -) -from litellm.llms.pass_through.guardrail_translation.handler import ( - LlmPassthroughRouteHandler, -) -from litellm.types.utils import CallTypes - - -class TestRegistry: - def test_allm_passthrough_route_registered(self): - """Regression: missing this mapping was the root cause of the bug.""" - assert CallTypes.allm_passthrough_route in guardrail_translation_mappings - - def test_allm_passthrough_route_maps_to_llm_passthrough_route_handler(self): - assert ( - guardrail_translation_mappings[CallTypes.allm_passthrough_route] - is LlmPassthroughRouteHandler - ) - - def test_pass_through_still_registered(self): - from litellm.llms.pass_through.guardrail_translation.handler import ( - PassThroughEndpointHandler, - ) - - assert ( - guardrail_translation_mappings[CallTypes.pass_through] - is PassThroughEndpointHandler - ) - diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py deleted file mode 100644 index 020b5de0a2a..00000000000 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py +++ /dev/null @@ -1,296 +0,0 @@ -""" -Test reasoning content preservation in Responses API transformation -""" - -from unittest.mock import AsyncMock - -from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, -) -from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, -) -from litellm.types.utils import ( - Choices, - Delta, - Message, - ModelResponse, - ModelResponseStream, - StreamingChoices, -) - - -class TestReasoningContentStreaming: - """Test reasoning content preservation during streaming""" - - def test_reasoning_content_in_delta(self): - """Test that reasoning content is preserved in streaming deltas""" - # Setup - chunk = ModelResponseStream( - id="test-id", - created=1234567890, - model="test-model", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - content="", - role="assistant", - reasoning_content="Let me think about this problem...", - ), - ) - ], - ) - - mock_stream = AsyncMock() - - iterator = LiteLLMCompletionStreamingIterator( - model="test-model", - litellm_custom_stream_wrapper=mock_stream, - request_input="Test input", - responses_api_request={}, - ) - - # Execute - transformed_chunk = ( - iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk) - ) - - # Assert - assert transformed_chunk.delta == "Let me think about this problem..." - assert transformed_chunk.type == "response.reasoning_summary_text.delta" - - def test_mixed_content_and_reasoning(self): - """Test handling of both content and reasoning content""" - # Setup - chunk = ModelResponseStream( - id="test-id", - created=1234567890, - model="test-model", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - content="Here is the answer", - role="assistant", - reasoning_content="First, let me analyze...", - ), - ) - ], - ) - - mock_stream = AsyncMock() - iterator = LiteLLMCompletionStreamingIterator( - model="test-model", - litellm_custom_stream_wrapper=mock_stream, - request_input="Test input", - responses_api_request={}, - ) - - # Execute - transformed_chunk = ( - iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk) - ) - - # Assert - assert transformed_chunk.delta == "First, let me analyze..." - assert transformed_chunk.type == "response.reasoning_summary_text.delta" - - def test_no_reasoning_content(self): - """Test handling when no reasoning content is present""" - # Setup - chunk = ModelResponseStream( - id="test-id", - created=1234567890, - model="test-model", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - content="Regular content only", - role="assistant", - ), - ) - ], - ) - - mock_stream = AsyncMock() - iterator = LiteLLMCompletionStreamingIterator( - model="test-model", - litellm_custom_stream_wrapper=mock_stream, - request_input="Test input", - responses_api_request={}, - ) - - # Execute - transformed_chunk = ( - iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk) - ) - - # Assert - assert transformed_chunk.delta == "Regular content only" - assert transformed_chunk.type == "response.output_text.delta" - - -class TestReasoningContentFinalResponse: - """Test reasoning content preservation in final response transformation""" - - def test_reasoning_content_in_final_response(self): - """Test that reasoning content is included in final response""" - # Setup - response = ModelResponse( - id="test-id", - created=1234567890, - model="test-model", - object="chat.completion", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="Here is my answer", - role="assistant", - reasoning_content="Let me think step by step about this problem...", - ), - ) - ], - ) - - # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Test input", - responses_api_request={}, - chat_completion_response=response, - ) - - # Assert - assert hasattr(responses_api_response, "output") - assert len(responses_api_response.output) > 0 - - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] - assert len(reasoning_items) > 0, "No reasoning item found in output" - - reasoning_item = reasoning_items[0] - assert ( - reasoning_item.content[0].text - == "Let me think step by step about this problem..." - ) - - def test_no_reasoning_content_in_response(self): - """Test handling when no reasoning content in response""" - # Setup - response = ModelResponse( - id="test-id", - created=1234567890, - model="test-model", - object="chat.completion", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="Simple answer", - role="assistant", - ), - ) - ], - ) - - # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Test input", - responses_api_request={}, - chat_completion_response=response, - ) - - # Assert - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] - assert ( - len(reasoning_items) == 0 - ), "Should have no reasoning items when no reasoning content present" - - def test_multiple_choices_with_reasoning(self): - """Test handling multiple choices, first with reasoning content""" - # Setup - response = ModelResponse( - id="test-id", - created=1234567890, - model="test-model", - object="chat.completion", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="First answer", - role="assistant", - reasoning_content="Reasoning for first answer", - ), - ), - Choices( - finish_reason="stop", - index=1, - message=Message( - content="Second answer", - role="assistant", - reasoning_content="Reasoning for second answer", - ), - ), - ], - ) - - # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Test input", - responses_api_request={}, - chat_completion_response=response, - ) - - # Assert - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] - assert len(reasoning_items) == 1, "Should have exactly one reasoning item" - assert reasoning_items[0].content[0].text == "Reasoning for first answer" - - -def test_streaming_chunk_id_raw(): - """Test that streaming chunk IDs are raw (not encoded) to match OpenAI format""" - chunk = ModelResponseStream( - id="chunk-123", - created=1234567890, - model="test-model", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta(content="Hello", role="assistant"), - ) - ], - ) - - iterator = LiteLLMCompletionStreamingIterator( - model="test-model", - litellm_custom_stream_wrapper=AsyncMock(), - request_input="Test input", - responses_api_request={}, - custom_llm_provider="openai", - litellm_metadata={"model_info": {"id": "gpt-4"}}, - ) - - result = iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk) - - # Streaming chunk IDs should be raw (like OpenAI's msg_xxx format) - assert result.item_id == "chunk-123" # Should be raw, not encoded - assert not result.item_id.startswith("resp_") # Should NOT have resp_ prefix diff --git a/tests/test_litellm/test_azure_video_router.py b/tests/test_litellm/test_azure_video_router.py deleted file mode 100644 index e7e2e0a01ea..00000000000 --- a/tests/test_litellm/test_azure_video_router.py +++ /dev/null @@ -1,53 +0,0 @@ -""" -Test suite for Azure video router functionality. -Tests that the router method gets called correctly for Azure video generation. -""" - -import pytest -from unittest.mock import Mock, patch, MagicMock -import litellm - - -class TestAzureVideoRouter: - """Test suite for Azure video router functionality""" - - def setup_method(self): - """Setup test fixtures""" - self.model = "azure/sora-2" - self.prompt = "A beautiful sunset over mountains" - self.seconds = "5" - self.size = "1280x720" - - @patch("litellm.videos.main.base_llm_http_handler") - def test_azure_video_generation_router_call_mock(self, mock_handler): - """Test that Azure video generation calls the router method with mock response""" - # Setup mock response - mock_response = { - "id": "video_123", - "model": "sora-2", - "object": "video", - "status": "processing", - "created_at": 1234567890, - "progress": 0, - } - - # Configure the mock handler - mock_handler.video_generation_handler.return_value = mock_response - - # Call the video generation function with mock response - result = litellm.video_generation( - prompt=self.prompt, - model=self.model, - seconds=self.seconds, - size=self.size, - custom_llm_provider="azure", - mock_response=mock_response, - ) - - # Verify the result is a VideoObject with the expected data - assert result.id == mock_response["id"] - assert result.model == mock_response["model"] - assert result.object == mock_response["object"] - assert result.status == mock_response["status"] - assert result.created_at == mock_response["created_at"] - assert result.progress == mock_response["progress"] From e1afe2e29cee700710faa85063a9a0f7927104f6 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 16:42:15 -0700 Subject: [PATCH 007/576] test(e2e): bound the post-/model/new servable wait at 40s _await_model_servable used poll_timeout (120s), the spend/log read-back budget. A stuck model reload therefore stalled every suite that creates a deployment for two minutes before failing Give create_model a fixed harness middle ground: model_servable_timeout=40s, polled every 2s, with each /v1/models call capped at 5s and clamped to the remaining deadline so one slow GET cannot overrun the wait. Happy path still returns on the first listing. Not derived from proxy general_settings or env Transport.get accepts an optional per-call timeout for that clamp. Unit tests cover the deadline arithmetic and clamp without a live proxy (cherry picked from commit c082a0e6488f50978bf5255f6b5298ba7e8fd8da) --- tests/e2e/proxy_client.py | 134 ++++++++++-- tests/e2e/test_proxy_client_model_servable.py | 190 ++++++++++++++++++ tests/e2e/transport.py | 13 +- 3 files changed, 313 insertions(+), 24 deletions(-) create mode 100644 tests/e2e/test_proxy_client_model_servable.py diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 6c6b948e29c..87693175e62 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -75,12 +75,95 @@ from transport import HttpTransport, SplitTransport, Transport RowsPredicate = Callable[[list[SpendLogRow]], bool] +# After /model/new, poll the data plane until the model is listed (or fail). +# Shorter than poll_timeout (spend/log read-backs ~120s); longer than a single +# request. 40s is the harness middle ground: happy path returns on the first +# poll, a stuck reload fails in under a minute instead of two. +MODEL_SERVABLE_TIMEOUT = 40.0 +MODEL_SERVABLE_INTERVAL = 2.0 +# Cap each /v1/models poll so one slow request cannot outlast the budget. +# Clamped further to remaining deadline inside await_servable. +MODEL_SERVABLE_REQUEST_TIMEOUT = 5.0 + + +@dataclass(frozen=True, slots=True) +class Servable: + """The data plane listed the model within the deadline.""" + + +@dataclass(frozen=True, slots=True) +class NotServable: + """The deadline passed without the data plane listing the model. + + `last_result` is the final /v1/models read, so the caller can tell "the proxy + answered but omitted the model" (propagation) from "the read itself failed" + (network/auth) when reporting.""" + + last_result: Result[ModelsListResponse] | None + + +ServableOutcome = Servable | NotServable + + +def await_servable( + list_models: Callable[[float], Result[ModelsListResponse]], + *, + model_name: str, + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> ServableOutcome: + """Poll `list_models` until the data plane lists `model_name` or `timeout` elapses. + + `list_models` receives the per-poll request timeout, clamped to the remaining + deadline so a slow final poll cannot overrun the overall budget. Clock and sleep + are injected so this is exercised without wall-clock waits. Always polls at least + once when the loop starts with a positive budget.""" + deadline = now() + timeout + last_result: Result[ModelsListResponse] | None = None + while True: + remaining = deadline - now() + if remaining <= 0 and last_result is not None: + return NotServable(last_result=last_result) + poll_timeout = min(request_timeout, remaining) if remaining > 0 else request_timeout + last_result = list_models(poll_timeout) + if isinstance(last_result, Success) and any( + entry.id == model_name for entry in last_result.data.data + ): + return Servable() + if now() + interval >= deadline: + return NotServable(last_result=last_result) + sleep(interval) + + +def servable_timeout_message( + *, + model_name: str, + timeout: float, + last_result: Result[ModelsListResponse] | None, +) -> str: + last_error = ( + f"; last /v1/models poll did not succeed: {last_result}" + if last_result is not None and not isinstance(last_result, Success) + else "" + ) + return ( + f"model {model_name!r} was created but never became servable on the data " + f"plane within {timeout}s of /model/new (control/data-plane propagation or " + f"STORE_MODEL_IN_DB reload issue){last_error}" + ) + @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport poll_timeout: float = 120.0 poll_interval: float = 5.0 + model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT + model_servable_interval: float = MODEL_SERVABLE_INTERVAL + model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT # ---- keys / customers (satisfies lifecycle.ResourceClient) ---------- @@ -167,7 +250,12 @@ class ProxyClient: this returns can race the reload and 400 with "Invalid model name passed". We therefore poll the data-plane /v1/models until the model appears before handing back, so callers can invoke it immediately. In the monolithic case - it is already present on the first poll, so this adds one request.""" + it is already present on the first poll, so this adds one request. + + The wait is bounded by `model_servable_timeout` rather than the much longer + `poll_timeout` used for batched read-backs, so a stuck reload fails in under + a minute instead of two. Happy path still returns as soon as /v1/models lists + the model (usually the first poll).""" model_id = unwrap( self.transport.post( "/model/new", @@ -185,32 +273,34 @@ class ProxyClient: def _await_model_servable(self, model_name: str) -> None: """Block until the data plane lists `model_name`, or fail loudly if it does - not within poll_timeout (a real propagation/config problem, surfaced here - instead of as a downstream "Invalid model name passed").""" - deadline = time.monotonic() + self.poll_timeout - last_result: Result[ModelsListResponse] | None = None - while time.monotonic() < deadline: - last_result = self.transport.get( + not within model_servable_timeout (a real propagation/config problem, + surfaced here instead of as a downstream "Invalid model name passed").""" + outcome = await_servable( + lambda poll_timeout: self.transport.get( "/v1/models", headers=self.transport.master, params=NoBody(), response_type=ModelsListResponse, - ) - if isinstance(last_result, Success) and any( - entry.id == model_name for entry in last_result.data.data - ): + timeout=poll_timeout, + ), + model_name=model_name, + timeout=self.model_servable_timeout, + interval=self.model_servable_interval, + request_timeout=self.model_servable_request_timeout, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case Servable(): return - time.sleep(self.poll_interval) - last_error = ( - f"; last /v1/models poll did not succeed: {last_result}" - if last_result is not None and not isinstance(last_result, Success) - else "" - ) - raise AssertionError( - f"model {model_name!r} was created but never became servable on the data " - f"plane within {self.poll_timeout}s of /model/new (control/data-plane " - f"propagation or STORE_MODEL_IN_DB reload issue){last_error}" - ) + case NotServable(last_result=last_result): + raise AssertionError( + servable_timeout_message( + model_name=model_name, + timeout=self.model_servable_timeout, + last_result=last_result, + ) + ) def update_model(self, model_id: str, litellm_params: LiteLLMParamsBody) -> None: """Merge `litellm_params` over the deployment `model_id`'s stored params via diff --git a/tests/e2e/test_proxy_client_model_servable.py b/tests/e2e/test_proxy_client_model_servable.py new file mode 100644 index 00000000000..0cc63f882e2 --- /dev/null +++ b/tests/e2e/test_proxy_client_model_servable.py @@ -0,0 +1,190 @@ +"""Harness coverage for the bounded wait after /model/new (no live proxy). + +Model propagation is polled to a deadline so a stuck control/data-plane reload fails +fast instead of stalling every test that creates a model. The clock and sleep are +injected, so these assert the deadline arithmetic without waiting. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from e2e_http import NetworkError, Result, Success +from models import ModelListEntry, ModelsListResponse +from proxy_client import ( + MODEL_SERVABLE_REQUEST_TIMEOUT, + MODEL_SERVABLE_TIMEOUT, + NotServable, + Servable, + await_servable, + servable_timeout_message, +) + + +def _listing(*model_names: str) -> Result[ModelsListResponse]: + return Success( + status_code=200, + data=ModelsListResponse(data=tuple(ModelListEntry(id=name) for name in model_names)), + ) + + +@dataclass(slots=True) +class FakeClock: + """A clock that only advances when the code under test sleeps or a slow poll runs.""" + + seconds: float = 0.0 + slept: list[float] = field(default_factory=list) # mutable-ok: records calls for assertions + + def now(self) -> float: + return self.seconds + + def sleep(self, duration: float) -> None: + self.slept.append(duration) + self.seconds += duration + + +@dataclass(slots=True) +class FakeModelList: + """Returns each queued /v1/models read in turn, repeating the last forever.""" + + responses: tuple[Result[ModelsListResponse], ...] + calls: int = 0 + timeouts: list[float] = field(default_factory=list) # mutable-ok: records call timeouts + + def __call__(self, request_timeout: float) -> Result[ModelsListResponse]: + self.timeouts.append(request_timeout) + response = self.responses[min(self.calls, len(self.responses) - 1)] + self.calls += 1 + return response + + +def test_returns_servable_on_first_listing_without_sleeping() -> None: + clock = FakeClock() + list_models = FakeModelList(responses=(_listing("my-model"),)) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=40.0, + interval=2.0, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == Servable() + assert list_models.calls == 1 + assert list_models.timeouts == [5.0] + assert clock.slept == [] + + +def test_polls_until_the_model_appears() -> None: + clock = FakeClock() + list_models = FakeModelList(responses=(_listing("other"), _listing("other"), _listing("other", "my-model"))) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=40.0, + interval=2.0, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == Servable() + assert list_models.calls == 3 + assert clock.seconds == 4.0 + + +def test_gives_up_at_the_deadline_rather_than_polling_forever() -> None: + clock = FakeClock() + list_models = FakeModelList(responses=(_listing("other"),)) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=10.0, + interval=2.0, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert isinstance(outcome, NotServable) + assert clock.seconds == 8.0 + assert list_models.calls == 5 + assert list_models.timeouts == [5.0, 5.0, 5.0, 4.0, 2.0] + + +def test_does_not_wait_past_the_overall_budget() -> None: + clock = FakeClock() + + outcome = await_servable( + FakeModelList(responses=(_listing("other"),)), + model_name="my-model", + timeout=MODEL_SERVABLE_TIMEOUT, + interval=2.0, + request_timeout=MODEL_SERVABLE_REQUEST_TIMEOUT, + now=clock.now, + sleep=clock.sleep, + ) + + assert isinstance(outcome, NotServable) + assert clock.seconds <= MODEL_SERVABLE_TIMEOUT + + +def test_clamps_request_timeout_to_remaining_deadline() -> None: + """A slow final poll must not receive the full request cap when less budget remains. + + Without the clamp, remaining=3 and cap=5 lets the transport block for 5s and the + overall wait overruns model_servable_timeout by up to ~cap seconds. + """ + clock = FakeClock() + timeouts: list[float] = [] + + def list_models(request_timeout: float) -> Result[ModelsListResponse]: + timeouts.append(request_timeout) + clock.seconds += request_timeout + return _listing("other") + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=10.0, + interval=2.0, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert isinstance(outcome, NotServable) + assert timeouts[0] == 5.0 + assert any(timeout < 5.0 for timeout in timeouts) + assert timeouts[-1] == 3.0 + assert clock.seconds <= 10.0 + + +def test_reports_a_failed_read_distinctly_from_a_missing_model() -> None: + clock = FakeClock() + unreachable: Result[ModelsListResponse] = NetworkError(message="connection refused") + + outcome = await_servable( + FakeModelList(responses=(unreachable,)), + model_name="my-model", + timeout=1.0, + interval=0.5, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == NotServable(last_result=unreachable) + message = servable_timeout_message(model_name="my-model", timeout=1.0, last_result=unreachable) + assert "connection refused" in message + + listed_without_model = _listing("other") + propagation_message = servable_timeout_message( + model_name="my-model", timeout=1.0, last_result=listed_without_model + ) + assert "did not succeed" not in propagation_message diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index a6adf83ed1f..27b11befc8e 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -58,6 +58,7 @@ class Transport(Protocol): headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: ... def delete[R: BaseModel]( @@ -136,13 +137,16 @@ class HttpTransport: headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: + """`timeout` overrides the transport-wide request_timeout for this call, for + pollers whose own deadline is shorter than it.""" return e2e_http.get( self._url(path), headers=headers, params=params, response_type=response_type, - timeout=self.request_timeout, + timeout=self.request_timeout if timeout is None else timeout, ) def delete[R: BaseModel]( @@ -336,9 +340,14 @@ class SplitTransport: headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: return self._route(path).get( - path, headers=headers, params=params, response_type=response_type + path, + headers=headers, + params=params, + response_type=response_type, + timeout=timeout, ) def delete[R: BaseModel]( From 5aa66ea33e6c194f66d360e22292f22eddac38e7 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 17:09:20 -0700 Subject: [PATCH 008/576] fix(e2e): wait one default DB reload interval of continuous listing create_model returned after the first /v1/models hit that listed the model, so chat could still land on a cold gateway worker (numWorkers>1 / peer pod) and 400 Invalid model name. Require continuous listing for the product default add_deployment interval (30s) after first sight so every worker has synced from the DB; first listing still bounded at 40s (cherry picked from commit 7d1ee2ff861b970f6de3f6759ff015947af9d2a1) --- tests/e2e/proxy_client.py | 85 ++++++++---- tests/e2e/test_proxy_client_model_servable.py | 125 ++++++++++++------ 2 files changed, 148 insertions(+), 62 deletions(-) diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 87693175e62..dfbb90ac08e 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -75,14 +75,17 @@ from transport import HttpTransport, SplitTransport, Transport RowsPredicate = Callable[[list[SpendLogRow]], bool] -# After /model/new, poll the data plane until the model is listed (or fail). -# Shorter than poll_timeout (spend/log read-backs ~120s); longer than a single -# request. 40s is the harness middle ground: happy path returns on the first -# poll, a stuck reload fails in under a minute instead of two. +# After /model/new, the control-plane writer reloads itself immediately, but every +# other gateway worker (and peer pod) only picks the model up on its add_deployment +# job. That job runs every proxy_config_reload_interval_seconds (product default 30). +# A single /v1/models hit can land on a hot worker while the next /chat hits a cold +# one ("Invalid model name"). Wait for first listing within MODEL_SERVABLE_TIMEOUT, +# then require continuous listing for MODEL_SERVABLE_DB_SYNC_SECONDS (the default +# reload interval) so every worker has had a chance to sync from the DB. MODEL_SERVABLE_TIMEOUT = 40.0 +MODEL_SERVABLE_DB_SYNC_SECONDS = 30.0 MODEL_SERVABLE_INTERVAL = 2.0 -# Cap each /v1/models poll so one slow request cannot outlast the budget. -# Clamped further to remaining deadline inside await_servable. +# Cap each /v1/models poll so one slow request cannot outlast the remaining budget. MODEL_SERVABLE_REQUEST_TIMEOUT = 5.0 @@ -112,29 +115,55 @@ def await_servable( timeout: float, interval: float, request_timeout: float, + db_sync_seconds: float, now: Callable[[], float], sleep: Callable[[float], None], ) -> ServableOutcome: - """Poll `list_models` until the data plane lists `model_name` or `timeout` elapses. + """Poll until `model_name` is listed long enough for every worker to DB-sync. - `list_models` receives the per-poll request timeout, clamped to the remaining - deadline so a slow final poll cannot overrun the overall budget. Clock and sleep - are injected so this is exercised without wall-clock waits. Always polls at least - once when the loop starts with a positive budget.""" - deadline = now() + timeout + First listing must happen within `timeout`. After that, the model must stay + listed continuously for `db_sync_seconds` (any miss resets the continuous + window). `db_sync_seconds=0` returns on the first listing. Each poll's request + timeout is clamped to the remaining budget. Clock and sleep are injected.""" + started = now() + first_seen_at: float | None = None last_result: Result[ModelsListResponse] | None = None while True: - remaining = deadline - now() + t = now() + if first_seen_at is None: + deadline = started + timeout + else: + deadline = first_seen_at + db_sync_seconds + remaining = deadline - t if remaining <= 0 and last_result is not None: + if first_seen_at is not None and db_sync_seconds <= 0: + return Servable() + if first_seen_at is not None and t - first_seen_at >= db_sync_seconds: + return Servable() return NotServable(last_result=last_result) poll_timeout = min(request_timeout, remaining) if remaining > 0 else request_timeout last_result = list_models(poll_timeout) - if isinstance(last_result, Success) and any( + listed = isinstance(last_result, Success) and any( entry.id == model_name for entry in last_result.data.data - ): + ) + t = now() + if not listed: + first_seen_at = None + elif first_seen_at is None: + first_seen_at = t + if db_sync_seconds <= 0: + return Servable() + elif t - first_seen_at >= db_sync_seconds: return Servable() - if now() + interval >= deadline: - return NotServable(last_result=last_result) + if first_seen_at is None: + if now() + interval >= started + timeout: + return NotServable(last_result=last_result) + elif now() + interval >= first_seen_at + db_sync_seconds: + # Final stretch: sleep only the remainder of the continuous window. + remainder = first_seen_at + db_sync_seconds - now() + if remainder > 0: + sleep(remainder) + continue sleep(interval) @@ -142,6 +171,7 @@ def servable_timeout_message( *, model_name: str, timeout: float, + db_sync_seconds: float, last_result: Result[ModelsListResponse] | None, ) -> str: last_error = ( @@ -151,7 +181,8 @@ def servable_timeout_message( ) return ( f"model {model_name!r} was created but never became servable on the data " - f"plane within {timeout}s of /model/new (control/data-plane propagation or " + f"plane within {timeout}s of first listing (plus {db_sync_seconds}s continuous " + f"DB sync) after /model/new (control/data-plane propagation or " f"STORE_MODEL_IN_DB reload issue){last_error}" ) @@ -162,6 +193,7 @@ class ProxyClient: poll_timeout: float = 120.0 poll_interval: float = 5.0 model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT + model_servable_db_sync_seconds: float = MODEL_SERVABLE_DB_SYNC_SECONDS model_servable_interval: float = MODEL_SERVABLE_INTERVAL model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT @@ -252,10 +284,10 @@ class ProxyClient: handing back, so callers can invoke it immediately. In the monolithic case it is already present on the first poll, so this adds one request. - The wait is bounded by `model_servable_timeout` rather than the much longer - `poll_timeout` used for batched read-backs, so a stuck reload fails in under - a minute instead of two. Happy path still returns as soon as /v1/models lists - the model (usually the first poll).""" + First listing must arrive within `model_servable_timeout` (not the longer + spend `poll_timeout`). The model must then stay listed for + `model_servable_db_sync_seconds` (product default DB reload interval) so every + gateway worker has run add_deployment before callers use the model.""" model_id = unwrap( self.transport.post( "/model/new", @@ -272,9 +304,10 @@ class ProxyClient: return model_id def _await_model_servable(self, model_name: str) -> None: - """Block until the data plane lists `model_name`, or fail loudly if it does - not within model_servable_timeout (a real propagation/config problem, - surfaced here instead of as a downstream "Invalid model name passed").""" + """Block until the data plane lists `model_name` long enough for DB sync. + + Fails if first listing misses model_servable_timeout, or if continuous listing + for model_servable_db_sync_seconds never holds (multi-worker / peer reload).""" outcome = await_servable( lambda poll_timeout: self.transport.get( "/v1/models", @@ -287,6 +320,7 @@ class ProxyClient: timeout=self.model_servable_timeout, interval=self.model_servable_interval, request_timeout=self.model_servable_request_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, now=time.monotonic, sleep=time.sleep, ) @@ -298,6 +332,7 @@ class ProxyClient: servable_timeout_message( model_name=model_name, timeout=self.model_servable_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, last_result=last_result, ) ) diff --git a/tests/e2e/test_proxy_client_model_servable.py b/tests/e2e/test_proxy_client_model_servable.py index 0cc63f882e2..c9edd148c2b 100644 --- a/tests/e2e/test_proxy_client_model_servable.py +++ b/tests/e2e/test_proxy_client_model_servable.py @@ -1,8 +1,9 @@ """Harness coverage for the bounded wait after /model/new (no live proxy). -Model propagation is polled to a deadline so a stuck control/data-plane reload fails -fast instead of stalling every test that creates a model. The clock and sleep are -injected, so these assert the deadline arithmetic without waiting. +create_model must wait for the product default DB reload interval of continuous +listing so multi-worker gateways finish add_deployment before callers use the +model. Clock and sleep are injected so these assert the deadline arithmetic +without wall-clock waits. """ from __future__ import annotations @@ -12,6 +13,7 @@ from dataclasses import dataclass, field from e2e_http import NetworkError, Result, Success from models import ModelListEntry, ModelsListResponse from proxy_client import ( + MODEL_SERVABLE_DB_SYNC_SECONDS, MODEL_SERVABLE_REQUEST_TIMEOUT, MODEL_SERVABLE_TIMEOUT, NotServable, @@ -30,8 +32,6 @@ def _listing(*model_names: str) -> Result[ModelsListResponse]: @dataclass(slots=True) class FakeClock: - """A clock that only advances when the code under test sleeps or a slow poll runs.""" - seconds: float = 0.0 slept: list[float] = field(default_factory=list) # mutable-ok: records calls for assertions @@ -45,8 +45,6 @@ class FakeClock: @dataclass(slots=True) class FakeModelList: - """Returns each queued /v1/models read in turn, repeating the last forever.""" - responses: tuple[Result[ModelsListResponse], ...] calls: int = 0 timeouts: list[float] = field(default_factory=list) # mutable-ok: records call timeouts @@ -58,7 +56,7 @@ class FakeModelList: return response -def test_returns_servable_on_first_listing_without_sleeping() -> None: +def test_returns_on_first_listing_when_db_sync_is_zero() -> None: clock = FakeClock() list_models = FakeModelList(responses=(_listing("my-model"),)) @@ -68,17 +66,64 @@ def test_returns_servable_on_first_listing_without_sleeping() -> None: timeout=40.0, interval=2.0, request_timeout=5.0, + db_sync_seconds=0.0, now=clock.now, sleep=clock.sleep, ) assert outcome == Servable() assert list_models.calls == 1 - assert list_models.timeouts == [5.0] assert clock.slept == [] -def test_polls_until_the_model_appears() -> None: +def test_requires_continuous_listing_for_default_db_sync_interval() -> None: + clock = FakeClock() + list_models = FakeModelList(responses=(_listing("my-model"),)) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=40.0, + interval=2.0, + request_timeout=5.0, + db_sync_seconds=MODEL_SERVABLE_DB_SYNC_SECONDS, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == Servable() + assert clock.seconds >= MODEL_SERVABLE_DB_SYNC_SECONDS + assert list_models.calls >= 2 + + +def test_resets_db_sync_window_when_a_poll_misses() -> None: + clock = FakeClock() + list_models = FakeModelList( + responses=( + _listing("my-model"), + _listing("my-model"), + _listing("other"), + _listing("my-model"), + ) + ) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=40.0, + interval=2.0, + request_timeout=5.0, + db_sync_seconds=6.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == Servable() + assert list_models.calls >= 4 + assert clock.seconds >= 6.0 + + +def test_polls_until_the_model_first_appears() -> None: clock = FakeClock() list_models = FakeModelList(responses=(_listing("other"), _listing("other"), _listing("other", "my-model"))) @@ -88,6 +133,7 @@ def test_polls_until_the_model_appears() -> None: timeout=40.0, interval=2.0, request_timeout=5.0, + db_sync_seconds=0.0, now=clock.now, sleep=clock.sleep, ) @@ -97,7 +143,7 @@ def test_polls_until_the_model_appears() -> None: assert clock.seconds == 4.0 -def test_gives_up_at_the_deadline_rather_than_polling_forever() -> None: +def test_gives_up_if_first_listing_never_arrives() -> None: clock = FakeClock() list_models = FakeModelList(responses=(_listing("other"),)) @@ -107,6 +153,7 @@ def test_gives_up_at_the_deadline_rather_than_polling_forever() -> None: timeout=10.0, interval=2.0, request_timeout=5.0, + db_sync_seconds=30.0, now=clock.now, sleep=clock.sleep, ) @@ -114,32 +161,9 @@ def test_gives_up_at_the_deadline_rather_than_polling_forever() -> None: assert isinstance(outcome, NotServable) assert clock.seconds == 8.0 assert list_models.calls == 5 - assert list_models.timeouts == [5.0, 5.0, 5.0, 4.0, 2.0] - - -def test_does_not_wait_past_the_overall_budget() -> None: - clock = FakeClock() - - outcome = await_servable( - FakeModelList(responses=(_listing("other"),)), - model_name="my-model", - timeout=MODEL_SERVABLE_TIMEOUT, - interval=2.0, - request_timeout=MODEL_SERVABLE_REQUEST_TIMEOUT, - now=clock.now, - sleep=clock.sleep, - ) - - assert isinstance(outcome, NotServable) - assert clock.seconds <= MODEL_SERVABLE_TIMEOUT def test_clamps_request_timeout_to_remaining_deadline() -> None: - """A slow final poll must not receive the full request cap when less budget remains. - - Without the clamp, remaining=3 and cap=5 lets the transport block for 5s and the - overall wait overruns model_servable_timeout by up to ~cap seconds. - """ clock = FakeClock() timeouts: list[float] = [] @@ -154,6 +178,7 @@ def test_clamps_request_timeout_to_remaining_deadline() -> None: timeout=10.0, interval=2.0, request_timeout=5.0, + db_sync_seconds=30.0, now=clock.now, sleep=clock.sleep, ) @@ -161,10 +186,27 @@ def test_clamps_request_timeout_to_remaining_deadline() -> None: assert isinstance(outcome, NotServable) assert timeouts[0] == 5.0 assert any(timeout < 5.0 for timeout in timeouts) - assert timeouts[-1] == 3.0 assert clock.seconds <= 10.0 +def test_does_not_wait_past_first_listing_budget_when_missing() -> None: + clock = FakeClock() + + outcome = await_servable( + FakeModelList(responses=(_listing("other"),)), + model_name="my-model", + timeout=MODEL_SERVABLE_TIMEOUT, + interval=2.0, + request_timeout=MODEL_SERVABLE_REQUEST_TIMEOUT, + db_sync_seconds=MODEL_SERVABLE_DB_SYNC_SECONDS, + now=clock.now, + sleep=clock.sleep, + ) + + assert isinstance(outcome, NotServable) + assert clock.seconds <= MODEL_SERVABLE_TIMEOUT + + def test_reports_a_failed_read_distinctly_from_a_missing_model() -> None: clock = FakeClock() unreachable: Result[ModelsListResponse] = NetworkError(message="connection refused") @@ -175,16 +217,25 @@ def test_reports_a_failed_read_distinctly_from_a_missing_model() -> None: timeout=1.0, interval=0.5, request_timeout=5.0, + db_sync_seconds=0.0, now=clock.now, sleep=clock.sleep, ) assert outcome == NotServable(last_result=unreachable) - message = servable_timeout_message(model_name="my-model", timeout=1.0, last_result=unreachable) + message = servable_timeout_message( + model_name="my-model", + timeout=1.0, + db_sync_seconds=0.0, + last_result=unreachable, + ) assert "connection refused" in message listed_without_model = _listing("other") propagation_message = servable_timeout_message( - model_name="my-model", timeout=1.0, last_result=listed_without_model + model_name="my-model", + timeout=1.0, + db_sync_seconds=30.0, + last_result=listed_without_model, ) assert "did not succeed" not in propagation_message From 5953a66eab8212beac24d3265cd118b7f4a51176 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 17:10:04 -0700 Subject: [PATCH 009/576] test(e2e): drop proxy_client model-servable unit tests Keep the create_model DB-sync wait in the harness; the pure-function unit file is not needed for this PR (cherry picked from commit 89204651d1a4537c6f21550c5ab85448ae0923f8) --- tests/e2e/test_proxy_client_model_servable.py | 241 ------------------ 1 file changed, 241 deletions(-) delete mode 100644 tests/e2e/test_proxy_client_model_servable.py diff --git a/tests/e2e/test_proxy_client_model_servable.py b/tests/e2e/test_proxy_client_model_servable.py deleted file mode 100644 index c9edd148c2b..00000000000 --- a/tests/e2e/test_proxy_client_model_servable.py +++ /dev/null @@ -1,241 +0,0 @@ -"""Harness coverage for the bounded wait after /model/new (no live proxy). - -create_model must wait for the product default DB reload interval of continuous -listing so multi-worker gateways finish add_deployment before callers use the -model. Clock and sleep are injected so these assert the deadline arithmetic -without wall-clock waits. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field - -from e2e_http import NetworkError, Result, Success -from models import ModelListEntry, ModelsListResponse -from proxy_client import ( - MODEL_SERVABLE_DB_SYNC_SECONDS, - MODEL_SERVABLE_REQUEST_TIMEOUT, - MODEL_SERVABLE_TIMEOUT, - NotServable, - Servable, - await_servable, - servable_timeout_message, -) - - -def _listing(*model_names: str) -> Result[ModelsListResponse]: - return Success( - status_code=200, - data=ModelsListResponse(data=tuple(ModelListEntry(id=name) for name in model_names)), - ) - - -@dataclass(slots=True) -class FakeClock: - seconds: float = 0.0 - slept: list[float] = field(default_factory=list) # mutable-ok: records calls for assertions - - def now(self) -> float: - return self.seconds - - def sleep(self, duration: float) -> None: - self.slept.append(duration) - self.seconds += duration - - -@dataclass(slots=True) -class FakeModelList: - responses: tuple[Result[ModelsListResponse], ...] - calls: int = 0 - timeouts: list[float] = field(default_factory=list) # mutable-ok: records call timeouts - - def __call__(self, request_timeout: float) -> Result[ModelsListResponse]: - self.timeouts.append(request_timeout) - response = self.responses[min(self.calls, len(self.responses) - 1)] - self.calls += 1 - return response - - -def test_returns_on_first_listing_when_db_sync_is_zero() -> None: - clock = FakeClock() - list_models = FakeModelList(responses=(_listing("my-model"),)) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=40.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=0.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == Servable() - assert list_models.calls == 1 - assert clock.slept == [] - - -def test_requires_continuous_listing_for_default_db_sync_interval() -> None: - clock = FakeClock() - list_models = FakeModelList(responses=(_listing("my-model"),)) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=40.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=MODEL_SERVABLE_DB_SYNC_SECONDS, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == Servable() - assert clock.seconds >= MODEL_SERVABLE_DB_SYNC_SECONDS - assert list_models.calls >= 2 - - -def test_resets_db_sync_window_when_a_poll_misses() -> None: - clock = FakeClock() - list_models = FakeModelList( - responses=( - _listing("my-model"), - _listing("my-model"), - _listing("other"), - _listing("my-model"), - ) - ) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=40.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=6.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == Servable() - assert list_models.calls >= 4 - assert clock.seconds >= 6.0 - - -def test_polls_until_the_model_first_appears() -> None: - clock = FakeClock() - list_models = FakeModelList(responses=(_listing("other"), _listing("other"), _listing("other", "my-model"))) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=40.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=0.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == Servable() - assert list_models.calls == 3 - assert clock.seconds == 4.0 - - -def test_gives_up_if_first_listing_never_arrives() -> None: - clock = FakeClock() - list_models = FakeModelList(responses=(_listing("other"),)) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=10.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=30.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert isinstance(outcome, NotServable) - assert clock.seconds == 8.0 - assert list_models.calls == 5 - - -def test_clamps_request_timeout_to_remaining_deadline() -> None: - clock = FakeClock() - timeouts: list[float] = [] - - def list_models(request_timeout: float) -> Result[ModelsListResponse]: - timeouts.append(request_timeout) - clock.seconds += request_timeout - return _listing("other") - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=10.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=30.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert isinstance(outcome, NotServable) - assert timeouts[0] == 5.0 - assert any(timeout < 5.0 for timeout in timeouts) - assert clock.seconds <= 10.0 - - -def test_does_not_wait_past_first_listing_budget_when_missing() -> None: - clock = FakeClock() - - outcome = await_servable( - FakeModelList(responses=(_listing("other"),)), - model_name="my-model", - timeout=MODEL_SERVABLE_TIMEOUT, - interval=2.0, - request_timeout=MODEL_SERVABLE_REQUEST_TIMEOUT, - db_sync_seconds=MODEL_SERVABLE_DB_SYNC_SECONDS, - now=clock.now, - sleep=clock.sleep, - ) - - assert isinstance(outcome, NotServable) - assert clock.seconds <= MODEL_SERVABLE_TIMEOUT - - -def test_reports_a_failed_read_distinctly_from_a_missing_model() -> None: - clock = FakeClock() - unreachable: Result[ModelsListResponse] = NetworkError(message="connection refused") - - outcome = await_servable( - FakeModelList(responses=(unreachable,)), - model_name="my-model", - timeout=1.0, - interval=0.5, - request_timeout=5.0, - db_sync_seconds=0.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == NotServable(last_result=unreachable) - message = servable_timeout_message( - model_name="my-model", - timeout=1.0, - db_sync_seconds=0.0, - last_result=unreachable, - ) - assert "connection refused" in message - - listed_without_model = _listing("other") - propagation_message = servable_timeout_message( - model_name="my-model", - timeout=1.0, - db_sync_seconds=30.0, - last_result=listed_without_model, - ) - assert "did not succeed" not in propagation_message From 38d03fd341bf5b7030b2ebbc45c2555a7f1ee1e6 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 17:16:30 -0700 Subject: [PATCH 010/576] fix(e2e): never skip the final deadline-clamped model-servable poll When less than one full poll interval remained in the first-listing budget, the pre-sleep check returned NotServable without another /v1/models call. Sleep only min(interval, time left) so a model that becomes listable in the last seconds of the timeout still gets a clamped final poll (cherry picked from commit 8439195922c913d118cb146409c75cc081d23e6e) --- tests/e2e/proxy_client.py | 43 ++++++++++++++++++++------------------- 1 file changed, 22 insertions(+), 21 deletions(-) diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index dfbb90ac08e..36dada5770a 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -124,24 +124,28 @@ def await_servable( First listing must happen within `timeout`. After that, the model must stay listed continuously for `db_sync_seconds` (any miss resets the continuous window). `db_sync_seconds=0` returns on the first listing. Each poll's request - timeout is clamped to the remaining budget. Clock and sleep are injected.""" + timeout is clamped to the remaining budget. Sleeps only min(interval, time left) + so a final deadline-clamped poll is never skipped just because a full interval + does not fit. Clock and sleep are injected.""" started = now() first_seen_at: float | None = None last_result: Result[ModelsListResponse] | None = None while True: t = now() - if first_seen_at is None: - deadline = started + timeout - else: - deadline = first_seen_at + db_sync_seconds - remaining = deadline - t - if remaining <= 0 and last_result is not None: - if first_seen_at is not None and db_sync_seconds <= 0: - return Servable() - if first_seen_at is not None and t - first_seen_at >= db_sync_seconds: + phase_deadline = ( + started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds + ) + remaining = phase_deadline - t + if remaining <= 0: + if ( + last_result is not None + and first_seen_at is not None + and (db_sync_seconds <= 0 or t - first_seen_at >= db_sync_seconds) + ): return Servable() return NotServable(last_result=last_result) - poll_timeout = min(request_timeout, remaining) if remaining > 0 else request_timeout + + poll_timeout = min(request_timeout, remaining) last_result = list_models(poll_timeout) listed = isinstance(last_result, Success) and any( entry.id == model_name for entry in last_result.data.data @@ -155,16 +159,13 @@ def await_servable( return Servable() elif t - first_seen_at >= db_sync_seconds: return Servable() - if first_seen_at is None: - if now() + interval >= started + timeout: - return NotServable(last_result=last_result) - elif now() + interval >= first_seen_at + db_sync_seconds: - # Final stretch: sleep only the remainder of the continuous window. - remainder = first_seen_at + db_sync_seconds - now() - if remainder > 0: - sleep(remainder) - continue - sleep(interval) + + phase_deadline = ( + started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds + ) + wait = min(interval, phase_deadline - now()) + if wait > 0: + sleep(wait) def servable_timeout_message( From 87be33f9354f399100c7ededdfba31b0e831d225 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 17:35:13 -0700 Subject: [PATCH 011/576] fix(e2e): reject first listing that returns after the 40s deadline A poll may start with remaining budget and still return after started+timeout if the transport overruns its clamp. Recheck the first-listing deadline after the response so a late listing does not open the continuous DB-sync phase (cherry picked from commit 7ff2bcbf1498ee82f7dbe0b4330c1ab48927ed01) --- tests/e2e/proxy_client.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 36dada5770a..b3fc8538322 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -154,6 +154,8 @@ def await_servable( if not listed: first_seen_at = None elif first_seen_at is None: + if t > started + timeout: + return NotServable(last_result=last_result) first_seen_at = t if db_sync_seconds <= 0: return Servable() From 82fa66908bb72748efc574a8a1a8750155a85c14 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 22:15:21 -0700 Subject: [PATCH 012/576] test(e2e): poll MCP tools across multi-worker lag (#35047) * fix(mcp): resolve call_tool by registry without requiring tool map Multi-worker reloads put MCP servers in the registry from the DB but do not re-run tools/list on every process. Gating call_tool on tool_name_to_mcp_server_name_mapping made cold workers 500 with Tool not found after another worker had already listed the tool. Treat a registry match on server id/name/alias as enough; upstream rejects unknown tools * test(e2e): poll MCP register, tools/list, and tools/call across multi-worker lag Stage multi-worker gateways only load MCP servers and tool maps on the process that handled the request. Poll until the server is listed, the tool appears on tools/list, and tools/call is not a cold-worker 500 so key-access and Datadog MCP e2e stop racing the LB * Revert "fix(mcp): resolve call_tool by registry without requiring tool map" This reverts commit 8b56e51e39b876d13d1112efa4130554ddf5f173. * test(e2e): tighten MCP multi-worker lag classifier Only retry tools/call on gateway shapes Tool not found and server_not_found, not any 500 that mentions tool/server not found, so upstream failures are not retried until the poll deadline * test(e2e): drop unit file for MCP lag classifier The live await_call_tool polls already cover multi-worker lag; a separate string-match unit module is not worth keeping (cherry picked from commit c274cf321c5c35c629220a89bb497d15b56f870f) --- tests/e2e/mcp/mcp_client.py | 91 +++++++++++++++++++++- tests/e2e/mcp/test_mcp_access_group_e2e.py | 1 + tests/e2e/mcp/test_mcp_datadog_e2e.py | 28 ++++--- tests/e2e/mcp/test_mcp_key_access_e2e.py | 15 ++-- 4 files changed, 111 insertions(+), 24 deletions(-) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 6d0f6ddc760..33ec557c339 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -11,13 +11,14 @@ request/response bodies are co-located here because only this suite speaks MCP. from __future__ import annotations +import re import time from collections.abc import Mapping from dataclasses import dataclass from pydantic import BaseModel, ConfigDict, Field, RootModel -from e2e_http import Headers, NoBody, Result, Success, unwrap +from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap from models import KeyGenerateBody, ObjectPermission from proxy_client import ProxyClient @@ -270,6 +271,60 @@ class McpClient: ) time.sleep(self.proxy.poll_interval) + def await_call_tool( + self, + key: str, + *, + server_id: str, + name: str, + arguments: McpToolArguments, + ) -> McpCallToolResponse: + """Poll tools/call until the result is not a multi-worker registry miss. + + Retries only on the gateway's own cold-worker 500 shapes (Tool + not found / server_not_found). Upstream tool errors and other 500s fail + immediately so non-idempotent calls are not repeated. + """ + deadline = time.monotonic() + self.proxy.poll_timeout + last: Result[McpCallToolResponse] | None = None + while True: + last = self.call_tool(key, server_id=server_id, name=name, arguments=arguments) + if not _is_mcp_not_synced(last, tool_name=name): + return unwrap(last) + if time.monotonic() >= deadline: + raise AssertionError( + f"tools/call for {name!r} on server {server_id} still missing on the " + f"data plane after {self.proxy.poll_timeout}s (multi-worker registry lag); " + f"last result: {last}" + ) + time.sleep(self.proxy.poll_interval) + + def await_call_tool_denied( + self, + key: str, + *, + server_id: str, + name: str, + arguments: McpToolArguments, + ) -> UnknownApiError: + """Poll tools/call until a cold-worker miss clears and the call is 403 access_denied.""" + deadline = time.monotonic() + self.proxy.poll_timeout + last: Result[McpCallToolResponse] | None = None + while True: + last = self.call_tool(key, server_id=server_id, name=name, arguments=arguments) + if isinstance(last, UnknownApiError) and last.status_code == 403: + return last + if not _is_mcp_not_synced(last, tool_name=name): + raise AssertionError( + f"ungranted key's tools/call was not 403 access_denied: {last}" + ) + if time.monotonic() >= deadline: + raise AssertionError( + f"ungranted key never got 403 for {name!r} within {self.proxy.poll_timeout}s; " + f"last result: {last}" + ) + time.sleep(self.proxy.poll_interval) + def register_mcp_content_filter(self, *, name: str, blocked_keyword: str) -> str: """Register a default-on content-filter guardrail that runs on the MCP tool-call hook (pre_mcp_call) and blocks a single keyword. The keyword is @@ -317,5 +372,39 @@ class McpClient: ) +def _is_mcp_not_synced( + result: Result[McpCallToolResponse], + *, + tool_name: str | None = None, +) -> bool: + """True only for gateway multi-worker registry misses, not upstream errors. + + Matches the proxy's own shapes: + - ValueError ``Tool not found`` wrapped as HTTP 500 (cold tool map / + unresolved server on this process) + - REST ``server_not_found`` when this worker has not loaded the MCP server row + + Does not treat arbitrary 500 bodies that merely mention "tool" and "not found" + (e.g. upstream MCP payload text) as lag, so await_call_tool does not retry + real failures or non-idempotent calls. + """ + if not isinstance(result, UnknownApiError) or result.status_code != 500: + return False + body = result.body + body_l = body.lower() + + if "server_not_found" in body_l: + return True + if re.search(r"mcp server ['\"][^'\"]+['\"] was not found", body_l): + return True + + # Gateway: "Tool search_datadog_logs not found" (optionally inside a longer message) + if tool_name is not None: + return ( + re.search(rf"\btool\s+{re.escape(tool_name)}\s+not found\b", body_l) is not None + ) + return re.search(r"\btool\s+\S+\s+not found\b", body_l) is not None + + def build_client(proxy: ProxyClient) -> McpClient: return McpClient(proxy=proxy) diff --git a/tests/e2e/mcp/test_mcp_access_group_e2e.py b/tests/e2e/mcp/test_mcp_access_group_e2e.py index 1b53d1ca0b4..f72b75fd43d 100644 --- a/tests/e2e/mcp/test_mcp_access_group_e2e.py +++ b/tests/e2e/mcp/test_mcp_access_group_e2e.py @@ -29,6 +29,7 @@ class TestMcpAccessGroupToolSelection: ) -> None: group = f"e2e-mcp-grp-{unique_marker()}" server_id = register_datadog_mcp(client, resources, mcp_access_groups=[group]) + client.await_registered(server_id) granted = client.generate_key( user_id=f"e2e-mcp-ag-granted-{unique_marker()}", diff --git a/tests/e2e/mcp/test_mcp_datadog_e2e.py b/tests/e2e/mcp/test_mcp_datadog_e2e.py index 8a539b86bff..d093e307f99 100644 --- a/tests/e2e/mcp/test_mcp_datadog_e2e.py +++ b/tests/e2e/mcp/test_mcp_datadog_e2e.py @@ -60,6 +60,7 @@ class TestDatadogMcpRoundTrip: _assert_datadog_logger_active(client.proxy) server_id = register_datadog_mcp(client, resources) + client.await_registered(server_id) marker = f"{MARKER_PREFIX}{unique_marker()}" key = client.generate_key( @@ -78,22 +79,19 @@ class TestDatadogMcpRoundTrip: ) tool_name = client.await_tool(key, server_id, SEARCH_LOGS_TOOL) - - call = unwrap( - client.call_tool( - key, - server_id=server_id, - name=tool_name, - arguments={ - "query": marker, - "from": DD_SEARCH_FROM, - "to": "now", - "max_tokens": 5000, - "telemetry": { - "intent": "e2e assert seeded litellm completion log is searchable via MCP" - }, + call = client.await_call_tool( + key, + server_id=server_id, + name=tool_name, + arguments={ + "query": marker, + "from": DD_SEARCH_FROM, + "to": "now", + "max_tokens": 5000, + "telemetry": { + "intent": "e2e assert seeded litellm completion log is searchable via MCP" }, - ) + }, ) assert call.is_error is not True, f"search_datadog_logs errored: {call}" body = call.all_text diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index 35c864c07d8..678424e36d1 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -16,7 +16,7 @@ import pytest from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp from e2e_config import DD_SEARCH_FROM, unique_marker -from e2e_http import UnknownApiError, unwrap +from e2e_http import unwrap from lifecycle import ResourceManager from mcp_client import McpClient @@ -72,13 +72,12 @@ class TestMcpKeyWithoutAccessIsDenied: "max_tokens": 1000, "telemetry": {"intent": "e2e control call proving granted key can invoke Datadog MCP"}, } - permitted_call = unwrap( - client.call_tool(permitted_key, server_id=server_id, name=tool_name, arguments=search_args) + permitted_call = client.await_call_tool( + permitted_key, server_id=server_id, name=tool_name, arguments=search_args ) assert permitted_call.is_error is not True, f"granted key's tool call errored: {permitted_call}" - match client.call_tool(denied_key, server_id=server_id, name=tool_name, arguments=search_args): - case UnknownApiError(status_code=403, body=body): - assert "access_denied" in body, f"403 was not an MCP access denial: {body}" - case other: - pytest.fail(f"ungranted key's tool call was not refused with 403 access_denied: {other}") + denied = client.await_call_tool_denied( + denied_key, server_id=server_id, name=tool_name, arguments=search_args + ) + assert "access_denied" in denied.body, f"403 was not an MCP access denial: {denied.body}" From 6cfcb6cd839c1d23c7a59f247b1900346b9b2cab Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 14:08:18 +0000 Subject: [PATCH 013/576] fix(vertex_ai): translate /v1/embeddings batch rows to Gemini embedding shape Vertex batch files sent every jsonl line through the generateContent transform, so embeddings rows went out as {"request": {"contents": [...]}} and Vertex rejected each one with "no such field: 'contents'"; the OpenAI "input" was dropped along the way too. Route lines by their own url: embeddings lines now emit the EmbedContentRequest shape (singular content, embed_content_config sibling, custom_id round-tripping through the top-level key), and matching output rows come back as OpenAI embeddings responses. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 253 +++++++++++++--- .../test_vertex_ai_files_transformation.py | 277 ++++++++++++++++++ 2 files changed, 491 insertions(+), 39 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index dd877b52eb8..516a3ca7184 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -5,6 +5,7 @@ import json import os import re import time +from collections.abc import Mapping from typing import ( Any, Callable, @@ -16,6 +17,7 @@ from typing import ( Tuple, Union, ) +from urllib.parse import unquote import httpx from httpx import Headers, Response @@ -51,6 +53,9 @@ from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) +from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( + transform_openai_input_gemini_embed_content, +) from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, @@ -62,13 +67,26 @@ from litellm.types.llms.openai import ( ) from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.vertex_ai import GcsBucketResponse -from litellm.types.utils import LlmProviders, ModelResponse +from litellm.types.utils import ( + Embedding, + EmbeddingResponse, + LlmProviders, + ModelResponse, + Usage, +) from ..common_utils import VertexAIError from ..vertex_llm_base import VertexBase _GCP_LABEL_VALUE_MAX_LEN = 63 _CUSTOM_ID_RAW_LABEL_PREFIX = "b32_" +_VERTEX_BATCH_KEY_FIELD = "key" +_MANAGED_GCS_MODEL_PATH_PATTERN = re.compile(r"publishers/[^/]+/models/([^/?]+)") +_EMBED_CONTENT_CONFIG_FIELD_BY_GEMINI_PARAM = { + "outputDimensionality": "output_dimensionality", + "taskType": "task_type", + "title": "title", +} def _sanitize_gcp_label_value(value: str) -> str: @@ -131,6 +149,21 @@ def _set_litellm_batch_custom_id_labels(labels: Dict[str, str], custom_id: Any) labels[f"litellm_custom_id_raw_{index}"] = raw_label_chunk +def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, Any]) -> str: + """ + Resolve the OpenAI `custom_id` for a Vertex batch output row. + + Embedding rows carry it in the top-level `key` field that Vertex echoes back; + `generateContent` rows have no such field, so it is smuggled through request + labels instead (see `_set_litellm_batch_custom_id_labels`). + """ + key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) + if key is not None: + return str(key) + request_data = vertex_output_row.get("request") or {} + return _get_litellm_batch_custom_id_from_labels(request_data.get("labels") or {}) + + def _get_litellm_batch_custom_id_from_labels(labels: Dict[str, Any]) -> str: """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" raw = labels.get("litellm_custom_id_raw") @@ -149,10 +182,156 @@ def _get_litellm_batch_custom_id_from_labels(labels: Dict[str, Any]) -> str: return str(labels.get("litellm_custom_id", "unknown")) +def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, Any]) -> bool: + """ + Whether a Vertex batch output row came from an `EmbedContentRequest`. + + Successful rows hold the vector under `response.embedding.values`; failed rows only + carry `status`, so they are recognized from the singular `content` that the + embeddings request shape echoes back. + """ + if "request" not in vertex_output_row: + return False + response = vertex_output_row.get("response") + if isinstance(response, dict) and isinstance(response.get("embedding"), dict): + return True + request_data = vertex_output_row.get("request") + return bool(vertex_output_row.get("status")) and isinstance(request_data, dict) and "content" in request_data + + +def _openai_batch_output_row( + custom_id: str, + body: Mapping[str, Any] | None = None, + error: Mapping[str, str] | None = None, +) -> Mapping[str, Any]: + """ + One row of an OpenAI batch output file. Per the OpenAI Batch spec, failed rows set + `response` to null and populate `error` instead. + """ + return { + "id": f"batch_req_{uuid.uuid4()}", + "custom_id": custom_id, + "response": None + if body is None + else { + "status_code": 200, + "request_id": body.get("id", ""), + "body": body, + }, + "error": error, + } + + +def _transform_vertex_embeddings_batch_output_row_to_openai( + vertex_output_row: Mapping[str, Any], + model: str | None, +) -> Mapping[str, Any]: + """ + Transforms one Vertex Gemini Embedding batch output row into an OpenAI batch + output row holding an `/v1/embeddings` response body. + + Example Vertex jsonl + {"key": "id_1", "request": {...}, "response": {"tokenCount": "2", "embedding": {"values": [-0.015, 0.024]}}} + + `tokenCount` is serialized as a string by Vertex (int64 proto field), and the row + carries no `modelVersion`, so the model comes from the batch the row belongs to. + """ + custom_id = _get_litellm_batch_custom_id(vertex_output_row) + status = vertex_output_row.get("status", "") + if status: + return _openai_batch_output_row( + custom_id=custom_id, + error={"code": "vertex_ai_error", "message": status}, + ) + + vertex_response = vertex_output_row.get("response") or {} + token_count = int(vertex_response.get("tokenCount") or 0) + body = EmbeddingResponse( + model=model or "", + data=[ + Embedding( + embedding=vertex_response["embedding"]["values"], + index=0, + object="embedding", + ) + ], + usage=Usage(prompt_tokens=token_count, total_tokens=token_count), + ).model_dump() + return _openai_batch_output_row(custom_id=custom_id, body=body) + + +def _model_from_managed_gcs_url(url: str) -> str | None: + """ + Extracts the model from a LiteLLM-managed Vertex batch GCS url. + + Batch inputs and their sibling outputs are stored under + `.../publishers/google/models//...`, which is the only place the model of an + embeddings batch output row can be recovered from; unlike `generateContent` + responses, embedding rows carry no `modelVersion`. + """ + match = _MANAGED_GCS_MODEL_PATH_PATTERN.search(unquote(url)) + return match.group(1) if match else None + + +def _is_embeddings_batch_entry(openai_entry: Mapping[str, Any]) -> bool: + """ + Whether an OpenAI batch JSONL line targets the embeddings endpoint. + + OpenAI puts the target route on each line's `url` (e.g. `/v1/embeddings`); Vertex + has no equivalent per-line field, so the route decides which Vertex request shape + the line has to be translated into. + """ + url = openai_entry.get("url") + if not isinstance(url, str): + return False + path = url.split("?")[0].rstrip("/") + return path == "embeddings" or path.endswith("/embeddings") + + +def _openai_batch_jsonl_entry_to_vertex_embeddings_row( + openai_entry: Mapping[str, Any], +) -> Mapping[str, Any]: + """ + Transforms a single OpenAI `/v1/embeddings` batch entry into a Vertex Gemini + Embedding batch row. + + Example Vertex jsonl + {"key": "id_1", "request": {"content": {"parts": [{"text": "Hello World"}]}}, "embed_content_config": {"output_dimensionality": 768, "task_type": "RETRIEVAL_DOCUMENT"}} + + Note that `content` is singular (an `EmbedContentRequest`, not a + `GenerateContentRequest`), the per-row config is a sibling of `request` rather than + part of it, and the `custom_id` round-trips through the top-level `key`. + + API Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings + """ + openai_request_body = openai_entry.get("body") or {} + embedding_input = openai_request_body.get("input") + if embedding_input is None: + raise ValueError("`input` is required on /v1/embeddings batch requests, but was not provided") + + embed_content_request = transform_openai_input_gemini_embed_content( + input=embedding_input, + model=openai_request_body.get("model", ""), + optional_params=openai_request_body, + ) + embed_content_config = { + config_field: embed_content_request[gemini_param] + for gemini_param, config_field in _EMBED_CONTENT_CONFIG_FIELD_BY_GEMINI_PARAM.items() + if gemini_param in embed_content_request + } + + custom_id = openai_entry.get("custom_id") + return { + **({_VERTEX_BATCH_KEY_FIELD: str(custom_id)} if custom_id is not None else {}), + "request": {"content": embed_content_request["content"]}, + **({"embed_content_config": embed_content_config} if embed_content_config else {}), + } + + def _openai_batch_jsonl_entry_to_vertex_wrapped_request( openai_entry: Dict[str, Any], map_openai_to_vertex_params: Callable[[Dict[str, Any]], Dict[str, Any]], -) -> Dict[str, Any]: +) -> Mapping[str, Any]: """ Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request. @@ -160,6 +339,9 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request( Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} """ + if _is_embeddings_batch_entry(openai_entry): + return _openai_batch_jsonl_entry_to_vertex_embeddings_row(openai_entry) + openai_request_body = openai_entry.get("body") or {} vertex_request_body = _transform_request_body( messages=openai_request_body.get("messages", []), @@ -629,6 +811,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): transformed_content = self._try_transform_vertex_batch_output_to_openai( content=content, logging_obj=logging_obj, + model=_model_from_managed_gcs_url(str(raw_response.request.url)), ) if transformed_content != content: # Create a new response with transformed content and updated Content-Length @@ -650,7 +833,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): return HttpxBinaryResponseContent(response=raw_response) def _try_transform_vertex_batch_output_to_openai( - self, content: bytes, logging_obj: Optional[LiteLLMLoggingObj] = None + self, + content: bytes, + logging_obj: LiteLLMLoggingObj | None = None, + model: str | None = None, ) -> bytes: """ Try to transform Vertex AI batch output to OpenAI format. @@ -692,7 +878,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # first line is not valid UTF-8/JSON) raises and falls through to the # passthrough below, leaving the content untouched. first_row = json.loads(first_line) - is_vertex_batch_output = ( + is_vertex_batch_output = _is_vertex_embeddings_batch_output_row(first_row) or ( "request" in first_row and "response" in first_row and "processed_time" in first_row @@ -731,11 +917,19 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): output = bytearray() for line in itertools.chain([first_line], lines): try: - openai_output = self._transform_single_vertex_batch_output_to_openai( - vertex_output=json.loads(line), - vertex_gemini_config=vertex_gemini_config, - logging_obj=batch_transform_logging_obj, - mock_httpx_response=mock_httpx_response, + vertex_output_row = json.loads(line) + openai_output = ( + _transform_vertex_embeddings_batch_output_row_to_openai( + vertex_output_row=vertex_output_row, + model=model, + ) + if _is_vertex_embeddings_batch_output_row(vertex_output_row) + else self._transform_single_vertex_batch_output_to_openai( + vertex_output=vertex_output_row, + vertex_gemini_config=vertex_gemini_config, + logging_obj=batch_transform_logging_obj, + mock_httpx_response=mock_httpx_response, + ) ) except Exception: return content @@ -755,30 +949,22 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): vertex_gemini_config: VertexGeminiConfig, logging_obj: Logging, mock_httpx_response: httpx.Response, - ) -> Dict[str, Any]: + ) -> Mapping[str, Any]: """ Transform a single Vertex AI batch output line to OpenAI format. Uses the existing VertexGeminiConfig transformation for the response. """ - # Extract custom_id from request labels (prefer raw for OpenAI round-trip) - request_data = vertex_output.get("request", {}) - labels = request_data.get("labels", {}) or {} - custom_id = _get_litellm_batch_custom_id_from_labels(labels) + custom_id = _get_litellm_batch_custom_id(vertex_output) # Check if there's an error status = vertex_output.get("status", "") has_error = bool(status) if has_error: - return { - "id": f"batch_req_{uuid.uuid4()}", - "custom_id": custom_id, - "response": None, - "error": { - "code": "vertex_ai_error", - "message": status, - }, - } + return _openai_batch_output_row( + custom_id=custom_id, + error={"code": "vertex_ai_error", "message": status}, + ) # Transform successful response using existing transformation vertex_response = vertex_output.get("response", {}) @@ -804,24 +990,13 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): response_dict = transformed_response.model_dump() # Return in OpenAI batch format - return { - "id": f"batch_req_{uuid.uuid4()}", - "custom_id": custom_id, - "response": { - "status_code": 200, - "request_id": response_dict.get("id", ""), - "body": response_dict, - }, - "error": None, - } + return _openai_batch_output_row(custom_id=custom_id, body=response_dict) except Exception as e: - return { - "id": f"batch_req_{uuid.uuid4()}", - "custom_id": custom_id, - "response": None, - "error": { + return _openai_batch_output_row( + custom_id=custom_id, + error={ "code": "transformation_error", "message": f"Failed to transform response: {str(e)}", }, - } + ) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 8c5305ee67b..636a3106617 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -1318,3 +1318,280 @@ class TestConfiguredBucketNameResolution: assert "bucket_name" in OPTIONAL_KWARGS_KEYS params = get_litellm_params(bucket_name="my-legacy-bucket") assert params.get("bucket_name") == "my-legacy-bucket" + + +def _embeddings_entry(**overrides): + entry = { + "custom_id": "request-1", + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": "gemini-embedding-2", "input": "hello world"}, + } + entry.update(overrides) + return entry + + +class TestVertexEmbeddingsBatchInputTranslation: + """ + /v1/embeddings batch lines must be translated to Vertex's Gemini Embedding batch + shape, not the generateContent shape. + + Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings + """ + + def test_should_emit_embed_content_request_shape(self): + (row,) = _wrap_entries([_embeddings_entry()]) + + assert row["request"] == {"content": {"parts": [{"text": "hello world"}]}} + assert "contents" not in row["request"] + assert "labels" not in row["request"] + + def test_should_round_trip_custom_id_through_top_level_key(self): + (row,) = _wrap_entries([_embeddings_entry(custom_id="MyRequest-1")]) + + assert row["key"] == "MyRequest-1" + + def test_should_omit_key_when_no_custom_id(self): + entry = _embeddings_entry() + del entry["custom_id"] + + (row,) = _wrap_entries([entry]) + + assert "key" not in row + + def test_should_map_openai_params_to_embed_content_config_sibling(self): + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-001", + "input": "hello world", + "dimensions": 768, + "task_type": "RETRIEVAL_DOCUMENT", + "title": "some_title", + } + ) + ] + ) + + assert row["embed_content_config"] == { + "output_dimensionality": 768, + "task_type": "RETRIEVAL_DOCUMENT", + "title": "some_title", + } + assert "output_dimensionality" not in row["request"] + + def test_should_omit_embed_content_config_when_no_params_given(self): + (row,) = _wrap_entries([_embeddings_entry()]) + + assert "embed_content_config" not in row + + def test_should_translate_multimodal_gcs_input(self): + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-2", + "input": "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + } + ) + ] + ) + + assert row["request"]["content"]["parts"] == [ + { + "file_data": { + "mime_type": "image/jpeg", + "file_uri": "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + } + } + ] + + @pytest.mark.parametrize("url", ["/v1/embeddings", "v1/embeddings", "/v1/embeddings/"]) + def test_should_detect_embeddings_route_variants(self, url): + (row,) = _wrap_entries([_embeddings_entry(url=url)]) + + assert "content" in row["request"] + + def test_should_raise_when_input_missing(self): + with pytest.raises(ValueError, match="`input` is required"): + _wrap_entries([_embeddings_entry(body={"model": "gemini-embedding-2"})]) + + def test_should_keep_chat_completions_lines_on_generate_content_path(self): + (row,) = _wrap_entries( + [ + { + "custom_id": "request-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.0-flash-001", + "messages": [{"role": "user", "content": "Hello"}], + }, + } + ] + ) + + assert row["request"]["contents"] == [ + {"role": "user", "parts": [{"text": "Hello"}]} + ] + assert row["request"]["labels"]["litellm_custom_id"] == "request-1" + assert "key" not in row + + def test_should_translate_each_line_by_its_own_url(self): + chat_row, embeddings_row = _wrap_entries( + [ + { + "custom_id": "chat-1", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.0-flash-001", + "messages": [{"role": "user", "content": "Hello"}], + }, + }, + _embeddings_entry(custom_id="embed-1"), + ] + ) + + assert "contents" in chat_row["request"] + assert "content" in embeddings_row["request"] + + +class TestVertexEmbeddingsBatchOutputTranslation: + """Vertex Gemini Embedding batch output rows must come back as OpenAI batch rows.""" + + def _vertex_embeddings_output_row(self, **overrides): + row = { + "key": "request-1", + "request": {"content": {"parts": [{"text": "hello world"}]}}, + "response": { + "tokenCount": "2", + "embedding": {"values": [-0.015, 0.024]}, + }, + } + row.update(overrides) + return row + + def _transform(self, config, rows, url="https://example.com"): + content = "\n".join(json.dumps(row) for row in rows).encode("utf-8") + result = config.transform_file_content_response( + raw_response=httpx.Response( + status_code=200, + content=content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request("GET", url), + ), + logging_obj=MagicMock(), + litellm_params={}, + ) + return [ + json.loads(line) + for line in result.response.content.decode("utf-8").split("\n") + ] + + def test_should_transform_embeddings_output_to_openai_batch_row(self, config): + (result,) = self._transform(config, [self._vertex_embeddings_output_row()]) + + assert result["custom_id"] == "request-1" + assert result["error"] is None + assert result["response"]["status_code"] == 200 + body = result["response"]["body"] + assert body["object"] == "list" + assert body["data"] == [ + {"embedding": [-0.015, 0.024], "index": 0, "object": "embedding"} + ] + assert body["usage"]["prompt_tokens"] == 2 + assert body["usage"]["total_tokens"] == 2 + + def test_should_resolve_model_from_managed_gcs_object_path(self, config): + object_path = urllib.parse.quote( + "litellm-vertex-files/publishers/google/models/gemini-embedding-2/" + "prediction-model-2026-07-29T05:55:52Z/predictions.jsonl", + safe="", + ) + url = f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{object_path}?alt=media" + + (result,) = self._transform( + config, [self._vertex_embeddings_output_row()], url=url + ) + + assert result["response"]["body"]["model"] == "gemini-embedding-2" + + def test_should_surface_failed_embeddings_row_as_error(self, config): + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row( + status="Failed to parse JSON into proto", response={} + ) + ], + ) + + assert result["custom_id"] == "request-1" + assert result["response"] is None + assert result["error"]["code"] == "vertex_ai_error" + assert "Failed to parse JSON into proto" in result["error"]["message"] + + def test_should_transform_every_row_of_a_multi_row_file(self, config): + results = self._transform( + config, + [ + self._vertex_embeddings_output_row(key=f"request-{index}") + for index in range(3) + ], + ) + + assert [result["custom_id"] for result in results] == [ + "request-0", + "request-1", + "request-2", + ] + + def test_should_end_to_end_round_trip_openai_embeddings_batch(self, config): + (vertex_row,) = _wrap_entries( + [ + _embeddings_entry( + custom_id="MyRequest-1", + body={ + "model": "gemini-embedding-2", + "input": "hello world", + "dimensions": 2, + }, + ) + ] + ) + + (result,) = self._transform( + config, + [ + { + **vertex_row, + "status": "", + "processed_time": "2026-07-29T05:55:52.379528Z", + "response": { + "tokenCount": "2", + "embedding": {"values": [-0.015, 0.024]}, + }, + } + ], + ) + + assert result["custom_id"] == "MyRequest-1" + assert result["response"]["body"]["data"][0]["embedding"] == [-0.015, 0.024] + + def test_should_leave_legacy_predict_embeddings_output_untouched(self, config): + legacy_row = { + "instance": {"content": "hello world"}, + "predictions": [ + { + "embeddings": { + "statistics": {"token_count": 2, "truncated": False}, + "values": [0.2], + } + } + ], + "status": "", + } + content = json.dumps(legacy_row).encode("utf-8") + + assert config._try_transform_vertex_batch_output_to_openai(content) == content From e96614a39f45d8b6ffde5a3e2a05e63f6d73541d Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 14:44:25 +0000 Subject: [PATCH 014/576] fix(vertex_ai): put embed config inside the request and read live usage A live Vertex batch run showed the documented "embed_content_config" sibling of "request" is rejected by the API ("unsupported type"), failing the whole job rather than the row; the same fields inside the EmbedContentRequest succeed and honor output_dimensionality. Real output rows also report usage under response.usageMetadata.promptTokenCount, not the documented response.tokenCount, so every row came back with zero tokens. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 29 +++++++------ .../test_vertex_ai_files_transformation.py | 42 ++++++++++++++----- 2 files changed, 48 insertions(+), 23 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 516a3ca7184..1a5d40dabb2 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -82,7 +82,7 @@ _GCP_LABEL_VALUE_MAX_LEN = 63 _CUSTOM_ID_RAW_LABEL_PREFIX = "b32_" _VERTEX_BATCH_KEY_FIELD = "key" _MANAGED_GCS_MODEL_PATH_PATTERN = re.compile(r"publishers/[^/]+/models/([^/?]+)") -_EMBED_CONTENT_CONFIG_FIELD_BY_GEMINI_PARAM = { +_EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = { "outputDimensionality": "output_dimensionality", "taskType": "task_type", "title": "title", @@ -231,10 +231,11 @@ def _transform_vertex_embeddings_batch_output_row_to_openai( output row holding an `/v1/embeddings` response body. Example Vertex jsonl - {"key": "id_1", "request": {...}, "response": {"tokenCount": "2", "embedding": {"values": [-0.015, 0.024]}}} + {"key": "id_1", "request": {...}, "response": {"embedding": {"values": [-0.015, 0.024]}, "usageMetadata": {"promptTokenCount": 2}}} - `tokenCount` is serialized as a string by Vertex (int64 proto field), and the row - carries no `modelVersion`, so the model comes from the batch the row belongs to. + Live rows report usage under `usageMetadata`; the documented `tokenCount` is kept as + a fallback. The row carries no `modelVersion`, so the model comes from the batch it + belongs to. """ custom_id = _get_litellm_batch_custom_id(vertex_output_row) status = vertex_output_row.get("status", "") @@ -245,7 +246,8 @@ def _transform_vertex_embeddings_batch_output_row_to_openai( ) vertex_response = vertex_output_row.get("response") or {} - token_count = int(vertex_response.get("tokenCount") or 0) + usage_metadata = vertex_response.get("usageMetadata") or {} + token_count = int(usage_metadata.get("promptTokenCount") or vertex_response.get("tokenCount") or 0) body = EmbeddingResponse( model=model or "", data=[ @@ -296,11 +298,13 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_row( Embedding batch row. Example Vertex jsonl - {"key": "id_1", "request": {"content": {"parts": [{"text": "Hello World"}]}}, "embed_content_config": {"output_dimensionality": 768, "task_type": "RETRIEVAL_DOCUMENT"}} + {"key": "id_1", "request": {"content": {"parts": [{"text": "Hello World"}]}, "output_dimensionality": 768, "task_type": "RETRIEVAL_DOCUMENT"}} Note that `content` is singular (an `EmbedContentRequest`, not a - `GenerateContentRequest`), the per-row config is a sibling of `request` rather than - part of it, and the `custom_id` round-trips through the top-level `key`. + `GenerateContentRequest`) and that the `custom_id` round-trips through the top-level + `key`. The docs put the per-row config in an `embed_content_config` sibling of + `request`, but the API rejects that key outright and fails the whole batch job, so + the config fields go inside the `EmbedContentRequest` itself. API Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings """ @@ -314,17 +318,16 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_row( model=openai_request_body.get("model", ""), optional_params=openai_request_body, ) - embed_content_config = { - config_field: embed_content_request[gemini_param] - for gemini_param, config_field in _EMBED_CONTENT_CONFIG_FIELD_BY_GEMINI_PARAM.items() + embed_request_fields = { + request_field: embed_content_request[gemini_param] + for gemini_param, request_field in _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM.items() if gemini_param in embed_content_request } custom_id = openai_entry.get("custom_id") return { **({_VERTEX_BATCH_KEY_FIELD: str(custom_id)} if custom_id is not None else {}), - "request": {"content": embed_content_request["content"]}, - **({"embed_content_config": embed_content_config} if embed_content_config else {}), + "request": {"content": embed_content_request["content"], **embed_request_fields}, } diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 636a3106617..3eff3083220 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -1359,7 +1359,11 @@ class TestVertexEmbeddingsBatchInputTranslation: assert "key" not in row - def test_should_map_openai_params_to_embed_content_config_sibling(self): + def test_should_map_openai_params_into_the_embed_content_request(self): + """ + The docs put these in an `embed_content_config` sibling of `request`, but Vertex + rejects that key and fails the whole job, so they belong inside the request. + """ (row,) = _wrap_entries( [ _embeddings_entry( @@ -1374,17 +1378,20 @@ class TestVertexEmbeddingsBatchInputTranslation: ] ) - assert row["embed_content_config"] == { - "output_dimensionality": 768, - "task_type": "RETRIEVAL_DOCUMENT", - "title": "some_title", + assert row == { + "key": "request-1", + "request": { + "content": {"parts": [{"text": "hello world"}]}, + "output_dimensionality": 768, + "task_type": "RETRIEVAL_DOCUMENT", + "title": "some_title", + }, } - assert "output_dimensionality" not in row["request"] - def test_should_omit_embed_content_config_when_no_params_given(self): + def test_should_omit_config_fields_when_no_params_given(self): (row,) = _wrap_entries([_embeddings_entry()]) - assert "embed_content_config" not in row + assert set(row["request"]) == {"content"} def test_should_translate_multimodal_gcs_input(self): (row,) = _wrap_entries( @@ -1465,8 +1472,8 @@ class TestVertexEmbeddingsBatchOutputTranslation: "key": "request-1", "request": {"content": {"parts": [{"text": "hello world"}]}}, "response": { - "tokenCount": "2", "embedding": {"values": [-0.015, 0.024]}, + "usageMetadata": {"promptTokenCount": 2}, }, } row.update(overrides) @@ -1503,6 +1510,21 @@ class TestVertexEmbeddingsBatchOutputTranslation: assert body["usage"]["prompt_tokens"] == 2 assert body["usage"]["total_tokens"] == 2 + def test_should_fall_back_to_documented_token_count_field(self, config): + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row( + response={ + "embedding": {"values": [-0.015, 0.024]}, + "tokenCount": "2", + } + ) + ], + ) + + assert result["response"]["body"]["usage"]["prompt_tokens"] == 2 + def test_should_resolve_model_from_managed_gcs_object_path(self, config): object_path = urllib.parse.quote( "litellm-vertex-files/publishers/google/models/gemini-embedding-2/" @@ -1569,8 +1591,8 @@ class TestVertexEmbeddingsBatchOutputTranslation: "status": "", "processed_time": "2026-07-29T05:55:52.379528Z", "response": { - "tokenCount": "2", "embedding": {"values": [-0.015, 0.024]}, + "usageMetadata": {"promptTokenCount": 2}, }, } ], From 627d2755da56d01a6f5838bec967005027b2f35d Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 15:06:58 +0000 Subject: [PATCH 015/576] fix(vertex_ai): fan array embeddings input out into one vertex row per element Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 221 +++++++++++++----- .../files/test_vertex_ai_files_streaming.py | 5 +- .../test_vertex_ai_files_transformation.py | 180 +++++++++++++- 3 files changed, 339 insertions(+), 67 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 1a5d40dabb2..78a16b002e7 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -66,7 +66,7 @@ from litellm.types.llms.openai import ( PathLike, ) from litellm.types.files import StreamingMediaUploadConfig -from litellm.types.llms.vertex_ai import GcsBucketResponse +from litellm.types.llms.vertex_ai import GcsBucketResponse, GeminiEmbeddingInput from litellm.types.utils import ( Embedding, EmbeddingResponse, @@ -87,6 +87,7 @@ _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = { "taskType": "task_type", "title": "title", } +_VERTEX_BATCH_FANNED_OUT_KEY_PATTERN = re.compile(r"(?P.*)#(?P\d+)/(?P\d+)") def _sanitize_gcp_label_value(value: str) -> str: @@ -222,46 +223,93 @@ def _openai_batch_output_row( } -def _transform_vertex_embeddings_batch_output_row_to_openai( - vertex_output_row: Mapping[str, Any], +def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, int]: + """ + Resolve `(custom_id, index within that custom_id)` for a Vertex batch output row. + + A `/v1/embeddings` entry whose `input` is an array fans out into one Vertex row per + element, tagged `#/` (see `_vertex_batch_embeddings_key`), + so the rows can be reassembled into a single OpenAI response. + """ + key = _get_litellm_batch_custom_id(vertex_output_row) + match = _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN.fullmatch(key) + if match is None or int(match["total"]) < 2: + return key, 0 + return match["custom_id"], int(match["index"]) + + +def _vertex_embeddings_rows_to_openai_batch_output_row( + custom_id: str, + vertex_output_rows: tuple[Mapping[str, Any], ...], model: str | None, ) -> Mapping[str, Any]: """ - Transforms one Vertex Gemini Embedding batch output row into an OpenAI batch - output row holding an `/v1/embeddings` response body. + Transforms the Vertex Gemini Embedding batch output rows belonging to one OpenAI + batch entry into an OpenAI batch output row holding an `/v1/embeddings` response. Example Vertex jsonl {"key": "id_1", "request": {...}, "response": {"embedding": {"values": [-0.015, 0.024]}, "usageMetadata": {"promptTokenCount": 2}}} - Live rows report usage under `usageMetadata`; the documented `tokenCount` is kept as - a fallback. The row carries no `modelVersion`, so the model comes from the batch it - belongs to. + An entry that asked for several embeddings at once maps to several rows here, which + become the indexed elements of a single `data` array. One failed element fails the + whole entry, since an OpenAI batch row is either a response or an error. Live rows + report usage under `usageMetadata`; the documented `tokenCount` is kept as a + fallback. Rows carry no `modelVersion`, so the model comes from the batch they + belong to. """ - custom_id = _get_litellm_batch_custom_id(vertex_output_row) - status = vertex_output_row.get("status", "") + status = next((row["status"] for row in vertex_output_rows if row.get("status")), "") if status: return _openai_batch_output_row( custom_id=custom_id, error={"code": "vertex_ai_error", "message": status}, ) - vertex_response = vertex_output_row.get("response") or {} - usage_metadata = vertex_response.get("usageMetadata") or {} - token_count = int(usage_metadata.get("promptTokenCount") or vertex_response.get("tokenCount") or 0) + responses = tuple(row.get("response") or {} for row in vertex_output_rows) + token_count = sum( + int((response.get("usageMetadata") or {}).get("promptTokenCount") or response.get("tokenCount") or 0) + for response in responses + ) body = EmbeddingResponse( model=model or "", data=[ Embedding( - embedding=vertex_response["embedding"]["values"], - index=0, + embedding=response["embedding"]["values"], + index=index, object="embedding", ) + for index, response in enumerate(responses) ], usage=Usage(prompt_tokens=token_count, total_tokens=token_count), ).model_dump() return _openai_batch_output_row(custom_id=custom_id, body=body) +def _transform_vertex_embeddings_batch_output_to_openai( + vertex_output_rows: Iterable[Mapping[str, Any]], + model: str | None, +) -> tuple[Mapping[str, Any], ...]: + """ + Transforms a whole Vertex Gemini Embedding batch output into OpenAI batch output + rows, one per OpenAI batch entry, in the order the entries first appear. + + Rows are grouped rather than mapped one to one because a single entry can fan out + into several Vertex rows, and Vertex returns them in arbitrary order. + """ + keyed_rows = tuple((_split_vertex_batch_key(row), row) for row in vertex_output_rows) + grouped_rows = { + custom_id: tuple(row for _, row in group) + for custom_id, group in itertools.groupby(sorted(keyed_rows, key=lambda kr: kr[0]), key=lambda kr: kr[0][0]) + } + return tuple( + _vertex_embeddings_rows_to_openai_batch_output_row( + custom_id=custom_id, + vertex_output_rows=grouped_rows[custom_id], + model=model, + ) + for custom_id in dict.fromkeys(custom_id for (custom_id, _), _ in keyed_rows) + ) + + def _model_from_managed_gcs_url(url: str) -> str | None: """ Extracts the model from a LiteLLM-managed Vertex batch GCS url. @@ -290,21 +338,49 @@ def _is_embeddings_batch_entry(openai_entry: Mapping[str, Any]) -> bool: return path == "embeddings" or path.endswith("/embeddings") -def _openai_batch_jsonl_entry_to_vertex_embeddings_row( - openai_entry: Mapping[str, Any], -) -> Mapping[str, Any]: +def _openai_embedding_input_elements( + embedding_input: GeminiEmbeddingInput, +) -> tuple[Union[str, List[str]], ...]: """ - Transforms a single OpenAI `/v1/embeddings` batch entry into a Vertex Gemini - Embedding batch row. + Split an OpenAI `input` into the elements that each get their own embedding. + + A string is one embedding, a flat array is one embedding per element, and a nested + array is one combined embedding per inner array, matching the online + `batchEmbedContents` path. + """ + if isinstance(embedding_input, list): + return tuple(embedding_input) + return (embedding_input,) + + +def _vertex_batch_embeddings_key(custom_id: str, index: int, total: int) -> str: + """ + The top-level `key` Vertex echoes back on an embeddings row. + + An entry asking for several embeddings needs several Vertex rows, so its key also + carries the element index and the group size; `_split_vertex_batch_key` reads them + back out. Entries asking for a single embedding keep their bare `custom_id`. + """ + return custom_id if total < 2 else f"{custom_id}#{index}/{total}" + + +def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( + openai_entry: Mapping[str, Any], +) -> tuple[Mapping[str, Any], ...]: + """ + Transforms a single OpenAI `/v1/embeddings` batch entry into Vertex Gemini Embedding + batch rows, one per requested embedding. Example Vertex jsonl {"key": "id_1", "request": {"content": {"parts": [{"text": "Hello World"}]}, "output_dimensionality": 768, "task_type": "RETRIEVAL_DOCUMENT"}} Note that `content` is singular (an `EmbedContentRequest`, not a `GenerateContentRequest`) and that the `custom_id` round-trips through the top-level - `key`. The docs put the per-row config in an `embed_content_config` sibling of - `request`, but the API rejects that key outright and fails the whole batch job, so - the config fields go inside the `EmbedContentRequest` itself. + `key`. An `EmbedContentRequest` returns exactly one vector, so an entry whose `input` + is an array fans out into one row per element and is reassembled on the way back. + The docs put the per-row config in an `embed_content_config` sibling of `request`, + but the API rejects that key outright and fails the whole batch job, so the config + fields go inside the `EmbedContentRequest` itself. API Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings """ @@ -313,37 +389,58 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_row( if embedding_input is None: raise ValueError("`input` is required on /v1/embeddings batch requests, but was not provided") - embed_content_request = transform_openai_input_gemini_embed_content( - input=embedding_input, - model=openai_request_body.get("model", ""), - optional_params=openai_request_body, + elements = _openai_embedding_input_elements(embedding_input) + if not elements: + raise ValueError("`input` on /v1/embeddings batch requests must not be empty") + + embed_content_requests = tuple( + transform_openai_input_gemini_embed_content( + input=element, + model=openai_request_body.get("model", ""), + optional_params=openai_request_body, + ) + for element in elements ) - embed_request_fields = { - request_field: embed_content_request[gemini_param] - for gemini_param, request_field in _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM.items() - if gemini_param in embed_content_request - } - custom_id = openai_entry.get("custom_id") - return { - **({_VERTEX_BATCH_KEY_FIELD: str(custom_id)} if custom_id is not None else {}), - "request": {"content": embed_content_request["content"], **embed_request_fields}, - } + return tuple( + { + **( + {} + if custom_id is None + else { + _VERTEX_BATCH_KEY_FIELD: _vertex_batch_embeddings_key( + custom_id=str(custom_id), + index=index, + total=len(embed_content_requests), + ) + } + ), + "request": { + "content": embed_content_request["content"], + **{ + request_field: embed_content_request[gemini_param] + for gemini_param, request_field in _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM.items() + if gemini_param in embed_content_request + }, + }, + } + for index, embed_content_request in enumerate(embed_content_requests) + ) -def _openai_batch_jsonl_entry_to_vertex_wrapped_request( +def _openai_batch_jsonl_entry_to_vertex_rows( openai_entry: Dict[str, Any], map_openai_to_vertex_params: Callable[[Dict[str, Any]], Dict[str, Any]], -) -> Mapping[str, Any]: +) -> tuple[Mapping[str, Any], ...]: """ - Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request. + Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to. jsonl body for vertex is {"request": } Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} """ if _is_embeddings_batch_entry(openai_entry): - return _openai_batch_jsonl_entry_to_vertex_embeddings_row(openai_entry) + return _openai_batch_jsonl_entry_to_vertex_embeddings_rows(openai_entry) openai_request_body = openai_entry.get("body") or {} vertex_request_body = _transform_request_body( @@ -361,7 +458,7 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request( vertex_request_body["labels"] = {} _set_litellm_batch_custom_id_labels(vertex_request_body["labels"], custom_id) - return {"request": vertex_request_body} + return ({"request": vertex_request_body},) def _iter_stripped_lines(raw_lines: Iterable[Union[str, bytes]]) -> Iterator[str]: @@ -459,10 +556,10 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): def _iter_vertex_jsonl_chunks(self) -> Iterator[bytes]: first = True for entry in _iter_openai_jsonl_entries(self._openai_file_content): - wrapped = _openai_batch_jsonl_entry_to_vertex_wrapped_request(entry, self._map_openai_to_vertex_params) - prefix = b"" if first else b"\n" - first = False - yield prefix + json.dumps(wrapped).encode("utf-8") + for wrapped in _openai_batch_jsonl_entry_to_vertex_rows(entry, self._map_openai_to_vertex_params): + prefix = b"" if first else b"\n" + first = False + yield prefix + json.dumps(wrapped).encode("utf-8") def iter_bytes(self) -> Iterator[bytes]: return self._iter_vertex_jsonl_chunks() @@ -914,25 +1011,29 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): request=httpx.Request(method="POST", url="https://example.com"), ) + all_lines = itertools.chain([first_line], lines) + + # Embedding rows are grouped by `custom_id` rather than transformed one at a + # time, since an entry that asked for several embeddings comes back as + # several rows, in arbitrary order. + if _is_vertex_embeddings_batch_output_row(first_row): + openai_outputs = _transform_vertex_embeddings_batch_output_to_openai( + vertex_output_rows=(json.loads(line) for line in all_lines), + model=model, + ) + return b"\n".join(json.dumps(openai_output).encode("utf-8") for openai_output in openai_outputs) + # Transform each row straight into the output buffer, so peak memory # stays at ~one row plus the output. If any row fails, return the # original content unchanged. output = bytearray() - for line in itertools.chain([first_line], lines): + for line in all_lines: try: - vertex_output_row = json.loads(line) - openai_output = ( - _transform_vertex_embeddings_batch_output_row_to_openai( - vertex_output_row=vertex_output_row, - model=model, - ) - if _is_vertex_embeddings_batch_output_row(vertex_output_row) - else self._transform_single_vertex_batch_output_to_openai( - vertex_output=vertex_output_row, - vertex_gemini_config=vertex_gemini_config, - logging_obj=batch_transform_logging_obj, - mock_httpx_response=mock_httpx_response, - ) + openai_output = self._transform_single_vertex_batch_output_to_openai( + vertex_output=json.loads(line), + vertex_gemini_config=vertex_gemini_config, + logging_obj=batch_transform_logging_obj, + mock_httpx_response=mock_httpx_response, ) except Exception: return content diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index 2e3280c0ed1..957fc7dbcf4 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -37,7 +37,7 @@ from litellm.llms.vertex_ai.files.transformation import ( _get_litellm_batch_custom_id_from_labels, _iter_openai_jsonl_entries, _iter_openai_jsonl_lines, - _openai_batch_jsonl_entry_to_vertex_wrapped_request, + _openai_batch_jsonl_entry_to_vertex_rows, ) from litellm.types.llms.openai import CreateFileRequest @@ -84,8 +84,9 @@ def _reference_vertex_jsonl_string(cfg: VertexAIFilesConfig, content: str) -> st transform, so the streaming path can be checked against it for parity.""" entries = [json.loads(line) for line in content.splitlines() if line.strip()] return "\n".join( - json.dumps(_openai_batch_jsonl_entry_to_vertex_wrapped_request(entry, cfg._map_openai_to_vertex_params)) + json.dumps(row) for entry in entries + for row in _openai_batch_jsonl_entry_to_vertex_rows(entry, cfg._map_openai_to_vertex_params) ) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 3eff3083220..73d1d6eeb5a 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -15,7 +15,7 @@ from unittest.mock import MagicMock from litellm.llms.vertex_ai.files.transformation import ( VertexAIFilesConfig, _get_litellm_batch_custom_id_from_labels, - _openai_batch_jsonl_entry_to_vertex_wrapped_request, + _openai_batch_jsonl_entry_to_vertex_rows, _sanitize_gcp_label_value, ) from litellm.types.llms.openai import OpenAIFileObject, HttpxBinaryResponseContent @@ -1054,14 +1054,15 @@ class TestTryTransformDoesNotMutateCallerLoggingObj: def _wrap_entries(openai_jsonl_content): - """Vertex-wrapped requests for a list of OpenAI batch entries, built via the - live single-entry transform that the streaming upload path uses.""" + """Vertex rows for a list of OpenAI batch entries, built via the live + single-entry transform that the streaming upload path uses.""" cfg = VertexAIFilesConfig() return [ - _openai_batch_jsonl_entry_to_vertex_wrapped_request( + row + for entry in openai_jsonl_content + for row in _openai_batch_jsonl_entry_to_vertex_rows( entry, cfg._map_openai_to_vertex_params ) - for entry in openai_jsonl_content ] @@ -1424,6 +1425,86 @@ class TestVertexEmbeddingsBatchInputTranslation: with pytest.raises(ValueError, match="`input` is required"): _wrap_entries([_embeddings_entry(body={"model": "gemini-embedding-2"})]) + def test_should_raise_when_input_empty(self): + with pytest.raises(ValueError, match="must not be empty"): + _wrap_entries( + [_embeddings_entry(body={"model": "gemini-embedding-2", "input": []})] + ) + + def test_should_fan_an_input_array_out_into_one_row_per_element(self): + """ + An `EmbedContentRequest` returns exactly one vector, so an OpenAI entry asking + for several embeddings needs several Vertex rows. + """ + rows = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-001", + "input": ["first", "second"], + "dimensions": 768, + } + ) + ] + ) + + assert rows == [ + { + "key": "request-1#0/2", + "request": { + "content": {"parts": [{"text": "first"}]}, + "output_dimensionality": 768, + }, + }, + { + "key": "request-1#1/2", + "request": { + "content": {"parts": [{"text": "second"}]}, + "output_dimensionality": 768, + }, + }, + ] + + def test_should_keep_the_bare_custom_id_for_single_element_arrays(self): + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={"model": "gemini-embedding-2", "input": ["only one"]} + ) + ] + ) + + assert row["key"] == "request-1" + + def test_should_combine_a_nested_input_into_one_multipart_row(self): + """Nested arrays are the combined-embedding shape, as on the online path.""" + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-2", + "input": [ + [ + "a caption", + "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + ] + ], + } + ) + ] + ) + + assert row["key"] == "request-1" + assert row["request"]["content"]["parts"] == [ + {"text": "a caption"}, + { + "file_data": { + "mime_type": "image/jpeg", + "file_uri": "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + } + }, + ] + def test_should_keep_chat_completions_lines_on_generate_content_path(self): (row,) = _wrap_entries( [ @@ -1569,6 +1650,95 @@ class TestVertexEmbeddingsBatchOutputTranslation: "request-2", ] + def test_should_reassemble_a_fanned_out_input_array_into_one_row(self, config): + """Vertex returns the rows of one entry in arbitrary order.""" + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row( + key="request-1#1/2", + response={ + "embedding": {"values": [0.3, 0.4]}, + "usageMetadata": {"promptTokenCount": 5}, + }, + ), + self._vertex_embeddings_output_row( + key="request-1#0/2", + response={ + "embedding": {"values": [0.1, 0.2]}, + "usageMetadata": {"promptTokenCount": 3}, + }, + ), + ], + ) + + assert result["custom_id"] == "request-1" + assert result["response"]["body"]["data"] == [ + {"embedding": [0.1, 0.2], "index": 0, "object": "embedding"}, + {"embedding": [0.3, 0.4], "index": 1, "object": "embedding"}, + ] + assert result["response"]["body"]["usage"]["prompt_tokens"] == 8 + + def test_should_keep_fanned_out_entries_apart_and_in_file_order(self, config): + results = self._transform( + config, + [ + self._vertex_embeddings_output_row(key="request-2#0/2"), + self._vertex_embeddings_output_row(key="request-1"), + self._vertex_embeddings_output_row(key="request-2#1/2"), + ], + ) + + assert [result["custom_id"] for result in results] == ["request-2", "request-1"] + assert len(results[0]["response"]["body"]["data"]) == 2 + assert len(results[1]["response"]["body"]["data"]) == 1 + + def test_should_fail_the_whole_entry_when_one_of_its_rows_failed(self, config): + """An OpenAI batch row is either a response or an error, never both.""" + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row(key="request-1#0/2"), + self._vertex_embeddings_output_row( + key="request-1#1/2", status="Quota exceeded", response={} + ), + ], + ) + + assert result["custom_id"] == "request-1" + assert result["response"] is None + assert result["error"]["message"] == "Quota exceeded" + + def test_should_end_to_end_round_trip_a_fanned_out_embeddings_batch(self, config): + first_row, second_row = _wrap_entries( + [ + _embeddings_entry( + custom_id="MyRequest-1", + body={ + "model": "gemini-embedding-2", + "input": ["hello world", "goodbye world"], + }, + ) + ] + ) + + (result,) = self._transform( + config, + [ + { + **row, + "status": "", + "response": {"embedding": {"values": values}}, + } + for row, values in ((second_row, [0.3]), (first_row, [0.1])) + ], + ) + + assert result["custom_id"] == "MyRequest-1" + assert [ + embedding["embedding"] for embedding in result["response"]["body"]["data"] + ] == [[0.1], [0.3]] + def test_should_end_to_end_round_trip_openai_embeddings_batch(self, config): (vertex_row,) = _wrap_entries( [ From 3c979f0b471214929b38de0fe61573fc864b7906 Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 15:23:14 +0000 Subject: [PATCH 016/576] test(vertex_ai): cover batch lines without a url staying on the chat path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_vertex_ai_files_transformation.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 73d1d6eeb5a..aef7f8b4684 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -1526,6 +1526,24 @@ class TestVertexEmbeddingsBatchInputTranslation: assert row["request"]["labels"]["litellm_custom_id"] == "request-1" assert "key" not in row + def test_should_keep_lines_without_a_url_on_generate_content_path(self): + """`url` is optional on a batch line, and chat is the shape LiteLLM has always assumed.""" + (row,) = _wrap_entries( + [ + { + "custom_id": "request-1", + "body": { + "model": "gemini-2.0-flash-001", + "messages": [{"role": "user", "content": "Hello"}], + }, + } + ] + ) + + assert row["request"]["contents"] == [ + {"role": "user", "parts": [{"text": "Hello"}]} + ] + def test_should_translate_each_line_by_its_own_url(self): chat_row, embeddings_row = _wrap_entries( [ From bf723fa9c167f48731f68ebe6b3bcab7351f5a83 Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 20:50:27 +0000 Subject: [PATCH 017/576] fix(vertex_ai): percent-encode the custom_id in fanned-out vertex batch keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 29 ++++--- .../test_vertex_ai_files_transformation.py | 75 +++++++++++++++++++ 2 files changed, 92 insertions(+), 12 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 78a16b002e7..90fc5fae082 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -17,7 +17,7 @@ from typing import ( Tuple, Union, ) -from urllib.parse import unquote +from urllib.parse import quote, unquote import httpx from httpx import Headers, Response @@ -87,7 +87,7 @@ _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = { "taskType": "task_type", "title": "title", } -_VERTEX_BATCH_FANNED_OUT_KEY_PATTERN = re.compile(r"(?P.*)#(?P\d+)/(?P\d+)") +_VERTEX_BATCH_FANNED_OUT_KEY_PATTERN = re.compile(r"(?P[^#]*)#(?P\d+)/(?P\d+)") def _sanitize_gcp_label_value(value: str) -> str: @@ -160,7 +160,7 @@ def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, Any]) -> str: """ key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) if key is not None: - return str(key) + return unquote(str(key)) request_data = vertex_output_row.get("request") or {} return _get_litellm_batch_custom_id_from_labels(request_data.get("labels") or {}) @@ -228,14 +228,17 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, Resolve `(custom_id, index within that custom_id)` for a Vertex batch output row. A `/v1/embeddings` entry whose `input` is an array fans out into one Vertex row per - element, tagged `#/` (see `_vertex_batch_embeddings_key`), - so the rows can be reassembled into a single OpenAI response. + element, tagged `#/` (see + `_vertex_batch_embeddings_key`), so the rows can be reassembled into a single OpenAI + response. """ - key = _get_litellm_batch_custom_id(vertex_output_row) - match = _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN.fullmatch(key) - if match is None or int(match["total"]) < 2: - return key, 0 - return match["custom_id"], int(match["index"]) + key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) + if key is None: + return _get_litellm_batch_custom_id(vertex_output_row), 0 + match = _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN.fullmatch(str(key)) + if match is None: + return unquote(str(key)), 0 + return unquote(match["custom_id"]), int(match["index"]) def _vertex_embeddings_rows_to_openai_batch_output_row( @@ -359,9 +362,11 @@ def _vertex_batch_embeddings_key(custom_id: str, index: int, total: int) -> str: An entry asking for several embeddings needs several Vertex rows, so its key also carries the element index and the group size; `_split_vertex_batch_key` reads them - back out. Entries asking for a single embedding keep their bare `custom_id`. + back out. The `custom_id` is percent-encoded so that a customer one ending in + `#/` cannot be mistaken for that tag, which would merge two entries. """ - return custom_id if total < 2 else f"{custom_id}#{index}/{total}" + encoded_custom_id = quote(custom_id, safe="") + return encoded_custom_id if total < 2 else f"{encoded_custom_id}#{index}/{total}" def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index aef7f8b4684..f95a63e4421 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -1476,6 +1476,19 @@ class TestVertexEmbeddingsBatchInputTranslation: assert row["key"] == "request-1" + def test_should_encode_a_custom_id_that_looks_like_a_fan_out_tag(self): + """A customer custom_id ending in `#/` must not read back as fan-out metadata.""" + (row,) = _wrap_entries( + [ + _embeddings_entry( + custom_id="request-1#0/2", + body={"model": "gemini-embedding-2", "input": "hello world"}, + ) + ] + ) + + assert row["key"] == "request-1%230%2F2" + def test_should_combine_a_nested_input_into_one_multipart_row(self): """Nested arrays are the combined-embedding shape, as on the online path.""" (row,) = _wrap_entries( @@ -1711,6 +1724,68 @@ class TestVertexEmbeddingsBatchOutputTranslation: assert len(results[0]["response"]["body"]["data"]) == 2 assert len(results[1]["response"]["body"]["data"]) == 1 + def test_should_not_merge_an_entry_whose_custom_id_looks_like_a_fan_out_tag(self, config): + """`request-1#0/2` is a legal custom_id, and a distinct entry from `request-1`.""" + lookalike_row, plain_row = _wrap_entries( + [ + _embeddings_entry( + custom_id="request-1#0/2", + body={"model": "gemini-embedding-2", "input": "lookalike"}, + ), + _embeddings_entry( + custom_id="request-1", + body={"model": "gemini-embedding-2", "input": "plain"}, + ), + ] + ) + + results = self._transform( + config, + [ + {**row, "status": "", "response": {"embedding": {"values": values}}} + for row, values in ((lookalike_row, [0.1]), (plain_row, [0.2])) + ], + ) + + assert [result["custom_id"] for result in results] == [ + "request-1#0/2", + "request-1", + ] + assert [ + result["response"]["body"]["data"][0]["embedding"] for result in results + ] == [[0.1], [0.2]] + + def test_should_round_trip_a_fan_out_of_a_custom_id_holding_the_separator(self, config): + rows = _wrap_entries( + [ + _embeddings_entry( + custom_id="request#1/1", + body={ + "model": "gemini-embedding-2", + "input": ["first", "second"], + }, + ) + ] + ) + + assert [row["key"] for row in rows] == [ + "request%231%2F1#0/2", + "request%231%2F1#1/2", + ] + + (result,) = self._transform( + config, + [ + {**row, "status": "", "response": {"embedding": {"values": values}}} + for row, values in zip(reversed(rows), ([0.3], [0.1])) + ], + ) + + assert result["custom_id"] == "request#1/1" + assert [ + embedding["embedding"] for embedding in result["response"]["body"]["data"] + ] == [[0.1], [0.3]] + def test_should_fail_the_whole_entry_when_one_of_its_rows_failed(self, config): """An OpenAI batch row is either a response or an error, never both.""" (result,) = self._transform( From b93030f84e7a414d2106528114b09f1fca1ad1aa Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 21:20:25 +0000 Subject: [PATCH 018/576] fix(vertex_ai): surface real error/status on vertex batch create instead of IndexError 500 --- litellm/llms/vertex_ai/batches/handler.py | 44 ++++++++++++---- .../llms/vertex_ai/batches/transformation.py | 46 ++++++++++++++--- .../llms/vertex_ai/batches/test_handler.py | 50 +++++++++++++++---- .../vertex_ai/batches/test_transformation.py | 43 +++++++++++++++- 4 files changed, 152 insertions(+), 31 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index ada1356fb6b..f0fd5480c75 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -13,7 +13,7 @@ from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, ) -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.types.llms.openai import CreateBatchRequest from litellm.types.llms.vertex_ai import ( @@ -98,7 +98,9 @@ class VertexAIBatchPrediction(VertexLLM): ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -130,7 +132,9 @@ class VertexAIBatchPrediction(VertexLLM): ) raise if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -242,7 +246,9 @@ class VertexAIBatchPrediction(VertexLLM): ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -292,7 +298,9 @@ class VertexAIBatchPrediction(VertexLLM): headers=headers, ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -365,7 +373,9 @@ class VertexAIBatchPrediction(VertexLLM): ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = ( @@ -390,7 +400,9 @@ class VertexAIBatchPrediction(VertexLLM): params=params, ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = ( @@ -475,7 +487,9 @@ class VertexAIBatchPrediction(VertexLLM): raise if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) # HTTPHandler.get() does not accept a timeout parameter retrieve_response = sync_handler.get( @@ -488,7 +502,10 @@ class VertexAIBatchPrediction(VertexLLM): retrieve_response.status_code, retrieve_response.text[:1000], ) - raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}") + raise VertexAIError( + status_code=retrieve_response.status_code, + message=f"Error: {retrieve_response.status_code} {retrieve_response.text}", + ) _json_response = retrieve_response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -521,7 +538,9 @@ class VertexAIBatchPrediction(VertexLLM): ) raise if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) # AsyncHTTPHandler.get() does not accept a timeout parameter retrieve_response = await client.get( @@ -534,7 +553,10 @@ class VertexAIBatchPrediction(VertexLLM): retrieve_response.status_code, retrieve_response.text[:1000], ) - raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}") + raise VertexAIError( + status_code=retrieve_response.status_code, + message=f"Error: {retrieve_response.status_code} {retrieve_response.text}", + ) _json_response = retrieve_response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index df903ba7ef0..e4299bcf2a0 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -1,7 +1,9 @@ from typing import Any, Dict, Optional +from urllib.parse import unquote from litellm._uuid import uuid from litellm.llms.vertex_ai.common_utils import ( + VertexAIError, _convert_vertex_datetime_to_openai_datetime, ) from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest @@ -199,16 +201,40 @@ class VertexAIBatchTransformation: gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8 returns: "publishers/google/models/gemini-1.5-flash-001" + + Raises a 400 `VertexAIError` when the uri carries no parseable model path. """ - from urllib.parse import unquote - - decoded_uri = unquote(gcs_file_uri) - - model_path = decoded_uri.split("publishers/")[1] - parts = model_path.split("/") - model = f"publishers/{'/'.join(parts[:3])}" + model = cls._parse_model_from_gcs_file(gcs_file_uri) + if model is None: + raise VertexAIError( + status_code=400, + message=( + "Vertex AI batch creation requires the model to be part of `input_file_id`, but " + f"'{gcs_file_uri}' contains no 'publishers//models/' path segment. " + "Either upload the input file through LiteLLM (POST /v1/files with " + "custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or " + "pass a uri of the form " + "gs:////publishers//models//" + ), + ) return model + @classmethod + def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None: + """ + Returns the `publishers//models/` path from a gcs uri, or None if the uri + does not contain one. + """ + _, separator, model_path = unquote(gcs_file_uri).partition("publishers/") + if not separator: + return None + + parts = model_path.split("/") + if len(parts) < 3 or parts[1] != "models" or not parts[2]: + return None + + return f"publishers/{'/'.join(parts[:3])}" + @classmethod def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: Optional[str]) -> bool: """ @@ -216,7 +242,11 @@ class VertexAIBatchTransformation: LiteLLM-managed unified file id) with a `publishers/` model path that `_get_model_from_gcs_file` can parse. """ - return input_file_id is not None and input_file_id.startswith("gs://") and "publishers/" in input_file_id + return ( + input_file_id is not None + and input_file_id.startswith("gs://") + and cls._parse_model_from_gcs_file(input_file_id) is not None + ) @classmethod def get_bare_model_name_from_gcs_file(cls, gcs_file_uri: str) -> str: diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index cacea234777..b9fb5dfe3c5 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -40,6 +40,7 @@ sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.llms.vertex_ai.batches.handler import ( # noqa: E402 VertexAIBatchPrediction, ) +from litellm.llms.vertex_ai.common_utils import VertexAIError # noqa: E402 from litellm.types.utils import LiteLLMBatch # noqa: E402 HMOD = "litellm.llms.vertex_ai.batches.handler" @@ -184,7 +185,7 @@ def test_create_batch_sync_non_200_raises(): client.post.return_value = _http_response(status_code=500) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500") as exc_info: h.create_batch( _is_async=False, create_batch_data=CREATE_DATA, @@ -196,6 +197,32 @@ def test_create_batch_sync_non_200_raises(): max_retries=None, ) + assert exc_info.value.status_code == 500 + assert "error text" in str(exc_info.value) + + +def test_create_batch_input_file_id_without_model_raises_400_before_post(): + """A gs:// uri with no publishers//models/ path is a 400, not a bare 500.""" + h = _make_handler() + client = MagicMock() + + with patch(f"{HMOD}._get_httpx_client", return_value=client): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data={"input_file_id": "gs://bucket/batch-input.jsonl"}, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert exc_info.value.status_code == 400 + assert "gs://bucket/batch-input.jsonl" in str(exc_info.value) + client.post.assert_not_called() + def test_create_batch_async_non_200_raises(): h = _make_handler() @@ -216,9 +243,12 @@ def test_create_batch_async_non_200_raises(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 403"): + with pytest.raises(VertexAIError, match="Error: 403") as exc_info: _run(coro) + assert exc_info.value.status_code == 403 + assert "error text" in str(exc_info.value) + # =========================================================================== # # retrieve_batch @@ -292,7 +322,7 @@ def test_retrieve_batch_sync_non_200_raises(): patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), patch(f"{HMOD}.safe_get", return_value=_http_response(status_code=404)), ): - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): h.retrieve_batch( _is_async=False, batch_id=BATCH_ID, @@ -438,7 +468,7 @@ def test_list_batches_sync_non_200_raises(): client.get.return_value = _http_response(status_code=500) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): h.list_batches( _is_async=False, after=None, @@ -530,7 +560,7 @@ def test_cancel_batch_sync_cancel_post_non_200_raises(): client.post.return_value = _http_response(status_code=500) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): h.cancel_batch( _is_async=False, batch_id=BATCH_ID, @@ -552,7 +582,7 @@ def test_cancel_batch_sync_retrieve_non_200_raises(): client.get.return_value = _http_response(status_code=404) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): h.cancel_batch( _is_async=False, batch_id=BATCH_ID, @@ -672,7 +702,7 @@ def test_async_retrieve_batch_non_200_raises(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): _run(coro) @@ -726,7 +756,7 @@ def test_async_list_batches_non_200_raises(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): _run(coro) @@ -779,7 +809,7 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): _run(coro) async_client_post500.get.assert_not_awaited() @@ -801,5 +831,5 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): _run(coro) diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index 1b37ade6b30..8352ec16389 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -25,6 +25,7 @@ from litellm.llms.vertex_ai.batches.transformation import ( # noqa: E402 VertexAIBatchTransformation, ) from litellm.llms.vertex_ai.common_utils import ( # noqa: E402 + VertexAIError, _convert_vertex_datetime_to_openai_datetime, ) from litellm.types.utils import LiteLLMBatch # noqa: E402 @@ -69,6 +70,24 @@ def test_transform_openai_request_missing_input_file_id_raises(): T.transform_openai_batch_request_to_vertex_ai_batch_request({}) +@pytest.mark.parametrize( + "input_file_id", + [ + "gs://bucket/no-model-here.jsonl", + "gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", + "gs://bucket/publishers/google/models", + "gs://bucket/publishers/google/models//file-uuid", + ], +) +def test_transform_openai_request_unparseable_model_raises_400(input_file_id: str): + """An input_file_id with no parseable model path is a client error, not an IndexError -> 500.""" + with pytest.raises(VertexAIError) as exc_info: + T.transform_openai_batch_request_to_vertex_ai_batch_request({"input_file_id": input_file_id}) + + assert exc_info.value.status_code == 400 + assert input_file_id in str(exc_info.value) + + # =========================================================================== # # transform_vertex_ai_batch_response_to_openai_batch_response # =========================================================================== # @@ -299,9 +318,29 @@ def test_get_model_from_gcs_file_url_encoded(): assert T._get_model_from_gcs_file(encoded) == "publishers/google/models/gemini-1.5-flash-001" -def test_get_model_from_gcs_file_no_publishers_raises(): - with pytest.raises(IndexError): +def test_get_model_from_gcs_file_no_publishers_raises_400(): + with pytest.raises(VertexAIError) as exc_info: T._get_model_from_gcs_file("gs://bucket/no-model-here.jsonl") + assert exc_info.value.status_code == 400 + + +# =========================================================================== # +# is_unmanaged_gcs_batch_input_file_id +# =========================================================================== # + + +@pytest.mark.parametrize( + "input_file_id, expected", + [ + (INPUT_FILE, True), + (None, False), + ("file-abc123", False), + ("gs://bucket/no-model-here.jsonl", False), + ("gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", False), + ], +) +def test_is_unmanaged_gcs_batch_input_file_id(input_file_id, expected): + assert T.is_unmanaged_gcs_batch_input_file_id(input_file_id) is expected # =========================================================================== # From 18d9c7aa21e1308c5ecf05254b29bc0715965bde Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 22:09:49 +0000 Subject: [PATCH 019/576] fix(bedrock): pass SSE-KMS key through to the batch input-file S3 upload --- .../llms/bedrock/batches/transformation.py | 8 ++- litellm/llms/bedrock/common_utils.py | 19 ++++- litellm/llms/bedrock/files/transformation.py | 16 ++++- litellm/types/router.py | 1 + .../bedrock/batches/test_transformation.py | 2 +- .../test_bedrock_files_transformation.py | 72 ++++++++++++++++++- tests/test_litellm/test_router.py | 31 ++++++++ 7 files changed, 141 insertions(+), 8 deletions(-) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index a4ff1c78467..7500531b81a 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -12,7 +12,6 @@ from litellm.litellm_core_utils.cloud_storage_security import ( from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.secret_managers.main import get_secret_str from litellm.types.llms.bedrock import ( BedrockCreateBatchRequest, BedrockCreateBatchResponse, @@ -29,7 +28,7 @@ from litellm.types.llms.openai import ( from litellm.types.utils import LiteLLMBatch, LlmProviders from ..base_aws_llm import BaseAWSLLM -from ..common_utils import CommonBatchFilesUtils +from ..common_utils import CommonBatchFilesUtils, resolve_s3_encryption_key_id # Bedrock batch input files are uploaded as # s3://bucket/litellm-bedrock-files-{model, ":" -> "-"}-{uuid4}.jsonl (see @@ -200,7 +199,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): ) # Add optional KMS encryption key ID if provided - s3_encryption_key_id = litellm_params.get("s3_encryption_key_id") or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") + s3_encryption_key_id = resolve_s3_encryption_key_id( + litellm_params=litellm_params, + optional_params=optional_params, + ) if s3_encryption_key_id: s3_output_config["s3EncryptionKeyId"] = s3_encryption_key_id diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 5114677ffc0..9d427fa6f12 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -35,7 +35,7 @@ from litellm.llms.base_llm.anthropic_messages.transformation import ( ) from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.secret_managers.main import get_secret +from litellm.secret_managers.main import get_secret, get_secret_str if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues @@ -1313,6 +1313,23 @@ def get_anthropic_beta_from_headers(headers: dict) -> List[str]: return [] +def resolve_s3_encryption_key_id( + litellm_params: Mapping[str, Any], + optional_params: Mapping[str, Any] | None = None, +) -> str | None: + """ + Resolve the SSE-KMS key configured for Bedrock batch/file S3 objects. + + Precedence: `s3_encryption_key_id` in litellm_params, then optional_params + (client-side / request params), then the AWS_S3_ENCRYPTION_KEY_ID env var. + """ + for source in (litellm_params, optional_params or {}): + value = source.get("s3_encryption_key_id") + if isinstance(value, str) and value: + return value + return get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") + + class CommonBatchFilesUtils: """ Common utilities for Bedrock batch and file operations. diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index d4865a1c87a..d1674b260b4 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -53,7 +53,7 @@ from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums from litellm.utils import get_llm_provider from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockError +from ..common_utils import BedrockError, resolve_s3_encryption_key_id # litellm_params key used to hand the SigV4-signed GET headers from # `transform_file_content_request` to `validate_environment` (the only hook @@ -741,6 +741,10 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): content=file_content, api_base=api_base, optional_params=optional_params, + s3_encryption_key_id=resolve_s3_encryption_key_id( + litellm_params=litellm_params, + optional_params=optional_params, + ), ) litellm_params["upload_url"] = api_base @@ -758,6 +762,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): content: str, api_base: str, optional_params: dict, + s3_encryption_key_id: str | None = None, ) -> Tuple[dict, str]: """ Sign S3 PUT request using the same proven logic as S3Logger. @@ -790,11 +795,20 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): content_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() # Prepare headers with required S3 headers (same as s3_v2.py) + sse_headers = ( + { + "x-amz-server-side-encryption": "aws:kms", + "x-amz-server-side-encryption-aws-kms-key-id": s3_encryption_key_id, + } + if s3_encryption_key_id + else {} + ) request_headers = { "Content-Type": "application/json", # JSONL files are JSON content "x-amz-content-sha256": content_hash, # REQUIRED by S3 "Content-Language": "en", "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", + **sse_headers, } # Use requests.Request to prepare the request (same pattern as s3_v2.py) diff --git a/litellm/types/router.py b/litellm/types/router.py index 28e4a8272e8..c4d679a2500 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -211,6 +211,7 @@ class CredentialLiteLLMParams(BaseModel): aws_bedrock_runtime_endpoint: Optional[str] = None aws_bedrock_project_id: Optional[str] = None s3_bucket_name: Optional[str] = None + s3_encryption_key_id: Optional[str] = None ## IBM WATSONX ## watsonx_region_name: Optional[str] = None diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 3681daffe5e..01420eb10df 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -172,7 +172,7 @@ def test_create_request_omits_kms_key_when_absent(config): "generate_unique_job_name", return_value="litellm-batch-1", ), patch.object(config.common_utils, "sign_aws_request") as mock_sign, patch( - "litellm.llms.bedrock.batches.transformation.get_secret_str", + "litellm.llms.bedrock.common_utils.get_secret_str", return_value=None, ): mock_sign.return_value = ({}, b"{}") diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index c548fe53e15..a57e5801327 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -442,7 +442,7 @@ class TestBedrockFilesTransformation: captured_optional_params: dict = {} - def fake_sign(content, api_base, optional_params): + def fake_sign(content, api_base, optional_params, s3_encryption_key_id=None): captured_optional_params.update(optional_params) return {"Authorization": "fake"}, content @@ -498,7 +498,7 @@ class TestBedrockFilesTransformation: captured_optional_params: dict = {} - def fake_sign(content, api_base, optional_params): + def fake_sign(content, api_base, optional_params, s3_encryption_key_id=None): captured_optional_params.update(optional_params) return {"Authorization": "fake"}, content @@ -514,6 +514,74 @@ class TestBedrockFilesTransformation: captured_optional_params.get("aws_region_name") == "us-gov-west-1" ), "s3_region_name must override aws_region_name for SigV4 signing" + def _signed_upload_request(self, litellm_params: dict) -> dict: + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + jsonl_content = json.dumps( + { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/amazon.nova-pro-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 10, + }, + } + ).encode() + + request = config.transform_create_file_request( + model="amazon.nova-pro-v1:0", + create_file_data={ + "file": ("batch.jsonl", jsonl_content, "application/jsonl"), + "purpose": "batch", + }, + optional_params={ + "aws_access_key_id": "test-key-id", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + }, + litellm_params={"s3_bucket_name": "litellm-batch-bucket", **litellm_params}, + ) + assert isinstance(request, dict) + return request + + def test_upload_signs_sse_kms_headers_when_key_configured(self, monkeypatch): + """ + Buckets whose policy requires SSE-KMS reject the batch input-file PutObject + unless the upload carries the aws:kms encryption headers; they must also be + covered by SigV4 SignedHeaders or S3 answers SignatureDoesNotMatch. + """ + monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False) + kms_key = "arn:aws:kms:us-west-2:1234:key/abcd" + + request = self._signed_upload_request({"s3_encryption_key_id": kms_key}) + + headers = {key.lower(): value for key, value in request["headers"].items()} + assert headers["x-amz-server-side-encryption"] == "aws:kms" + assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == kms_key + signed_headers = headers["authorization"].split("SignedHeaders=")[1].split(",")[0] + assert "x-amz-server-side-encryption" in signed_headers + assert "x-amz-server-side-encryption-aws-kms-key-id" in signed_headers + + def test_upload_reads_sse_kms_key_from_env(self, monkeypatch): + monkeypatch.setenv("AWS_S3_ENCRYPTION_KEY_ID", "env-kms-key") + + request = self._signed_upload_request({}) + + headers = {key.lower(): value for key, value in request["headers"].items()} + assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == "env-kms-key" + + def test_upload_omits_sse_headers_when_no_key_configured(self, monkeypatch): + monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False) + + request = self._signed_upload_request({}) + + headers = {key.lower() for key in request["headers"]} + assert "x-amz-server-side-encryption" not in headers + assert "x-amz-server-side-encryption-aws-kms-key-id" not in headers + def test_openai_passthrough_still_works(self): """ Regression test: ensure OpenAI-compatible models (e.g. gpt-oss) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 46b5ce65c3f..fa047d7ee46 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3666,6 +3666,37 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): assert credentials["custom_llm_provider"] == "vertex_ai" +def test_get_deployment_credentials_with_provider_includes_s3_encryption_key_id(): + """ + Regression: s3_encryption_key_id must survive the CredentialLiteLLMParams filter, + otherwise the Bedrock batch input-file upload loses the SSE-KMS key and S3 rejects + the PutObject on buckets whose policy requires aws:kms encryption. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-batch", + "litellm_params": { + "model": "bedrock/anthropic.claude-sonnet-4-20250514-v1:0", + "aws_region_name": "us-west-2", + "s3_bucket_name": "my-batch-bucket", + "s3_encryption_key_id": "arn:aws:kms:us-west-2:1234:key/abcd", + }, + } + ], + ) + + credentials = router.get_deployment_credentials_with_provider( + model_id="bedrock-batch" + ) + + assert credentials is not None + assert ( + credentials["s3_encryption_key_id"] + == "arn:aws:kms:us-west-2:1234:key/abcd" + ) + + def test_get_deployment_credentials_with_provider_resolves_credential_name(): """ Test that get_deployment_credentials_with_provider correctly resolves From 7c56317edf153d61b395f4257476aefdd02f2236 Mon Sep 17 00:00:00 2001 From: Yaroslav Date: Thu, 30 Jul 2026 21:37:28 +0300 Subject: [PATCH 020/576] fix(bedrock): drop toolSpec.strict for Claude Sonnet 5 on Converse (#33196) Bedrock routes Claude Sonnet 5 through the same Anthropic-compatible validator as Opus 4.7/4.8 and Sonnet 4, which rejects toolSpec.strict with 'tools.0.custom.strict: Extra inputs are not permitted'. Set bedrock_converse_supports_strict_tools: false on all six Sonnet 5 entries so the existing gate strips the field, matching the fix shape of #31582 Co-authored-by: Yaroslav Budyanskiy --- ...odel_prices_and_context_window_backup.json | 6 +++++ model_prices_and_context_window.json | 6 +++++ ...edrock_converse_strict_tools_opus_47_48.py | 27 +++++++++++++++---- 3 files changed, 34 insertions(+), 5 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e5afc81b641..5e21a868729 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1759,6 +1759,7 @@ "prompt_cache_min_tokens": 2048 }, "anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -1795,6 +1796,7 @@ "prompt_cache_min_tokens": 1024 }, "global.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -1831,6 +1833,7 @@ "prompt_cache_min_tokens": 1024 }, "us.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -1867,6 +1870,7 @@ "prompt_cache_min_tokens": 1024 }, "eu.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -1903,6 +1907,7 @@ "prompt_cache_min_tokens": 1024 }, "au.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -1939,6 +1944,7 @@ "prompt_cache_min_tokens": 1024 }, "jp.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5cf99ba8bac..7587f71bffc 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1759,6 +1759,7 @@ "prompt_cache_min_tokens": 2048 }, "anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -1795,6 +1796,7 @@ "prompt_cache_min_tokens": 1024 }, "global.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -1831,6 +1833,7 @@ "prompt_cache_min_tokens": 1024 }, "us.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -1867,6 +1870,7 @@ "prompt_cache_min_tokens": 1024 }, "eu.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -1903,6 +1907,7 @@ "prompt_cache_min_tokens": 1024 }, "au.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -1939,6 +1944,7 @@ "prompt_cache_min_tokens": 1024 }, "jp.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py index 791982fc3dc..b02324af0a5 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py @@ -1,9 +1,9 @@ """Regression tests for Bedrock Converse ``toolSpec.strict`` forwarding. -Bedrock Converse routes Claude Opus 4.7/4.8 and Claude Sonnet 4 through an -Anthropic-compatible validator that rejects ``toolSpec.strict`` even though -Anthropic's native API accepts ``strict`` as a top-level tool field. See -BerriAI/litellm#31582. +Bedrock Converse routes Claude Opus 4.7/4.8, Claude Sonnet 4 and Claude +Sonnet 5 through an Anthropic-compatible validator that rejects +``toolSpec.strict`` even though Anthropic's native API accepts ``strict`` +as a top-level tool field. See BerriAI/litellm#31582. """ import pytest @@ -48,12 +48,18 @@ _STRICT_TOOL = [ "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0", "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0", "bedrock/apac.anthropic.claude-sonnet-4-20250514-v1:0", + "anthropic.claude-sonnet-5", + "bedrock/global.anthropic.claude-sonnet-5", + "bedrock/us.anthropic.claude-sonnet-5", + "bedrock/eu.anthropic.claude-sonnet-5", + "bedrock/au.anthropic.claude-sonnet-5", + "bedrock/jp.anthropic.claude-sonnet-5", ], ) def test_bedrock_tools_pt_strict_dropped_for_strict_unsupported_models( model_id: str, ) -> None: - """Opus 4.7/4.8 and Sonnet 4 reject toolSpec.strict and additionalProperties.""" + """Opus 4.7/4.8, Sonnet 4 and Sonnet 5 reject toolSpec.strict and additionalProperties.""" result = _bedrock_tools_pt(_STRICT_TOOL, model=model_id) tool_spec = result[0]["toolSpec"] assert ( @@ -129,6 +135,11 @@ def test_bedrock_converse_supports_strict_tools_helper() -> None: ) is False ) + assert bedrock_converse_supports_strict_tools("anthropic.claude-sonnet-5") is False + assert ( + bedrock_converse_supports_strict_tools("bedrock/us.anthropic.claude-sonnet-5") + is False + ) @pytest.mark.parametrize( @@ -143,6 +154,12 @@ def test_bedrock_converse_supports_strict_tools_helper() -> None: "us.anthropic.claude-sonnet-4-20250514-v1:0", "eu.anthropic.claude-sonnet-4-20250514-v1:0", "apac.anthropic.claude-sonnet-4-20250514-v1:0", + "anthropic.claude-sonnet-5", + "global.anthropic.claude-sonnet-5", + "us.anthropic.claude-sonnet-5", + "eu.anthropic.claude-sonnet-5", + "au.anthropic.claude-sonnet-5", + "jp.anthropic.claude-sonnet-5", ], ) def test_strict_tools_flag_set_in_model_cost_map(cost_map_key: str) -> None: From 3e4669dbc5d31a261e65af8e021d2ad12a2d15c4 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 30 Jul 2026 23:03:19 +0000 Subject: [PATCH 021/576] fix(cost): track OpenAI/Azure web search tool cost per call Adds search_context_cost_per_query pricing for the 82 OpenAI/Azure models that advertise supports_web_search but had none (gpt-5 family, o-series, deep-research at $0.01/call; gpt-4.1 at $0.025/call), so built-in web search is no longer billed as $0. Also counts web_search_call items in Responses output so N searches bill N times instead of once; usage-count providers (gemini, anthropic, xai, vertex) still route through get_cost_for_web_search_request and are unaffected. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_cost_calc/tool_call_cost_tracking.py | 20 +- ...odel_prices_and_context_window_backup.json | 410 ++++++++++++++++++ model_prices_and_context_window.json | 410 ++++++++++++++++++ .../test_tool_call_cost_tracking.py | 85 ++++ 4 files changed, 924 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 221b1ae6eab..1d8f2a6c965 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -117,10 +117,28 @@ class StandardBuiltInToolCostTracking: if result is not None: return result - return StandardBuiltInToolCostTracking.get_cost_for_web_search( + per_call_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search( web_search_options=standard_built_in_tools_params.get("web_search_options", None), model_info=model_info, ) + return per_call_cost * StandardBuiltInToolCostTracking._count_web_search_calls(response_object) + + @staticmethod + def _count_web_search_calls(response_object: object) -> int: + """ + Number of web searches to bill for on the per-call pricing path. + + Providers that report a request count in usage (gemini, anthropic, xai, vertex) are handled by + get_cost_for_web_search_request and never reach here. This path prices per call, so it must count + the web_search_call items. Chat-completions responses only expose url_citation annotations with no + count, so they floor to a single billable search. + """ + if isinstance(response_object, ResponsesAPIResponse): + count = sum( + 1 for output_item in response_object.output if getattr(output_item, "type", None) == "web_search_call" + ) + return max(count, 1) + return 1 @staticmethod def _handle_file_search_cost( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5e21a868729..193d2bc2025 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5760,6 +5760,11 @@ "supports_vision": true }, "azure/gpt-5.2-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5791,6 +5796,11 @@ "supports_web_search": true }, "azure/gpt-5.2-pro-2025-12-11": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -6044,6 +6054,11 @@ "supports_vision": true }, "azure/gpt-5.4-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6079,6 +6094,11 @@ "supports_web_search": true }, "azure/gpt-5.4-pro-2026-03-05": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6114,6 +6134,11 @@ "supports_web_search": true }, "azure/gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6159,6 +6184,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6204,6 +6234,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -6249,6 +6284,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_above_272k_tokens": 2e-07, "cache_read_input_token_cost_priority": 2e-07, @@ -6294,6 +6334,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6336,6 +6381,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6378,6 +6428,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "cache_read_input_token_cost_priority": 6.875e-07, @@ -6420,6 +6475,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.1e-07, "cache_read_input_token_cost_above_272k_tokens": 2.2e-07, "cache_read_input_token_cost_priority": 2.75e-07, @@ -6462,6 +6522,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6504,6 +6569,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6546,6 +6616,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "cache_read_input_token_cost_priority": 6.875e-07, @@ -6588,6 +6663,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.1e-07, "cache_read_input_token_cost_above_272k_tokens": 2.2e-07, "cache_read_input_token_cost_priority": 2.75e-07, @@ -6630,6 +6710,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6675,6 +6760,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6717,6 +6807,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6759,6 +6854,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6801,6 +6901,11 @@ "supports_web_search": true }, "azure/us/gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6840,6 +6945,11 @@ "supports_web_search": true }, "azure/eu/gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6879,6 +6989,11 @@ "supports_web_search": true }, "azure/gpt-5.5-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6918,6 +7033,11 @@ "supports_low_reasoning_effort": false }, "azure/gpt-5.5-pro-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6953,6 +7073,11 @@ "supports_web_search": true }, "azure/gpt-5.4-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 7.5e-07, "litellm_provider": "azure", @@ -6988,6 +7113,11 @@ "supports_xhigh_reasoning_effort": false }, "azure/gpt-5.4-mini-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 7.5e-07, "litellm_provider": "azure", @@ -7023,6 +7153,11 @@ "supports_xhigh_reasoning_effort": false }, "azure/gpt-5.4-nano": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "azure", @@ -7058,6 +7193,11 @@ "supports_xhigh_reasoning_effort": false }, "azure/gpt-5.4-nano-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "azure", @@ -7521,6 +7661,11 @@ "supports_vision": true }, "azure/o3-deep-research": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_token": 1e-05, "litellm_provider": "azure", @@ -21452,6 +21597,11 @@ "supports_tool_choice": true }, "gpt-4.1": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_priority": 8.75e-07, "input_cost_per_token": 2e-06, @@ -21489,6 +21639,11 @@ "supports_web_search": true }, "gpt-4.1-2025-04-14": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -21523,6 +21678,11 @@ "supports_web_search": true }, "gpt-4.1-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_priority": 1.75e-07, "input_cost_per_token": 4e-07, @@ -21560,6 +21720,11 @@ "supports_web_search": true }, "gpt-4.1-mini-2025-04-14": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 4e-07, "input_cost_per_token_batches": 2e-07, @@ -22706,6 +22871,11 @@ "supports_pdf_input": true }, "gpt-5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, @@ -22748,6 +22918,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -22787,6 +22962,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-2025-11-13": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -22826,6 +23006,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-chat-latest": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -22865,6 +23050,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -22905,6 +23095,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-2025-12-11": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -22945,6 +23140,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-chat-latest": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -22983,6 +23183,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.3-chat-latest": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -23021,6 +23226,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -23055,6 +23265,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-pro-2025-12-11": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -23089,6 +23304,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, "cache_creation_input_token_cost_flex": 3.125e-06, @@ -23142,6 +23362,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, "cache_creation_input_token_cost_flex": 3.125e-06, @@ -23195,6 +23420,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 3.125e-06, "cache_creation_input_token_cost_above_272k_tokens": 6.25e-06, "cache_creation_input_token_cost_flex": 1.5625e-06, @@ -23248,6 +23478,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, "cache_creation_input_token_cost_flex": 6.25e-07, @@ -23301,6 +23536,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_flex": 2.5e-07, @@ -23350,6 +23590,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_flex": 2.5e-07, @@ -23399,6 +23644,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.5-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23444,6 +23694,11 @@ "supports_low_reasoning_effort": false }, "gpt-5.5-pro-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23582,6 +23837,11 @@ "supports_vision": true }, "gpt-5.4-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23626,6 +23886,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.4-pro-2026-03-05": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23670,6 +23935,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.4-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "cache_read_input_token_cost_priority": 1.5e-07, @@ -23716,6 +23986,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.4-mini-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "cache_read_input_token_cost_priority": 1.5e-07, @@ -23762,6 +24037,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.4-nano": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_flex": 1e-08, "input_cost_per_token": 2e-07, @@ -23805,6 +24085,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.4-nano-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_flex": 1e-08, "input_cost_per_token": 2e-07, @@ -23848,6 +24133,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", @@ -23884,6 +24174,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-pro-2025-10-06": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", @@ -23920,6 +24215,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-2025-08-07": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, @@ -24032,6 +24332,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", @@ -24066,6 +24371,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -24103,6 +24413,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-codex-max": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", @@ -24137,6 +24452,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-codex-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_priority": 4.5e-08, "input_cost_per_token": 2.5e-07, @@ -24174,6 +24494,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -24211,6 +24536,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.3-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -24248,6 +24578,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -24290,6 +24625,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-mini-2025-08-07": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -24332,6 +24672,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-nano": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_flex": 2.5e-09, "input_cost_per_token": 5e-08, @@ -24372,6 +24717,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-nano-2025-08-07": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_flex": 2.5e-09, "input_cost_per_token": 5e-08, @@ -28548,6 +28898,11 @@ "supports_vision": true }, "o3": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 8.75e-07, @@ -28586,6 +28941,11 @@ "supports_web_search": true }, "o3-2025-04-16": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "openai", @@ -28618,6 +28978,11 @@ "supports_web_search": true }, "o3-deep-research": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_token": 1e-05, "input_cost_per_token_batches": 5e-06, @@ -28652,6 +29017,11 @@ "supports_web_search": true }, "o3-deep-research-2025-06-26": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_token": 1e-05, "input_cost_per_token_batches": 5e-06, @@ -28720,6 +29090,11 @@ "supports_vision": false }, "o3-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "openai", @@ -28751,6 +29126,11 @@ "supports_web_search": true }, "o3-pro-2025-06-10": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "openai", @@ -28782,6 +29162,11 @@ "supports_web_search": true }, "o4-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_flex": 1.375e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -28807,6 +29192,11 @@ "supports_web_search": true }, "o4-mini-2025-04-16": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "openai", @@ -28826,6 +29216,11 @@ "supports_web_search": true }, "o4-mini-deep-research": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -28860,6 +29255,11 @@ "supports_web_search": true }, "o4-mini-deep-research-2025-06-26": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -43526,6 +43926,11 @@ ] }, "gpt-5-search-api": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", @@ -43548,6 +43953,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-search-api-2025-10-14": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7587f71bffc..35bf17675cb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5760,6 +5760,11 @@ "supports_vision": true }, "azure/gpt-5.2-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5791,6 +5796,11 @@ "supports_web_search": true }, "azure/gpt-5.2-pro-2025-12-11": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -6044,6 +6054,11 @@ "supports_vision": true }, "azure/gpt-5.4-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6079,6 +6094,11 @@ "supports_web_search": true }, "azure/gpt-5.4-pro-2026-03-05": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6114,6 +6134,11 @@ "supports_web_search": true }, "azure/gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6159,6 +6184,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6204,6 +6234,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -6249,6 +6284,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_above_272k_tokens": 2e-07, "cache_read_input_token_cost_priority": 2e-07, @@ -6294,6 +6334,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6336,6 +6381,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6378,6 +6428,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "cache_read_input_token_cost_priority": 6.875e-07, @@ -6420,6 +6475,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.1e-07, "cache_read_input_token_cost_above_272k_tokens": 2.2e-07, "cache_read_input_token_cost_priority": 2.75e-07, @@ -6462,6 +6522,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6504,6 +6569,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6546,6 +6616,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "cache_read_input_token_cost_priority": 6.875e-07, @@ -6588,6 +6663,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.1e-07, "cache_read_input_token_cost_above_272k_tokens": 2.2e-07, "cache_read_input_token_cost_priority": 2.75e-07, @@ -6630,6 +6710,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6675,6 +6760,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6717,6 +6807,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6759,6 +6854,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6801,6 +6901,11 @@ "supports_web_search": true }, "azure/us/gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6840,6 +6945,11 @@ "supports_web_search": true }, "azure/eu/gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6879,6 +6989,11 @@ "supports_web_search": true }, "azure/gpt-5.5-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6918,6 +7033,11 @@ "supports_low_reasoning_effort": false }, "azure/gpt-5.5-pro-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6953,6 +7073,11 @@ "supports_web_search": true }, "azure/gpt-5.4-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 7.5e-07, "litellm_provider": "azure", @@ -6988,6 +7113,11 @@ "supports_xhigh_reasoning_effort": false }, "azure/gpt-5.4-mini-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 7.5e-07, "litellm_provider": "azure", @@ -7023,6 +7153,11 @@ "supports_xhigh_reasoning_effort": false }, "azure/gpt-5.4-nano": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "azure", @@ -7058,6 +7193,11 @@ "supports_xhigh_reasoning_effort": false }, "azure/gpt-5.4-nano-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "azure", @@ -7521,6 +7661,11 @@ "supports_vision": true }, "azure/o3-deep-research": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_token": 1e-05, "litellm_provider": "azure", @@ -21527,6 +21672,11 @@ "supports_tool_choice": true }, "gpt-4.1": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_priority": 8.75e-07, "input_cost_per_token": 2e-06, @@ -21564,6 +21714,11 @@ "supports_web_search": true }, "gpt-4.1-2025-04-14": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -21598,6 +21753,11 @@ "supports_web_search": true }, "gpt-4.1-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_priority": 1.75e-07, "input_cost_per_token": 4e-07, @@ -21635,6 +21795,11 @@ "supports_web_search": true }, "gpt-4.1-mini-2025-04-14": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 4e-07, "input_cost_per_token_batches": 2e-07, @@ -22781,6 +22946,11 @@ "supports_pdf_input": true }, "gpt-5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, @@ -22823,6 +22993,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -22862,6 +23037,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-2025-11-13": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -22901,6 +23081,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-chat-latest": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -22940,6 +23125,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -22980,6 +23170,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-2025-12-11": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -23020,6 +23215,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-chat-latest": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -23058,6 +23258,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.3-chat-latest": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -23096,6 +23301,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -23130,6 +23340,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-pro-2025-12-11": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -23164,6 +23379,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, "cache_creation_input_token_cost_flex": 3.125e-06, @@ -23217,6 +23437,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, "cache_creation_input_token_cost_flex": 3.125e-06, @@ -23270,6 +23495,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 3.125e-06, "cache_creation_input_token_cost_above_272k_tokens": 6.25e-06, "cache_creation_input_token_cost_flex": 1.5625e-06, @@ -23323,6 +23553,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, "cache_creation_input_token_cost_flex": 6.25e-07, @@ -23376,6 +23611,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_flex": 2.5e-07, @@ -23425,6 +23665,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_flex": 2.5e-07, @@ -23474,6 +23719,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.5-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23519,6 +23769,11 @@ "supports_low_reasoning_effort": false }, "gpt-5.5-pro-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23657,6 +23912,11 @@ "supports_vision": true }, "gpt-5.4-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23701,6 +23961,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.4-pro-2026-03-05": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23745,6 +24010,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.4-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "cache_read_input_token_cost_priority": 1.5e-07, @@ -23791,6 +24061,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.4-mini-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "cache_read_input_token_cost_priority": 1.5e-07, @@ -23837,6 +24112,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.4-nano": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_flex": 1e-08, "input_cost_per_token": 2e-07, @@ -23880,6 +24160,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.4-nano-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_flex": 1e-08, "input_cost_per_token": 2e-07, @@ -23923,6 +24208,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", @@ -23959,6 +24249,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-pro-2025-10-06": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", @@ -23995,6 +24290,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-2025-08-07": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, @@ -24107,6 +24407,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", @@ -24141,6 +24446,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -24178,6 +24488,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-codex-max": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", @@ -24212,6 +24527,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-codex-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_priority": 4.5e-08, "input_cost_per_token": 2.5e-07, @@ -24249,6 +24569,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -24286,6 +24611,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.3-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -24323,6 +24653,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -24365,6 +24700,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-mini-2025-08-07": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -24407,6 +24747,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-nano": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_flex": 2.5e-09, "input_cost_per_token": 5e-08, @@ -24447,6 +24792,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-nano-2025-08-07": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_flex": 2.5e-09, "input_cost_per_token": 5e-08, @@ -28623,6 +28973,11 @@ "supports_vision": true }, "o3": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 8.75e-07, @@ -28661,6 +29016,11 @@ "supports_web_search": true }, "o3-2025-04-16": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "openai", @@ -28693,6 +29053,11 @@ "supports_web_search": true }, "o3-deep-research": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_token": 1e-05, "input_cost_per_token_batches": 5e-06, @@ -28727,6 +29092,11 @@ "supports_web_search": true }, "o3-deep-research-2025-06-26": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_token": 1e-05, "input_cost_per_token_batches": 5e-06, @@ -28795,6 +29165,11 @@ "supports_vision": false }, "o3-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "openai", @@ -28826,6 +29201,11 @@ "supports_web_search": true }, "o3-pro-2025-06-10": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "openai", @@ -28857,6 +29237,11 @@ "supports_web_search": true }, "o4-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_flex": 1.375e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -28882,6 +29267,11 @@ "supports_web_search": true }, "o4-mini-2025-04-16": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "openai", @@ -28901,6 +29291,11 @@ "supports_web_search": true }, "o4-mini-deep-research": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -28935,6 +29330,11 @@ "supports_web_search": true }, "o4-mini-deep-research-2025-06-26": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -43647,6 +44047,11 @@ ] }, "gpt-5-search-api": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", @@ -43669,6 +44074,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-search-api-2025-10-14": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 24fd3c94ee3..f194db7676a 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -602,5 +602,90 @@ def test_web_search_provider_prefix_fallback_does_not_misprice_non_gemini_model( ) +def _openai_responses_with_web_search_calls(model, num_calls): + from litellm.types.llms.openai import ResponsesAPIResponse + from openai.types.responses.response_function_web_search import ( + ActionSearch, + ResponseFunctionWebSearch, + ) + + output = [ + ResponseFunctionWebSearch( + id=f"ws_{i}", + type="web_search_call", + status="completed", + action=ActionSearch(type="search", query="latest news"), + ) + for i in range(num_calls) + ] + return ResponsesAPIResponse( + id="resp_1", + created_at=0, + model=model, + object="response", + output=output, + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + + +def test_openai_responses_web_search_priced_per_call(local_model_cost_map): + """ + Regression for LIT-5013 bug 1: OpenAI reasoning models (gpt-5 family, o-series, deep-research) + carry supports_web_search but had no search_context_cost_per_query, so get_cost_for_web_search_request + (no openai branch) returned None and the default fallback billed web search as $0. gpt-5-nano now + prices at $0.01 per call, and two web_search_call items in the Responses output must bill 2 x $0.01. + """ + from litellm.types.utils import Usage + + model = "gpt-5-nano" + per_call = litellm.get_model_info(model)["search_context_cost_per_query"][ + "search_context_size_medium" + ] + assert per_call == 0.01 + + response = _openai_responses_with_web_search_calls(model, num_calls=2) + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=response, + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + assert cost == pytest.approx(2 * per_call), ( + f"gpt-5-nano web search must bill 2 x ${per_call}, got ${cost}" + ) + + +def test_openai_responses_web_search_multiplied_by_call_count(local_model_cost_map): + """ + Regression for LIT-5013 bug 2: web_search_call detection was binary, so a Responses output with + multiple web searches was charged once. gpt-4o-search-preview carries per-call pricing; N calls + must bill N times, and a single call must still bill exactly once. + """ + from litellm.types.utils import Usage + + model = "gpt-4o-search-preview" + per_call = litellm.get_model_info(model)["search_context_cost_per_query"][ + "search_context_size_medium" + ] + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + + for num_calls in (1, 3): + response = _openai_responses_with_web_search_calls(model, num_calls=num_calls) + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=response, + usage=usage, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + assert cost == pytest.approx(num_calls * per_call), ( + f"{num_calls} web searches must bill {num_calls} x ${per_call}, got ${cost}" + ) + + # Note: File search integration test removed due to complex annotation detection logic # The unit tests in test_azure_assistant_cost_tracking.py provide comprehensive coverage From 0b809cf7d68e368b7a7ee90d7134c4841c944e3c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 31 Jul 2026 03:21:20 +0000 Subject: [PATCH 022/576] fix(anthropic adapter): stop indexing choices[0] on choiceless streaming chunks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../adapters/streaming_iterator.py | 30 +++++++ .../test_streaming_iterator_empty_choices.py | 87 +++++++++++++++++++ 2 files changed, 117 insertions(+) create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index d9bcfa19a7f..194cbcc9327 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -329,6 +329,26 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): merged_chunk["context_management"] = ContextManagementResponse(applied_edits=list(self.applied_edits)) return self._augment_message_delta_usage(merged_chunk) + def _handle_choiceless_chunk(self, chunk: Any) -> bool: + """Consume an OpenAI-compatible chunk that carries no ``choices``. + + ``choices`` is legitimately empty on metadata-only chunks; the final + usage chunk emitted when ``stream_options.include_usage`` is set is the + common case (vLLM and other OpenAI-compatible servers do this). Such a + chunk carries no content-block information, so the caller must not run + the content-block state machine over it. + + Returns True when a merged ``message_delta`` was queued (usage folded + into the held stop-reason chunk); False when the chunk should be + skipped entirely. + """ + if self.holding_stop_reason_chunk is not None and getattr(chunk, "usage", None) is not None: + self.chunk_queue.append(self._merge_usage_into_held_stop_reason_chunk(chunk)) + self.queued_usage_chunk = True + self.holding_stop_reason_chunk = None + return True + return False + def _ensure_context_management_attached(self, message_delta_chunk: Dict[str, Any]) -> Dict[str, Any]: """Attach ``context_management`` to a ``message_delta`` chunk if ``self.applied_edits`` is non-empty and the chunk does not already @@ -490,6 +510,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if chunk == "None" or chunk is None: raise Exception + if not getattr(chunk, "choices", None): + if self._handle_choiceless_chunk(chunk): + return self.chunk_queue.popleft() + continue + should_start_new_block = self._should_start_new_content_block(chunk) is_opening_first_block = self.sent_content_block_start is False if is_opening_first_block and self._is_blank_delta(chunk): @@ -713,6 +738,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if chunk == "None" or chunk is None: raise Exception + if not getattr(chunk, "choices", None): + if self._handle_choiceless_chunk(chunk): + return self.chunk_queue.popleft() + continue + should_start_new_block = self._should_start_new_content_block(chunk) is_opening_first_block = self.sent_content_block_start is False if is_opening_first_block and self._is_blank_delta(chunk): diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py new file mode 100644 index 00000000000..3e85872f1e5 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py @@ -0,0 +1,87 @@ +""" +Regression tests for OpenAI-compatible chunks with an empty ``choices`` list. + +``choices: []`` is valid OpenAI-compatible streaming: vLLM (and OpenAI itself, +when ``stream_options.include_usage`` is set) emits a final usage chunk with no +choices, and some gateways emit metadata-only chunks mid-stream. The adapter +used to index ``chunk.choices[0]`` unconditionally, so such a chunk raised +``IndexError: list index out of range`` and killed the ``/v1/messages`` stream. +""" + +import asyncio +import json +from typing import Any, AsyncIterator, Dict, List, Optional + +from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( + AnthropicStreamWrapper, +) +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage + + +def _text_chunk(text: str) -> ModelResponseStream: + return ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=text), finish_reason=None)] + ) + + +def _finish_chunk() -> ModelResponseStream: + return ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")]) + + +def _empty_choices_chunk(usage: Optional[Usage] = None) -> ModelResponseStream: + return ModelResponseStream(choices=[], usage=usage) + + +def _collect_async(wrapper: AnthropicStreamWrapper) -> str: + async def _run() -> str: + return "".join( + [raw.decode() if isinstance(raw, bytes) else raw async for raw in wrapper.async_anthropic_sse_wrapper()] + ) + + return asyncio.run(_run()) + + +def _message_delta(sse: str) -> Dict[str, Any]: + return next( + json.loads(line[len("data: ") :]) + for block in sse.split("\n\n") + for line in block.splitlines() + if line.startswith("data: ") and '"message_delta"' in line + ) + + +def test_leading_metadata_chunk_without_choices_does_not_kill_stream(): + """A metadata-only chunk before any content must be skipped, not indexed.""" + chunks: List[ModelResponseStream] = [ + _empty_choices_chunk(), + _text_chunk("Hello"), + _text_chunk(" there"), + _finish_chunk(), + ] + wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="mock-model") + events = list(wrapper) + + text = "".join( + event["delta"]["text"] for event in events if event.get("type") == "content_block_delta" + ) + assert text == "Hello there" + assert events[-1]["type"] == "message_stop" + + +def test_final_usage_chunk_without_choices_is_merged_into_message_delta(): + """The vLLM/OpenAI final usage chunk carries no choices; its usage must + still land on the Anthropic ``message_delta``.""" + usage = Usage(prompt_tokens=10, completion_tokens=3, total_tokens=13) + + async def _aiter() -> "AsyncIterator[ModelResponseStream]": + for chunk in [_text_chunk("Hi"), _finish_chunk(), _empty_choices_chunk(usage)]: + yield chunk + + sse = _collect_async(AnthropicStreamWrapper(completion_stream=_aiter(), model="mock-model")) + + message_delta = _message_delta(sse) + assert message_delta["delta"]["stop_reason"] == "end_turn" + assert message_delta["usage"]["input_tokens"] == 10 + assert message_delta["usage"]["output_tokens"] == 3 + assert "Hi" in sse + assert "message_stop" in sse From bae58eb4e0aa015d5085264e8b0d9a342f163ce5 Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Fri, 31 Jul 2026 13:03:13 -0400 Subject: [PATCH 023/576] fix(anthropic): preserve mid-turn system messages Generated with AI Co-Authored-By: Claude Code --- .../chat/guardrail_translation/handler.py | 209 ++++-- .../adapters/transformation.py | 48 +- .../responses_adapters/transformation.py | 49 +- litellm/types/guardrails.py | 47 +- litellm/types/llms/anthropic.py | 12 + .../test_anthropic_guardrail_handler.py | 699 ++++++++++++++++++ ...al_pass_through_adapters_transformation.py | 218 ++++++ .../context_management/test_compact.py | 54 ++ .../test_responses_adapters_transformation.py | 138 ++++ 9 files changed, 1380 insertions(+), 94 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index a549db94224..82707e741c0 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,6 +13,7 @@ Pattern Overview: """ import json +from copy import deepcopy from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast from litellm._logging import verbose_proxy_logger @@ -24,7 +25,6 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTra from litellm.llms.base_llm.guardrail_translation.utils import ( effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, - openai_messages_without_system, openai_messages_without_tool, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( @@ -59,14 +59,10 @@ if TYPE_CHECKING: class AnthropicMessagesHandler(BaseTranslation): - """ - Handler for processing Anthropic messages with guardrails. + """Process Anthropic messages with guardrails. - This class provides methods to: - 1. Process input messages (pre-call hook) - 2. Process output responses (post-call hook) - - Methods can be overridden to customize behavior for different message formats. + In-sequence system entries are untrusted client input. This handler scans and preserves + them through guardrail rewrites; downstream provider handling is out of scope. """ def __init__(self): @@ -279,14 +275,26 @@ class AnthropicMessagesHandler(BaseTranslation): skip_system = effective_skip_system_message_for_guardrail(guardrail_to_apply) skip_tool = effective_skip_tool_message_for_guardrail(guardrail_to_apply) - chat_completion_compatible_request = self._translate_to_openai(data) + # Exclude only the trusted top-level prompt. In-sequence system entries are untrusted + # and must stay aligned with texts_to_check for positional masking. When the top-level + # prompt is included, the pre-existing count mismatch disables positional masking. + translation_source = { # mutable-ok: API message payload + key: value for key, value in data.items() if key != "system" + } # mutable-ok: API message payload + chat_completion_compatible_request = self._translate_to_openai(translation_source) structured_messages = cast( List[AllMessageValues], chat_completion_compatible_request.get("messages", []), ) - if skip_system: - structured_messages = openai_messages_without_system(structured_messages) + has_midturn_system_message = any( + str(message.get("role") or "").lower() == "system" for message in structured_messages + ) + hoisted_system_message: AllMessageValues | None = None + if not skip_system: + hoisted_system_message = self._hoisted_top_level_system_message(data) + if hoisted_system_message is not None: + structured_messages.insert(0, hoisted_system_message) if skip_tool: structured_messages = openai_messages_without_tool(structured_messages) @@ -346,7 +354,12 @@ class AnthropicMessagesHandler(BaseTranslation): guardrailed_structured_messages is not None and guardrailed_structured_messages is not original_structured_messages ): - self._write_back_structured_messages(data, guardrailed_structured_messages) + self._write_back_structured_messages( + data, + guardrailed_structured_messages, + hoisted_system_message=hoisted_system_message, + preserve_system_messages=has_midturn_system_message, + ) else: # Step 3: Map guardrail responses back to original message structure await self._apply_guardrail_responses_to_input( @@ -359,36 +372,120 @@ class AnthropicMessagesHandler(BaseTranslation): return data - @staticmethod - def _write_back_structured_messages(data: dict, structured_messages: list) -> None: - """Convert compressed structured_messages back to Anthropic format and write to data. + def _hoisted_top_level_system_message( + self, data: dict + ) -> AllMessageValues | None: # mutable-ok: API message payload + """Return the system message produced by translating the top-level prompt.""" + system = data.get("system") + if not system: + return None + probe = self._translate_to_openai( + { # mutable-ok: API message payload + "model": data.get("model") or "", + "messages": [], # mutable-ok: API message payload + "system": system, + } + ) + hoisted = probe.get("messages") or [] # mutable-ok: API message payload + return hoisted[0] if hoisted else None - ``anthropic_messages_pt`` merges every run of consecutive user/tool rows - into a single message, so a turn carrying only tool results and the user - turn that follows it come back fused, and the request the model sees no - longer has the boundaries the client sent. Converting a row at a time - would keep them apart but breaks tool pairing: an assistant row whose - tool results sit outside its own call reads as an orphaned tool call, - and under ``modify_params`` the sanitizer answers it with a synthetic - "tool execution skipped" result and drops the real one. Converting each - assistant row together with the tool rows that answer it, and every - other row on its own, satisfies both. - """ + @staticmethod + def _openai_system_message_to_anthropic( + message: dict[str, Any], + ) -> dict[str, Any] | None: # mutable-ok: API message payload + """Convert an OpenAI system message to the client's Anthropic-shaped entry.""" + content = message.get("content") + if isinstance(content, str): + return ( + {"role": "system", "content": content} if content else None # mutable-ok: API message payload + ) # mutable-ok: API message payload + if not isinstance(content, list): + return None + blocks: list[dict[str, Any]] = [] # mutable-ok: API message payload + for block in content: + if not isinstance(block, dict) or block.get("type") != "text": + continue + text = block.get("text") + if not isinstance(text, str) or not text: + continue + anthropic_block: dict[str, Any] = { # mutable-ok: API message payload + "type": "text", + "text": text, + } # mutable-ok: API message payload + cache_control = block.get("cache_control") + if cache_control: + anthropic_block["cache_control"] = deepcopy(cache_control) + blocks.append(anthropic_block) + return ( + {"role": "system", "content": blocks} if blocks else None # mutable-ok: API message payload + ) # mutable-ok: API message payload + + @staticmethod + def _is_hoisted_top_level_system(message: Any, hoisted_system_message: Any) -> bool: + """Match the hoisted prompt by identity, or by value after serialization.""" + if hoisted_system_message is None: + return False + if message is hoisted_system_message: + return True + return ( + isinstance(message, dict) and isinstance(hoisted_system_message, dict) and message == hoisted_system_message + ) + + @staticmethod + def _write_back_structured_messages( + data: dict, # mutable-ok: API message payload + structured_messages: list, # mutable-ok: API message payload + hoisted_system_message: Any = None, + preserve_system_messages: bool = False, + ) -> None: + """Write a guardrail's structured-message rewrite back without losing corrections.""" from litellm.litellm_core_utils.prompt_templates.factory import ( anthropic_messages_pt, group_tool_exchanges, ) + def _is_system(message: Any) -> bool: + return isinstance(message, dict) and str(message.get("role") or "").lower() == "system" + model = str(data.get("model") or "") - non_system = [m for m in structured_messages if m.get("role") != "system"] - groups = tuple([non_system[index] for index in group] for group in group_tool_exchanges(non_system)) or ( - non_system, - ) - converted = [ - message - for group in groups - for message in anthropic_messages_pt(messages=group, model=model, llm_provider="anthropic") - ] + converted: list = [] # mutable-ok: API message payload + + def _convert_run(run: list) -> None: # mutable-ok: API message payload + for group in group_tool_exchanges(run): + converted.extend( + anthropic_messages_pt( + messages=[ # mutable-ok: API message payload + run[index] for index in group + ], # mutable-ok: API message payload + model=model, + llm_provider="anthropic", + ) + ) + + run: list = [] # mutable-ok: API message payload + hoisted_dropped = False + for message in structured_messages: + if not _is_system(message): + run.append(message) + continue + _convert_run(run) + run = [] # mutable-ok: API message payload + if not hoisted_dropped and AnthropicMessagesHandler._is_hoisted_top_level_system( + message, hoisted_system_message + ): + hoisted_dropped = True + continue + if preserve_system_messages: + anthropic_system = AnthropicMessagesHandler._openai_system_message_to_anthropic(message) + if anthropic_system is not None: + converted.append(anthropic_system) + _convert_run(run) + if not any(not _is_system(message) for message in converted): + converted.extend( + anthropic_messages_pt( + messages=[], model=model, llm_provider="anthropic" + ) # mutable-ok: API message payload + ) # mutable-ok: API message payload for msg in converted: content = msg.get("content") if isinstance(content, list): @@ -397,6 +494,29 @@ class AnthropicMessagesHandler(BaseTranslation): block.pop("cache_control", None) data["messages"] = converted + @staticmethod + def _extract_midturn_system_text( + message: dict[str, Any], # mutable-ok: API message payload + msg_idx: int, + texts_to_check: list[str], # mutable-ok: API message payload + task_mappings: list[tuple[int, int | None]], # mutable-ok: API message payload + ) -> None: + content = message.get("content") + if isinstance(content, str): + if content: + texts_to_check.append(content) + task_mappings.append((msg_idx, None)) + return + if not isinstance(content, list): + return + for content_idx, content_item in enumerate(content): + if not isinstance(content_item, dict) or content_item.get("type") != "text": + continue + text_str = content_item.get("text") + if isinstance(text_str, str) and text_str: + texts_to_check.append(text_str) + task_mappings.append((msg_idx, content_idx)) + def extract_request_tool_names(self, data: dict) -> List[str]: """Extract tool names from Anthropic messages request (tools[].name).""" names: List[str] = [] @@ -415,15 +535,18 @@ class AnthropicMessagesHandler(BaseTranslation): skip_system_message: bool = False, skip_tool_message: bool = False, ) -> None: - """ - Extract text content and images from a message. - - Override this method to customize text/image extraction logic. - """ - role = str(message.get("role") or "").lower() - if skip_system_message and role == "system": + """Extract text content and images from a message.""" + role = str(message.get("role") or "") + if role == "system": + # Match the adapter's filtering so positional guardrail write-back stays aligned. + self._extract_midturn_system_text( + message=message, + msg_idx=msg_idx, + texts_to_check=texts_to_check, + task_mappings=task_mappings, + ) return - if skip_tool_message and role == "tool": + if skip_tool_message and role.lower() == "tool": return content = message.get("content", None) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 86c9c1db481..707d53e9006 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -85,12 +85,12 @@ from litellm.llms.anthropic.experimental_pass_through.context_management import ) from litellm.types.llms.anthropic import ( ANTHROPIC_HOSTED_TOOLS, + AllAnthropicPassThroughMessageValues, AllAnthropicToolsValues, - AnthopicMessagesAssistantMessageParam, AnthropicFinishReason, AnthropicMessagesRequest, + AnthropicMessagesSystemMessageParam, AnthropicMessagesToolChoice, - AnthropicMessagesUserMessageParam, AnthropicResponseContentBlockRedactedThinking, AnthropicResponseContentBlockText, AnthropicResponseContentBlockThinking, @@ -354,12 +354,7 @@ class LiteLLMAnthropicMessagesAdapter: def translate_anthropic_messages_to_openai( self, - messages: List[ - Union[ - AnthropicMessagesUserMessageParam, - AnthopicMessagesAssistantMessageParam, - ] - ], + messages: List[AllAnthropicPassThroughMessageValues], # mutable-ok: API message payload model: Optional[str] = None, ) -> List: new_messages: List[AllMessageValues] = [] @@ -367,6 +362,11 @@ class LiteLLMAnthropicMessagesAdapter: user_message: Optional[ChatCompletionUserMessage] = None tool_message_list: List[ChatCompletionToolMessage] = [] new_user_content_list: List[Union[ChatCompletionTextObject, ChatCompletionImageObject]] = [] + if m["role"] == "system": + system_message = self._translate_midturn_system_message_to_openai(m, model) + if system_message is not None: + new_messages.append(system_message) + continue ## USER MESSAGE ## if m["role"] == "user": ## translate user message @@ -867,6 +867,29 @@ class LiteLLMAnthropicMessagesAdapter: for def_schema in schema[key].values(): LiteLLMAnthropicMessagesAdapter._add_additional_properties_false(def_schema) + def _translate_midturn_system_message_to_openai( + self, + message: AnthropicMessagesSystemMessageParam, + model: str | None, + ) -> ChatCompletionSystemMessage | None: + """Translate an in-sequence system entry without changing its role or position.""" + content = message.get("content") + if isinstance(content, str): + return ChatCompletionSystemMessage(role="system", content=content) if content else None + if not isinstance(content, list): + return None + text_parts: list[ChatCompletionTextObject] = [] # mutable-ok: API message payload + for block in content: + if not isinstance(block, dict) or block.get("type") != "text": + continue + text = block.get("text") + if not text: + continue + text_obj = ChatCompletionTextObject(type="text", text=text) + self._add_cache_control_if_applicable(block, text_obj, model) + text_parts.append(text_obj) + return ChatCompletionSystemMessage(role="system", content=text_parts) if text_parts else None + def _add_system_message_to_messages( self, new_messages: List[AllMessageValues], @@ -1068,13 +1091,8 @@ class LiteLLMAnthropicMessagesAdapter: tool_name_mapping: Dict[str, str] = {} ## CONVERT ANTHROPIC MESSAGES TO OPENAI - messages_list: List[Union[AnthropicMessagesUserMessageParam, AnthopicMessagesAssistantMessageParam]] = cast( - List[ - Union[ - AnthropicMessagesUserMessageParam, - AnthopicMessagesAssistantMessageParam, - ] - ], + messages_list = cast( + List[AllAnthropicPassThroughMessageValues], anthropic_message_request["messages"], ) new_messages = self.translate_anthropic_messages_to_openai( diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 172e54de98e..cbe36100eb0 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -6,6 +6,7 @@ path used for OpenAI and Azure models. """ import json +from collections.abc import Iterable from typing import Any, Dict, List, Optional, Union, cast from litellm.litellm_core_utils.reasoning_effort_utils import ( @@ -15,15 +16,15 @@ from litellm.llms.anthropic.experimental_pass_through.utils import ( is_reasoning_auto_summary_enabled, ) from litellm.types.llms.anthropic import ( + AllAnthropicPassThroughMessageValues, AllAnthropicToolsValues, - AnthopicMessagesAssistantMessageParam, AnthropicFinishReason, AnthropicMessagesRequest, AnthropicMessagesToolChoice, - AnthropicMessagesUserMessageParam, AnthropicResponseContentBlockText, AnthropicResponseContentBlockThinking, AnthropicResponseContentBlockToolUse, + AnthropicSystemMessageContent, ) from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, @@ -54,19 +55,32 @@ class LiteLLMAnthropicToResponsesAPIAdapter: return source.get("url") return None + @staticmethod + def _translate_midturn_system_content_to_responses( + content: Union[str, Iterable[AnthropicSystemMessageContent]], + ) -> list[dict[str, str]]: # mutable-ok: API message payload + """Convert in-sequence system content to Responses input-text parts.""" + if isinstance(content, str): + return ( + [{"type": "input_text", "text": content}] if content else [] # mutable-ok: API message payload + ) # mutable-ok: API message payload + if not isinstance(content, list): + return [] # mutable-ok: API message payload + return [ # mutable-ok: API message payload + {"type": "input_text", "text": text} # mutable-ok: API message payload + for block in content + if isinstance(block, dict) and block.get("type") == "text" and (text := block.get("text")) + ] + def translate_messages_to_responses_input( self, - messages: List[ - Union[ - AnthropicMessagesUserMessageParam, - AnthopicMessagesAssistantMessageParam, - ] - ], + messages: List[AllAnthropicPassThroughMessageValues], # mutable-ok: API message payload ) -> List[Dict[str, Any]]: """ Convert Anthropic messages list to Responses API `input` items. Mapping: + system text -> message(role=system, input_text) user text -> message(role=user, input_text) user image -> message(role=user, input_image) user tool_result -> function_call_output @@ -76,6 +90,18 @@ class LiteLLMAnthropicToResponsesAPIAdapter: input_items: List[Dict[str, Any]] = [] for m in messages: + if m["role"] == "system": + system_parts = self._translate_midturn_system_content_to_responses(m.get("content")) + if system_parts: + input_items.append( + { # mutable-ok: API message payload + "type": "message", + "role": "system", + "content": system_parts, + } + ) + continue + role = m["role"] content = m.get("content") @@ -287,12 +313,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: """ model: str = anthropic_request["model"] messages_list = cast( - List[ - Union[ - AnthropicMessagesUserMessageParam, - AnthopicMessagesAssistantMessageParam, - ] - ], + List[AllAnthropicPassThroughMessageValues], anthropic_request["messages"], ) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index af419d8cb6f..2e0da24ccda 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -11,12 +11,24 @@ from litellm.types.proxy.guardrails.guardrail_hooks.akto import ( from litellm.types.proxy.guardrails.guardrail_hooks.block_code_execution import ( BlockCodeExecutionGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.cisco_ai_defense import ( + CiscoAIDefenseGuardrailConfigModel, +) +from litellm.types.proxy.guardrails.guardrail_hooks.compresr import ( + CompresrGuardrailConfigModel, +) from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import ( EnkryptAIGuardrailConfigs, ) from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import ( GraySwanGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( + HeadroomGuardrailConfigModel, +) +from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import ( + HiddenlayerGuardrailConfigModel, +) from litellm.types.proxy.guardrails.guardrail_hooks.ibm import ( IBMGuardrailsBaseConfigModel, ) @@ -29,38 +41,26 @@ from litellm.types.proxy.guardrails.guardrail_hooks.ovalix import ( from litellm.types.proxy.guardrails.guardrail_hooks.promptguard import ( PromptGuardConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.xecguard import ( - XecGuardConfigModel, +from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( + QostodianNexusConfigModel, ) from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import ( QualifireGuardrailConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( - ToolPermissionGuardrailConfigModel, -) -from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import ( - HiddenlayerGuardrailConfigModel, -) -from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( - QostodianNexusConfigModel, -) from litellm.types.proxy.guardrails.guardrail_hooks.repelloai import ( RepelloAIGuardrailConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import ( - VigilGuardGuardrailConfigModel, -) -from litellm.types.proxy.guardrails.guardrail_hooks.cisco_ai_defense import ( - CiscoAIDefenseGuardrailConfigModel, -) from litellm.types.proxy.guardrails.guardrail_hooks.singulr import ( SingulrGuardrailConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( - HeadroomGuardrailConfigModel, +from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( + ToolPermissionGuardrailConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.compresr import ( - CompresrGuardrailConfigModel, +from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import ( + VigilGuardGuardrailConfigModel, +) +from litellm.types.proxy.guardrails.guardrail_hooks.xecguard import ( + XecGuardConfigModel, ) """ @@ -743,7 +743,10 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up "When True, unified guardrails skip system-role messages when building " "evaluation inputs (texts and structured_messages). When False, system " "messages are included even if litellm_settings sets a global skip. When " - "None, use the global litellm.skip_system_message_in_guardrail setting." + "None, use the global litellm.skip_system_message_in_guardrail setting. " + "For Anthropic /v1/messages, the flag applies only to the trusted top-level " + "system prompt. In-sequence system entries are untrusted client input and remain " + "in texts and structured_messages." ), ) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index c24d072217a..29faf500b3d 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -365,8 +365,20 @@ class AnthropicSystemMessageContent(TypedDict, total=False): cache_control: Optional[Union[dict, ChatCompletionCachedContent]] +class AnthropicMessagesSystemMessageParam(TypedDict, total=False): + role: Required[Literal["system"]] + content: Required[Union[str, Iterable[AnthropicSystemMessageContent]]] + + AllAnthropicMessageValues = Union[AnthropicMessagesUserMessageParam, AnthopicMessagesAssistantMessageParam] +# System is not a native Anthropic message role; only pass-through adapters use this union. +AllAnthropicPassThroughMessageValues = Union[ + AnthropicMessagesUserMessageParam, + AnthopicMessagesAssistantMessageParam, + AnthropicMessagesSystemMessageParam, +] + class AnthropicMessagesRequestOptionalParams(TypedDict, total=False): max_tokens: Optional[int] diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 48acdd348e9..7757b0fa5a4 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -75,6 +75,51 @@ class MockRecordingGuardrail(CustomGuardrail): return inputs +class MockMaskingGuardrail(CustomGuardrail): + """Capture request inputs and mask one known prohibited value.""" + + def __init__(self, skip_system_message_in_guardrail: Optional[bool] = True): + super().__init__(guardrail_name="masking-test") + self.skip_system_message_in_guardrail = skip_system_message_in_guardrail + self.inputs: Optional[GenericGuardrailAPIInputs] = None + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.inputs = inputs.copy() + masked_inputs = inputs.copy() + masked_inputs["texts"] = [ + "[MASKED]" if text == "prohibited correction" else text for text in inputs.get("texts", []) + ] + return masked_inputs + + +class MockCompactingGuardrail(CustomGuardrail): + """Stand in for a compaction guardrail that rewrites `structured_messages` wholesale.""" + + def __init__(self, replacement_messages: list): + super().__init__(guardrail_name="compacting-test") + self.replacement_messages = replacement_messages + self.inputs: Optional[GenericGuardrailAPIInputs] = None + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.inputs = inputs.copy() + rewritten = inputs.copy() + # A new list object -- this is what signals a rewrite to the handler. + rewritten["structured_messages"] = list(self.replacement_messages) + return rewritten + + class TestAnthropicMessagesHandlerStreamingRequestData: """Post-call guardrails on streaming /v1/messages receive the response and identity metadata""" @@ -210,6 +255,660 @@ class TestAnthropicMessagesHandlerInputProcessing: assert data.get("litellm_metadata", {}).get("guardrails") assert guardrail.dynamic_params == {"policy_id": "policy-123"} + @pytest.mark.asyncio + async def test_midturn_system_correction_is_guardrailed_when_top_level_system_is_skipped( + self, + ): + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "user", "content": "safe text"}, + { + "role": "system", + "content": [ + {"type": "unsupported", "text": "discarded text"}, + {"type": "text", "text": "prohibited correction"}, + ], + }, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + assert guardrail.inputs["texts"] == ["safe text", "prohibited correction"] + assert "trusted top-level system prompt" not in guardrail.inputs["texts"] + assert data["messages"][1]["content"][0]["text"] == "discarded text" + assert data["messages"][1]["content"][1]["text"] == "[MASKED]" + + @pytest.mark.asyncio + async def test_string_midturn_system_correction_is_guardrailed(self): + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [{"role": "system", "content": "prohibited correction"}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + assert guardrail.inputs["texts"] == ["prohibited correction"] + assert data["messages"][0]["content"] == "[MASKED]" + + @pytest.mark.asyncio + async def test_unsupported_midturn_system_content_is_not_guardrailed(self): + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + { + "role": "system", + "content": [{"type": "image", "source": {"type": "url"}}], + } + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is None + + @pytest.mark.asyncio + async def test_skip_system_message_excludes_only_hoisted_top_level_system(self): + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "user", "content": "safe text"}, + {"role": "system", "content": "prohibited correction"}, + {"role": "user", "content": "continue"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + structured = guardrail.inputs["structured_messages"] + assert [m["role"] for m in structured] == ["user", "system", "user"] + assert structured[1]["content"] == "prohibited correction" + + @pytest.mark.asyncio + async def test_default_skip_false_scans_midturn_system_and_hoists_top_level_system( + self, + ): + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail(skip_system_message_in_guardrail=None) + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "user", "content": "safe text"}, + {"role": "system", "content": "prohibited correction"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + assert guardrail.inputs["texts"] == ["safe text", "prohibited correction"] + structured = guardrail.inputs["structured_messages"] + assert [m["role"] for m in structured] == ["system", "user", "system"] + assert structured[0]["content"] == "trusted top-level system prompt" + assert data["messages"][1]["content"] == "[MASKED]" + + @pytest.mark.asyncio + async def test_bedrock_masking_slice_is_unavailable_when_top_level_system_is_included( + self, + ): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + ) + + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail(skip_system_message_in_guardrail=None) + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "user", "content": "safe text"}, + {"role": "system", "content": "prohibited correction"}, + {"role": "user", "content": "latest question"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + texts = guardrail.inputs["texts"] + structured = guardrail.inputs["structured_messages"] + + bedrock = BedrockGuardrail(guardrailIdentifier="gi", guardrailVersion="1") + assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts) + 1 + latest_user_index = bedrock._find_latest_message_index(structured, target_role="user") + assert ( + bedrock._locate_message_texts_slice( + structured_messages=structured, + target_index=latest_user_index, + texts=texts, + ) + is None + ) + assert ( + bedrock._merge_masked_texts( + masked_texts=["{MASKED}"], + texts=texts, + scanned_slice=None, + scanned_role_subset=True, + ) + == texts + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("skip_system_message_in_guardrail", [True, None]) + async def test_midturn_system_text_extraction_matches_translation_in_both_skip_modes( + self, + skip_system_message_in_guardrail: Optional[bool], + ): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + ) + + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail(skip_system_message_in_guardrail=skip_system_message_in_guardrail) + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "safe text"}, + { + "role": "system", + "content": [ + {"type": "text", "text": ""}, + {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}, + {"type": "text", "text": "prohibited correction"}, + ], + }, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + texts = guardrail.inputs["texts"] + structured = guardrail.inputs["structured_messages"] + assert texts == ["safe text", "prohibited correction"] + bedrock = BedrockGuardrail(guardrailIdentifier="gi", guardrailVersion="1") + assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts) + assert data["messages"][1]["content"][2]["text"] == "[MASKED]" + + @pytest.mark.asyncio + async def test_bedrock_masking_slice_stays_aligned_with_midturn_system(self): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + ) + + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "user", "content": "safe text"}, + { + "role": "system", + "content": [ + {"type": "text", "text": "prohibited correction"}, + {"type": "text", "text": "second correction"}, + ], + }, + {"role": "user", "content": "latest question"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + texts = guardrail.inputs["texts"] + structured = guardrail.inputs["structured_messages"] + + bedrock = BedrockGuardrail(guardrailIdentifier="gi", guardrailVersion="1") + total = sum(bedrock._count_message_texts(m) for m in structured) + assert total == len(texts) + + latest_user_index = bedrock._find_latest_message_index(structured, target_role="user") + assert latest_user_index == 2 + scanned_slice = bedrock._locate_message_texts_slice( + structured_messages=structured, + target_index=latest_user_index, + texts=texts, + ) + assert scanned_slice == (3, 1) + + merged = bedrock._merge_masked_texts( + masked_texts=["{MASKED}"], + texts=texts, + scanned_slice=scanned_slice, + scanned_role_subset=True, + ) + assert merged == [ + "safe text", + "prohibited correction", + "second correction", + "{MASKED}", + ] + + @pytest.mark.asyncio + async def test_compaction_rewrite_keeps_midturn_system_messages(self): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "user", "content": "compacted history"}, + { + "role": "system", + "content": [{"type": "text", "text": "use the corrected result"}], + }, + {"role": "user", "content": "continue"}, + ] + ) + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "continue"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["user", "system", "user"] + assert data["messages"][1]["content"] == [{"type": "text", "text": "use the corrected result"}] + assert data["messages"][0]["content"] == [{"type": "text", "text": "compacted history"}] + assert data["messages"][2]["content"] == [{"type": "text", "text": "continue"}] + assert data["system"] == "trusted top-level system prompt" + + @pytest.mark.asyncio + async def test_compaction_rewrite_does_not_duplicate_hoisted_top_level_system(self): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "system", "content": "trusted top-level system prompt"}, + {"role": "user", "content": "compacted history"}, + {"role": "system", "content": "use the corrected result"}, + ] + ) + guardrail.skip_system_message_in_guardrail = None + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "use the corrected result"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["user", "system"] + assert data["messages"][1]["content"] == "use the corrected result" + assert data["system"] == "trusted top-level system prompt" + + @pytest.mark.asyncio + async def test_compaction_rewrite_keeps_leading_midturn_system_when_system_is_skipped( + self, + ): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "compacted history"}, + ] + ) + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "trusted top-level system prompt", + "messages": [ + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "original history"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["system", "user"] + assert data["messages"][0]["content"] == "use the corrected result" + + @pytest.mark.asyncio + async def test_compaction_rewrite_keeps_leading_correction_when_top_level_system_hoists_nothing( + self, + ): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "compacted history"}, + ] + ) + guardrail.skip_system_message_in_guardrail = None + data = { + "model": "claude-3-5-sonnet-20241022", + "system": [{"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}], + "messages": [ + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "original history"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["system", "user"] + assert data["messages"][0]["content"] == "use the corrected result" + + @pytest.mark.asyncio + async def test_compaction_rewrite_keeps_leading_correction_when_hoisted_prompt_is_dropped( + self, + ): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "system", "content": "CLIENT CORRECTION"}, + {"role": "user", "content": "compacted history"}, + ] + ) + guardrail.skip_system_message_in_guardrail = None + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "TRUSTED", + "messages": [ + {"role": "system", "content": "CLIENT CORRECTION"}, + {"role": "user", "content": "original history"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + assert guardrail.inputs["structured_messages"][0] == { + "role": "system", + "content": "TRUSTED", + } + assert [m["role"] for m in data["messages"]] == ["system", "user"] + assert data["messages"][0]["content"] == "CLIENT CORRECTION" + assert data["system"] == "TRUSTED" + + @pytest.mark.asyncio + async def test_compaction_rewrite_drops_hoisted_prompt_matched_by_content_copy(self): + import json + + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + json.loads(json.dumps({"role": "system", "content": "TRUSTED"})), + {"role": "user", "content": "compacted history"}, + {"role": "system", "content": "CLIENT CORRECTION"}, + ] + ) + guardrail.skip_system_message_in_guardrail = None + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "TRUSTED", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "CLIENT CORRECTION"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["user", "system"] + assert data["messages"][1]["content"] == "CLIENT CORRECTION" + assert data["system"] == "TRUSTED" + + @pytest.mark.asyncio + async def test_compaction_rewrite_preserves_cache_control_on_system_blocks(self): + """ + `cache_control` on an in-sequence system text block survives the write-back, and is + copied rather than aliased into the guardrail's own returned list. + """ + handler = AnthropicMessagesHandler() + source_cache_control = {"type": "ephemeral"} + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "user", "content": "compacted history"}, + { + "role": "system", + "content": [ + { + "type": "text", + "text": "use the corrected result", + "cache_control": source_cache_control, + } + ], + }, + ] + ) + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "use the corrected result"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["messages"][1]["content"] == [ + { + "type": "text", + "text": "use the corrected result", + "cache_control": {"type": "ephemeral"}, + } + ] + assert data["messages"][1]["content"][0]["cache_control"] is not source_cache_control + + @pytest.mark.asyncio + async def test_compaction_rewrite_rstrips_trailing_assistant_in_each_run(self): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "user", "content": "compacted history"}, + {"role": "assistant", "content": "earlier "}, + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "continue"}, + {"role": "assistant", "content": "prefill "}, + ] + ) + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "use the corrected result"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == [ + "user", + "assistant", + "system", + "user", + "assistant", + ] + assert data["messages"][1]["content"] == [{"type": "text", "text": "earlier"}] + assert data["messages"][-1]["content"] == [{"type": "text", "text": "prefill"}] + + @pytest.mark.asyncio + async def test_compaction_rewrite_drops_text_free_system_message(self): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[ + {"role": "user", "content": "compacted history"}, + {"role": "system", "content": [{"type": "text", "text": ""}]}, + {"role": "system", "content": ""}, + {"role": "user", "content": "continue"}, + ] + ) + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "continue"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["user", "user"] + assert data["messages"][0]["content"] == [{"type": "text", "text": "compacted history"}] + assert data["messages"][1]["content"] == [{"type": "text", "text": "continue"}] + + @pytest.mark.asyncio + async def test_noncanonical_system_role_casing_is_still_scanned(self): + handler = AnthropicMessagesHandler() + guardrail = MockMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "safe text"}, + {"role": "System", "content": "prohibited correction"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.inputs is not None + assert "prohibited correction" in guardrail.inputs["texts"] + assert data["messages"][1]["content"] == "[MASKED]" + + @pytest.mark.asyncio + async def test_tool_result_turns_have_a_preexisting_alignment_gap(self): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + ) + + handler = AnthropicMessagesHandler() + bedrock = BedrockGuardrail(guardrailIdentifier="gi", guardrailVersion="1") + tool_loop = [ + {"role": "user", "content": "call the tool"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tu_1", "name": "get", "input": {"a": 1}}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tu_1", + "content": [{"type": "text", "text": "tool output"}], + } + ], + }, + ] + + async def _slice_for(messages: list): + guardrail = MockMaskingGuardrail() + data = {"model": "claude-3-5-sonnet-20241022", "messages": messages} + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + assert guardrail.inputs is not None + texts = guardrail.inputs["texts"] + structured = guardrail.inputs["structured_messages"] + target_index = bedrock._find_latest_message_index(structured, target_role="user") + return ( + sum(bedrock._count_message_texts(m) for m in structured) - len(texts), + bedrock._locate_message_texts_slice( + structured_messages=structured, + target_index=target_index, + texts=texts, + ), + ) + + with_system = await _slice_for( + tool_loop + + [ + {"role": "system", "content": "use the corrected result"}, + {"role": "user", "content": "latest question"}, + ] + ) + without_system = await _slice_for(tool_loop + [{"role": "user", "content": "latest question"}]) + + assert with_system == without_system == (1, None) + + @pytest.mark.asyncio + async def test_compaction_rewrite_to_only_system_messages_is_rejected(self): + import litellm + + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[{"role": "system", "content": "use the corrected result"}] + ) + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "use the corrected result"}, + ], + } + + with patch.object(litellm, "modify_params", False): + with pytest.raises(litellm.BadRequestError, match="at least one non-system message"): + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + @pytest.mark.asyncio + async def test_compaction_rewrite_to_only_system_messages_repaired_with_modify_params( + self, + ): + import litellm + + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail( + replacement_messages=[{"role": "system", "content": "use the corrected result"}] + ) + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "original history"}, + {"role": "system", "content": "use the corrected result"}, + ], + } + + with patch.object(litellm, "modify_params", True): + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["system", "user"] + assert data["messages"][0]["content"] == "use the corrected result" + + @pytest.mark.asyncio + async def test_compaction_rewrite_without_system_messages_is_unchanged(self): + handler = AnthropicMessagesHandler() + guardrail = MockCompactingGuardrail(replacement_messages=[{"role": "user", "content": "compacted history"}]) + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "user", "content": "a"}, + {"role": "assistant", "content": "b"}, + {"role": "user", "content": "c"}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["messages"] == [{"role": "user", "content": [{"type": "text", "text": "compacted history"}]}] + @pytest.mark.asyncio async def test_process_output_streaming_response_empty_choices(self): """Test that streaming response with empty choices doesn't raise IndexError diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index c0c6e315b5b..b72620f9918 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -413,6 +413,224 @@ def test_translate_anthropic_messages_to_openai_tool_message_placement(): ), "Tool message should be placed before user message" +@pytest.mark.parametrize( + ("system_content", "expected_content"), + [ + ("Use the corrected result.", "Use the corrected result."), + ( + [{"type": "text", "text": "Use the corrected result."}], + [{"type": "text", "text": "Use the corrected result."}], + ), + ( + [ + { + "type": "image", + "source": {"type": "url", "url": "https://example.com/a.png"}, + }, + {"type": "text", "text": "Use the corrected result."}, + ], + [{"type": "text", "text": "Use the corrected result."}], + ), + ( + [ + {"type": "text", "text": "First correction."}, + {"type": "text", "text": "Second correction."}, + ], + [ + {"type": "text", "text": "First correction."}, + {"type": "text", "text": "Second correction."}, + ], + ), + ], +) +def test_translate_anthropic_messages_to_openai_preserves_midturn_system_correction( + system_content: object, + expected_content: object, +): + messages = [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_01234", + "name": "get_weather", + "input": {"location": "Boston"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01234", + "content": "Rainy, 55°F", + } + ], + }, + {"role": "system", "content": system_content}, + {"role": "user", "content": "Continue."}, + ] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=messages, + model="claude-3-5-sonnet-20240620", + ) + + assert result == [ + { + "role": "assistant", + "content": None, + "thinking_blocks": None, + "tool_calls": [ + { + "id": "toolu_01234", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "Boston"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "toolu_01234", + "content": "Rainy, 55°F", + }, + {"role": "system", "content": expected_content}, + {"role": "user", "content": "Continue."}, + ] + + +def test_translate_anthropic_messages_to_openai_preserves_midturn_system_cache_control(): + """ + `cache_control` on an in-sequence system text block survives, matching how the + hoisted top-level `system` prompt and user text blocks are already handled. + """ + messages = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "Use the corrected result.", + "cache_control": {"type": "ephemeral"}, + } + ], + } + ] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=messages, + model="claude-3-5-sonnet-20240620", + ) + + assert result == [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "Use the corrected result.", + "cache_control": {"type": "ephemeral"}, + } + ], + } + ] + + +def test_translate_anthropic_messages_to_openai_drops_midturn_system_cache_control_for_non_claude(): + """ + `cache_control` goes through the same `_add_cache_control_if_applicable` gate as the + hoisted top-level prompt and user text blocks, so a non-Claude *requested model name* + does not get it. That gate is a best-effort check of the requested name before routing + (behind the proxy it is often a public alias), not a guarantee about the backend that + ultimately serves the request. + """ + messages = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "Use the corrected result.", + "cache_control": {"type": "ephemeral"}, + } + ], + } + ] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=messages, + model="gpt-4o", + ) + + assert result == [ + { + "role": "system", + "content": [{"type": "text", "text": "Use the corrected result."}], + } + ] + + +@pytest.mark.parametrize( + "system_content", + [ + "", + [{"type": "text", "text": ""}], + [ + { + "type": "image", + "source": {"type": "url", "url": "https://example.com/a.png"}, + } + ], + None, + ], +) +def test_translate_anthropic_messages_to_openai_drops_empty_midturn_system( + system_content: object, +): + messages = [{"role": "system", "content": system_content}] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=messages, + model="claude-3-5-sonnet-20240620", + ) + + assert result == [] + + +def test_translate_anthropic_to_openai_orders_top_level_and_midturn_system(): + """ + Request level: the trusted top-level prompt is hoisted to index 0 exactly once and the + in-sequence correction keeps its own position and `role: "system"` -- no duplication of + either, and no reordering of the surrounding turns. + """ + openai_request, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + anthropic_message_request={ + "model": "claude-3-5-sonnet-20240620", + "max_tokens": 100, + "system": "Trusted top-level prompt.", + "messages": [ + {"role": "user", "content": "First question."}, + {"role": "assistant", "content": "First answer."}, + {"role": "system", "content": "Use the corrected result."}, + {"role": "user", "content": "Continue."}, + ], + } + ) + + assert openai_request["messages"] == [ + {"role": "system", "content": "Trusted top-level prompt."}, + {"role": "user", "content": "First question."}, + {"role": "assistant", "content": "First answer.", "thinking_blocks": None}, + {"role": "system", "content": "Use the corrected result."}, + {"role": "user", "content": "Continue."}, + ] + + def test_translate_openai_content_to_anthropic_empty_function_arguments(): """Test that empty function arguments are handled safely and don't cause JSON parsing errors.""" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py index 9c8df1c79f9..6cc1d9e5add 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -2475,3 +2475,57 @@ def test_endpoint_runs_failure_hook_on_500_context_management_error(): body = response.json() assert body["type"] == "error" failure_hook.assert_awaited_once() + + +def test_count_effective_tokens_counts_midturn_system_correction(): + from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import ( + _count_effective_tokens, + ) + + base: List[Dict[str, Any]] = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + ] + correction = { + "role": "system", + "content": [{"type": "text", "text": "use the corrected result " * 20}], + } + + without_correction = _count_effective_tokens( + model=MODEL, effective_messages=base, compaction_block=None, tools=None + ) + with_correction = _count_effective_tokens( + model=MODEL, + effective_messages=base + [correction], + compaction_block=None, + tools=None, + ) + + assert with_correction > without_correction + + +def test_build_summary_messages_keeps_midturn_system_correction_in_place(): + from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import ( + _build_summary_messages, + ) + + summary_messages = _build_summary_messages( + effective_messages=[ + {"role": "user", "content": "original question"}, + {"role": "system", "content": "use the corrected result"}, + {"role": "assistant", "content": "acknowledged"}, + ], + prompt="summarize the conversation", + system="caller system prompt", + ) + + assert [m["role"] for m in summary_messages] == [ + "system", + "user", + "system", + "assistant", + "user", + ] + assert summary_messages[0]["content"] == "caller system prompt" + assert summary_messages[2]["content"] == "use the corrected result" + assert summary_messages[-1]["content"] == "summarize the conversation" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 606ff39b35e..8963012ecd5 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -9,6 +9,8 @@ import sys from typing import Any, Dict, List from unittest.mock import MagicMock +import pytest + sys.path.insert(0, os.path.abspath("../../../../../../..")) from litellm.constants import ( @@ -221,6 +223,106 @@ class TestTranslateMessagesToResponsesInput: {"type": "input_text", "text": "Second part."}, ] + @pytest.mark.parametrize( + "system_content", + [ + "Use the corrected result.", + [{"type": "text", "text": "Use the corrected result."}], + [ + {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}, + {"type": "text", "text": "Use the corrected result."}, + ], + ], + ) + def test_midturn_system_correction_stays_system_in_sequence(self, system_content: object): + messages = [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_01234", + "name": "get_weather", + "input": {"location": "Boston"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01234", + "content": "Rainy, 55°F", + } + ], + }, + {"role": "system", "content": system_content}, + {"role": "user", "content": "Continue."}, + ] + + result = _translate_messages(messages) + + assert result == [ + { + "type": "function_call", + "call_id": "toolu_01234", + "name": "get_weather", + "arguments": '{"location": "Boston"}', + }, + { + "type": "function_call_output", + "call_id": "toolu_01234", + "output": "Rainy, 55°F", + }, + { + "type": "message", + "role": "system", + "content": [{"type": "input_text", "text": "Use the corrected result."}], + }, + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Continue."}], + }, + ] + + def test_midturn_system_correction_keeps_multiple_text_blocks(self): + messages = [ + { + "role": "system", + "content": [ + {"type": "text", "text": "First correction."}, + {"type": "text", "text": "Second correction."}, + ], + } + ] + + assert _translate_messages(messages) == [ + { + "type": "message", + "role": "system", + "content": [ + {"type": "input_text", "text": "First correction."}, + {"type": "input_text", "text": "Second correction."}, + ], + } + ] + + @pytest.mark.parametrize( + "system_content", + [ + "", + [{"type": "text", "text": ""}], + [{"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}], + None, + ], + ) + def test_empty_or_unsupported_midturn_system_correction_is_dropped(self, system_content: object): + messages = [{"role": "system", "content": system_content}] + + assert _translate_messages(messages) == [] + def test_user_base64_image(self): """User message with base64 image source becomes input_image with data URL.""" messages = [ @@ -722,6 +824,42 @@ class TestTranslateRequestBroaderCoverage: kwargs = _ADAPTER.translate_request(req) assert kwargs["instructions"] == "You are a helpful assistant." + def test_top_level_system_and_midturn_correction_are_not_duplicated(self): + """ + Request level: the trusted top-level prompt goes to `instructions` only, and the + in-sequence correction stays a `role: "system"` input item in its original position. + Neither appears twice, and the surrounding turns keep their order. + """ + req = _make_request( + system="Trusted top-level prompt.", + messages=[ + {"role": "user", "content": "First question."}, + {"role": "system", "content": "Use the corrected result."}, + {"role": "user", "content": "Continue."}, + ], + ) + + kwargs = _ADAPTER.translate_request(req) + + assert kwargs["instructions"] == "Trusted top-level prompt." + assert kwargs["input"] == [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "First question."}], + }, + { + "type": "message", + "role": "system", + "content": [{"type": "input_text", "text": "Use the corrected result."}], + }, + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Continue."}], + }, + ] + def test_system_list_of_text_blocks_joined(self): req = _make_request( system=[ From c5c5a276790529e2de3378654864fd847530c5a0 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 2 Aug 2026 00:19:42 +0000 Subject: [PATCH 024/576] fix(files): enforce require_managed_files on file retrieve, content and delete Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../openai_files_endpoints/common_utils.py | 30 +++++ .../openai_files_endpoints/files_endpoints.py | 7 ++ .../test_files_endpoint.py | 116 ++++++++++++++++++ 3 files changed, 153 insertions(+) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 87514b46dbd..3eef3868c94 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -866,6 +866,36 @@ def validate_managed_files_requirement( ) +def validate_managed_file_id_requirement(file_id: str) -> None: + """ + Enforce proxy-level managed files on the file read/delete routes when + ``litellm.require_managed_files`` is enabled. + + Ownership is only recorded for LiteLLM managed files, so a raw provider file id sent to + retrieve/content/delete is forwarded to the provider under shared credentials without any + tenant check; knowing another tenant's provider file id would be enough to read or delete it. + + Raises: + HTTPException: 400 if ``file_id`` is not a LiteLLM managed file id. + """ + import litellm + from fastapi import HTTPException + + if litellm.require_managed_files is not True: + return + + if _is_base64_encoded_unified_file_id(file_id): + return + + raise HTTPException( + status_code=400, + detail=( + "Raw provider file ids cannot be used when require_managed_files is enabled in " + "litellm_settings. Use the LiteLLM managed file id returned when the file was created." + ), + ) + + def _extract_model_param(request: "Request", request_body: dict) -> str | None: """ Extract model parameter from request. diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 4e4718272bd..37f1ced6996 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -49,6 +49,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( get_credentials_for_model, handle_model_based_routing, prepare_data_with_credentials, + validate_managed_file_id_requirement, validate_managed_files_requirement, ) from litellm.proxy.utils import ProxyLogging, is_known_model @@ -612,6 +613,8 @@ async def get_file_content( data: dict = {"file_id": file_id} try: + validate_managed_file_id_requirement(file_id=file_id) + # Include original request and headers in the data base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) ( @@ -908,6 +911,8 @@ async def get_file( data: dict = {"file_id": file_id} try: + validate_managed_file_id_requirement(file_id=file_id) + custom_llm_provider = ( provider or get_custom_llm_provider_from_request_headers(request=request) @@ -1098,6 +1103,8 @@ async def delete_file( data: dict = {"file_id": file_id} try: + validate_managed_file_id_requirement(file_id=file_id) + custom_llm_provider = ( provider or get_custom_llm_provider_from_request_headers(request=request) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index ac01c6ae1d1..24b814bae1f 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -3051,3 +3051,119 @@ def test_list_files_key_allowed_openai_model_still_resolves_team_credentials( mocker, monkeypatch, _team_openai_plus_global_anthropic_router(), ["team-gpt"] ) assert captured_kwargs.get("api_key") == "team-openai-key" + + +@pytest.mark.parametrize( + "http_method, url, patched_litellm_call", + [ + ("get", "/v1/files/file-victim-abc123", "litellm.afile_retrieve"), + ("get", "/v1/files/file-victim-abc123/content", "litellm.afile_content"), + ("delete", "/v1/files/file-victim-abc123", "litellm.afile_delete"), + ], +) +def test_require_managed_files_rejects_raw_provider_file_id( + mocker: MockerFixture, + monkeypatch, + llm_router: Router, + http_method: str, + url: str, + patched_litellm_call: str, +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr("litellm.require_managed_files", True) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + mock_call = mocker.patch(patched_litellm_call, new=mocker.AsyncMock()) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="attacker-user" + ) + + try: + response = getattr(client, http_method)( + url, headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + monkeypatch.setattr("litellm.require_managed_files", False) + + assert response.status_code == 400, response.text + mock_call.assert_not_called() + + +def _unified_managed_file_id() -> str: + import base64 + + from litellm.types.utils import SpecialEnums + + unified_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "victim-unified-id", "gpt-3.5-turbo", "file-victim-abc123", "gpt-3.5-turbo-id" + ) + return base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=") + + +def test_require_managed_files_allows_unified_managed_file_id(monkeypatch): + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_file_id_requirement, + ) + + monkeypatch.setattr("litellm.require_managed_files", True) + + validate_managed_file_id_requirement(file_id=_unified_managed_file_id()) + + +def test_managed_file_id_requirement_is_opt_in(monkeypatch): + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_file_id_requirement, + ) + + monkeypatch.setattr("litellm.require_managed_files", False) + + validate_managed_file_id_requirement(file_id="file-victim-abc123") + + +def test_raw_provider_file_id_retrieve_allowed_when_managed_files_not_required( + mocker: MockerFixture, monkeypatch, llm_router: Router +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr("litellm.require_managed_files", False) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + mock_retrieve = mocker.patch( + "litellm.afile_retrieve", + new=mocker.AsyncMock( + return_value=OpenAIFileObject( + id="file-victim-abc123", + object="file", + bytes=3, + created_at=1234567890, + filename="test.txt", + purpose="user_data", + status="uploaded", + ) + ), + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="some-user" + ) + + try: + response = client.get( + "/v1/files/file-victim-abc123", headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + mock_retrieve.assert_called_once() From 81a80b8c632a9b5d942193bb71eda601630e9906 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 2 Aug 2026 01:02:27 +0000 Subject: [PATCH 025/576] fix(anthropic): preserve speed=fast in usage for /v1/messages and pass-through Fast mode is priced with a provider-specific multiplier applied off usage.speed, but only chat completions kept that field. The Messages route rebuilt usage with empty optional params, stream reassembly dropped speed and inference_geo, and the pass-through handler never read speed off the request body, so fast-mode spend was logged at the standard rate. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 2 +- .../streaming_chunk_builder_utils.py | 16 +++ .../anthropic_passthrough_logging_handler.py | 24 +++- .../streaming_chunk_builder_utils.py | 2 + .../test_litellm_logging.py | 35 ++++++ .../test_streaming_chunk_builder_utils.py | 39 ++++++ ...t_anthropic_passthrough_logging_handler.py | 113 ++++++++++++++++++ 7 files changed, 228 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c2dc7189934..15330a2a910 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3278,7 +3278,7 @@ class Logging(LiteLLMLoggingBaseClass): model=self.model, messages=[], logging_obj=self, - optional_params={}, + optional_params=self.optional_params or {}, api_key="", request_data={}, encoding=litellm.encoding, diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index d52d9849310..5ceb64547a0 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -569,6 +569,8 @@ class ChunkProcessor: # lost and 1h cache writes get billed at the 5m rate. cache_creation_token_details: Optional[CacheCreationTokenDetails] = None cost: Optional[float] = None + inference_geo: Optional[str] = None + speed: Optional[str] = None for chunk in chunks: usage_chunk = self._extract_usage_chunk(chunk) @@ -627,6 +629,13 @@ class ChunkProcessor: if usage_chunk_dict["cost"] is not None: cost = usage_chunk_dict["cost"] + chunk_inference_geo = getattr(usage_chunk, "inference_geo", None) + if isinstance(chunk_inference_geo, str): + inference_geo = chunk_inference_geo + chunk_speed = getattr(usage_chunk, "speed", None) + if isinstance(chunk_speed, str): + speed = chunk_speed + prompt_tokens_details = self._attach_cache_creation_token_details( prompt_tokens_details, cache_creation_token_details ) @@ -647,6 +656,8 @@ class ChunkProcessor: completion_tokens_details=completion_tokens_details, prompt_tokens_details=prompt_tokens_details, cost=cost, + inference_geo=inference_geo, + speed=speed, ) @staticmethod @@ -806,6 +817,11 @@ class ChunkProcessor: if cost is not None: setattr(returned_usage, "cost", cost) + if calculated_usage_per_chunk["inference_geo"] is not None: + setattr(returned_usage, "inference_geo", calculated_usage_per_chunk["inference_geo"]) + if calculated_usage_per_chunk["speed"] is not None: + setattr(returned_usage, "speed", calculated_usage_per_chunk["speed"]) + # Return a new usage object with the new values returned_usage = Usage(**returned_usage.model_dump()) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 50e90699194..112780581de 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -80,7 +80,9 @@ class AnthropicPassthroughLoggingHandler: model=model, messages=[], logging_obj=logging_obj, - optional_params={}, + optional_params=AnthropicPassthroughLoggingHandler._cost_relevant_request_params( + request_body or kwargs.get("request_body") + ), api_key="", request_data={}, encoding=litellm.encoding, @@ -102,6 +104,15 @@ class AnthropicPassthroughLoggingHandler: "kwargs": kwargs, } + @staticmethod + def _cost_relevant_request_params(request_body: Optional[dict]) -> dict: + """ + Request params that change how the response is priced, and so must reach the + usage-building paths. Anthropic's ``speed=fast`` multiplies non-cache token cost. + """ + speed = (request_body or {}).get("speed") + return {"speed": speed} if isinstance(speed, str) else {} + @staticmethod def _get_user_from_metadata( passthrough_logging_payload: PassthroughStandardLoggingPayload, @@ -315,6 +326,7 @@ class AnthropicPassthroughLoggingHandler: - Logs in litellm callbacks """ + speed = AnthropicPassthroughLoggingHandler._cost_relevant_request_params(request_body).get("speed") model = request_body.get("model", "") # Check if it's available in the logging object if ( @@ -334,6 +346,7 @@ class AnthropicPassthroughLoggingHandler: all_chunks=all_chunks, litellm_logging_obj=litellm_logging_obj, model=model, + speed=speed, ) except Exception as e: # stream_chunk_builder re-raises assembly failures (as litellm.APIError) @@ -355,6 +368,7 @@ class AnthropicPassthroughLoggingHandler: complete_streaming_response = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( all_chunks=all_chunks, model=model, + speed=speed, ) except Exception as e: verbose_proxy_logger.warning( @@ -419,6 +433,7 @@ class AnthropicPassthroughLoggingHandler: all_chunks: Sequence[Union[str, bytes]], litellm_logging_obj: LiteLLMLoggingObj, model: str, + speed: Optional[str] = None, ) -> Optional[Union[ModelResponse, TextCompletionResponse]]: """ Builds complete response from raw Anthropic chunks. @@ -443,11 +458,13 @@ class AnthropicPassthroughLoggingHandler: all_chunks=collapsed, litellm_logging_obj=litellm_logging_obj, model=model, + speed=speed, ) return AnthropicPassthroughLoggingHandler._build_complete_streaming_response_legacy( all_chunks=all_chunks, litellm_logging_obj=litellm_logging_obj, model=model, + speed=speed, ) # Anthropic SSE block/delta types that the fast path is NOT allowed to @@ -575,6 +592,7 @@ class AnthropicPassthroughLoggingHandler: all_chunks: Sequence[Union[str, bytes]], litellm_logging_obj: LiteLLMLoggingObj, model: str, + speed: Optional[str] = None, ) -> Optional[Union[ModelResponse, TextCompletionResponse]]: """ Original reconstruction: convert every SSE event to a generic chunk @@ -590,6 +608,7 @@ class AnthropicPassthroughLoggingHandler: anthropic_model_response_iterator = AnthropicModelResponseIterator( streaming_response=None, sync_stream=False, + speed=speed, ) all_openai_chunks = [] @@ -649,6 +668,7 @@ class AnthropicPassthroughLoggingHandler: def _build_usage_only_response_from_chunks( all_chunks: Sequence[Union[str, bytes]], model: str, + speed: Optional[str] = None, ) -> Optional[ModelResponse]: """ Build a usage-bearing ModelResponse from Anthropic SSE token-usage events, for @@ -742,7 +762,7 @@ class AnthropicPassthroughLoggingHandler: usage_object["server_tool_use"] = _server_tool_use if inference_geo is not None: usage_object["inference_geo"] = inference_geo - usage_obj = AnthropicConfig().calculate_usage(usage_object=usage_object, reasoning_content=None) + usage_obj = AnthropicConfig().calculate_usage(usage_object=usage_object, reasoning_content=None, speed=speed) return ModelResponse( model=resolved_model, choices=[ diff --git a/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py index c9f9d4e6baa..4ff47171e32 100644 --- a/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py @@ -15,3 +15,5 @@ class UsagePerChunk(TypedDict): completion_tokens_details: Optional[CompletionTokensDetails] prompt_tokens_details: Optional[PromptTokensDetailsWrapper] cost: Optional[float] + inference_geo: Optional[str] + speed: Optional[str] diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index eaa4bd3e3fc..1cd4174c57c 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -4153,3 +4153,38 @@ def test_pre_call_does_not_pin_request_in_module_state(logging_obj): logging_obj.post_call(original_response='{"ok": true}', input=big_input, api_key="sk-test") assert litellm.error_logs == {} + + +def test_handle_anthropic_messages_response_logging_preserves_fast_mode_speed(): + """/v1/messages non-streaming rebuilds usage by re-transforming the raw Anthropic + response. Anthropic's fast-mode multiplier is applied off ``usage.speed``, which the + response body never carries, so the request's optional params have to be passed in or + fast-mode spend is logged at the standard rate.""" + import httpx + + logging_obj = LitellmLogging( + model="claude-opus-4-8", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="anthropic_messages", + start_time=time.time(), + litellm_call_id="lit-5115", + function_id="lit-5115", + ) + logging_obj.optional_params = {"speed": "fast"} + logging_obj.model_call_details["httpx_response"] = httpx.Response( + status_code=200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1000, "cache_read_input_tokens": 200, "output_tokens": 100}, + }, + ) + + result = logging_obj._handle_anthropic_messages_response_logging(result=None) + + assert result.usage.speed == "fast" # type: ignore[attr-defined] diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index be8c5a05601..72bbad7ae49 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -992,3 +992,42 @@ def test_cost_field_in_usage_chunks(): assert usage.cost == 0.00025 assert usage.prompt_tokens == 10 assert usage.completion_tokens == 5 + + +def test_anthropic_speed_and_geo_survive_stream_assembly(): + """Anthropic prices fast mode and non-global regions with a multiplier read off + ``usage.speed`` / ``usage.inference_geo``. Dropping them while reassembling a stream + bills streamed fast-mode calls at the standard rate.""" + from litellm.llms.anthropic.cost_calculation import cost_per_token + + def _usage(**extra): + usage = Usage(completion_tokens=100, prompt_tokens=1000, total_tokens=1100) + for key, value in extra.items(): + setattr(usage, key, value) + return usage + + def _chunk(usage): + return ModelResponseStream( + id="chatcmpl-1", + created=1745513206, + model="claude-opus-4-8", + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="Hi"))], + usage=usage, + ) + + fast_chunk = _chunk(_usage(speed="fast", inference_geo="global")) + fast_usage = ChunkProcessor(chunks=[fast_chunk]).calculate_usage( + chunks=[fast_chunk], model="claude-opus-4-8", completion_output="Hi" + ) + standard_chunk = _chunk(_usage(inference_geo="global")) + standard_usage = ChunkProcessor(chunks=[standard_chunk]).calculate_usage( + chunks=[standard_chunk], model="claude-opus-4-8", completion_output="Hi" + ) + + assert fast_usage.speed == "fast" + assert fast_usage.inference_geo == "global" + assert getattr(standard_usage, "speed", None) is None + + fast_cost = sum(cost_per_token(model="claude-opus-4-8", usage=fast_usage)) + standard_cost = sum(cost_per_token(model="claude-opus-4-8", usage=standard_usage)) + assert fast_cost == pytest.approx(standard_cost * 2.0) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 947a7a64beb..ffbdb3485ae 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -2192,3 +2192,116 @@ class TestAnthropicResponseCostRecordedOnModelCallDetails: logging_obj.model_call_details["response_cost"] == kwargs["response_cost"] ) assert logging_obj.model_call_details["response_cost"] > 0 + + +class TestAnthropicPassthroughFastMode: + """Anthropic charges a provider-specific multiplier for ``speed=fast``, and the + multiplier is applied off ``usage.speed``. The pass-through handler only sees the + speed in the request body, so it has to thread it into every usage-building path or + fast-mode pass-through spend is under-reported.""" + + MODEL = "claude-opus-4-8" + STREAM_CHUNKS = [ + 'event: message_start', + 'data: {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant",' + ' "model": "claude-opus-4-8", "content": [], "stop_reason": null,' + ' "usage": {"input_tokens": 1000, "cache_read_input_tokens": 200, "output_tokens": 0}}}', + 'event: content_block_start', + 'data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}', + 'event: content_block_delta', + 'data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}', + 'event: content_block_stop', + 'data: {"type": "content_block_stop", "index": 0}', + 'event: message_delta', + 'data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"},' + ' "usage": {"input_tokens": 1000, "cache_read_input_tokens": 200, "output_tokens": 100}}', + 'event: message_stop', + 'data: {"type": "message_stop"}', + ] + + def _logging_obj(self) -> LiteLLMLoggingObj: + return LiteLLMLoggingObj( + model=self.MODEL, + messages=[], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="fast-mode", + function_id="fast-mode", + ) + + def _cost(self, response) -> float: + import litellm + + return litellm.completion_cost(completion_response=response, model=f"anthropic/{self.MODEL}") + + def _expected_fast_cost(self, standard_cost: float) -> float: + import litellm + + model_info = litellm.get_model_info(model=self.MODEL, custom_llm_provider="anthropic") + cache_read_cost = 200 * (model_info.get("cache_read_input_token_cost") or 0.0) + return (standard_cost - cache_read_cost) * 2.0 + cache_read_cost + + def test_non_streaming_applies_fast_multiplier(self): + import httpx + + response_body = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": self.MODEL, + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1000, "cache_read_input_tokens": 200, "output_tokens": 100}, + } + + def _handle(request_body): + logging_obj = self._logging_obj() + logging_obj.model_call_details["stream"] = False + return AnthropicPassthroughLoggingHandler.anthropic_passthrough_handler( + httpx_response=httpx.Response(status_code=200, json=response_body), + response_body=response_body, + logging_obj=logging_obj, + url_route="https://api.anthropic.com/v1/messages", + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body=request_body, + ) + + fast = _handle({"model": self.MODEL, "speed": "fast"}) + standard = _handle({"model": self.MODEL}) + + assert fast["result"].usage.speed == "fast" + assert self._cost(fast["result"]) == pytest.approx(self._expected_fast_cost(self._cost(standard["result"]))) + + def test_streaming_reconstruction_applies_fast_multiplier(self): + fast = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=self.STREAM_CHUNKS, + litellm_logging_obj=self._logging_obj(), + model=self.MODEL, + speed="fast", + ) + standard = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=self.STREAM_CHUNKS, + litellm_logging_obj=self._logging_obj(), + model=self.MODEL, + ) + + assert fast.usage.speed == "fast" + assert self._cost(fast) == pytest.approx(self._expected_fast_cost(self._cost(standard))) + + def test_usage_only_fallback_applies_fast_multiplier(self): + fast = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=self.STREAM_CHUNKS, + model=self.MODEL, + speed="fast", + ) + standard = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=self.STREAM_CHUNKS, + model=self.MODEL, + ) + + assert fast.usage.speed == "fast" + assert self._cost(fast) == pytest.approx(self._expected_fast_cost(self._cost(standard))) From ef614b7b5bcdf94472b876bea64ad17e5dfd4282 Mon Sep 17 00:00:00 2001 From: milan Date: Mon, 3 Aug 2026 14:35:47 +0000 Subject: [PATCH 026/576] refactor(vertex_ai): keep the batch embeddings translation within the LIT002 ceiling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 118 +++++++++++------- 1 file changed, 71 insertions(+), 47 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 9363540fe1b..bbb97a1edc2 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -74,11 +74,11 @@ _GCP_LABEL_VALUE_MAX_LEN = 63 _CUSTOM_ID_RAW_LABEL_PREFIX = "b32_" _VERTEX_BATCH_KEY_FIELD = "key" _MANAGED_GCS_MODEL_PATH_PATTERN = re.compile(r"publishers/[^/]+/models/([^/?]+)") -_EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = { - "outputDimensionality": "output_dimensionality", - "taskType": "task_type", - "title": "title", -} +_EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = ( + ("outputDimensionality", "output_dimensionality"), + ("taskType", "task_type"), + ("title", "title"), +) _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN = re.compile(r"(?P[^#]*)#(?P\d+)/(?P\d+)") @@ -153,12 +153,15 @@ def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, Any]) -> str: key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) if key is not None: return unquote(str(key)) - request_data = vertex_output_row.get("request") or {} - return _get_litellm_batch_custom_id_from_labels(request_data.get("labels") or {}) + request_data = vertex_output_row.get("request") + labels = request_data.get("labels") if isinstance(request_data, Mapping) else None + return _get_litellm_batch_custom_id_from_labels(labels) -def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str: +def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, Any] | None) -> str: """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" + if not labels: + return "unknown" raw = labels.get("litellm_custom_id_raw") if raw: raw_chunks = [str(raw)] @@ -195,7 +198,8 @@ def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, Any]) def _openai_batch_output_row( custom_id: str, body: Mapping[str, Any] | None = None, - error: Mapping[str, str] | None = None, + error_code: str | None = None, + error_message: str = "", ) -> Mapping[str, Any]: """ One row of an OpenAI batch output file. Per the OpenAI Batch spec, failed rows set @@ -211,7 +215,7 @@ def _openai_batch_output_row( "request_id": body.get("id", ""), "body": body, }, - "error": error, + "error": None if error_code is None else {"code": error_code, "message": error_message}, } @@ -233,6 +237,19 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, return unquote(match["custom_id"]), int(match["index"]) +def _embedding_prompt_token_count(vertex_response: Mapping[str, Any]) -> int: + """ + Prompt tokens billed for one Vertex Gemini Embedding batch row. + + Live rows report usage under `usageMetadata`; the documented `tokenCount` is kept as + a fallback. + """ + usage_metadata = vertex_response.get("usageMetadata") + if isinstance(usage_metadata, Mapping): + return int(usage_metadata.get("promptTokenCount") or 0) + return int(vertex_response.get("tokenCount") or 0) + + def _vertex_embeddings_rows_to_openai_batch_output_row( custom_id: str, vertex_output_rows: tuple[Mapping[str, Any], ...], @@ -247,23 +264,19 @@ def _vertex_embeddings_rows_to_openai_batch_output_row( An entry that asked for several embeddings at once maps to several rows here, which become the indexed elements of a single `data` array. One failed element fails the - whole entry, since an OpenAI batch row is either a response or an error. Live rows - report usage under `usageMetadata`; the documented `tokenCount` is kept as a - fallback. Rows carry no `modelVersion`, so the model comes from the batch they - belong to. + whole entry, since an OpenAI batch row is either a response or an error. Rows carry + no `modelVersion`, so the model comes from the batch they belong to. """ status = next((row["status"] for row in vertex_output_rows if row.get("status")), "") if status: return _openai_batch_output_row( custom_id=custom_id, - error={"code": "vertex_ai_error", "message": status}, + error_code="vertex_ai_error", + error_message=status, ) - responses = tuple(row.get("response") or {} for row in vertex_output_rows) - token_count = sum( - int((response.get("usageMetadata") or {}).get("promptTokenCount") or response.get("tokenCount") or 0) - for response in responses - ) + responses = tuple(row["response"] for row in vertex_output_rows) + token_count = sum(_embedding_prompt_token_count(response) for response in responses) body = EmbeddingResponse( model=model or "", data=[ @@ -361,6 +374,27 @@ def _vertex_batch_embeddings_key(custom_id: str, index: int, total: int) -> str: return encoded_custom_id if total < 2 else f"{encoded_custom_id}#{index}/{total}" +def _vertex_embeddings_row(key: str | None, embed_content_request: Mapping[str, Any]) -> Mapping[str, Any]: + """ + One Vertex Gemini Embedding batch input row. + + The config fields live inside the `EmbedContentRequest` under their snake_case batch + names, and the OpenAI `custom_id` rides along in the top-level `key` that Vertex + echoes back. + """ + request = { + "content": embed_content_request["content"], + **{ + request_field: embed_content_request[gemini_param] + for gemini_param, request_field in _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM + if gemini_param in embed_content_request + }, + } + if key is None: + return {"request": request} + return {_VERTEX_BATCH_KEY_FIELD: key, "request": request} + + def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( openai_entry: Mapping[str, Any], ) -> tuple[Mapping[str, Any], ...]: @@ -381,7 +415,9 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( API Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings """ - openai_request_body = openai_entry.get("body") or {} + openai_request_body = openai_entry.get("body") + if not isinstance(openai_request_body, dict): + raise ValueError("`body` is required on /v1/embeddings batch requests, but was not provided") embedding_input = openai_request_body.get("input") if embedding_input is None: raise ValueError("`input` is required on /v1/embeddings batch requests, but was not provided") @@ -400,27 +436,16 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( ) custom_id = openai_entry.get("custom_id") return tuple( - { - **( - {} - if custom_id is None - else { - _VERTEX_BATCH_KEY_FIELD: _vertex_batch_embeddings_key( - custom_id=str(custom_id), - index=index, - total=len(embed_content_requests), - ) - } + _vertex_embeddings_row( + key=None + if custom_id is None + else _vertex_batch_embeddings_key( + custom_id=str(custom_id), + index=index, + total=len(embed_content_requests), ), - "request": { - "content": embed_content_request["content"], - **{ - request_field: embed_content_request[gemini_param] - for gemini_param, request_field in _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM.items() - if gemini_param in embed_content_request - }, - }, - } + embed_content_request=embed_content_request, + ) for index, embed_content_request in enumerate(embed_content_requests) ) @@ -1008,7 +1033,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): request=httpx.Request(method="POST", url="https://example.com"), ) - all_lines = itertools.chain([first_line], lines) + all_lines = itertools.chain((first_line,), lines) # Embedding rows are grouped by `custom_id` rather than transformed one at a # time, since an entry that asked for several embeddings comes back as @@ -1064,7 +1089,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if has_error: return _openai_batch_output_row( custom_id=custom_id, - error={"code": "vertex_ai_error", "message": status}, + error_code="vertex_ai_error", + error_message=status, ) # Transform successful response using existing transformation @@ -1096,8 +1122,6 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): except Exception as e: return _openai_batch_output_row( custom_id=custom_id, - error={ - "code": "transformation_error", - "message": f"Failed to transform response: {e!s}", - }, + error_code="transformation_error", + error_message=f"Failed to transform response: {e!s}", ) From a7250f4eeab560299215773b50ba32405f597a1c Mon Sep 17 00:00:00 2001 From: milan Date: Mon, 3 Aug 2026 17:11:50 +0000 Subject: [PATCH 027/576] fix(bedrock): normalize /v1/completions and /v1/responses batch records Bedrock managed-batch file upload read `messages` unconditionally, so a JSONL record shaped for /v1/completions (`prompt`) or /v1/responses (`input`) reached the per-provider transform with an empty message list. Anthropic and Nova rejected it at POST /v1/files, and the passthrough providers shipped an empty conversation to AWS. Classify each record by its OpenAI batch `url`, then normalize the non-embedding shapes to chat completions before the Bedrock transforms: `prompt` wraps into user messages the way litellm.text_completion does in real time, and `input` goes through the existing Responses-to-Chat bridge. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../prompt_templates/common_utils.py | 22 +- litellm/llms/bedrock/files/transformation.py | 194 ++++++++-- litellm/types/llms/bedrock.py | 15 + ...ore_utils_prompt_templates_common_utils.py | 43 +++ .../test_bedrock_files_transformation.py | 347 ++++++++++++++++-- 5 files changed, 559 insertions(+), 62 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 639c93dfb80..777ba398d5a 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -6,7 +6,7 @@ import io import json import mimetypes import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from os import PathLike from pathlib import Path from typing import ( @@ -1742,3 +1742,23 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, Any]]: idx = end_idx return results + + +def text_completion_prompt_to_messages(prompt: str | Sequence[str]) -> tuple[AllMessageValues, ...]: + """ + Wrap an OpenAI ``/v1/completions`` ``prompt`` into Chat Completion messages. + + Mirrors what ``litellm.text_completion`` does on the real-time path: a + string becomes a single user message, and a list of strings becomes one + user message per element. Pre-tokenized prompts (``list[int]`` / + ``list[list[int]]``) are only meaningful for the OpenAI-family text + endpoints, so they are rejected here rather than silently forwarded, as is + an empty prompt, which every chat-shaped provider rejects downstream. + """ + if isinstance(prompt, str) and prompt: + return (ChatCompletionUserMessage(role="user", content=prompt),) + if isinstance(prompt, Sequence) and prompt and all(isinstance(entry, str) and entry for entry in prompt): + return tuple(ChatCompletionUserMessage(role="user", content=entry) for entry in prompt) + raise ValueError( + f"`prompt` must be a non-empty string or a non-empty list of strings. Got: {type(prompt).__name__}." + ) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 3656088cb9d..0aa832780e5 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -2,7 +2,9 @@ import base64 import json import os import time -from collections.abc import Mapping, MutableMapping +from collections.abc import Iterable, Mapping, MutableMapping +from functools import cache +from itertools import chain from types import MappingProxyType from typing import ( Any, @@ -12,7 +14,7 @@ from urllib.parse import unquote import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -28,12 +30,16 @@ from litellm.litellm_core_utils.cloud_storage_security import ( split_configured_cloud_bucket_name, validate_managed_cloud_file_id, ) -from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + text_completion_prompt_to_messages, +) from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, ) +from litellm.types.llms.bedrock import BedrockBatchRecordKind from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, @@ -43,6 +49,8 @@ from litellm.types.llms.openai import ( OpenAICreateFileRequestOptionalParams, OpenAIFileObject, PathLike, + ResponseInputParam, + ResponsesAPIOptionalRequestParams, ) from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums from litellm.utils import get_llm_provider @@ -57,6 +65,26 @@ from ..common_utils import BedrockError S3_SIGNED_GET_HEADERS_PARAM = "_s3_signed_get_headers" +def _frozen_mapping(items: Iterable[tuple[str, Any]]) -> Mapping[str, Any]: + return MappingProxyType(dict(items)) + + +# JSONL batch records are untyped json, so the `/v1/responses` fields are +# validated into their concrete Responses API types before being handed to the +# Responses-to-Chat bridge. Both adapters drop keys the Responses API doesn't +# define, which is what the bridge would ignore anyway. Built on first use +# rather than at import: `ResponseInputParam` is a deep union and only batch +# files carrying `/v1/responses` records need it. +@cache +def _responses_input_adapter() -> TypeAdapter[str | ResponseInputParam]: + return TypeAdapter(str | ResponseInputParam) + + +@cache +def _responses_request_adapter() -> TypeAdapter[ResponsesAPIOptionalRequestParams]: + return TypeAdapter(ResponsesAPIOptionalRequestParams) + + class _BedrockS3RequestParams(BaseModel): """Typed view of the credential/region params the S3 GetObject path reads.""" @@ -305,41 +333,55 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # example; add others here as they adopt the same schema. CONVERSE_INVOKE_PROVIDERS = ("nova",) - # OpenAI batch URL that signals an embedding request. Per OpenAI Batch API - # spec, every JSONL record carries a `url` field; we use it as the - # authoritative signal to route the line to the embedding code path - # instead of inferring from the presence of `input` vs `messages`. + # OpenAI batch URLs that select which request shape a JSONL line carries. + # Per the OpenAI Batch API spec every record carries a `url`, so we use it + # as the authoritative routing signal instead of inferring from the + # presence of `input` vs `prompt` vs `messages`. OPENAI_EMBEDDINGS_URL = "/v1/embeddings" + OPENAI_TEXT_COMPLETIONS_URL = "/v1/completions" + OPENAI_RESPONSES_URL = "/v1/responses" @staticmethod - def _is_embedding_record(openai_jsonl_record: dict[str, Any]) -> bool: + def _classify_batch_record(openai_jsonl_record: Mapping[str, Any]) -> BedrockBatchRecordKind: """ - Decide whether an OpenAI batch JSONL line is an embedding request. + Decide which OpenAI endpoint shape an OpenAI batch JSONL line carries. - Precedence (strict - any explicit `url` short-circuits): - 1. `url == "/v1/embeddings"` -> embedding. Authoritative per the - OpenAI Batch API spec. - 2. Any other non-empty `url` (e.g. `/v1/chat/completions`) -> NOT - embedding. We trust the caller's explicit signal even if the - body would otherwise suggest embedding; misrouting a chat - record into the embedding transformer would corrupt the - modelInput, while a chat-shaped body sent to the chat path - either succeeds or fails cleanly inside that transformer. - 3. `url` missing/empty -> fall back to body shape. Requires - `input` present AND `messages` absent so a malformed record - carrying both keys routes to the chat path (safer default: - Anthropic transforms ignore unknown top-level keys, whereas - the embedding transformer would silently drop the messages). + Precedence (strict - any recognized `url` short-circuits): + 1. A `url` matching a supported endpoint wins. Authoritative per the + OpenAI Batch API spec, which requires it on every record. + 2. Any other non-empty `url` -> chat. We trust the caller's explicit + signal rather than re-deriving it from the body, and an + unexpectedly-shaped body fails cleanly inside the chat + transformer instead of being silently misrouted. + 3. `url` missing/empty -> fall back to body shape. `messages` wins + over the other keys so a malformed record carrying several of + them keeps its conversation instead of having it dropped, and a + bare `input` stays an embedding for backwards compatibility + (that ambiguity with `/v1/responses` is only resolvable from + `url`). """ - url = openai_jsonl_record.get("url") - if url == BedrockFilesConfig.OPENAI_EMBEDDINGS_URL: - return True - if url: - return False - body = openai_jsonl_record.get("body", {}) - if not isinstance(body, dict): - return False - return "input" in body and "messages" not in body + match openai_jsonl_record.get("url"): + case BedrockFilesConfig.OPENAI_EMBEDDINGS_URL: + return BedrockBatchRecordKind.EMBEDDING + case BedrockFilesConfig.OPENAI_TEXT_COMPLETIONS_URL: + return BedrockBatchRecordKind.TEXT_COMPLETION + case BedrockFilesConfig.OPENAI_RESPONSES_URL: + return BedrockBatchRecordKind.RESPONSES + case None | "": + pass + case _: + return BedrockBatchRecordKind.CHAT + + body = openai_jsonl_record.get("body") + if not isinstance(body, Mapping): + return BedrockBatchRecordKind.CHAT + if "messages" in body: + return BedrockBatchRecordKind.CHAT + if "prompt" in body: + return BedrockBatchRecordKind.TEXT_COMPLETION + if "input" in body: + return BedrockBatchRecordKind.EMBEDDING + return BedrockBatchRecordKind.CHAT # Identifier for the Bedrock Titan v2 InvokeModel body schema as stored # in `model_prices_and_context_window.json`. Centralized so future @@ -546,9 +588,83 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) return dict(titan_config._transform_request(input=input_text, inference_params=inference_params)) + @staticmethod + def _transform_text_completion_body_to_chat_body(openai_request_body: Mapping[str, Any]) -> Mapping[str, Any]: + """ + Rewrite an OpenAI `/v1/completions` batch body as a Chat Completions body. + + Bedrock batch `modelInput` is the model's InvokeModel/Converse body, and + no Bedrock batch model takes a bare `prompt`, so the wrapping that + `litellm.text_completion` does in real time has to happen here too. + """ + prompt = openai_request_body.get("prompt") + if prompt is None: + raise ValueError( + "Batch record for /v1/completions is missing required `prompt` field: " + f"model={openai_request_body.get('model', '')}" + ) + return _frozen_mapping( + chain( + ((key, value) for key, value in openai_request_body.items() if key != "prompt"), + (("messages", text_completion_prompt_to_messages(prompt)),), + ) + ) + + @staticmethod + def _transform_responses_body_to_chat_body(openai_request_body: Mapping[str, Any]) -> Mapping[str, Any]: + """ + Rewrite an OpenAI `/v1/responses` batch body as a Chat Completions body. + + Delegates to the same Responses-to-Chat bridge the real-time path uses + for providers without a native Responses API (which is every Bedrock + model), so `input`, `instructions`, `max_output_tokens` and the tool + params translate identically in batch and real time. The bridge always + emits a `tools` key; an empty one is dropped rather than shipped as an + empty array inside `modelInput`. + """ + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + responses_input = openai_request_body.get("input") + if responses_input is None: + raise ValueError( + "Batch record for /v1/responses is missing required `input` field: " + f"model={openai_request_body.get('model', '')}" + ) + chat_body = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=openai_request_body.get("model", ""), + input=_responses_input_adapter().validate_python(responses_input), + responses_api_request=_responses_request_adapter().validate_python( + _frozen_mapping( + (key, value) for key, value in openai_request_body.items() if key not in ("model", "input") + ) + ), + ) + return _frozen_mapping((key, value) for key, value in chat_body.items() if key != "tools" or value) + + @staticmethod + def _transform_batch_body_to_chat_body( + openai_request_body: Mapping[str, Any], + record_kind: BedrockBatchRecordKind, + ) -> Mapping[str, Any]: + """ + Normalize a non-embedding batch body to the Chat Completions shape the + per-provider Bedrock transformations expect. + """ + match record_kind: + case BedrockBatchRecordKind.TEXT_COMPLETION: + return BedrockFilesConfig._transform_text_completion_body_to_chat_body(openai_request_body) + case BedrockBatchRecordKind.RESPONSES: + return BedrockFilesConfig._transform_responses_body_to_chat_body(openai_request_body) + case BedrockBatchRecordKind.CHAT: + return openai_request_body + case BedrockBatchRecordKind.EMBEDDING: + raise ValueError("Embedding batch records do not have a chat-completion equivalent") + def _map_openai_to_bedrock_params( self, - openai_request_body: dict[str, Any], + openai_request_body: Mapping[str, Any], provider: str | None = None, ) -> dict[str, Any]: """ @@ -659,14 +775,18 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): provider = self.get_bedrock_invoke_provider(model) # Route to the embedding transformer when the OpenAI batch line - # targets /v1/embeddings; otherwise fall back to the existing - # chat-completion path. We branch here (rather than inside + # targets /v1/embeddings; every other endpoint shape is normalized + # to chat completions first. We branch here (rather than inside # `_map_openai_to_bedrock_params`) so the chat helper keeps its # narrow contract and the embedding helper can evolve independently. - if self._is_embedding_record(_openai_jsonl_content): + record_kind = self._classify_batch_record(_openai_jsonl_content) + if record_kind is BedrockBatchRecordKind.EMBEDDING: model_input = self._map_openai_embedding_to_bedrock_params(openai_request_body=openai_body) else: - model_input = self._map_openai_to_bedrock_params(openai_request_body=openai_body, provider=provider) + model_input = self._map_openai_to_bedrock_params( + openai_request_body=self._transform_batch_body_to_chat_body(openai_body, record_kind), + provider=provider, + ) # Create Bedrock batch record record_id = _openai_jsonl_content.get("custom_id", f"CALL{str(idx).zfill(7)}") diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index d9f8229dbed..f7a4682cb03 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1,4 +1,5 @@ import json +from enum import Enum from typing import Any, Dict, List, Literal, Optional, Union from typing_extensions import TYPE_CHECKING, Required, TypedDict, override @@ -1100,3 +1101,17 @@ class BedrockInvokeAnthropicMessagesRequest(TypedDict, total=False): # supported subset and strips the field entirely when nothing remains, so # other edit types (e.g. `clear_thinking_20251015`) never reach Bedrock. context_management: dict + + +class BedrockBatchRecordKind(Enum): + """ + Which OpenAI endpoint shape a line of a Bedrock managed-batch JSONL file + carries. Bedrock batch `modelInput` is always the model's InvokeModel / + Converse body, so every non-embedding shape is normalized to Chat + Completions before being handed to the per-provider transformation. + """ + + CHAT = "chat" + TEXT_COMPLETION = "text_completion" + RESPONSES = "responses" + EMBEDDING = "embedding" diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 1b1db634ed2..d10ccf77703 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -721,3 +721,46 @@ class TestUnpackLegacyDefs: out = unpack_legacy_defs(schema) assert "components" not in out assert out["properties"]["r0"]["properties"]["p0"] == {"type": "string"} + + +class TestTextCompletionPromptToMessages: + """`/v1/completions` prompt wrapping, shared by the real-time and batch paths.""" + + def test_string_prompt_becomes_single_user_message(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + text_completion_prompt_to_messages, + ) + + assert text_completion_prompt_to_messages("summarize this") == ( + {"role": "user", "content": "summarize this"}, + ) + + def test_list_of_strings_becomes_one_message_each(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + text_completion_prompt_to_messages, + ) + + assert text_completion_prompt_to_messages(["first", "second"]) == ( + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + ) + + @pytest.mark.parametrize( + "prompt", + [ + [1, 2, 3], + [[1, 2], [3, 4]], + ["ok", 7], + [], + "", + None, + {"role": "user"}, + ], + ) + def test_unsupported_prompt_shapes_raise(self, prompt): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + text_completion_prompt_to_messages, + ) + + with pytest.raises(ValueError, match="non-empty string or a non-empty list of strings"): + text_completion_prompt_to_messages(prompt) diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index c548fe53e15..87b03b02e1a 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1072,21 +1072,34 @@ class TestBedrockFilesEmbeddingTransformation: is None ) - def test_is_embedding_record_helper(self): - """Helper detects embeddings via `url` first, then by body shape.""" + def test_classify_batch_record_helper(self): + """Helper classifies by `url` first, then by body shape.""" from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind - assert BedrockFilesConfig._is_embedding_record( - {"url": "/v1/embeddings", "body": {"input": "x"}} + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/embeddings", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.EMBEDDING ) # body-only fallback - assert BedrockFilesConfig._is_embedding_record({"body": {"input": "x"}}) - # chat shape - assert not BedrockFilesConfig._is_embedding_record( - {"url": "/v1/chat/completions", "body": {"messages": []}} + assert ( + BedrockFilesConfig._classify_batch_record({"body": {"input": "x"}}) + is BedrockBatchRecordKind.EMBEDDING + ) + # chat shape + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/chat/completions", "body": {"messages": []}} + ) + is BedrockBatchRecordKind.CHAT + ) + # ambiguous body without any recognized key is treated as chat + assert ( + BedrockFilesConfig._classify_batch_record({"body": {}}) + is BedrockBatchRecordKind.CHAT ) - # ambiguous body without `input` is treated as not-embedding - assert not BedrockFilesConfig._is_embedding_record({"body": {}}) def test_explicit_chat_url_with_input_body_short_circuits_to_chat(self): """Explicit url=/v1/chat/completions wins even if body looks like embedding. @@ -1097,15 +1110,20 @@ class TestBedrockFilesEmbeddingTransformation: """ from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind + # Direct helper assertion - assert not BedrockFilesConfig._is_embedding_record( - { - "url": "/v1/chat/completions", - "body": { - "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - "input": "this would mis-route under the old precedence", - }, - } + assert ( + BedrockFilesConfig._classify_batch_record( + { + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "input": "this would mis-route under the old precedence", + }, + } + ) + is BedrockBatchRecordKind.CHAT ) # End-to-end: a record like this routes through the chat path. We @@ -1164,20 +1182,301 @@ class TestBedrockFilesEmbeddingTransformation: with pytest.raises(ValueError, match="must be a string"): BedrockFilesConfig._coerce_embedding_input_to_string({"unsupported": True}) - def test_other_non_embedding_urls_route_to_chat(self): - """Any non-/v1/embeddings url short-circuits to chat path.""" + def test_other_non_embedding_urls_do_not_route_to_embeddings(self): + """An `input` body only means "embedding" when the url says so.""" from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind # /v1/completions (legacy completions endpoint) - assert not BedrockFilesConfig._is_embedding_record( - {"url": "/v1/completions", "body": {"input": "x"}} + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/completions", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.TEXT_COMPLETION + ) + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/responses", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.RESPONSES ) # Arbitrary unknown url - caller's explicit signal still wins - assert not BedrockFilesConfig._is_embedding_record( - {"url": "/v1/responses", "body": {"input": "x"}} + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/moderations", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.CHAT ) +class TestBedrockBatchNonChatEndpointRecords: + """`/v1/completions` and `/v1/responses` JSONL records (issue #35639). + + Bedrock batch `modelInput` is always the model's InvokeModel/Converse body, + so a record shaped for another OpenAI endpoint has to be normalized to chat + completions first. Before this normalization every record below either + raised `BadRequestError` at `POST /v1/files` (Anthropic, Nova) or silently + shipped an empty `messages` list to AWS (passthrough providers). + """ + + ANTHROPIC_MODEL = "bedrock/us.anthropic.claude-sonnet-4-6" + NOVA_MODEL = "bedrock/us.amazon.nova-pro-v1:0" + PASSTHROUGH_MODEL = "bedrock/openai.gpt-oss-120b-1:0" + + def _transform(self, record: dict) -> dict: + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + result = BedrockFilesConfig()._transform_openai_jsonl_content_to_bedrock_jsonl_content([record]) + assert len(result) == 1 + assert result[0]["recordId"] == record["custom_id"] + return result[0]["modelInput"] + + def test_anthropic_text_completion_record_wraps_prompt(self): + model_input = self._transform( + { + "custom_id": "1", + "method": "POST", + "url": "/v1/completions", + "body": { + "model": self.ANTHROPIC_MODEL, + "prompt": "Summarize the following call transcript", + "max_tokens": 64, + }, + } + ) + + assert model_input == { + "messages": [ + { + "role": "user", + "content": [{"type": "text", "text": "Summarize the following call transcript"}], + } + ], + "max_tokens": 64, + "anthropic_version": "bedrock-2023-05-31", + } + + def test_anthropic_text_completion_record_keeps_every_prompt_in_a_list(self): + model_input = self._transform( + { + "custom_id": "2", + "method": "POST", + "url": "/v1/completions", + "body": { + "model": self.ANTHROPIC_MODEL, + "prompt": ["first prompt", "second prompt"], + "max_tokens": 8, + }, + } + ) + + # Consecutive user messages are merged by the Anthropic transform, the + # same way they are on the real-time path. + assert model_input["messages"] == [ + { + "role": "user", + "content": [ + {"type": "text", "text": "first prompt"}, + {"type": "text", "text": "second prompt"}, + ], + } + ] + assert "prompt" not in model_input + + def test_anthropic_responses_record_wraps_string_input(self): + model_input = self._transform( + { + "custom_id": "3", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": self.ANTHROPIC_MODEL, + "input": "hi", + "max_output_tokens": 16, + }, + } + ) + + assert model_input == { + "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}], + "max_tokens": 16, + "anthropic_version": "bedrock-2023-05-31", + } + assert "tools" not in model_input, "an empty tools array must not be shipped to Bedrock" + + def test_anthropic_responses_record_maps_instructions_and_input_items(self): + """The Responses-specific params go through the same bridge as real time.""" + model_input = self._transform( + { + "custom_id": "4", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": self.ANTHROPIC_MODEL, + "instructions": "be terse", + "input": [ + {"role": "user", "content": "what is 2+2?"}, + {"role": "assistant", "content": "4"}, + {"role": "user", "content": "and 3+3?"}, + ], + "max_output_tokens": 32, + "temperature": 0.2, + }, + } + ) + + assert model_input["system"] == [{"type": "text", "text": "be terse"}] + assert model_input["max_tokens"] == 32 + assert model_input["temperature"] == 0.2 + assert [message["role"] for message in model_input["messages"]] == [ + "user", + "assistant", + "user", + ] + assert model_input["messages"][-1]["content"] == [{"type": "text", "text": "and 3+3?"}] + assert "input" not in model_input + assert "max_output_tokens" not in model_input + + @pytest.mark.parametrize( + "body", + [ + {"prompt": "hi"}, + {"input": "hi"}, + ], + ids=["prompt", "input"], + ) + def test_nova_converse_record_wraps_prompt_and_input(self, body): + url = "/v1/completions" if "prompt" in body else "/v1/responses" + model_input = self._transform( + { + "custom_id": "5", + "method": "POST", + "url": url, + "body": {"model": self.NOVA_MODEL, **body}, + } + ) + + assert model_input["messages"] == [{"role": "user", "content": [{"text": "hi"}]}] + + @pytest.mark.parametrize( + "body", + [ + {"prompt": "hi"}, + {"input": "hi"}, + ], + ids=["prompt", "input"], + ) + def test_passthrough_provider_record_no_longer_emits_empty_messages(self, body): + """The passthrough branch used to emit `{"messages": [], "prompt": ...}`. + + That shape is accepted by `POST /v1/files`, so the whole batch job was + submitted to AWS and only failed there. + """ + url = "/v1/completions" if "prompt" in body else "/v1/responses" + model_input = self._transform( + { + "custom_id": "6", + "method": "POST", + "url": url, + "body": {"model": self.PASSTHROUGH_MODEL, **body}, + } + ) + + # Asserted on the serialized form, since the passthrough branch hands + # `messages` straight to S3 without a per-provider transform. + assert json.loads(json.dumps(model_input)) == {"messages": [{"role": "user", "content": "hi"}]} + + def test_mixed_endpoints_in_one_file_keep_their_own_shapes(self): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + result = BedrockFilesConfig()._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "chat", + "url": "/v1/chat/completions", + "body": { + "model": self.ANTHROPIC_MODEL, + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 4, + }, + }, + { + "custom_id": "text", + "url": "/v1/completions", + "body": {"model": self.ANTHROPIC_MODEL, "prompt": "hi", "max_tokens": 4}, + }, + { + "custom_id": "responses", + "url": "/v1/responses", + "body": {"model": self.ANTHROPIC_MODEL, "input": "hi", "max_output_tokens": 4}, + }, + { + "custom_id": "embedding", + "url": "/v1/embeddings", + "body": {"model": "bedrock/amazon.titan-embed-text-v2:0", "input": "hi"}, + }, + ] + ) + + assert [record["recordId"] for record in result] == [ + "chat", + "text", + "responses", + "embedding", + ] + for record in result[:3]: + assert record["modelInput"]["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]} + ] + assert result[3]["modelInput"] == {"inputText": "hi"} + + @pytest.mark.parametrize( + ("url", "expected_message"), + [ + ("/v1/completions", "missing required `prompt` field"), + ("/v1/responses", "missing required `input` field"), + ], + ) + def test_missing_required_field_raises_actionable_error(self, url, expected_message): + with pytest.raises(ValueError, match=expected_message): + self._transform( + { + "custom_id": "7", + "method": "POST", + "url": url, + "body": {"model": self.ANTHROPIC_MODEL, "max_tokens": 4}, + } + ) + + def test_prompt_body_without_url_is_still_wrapped(self): + """A record can omit `url`; the body shape then decides.""" + model_input = self._transform( + { + "custom_id": "8", + "body": {"model": self.ANTHROPIC_MODEL, "prompt": "hi", "max_tokens": 4}, + } + ) + + assert model_input["messages"] == [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + + def test_messages_win_over_prompt_when_url_is_absent(self): + model_input = self._transform( + { + "custom_id": "9", + "body": { + "model": self.ANTHROPIC_MODEL, + "messages": [{"role": "user", "content": "from messages"}], + "prompt": "from prompt", + "max_tokens": 4, + }, + } + ) + + assert model_input["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "from messages"}]} + ] + + class TestBedrockFileContentTransformation: """SigV4-signed S3 GetObject retrieval of Bedrock batch output files.""" From 6def61e672d95c4b145613009cd3064d0a133475 Mon Sep 17 00:00:00 2001 From: milan Date: Mon, 3 Aug 2026 17:22:46 +0000 Subject: [PATCH 028/576] fix(bedrock): keep /v1/responses batch metadata through the chat bridge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/files/transformation.py | 1 + .../files/test_bedrock_files_transformation.py | 17 +++++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 0aa832780e5..baa2630556a 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -640,6 +640,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): (key, value) for key, value in openai_request_body.items() if key not in ("model", "input") ) ), + metadata=openai_request_body.get("metadata"), ) return _frozen_mapping((key, value) for key, value in chat_body.items() if key != "tools" or value) diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 87b03b02e1a..09aa6b1cf2d 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1337,6 +1337,23 @@ class TestBedrockBatchNonChatEndpointRecords: assert "input" not in model_input assert "max_output_tokens" not in model_input + def test_responses_record_keeps_metadata(self): + """`metadata` reaches the bridge, which reads it as its own kwarg.""" + model_input = self._transform( + { + "custom_id": "4b", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": self.PASSTHROUGH_MODEL, + "input": "hi", + "metadata": {"tenant": "acct-1"}, + }, + } + ) + + assert model_input["metadata"] == {"tenant": "acct-1"} + @pytest.mark.parametrize( "body", [ From 0c0e1e8374d7e956e65d275ebf5f2f832ec374b9 Mon Sep 17 00:00:00 2001 From: Miles Adkins Date: Wed, 5 Aug 2026 13:21:15 -0500 Subject: [PATCH 029/576] feat(fireworks_ai): translate NIM/vLLM extra params to Fireworks-native args Requests migrated from NIM/vLLM servers carry extras that flow through the extra_body passthrough verbatim, but the Fireworks chat completions API either names them differently or does not accept them at all. Add FireworksAIConfig.map_extra_body_params, invoked from the fireworks chat dispatch, which renames truncate_prompt_tokens to prompt_truncate_len, maps chat_template_kwargs.enable_thinking to reasoning_effort, converts guided_json/guided_grammar/guided_choice to response_format, and drops the remaining extras (min_tokens, stop_token_ids, skip_special_tokens, guided_regex, etc.) with a debug log. Alias and competing-constraint combinations raise BadRequestError. Unrecognized extras keep passing through untouched, as do fireworks-native params like top_k. --- .../llms/fireworks_ai/chat/transformation.py | 162 +++++++++++- litellm/main.py | 6 +- .../test_fireworks_ai_chat_transformation.py | 249 ++++++++++++++++++ 3 files changed, 415 insertions(+), 2 deletions(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index a796aa47b70..4740e84d513 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -1,5 +1,5 @@ import json -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Mapping from typing import Any, Final, Literal, cast import httpx @@ -61,6 +61,36 @@ def _extract_fireworks_hidden_params(payload: dict) -> dict: return {**top_level, **per_choice} +def _json_schema_response_format(schema: object) -> Mapping[str, object]: + return {"type": "json_schema", "json_schema": {"schema": schema}} # mutable-ok: JSON request body + + +_NIM_VLLM_STRIP_PARAMS: Final = frozenset( + { + "min_tokens", + "stop_token_ids", + "include_stop_str_in_output", + "skip_special_tokens", + "spaces_between_special_tokens", + "best_of", + "use_beam_search", + "guided_decoding_backend", + "guided_regex", + "add_generation_prompt", + "continue_final_message", + "add_special_tokens", + "detokenize", + "allowed_token_ids", + "bad_words", + } +) + +_EXTRA_BODY_CONSUMED_PARAMS: Final = ( + frozenset({"truncate_prompt_tokens", "chat_template_kwargs", "guided_json", "guided_grammar", "guided_choice"}) + | _NIM_VLLM_STRIP_PARAMS +) + + class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): """ Reference: https://docs.fireworks.ai/api-reference/post-chatcompletions @@ -273,6 +303,136 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): return optional_params + def map_extra_body_params(self, optional_params: Mapping[str, object], model: str) -> dict: # noqa: LIT001 # http handler pops extra_body off the returned dict + extra_body: Final = optional_params.get("extra_body") + if not isinstance(extra_body, dict): + return dict(optional_params) # mutable-ok: JSON request body + + self._validate_extra_body_conflicts(extra_body=extra_body, optional_params=optional_params, model=model) + stripped: Final = tuple(sorted(k for k in extra_body if k in _NIM_VLLM_STRIP_PARAMS)) + if stripped: + verbose_logger.debug( + "fireworks_ai does not support NIM/vLLM params %s for model=%s; dropping them from the request.", + stripped, + model, + ) + promoted: Final = ( + *self._translate_truncate_prompt_tokens(extra_body), + *self._translate_chat_template_kwargs(extra_body, model), + *self._translate_guided_params(extra_body), + ) + remaining: Final = tuple((k, v) for k, v in extra_body.items() if k not in _EXTRA_BODY_CONSUMED_PARAMS) + base: Final = {k: v for k, v in optional_params.items() if k != "extra_body"} # mutable-ok: JSON request body + return { # mutable-ok: JSON request body + **base, + **dict(promoted), # mutable-ok: JSON request body + **({"extra_body": dict(remaining)} if remaining else {}), # mutable-ok: JSON request body + } + + def _validate_extra_body_conflicts( + self, extra_body: Mapping[str, object], optional_params: Mapping[str, object], model: str + ) -> None: + if "truncate_prompt_tokens" in extra_body and ( + "prompt_truncate_len" in extra_body or "prompt_truncate_len" in optional_params + ): + raise litellm.BadRequestError( + message=( + "Fireworks AI chat completions received both `truncate_prompt_tokens` and " + "`prompt_truncate_len`; they are aliases, send only one." + ), + model=model, + llm_provider="fireworks_ai", + ) + chat_template_kwargs: Final = extra_body.get("chat_template_kwargs") + if ( + isinstance(chat_template_kwargs, dict) + and "enable_thinking" in chat_template_kwargs + and ("reasoning_effort" in optional_params or "thinking" in optional_params) + ): + raise litellm.BadRequestError( + message=( + "Fireworks AI chat completions does not support specifying both " + "`chat_template_kwargs.enable_thinking` and `reasoning_effort`/`thinking` in the same request." + ), + model=model, + llm_provider="fireworks_ai", + ) + guided_params: Final = tuple( + k for k in ("guided_json", "guided_grammar", "guided_choice") if extra_body.get(k) is not None + ) + if len(guided_params) > 1: + raise litellm.BadRequestError( + message=( + f"Fireworks AI chat completions received multiple guided decoding params " + f"{guided_params}; send only one." + ), + model=model, + llm_provider="fireworks_ai", + ) + if guided_params and "response_format" in optional_params: + raise litellm.BadRequestError( + message=( + f"Fireworks AI chat completions received both `{guided_params[0]}` and " + "`response_format`; they are competing output constraints, send only one." + ), + model=model, + llm_provider="fireworks_ai", + ) + + @staticmethod + def _translate_truncate_prompt_tokens(extra_body: Mapping[str, object]) -> tuple[tuple[str, object], ...]: + if extra_body.get("truncate_prompt_tokens") is None: + return () + return (("prompt_truncate_len", extra_body["truncate_prompt_tokens"]),) + + def _translate_chat_template_kwargs( + self, extra_body: Mapping[str, object], model: str + ) -> tuple[tuple[str, object], ...]: + chat_template_kwargs: Final = extra_body.get("chat_template_kwargs") + if chat_template_kwargs is None: + return () + if not isinstance(chat_template_kwargs, dict): + raise litellm.BadRequestError( + message="Fireworks AI chat completions expects `chat_template_kwargs` to be an object.", + model=model, + llm_provider="fireworks_ai", + ) + other_keys: Final = tuple(sorted(k for k in chat_template_kwargs if k != "enable_thinking")) + if other_keys: + verbose_logger.debug( + "fireworks_ai does not support chat_template_kwargs keys %s for model=%s; dropping them.", + other_keys, + model, + ) + if "enable_thinking" not in chat_template_kwargs: + return () + if not supports_reasoning(model=model, custom_llm_provider="fireworks_ai"): + verbose_logger.debug( + "fireworks_ai model %r does not support reasoning; dropping chat_template_kwargs.enable_thinking.", + model, + ) + return () + effort: Final = "medium" if chat_template_kwargs["enable_thinking"] else "none" + return (("reasoning_effort", effort),) + + @staticmethod + def _translate_guided_params(extra_body: Mapping[str, object]) -> tuple[tuple[str, object], ...]: + if extra_body.get("guided_json") is not None: + return (("response_format", _json_schema_response_format(extra_body["guided_json"])),) + if extra_body.get("guided_grammar") is not None: + grammar_response_format: Final = { # mutable-ok: JSON request body + "type": "grammar", + "grammar": extra_body["guided_grammar"], + } + return (("response_format", grammar_response_format),) + if extra_body.get("guided_choice") is not None: + choice_schema: Final = { # mutable-ok: JSON request body + "type": "string", + "enum": extra_body["guided_choice"], + } + return (("response_format", _json_schema_response_format(choice_schema)),) + return () + def _transform_tools(self, tools: list[OpenAIChatCompletionToolParam]) -> list[OpenAIChatCompletionToolParam]: for tool in tools: if tool.get("type") != "function": diff --git a/litellm/main.py b/litellm/main.py index f906c78f9ae..660e024b113 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1711,11 +1711,15 @@ def _complete_fireworks_ai( messages: Final = ctx.messages model: Final = ctx.model model_response: Final = ctx.model_response - optional_params: Final = ctx.optional_params provider_config: Final = ctx.provider_config shared_session: Final = ctx.shared_session stream: Final = ctx.stream timeout: Final = ctx.timeout + optional_params: Final = ( + provider_config.map_extra_body_params(optional_params=ctx.optional_params, model=model) + if isinstance(provider_config, litellm.FireworksAIConfig) + else ctx.optional_params + ) try: response: Final = base_llm_http_handler.completion( diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 94945ed4bfb..3fbbc70916a 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1282,3 +1282,252 @@ def test_streaming_surfaces_fireworks_response_fields(): assert surfaced["fireworks_raw_outputs"] == [raw_output] assert surfaced["fireworks_perf_metrics"] == {"prompt-tokens": 5} assert surfaced["fireworks_prompt_token_ids"] == [1, 2, 3] + + +def test_map_extra_body_params_translates_truncate_prompt_tokens(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"truncate_prompt_tokens": 4096}}, _REASONING_MODEL + ) + assert result == {"prompt_truncate_len": 4096} + + +def test_map_extra_body_params_truncate_prompt_tokens_conflicts_with_alias(): + config = FireworksAIConfig() + with pytest.raises(litellm.BadRequestError, match="aliases"): + config.map_extra_body_params( + {"prompt_truncate_len": 2048, "extra_body": {"truncate_prompt_tokens": 4096}}, + _REASONING_MODEL, + ) + with pytest.raises(litellm.BadRequestError, match="aliases"): + config.map_extra_body_params( + {"extra_body": {"truncate_prompt_tokens": 4096, "prompt_truncate_len": 2048}}, + _REASONING_MODEL, + ) + + +def test_map_extra_body_params_chat_template_kwargs_enable_thinking(): + config = FireworksAIConfig() + disabled = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"enable_thinking": False}}}, + _REASONING_MODEL, + ) + assert disabled == {"reasoning_effort": "none"} + + enabled = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"enable_thinking": True}}}, + _REASONING_MODEL, + ) + assert enabled == {"reasoning_effort": "medium"} + + +def test_map_extra_body_params_chat_template_kwargs_conflicts_with_reasoning_effort(): + config = FireworksAIConfig() + with pytest.raises(litellm.BadRequestError, match="enable_thinking"): + config.map_extra_body_params( + { + "reasoning_effort": "high", + "extra_body": {"chat_template_kwargs": {"enable_thinking": False}}, + }, + _REASONING_MODEL, + ) + + +def test_map_extra_body_params_chat_template_kwargs_conflicts_with_thinking(): + config = FireworksAIConfig() + with pytest.raises(litellm.BadRequestError, match="enable_thinking"): + config.map_extra_body_params( + { + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "extra_body": {"chat_template_kwargs": {"enable_thinking": True}}, + }, + _REASONING_MODEL, + ) + + +def test_map_extra_body_params_chat_template_kwargs_dropped_for_non_reasoning_model(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"enable_thinking": False, "custom_flag": 1}}}, + _NON_REASONING_MODEL, + ) + assert result == {} + + +def test_map_extra_body_params_guided_json(): + config = FireworksAIConfig() + schema = {"type": "object", "properties": {"x": {"type": "string"}}} + result = config.map_extra_body_params( + {"extra_body": {"guided_json": schema}}, _REASONING_MODEL + ) + assert result == { + "response_format": {"type": "json_schema", "json_schema": {"schema": schema}} + } + + +def test_map_extra_body_params_guided_grammar_and_choice(): + config = FireworksAIConfig() + grammar = config.map_extra_body_params( + {"extra_body": {"guided_grammar": "root ::= 'hello'"}}, _REASONING_MODEL + ) + assert grammar == { + "response_format": {"type": "grammar", "grammar": "root ::= 'hello'"} + } + + choice = config.map_extra_body_params( + {"extra_body": {"guided_choice": ["yes", "no"]}}, _REASONING_MODEL + ) + assert choice == { + "response_format": { + "type": "json_schema", + "json_schema": {"schema": {"type": "string", "enum": ["yes", "no"]}}, + } + } + + +def test_map_extra_body_params_guided_conflicts_with_response_format(): + config = FireworksAIConfig() + with pytest.raises(litellm.BadRequestError, match="response_format"): + config.map_extra_body_params( + { + "response_format": {"type": "json_object"}, + "extra_body": {"guided_json": {"type": "object"}}, + }, + _REASONING_MODEL, + ) + + +def test_map_extra_body_params_multiple_guided_params_rejected(): + config = FireworksAIConfig() + with pytest.raises(litellm.BadRequestError, match="multiple guided decoding params"): + config.map_extra_body_params( + {"extra_body": {"guided_json": {"type": "object"}, "guided_grammar": "root ::= 'x'"}}, + _REASONING_MODEL, + ) + + +@pytest.mark.parametrize( + "param,value", + [ + ("min_tokens", 10), + ("stop_token_ids", [1, 2]), + ("include_stop_str_in_output", True), + ("skip_special_tokens", False), + ("spaces_between_special_tokens", True), + ("best_of", 2), + ("use_beam_search", True), + ("guided_decoding_backend", "outlines"), + ("guided_regex", "[0-9]+"), + ("add_generation_prompt", True), + ("continue_final_message", True), + ("add_special_tokens", False), + ("detokenize", True), + ("allowed_token_ids", [1]), + ("bad_words", ["foo"]), + ], +) +def test_map_extra_body_params_strips_unsupported_nim_vllm_params(param, value, caplog): + import logging + + config = FireworksAIConfig() + with caplog.at_level(logging.DEBUG): + result = config.map_extra_body_params( + {"extra_body": {param: value}}, _REASONING_MODEL + ) + assert result == {} + assert param in caplog.text + + +def test_map_extra_body_params_preserves_unknown_passthrough(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"top_k": 40, "some_future_param": "x", "truncate_prompt_tokens": 100}}, + _REASONING_MODEL, + ) + assert result == { + "prompt_truncate_len": 100, + "extra_body": {"top_k": 40, "some_future_param": "x"}, + } + + +def test_map_extra_body_params_no_extra_body(): + config = FireworksAIConfig() + assert config.map_extra_body_params({}, _REASONING_MODEL) == {} + unchanged = {"temperature": 0.5, "extra_body": None} + assert config.map_extra_body_params(unchanged, _REASONING_MODEL) == unchanged + + +def test_nim_vllm_extras_translated_end_to_end_in_request_body(): + """ + Passing NIM/vLLM extras to litellm.completion must reach the Fireworks + request body translated, not verbatim: truncate_prompt_tokens becomes + prompt_truncate_len, chat_template_kwargs.enable_thinking becomes + reasoning_effort, min_tokens is dropped, and fireworks-native top_k still + passes through. Asserts on the actual JSON posted to the API, so a revert + of the _complete_fireworks_ai wiring fails this test. + """ + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + model = "accounts/fireworks/models/glm-5p1" + body = { + "id": "chat-1", + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hi"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + raw_response = MagicMock() + raw_response.status_code = 200 + raw_response.headers = {} + raw_response.text = json.dumps(body) + raw_response.json = lambda: body + + client = HTTPHandler() + with patch.object(client, "post", return_value=raw_response) as mock_post: + litellm.completion( + model=f"fireworks_ai/{model}", + messages=[{"role": "user", "content": "hi"}], + api_key="fw-test-key", + client=client, + truncate_prompt_tokens=4096, + chat_template_kwargs={"enable_thinking": False}, + min_tokens=10, + top_k=40, + ) + + request_body = json.loads(mock_post.call_args.kwargs["data"]) + assert request_body["prompt_truncate_len"] == 4096 + assert "truncate_prompt_tokens" not in request_body + assert request_body["reasoning_effort"] == "none" + assert "chat_template_kwargs" not in request_body + assert "min_tokens" not in request_body + assert request_body["top_k"] == 40 + + +def test_in_schema_unsupported_params_still_raise(): + """ + The extras translation channel does not weaken the supported-params gate + for in-schema OpenAI params: store is still rejected with drop_params=False + and dropped with drop_params=True. + """ + with pytest.raises(litellm.UnsupportedParamsError): + litellm.get_optional_params( + model="accounts/fireworks/models/llama-v3-70b-instruct", + custom_llm_provider="fireworks_ai", + drop_params=False, + store=True, + ) + optional_params = litellm.get_optional_params( + model="accounts/fireworks/models/llama-v3-70b-instruct", + custom_llm_provider="fireworks_ai", + drop_params=True, + store=True, + ) + assert "store" not in optional_params From 599283584f2d16448c25a1ee4fbfdda2062eecf3 Mon Sep 17 00:00:00 2001 From: Miles Adkins Date: Wed, 5 Aug 2026 15:09:16 -0500 Subject: [PATCH 030/576] feat(fireworks_ai): drop reasoning_effort=auto to the model default Fireworks rejects reasoning_effort="auto" (accepted set: low, medium, high, xhigh, max, none, adaptive), so OpenAI-compatible clients sending it 400. Omitting the param means model default on Fireworks, which is exactly what auto means on OpenAI's side, so skip it in map_openai_params instead of forwarding. --- litellm/llms/fireworks_ai/chat/transformation.py | 6 ++++-- .../test_fireworks_ai_chat_transformation.py | 16 ++++++++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 4740e84d513..a05b160413e 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -295,7 +295,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): optional_params["reasoning_effort"] = "medium" elif value is False: optional_params["reasoning_effort"] = "none" - else: + elif value != "auto": optional_params["reasoning_effort"] = value elif param in supported_openai_params: if value is not None: @@ -303,7 +303,9 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): return optional_params - def map_extra_body_params(self, optional_params: Mapping[str, object], model: str) -> dict: # noqa: LIT001 # http handler pops extra_body off the returned dict + def map_extra_body_params( + self, optional_params: Mapping[str, object], model: str + ) -> dict: # mutable-ok: http handler pops extra_body off the returned dict extra_body: Final = optional_params.get("extra_body") if not isinstance(extra_body, dict): return dict(optional_params) # mutable-ok: JSON request body diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 3fbbc70916a..bbb7fb197d0 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1153,6 +1153,22 @@ def test_reasoning_effort_integer_passthrough(): assert isinstance(result["reasoning_effort"], int) +def test_reasoning_effort_auto_dropped_to_model_default(): + """ + Fireworks rejects reasoning_effort="auto" (accepted set: low/medium/high/ + xhigh/max/none/adaptive). Omitting the param is the model default, which is + exactly what "auto" means on OpenAI's side, so it must not reach the request. + """ + config = FireworksAIConfig() + result = config.map_openai_params( + {"reasoning_effort": "auto"}, + {}, + _REASONING_MODEL, + drop_params=False, + ) + assert "reasoning_effort" not in result + + def test_transform_response_captures_perf_metrics(): body = { **_BASE_CHAT_COMPLETION_RESPONSE, From 6d80d0509976feb702aad744cb7aff5fa81d3f54 Mon Sep 17 00:00:00 2001 From: Miles Adkins Date: Wed, 5 Aug 2026 15:56:33 -0500 Subject: [PATCH 031/576] fix(fireworks_ai): align extras translation with the API gateway matrix min_tokens is accepted natively by the Fireworks API (verified live), so stop stripping it and let it pass through extra_body. Add the NIM-specific include_reasoning and nvext keys to the strip set. enable_thinking=true now omits reasoning_effort (model default) instead of forcing medium, matching the gateway translation and preserving default-off models' behavior; enable_thinking=false still maps to none. --- litellm/llms/fireworks_ai/chat/transformation.py | 8 +++++--- .../test_fireworks_ai_chat_transformation.py | 16 ++++++++++------ 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index a05b160413e..b5c82129f45 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -67,7 +67,6 @@ def _json_schema_response_format(schema: object) -> Mapping[str, object]: _NIM_VLLM_STRIP_PARAMS: Final = frozenset( { - "min_tokens", "stop_token_ids", "include_stop_str_in_output", "skip_special_tokens", @@ -82,6 +81,8 @@ _NIM_VLLM_STRIP_PARAMS: Final = frozenset( "detokenize", "allowed_token_ids", "bad_words", + "include_reasoning", + "nvext", } ) @@ -414,8 +415,9 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): model, ) return () - effort: Final = "medium" if chat_template_kwargs["enable_thinking"] else "none" - return (("reasoning_effort", effort),) + if chat_template_kwargs["enable_thinking"]: + return () + return (("reasoning_effort", "none"),) @staticmethod def _translate_guided_params(extra_body: Mapping[str, object]) -> tuple[tuple[str, object], ...]: diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index bbb7fb197d0..d25bbd69b91 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1334,7 +1334,7 @@ def test_map_extra_body_params_chat_template_kwargs_enable_thinking(): {"extra_body": {"chat_template_kwargs": {"enable_thinking": True}}}, _REASONING_MODEL, ) - assert enabled == {"reasoning_effort": "medium"} + assert enabled == {} def test_map_extra_body_params_chat_template_kwargs_conflicts_with_reasoning_effort(): @@ -1425,7 +1425,6 @@ def test_map_extra_body_params_multiple_guided_params_rejected(): @pytest.mark.parametrize( "param,value", [ - ("min_tokens", 10), ("stop_token_ids", [1, 2]), ("include_stop_str_in_output", True), ("skip_special_tokens", False), @@ -1440,6 +1439,8 @@ def test_map_extra_body_params_multiple_guided_params_rejected(): ("detokenize", True), ("allowed_token_ids", [1]), ("bad_words", ["foo"]), + ("include_reasoning", False), + ("nvext", {"verbosity": 1}), ], ) def test_map_extra_body_params_strips_unsupported_nim_vllm_params(param, value, caplog): @@ -1478,9 +1479,10 @@ def test_nim_vllm_extras_translated_end_to_end_in_request_body(): Passing NIM/vLLM extras to litellm.completion must reach the Fireworks request body translated, not verbatim: truncate_prompt_tokens becomes prompt_truncate_len, chat_template_kwargs.enable_thinking becomes - reasoning_effort, min_tokens is dropped, and fireworks-native top_k still - passes through. Asserts on the actual JSON posted to the API, so a revert - of the _complete_fireworks_ai wiring fails this test. + reasoning_effort, include_reasoning is dropped, and min_tokens and + fireworks-native top_k still pass through. Asserts on the actual JSON + posted to the API, so a revert of the _complete_fireworks_ai wiring + fails this test. """ from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -1515,6 +1517,7 @@ def test_nim_vllm_extras_translated_end_to_end_in_request_body(): truncate_prompt_tokens=4096, chat_template_kwargs={"enable_thinking": False}, min_tokens=10, + include_reasoning=False, top_k=40, ) @@ -1523,7 +1526,8 @@ def test_nim_vllm_extras_translated_end_to_end_in_request_body(): assert "truncate_prompt_tokens" not in request_body assert request_body["reasoning_effort"] == "none" assert "chat_template_kwargs" not in request_body - assert "min_tokens" not in request_body + assert "include_reasoning" not in request_body + assert request_body["min_tokens"] == 10 assert request_body["top_k"] == 40 From 431f61b4f7c20b8f722f30c42c279edd19fe6a2d Mon Sep 17 00:00:00 2001 From: Miles Adkins Date: Wed, 5 Aug 2026 16:02:47 -0500 Subject: [PATCH 032/576] fix(fireworks_ai): prefer native values silently on extras conflicts Align with the API gateway translation: instead of raising BadRequestError on alias or competing-constraint conflicts, the explicit Fireworks-native param wins and the NIM/vLLM extra is dropped with a debug log. Covers truncate_prompt_tokens vs prompt_truncate_len, chat_template_kwargs enable_thinking vs reasoning_effort/thinking, guided_* vs response_format (including response_format nested in an explicit extra_body, which the previous conflict check missed), and multiple guided_* params (priority order json, grammar, choice). Malformed non-object chat_template_kwargs is also dropped with a log instead of raising. --- .../llms/fireworks_ai/chat/transformation.py | 108 +++++++---------- .../test_fireworks_ai_chat_transformation.py | 111 +++++++++++------- 2 files changed, 107 insertions(+), 112 deletions(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index b5c82129f45..3bacb3cd28e 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -311,7 +311,6 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): if not isinstance(extra_body, dict): return dict(optional_params) # mutable-ok: JSON request body - self._validate_extra_body_conflicts(extra_body=extra_body, optional_params=optional_params, model=model) stripped: Final = tuple(sorted(k for k in extra_body if k in _NIM_VLLM_STRIP_PARAMS)) if stripped: verbose_logger.debug( @@ -320,9 +319,9 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): model, ) promoted: Final = ( - *self._translate_truncate_prompt_tokens(extra_body), - *self._translate_chat_template_kwargs(extra_body, model), - *self._translate_guided_params(extra_body), + *self._translate_truncate_prompt_tokens(extra_body, optional_params), + *self._translate_chat_template_kwargs(extra_body, optional_params, model), + *self._translate_guided_params(extra_body, optional_params), ) remaining: Final = tuple((k, v) for k, v in extra_body.items() if k not in _EXTRA_BODY_CONSUMED_PARAMS) base: Final = {k: v for k, v in optional_params.items() if k != "extra_body"} # mutable-ok: JSON request body @@ -332,74 +331,32 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): **({"extra_body": dict(remaining)} if remaining else {}), # mutable-ok: JSON request body } - def _validate_extra_body_conflicts( - self, extra_body: Mapping[str, object], optional_params: Mapping[str, object], model: str - ) -> None: - if "truncate_prompt_tokens" in extra_body and ( - "prompt_truncate_len" in extra_body or "prompt_truncate_len" in optional_params - ): - raise litellm.BadRequestError( - message=( - "Fireworks AI chat completions received both `truncate_prompt_tokens` and " - "`prompt_truncate_len`; they are aliases, send only one." - ), - model=model, - llm_provider="fireworks_ai", - ) - chat_template_kwargs: Final = extra_body.get("chat_template_kwargs") - if ( - isinstance(chat_template_kwargs, dict) - and "enable_thinking" in chat_template_kwargs - and ("reasoning_effort" in optional_params or "thinking" in optional_params) - ): - raise litellm.BadRequestError( - message=( - "Fireworks AI chat completions does not support specifying both " - "`chat_template_kwargs.enable_thinking` and `reasoning_effort`/`thinking` in the same request." - ), - model=model, - llm_provider="fireworks_ai", - ) - guided_params: Final = tuple( - k for k in ("guided_json", "guided_grammar", "guided_choice") if extra_body.get(k) is not None - ) - if len(guided_params) > 1: - raise litellm.BadRequestError( - message=( - f"Fireworks AI chat completions received multiple guided decoding params " - f"{guided_params}; send only one." - ), - model=model, - llm_provider="fireworks_ai", - ) - if guided_params and "response_format" in optional_params: - raise litellm.BadRequestError( - message=( - f"Fireworks AI chat completions received both `{guided_params[0]}` and " - "`response_format`; they are competing output constraints, send only one." - ), - model=model, - llm_provider="fireworks_ai", - ) - @staticmethod - def _translate_truncate_prompt_tokens(extra_body: Mapping[str, object]) -> tuple[tuple[str, object], ...]: + def _translate_truncate_prompt_tokens( + extra_body: Mapping[str, object], optional_params: Mapping[str, object] + ) -> tuple[tuple[str, object], ...]: if extra_body.get("truncate_prompt_tokens") is None: return () + if "prompt_truncate_len" in extra_body or "prompt_truncate_len" in optional_params: + verbose_logger.debug( + "fireworks_ai ignoring truncate_prompt_tokens; explicit prompt_truncate_len takes precedence." + ) + return () return (("prompt_truncate_len", extra_body["truncate_prompt_tokens"]),) def _translate_chat_template_kwargs( - self, extra_body: Mapping[str, object], model: str + self, extra_body: Mapping[str, object], optional_params: Mapping[str, object], model: str ) -> tuple[tuple[str, object], ...]: chat_template_kwargs: Final = extra_body.get("chat_template_kwargs") if chat_template_kwargs is None: return () if not isinstance(chat_template_kwargs, dict): - raise litellm.BadRequestError( - message="Fireworks AI chat completions expects `chat_template_kwargs` to be an object.", - model=model, - llm_provider="fireworks_ai", + verbose_logger.debug( + "fireworks_ai dropping chat_template_kwargs for model=%s; expected an object, got %s.", + model, + type(chat_template_kwargs).__name__, ) + return () other_keys: Final = tuple(sorted(k for k in chat_template_kwargs if k != "enable_thinking")) if other_keys: verbose_logger.debug( @@ -409,6 +366,11 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): ) if "enable_thinking" not in chat_template_kwargs: return () + if "reasoning_effort" in optional_params or "thinking" in optional_params: + verbose_logger.debug( + "fireworks_ai ignoring chat_template_kwargs.enable_thinking; explicit reasoning_effort/thinking takes precedence." + ) + return () if not supports_reasoning(model=model, custom_llm_provider="fireworks_ai"): verbose_logger.debug( "fireworks_ai model %r does not support reasoning; dropping chat_template_kwargs.enable_thinking.", @@ -420,7 +382,19 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): return (("reasoning_effort", "none"),) @staticmethod - def _translate_guided_params(extra_body: Mapping[str, object]) -> tuple[tuple[str, object], ...]: + def _translate_guided_params( + extra_body: Mapping[str, object], optional_params: Mapping[str, object] + ) -> tuple[tuple[str, object], ...]: + has_guided: Final = any( + extra_body.get(key) is not None for key in ("guided_json", "guided_grammar", "guided_choice") + ) + if not has_guided: + return () + if "response_format" in optional_params or "response_format" in extra_body: + verbose_logger.debug( + "fireworks_ai ignoring guided decoding params; explicit response_format takes precedence." + ) + return () if extra_body.get("guided_json") is not None: return (("response_format", _json_schema_response_format(extra_body["guided_json"])),) if extra_body.get("guided_grammar") is not None: @@ -429,13 +403,11 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): "grammar": extra_body["guided_grammar"], } return (("response_format", grammar_response_format),) - if extra_body.get("guided_choice") is not None: - choice_schema: Final = { # mutable-ok: JSON request body - "type": "string", - "enum": extra_body["guided_choice"], - } - return (("response_format", _json_schema_response_format(choice_schema)),) - return () + choice_schema: Final = { # mutable-ok: JSON request body + "type": "string", + "enum": extra_body["guided_choice"], + } + return (("response_format", _json_schema_response_format(choice_schema)),) def _transform_tools(self, tools: list[OpenAIChatCompletionToolParam]) -> list[OpenAIChatCompletionToolParam]: for tool in tools: diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index d25bbd69b91..48d868b5846 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1308,18 +1308,19 @@ def test_map_extra_body_params_translates_truncate_prompt_tokens(): assert result == {"prompt_truncate_len": 4096} -def test_map_extra_body_params_truncate_prompt_tokens_conflicts_with_alias(): +def test_map_extra_body_params_truncate_prompt_tokens_native_wins(): config = FireworksAIConfig() - with pytest.raises(litellm.BadRequestError, match="aliases"): - config.map_extra_body_params( - {"prompt_truncate_len": 2048, "extra_body": {"truncate_prompt_tokens": 4096}}, - _REASONING_MODEL, - ) - with pytest.raises(litellm.BadRequestError, match="aliases"): - config.map_extra_body_params( - {"extra_body": {"truncate_prompt_tokens": 4096, "prompt_truncate_len": 2048}}, - _REASONING_MODEL, - ) + top_level = config.map_extra_body_params( + {"prompt_truncate_len": 2048, "extra_body": {"truncate_prompt_tokens": 4096}}, + _REASONING_MODEL, + ) + assert top_level == {"prompt_truncate_len": 2048} + + nested = config.map_extra_body_params( + {"extra_body": {"truncate_prompt_tokens": 4096, "prompt_truncate_len": 2048}}, + _REASONING_MODEL, + ) + assert nested == {"extra_body": {"prompt_truncate_len": 2048}} def test_map_extra_body_params_chat_template_kwargs_enable_thinking(): @@ -1337,28 +1338,29 @@ def test_map_extra_body_params_chat_template_kwargs_enable_thinking(): assert enabled == {} -def test_map_extra_body_params_chat_template_kwargs_conflicts_with_reasoning_effort(): +def test_map_extra_body_params_chat_template_kwargs_native_reasoning_effort_wins(): config = FireworksAIConfig() - with pytest.raises(litellm.BadRequestError, match="enable_thinking"): - config.map_extra_body_params( - { - "reasoning_effort": "high", - "extra_body": {"chat_template_kwargs": {"enable_thinking": False}}, - }, - _REASONING_MODEL, - ) + result = config.map_extra_body_params( + { + "reasoning_effort": "high", + "extra_body": {"chat_template_kwargs": {"enable_thinking": False}}, + }, + _REASONING_MODEL, + ) + assert result == {"reasoning_effort": "high"} -def test_map_extra_body_params_chat_template_kwargs_conflicts_with_thinking(): +def test_map_extra_body_params_chat_template_kwargs_native_thinking_wins(): config = FireworksAIConfig() - with pytest.raises(litellm.BadRequestError, match="enable_thinking"): - config.map_extra_body_params( - { - "thinking": {"type": "enabled", "budget_tokens": 4096}, - "extra_body": {"chat_template_kwargs": {"enable_thinking": True}}, - }, - _REASONING_MODEL, - ) + thinking = {"type": "enabled", "budget_tokens": 4096} + result = config.map_extra_body_params( + { + "thinking": thinking, + "extra_body": {"chat_template_kwargs": {"enable_thinking": True}}, + }, + _REASONING_MODEL, + ) + assert result == {"thinking": thinking} def test_map_extra_body_params_chat_template_kwargs_dropped_for_non_reasoning_model(): @@ -1370,6 +1372,15 @@ def test_map_extra_body_params_chat_template_kwargs_dropped_for_non_reasoning_mo assert result == {} +def test_map_extra_body_params_non_dict_chat_template_kwargs_dropped(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": "enable_thinking"}}, + _REASONING_MODEL, + ) + assert result == {} + + def test_map_extra_body_params_guided_json(): config = FireworksAIConfig() schema = {"type": "object", "properties": {"x": {"type": "string"}}} @@ -1401,25 +1412,37 @@ def test_map_extra_body_params_guided_grammar_and_choice(): } -def test_map_extra_body_params_guided_conflicts_with_response_format(): +def test_map_extra_body_params_guided_native_response_format_wins(): config = FireworksAIConfig() - with pytest.raises(litellm.BadRequestError, match="response_format"): - config.map_extra_body_params( - { - "response_format": {"type": "json_object"}, - "extra_body": {"guided_json": {"type": "object"}}, - }, - _REASONING_MODEL, - ) + top_level = config.map_extra_body_params( + { + "response_format": {"type": "json_object"}, + "extra_body": {"guided_json": {"type": "object"}}, + }, + _REASONING_MODEL, + ) + assert top_level == {"response_format": {"type": "json_object"}} + + nested_format = {"type": "json_object"} + nested = config.map_extra_body_params( + {"extra_body": {"guided_json": {"type": "object"}, "response_format": nested_format}}, + _REASONING_MODEL, + ) + assert nested == {"extra_body": {"response_format": nested_format}} -def test_map_extra_body_params_multiple_guided_params_rejected(): +def test_map_extra_body_params_multiple_guided_params_priority_order(): config = FireworksAIConfig() - with pytest.raises(litellm.BadRequestError, match="multiple guided decoding params"): - config.map_extra_body_params( - {"extra_body": {"guided_json": {"type": "object"}, "guided_grammar": "root ::= 'x'"}}, - _REASONING_MODEL, - ) + result = config.map_extra_body_params( + {"extra_body": {"guided_grammar": "root ::= 'x'", "guided_json": {"type": "object"}}}, + _REASONING_MODEL, + ) + assert result == { + "response_format": { + "type": "json_schema", + "json_schema": {"schema": {"type": "object"}}, + } + } @pytest.mark.parametrize( From 3d673f9534f961c7f709b0a70063f349ab7cfd2c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:58:20 +0000 Subject: [PATCH 033/576] fix(managed_files): skip unparseable rows when listing managed files get_user_created_file_ids validated every row's file_object without a guard, so a single row failing OpenAIFileObject validation raised ValidationError and turned the whole GET /v1/files response into a 500. #35365 covered the null case only, leaving malformed or partial rows able to take the entire listing down. Rows now parse through a helper that returns None on failure and logs a warning, matching how list_user_batches already tolerates rows it cannot parse, so one bad row costs its own entry instead of the caller's whole listing. Null rows stay silent since the batch cost poller registers those legitimately. Refs #35361 --- .../proxy/hooks/managed_files.py | 23 +++++++++++++++++-- .../proxy/test_managed_files_hook.py | 23 +++++++++++++++++++ 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index ec47b6ac0e6..2349b618a28 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -73,6 +73,20 @@ else: PrismaClient = Any +def _parse_managed_file_object( + raw_file_object: object, unified_file_id: str +) -> Optional[OpenAIFileObject]: + if raw_file_object is None: + return None + try: + return OpenAIFileObject.model_validate(raw_file_object) + except Exception as e: + verbose_logger.warning( + f"Failed to parse managed file object {unified_file_id}: {e}" + ) + return None + + class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Class variables or attributes def __init__( @@ -383,9 +397,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): } ) return [ - OpenAIFileObject.model_validate(file_object.file_object) + parsed_file_object for file_object in file_ids - if file_object.file_object is not None + if ( + parsed_file_object := _parse_managed_file_object( + file_object.file_object, file_object.unified_file_id + ) + ) + is not None ] async def check_managed_file_id_access( diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 4a4aa7aa5ea..4da6de6353f 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -154,6 +154,29 @@ async def test_get_user_created_file_ids_skips_rows_without_file_object(): assert [file.id for file in files] == ["file-output-abc"] +@pytest.mark.asyncio +async def test_get_user_created_file_ids_skips_unparseable_rows(): + managed_files = _make_managed_files_instance() + managed_files.prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=[ + MagicMock( + file_object={"id": "file-corrupt", "object": "file"}, + unified_file_id="unified-corrupt", + ), + MagicMock( + file_object=_make_file_object().model_dump(), + unified_file_id="unified-valid", + ), + ] + ) + + files = await managed_files.get_user_created_file_ids( + _make_user_api_key_dict(), ["file-output-abc"] + ) + + assert [file.id for file in files] == ["file-output-abc"] + + @pytest.mark.asyncio async def test_should_fallback_when_no_router(): """ From 1b6f3cebf1a4a4804a9bd9a0c3287cfc0d07c971 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 00:16:05 +0000 Subject: [PATCH 034/576] fix(managed_files): log sanitized validation errors when skipping rows The skip warning interpolated the full pydantic ValidationError, whose string embeds input_value with the rejected row's contents. Managed-file rows carry a caller-supplied filename, so a malformed row copied that into operational logs. Log the error locations, types, and messages via errors() with input, url, and context excluded, keeping the field-level diagnostics without the values. Non-validation failures fall back to the exception type. --- .../proxy/hooks/managed_files.py | 9 ++++++++- .../proxy/test_managed_files_hook.py | 19 +++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 2349b618a28..688ffb35ff7 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -7,6 +7,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast from fastapi import HTTPException +from pydantic import ValidationError import litellm from litellm import Router, verbose_logger @@ -80,9 +81,15 @@ def _parse_managed_file_object( return None try: return OpenAIFileObject.model_validate(raw_file_object) + except ValidationError as e: + verbose_logger.warning( + f"Failed to parse managed file object {unified_file_id}: " + f"{e.errors(include_input=False, include_url=False, include_context=False)}" + ) + return None except Exception as e: verbose_logger.warning( - f"Failed to parse managed file object {unified_file_id}: {e}" + f"Failed to parse managed file object {unified_file_id}: {type(e).__name__}" ) return None diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 4da6de6353f..6397e0be247 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -6,6 +6,7 @@ async_post_call_success_hook when processing completed batch responses. """ import json +import logging import pytest from typing import Optional @@ -154,6 +155,24 @@ async def test_get_user_created_file_ids_skips_rows_without_file_object(): assert [file.id for file in files] == ["file-output-abc"] +@pytest.mark.asyncio +async def test_parse_managed_file_object_warning_omits_rejected_values(caplog): + from litellm_enterprise.proxy.hooks.managed_files import ( + _parse_managed_file_object, + ) + + with caplog.at_level(logging.WARNING): + parsed = _parse_managed_file_object( + {"id": "file-corrupt", "object": "file", "filename": "confidential.jsonl"}, + "unified-corrupt", + ) + + assert parsed is None + assert "unified-corrupt" in caplog.text + assert "bytes" in caplog.text + assert "confidential.jsonl" not in caplog.text + + @pytest.mark.asyncio async def test_get_user_created_file_ids_skips_unparseable_rows(): managed_files = _make_managed_files_instance() From 7d00f9d019f84be709a7515094fed4ce7bbee900 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:46:00 -0700 Subject: [PATCH 035/576] fix(managed_files): return unified output file ids from GET /batches list_user_batches parsed each stored batch blob and returned it as-is, so any row whose blob still carried raw provider file ids (for example a batch that reached a terminal state through the cost poller, or rows written before output registration existed) leaked raw output_file_id and error_file_id values that clients cannot fetch through the proxy. The list path now runs each row through ensure_batch_response_managed_file_ids, which swaps in existing managed ids and registers missing ones under the batch owner's identity, matching what GET /batches/{id} already does --- .../proxy/hooks/managed_files.py | 12 ++ .../proxy/hooks/test_managed_files.py | 142 ++++++++++++++++++ 2 files changed, 154 insertions(+) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 0036603bcd1..07a1f959940 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -31,6 +31,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + ensure_batch_response_managed_file_ids, get_batch_id_from_unified_batch_id, get_content_type_from_file_object, get_model_id_from_unified_batch_id, @@ -352,6 +353,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) batch_obj = LiteLLMBatch.model_validate(batch_data) batch_obj.id = batch.unified_object_id + await ensure_batch_response_managed_file_ids( + response=batch_obj, + managed_files_obj=self, + prisma_client=self.prisma_client, + verbose_proxy_logger=verbose_logger, + user_api_key_dict=user_api_key_dict, + db_batch_object=batch, + unified_batch_id=_is_base64_encoded_unified_file_id( + batch.unified_object_id + ), + ) batch_objects.append(batch_obj) except Exception as e: diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 50af6465d06..fc10a1257e1 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -1813,6 +1813,148 @@ def _create_unified_batch_id(model_id: str, batch_id: str) -> str: return base64.urlsafe_b64encode(unified_str.encode()).decode().rstrip("=") +def _decode_unified_id(b64_id: str) -> str: + return base64.urlsafe_b64decode(b64_id + "=" * (-len(b64_id) % 4)).decode() + + +def _terminal_batch_record( + unified_batch_uid: str, + raw_input_file_id: str, + raw_output_file_id: str, + raw_error_file_id: str, +): + record = MagicMock() + record.unified_object_id = unified_batch_uid + record.created_by = "owner-user" + record.team_id = "owner-team" + record.status = "cancelled" + record.file_object = json.dumps( + { + "id": "batch-raw-456", + "object": "batch", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "status": "cancelled", + "created_at": 1234567890, + "input_file_id": raw_input_file_id, + "output_file_id": raw_output_file_id, + "error_file_id": raw_error_file_id, + } + ) + return record + + +@pytest.mark.asyncio +async def test_list_batches_registers_and_returns_unified_output_file_ids(): + """A stored batch blob with raw provider file IDs (e.g. persisted by the cost + poller for a cancelled batch) must be listed with unified managed IDs, and the + output/error files must be registered in the managed file table so GET + /files/{id}/content can route them.""" + from litellm.proxy._types import UserAPIKeyAuth + + unified_batch_uid = _create_unified_batch_id("model-123", "batch-456") + raw_input_file_id = "file-list-in-1" + raw_output_file_id = "file-list-out-1" + raw_error_file_id = "file-list-err-1" + unified_input_file_id = base64.urlsafe_b64encode( + b"litellm_proxy:application/octet-stream;unified_id,in-1;target_model_names,gpt-5-batch" + ).decode() + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [ + _terminal_batch_record( + unified_batch_uid, raw_input_file_id, raw_output_file_id, raw_error_file_id + ) + ] + + input_file_row = MagicMock() + input_file_row.unified_file_id = unified_input_file_id + + def find_managed_file(where): + if where["flat_model_file_ids"]["has"] == raw_input_file_id: + return input_file_row + return None + + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( + side_effect=find_managed_file + ) + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), + limit=10, + ) + + listed = result["data"][0] + assert listed.id == unified_batch_uid + assert listed.input_file_id == unified_input_file_id + + decoded_output = _decode_unified_id(listed.output_file_id) + assert decoded_output.startswith("litellm_proxy") + assert f"llm_output_file_id,{raw_output_file_id}" in decoded_output + assert "llm_output_file_model_id,model-123" in decoded_output + assert "target_model_names,gpt-5-batch" in decoded_output + + decoded_error = _decode_unified_id(listed.error_file_id) + assert f"llm_output_file_id,{raw_error_file_id}" in decoded_error + + upsert_calls = prisma_client.db.litellm_managedfiletable.upsert.await_args_list + stored_raw_ids = { + c.kwargs["data"]["create"]["flat_model_file_ids"][0] for c in upsert_calls + } + assert stored_raw_ids == {raw_output_file_id, raw_error_file_id} + for c in upsert_calls: + assert c.kwargs["data"]["create"]["created_by"] == "owner-user" + assert c.kwargs["data"]["create"]["team_id"] == "owner-team" + + +@pytest.mark.asyncio +async def test_list_batches_resolves_existing_managed_rows_without_minting(): + """When the raw provider file IDs already have managed file rows, listing must + swap in the existing unified IDs and must not upsert duplicate rows.""" + from litellm.proxy._types import UserAPIKeyAuth + + unified_batch_uid = _create_unified_batch_id("model-123", "batch-456") + raw_output_file_id = "file-list-out-existing" + existing_unified_output_id = base64.urlsafe_b64encode( + f"litellm_proxy:application/json;unified_id,u-9;llm_output_file_id,{raw_output_file_id}".encode() + ).decode() + + record = _terminal_batch_record( + unified_batch_uid, "file-list-in-9", raw_output_file_id, "" + ) + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [record] + + existing_row = MagicMock() + existing_row.unified_file_id = existing_unified_output_id + + def find_managed_file(where): + if where["flat_model_file_ids"]["has"] == raw_output_file_id: + return existing_row + return None + + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( + side_effect=find_managed_file + ) + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), + limit=10, + ) + + assert result["data"][0].output_file_id == existing_unified_output_id + prisma_client.db.litellm_managedfiletable.upsert.assert_not_awaited() + + @pytest.mark.asyncio async def test_list_batches_from_managed_objects_table_provider_filter_raises_exception(): from litellm.proxy._types import UserAPIKeyAuth From 59041240f036fe80776b297b36757c48d85f7978 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 02:23:32 -0700 Subject: [PATCH 036/576] fix(managed_files): cap batch list page size at 100 and bulk-resolve raw file ids in one query --- .../proxy/hooks/managed_files.py | 104 +++++++++++++----- .../openai_files_endpoints/common_utils.py | 30 +++++ .../proxy/hooks/test_managed_files.py | 94 +++++++++++----- 3 files changed, 172 insertions(+), 56 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 07a1f959940..6fa6ef46ad4 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -3,6 +3,7 @@ import base64 import json +from collections.abc import Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Union, cast from uuid import NAMESPACE_URL, uuid5 @@ -31,10 +32,12 @@ from litellm.proxy._types import ( ) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + apply_unified_file_ids, ensure_batch_response_managed_file_ids, get_batch_id_from_unified_batch_id, get_content_type_from_file_object, get_model_id_from_unified_batch_id, + map_raw_file_ids_to_unified, normalize_mime_type_for_provider, resolve_managed_output_file_model_name, ) @@ -62,6 +65,9 @@ if TYPE_CHECKING: if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from prisma.models import ( + LiteLLM_ManagedObjectTable as PrismaManagedObjectRow, + ) from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache from litellm.proxy.utils import PrismaClient as _PrismaClient @@ -75,6 +81,20 @@ else: PrismaClient = Any +def _decode_json_blob(blob: object) -> object: + return json.loads(blob) if isinstance(blob, str) else blob + + +def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> Optional[LiteLLMBatch]: + try: + batch_obj: Final = LiteLLMBatch.model_validate(_decode_json_blob(row.file_object)) + except Exception as e: + verbose_logger.warning(f"Failed to parse batch object {row.unified_object_id}: {e}") + return None + batch_obj.id = row.unified_object_id + return batch_obj + + class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Class variables or attributes def __init__( @@ -329,7 +349,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): detail=f"Invalid 'after' cursor: no batch found with id '{after}'.", ) - page_size = limit or 20 + page_size: Final = min(limit or 20, 100) cursor_args: Dict[str, Any] = ( {"cursor": {"unified_object_id": after}, "skip": 1} if after else {} ) @@ -343,36 +363,60 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): has_more = len(batches) > page_size - batch_objects: List[LiteLLMBatch] = [] - for batch in batches[:page_size]: - try: - batch_data = ( - json.loads(batch.file_object) - if isinstance(batch.file_object, str) - else batch.file_object - ) - batch_obj = LiteLLMBatch.model_validate(batch_data) - batch_obj.id = batch.unified_object_id - await ensure_batch_response_managed_file_ids( - response=batch_obj, - managed_files_obj=self, - prisma_client=self.prisma_client, - verbose_proxy_logger=verbose_logger, - user_api_key_dict=user_api_key_dict, - db_batch_object=batch, - unified_batch_id=_is_base64_encoded_unified_file_id( - batch.unified_object_id - ), - ) - batch_objects.append(batch_obj) + parsed_rows: Final = tuple( + (row, batch_obj) + for row in batches[:page_size] + if (batch_obj := _parse_managed_batch_row(row)) is not None + ) + unified_id_by_raw_id: Final = await map_raw_file_ids_to_unified( + raw_file_ids=frozenset( + file_id + for _, batch_obj in parsed_rows + for file_id in (batch_obj.input_file_id, batch_obj.output_file_id, batch_obj.error_file_id) + if file_id and not _is_base64_encoded_unified_file_id(file_id) + ), + prisma_client=self.prisma_client, + ) + resolved_batches: Final = [ + await self._resolve_listed_batch( + row=row, + batch_obj=batch_obj, + unified_id_by_raw_id=unified_id_by_raw_id, + user_api_key_dict=user_api_key_dict, + ) + for row, batch_obj in parsed_rows + ] + return build_list_page( + [batch_obj for batch_obj in resolved_batches if batch_obj is not None], + has_more=has_more, + ) - except Exception as e: - verbose_logger.warning( - f"Failed to parse batch object {batch.unified_object_id}: {e}" - ) - continue - - return build_list_page(batch_objects, has_more=has_more) + async def _resolve_listed_batch( + self, + row: "PrismaManagedObjectRow", + batch_obj: LiteLLMBatch, + unified_id_by_raw_id: Mapping[str, str], + user_api_key_dict: UserAPIKeyAuth, + ) -> Optional[LiteLLMBatch]: + apply_unified_file_ids(batch_obj, unified_id_by_raw_id) + try: + await ensure_batch_response_managed_file_ids( + response=batch_obj, + managed_files_obj=self, + prisma_client=self.prisma_client, + verbose_proxy_logger=verbose_logger, + user_api_key_dict=user_api_key_dict, + db_batch_object=row, + unified_batch_id=_is_base64_encoded_unified_file_id( + row.unified_object_id + ), + ) + except Exception as e: + verbose_logger.warning( + f"Failed to resolve managed file ids for batch {row.unified_object_id}: {e}" + ) + return None + return batch_obj async def get_user_created_file_ids( self, user_api_key_dict: UserAPIKeyAuth, model_object_ids: List[str] diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 080b8b80ae4..bf83a7cf25c 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1,6 +1,7 @@ import base64 import mimetypes import re +from collections.abc import Mapping from dataclasses import dataclass, field from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Optional @@ -16,6 +17,7 @@ if TYPE_CHECKING: from prisma.models import LiteLLM_ManagedObjectTable from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import PrismaClient from litellm.router import Router from litellm.types.utils import LiteLLMBatch @@ -1002,6 +1004,34 @@ async def resolve_output_file_ids_to_unified(response, prisma_client) -> None: pass +async def map_raw_file_ids_to_unified( + raw_file_ids: frozenset[str], prisma_client: "PrismaClient | None" +) -> Mapping[str, str]: + if not raw_file_ids or not prisma_client: + return MappingProxyType({}) + managed_files: Final = await ManagedFileRepository(prisma_client).table.find_many( + where={"flat_model_file_ids": {"hasSome": sorted(raw_file_ids)}} # mutable-ok: prisma where is a plain dict + ) + return MappingProxyType( + { + raw_id: managed_file.unified_file_id + for managed_file in managed_files + for raw_id in managed_file.flat_model_file_ids + if raw_id in raw_file_ids + } + ) + + +def apply_unified_file_ids(response: "LiteLLMBatch", unified_id_by_raw_id: Mapping[str, str]) -> None: + for file_attr, raw_id in ( + ("input_file_id", getattr(response, "input_file_id", None)), + ("output_file_id", getattr(response, "output_file_id", None)), + ("error_file_id", getattr(response, "error_file_id", None)), + ): + if isinstance(raw_id, str) and raw_id in unified_id_by_raw_id: + setattr(response, file_attr, unified_id_by_raw_id[raw_id]) + + async def ensure_batch_response_managed_file_ids( response, managed_files_obj, diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index fc10a1257e1..e1e5cc6c532 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -1869,15 +1869,12 @@ async def test_list_batches_registers_and_returns_unified_output_file_ids(): input_file_row = MagicMock() input_file_row.unified_file_id = unified_input_file_id + input_file_row.flat_model_file_ids = [raw_input_file_id] - def find_managed_file(where): - if where["flat_model_file_ids"]["has"] == raw_input_file_id: - return input_file_row - return None - - prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( - side_effect=find_managed_file + prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=[input_file_row] ) + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client @@ -1892,6 +1889,13 @@ async def test_list_batches_registers_and_returns_unified_output_file_ids(): assert listed.id == unified_batch_uid assert listed.input_file_id == unified_input_file_id + bulk_lookup = prisma_client.db.litellm_managedfiletable.find_many.await_args + assert set(bulk_lookup.kwargs["where"]["flat_model_file_ids"]["hasSome"]) == { + raw_input_file_id, + raw_output_file_id, + raw_error_file_id, + } + decoded_output = _decode_unified_id(listed.output_file_id) assert decoded_output.startswith("litellm_proxy") assert f"llm_output_file_id,{raw_output_file_id}" in decoded_output @@ -1914,33 +1918,43 @@ async def test_list_batches_registers_and_returns_unified_output_file_ids(): @pytest.mark.asyncio async def test_list_batches_resolves_existing_managed_rows_without_minting(): """When the raw provider file IDs already have managed file rows, listing must - swap in the existing unified IDs and must not upsert duplicate rows.""" + swap in the existing unified IDs via one bulk lookup for the whole page, with + no per-row queries and no duplicate upserts.""" from litellm.proxy._types import UserAPIKeyAuth - unified_batch_uid = _create_unified_batch_id("model-123", "batch-456") - raw_output_file_id = "file-list-out-existing" - existing_unified_output_id = base64.urlsafe_b64encode( - f"litellm_proxy:application/json;unified_id,u-9;llm_output_file_id,{raw_output_file_id}".encode() + unified_input_file_id = base64.urlsafe_b64encode( + b"litellm_proxy:application/octet-stream;unified_id,in-9;target_model_names,gpt-5-batch" ).decode() + raw_output_file_ids = ["file-list-out-existing-1", "file-list-out-existing-2"] + existing_unified_output_ids = [ + base64.urlsafe_b64encode( + f"litellm_proxy:application/json;unified_id,u-{i};llm_output_file_id,{raw_id}".encode() + ).decode() + for i, raw_id in enumerate(raw_output_file_ids) + ] - record = _terminal_batch_record( - unified_batch_uid, "file-list-in-9", raw_output_file_id, "" - ) + records = [ + _terminal_batch_record( + _create_unified_batch_id("model-123", f"batch-{i}"), + unified_input_file_id, + raw_id, + "", + ) + for i, raw_id in enumerate(raw_output_file_ids) + ] prisma_client = AsyncMock() - prisma_client.db.litellm_managedobjecttable.find_many.return_value = [record] + prisma_client.db.litellm_managedobjecttable.find_many.return_value = records - existing_row = MagicMock() - existing_row.unified_file_id = existing_unified_output_id + existing_rows = [ + MagicMock(unified_file_id=unified_id, flat_model_file_ids=[raw_id]) + for raw_id, unified_id in zip(raw_output_file_ids, existing_unified_output_ids) + ] - def find_managed_file(where): - if where["flat_model_file_ids"]["has"] == raw_output_file_id: - return existing_row - return None - - prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( - side_effect=find_managed_file + prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=existing_rows ) + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock() proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client @@ -1951,10 +1965,38 @@ async def test_list_batches_resolves_existing_managed_rows_without_minting(): limit=10, ) - assert result["data"][0].output_file_id == existing_unified_output_id + assert [b.output_file_id for b in result["data"]] == existing_unified_output_ids + prisma_client.db.litellm_managedfiletable.find_many.assert_awaited_once() + prisma_client.db.litellm_managedfiletable.find_first.assert_not_awaited() prisma_client.db.litellm_managedfiletable.upsert.assert_not_awaited() +@pytest.mark.asyncio +async def test_list_batches_caps_page_size_at_100(): + """The list page size must be capped at 100 rows (matching OpenAI's limit) + even when the caller asks for more, so one request cannot fan out into an + unbounded scan.""" + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [] + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), + limit=100000, + ) + + assert ( + prisma_client.db.litellm_managedobjecttable.find_many.await_args.kwargs["take"] + == 101 + ) + assert result["data"] == [] + + @pytest.mark.asyncio async def test_list_batches_from_managed_objects_table_provider_filter_raises_exception(): from litellm.proxy._types import UserAPIKeyAuth From 5a5bb8c9d844870c25684e169960f1571d08e5ce Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:29:39 +0000 Subject: [PATCH 037/576] fix(proxy): stop /{provider}/v1/files from capturing /openai_passthrough The native files and batches routes declare /{provider}/v1/... and their routers are mounted before the passthrough router, so /openai_passthrough/v1/files and /openai_passthrough/v1/batches matched them with provider="openai_passthrough" and 500'd on the LlmProviders lookup instead of reaching openai_proxy_route. Move the dedicated /openai_passthrough prefix onto its own router mounted ahead of the batches and files routers. /openai/... and every other provider prefix keep their current behavior. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_passthrough_endpoints.py | 3 +- litellm/proxy/proxy_server.py | 2 + .../test_llm_pass_through_endpoints.py | 57 +++++++++++++++++++ 3 files changed, 61 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 38da00a3bb9..baa74c19182 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -60,6 +60,7 @@ from .passthrough_endpoint_router import PassthroughEndpointRouter vertex_llm_base: Final = VertexBase() router: Final = APIRouter() +openai_passthrough_router: Final = APIRouter() default_vertex_config: Final = None passthrough_endpoint_router: Final = PassthroughEndpointRouter() @@ -1875,7 +1876,7 @@ async def vertex_proxy_route( ) -@router.api_route( +@openai_passthrough_router.api_route( "/openai_passthrough/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], tags=["OpenAI Pass-through", "pass-through"], diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fb9c4e67aad..e75277e7f0a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -522,6 +522,7 @@ from litellm.proxy.openai_files_endpoints.files_endpoints import ( set_files_config, ) from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + openai_passthrough_router, passthrough_endpoint_router, vertex_ai_live_websocket_passthrough, ) @@ -16433,6 +16434,7 @@ app.include_router(search_router) app.include_router(image_router) app.include_router(fine_tuning_router) app.include_router(credential_router) +app.include_router(openai_passthrough_router) app.include_router(batches_router) app.include_router(openai_files_router) app.include_router(llm_passthrough_router) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 181846fe289..27d6e4c8585 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -2814,6 +2814,63 @@ class TestOpenAIPassthroughRoute: assert result == {"id": "asst_123", "object": "assistant"} +def _resolve_route_name(method: str, path: str) -> str | None: + from starlette.routing import Match + + from litellm.proxy.proxy_server import app + + scope = { + "type": "http", + "method": method, + "path": path, + "headers": [], + "query_string": b"", + "root_path": "", + } + for route in app.router.routes: + if route.matches(scope)[0] == Match.FULL: + return getattr(route, "name", None) + return None + + +@pytest.mark.parametrize( + "method, path", + [ + ("POST", "/openai_passthrough/v1/files"), + ("GET", "/openai_passthrough/v1/files"), + ("GET", "/openai_passthrough/v1/files/file-abc123"), + ("DELETE", "/openai_passthrough/v1/files/file-abc123"), + ("GET", "/openai_passthrough/v1/files/file-abc123/content"), + ("POST", "/openai_passthrough/v1/batches"), + ("GET", "/openai_passthrough/v1/batches"), + ("GET", "/openai_passthrough/v1/batches/batch_abc123"), + ("POST", "/openai_passthrough/v1/batches/batch_abc123/cancel"), + ("POST", "/openai_passthrough/v1/responses"), + ], +) +def test_openai_passthrough_prefix_wins_over_native_provider_routes(method, path): + """ + /openai_passthrough exists to guarantee passthrough, so the native + /{provider}/v1/files and /{provider}/v1/batches routes must never capture it + with provider="openai_passthrough" (which 500s on the LlmProviders lookup). + """ + assert _resolve_route_name(method, path) == "openai_proxy_route" + + +@pytest.mark.parametrize( + "method, path, expected_name", + [ + ("POST", "/openai/v1/files", "create_file"), + ("GET", "/azure/v1/files", "list_files"), + ("POST", "/v1/files", "create_file"), + ("POST", "/v1/batches", "create_batch"), + ("POST", "/openai/v1/chat/completions", "openai_proxy_route"), + ], +) +def test_native_provider_routes_are_unchanged(method, path, expected_name): + assert _resolve_route_name(method, path) == expected_name + + class TestCursorProxyRoute: """Tests for the Cursor Cloud Agents pass-through route.""" From 357f90fa39d18c9a158a978ebd1ed0fecac6044d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:33:40 +0000 Subject: [PATCH 038/576] fix(proxy): scope file list pagination cursors to the caller GET /v1/files filters data down to the caller's own managed files but left first_id and last_id as the upstream page's, so a non-owner got back file ids belonging to other users even with an empty data array --- .../proxy/hooks/managed_files.py | 15 +++ .../proxy/hooks/test_managed_files.py | 100 ++++++++++++++++++ 2 files changed, 115 insertions(+) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 0036603bcd1..851e202e2fb 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1270,10 +1270,25 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) ## Filter the response to only include the files created by the user response.data = user_created_file_ids # type: ignore + self._scope_list_page_cursors(response, user_created_file_ids) return response return response return response + @staticmethod + def _scope_list_page_cursors(response: AsyncCursorPage, data: List[OpenAIFileObject]) -> None: + """Rebuild ``first_id`` / ``last_id`` from the caller-scoped page. + + The upstream cursors point at rows that were just filtered out, so + leaving them in place discloses other callers' file ids. + """ + if hasattr(response, "first_id"): + response.first_id = data[0].id if data else None + if hasattr(response, "last_id"): + response.last_id = data[-1].id if data else None + if not data and hasattr(response, "has_more"): + response.has_more = False + async def afile_retrieve( self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router=None ) -> OpenAIFileObject: diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 50af6465d06..3384c553740 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -2861,3 +2861,103 @@ async def test_same_user_different_keys_can_access_batch(): assert "batch_id" in result2 # Both keys should get the same result assert result1["batch_id"] == result2["batch_id"] + + +@pytest.mark.asyncio +async def test_file_list_cursors_are_scoped_to_the_caller(): + """A non-owner must not learn other callers' file ids through the page cursors.""" + from openai.pagination import AsyncCursorPage + from openai.types import FileObject + + from litellm.proxy._types import UserAPIKeyAuth + + owner_file = FileObject( + id="file-owner-1", + bytes=100, + created_at=1, + filename="owner.jsonl", + object="file", + purpose="batch", + status="processed", + ) + upstream_page = AsyncCursorPage[FileObject].construct( + data=[owner_file], + has_more=True, + first_id=owner_file.id, + last_id=owner_file.id, + object="list", + ) + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedfiletable.find_many.return_value = [] + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + response = await proxy_managed_files.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth( + user_id="other-user", team_id="other-team", parent_otel_span=MagicMock() + ), + response=upstream_page, + ) + + assert response.data == [] + assert response.first_id is None + assert response.last_id is None + assert response.has_more is False + + +@pytest.mark.asyncio +async def test_file_list_cursors_follow_the_owner_scoped_page(): + from openai.pagination import AsyncCursorPage + from openai.types import FileObject + + from litellm.proxy._types import UserAPIKeyAuth + + def _raw_file(file_id: str) -> FileObject: + return FileObject( + id=file_id, + bytes=100, + created_at=1, + filename=f"{file_id}.jsonl", + object="file", + purpose="batch", + status="processed", + ) + + upstream_page = AsyncCursorPage[FileObject].construct( + data=[_raw_file("file-someone-else"), _raw_file("file-mine")], + has_more=False, + first_id="file-someone-else", + last_id="file-mine", + object="list", + ) + + managed_row = MagicMock() + managed_row.file_object = { + "id": "litellm_proxy:mine", + "bytes": 100, + "created_at": 1, + "filename": "mine.jsonl", + "object": "file", + "purpose": "batch", + "status": "processed", + } + prisma_client = AsyncMock() + prisma_client.db.litellm_managedfiletable.find_many.return_value = [managed_row] + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + response = await proxy_managed_files.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth( + user_id="mine-user", parent_otel_span=MagicMock() + ), + response=upstream_page, + ) + + assert [file_object.id for file_object in response.data] == ["litellm_proxy:mine"] + assert response.first_id == "litellm_proxy:mine" + assert response.last_id == "litellm_proxy:mine" From 4a601c49a60d34d12810bd0372b062dca57d34b7 Mon Sep 17 00:00:00 2001 From: Miles Adkins Date: Thu, 6 Aug 2026 10:46:48 -0500 Subject: [PATCH 039/576] feat(fireworks_ai): full chat_template_kwargs parity with the gateway Map the remaining gateway-documented effort keys: thinking as an alias for enable_thinking (enable_thinking wins when both are present), reasoning_budget to an integer reasoning_effort (skipped when thinking is explicitly off), and low_effort=true to reasoning_effort=low (budget wins when both are set). guided_json and guided_choice response_format wrappers now include the name field (response and choice) to match the gateway wire shape. --- .../llms/fireworks_ai/chat/transformation.py | 47 ++++++++---- .../test_fireworks_ai_chat_transformation.py | 72 ++++++++++++++++++- 2 files changed, 104 insertions(+), 15 deletions(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 3bacb3cd28e..6b763be0bfe 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -61,8 +61,32 @@ def _extract_fireworks_hidden_params(payload: dict) -> dict: return {**top_level, **per_choice} -def _json_schema_response_format(schema: object) -> Mapping[str, object]: - return {"type": "json_schema", "json_schema": {"schema": schema}} # mutable-ok: JSON request body +def _json_schema_response_format(schema: object, name: str) -> Mapping[str, object]: + return {"type": "json_schema", "json_schema": {"name": name, "schema": schema}} # mutable-ok: JSON request body + + +_EFFORT_KWARG_KEYS: Final = frozenset({"enable_thinking", "thinking", "reasoning_budget", "low_effort"}) + + +def _bool_from_kwargs(kwargs: Mapping[str, object], keys: tuple[str, ...]) -> bool | None: + for key in keys: + value = kwargs.get(key) + if isinstance(value, bool): + return value + return None + + +def _effort_from_chat_template_kwargs(kwargs: Mapping[str, object]) -> object: + enable_thinking: Final = _bool_from_kwargs(kwargs, ("enable_thinking", "thinking")) + if enable_thinking is False: + return "none" + budget: Final = kwargs.get("reasoning_budget") + if isinstance(budget, (int, float)) and not isinstance(budget, bool) and budget > 0: + return int(budget) + low_effort: Final = _bool_from_kwargs(kwargs, ("low_effort",)) + if low_effort is True: + return "low" + return None _NIM_VLLM_STRIP_PARAMS: Final = frozenset( @@ -357,29 +381,28 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): type(chat_template_kwargs).__name__, ) return () - other_keys: Final = tuple(sorted(k for k in chat_template_kwargs if k != "enable_thinking")) + other_keys: Final = tuple(sorted(k for k in chat_template_kwargs if k not in _EFFORT_KWARG_KEYS)) if other_keys: verbose_logger.debug( "fireworks_ai does not support chat_template_kwargs keys %s for model=%s; dropping them.", other_keys, model, ) - if "enable_thinking" not in chat_template_kwargs: - return () if "reasoning_effort" in optional_params or "thinking" in optional_params: verbose_logger.debug( - "fireworks_ai ignoring chat_template_kwargs.enable_thinking; explicit reasoning_effort/thinking takes precedence." + "fireworks_ai ignoring chat_template_kwargs; explicit reasoning_effort/thinking takes precedence." ) return () + effort: Final = _effort_from_chat_template_kwargs(chat_template_kwargs) + if effort is None: + return () if not supports_reasoning(model=model, custom_llm_provider="fireworks_ai"): verbose_logger.debug( - "fireworks_ai model %r does not support reasoning; dropping chat_template_kwargs.enable_thinking.", + "fireworks_ai model %r does not support reasoning; dropping chat_template_kwargs effort keys.", model, ) return () - if chat_template_kwargs["enable_thinking"]: - return () - return (("reasoning_effort", "none"),) + return (("reasoning_effort", effort),) @staticmethod def _translate_guided_params( @@ -396,7 +419,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): ) return () if extra_body.get("guided_json") is not None: - return (("response_format", _json_schema_response_format(extra_body["guided_json"])),) + return (("response_format", _json_schema_response_format(extra_body["guided_json"], "response")),) if extra_body.get("guided_grammar") is not None: grammar_response_format: Final = { # mutable-ok: JSON request body "type": "grammar", @@ -407,7 +430,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): "type": "string", "enum": extra_body["guided_choice"], } - return (("response_format", _json_schema_response_format(choice_schema)),) + return (("response_format", _json_schema_response_format(choice_schema, "choice")),) def _transform_tools(self, tools: list[OpenAIChatCompletionToolParam]) -> list[OpenAIChatCompletionToolParam]: for tool in tools: diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 48d868b5846..e1b5d457205 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1338,6 +1338,66 @@ def test_map_extra_body_params_chat_template_kwargs_enable_thinking(): assert enabled == {} +def test_map_extra_body_params_chat_template_kwargs_thinking_alias(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"thinking": False}}}, + _REASONING_MODEL, + ) + assert result == {"reasoning_effort": "none"} + + +def test_map_extra_body_params_chat_template_kwargs_enable_thinking_wins_over_thinking(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"enable_thinking": True, "thinking": False}}}, + _REASONING_MODEL, + ) + assert result == {} + + +def test_map_extra_body_params_chat_template_kwargs_reasoning_budget(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"reasoning_budget": 512}}}, + _REASONING_MODEL, + ) + assert result == {"reasoning_effort": 512} + + +def test_map_extra_body_params_chat_template_kwargs_budget_ignored_when_thinking_off(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"enable_thinking": False, "reasoning_budget": 512}}}, + _REASONING_MODEL, + ) + assert result == {"reasoning_effort": "none"} + + +def test_map_extra_body_params_chat_template_kwargs_low_effort(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"low_effort": True}}}, + _REASONING_MODEL, + ) + assert result == {"reasoning_effort": "low"} + + budget_wins = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"low_effort": True, "reasoning_budget": 256}}}, + _REASONING_MODEL, + ) + assert budget_wins == {"reasoning_effort": 256} + + +def test_map_extra_body_params_chat_template_kwargs_effort_keys_dropped_for_non_reasoning_model(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"reasoning_budget": 512, "low_effort": True}}}, + _NON_REASONING_MODEL, + ) + assert result == {} + + def test_map_extra_body_params_chat_template_kwargs_native_reasoning_effort_wins(): config = FireworksAIConfig() result = config.map_extra_body_params( @@ -1388,7 +1448,10 @@ def test_map_extra_body_params_guided_json(): {"extra_body": {"guided_json": schema}}, _REASONING_MODEL ) assert result == { - "response_format": {"type": "json_schema", "json_schema": {"schema": schema}} + "response_format": { + "type": "json_schema", + "json_schema": {"name": "response", "schema": schema}, + } } @@ -1407,7 +1470,10 @@ def test_map_extra_body_params_guided_grammar_and_choice(): assert choice == { "response_format": { "type": "json_schema", - "json_schema": {"schema": {"type": "string", "enum": ["yes", "no"]}}, + "json_schema": { + "name": "choice", + "schema": {"type": "string", "enum": ["yes", "no"]}, + }, } } @@ -1440,7 +1506,7 @@ def test_map_extra_body_params_multiple_guided_params_priority_order(): assert result == { "response_format": { "type": "json_schema", - "json_schema": {"schema": {"type": "object"}}, + "json_schema": {"name": "response", "schema": {"type": "object"}}, } } From b5823d5894d28130b1a8748c9edea898d4055452 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 6 Aug 2026 09:49:13 -0700 Subject: [PATCH 040/576] feat(terraform): sync provider 0.3.0 from mirror and cut 0.4.0 The provider's release gate in project-releaser publishes only when the topmost released heading in terraform/provider/CHANGELOG.md moves past the tag the mirror already carries. That heading has been 0.2.2 since 2026-05-13, so every stable release since has correctly decided there was nothing to publish and the registry has gone stale. Two things were blocking a release: 1. The mirror shipped 0.3.0 out-of-band on 2026-07-13 (pricing_base_model, BerriAI/terraform-provider-litellm#47) after the source move, so that code exists only in the mirror. The publish rsyncs monorepo -> mirror with --delete, so publishing without this port would have deleted a released feature from the registry. 2. Nothing here declared a new version. Port #47 verbatim (resource_model.go and resource_model_crud.go are now byte-identical to the mirror's released files), backfill the 0.3.0 changelog entry it shipped under, and cut 0.4.0 covering the changes made here since the source move. 0.3.0 is not reusable as the next version -- the mirror holds that tag and the publish workflow's tag guard rejects it. --- terraform/provider/CHANGELOG.md | 13 ++++++++++++ terraform/provider/docs/resources/model.md | 2 ++ terraform/provider/litellm/resource_model.go | 8 +++++++ .../provider/litellm/resource_model_crud.go | 21 +++++++++++++++++-- 4 files changed, 42 insertions(+), 2 deletions(-) diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 101519c0b08..7c744f04064 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -7,13 +7,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [0.4.0] - 2026-08-06 + ### Fixed - **organization**: Send `PATCH` instead of `POST` to `/organization/update` and `/organization/member_update`, matching the methods the LiteLLM proxy serves; organization and organization member updates previously failed with a 405 +- **team_member**: Include `role` in the update payload so a role change on an existing `litellm_team_member` is applied instead of being silently dropped ### Changed - The provider source of truth moved to `terraform/provider/` in [BerriAI/litellm](https://github.com/BerriAI/litellm); this repository is now a release mirror. CI in the monorepo statically audits every endpoint the provider calls against the proxy's OpenAPI schema on every change +- **mcp_server**, **vector_store**: `env` and `litellm_params` are now marked sensitive, so they are redacted from plan/apply output, and they are no longer read back from the API into state — the configured value is authoritative. If the proxy returns values that differ from the configuration, that drift is no longer surfaced on refresh +- Dependency updates: `grpc` and `golang.org/x` modules + +## [0.3.0] - 2026-07-13 + +Released from the mirror repository before the source move was complete; this entry backfills it in the monorepo changelog. + +### Added + +- **model**: Add optional `pricing_base_model` attribute that sets `model_info.base_model` (the cost-map lookup key) independently of routing. Deployments whose routing name differs from the pricing key (for example Azure Data Zone, routed as `azure/gpt-4.1` but priced via `us/gpt-4.1-2025-04-14`) can now be billed correctly without breaking routing. When unset, behavior is unchanged and `base_model` continues to drive both routing and pricing (#47) ## [0.2.2] - 2026-05-13 diff --git a/terraform/provider/docs/resources/model.md b/terraform/provider/docs/resources/model.md index 5a46fe2f073..0409b48b391 100644 --- a/terraform/provider/docs/resources/model.md +++ b/terraform/provider/docs/resources/model.md @@ -118,6 +118,8 @@ The following arguments are supported: * `base_model` - (Required) string. The actual model identifier from the provider (e.g., "gpt-4", "claude-2"). +* `pricing_base_model` - (Optional) string. A pricing key fed to `model_info.base_model` **independently of routing**. When set, `litellm_params.model` still routes via `base_model`, but LiteLLM looks up cost against this key. Useful when the routing/deployment name differs from the cost-map key — e.g. an Azure deployment routed as `azure/gpt-4.1` whose real tier is Data Zone: set `pricing_base_model = "us/gpt-4.1-2025-04-14"` so it is billed at the Data Zone rate. When unset, `base_model` drives pricing as before. + * `litellm_credential_name` - (Optional) string. Name of a LiteLLM credential to use for this model. * `tier` - (Optional) string. The usage tier for this model. Valid values are `"free"` or `"paid"`. Default: `"free"`. diff --git a/terraform/provider/litellm/resource_model.go b/terraform/provider/litellm/resource_model.go index 2858b6e763d..4bad057871d 100644 --- a/terraform/provider/litellm/resource_model.go +++ b/terraform/provider/litellm/resource_model.go @@ -73,6 +73,14 @@ func resourceLiteLLMModel() *schema.Resource { Type: schema.TypeString, Required: true, }, + "pricing_base_model": { + // Optional pricing key fed to model_info.base_model, DECOUPLED + // from routing. When set, litellm_params.model still routes via + // base_model, but cost is looked up against this key (e.g. + // "us/gpt-4.1-2025-04-14" for Azure Data Zone pricing). + Type: schema.TypeString, + Optional: true, + }, "tier": { Type: schema.TypeString, Optional: true, diff --git a/terraform/provider/litellm/resource_model_crud.go b/terraform/provider/litellm/resource_model_crud.go index 40766c8e312..fc5d5b09dd5 100644 --- a/terraform/provider/litellm/resource_model_crud.go +++ b/terraform/provider/litellm/resource_model_crud.go @@ -68,6 +68,14 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e baseModel := d.Get("base_model").(string) modelName := fmt.Sprintf("%s/%s", customLLMProvider, baseModel) + // Pricing base_model, decoupled from routing. When pricing_base_model is + // set it feeds model_info.base_model (the cost-lookup key) WITHOUT changing + // the routing string above; otherwise base_model drives pricing as before. + pricingBaseModel := baseModel + if v, ok := d.GetOk("pricing_base_model"); ok && v.(string) != "" { + pricingBaseModel = v.(string) + } + // Generate a UUID for new models modelID := d.Id() if !isUpdate { @@ -240,7 +248,7 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e ModelInfo: ModelInfo{ ID: modelID, DBModel: true, - BaseModel: baseModel, + BaseModel: pricingBaseModel, Tier: d.Get("tier").(string), Mode: d.Get("mode").(string), TeamID: d.Get("team_id").(string), @@ -306,7 +314,16 @@ func resourceLiteLLMModelRead(d *schema.ResourceData, m interface{}) error { d.Set("rpm", GetIntValue(modelResp.LiteLLMParams.RPM, d.Get("rpm").(int))) d.Set("model_api_base", GetStringValue(modelResp.LiteLLMParams.APIBase, d.Get("model_api_base").(string))) d.Set("api_version", GetStringValue(modelResp.LiteLLMParams.APIVersion, d.Get("api_version").(string))) - d.Set("base_model", GetStringValue(modelResp.ModelInfo.BaseModel, d.Get("base_model").(string))) + // base_model / pricing_base_model read-back. When pricing_base_model is + // configured, model_info.base_model holds the PRICING key, so recover the + // routing base_model from state (not returned by the API) and read + // pricing_base_model from model_info. + if pbm, ok := d.GetOk("pricing_base_model"); ok && pbm.(string) != "" { + d.Set("base_model", d.Get("base_model").(string)) + d.Set("pricing_base_model", GetStringValue(modelResp.ModelInfo.BaseModel, pbm.(string))) + } else { + d.Set("base_model", GetStringValue(modelResp.ModelInfo.BaseModel, d.Get("base_model").(string))) + } d.Set("tier", GetStringValue(modelResp.ModelInfo.Tier, d.Get("tier").(string))) d.Set("mode", GetStringValue(modelResp.ModelInfo.Mode, d.Get("mode").(string))) d.Set("team_id", GetStringValue(modelResp.ModelInfo.TeamID, d.Get("team_id").(string))) From 845680ed1dc1e2f4b6c4493a00289e2f9422bbf0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 10:50:09 -0700 Subject: [PATCH 041/576] test(proxy): unit test batch file id mapping helpers directly --- .../test_common_utils.py | 97 +++++++++++++++++++ 1 file changed, 97 insertions(+) create mode 100644 tests/test_litellm/proxy/openai_files_endpoint/test_common_utils.py diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_common_utils.py new file mode 100644 index 00000000000..4a021627c3e --- /dev/null +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_common_utils.py @@ -0,0 +1,97 @@ +import os +import sys +from types import MappingProxyType +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy.openai_files_endpoints.common_utils import ( + apply_unified_file_ids, + map_raw_file_ids_to_unified, +) +from litellm.types.utils import LiteLLMBatch + + +def _batch(input_file_id, output_file_id, error_file_id) -> LiteLLMBatch: + return LiteLLMBatch( + id="batch-1", + completion_window="24h", + created_at=1234567890, + endpoint="/v1/chat/completions", + input_file_id=input_file_id, + object="batch", + status="cancelled", + output_file_id=output_file_id, + error_file_id=error_file_id, + ) + + +@pytest.mark.asyncio +async def test_map_raw_file_ids_to_unified_empty_ids_skips_db(): + prisma_client = MagicMock() + + assert await map_raw_file_ids_to_unified(frozenset(), prisma_client) == {} + + prisma_client.db.litellm_managedfiletable.find_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_map_raw_file_ids_to_unified_no_prisma_client_returns_empty(): + assert await map_raw_file_ids_to_unified(frozenset({"file-raw-1"}), None) == {} + + +@pytest.mark.asyncio +async def test_map_raw_file_ids_to_unified_bulk_queries_and_filters_to_requested_ids(): + row_a = MagicMock( + unified_file_id="unified-a", + flat_model_file_ids=["file-raw-a", "file-raw-other"], + ) + row_b = MagicMock(unified_file_id="unified-b", flat_model_file_ids=["file-raw-b"]) + prisma_client = MagicMock() + prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(return_value=[row_a, row_b]) + + mapping = await map_raw_file_ids_to_unified( + frozenset({"file-raw-b", "file-raw-a", "file-raw-missing"}), prisma_client + ) + + prisma_client.db.litellm_managedfiletable.find_many.assert_awaited_once_with( + where={"flat_model_file_ids": {"hasSome": ["file-raw-a", "file-raw-b", "file-raw-missing"]}} + ) + assert dict(mapping) == {"file-raw-a": "unified-a", "file-raw-b": "unified-b"} + + +def test_apply_unified_file_ids_swaps_only_mapped_ids(): + batch = _batch(input_file_id="file-raw-in", output_file_id="file-raw-out", error_file_id=None) + + apply_unified_file_ids(batch, MappingProxyType({"file-raw-out": "unified-out"})) + + assert batch.input_file_id == "file-raw-in" + assert batch.output_file_id == "unified-out" + assert batch.error_file_id is None + + +def test_apply_unified_file_ids_swaps_all_three_ids(): + batch = _batch( + input_file_id="file-raw-in", + output_file_id="file-raw-out", + error_file_id="file-raw-err", + ) + + apply_unified_file_ids( + batch, + MappingProxyType( + { + "file-raw-in": "unified-in", + "file-raw-out": "unified-out", + "file-raw-err": "unified-err", + } + ), + ) + + assert (batch.input_file_id, batch.output_file_id, batch.error_file_id) == ( + "unified-in", + "unified-out", + "unified-err", + ) From 495eb7e7f428a64ebfb9b57004026dc7739dcbc1 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 6 Aug 2026 11:26:32 -0700 Subject: [PATCH 042/576] test(router): assert the auto-router max_input_chars kwarg PR #35956 added the max_input_chars passthrough to the AutoRouter constructor but left this mock assertion in tests/router_unit_tests unchanged, so test_init_auto_router_deployment_success has been failing on litellm_internal_staging ever since. The passthrough itself is intentional and its behaviour is already covered by TestAutoRouterMaxInputCharsWiring in tests/test_litellm, so only the stale expected kwargs need updating. Assert the shared constant rather than the literal 2000 so tuning the default does not break this test again. --- tests/router_unit_tests/test_router_helper_utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index bcc70fae67c..0655763d41b 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -15,6 +15,7 @@ from unittest.mock import patch, MagicMock, AsyncMock from create_mock_standard_logging_payload import create_standard_logging_payload from litellm.types.utils import StandardLoggingPayload from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo +from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS @pytest.fixture @@ -1816,6 +1817,7 @@ def test_init_auto_router_deployment_success(mock_auto_router, model_list): default_model="gpt-5-mini", embedding_model="text-embedding-3-small", litellm_router_instance=router, + max_input_chars=DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ) # Verify the auto-router was added to the router's auto_routers dict From f5d98c0b8ce15164f25b880258fb88c24f03baeb Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 6 Aug 2026 14:25:11 -0700 Subject: [PATCH 043/576] feat(ui): migrate playground chat controls toward shadcn Continue the Playground Chat Ant Design/Tremor migration: shared MultiSelect, upload validation with semantic file inputs, collapsible message widgets, and AdditionalModelSettings on Base UI controls --- .../components/chat_ui/A2AMetrics.tsx | 237 +++++++------ .../chat_ui/AdditionalModelSettings.tsx | 224 +++++++----- .../components/chat_ui/ChatImageUpload.tsx | 87 +++-- .../components/chat_ui/ChatMessageBubble.tsx | 8 +- .../playground/components/chat_ui/ChatUI.tsx | 251 +++++++------- .../chat_ui/CodeInterpreterOutput.tsx | 215 ++++++------ .../chat_ui/CodeInterpreterTool.tsx | 27 +- .../components/chat_ui/EndpointSelector.tsx | 14 +- .../components/chat_ui/FilePreviewCard.tsx | 17 +- .../chat_ui/ResponsesImageUpload.tsx | 83 +++-- .../chat_ui/SearchResultsDisplay.tsx | 168 ++++----- .../components/chat_ui/SessionManagement.tsx | 73 ++-- .../chat_ui/uploadValidation.test.ts | 78 +++++ .../components/chat_ui/uploadValidation.ts | 97 ++++++ .../src/app/(dashboard)/playground/page.tsx | 10 +- .../components/chat_ui/MCPEventsDisplay.tsx | 323 ++++++++---------- .../components/chat_ui/ReasoningContent.tsx | 117 +++---- .../components/chat_ui/ResponseMetrics.tsx | 131 +++---- .../guardrails/GuardrailSelector.tsx | 13 +- .../src/components/llm_calls/fetch_models.tsx | 20 +- .../components/policies/PolicySelector.tsx | 13 +- .../src/components/shared/MultiSelect.tsx | 119 +++++++ .../src/components/shared/SearchSelect.tsx | 4 +- .../components/tag_management/TagSelector.tsx | 17 +- .../src/components/ui/combobox.tsx | 7 +- .../VectorStoreSelector.tsx | 17 +- 26 files changed, 1414 insertions(+), 956 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/uploadValidation.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/uploadValidation.ts create mode 100644 ui/litellm-dashboard/src/components/shared/MultiSelect.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/A2AMetrics.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/A2AMetrics.tsx index 004a513f061..6ddfe1442f5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/A2AMetrics.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/A2AMetrics.tsx @@ -1,17 +1,19 @@ import React, { useState } from "react"; -import { Tooltip, Button } from "antd"; import { - CheckCircleOutlined, - ClockCircleOutlined, - LoadingOutlined, - ExclamationCircleOutlined, - CopyOutlined, - DownOutlined, - RightOutlined, - LinkOutlined, - FileTextOutlined, - RobotOutlined, -} from "@ant-design/icons"; + Bot, + CheckCircle, + ChevronDown, + ChevronRight, + CircleAlert, + Clock, + Copy, + FileText, + Link, + LoaderCircle, +} from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; export interface A2ATaskMetadata { taskId?: string; @@ -21,7 +23,7 @@ export interface A2ATaskMetadata { timestamp?: string; message?: string; }; - metadata?: Record; + metadata?: Record; } interface A2AMetricsProps { @@ -33,15 +35,15 @@ interface A2AMetricsProps { const getStatusIcon = (state?: string) => { switch (state) { case "completed": - return ; + return ; case "working": case "submitted": - return ; + return ; case "failed": case "canceled": - return ; + return ; default: - return ; + return ; } }; @@ -91,7 +93,7 @@ const A2AMetrics: React.FC = ({ a2aMetadata, timeToFirstToken,
{/* A2A Metadata Header */}
- + A2A Metadata
@@ -109,28 +111,33 @@ const A2AMetrics: React.FC = ({ a2aMetadata, timeToFirstToken, {/* Timestamp */} {formattedTime && ( - - - + + }> + {formattedTime} - + + {status?.timestamp} )} {/* Latency */} {totalLatency !== undefined && ( - - - + + }> + {(totalLatency / 1000).toFixed(2)}s - + + Total latency )} {/* Time to first token */} {timeToFirstToken !== undefined && ( - - TTFT: {(timeToFirstToken / 1000).toFixed(2)}s + + }> + TTFT: {(timeToFirstToken / 1000).toFixed(2)}s + + Time to first token )}
@@ -139,95 +146,133 @@ const A2AMetrics: React.FC = ({ a2aMetadata, timeToFirstToken,
{/* Task ID */} {taskId && ( - - copyToClipboard(taskId)} + + copyToClipboard(taskId)} + aria-label={`Copy task ID ${taskId}`} + /> + } > - + Task: {truncateId(taskId)} - - + + + Click to copy: {taskId} )} {/* Context/Session ID */} {contextId && ( - - copyToClipboard(contextId)} + + copyToClipboard(contextId)} + aria-label={`Copy session ID ${contextId}`} + /> + } > - + Session: {truncateId(contextId)} - - + + + Click to copy: {contextId} )} {/* Details toggle */} {(metadata || status?.message) && ( - + + + } + > + {showDetails ? : } + Details + + )}
{/* Expandable details panel */} - {showDetails && ( -
- {/* Status message */} - {status?.message && ( -
- Status Message: - {status.message} -
- )} + + +
+ {/* Status message */} + {status?.message && ( +
+ Status Message: + {status.message} +
+ )} - {/* Full IDs */} - {taskId && ( -
- Task ID: - - {taskId} - - copyToClipboard(taskId)} - /> -
- )} + {/* Full IDs */} + {taskId && ( +
+ Task ID: + + {taskId} + + +
+ )} - {contextId && ( -
- Session ID: - - {contextId} - - copyToClipboard(contextId)} - /> -
- )} + {contextId && ( +
+ Session ID: + + {contextId} + + +
+ )} - {/* Metadata fields */} - {metadata && Object.keys(metadata).length > 0 && ( -
- Custom Metadata: - 
-                {JSON.stringify(metadata, null, 2)}
-
-
- )} -
- )} + {/* Metadata fields */} + {metadata && Object.keys(metadata).length > 0 && ( +
+ Custom Metadata: + 
+                  {JSON.stringify(metadata, null, 2)}
+
+
+ )} +
+ + ); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx index d4320110c4c..4deb7051954 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx @@ -1,7 +1,10 @@ -import { InfoCircleOutlined } from "@ant-design/icons"; -import { Text } from "@tremor/react"; -import { Checkbox, InputNumber, Popover, Slider, Tooltip, Typography } from "antd"; -import React, { useEffect, useState } from "react"; +import { Info } from "lucide-react"; +import React, { useEffect, useId, useState } from "react"; +import { Checkbox } from "@/components/ui/checkbox"; +import { Input } from "@/components/ui/input"; +import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { cn } from "@/lib/cva.config"; interface AdditionalModelSettingsProps { temperature?: number; @@ -17,6 +20,10 @@ interface AdditionalModelSettingsProps { showAdvancedParams?: boolean; } +function clamp(value: number, min: number, max: number): number { + return Math.min(max, Math.max(min, value)); +} + const AdditionalModelSettings: React.FC = ({ temperature = 1.0, maxTokens = 2048, @@ -36,7 +43,12 @@ const AdditionalModelSettings: React.FC = ({ const [localTemperature, setLocalTemperature] = useState(temperature); const [localMaxTokens, setLocalMaxTokens] = useState(maxTokens); - // Sync local state with props when they change + const streamingId = useId(); + const advancedId = useId(); + const fallbacksId = useId(); + const temperatureId = useId(); + const maxTokensId = useId(); + useEffect(() => { setLocalTemperature(temperature); }, [temperature]); @@ -45,21 +57,18 @@ const AdditionalModelSettings: React.FC = ({ setLocalMaxTokens(maxTokens); }, [maxTokens]); - const handleTemperatureChange = (value: number | null) => { - const newValue = value ?? 1.0; + const handleTemperatureChange = (value: number) => { + const newValue = clamp(Number.isFinite(value) ? value : 1.0, 0, 2); setLocalTemperature(newValue); onTemperatureChange?.(newValue); }; - const handleMaxTokensChange = (value: number | null) => { - const newValue = value ?? 1000; + const handleMaxTokensChange = (value: number) => { + const newValue = clamp(Number.isFinite(value) ? Math.round(value) : 1000, 1, 32768); setLocalMaxTokens(newValue); onMaxTokensChange?.(newValue); }; - const disabledOpacity = useAdvancedParams ? 1 : 0.4; - const disabledTextColor = useAdvancedParams ? "text-gray-700" : "text-gray-400"; - const handleUseAdvancedParamsChange = (checked: boolean) => { if (onUseAdvancedParamsChange) { onUseAdvancedParamsChange(checked); @@ -68,129 +77,176 @@ const AdditionalModelSettings: React.FC = ({ } }; + const disabledTextColor = useAdvancedParams ? "text-gray-700" : "text-gray-400"; + return ( -
+
{onStreamingChange && ( -
- onStreamingChange(e.target.checked)}> - Stream responses - - - +
+ onStreamingChange(checked === true)} + aria-label="Stream responses" + /> + + + + + + + Streams the answer token by token. Uncheck to send a non-streaming request and render the full response at + once. +
)} {showAdvancedParams && ( - handleUseAdvancedParamsChange(e.target.checked)}> - Use Advanced Parameters - +
+ handleUseAdvancedParamsChange(checked === true)} + aria-label="Use Advanced Parameters" + /> + +
)} {onMockTestFallbacksChange && ( -
- onMockTestFallbacksChange(e.target.checked)}> - Simulate failure to test fallbacks - - - - Causes the first request to fail so the router tries fallbacks (if configured). Use this to verify - your fallback setup. - - - Behavior can differ when keys, teams, or router settings are configured.{" "} - - Learn more - - -
- } - > - +
+ onMockTestFallbacksChange(checked === true)} + aria-label="Simulate failure to test fallbacks" + /> + + + + + + +

+ Causes the first request to fail so the router tries fallbacks (if configured). Use this to verify your + fallback setup. +

+

+ Behavior can differ when keys, teams, or router settings are configured.{" "} + + Learn more + +

+
)} {showAdvancedParams && ( -
+
-
+
- Temperature - - + + + + + + + Controls randomness. Lower values make output more deterministic, higher values more creative. +
- handleTemperatureChange(Number(event.target.value))} />
- handleTemperatureChange(Number(event.target.value))} /> +
+ 0 + 1.0 + 2.0 +
-
+
- Max Tokens - - + + + + + + + Maximum number of tokens to generate in the response. +
- handleMaxTokensChange(Number(event.target.value))} />
- handleMaxTokensChange(Number(event.target.value))} /> +
+ 1 + 32768 +
)} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUpload.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUpload.tsx index 55527d997ac..6f210118281 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUpload.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUpload.tsx @@ -1,43 +1,70 @@ -import React from "react"; -import { Upload, Tooltip } from "antd"; -import { PaperClipOutlined } from "@ant-design/icons"; - -const { Dragger } = Upload; +import React, { useId, useRef } from "react"; +import { Paperclip } from "lucide-react"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { Button } from "@/components/ui/button"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { CHAT_ATTACHMENT_ACCEPT, validateChatAttachment } from "./uploadValidation"; interface ChatImageUploadProps { chatUploadedImage: File | null; chatImagePreviewUrl: string | null; - onImageUpload: (file: File) => false; + onImageUpload: (file: File) => void; onRemoveImage: () => void; + disabled?: boolean; } -const ChatImageUpload: React.FC = ({ - chatUploadedImage, - chatImagePreviewUrl, - onImageUpload, - onRemoveImage, -}) => { +const ChatImageUpload: React.FC = ({ chatUploadedImage, onImageUpload, disabled = false }) => { + const inputRef = useRef(null); + const inputId = useId(); + + if (chatUploadedImage) { + return null; + } + + const handleFileChange = (event: React.ChangeEvent) => { + const file = event.target.files?.[0]; + event.target.value = ""; + if (!file) { + return; + } + const result = validateChatAttachment(file); + if (!result.ok) { + NotificationsManager.error(result.error); + return; + } + onImageUpload(file); + }; + return ( <> - {/* Subtle upload button - only show when no image */} - {!chatUploadedImage && ( - - - - - - )} + variant="ghost" + size="icon-sm" + disabled={disabled} + aria-label="Attach image or PDF" + className="text-gray-400 hover:text-gray-600" + onClick={() => inputRef.current?.click()} + /> + } + > + + + Attach image or PDF + ); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx index 8e71017a7b5..c438b4982bf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx @@ -41,9 +41,9 @@ function ChatMessageBubble({ const isUser = message.role === "user"; return ( -
+
{/* Header: role icon + name + model badge */} -
+
{message.role} {message.role === "assistant" && message.model && ( - + {message.model} )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index 57ff7906eda..5edcbe84aa8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -66,7 +66,16 @@ import { MessageType } from "@/components/chat_ui/types"; import { useCodeInterpreter } from "../../hooks/useCodeInterpreter"; import { useChatHistory } from "../../hooks/useChatHistory"; import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; +import { SearchSelect } from "@/components/shared/SearchSelect"; +import { Select as ShadcnSelect, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; +import { + AUDIO_ACCEPT, + IMAGE_EDIT_ACCEPT, + validateAudioFile, + validateChatAttachment, + validateImageEditFile, +} from "./uploadValidation"; const { TextArea } = Input; const { Dragger } = Upload; @@ -177,6 +186,8 @@ const ChatUI: React.FC = ({ const [selectedModel, setSelectedModel] = useState(simplified ? fixedModel : undefined); const [showCustomModelInput, setShowCustomModelInput] = useState(false); const [modelInfo, setModelInfo] = useState([]); + const [isLoadingModels, setIsLoadingModels] = useState(false); + const [modelLoadError, setModelLoadError] = useState(false); const [agentInfo, setAgentInfo] = useState([]); const [selectedAgent, setSelectedAgent] = useState(undefined); const debouncedSetSelectedModel = useDebouncedCallback((value: string) => setSelectedModel(value), { @@ -388,17 +399,17 @@ const ChatUI: React.FC = ({ ]); useEffect(() => { - let userApiKey = apiKeySource === "session" ? accessToken : apiKey; - if (!userApiKey || !token || !userRole || !userID) { + const userApiKey = apiKeySource === "session" ? accessToken : apiKey.trim(); + if (!userApiKey) { + setModelInfo([]); + setModelLoadError(false); return; } - // Fetch model info and set the default selected model (skip in simplified mode; we use fixedModel) const loadModels = async () => { + setIsLoadingModels(true); + setModelLoadError(false); try { - if (!userApiKey) { - return; - } const uniqueModels = await fetchAvailableModels(userApiKey); setModelInfo(uniqueModels); @@ -412,6 +423,10 @@ const ChatUI: React.FC = ({ } } catch (error) { console.error("Error fetching model info:", error); + setModelInfo([]); + setModelLoadError(true); + } finally { + setIsLoadingModels(false); } }; @@ -419,7 +434,7 @@ const ChatUI: React.FC = ({ loadModels(); } loadMCPServers(); - }, [accessToken, userID, userRole, apiKeySource, apiKey, token, simplified]); + }, [accessToken, apiKeySource, apiKey, simplified]); // Load tools when MCP direct mode has a server (or toolset) selected useEffect(() => { @@ -494,13 +509,35 @@ const ChatUI: React.FC = ({ } }; - const handleImageUpload = (file: File) => { - setUploadedImages((prev) => [...prev, file]); + const createBlobPreviewUrl = (file: File): string => { const rawPreviewUrl = URL.createObjectURL(file); - // Sanitize: only allow blob: URLs to prevent XSS via img src injection. - const previewUrl = rawPreviewUrl.startsWith("blob:") ? rawPreviewUrl : ""; - setImagePreviewUrls((prev) => [...prev, previewUrl]); - return false; // Prevent default upload behavior + return rawPreviewUrl.startsWith("blob:") ? rawPreviewUrl : ""; + }; + + const handleImageFiles = (files: File[]) => { + let nextCount = uploadedImages.length; + const accepted: File[] = []; + const previews: string[] = []; + for (const file of files) { + const result = validateImageEditFile(file, nextCount); + if (!result.ok) { + NotificationsManager.error(result.error); + continue; + } + accepted.push(file); + previews.push(createBlobPreviewUrl(file)); + nextCount += 1; + } + if (accepted.length === 0) { + return; + } + setUploadedImages((prev) => [...prev, ...accepted]); + setImagePreviewUrls((prev) => [...prev, ...previews]); + }; + + const handleImageUpload = (file: File): false => { + handleImageFiles([file]); + return false; }; const handleRemoveImage = (index: number) => { @@ -519,11 +556,14 @@ const ChatUI: React.FC = ({ setImagePreviewUrls([]); }; - const handleResponsesImageUpload = (file: File): false => { + const handleResponsesImageUpload = (file: File): void => { + const result = validateChatAttachment(file); + if (!result.ok) { + NotificationsManager.error(result.error); + return; + } setResponsesUploadedImage(file); - const previewUrl = URL.createObjectURL(file); - setResponsesImagePreviewUrl(previewUrl); - return false; // Prevent default upload behavior + setResponsesImagePreviewUrl(createBlobPreviewUrl(file)); }; const handleRemoveResponsesImage = () => { @@ -534,11 +574,14 @@ const ChatUI: React.FC = ({ setResponsesImagePreviewUrl(null); }; - const handleChatImageUpload = (file: File): false => { + const handleChatImageUpload = (file: File): void => { + const result = validateChatAttachment(file); + if (!result.ok) { + NotificationsManager.error(result.error); + return; + } setChatUploadedImage(file); - const previewUrl = URL.createObjectURL(file); - setChatImagePreviewUrl(previewUrl); - return false; // Prevent default upload behavior + setChatImagePreviewUrl(createBlobPreviewUrl(file)); }; const handleRemoveChatImage = () => { @@ -550,8 +593,13 @@ const ChatUI: React.FC = ({ }; const handleAudioUpload = (file: File): false => { + const result = validateAudioFile(file); + if (!result.ok) { + NotificationsManager.error(result.error); + return false; + } setUploadedAudio(file); - return false; // Prevent default upload behavior + return false; }; const handleRemoveAudio = () => { @@ -1002,8 +1050,12 @@ const ChatUI: React.FC = ({ const onModelChange = (value: string) => { setSelectedModel(value); - setShowCustomModelInput(value === "custom"); + + const model = modelInfo.find((option) => option.model_group === value); + if (model?.mode) { + setEndpointType(getEndpointType(model.mode)); + } }; // Check if the selected model is a chat model @@ -1020,35 +1072,43 @@ const ChatUI: React.FC = ({ }; const supportsStreamingToggle = endpointType === EndpointType.CHAT || endpointType === EndpointType.RESPONSES; + let modelEmptyText = "No models available for this key"; + if (modelLoadError) { + modelEmptyText = "Unable to load models for this key"; + } else if (apiKeySource === "custom" && !apiKey.trim()) { + modelEmptyText = "Enter a Virtual Key to load models"; + } const antIcon = ; return ( -
- -
+
+ +
{/* Left Sidebar with Controls - hidden in simplified mode */} {!simplified && ( -
+
Configurations
Virtual Key Source - { + onValueChange={(value) => { setSelectedVoice(value); sessionStorage.setItem("selectedVoice", value); }} - style={{ width: "100%" }} - className="rounded-md" - options={OPEN_AI_VOICE_SELECT_OPTIONS} - /> + > + + + + + {OPEN_AI_VOICE_SELECT_OPTIONS.map((voice) => ( + + {voice.label} + + ))} + +
)} @@ -1212,46 +1280,20 @@ const ChatUI: React.FC = ({ )} - setSelectedAgent(value)} + onValueChange={(value) => setSelectedAgent(value)} options={agentInfo.map((agent) => ({ value: agent.agent_name, label: agent.agent_name || agent.agent_id, - key: agent.agent_id, + sublabel: agent.agent_card_params?.description, }))} - style={{ width: "100%" }} - showSearch={true} - className="rounded-md" - optionLabelProp="label" - > - {agentInfo.map((agent) => ( - -
- {agent.agent_name || agent.agent_id} - {agent.agent_card_params?.description && ( - {agent.agent_card_params.description} - )} -
- - ))} - + /> {agentInfo.length === 0 && ( No agents found. Create agents via /v1/agents endpoint. @@ -1697,7 +1720,7 @@ const ChatUI: React.FC = ({ )} {/* Main Chat Area */} -
+
{endpointType === EndpointType.REALTIME ? ( = ({ /> ) : ( <> -
+
{simplified ? "Chat" : "Test Key"} -
+
= ({ )}
-
+
{chatHistory.length === 0 && (
@@ -1788,18 +1811,18 @@ const ChatUI: React.FC = ({
-
+
{/* Image Upload Section for Image Edits */} {endpointType === EndpointType.IMAGE_EDITS && (
{uploadedImages.length === 0 ? ( - +

Click or drag images to upload

- Support for PNG, JPG, JPEG formats. Multiple images supported. + Support for PNG, JPG, JPEG, GIF, WebP. Multiple images supported.

) : ( @@ -1840,12 +1863,12 @@ const ChatUI: React.FC = ({ { - const files = Array.from(e.target.files || []); - files.forEach((file) => handleImageUpload(file)); + handleImageFiles(Array.from(e.target.files || [])); + e.target.value = ""; }} />
@@ -1858,11 +1881,7 @@ const ChatUI: React.FC = ({ {endpointType === EndpointType.TRANSCRIPTION && (
{!uploadedAudio ? ( - +

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx index c27273a116d..8dc503f369c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx @@ -1,15 +1,10 @@ -import React, { useState, useEffect } from "react"; -import { Collapse, Spin } from "antd"; -import { - CodeOutlined, - DownloadOutlined, - FileImageOutlined, - FileTextOutlined, - LoadingOutlined, -} from "@ant-design/icons"; +import React, { useEffect, useState } from "react"; +import { Code, Download, FileImage, FileText, Loader2 } from "lucide-react"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; +import { Button } from "@/components/ui/button"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; interface ContainerFileCitation { type: "container_file_citation"; @@ -27,48 +22,60 @@ interface CodeInterpreterOutputProps { accessToken: string; } -const CodeInterpreterOutput: React.FC = ({ - code, - containerId, - annotations = [], - accessToken, -}) => { +const IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".gif"] as const; + +function isImageFilename(filename: string | undefined): boolean { + if (!filename) { + return false; + } + const lower = filename.toLowerCase(); + return IMAGE_EXTENSIONS.some((ext) => lower.endsWith(ext)); +} + +const CodeInterpreterOutput: React.FC = ({ code, annotations = [], accessToken }) => { const [imageUrls, setImageUrls] = useState>({}); const [loadingImages, setLoadingImages] = useState>({}); + const [codeOpen, setCodeOpen] = useState(false); const proxyBaseUrl = getProxyBaseUrl(); - // Fetch images from container files API useEffect(() => { + const createdUrls: string[] = []; + let cancelled = false; + const fetchImages = async () => { for (const annotation of annotations) { - const isImage = - annotation.filename?.toLowerCase().endsWith(".png") || - annotation.filename?.toLowerCase().endsWith(".jpg") || - annotation.filename?.toLowerCase().endsWith(".jpeg") || - annotation.filename?.toLowerCase().endsWith(".gif"); + if (!isImageFilename(annotation.filename) || !annotation.container_id || !annotation.file_id) { + continue; + } - if (isImage && annotation.container_id && annotation.file_id) { + if (!cancelled) { setLoadingImages((prev) => ({ ...prev, [annotation.file_id]: true })); + } - try { - // Fetch image content from container files API - const response = await fetch( - `${proxyBaseUrl}/v1/containers/${annotation.container_id}/files/${annotation.file_id}/content`, - { - headers: { - [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, - }, + try { + const response = await fetch( + `${proxyBaseUrl}/v1/containers/${annotation.container_id}/files/${annotation.file_id}/content`, + { + headers: { + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, }, - ); + }, + ); - if (response.ok) { - const blob = await response.blob(); - const url = URL.createObjectURL(blob); + if (response.ok) { + const blob = await response.blob(); + const url = URL.createObjectURL(blob); + createdUrls.push(url); + if (!cancelled) { setImageUrls((prev) => ({ ...prev, [annotation.file_id]: url })); + } else { + URL.revokeObjectURL(url); } - } catch (error) { - console.error("Error fetching image:", error); - } finally { + } + } catch (error) { + console.error("Error fetching image:", error); + } finally { + if (!cancelled) { setLoadingImages((prev) => ({ ...prev, [annotation.file_id]: false })); } } @@ -76,12 +83,12 @@ const CodeInterpreterOutput: React.FC = ({ }; if (annotations.length > 0 && accessToken) { - fetchImages(); + void fetchImages(); } - // Cleanup URLs on unmount return () => { - Object.values(imageUrls).forEach((url) => URL.revokeObjectURL(url)); + cancelled = true; + createdUrls.forEach((url) => URL.revokeObjectURL(url)); }; }, [annotations, accessToken, proxyBaseUrl]); @@ -112,22 +119,8 @@ const CodeInterpreterOutput: React.FC = ({ } }; - // Separate images and other files - const imageAnnotations = annotations.filter( - (a) => - a.filename?.toLowerCase().endsWith(".png") || - a.filename?.toLowerCase().endsWith(".jpg") || - a.filename?.toLowerCase().endsWith(".jpeg") || - a.filename?.toLowerCase().endsWith(".gif"), - ); - - const fileAnnotations = annotations.filter( - (a) => - !a.filename?.toLowerCase().endsWith(".png") && - !a.filename?.toLowerCase().endsWith(".jpg") && - !a.filename?.toLowerCase().endsWith(".jpeg") && - !a.filename?.toLowerCase().endsWith(".gif"), - ); + const imageAnnotations = annotations.filter((a) => isImageFilename(a.filename)); + const fileAnnotations = annotations.filter((a) => !isImageFilename(a.filename)); if (!code && annotations.length === 0) { return null; @@ -135,44 +128,46 @@ const CodeInterpreterOutput: React.FC = ({ return (
- {/* Executed Code - Collapsible */} {code && ( - - Python Code Executed - - ), - children: ( - - {code} - - ), - }, - ]} - /> + + + } + > + + Python Code Executed + + +
+ + {code} + +
+ + )} - {/* Generated Images */} {imageAnnotations.map((annotation) => ( -
+
{loadingImages[annotation.file_id] ? ( -
- } /> +
+
) : imageUrls[annotation.file_id] ? ( @@ -180,42 +175,48 @@ const CodeInterpreterOutput: React.FC = ({ {annotation.filename -
- - {annotation.filename} +
+ + - + + Download +
) : ( -
+
Image not available
)}
))} - {/* Download Links for Other Files */} {fileAnnotations.length > 0 && (
{fileAnnotations.map((annotation) => ( - +
)} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterTool.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterTool.tsx index d2682e3a7a8..e5744ac8e38 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterTool.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterTool.tsx @@ -1,8 +1,8 @@ import React from "react"; -import { Switch, Tooltip } from "antd"; import MessageManager from "@/components/molecules/message_manager"; -import { CodeOutlined, InfoCircleOutlined, ExclamationCircleOutlined } from "@ant-design/icons"; -import { Text } from "@tremor/react"; +import { Code, Info, TriangleAlert } from "lucide-react"; +import { Switch } from "@/components/ui/switch"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; interface CodeInterpreterToolProps { accessToken: string; @@ -49,25 +49,30 @@ const CodeInterpreterTool: React.FC = ({
- - Code Interpreter - - + + Code Interpreter + + + + + + Run Python code to generate files, charts, and analyze data. Container is created automatically. +
{!isOpenAI && (
- +
Code Interpreter is currently only supported for OpenAI models. = ({ endpointType, onEndpointChange, className }) => { return (
- + { return { label: `${guardrail.guardrail_name}`, value: guardrail.guardrail_name, }; })} - optionFilterProp="label" - showSearch - style={{ width: "100%" }} />
); diff --git a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx index 0de98330c2e..24f1e038f85 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx @@ -7,6 +7,13 @@ export interface ModelGroup { mode?: string; } +interface AvailableModel { + model_group?: string | null; + model_name?: string | null; + id?: string | null; + mode?: string | null; +} + /** * Fetches available models using modelHubCall and formats them for the selection dropdown. */ @@ -15,14 +22,15 @@ export const fetchAvailableModels = async (accessToken: string): Promise 0) { - const models: ModelGroup[] = fetchedModels.data.map((item: any) => ({ - model_group: item.model_group, // Display the model_group to the user - mode: item?.mode, // Save the mode for auto-selection of endpoint type - })); + const models: ModelGroup[] = fetchedModels.data + .map((item: AvailableModel) => ({ + model_group: item.model_group || item.id || item.model_name || "", + mode: item.mode || undefined, + })) + .filter((model: ModelGroup) => model.model_group !== ""); - // Sort models alphabetically by label models.sort((a, b) => a.model_group.localeCompare(b.model_group)); - return models; + return Array.from(new Map(models.map((model) => [model.model_group, model])).values()); } return []; } catch (error) { diff --git a/ui/litellm-dashboard/src/components/policies/PolicySelector.tsx b/ui/litellm-dashboard/src/components/policies/PolicySelector.tsx index 132538d439f..1e565d938d7 100644 --- a/ui/litellm-dashboard/src/components/policies/PolicySelector.tsx +++ b/ui/litellm-dashboard/src/components/policies/PolicySelector.tsx @@ -1,7 +1,7 @@ import React, { useEffect, useState } from "react"; -import { Select } from "antd"; import { Policy } from "./types"; import { getPoliciesList } from "../networking"; +import { MultiSelect } from "@/components/shared/MultiSelect"; /** Prefix for policy version IDs in request body; must match backend POLICY_VERSION_ID_PREFIX. */ export const POLICY_VERSION_ID_PREFIX = "policy_"; @@ -80,22 +80,17 @@ const PolicySelector: React.FC = ({ }; return ( -
- ({ label: tag.name, value: tag.name, - title: tag.description || tag.name, + description: tag.description || undefined, }))} - optionFilterProp="label" - tokenSeparators={[","]} - maxTagCount="responsive" - allowClear - style={{ width: "100%" }} /> ); }; diff --git a/ui/litellm-dashboard/src/components/ui/combobox.tsx b/ui/litellm-dashboard/src/components/ui/combobox.tsx index 2854928140e..541ad8bb25c 100644 --- a/ui/litellm-dashboard/src/components/ui/combobox.tsx +++ b/ui/litellm-dashboard/src/components/ui/combobox.tsx @@ -83,10 +83,14 @@ function ComboboxContent({ sideOffset = 6, align = "start", alignOffset = 0, + collisionAvoidance, anchor, ...props }: ComboboxPrimitive.Popup.Props & - Pick) { + Pick< + ComboboxPrimitive.Positioner.Props, + "side" | "align" | "sideOffset" | "alignOffset" | "collisionAvoidance" | "anchor" + >) { return ( diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.tsx index 2642b74e492..d80996f6a28 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.tsx @@ -1,7 +1,7 @@ import React, { useEffect, useState } from "react"; -import { Select } from "antd"; import { VectorStore } from "./types"; import { vectorStoreListCall } from "../networking"; +import { MultiSelect } from "@/components/shared/MultiSelect"; interface VectorStoreSelectorProps { onChange: (selectedVectorStores: string[]) => void; value?: string[]; @@ -43,24 +43,19 @@ const VectorStoreSelector: React.FC = ({ }, [accessToken]); return ( -
- setApiKey(event.target.value)} + value={apiKey} + /> +
)}
-
- - Custom Proxy Base URL - +
+ {proxySettings?.LITELLM_UI_API_DOC_BASE_URL && !customProxyBaseUrl && ( )} {customProxyBaseUrl && ( )}
- { - setCustomProxyBaseUrl(value); - sessionStorage.setItem("customProxyBaseUrl", value); - }} - value={customProxyBaseUrl} - icon={ApiOutlined} - /> +
+ + { + setCustomProxyBaseUrl(event.target.value); + sessionStorage.setItem("customProxyBaseUrl", event.target.value); + }} + /> +
{customProxyBaseUrl && ( - API calls will be sent to: {customProxyBaseUrl} +

API calls will be sent to: {customProxyBaseUrl}

)}
- - Endpoint Type - + { setEndpointType(value); - // Clear model/agent selection when switching endpoint type setSelectedModel(undefined); setSelectedAgent(undefined); setShowCustomModelInput(false); setSelectedMCPDirectTool(undefined); - // For MCP direct mode, require single server (clear __all__ or multiple) if (value === EndpointType.MCP) { setSelectedMCPServers((prev) => (prev.length === 1 && prev[0] !== "__all__" ? prev : [])); } @@ -1194,13 +1279,12 @@ const ChatUI: React.FC = ({ className="mb-4" /> - {/* Voice Selector for Speech Endpoint */} {endpointType === EndpointType.SPEECH && (
- - + + { @@ -1222,7 +1306,6 @@ const ChatUI: React.FC = ({
)} - {/* Session Management Component */} = ({ />
- {/* Model Selector - shown when NOT using A2A Agents or MCP direct mode */} {endpointType !== EndpointType.A2A_AGENTS && endpointType !== EndpointType.MCP && (
- +
- Select Model + {isChatModel() || supportsStreamingToggle ? ( - + + } + > + + + +
Model Settings
= ({ streamingEnabled={streamingEnabled} onStreamingChange={supportsStreamingToggle ? setStreamingEnabled : undefined} /> - } - title="Model Settings" - trigger="click" - placement="right" - > -
= ({ ]} /> {showCustomModelInput && ( - debouncedSetSelectedModel(event.target.value)} /> )}
)} - {/* Agent Selector - shown ONLY for A2A Agents endpoint */} {endpointType === EndpointType.A2A_AGENTS && (
- - Select Agent - + = ({ }))} /> {agentInfo.length === 0 && ( - +

No agents found. Create agents via /v1/agents endpoint. - +

)}
)}
- - Tags - + = ({ />
- {/* MCP Server Selection */}
- - +
+
)} - {/* BYOK credential status for selected servers */} {selectedMCPServers.length > 0 && !selectedMCPServers.includes("__all__") && selectedMCPServers.some((serverId) => { @@ -1593,28 +1571,31 @@ const ChatUI: React.FC = ({ return (
- {serverName} requires your API key +

{serverName} requires your API key

{server.has_user_credential ? (
- - Connected + + Connected
) : ( - + )}
); @@ -1624,23 +1605,21 @@ const ChatUI: React.FC = ({
- - Vector Store - - Select vector store(s) to use for this LLM API call. You can set up your vector store{" "} - - here - - . - - } - > - +
+
= ({
- - Guardrails - - Select guardrail(s) to use for this LLM API call. You can set up your guardrails{" "} - - here - - . - - } - > - +
+
= ({
- - Policies - - Select policy/policies to apply to this LLM API call. Policies define which guardrails are - applied based on conditions. You can set up your policies{" "} - - here - - . - - } - > - +
+
= ({ />
- {/* Code Interpreter Toggle - Only for Responses endpoint */} {endpointType === EndpointType.RESPONSES && (
= ({
)} - {/* Main Chat Area */}
{endpointType === EndpointType.REALTIME ? ( = ({ ) : ( <>
- {simplified ? "Chat" : "Test Key"} +

{simplified ? "Chat" : "Test Key"}

- + {!simplified && ( - setIsGetCodeModalVisible(true)} - className="bg-gray-100 hover:bg-gray-200 text-gray-700 border-gray-300" - icon={CodeOutlined} - > + )}
{chatHistory.length === 0 && ( -
- - Start a conversation, generate an image, or handle audio +
+
)} @@ -1772,29 +1739,26 @@ const ChatUI: React.FC = ({
))} - {/* Show MCP events during loading if no assistant message exists yet */} {isLoading && mcpEvents.length > 0 && (endpointType === EndpointType.RESPONSES || endpointType === EndpointType.CHAT) && chatHistory.length > 0 && chatHistory[chatHistory.length - 1].role === "user" && ( -
+
-
+
- +
Assistant
@@ -1804,27 +1768,34 @@ const ChatUI: React.FC = ({ )} {isLoading && ( -
- +
+
)}
- {/* Image Upload Section for Image Edits */} {endpointType === EndpointType.IMAGE_EDITS && (
{uploadedImages.length === 0 ? ( - -

- -

-

Click or drag images to upload

-

+

Click or drag images to upload

+

Support for PNG, JPG, JPEG, GIF, WebP. Multiple images supported.

- + { + handleImageFiles(Array.from(event.target.files || [])); + event.target.value = ""; + }} + /> + ) : (
{uploadedImages.map((file, index) => ( @@ -1841,76 +1812,83 @@ const ChatUI: React.FC = ({ } })()} alt={`Upload preview ${index + 1}`} - className="max-w-32 max-h-32 rounded-md border border-gray-200 object-cover" + className="max-h-32 max-w-32 rounded-md border border-gray-200 object-cover" /> - + +
))} - {/* Add more images button */} -
document.getElementById("additional-image-upload")?.click()} - > -
- -

Add more

-
+
+
)}
)} - {/* Audio Upload Section for Transcriptions */} {endpointType === EndpointType.TRANSCRIPTION && (
{!uploadedAudio ? ( - -

- -

-

Click or drag audio file to upload

-

+

Click or drag audio file to upload

+

Support for MP3, MP4, MPEG, MPGA, M4A, WAV, WEBM formats. Max file size: 25 MB.

- + { + const file = event.target.files?.[0]; + if (file) handleAudioUpload(file); + event.target.value = ""; + }} + /> + ) : ( -
-
- +
+
+
- + + Remove +
)}
)} - {/* Show file previews above input when files are uploaded */} {endpointType === EndpointType.RESPONSES && responsesUploadedImage && ( = ({ /> )} - {/* Code Interpreter indicator and sample prompts when enabled */} {endpointType === EndpointType.RESPONSES && codeInterpreter.enabled && (
-
+
{isLoading ? ( <> - - Running Python code... +
- {/* Sample prompts - only show when not loading */} {!isLoading && (
{[ @@ -1961,7 +1938,8 @@ const ChatUI: React.FC = ({ ].map((prompt, idx) => (
)} - {/* Suggested prompts - show when chat is empty and not loading (skip for MCP - uses structured form) */} {chatHistory.length === 0 && !isLoading && endpointType !== EndpointType.MCP && ( -
+
{(endpointType === EndpointType.A2A_AGENTS ? ["What can you help me with?", "Tell me about yourself", "What tasks can you perform?"] : ["Write me a poem", "Explain quantum computing", "Draft a polite email requesting a meeting"] @@ -1982,7 +1959,7 @@ const ChatUI: React.FC = ({ + + + + {codeInterpreter.enabled + ? "Code Interpreter enabled (click to disable)" + : "Enable Code Interpreter"} + )}
- {/* Middle: input field or MCP structured form */} {endpointType === EndpointType.MCP && selectedMCPServers.length === 1 && selectedMCPServers[0] !== "__all__" && selectedMCPDirectTool ? ( -
+
{(() => { const rawSel = selectedMCPServers[0]; - let toolPool: any[] = []; + let toolPool: { name: string }[] = []; if (rawSel.startsWith("toolset:")) { const toolsetId = rawSel.slice("toolset:".length); const toolset = mcpToolsets.find((t) => t.toolset_id === toolsetId); @@ -2060,82 +2043,51 @@ const ChatUI: React.FC = ({ } else { toolPool = serverToolsMap[rawSel] || []; } - const mcpTool = toolPool.find((t: any) => t.name === selectedMCPDirectTool); + const mcpTool = toolPool.find((t) => t.name === selectedMCPDirectTool); return mcpTool ? ( ) : ( -
+
Loading tool schema...
); })()}
) : ( -