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 8827c5cb..6c516aff 100644 --- a/src/haiku/rag/qa/anthropic.py +++ b/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( { diff --git a/src/haiku/rag/qa/base.py b/src/haiku/rag/qa/base.py index 0ff2a55b..eea670df 100644 --- a/src/haiku/rag/qa/base.py +++ b/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", + }, + }, + }, + }, + }, }, } ] diff --git a/src/haiku/rag/qa/ollama.py b/src/haiku/rag/qa/ollama.py index 7521d8bf..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) @@ -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", diff --git a/src/haiku/rag/qa/openai.py b/src/haiku/rag/qa/openai.py index 24f58cf9..27396d7a 100644 --- a/src/haiku/rag/qa/openai.py +++ b/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( diff --git a/src/haiku/rag/qa/prompts.py b/src/haiku/rag/qa/prompts.py index af490fd7..d55131be 100644 --- a/src/haiku/rag/qa/prompts.py +++ b/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. +""" 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?")