chat_config.py

"""Conversational assistant that turns free-form instructions into
extraction settings (categories, notification keywords, value fields,
export filter) using a Hugging Face-hosted chat model.
"""
import json
import os
import re

from huggingface_hub import InferenceClient

import extractors

CATEGORY_CHOICES = list(extractors.KEYWORD_SETS.keys())

# Tried in order until one is accepted by an inference provider. An
# HF_CHAT_MODEL env var/Space variable always takes first priority.
CANDIDATE_MODELS = [
    m for m in [
        os.environ.get("HF_CHAT_MODEL"),
        "Qwen/Qwen2.5-7B-Instruct",
        "mistralai/Mistral-7B-Instruct-v0.3",
        "Qwen/Qwen2.5-72B-Instruct",
    ] if m
]

DEFAULT_MODEL = CANDIDATE_MODELS[0]

# Set to whichever model actually served the most recent successful call,
# so the run-metadata audit trail records the real model used.
ACTIVE_MODEL = DEFAULT_MODEL

SYSTEM_PROMPT = f"""You are a configuration assistant for a document OCR/data-extraction tool.

The tool distinguishes two different kinds of things a user can ask for:

1. NOTIFICATIONS - keyword/phrase hits that flag an event worth a reviewer's
   attention (a boolean "this happened somewhere on the page"). Built-in
   category keyword sets:
   - construction: change order, delay notice, punch list, non-conformance,
     RFI, submittal rejected, back charge, stop work, safety violation,
     schedule slip, etc.
   - compliance: non-compliant, violation, deficiency, corrective action,
     audit finding, expired, recall, citation, penalty, failed inspection.
   - finance: past due, overdue, payment rejected, invoice disputed, credit
     hold, chargeback, insufficient funds, write-off, collections, late fee.
   - general: urgent, cancelled, rejected, approved, pending approval, on hold.
   Users can also name their own free-text notification phrases.

2. FIELDS - labeled values to pull out of the document as data, not just
   flag. Two kinds:
   - Standard fields, auto-scoped to the categories selected above: PO
     numbers + SAP document numbers (construction), certification codes
     ISO/OSHA/ASTM/ANSI (compliance), invoice numbers + percentages
     (finance), and dates/dollar amounts wherever relevant.
   - Custom fields: any label the user names that appears in the document
     followed by a value, e.g. "reference number", "gross weight",
     "delivery date", "customer PO". These are looked up generically by
     label text, so use the user's own wording for the label.

The user will describe in plain language what they want flagged vs. pulled
as data. Given their message and the current settings, respond with:
1. A short (1-3 sentence) conversational reply confirming what you changed,
   or asking a clarifying question if the request is ambiguous.
2. On its own line, a fenced json block with the FULL updated settings:
```json
{{"categories": ["finance"], "custom_keywords": ["warranty claim"], "custom_fields": ["reference number", "gross weight"], "only_matches": true}}
```

Rules:
- "categories" must only contain values from {CATEGORY_CHOICES}.
- "custom_keywords" is a list of extra free-text NOTIFICATION phrases to
  flag (empty list if none requested).
- "custom_fields" is a list of label names to extract as VALUES, in the
  user's own wording (empty list if none requested).
- "only_matches" is true if only pages with a hit/field should be
  exported, false to export every page regardless.
- Always include all four keys, carrying over prior values the user
  didn't ask to change.
- If the user names something ambiguous, prefer treating it as a
  custom_field when they clearly want a value (a number, date, code) back,
  and as a custom_keyword when they clearly want an alert/flag instead.
- Never invent a category or capability that isn't listed above.
"""


def _extract_json_block(text):
    match = re.search(r"```json\s*(\{.*?\})\s*```", text, re.DOTALL)
    if not match:
        match = re.search(r"(\{.*\})", text, re.DOTALL)
    if not match:
        return None
    try:
        return json.loads(match.group(1))
    except json.JSONDecodeError:
        return None


def _sanitize_settings(parsed, current_settings):
    categories = current_settings.get("categories", [])
    custom_keywords = current_settings.get("custom_keywords", [])
    custom_fields = current_settings.get("custom_fields", [])
    only_matches = current_settings.get("only_matches", True)

    if isinstance(parsed, dict):
        if isinstance(parsed.get("categories"), list):
            categories = [c for c in parsed["categories"] if c in CATEGORY_CHOICES]
        if isinstance(parsed.get("custom_keywords"), list):
            custom_keywords = [str(k).strip() for k in parsed["custom_keywords"] if str(k).strip()]
        if isinstance(parsed.get("custom_fields"), list):
            custom_fields = [str(k).strip() for k in parsed["custom_fields"] if str(k).strip()]
        if isinstance(parsed.get("only_matches"), bool):
            only_matches = parsed["only_matches"]

    return {
        "categories": categories,
        "custom_keywords": custom_keywords,
        "custom_fields": custom_fields,
        "only_matches": only_matches,
    }


def chat_update(history, user_message, current_settings, hf_token=None):
    """Send the conversation to the HF-hosted model and return (reply, new_settings).

    `history` is a list of {"role": ..., "content": ...} dicts (gr.Chatbot's
    "messages" format). On any failure, the reply explains the error and
    current_settings are returned unchanged.
    """
    token = hf_token or os.environ.get("HF_TOKEN")
    if not token:
        return (
            "I need a Hugging Face token to reach the chat model. Set the "
            "HF_TOKEN secret on this Space, or paste a token in the field above.",
            current_settings,
        )

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    messages.append({
        "role": "system",
        "content": f"Current settings: {json.dumps(current_settings)}",
    })
    messages.extend(history)
    messages.append({"role": "user", "content": user_message})

    global ACTIVE_MODEL
    client = InferenceClient(token=token, provider="auto")
    content = None
    last_exc = None
    for model in CANDIDATE_MODELS:
        try:
            response = client.chat_completion(
                messages=messages, model=model, max_tokens=400, temperature=0.2
            )
            content = response.choices[0].message.content
            ACTIVE_MODEL = model
            break
        except Exception as exc:
            last_exc = exc
    if content is None:
        return (
            f"The chat model call failed for every candidate model "
            f"({', '.join(CANDIDATE_MODELS)}). Last error: {last_exc}. "
            "Check that HF_TOKEN is valid and has Inference Providers "
            "permission, or set the HF_CHAT_MODEL env var to a model your "
            "account can reach.",
            current_settings,
        )

    parsed = _extract_json_block(content)
    new_settings = _sanitize_settings(parsed, current_settings)
    reply_text = re.sub(r"```json.*?```", "", content, flags=re.DOTALL).strip()
    if not reply_text:
        reply_text = "Updated the extraction settings."
    return reply_text, new_settings