Update agent.py

agent.py

@@ -7,10 +7,14 @@ 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
@@ -41,6 +45,25 @@ _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
 # ──────────────────────────────────────────────────────────────────────────────
@@ -509,51 +532,105 @@ class BERTopicAgent:
         stream: Any,
         history: list,
         user_msg: str,
+        _attempt: int = 0,
     ) -> Generator[tuple, None, None]:
         """
         Consume a LangGraph stream, yielding Gradio 5-tuples incrementally.
-        Handles both streaming token chunks and full message events.
+
+        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 = ""
 
-        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
-
+        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 (
-                updated_history,
+                history + [[user_msg, accumulated]],
                 self.phase,
                 dict(self._charts),
                 list(self._downloads),
-                [],          # topic_rows populated in final yield
+                topic_rows,
             )
 
-        # Final yield with review-table rows (if artefacts exist)
-        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."""
@@ -624,37 +701,81 @@ class BERTopicAgent:
             f"Please continue to the next B&C phase now."
         )
 
-        stream = self._graph.stream(
-            {"messages": [{"role": "user", "content": review_message}]},
-            config=self._config(),
-            stream_mode="updates",
-        )
-
-        accumulated = ""
-        for chunk in 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),
+        def _make_review_stream() -> Any:
+            return self._graph.stream(
+                {"messages": [{"role": "user", "content": review_message}]},
+                config=self._config(),
+                stream_mode="updates",
             )
 
-        yield (
-            history + [["[Review submitted]", accumulated]],
-            self.phase,
-            dict(self._charts),
-            list(self._downloads),
-        )
+        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