diff --git a/openrag/api/routers/admin/partitions.py b/openrag/api/routers/admin/partitions.py index 7196d540d..2136eeac1 100644 --- a/openrag/api/routers/admin/partitions.py +++ b/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"])), diff --git a/openrag/services/orchestrators/partition_service.py b/openrag/services/orchestrators/partition_service.py index c4d4d8869..7a972f20b 100644 --- a/openrag/services/orchestrators/partition_service.py +++ b/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] = {} diff --git a/tests/unit/services/orchestrators/test_partition_service.py b/tests/unit/services/orchestrators/test_partition_service.py index 6089a1387..d09586ce7 100644 --- a/tests/unit/services/orchestrators/test_partition_service.py +++ b/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 # --------------------------------------------------------------------------- #