Merge pull request #628 from yangcao77/shield-endpoint

[RHDHPAI-1150] create /v1/shields endpoint

Author: tisnik

docs/openapi.json

@@ -123,6 +123,40 @@
                 }
             }
         },
+        "/v1/shields": {
+            "get": {
+                "tags": [
+                    "shields"
+                ],
+                "summary": "Shields Endpoint Handler",
+                "description": "Handle requests to the /shields endpoint.\n\nProcess GET requests to the /shields endpoint, returning a list of available\nshields from the Llama Stack service.\n\nRaises:\n    HTTPException: If unable to connect to the Llama Stack server or if\n    shield retrieval fails for any reason.\n\nReturns:\n    ShieldsResponse: An object containing the list of available shields.",
+                "operationId": "shields_endpoint_handler_v1_shields_get",
+                "responses": {
+                    "200": {
+                        "description": "Successful Response",
+                        "content": {
+                            "application/json": {
+                                "schema": {
+                                    "$ref": "#/components/schemas/ShieldsResponse"
+                                }
+                            }
+                        },
+                        "shields": [
+                            {
+                                "identifier": "lightspeed_question_validity-shield",
+                                "provider_resource_id": "lightspeed_question_validity-shield",
+                                "provider_id": "lightspeed_question_validity",
+                                "type": "shield",
+                                "params": {}
+                            }
+                        ]
+                    },
+                    "500": {
+                        "description": "Connection to Llama Stack is broken"
+                    }
+                }
+            }
+        },
         "/v1/query": {
             "post": {
                 "tags": [
@@ -1082,6 +1116,7 @@
                     "delete_conversation",
                     "feedback",
                     "get_models",
+                    "get_shields",
                     "get_metrics",
                     "get_config",
                     "info",
@@ -2990,6 +3025,34 @@
                 "title": "ServiceConfiguration",
                 "description": "Service configuration."
             },
+            "ShieldsResponse": {
+                "properties": {
+                    "shields": {
+                        "items": {
+                            "additionalProperties": true,
+                            "type": "object"
+                        },
+                        "type": "array",
+                        "title": "Shields",
+                        "description": "List of shields available",
+                        "examples": [
+                            {
+                                "identifier": "lightspeed_question_validity-shield",
+                                "params": {},
+                                "provider_id": "lightspeed_question_validity",
+                                "provider_resource_id": "lightspeed_question_validity-shield",
+                                "type": "shield"
+                            }
+                        ]
+                    }
+                },
+                "type": "object",
+                "required": [
+                    "shields"
+                ],
+                "title": "ShieldsResponse",
+                "description": "Model representing a response to shields request."
+            },
             "StatusResponse": {
                 "properties": {
                     "functionality": {

src/app/endpoints/shields.py

@@ -0,0 +1,96 @@
+"""Handler for REST API call to list available shields."""
+
+import logging
+from typing import Annotated, Any
+
+from fastapi import APIRouter, HTTPException, Request, status
+from fastapi.params import Depends
+from llama_stack_client import APIConnectionError
+
+from authentication import get_auth_dependency
+from authentication.interface import AuthTuple
+from authorization.middleware import authorize
+from client import AsyncLlamaStackClientHolder
+from configuration import configuration
+from models.config import Action
+from models.responses import ShieldsResponse
+from utils.endpoints import check_configuration_loaded
+
+logger = logging.getLogger(__name__)
+router = APIRouter(tags=["shields"])
+
+
+shields_responses: dict[int | str, dict[str, Any]] = {
+    200: {
+        "shields": [
+            {
+                "identifier": "lightspeed_question_validity-shield",
+                "provider_resource_id": "lightspeed_question_validity-shield",
+                "provider_id": "lightspeed_question_validity",
+                "type": "shield",
+                "params": {},
+            }
+        ]
+    },
+    500: {"description": "Connection to Llama Stack is broken"},
+}
+
+
+@router.get("/shields", responses=shields_responses)
+@authorize(Action.GET_SHIELDS)
+async def shields_endpoint_handler(
+    request: Request,
+    auth: Annotated[AuthTuple, Depends(get_auth_dependency())],
+) -> ShieldsResponse:
+    """
+    Handle requests to the /shields endpoint.
+
+    Process GET requests to the /shields endpoint, returning a list of available
+    shields from the Llama Stack service.
+
+    Raises:
+        HTTPException: If unable to connect to the Llama Stack server or if
+        shield retrieval fails for any reason.
+
+    Returns:
+        ShieldsResponse: An object containing the list of available shields.
+    """
+    # Used only by the middleware
+    _ = auth
+
+    # Nothing interesting in the request
+    _ = request
+
+    check_configuration_loaded(configuration)
+
+    llama_stack_configuration = configuration.llama_stack_configuration
+    logger.info("Llama stack config: %s", llama_stack_configuration)
+
+    try:
+        # try to get Llama Stack client
+        client = AsyncLlamaStackClientHolder().get_client()
+        # retrieve shields
+        shields = await client.shields.list()
+        s = [dict(s) for s in shields]
+        return ShieldsResponse(shields=s)
+
+    # connection to Llama Stack server
+    except APIConnectionError as e:
+        logger.error("Unable to connect to Llama Stack: %s", e)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={
+                "response": "Unable to connect to Llama Stack",
+                "cause": str(e),
+            },
+        ) from e
+    # any other exception that can occur during shield listing
+    except Exception as e:
+        logger.error("Unable to retrieve list of shields: %s", e)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={
+                "response": "Unable to retrieve list of shields",
+                "cause": str(e),
+            },
+        ) from e

src/app/routers.py

@@ -5,6 +5,7 @@
 from app.endpoints import (
     info,
     models,
+    shields,
     root,
     query,
     health,
@@ -27,6 +28,7 @@ def include_routers(app: FastAPI) -> None:
     app.include_router(root.router)
     app.include_router(info.router, prefix="/v1")
     app.include_router(models.router, prefix="/v1")
+    app.include_router(shields.router, prefix="/v1")
     app.include_router(query.router, prefix="/v1")
     app.include_router(streaming_query.router, prefix="/v1")
     app.include_router(config.router, prefix="/v1")

src/models/config.py

@@ -358,6 +358,7 @@ class Action(str, Enum):
     DELETE_CONVERSATION = "delete_conversation"
     FEEDBACK = "feedback"
     GET_MODELS = "get_models"
+    GET_SHIELDS = "get_shields"
     GET_METRICS = "get_metrics"
     GET_CONFIG = "get_config"
 

src/models/responses.py

@@ -36,6 +36,24 @@ class ModelsResponse(BaseModel):
     )
 
 
+class ShieldsResponse(BaseModel):
+    """Model representing a response to shields request."""
+
+    shields: list[dict[str, Any]] = Field(
+        ...,
+        description="List of shields available",
+        examples=[
+            {
+                "identifier": "lightspeed_question_validity-shield",
+                "provider_resource_id": "lightspeed_question_validity-shield",
+                "provider_id": "lightspeed_question_validity",
+                "type": "shield",
+                "params": {},
+            }
+        ],
+    )
+
+
 class RAGChunk(BaseModel):
     """Model representing a RAG chunk used in the response."""