workbench_codes.py

# ============================================================================
# workbench_codes.py — Sentence-level coding (Layer 1)
# ============================================================================
"""Layer 1 coder — assigns one short, POS-bound code to every sentence.

Purpose
-------
Implement the Braun & Clarke 2006 Phase 2 ("Generating initial codes") /
Grounded Theory open-coding step in code. Every sentence in the corpus
receives a 2-3 word noun-phrase code following the same POS-rule grammar
as the cluster-level theme labels (Layer 2). Same rule at both layers =
researcher-readable, repeatable, and inter-rater consistent.

The grammar
-----------
Every code MUST be:
  - 2 OR 3 words (never 1, never 4+)
  - [ADJ] + N or N + N pattern
  - Title Case (every word capitalised)
  - Singular noun head (unless the concept is inherently plural)
  - No verb as head, no pronoun, no article ("a", "an", "the")
  - No quotes, no punctuation, no preamble

Inputs
------
    sentences      : list[str]
    llm_provider   : str (workbench provider key)
    llm_key        : str (API key)
    batch_size     : int — sentences per LLM round-trip (default 20)

Outputs
-------
    list[str] codes, one per input sentence (parallel order).

Side effects
------------
    None. Pure function except for the LLM HTTP call.

Contract
--------
    code_sentences(sentences, llm_provider, llm_key, batch_size) -> list[str]
"""

from __future__ import annotations

import re
from typing import Iterable

import providers


# ----------------------------------------------------------------
# The codified rule — sent verbatim to the LLM
# ----------------------------------------------------------------
RULE_PROMPT_HEADER = (
    "You are coding sentences for a thematic-analysis codebook (Braun & Clarke "
    "2006 Phase 2; Saldaña 2016, Ch. 3).\n"
    "For each sentence below, produce TWO candidate codes (option 01 and "
    "option 02) so the researcher can pick the better one (closed card-sort "
    "methodology, Krippendorff 2018 §7).\n"
    "\n"
    "RULES every code must follow:\n"
    "  RULE 1 — Format: ADJECTIVE + NOUN or NOUN + NOUN, 2 OR 3 words total. "
    "Never 1, never 4+.\n"
    "  RULE 2 — Title Case. Every word capitalised.\n"
    "  RULE 3 — Singular noun head (unless inherently plural).\n"
    "  RULE 4 — No verbs as head, no pronouns, no articles, no quotes, no "
    "punctuation.\n"
    "  RULE 5 — Use ONLY Nouns and Adjectives.\n"
    "  RULE 6 — Option 01 and Option 02 MUST use DIFFERENT POS patterns (e.g. "
    "01 = N+N, 02 = A+N) so the researcher gets a real choice. If only one "
    "valid pattern fits the sentence, change the noun-head wording for option 02.\n"
    "\n"
    "POS letters: N = Noun, A = Adjective. Valid patterns: N+N, A+N, N+N+N, "
    "A+N+N, A+A+N, N+A+N. Forbidden: V (verb), D (determiner), P (pronoun), "
    "R (adverb), C (conjunction).\n"
    "\n"
    "OUTPUT FORMAT — one line per input sentence, in order:\n"
    "  <option_01_code> | <option_01_pattern> || <option_02_code> | <option_02_pattern>\n"
    "Separator between options is the double-pipe '||'. No numbering, no "
    "bullets, no preamble, no commentary.\n"
    "\n"
    "EXAMPLES of valid output lines:\n"
    "  Stolen Card Purchase | A+N+N || Fraudulent Purchase | A+N\n"
    "  Defective Product | A+N || Product Defect | N+N\n"
    "  Poor Hotel Condition | A+N+N || Smoke Smell | N+N\n"
    "  Bank Fraud | N+N || Fraudulent Bank | A+N\n"
    "  Refund Delay | N+N || Late Refund | A+N\n"
    "\n"
    "Sentences to code (one '01 || 02' line per sentence, in order):\n"
)


# Pronouns and articles we reject during post-processing.
_PRONOUNS: frozenset[str] = frozenset({
    "i", "me", "my", "mine", "myself",
    "we", "us", "our", "ours", "ourselves",
    "you", "your", "yours", "yourself", "yourselves",
    "he", "him", "his", "himself",
    "she", "her", "hers", "herself",
    "it", "its", "itself",
    "they", "them", "their", "theirs", "themselves",
    "this", "that", "these", "those",
})
_ARTICLES: frozenset[str] = frozenset({"a", "an", "the"})

# A simple "looks like a verb" filter for the head word. spaCy would be
# more accurate; this lightweight heuristic catches the common offenders
# without an extra dependency. The full spaCy validator is a follow-up.
_VERB_SUFFIXES: tuple[str, ...] = ("ing",)
_KNOWN_VERB_HEADS: frozenset[str] = frozenset({
    "stealing", "buying", "selling", "running", "working", "making",
    "breaking", "hiding", "creating", "issuing", "forging",
})


# ----------------------------------------------------------------
# Validator — applies the POS rule programmatically
# ----------------------------------------------------------------
_VALID_POS_LETTERS: frozenset[str] = frozenset({"N", "A"})
_VALID_POS_PATTERNS: frozenset[str] = frozenset({
    "N+N", "A+N", "N+N+N", "A+N+N", "A+A+N", "N+A+N",
})


def _split_code_and_pattern(raw: str) -> tuple[str, str]:
    """Split an LLM output line into ``(code, pattern)``.

    Accepts ``Bank Fraud | N+N`` or just ``Bank Fraud`` (no pattern). When
    the pattern is missing we return an empty string for it; the validator
    will then reject the line.
    """
    s = (raw or "").strip()
    if "|" in s:
        left, right = s.split("|", 1)
        return left.strip(), right.strip().upper()
    return s, ""


def _validate_code(raw: str) -> tuple[bool, str]:
    """Validate just the code text (used for researcher overrides).

    Args:
        raw: A code string (may include preambles or punctuation).

    Returns:
        ``(True, cleaned_code)`` or ``(False, reason)``.
    """
    s = (raw or "").strip()
    # Strip common preambles the LLM or researcher occasionally adds.
    for prefix in ("Code:", "Theme:", "Label:", "-", "*"):
        if s.lower().startswith(prefix.lower()):
            s = s[len(prefix):].strip()
    s = s.strip("\"'`.,;:!?()[]{}")
    if not s:
        return False, "empty"

    words = s.split()
    if len(words) not in (2, 3):
        return False, f"wrong word count ({len(words)})"
    if any(w.lower() in _PRONOUNS for w in words):
        return False, "contains pronoun"
    if any(w.lower() in _ARTICLES for w in words):
        return False, "contains article"
    head = words[-1].lower()
    if head.endswith(_VERB_SUFFIXES) or head in _KNOWN_VERB_HEADS:
        return False, "verb head (-ing or known verb)"
    cleaned = " ".join(w[:1].upper() + w[1:].lower() if w else w for w in words)
    return True, cleaned


def _validate_pattern(pattern: str, n_words: int) -> tuple[bool, str]:
    """Validate the LLM's POS pattern (e.g. "A+N", "N+N+N").

    Rules:
      - non-empty
      - made up of '+' separated 1-letter tokens, each in {N, A}
      - token count must equal the code's word count
      - the joined pattern must be in the closed list of allowed patterns

    Returns:
        ``(True, normalised_pattern)`` or ``(False, reason)``.
    """
    if not pattern:
        return False, "missing pattern"
    tokens = [t.strip().upper() for t in pattern.split("+") if t.strip()]
    if not tokens:
        return False, "empty pattern"
    if len(tokens) != n_words:
        return False, f"pattern length {len(tokens)} ≠ code word count {n_words}"
    if any(t not in _VALID_POS_LETTERS for t in tokens):
        return False, f"contains forbidden POS letter (only N and A allowed): {pattern}"
    joined = "+".join(tokens)
    if joined not in _VALID_POS_PATTERNS:
        return False, f"pattern {joined!r} not in allowed set"
    return True, joined


def _validate_line(raw: str) -> tuple[bool, str, str]:
    """Validate a full ``<code> | <pattern>`` line from the LLM.

    Args:
        raw: One line of LLM output.

    Returns:
        ``(ok, code_or_reason, pattern_or_reason)``.
    """
    code_raw, pattern_raw = _split_code_and_pattern(raw)
    ok_code, code_value = _validate_code(code_raw)
    if not ok_code:
        return False, code_value, ""
    n_words = len(code_value.split())
    ok_pat, pat_value = _validate_pattern(pattern_raw, n_words)
    if not ok_pat:
        return False, code_value, pat_value
    return True, code_value, pat_value


def _split_two_candidates(raw: str) -> tuple[str, str]:
    """Split a single LLM output line into the two candidates.

    The expected separator is a double-pipe '||'. Returns
    ``(option_01_raw, option_02_raw)``. If the LLM didn't produce a
    second candidate (or used a different separator), the second
    return value is empty so the caller's validator marks option 02
    as missing and we fall back to option 01 only.
    """
    s = (raw or "").strip()
    if "||" in s:
        a, b = s.split("||", 1)
        return a.strip(), b.strip()
    return s, ""


def _validate_two_candidates(raw: str) -> dict:
    """Validate one LLM line that contains two candidate codes.

    Returns a dict with the parsed and validated fields:
        option_01_code, option_01_pattern,
        option_02_code, option_02_pattern,
        ok_01, ok_02   (booleans)
    Failed candidates collapse to ``("Unlabeled Code", "")``.
    """
    raw_01, raw_02 = _split_two_candidates(raw)

    ok_01, code_01, pat_01 = _validate_line(raw_01)
    ok_02, code_02, pat_02 = _validate_line(raw_02) if raw_02 else (False, "", "")

    if not ok_01:
        code_01, pat_01 = "Unlabeled Code", ""
    if not ok_02:
        code_02, pat_02 = "Unlabeled Code", ""

    return {
        "option_01_code": code_01,
        "option_01_pattern": pat_01,
        "option_02_code": code_02,
        "option_02_pattern": pat_02,
        "ok_01": ok_01,
        "ok_02": ok_02,
    }


# ----------------------------------------------------------------
# Batch coder — the heavy LLM loop
# ----------------------------------------------------------------
def _chunked(seq: list[str], n: int) -> Iterable[list[str]]:
    """Yield successive chunks of ``seq`` of size ``n``."""
    for i in range(0, len(seq), n):
        yield seq[i:i + n]


def _call_llm_batch(
    batch: list[str], llm_provider: str, llm_key: str,
) -> list[str]:
    """Send one batch of sentences and parse the LLM's per-line response.

    Args:
        batch:        The sentences in this batch (1..N strings).
        llm_provider: Workbench provider key.
        llm_key:      API key for the provider.

    Returns:
        A list of raw code strings, one per input sentence. If the LLM
        returns fewer or more lines than expected, the list is padded
        with empty strings (which the validator will then reject).

    Raises:
        RuntimeError: when the provider client raises.
    """
    client = providers.get_llm_client(llm_provider or "Mistral", llm_key or "")
    model = providers.get_llm_model(llm_provider or "Mistral")
    prompt = RULE_PROMPT_HEADER + "\n".join(
        f"{i+1}. {s}" for i, s in enumerate(batch)
    )
    try:
        resp = client.chat.complete(
            model=model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.0,
            # Generous: each code ~3 words ~6 tokens × 20 sentences = 120 tokens.
            max_tokens=64 * len(batch),
        )
    except Exception as exc:
        raise RuntimeError(f"Layer 1 LLM call failed: {exc}") from exc

    text = (resp.choices[0].message.content or "").strip()
    # Parse "1. Foo Bar" / "Foo Bar" / "- Foo Bar" — strip leading numbering.
    lines = [
        re.sub(r"^\s*(?:\d+[.)]|-|\*)\s*", "", ln).strip()
        for ln in text.splitlines() if ln.strip()
    ]
    # Pad / truncate to len(batch) so caller can rely on parallel indexing.
    if len(lines) < len(batch):
        lines = lines + [""] * (len(batch) - len(lines))
    return lines[:len(batch)]


# ----------------------------------------------------------------
# PUBLIC: code_sentences
# ----------------------------------------------------------------
def code_sentences(
    sentences: list[str],
    llm_provider: str = "Mistral",
    llm_key: str = "",
    batch_size: int = 20,
) -> list[dict]:
    """Assign TWO candidate POS-bound codes to every sentence.

    Implements forced-choice closed card-sort coding (Krippendorff 2018
    §7) — the LLM produces two candidates with different POS patterns,
    the researcher picks one (default option 01) or overrides with their
    own (agreement = N).

    Args:
        sentences:    The corpus.
        llm_provider: Workbench provider key. Default Mistral.
        llm_key:      API key for the provider.
        batch_size:   Sentences per LLM round-trip. Default 20.

    Returns:
        A list of dicts parallel to ``sentences``, each containing:
            option_01_code, option_01_pattern,
            option_02_code, option_02_pattern,
            ok_01, ok_02

        When a candidate fails POS validation it collapses to
        ``"Unlabeled Code"`` with empty pattern, and ``ok_NN`` is False
        so the UI can flag it.

    Raises:
        RuntimeError: bubbled up from the LLM call.

    Example:
        >>> code_sentences(["He used a stolen credit card."])[0]
        {'option_01_code': 'Stolen Card Purchase', 'option_01_pattern': 'A+N+N',
         'option_02_code': 'Fraudulent Purchase',  'option_02_pattern': 'A+N',
         'ok_01': True, 'ok_02': True}
    """
    if not sentences:
        return []

    out: list[dict] = []
    for batch in _chunked(list(sentences), max(1, int(batch_size))):
        raw_lines = _call_llm_batch(batch, llm_provider, llm_key)
        for raw in raw_lines:
            out.append(_validate_two_candidates(raw))
    return out


def resolve_final_code(
    option_01_code: str,
    option_02_code: str,
    researcher_choice: str,
    researcher_code_override: str = "",
) -> tuple[str, str]:
    """Resolve the final_code given the researcher's choice.

    Args:
        option_01_code:           First LLM candidate.
        option_02_code:           Second LLM candidate.
        researcher_choice:        "01", "02", or "override".
        researcher_code_override: Researcher's typed override (used only
                                  when researcher_choice == "override").

    Returns:
        ``(final_code, researcher_agreement)`` where agreement is "Y"
        when choice ∈ {01, 02} and "N" when choice == "override".

    Override codes are POS-validated; a failed override becomes
    "Unlabeled Code".
    """
    choice = (researcher_choice or "01").strip().lower()
    if choice == "02":
        return option_02_code, "Y"
    if choice in ("override", "n"):
        if not researcher_code_override:
            return option_01_code, "Y"
        ok, value = _validate_code(researcher_code_override)
        return (value if ok else "Unlabeled Code"), "N"
    # default
    return option_01_code, "Y"