app/services/entities.py

"""NLP entity extraction — no LLM, pure regex/heuristic.

Extracts: parties, dates, monetary amounts, durations, jurisdictions, key obligations.
Runs synchronously on the parsed document text before any LLM calls.
"""

from __future__ import annotations

import re

from app.schemas import DocumentEntities

# --- Pattern sets ---

_PARTY_PATTERNS = [
    r'\b(?:the\s+)?(?:Company|Employer|Client|Contractor|Vendor|Provider|Lessor|Landlord|'
    r'Licensor|Disclosing\s+Party|Receiving\s+Party|Service\s+Provider)\b',
    r'\b(?:the\s+)?(?:Employee|Consultant|Freelancer|Tenant|Lessee|Licensee|'
    r'Customer|User|Subscriber)\b',
]

_DATE_PATTERNS = [
    r'\b(?:January|February|March|April|May|June|July|August|September|October|November|December)'
    r'\s+\d{1,2},?\s+\d{4}\b',
    r'\b\d{1,2}[\/\-]\d{1,2}[\/\-]\d{2,4}\b',
    r'\beffective\s+(?:date|as\s+of)\s+[^\.,;]{3,30}',
]

_MONEY_PATTERNS = [
    r'\$\s*[\d,]+(?:\.\d{2})?(?:\s*(?:million|billion|thousand|USD|CAD|EUR|GBP))?',
    r'\b[\d,]+(?:\.\d{2})?\s*(?:US\s*)?dollars?\b',
    r'\b(?:USD|EUR|GBP|CAD)\s*[\d,]+(?:\.\d{2})?\b',
]

_DURATION_PATTERNS = [
    r'\b(?:zero\s*\(0\)|one\s*\(1\)|two\s*\(2\)|three\s*\(3\)|four\s*\(4\)|'
    r'five\s*\(5\)|six\s*\(6\)|seven\s*\(7\)|eight\s*\(8\)|nine\s*\(9\)|'
    r'ten\s*\(10\)|twelve\s*\(12\)|eighteen\s*\(18\)|twenty-?four\s*\(24\)|'
    r'thirty\s*\(30\)|sixty\s*\(60\)|ninety\s*\(90\)|one\s+hundred\s+eighty\s*\(180\))\s+'
    r'(?:calendar\s+)?(?:days?|weeks?|months?|years?)\b',
    r'\b\d+\s*(?:calendar\s+)?(?:days?|weeks?|months?|years?)\b',
    r'\b(?:perpetual|irrevocable|indefinite(?:ly)?|forever|permanent(?:ly)?)\b',
]

_JURISDICTION_PATTERNS = [
    r'\bState\s+of\s+[A-Z][a-z]+(?:\s+[A-Z][a-z]+)?\b',
    r'\b(?:California|Delaware|New\s+York|Texas|Florida|Illinois|Washington|'
    r'Colorado|Massachusetts|Georgia)\b',
    r'\b(?:United\s+States|federal|U\.S\.|EU|European\s+Union|GDPR|CCPA)\b',
]

_OBLIGATION_KEYWORDS = [
    (r'\bshall\s+not\b', "Prohibition"),
    (r'\bmust\s+not\b', "Prohibition"),
    (r'\bhereby\s+assigns?\b', "Assignment of rights"),
    (r'\bwaives?\s+(?:all|any)\b', "Waiver of rights"),
    (r'\birrevocably\b', "Irrevocable commitment"),
    (r'\bin\s+perpetuity\b', "Perpetual obligation"),
    (r'\bsole\s+(?:discretion|remedy)\b', "Unilateral discretion"),
    (r'\bclass\s+action\s+waiver\b', "Class action waiver"),
    (r'\bbinding\s+arbitration\b', "Mandatory arbitration"),
    (r'\bat.?will\b', "At-will relationship"),
    (r'\bauto.?renew(?:al|s)?\b', "Auto-renewal clause"),
]


def _dedupe(items: list[str]) -> list[str]:
    seen: set[str] = set()
    out: list[str] = []
    for item in items:
        key = item.strip().lower()
        if key not in seen and len(key) > 1:
            seen.add(key)
            out.append(item.strip())
    return out


def extract_entities(text: str) -> DocumentEntities:
    parties: list[str] = []
    for pat in _PARTY_PATTERNS:
        parties += re.findall(pat, text, re.IGNORECASE)
    parties = _dedupe([p.replace("the ", "").strip().title() for p in parties])[:8]

    dates: list[str] = []
    for pat in _DATE_PATTERNS:
        dates += re.findall(pat, text, re.IGNORECASE)
    dates = _dedupe(dates)[:6]

    amounts: list[str] = []
    for pat in _MONEY_PATTERNS:
        amounts += re.findall(pat, text, re.IGNORECASE)
    amounts = _dedupe(amounts)[:8]

    durations: list[str] = []
    for pat in _DURATION_PATTERNS:
        durations += re.findall(pat, text, re.IGNORECASE)
    durations = _dedupe(durations)[:8]

    jurisdictions: list[str] = []
    for pat in _JURISDICTION_PATTERNS:
        jurisdictions += re.findall(pat, text, re.IGNORECASE)
    jurisdictions = _dedupe(jurisdictions)[:6]

    obligations: list[str] = []
    for pat, label in _OBLIGATION_KEYWORDS:
        if re.search(pat, text, re.IGNORECASE):
            obligations.append(label)
    obligations = _dedupe(obligations)[:6]

    return DocumentEntities(
        parties=parties,
        dates=dates,
        monetary_amounts=amounts,
        durations=durations,
        jurisdictions=jurisdictions,
        key_obligations=obligations,
    )