Merge pull request #36 from ggozad/feat/citations

Support for requesting citations in the QA agent and ask() methods.

README.md

@@ -33,6 +33,9 @@ haiku-rag search "query"
 # Ask questions
 haiku-rag ask "Who is the author of haiku.rag?"
 
+# Ask questions with citations
+haiku-rag ask "Who is the author of haiku.rag?" --cite
+
 # Rebuild database (re-chunk and re-embed all documents)
 haiku-rag rebuild
 
@@ -58,6 +61,10 @@ async with HaikuRAG("database.db") as client:
     # Ask questions
     answer = await client.ask("Who is the author of haiku.rag?")
     print(answer)
+
+    # Ask questions with citations
+    answer = await client.ask("Who is the author of haiku.rag?", cite=True)
+    print(answer)
 ```
 
 ## MCP Server

docs/cli.md

@@ -64,7 +64,12 @@ Ask questions about your documents:
 haiku-rag ask "Who is the author of haiku.rag?"
 ```
 
-The QA agent will search your documents for relevant information and provide a comprehensive answer.
+Ask questions with citations showing source documents:
+```bash
+haiku-rag ask "Who is the author of haiku.rag?" --cite
+```
+
+The QA agent will search your documents for relevant information and provide a comprehensive answer. With `--cite`, responses include citations showing which documents were used.
 
 ## Configuration
 

docs/python.md

@@ -139,6 +139,13 @@ answer = await client.ask("Who is the author of haiku.rag?")
 print(answer)
 ```
 
-The QA agent will search your documents for relevant information and use the configured LLM to generate a comprehensive answer.
+Ask questions with citations showing source documents:
+
+```python
+answer = await client.ask("Who is the author of haiku.rag?", cite=True)
+print(answer)
+```
+
+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.
 
 The QA provider and model can be configured via environment variables (see [Configuration](configuration.md)).

src/haiku/rag/app.py

@@ -62,10 +62,10 @@ class HaikuRAGApp:
             for chunk, score in results:
                 self._rich_print_search_result(chunk, score)
 
-    async def ask(self, question: str):
+    async def ask(self, question: str, cite: bool = False):
         async with HaikuRAG(db_path=self.db_path) as self.client:
             try:
-                answer = await self.client.ask(question)
+                answer = await self.client.ask(question, cite=cite)
                 self.console.print(f"[bold blue]Question:[/bold blue] {question}")
                 self.console.print()
                 self.console.print("[bold green]Answer:[/bold green]")

src/haiku/rag/cli.py

@@ -160,9 +160,14 @@ def ask(
         "--db",
         help="Path to the SQLite database file",
     ),
+    cite: bool = typer.Option(
+        False,
+        "--cite",
+        help="Include citations in the response",
+    ),
 ):
     app = HaikuRAGApp(db_path=db)
-    asyncio.run(app.ask(question=question))
+    asyncio.run(app.ask(question=question, cite=cite))
 
 
 @cli.command("settings", help="Display current configuration settings")

src/haiku/rag/client.py

@@ -348,18 +348,19 @@ class HaikuRAG:
         # Return reranked results with scores from reranker
         return reranked_results
 
-    async def ask(self, question: str) -> str:
+    async def ask(self, question: str, cite: bool = False) -> str:
         """Ask a question using the configured QA agent.
 
         Args:
             question: The question to ask.
+            cite: Whether to include citations in the response.
 
         Returns:
             The generated answer as a string.
         """
         from haiku.rag.qa import get_qa_agent
 
-        qa_agent = get_qa_agent(self)
+        qa_agent = get_qa_agent(self, use_citations=cite)
         return await qa_agent.answer(question)
 
     async def rebuild_database(self) -> AsyncGenerator[int, None]:

src/haiku/rag/qa/__init__.py

@@ -4,12 +4,16 @@ from haiku.rag.qa.base import QuestionAnswerAgentBase
 from haiku.rag.qa.ollama import QuestionAnswerOllamaAgent
 
 
-def get_qa_agent(client: HaikuRAG, model: str = "") -> QuestionAnswerAgentBase:
+def get_qa_agent(
+    client: HaikuRAG, model: str = "", use_citations: bool = False
+) -> QuestionAnswerAgentBase:
     """
     Factory function to get the appropriate QA agent based on the configuration.
     """
     if Config.QA_PROVIDER == "ollama":
-        return QuestionAnswerOllamaAgent(client, model or Config.QA_MODEL)
+        return QuestionAnswerOllamaAgent(
+            client, model or Config.QA_MODEL, use_citations
+        )
 
     if Config.QA_PROVIDER == "openai":
         try:
@@ -20,7 +24,9 @@ def get_qa_agent(client: HaikuRAG, model: str = "") -> QuestionAnswerAgentBase:
                 "Please install haiku.rag with the 'openai' extra:"
                 "uv pip install haiku.rag[openai]"
             )
-        return QuestionAnswerOpenAIAgent(client, model or Config.QA_MODEL)
+        return QuestionAnswerOpenAIAgent(
+            client, model or Config.QA_MODEL, use_citations
+        )
 
     if Config.QA_PROVIDER == "anthropic":
         try:
@@ -31,6 +37,8 @@ def get_qa_agent(client: HaikuRAG, model: str = "") -> QuestionAnswerAgentBase:
                 "Please install haiku.rag with the 'anthropic' extra:"
                 "uv pip install haiku.rag[anthropic]"
             )
-        return QuestionAnswerAnthropicAgent(client, model or Config.QA_MODEL)
+        return QuestionAnswerAnthropicAgent(
+            client, model or Config.QA_MODEL, use_citations
+        )
 
     raise ValueError(f"Unsupported QA provider: {Config.QA_PROVIDER}")

src/haiku/rag/qa/anthropic.py

@@ -1,19 +1,29 @@
 from collections.abc import Sequence
 
 try:
-    from anthropic import AsyncAnthropic
-    from anthropic.types import MessageParam, TextBlock, ToolParam, ToolUseBlock
+    from anthropic import AsyncAnthropic  # type: ignore
+    from anthropic.types import (  # type: ignore
+        MessageParam,
+        TextBlock,
+        ToolParam,
+        ToolUseBlock,
+    )
 
     from haiku.rag.client import HaikuRAG
     from haiku.rag.qa.base import QuestionAnswerAgentBase
 
     class QuestionAnswerAnthropicAgent(QuestionAnswerAgentBase):
-        def __init__(self, client: HaikuRAG, model: str = "claude-3-5-haiku-20241022"):
-            super().__init__(client, model or self._model)
+        def __init__(
+            self,
+            client: HaikuRAG,
+            model: str = "claude-3-5-haiku-20241022",
+            use_citations: bool = False,
+        ):
+            super().__init__(client, model or self._model, use_citations)
             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": {
@@ -73,13 +83,7 @@ try:
                                     query, limit=limit
                                 )
 
-                                context_chunks = []
-                                for chunk, score in search_results:
-                                    context_chunks.append(
-                                        f"Content: {chunk.content}\nScore: {score:.4f}"
-                                    )
-
-                                context = "\n\n".join(context_chunks)
+                                context = self._format_search_results(search_results)
 
                                 tool_results.append(
                                     {

src/haiku/rag/qa/base.py

@@ -1,26 +1,44 @@
+import json
+
 from haiku.rag.client import HaikuRAG
-from haiku.rag.qa.prompts import SYSTEM_PROMPT
+from haiku.rag.qa.prompts import SYSTEM_PROMPT, SYSTEM_PROMPT_WITH_CITATIONS
 
 
 class QuestionAnswerAgentBase:
     _model: str = ""
     _system_prompt: str = SYSTEM_PROMPT
 
-    def __init__(self, client: HaikuRAG, model: str = ""):
+    def __init__(self, client: HaikuRAG, model: str = "", use_citations: bool = False):
         self._model = model
         self._client = client
+        self._system_prompt = (
+            SYSTEM_PROMPT_WITH_CITATIONS if use_citations else SYSTEM_PROMPT
+        )
 
     async def answer(self, question: str) -> str:
         raise NotImplementedError(
             "QABase is an abstract class. Please implement the answer method in a subclass."
         )
 
+    def _format_search_results(self, search_results) -> str:
+        """Format search results as JSON list of {content, score, document_uri}"""
+        formatted_results = []
+        for chunk, score in search_results:
+            formatted_results.append(
+                {
+                    "content": chunk.content,
+                    "score": score,
+                    "document_uri": chunk.document_uri,
+                }
+            )
+        return json.dumps(formatted_results, indent=2)
+
     tools = [
         {
             "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": {
@@ -36,6 +54,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/ollama.py

@@ -8,8 +8,13 @@ OLLAMA_OPTIONS = {"temperature": 0.0, "seed": 42, "num_ctx": 16384}
 
 
 class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase):
-    def __init__(self, client: HaikuRAG, model: str = Config.QA_MODEL):
-        super().__init__(client, model or self._model)
+    def __init__(
+        self,
+        client: HaikuRAG,
+        model: str = Config.QA_MODEL,
+        use_citations: bool = False,
+    ):
+        super().__init__(client, model or self._model, use_citations)
 
     async def answer(self, question: str) -> str:
         ollama_client = AsyncClient(host=Config.OLLAMA_BASE_URL)
@@ -41,14 +46,7 @@ class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase):
 
                         search_results = await self._client.search(query, limit=limit)
 
-                        context_chunks = []
-                        for chunk, score in search_results:
-                            context_chunks.append(
-                                f"Content: {chunk.content}\nScore: {score:.4f}"
-                            )
-
-                        context = "\n\n".join(context_chunks)
-
+                        context = self._format_search_results(search_results)
                         messages.append(
                             {
                                 "role": "tool",

src/haiku/rag/qa/openai.py

@@ -1,22 +1,29 @@
 from collections.abc import Sequence
 
 try:
-    from openai import AsyncOpenAI
-    from openai.types.chat import (
+    from openai import AsyncOpenAI  # type: ignore
+    from openai.types.chat import (  # type: ignore
         ChatCompletionAssistantMessageParam,
         ChatCompletionMessageParam,
         ChatCompletionSystemMessageParam,
         ChatCompletionToolMessageParam,
         ChatCompletionUserMessageParam,
     )
-    from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam
+    from openai.types.chat.chat_completion_tool_param import (  # type: ignore
+        ChatCompletionToolParam,
+    )
 
     from haiku.rag.client import HaikuRAG
     from haiku.rag.qa.base import QuestionAnswerAgentBase
 
     class QuestionAnswerOpenAIAgent(QuestionAnswerAgentBase):
-        def __init__(self, client: HaikuRAG, model: str = "gpt-4o-mini"):
-            super().__init__(client, model or self._model)
+        def __init__(
+            self,
+            client: HaikuRAG,
+            model: str = "gpt-4o-mini",
+            use_citations: bool = False,
+        ):
+            super().__init__(client, model or self._model, use_citations)
             self.tools: Sequence[ChatCompletionToolParam] = [
                 ChatCompletionToolParam(tool) for tool in self.tools
             ]
@@ -74,13 +81,7 @@ try:
                                 query, limit=limit
                             )
 
-                            context_chunks = []
-                            for chunk, score in search_results:
-                                context_chunks.append(
-                                    f"Content: {chunk.content}\nScore: {score:.4f}"
-                                )
-
-                            context = "\n\n".join(context_chunks)
+                            context = self._format_search_results(search_results)
 
                             messages.append(
                                 ChatCompletionToolMessageParam(

src/haiku/rag/qa/prompts.py

@@ -19,3 +19,40 @@ Guidelines:
 
 Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
 """
+
+SYSTEM_PROMPT_WITH_CITATIONS = """
+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. IMMEDIATELY call the search_documents tool with relevant keywords from the user's question
+2. Review the search results and their relevance scores
+3. 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. Always include citations for the sources used in your answer
+
+Guidelines:
+- Base your answers strictly on the provided document content
+- 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.
+"""

tests/test_app.py

@@ -75,7 +75,7 @@ async def test_add_document_from_source(app: HaikuRAGApp, monkeypatch):
     monkeypatch.setattr(app, "_rich_print_document", mock_rich_print)
     monkeypatch.setattr(app.console, "print", mock_print)
 
-    file_path = Path("test.txt")
+    file_path = "test.txt"
     with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
         await app.add_document_from_source(file_path)
 
@@ -206,3 +206,37 @@ async def test_serve(app: HaikuRAGApp, monkeypatch, transport):
         mock_server.run_http_async.assert_called_once_with("streamable-http")
 
     mock_task.cancel.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_ask_without_cite(app: HaikuRAGApp, monkeypatch):
+    """Test asking a question without citations."""
+    mock_answer = "Test answer"
+    mock_client = AsyncMock()
+    mock_client.ask.return_value = mock_answer
+    mock_client.__aenter__.return_value = mock_client
+
+    mock_print = MagicMock()
+    monkeypatch.setattr(app.console, "print", mock_print)
+
+    with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
+        await app.ask("test question")
+
+    mock_client.ask.assert_called_once_with("test question", cite=False)
+
+
+@pytest.mark.asyncio
+async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch):
+    """Test asking a question with citations."""
+    mock_answer = "Test answer with citations"
+    mock_client = AsyncMock()
+    mock_client.ask.return_value = mock_answer
+    mock_client.__aenter__.return_value = mock_client
+
+    mock_print = MagicMock()
+    monkeypatch.setattr(app.console, "print", mock_print)
+
+    with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
+        await app.ask("test question", cite=True)
+
+    mock_client.ask.assert_called_once_with("test question", cite=True)

tests/test_cli.py

@@ -126,4 +126,32 @@ def test_serve_stdio_and_sse():
         result = runner.invoke(cli, ["serve", "--stdio", "--sse"])
 
         assert result.exit_code == 1
-        assert "Error: Cannot use both --stdio and --http options" in result.stdout
+        assert "Error: Cannot use both --stdio and --http options" in result.stdout
+
+
+def test_ask():
+    with patch("haiku.rag.cli.HaikuRAGApp") as mock_app:
+        mock_app_instance = MagicMock()
+        mock_app_instance.ask = AsyncMock()
+        mock_app.return_value = mock_app_instance
+
+        result = runner.invoke(cli, ["ask", "What is Python?"])
+
+        assert result.exit_code == 0
+        mock_app_instance.ask.assert_called_once_with(
+            question="What is Python?", cite=False
+        )
+
+
+def test_ask_with_cite():
+    with patch("haiku.rag.cli.HaikuRAGApp") as mock_app:
+        mock_app_instance = MagicMock()
+        mock_app_instance.ask = AsyncMock()
+        mock_app.return_value = mock_app_instance
+
+        result = runner.invoke(cli, ["ask", "What is Python?", "--cite"])
+
+        assert result.exit_code == 0
+        mock_app_instance.ask.assert_called_once_with(
+            question="What is Python?", cite=True
+        )

tests/test_client.py

@@ -489,3 +489,33 @@ async def test_client_create_document_with_custom_chunks():
             assert (
                 chunk.metadata["custom"] == f"metadata{i + 1}"
             )  # Original metadata preserved
+
+
+@pytest.mark.asyncio
+async def test_client_ask_without_cite():
+    """Test asking questions without citations."""
+    async with HaikuRAG(":memory:") as client:
+        # Mock the QA agent
+        mock_qa_agent = AsyncMock()
+        mock_qa_agent.answer.return_value = "Test answer"
+
+        with patch("haiku.rag.qa.get_qa_agent", return_value=mock_qa_agent):
+            answer = await client.ask("What is Python?")
+
+        assert answer == "Test answer"
+        mock_qa_agent.answer.assert_called_once_with("What is Python?")
+
+
+@pytest.mark.asyncio
+async def test_client_ask_with_cite():
+    """Test asking questions with citations."""
+    async with HaikuRAG(":memory:") as client:
+        # Mock the QA agent
+        mock_qa_agent = AsyncMock()
+        mock_qa_agent.answer.return_value = "Test answer with citations [1]"
+
+        with patch("haiku.rag.qa.get_qa_agent", return_value=mock_qa_agent):
+            answer = await client.ask("What is Python?", cite=True)
+
+        assert answer == "Test answer with citations [1]"
+        mock_qa_agent.answer.assert_called_once_with("What is Python?")