main.py

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")

class PredictRequest(BaseModel):
    text: str

class PredictResponse(BaseModel):
    verdict: str
    confidence: float

class ExplainResponse(BaseModel):
    verdict: str
    confidence: float
    explanation: str
    retrievedChunks: dict
    extractedKeywords: list
    prompt: str

@app.get("/health")
async def healthCheck():
    return {"status": "ok"}

@app.post("/predict", response_model=PredictResponse)
async def predict(request: PredictRequest):
    try:
        verdictResult = predictVerdict(request.text)
        confidenceScore = getConfidence(request.text)
        return PredictResponse(verdict=verdictResult, confidence=confidenceScore)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/explain", response_model=ExplainResponse)
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=verdict,
            confidence=confidence,
            explanation=explanation,
            retrievedChunks=retrievedChunks,
            extractedKeywords=extractedKeywords,
            prompt=prompt
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))