From a7850e2210c98f34a9d87953a3847c6416fae9bd Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Fri, 20 Feb 2026 11:02:46 +0200
Subject: [PATCH] Assign stable index to citations. Use a strong prompt
 preamble for our app agents

---
 app/backend/main.py                    | 13 ++++++-
 haiku_rag_slim/haiku/rag/chat/app.py   | 12 ++++++-
 haiku_rag_slim/haiku/rag/skills/rag.py |  4 +++
 tests/skills/test_rag.py               | 50 +++++++++++++++++++++++++-
 4 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/app/backend/main.py b/app/backend/main.py
index 5d9f0067..2e53d854 100644
--- a/app/backend/main.py
+++ b/app/backend/main.py
@@ -65,9 +65,20 @@ def get_client() -> HaikuRAG:
 # Create skill, toolset, and agent
 skill = create_skill(db_path=db_path, config=Config)
 toolset = SkillToolset(skills=[skill])
+
+AGENT_PREAMBLE = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.
+
+CRITICAL RULES:
+1. For greetings or casual chat: respond directly WITHOUT using any tools
+2. NEVER make up information - always use tools to get facts from the knowledge base
+3. For questions: Use the "ask" tool - it handles search and citation automatically
+4. For searches: Use the "search" tool - copy the ENTIRE tool response to your output INCLUDING content snippets
+5. When you use the "ask" tool, summarize the key findings and always include citations in your response
+"""
+
 agent = Agent(
     os.getenv("HAIKU_CHAT_MODEL", "openai:gpt-4o"),
-    instructions=toolset.system_prompt,
+    instructions=AGENT_PREAMBLE + toolset.system_prompt,
     toolsets=[toolset],
 )
 
diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py
index dd789b74..e3bcd107 100644
--- a/haiku_rag_slim/haiku/rag/chat/app.py
+++ b/haiku_rag_slim/haiku/rag/chat/app.py
@@ -54,6 +54,16 @@ except ImportError:  # pragma: no cover
 
 RAG_STATE_NAMESPACE = "rag"
 
+AGENT_PREAMBLE = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.
+
+CRITICAL RULES:
+1. For greetings or casual chat: respond directly WITHOUT using any tools
+2. NEVER make up information - always use tools to get facts from the knowledge base
+3. For questions: Use the "ask" tool - it handles search and citation automatically
+4. For searches: Use the "search" tool - copy the ENTIRE tool response to your output INCLUDING content snippets
+5. When you use the "ask" tool, summarize the key findings and always include citations in your response
+"""
+
 
 class ChatApp(App):
     """Textual TUI for conversational RAG."""
@@ -158,7 +168,7 @@ class ChatApp(App):
         self._toolset = SkillToolset(skills=[self._skill])
         self._agent = Agent(
             self._model,
-            instructions=self._toolset.system_prompt,
+            instructions=AGENT_PREAMBLE + self._toolset.system_prompt,
             toolsets=[self._toolset],
         )
         self._state = self._toolset.build_state_snapshot()
diff --git a/haiku_rag_slim/haiku/rag/skills/rag.py b/haiku_rag_slim/haiku/rag/skills/rag.py
index ca76fb9f..8eaa9186 100644
--- a/haiku_rag_slim/haiku/rag/skills/rag.py
+++ b/haiku_rag_slim/haiku/rag/skills/rag.py
@@ -172,6 +172,10 @@ def create_skill(
             answer, citations = await rag.ask(question)
 
         if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
+            next_index = len(ctx.deps.state.citations) + 1
+            for citation in citations:
+                citation.index = next_index
+                next_index += 1
             ctx.deps.state.citations.extend(citations)
             ctx.deps.state.qa_history.append(
                 QAHistoryEntry(question=question, answer=answer, citations=citations)
diff --git a/tests/skills/test_rag.py b/tests/skills/test_rag.py
index 17c64f0c..1cd657f3 100644
--- a/tests/skills/test_rag.py
+++ b/tests/skills/test_rag.py
@@ -203,7 +203,55 @@ class TestAskTool:
         assert len(state.qa_history) == 1
         assert isinstance(state.qa_history[0], QAHistoryEntry)
         assert state.qa_history[0].question == "What is AI?"
-        assert state.qa_history[0].citations == citations
+
+    async def test_ask_assigns_citation_indices(self, rag_db, monkeypatch):
+        from haiku.rag.skills.rag import RAGState, create_skill
+
+        first_citations = [
+            Citation(
+                document_id="d1",
+                chunk_id="c1",
+                document_uri="test://doc1",
+                content="First.",
+            ),
+            Citation(
+                document_id="d2",
+                chunk_id="c2",
+                document_uri="test://doc2",
+                content="Second.",
+            ),
+        ]
+        second_citations = [
+            Citation(
+                document_id="d3",
+                chunk_id="c3",
+                document_uri="test://doc3",
+                content="Third.",
+            ),
+        ]
+
+        call_count = 0
+
+        async def mock_ask(self, question):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                return ("Answer 1", first_citations)
+            return ("Answer 2", second_citations)
+
+        monkeypatch.setattr(HaikuRAG, "ask", mock_ask)
+
+        skill = create_skill(db_path=rag_db)
+        ask = _get_tool(skill, "ask")
+        state = RAGState()
+        ctx = _make_ctx(state)
+
+        await ask(ctx, question="First question")
+        assert state.citations[0].index == 1
+        assert state.citations[1].index == 2
+
+        await ask(ctx, question="Second question")
+        assert state.citations[2].index == 3
 
 
 class TestAnalyzeTool: