Update ai_mapping.py

ai_mapping.py

@@ -31,20 +31,23 @@ def extract_key_values_with_layoutlm(page_data: list, pdf_path: str) -> Dict[str
         key_values = {}
         # Enhanced regex patterns with flexibility
         dates = re.findall(r'(Agreement\s+(?:Start|End)\s+Date(?:s)?)\s*[:\s]*(\d{1,2}[/-]\d{1,2}[/-]\d{4})', text_data, re.IGNORECASE)
-        # Capture date near "Effective Date" or "executed as of"
-        date_context = re.findall(r'(?:Effective\s+Date|executed\s+as\s+of)\s*[:\s]*(\d{1,2}[/-]\d{1,2}[/-]\d{4})', text_data, re.IGNORECASE)
+        # Capture date anywhere with context like "executed as of" or "Effective Date"
+        date_context = re.findall(r'(?:executed\s+as\s+of|Effective\s+Date)\s*[:\s]*(\d{1,2}[/-]\d{1,2}[/-]\d{4})', text_data, re.IGNORECASE)
         amounts = re.findall(r'\$\d{1,3}(?:,\d{3})*(?:\.\d{2})?', text_data)
-        # Heuristic for Agreement Name near Order Form or titles
-        name_context = re.findall(r'(?:Order\s+Form|Agreement)\s*[:\s]*([A-Za-z0-9\s-]+)(?=\s*(?:Product|Quantity|List|Net))', text_data, re.IGNORECASE)
+        # Refined Agreement Name near Order Form or document start
+        name_context = re.findall(r'(?:Order\s+Form|Agreement)\s*[:\s]*([A-Za-z0-9\s-]+)(?=\s*(?:Product|Quantity|List|Net|\Z))', text_data, re.IGNORECASE)
+        # Prioritize first meaningful name, avoiding procedural text
+        if name_context:
+            key_values["Agreement Name"] = next((name.strip() for name in name_context if len(name.split()) > 1 and not name.lower().startswith("no")), "Unknown")
         # Update key_values with matched fields
         for key, value in dates:
             key_values[key] = value
         if date_context and not key_values.get("Agreement Start Date"):
             key_values["Agreement Start Date"] = date_context[0]
+        if not key_values.get("Agreement End Date") and len(date_context) > 1:
+            key_values["Agreement End Date"] = date_context[1] if date_context[1:] else ""
         if amounts:
             key_values["Amount"] = amounts[0]
-        if name_context:
-            key_values["Agreement Name"] = name_context[0].strip()
 
         # Attempt LayoutLMv3 processing
         doc = fitz.open(pdf_path)
@@ -122,9 +125,9 @@ def extract_clauses(page_data: list) -> Dict[str, str]:
     """
     clauses = {}
     text_data = "\n".join([page["text"] for page in page_data])  # Use newlines for better segmentation
-    # Target "NO WAIVER" specifically under "General Provisions"
-    no_waiver_match = re.search(r'(?:General\s+Provisions\s*[\s\S]*?NO\s+WAIVER\s*[:\s]*)([\s\S]*?)(?=(?:General\s+Provisions|\Z))', text_data, re.IGNORECASE)
-    if no_waiver_match and len(no_waiver_match.group(1).strip().split()) > 5:
+    # Target exact "NO WAIVER" text only
+    no_waiver_match = re.search(r'NO\s+WAIVER\s*[:\s]*(.*?)(?=\n\n|\Z)', text_data, re.IGNORECASE)
+    if no_waiver_match:
         clauses["NO WAIVER"] = no_waiver_match.group(1).strip()
     return clauses if clauses else {}