feat: add pagination for custom components (#178)

Author: mathislucka

src/deepset_mcp/api/custom_components/models.py

@@ -19,11 +19,3 @@ class CustomComponentInstallation(BaseModel):
     logs: list[dict[str, Any]]
     organization_id: str
     user_info: DeepsetUser | None = None
-
-
-class CustomComponentInstallationList(BaseModel):
-    """Model representing a list of custom component installations."""
-
-    data: list[CustomComponentInstallation]
-    total: int
-    has_more: bool

src/deepset_mcp/api/custom_components/protocols.py

@@ -4,15 +4,16 @@
 
 from typing import Protocol
 
-from deepset_mcp.api.custom_components.models import CustomComponentInstallationList
+from deepset_mcp.api.custom_components.models import CustomComponentInstallation
+from deepset_mcp.api.shared_models import PaginatedResponse
 
 
 class CustomComponentsProtocol(Protocol):
     """Protocol defining the implementation for CustomComponentsResource."""
 
     async def list_installations(
-        self, limit: int = 20, page_number: int = 1, field: str = "created_at", order: str = "DESC"
-    ) -> CustomComponentInstallationList:
+        self, limit: int = 20, after: str | None = None, field: str = "created_at", order: str = "DESC"
+    ) -> PaginatedResponse[CustomComponentInstallation]:
         """List custom component installations."""
         ...
 

src/deepset_mcp/api/custom_components/resource.py

@@ -3,10 +3,13 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from typing import Any
+from urllib.parse import quote
 
-from deepset_mcp.api.custom_components.models import CustomComponentInstallationList
+from deepset_mcp.api.custom_components.models import CustomComponentInstallation
 from deepset_mcp.api.custom_components.protocols import CustomComponentsProtocol
+from deepset_mcp.api.exceptions import UnexpectedAPIError
 from deepset_mcp.api.protocols import AsyncClientProtocol
+from deepset_mcp.api.shared_models import PaginatedResponse
 from deepset_mcp.api.transport import raise_for_status
 
 
@@ -21,29 +24,52 @@ def __init__(self, client: AsyncClientProtocol) -> None:
         self._client = client
 
     async def list_installations(
-        self, limit: int = 20, page_number: int = 1, field: str = "created_at", order: str = "DESC"
-    ) -> CustomComponentInstallationList:
-        """List custom component installations.
+        self, limit: int = 20, after: str | None = None, field: str = "created_at", order: str = "DESC"
+    ) -> PaginatedResponse[CustomComponentInstallation]:
+        """Lists custom component installations and returns the first page of results.
 
-        :param limit: Maximum number of installations to return.
-        :param page_number: Page number for pagination.
+        The returned object can be iterated over to fetch subsequent pages.
+
+        :param limit: Maximum number of installations to return per page.
+        :param after: The cursor to fetch the next page of results.
         :param field: Field to sort by.
         :param order: Sort order (ASC or DESC).
-
-        :returns: List of custom component installations.
+        :returns: A `PaginatedResponse` object containing the first page of installations.
         """
+        # 1. Prepare arguments for the initial API call
+        # TODO: Pagination in the deepset API is currently implemented in an unintuitive way.
+        # TODO: The cursor is always time based (created_at) and after signifies installations older than the
+        # TODO: current cursor
+        # TODO: while 'before' signals installations younger than the current cursor.
+        # TODO: This is applied irrespective of any sort (e.g. name) that would conflict with this approach.
+        # TODO: Change this to 'after' once the behaviour is fixed on the deepset API
+        request_params = {"limit": limit, "field": field, "order": order, "before": after}
+        request_params = {k: v for k, v in request_params.items() if v is not None}
+
+        # 2. Make the first API call using a private, stateless method
+        page = await self._list_api_call(**request_params)
+
+        # 3. Inject the logic needed for subsequent fetches into the response object
+        page._inject_paginator(
+            fetch_func=self._list_api_call,
+            # Base args for the *next* fetch don't include initial cursors
+            base_args={"limit": limit, "field": field, "order": order},
+        )
+        return page
+
+    async def _list_api_call(self, **kwargs: Any) -> PaginatedResponse[CustomComponentInstallation]:
+        """A private, stateless method that performs the raw API call."""
+        params = "&".join([f"{key}={quote(str(value), safe='')}" for key, value in kwargs.items()])
         resp = await self._client.request(
-            endpoint=f"v2/custom_components?limit={limit}&page_number={page_number}&field={field}&order={order}",
+            endpoint=f"v2/custom_components?{params}",
             method="GET",
             response_type=dict[str, Any],
         )
-
         raise_for_status(resp)
-
         if resp.json is None:
-            return CustomComponentInstallationList(data=[], total=0, has_more=False)
+            raise UnexpectedAPIError(status_code=resp.status_code, message="Empty response", detail=None)
 
-        return CustomComponentInstallationList(**resp.json)
+        return PaginatedResponse[CustomComponentInstallation].create_with_cursor_field(resp.json, "custom_component_id")
 
     async def get_latest_installation_logs(self) -> str | None:
         """Get the logs from the latest custom component installation.

src/deepset_mcp/tools/custom_components.py

@@ -2,25 +2,28 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from deepset_mcp.api.custom_components.models import CustomComponentInstallationList
+from deepset_mcp.api.custom_components.models import CustomComponentInstallation
 from deepset_mcp.api.protocols import AsyncClientProtocol
+from deepset_mcp.api.shared_models import PaginatedResponse
 
 
 async def list_custom_component_installations(
-    *, client: AsyncClientProtocol, workspace: str
-) -> CustomComponentInstallationList | str:
+    *, client: AsyncClientProtocol, workspace: str, limit: int = 20, after: str | None = None
+) -> PaginatedResponse[CustomComponentInstallation] | str:
     """List custom component installations.
 
     :param client: The API client to use.
     :param workspace: The workspace to operate in.
+    :param limit: Maximum number of installations to return per page.
+    :param after: The cursor to fetch the next page of results.
 
     :returns: Custom component installations or error message.
     """
     custom_components = client.custom_components(workspace)
     users = client.users()
 
     try:
-        installations = await custom_components.list_installations()
+        installations = await custom_components.list_installations(limit=limit, after=after)
     except Exception as e:
         return f"Failed to retrieve custom component installations: {e}"
 

test/unit/api/custom_components/test_custom_components_resource.py

@@ -4,10 +4,8 @@
 
 import pytest
 
-from deepset_mcp.api.custom_components.models import (
-    CustomComponentInstallationList,
-)
 from deepset_mcp.api.custom_components.resource import CustomComponentsResource
+from deepset_mcp.api.shared_models import PaginatedResponse
 from test.unit.conftest import BaseFakeClient
 
 
@@ -29,9 +27,7 @@ async def test_list_installations() -> None:
         "has_more": False,
     }
 
-    fake_client = BaseFakeClient(
-        responses={"v2/custom_components?limit=20&page_number=1&field=created_at&order=DESC": mock_data}
-    )
+    fake_client = BaseFakeClient(responses={"v2/custom_components?limit=20&field=created_at&order=DESC": mock_data})
 
     def custom_components(workspace: str) -> CustomComponentsResource:
         return CustomComponentsResource(client=fake_client)
@@ -41,7 +37,7 @@ def custom_components(workspace: str) -> CustomComponentsResource:
     resource = fake_client.custom_components("test-workspace")
     result = await resource.list_installations()
 
-    assert isinstance(result, CustomComponentInstallationList)
+    assert isinstance(result, PaginatedResponse)
     assert len(result.data) == 1
     assert result.total == 1
     assert not result.has_more
@@ -63,9 +59,7 @@ async def test_list_installations_empty() -> None:
         "has_more": False,
     }
 
-    fake_client = BaseFakeClient(
-        responses={"v2/custom_components?limit=20&page_number=1&field=created_at&order=DESC": mock_data}
-    )
+    fake_client = BaseFakeClient(responses={"v2/custom_components?limit=20&field=created_at&order=DESC": mock_data})
 
     def custom_components(workspace: str) -> CustomComponentsResource:
         return CustomComponentsResource(client=fake_client)
@@ -75,7 +69,7 @@ def custom_components(workspace: str) -> CustomComponentsResource:
     resource = fake_client.custom_components("test-workspace")
     result = await resource.list_installations()
 
-    assert isinstance(result, CustomComponentInstallationList)
+    assert isinstance(result, PaginatedResponse)
     assert len(result.data) == 0
     assert result.total == 0
     assert not result.has_more
@@ -91,7 +85,7 @@ async def test_list_installations_with_params() -> None:
     }
 
     fake_client = BaseFakeClient(
-        responses={"v2/custom_components?limit=50&page_number=2&field=status&order=ASC": mock_data}
+        responses={"v2/custom_components?limit=50&field=status&order=ASC&before=cursor_123": mock_data}
     )
 
     def custom_components(workspace: str) -> CustomComponentsResource:
@@ -100,13 +94,13 @@ def custom_components(workspace: str) -> CustomComponentsResource:
     fake_client.custom_components = custom_components  # type: ignore[method-assign]
 
     resource = fake_client.custom_components("test-workspace")
-    await resource.list_installations(limit=50, page_number=2, field="status", order="ASC")
+    await resource.list_installations(limit=50, after="cursor_123", field="status", order="ASC")
 
     # Check that the request was made with the correct parameters
     assert len(fake_client.requests) == 1
     request = fake_client.requests[0]
     assert "limit=50" in request["endpoint"]
-    assert "page_number=2" in request["endpoint"]
+    assert "before=cursor_123" in request["endpoint"]  # after becomes before due to API quirk
     assert "field=status" in request["endpoint"]
     assert "order=ASC" in request["endpoint"]
 

test/unit/tools/test_custom_components.py

@@ -2,14 +2,15 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from typing import Any
+
 import pytest
 
 from deepset_mcp.api.custom_components.models import (
     CustomComponentInstallation,
-    CustomComponentInstallationList,
 )
 from deepset_mcp.api.exceptions import UnexpectedAPIError
-from deepset_mcp.api.shared_models import DeepsetUser
+from deepset_mcp.api.shared_models import DeepsetUser, PaginatedResponse
 from deepset_mcp.tools.custom_components import (
     get_latest_custom_component_installation_logs,
     list_custom_component_installations,
@@ -20,7 +21,7 @@
 class FakeCustomComponentsResource:
     def __init__(
         self,
-        installations_response: CustomComponentInstallationList | None = None,
+        installations_response: PaginatedResponse[CustomComponentInstallation] | None = None,
         latest_logs_response: str | None = None,
         exception: Exception | None = None,
     ):
@@ -29,8 +30,8 @@ def __init__(
         self._exception = exception
 
     async def list_installations(
-        self, limit: int = 20, page_number: int = 1, field: str = "created_at", order: str = "DESC"
-    ) -> CustomComponentInstallationList:
+        self, limit: int = 20, after: str | None = None, field: str = "created_at", order: str = "DESC"
+    ) -> PaginatedResponse[CustomComponentInstallation]:
         if self._exception:
             raise self._exception
         if self._installations_response is not None:
@@ -84,7 +85,7 @@ def users(self) -> FakeUserResource:
 @pytest.mark.asyncio
 async def test_list_custom_component_installations() -> None:
     """Test listing custom component installations."""
-    mock_installations = CustomComponentInstallationList(
+    mock_installations = PaginatedResponse[CustomComponentInstallation](
         data=[
             CustomComponentInstallation(
                 custom_component_id="comp_123",
@@ -134,7 +135,7 @@ async def test_list_custom_component_installations() -> None:
 
     result = await list_custom_component_installations(client=client, workspace="test-workspace")
 
-    assert isinstance(result, CustomComponentInstallationList)
+    assert isinstance(result, PaginatedResponse)
     assert len(result.data) == 2
     assert result.total == 2
     assert result.has_more is False
@@ -168,7 +169,7 @@ async def test_list_custom_component_installations() -> None:
 @pytest.mark.asyncio
 async def test_list_custom_component_installations_empty() -> None:
     """Test listing custom component installations when none exist."""
-    mock_installations = CustomComponentInstallationList(
+    mock_installations = PaginatedResponse[CustomComponentInstallation](
         data=[],
         total=0,
         has_more=False,
@@ -183,7 +184,7 @@ async def test_list_custom_component_installations_empty() -> None:
 
     result = await list_custom_component_installations(client=client, workspace="test-workspace")
 
-    assert isinstance(result, CustomComponentInstallationList)
+    assert isinstance(result, PaginatedResponse)
     assert len(result.data) == 0
     assert result.total == 0
     assert result.has_more is False
@@ -192,7 +193,7 @@ async def test_list_custom_component_installations_empty() -> None:
 @pytest.mark.asyncio
 async def test_list_custom_component_installations_user_fetch_error() -> None:
     """Test listing custom component installations when user fetch fails."""
-    mock_installations = CustomComponentInstallationList(
+    mock_installations = PaginatedResponse[CustomComponentInstallation](
         data=[
             CustomComponentInstallation(
                 custom_component_id="comp_123",
@@ -216,7 +217,7 @@ async def test_list_custom_component_installations_user_fetch_error() -> None:
 
     result = await list_custom_component_installations(client=client, workspace="test-workspace")
 
-    assert isinstance(result, CustomComponentInstallationList)
+    assert isinstance(result, PaginatedResponse)
     assert len(result.data) == 1
     assert result.data[0].created_by_user_id == "user_unknown"
     assert result.data[0].user_info is None  # User fetch failed, so user_info should be None
@@ -271,3 +272,62 @@ async def test_get_latest_custom_component_installation_logs_api_error() -> None
 
     result = await get_latest_custom_component_installation_logs(client=client, workspace="test-workspace")
     assert result == "Failed to retrieve latest installation logs: API Error (Status Code: 500)"
+
+
+@pytest.mark.asyncio
+async def test_list_custom_component_installations_with_pagination_params() -> None:
+    """Test listing custom component installations with pagination parameters."""
+    mock_installations = PaginatedResponse[CustomComponentInstallation](
+        data=[
+            CustomComponentInstallation(
+                custom_component_id="comp_123",
+                status="installed",
+                version="1.0.0",
+                created_by_user_id="user_123",
+                organization_id="org-123",
+                logs=[{"level": "INFO", "msg": "Installation complete"}],
+            )
+        ],
+        total=1,
+        has_more=False,
+    )
+
+    # Create a custom resource that tracks the parameters passed
+    class TrackingCustomComponentsResource(FakeCustomComponentsResource):
+        def __init__(self, installations_response: PaginatedResponse[CustomComponentInstallation]) -> None:
+            super().__init__(installations_response=installations_response)
+            self.called_with: dict[str, Any] = {}
+
+        async def list_installations(
+            self, limit: int = 20, after: str | None = None, field: str = "created_at", order: str = "DESC"
+        ) -> PaginatedResponse[CustomComponentInstallation]:
+            self.called_with = {"limit": limit, "after": after, "field": field, "order": order}
+            return await super().list_installations(limit, after, field, order)
+
+    custom_components_resource = TrackingCustomComponentsResource(installations_response=mock_installations)
+    user_resource = FakeUserResource(
+        users={
+            "user_123": DeepsetUser(
+                user_id="user_123", given_name="John", family_name="Doe", email="john.doe@example.com"
+            )
+        }
+    )
+    client = FakeClient(
+        custom_components_resource=custom_components_resource,
+        user_resource=user_resource,
+    )
+
+    # Test with custom parameters
+    result = await list_custom_component_installations(
+        client=client, workspace="test-workspace", limit=50, after="cursor_123"
+    )
+
+    # Verify the parameters were passed correctly
+    assert custom_components_resource.called_with["limit"] == 50
+    assert custom_components_resource.called_with["after"] == "cursor_123"
+    assert custom_components_resource.called_with["field"] == "created_at"  # default
+    assert custom_components_resource.called_with["order"] == "DESC"  # default
+
+    # Verify the result
+    assert isinstance(result, PaginatedResponse)
+    assert len(result.data) == 1