diff --git a/litellm/router.py b/litellm/router.py index 1ef68e60440..53aa5041784 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6504,6 +6504,7 @@ class Router: "first_id": None, "last_id": None, "has_more": False, + "next_page_token": None, } for result in results: @@ -6513,6 +6514,9 @@ class Router: final_results["first_id"] = getattr(result, "first_id") final_results["last_id"] = getattr(result, "last_id") final_results["data"].extend(result.data) + page_token = getattr(result, "next_page_token", None) + if page_token is not None: + final_results["next_page_token"] = page_token ## check 'has_more' if getattr(result, "has_more", False) is True: diff --git a/tests/unit/test_router_batch_list_pagination.py b/tests/unit/test_router_batch_list_pagination.py new file mode 100644 index 00000000000..87b956aea9d --- /dev/null +++ b/tests/unit/test_router_batch_list_pagination.py @@ -0,0 +1,35 @@ +from typing import Final + +import pytest + +import litellm +from litellm import Router +from litellm.types.utils import OpenAIBatchListResponse + +TOKEN_BY_KEY: Final = {"key-with-more-pages": "1", "key-on-last-page": None} + + +def _deployment(api_key: str) -> dict: + return { + "model_name": "mistral-ocr", + "litellm_params": {"model": "mistral/mistral-ocr-latest", "api_key": api_key}, + "model_info": {"id": api_key}, + } + + +async def _fake_alist_batches(**kwargs: object) -> OpenAIBatchListResponse: + token: Final = TOKEN_BY_KEY[str(kwargs["api_key"])] + return OpenAIBatchListResponse( + data=(), first_id=None, last_id=None, has_more=token is not None, next_page_token=token + ) + + +@pytest.mark.asyncio +async def test_alist_batches_keeps_the_page_token_a_deployment_returned(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "alist_batches", _fake_alist_batches) + router: Final = Router(model_list=[_deployment("key-with-more-pages"), _deployment("key-on-last-page")]) + + result: Final = await router.alist_batches(model="mistral-ocr", limit=3) + + assert result["has_more"] is True + assert result["next_page_token"] == "1"