AI-Powered Contract Intelligence System

import fitz  # PyMuPDF
import re
import pytesseract
from PIL import Image
from dataclasses import dataclass
from typing import List, Optional
import tiktoken


@dataclass
class ContractChunk:
    text: str
    section_title: str
    clause_number: Optional[str]
    page_number: int
    chunk_type: str  # "clause" | "definition" | "schedule" | "recital"
    token_count: int


class ContractParser:
    """Structure-aware contract parser with OCR fallback."""

    CLAUSE_PATTERN = re.compile(
        r'^(\d+\.?\d*\.?\d*)\s+([A-Z][A-Za-z\s&,]+)'
    )
    DEFINITION_PATTERN = re.compile(
        r'"([^"]+)"\s+(means|shall mean|refers to)'
    )
    MAX_CHUNK_TOKENS = 512

    def __init__(self):
        self.tokenizer = tiktoken.encoding_for_model("gpt-4")

    def parse_pdf(self, pdf_path: str) -> List[ContractChunk]:
        """Extract text from PDF with OCR fallback."""
        doc = fitz.open(pdf_path)
        chunks = []

        for page_num, page in enumerate(doc, 1):
            text = page.get_text("text")

            # OCR fallback for scanned pages
            if len(text.strip()) < 50:
                pix = page.get_pixmap(dpi=300)
                img = Image.frombytes(
                    "RGB", [pix.width, pix.height], pix.samples
                )
                text = pytesseract.image_to_string(img)

            page_chunks = self._structure_aware_chunk(text, page_num)
            chunks.extend(page_chunks)

        return self._merge_small_chunks(chunks)

    def _structure_aware_chunk(
        self, text: str, page_num: int
    ) -> List[ContractChunk]:
        """Split text by clause boundaries, not arbitrary windows."""
        lines = text.split('\n')
        chunks = []
        current_section = ""
        current_clause = None
        current_text = []

        for line in lines:
            clause_match = self.CLAUSE_PATTERN.match(line.strip())

            if clause_match:
                # Save previous chunk
                if current_text:
                    joined = '\n'.join(current_text)
                    tokens = len(self.tokenizer.encode(joined))
                    chunks.append(ContractChunk(
                        text=joined,
                        section_title=current_section,
                        clause_number=current_clause,
                        page_number=page_num,
                        chunk_type=self._classify_chunk(joined),
                        token_count=tokens,
                    ))

                current_clause = clause_match.group(1)
                current_section = clause_match.group(2).strip()
                current_text = [line]
            else:
                current_text.append(line)

        # Don't forget the last chunk
        if current_text:
            joined = '\n'.join(current_text)
            tokens = len(self.tokenizer.encode(joined))
            chunks.append(ContractChunk(
                text=joined,
                section_title=current_section,
                clause_number=current_clause,
                page_number=page_num,
                chunk_type=self._classify_chunk(joined),
                token_count=tokens,
            ))

        return chunks

    def _classify_chunk(self, text: str) -> str:
        """Classify chunk type based on content patterns."""
        if self.DEFINITION_PATTERN.search(text):
            return "definition"
        if re.search(r'SCHEDULE|EXHIBIT|APPENDIX', text, re.IGNORECASE):
            return "schedule"
        if re.search(r'WHEREAS|RECITAL', text, re.IGNORECASE):
            return "recital"
        return "clause"

    def _merge_small_chunks(
        self, chunks: List[ContractChunk]
    ) -> List[ContractChunk]:
        """Merge chunks smaller than 64 tokens with neighbors."""
        merged = []
        for chunk in chunks:
            if (merged and chunk.token_count < 64
                and merged[-1].token_count + chunk.token_count
                    < self.MAX_CHUNK_TOKENS):
                merged[-1].text += '\n' + chunk.text
                merged[-1].token_count += chunk.token_count
            else:
                merged.append(chunk)
        return merged

import anthropic
import json
from typing import List, Dict
from pydantic import BaseModel, Field
from enum import Enum


class ClauseType(Enum):
    INDEMNIFICATION = "indemnification"
    LIMITATION_OF_LIABILITY = "limitation_of_liability"
    TERMINATION = "termination"
    CONFIDENTIALITY = "confidentiality"
    INTELLECTUAL_PROPERTY = "intellectual_property"
    DATA_PROTECTION = "data_protection"
    PAYMENT_TERMS = "payment_terms"
    WARRANTY = "warranty"
    FORCE_MAJEURE = "force_majeure"
    NON_COMPETE = "non_compete"
    GOVERNING_LAW = "governing_law"
    ASSIGNMENT = "assignment"


class ExtractedClause(BaseModel):
    clause_type: ClauseType
    text: str
    section_ref: str
    obligations: List[str] = Field(default_factory=list)
    key_dates: List[str] = Field(default_factory=list)
    monetary_values: List[str] = Field(default_factory=list)
    risk_indicators: List[str] = Field(default_factory=list)


class ClauseExtractor:
    """LLM-powered clause extraction with structured output."""

    EXTRACTION_PROMPT = """You are a senior contract analyst. Extract all
legal clauses from this contract section. For each clause, identify:

1. clause_type: One of: {clause_types}
2. text: The exact clause text
3. section_ref: Section number reference (e.g., "3.2.1")
4. obligations: List of obligations imposed on each party
5. key_dates: Any dates, deadlines, or time periods mentioned
6. monetary_values: Dollar amounts, percentages, caps
7. risk_indicators: Red flags or unusual terms

Return valid JSON array of extracted clauses.

CONTRACT SECTION:
{contract_text}"""

    def __init__(self):
        self.client = anthropic.Anthropic()

    async def extract_clauses(
        self, contract_text: str
    ) -> List[ExtractedClause]:
        """Extract structured clause data using Claude."""

        prompt = self.EXTRACTION_PROMPT.format(
            clause_types=", ".join(
                [ct.value for ct in ClauseType]
            ),
            contract_text=contract_text,
        )

        response = self.client.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=4096,
            messages=[{"role": "user", "content": prompt}],
            # Force structured JSON output
            system="Respond only with a valid JSON array. No markdown.",
        )

        raw = json.loads(response.content[0].text)
        return [ExtractedClause(**clause) for clause in raw]

from dataclasses import dataclass
from typing import List, Dict
import anthropic


@dataclass
class RiskAssessment:
    overall_score: float       # 0.0 (low risk) to 1.0 (critical)
    category_scores: Dict[str, float]
    high_risk_clauses: List[Dict]
    missing_protections: List[str]
    recommendations: List[str]


class RiskAssessor:
    """Multi-dimensional risk scoring for contracts."""

    RISK_CATEGORIES = [
        "financial_exposure",
        "liability_scope",
        "ip_protection",
        "data_privacy",
        "termination_risk",
        "regulatory_compliance",
    ]

    RISK_PROMPT = """You are an enterprise legal risk analyst.
Assess the following contract clauses across these risk dimensions:
{categories}

For each dimension, provide:
- score (0.0 = no risk, 1.0 = critical risk)
- rationale (1-2 sentences)
- specific clause references

Also identify:
1. High-risk clauses that need immediate legal review
2. Standard protections MISSING from this contract
3. Specific recommendations to reduce risk

Contract clauses:
{clauses_json}

Respond with valid JSON."""

    def __init__(self):
        self.client = anthropic.Anthropic()

    async def assess_risk(
        self, clauses: List[Dict]
    ) -> RiskAssessment:
        """Run multi-dimensional risk assessment via LLM."""

        prompt = self.RISK_PROMPT.format(
            categories="\n".join(
                f"- {cat}" for cat in self.RISK_CATEGORIES
            ),
            clauses_json=json.dumps(clauses, indent=2),
        )

        response = self.client.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=4096,
            messages=[{"role": "user", "content": prompt}],
        )

        result = json.loads(response.content[0].text)

        # Compute weighted overall score
        weights = {"financial_exposure": 0.25, "liability_scope": 0.20,
                   "ip_protection": 0.15, "data_privacy": 0.20,
                   "termination_risk": 0.10, "regulatory_compliance": 0.10}

        overall = sum(
            result["category_scores"].get(cat, 0) * w
            for cat, w in weights.items()
        )

        return RiskAssessment(
            overall_score=round(overall, 3),
            category_scores=result["category_scores"],
            high_risk_clauses=result["high_risk_clauses"],
            missing_protections=result["missing_protections"],
            recommendations=result["recommendations"],
        )

import openai
import asyncpg
import numpy as np
from typing import List, Dict, Optional
from dataclasses import dataclass


@dataclass
class SimilarClause:
    clause_id: str
    contract_id: str
    text: str
    clause_type: str
    similarity_score: float
    contract_title: str
    jurisdiction: str


class ContractVectorSearch:
    """RAG-based clause similarity search using pgvector."""

    def __init__(self, db_url: str):
        self.db_url = db_url
        self.openai = openai.AsyncOpenAI()

    async def embed_text(self, text: str) -> List[float]:
        """Generate embedding via OpenAI."""
        response = await self.openai.embeddings.create(
            model="text-embedding-3-small",
            input=text,
        )
        return response.data[0].embedding

    async def find_similar_clauses(
        self,
        query_text: str,
        clause_type: Optional[str] = None,
        jurisdiction: Optional[str] = None,
        top_k: int = 10,
        min_similarity: float = 0.75,
    ) -> List[SimilarClause]:
        """Hybrid search: vector similarity + SQL filters."""

        query_embedding = await self.embed_text(query_text)

        # Build dynamic WHERE clause for filtering
        conditions = ["1 = 1"]
        params = [str(query_embedding), top_k]

        if clause_type:
            conditions.append(f"c.clause_type = ${len(params) + 1}")
            params.append(clause_type)

        if jurisdiction:
            conditions.append(f"ct.jurisdiction = ${len(params) + 1}")
            params.append(jurisdiction)

        where = " AND ".join(conditions)

        query = f"""
            SELECT
                c.id AS clause_id,
                c.contract_id,
                c.text,
                c.clause_type,
                1 - (c.embedding <=> $1::vector) AS similarity,
                ct.title AS contract_title,
                ct.jurisdiction
            FROM clause_embeddings c
            JOIN contracts ct ON ct.id = c.contract_id
            WHERE {where}
              AND 1 - (c.embedding <=> $1::vector) >= {min_similarity}
            ORDER BY c.embedding <=> $1::vector
            LIMIT $2
        """

        conn = await asyncpg.connect(self.db_url)
        try:
            rows = await conn.fetch(query, *params)
            return [
                SimilarClause(
                    clause_id=str(row["clause_id"]),
                    contract_id=str(row["contract_id"]),
                    text=row["text"],
                    clause_type=row["clause_type"],
                    similarity_score=float(row["similarity"]),
                    contract_title=row["contract_title"],
                    jurisdiction=row["jurisdiction"],
                )
                for row in rows
            ]
        finally:
            await conn.close()

    async def detect_deviations(
        self, clause_text: str, clause_type: str
    ) -> Dict:
        """Compare clause against corpus standard to find deviations."""

        # Find the most similar approved clauses
        similar = await self.find_similar_clauses(
            clause_text,
            clause_type=clause_type,
            top_k=5,
            min_similarity=0.70,
        )

        if not similar:
            return {
                "status": "no_precedent",
                "message": "No similar clauses found — manual review required",
            }

        avg_similarity = np.mean(
            [s.similarity_score for s in similar]
        )

        return {
            "status": "deviation" if avg_similarity < 0.85 else "standard",
            "avg_similarity": round(float(avg_similarity), 3),
            "closest_precedent": similar[0].text,
            "precedent_contract": similar[0].contract_title,
            "all_matches": len(similar),
        }

import asyncio
import asyncpg
import math
from typing import List, Dict, Tuple
from dataclasses import dataclass


@dataclass
class HybridResult:
    clause_id: str
    text: str
    cosine_score: float
    bm25_score: float
    rrf_score: float       # combined ranking score
    clause_type: str
    contract_title: str


class HybridContractSearch:
    """
    Hybrid retrieval: Cosine similarity (semantic) + BM25 (lexical).
    Uses Reciprocal Rank Fusion to combine rankings.

    Same approach used in chinnam.AI RAG pipeline for research papers.
    """

    RRF_K = 60  # RRF constant — controls rank vs score weighting

    def __init__(self, db_url: str):
        self.db_url = db_url
        self.openai = openai.AsyncOpenAI()

    async def search(
        self,
        query: str,
        clause_type: str | None = None,
        top_k: int = 10,
        cosine_weight: float = 0.6,
        bm25_weight: float = 0.4,
    ) -> List[HybridResult]:
        """
        Run cosine + BM25 in parallel, merge with RRF.

        Weights: 60% semantic, 40% lexical — tuned for legal text
        where exact terms matter more than general prose.
        """

        # Run both searches in parallel — max(latency) not sum(latency)
        cosine_results, bm25_results = await asyncio.gather(
            self._cosine_search(query, clause_type, top_k * 2),
            self._bm25_search(query, clause_type, top_k * 2),
        )

        # Reciprocal Rank Fusion — position-based, not score-based
        # This handles the fact that cosine ∈ [0,1] and BM25 is unbounded
        fused = self._reciprocal_rank_fusion(
            cosine_results, bm25_results,
            cosine_weight, bm25_weight,
        )

        return sorted(fused, key=lambda r: r.rrf_score, reverse=True)[:top_k]

    async def _cosine_search(
        self, query: str, clause_type: str | None, limit: int
    ) -> List[Tuple[str, float, Dict]]:
        """pgvector cosine similarity search."""
        embedding = await self._embed(query)
        conn = await asyncpg.connect(self.db_url)
        try:
            rows = await conn.fetch(f"""
                SELECT c.id, c.text, c.clause_type, ct.title,
                       1 - (c.embedding <=> $1::vector) AS score
                FROM clause_embeddings c
                JOIN contracts ct ON ct.id = c.contract_id
                WHERE ($3::text IS NULL OR c.clause_type = $3)
                  AND 1 - (c.embedding <=> $1::vector) >= 0.70
                ORDER BY c.embedding <=> $1::vector
                LIMIT $2
            """, str(embedding), limit, clause_type)
            return [
                (row["id"], float(row["score"]), dict(row))
                for row in rows
            ]
        finally:
            await conn.close()

    async def _bm25_search(
        self, query: str, clause_type: str | None, limit: int
    ) -> List[Tuple[str, float, Dict]]:
        """PostgreSQL full-text search with ts_rank (BM25-equivalent)."""
        conn = await asyncpg.connect(self.db_url)
        try:
            # plainto_tsquery handles multi-word queries automatically
            rows = await conn.fetch(f"""
                SELECT c.id, c.text, c.clause_type, ct.title,
                       ts_rank_cd(
                           to_tsvector('english', c.text),
                           plainto_tsquery('english', $1),
                           32  -- normalize by document length
                       ) AS score
                FROM clause_embeddings c
                JOIN contracts ct ON ct.id = c.contract_id
                WHERE to_tsvector('english', c.text)
                      @@ plainto_tsquery('english', $1)
                  AND ($3::text IS NULL OR c.clause_type = $3)
                ORDER BY score DESC
                LIMIT $2
            """, query, limit, clause_type)
            return [
                (row["id"], float(row["score"]), dict(row))
                for row in rows
            ]
        finally:
            await conn.close()

    def _reciprocal_rank_fusion(
        self,
        cosine_results: List,
        bm25_results: List,
        w_cosine: float,
        w_bm25: float,
    ) -> List[HybridResult]:
        """
        RRF merges two ranked lists without needing score normalization.

        Formula: RRF(d) = Σ  weight_i / (K + rank_i(d))

        K=60 is the standard constant that prevents top-ranked items
        from dominating too heavily.
        """
        scores: Dict[str, Dict] = {}
        K = self.RRF_K

        # Score from cosine rankings
        for rank, (cid, score, meta) in enumerate(cosine_results):
            scores[cid] = {
                "meta": meta,
                "cosine_score": score,
                "bm25_score": 0.0,
                "rrf": w_cosine / (K + rank + 1),
            }

        # Add BM25 rankings
        for rank, (cid, score, meta) in enumerate(bm25_results):
            if cid in scores:
                scores[cid]["bm25_score"] = score
                scores[cid]["rrf"] += w_bm25 / (K + rank + 1)
            else:
                scores[cid] = {
                    "meta": meta,
                    "cosine_score": 0.0,
                    "bm25_score": score,
                    "rrf": w_bm25 / (K + rank + 1),
                }

        return [
            HybridResult(
                clause_id=cid,
                text=data["meta"]["text"],
                cosine_score=data["cosine_score"],
                bm25_score=data["bm25_score"],
                rrf_score=data["rrf"],
                clause_type=data["meta"]["clause_type"],
                contract_title=data["meta"]["title"],
            )
            for cid, data in scores.items()
        ]

import re
import anthropic
from dataclasses import dataclass, field
from typing import List, Dict, Optional
from enum import Enum


class ComplianceFramework(Enum):
    GDPR = "gdpr"
    SOX = "sox"
    HIPAA = "hipaa"
    SOC2 = "soc2"
    CCPA = "ccpa"


@dataclass
class ComplianceRule:
    rule_id: str
    framework: ComplianceFramework
    description: str
    required_patterns: List[str]    # regex patterns that MUST exist
    prohibited_patterns: List[str]  # regex patterns that MUST NOT exist
    severity: str  # "critical" | "high" | "medium" | "low"


@dataclass
class ComplianceResult:
    rule_id: str
    passed: bool
    confidence: float
    source: str  # "rule_engine" | "llm" | "hybrid"
    details: str
    clause_ref: Optional[str] = None


class ComplianceValidator:
    """Hybrid rule-engine + LLM compliance checker."""

    # Deterministic rules for known requirements
    RULES: List[ComplianceRule] = [
        ComplianceRule(
            rule_id="GDPR-001",
            framework=ComplianceFramework.GDPR,
            description="Must include data processing agreement",
            required_patterns=[
                r"data\s+process(ing|or)\s+agreement",
                r"sub-?processor",
            ],
            prohibited_patterns=[],
            severity="critical",
        ),
        ComplianceRule(
            rule_id="GDPR-002",
            framework=ComplianceFramework.GDPR,
            description="Must specify data retention period",
            required_patterns=[
                r"retention\s+period",
                r"delet(e|ion)\s+.{0,30}(data|records)",
            ],
            prohibited_patterns=[
                r"retain\s+(indefinitely|forever|permanently)",
            ],
            severity="critical",
        ),
        ComplianceRule(
            rule_id="SOX-001",
            framework=ComplianceFramework.SOX,
            description="Must include audit rights clause",
            required_patterns=[
                r"audit\s+right",
                r"(inspect|examin)\w*\s+.{0,20}(books|records)",
            ],
            prohibited_patterns=[],
            severity="high",
        ),
    ]

    def __init__(self):
        self.client = anthropic.Anthropic()

    async def validate(
        self,
        contract_text: str,
        frameworks: List[ComplianceFramework],
    ) -> List[ComplianceResult]:
        """Run hybrid compliance validation."""
        results = []

        # Phase 1: Deterministic rule checks
        applicable_rules = [
            r for r in self.RULES
            if r.framework in frameworks
        ]

        for rule in applicable_rules:
            result = self._check_rule(rule, contract_text)
            results.append(result)

        # Phase 2: LLM for nuanced checks rules can't cover
        llm_results = await self._llm_compliance_check(
            contract_text, frameworks
        )
        results.extend(llm_results)

        return results

    def _check_rule(
        self, rule: ComplianceRule, text: str
    ) -> ComplianceResult:
        """Deterministic regex-based rule check."""
        text_lower = text.lower()

        # Check required patterns
        required_found = all(
            re.search(pattern, text_lower)
            for pattern in rule.required_patterns
        )

        # Check prohibited patterns
        prohibited_found = any(
            re.search(pattern, text_lower)
            for pattern in rule.prohibited_patterns
        )

        passed = required_found and not prohibited_found

        return ComplianceResult(
            rule_id=rule.rule_id,
            passed=passed,
            confidence=1.0,  # deterministic = full confidence
            source="rule_engine",
            details=rule.description,
        )

    async def _llm_compliance_check(
        self,
        contract_text: str,
        frameworks: List[ComplianceFramework],
    ) -> List[ComplianceResult]:
        """LLM-based check for nuanced compliance requirements."""

        prompt = f"""As a regulatory compliance expert, analyze this
contract for compliance with: {', '.join(f.value.upper() for f in frameworks)}

Focus on requirements that simple keyword matching would MISS:
1. Implied obligations without explicit language
2. Contradictory clauses that create compliance gaps
3. Missing standard protections for the framework
4. Ambiguous language that could be interpreted non-compliantly

Contract text:
{contract_text[:8000]}

Return JSON array with fields: rule_id, passed, confidence (0-1),
details, clause_ref"""

        response = self.client.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=4096,
            messages=[{"role": "user", "content": prompt}],
        )

        raw = json.loads(response.content[0].text)
        return [
            ComplianceResult(
                rule_id=r["rule_id"],
                passed=r["passed"],
                confidence=r["confidence"],
                source="llm",
                details=r["details"],
                clause_ref=r.get("clause_ref"),
            )
            for r in raw
        ]

import asyncio
import logging
from enum import Enum
from dataclasses import dataclass, field
from typing import Dict, Any, Optional, List
from datetime import datetime

logger = logging.getLogger(__name__)


class PipelineState(Enum):
    PENDING = "pending"
    EXTRACTING = "extracting"
    ASSESSING_RISK = "assessing_risk"
    CHECKING_COMPLIANCE = "checking_compliance"
    SUMMARIZING = "summarizing"
    COMPLETED = "completed"
    FAILED = "failed"
    NEEDS_REVIEW = "needs_review"


@dataclass
class PipelineContext:
    """Shared context passed between agents."""
    contract_id: str
    raw_text: str
    contract_type: str
    state: PipelineState = PipelineState.PENDING
    extracted_clauses: List[Dict] = field(default_factory=list)
    risk_assessment: Optional[Dict] = None
    compliance_results: List[Dict] = field(default_factory=list)
    summary: Optional[str] = None
    errors: List[Dict] = field(default_factory=list)
    started_at: datetime = field(default_factory=datetime.utcnow)
    audit_log: List[Dict] = field(default_factory=list)


class ContractOrchestrator:
    """State machine orchestrator for multi-agent pipeline."""

    # Valid state transitions
    TRANSITIONS = {
        PipelineState.PENDING: [PipelineState.EXTRACTING],
        PipelineState.EXTRACTING: [
            PipelineState.ASSESSING_RISK,
            PipelineState.FAILED,
        ],
        PipelineState.ASSESSING_RISK: [
            PipelineState.CHECKING_COMPLIANCE,
            PipelineState.NEEDS_REVIEW,
            PipelineState.FAILED,
        ],
        PipelineState.CHECKING_COMPLIANCE: [
            PipelineState.SUMMARIZING,
            PipelineState.NEEDS_REVIEW,
            PipelineState.FAILED,
        ],
        PipelineState.SUMMARIZING: [
            PipelineState.COMPLETED,
            PipelineState.FAILED,
        ],
    }

    MAX_RETRIES = 3

    def __init__(self, extraction_agent, risk_agent,
                 compliance_agent, summary_agent):
        self.agents = {
            PipelineState.EXTRACTING: extraction_agent,
            PipelineState.ASSESSING_RISK: risk_agent,
            PipelineState.CHECKING_COMPLIANCE: compliance_agent,
            PipelineState.SUMMARIZING: summary_agent,
        }

    async def run_pipeline(
        self, ctx: PipelineContext
    ) -> PipelineContext:
        """Execute the full agent pipeline with retry logic."""

        pipeline_order = [
            PipelineState.EXTRACTING,
            PipelineState.ASSESSING_RISK,
            PipelineState.CHECKING_COMPLIANCE,
            PipelineState.SUMMARIZING,
        ]

        for target_state in pipeline_order:
            # Validate transition
            allowed = self.TRANSITIONS.get(ctx.state, [])
            if target_state not in allowed:
                logger.error(
                    f"Invalid transition: {ctx.state} -> {target_state}"
                )
                ctx.state = PipelineState.FAILED
                break

            # Execute agent with retries
            success = await self._execute_with_retry(
                ctx, target_state
            )

            if not success:
                ctx.state = PipelineState.FAILED
                break

            # Check if high-risk triggers manual review
            if (target_state == PipelineState.ASSESSING_RISK
                and ctx.risk_assessment
                and ctx.risk_assessment.get("overall_score", 0) > 0.8):
                ctx.state = PipelineState.NEEDS_REVIEW
                self._log_audit(
                    ctx, "HIGH_RISK_ESCALATION",
                    "Risk score > 0.8 — routing to senior review",
                )
                break

        if ctx.state == PipelineState.SUMMARIZING:
            ctx.state = PipelineState.COMPLETED

        return ctx

    async def _execute_with_retry(
        self, ctx: PipelineContext, state: PipelineState
    ) -> bool:
        """Execute an agent with exponential backoff retry."""
        agent = self.agents[state]

        for attempt in range(self.MAX_RETRIES):
            try:
                ctx.state = state
                self._log_audit(
                    ctx, "AGENT_START",
                    f"Starting {state.value} (attempt {attempt + 1})",
                )

                await agent.execute(ctx)

                self._log_audit(
                    ctx, "AGENT_COMPLETE",
                    f"Completed {state.value}",
                )
                return True

            except Exception as e:
                wait = 2 ** attempt
                logger.warning(
                    f"Agent {state.value} failed (attempt {attempt + 1}): {e}"
                )
                ctx.errors.append({
                    "agent": state.value,
                    "attempt": attempt + 1,
                    "error": str(e),
                })
                if attempt < self.MAX_RETRIES - 1:
                    await asyncio.sleep(wait)

        return False

    def _log_audit(
        self, ctx: PipelineContext,
        event: str, details: str
    ):
        ctx.audit_log.append({
            "timestamp": datetime.utcnow().isoformat(),
            "contract_id": ctx.contract_id,
            "event": event,
            "state": ctx.state.value,
            "details": details,
        })

import asyncio
from datetime import datetime, timedelta
from typing import List, Dict
from dataclasses import dataclass
from enum import Enum
import aiohttp


class AlertPriority(Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"


@dataclass
class Alert:
    alert_type: str
    contract_id: str
    priority: AlertPriority
    message: str
    recipients: List[str]
    metadata: Dict


class ContractAlertService:
    """Monitors contracts and fires alerts via multiple channels."""

    EXPIRY_WINDOWS = [90, 60, 30, 7]  # days before expiry

    def __init__(self, db, kafka_producer, slack_client):
        self.db = db
        self.kafka = kafka_producer
        self.slack = slack_client

    async def scan_expiring_contracts(self):
        """Cron job: scan for contracts approaching expiry."""

        for window in self.EXPIRY_WINDOWS:
            target_date = datetime.utcnow() + timedelta(days=window)

            contracts = await self.db.fetch("""
                SELECT id, title, expiry_date, owner_email,
                       contract_value
                FROM contracts
                WHERE expiry_date::date = $1::date
                  AND status = 'executed'
                  AND auto_renew = FALSE
            """, target_date)

            for contract in contracts:
                alert = Alert(
                    alert_type="CONTRACT_EXPIRY",
                    contract_id=str(contract["id"]),
                    priority=(
                        AlertPriority.CRITICAL if window <= 7
                        else AlertPriority.HIGH
                    ),
                    message=(
                        f"Contract '{contract['title']}' expires "
                        f"in {window} days ({contract['expiry_date']})"
                    ),
                    recipients=[contract["owner_email"]],
                    metadata={
                        "days_until_expiry": window,
                        "contract_value": str(
                            contract["contract_value"]
                        ),
                    },
                )
                await self._dispatch_alert(alert)

    async def _dispatch_alert(self, alert: Alert):
        """Route alert to appropriate channels based on priority."""

        # Always publish to Kafka for audit trail
        await self.kafka.send(
            "contract-alerts",
            key=alert.contract_id,
            value=alert.__dict__,
        )

        # Email for all priorities
        await self._send_email(alert)

        # Slack for high and critical
        if alert.priority in (
            AlertPriority.HIGH, AlertPriority.CRITICAL
        ):
            await self._send_slack(alert)

        # PagerDuty for critical only
        if alert.priority == AlertPriority.CRITICAL:
            await self._trigger_pagerduty(alert)

-- Enable pgvector extension
CREATE EXTENSION IF NOT EXISTS vector;

-- ═══════════════════════════════════════════════
-- Core: Contracts
-- ═══════════════════════════════════════════════
CREATE TABLE contracts (
    id              UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    title           VARCHAR(500) NOT NULL,
    contract_type   VARCHAR(50) NOT NULL,  -- NDA, MSA, SOW, etc.
    status          VARCHAR(20) NOT NULL DEFAULT 'draft',
    version         INTEGER NOT NULL DEFAULT 1,
    parent_id       UUID REFERENCES contracts(id),  -- for amendments
    parties         JSONB NOT NULL,  -- [{name, role, contact}]
    effective_date  DATE,
    expiry_date     DATE,
    contract_value  DECIMAL(15,2),
    currency        VARCHAR(3) DEFAULT 'USD',
    jurisdiction    VARCHAR(100),
    governing_law   VARCHAR(100),
    auto_renew      BOOLEAN DEFAULT FALSE,
    renewal_term    INTERVAL,
    raw_text        TEXT,
    file_url        VARCHAR(1000),
    owner_id        UUID NOT NULL,
    owner_email     VARCHAR(255),
    created_at      TIMESTAMPTZ DEFAULT now(),
    updated_at      TIMESTAMPTZ DEFAULT now(),
    deleted_at      TIMESTAMPTZ  -- soft delete
);

CREATE INDEX idx_contracts_status ON contracts(status);
CREATE INDEX idx_contracts_expiry ON contracts(expiry_date)
    WHERE deleted_at IS NULL;
CREATE INDEX idx_contracts_type ON contracts(contract_type);

-- ═══════════════════════════════════════════════
-- Clauses (extracted by LLM)
-- ═══════════════════════════════════════════════
CREATE TABLE clauses (
    id              UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    contract_id     UUID NOT NULL REFERENCES contracts(id),
    clause_type     VARCHAR(50) NOT NULL,
    section_ref     VARCHAR(20),
    text            TEXT NOT NULL,
    page_number     INTEGER,
    risk_indicators JSONB DEFAULT '[]',
    key_dates       JSONB DEFAULT '[]',
    monetary_values JSONB DEFAULT '[]',
    created_at      TIMESTAMPTZ DEFAULT now()
);

CREATE INDEX idx_clauses_contract ON clauses(contract_id);
CREATE INDEX idx_clauses_type ON clauses(clause_type);

-- ═══════════════════════════════════════════════
-- Clause Embeddings (pgvector)
-- ═══════════════════════════════════════════════
CREATE TABLE clause_embeddings (
    id              UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    clause_id       UUID NOT NULL REFERENCES clauses(id),
    contract_id     UUID NOT NULL REFERENCES contracts(id),
    clause_type     VARCHAR(50),
    text            TEXT NOT NULL,
    embedding       vector(1536) NOT NULL,
    created_at      TIMESTAMPTZ DEFAULT now()
);

-- HNSW index for fast approximate nearest neighbor
CREATE INDEX idx_clause_embeddings_hnsw
    ON clause_embeddings
    USING hnsw (embedding vector_cosine_ops)
    WITH (m = 16, ef_construction = 200);

-- GIN index for BM25 full-text search (hybrid retrieval)
CREATE INDEX idx_clause_text_search
    ON clause_embeddings
    USING gin (to_tsvector('english', text));

-- ═══════════════════════════════════════════════
-- Obligations
-- ═══════════════════════════════════════════════
CREATE TABLE obligations (
    id              UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    contract_id     UUID NOT NULL REFERENCES contracts(id),
    clause_id       UUID REFERENCES clauses(id),
    description     TEXT NOT NULL,
    obligated_party VARCHAR(200),
    due_date        DATE,
    recurring       BOOLEAN DEFAULT FALSE,
    recurrence_rule VARCHAR(100),  -- RRULE format
    status          VARCHAR(20) DEFAULT 'pending',
    assigned_to     UUID,
    completed_at    TIMESTAMPTZ,
    created_at      TIMESTAMPTZ DEFAULT now()
);

CREATE INDEX idx_obligations_due ON obligations(due_date)
    WHERE status = 'pending';

-- ═══════════════════════════════════════════════
-- Compliance Checks
-- ═══════════════════════════════════════════════
CREATE TABLE compliance_checks (
    id              UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    contract_id     UUID NOT NULL REFERENCES contracts(id),
    framework       VARCHAR(20) NOT NULL,  -- GDPR, SOX, HIPAA
    rule_id         VARCHAR(50) NOT NULL,
    passed          BOOLEAN NOT NULL,
    confidence      DECIMAL(3,2),
    source          VARCHAR(20),  -- rule_engine | llm | hybrid
    details         TEXT,
    clause_ref      VARCHAR(20),
    reviewed_by     UUID,
    reviewed_at     TIMESTAMPTZ,
    created_at      TIMESTAMPTZ DEFAULT now()
);

-- ═══════════════════════════════════════════════
-- Risk Scores
-- ═══════════════════════════════════════════════
CREATE TABLE risk_scores (
    id              UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    contract_id     UUID NOT NULL REFERENCES contracts(id),
    overall_score   DECIMAL(4,3) NOT NULL,
    category_scores JSONB NOT NULL,
    high_risk_clauses JSONB,
    missing_protections JSONB,
    recommendations JSONB,
    assessed_by     VARCHAR(20) DEFAULT 'system',
    created_at      TIMESTAMPTZ DEFAULT now()
);

CREATE INDEX idx_risk_scores_contract ON risk_scores(contract_id);
CREATE INDEX idx_risk_scores_high ON risk_scores(overall_score)
    WHERE overall_score > 0.7;

-- ═══════════════════════════════════════════════
-- Audit Logs (immutable, append-only)
-- ═══════════════════════════════════════════════
CREATE TABLE audit_logs (
    id              BIGSERIAL PRIMARY KEY,
    contract_id     UUID REFERENCES contracts(id),
    actor_id        UUID,
    actor_type      VARCHAR(20),  -- user | system | agent
    action          VARCHAR(50) NOT NULL,
    details         JSONB,
    ip_address      INET,
    created_at      TIMESTAMPTZ DEFAULT now()
);

-- Partition audit logs by month for performance
CREATE INDEX idx_audit_contract ON audit_logs(contract_id);
CREATE INDEX idx_audit_created ON audit_logs(created_at);