Update main.py

main.py

@@ -6,7 +6,6 @@ from langchain_groq import ChatGroq
 from langchain_core.prompts import ChatPromptTemplate
 import os
 import torch
-import jwt
 import fitz  # PyMuPDF
 from dotenv import load_dotenv
 from typing import Dict, Optional
@@ -19,7 +18,7 @@ from database import db
 load_dotenv()
 app = FastAPI(title="LexGuard AI Core")
 
-# --- CORS CONFIGURATION (CRITICAL FOR VERCEL & REACT) ---
+# --- CORS CONFIGURATION ---
 origins = [
     "http://localhost:3000",
     "http://127.0.0.1:3000",
@@ -36,33 +35,43 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- MODEL PATHS ---
+# --- MODEL PATHS & API KEYS ---
 MODEL_DIR = "./lexguard_model"
 ABS_MODEL_PATH = os.path.abspath(MODEL_DIR)
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+
+# Retrieve and sanitize the Groq API key
+raw_groq_key = os.getenv("GROQ_API_KEY")
+GROQ_API_KEY = raw_groq_key.strip() if raw_groq_key else None
+
+if not GROQ_API_KEY:
+    print("⚠️ WARNING: GROQ_API_KEY is not detected in environment variables!")
+else:
+    print(f"🔑 GROQ_API_KEY detected (starts with: {GROQ_API_KEY[:8]}...)")
 
 # 2. Load Models (The Sniper & The Analyst)
 print(f"Loading Predictive Model from: {ABS_MODEL_PATH}")
 try:
-    # Force local loading to prevent Hugging Face URL errors
     model = AutoModelForSequenceClassification.from_pretrained(
         ABS_MODEL_PATH, local_files_only=True
     )
     tokenizer = AutoTokenizer.from_pretrained(
         ABS_MODEL_PATH, local_files_only=True
     )
-    # Automatically switch to CPU (-1) if running on a free-tier non-GPU Hugging Face Space
+    # Dynamically select GPU (0) if CUDA is available, otherwise CPU (-1)
     device = 0 if torch.cuda.is_available() else -1
     classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, device=device)
     print(f"✅ Sniper Model Loaded successfully on device: {'GPU' if device == 0 else 'CPU'}")
 except Exception as e:
-    print(f"❌ Error loading model: {e}")
-    exit()
+    print(f"❌ Error loading predictive model: {e}")
+    exit(1)
 
+# Initialize Groq LLM client
 llm = ChatGroq(
     temperature=0,
     model_name="llama-3.3-70b-versatile",
-    groq_api_key=GROQ_API_KEY
+    groq_api_key=GROQ_API_KEY,
+    request_timeout=60,
+    max_retries=2
 )
 
 # 3. Helper: PDF to Text Chunks
@@ -79,7 +88,7 @@ def extract_text_from_pdf(file_bytes):
         paragraphs = text.split('\n\n')
         for p in paragraphs:
             clean_p = " ".join(p.split()).strip()
-            # Filter out tiny page numbers or headers
+            # Filter out tiny page numbers or short headers
             if len(clean_p) > 50: 
                 text_chunks.append(clean_p)
     return text_chunks
@@ -90,7 +99,8 @@ def extract_text_from_pdf(file_bytes):
 def health_check():
     return {
         "status": "LexGuard API is Ready",
-        "gpu_active": torch.cuda.is_available()
+        "gpu_active": torch.cuda.is_available(),
+        "groq_configured": GROQ_API_KEY is not None
     }
 
 @app.post("/users/sync")
@@ -100,7 +110,6 @@ async def sync_user(user: UserSync):
     Syncs user from Clerk Frontend to MongoDB Backend.
     """
     user_data = user.model_dump()
-    # Upsert based on clerk_id
     result = await db.users.update_one(
         {"clerk_id": user.clerk_id},
         {"$set": user_data},
@@ -112,7 +121,7 @@ async def sync_user(user: UserSync):
 @app.post("/analyze_document/")
 async def analyze_document(
     file: UploadFile = File(...),
-    user_rule: str = Form(None)  # React sends this as FormData
+    user_rule: str = Form(None)
 ):
     """
     Main Endpoint: Receives a PDF file + Optional User Rule.
@@ -140,10 +149,7 @@ async def analyze_document(
     # B. The Sniper Pass (Batch Prediction)
     print(f"🔍 Scanning {len(clauses)} clauses...")
     
-    # Map model output to human labels
     label_map = {"LABEL_0": "Safe", "LABEL_1": "Termination", "LABEL_2": "Non-Compete"}
-    
-    # Run inference
     predictions = classifier(clauses, batch_size=8, truncation=True)
 
     # C. The Filter & Logic Pass
@@ -152,7 +158,6 @@ async def analyze_document(
         score = pred['score']
         risk_type = label_map.get(label_str, "Safe")
 
-        # LOGIC: We only keep it if it's RISKY OR if User has a specific rule to check
         is_risky = risk_type != "Safe"
         has_rule = user_rule is not None and len(user_rule.strip()) > 5
         
@@ -160,20 +165,17 @@ async def analyze_document(
             explanation = "Standard clause."
             
             # D. The Analyst Pass (GenAI)
-            if is_risky or (has_rule and i < 20):  # Limit rule checking to first 20 clauses for speed
-                
-                system_msg = f"""
-                You are a legal auditor.
-                Detected Risk Category: {risk_type} (Confidence: {score:.2f})
-                User's Constraint Rule: {user_rule if user_rule else "None"}
-                
-                Task:
-                1. Summarize what this clause says in plain English.
-                2. If a User Rule exists, explicitly state if this clause violates it.
-                3. If it is risky, suggest a 1-sentence edit to make it safer.
-                """
+            if is_risky or (has_rule and i < 20):
+                system_msg = f"""You are a legal auditor.
+Detected Risk Category: {risk_type} (Confidence: {score:.2f})
+User's Constraint Rule: {user_rule if user_rule else "None"}
+
+Task:
+1. Summarize what this clause says in plain English.
+2. If a User Rule exists, explicitly state if this clause violates it.
+3. If it is risky, suggest a 1-sentence edit to make it safer."""
                 
-                # Use parameter placeholder to prevent syntax errors on raw contract brackets
+                # Parameterized human prompt prevents template syntax errors on contract brackets
                 prompt = ChatPromptTemplate.from_messages([
                     ("system", system_msg),
                     ("human", "{clause_text}")
@@ -183,10 +185,10 @@ async def analyze_document(
                     ai_response = (prompt | llm).invoke({"clause_text": clause})
                     explanation = ai_response.content
                 except Exception as e:
-                    print(f"Groq API call error: {e}")
+                    cause = getattr(e, '__cause__', None)
+                    print(f"❌ Groq API call error: {e} | Underlying cause: {cause}")
                     explanation = "AI Analysis unavailable."
 
-            # Append to results
             results.append({
                 "id": i,
                 "text": clause,