Merge pull request #92 from ggozad/feat/customizable-system-prompt

Customizable system prompt for Q/A agent

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)).

src/evaluations/benchmark.py

@@ -15,6 +15,7 @@ from rich.progress import Progress
 from evaluations.config import DatasetSpec, RetrievalSample
 from evaluations.datasets import DATASETS
 from evaluations.llm_judge import ANSWER_EQUIVALENCE_RUBRIC
+from evaluations.prompts import WIX_SUPPORT_PROMPT
 from haiku.rag import logging  # noqa: F401
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
@@ -204,7 +205,8 @@ async def run_qa_benchmark(
         )
 
         async with HaikuRAG(spec.db_path) as rag:
-            qa = get_qa_agent(rag)
+            system_prompt = WIX_SUPPORT_PROMPT if spec.key == "wix" else None
+            qa = get_qa_agent(rag, system_prompt=system_prompt)
 
             async def answer_question(question: str) -> str:
                 return await qa.answer(question)

src/evaluations/prompts.py

@@ -0,0 +1,22 @@
+WIX_SUPPORT_PROMPT = """
+You are a WIX technical support expert helping users with questions about the WIX platform.
+
+Your process:
+1. When a user asks a question, use the search_documents tool to find relevant information
+2. Search with specific keywords and phrases from the user's question
+3. Review the search results and their relevance scores
+4. If you need additional context, perform follow-up searches with different keywords
+5. Provide a short and to the point comprehensive answer based only on the retrieved documents
+
+Guidelines:
+- Base your answers strictly on the provided document content
+- Quote or reference specific information when possible
+- If multiple documents contain relevant information, synthesize them coherently
+- Indicate when information is incomplete or when you need to search for additional context
+- If the retrieved documents don't contain sufficient information, clearly state: "I cannot find enough information in the knowledge base to answer this question."
+- For complex questions, consider breaking them down and performing multiple searches
+- Stick to the answer, do not ellaborate or provide context unless explicitly asked for it.
+
+Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
+/no_think
+"""

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]:

src/haiku/rag/qa/__init__.py

@@ -3,7 +3,11 @@ from haiku.rag.config import Config
 from haiku.rag.qa.agent import QuestionAnswerAgent
 
 
-def get_qa_agent(client: HaikuRAG, use_citations: bool = False) -> QuestionAnswerAgent:
+def get_qa_agent(
+    client: HaikuRAG,
+    use_citations: bool = False,
+    system_prompt: str | None = None,
+) -> QuestionAnswerAgent:
     provider = Config.QA_PROVIDER
     model_name = Config.QA_MODEL
 
@@ -12,4 +16,5 @@ def get_qa_agent(client: HaikuRAG, use_citations: bool = False) -> QuestionAnswe
         provider=provider,
         model=model_name,
         use_citations=use_citations,
+        system_prompt=system_prompt,
     )

src/haiku/rag/qa/agent.py

@@ -30,12 +30,14 @@ class QuestionAnswerAgent:
         model: str,
         use_citations: bool = False,
         q: float = 0.0,
+        system_prompt: str | None = None,
     ):
         self._client = client
 
-        system_prompt = (
-            QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT
-        )
+        if system_prompt is None:
+            system_prompt = (
+                QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT
+            )
         model_obj = self._get_model(provider, model)
 
         self._agent = Agent(