Fix follow-ups: explicit session-cached conversation + balanced decline rule

- Replace ChatInterface with a gr.Blocks chat that holds the conversation in the
Chatbot component (per-session cache) and explicitly passes the prior turns to
retrieval and the model — so multi-turn is controlled and verifiable (couldn't
confirm ChatInterface's history via the API, which strips it).
- Retrieval query now carries the last few user turns (not just one), so the
topic persists across several follow-ups.
- RULE 1 rebalanced: decline only when nothing relevant is present; answer when
it is (incl. follow-ups). Verified: maternity -> "financially how does it
change?" -> "and paternity leave?" all stay on topic and answer correctly;
car tax still declines; revenue/deficit correct.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

app/ollama_client.py

@@ -39,11 +39,12 @@ def _extra() -> dict:
 _SYSTEM = (
     "You are a precise analyst of the Republic of Mauritius National Budget 2026-2027, answering from "
     "the numbered budget excerpts in the latest message.\n"
-    "RULE 1 (most important): use an excerpt ONLY if it is specifically about the SAME subject as the "
-    "question. If none of the excerpts specifically address the asked subject, reply that the budget "
-    "documents provided do not appear to address it, and stop — it is better to decline than to answer "
-    "using an excerpt about a different subject (e.g. never answer about cars/vehicles using a "
-    "plastic-bottle, sugar or tobacco excerpt). NEVER invent, guess or estimate a figure.\n"
+    "RULE 1 (most important): answer only from information actually present in the excerpts. If the "
+    "excerpts contain nothing relevant to the question's topic, reply that the budget documents provided "
+    "do not appear to address it, and stop — never answer using an excerpt about an unrelated topic "
+    "(e.g. do not answer about cars/vehicles using a plastic-bottle, sugar or tobacco excerpt) and never "
+    "invent, guess or estimate a figure. But when relevant information IS present, answer it fully, "
+    "including for follow-up questions about a topic already discussed in the conversation.\n"
     "RULE 2: when the question IS covered, lead with the direct answer, cite the excerpts you use inline "
     "like [1], and do NOT repeat the excerpt headers (the 'Document: …' lines).\n"
     "RULE 3: be brief — answer in 1-2 short sentences. Quote figures exactly with their unit "

app/rag.py

@@ -75,11 +75,13 @@ def _clean_history(history: list | None) -> list:
 
 
 def _retrieval_query(message: str, history: list | None) -> str:
-    """Prepend the most recent user turn so a terse follow-up still retrieves the
-    right passages (e.g. 'financially how does it change' -> includes 'maternity')."""
+    """Prepend recent user turns so a terse follow-up still retrieves the right
+    passages (e.g. 'financially how does it change' -> keeps 'maternity'). Uses the
+    last few user turns so the topic persists across several follow-ups."""
     prev = [m["content"] for m in (history or [])
             if m.get("role") == "user" and isinstance(m.get("content"), str)]
-    return f"{prev[-1]} {message}".strip() if prev else message
+    recent = prev[-_HISTORY_TURNS:]
+    return (" ".join(recent + [message])).strip() if recent else message
 
 
 def answer(message: str, history: list | None = None) -> Iterator[str]:

app/ui.py

@@ -1,8 +1,10 @@
 """Gradio UI for the Mauritius Budget 2026-2027 RAG assistant.
 
-Pure-Python front end: a full-width, mobile-responsive streaming chat that keeps
-conversation context, with clickable/hoverable source citations and a collapsible
-list of indexed documents at the bottom."""
+A full-width, mobile-responsive chat built on explicit gr.Blocks state: the
+conversation lives in the Chatbot component (a per-session cache), and the bot
+handler passes the prior turns to retrieval + the model itself — so follow-up
+questions reliably keep their context. Citations are clickable/hoverable; the
+indexed-document list sits in a collapsible accordion at the bottom."""
 from __future__ import annotations
 
 import gradio as gr
@@ -13,13 +15,11 @@ EXAMPLES = [
     "What is the projected total revenue for 2026-2027?",
     "How much is allocated to education?",
     "How long is maternity leave being extended to?",
-    "What is the projected budget deficit and how is it financed?",
-    "Summarise the macroeconomic outlook and growth forecast.",
+    "What is the projected budget deficit?",
+    "Summarise the macroeconomic outlook.",
     "What are the main capital projects in the PSIP?",
 ]
 
-# Mobile-first: cap content width on desktop, let it go full-bleed on phones,
-# and make citation links stand out.
 CSS = """
 .gradio-container {max-width: 900px !important; margin: 0 auto !important}
 #intro h1 {font-size: 1.4rem; margin-bottom: .2rem}
@@ -44,12 +44,26 @@ def _docs_markdown() -> str:
     return "\n".join(lines)
 
 
-def _respond(message, history):
+def _add_user(message: str, history: list):
+    """Append the user's message to the conversation and clear the textbox."""
     message = (message or "").strip()
     if not message:
-        yield "Please type a question about the 2026-2027 budget."
+        return history, ""
+    return history + [{"role": "user", "content": message}], ""
+
+
+def _bot_reply(history: list):
+    """Stream the grounded answer, passing the PRIOR turns (the session cache)
+    to the RAG layer so follow-up questions keep their context."""
+    if not history or history[-1].get("role") != "user":
+        yield history
         return
-    yield from rag.answer(message, history)
+    message = history[-1]["content"]
+    prior = history[:-1]                      # everything before this question
+    history = history + [{"role": "assistant", "content": ""}]
+    for partial in rag.answer(message, prior):
+        history[-1]["content"] = partial
+        yield history
 
 
 def build_ui() -> gr.Blocks:
@@ -63,23 +77,28 @@ def build_ui() -> gr.Blocks:
             f"<sub>Model: <code>{model}</code> · local Ollama · retrieval-augmented (RAG)</sub>",
             elem_id="intro",
         )
-        gr.ChatInterface(
-            fn=_respond,
-            examples=EXAMPLES,
-            cache_examples=False,
-            chatbot=gr.Chatbot(
-                height=560,
-                render_markdown=True,
-                resizable=True,
-                latex_delimiters=[],
-                placeholder="Ask a question about the Mauritius 2026-2027 budget…",
-            ),
-            fill_height=True,
+        chatbot = gr.Chatbot(
+            height=520, render_markdown=True, resizable=True,
+            placeholder="Ask a question about the 2026-2027 budget…",
         )
+        with gr.Row():
+            msg = gr.Textbox(
+                placeholder="e.g. What is the projected total revenue for 2026-2027?",
+                scale=8, show_label=False, autofocus=True, submit_btn=True,
+            )
+            clear = gr.Button("Clear", scale=1, variant="secondary")
+
+        gr.Examples(examples=EXAMPLES, inputs=msg, label="Try one")
+
         with gr.Accordion("📄 Indexed documents", open=False):
             gr.Markdown(_docs_markdown())
             gr.Markdown(
                 "<sub>Source: Ministry of Finance, Economic Planning and Development, "
                 "Republic of Mauritius — National Budget 2026-2027.</sub>"
             )
+
+        # Submit -> show the user message, then stream the grounded reply.
+        msg.submit(_add_user, [msg, chatbot], [chatbot, msg], queue=False).then(
+            _bot_reply, chatbot, chatbot)
+        clear.click(lambda: ([], ""), None, [chatbot, msg], queue=False)
     return demo