mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
fix(vector-stores): paginate managed resource list by unified_resource_id cursor
`list_user_resources`, which backs the managed vector store listing, carried the same broken cursor as the managed batch listing: `after` was applied as `where id > after` against the table's random-uuid primary key, while the value a client holds is the `unified_resource_id` returned as each item's `id`. That compares a base64 cursor to unrelated uuids, so a page could re-select the rows it just returned or skip rows that should have come next. Ordering by non-unique `created_at` alone left the order partial as well, so rows sharing a timestamp could shift between requests. The listing now uses a Prisma cursor on the unique `unified_resource_id` with `unified_resource_id` as a secondary sort key, rejects an `after` that does not resolve to a resource the caller can list with a 400 rather than an empty page, and derives `has_more` from an extra fetched row instead of from whether the page came back full. `build_list_page` read `first_id` / `last_id` through attribute access only, which raised AttributeError on every non-empty page of this listing because it builds plain dicts. It now reads the id from either shape.
This commit is contained in:
parent
7b019cf152
commit
f913d42abb
4 changed files with 253 additions and 21 deletions
|
|
@ -542,35 +542,39 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]):
|
|||
Returns:
|
||||
Dictionary with list of resources and pagination info
|
||||
"""
|
||||
from fastapi import HTTPException
|
||||
|
||||
owner_filter = build_owner_filter(user_api_key_dict)
|
||||
if owner_filter is None:
|
||||
return build_list_page([])
|
||||
|
||||
where_clause: Dict[str, Any] = {**owner_filter}
|
||||
where_clause: Dict[str, Any] = {**owner_filter, **(additional_filters or {})}
|
||||
|
||||
table = getattr(self.prisma_client.db, self.table_name)
|
||||
|
||||
if after:
|
||||
where_clause["id"] = {"gt": after}
|
||||
cursor_row = await table.find_first(where={**where_clause, "unified_resource_id": after})
|
||||
if cursor_row is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Invalid 'after' cursor: no {self.resource_type} found with id '{after}'.",
|
||||
)
|
||||
|
||||
# Add additional filters
|
||||
if additional_filters:
|
||||
where_clause.update(additional_filters)
|
||||
page_size = limit or 20
|
||||
cursor_args: Dict[str, Any] = {"cursor": {"unified_resource_id": after}, "skip": 1} if after else {}
|
||||
|
||||
# Fetch resources
|
||||
fetch_limit = limit or 20
|
||||
table = getattr(self.prisma_client.db, self.table_name)
|
||||
resources = await table.find_many(
|
||||
where=where_clause,
|
||||
take=fetch_limit,
|
||||
order={"created_at": "desc"},
|
||||
take=page_size + 1,
|
||||
order=[{"created_at": "desc"}, {"unified_resource_id": "desc"}],
|
||||
**cursor_args,
|
||||
)
|
||||
|
||||
resource_objects: List[Any] = []
|
||||
for resource in resources:
|
||||
try:
|
||||
# Stop once we have enough
|
||||
if len(resource_objects) >= (limit or 20):
|
||||
break
|
||||
has_more = len(resources) > page_size
|
||||
|
||||
resource_objects: List[Any] = []
|
||||
for resource in resources[:page_size]:
|
||||
try:
|
||||
# Parse resource object
|
||||
resource_data = resource.resource_object
|
||||
if isinstance(resource_data, str):
|
||||
|
|
@ -590,4 +594,4 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]):
|
|||
)
|
||||
continue
|
||||
|
||||
return build_list_page(resource_objects, has_more=len(resource_objects) == (limit or 20))
|
||||
return build_list_page(resource_objects, has_more=has_more)
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ identifying ids are denied so an empty user_id can never select an
|
|||
unscoped query.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from litellm.proxy._types import (
|
||||
|
|
@ -17,15 +18,23 @@ from litellm.proxy._types import (
|
|||
)
|
||||
|
||||
|
||||
def _item_id(item: Any) -> str | None:
|
||||
if isinstance(item, Mapping):
|
||||
return item.get("id")
|
||||
return getattr(item, "id", None)
|
||||
|
||||
|
||||
def build_list_page(items: List[Any], has_more: bool = False) -> Dict[str, Any]:
|
||||
"""Build the OpenAI-style paginated list response shape used by managed
|
||||
file/batch/vector-store listings. ``first_id`` and ``last_id`` are
|
||||
sourced from each item's ``.id`` attribute."""
|
||||
file/batch/vector-store listings. ``first_id`` and ``last_id`` are sourced
|
||||
from each item's ``id``, which is an attribute on the pydantic objects the
|
||||
file and batch listings build and a key on the plain dicts the vector-store
|
||||
listing builds."""
|
||||
return {
|
||||
"object": "list",
|
||||
"data": items,
|
||||
"first_id": items[0].id if items else None,
|
||||
"last_id": items[-1].id if items else None,
|
||||
"first_id": _item_id(items[0]) if items else None,
|
||||
"last_id": _item_id(items[-1]) if items else None,
|
||||
"has_more": has_more,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -108,6 +108,184 @@ async def test_list_identity_less_caller_returns_empty_without_query():
|
|||
resource.prisma_client.db.litellm_test_resource_table.find_many.assert_not_awaited()
|
||||
|
||||
|
||||
def _make_row(index: int, created_by: str = "alice"):
|
||||
row = MagicMock()
|
||||
row.id = f"pk-{index:03d}"
|
||||
row.unified_resource_id = f"unified-resource-{index:03d}"
|
||||
row.created_at = 1_000_000 + index
|
||||
row.created_by = created_by
|
||||
row.resource_object = {
|
||||
"id": f"provider-vector-store-{index:03d}",
|
||||
"object": "vector_store",
|
||||
"name": f"store-{index:03d}",
|
||||
}
|
||||
return row
|
||||
|
||||
|
||||
def _row_matches(row, where) -> bool:
|
||||
return all(
|
||||
isinstance(value, dict) or getattr(row, field, None) == value
|
||||
for field, value in where.items()
|
||||
)
|
||||
|
||||
|
||||
def _make_paginating_resource(rows: List) -> _StubResource:
|
||||
async def find_many(where, take, order, cursor=None, skip=0):
|
||||
result = sorted(
|
||||
rows, key=lambda r: (r.created_at, r.unified_resource_id), reverse=True
|
||||
)
|
||||
id_filter = where.get("id")
|
||||
if isinstance(id_filter, dict) and "gt" in id_filter:
|
||||
result = [r for r in result if r.id > id_filter["gt"]]
|
||||
if cursor is not None:
|
||||
(cur_field, cur_val), = cursor.items()
|
||||
idx = next(
|
||||
(i for i, r in enumerate(result) if getattr(r, cur_field) == cur_val),
|
||||
None,
|
||||
)
|
||||
if idx is None:
|
||||
return []
|
||||
result = result[idx + skip:]
|
||||
return result[:take]
|
||||
|
||||
async def find_first(where):
|
||||
return next((r for r in rows if _row_matches(r, where)), None)
|
||||
|
||||
cache = MagicMock()
|
||||
cache.async_get_cache = AsyncMock(return_value=None)
|
||||
|
||||
prisma = MagicMock()
|
||||
table = MagicMock()
|
||||
table.find_many = AsyncMock(side_effect=find_many)
|
||||
table.find_first = AsyncMock(side_effect=find_first)
|
||||
prisma.db = MagicMock()
|
||||
setattr(prisma.db, "litellm_test_resource_table", table)
|
||||
|
||||
return _StubResource(internal_usage_cache=cache, prisma_client=prisma)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_resources_pagination_uses_unified_resource_id_cursor():
|
||||
"""The ``after`` cursor a client sends back is a resource's
|
||||
``unified_resource_id`` -- that is what the listing returns as each item's
|
||||
``id`` and as ``last_id``. Paginating must use a Prisma cursor on that
|
||||
unique column, not a ``where id > after`` filter against the random-uuid
|
||||
primary key, which compares the cursor to unrelated values and so loops or
|
||||
silently drops resources.
|
||||
"""
|
||||
rows = [_make_row(i) for i in range(3)]
|
||||
resource = _make_paginating_resource(rows)
|
||||
|
||||
await resource.list_user_resources(
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="alice"),
|
||||
limit=2,
|
||||
after=rows[2].unified_resource_id,
|
||||
)
|
||||
|
||||
table = resource.prisma_client.db.litellm_test_resource_table
|
||||
kwargs = table.find_many.await_args.kwargs
|
||||
assert kwargs["cursor"] == {"unified_resource_id": rows[2].unified_resource_id}
|
||||
assert kwargs["skip"] == 1
|
||||
assert kwargs["order"] == [
|
||||
{"created_at": "desc"},
|
||||
{"unified_resource_id": "desc"},
|
||||
]
|
||||
assert "id" not in kwargs["where"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_resources_pagination_walks_all_pages_without_loops_or_gaps():
|
||||
"""Walking pages the way a client does -- feeding ``last_id`` back as
|
||||
``after`` -- must return every resource exactly once, newest first."""
|
||||
rows = [_make_row(i) for i in range(5)]
|
||||
resource = _make_paginating_resource(rows)
|
||||
user = UserAPIKeyAuth(user_id="alice")
|
||||
|
||||
seen = []
|
||||
after = None
|
||||
for _ in range(len(rows) + 5):
|
||||
page = await resource.list_user_resources(
|
||||
user_api_key_dict=user, limit=2, after=after
|
||||
)
|
||||
seen.extend(item["id"] for item in page["data"])
|
||||
if not page["has_more"]:
|
||||
break
|
||||
assert page["last_id"] is not None, "has_more was true but there is no cursor"
|
||||
assert page["last_id"] != after, "cursor did not advance (pagination loop)"
|
||||
after = page["last_id"]
|
||||
|
||||
assert seen == [r.unified_resource_id for r in reversed(rows)]
|
||||
assert len(seen) == len(set(seen))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_resources_has_more_false_on_exactly_full_final_page():
|
||||
"""``has_more`` must mean "another row exists", not "this page is full",
|
||||
so a resource count that is an exact multiple of ``limit`` does not cost
|
||||
every client an extra empty request."""
|
||||
rows = [_make_row(i) for i in range(4)]
|
||||
resource = _make_paginating_resource(rows)
|
||||
user = UserAPIKeyAuth(user_id="alice")
|
||||
|
||||
first = await resource.list_user_resources(user_api_key_dict=user, limit=2)
|
||||
second = await resource.list_user_resources(
|
||||
user_api_key_dict=user, limit=2, after=first["last_id"]
|
||||
)
|
||||
|
||||
assert first["has_more"] is True
|
||||
assert second["has_more"] is False
|
||||
assert [item["id"] for item in second["data"]] == [
|
||||
rows[1].unified_resource_id,
|
||||
rows[0].unified_resource_id,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_resources_rejects_unknown_after_cursor():
|
||||
"""An ``after`` that does not resolve to a resource the caller can see is
|
||||
a client error. Returning an empty page instead is indistinguishable from
|
||||
the end of the list, so a stale cursor silently truncates the listing."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
resource = _make_paginating_resource([_make_row(0)])
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await resource.list_user_resources(
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="alice"),
|
||||
limit=2,
|
||||
after="does-not-exist-xyz",
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "does-not-exist-xyz" in str(exc_info.value.detail)
|
||||
resource.prisma_client.db.litellm_test_resource_table.find_many.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_resources_cursor_lookup_is_scoped_to_the_caller():
|
||||
"""The cursor lookup must run against the rows the caller can list. A
|
||||
Prisma cursor resolves by unique column regardless of the ``where``
|
||||
filter, so an unscoped lookup would let one user anchor their page window
|
||||
to another user's resource."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
resource = _make_paginating_resource([_make_row(0)])
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
await resource.list_user_resources(
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="alice"),
|
||||
limit=2,
|
||||
after="unified-resource-000",
|
||||
additional_filters={"created_by": "bob"},
|
||||
)
|
||||
|
||||
where = resource.prisma_client.db.litellm_test_resource_table.find_first.await_args.kwargs[
|
||||
"where"
|
||||
]
|
||||
assert where["created_by"] == "bob"
|
||||
assert where["unified_resource_id"] == "unified-resource-000"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"caller_team_id,expected",
|
||||
|
|
|
|||
|
|
@ -5,12 +5,53 @@ Tests for managed-resource tenant isolation helpers.
|
|||
import pytest
|
||||
|
||||
from litellm.llms.base_llm.managed_resources.isolation import (
|
||||
build_list_page,
|
||||
build_owner_filter,
|
||||
can_access_resource,
|
||||
)
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# build_list_page
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_list_page_reads_ids_from_dict_items():
|
||||
"""The vector-store listing builds plain dicts, so sourcing the cursor ids
|
||||
through attribute access alone raised AttributeError on every non-empty
|
||||
page."""
|
||||
page = build_list_page(
|
||||
[{"id": "resource-a"}, {"id": "resource-b"}], has_more=True
|
||||
)
|
||||
|
||||
assert page["first_id"] == "resource-a"
|
||||
assert page["last_id"] == "resource-b"
|
||||
assert page["has_more"] is True
|
||||
|
||||
|
||||
def test_list_page_reads_ids_from_object_items():
|
||||
class _Item:
|
||||
def __init__(self, id: str):
|
||||
self.id = id
|
||||
|
||||
page = build_list_page([_Item("batch-a"), _Item("batch-b")])
|
||||
|
||||
assert page["first_id"] == "batch-a"
|
||||
assert page["last_id"] == "batch-b"
|
||||
assert page["has_more"] is False
|
||||
|
||||
|
||||
def test_list_page_empty():
|
||||
assert build_list_page([]) == {
|
||||
"object": "list",
|
||||
"data": [],
|
||||
"first_id": None,
|
||||
"last_id": None,
|
||||
"has_more": False,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# build_owner_filter
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue