File size: 11,971 Bytes
2f02fd3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
"""
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}")