Update main.py

main.py

@@ -5,7 +5,7 @@ import torch
 from detoxify import Detoxify
 import asyncio
 from fastapi.concurrency import run_in_threadpool
-from typing import List
+from typing import List, Optional
 
 class Guardrail:
     def __init__(self):
@@ -80,6 +80,16 @@ class TopicBannerResult(BaseModel):
     labels: list
     scores: list
 
+class GuardrailsRequest(BaseModel):
+    prompt: str
+    guardrails: List[str]
+    labels: Optional[List[str]] = None
+
+class GuardrailsResponse(BaseModel):
+    prompt_injection: Optional[ClassificationResult] = None
+    toxicity: Optional[ToxicityResult] = None
+    topic_banner: Optional[TopicBannerResult] = None
+
 app = FastAPI()
 guardrail = Guardrail()
 toxicity_classifier = Detoxify('original')
@@ -123,6 +133,35 @@ async def classify_topic_banner(request: TopicBannerRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+@app.post("/api/guardrails", response_model=GuardrailsResponse)
+async def evaluate_guardrails(request: GuardrailsRequest):
+    tasks = []
+    response = GuardrailsResponse()
+
+    if "pi" in request.guardrails:
+        tasks.append(classify_text(TextPrompt(prompt=request.prompt)))
+    if "tox" in request.guardrails:
+        tasks.append(classify_toxicity(TextPrompt(prompt=request.prompt)))
+    if "top" in request.guardrails:
+        if not request.labels:
+            raise HTTPException(status_code=400, detail="Labels are required for topic banner classification")
+        tasks.append(classify_topic_banner(TopicBannerRequest(prompt=request.prompt, labels=request.labels)))
+
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+
+    for result, guardrail in zip(results, request.guardrails):
+        if isinstance(result, Exception):
+            # Handle the exception as needed
+            continue
+        if guardrail == "pi":
+            response.prompt_injection = result
+        elif guardrail == "tox":
+            response.toxicity = result
+        elif guardrail == "top":
+            response.topic_banner = result
+
+    return response
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)