Update main.py

main.py

@@ -21,20 +21,43 @@ class Guardrail:
     def guard(self, prompt):
         return self.classifier(prompt)
 
+    def determine_level(self, label, score):
+        if label == "SAFE":
+            return 0, "safe"
+        else:
+            if score > 0.9:
+                return 4, "high"
+            elif score > 0.75:
+                return 3, "medium"
+            elif score > 0.5:
+                return 2, "low"
+            else:
+                return 1, "very low"
+
 
 class TextPrompt(BaseModel):
     prompt: str
 
 
+class ClassificationResult(BaseModel):
+    label: str
+    score: float
+    level: int
+    severity_label: str
+
+
 app = FastAPI()
 guardrail = Guardrail()
 
 
-@app.post("/classify/")
+@app.post("/classify/", response_model=ClassificationResult)
 def classify_text(text_prompt: TextPrompt):
     try:
         result = guardrail.guard(text_prompt.prompt)
-        return result
+        label = result[0]['label']
+        score = result[0]['score']
+        level, severity_label = guardrail.determine_level(label, score)
+        return {"label": label, "score": score, "level": level, "severity_label": severity_label}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))