Contract Plain-Language Explainer Agent

import io
import json
import inspect
import sys
import locale
from typing import get_type_hints

# ── Encoding fixes ────────────────────────────────────────────────────────────
locale.getpreferredencoding = lambda do_set_locale=True: 'UTF-8'
if hasattr(sys.stdout, 'reconfigure'):
    sys.stdout.reconfigure(encoding='utf-8', errors='replace')
if hasattr(sys.stderr, 'reconfigure'):
    sys.stderr.reconfigure(encoding='utf-8', errors='replace')

try:
    import httpx.models as _httpx_models
    def _safe_normalize(value, encoding=None):
        if isinstance(value, bytes): return value
        if value is None: return b''
        return str(value).encode(encoding or 'ascii', errors='replace')
    _httpx_models.normalize_header_value = _safe_normalize
except Exception:
    pass

# ── Provider (injected by AgentHTMLGenerator) ─────────────────────────────────
PROVIDER = globals().get('PROVIDER', 'local')

# ── Contract state ────────────────────────────────────────────────────────────
contract_context = {
    'text':          '',
    'loaded':        False,
    'contract_type': 'employment',
    'your_role':     'employee',
    'jurisdiction':  'not specified',
}
conversation_history = []


# ─────────────────────────────────────────────────────────────────────────────
# Schema helpers
# ─────────────────────────────────────────────────────────────────────────────

def python_type_to_json_type(python_type):
    type_mapping = {
        'str': 'string', 'int': 'integer', 'float': 'number',
        'bool': 'boolean', 'list': 'array', 'dict': 'object',
    }
    type_name = python_type.__name__ if hasattr(python_type, '__name__') else str(python_type)
    return type_mapping.get(type_name, 'string')


def extract_function_schema(func):
    sig = inspect.signature(func)
    try:
        hints = get_type_hints(func)
    except Exception:
        hints = {}
    doc = inspect.getdoc(func) or f'Execute {func.__name__}'
    description = doc.split('\n')[0]
    parameters = {'type': 'object', 'properties': {}, 'required': []}
    param_descriptions = {}
    if 'Args:' in doc:
        args_section = doc.split('Args:')[1].split('Returns:')[0] if 'Returns:' in doc else doc.split('Args:')[1]
        for line in args_section.split('\n'):
            line = line.strip()
            if ':' in line:
                param_name = line.split(':')[0].strip()
                param_desc = line.split(':', 1)[1].strip()
                param_descriptions[param_name] = param_desc
    for param_name, param in sig.parameters.items():
        if param_name in ('self', 'cls'):
            continue
        param_type = hints.get(param_name, param.annotation)
        if param_type is inspect.Parameter.empty:
            param_type = str
        json_type = python_type_to_json_type(param_type)
        prop_schema = {
            'type': json_type,
            'description': param_descriptions.get(param_name, f'The {param_name} parameter'),
        }
        parameters['properties'][param_name] = prop_schema
        if param.default is inspect.Parameter.empty:
            parameters['required'].append(param_name)
    return {'name': func.__name__, 'description': description, 'parameters': parameters}


# ─────────────────────────────────────────────────────────────────────────────
# Tool: 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)}'


# ─────────────────────────────────────────────────────────────────────────────
# Tool: load_contract
# ─────────────────────────────────────────────────────────────────────────────

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

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

    Returns:
        Confirmation message with word count
    """
    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.'


# ─────────────────────────────────────────────────────────────────────────────
# Context window budget
#
# Hermes-3-Llama-3.1-8B has a 4096-token context window.
# Budget breakdown (tokens):
#   ~400  system prompt
#   ~100  metadata header (type / role / jurisdiction)
#   ~200  overhead (role tags, formatting)
#   ~800  assistant response headroom
#   ────
#   ~1500 remaining for contract text  →  ~6000 chars (≈4 chars/token)
#
# For follow-up turns the contract is NOT re-sent; only a short summary
# (~300 chars) is prepended so the model remembers what it read.
# ─────────────────────────────────────────────────────────────────────────────

_LOCAL_MAX_CHARS        = 6000   # first-turn full contract cap
_LOCAL_SUMMARY_CHARS    = 300    # follow-up: short reminder snippet
_HISTORY_ASSISTANT_CAP  = 600    # max chars kept per assistant turn in history
_MAX_HISTORY_TURNS      = 2      # how many past user/assistant pairs to include


def _truncate_for_local(text: str, max_chars: int = _LOCAL_MAX_CHARS) -> str:
    """Hard-cap contract text and append an omission notice."""
    if len(text) <= max_chars:
        return text
    cutoff = text.rfind('\n', 0, max_chars)
    if cutoff == -1:
        cutoff = max_chars
    omitted = len(text) - cutoff
    return (
        text[:cutoff]
        + f'\n\n[... {omitted:,} characters omitted to fit local model context window ...]'
    )


def _contract_summary_snippet() -> str:
    """Return a short reminder of the contract for follow-up turns."""
    text = contract_context['text']
    snippet = text[:_LOCAL_SUMMARY_CHARS].strip()
    if len(text) > _LOCAL_SUMMARY_CHARS:
        snippet += ' [...]'
    return snippet


def _build_history_block() -> str:
    """
    Build a compact conversation history string.
    Keeps only the last _MAX_HISTORY_TURNS pairs.
    Truncates long assistant replies to _HISTORY_ASSISTANT_CAP chars.
    """
    if not conversation_history:
        return ''
    # Keep last N pairs (each pair = 1 user + 1 assistant message)
    pairs = []
    history = conversation_history[:]
    while history and len(pairs) < _MAX_HISTORY_TURNS:
        # Walk backwards looking for assistant+user pairs
        if len(history) >= 2 and history[-1]['role'] == 'assistant' and history[-2]['role'] == 'user':
            pairs.insert(0, (history[-2], history[-1]))
            history = history[:-2]
        else:
            history = history[:-1]

    lines = []
    for user_msg, asst_msg in pairs:
        user_content = user_msg['content']
        # Strip the full contract block from history — keep only the question part
        if 'CONTRACT TEXT:' in user_content:
            user_content = user_content.split('CONTRACT TEXT:')[0].strip()
        asst_content = asst_msg['content']
        if len(asst_content) > _HISTORY_ASSISTANT_CAP:
            asst_content = asst_content[:_HISTORY_ASSISTANT_CAP] + ' [...]'
        lines.append(f'User: {user_content}')
        lines.append(f'Assistant: {asst_content}')
    return '\n'.join(lines)


# ─────────────────────────────────────────────────────────────────────────────
# Prompt builders
# ─────────────────────────────────────────────────────────────────────────────

def build_initial_prompt() -> str:
    """
    Build the structured first user message from loaded contract context.
    Truncates contract text automatically for local models.

    Returns:
        Formatted prompt string ready to pass to process_user_query
    """
    is_local = globals().get('PROVIDER', 'local') == 'local'
    body = (
        _truncate_for_local(contract_context['text'])
        if is_local
        else contract_context['text']
    )
    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{body}"
    )


def build_followup_prompt(question: str) -> str:
    """
    Build a follow-up question prompt WITHOUT re-sending the full contract.
    Includes a short contract snippet + compact history so the model has context.

    Args:
        question: The follow-up question from the user

    Returns:
        Compact prompt string safe for a 4096-token context window
    """
    snippet  = _contract_summary_snippet()
    history  = _build_history_block()
    c_type   = contract_context['contract_type']
    role     = contract_context['your_role']
    jur      = contract_context['jurisdiction']

    parts = [
        f"Contract type: {c_type} | Role: {role} | Jurisdiction: {jur}",
        f"Contract snippet for reference:\n{snippet}",
    ]
    if history:
        parts.append(f"Previous exchanges:\n{history}")
    parts.append(f"Follow-up question: {question}")
    return '\n\n'.join(parts)


# ─────────────────────────────────────────────────────────────────────────────
# Lawyer questions — pure logic, no LLM call
# ─────────────────────────────────────────────────────────────────────────────

def get_lawyer_questions() -> str:
    """
    Generate relevant lawyer questions based on contract type and role.

    Returns:
        Numbered list of questions as a plain string
    """
    c_type = contract_context.get('contract_type', 'employment')
    jur    = contract_context.get('jurisdiction',  'not specified')

    type_questions = {
        'employment': [
            'Is the non-compete or non-solicitation clause enforceable and proportionate to my role?',
            'What are my rights regarding intellectual property I create outside of work hours?',
            'What severance or notice period am I entitled to if terminated without cause?',
        ],
        'lease': [
            'What are my rights if the landlord wants to enter the property or sell it?',
            'Are there hidden costs or maintenance obligations not immediately obvious?',
            'What are the exact conditions under which I can be evicted?',
        ],
        'nda': [
            'Is the definition of "confidential information" too broad to be practical?',
            'How long does the confidentiality obligation last after the agreement ends?',
            'Are there exceptions to the confidentiality obligation I should know about?',
        ],
        'freelance': [
            'Who owns the intellectual property I create under this contract?',
            'What is the payment schedule and what happens if payment is delayed?',
            'Under what conditions can either party terminate the engagement early?',
        ],
        'saas': [
            'What data do you collect and how is it stored or shared?',
            'What happens to my data if I cancel the subscription?',
            'Are there automatic price increases or renewal lock-ins?',
        ],
        'purchase': [
            'What warranties or guarantees are included, and for how long?',
            'What are the return, refund, or dispute resolution procedures?',
            'Are there any ongoing obligations or fees after the initial purchase?',
        ],
    }

    base_questions = [
        f'Are there any clauses unenforceable under {jur} law?',
        'What is the dispute resolution process if something goes wrong?',
        'Are there automatic renewal or lock-in clauses I should be aware of?',
        'What happens to my rights if this contract is assigned to a third party?',
    ]

    specific = type_questions.get(c_type, [
        'What are my core obligations under this contract?',
        'What happens if either party fails to meet their obligations?',
        'Are there any unusual or one-sided clauses I should negotiate?',
    ])

    all_questions = specific + base_questions
    return '\n'.join(f'{i+1}. {q}' for i, q in enumerate(all_questions))


# ─────────────────────────────────────────────────────────────────────────────
# Main query processor
# ─────────────────────────────────────────────────────────────────────────────

async def process_user_query(query: str) -> str:
    """
    Process a user question about the loaded contract.
    Routes to local WebLLM or cloud provider (OpenAI / Anthropic).
    First turn receives the full contract; follow-up turns get a compact prompt.

    Args:
        query: The prompt string built by build_initial_prompt() or build_followup_prompt()

    Returns:
        AI response as a plain string
    """
    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.'
        )

    conversation_history.append({'role': 'user', 'content': query})

    # ── Local WebLLM ──────────────────────────────────────────────────────────
    if provider == 'local':
        try:
            system_prompt = globals().get('TEMPLATE_SYSTEM_PROMPT', '')
            result = await process_user_query_wllama(query, 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 ───────────────────────────────────────────────────────
    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}'

        system_prompt     = globals().get('TEMPLATE_SYSTEM_PROMPT', '')
        few_shot_examples = globals().get('AGENT_FEW_SHOT_EXAMPLES', [])

        messages = [SystemMessage(content=system_prompt)]
        for ex in few_shot_examples:
            if ex.get('role') == 'user':
                messages.append(HumanMessage(content=ex['content']))
            else:
                messages.append(AIMessage(content=ex['content']))

        # For cloud: include full history — no token budget concern
        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)}'


# ─────────────────────────────────────────────────────────────────────────────
# Tool schema export
# ─────────────────────────────────────────────────────────────────────────────

def get_tool_schemas() -> str:
    """
    Export all tool function schemas in OpenAI format.
    Called by PyodideToolBridge in JS to discover available tools.

    Returns:
        JSON string with OpenAI-format tool schemas
    """
    tool_functions = [load_contract]
    schemas = []
    for func in tool_functions:
        try:
            function_schema = extract_function_schema(func)
            openai_schema = {'type': 'function', 'function': function_schema}
            schemas.append(openai_schema)
        except Exception as e:
            print(f'Warning: Failed to extract schema for {func.__name__}: {e}')
    return json.dumps(schemas, indent=2)


print(f'✅ Contract Plain Language Explainer initialized (provider: {PROVIDER})')