diff --git a/README.md b/README.md index bc3a9567..873364f4 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,7 @@ Retrieval-Augmented Generation (RAG) library on SQLite. - **Local SQLite**: No external servers required - **Multiple embedding providers**: Ollama, VoyageAI, OpenAI +- **Multiple QA providers**: Ollama, OpenAI, Anthropic - **Hybrid search**: Vector + full-text search with Reciprocal Rank Fusion - **Question answering**: Built-in QA agents on your documents - **File monitoring**: Auto-index files when run as server diff --git a/docs/configuration.md b/docs/configuration.md index cae8a1dd..a9846508 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -84,6 +84,22 @@ QA_MODEL="gpt-4o-mini" # or gpt-4, gpt-3.5-turbo, etc. OPENAI_API_KEY="your-api-key" ``` +### Anthropic + +For Anthropic QA, you need to install haiku.rag with Anthropic extras: + +```bash +uv pip install haiku.rag --extra anthropic +``` + +Then configure: + +```bash +QA_PROVIDER="anthropic" +QA_MODEL="claude-3-5-haiku-20241022" # or claude-3-5-sonnet-20241022, etc. +ANTHROPIC_API_KEY="your-api-key" +``` + ## Other Settings ### Database and Storage diff --git a/docs/index.md b/docs/index.md index 66117c63..19da8d6d 100644 --- a/docs/index.md +++ b/docs/index.md @@ -8,7 +8,7 @@ - **Local SQLite**: No need to run additional servers - **Support for various embedding providers**: Ollama, VoyageAI, OpenAI or add your own - **Hybrid Search**: Vector search using `sqlite-vec` combined with full-text search `FTS5`, using Reciprocal Rank Fusion -- **Question Answering**: Built-in QA agents using Ollama or OpenAI. +- **Question Answering**: Built-in QA agents using Ollama, OpenAI, or Anthropic. - **File monitoring**: Automatically index files when run as a server - **Extended file format support**: Parse 40+ file formats including PDF, DOCX, HTML, Markdown, audio and more. Or add a URL! - **MCP server**: Exposes functionality as MCP tools diff --git a/docs/installation.md b/docs/installation.md index b3da3847..fd5eb509 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -24,6 +24,12 @@ uv pip install haiku.rag --extra voyageai uv pip install haiku.rag --extra openai ``` +### Anthropic + +```bash +uv pip install haiku.rag --extra anthropic +``` + ## Requirements - Python 3.10+ diff --git a/pyproject.toml b/pyproject.toml index 3e33e1a8..f0b0b457 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ dependencies = [ [project.optional-dependencies] voyageai = ["voyageai>=0.3.2"] openai = ["openai>=1.0.0"] +anthropic = ["anthropic>=0.56.0"] [project.scripts] haiku-rag = "haiku.rag.cli:cli" diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index c388064f..e1552873 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -30,6 +30,7 @@ class AppConfig(BaseModel): # Provider keys VOYAGE_API_KEY: str = "" OPENAI_API_KEY: str = "" + ANTHROPIC_API_KEY: str = "" @field_validator("MONITOR_DIRECTORIES", mode="before") @classmethod @@ -49,3 +50,5 @@ if Config.OPENAI_API_KEY: os.environ["OPENAI_API_KEY"] = Config.OPENAI_API_KEY if Config.VOYAGE_API_KEY: os.environ["VOYAGE_API_KEY"] = Config.VOYAGE_API_KEY +if Config.ANTHROPIC_API_KEY: + os.environ["ANTHROPIC_API_KEY"] = Config.ANTHROPIC_API_KEY diff --git a/src/haiku/rag/qa/__init__.py b/src/haiku/rag/qa/__init__.py index 69ff4523..bcf53380 100644 --- a/src/haiku/rag/qa/__init__.py +++ b/src/haiku/rag/qa/__init__.py @@ -23,4 +23,17 @@ def get_qa_agent(client: HaikuRAG, model: str = "") -> QuestionAnswerAgentBase: ) return QuestionAnswerOpenAIAgent(client, model or "gpt-4o-mini") + if Config.QA_PROVIDER == "anthropic": + try: + from haiku.rag.qa.anthropic import QuestionAnswerAnthropicAgent + except ImportError: + raise ImportError( + "Anthropic QA agent requires the 'anthropic' package. " + "Please install haiku.rag with the 'anthropic' extra:" + "uv pip install haiku.rag --extra anthropic" + ) + return QuestionAnswerAnthropicAgent( + client, model or "claude-3-5-haiku-20241022" + ) + 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 new file mode 100644 index 00000000..5b4479b3 --- /dev/null +++ b/src/haiku/rag/qa/anthropic.py @@ -0,0 +1,112 @@ +from collections.abc import Sequence + +try: + from anthropic import AsyncAnthropic + from anthropic.types import 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) + self.tools: Sequence[ToolParam] = [ + ToolParam( + name="search_documents", + description="Search the knowledge base for relevant documents", + input_schema={ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query to find relevant documents", + }, + "limit": { + "type": "integer", + "description": "Maximum number of results to return", + "default": 3, + }, + }, + "required": ["query"], + }, + ) + ] + + async def answer(self, question: str) -> str: + anthropic_client = AsyncAnthropic() + + messages: list[MessageParam] = [{"role": "user", "content": question}] + + response = await anthropic_client.messages.create( + model=self._model, + max_tokens=4096, + system=self._system_prompt, + messages=messages, + tools=self.tools, + temperature=0.0, + ) + + if response.stop_reason == "tool_use": + messages.append({"role": "assistant", "content": response.content}) + + # Process tool calls + tool_results = [] + for content_block in response.content: + if isinstance(content_block, ToolUseBlock): + if content_block.name == "search_documents": + args = content_block.input + query = ( + args.get("query", question) + if isinstance(args, dict) + else question + ) + limit = ( + int(args.get("limit", 3)) + if isinstance(args, dict) + else 3 + ) + + 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) + + tool_results.append( + { + "type": "tool_result", + "tool_use_id": content_block.id, + "content": context, + } + ) + + if tool_results: + messages.append({"role": "user", "content": tool_results}) + + final_response = await anthropic_client.messages.create( + model=self._model, + max_tokens=4096, + system=self._system_prompt, + messages=messages, + temperature=0.0, + ) + if final_response.content: + first_content = final_response.content[0] + if isinstance(first_content, TextBlock): + return first_content.text + return "" + + if response.content: + first_content = response.content[0] + if isinstance(first_content, TextBlock): + return first_content.text + return "" + +except ImportError: + pass diff --git a/tests/test_qa.py b/tests/test_qa.py index 41312a75..686fb53f 100644 --- a/tests/test_qa.py +++ b/tests/test_qa.py @@ -12,11 +12,19 @@ except ImportError: QuestionAnswerOpenAIAgent = None OPENAI_AVAILABLE = False +try: + from haiku.rag.qa.anthropic import QuestionAnswerAnthropicAgent + + ANTHROPIC_AVAILABLE = True +except ImportError: + QuestionAnswerAnthropicAgent = None + ANTHROPIC_AVAILABLE = False + from .llm_judge import LLMJudge @pytest.mark.asyncio -async def test_qa_ollama_with_dataset_question(qa_corpus: Dataset): +async def test_qa_ollama(qa_corpus: Dataset): """Test QA with actual question from the dataset using LLM judge.""" client = HaikuRAG(":memory:") qa = QuestionAnswerOllamaAgent(client) @@ -40,7 +48,7 @@ async def test_qa_ollama_with_dataset_question(qa_corpus: Dataset): @pytest.mark.asyncio @pytest.mark.skipif(not OPENAI_AVAILABLE, reason="OpenAI not available") -async def test_qa_openai_basic(qa_corpus: Dataset): +async def test_qa_openai(qa_corpus: Dataset): """Test OpenAI QA basic functionality.""" client = HaikuRAG(":memory:") qa = QuestionAnswerOpenAIAgent(client) # type: ignore @@ -60,3 +68,27 @@ async def test_qa_openai_basic(qa_corpus: Dataset): assert is_equivalent, ( f"Generated answer not equivalent to expected answer.\nQuestion: {question}\nGenerated: {answer}\nExpected: {expected_answer}" ) + + +@pytest.mark.asyncio +@pytest.mark.skipif(not ANTHROPIC_AVAILABLE, reason="Anthropic not available") +async def test_qa_anthropic(qa_corpus: Dataset): + """Test Anthropic QA basic functionality.""" + client = HaikuRAG(":memory:") + qa = QuestionAnswerAnthropicAgent(client) # type: ignore + llm_judge = LLMJudge() + + doc = qa_corpus[1] + await client.create_document( + content=doc["document_extracted"], uri=doc["document_id"] + ) + + question = doc["question"] + expected_answer = doc["answer"] + + answer = await qa.answer(question) + is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer) + + assert is_equivalent, ( + f"Generated answer not equivalent to expected answer.\nQuestion: {question}\nGenerated: {answer}\nExpected: {expected_answer}" + ) diff --git a/uv.lock b/uv.lock index 176dc658..d12ce5a2 100644 --- a/uv.lock +++ b/uv.lock @@ -133,6 +133,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, ] +[[package]] +name = "anthropic" +version = "0.56.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "distro" }, + { name = "httpx" }, + { name = "jiter" }, + { name = "pydantic" }, + { name = "sniffio" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4d/40/0c4eb5728466849803782c8a86eb315af1a6eb0efea6a751de120ab845c9/anthropic-0.56.0.tar.gz", hash = "sha256:56fa9eb61afa004a1664bc85eed071e77b96c579b77395e9cc893097e599f72e", size = 421538, upload-time = "2025-07-01T19:39:10.805Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e5/90/7f4d4084f9c35c3ea3e784646ec12f9b2c8cf8743b2bb5489252659b5bda/anthropic-0.56.0-py3-none-any.whl", hash = "sha256:91f1f74abdcf0958d3296b657304588cc244b1107b89f973ff6f511afdacfc56", size = 289603, upload-time = "2025-07-01T19:39:08.794Z" }, +] + [[package]] name = "anyio" version = "4.9.0" @@ -815,6 +833,9 @@ dependencies = [ ] [package.optional-dependencies] +anthropic = [ + { name = "anthropic" }, +] openai = [ { name = "openai" }, ] @@ -837,6 +858,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "anthropic", marker = "extra == 'anthropic'", specifier = ">=0.56.0" }, { name = "fastmcp", specifier = ">=2.8.1" }, { name = "httpx", specifier = ">=0.28.1" }, { name = "markitdown", extras = ["audio-transcription", "docx", "pdf", "pptx", "xlsx"], specifier = ">=0.1.2" }, @@ -851,7 +873,7 @@ requires-dist = [ { name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.2" }, { name = "watchfiles", specifier = ">=1.1.0" }, ] -provides-extras = ["voyageai", "openai"] +provides-extras = ["voyageai", "openai", "anthropic"] [package.metadata.requires-dev] dev = [