Change prompt and tool definition to to include citations in the QA agent.

src/haiku/rag/qa/anthropic.py

@@ -18,7 +18,7 @@ try:
             self.tools: Sequence[ToolParam] = [
                 ToolParam(
                     name="search_documents",
-                    description="Search the knowledge base for relevant documents",
+                    description="Search the knowledge base for relevant documents. Returns a JSON array with content, score, and document_uri for each result.",
                     input_schema={
                         "type": "object",
                         "properties": {

src/haiku/rag/qa/base.py

@@ -35,7 +35,7 @@ class QuestionAnswerAgentBase:
             "type": "function",
             "function": {
                 "name": "search_documents",
-                "description": "Search the knowledge base for relevant documents",
+                "description": "Search the knowledge base for relevant documents. Returns a JSON array of search results.",
                 "parameters": {
                     "type": "object",
                     "properties": {
@@ -51,6 +51,30 @@ class QuestionAnswerAgentBase:
                     },
                     "required": ["query"],
                 },
+                "returns": {
+                    "type": "string",
+                    "description": "JSON array of search results",
+                    "schema": {
+                        "type": "array",
+                        "items": {
+                            "type": "object",
+                            "properties": {
+                                "content": {
+                                    "type": "string",
+                                    "description": "The document text content",
+                                },
+                                "score": {
+                                    "type": "number",
+                                    "description": "Relevance score (higher is more relevant)",
+                                },
+                                "document_uri": {
+                                    "type": "string",
+                                    "description": "Source URI/path of the document",
+                                },
+                            },
+                        },
+                    },
+                },
             },
         }
     ]

src/haiku/rag/qa/prompts.py

@@ -1,21 +1,36 @@
 SYSTEM_PROMPT = """
 You are a knowledgeable assistant that helps users find information from a document knowledge base.
 
+IMPORTANT: You MUST use the search_documents tool for every question. Do not answer any question without first searching the knowledge base.
+
 Your process:
-1. When a user asks a question, use the search_documents tool to find relevant information
+1. IMMEDIATELY call the search_documents tool with relevant keywords from the user's question
-2. Search with specific keywords and phrases from the user's question
+2. Review the search results and their relevance scores
-3. Review the search results and their relevance scores
+3. If you need additional context, perform follow-up searches with different keywords
-4. If you need additional context, perform follow-up searches with different keywords
+4. Provide a short and to the point comprehensive answer based only on the retrieved documents
-5. Provide a short and to the point comprehensive answer based only on the retrieved documents
+5. Always include citations for the sources used in your answer
 
 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.
+- ALWAYS include citations at the end of your response using the format below
+
+Citation Format:
+After your answer, include a "Citations:" section that lists:
+- The document URI from each search result used
+- A brief excerpt (first 50-100 characters) of the content that supported your answer
+- Format: "Citations:\n- [document_uri]: [content_excerpt]..."
+
+Example response format:
+[Your answer here]
+
+Citations:
+- /path/to/document1.pdf: "This document explains that AFMAN stands for Air Force Manual..."
+- /path/to/document2.pdf: "The manual provides guidance on military procedures and..."
 
 Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
 """