diff --git a/CHANGELOG.md b/CHANGELOG.md index 4874a4d..8a80a64 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -165,6 +165,18 @@ This project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html). `AMP_SCORING_PATCH_WINDOW_SECONDS`, disabled with `0`, and advertised at `GET /spec` as `scoring_patch_limit` (`null` when off). Counters are per process. +- **Search pages, like the listing endpoints.** `POST /memories/search` gained + `offset`, and both endpoints now share one page ceiling: `MAX_PAGE_SIZE` lives + beside the request models because it is part of the request contract, and two + hard-coded `100`s - one in `SearchRequest`, one in the listing routes - is how a + client ends up told one ceiling by one endpoint and refused by the next. + `offset` windows the results the caller may *read*, the same measure the listing + endpoints use, so a page means the same thing wherever it comes from. The + response gained `has_more` and echoes `offset` / `limit`, since a short page and + a final page are indistinguishable without one of them. +- **A stray `total` in the search example** in `docs/api-reference.md` - left from + the field's rename, and still documenting the count that was never computed. + ### Changed - **Every endpoint returns one error shape.** `PATCH /memories/{id}` answered a conflict with `{"detail": ...}` while `DELETE` answered with diff --git a/docs/api-reference.md b/docs/api-reference.md index 922898b..5f12f81 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -539,7 +539,8 @@ curl -X POST http://localhost:8765/amp/v1/memories/search \ | `owner_id` | string | Yes | Only return cells belonging to this owner | | `types` | string[] | No | Filter to specific memory types. Omit to search all types | | `status` | string[] | No | Lifecycle statuses to include (default: `["active"]`) | -| `limit` | int [1-100] | No | Maximum number of results to return (default: `10`) | +| `limit` | int [1-100] | No | Maximum number of results in this page (default: `10`, ceiling `max_page_size` from [`GET /spec`](#get-spec)). A larger value is refused with `422`, not clamped | +| `offset` | int ≥ 0 | No | Skip this many results the caller may read (default: `0`). Pages the ranked results, so page 2 is page 2 of what this caller can see | | `include_stale` | bool | No | Shorthand to add `"stale"` to the status filter (default: `false`) | **Response `200 OK`** @@ -589,7 +590,10 @@ curl -X POST http://localhost:8765/amp/v1/memories/search \ } } ], - "total": 1, + "returned": 1, + "has_more": true, + "offset": 0, + "limit": 5, "query": "what programming languages does the user know?" } ``` @@ -600,13 +604,15 @@ curl -X POST http://localhost:8765/amp/v1/memories/search \ |-------|-------------| | `results` | Array of matching `MemoryCell` objects, ordered by relevance | | `returned` | The number of cells in this page, bounded by the request's `limit`. It is not a count of everything that matched - the server does not compute one. | +| `has_more` | Whether another page holds anything. `false` means this was the last one, which `returned` alone cannot say: a short page and a final page look identical | +| `offset`, `limit` | The window this page used, echoed back so a client can advance without keeping its own count | | `query` | The query string echoed back | **Error responses** | Status | `error.code` | Cause | |--------|-------------|-------| -| `422` | `VALIDATION_ERROR` | Missing `query` or `owner_id`, or `limit` out of range | +| `422` | `VALIDATION_ERROR` | Missing `query` or `owner_id`, or `limit` / `offset` out of range | --- diff --git a/server/amp_server/main.py b/server/amp_server/main.py index 7200dfd..83bc834 100644 --- a/server/amp_server/main.py +++ b/server/amp_server/main.py @@ -27,6 +27,7 @@ from amp_server.lifecycle import LifecycleEngine from amp_server.limits import MAX_CELL_SIZE_BYTES from amp_server.models import ( + MAX_PAGE_SIZE, ErrorResponse, LifecycleStatus, MemoryAccessPolicy, @@ -41,7 +42,7 @@ SearchRequest, SearchResponse, ) -from amp_server.paging import DEFAULT_PAGE_SIZE, MAX_PAGE_SIZE, readable_page +from amp_server.paging import DEFAULT_PAGE_SIZE, readable_page from amp_server.ratelimit import ( ScoringPatchLimiter, limit_from_env, @@ -583,8 +584,24 @@ async def search_memories( raise missing_agent_id() storage = get_storage() - results = await storage.search(body, agent_id=x_amp_agent_id) - response = SearchResponse(results=results, returned=len(results), query=body.query) + + # One cell past the page, which is what makes `has_more` an answer rather than + # an inference: the adapter ranks the whole candidate set before slicing, so + # the extra cell costs one comparison and no extra query. `model_copy` does + # not re-validate, so asking for `limit + 1` cannot trip the ceiling that + # `limit` itself is bounded by. + window = body.model_copy(update={"limit": body.limit + 1}) + cells = await storage.search(window, agent_id=x_amp_agent_id) + + page = cells[: body.limit] + response = SearchResponse( + results=page, + returned=len(page), + has_more=len(cells) > body.limit, + offset=body.offset, + limit=body.limit, + query=body.query, + ) return response.model_dump(mode="json") diff --git a/server/amp_server/models.py b/server/amp_server/models.py index 4e93256..40800aa 100644 --- a/server/amp_server/models.py +++ b/server/amp_server/models.py @@ -161,6 +161,13 @@ class MemoryCellUpdate(BaseModel): # --- Search Models --- +#: The largest page any endpoint returns. It lives beside the request models +#: because it is part of the request contract: the search model bounds `limit` +#: with it and the listing routes do too, so no endpoint can promise a client one +#: ceiling while the next refuses it. Two hard-coded 100s is how that drifts. +MAX_PAGE_SIZE = 100 + + class SearchRequest(BaseModel): query: str owner_id: str | None = None @@ -168,22 +175,30 @@ class SearchRequest(BaseModel): status: list[LifecycleStatus] = Field( default_factory=lambda: [LifecycleStatus.ACTIVE] ) - limit: int = Field(default=10, ge=1, le=100) + # Default 10 rather than the listing endpoints' 20: that is what search has + # always used, and a default is not worth breaking clients over. + limit: int = Field(default=10, ge=1, le=MAX_PAGE_SIZE) + #: Skips this many results *the caller may read*, so page 2 is page 2 of the + #: caller's own view, exactly as `offset` on the listing endpoints is. + offset: int = Field(default=0, ge=0) include_stale: bool = False class SearchResponse(BaseModel): """Search results for one page. - `returned` is the size of this page, not the number of cells that matched: - the field was called `total` and documented as "may exceed `limit`", which it - never could, because the value is the length of the list beside it. A name - that promises a count the server does not compute is worse than no count. - Paging a search is a v0.2 question; the adapter returns one page. + The fields say what the page is and what it is not. `returned` is the size of + this page - the field was called `total` and documented as "may exceed + `limit`", which it never could, because the value was the length of the list + beside it. `has_more` is the honest answer to "is there more", and `offset` / + `limit` echo the window so a client can page without keeping its own count. """ results: list[MemoryCell] returned: int + has_more: bool + offset: int + limit: int query: str diff --git a/server/amp_server/paging.py b/server/amp_server/paging.py index 2cac4f2..f7d4221 100644 --- a/server/amp_server/paging.py +++ b/server/amp_server/paging.py @@ -15,11 +15,16 @@ from __future__ import annotations from amp_server.access_control import check_read_access -from amp_server.models import LifecycleStatus, MemoryCell, MemoryType +from amp_server.models import ( + MAX_PAGE_SIZE, + LifecycleStatus, + MemoryCell, + MemoryType, +) from amp_server.storage.base import StorageAdapter -#: The most cells one request may ask for. -MAX_PAGE_SIZE = 100 +__all__ = ["DEFAULT_PAGE_SIZE", "MAX_PAGE_SIZE", "SCAN_CEILING", "readable_page"] + #: The page size when the caller does not choose one. DEFAULT_PAGE_SIZE = 20 #: How many candidates one request may examine. A filter that matches almost diff --git a/server/amp_server/storage/chroma.py b/server/amp_server/storage/chroma.py index 1d8efa2..efed955 100644 --- a/server/amp_server/storage/chroma.py +++ b/server/amp_server/storage/chroma.py @@ -238,7 +238,13 @@ async def search(self, request: SearchRequest, agent_id: str) -> list[MemoryCell scored.append((combined_score(similarity, cell), cell)) scored.sort(key=lambda pair: pair[0], reverse=True) - return [cell for _, cell in scored[: request.limit]] + # `offset` windows the cells the caller may read, the same measure + # `query` uses, so a page means the same thing on both endpoints and both + # backends. The access filter above has already run, so the window counts + # readable matches rather than candidates. + return [ + cell for _, cell in scored[request.offset : request.offset + request.limit] + ] async def query( self, diff --git a/server/amp_server/storage/postgres.py b/server/amp_server/storage/postgres.py index f74d096..91631aa 100644 --- a/server/amp_server/storage/postgres.py +++ b/server/amp_server/storage/postgres.py @@ -286,7 +286,13 @@ async def search(self, request: SearchRequest, agent_id: str) -> list[MemoryCell scored.append((combined_score(similarity, cell), cell)) scored.sort(key=lambda pair: pair[0], reverse=True) - return [cell for _, cell in scored[: request.limit]] + # `offset` windows the cells the caller may read, the same measure + # `query` uses, so a page means the same thing on both endpoints and both + # backends. The access filter above has already run, so the window counts + # readable matches rather than candidates. + return [ + cell for _, cell in scored[request.offset : request.offset + request.limit] + ] async def query( self, diff --git a/server/tests/test_adapter_contract.py b/server/tests/test_adapter_contract.py index 83b6a2c..45df8a4 100644 --- a/server/tests/test_adapter_contract.py +++ b/server/tests/test_adapter_contract.py @@ -430,6 +430,61 @@ async def test_search_can_include_stale_cells(adapter): assert len(with_stale) == 1 +@pytest.mark.asyncio +async def test_search_windows_the_ranked_result(adapter): + """`offset` pages the *readable* ranked matches, on every backend. + + Windows must mean the same thing here as in `query`: a client that pages one + backend and then another has to see the same cells in the same order. + """ + ids = [] + for index in range(4): + cell = make_cell(owner_id=_OWNER, created_by=_CREATOR, text=f"invoice {index}") + await adapter.save(cell) + ids.append(cell.id) + + request = SearchRequest(query="invoice", owner_id=_OWNER, limit=2) + first = await adapter.search(request, agent_id=_OWNER) + + request = SearchRequest(query="invoice", owner_id=_OWNER, limit=2, offset=2) + second = await adapter.search(request, agent_id=_OWNER) + + request = SearchRequest(query="invoice", owner_id=_OWNER, limit=2, offset=99) + beyond = await adapter.search(request, agent_id=_OWNER) + + assert [len(first), len(second)] == [2, 2] + assert not ({c.id for c in first} & {c.id for c in second}) + assert {c.id for c in first} | {c.id for c in second} == set(ids) + assert beyond == [] + + +@pytest.mark.asyncio +async def test_search_windows_after_the_access_filter(adapter): + """A page counts cells the caller may read, not cells ranked ahead of them.""" + await adapter.save( + make_cell( + owner_id=_OWNER, + created_by=_CREATOR, + text="invoice for someone else", + readable_by=["agent_elsewhere*"], + ) + ) + for index in range(2): + await adapter.save( + make_cell( + owner_id=_OWNER, + created_by=_CREATOR, + text=f"invoice {index}", + readable_by=["agent_reader*"], + ) + ) + + request = SearchRequest(query="invoice", owner_id=_OWNER, limit=2) + results = await adapter.search(request, agent_id="agent_reader_1") + + assert len(results) == 2, "the unreadable cell was counted against the limit" + + @pytest.mark.asyncio async def test_search_honours_the_limit(adapter): for index in range(5): diff --git a/server/tests/test_paging.py b/server/tests/test_paging.py index e8ce31e..b99415f 100644 --- a/server/tests/test_paging.py +++ b/server/tests/test_paging.py @@ -14,10 +14,9 @@ import pytest from conftest import make_cell -from httpx import ASGITransport, AsyncClient +from httpx import ASGITransport, AsyncClient, Response -from amp_server.models import LifecycleStatus -from amp_server.paging import MAX_PAGE_SIZE +from amp_server.models import MAX_PAGE_SIZE, LifecycleStatus _OWNER = "user-paging" _READER = "agent-paging-reader" @@ -222,3 +221,112 @@ async def test_spec_advertises_the_page_ceiling(): assert capabilities["max_page_size"] == MAX_PAGE_SIZE assert MAX_PAGE_SIZE >= 1 + + +# --------------------------------------------------------------------------- +# Search pages the same way +# --------------------------------------------------------------------------- + + +async def _search(payload: dict) -> Response: + async with await _client() as client: + return await client.post( + "/amp/v1/memories/search", json=payload, headers=_headers() + ) + + +def _query(**overrides: object) -> dict: + payload: dict[str, object] = {"query": "page cell", "owner_id": _OWNER} + payload.update(overrides) + return payload + + +@pytest.mark.asyncio +async def test_search_pages_with_offset_and_reports_has_more(): + _install_app() + await _seed(3, readable_by=[_READER]) + + first = (await _search(_query(limit=2))).json() + second = (await _search(_query(limit=2, offset=2))).json() + + assert first["returned"] == 2 + assert first["has_more"] is True + assert second["returned"] == 1 + assert second["has_more"] is False + + +@pytest.mark.asyncio +async def test_search_echoes_the_window_it_used(): + _install_app() + await _seed(2, readable_by=[_READER]) + + body = (await _search(_query(limit=1, offset=1))).json() + + assert body["offset"] == 1 + assert body["limit"] == 1 + assert body["returned"] == 1 + + +@pytest.mark.asyncio +async def test_search_pages_do_not_overlap_and_cover_everything(): + _install_app() + stored = await _seed(5, readable_by=[_READER]) + + seen: list[str] = [] + async with await _client() as client: + offset = 0 + while True: + page = ( + await client.post( + "/amp/v1/memories/search", + json=_query(limit=2, offset=offset), + headers=_headers(), + ) + ).json() + seen.extend(cell["id"] for cell in page["results"]) + if not page["has_more"]: + break + offset += page["returned"] + + assert sorted(seen) == sorted(stored) + assert len(seen) == len(set(seen)), "a cell appeared on two pages" + + +@pytest.mark.asyncio +async def test_the_search_window_counts_results_the_caller_may_read(): + """Same rule as the listing endpoints: read first, then window.""" + _install_app() + await _seed(4, readable_by=[_STRANGER]) + readable = await _seed(3, readable_by=[_READER]) + + page = (await _search(_query(limit=3))).json() + + assert page["returned"] == 3 + assert page["has_more"] is False + assert {cell["id"] for cell in page["results"]} == set(readable) + + +@pytest.mark.asyncio +async def test_a_search_window_outside_the_bounds_is_refused(): + _install_app() + await _seed(1, readable_by=[_READER]) + + assert (await _search(_query(limit=MAX_PAGE_SIZE + 1))).status_code == 422 + assert (await _search(_query(limit=0))).status_code == 422 + assert (await _search(_query(offset=-1))).status_code == 422 + + +@pytest.mark.asyncio +async def test_the_search_ceiling_is_the_advertised_one(): + """One number for both endpoints, so a client cannot be told two.""" + _install_app() + + async with await _client() as client: + advertised = (await client.get("/amp/v1/spec")).json()["capabilities"][ + "max_page_size" + ] + + assert advertised == MAX_PAGE_SIZE + await _seed(1, readable_by=[_READER]) + assert (await _search(_query(limit=advertised))).status_code == 200 + assert (await _search(_query(limit=advertised + 1))).status_code == 422 diff --git a/spec/v0.1.0/openapi.json b/spec/v0.1.0/openapi.json index a73b0b0..c60898a 100644 --- a/spec/v0.1.0/openapi.json +++ b/spec/v0.1.0/openapi.json @@ -412,6 +412,12 @@ "title": "Limit", "type": "integer" }, + "offset": { + "default": 0, + "minimum": 0.0, + "title": "Offset", + "type": "integer" + }, "owner_id": { "anyOf": [ {