Make citations optional, through cite parameter

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

@@ -13,8 +13,13 @@ try:
     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",

src/haiku/rag/qa/base.py

@@ -1,16 +1,19 @@
 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(

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)

src/haiku/rag/qa/openai.py

@@ -17,8 +17,13 @@ try:
     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
             ]

src/haiku/rag/qa/prompts.py

@@ -1,6 +1,28 @@
 SYSTEM_PROMPT = """
 You are a knowledgeable assistant that helps users find information from a document knowledge base.
 
+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.
+"""
+
+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:

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?")