From 436c7aa75143b2e1091a8215cc74c721bd9611db Mon Sep 17 00:00:00 2001 From: Rohan Awhad <30470101+RohanAwhad@users.noreply.github.com> Date: Mon, 16 Jun 2025 05:19:48 -0400 Subject: [PATCH] =?UTF-8?q?feat:=20Add=20url=20field=20to=20PaginatedRespo?= =?UTF-8?q?nse=20and=20populate=20it=20using=20route=20=E2=80=A6=20(#2419)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit …path # What does this PR do? Closes #1847 Changes: - llama_stack/apis/common/responses.py: adds optional `url` field to PaginatedResponse - llama_stack/distribution/server/server.py: automatically populate the URL field with route path ## Test Plan - Built and ran llama stack server using the following cmds: ```bash export INFERENCE_MODEL=llama3.1:8b llama stack build --run --template ollama --image-type container llama stack run llama_stack/templates/ollama/run.yaml ``` - Ran `curl` to test if we are seeing the `url` param in response: ```bash curl -X 'GET' \ 'http://localhost:8321/v1/agents' \ -H 'accept: application/json' ``` - Expected and Received Output: `{"data":[],"has_more":false,"url":"/v1/agents"}` --------- Co-authored-by: Rohan Awhad --- docs/_static/llama-stack-spec.html | 4 ++++ docs/_static/llama-stack-spec.yaml | 3 +++ llama_stack/apis/common/responses.py | 2 ++ llama_stack/distribution/server/server.py | 6 +++++- tests/unit/server/test_sse.py | 25 ++++++++++++++++++++++- 5 files changed, 38 insertions(+), 2 deletions(-) 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"