agent.py

"""
agent.py — BERTopic Thematic Analysis Agent
Braun & Clarke (2006) six-phase methodology implemented as a ReAct agent
using LangGraph, ChatMistralAI, and MemorySaver.
"""

from __future__ import annotations

import json
import logging
import re
import time
from pathlib import Path
from typing import Any, Generator

logger = logging.getLogger(__name__)

from langchain_mistralai import ChatMistralAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent

from tools import (
    load_scopus_csv,
    run_bertopic_discovery,
    label_topics_with_llm,
    consolidate_into_themes,
    compare_with_taxonomy,
    generate_comparison_csv,
    export_narrative,
)

# ──────────────────────────────────────────────────────────────────────────────
# Artifact paths (shared across phases)
# ──────────────────────────────────────────────────────────────────────────────

ARTIFACTS_DIR   = Path("artifacts")
ARTIFACTS_DIR.mkdir(exist_ok=True)

_LOADED_DATA    = str(ARTIFACTS_DIR / "loaded_data.json")
_SUMMARIES      = str(ARTIFACTS_DIR / "summaries.json")
_EMB            = str(ARTIFACTS_DIR / "emb.npy")
_LABELS         = str(ARTIFACTS_DIR / "topic_labels.json")
_THEMES         = str(ARTIFACTS_DIR / "themes.json")
_TAXONOMY       = str(ARTIFACTS_DIR / "taxonomy_mapping.json")
_COMPARISON_CSV = str(ARTIFACTS_DIR / "abstract_vs_title_comparison.csv")
_NARRATIVE      = str(ARTIFACTS_DIR / "section7_narrative.txt")

# ──────────────────────────────────────────────────────────────────────────────
# Rate-limit retry config
# ──────────────────────────────────────────────────────────────────────────────

_RL_MAX_RETRIES  = 4                        # max automatic retries on 429
_RL_BACKOFF_SECS = [15, 30, 60, 120]        # wait before each retry attempt


def _is_rate_limit(exc: Exception) -> bool:
    """Return True when *exc* is an HTTP 429 rate-limit error from any client."""
    # httpx.HTTPStatusError carries a .response attribute
    resp = getattr(exc, "response", None)
    if resp is not None and getattr(resp, "status_code", None) == 429:
        return True
    # Fallback: inspect the string representation (handles wrapped exceptions)
    s = str(exc).lower()
    return "429" in s and ("rate limit" in s or "rate_limit" in s or "rate_limited" in s)


# ──────────────────────────────────────────────────────────────────────────────
# System Prompt
# ──────────────────────────────────────────────────────────────────────────────

SYSTEM_PROMPT = """
╔══════════════════════════════════════════════════════════════════════════════╗
║          COMPUTATIONAL THEMATIC ANALYSIS AGENT  —  SYSTEM PROMPT           ║
╚══════════════════════════════════════════════════════════════════════════════╝

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
ROLE
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
You are a computational thematic analysis expert trained in the Braun & Clarke
(2006) six-phase framework for rigorous qualitative and mixed-methods research.
You specialise in applying BERTopic-based semantic clustering to academic
literature corpora, with deep expertise in:

  • Systematic literature review methodology
  • Sentence-level semantic embedding and agglomerative clustering
  • LLM-assisted topic labelling and theme consolidation
  • PAJAIS (Pacific-Asia Journal of the Association for Information Systems)
    25-category research taxonomy alignment
  • Transparent, reproducible, human-in-the-loop analytical pipelines

Your outputs are used in peer-reviewed academic research. Precision,
methodological rigour, and faithful adherence to the B&C (2006) phases are
non-negotiable.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
CRITICAL RULES  (must be followed without exception)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
1.  ONE PHASE PER MESSAGE. Complete exactly one B&C phase per conversational
    turn. Never skip ahead or combine phases in a single response.

2.  ALL APPROVALS VIA REVIEW TABLE — NEVER VIA CHAT. You must NEVER ask the
    user to approve, reject, or rename topics in free-text chat. Every approval
    workflow must go through the Gradio review table. After populating the table
    you must STOP and wait for the user to click "Submit Review".

3.  STOP GATES ARE MANDATORY. At the end of Phases 2, 3, 4, and 5.5 you must
    output the exact STOP phrase:
        ⏸ STOP GATE — awaiting your review table submission to continue.
    Do not proceed until the user's next message contains review data.

4.  NEVER HALLUCINATE TOOL RESULTS. If a tool call fails, report the exact
    error verbatim and ask the user how to proceed. Do not invent file paths,
    cluster counts, or topic labels.

5.  COLUMN DISCIPLINE. Never include "Author Keywords" in any clustering run.
    Use only the columns specified in RUN_CONFIGS: Abstract (abstract run) or
    Title (title run).

6.  ARTEFACT HYGIENE. Every tool saves files to the artifacts/ directory.
    Always pass the exact saved_path returned by a previous tool to the next
    tool. Never guess or construct file paths manually.

7.  STREAMING DISCIPLINE. Yield one streamed chunk per reasoning step so the
    Gradio UI can update the phase progress bar in real time.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
TOOLS  (7 available)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
1.  load_scopus_csv(csv_path, run_mode)
      Load a Scopus-exported CSV. Counts papers and sentences. Applies
      boilerplate regex filter. Saves loaded_data.json. Use in Phase 1.

2.  run_bertopic_discovery(loaded_data_path)
      Embeds sentences with all-MiniLM-L6-v2 (normalize_embeddings=True).
      Clusters with AgglomerativeClustering(metric=cosine, threshold=0.7).
      No UMAP. Finds 5 nearest centroid sentences per cluster. Generates
      4 Plotly charts. Saves summaries.json + emb.npy. Use in Phase 2.

3.  label_topics_with_llm(summaries_path, top_n)
      Sends top-N topics (max 100) to Mistral via PromptTemplate +
      JsonOutputParser. Returns short labels and descriptions. Saves
      topic_labels.json. Use in Phase 2 after discovery.

4.  consolidate_into_themes(labels_path, summaries_path, emb_path, approved_groups)
      Merges approved topic groups into named themes. Recomputes centroids.
      Saves themes.json. Use in Phase 3 after the review table is submitted.

5.  compare_with_taxonomy(themes_path)
      Maps consolidated themes to PAJAIS 25 categories via Mistral.
      Returns confidence scores and rationale. Saves taxonomy_mapping.json.
      Use in Phase 5.5.

6.  generate_comparison_csv(csv_path, taxonomy_path)
      Produces abstract vs title side-by-side CSV with PAJAIS categories
      and confidence scores. Saves abstract_vs_title_comparison.csv.
      Use in Phase 6.

7.  export_narrative(taxonomy_path)
      Generates a ~500-word Section 7 (Discussion & Implications) as
      flowing academic prose via Mistral. Saves section7_narrative.txt.
      Use in Phase 6.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
BRAUN & CLARKE (2006) SIX-PHASE PROTOCOL
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

┌─────────────────────────────────────────────────────────────────────────────┐
│  PHASE 1 — FAMILIARISATION WITH THE DATA                                   │
└─────────────────────────────────────────────────────────────────────────────┘
Objective: Immerse in the corpus. Understand its scope, structure, and quality.

Instructions:
  a. Call load_scopus_csv(csv_path=<user_provided>, run_mode=<"abstract"|"title">).
  b. Display the returned statistics in a clear summary:
       • Total papers loaded
       • Total sentences extracted
       • Sentences remaining after boilerplate filtering
       • Column(s) used
       • Run mode (abstract / title)
  c. Comment briefly on data quality: density, likely noise level, any
     column mapping issues detected.
  d. STOP. Do not proceed to Phase 2 until the user explicitly confirms
     they are satisfied with the loaded data.

Output template:
  📂 Phase 1 Complete — Familiarisation
  ─────────────────────────────────────
  Papers:              {N}
  Sentences extracted: {S}
  After filtering:     {F}
  Column used:         {C}
  Run mode:            {M}
  [Quality commentary]
  ✅ Ready for Phase 2. Reply "proceed" to start Initial Coding.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

┌─────────────────────────────────────────────────────────────────────────────┐
│  PHASE 2 — GENERATING INITIAL CODES                                        │
└─────────────────────────────────────────────────────────────────────────────┘
Objective: Produce a full set of atomic semantic codes from the corpus.

Instructions:
  a. Call run_bertopic_discovery(loaded_data_path=artifacts/loaded_data.json).
     Report: number of clusters found, total sentences clustered, chart paths.
  b. Call label_topics_with_llm(summaries_path=artifacts/summaries.json, top_n=100).
     Report: number of topics labelled.
  c. Populate the Gradio review table with ALL labelled topics. Each row must
     contain:
       • #           — topic_id (integer)
       • Topic Label — LLM-generated label
       • Top Evidence — first centroid sentence (truncated to 120 chars)
       • Sentences   — cluster size
       • Papers      — estimated paper count (size ÷ avg sentences per paper)
       • Approve     — default True
       • Rename To   — empty (user fills)
       • Reasoning   — empty (user fills)
  d. Present the 4 Plotly charts by referencing their file paths.
  e. Explain to the user:
       • Check "Approve" for topics to keep; uncheck to discard.
       • Fill "Rename To" with a preferred label; leave blank to keep LLM label.
       • Optionally note merging intentions in "Reasoning".
       • Topics with the same "Reasoning" group tag will be merged in Phase 3.

⏸ STOP GATE — awaiting your review table submission to continue.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

┌─────────────────────────────────────────────────────────────────────────────┐
│  PHASE 3 — SEARCHING FOR THEMES                                            │
└─────────────────────────────────────────────────────────────────────────────┘
Objective: Collate approved codes into candidate themes.

Instructions:
  a. Parse the submitted review table. Extract:
       • Approved topic IDs (Approve == True)
       • Rename mappings (Rename To != "")
       • Merge groups (topics sharing the same Reasoning tag)
  b. Construct approved_groups: a JSON list of lists, where each inner list
     contains the topic_ids belonging to one theme. Topics with a shared
     Reasoning tag form one group. Approved topics with no Reasoning tag
     each form a singleton group.
  c. Call consolidate_into_themes(
         labels_path=artifacts/topic_labels.json,
         summaries_path=artifacts/summaries.json,
         emb_path=artifacts/emb.npy,
         approved_groups=<constructed JSON string>
     ).
  d. Display a theme summary table:
       Theme # | Theme Label | Topics Merged | Total Sentences

⏸ STOP GATE — awaiting your review table submission to continue.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

┌─────────────────────────────────────────────────────────────────────────────┐
│  PHASE 4 — REVIEWING THEMES / SATURATION CHECK                             │
└─────────────────────────────────────────────────────────────────────────────┘
Objective: Assess whether themes are internally coherent and collectively
exhaustive. Check corpus coverage.

Instructions:
  a. Load artifacts/themes.json (already created by Phase 3).
  b. Compute and display a saturation report:
       • Total sentences covered by all themes vs. total corpus sentences
       • Coverage percentage
       • Theme coherence flag: warn if any theme covers < 1 % of corpus
       • Overlap flag: warn if any two themes share > 30 % vocabulary
       (Vocabulary overlap is approximated by comparing top-evidence sentences
        using word-level Jaccard similarity — compute in Python, no tool call.)
  c. Populate the review table again with the THEME list (not topic list):
       • #           — theme_id
       • Topic Label — current theme_label
       • Top Evidence — first top_evidence sentence (120 chars)
       • Sentences   — total_size
       • Papers      — estimated
       • Approve     — default True
       • Rename To   — user may provide final name
       • Reasoning   — any split/merge instructions
  d. Ask the user to confirm themes, request splits/merges, or rename.

⏸ STOP GATE — awaiting your review table submission to continue.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

┌─────────────────────────────────────────────────────────────────────────────┐
│  PHASE 5 — DEFINING AND NAMING THEMES                                      │
└─────────────────────────────────────────────────────────────────────────────┘
Objective: Produce final, publication-ready theme names and definitions.

Instructions:
  a. Apply all renames from the Phase 4 review table to artifacts/themes.json
     in memory (update theme_label field for each theme_id where Rename To
     is non-empty).
  b. For each finalised theme, generate a two-sentence academic definition
     grounded in the top_evidence sentences. Output this as a numbered list.
  c. Confirm the final theme set to the user in a clean summary:
       Theme # | Final Name | Definition (2 sentences) | Sentence Count
  d. STOP. Ask the user to confirm the final names before PAJAIS mapping.
     ✅ Reply "proceed to taxonomy" to continue to Phase 5.5.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

┌─────────────────────────────────────────────────────────────────────────────┐
│  PHASE 5.5 — PAJAIS TAXONOMY ALIGNMENT                                     │
└─────────────────────────────────────────────────────────────────────────────┘
Objective: Map each finalised theme to the PAJAIS 25-category taxonomy.

Instructions:
  a. Call compare_with_taxonomy(themes_path=artifacts/themes.json).
  b. Display the mapping results in a structured table:
       Theme Name | PAJAIS Category | Confidence | Rationale
  c. Highlight any themes with confidence < 0.5 as requiring manual review.
  d. Note any PAJAIS categories not covered by the corpus (research gaps).
  e. Populate the review table with the mapping results:
       • #           — theme_id
       • Topic Label — theme_label → PAJAIS category
       • Top Evidence — rationale (truncated)
       • Approve     — default True (uncheck to override mapping)
       • Rename To   — alternative PAJAIS category if user disagrees
       • Reasoning   — free notes

⏸ STOP GATE — awaiting your review table submission to continue.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

┌─────────────────────────────────────────────────────────────────────────────┐
│  PHASE 6 — PRODUCING THE REPORT                                            │
└─────────────────────────────────────────────────────────────────────────────┘
Objective: Generate all final deliverables.

Instructions:
  a. Apply any PAJAIS category overrides from the Phase 5.5 review table
     to artifacts/taxonomy_mapping.json in memory.
  b. Call generate_comparison_csv(
         csv_path=<original CSV path>,
         taxonomy_path=artifacts/taxonomy_mapping.json
     ).
     Report: row count, file path.
  c. Call export_narrative(taxonomy_path=artifacts/taxonomy_mapping.json).
     Report: word count, file path, first 150 chars of preview.
  d. Present a final deliverables checklist:
       ✅ artifacts/summaries.json            — raw cluster summaries
       ✅ artifacts/topic_labels.json         — LLM-generated labels
       ✅ artifacts/themes.json               — consolidated themes
       ✅ artifacts/taxonomy_mapping.json     — PAJAIS alignment
       ✅ artifacts/abstract_vs_title_comparison.csv
       ✅ artifacts/section7_narrative.txt    — ~500-word Section 7
       ✅ artifacts/chart_cluster_sizes.html
       ✅ artifacts/chart_pca_scatter.html
       ✅ artifacts/chart_top10_pie.html
       ✅ artifacts/chart_centroid_heatmap.html
  e. Congratulate the user and offer to re-run in title mode for comparison.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
END OF SYSTEM PROMPT
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
""".strip()

# ──────────────────────────────────────────────────────────────────────────────
# Tool registry
# ──────────────────────────────────────────────────────────────────────────────

TOOLS = [
    load_scopus_csv,
    run_bertopic_discovery,
    label_topics_with_llm,
    consolidate_into_themes,
    compare_with_taxonomy,
    generate_comparison_csv,
    export_narrative,
]

# ──────────────────────────────────────────────────────────────────────────────
# Phase detection helpers
# ──────────────────────────────────────────────────────────────────────────────

_PHASE_PATTERNS = {
    "loading":    re.compile(r"load_scopus_csv|loaded_data", re.I),
    "embedding":  re.compile(r"run_bertopic_discovery|embedding", re.I),
    "clustering": re.compile(r"summaries\.json|n_topics|clusters found", re.I),
    "labelling":  re.compile(r"label_topics_with_llm|topic_labels", re.I),
    "review":     re.compile(r"STOP GATE|review table|submit review", re.I),
    "done":       re.compile(r"section7_narrative|deliverables checklist", re.I),
}


def _detect_phase(text: str) -> str:
    """Return the most specific pipeline phase detectable from agent output."""
    matched = list(filter(
        lambda kv: kv[1].search(text),
        _PHASE_PATTERNS.items(),
    ))
    return matched[-1][0] if matched else "idle"


# ──────────────────────────────────────────────────────────────────────────────
# Review-table row builder
# ──────────────────────────────────────────────────────────────────────────────

_TABLE_COLS = ["#", "Topic Label", "Top Evidence", "Sentences", "Papers", "Approve", "Rename To", "Reasoning"]


def _topic_to_row(topic: dict, papers_per_sent: float = 0.2) -> list:
    """Convert a topic/theme dict to a review-table row."""
    evidence = (topic.get("top_evidence") or [""])[0]
    return [
        topic.get("topic_id", topic.get("theme_id", 0)),
        topic.get("label", topic.get("theme_label", "")),
        evidence[:120],
        topic.get("size", topic.get("total_size", 0)),
        round(topic.get("size", topic.get("total_size", 0)) * papers_per_sent),
        True,
        "",
        "",
    ]


def _build_review_rows(path: str, id_key: str = "topic_id") -> list[list]:
    """Load a JSON artefact and convert every entry to a review-table row."""
    records = json.loads(Path(path).read_text())
    return list(map(_topic_to_row, records))


# ──────────────────────────────────────────────────────────────────────────────
# Approved-groups extractor  (called in handle_review for Phase 2 → 3)
# ──────────────────────────────────────────────────────────────────────────────

def _extract_approved_groups(rows: list[list]) -> str:
    """
    Parse review-table rows into approved_groups JSON string.

    Groups are formed by the Reasoning field value:
      • Rows sharing a non-empty Reasoning tag → merged into one group
      • Approved rows with empty Reasoning    → singleton group each
      • Unapproved rows (Approve == False)    → discarded
    """
    approved = list(filter(lambda r: r[5] is True or r[5] == "True" or r[5] == 1, rows))

    tagged    = list(filter(lambda r: str(r[7]).strip(), approved))
    untagged  = list(filter(lambda r: not str(r[7]).strip(), approved))

    # Group tagged rows by their Reasoning value
    reasoning_vals = list(set(map(lambda r: str(r[7]).strip(), tagged)))

    tagged_groups = list(map(
        lambda tag: list(map(
            lambda r: int(r[0]),
            filter(lambda r: str(r[7]).strip() == tag, tagged),
        )),
        reasoning_vals,
    ))

    singleton_groups = list(map(lambda r: [int(r[0])], untagged))

    all_groups = tagged_groups + singleton_groups
    return json.dumps(all_groups)


# ──────────────────────────────────────────────────────────────────────────────
# BERTopicAgent  — the class consumed by app.py
# ──────────────────────────────────────────────────────────────────────────────

class BERTopicAgent:
    """
    Wraps a LangGraph ReAct agent and exposes the two generator methods
    expected by the Gradio front-end:

        handle_message(message, history, csv_path) → yields 5-tuple
        handle_review(table_rows, history)          → yields 4-tuple
    """

    # Gradio 5-tuple: (history, phase, charts_dict, downloads_list, topic_rows)
    # Gradio 4-tuple: (history, phase, charts_dict, downloads_list)

    def __init__(self) -> None:
        self.phase       = "idle"
        self._charts:    dict[str, str] = {}
        self._downloads: list[str]      = []
        self._csv_path:  str | None     = None

        self._llm       = ChatMistralAI(
            model="mistral-large-latest",
            temperature=0.2,
            streaming=True,
        )
        self._memory    = MemorySaver()

        # handle_tool_error is no longer a @tool() decorator argument in
        # newer LangChain versions — set it directly on each tool object.
        for t in TOOLS:
            t.handle_tool_error = True

        self._graph     = create_react_agent(
            model=self._llm,
            tools=TOOLS,
            checkpointer=self._memory,
            prompt=SYSTEM_PROMPT,
        )
        self._thread_id = "bc2006-session-1"

    # ── internal ──────────────────────────────────────────────────────────────

    def _config(self) -> dict:
        return {"configurable": {"thread_id": self._thread_id}}

    def _update_charts(self, text: str) -> None:
        """Scan agent text for chart file paths and register them."""
        found = re.findall(r'artifacts/chart_[a-z_]+\.html', text)
        label_map = {
            "chart_cluster_sizes":    "Cluster Sizes",
            "chart_pca_scatter":      "PCA Scatter",
            "chart_top10_pie":        "Top-10 Pie",
            "chart_centroid_heatmap": "Centroid Heatmap",
        }
        list(map(
            lambda p: self._charts.__setitem__(
                label_map.get(Path(p).stem, Path(p).stem), p
            ),
            found,
        ))

    def _update_downloads(self, text: str) -> None:
        """Scan agent text for downloadable artefact paths."""
        # Added ?: to make it a non-capturing group
        found = re.findall(r'artifacts/[\w_]+\.(?:json|csv|txt|npy|html)', text)
        new   = list(filter(lambda p: p not in self._downloads, found))
        self._downloads.extend(new)

    def _accumulate_stream(
        self,
        stream: Any,
        history: list,
        user_msg: str,
        _attempt: int = 0,
    ) -> Generator[tuple, None, None]:
        """
        Consume a LangGraph stream, yielding Gradio 5-tuples incrementally.

        Automatically retries on HTTP 429 (rate-limit) errors with exponential
        back-off up to _RL_MAX_RETRIES times.  All other exceptions surface a
        friendly error message in the chat rather than crashing the generator.
        """
        accumulated = ""

        try:
            for chunk in stream:
                # LangGraph yields dicts keyed by node name
                node_output = (
                    chunk.get("agent") or
                    chunk.get("tools") or
                    {}
                )
                messages = node_output.get("messages", [])

                text_delta = "".join(list(map(
                    lambda m: getattr(m, "content", "") if hasattr(m, "content") else "",
                    messages,
                )))

                accumulated += text_delta
                self._update_charts(accumulated)
                self._update_downloads(accumulated)
                self.phase = _detect_phase(accumulated)

                updated_history = history + [[user_msg, accumulated]] if accumulated else history

                yield (
                    updated_history,
                    self.phase,
                    dict(self._charts),
                    list(self._downloads),
                    [],          # topic_rows populated in final yield
                )

            # ── Success: final yield with review-table rows ───────────────────
            topic_rows = self._latest_review_rows()
            yield (
                history + [[user_msg, accumulated]],
                self.phase,
                dict(self._charts),
                list(self._downloads),
                topic_rows,
            )

        except Exception as exc:  # noqa: BLE001
            if _is_rate_limit(exc) and _attempt < _RL_MAX_RETRIES:
                # ── Rate-limit: back off then restart the stream ──────────────
                wait = _RL_BACKOFF_SECS[_attempt]
                notice = (
                    f"\n\n⏳ **Mistral rate limit hit** — waiting **{wait}s** "
                    f"then retrying automatically "
                    f"(attempt {_attempt + 1}/{_RL_MAX_RETRIES})…"
                )
                logger.warning("Rate limit 429 on attempt %d; sleeping %ds", _attempt, wait)
                yield (
                    history + [[user_msg, accumulated + notice]],
                    "idle",
                    dict(self._charts),
                    list(self._downloads),
                    [],
                )
                time.sleep(wait)

                # Rebuild the stream — MemorySaver resumes from last checkpoint
                new_stream = self._graph.stream(
                    {"messages": [{"role": "user", "content": user_msg}]},
                    config=self._config(),
                    stream_mode="updates",
                )
                yield from self._accumulate_stream(
                    new_stream, history, user_msg, _attempt=_attempt + 1
                )

            else:
                # ── Non-retryable error: surface gracefully in chat ───────────
                if _is_rate_limit(exc):
                    err_header = (
                        f"❌ **Rate limit persists after {_RL_MAX_RETRIES} retries.**\n"
                        "Please wait a few minutes before sending another message."
                    )
                else:
                    err_header = f"❌ **API / tool error:** `{type(exc).__name__}: {exc}`"

                logger.exception("Unhandled error in _accumulate_stream (attempt %d)", _attempt)
                self.phase = "idle"
                yield (
                    history + [[user_msg, accumulated + f"\n\n{err_header}"]],
                    "idle",
                    dict(self._charts),
                    list(self._downloads),
                    self._latest_review_rows(),   # keep existing table intact
                )

    def _latest_review_rows(self) -> list[list]:
        """Return review rows from the most recently produced artefact."""
        candidates = [
            (_THEMES,   "theme_id"),
            (_LABELS,   "topic_id"),
            (_SUMMARIES,"topic_id"),
        ]
        existing = list(filter(lambda t: Path(t[0]).exists(), candidates))
        return _build_review_rows(*existing[0]) if existing else []

    # ── public API ────────────────────────────────────────────────────────────

    def handle_message(
        self,
        message:  str,
        history:  list,
        csv_path: str | None = None,
    ) -> Generator[tuple, None, None]:
        """
        Send a user message to the ReAct agent and stream back Gradio 5-tuples.

        Yields: (history, phase, charts_dict, downloads_list, topic_rows)
        """
        self._csv_path = csv_path or self._csv_path

        # Inject CSV path into message so the agent can reference it
        enriched = (
            f"{message}\n\n[SYSTEM CONTEXT] CSV path: {self._csv_path}"
            if self._csv_path and "csv" not in message.lower()
            else message
        )

        stream = self._graph.stream(
            {"messages": [{"role": "user", "content": enriched}]},
            config=self._config(),
            stream_mode="updates",
        )

        yield from self._accumulate_stream(stream, history, message)

    def handle_review(
        self,
        table_data: list,
        history:    list,
    ) -> Generator[tuple, None, None]:
        """
        Process a submitted review table and advance to the next B&C phase.

        The table rows are serialised to JSON and injected as a structured
        user message so the agent can parse approvals, renames, and groups.

        Yields: (history, phase, charts_dict, downloads_list)
        """
        approved_groups = _extract_approved_groups(table_data)

        review_payload = json.dumps({
            "event":           "review_submitted",
            "rows":            table_data,
            "approved_groups": json.loads(approved_groups),
            "approved_count":  len(json.loads(approved_groups)),
        }, ensure_ascii=False, indent=2)

        review_message = (
            f"The user has submitted the review table. "
            f"Approved groups: {approved_groups}. "
            f"Full payload:\n{review_payload}\n\n"
            f"Please continue to the next B&C phase now."
        )

        def _make_review_stream() -> Any:
            return self._graph.stream(
                {"messages": [{"role": "user", "content": review_message}]},
                config=self._config(),
                stream_mode="updates",
            )

        accumulated = ""
        attempt     = 0

        while True:
            current_stream = _make_review_stream()
            try:
                for chunk in current_stream:
                    node_output = chunk.get("agent") or chunk.get("tools") or {}
                    messages    = node_output.get("messages", [])

                    text_delta = "".join(list(map(
                        lambda m: getattr(m, "content", "") if hasattr(m, "content") else "",
                        messages,
                    )))

                    accumulated += text_delta
                    self._update_charts(accumulated)
                    self._update_downloads(accumulated)
                    self.phase = _detect_phase(accumulated)

                    yield (
                        history + [["[Review submitted]", accumulated]],
                        self.phase,
                        dict(self._charts),
                        list(self._downloads),
                    )

                # ── Success ───────────────────────────────────────────────────
                yield (
                    history + [["[Review submitted]", accumulated]],
                    self.phase,
                    dict(self._charts),
                    list(self._downloads),
                )
                break   # exit retry loop

            except Exception as exc:  # noqa: BLE001
                if _is_rate_limit(exc) and attempt < _RL_MAX_RETRIES:
                    wait = _RL_BACKOFF_SECS[attempt]
                    notice = (
                        f"\n\n⏳ **Rate limit hit** — waiting **{wait}s** "
                        f"then retrying (attempt {attempt + 1}/{_RL_MAX_RETRIES})…"
                    )
                    logger.warning("Rate limit 429 in handle_review attempt %d; sleeping %ds", attempt, wait)
                    yield (
                        history + [["[Review submitted]", accumulated + notice]],
                        "idle",
                        dict(self._charts),
                        list(self._downloads),
                    )
                    time.sleep(wait)
                    attempt += 1
                    # Loop re-creates the stream via _make_review_stream()
                else:
                    if _is_rate_limit(exc):
                        err = (
                            f"❌ **Rate limit persists after {_RL_MAX_RETRIES} retries.**\n"
                            "Please wait a few minutes before trying again."
                        )
                    else:
                        err = f"❌ **Error processing review:** `{type(exc).__name__}: {exc}`"

                    logger.exception("Unhandled error in handle_review (attempt %d)", attempt)
                    self.phase = "idle"
                    yield (
                        history + [["[Review submitted]", accumulated + f"\n\n{err}"]],
                        "idle",
                        dict(self._charts),
                        list(self._downloads),
                    )
                    break