Use custom prompt when evaluating with the Wix dataset to make the QA agent act as an assistant

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/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(