Update app/services/risk_analyzer.py

app/services/risk_analyzer.py

@@ -3,30 +3,27 @@ from app.config import settings
 from app.services.risk_scorer import RISK_DEFINITIONS
 import json
 import re
-from typing import Dict, List
+from typing import Dict
 
 
-# Configure the Gemini API
+# Configure Gemini API
 genai.configure(api_key=settings.GEMINI_API_KEY)
 
 
 def analyze_clause_with_gemini(clause_text: str) -> Dict:
     """
-    Analyze a contract clause using Google Gemini AI for risk identification.
-
+    Analyze a contract (clause or full text) using Google Gemini AI for risk identification.
     Args:
-        clause_text: The text of the clause to analyze
-
+        clause_text: The text of the clause or contract to analyze
     Returns:
         Dictionary containing identified risk IDs and suggestions
     """
 
-    # Create the detailed prompt for Gemini
     prompt = f"""
 You are an expert Indian legal consultant specializing in contract analysis and risk assessment. 
-Analyze the following contract clause and identify any legal risks based on the predefined risk categories.
+Analyze the following contract text and identify any legal risks based on the predefined categories.
 
-CONTRACT CLAUSE TO ANALYZE:
+CONTRACT TEXT TO ANALYZE:
 {clause_text}
 
 RISK CATEGORIES TO CHECK FOR:
@@ -36,64 +33,59 @@ RISK CATEGORIES TO CHECK FOR:
 4. DPDP_NON_COMPLIANCE: Data protection clause may not comply with the DPDP Act 2023
 
 INSTRUCTIONS:
-1. Carefully read the clause text
-2. Identify which of the above risk categories apply to this clause
-3. For each identified risk, provide a brief explanation
-4. Suggest a compliant alternative or modification for any identified risks
+1. Carefully read the contract text
+2. Identify which of the above risk categories apply
+3. For each identified risk, provide a **short explanation (max 1–2 lines, no long paragraphs)**
+4. Suggest a **brief compliant alternative (1–2 lines)**
 5. If no risks are found, respond with "No risks identified"
 
-RESPONSE FORMAT (JSON):
+RESPONSE FORMAT (strict JSON only):
 {{
     "risks": [
         {{
             "risk_id": "RISK_CATEGORY_ID",
-            "explanation": "Brief explanation of why this risk applies"
+            "explanation": "1–2 line explanation of why this risk applies"
         }}
     ],
-    "suggestion": "Compliant alternative or modification suggestion"
+    "suggestion": "1–2 line compliant alternative or modification suggestion"
 }}
 
 If no risks are identified, return:
 {{
     "risks": [],
-    "suggestion": "No risks identified - clause appears compliant"
+    "suggestion": "No risks identified - contract appears compliant"
 }}
 """
 
     try:
-        # Initialize the Gemini model
-        model = genai.GenerativeModel('gemini-2.5-flash-lite')
+        # Initialize Gemini model
+        model = genai.GenerativeModel("gemini-2.5-flash-lite")
 
         # Generate response
         response = model.generate_content(prompt)
-
-        # Extract the text response
         response_text = response.text.strip()
 
         # Try to parse JSON from the response
         try:
-            # Look for JSON in the response (sometimes Gemini includes extra text)
-            json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
+            json_match = re.search(r"\{.*\}", response_text, re.DOTALL)
             if json_match:
                 json_str = json_match.group()
                 result = json.loads(json_str)
             else:
-                # Fallback: try to parse the entire response as JSON
                 result = json.loads(response_text)
         except json.JSONDecodeError:
-            # If JSON parsing fails, create a fallback response
             result = {
                 "risks": [],
                 "suggestion": "Unable to parse AI response - manual review recommended"
             }
 
-        # Validate and clean the response
+        # Ensure required keys exist
         if "risks" not in result:
             result["risks"] = []
         if "suggestion" not in result:
             result["suggestion"] = "No suggestion provided"
 
-        # Validate risk IDs
+        # Validate risk IDs against known definitions
         valid_risks = []
         for risk in result["risks"]:
             if isinstance(risk, dict) and "risk_id" in risk:
@@ -102,7 +94,6 @@ If no risks are identified, return:
                     valid_risks.append(risk)
 
         result["risks"] = valid_risks
-
         return result
 
     except Exception as e: