From 0991df6d9e3bd99002aca65d45af2b54ea393dc2 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Tue, 7 Oct 2025 13:15:58 +0300
Subject: [PATCH] Make it possible to override the QA agent system prompt using
 the client

---
 docs/python.md          | 13 +++++++++++++
 src/haiku/rag/client.py |  7 +++++--
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/docs/python.md b/docs/python.md
index 78649480..517599d0 100644
--- a/docs/python.md
+++ b/docs/python.md
@@ -207,6 +207,19 @@ answer = await client.ask("Who is the author of haiku.rag?", cite=True)
 print(answer)
 ```
 
+Customize the QA agent's behavior with a custom system prompt:
+
+```python
+custom_prompt = """You are a technical support expert for WIX.
+Answer questions based on the knowledge base documents provided.
+Be concise and helpful."""
+
+answer = await client.ask(
+    "How do I create a blog?",
+    system_prompt=custom_prompt
+)
+```
+
 The QA agent will search your documents for relevant information and use the configured LLM to generate a comprehensive answer. With `cite=True`, responses include citations showing which documents were used as sources. Citations prefer the document title when present, otherwise they use the URI.
 
 The QA provider and model can be configured via environment variables (see [Configuration](configuration.md)).
diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py
index e4107b86..193025b9 100644
--- a/src/haiku/rag/client.py
+++ b/src/haiku/rag/client.py
@@ -525,19 +525,22 @@ class HaikuRAG:
         merged.append(current)
         return merged
 
-    async def ask(self, question: str, cite: bool = False) -> str:
+    async def ask(
+        self, question: str, cite: bool = False, system_prompt: str | None = None
+    ) -> str:
         """Ask a question using the configured QA agent.
 
         Args:
             question: The question to ask.
             cite: Whether to include citations in the response.
+            system_prompt: Optional custom system prompt for the QA agent.
 
         Returns:
             The generated answer as a string.
         """
         from haiku.rag.qa import get_qa_agent
 
-        qa_agent = get_qa_agent(self, use_citations=cite)
+        qa_agent = get_qa_agent(self, use_citations=cite, system_prompt=system_prompt)
         return await qa_agent.answer(question)
 
     async def rebuild_database(self) -> AsyncGenerator[str, None]: