diff --git a/README.md b/README.md index dc4ac927..be1a33b2 100644 --- a/README.md +++ b/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 diff --git a/docs/cli.md b/docs/cli.md index 3c857cd5..0dd19b0b 100644 --- a/docs/cli.md +++ b/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 diff --git a/docs/python.md b/docs/python.md index 8368bcc0..a498c22a 100644 --- a/docs/python.md +++ b/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)). diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 590966db..c75128cb 100644 --- a/src/haiku/rag/app.py +++ b/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]") diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index 8855d888..b0c282e0 100644 --- a/src/haiku/rag/cli.py +++ b/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") diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 51694138..c5a6eb96 100644 --- a/src/haiku/rag/client.py +++ b/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]: diff --git a/src/haiku/rag/qa/__init__.py b/src/haiku/rag/qa/__init__.py index 83e2451f..0d578cd6 100644 --- a/src/haiku/rag/qa/__init__.py +++ b/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}") diff --git a/src/haiku/rag/qa/anthropic.py b/src/haiku/rag/qa/anthropic.py index f6671ee4..6c516aff 100644 --- a/src/haiku/rag/qa/anthropic.py +++ b/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", diff --git a/src/haiku/rag/qa/base.py b/src/haiku/rag/qa/base.py index 21513a60..eea670df 100644 --- a/src/haiku/rag/qa/base.py +++ b/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( diff --git a/src/haiku/rag/qa/ollama.py b/src/haiku/rag/qa/ollama.py index f49315c0..c993b510 100644 --- a/src/haiku/rag/qa/ollama.py +++ b/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) diff --git a/src/haiku/rag/qa/openai.py b/src/haiku/rag/qa/openai.py index 54b2fc3e..27396d7a 100644 --- a/src/haiku/rag/qa/openai.py +++ b/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 ] diff --git a/src/haiku/rag/qa/prompts.py b/src/haiku/rag/qa/prompts.py index a1b38246..d55131be 100644 --- a/src/haiku/rag/qa/prompts.py +++ b/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: diff --git a/tests/test_app.py b/tests/test_app.py index f432d5ff..ca096fef 100644 --- a/tests/test_app.py +++ b/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) diff --git a/tests/test_cli.py b/tests/test_cli.py index 36289558..7f51af99 100644 --- a/tests/test_cli.py +++ b/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 \ No newline at end of file + 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 + ) diff --git a/tests/test_client.py b/tests/test_client.py index e2bb3b14..7c20c517 100644 --- a/tests/test_client.py +++ b/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?")