astraea/documents.py

# -*- coding: utf-8 -*-
"""Astraea document generation — prep Washington divorce paperwork on demand.
Each doc type carries its real WA form structure; the LLM fills it with the user's facts.
"""
import urllib.request
import json

OLLAMA_URL = "http://10.30.20.29:11434"

# WA family-law paperwork the specialist can draft. `sections` mirrors the real form.
DOC_TYPES = {
    "petition": {
        "name": "Petition for Dissolution of Marriage",
        "form": "FL Divorce 201",
        "description": "The filing that starts a divorce — party info, marriage details, grounds, and requests.",
        "sections": [
            "Caption (Superior Court of Washington, County of ___)",
            "Parties (Petitioner / Respondent names)",
            "Marriage information (date + place of marriage)",
            "Residency basis (RCW 26.09.030)",
            "Grounds: irretrievably broken",
            "Children of the marriage (names, ages)",
            "Requests: dissolution, property division, maintenance, parenting plan, fees",
        ],
    },
    "summons": {
        "name": "Summons (Dissolution)",
        "form": "FL Divorce 200",
        "description": "Formal notice to the other spouse that a case has been filed and they must respond.",
        "sections": [
            "Caption", "Notice of case filing", "Response deadline (20/60 days)",
            "Warning: default judgment if no response", "Certificate of service",
        ],
    },
    "parenting-plan": {
        "name": "Parenting Plan",
        "form": "FL All Family 140",
        "description": "The residential schedule and decision-making plan (RCW 26.09.184).",
        "sections": [
            "Children covered", "Residential schedule (school year / holidays / summer)",
            "Decision-making (joint or sole, RCW 26.09.187)", "Dispute resolution",
            "Relocation restrictions (RCW 26.09.405-.560)", "Safety provisions if DV (RCW 26.09.191)",
        ],
    },
    "child-support": {
        "name": "Child Support Order + Worksheet",
        "form": "WSCSS / FL All Family 130",
        "description": "The income-shares support calculation under RCW 26.19.",
        "sections": [
            "Combined monthly net income (each parent)", "Economic table lookup (RCW 26.19.020)",
            "Basic support obligation", "Proportionate share split",
            "Transfer payment amount", "Healthcare + special expenses", "Deviations (if any)",
        ],
    },
    "financial": {
        "name": "Financial Declaration",
        "form": "FL All Family 131",
        "description": "Full disclosure of income, assets, debts, and expenses.",
        "sections": ["Income (all sources)", "Assets (real + personal)", "Debts/liabilities",
                    "Monthly living expenses", "Separate property claims"],
    },
    "final-order": {
        "name": "Findings of Fact, Conclusions of Law & Decree",
        "form": "FL Divorce 241",
        "description": "The final order that ends the marriage and sets the terms.",
        "sections": [
            "Findings of fact (residency, marriage, children)", "Conclusions of law (grounds)",
            "Property + debt division (RCW 26.09.080)", "Maintenance award (RCW 26.09.090)",
            "Parenting plan incorporation", "Child support order", "Name restoration",
            "Attorney fees (RCW 26.09.140)",
        ],
    },
    "settlement": {
        "name": "Separation / Settlement Agreement",
        "form": "RCW 26.09.070",
        "description": "A binding written agreement resolving all issues without trial.",
        "sections": ["Recitals", "Property + debt division", "Maintenance",
                    "Parenting + support", "Disclosure of assets", "Mutual release",
                    "Enforcement + severability"],
    },
    "confidential": {
        "name": "Confidential Information Form",
        "form": "FL All Family 001",
        "description": "Sealed identifying info (SSNs, DOBs) filed with every case.",
        "sections": ["Party names + DOB + SSN", "Children (names/DOB)", "Protective address (if DV)"],
    },
    "letter": {
        "name": "Formal Correspondence",
        "form": "custom",
        "description": "A professional letter to a spouse, the court, or counsel.",
        "sections": ["Date + address", "Salutation", "Body (purpose)", "Closing + signature"],
    },
}


def build_prompt(doc_type, about_me, instructions, context):
    dt = DOC_TYPES.get(doc_type)
    if not dt:
        return None, None
    sections = "\n".join(f"- {s}" for s in dt["sections"])
    system = (
        "You are a meticulous Washington State family-law paralegal preparing court documents. "
        "Draft the requested document as a clean, complete, court-ready text. Use the user's "
        "ABOUT ME details and any reference context to fill in every field you can; where a "
        "fact is unknown, insert a blank '[____]' placeholder — never invent names, dates, or "
        "numbers. Follow the real Washington form structure precisely. Do not add legal advice "
        "to the document itself."
    )
    user = (
        f"DOCUMENT: {dt['name']} (form {dt['form']})\n"
        f"PURPOSE: {dt['description']}\n\n"
        f"REQUIRED SECTIONS:\n{sections}\n\n"
        f"ABOUT ME (the client):\n{about_me or '(none provided)'}\n\n"
        f"SPECIFIC INSTRUCTIONS:\n{instructions or '(draft a complete standard version)'}\n\n"
        f"REFERENCE (Washington law):\n{context}\n\n"
        "Draft the full document now, formatted for print."
    )
    return system, user


def generate(doc_type, about_me, instructions, context):
    system, user = build_prompt(doc_type, about_me, instructions, context)
    if not system:
        return None
    payload = {
        "model": "ornith-1.5:9b",
        "messages": [{"role": "system", "content": system},
                     {"role": "user", "content": user}],
        "stream": False, "think": False,
        "options": {"temperature": 0.2, "num_predict": 3000, "num_ctx": 24000},
    }
    req = urllib.request.Request(f"{OLLAMA_URL}/api/chat",
                                 data=json.dumps(payload).encode(),
                                 headers={"Content-Type": "application/json"})
    op = urllib.request.build_opener(urllib.request.ProxyHandler({}))
    with op.open(req, timeout=360) as r:
        d = json.loads(r.read().decode("utf-8"))
    m = d.get("message", {})
    return (m.get("content") or m.get("thinking") or "").strip()