From e66c160055065d74fd96378a011784485013b74d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Jul 2025 11:48:19 +0300 Subject: [PATCH] Support for anthropic in Question/Answering --- pyproject.toml | 1 + src/haiku/rag/config.py | 3 + src/haiku/rag/qa/__init__.py | 13 ++++ src/haiku/rag/qa/anthropic.py | 112 ++++++++++++++++++++++++++++++++++ tests/test_qa.py | 36 ++++++++++- uv.lock | 24 +++++++- 6 files changed, 186 insertions(+), 3 deletions(-) create mode 100644 src/haiku/rag/qa/anthropic.py 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 = [