Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 63 additions & 0 deletions docs/openapi.json
Original file line number Diff line number Diff line change
Expand Up @@ -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": [
Expand Down Expand Up @@ -1082,6 +1116,7 @@
"delete_conversation",
"feedback",
"get_models",
"get_shields",
"get_metrics",
"get_config",
"info",
Expand Down Expand Up @@ -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": {
Expand Down
96 changes: 96 additions & 0 deletions src/app/endpoints/shields.py
Original file line number Diff line number Diff line change
@@ -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
Comment on lines +6 to +8
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion | 🟠 Major

Import Depends from fastapi, not fastapi.params.

Aligns with project guidelines for FastAPI imports.

Apply this diff:

-from fastapi import APIRouter, HTTPException, Request, status
-from fastapi.params import Depends
+from fastapi import APIRouter, HTTPException, Request, status, Depends

As per coding guidelines

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
from fastapi import APIRouter, HTTPException, Request, status
from fastapi.params import Depends
from llama_stack_client import APIConnectionError
from fastapi import APIRouter, HTTPException, Request, status, Depends
from llama_stack_client import APIConnectionError
🤖 Prompt for AI Agents
In src/app/endpoints/shields.py around lines 6 to 8, the file currently imports
Depends from fastapi.params; update the import to bring Depends from fastapi
instead (i.e., remove fastapi.params and add Depends to the existing fastapi
import list) so all FastAPI symbols are imported consistently from fastapi per
project guidelines.


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
2 changes: 2 additions & 0 deletions src/app/routers.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from app.endpoints import (
info,
models,
shields,
root,
query,
health,
Expand All @@ -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")
Expand Down
1 change: 1 addition & 0 deletions src/models/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down
18 changes: 18 additions & 0 deletions src/models/responses.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""

Expand Down
Loading
Loading