Upload drafting_engine.py

drafting_engine.py

@@ -0,0 +1,296 @@
+"""
+Contract Drafting Engine.
+Orchestrates clause retrieval, playbook rules, fallback positions,
+risk flags, drafting checklist, and verifier pass.
+"""
+
+import json
+from typing import List, Dict, Optional, Any
+from dataclasses import dataclass, asdict
+
+from playbook import (
+    get_required_clauses,
+    get_fallback_position,
+    get_risk_flags,
+    get_checklist,
+)
+from clause_retriever import ClauseRetriever
+
+
+@dataclass
+class DraftingContext:
+    contract_type: str
+    party_position: str  # pro_company, pro_counterparty, balanced
+    deal_context: str
+    business_constraints: List[str]
+    governing_law: Optional[str] = None
+    counterparty_name: Optional[str] = None
+    company_name: Optional[str] = None
+    deal_value: Optional[str] = None
+    term_length: Optional[str] = None
+
+
+@dataclass
+class DraftedClause:
+    clause_name: str
+    clause_text: str
+    source: str
+    fallback_applied: bool
+    risk_flags: List[Dict[str, str]]
+    checklist_items: List[str]
+    retrieved_clauses: List[Dict]
+
+
+@dataclass
+class DraftedContract:
+    contract_type: str
+    context: DraftingContext
+    clauses: List[DraftedClause]
+    risk_flags: List[Dict[str, Any]]
+    checklist: List[Dict[str, Any]]
+    verifier_notes: List[str]
+
+
+class ContractDraftingEngine:
+    def __init__(
+        self,
+        retriever: Optional[ClauseRetriever] = None,
+        generator_model_name: Optional[str] = None,
+    ):
+        self.retriever = retriever or ClauseRetriever()
+        self.generator_model_name = generator_model_name
+        self.generator = None
+        if generator_model_name:
+            try:
+                from transformers import pipeline
+                self.generator = pipeline(
+                    "text-generation",
+                    model=generator_model_name,
+                    device_map="auto",
+                )
+            except Exception as e:
+                print(f"Warning: could not load generator {generator_model_name}: {e}")
+
+    def draft(self, context: DraftingContext) -> DraftedContract:
+        required = get_required_clauses(context.contract_type)
+        checklist = get_checklist(context.contract_type)
+        drafted_clauses: List[DraftedClause] = []
+        all_risk_flags: List[Dict[str, Any]] = []
+
+        for clause_name in required:
+            # Retrieve precedent clauses
+            query = f"{clause_name.replace('_', ' ')} clause for {context.contract_type.replace('_', ' ')}"
+            retrieved = self.retriever.retrieve(
+                query=query,
+                clause_type=clause_name,
+                top_k=3,
+            )
+
+            # Get fallback position
+            fallback = get_fallback_position(clause_name, context.party_position)
+
+            # Generate clause text
+            clause_text = self._generate_clause(
+                clause_name=clause_name,
+                context=context,
+                retrieved_clauses=retrieved,
+                fallback=fallback,
+            )
+
+            # Get risk flags
+            flags = get_risk_flags(clause_name)
+            # Apply contextual risk analysis
+            active_flags = self._evaluate_risk_flags(clause_text, flags, context)
+            all_risk_flags.extend([
+                {"clause": clause_name, **f} for f in active_flags
+            ])
+
+            # Checklist items for this clause
+            clause_checklist = [
+                c["item"] for c in checklist
+                if clause_name.replace("_", " ") in c["item"].lower()
+                or c["category"] in clause_name
+            ]
+
+            drafted_clauses.append(DraftedClause(
+                clause_name=clause_name,
+                clause_text=clause_text,
+                source="retrieved+generated",
+                fallback_applied=fallback is not None,
+                risk_flags=active_flags,
+                checklist_items=clause_checklist,
+                retrieved_clauses=retrieved,
+            ))
+
+        # Verifier pass
+        verifier_notes = self._verifier_pass(drafted_clauses, context)
+
+        return DraftedContract(
+            contract_type=context.contract_type,
+            context=context,
+            clauses=drafted_clauses,
+            risk_flags=all_risk_flags,
+            checklist=[{"item": c["item"], "category": c["category"], "checked": False} for c in checklist],
+            verifier_notes=verifier_notes,
+        )
+
+    def _generate_clause(
+        self,
+        clause_name: str,
+        context: DraftingContext,
+        retrieved_clauses: List[Dict],
+        fallback: Optional[Dict[str, str]],
+    ) -> str:
+        # Build prompt from retrieved clauses + playbook
+        prompt_parts = [
+            f"Draft a {clause_name.replace('_', ' ')} clause for a {context.contract_type.replace('_', ' ')}.",
+            f"Party position: {context.party_position}.",
+            f"Deal context: {context.deal_context}",
+        ]
+        if fallback:
+            prompt_parts.append(f"Fallback position: {json.dumps(fallback)}")
+        if retrieved_clauses:
+            prompt_parts.append("Precedent clauses:")
+            for rc in retrieved_clauses:
+                prompt_parts.append(f"- {rc['clause_text'][:500]}")
+        prompt = "\n".join(prompt_parts)
+
+        if self.generator:
+            try:
+                out = self.generator(
+                    prompt,
+                    max_new_tokens=512,
+                    do_sample=True,
+                    temperature=0.3,
+                )
+                return out[0]["generated_text"][len(prompt):].strip()
+            except Exception as e:
+                print(f"Generation failed for {clause_name}: {e}")
+
+        # Fallback: template-based generation
+        return self._template_clause(clause_name, context, fallback)
+
+    def _template_clause(
+        self,
+        clause_name: str,
+        context: DraftingContext,
+        fallback: Optional[Dict[str, str]],
+    ) -> str:
+        """Simple template-based clause generation when no LLM is available."""
+        templates = {
+            "limitation_of_liability": (
+                "LIMITATION OF LIABILITY. "
+                "Except for breaches of confidentiality, IP infringement, or gross negligence, "
+                "each party's aggregate liability arising out of this agreement shall not exceed "
+                f"{fallback.get('cap', 'the fees paid in the 12 months preceding the claim') if fallback else 'the fees paid in the 12 months preceding the claim'}."
+            ),
+            "indemnification": (
+                "INDEMNIFICATION. "
+                f"{context.company_name or 'Company'} shall indemnify {context.counterparty_name or 'Counterparty'} against third-party claims arising from "
+                f"{fallback.get('scope', 'IP infringement and breach of confidentiality') if fallback else 'IP infringement and breach of confidentiality'}."
+            ),
+            "data_protection": (
+                "DATA PROTECTION. "
+                "Each party shall process personal data in accordance with applicable data protection laws. "
+                f"{fallback.get('role', 'The parties shall act as independent controllers') if fallback else 'The parties shall act as independent controllers'}."
+            ),
+            "termination": (
+                "TERMINATION. "
+                f"Either party may terminate this agreement {fallback.get('for_convenience', 'for convenience with 60 days notice') if fallback else 'for convenience with 60 days notice'}. "
+                "Upon termination, all fees owed survive and data shall be returned within 30 days."
+            ),
+            "intellectual_property": (
+                "INTELLECTUAL PROPERTY. "
+                f"{fallback.get('ownership', 'Each party retains its pre-existing IP') if fallback else 'Each party retains its pre-existing IP'}. "
+                "All custom deliverables shall be owned as specified in the applicable SOW."
+            ),
+            "confidentiality": (
+                "CONFIDENTIALITY. "
+                "Each party agrees to hold all Confidential Information in strict confidence and not disclose it to any third parties except as required by law."
+            ),
+            "governing_law": (
+                f"GOVERNING LAW. This agreement shall be governed by the laws of {context.governing_law or 'the State of Delaware'}, without regard to conflict of laws principles."
+            ),
+        }
+        return templates.get(
+            clause_name,
+            f"[{clause_name.replace('_', ' ').title()}]. [Placeholder clause for {context.contract_type}.]"
+        )
+
+    def _evaluate_risk_flags(
+        self,
+        clause_text: str,
+        flags: List[Dict[str, str]],
+        context: DraftingContext,
+    ) -> List[Dict[str, str]]:
+        active = []
+        text_lower = clause_text.lower()
+        for flag in flags:
+            if flag["flag"] == "NO_CAP" and "cap" not in text_lower and "limited" not in text_lower:
+                active.append(flag)
+            elif flag["flag"] == "NO_IP_CARVEOUT" and "intellectual property" not in text_lower and "ip" not in text_lower:
+                active.append(flag)
+            elif flag["flag"] == "NO_DPA" and "data processing" not in text_lower and "dpa" not in text_lower:
+                active.append(flag)
+            elif flag["flag"] == "NO_CURE_PERIOD" and "cure" not in text_lower:
+                active.append(flag)
+            elif flag["flag"] == "NO_DATA_RETURN" and "return" not in text_lower and "delete" not in text_lower:
+                active.append(flag)
+            elif flag["flag"] == "NO_MUTUALITY" and "mutual" not in text_lower:
+                active.append(flag)
+        return active
+
+    def _verifier_pass(
+        self,
+        clauses: List[DraftedClause],
+        context: DraftingContext,
+    ) -> List[str]:
+        notes = []
+        clause_names = {c.clause_name for c in clauses}
+        required = set(get_required_clauses(context.contract_type))
+        missing = required - clause_names
+        if missing:
+            notes.append(f"MISSING CLAUSES: {', '.join(missing)}")
+
+        # Check internal consistency
+        has_limitation = any(c.clause_name == "limitation_of_liability" for c in clauses)
+        has_indemnity = any(c.clause_name == "indemnification" for c in clauses)
+        if has_limitation and has_indemnity:
+            notes.append("PASS: Limitation and indemnification both present.")
+        if not has_limitation:
+            notes.append("WARNING: No limitation of liability clause.")
+
+        # Check for invented terms (basic heuristic)
+        for c in clauses:
+            if "[Placeholder" in c.clause_text:
+                notes.append(f"WARNING: {c.clause_name} contains placeholder text.")
+
+        return notes
+
+    def export(self, contract: DraftedContract, fmt: str = "json") -> str:
+        if fmt == "json":
+            return json.dumps(asdict(contract), indent=2)
+        elif fmt == "markdown":
+            lines = [
+                f"# {contract.contract_type.replace('_', ' ').title()}",
+                "",
+                "## Context",
+                f"- Party position: {contract.context.party_position}",
+                f"- Deal context: {contract.context.deal_context}",
+                "",
+                "## Clauses",
+            ]
+            for c in contract.clauses:
+                lines.append(f"### {c.clause_name.replace('_', ' ').title()}")
+                lines.append(c.clause_text)
+                lines.append("")
+            lines.append("## Risk Flags")
+            for rf in contract.risk_flags:
+                lines.append(f"- **{rf['severity']}** [{rf['clause']}]: {rf['description']}")
+            lines.append("")
+            lines.append("## Verifier Notes")
+            for note in contract.verifier_notes:
+                lines.append(f"- {note}")
+            return "\n".join(lines)
+        else:
+            raise ValueError(f"Unknown format: {fmt}")