Contract Plain-Language Explainer Agent

import io
import json

# ============================================================================
# Contract State
# ============================================================================

contract_context = {
    "text": "",
    "loaded": False,
    "contract_type": "employment",
    "your_role": "employee",
    "jurisdiction": "not specified"
}

conversation_history = []

# ============================================================================
# System Prompt & Few-Shot Examples
# ============================================================================

TEMPLATE_SYSTEM_PROMPT = """You are a plain-language contract writing assistant. You will receive pre-computed contract analysis results. Your only job is to write clear, readable output.

Rules:
- Follow the output structure exactly as given in the user message
- Write at a Grade 8 reading level — explain every legal term
- Always quote contract text in _italics_ before explaining it
- Use [HIGH RISK] / [MEDIUM RISK] / [LOW RISK] labels exactly as shown
- Never fabricate clauses, risks, or quotes not present in the analysis results
- Never call tools — all analysis is already done
- Keep each clause explanation to 1-2 sentences maximum"""

FEW_SHOT_EXAMPLES = [
    {
        "role": "user",
        "content": (
            "Contract type: employment\n"
            "My role: employee\n"
            "Jurisdiction: Netherlands\n\n"
            "CONTRACT TEXT:\n"
            "The Employee shall not, for a period of 12 months following termination, "
            "directly or indirectly engage in any business that competes with the Company "
            "within the Netherlands."
        )
    },
    {
        "role": "assistant",
        "content": (
            "## \U0001f4cb Plain-Language Summary\n"
            "This is a non-compete clause from an employment contract. As an employee in the Netherlands, "
            "you are agreeing not to work for a competing business for 12 months after you leave. "
            "The overall commitment is medium-to-high: Dutch law does allow non-competes but courts frequently limit their scope.\n\n"
            "## \u2696\ufe0f Clauses Found \u2014 Plain English\n"
            "**Non-compete** \u2014 _\u2018shall not...directly or indirectly engage in any business that competes\u2019_ "
            "\u2014 After leaving this job, you cannot work for a competitor or start a competing business for 12 months inside the Netherlands.\n\n"
            "## \U0001f6a8 Risk Assessment\n"
            "[HIGH RISK] **Non-Compete** \u2014 For an employee, a broad non-compete can prevent you from working in your field for a year, directly impacting your income.\n\n"
            "## \U0001f50d Missing Clauses\n"
            "- **Severance / compensation** \u2014 Dutch law (Article 7:653 BW) may require financial compensation for a non-compete clause; its absence may make this clause unenforceable.\n"
            "- **Governing law** \u2014 Not explicitly stated; important for knowing which court handles disputes.\n\n"
            "## \u2753 Questions for Your Lawyer\n"
            "1. Is this non-compete enforceable under Dutch law (Art. 7:653 BW) without a compensation clause?\n"
            "2. Is the geographic scope (Netherlands only) and duration (12 months) proportionate to my role?\n"
            "3. Does this prevent me from freelancing in the same industry, or only working for direct competitors?\n"
            "4. Are there any clauses here that would be unenforceable under local consumer/employment law?\n\n"
            "---\n"
            "\u26a0\ufe0f **Not legal advice.** Always consult a qualified lawyer before signing."
        )
    },
    {
        "role": "user",
        "content": "What does the liability cap clause mean for me?"
    },
    {
        "role": "assistant",
        "content": (
            "Based on the contract we analyzed, the liability cap clause limits the maximum amount "
            "either party can claim from the other in a dispute.\n\n"
            "In plain terms: even if you suffer a large loss caused by the other party, you can only "
            "claim up to the capped amount \u2014 typically 3 months of fees or the total contract value, whichever is lower.\n\n"
            "**Why it matters for you as a contractor [MEDIUM RISK]:**\n"
            "- If a client\u2019s project failure causes you significant damage (e.g. loss of future income), "
            "you may be unable to recover the full amount.\n"
            "- The cap also protects *you* \u2014 if a client claims you caused them a loss, your exposure is limited.\n\n"
            "**Key question to ask your lawyer:**\n"
            "> Is the liability cap in this contract proportionate to the value and risk of this engagement, "
            "and does it comply with applicable consumer/commercial law?\n\n"
            "---\n"
            "\u26a0\ufe0f **Not legal advice.** Always consult a qualified lawyer before signing."
        )
    }
]

# ============================================================================
# PDF Extraction
# ============================================================================

async def extract_pdf_text(pdf_bytes: bytes) -> str:
    """
    Extract all text from a PDF file given its raw bytes.

    Args:
        pdf_bytes: Raw bytes of the PDF file

    Returns:
        Extracted plain text from all pages
    """
    try:
        import micropip
        await micropip.install("pypdf")
        from pypdf import PdfReader
        reader = PdfReader(io.BytesIO(pdf_bytes))
        pages_text = []
        for i, page in enumerate(reader.pages):
            text = page.extract_text()
            if text and text.strip():
                pages_text.append(f"--- Page {i+1} ---\n{text.strip()}")
        if not pages_text:
            return "⚠️ No readable text found in this PDF. It may be a scanned/image-based PDF."
        return "\n\n".join(pages_text)
    except Exception as e:
        return f"❌ Error extracting PDF text: {str(e)}"


# ============================================================================
# Contract Loading
# ============================================================================

def load_contract(text: str, contract_type: str, your_role: str, jurisdiction: str) -> str:
    """
    Load contract text and metadata into context.

    Args:
        text: Full contract text
        contract_type: Type of contract e.g. employment, lease, NDA
        your_role: User role e.g. employee, tenant, contractor
        jurisdiction: Legal jurisdiction e.g. Netherlands, UK, US-CA

    Returns:
        Confirmation message
    """
    contract_context["text"] = text.strip()
    contract_context["loaded"] = True
    contract_context["contract_type"] = contract_type or "employment"
    contract_context["your_role"] = your_role or "employee"
    contract_context["jurisdiction"] = jurisdiction or "not specified"
    conversation_history.clear()
    word_count = len(text.split())
    return f"✅ Contract loaded ({word_count} words). Ready for analysis."


# ============================================================================
# Prompt Builder
# ============================================================================

def _build_initial_prompt() -> str:
    """Build the structured first user message using the exact prompt template."""
    return (
        f"Contract type: {contract_context['contract_type']}\n"
        f"My role in this contract: {contract_context['your_role']}\n"
        f"Jurisdiction: {contract_context['jurisdiction']}\n\n"
        f"CONTRACT TEXT:\n{contract_context['text']}"
    )


# ============================================================================
# Query Processing
# ============================================================================

async def process_user_query(query: str) -> str:
    """
    Process a user question about the loaded contract.
    Maintains full conversation history for follow-up questions.

    Args:
        query: User question, or 'INITIAL_ANALYSIS' to trigger first summary

    Returns:
        AI response grounded in the contract
    """
    provider = globals().get('PROVIDER', 'local')

    if not contract_context["loaded"]:
        return "⚠️ No contract loaded yet. Please upload a PDF or paste contract text and click **Load Contract** first."

    # Use structured prompt template for initial analysis
    actual_query = _build_initial_prompt() if query == "INITIAL_ANALYSIS" else query

    conversation_history.append({"role": "user", "content": actual_query})

    # ── Local WebLLM ──────────────────────────────────────────────────────────
    if provider == 'local':
        try:
            # Build prior conversation as plain text for WebLLM context window
            history_text = ""
            for msg in conversation_history[:-1]:
                role = "User" if msg["role"] == "user" else "Assistant"
                history_text += f"{role}: {msg['content']}\n\n"

            full_query = actual_query
            if history_text:
                full_query = f"[Conversation so far]\n{history_text}[Current question]\n{actual_query}"

            result = await process_user_query_webllm(full_query, TEMPLATE_SYSTEM_PROMPT)
            conversation_history.append({"role": "assistant", "content": result})
            return result
        except Exception as e:
            import traceback
            traceback.print_exc()
            conversation_history.pop()
            return f"❌ Error using WebLLM: {str(e)}"

    # ── Cloud Providers (OpenAI / Anthropic) ──────────────────────────────────
    else:
        try:
            from langchain_core.messages import HumanMessage, AIMessage, SystemMessage

            api_key = globals().get('current_api_key', '') or globals().get('api_key', '')
            if not api_key:
                conversation_history.pop()
                return "⚠️ Please enter an API key to use this agent."

            if provider == 'openai':
                from langchain_openai import ChatOpenAI
                llm = ChatOpenAI(
                    model="gpt-4o-mini",
                    api_key=api_key,
                    temperature=0.3
                )
            elif provider == 'anthropic':
                from langchain_anthropic import ChatAnthropic
                llm = ChatAnthropic(
                    model="claude-3-5-sonnet-20241022",
                    api_key=api_key,
                    temperature=0.3
                )
            else:
                conversation_history.pop()
                return f"❌ Unknown provider: {provider}"

            # Build full message list:
            # system prompt → few-shot examples → live conversation history
            messages = [SystemMessage(content=TEMPLATE_SYSTEM_PROMPT)]

            for ex in FEW_SHOT_EXAMPLES:
                if ex["role"] == "user":
                    messages.append(HumanMessage(content=ex["content"]))
                else:
                    messages.append(AIMessage(content=ex["content"]))

            for msg in conversation_history:
                if msg["role"] == "user":
                    messages.append(HumanMessage(content=msg["content"]))
                else:
                    messages.append(AIMessage(content=msg["content"]))

            response = llm.invoke(messages)
            assistant_reply = response.content
            conversation_history.append({"role": "assistant", "content": assistant_reply})
            return assistant_reply

        except Exception as e:
            import traceback
            traceback.print_exc()
            conversation_history.pop()
            return f"❌ Error: {str(e)}"


print("✅ Contract Plain Language Explainer initialized")