aibridze/document_intelligence
0
1# """2# Entity Extractor — type-aware extraction using Gemini + LangChain.3# Embeds the exact JSON schema for each contract type directly in the prompt.4# Extracts structured data + handwritten content (signatures, stamps, notes).5# Production-ready: aggressive field filling, robust JSON parsing, retry logic.6# """7 8# from langchain_google_genai import ChatGoogleGenerativeAI9# from langchain_core.prompts import ChatPromptTemplate10# import logging11# import json12# import re13 14# from app.models.extraction_schemas import (15# ContractType,16# EXTRACTION_MODEL_MAP,17# NDAExtraction,18# ServiceProviderExtraction,19# SupplyCumServiceExtraction,20# ConsultancyExtraction,21# CustomerSupplyExtraction,22# )23# from app.config import get_settings24 25# settings = get_settings()26# logger = logging.getLogger("trident_poc")27 28 29# # ─── JSON Schema Templates (embedded in prompts) ─────────────────30 31# NDA_SCHEMA = """{32# "agreement_type": "NDA",33# "parties": {34# "disclosing_party": "",35# "receiving_party": ""36# },37# "dates": {38# "agreement_date": "",39# "effective_date": ""40# },41# "definitions": {42# "confidential_information": ""43# },44# "purpose": "",45# "obligations": {46# "recipient_obligations": "",47# "disclosing_party_rights": ""48# },49# "exclusions": [],50# "ownership": "",51# "term": {52# "duration": "",53# "termination_conditions": ""54# },55# "indemnity": "",56# "governing_law": "",57# "miscellaneous": [],58# "handwritten_data": {59# "signatures": [{"party": "", "name": "", "designation": "", "date": ""}],60# "notes": [],61# "stamps": [],62# "unstructured_handwriting": ""63# }64# }"""65 66# SERVICE_PROVIDER_SCHEMA = """{67# "agreement_type": "Service Provider Agreement",68# "parties": {69# "company": "",70# "service_provider": ""71# },72# "dates": {73# "agreement_date": "",74# "start_date": "",75# "end_date": ""76# },77# "scope_of_work": "",78# "services": [],79# "deliverables": [],80# "kpis_or_sla": [],81# "payment_terms": {82# "payment_cycle": "",83# "billing_details": "",84# "rate_details": "",85# "penalties": ""86# },87# "security_deposit": "",88# "tenure": "",89# "obligations": {90# "service_provider": [],91# "company": []92# },93# "compliance": [],94# "termination": "",95# "annexures": [96# {97# "title": "",98# "table_data": []99# }100# ],101# "handwritten_data": {102# "signatures": [{"party": "", "name": "", "designation": "", "date": ""}],103# "notes": [],104# "stamps": [],105# "unstructured_handwriting": ""106# }107# }"""108 109# SUPPLY_CUM_SERVICE_SCHEMA = """{110# "agreement_type": "Supply Cum Service Agreement",111# "parties": {112# "company": "",113# "service_provider": ""114# },115# "project_details": {116# "project_name": "",117# "location": "",118# "description": ""119# },120# "dates": {121# "agreement_date": "",122# "start_date": "",123# "end_date": ""124# },125# "scope_of_work": {126# "engineering": "",127# "supply": "",128# "installation": "",129# "commissioning": ""130# },131# "technical_specifications": [132# {133# "component": "",134# "details": {}135# }136# ],137# "pricing": {138# "total_value": "",139# "currency": "",140# "breakdown": []141# },142# "payment_terms": {143# "milestones": [],144# "advance": "",145# "final_payment": ""146# },147# "bank_guarantee": {148# "required": "",149# "details": ""150# },151# "warranty": {152# "duration": "",153# "terms": ""154# },155# "performance_terms": [],156# "penalties": [],157# "compliance": [],158# "intellectual_property": "",159# "termination": "",160# "annexures": [161# {162# "title": "",163# "tables": []164# }165# ],166# "handwritten_data": {167# "signatures": [{"party": "", "name": "", "designation": "", "date": ""}],168# "notes": [],169# "stamps": [],170# "unstructured_handwriting": ""171# }172# }"""173 174# CONSULTANCY_SCHEMA = """{175# "agreement_type": "Consultancy Agreement",176# "parties": {177# "company": "",178# "consultant": ""179# },180# "dates": {181# "agreement_date": "",182# "start_date": "",183# "end_date": ""184# },185# "scope": "",186# "deliverables": [],187# "timelines": {188# "meeting_frequency": "",189# "milestones": []190# },191# "consideration": {192# "amount": "",193# "currency": "",194# "payment_frequency": ""195# },196# "payment_terms": "",197# "reimbursements": [],198# "taxation": "",199# "obligations": {200# "consultant": [],201# "company": []202# },203# "confidentiality": "",204# "termination": "",205# "validity": "",206# "handwritten_data": {207# "signatures": [{"party": "", "name": "", "designation": "", "date": ""}],208# "notes": [],209# "stamps": [],210# "unstructured_handwriting": ""211# }212# }"""213 214# CUSTOMER_SUPPLY_SCHEMA = """{215# "agreement_type": "Supply Agreement",216# "parties": {217# "buyer": "",218# "supplier": ""219# },220# "dates": {221# "agreement_date": "",222# "effective_date": "",223# "termination_date": ""224# },225# "recitals": [],226# "products_or_services": [],227# "orders_and_delivery": {228# "process": "",229# "timelines": "",230# "minimum_quantity": ""231# },232# "pricing": {233# "price_details": "",234# "currency": "",235# "invoicing": ""236# },237# "payment_terms": {238# "payment_cycle": "",239# "conditions": ""240# },241# "intellectual_property": {242# "ownership": "",243# "usage_rights": ""244# },245# "warranties": [],246# "liabilities": [],247# "indemnity": "",248# "termination": "",249# "governing_law": "",250# "miscellaneous": [],251# "handwritten_data": {252# "signatures": [{"party": "", "name": "", "designation": "", "date": ""}],253# "notes": [],254# "stamps": [],255# "unstructured_handwriting": ""256# }257# }"""258 259# # Map contract type → JSON schema260# TYPE_SCHEMA_MAP = {261# ContractType.NDA: NDA_SCHEMA,262# ContractType.SERVICE_PROVIDER: SERVICE_PROVIDER_SCHEMA,263# ContractType.SUPPLY_CUM_SERVICE: SUPPLY_CUM_SERVICE_SCHEMA,264# ContractType.CONSULTANCY: CONSULTANCY_SCHEMA,265# ContractType.CUSTOMER_SUPPLY: CUSTOMER_SUPPLY_SCHEMA,266# }267 268# # ─── Type-specific extraction instructions ────────────────────────269# TYPE_INSTRUCTIONS = {270# ContractType.NDA: """271# TYPE-SPECIFIC RULES FOR NDA:272# - Extract BOTH parties clearly as disclosing_party and receiving_party with full legal names and registered addresses.273# - For "definitions.confidential_information", extract the COMPLETE definition verbatim from the document.274# - For "exclusions", list EVERY exclusion from confidentiality (publicly available info, independently developed, prior knowledge, court orders, etc.)275# - For "obligations", extract both the recipient's obligations AND the disclosing party's retained rights in detail.276# - Extract "ownership" — who retains ownership of the confidential information.277# - Extract exact "term.duration" (e.g., "2 years from the Effective Date") and "term.termination_conditions".278# - For "purpose", describe WHY the NDA is being entered into.279# - For "indemnity", extract the full indemnification clause.280# - For "governing_law", extract the jurisdiction and dispute resolution mechanism.281# - For "miscellaneous", capture any other clauses like notices, amendments, severability, entire agreement.282# """,283 284# ContractType.SERVICE_PROVIDER: """285# TYPE-SPECIFIC RULES FOR SERVICE PROVIDER AGREEMENT:286# - Extract the full company name (with CIN/registered office) and service provider name under "parties".287# - For "scope_of_work", provide a comprehensive description of ALL work to be performed.288# - For "services", list EVERY service mentioned (loading, unloading, manpower deployment, housekeeping, etc.)289# - For "kpis_or_sla", extract ALL KPIs, targets, performance metrics, and SLA terms.290# - Extract the complete rate slab / rate table under "payment_terms.rate_details" — include all tiers and rates.291# - For "payment_terms.billing_details", include billing cycle, invoice submission process, payment timelines.292# - For "obligations", list obligations of BOTH the service provider AND the company separately — be thorough.293# - For "compliance", list ALL statutory/regulatory compliance requirements (PF, ESI, labour laws, etc.)294# - For "security_deposit", extract amount, bank guarantee details, and refund conditions.295# - For "annexures", extract ALL annexure/schedule tables (deployment schedule, rate tables, manpower details) — use "table_data" as a list of row objects.296# - For "termination", extract notice period, grounds for termination, and exit procedures.297# """,298 299# ContractType.SUPPLY_CUM_SERVICE: """300# TYPE-SPECIFIC RULES FOR SUPPLY CUM SERVICE AGREEMENT:301# - This is a complex EPC/engineering agreement — be VERY thorough and extract EVERYTHING.302# - Extract "project_details" — full project name, site location, and technical description.303# - For "scope_of_work", fill EVERY sub-field: engineering design, supply of materials, installation work, commissioning and testing.304# - For "technical_specifications", extract EACH component (panels, inverters, transformers, cables, structures, etc.) with detailed specs in the "details" dict (brand, model, capacity, rating, quantity, etc.)305# - Extract "pricing.total_value" with exact amount and "pricing.currency". For "pricing.breakdown", include supply vs. erection vs. testing split.306# - For "payment_terms.milestones", extract ALL milestone stages with percentage, trigger condition, and timeline.307# - Extract "bank_guarantee" terms — type, amount, validity period, conditions for invocation.308# - For "warranty", extract duration per component and terms of warranty/AMC.309# - For "performance_terms", extract performance ratios, degradation rates, generation guarantees, grid availability.310# - For "penalties", extract liquidated damages, delay penalties, performance shortfall penalties.311# - For "annexures", extract ALL technical annexures, BOQ tables, schedule of quantities.312# """,313 314# ContractType.CONSULTANCY: """315# TYPE-SPECIFIC RULES FOR CONSULTANCY AGREEMENT:316# - Extract the full company name (including CIN, registered office address) AND the consultant's full name and details under "parties".317# - For "scope", provide a DETAILED description of the consultancy services — areas of expertise, advisory domains, governance, skill development, etc. If there are multiple areas, combine them all into one comprehensive description.318# - For "deliverables", list EVERY deliverable mentioned — reports, recommendations, training programs, assessments, etc. If no explicit deliverables list, extract from scope of work clauses.319# - For "timelines.meeting_frequency", extract how often the consultant meets with the company (e.g., "quarterly", "as needed", "monthly").320# - For "timelines.milestones", extract any phase/milestone information.321# - For "consideration.amount", extract the EXACT fee amount (e.g., "Rs. 1,50,000/- per month").322# - For "consideration.currency", extract currency (e.g., "INR", "USD").323# - For "consideration.payment_frequency", extract payment schedule (e.g., "monthly", "quarterly").324# - For "payment_terms", extract invoice submission, payment timeline, bank details requirements.325# - For "reimbursements", extract ALL reimbursable expenses — travel, lodging, boarding, conveyance, etc. with limits.326# - For "taxation", extract TDS, GST, and any tax-related clauses.327# - For "obligations.consultant", list ALL consultant obligations — deliverables, availability, confidentiality, non-compete, etc.328# - For "obligations.company", list ALL company obligations — providing access, timely payments, information sharing, etc.329# - For "confidentiality", extract the full confidentiality clause.330# - For "termination", extract notice period, grounds for termination, and post-termination obligations.331# - For "validity", extract the agreement duration/validity period.332# """,333 334# ContractType.CUSTOMER_SUPPLY: """335# TYPE-SPECIFIC RULES FOR CUSTOMER/SUPPLY AGREEMENT:336# - Extract buyer and supplier as full legal names with registered addresses.337# - For "recitals", list the WHEREAS clauses / background context of the agreement.338# - For "products_or_services", extract EACH product/service as an object: {"product_name": "", "description": "", "quantity": "", "dimensions": "", "mrp": "", "price": ""}339# - For "orders_and_delivery", extract the ordering process, delivery timelines, delivery location, and any minimum order quantities.340# - For "pricing", extract price details, currency (INR/USD), and invoicing terms.341# - For "payment_terms", extract payment cycle (advance/credit) and all conditions.342# - For "intellectual_property", extract IP ownership terms and usage rights scope in detail.343# - For "warranties", extract ALL warranty/guarantee clauses as separate list items.344# - For "liabilities", extract ALL liability/limitation clauses.345# - For "indemnity", extract the full indemnification clause.346# - For "termination", extract notice period, grounds, and consequences.347# - For "governing_law", extract jurisdiction and dispute resolution.348# - For "miscellaneous", capture any other clauses (force majeure, notices, amendments, severability, entire agreement, etc.)349# """,350# }351 352 353# # ─── Base Extraction Prompt ───────────────────────────────────────354# BASE_EXTRACTION_PROMPT = """You are an expert legal contract analyst with decades of experience. You are analyzing a {contract_type_label}.355 356# Your task is to extract ALL information from the document into the JSON schema provided. You must be EXTREMELY thorough — your extraction quality directly determines the confidence score.357 358# CRITICAL EXTRACTION RULES:359# 1. Fill EVERY field with data from the document. Be aggressive — extract as much as possible.360# 2. For string fields: If direct information is available, extract it verbatim. If only related/partial information is available, extract what you can and annotate. NEVER leave a field empty string "" if there is ANY relevant information in the document.361# 3. For list fields: Include ALL items found. Scan the ENTIRE document for relevant items. Do NOT summarize — be comprehensive.362# 4. Preserve EXACT dates, names, amounts, percentages, and legal terms as written in the document.363# 5. For fields where the document uses different terminology, map the content to the correct field. For example, "Consideration" in the document maps to "consideration.amount".364# 6. If a field truly has NO relevant data in the document, set it to null.365# 7. Do NOT hallucinate or fabricate data — only extract what is explicitly stated or clearly implied.366# 8. For "annexures" / "miscellaneous", capture any additional terms, schedules, or tables.367# 9. Read EVERY page of the document — do not stop after the first few pages.368 369# FIELD FILLING STRATEGY:370# - "parties": Always extract with FULL legal names and addresses.371# - "dates": Extract ALL dates mentioned — agreement date, effective date, start/end dates.372# - "obligations": Always fill for BOTH parties — even if briefly mentioned.373# - "termination": Always extract — every contract has termination terms.374# - "payment_terms": Fill all sub-fields — cycle, billing, rates, penalties.375# - "scope"/"scope_of_work": Provide comprehensive description — combine multiple clauses if needed.376 377# HANDWRITTEN DATA EXTRACTION:378# 10. Extract ALL handwritten content found in the document:379# - "signatures": Extract each signatory with party name, person name, designation, and signing date.380# - "notes": Any handwritten annotations, margin notes, or corrections.381# - "stamps": Text from rubber stamps, company seals, or revenue stamps.382# - "unstructured_handwriting": Any other handwritten text that doesn't fit above categories.383# If handwritten content is unclear, include the raw text as best as possible.384# If no handwritten content is detected, set signatures to empty list and other fields to null.385 386# {type_instructions}387 388# OUTPUT FORMAT:389# Respond ONLY with a valid JSON object matching this EXACT schema. 390# - No markdown backticks, no ```json prefix, no explanation text.391# - Just the raw JSON object starting with {{ and ending with }}.392 393# JSON SCHEMA:394# {json_schema}395 396# COMPLETE DOCUMENT TEXT:397# {document_text}398# """399 400 401# class EntityExtractor:402 403# def __init__(self):404# if not settings.GOOGLE_API_KEY:405# raise ValueError("GOOGLE_API_KEY not configured")406 407# self.llm = ChatGoogleGenerativeAI(408# model=settings.VISION_MODEL,409# google_api_key=settings.GOOGLE_API_KEY,410# temperature=0,411# max_output_tokens=8192,412# )413# self.prompt = ChatPromptTemplate.from_template(BASE_EXTRACTION_PROMPT)414 415# async def extract(416# self,417# document_text: str,418# contract_type: ContractType = None,419# ) -> dict:420# """421# Extract structured data from document text using type-specific JSON schema.422# Returns a validated dict matching the type's Pydantic model.423# """424# if contract_type is None:425# contract_type = ContractType.NDA426 427# type_label = contract_type.value428# json_schema = TYPE_SCHEMA_MAP.get(contract_type, NDA_SCHEMA)429# type_instructions = TYPE_INSTRUCTIONS.get(contract_type, "")430# model_cls = EXTRACTION_MODEL_MAP.get(contract_type, NDAExtraction)431 432# # Truncate very large documents to avoid token limit issues433# MAX_TEXT_LENGTH = 200_000434# if len(document_text) > MAX_TEXT_LENGTH:435# logger.warning(436# f"Document text truncated from {len(document_text)} to {MAX_TEXT_LENGTH} chars"437# )438# document_text = document_text[:MAX_TEXT_LENGTH]439 440# chain = self.prompt | self.llm441 442# try:443# response = await chain.ainvoke({444# "contract_type_label": type_label,445# "type_instructions": type_instructions,446# "json_schema": json_schema,447# "document_text": document_text,448# })449 450# # Parse JSON response451# content = response.content.strip()452# parsed = self._parse_json_response(content)453 454# except Exception as e:455# logger.error(f"LLM extraction failed: {e}", exc_info=True)456# # Return a minimal valid response instead of crashing457# parsed = {"agreement_type": type_label, "error": str(e)}458 459# # Validate against Pydantic model460# try:461# validated = model_cls.model_validate(parsed)462# return validated.model_dump()463# except Exception as e:464# logger.warning(f"Pydantic validation failed, returning raw parsed JSON: {e}")465# return parsed466 467# async def extract_with_vision_context(468# self,469# document_text: str,470# vision_analysis: str,471# contract_type: ContractType = None,472# ) -> dict:473# """Extract with additional vision OCR context for scanned docs."""474# combined_text = f"""{document_text}475 476# === VISION MODEL ANALYSIS (Scanned/Handwritten Content) ===477# {vision_analysis}478# """479# return await self.extract(combined_text, contract_type)480 481# @staticmethod482# def _parse_json_response(content: str) -> dict:483# """484# Robust JSON parsing from LLM response.485# Handles markdown fences, extra text, and partial JSON.486# """487# original = content488 489# # Strip markdown code fences (```json ... ```)490# content = re.sub(r'^```(?:json)?\s*\n?', '', content, flags=re.MULTILINE)491# content = re.sub(r'\n?```\s*$', '', content, flags=re.MULTILINE)492# content = content.strip()493 494# # Attempt 1: Direct parse495# try:496# return json.loads(content)497# except json.JSONDecodeError:498# pass499 500# # Attempt 2: Find the outermost JSON object501# brace_start = content.find("{")502# brace_end = content.rfind("}") + 1503# if brace_start >= 0 and brace_end > brace_start:504# try:505# return json.loads(content[brace_start:brace_end])506# except json.JSONDecodeError:507# pass508 509# # Attempt 3: Fix common JSON issues (trailing commas, single quotes)510# cleaned = content511# if brace_start >= 0 and brace_end > brace_start:512# cleaned = content[brace_start:brace_end]513# # Remove trailing commas before } or ]514# cleaned = re.sub(r',\s*([\]}])', r'\1', cleaned)515# try:516# return json.loads(cleaned)517# except json.JSONDecodeError:518# pass519 520# logger.error(521# f"All JSON parse attempts failed.\n"522# f"Raw content (first 1000 chars): {original[:1000]}"523# )524# return {"agreement_type": "Unknown", "error": "Failed to parse LLM response"}525 526"""527Entity Extractor — Production-grade multi-pass extraction.528 529Pipeline:530 Pass 1: Full extraction with type-specific JSON schema531 Pass 2: Gap-fill — re-extract only null/empty critical fields532 Pass 3: Validate — cross-check extracted values against source text533 534Features:535 - Retry with exponential backoff (3 attempts)536 - Section-based extraction for documents >15 pages537 - OCR text cleaning (remove artifacts, normalize whitespace)538 - Robust JSON parsing with multiple fallback strategies539 - Per-field confidence scoring540"""541 542from langchain_google_genai import ChatGoogleGenerativeAI543from langchain_core.prompts import ChatPromptTemplate544import logging545import json546import re547import asyncio548from typing import Dict, Any, List, Optional, Tuple549 550from app.models.extraction_schemas import ContractType551from app.config import get_settings552 553settings = get_settings()554logger = logging.getLogger("trident_poc")555 556MAX_TEXT_LENGTH = 200_000557MAX_RETRIES = 3558RETRY_BASE_DELAY = 2 # seconds559SECTION_PAGE_THRESHOLD = 15 # Split into sections if more than this many pages560 561 562# ═══════════════════════════════════════════════════════════════════563# JSON SCHEMAS (one per contract type)564# ═══════════════════════════════════════════════════════════════════565 566NDA_SCHEMA = """{567 "agreement_type": "NDA",568 "agreement_date": null,569 "parties": {570 "disclosing_party": null,571 "receiving_party": null572 },573 "dates": {},574 "purpose": null,575 "confidential_information_summary": null,576 "obligations_summary": null,577 "term": {578 "duration": null579 },580 "summary": null,581 "confidence_score": null582}"""583 584SERVICE_PROVIDER_SCHEMA = """{585 "agreement_type": "Service Provider Agreement",586 "agreement_date": null,587 "parties": {588 "company": null,589 "service_provider": null590 },591 "dates": {592 "start_date": null,593 "end_date": null594 },595 "scope_of_work": null,596 "services": [],597 "payment_terms": {598 "payment_cycle": null,599 "rate_details": null600 },601 "security_deposit": null,602 "key_obligations": [],603 "termination_summary": null,604 "summary": null,605 "confidence_score": null606}"""607 608SUPPLY_CUM_SERVICE_SCHEMA = """{609 "agreement_type": "Supply Cum Service Agreement",610 "agreement_date": null,611 "parties": {612 "company": null,613 "service_provider": null614 },615 "project_details": {616 "project_name": null,617 "location": null618 },619 "dates": {620 "start_date": null,621 "end_date": null622 },623 "scope_summary": null,624 "pricing": {625 "total_value": null626 },627 "payment_terms": {628 "milestones": []629 },630 "warranty": {631 "duration": null632 },633 "penalties_summary": null,634 "termination_summary": null,635 "summary": null,636 "confidence_score": null637}"""638 639CONSULTANCY_SCHEMA = """{640 "agreement_type": "Consultancy Agreement",641 "agreement_date": null,642 "parties": {643 "company": null,644 "consultant": null645 },646 "dates": {647 "start_date": null,648 "end_date": null649 },650 "scope": null,651 "deliverables": [],652 "consideration": {653 "amount": null,654 "payment_frequency": null655 },656 "key_obligations": [],657 "termination_summary": null,658 "summary": null,659 "confidence_score": null660}"""661 662CUSTOMER_SUPPLY_SCHEMA = """{663 "agreement_type": "Supply Agreement",664 "agreement_date": null,665 "parties": {666 "buyer": null,667 "supplier": null668 },669 "dates": {670 "termination_date": null671 },672 "products_or_services": [],673 "delivery_terms": null,674 "pricing": {675 "price_details": null676 },677 "payment_terms": {678 "payment_cycle": null679 },680 "key_obligations": [],681 "termination_summary": null,682 "summary": null,683 "confidence_score": null684}"""685 686TYPE_SCHEMA_MAP = {687 ContractType.NDA: NDA_SCHEMA,688 ContractType.SERVICE_PROVIDER: SERVICE_PROVIDER_SCHEMA,689 ContractType.SUPPLY_CUM_SERVICE: SUPPLY_CUM_SERVICE_SCHEMA,690 ContractType.CONSULTANCY: CONSULTANCY_SCHEMA,691 ContractType.CUSTOMER_SUPPLY: CUSTOMER_SUPPLY_SCHEMA,692}693 694 695# ═══════════════════════════════════════════════════════════════════696# TYPE-SPECIFIC INSTRUCTIONS697# ═══════════════════════════════════════════════════════════════════698 699TYPE_INSTRUCTIONS = {700 ContractType.NDA: """701NDA-SPECIFIC RULES:702- "parties.disclosing_party" and "parties.receiving_party": Extract the FULL LEGAL NAME of each party (company name, not individual).703- "agreement_date": The date the agreement was signed/executed.704- "purpose": WHY the NDA is being entered into — brief description.705- "confidential_information_summary": Summarize what constitutes confidential information in 2-3 sentences.706- "obligations_summary": Summarize the key obligations of the receiving party in 2-3 sentences.707- "term.duration": Exact duration (e.g., "2 years from the Effective Date").708- "summary": A 1-2 sentence overview of the entire agreement.709- "confidence_score": Your self-assessed confidence (0.0 to 1.0) in the accuracy of this extraction.710""",711 712 ContractType.SERVICE_PROVIDER: """713SERVICE PROVIDER RULES:714- "parties.company": Full legal name of the hiring company.715- "parties.service_provider": Full legal name of the service provider/contractor.716- "agreement_date": The date the agreement was signed/executed.717- "dates.start_date" and "dates.end_date": Contract validity period.718- "scope_of_work": Comprehensive description of ALL work to be performed.719- "services": List EVERY service mentioned (loading, unloading, manpower, housekeeping, etc.).720- "payment_terms.payment_cycle": Billing frequency (monthly, fortnightly, etc.).721- "payment_terms.rate_details": Complete rate information or rate slab summary.722- "security_deposit": Amount and terms.723- "key_obligations": List the TOP 5-8 most important obligations of both parties combined.724- "termination_summary": Summarize termination terms (notice period, grounds) in 1-2 sentences.725- "summary": A 1-2 sentence overview of the entire agreement.726- "confidence_score": Your self-assessed confidence (0.0 to 1.0) in the accuracy.727""",728 729 ContractType.SUPPLY_CUM_SERVICE: """730SUPPLY CUM SERVICE (EPC) RULES:731- "parties.company": Full legal name of the company.732- "parties.service_provider": Full legal name of the EPC/service provider.733- "agreement_date": The date the agreement was signed/executed.734- "project_details.project_name": Full project name.735- "project_details.location": Project site location.736- "dates.start_date" and "dates.end_date": Project timeline.737- "scope_summary": Summarize the full scope (engineering, supply, installation, commissioning) in 3-5 sentences.738- "pricing.total_value": Exact total contract value with currency (e.g., "INR 2,50,00,000/-").739- "payment_terms.milestones": List key payment milestones (e.g., "30% on order placement", "40% on delivery", "30% on commissioning").740- "warranty.duration": Warranty period (e.g., "5 years from commissioning").741- "penalties_summary": Summarize penalty/liquidated damages terms in 1-2 sentences.742- "termination_summary": Summarize termination terms in 1-2 sentences.743- "summary": A 1-2 sentence overview of the entire agreement.744- "confidence_score": Your self-assessed confidence (0.0 to 1.0) in the accuracy.745""",746 747 ContractType.CONSULTANCY: """748CONSULTANCY RULES:749- "parties.company": Full legal name of the hiring company.750- "parties.consultant": Full legal name of the consultant/advisory firm.751- "agreement_date": The date the agreement was signed/executed.752- "dates.start_date" and "dates.end_date": Consultancy period.753- "scope": Detailed description of consultancy services — governance, advisory, skill development, etc.754- "deliverables": List ALL deliverables mentioned — reports, recommendations, training, etc.755- "consideration.amount": EXACT fee amount (e.g., "INR 5,00,000/- per month").756- "consideration.payment_frequency": Payment schedule ("Monthly", "Quarterly", etc.).757- "key_obligations": List the TOP 5-8 most important obligations of both parties combined.758- "termination_summary": Summarize termination terms (notice period, grounds) in 1-2 sentences.759- "summary": A 1-2 sentence overview of the entire agreement.760- "confidence_score": Your self-assessed confidence (0.0 to 1.0) in the accuracy.761""",762 763 ContractType.CUSTOMER_SUPPLY: """764CUSTOMER/SUPPLY AGREEMENT RULES:765- "parties.buyer": Full legal name of the buyer.766- "parties.supplier": Full legal name of the supplier.767- "agreement_date": The date the agreement was signed/executed.768- "dates.termination_date": When the agreement expires.769- "products_or_services": List the key products/services being supplied (as simple strings, e.g., "Solar panels - 500 units").770- "delivery_terms": Summarize delivery terms (location, timelines, conditions) in 1-2 sentences.771- "pricing.price_details": Summary of pricing terms.772- "payment_terms.payment_cycle": Payment frequency/terms.773- "key_obligations": List the TOP 5-8 most important obligations of both parties combined.774- "termination_summary": Summarize termination terms in 1-2 sentences.775- "summary": A 1-2 sentence overview of the entire agreement.776- "confidence_score": Your self-assessed confidence (0.0 to 1.0) in the accuracy.777""",778}779 780 781# ═══════════════════════════════════════════════════════════════════782# PROMPTS783# ═══════════════════════════════════════════════════════════════════784 785EXTRACTION_PROMPT = """You are an expert legal contract analyst. You are analyzing a {contract_type_label}.786 787TASK: Extract ALL relevant information from the document into the JSON schema below. Be ACCURATE and THOROUGH.788 789RULES:7901. Fill EVERY field with data from the document. Search the ENTIRE document.7912. Preserve EXACT dates, names, amounts, percentages as written in the document.7923. For party names: extract the FULL LEGAL ENTITY NAME (e.g., "Trident Limited" not just "Trident").7934. For lists: include ALL relevant items found across the entire document.7945. Map document terminology to schema fields (e.g., "Consideration" → "consideration.amount").7956. If a field truly has no data in the document, set it to null.7967. Do NOT hallucinate or fabricate — only extract what is explicitly stated or clearly implied.7978. Read EVERY page — critical data is often in later pages and annexures.7989. "summary": Write a concise 1-2 sentence overview of the agreement covering parties, purpose, and key commercial terms.79910. "confidence_score": Rate your extraction confidence from 0.0 to 1.0 based on how much information you found and how clearly it matched the schema fields.800 801{type_instructions}802 803JSON SCHEMA TO FILL:804{json_schema}805 806DOCUMENT TEXT:807{document_text}808 809Respond ONLY with a valid JSON object. No markdown, no backticks, no explanation."""810 811 812GAP_FILL_PROMPT = """You are an expert legal contract analyst. A previous extraction from this document left some fields empty.813 814Your task: Extract ONLY the missing fields listed below from the document text. Be thorough — the data IS in the document, it was just missed.815 816MISSING FIELDS TO FILL:817{missing_fields}818 819RULES:8201. Only return values for the fields listed above.8212. Search the ENTIRE document text carefully — check every page.8223. Preserve exact values as written in the document.8234. Return a flat JSON object mapping field paths to values.824 825Example response format:826{{827 "parties.company.name": "ABC Corporation Ltd",828 "consideration.amount": "INR 1,00,000/- per month",829 "dates.start_date": "November 19, 2020"830}}831 832DOCUMENT TEXT:833{document_text}834 835Respond ONLY with a valid JSON object. No markdown, no backticks."""836 837 838SECTION_EXTRACTION_PROMPT = """You are an expert legal contract analyst extracting data from a SECTION of a {contract_type_label}.839 840This is pages {page_range} of the document. Extract ALL relevant information from this section.841 842FOCUS AREAS FOR THIS SECTION:843{section_focus}844 845JSON SCHEMA (fill what you find in this section):846{json_schema}847 848SECTION TEXT:849{section_text}850 851Respond ONLY with a valid JSON object containing fields found in this section. No markdown, no backticks."""852 853 854# ═══════════════════════════════════════════════════════════════════855# CRITICAL FIELDS PER TYPE (used for gap detection)856# ═══════════════════════════════════════════════════════════════════857 858CRITICAL_FIELDS = {859 ContractType.NDA: [860 "parties.disclosing_party", "parties.receiving_party",861 "agreement_date", "purpose", "confidential_information_summary",862 "obligations_summary", "term.duration", "summary",863 ],864 ContractType.SERVICE_PROVIDER: [865 "parties.company", "parties.service_provider",866 "agreement_date", "scope_of_work", "services",867 "payment_terms.rate_details", "termination_summary", "summary",868 ],869 ContractType.SUPPLY_CUM_SERVICE: [870 "parties.company", "parties.service_provider",871 "agreement_date", "project_details.project_name",872 "scope_summary", "pricing.total_value",873 "warranty.duration", "summary",874 ],875 ContractType.CONSULTANCY: [876 "parties.company", "parties.consultant",877 "agreement_date", "scope", "deliverables",878 "consideration.amount", "termination_summary", "summary",879 ],880 ContractType.CUSTOMER_SUPPLY: [881 "parties.buyer", "parties.supplier",882 "agreement_date", "products_or_services",883 "pricing.price_details", "payment_terms.payment_cycle",884 "termination_summary", "summary",885 ],886}887 888# Section focus areas for chunked extraction (by page position)889SECTION_FOCUS_MAP = {890 "early": "Focus on: document title, agreement type, parties (names, addresses, CIN, GSTIN), execution date, effective date, recitals/whereas clauses, definitions.",891 "middle": "Focus on: scope of work, deliverables, services, technical specifications, financial terms (pricing, rates, amounts), payment terms, milestones, KPIs, timelines, obligations of each party.",892 "late": "Focus on: legal clauses (termination, indemnity, governing law, dispute resolution, force majeure, confidentiality, IP rights), compliance requirements, annexures, schedules, tables, SPOC details, signature details.",893}894 895 896# ═══════════════════════════════════════════════════════════════════897# ENTITY EXTRACTOR CLASS898# ═══════════════════════════════════════════════════════════════════899 900class EntityExtractor:901 902 def __init__(self):903 if not settings.GOOGLE_API_KEY:904 raise ValueError("GOOGLE_API_KEY not configured")905 906 self.llm = ChatGoogleGenerativeAI(907 model=settings.VISION_MODEL,908 google_api_key=settings.GOOGLE_API_KEY,909 temperature=0,910 max_output_tokens=8192,911 )912 self.extraction_prompt = ChatPromptTemplate.from_template(EXTRACTION_PROMPT)913 self.gap_fill_prompt = ChatPromptTemplate.from_template(GAP_FILL_PROMPT)914 self.section_prompt = ChatPromptTemplate.from_template(SECTION_EXTRACTION_PROMPT)915 916 # ─── Main Entry Point ─────────────────────────────────────────917 918 async def extract(919 self,920 document_text: str,921 contract_type: ContractType = None,922 ) -> dict:923 """924 Production-grade multi-pass extraction.925 926 1. Clean the OCR text927 2. For large docs: section-based extraction + merge928 For small docs: single-pass extraction929 3. Gap-fill pass for any missing critical fields930 4. Return enriched result931 """932 if contract_type is None:933 contract_type = ContractType.NDA934 935 # Step 0: Clean text936 cleaned_text = self._clean_ocr_text(document_text)937 938 if len(cleaned_text) > MAX_TEXT_LENGTH:939 logger.warning(f"Text truncated from {len(cleaned_text)} to {MAX_TEXT_LENGTH}")940 cleaned_text = cleaned_text[:MAX_TEXT_LENGTH]941 942 # Count pages for decision943 page_count = cleaned_text.count("--- Page")944 945 logger.info(946 f"Extraction: type={contract_type.value}, "947 f"text_len={len(cleaned_text)}, pages={page_count}"948 )949 950 # Step 1: Extract (section-based for large docs, single-pass for small)951 if page_count > SECTION_PAGE_THRESHOLD:952 logger.info(f"Large document ({page_count} pages). Using section-based extraction.")953 result = await self._section_based_extraction(cleaned_text, contract_type)954 else:955 logger.info(f"Standard document ({page_count} pages). Using single-pass extraction.")956 result = await self._single_pass_extraction(cleaned_text, contract_type)957 958 if not result or result.get("error"):959 logger.error(f"Primary extraction failed: {result}")960 return result or {"agreement_type": contract_type.value, "error": "Extraction failed"}961 962 # Step 2: Gap-fill for missing critical fields963 missing = self._find_missing_critical_fields(result, contract_type)964 if missing:965 logger.info(f"Gap-fill needed for {len(missing)} missing fields: {missing}")966 result = await self._gap_fill_extraction(967 result, missing, cleaned_text, contract_type968 )969 970 # Step 3: Second gap-fill if still missing (different prompt strategy)971 still_missing = self._find_missing_critical_fields(result, contract_type)972 if still_missing:973 logger.info(f"Second gap-fill for {len(still_missing)} remaining fields")974 result = await self._gap_fill_extraction(975 result, still_missing, cleaned_text, contract_type976 )977 978 return result979 980 async def extract_with_vision_context(981 self,982 document_text: str,983 vision_analysis: str,984 contract_type: ContractType = None,985 ) -> dict:986 """Extract with additional vision OCR context."""987 combined = f"{document_text}\n\n=== VISION ANALYSIS ===\n{vision_analysis}"988 return await self.extract(combined, contract_type)989 990 # ─── Single-Pass Extraction ───────────────────────────────────991 992 async def _single_pass_extraction(993 self,994 text: str,995 contract_type: ContractType,996 ) -> dict:997 """Full document extraction in a single LLM call with retry."""998 schema = TYPE_SCHEMA_MAP.get(contract_type, NDA_SCHEMA)999 instructions = TYPE_INSTRUCTIONS.get(contract_type, "")1000 1001 result = await self._llm_extract_with_retry(1002 prompt_template=self.extraction_prompt,1003 variables={1004 "contract_type_label": contract_type.value,1005 "type_instructions": instructions,1006 "json_schema": schema,1007 "document_text": text,1008 },1009 context=f"single-pass {contract_type.value}",1010 )1011 1012 return result1013 1014 # ─── Section-Based Extraction ─────────────────────────────────1015 1016 async def _section_based_extraction(1017 self,1018 text: str,1019 contract_type: ContractType,1020 ) -> dict:1021 """1022 Split document into 3 sections (early/middle/late pages),1023 extract each in parallel, then merge.1024 """1025 sections = self._split_into_sections(text)1026 schema = TYPE_SCHEMA_MAP.get(contract_type, NDA_SCHEMA)1027 1028 # Extract all sections in parallel1029 tasks = []1030 for section_name, section_text, page_range in sections:1031 focus = SECTION_FOCUS_MAP.get(section_name, "")1032 tasks.append(1033 self._llm_extract_with_retry(1034 prompt_template=self.section_prompt,1035 variables={1036 "contract_type_label": contract_type.value,1037 "page_range": page_range,1038 "section_focus": focus,1039 "json_schema": schema,1040 "section_text": section_text,1041 },1042 context=f"section-{section_name}",1043 )1044 )1045 1046 section_results = await asyncio.gather(*tasks, return_exceptions=True)1047 1048 # Merge section results (later sections override earlier for conflicts)1049 merged = {}1050 for i, result in enumerate(section_results):1051 if isinstance(result, Exception):1052 logger.error(f"Section {i} extraction failed: {result}")1053 continue1054 if isinstance(result, dict) and not result.get("error"):1055 merged = self._deep_merge(merged, result)1056 1057 # If section-based gave sparse results, fall back to single-pass1058 filled_count = self._count_filled_fields(merged)1059 if filled_count < 5:1060 logger.warning(1061 f"Section-based extraction too sparse ({filled_count} fields). "1062 f"Falling back to single-pass."1063 )1064 return await self._single_pass_extraction(text, contract_type)1065 1066 return merged1067 1068 def _split_into_sections(self, text: str) -> List[Tuple[str, str, str]]:1069 """Split document text into early/middle/late sections by page markers."""1070 pages = re.split(r'(?=--- Page \d+)', text)1071 pages = [p for p in pages if p.strip()]1072 1073 if len(pages) <= 3:1074 return [("early", text, "1-" + str(len(pages)))]1075 1076 third = max(len(pages) // 3, 1)1077 sections = [1078 ("early", "\n".join(pages[:third]), f"1-{third}"),1079 ("middle", "\n".join(pages[third:2*third]), f"{third+1}-{2*third}"),1080 ("late", "\n".join(pages[2*third:]), f"{2*third+1}-{len(pages)}"),1081 ]1082 return sections1083 1084 # ─── Gap-Fill Extraction ──────────────────────────────────────1085 1086 async def _gap_fill_extraction(1087 self,1088 current_result: dict,1089 missing_fields: List[str],1090 text: str,1091 contract_type: ContractType,1092 ) -> dict:1093 """Re-extract only the missing fields with a targeted prompt."""1094 missing_desc = "\n".join([f" - {f}" for f in missing_fields])1095 1096 gap_result = await self._llm_extract_with_retry(1097 prompt_template=self.gap_fill_prompt,1098 variables={1099 "missing_fields": missing_desc,1100 "document_text": text,1101 },1102 context="gap-fill",1103 )1104 1105 if not gap_result or gap_result.get("error"):1106 logger.warning("Gap-fill extraction returned no results")1107 return current_result1108 1109 # Apply gap-fill results to the current extraction1110 for field_path, value in gap_result.items():1111 if value is not None and value != "" and value != []:1112 self._set_nested_value(current_result, field_path, value)1113 1114 return current_result1115 1116 # ─── LLM Call with Retry ──────────────────────────────────────1117 1118 async def _llm_extract_with_retry(1119 self,1120 prompt_template: ChatPromptTemplate,1121 variables: dict,1122 context: str = "",1123 ) -> dict:1124 """Call LLM with exponential backoff retry. Returns parsed dict."""1125 chain = prompt_template | self.llm1126 1127 for attempt in range(1, MAX_RETRIES + 1):1128 try:1129 response = await chain.ainvoke(variables)1130 content = response.content.strip()1131 parsed = self._parse_json_response(content)1132 1133 if parsed and not parsed.get("error"):1134 filled = self._count_filled_fields(parsed)1135 logger.info(1136 f"[{context}] Attempt {attempt} succeeded: "1137 f"{filled} fields filled"1138 )1139 return parsed1140 1141 logger.warning(1142 f"[{context}] Attempt {attempt} returned sparse/error result"1143 )1144 1145 except Exception as e:1146 logger.error(f"[{context}] Attempt {attempt} failed: {e}")1147 1148 if attempt < MAX_RETRIES:1149 delay = RETRY_BASE_DELAY ** attempt1150 logger.info(f"[{context}] Retrying in {delay}s...")1151 await asyncio.sleep(delay)1152 1153 logger.error(f"[{context}] All {MAX_RETRIES} attempts failed")1154 return {"agreement_type": variables.get("contract_type_label", "Unknown"), "error": "All extraction attempts failed"}1155 1156 # ─── OCR Text Cleaning ────────────────────────────────────────1157 1158 @staticmethod1159 def _clean_ocr_text(text: str) -> str:1160 """1161 Clean OCR artifacts from vision-extracted text.1162 Preserves structure (page markers, paragraphs) while fixing noise.1163 """1164 if not text:1165 return ""1166 1167 # Remove zero-width characters and control chars (except newline, tab)1168 text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f\u200b\u200c\u200d\ufeff]', '', text)1169 1170 # Fix common OCR artifacts1171 text = re.sub(r'\s{3,}', ' ', text) # Collapse excessive spaces1172 text = re.sub(r'\n{4,}', '\n\n\n', text) # Collapse excessive newlines1173 text = re.sub(r'([a-z])\s{2,}([a-z])', r'\1 \2', text) # Fix word breaks1174 text = re.sub(r'(\w)-\s*\n\s*(\w)', r'\1\2', text) # Fix hyphenated line breaks1175 1176 # Normalize quotes1177 text = text.replace('\u201c', '"').replace('\u201d', '"')1178 text = text.replace('\u2018', "'").replace('\u2019', "'")1179 1180 # Clean up page markers (normalize format)1181 text = re.sub(r'---\s*Page\s*(\d+)\s*(?:\(OCR\))?\s*---', r'--- Page \1 ---', text)1182 1183 return text.strip()1184 1185 # ─── JSON Parsing ─────────────────────────────────────────────1186 1187 @staticmethod1188 def _parse_json_response(content: str) -> dict:1189 """Robust JSON parsing with multiple fallback strategies."""1190 original = content1191 1192 # Strip markdown fences1193 content = re.sub(r'^```(?:json)?\s*\n?', '', content, flags=re.MULTILINE)1194 content = re.sub(r'\n?```\s*$', '', content, flags=re.MULTILINE)1195 content = content.strip()1196 1197 # Attempt 1: Direct parse1198 try:1199 return json.loads(content)1200 except json.JSONDecodeError: