v4.0: Backend API — add /api/redline, /api/chat, /api/ocr endpoints

api/main.py

@@ -1,19 +1,19 @@
 """
-ClauseGuard — FastAPI Backend v3.0
+ClauseGuard — FastAPI Backend v4.0
 ══════════════════════════════════
-FIXED in v3.0:
-  • Imports shared modules (no code duplication)
-  • Fixed API schema to accept both {text} and {clauses} from extension
-  • Added rate limiting
-  • Added max text length validation
-  • Fixed CORS (removed wildcard)
-  • Added proper error responses
+New in v4.0:
+  • /api/redline — clause redlining suggestions
+  • /api/chat — RAG chatbot (streaming)
+  • /api/ocr — OCR scanned PDF extraction
+  • Updated analysis to include redlining data
 """
 
 import os
 import re
 import json
 import time
+import uuid
+import tempfile
 from contextlib import asynccontextmanager
 from typing import Optional
 from collections import defaultdict
@@ -21,14 +21,14 @@ from datetime import datetime
 
 import httpx
 import numpy as np
-from fastapi import FastAPI, HTTPException, Depends, Body, Request
+from fastapi import FastAPI, HTTPException, Depends, Body, Request, UploadFile, File as FastAPIFile
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, Field
 
 from auth import get_current_user, require_auth
 
 # ── Import shared modules ──
-# When deployed, these must be in the same directory or on PYTHONPATH
 import sys
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
@@ -36,29 +36,32 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 try:
     from app import (
         split_clauses, classify_cuad, extract_entities,
-        detect_contradictions, compute_risk_score,
+        detect_contradictions, compute_risk_score, analyze_contract,
         CUAD_LABELS, RISK_MAP, DESC_MAP, _model_status,
         cuad_model, cuad_tokenizer
     )
     from obligations import extract_obligations
     from compliance import check_compliance
     from compare import compare_contracts
+    from redlining import generate_redlines
+    from chatbot import index_contract, chat_respond
+    from ocr_engine import parse_pdf_smart, get_ocr_status
     _SHARED_MODULES = True
-except ImportError:
+except ImportError as e:
     _SHARED_MODULES = False
-    print("[API] WARNING: Could not import shared modules, using inline fallbacks")
+    print(f"[API] WARNING: Could not import shared modules: {e}")
 
 # ─── Config ───
 SUPABASE_URL = os.environ.get("SUPABASE_URL", "")
 SUPABASE_SERVICE_KEY = os.environ.get("SUPABASE_SERVICE_ROLE_KEY", "")
 HF_API_TOKEN = os.environ.get("HF_API_TOKEN", "")
 SAULLM_ENDPOINT = os.environ.get("SAULLM_ENDPOINT", "")
-MAX_TEXT_LENGTH = int(os.environ.get("MAX_TEXT_LENGTH", "100000"))  # 100KB default
+MAX_TEXT_LENGTH = int(os.environ.get("MAX_TEXT_LENGTH", "100000"))
 
 # ─── Rate Limiting ───
-_rate_limits = {}  # ip -> (count, window_start)
+_rate_limits = {}
 RATE_LIMIT_REQUESTS = 30
-RATE_LIMIT_WINDOW = 60  # seconds
+RATE_LIMIT_WINDOW = 60
 
 def _check_rate_limit(client_ip: str) -> bool:
     now = time.time()
@@ -113,25 +116,16 @@ async def supabase_query(table: str, params: dict, headers_extra: dict = {}):
     except Exception:
         return []
 
+# ─── In-memory RAG session store ───
+_rag_sessions: dict = {}
+_RAG_SESSION_MAX = 100
+
 # ─── Request/Response Models ───
 class AnalyzeRequest(BaseModel):
     text: Optional[str] = Field(None, min_length=50)
-    clauses: Optional[list] = None  # FIXED: accept clauses array from extension
+    clauses: Optional[list] = None
     source_url: Optional[str] = None
 
-class AnalyzeResponse(BaseModel):
-    risk_score: int
-    grade: str
-    total_clauses: int
-    flagged_count: int
-    results: list[dict]
-    entities: list[dict]
-    contradictions: list[dict]
-    obligations: list[dict]
-    compliance: dict
-    model: str
-    latency_ms: int
-
 class CompareRequest(BaseModel):
     text_a: str = Field(..., min_length=50)
     text_b: str = Field(..., min_length=50)
@@ -147,21 +141,28 @@ class ExplainResponse(BaseModel):
     legal_basis: str
     recommendation: str
 
+class ChatRequest(BaseModel):
+    message: str = Field(..., min_length=1, max_length=2000)
+    session_id: str
+    history: Optional[list[dict]] = None
+
+class RedlineRequest(BaseModel):
+    session_id: Optional[str] = None
+    text: Optional[str] = None
+    use_llm: bool = True
+
 # ─── App ───
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Models are loaded when app.py is imported
     yield
 
-app = FastAPI(title="ClauseGuard API", version="3.0.0", lifespan=lifespan)
+app = FastAPI(title="ClauseGuard API", version="4.0.0", lifespan=lifespan)
 
-# FIXED: No wildcard CORS
 ALLOWED_ORIGINS = [
     "https://clauseguardweb.netlify.app",
     "http://localhost:3000",
     "http://localhost:3001",
 ]
-# Allow chrome extensions
 app.add_middleware(
     CORSMiddleware,
     allow_origins=ALLOWED_ORIGINS,
@@ -174,36 +175,36 @@ app.add_middleware(
 @app.get("/health")
 async def health():
     model_status = "ml" if _SHARED_MODULES and cuad_model else "regex"
+    ocr_status = get_ocr_status() if _SHARED_MODULES else "unavailable"
     return {
         "status": "ok",
         "model": model_status,
-        "version": "3.0.0",
+        "version": "4.0.0",
         "shared_modules": _SHARED_MODULES,
+        "ocr": ocr_status,
+        "features": ["analyze", "compare", "redline", "chat", "ocr"],
     }
 
-@app.post("/api/analyze", response_model=AnalyzeResponse)
+@app.post("/api/analyze")
 async def analyze(req: AnalyzeRequest, request: Request, user: Optional[dict] = Depends(get_current_user)):
-    # Rate limiting
     client_ip = request.client.host if request.client else "unknown"
     if not _check_rate_limit(client_ip):
-        raise HTTPException(status_code=429, detail="Rate limit exceeded. Try again in 60 seconds.")
+        raise HTTPException(status_code=429, detail="Rate limit exceeded.")
 
-    # FIXED: Accept either text or clauses from extension
     text = req.text
     if not text and req.clauses:
         text = "\n\n".join(req.clauses) if isinstance(req.clauses, list) else str(req.clauses)
 
     if not text or len(text.strip()) < 50:
         raise HTTPException(status_code=400, detail="Text too short (minimum 50 characters)")
-
-    # Max length check
     if len(text) > MAX_TEXT_LENGTH:
-        raise HTTPException(status_code=400, detail=f"Text too long (maximum {MAX_TEXT_LENGTH} characters)")
+        raise HTTPException(status_code=400, detail=f"Text too long (max {MAX_TEXT_LENGTH} chars)")
 
     start = time.time()
+
     clauses = split_clauses(text)
     if not clauses:
-        raise HTTPException(status_code=400, detail="No clauses detected in document")
+        raise HTTPException(status_code=400, detail="No clauses detected")
 
     clause_results = []
     for clause in clauses:
@@ -224,6 +225,15 @@ async def analyze(req: AnalyzeRequest, request: Request, user: Optional[dict] =
     risk, grade, sev_counts = compute_risk_score(clause_results, len(clauses))
     obligations = extract_obligations(text)
     compliance = check_compliance(text)
+
+    # v4.0: Redlining
+    analysis_for_redline = {"clauses": clause_results}
+    redlines = []
+    try:
+        redlines = generate_redlines(analysis_for_redline, use_llm=True)
+    except Exception as e:
+        print(f"[API] Redlining error: {e}")
+
     latency = int((time.time() - start) * 1000)
 
     results_for_db = []
@@ -238,6 +248,29 @@ async def analyze(req: AnalyzeRequest, request: Request, user: Optional[dict] =
             }],
         })
 
+    # v4.0: RAG indexing
+    session_id = None
+    try:
+        chunks, embeddings, _status = index_contract(text)
+        if chunks and embeddings is not None:
+            session_id = uuid.uuid4().hex[:12]
+            if len(_rag_sessions) >= _RAG_SESSION_MAX:
+                oldest = next(iter(_rag_sessions))
+                del _rag_sessions[oldest]
+            _rag_sessions[session_id] = {
+                "chunks": chunks,
+                "embeddings": embeddings,
+                "analysis": {
+                    "risk": {"score": risk, "grade": grade, "breakdown": sev_counts},
+                    "metadata": {"total_clauses": len(clauses), "flagged_clauses": len(clause_results)},
+                    "clauses": clause_results[:30],
+                    "entities": entities[:30],
+                    "contradictions": contradictions,
+                },
+            }
+    except Exception as e:
+        print(f"[API] RAG indexing error: {e}")
+
     if user:
         await supabase_insert("analyses", {
             "user_id": user["id"],
@@ -253,46 +286,120 @@ async def analyze(req: AnalyzeRequest, request: Request, user: Optional[dict] =
             "compliance": compliance,
         })
 
-    return AnalyzeResponse(
-        risk_score=risk,
-        grade=grade,
-        total_clauses=len(clauses),
-        flagged_count=len(set(cr["text"] for cr in clause_results)),
-        results=results_for_db,
-        entities=entities,
-        contradictions=contradictions,
-        obligations=obligations,
-        compliance=compliance,
-        model="ml" if cuad_model else "regex",
-        latency_ms=latency,
-    )
+    return {
+        "risk_score": risk,
+        "grade": grade,
+        "total_clauses": len(clauses),
+        "flagged_count": len(set(cr["text"] for cr in clause_results)),
+        "results": results_for_db,
+        "entities": entities,
+        "contradictions": contradictions,
+        "obligations": obligations,
+        "compliance": compliance,
+        "redlines": redlines,
+        "model": "ml" if cuad_model else "regex",
+        "latency_ms": latency,
+        "session_id": session_id,
+    }
 
 @app.post("/api/compare")
 async def compare(req: CompareRequest, request: Request):
     client_ip = request.client.host if request.client else "unknown"
     if not _check_rate_limit(client_ip):
         raise HTTPException(status_code=429, detail="Rate limit exceeded.")
-    result = compare_contracts(req.text_a, req.text_b)
-    return result
+    return compare_contracts(req.text_a, req.text_b)
+
+@app.post("/api/redline")
+async def redline(req: RedlineRequest, request: Request):
+    client_ip = request.client.host if request.client else "unknown"
+    if not _check_rate_limit(client_ip):
+        raise HTTPException(status_code=429, detail="Rate limit exceeded.")
+
+    if req.session_id and req.session_id in _rag_sessions:
+        analysis = _rag_sessions[req.session_id]["analysis"]
+    elif req.text:
+        result, error = analyze_contract(req.text)
+        if error:
+            raise HTTPException(status_code=400, detail=error)
+        analysis = result
+    else:
+        raise HTTPException(status_code=400, detail="Provide session_id or text")
+
+    redlines = generate_redlines(analysis, use_llm=req.use_llm)
+    return {"redlines": redlines, "count": len(redlines)}
+
+@app.post("/api/chat")
+async def chat(req: ChatRequest, request: Request):
+    client_ip = request.client.host if request.client else "unknown"
+    if not _check_rate_limit(client_ip):
+        raise HTTPException(status_code=429, detail="Rate limit exceeded.")
+
+    if req.session_id not in _rag_sessions:
+        raise HTTPException(status_code=404, detail="Session not found. Analyze a contract first.")
+
+    session = _rag_sessions[req.session_id]
+    response_text = ""
+    for partial in chat_respond(req.message, req.history or [],
+                                 session["chunks"], session["embeddings"], session["analysis"]):
+        response_text = partial
+
+    return {"response": response_text, "session_id": req.session_id}
+
+@app.post("/api/chat/stream")
+async def chat_stream(req: ChatRequest, request: Request):
+    client_ip = request.client.host if request.client else "unknown"
+    if not _check_rate_limit(client_ip):
+        raise HTTPException(status_code=429, detail="Rate limit exceeded.")
+
+    if req.session_id not in _rag_sessions:
+        raise HTTPException(status_code=404, detail="Session not found.")
+
+    session = _rag_sessions[req.session_id]
+
+    async def generate():
+        last = ""
+        for partial in chat_respond(
+            req.message, req.history or [],
+            session["chunks"], session["embeddings"], session["analysis"]
+        ):
+            delta = partial[len(last):]
+            last = partial
+            if delta:
+                yield f"data: {json.dumps({'delta': delta})}\n\n"
+        yield "data: [DONE]\n\n"
+
+    return StreamingResponse(generate(), media_type="text/event-stream")
+
+@app.post("/api/ocr")
+async def ocr_endpoint(file: UploadFile = FastAPIFile(...)):
+    if not file.filename or not file.filename.lower().endswith(".pdf"):
+        raise HTTPException(status_code=400, detail="Only PDF files supported")
+
+    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
+        content = await file.read()
+        tmp.write(content)
+        tmp_path = tmp.name
+
+    try:
+        text, error, method = parse_pdf_smart(tmp_path)
+        if error:
+            raise HTTPException(status_code=400, detail=error)
+        return {"text": text, "method": method, "chars": len(text) if text else 0, "filename": file.filename}
+    finally:
+        os.unlink(tmp_path)
 
 @app.post("/api/explain", response_model=ExplainResponse)
 async def explain(req: ExplainRequest, user: dict = Depends(require_auth)):
     desc = DESC_MAP.get(req.category, "Unknown category.")
     legal = "Consult local consumer protection laws."
-    recommendation = "Review this clause carefully. Consider negotiating or seeking legal advice before agreeing."
+    recommendation = "Review this clause carefully."
 
     if SAULLM_ENDPOINT and HF_API_TOKEN:
         try:
             prompt = (
-                f"You are a consumer protection legal analyst. Analyze this contract clause "
-                f"and explain why it may be unfair or risky.\n\n"
-                f"Clause: \"{req.clause}\"\n"
-                f"Category: {req.category}\n\n"
-                f"Provide:\n"
-                f"1. A plain-English explanation of what this clause means\n"
-                f"2. The specific legal basis or consumer protection concern\n"
-                f"3. A practical recommendation\n\n"
-                f"Be concise. 3-4 sentences per section."
+                f"Analyze this contract clause and explain why it may be risky.\n\n"
+                f"Clause: \"{req.clause}\"\nCategory: {req.category}\n\n"
+                f"Provide: 1) Plain-English explanation 2) Legal basis 3) Recommendation"
             )
             async with httpx.AsyncClient(timeout=30.0) as client:
                 resp = await client.post(
@@ -311,27 +418,16 @@ async def explain(req: ExplainRequest, user: dict = Depends(require_auth)):
         except Exception:
             pass
 
-    return ExplainResponse(
-        clause=req.clause,
-        category=req.category,
-        explanation=desc,
-        legal_basis=legal,
-        recommendation=recommendation,
-    )
+    return ExplainResponse(clause=req.clause, category=req.category,
+                           explanation=desc, legal_basis=legal, recommendation=recommendation)
 
 @app.get("/api/history")
 async def history(user: dict = Depends(require_auth), limit: int = 20, offset: int = 0):
     limit = min(limit, 100)
-    data = await supabase_query(
-        "analyses",
-        {
-            "user_id": f"eq.{user['id']}",
-            "select": "*",
-            "order": "created_at.desc",
-            "limit": str(limit),
-            "offset": str(offset),
-        },
-    )
+    data = await supabase_query("analyses", {
+        "user_id": f"eq.{user['id']}", "select": "*",
+        "order": "created_at.desc", "limit": str(limit), "offset": str(offset),
+    })
     return {"analyses": data, "limit": limit, "offset": offset}
 
 if __name__ == "__main__":