Automated Legal Contract Review with LLMs: Clause Extraction, Risk Scoring, and Redlines

import io
from typing import Optional, Tuple
import docx
from pypdf import PdfReader
import pandas as pd

class ContractIngestor:
    """Parses contract files, preserving basic structure for clause analysis."""

    def ingest(self, file_bytes: bytes, filename: str) -> Tuple[str, dict]:
        """
        Returns (extracted_text, metadata_dict).
        Metadata includes page count, word count, and detected sections.
        """
        metadata = {"filename": filename, "word_count": 0, "sections": []}
        text = ""

        if filename.lower().endswith('.pdf'):
            # PDF Parsing
            pdf_file = io.BytesIO(file_bytes)
            reader = PdfReader(pdf_file)
            metadata["page_count"] = len(reader.pages)
            pages_text = []
            for page_num, page in enumerate(reader.pages):
                page_text = page.extract_text()
                if page_text.strip():
                    # Add a section marker for each page (simple heuristic)
                    pages_text.append(f"\n--- PAGE {page_num+1} ---\n{page_text}")
            text = "\n".join(pages_text)

        elif filename.lower().endswith('.docx'):
            # DOCX Parsing
            doc_file = io.BytesIO(file_bytes)
            doc = docx.Document(doc_file)
            metadata["page_count"] = len(doc.paragraphs) // 50  # rough estimate
            for para in doc.paragraphs:
                if para.text.strip():
                    text += para.text + "\n"
                    # Heuristic: Bold and large text might be a section header
                    if para.runs and (para.runs[0].bold or para.runs[0].font.size > 14):
                        metadata["sections"].append(para.text.strip())

        # Basic word count
        words = text.split()
        metadata["word_count"] = len(words)

        return text, metadata


@app.post("/ingest")
async def ingest_contract(file: UploadFile = File(...)):
    ingestor = ContractIngestor()
    file_bytes = await file.read()
    text, metadata = ingestor.ingest(file_bytes, file.filename)
    # Queue for processing
    task = process_contract.delay(text, metadata)
    return {"task_id": task.id, "metadata": metadata}

from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain.output_parsers import PydanticOutputParser
import json

# Define the exact structure we want
class ExtractedClause(BaseModel):
    clause_name: str = Field(description="The standardized name of the clause, e.g., GOVERNING_LAW, LIABILITY_CAP, TERM, CONFIDENTIALITY_DEFINITION")
    original_text: str = Field(description="The exact text from the contract for this clause.")
    confidence: float = Field(description="Confidence score from 0.0 to 1.0")
    found_page: Optional[int] = Field(description="Page number where the clause was found, if detectable.")

class ContractClauses(BaseModel):
    clauses: list[ExtractedClause] = Field(description="List of extracted clauses")

def extract_clauses_with_llm(contract_text: str, metadata: dict) -> ContractClauses:
    """Send to LLM for structured extraction."""
    parser = PydanticOutputParser(pydantic_object=ContractClauses)
    llm = ChatOpenAI(model="gpt-4o", temperature=0)  # Zero temperature for consistency

    prompt = ChatPromptTemplate.from_template("""
    You are a legal contract analyst. Extract the following specific clauses from the contract text below.

    **Instructions:**
    1. Only extract clauses you can identify with high confidence.
    2. Use the EXACT clause names provided.
    3. Return the `original_text` verbatim from the contract.
    4. Estimate `found_page` based on page markers like '--- PAGE X ---' in the text.

    **Clause Names to Find:**
    - GOVERNING_LAW
    - LIABILITY_CAP
    - TERM
    - CONFIDENTIALITY_DEFINITION
    - INDEMNIFICATION
    - NOTICE_PERIOD

    Contract Text:
    {contract_text}

    {format_instructions}
    """)

    chain = prompt | llm | parser
    # Error handling is critical here
    try:
        result = chain.invoke({
            "contract_text": contract_text[:15000],  # Token limit safety
            "format_instructions": parser.get_format_instructions()
        })
    except openai.RateLimitError:
        # Real Error & Fix: You exceeded your current quota
        # Fix: implement per-tenant rate limiting with Redis token bucket
        # For now, we log and re-raise, but in production:
        #   bucket = redis.get(f"tenant:{tenant_id}:tokens")
        #   if bucket and int(bucket) < 1: raise RateLimitError
        raise
    return result

# The output is a validated Pydantic object. Access: result.clauses[0].original_text

class StandardClause(BaseModel):
    name: str
    standard_text: str
    risk_weight: float  # 1.0 (low) to 5.0 (critical)
    allowed_variations: list[str]  # e.g., ["Delaware", "New York"] for governing law

class RiskScorer:
    def __init__(self):
        # In reality, load from a database
        self.standards = {
            "GOVERNING_LAW": StandardClause(
                name="GOVERNING_LAW",
                standard_text="This Agreement shall be governed by the laws of the State of Delaware.",
                risk_weight=3.0,
                allowed_variations=["Delaware", "DE"]
            ),
            "LIABILITY_CAP": StandardClause(
                name="LIABILITY_CAP",
                standard_text="Liability cap: the greater of $100,000 or fees paid in the 12 months preceding the claim.",
                risk_weight=5.0,
                allowed_variations=["100,000", "100k"]
            )
        }

    def score_clause(self, extracted_clause: ExtractedClause) -> dict:
        """Returns a risk score and deviation analysis for a single clause."""
        standard = self.standards.get(extracted_clause.clause_name)
        if not standard:
            return {"risk_score": 0, "flag": "NO_STANDARD", "suggestion": ""}

        # Simple semantic comparison: check if standard keywords are present
        deviation_detected = False
        suggestion = ""
        original_lower = extracted_clause.original_text.lower()

        # Check for allowed variations (simplified logic)
        if extracted_clause.clause_name == "GOVERNING_LAW":
            if not any(var.lower() in original_lower for var in standard.allowed_variations):
                deviation_detected = True
                suggestion = f"Consider requesting change to: {standard.standard_text}"

        # More complex logic for liability caps could use regex to extract amounts
        if deviation_detected:
            risk_score = standard.risk_weight * (1 - extracted_clause.confidence)
            flag = "NON_STANDARD"
        else:
            risk_score = 0
            flag = "STANDARD"

        return {
            "clause_name": extracted_clause.clause_name,
            "risk_score": round(risk_score, 2),
            "flag": flag,
            "suggestion": suggestion,
            "original_text_snippet": extracted_clause.original_text[:150] + "..."
        }

    def score_contract(self, extracted_clauses: ContractClauses) -> pd.DataFrame:
        """Score all clauses and return a DataFrame for reporting."""
        scores = []
        for clause in extracted_clauses.clauses:
            scores.append(self.score_clause(clause))
        df = pd.DataFrame(scores)
        df = df.sort_values("risk_score", ascending=False)
        return df

from docx import Document
from docx.shared import RGBColor
from docx.opc.constants import RELATIONSHIP_TYPE as RT

def generate_redline_docx(original_text: str, risk_df: pd.DataFrame, metadata: dict) -> Document:
    """Creates a DOCX with risky clauses highlighted and suggestions in comments."""
    doc = Document()
    doc.add_heading(f'Contract Review: {metadata.get("filename", "Document")}', 0)

    # Add summary table
    doc.add_heading('Risk Summary', level=1)
    table = doc.add_table(rows=1, cols=4)
    hdr_cells = table.rows[0].cells
    hdr_cells[0].text = 'Clause'
    hdr_cells[1].text = 'Risk Score'
    hdr_cells[2].text = 'Flag'
    hdr_cells[3].text = 'Suggestion'
    for _, row in risk_df[risk_df['risk_score'] > 0].iterrows():
        row_cells = table.add_row().cells
        row_cells[0].text = row['clause_name']
        row_cells[1].text = str(row['risk_score'])
        row_cells[2].text = row['flag']

    # Add the contract text with highlights
    doc.add_heading('Reviewed Text', level=1)
    # Simplified: We add the text as a single paragraph. A robust solution would map back to original positions.
    p = doc.add_paragraph(original_text[:5000])  # Truncate for demo

    # Highlight risky terms (simplified mapping)
    for _, row in risk_df[risk_df['risk_score'] > 0].iterrows():
        if row['original_text_snippet'] in original_text:
            # This is naive; a real impl would use character offsets from the LLM
            for run in p.runs:
                if row['original_text_snippet'] in run.text:
                    run.font.highlight_color = RGBColor(255, 255, 0)  # Yellow highlight
                    # Adding a comment is more complex and requires relationship management
                    # doc_part = doc.part
                    # comment_part = doc_part.add_comment_part()
                    # comment = comment_part.add_comment("Reviewer", row['suggestion'])
                    # run.add_comment(comment, "AI", "Suggestion: " + row['suggestion'])

    return doc

# Save the document
# doc.save(f"review_{metadata['filename']}.docx")

# Real Error & Fix: PII detected in prompt
# Fix: run presidio analyzer before sending to LLM, redact then re-inject
from presidio_analyzer import AnalyzerEngine
from presidio_anonymizer import AnonymizerEngine

def sanitize_text(text: str) -> str:
    analyzer = AnalyzerEngine()
    anonymizer = AnonymizerEngine()
    results = analyzer.analyze(text=text, language='en')
    anonymized = anonymizer.anonymize(text=text, analyzer_results=results)
    return anonymized.text
# Store mapping to re-inject later if needed for redlining.