Update main.py

main.py

@@ -2,6 +2,7 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from model_loader import predictVerdict, getConfidence
 from rag_service import evaluateCase
+import uvicorn
 
 app = FastAPI(title="Legal RAG Backend", version="1.0.0")
 
@@ -17,6 +18,8 @@ class ExplainResponse(BaseModel):
     confidence: float
     explanation: str
     retrievedChunks: dict
+    extractedKeywords: list
+    prompt: str
 
 @app.get("/health")
 async def healthCheck():
@@ -36,12 +39,45 @@ async def explain(request: PredictRequest):
     try:
         result = evaluateCase(request.text)
 
+        # SAFELY extract fields (prevents KeyError)
+        verdict = (
+            result.get("verdict")
+            or result.get("finalVerdictByGemini")
+            or "unknown"
+        )
+
+        confidence = result.get("confidence", 0.0)
+
+        explanation = (
+            result.get("explanation")
+            or result.get("geminiOutput")
+            or "No explanation generated."
+        )
+
+        retrievedChunks = (
+            result.get("retrievedChunks")
+            or result.get("support")
+            or {}
+        )
+
+        extractedKeywords = result.get("extractedKeywords", [])
+
+        prompt = (
+            result.get("prompt")
+            or result.get("promptToGemini")
+            or ""
+        )
+
         return ExplainResponse(
-            verdict=result["finalVerdictByGemini"],
-            confidence=result["confidence"],
-            explanation=result["geminiOutput"],
-            retrievedChunks=result["support"]
+            verdict=verdict,
+            confidence=confidence,
+            explanation=explanation,
+            retrievedChunks=retrievedChunks,
+            extractedKeywords=extractedKeywords,
+            prompt=prompt
         )
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
+
+