diff --git a/docs/_static/llama-stack-spec.html b/docs/_static/llama-stack-spec.html index fddce0c57..486068258 100644 --- a/docs/_static/llama-stack-spec.html +++ b/docs/_static/llama-stack-spec.html @@ -11157,6 +11157,10 @@ "has_more": { "type": "boolean", "description": "Whether there are more items available after this set" + }, + "url": { + "type": "string", + "description": "The URL for accessing this list" } }, "additionalProperties": false, diff --git a/docs/_static/llama-stack-spec.yaml b/docs/_static/llama-stack-spec.yaml index 49388939f..72f60032e 100644 --- a/docs/_static/llama-stack-spec.yaml +++ b/docs/_static/llama-stack-spec.yaml @@ -7818,6 +7818,9 @@ components: type: boolean description: >- Whether there are more items available after this set + url: + type: string + description: The URL for accessing this list additionalProperties: false required: - data diff --git a/llama_stack/apis/common/responses.py b/llama_stack/apis/common/responses.py index 5cb41e23d..e4cf21a54 100644 --- a/llama_stack/apis/common/responses.py +++ b/llama_stack/apis/common/responses.py @@ -23,7 +23,9 @@ class PaginatedResponse(BaseModel): :param data: The list of items for the current page :param has_more: Whether there are more items available after this set + :param url: The URL for accessing this list """ data: list[dict[str, Any]] has_more: bool + url: str | None = None diff --git a/llama_stack/distribution/server/server.py b/llama_stack/distribution/server/server.py index 4f2427a55..88b64ef2e 100644 --- a/llama_stack/distribution/server/server.py +++ b/llama_stack/distribution/server/server.py @@ -30,6 +30,7 @@ from fastapi.responses import JSONResponse, StreamingResponse from openai import BadRequestError from pydantic import BaseModel, ValidationError +from llama_stack.apis.common.responses import PaginatedResponse from llama_stack.distribution.datatypes import AuthenticationRequiredError, LoggingConfig, StackRunConfig from llama_stack.distribution.distribution import builtin_automatically_routed_apis from llama_stack.distribution.request_headers import PROVIDER_DATA_VAR, User, request_provider_data_context @@ -230,7 +231,10 @@ def create_dynamic_typed_route(func: Any, method: str, route: str) -> Callable: return StreamingResponse(gen, media_type="text/event-stream") else: value = func(**kwargs) - return await maybe_await(value) + result = await maybe_await(value) + if isinstance(result, PaginatedResponse) and result.url is None: + result.url = route + return result except Exception as e: logger.exception(f"Error executing endpoint {route=} {method=}") raise translate_exception(e) from e diff --git a/tests/unit/server/test_sse.py b/tests/unit/server/test_sse.py index c78122294..60e9f4609 100644 --- a/tests/unit/server/test_sse.py +++ b/tests/unit/server/test_sse.py @@ -5,10 +5,12 @@ # the root directory of this source tree. import asyncio +from unittest.mock import AsyncMock, MagicMock import pytest -from llama_stack.distribution.server.server import create_sse_event, sse_generator +from llama_stack.apis.common.responses import PaginatedResponse +from llama_stack.distribution.server.server import create_dynamic_typed_route, create_sse_event, sse_generator @pytest.mark.asyncio @@ -89,3 +91,24 @@ async def test_sse_generator_error_before_response_starts(): # We should have 1 error event assert len(seen_events) == 1 assert 'data: {"error":' in seen_events[0] + + +@pytest.mark.asyncio +async def test_paginated_response_url_setting(): + """Test that PaginatedResponse gets url set to route path.""" + + async def mock_api_method(): + return PaginatedResponse(data=[], has_more=False, url=None) + + route_handler = create_dynamic_typed_route(mock_api_method, "get", "/test/route") + + # Mock minimal request + request = MagicMock() + request.scope = {"user_attributes": {}, "principal": ""} + request.headers = {} + request.body = AsyncMock(return_value=b"") + + result = await route_handler(request) + + assert isinstance(result, PaginatedResponse) + assert result.url == "/test/route"