feat(partitions): scope chunk listing to a single file (fix O(partition) detail view)

openrag/api/routers/admin/partitions.py

@@ -20,7 +20,7 @@
 from api.schemas.admin.partition_schemas import PartitionDetailResponse, UpdatePartitionRequest
 from core.utils.logging import get_logger
 from di.providers import get_partition_service
-from fastapi import APIRouter, Depends, Form, HTTPException, Request, Response, status
+from fastapi import APIRouter, Depends, Form, HTTPException, Query, Request, Response, status
 from fastapi.responses import JSONResponse
 
 logger = get_logger()
@@ -166,7 +166,7 @@ async def get_file(
     request: Request,
     partition: str,
     file_id: str,
-    limit: int = 2000,
+    limit: int = Query(default=2000, ge=0),
     partition_viewer=Depends(require_partition_viewer),
     service=Depends(get_partition_service),
 ):
@@ -188,11 +188,14 @@ async def get_file(
 
 @router.get(
     "/{partition}/chunks",
-    description="""List all document chunks in a partition.
+    description="""List document chunks in a partition.
 
 **Parameters:**
 - `partition`: The partition name
 - `include_embedding`: Include vector embeddings in response (default: true)
+- `file_id`: Restrict to a single file's chunks (filtered server-side; recommended
+  for the document detail view to avoid loading the whole partition)
+- `limit`: Maximum number of chunks to return (default: unbounded)
 
 **Response:**
 Returns all chunks with:
@@ -211,11 +214,18 @@ async def list_all_chunks(
     request: Request,
     partition: str,
     include_embedding: bool = True,
+    file_id: str | None = None,
+    limit: int | None = Query(default=None, ge=0),
     partition_viewer=Depends(require_partition_viewer),
     service=Depends(get_partition_service),
 ):
-    """List all chunks in a partition."""
-    items = await service.list_all_chunks(partition=partition, include_embedding=include_embedding)
+    """List chunks in a partition, optionally scoped to a single file."""
+    items = await service.list_all_chunks(
+        partition=partition,
+        include_embedding=include_embedding,
+        file_id=file_id,
+        limit=limit,
+    )
     chunks = [
         {
             "link": str(request.url_for("get_extract", extract_id=it["metadata"]["_id"])),

openrag/services/orchestrators/partition_service.py

@@ -51,6 +51,16 @@
 logger = get_logger()
 
 
+def _validate_limit(limit: int | None) -> None:
+    """Reject negative ``limit`` values before they reach ``rows[:limit]``.
+
+    A negative bound would silently drop tail rows (e.g. ``-1`` returns all but
+    the last chunk) instead of capping the result, so treat it as a 422.
+    """
+    if limit is not None and limit < 0:
+        raise ValidationError("`limit` must be greater than or equal to 0.", code="INVALID_LIMIT")
+
+
 class PartitionService:
     """Partition lifecycle, membership and read-through orchestration."""
 
@@ -341,6 +351,7 @@ async def get_file_chunks(self, partition: str, file_id: str, limit: int = 2000)
         The router builds the extract links and strips ``_id`` from the
         surfaced metadata, exactly as before.
         """
+        _validate_limit(limit)
         if not await self.file_exists(file_id, partition):
             raise NotFoundError(
                 f"'{file_id}' not found in partition '{partition}'",
@@ -355,16 +366,34 @@ async def get_file_chunks(self, partition: str, file_id: str, limit: int = 2000)
             rows = rows[:limit]
         return [{k: v for k, v in row.items() if k != "text"} for row in rows]
 
-    async def list_all_chunks(self, partition: str, include_embedding: bool = True) -> list[dict]:
-        """Return ``{"content", "metadata"}`` dicts for every chunk."""
+    async def list_all_chunks(
+        self,
+        partition: str,
+        include_embedding: bool = True,
+        file_id: str | None = None,
+        limit: int | None = None,
+    ) -> list[dict]:
+        """Return ``{"content", "metadata"}`` dicts for chunks in a partition.
+
+        ``file_id`` scopes the query to a single file, pushing the filter down
+        to the vector store so the document detail view costs O(file) instead of
+        O(partition). ``limit`` caps the number of chunks returned (a defensive
+        bound for pathologically large files).
+        """
+        _validate_limit(limit)
         await self._ensure_partition(partition)
         excluded = {"text"} if include_embedding else {"text", "vector"}
         output_fields = ["*", "vector"] if include_embedding else ["*"]
+        filters: dict[str, Any] = {"partition": partition}
+        if file_id is not None:
+            filters["file_id"] = file_id
         rows = await self._vector_store.query_chunks_by_filter(
             self._collection,
-            {"partition": partition},
+            filters,
             output_fields=output_fields,
         )
+        if limit is not None and len(rows) > limit:
+            rows = rows[:limit]
 
         def _meta(row: dict[str, Any]) -> dict[str, Any]:
             meta: dict[str, Any] = {}

tests/unit/services/orchestrators/test_partition_service.py

@@ -79,6 +79,7 @@ def __init__(self, ids=None, rows=None):
         self._ids = ids or []
         self._rows = rows or []
         self.deleted_ids: list[str] = []
+        self.last_chunk_filters: dict | None = None
 
     async def query_ids_by_filter(self, collection, filters):
         return list(self._ids)
@@ -88,6 +89,7 @@ async def delete(self, ids, collection="default") -> int:
         return len(ids)
 
     async def query_chunks_by_filter(self, collection, filters, output_fields=None):
+        self.last_chunk_filters = dict(filters)
         return list(self._rows)
 
 
@@ -264,6 +266,52 @@ async def test_list_all_chunks_stringifies_vector_when_included():
     assert isinstance(out[0]["metadata"]["vector"], str)
 
 
+@pytest.mark.asyncio
+async def test_list_all_chunks_without_file_id_filters_partition_only():
+    vstore = FakeVectorStore(rows=[{"text": "t", "_id": "1"}])
+    svc = _svc(prepo=FakePartitionRepo({"p"}), vstore=vstore)
+    await svc.list_all_chunks("p", include_embedding=False)
+    assert vstore.last_chunk_filters == {"partition": "p"}
+
+
+@pytest.mark.asyncio
+async def test_list_all_chunks_scopes_to_file_id_when_given():
+    """file_id is pushed down to the vector store so the detail view is O(file)."""
+    vstore = FakeVectorStore(rows=[{"text": "t", "_id": "1"}])
+    svc = _svc(prepo=FakePartitionRepo({"p"}), vstore=vstore)
+    await svc.list_all_chunks("p", include_embedding=False, file_id="f-123")
+    assert vstore.last_chunk_filters == {"partition": "p", "file_id": "f-123"}
+
+
+@pytest.mark.asyncio
+async def test_list_all_chunks_applies_limit():
+    rows = [{"text": str(i), "_id": str(i)} for i in range(5)]
+    svc = _svc(prepo=FakePartitionRepo({"p"}), vstore=FakeVectorStore(rows=rows))
+    out = await svc.list_all_chunks("p", include_embedding=False, limit=2)
+    assert len(out) == 2
+
+
+@pytest.mark.asyncio
+async def test_list_all_chunks_rejects_negative_limit():
+    rows = [{"text": str(i), "_id": str(i)} for i in range(5)]
+    svc = _svc(prepo=FakePartitionRepo({"p"}), vstore=FakeVectorStore(rows=rows))
+    with pytest.raises(ValidationError) as ei:
+        await svc.list_all_chunks("p", include_embedding=False, limit=-1)
+    assert ei.value.status_code == 422
+
+
+@pytest.mark.asyncio
+async def test_get_file_chunks_rejects_negative_limit():
+    rows = [{"_id": str(i), "text": "body"} for i in range(5)]
+    svc = _svc(
+        drepo=FakeDocumentRepo(files={("f", "p")}),
+        vstore=FakeVectorStore(rows=rows),
+    )
+    with pytest.raises(ValidationError) as ei:
+        await svc.get_file_chunks("p", "f", limit=-1)
+    assert ei.value.status_code == 422
+
+
 # --------------------------------------------------------------------------- #
 # membership
 # --------------------------------------------------------------------------- #