Update ai_mapping.py

ai_mapping.py

@@ -35,26 +35,32 @@ def extract_key_values_with_layoutlm(page_data: list, pdf_path: str) -> Dict[str
         text_data = " ".join([page["text"] for page in page_data])
         
         # Refined regex patterns for required fields, avoiding record type as Agreement Name
-        name_context = re.findall(r'(?:Agreement\s+Name|Contract\s+Title)\s*[:\s]*([A-Za-z0-9\s]+?)(?=\s*(?:Exhibit|\n\n|\Z))', text_data, re.IGNORECASE)
+        name_context = re.findall(r'(?:Agreement\s+Name|Contract\s+Title|Agreement\s+Title)\s*[:\s]*([A-Za-z0-9\s]+?)(?=\s*(?:Exhibit|\n\n|\Z))', text_data, re.IGNORECASE)
         if name_context:
             key_values["Agreement Name"] = next((name.strip() for name in name_context if len(name.split()) > 1 and "MASTER SUBSCRIPTION AGREEMENT" not in name.upper()), "Unknown")
-        
-        # Enhanced date patterns to capture context like "executed as of" or specific date labels
+        else:
+            # Fallback to infer name from context if no explicit title
+            party_match = re.search(r'(?:between\s+([A-Za-z\s]+)\s+and)', text_data, re.IGNORECASE)
+            if party_match:
+                key_values["Agreement Name"] = party_match.group(1).strip() or "Unknown"
+
+        # Enhanced date patterns to capture "executed as of" and other date contexts
         date_patterns = [
-            r'(?:Agreement\s+Start\s+Date|Effective\s+Date|executed\s+as\s+of)\s*[:\s]*(\d{1,2}[/-]\d{1,2}[/-]\d{2,4})',
-            r'(?:Agreement\s+End\s+Date|Termination\s+Date)\s*[:\s]*(\d{1,2}[/-]\d{1,2}[/-]\d{2,4})'
+            r'(?:Agreement\s+Start\s+Date|Effective\s+Date|executed\s+as\s+of)\s*[:\s]*(\d{1,2}/\d{1,2}/\d{2,4})',
+            r'(?:Agreement\s+End\s+Date|Termination\s+Date)\s*[:\s]*(\d{1,2}/\d{1,2}/\d{2,4})'
         ]
         for pattern in date_patterns:
             matches = re.findall(pattern, text_data, re.IGNORECASE)
-            for key, value in [("Agreement Start Date", matches[0][0]) if "start" in pattern.lower() or "effective" in pattern.lower() else ("Agreement End Date", matches[0][0]) for matches in [m for m in [re.findall(pattern, text_data, re.IGNORECASE)] if m]]:
+            if matches:
+                key, value = ("Agreement Start Date", matches[0]) if "start" in pattern.lower() or "effective" in pattern.lower() or "executed" in pattern.lower() else ("Agreement End Date", matches[0])
                 if value and not key_values.get(key):
                     key_values[key] = value
 
         # Improved amount pattern to capture total value context
-        amount_pattern = r'(?:Total\s+Agreement\s+Value|Total\s+Amount|Contract\s+Value)\s*[:\s]*\$?\d{1,3}(?:,\d{3})*(?:\.\d{2})?'
+        amount_pattern = r'(?:Total\s+Agreement\s+Value|Total\s+Amount|Contract\s+Value|List\s+Price)\s*[:\s]*\$?\d{1,3}(?:,\d{3})*(?:\.\d{2})?'
         amounts = re.findall(amount_pattern, text_data, re.IGNORECASE)
         if amounts:
-            key_values["Total Agreement Value"] = next((amt.split(":")[-1].strip() if ":" in amt else amt.strip() for amt in amounts if any(k.lower() in amt.lower() for k in ["total", "value"])), "")
+            key_values["Total Agreement Value"] = next((amt.split(":")[-1].strip() if ":" in amt else amt.strip() for amt in amounts if any(k.lower() in amt.lower() for k in ["total", "value", "price"])), "")
 
         # Attempt LayoutLMv3 processing for enhanced extraction
         doc = fitz.open(pdf_path)
@@ -108,11 +114,11 @@ def extract_key_values_with_layoutlm(page_data: list, pdf_path: str) -> Dict[str
                         key = " ".join(current_value).strip()
                         if "agreement name" in current_key.lower() and "MASTER SUBSCRIPTION AGREEMENT" not in key.upper():
                             key_values["Agreement Name"] = key
-                        elif "start date" in current_key.lower() or "effective date" in current_key.lower():
+                        elif "start date" in current_key.lower() or "effective date" in current_key.lower() or "executed as of" in current_key.lower():
                             key_values["Agreement Start Date"] = key
                         elif "end date" in current_key.lower() or "termination date" in current_key.lower():
                             key_values["Agreement End Date"] = key
-                        elif "total agreement value" in current_key.lower() or "amount" in current_key.lower():
+                        elif "total agreement value" in current_key.lower() or "amount" in current_key.lower() or "price" in current_key.lower():
                             key_values["Total Agreement Value"] = key
                     current_key = token
                     current_value = []
@@ -122,11 +128,11 @@ def extract_key_values_with_layoutlm(page_data: list, pdf_path: str) -> Dict[str
                 key = " ".join(current_value).strip()
                 if "agreement name" in current_key.lower() and "MASTER SUBSCRIPTION AGREEMENT" not in key.upper():
                     key_values["Agreement Name"] = key
-                elif "start date" in current_key.lower() or "effective date" in current_key.lower():
+                elif "start date" in current_key.lower() or "effective date" in current_key.lower() or "executed as of" in current_key.lower():
                     key_values["Agreement Start Date"] = key
                 elif "end date" in current_key.lower() or "termination date" in current_key.lower():
                     key_values["Agreement End Date"] = key
-                elif "total agreement value" in current_key.lower() or "amount" in current_key.lower():
+                elif "total agreement value" in current_key.lower() or "amount" in current_key.lower() or "price" in current_key.lower():
                     key_values["Total Agreement Value"] = key
 
             # Clean up temporary image
@@ -140,7 +146,7 @@ def extract_key_values_with_layoutlm(page_data: list, pdf_path: str) -> Dict[str
 
 def extract_clauses(page_data: list) -> Dict[str, str]:
     """
-    Extract clauses from PDF text based on keywords, focusing on key clauses like NO WAIVER.
+    Extract clauses from PDF text based on keywords, focusing on key clauses like NO WAIVER and Termination.
     Args:
         page_data (list): List of dictionaries with 'text' (str) per page.
     Returns:
@@ -155,9 +161,9 @@ def extract_clauses(page_data: list) -> Dict[str, str]:
         clause_text = no_waiver_match.group(1).strip()
         clauses["NO WAIVER"] = clause_text if clause_text else "NO WAIVER clause found but no content extracted"
     elif "NO WAIVER" in text_data.upper():
-        clauses["NO WAIVER"] = "NO WAIVER clause identified but no detailed content extracted"
+        clauses["NO WAIVER"] = re.search(r'(NO\s+WAIVER\s*[:\s]*[\s\S]*?)(?=\n\n|\Z)', text_data, re.IGNORECASE).group(1).strip() if re.search(r'(NO\s+WAIVER\s*[:\s]*[\s\S]*?)(?=\n\n|\Z)', text_data, re.IGNORECASE) else "NO WAIVER clause identified but no detailed content extracted"
     
-    # Add Termination clause
+    # Search for Termination clause
     termination_match = re.search(r'(?:Termination\s*[:\s]*)([\s\S]*?)(?=\n\n|\Z)', text_data, re.IGNORECASE)
     if termination_match:
         clauses["Termination"] = termination_match.group(1).strip()