diff --git a/src/haiku/rag/a2a.py b/src/haiku/rag/a2a.py deleted file mode 100644 index 49cefadb..00000000 --- a/src/haiku/rag/a2a.py +++ /dev/null @@ -1,664 +0,0 @@ -import logging -import uuid -from collections import OrderedDict -from contextlib import asynccontextmanager -from pathlib import Path - -import logfire -from pydantic import BaseModel, Field, TypeAdapter -from pydantic_ai import Agent, RunContext -from pydantic_ai.messages import ModelMessage -from pydantic_core import to_jsonable_python - -from haiku.rag.client import HaikuRAG -from haiku.rag.config import Config -from haiku.rag.graph.common import get_model -from haiku.rag.qa.deep.dependencies import DeepQAContext -from haiku.rag.qa.deep.graph import build_deep_qa_graph -from haiku.rag.qa.deep.nodes import DeepQAPlanNode -from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState - -logger = logging.getLogger(__name__) - -try: - from fasta2a import FastA2A, Worker # type: ignore - from fasta2a.broker import InMemoryBroker # type: ignore - from fasta2a.schema import ( # type: ignore - Artifact, - DataPart, - Message, - Skill, - TaskIdParams, - TaskSendParams, - TaskState, - TextPart, - ) - from fasta2a.storage import InMemoryStorage, Storage # type: ignore -except ImportError as e: - raise ImportError( - "A2A support requires the 'a2a' extra. " - "Install with: uv pip install 'haiku.rag[a2a]'" - ) from e - -logfire.configure(send_to_logfire="if-token-present", service_name="a2a") -logfire.instrument_pydantic_ai() - -ModelMessagesTypeAdapter = TypeAdapter(list[ModelMessage]) - - -class SearchResult(BaseModel): - """Search result with both title and URI for A2A agent.""" - - content: str = Field(description="The document text content") - score: float = Field(description="Relevance score (higher is more relevant)") - document_title: str | None = Field( - description="Human-readable document title", default=None - ) - document_uri: str = Field(description="Document URI/path for get_full_document") - - -class AgentDependencies(BaseModel): - """Dependencies for the A2A conversational agent.""" - - model_config = {"arbitrary_types_allowed": True} - client: HaikuRAG - - -A2A_SYSTEM_PROMPT = """You are Haiku.rag, an AI assistant that helps users find information from a document knowledge base. - -IMPORTANT: You are NOT any person mentioned in the documents. You retrieve and present information about them. - -Tools available: -- search_documents: Query for relevant text chunks -- get_full_document: Get complete document content by document_uri -- list_documents: Show available documents - -Your process: -1. Search phase: For straightforward questions use one search, for complex questions search multiple times with different queries -2. Synthesis phase: Combine the search results into a comprehensive answer -3. When user requests full document: use get_full_document with the exact document_uri from Sources - -Critical rules: -- ONLY answer based on information found via search_documents -- NEVER fabricate or assume information -- If not found, say: "I cannot find information about this in the knowledge base." -- For follow-ups, understand context (pronouns like "he", "it") but always search for facts -- ALWAYS include citations at the end showing document URIs used -- Be concise and direct - -Citation Format: -After your answer, include a "Sources:" section listing documents from search results. -Show both title and URI if available, otherwise just the URI. -Format: "Sources:\n- [document_title] ([document_uri])" or "Sources:\n- [document_uri]" - -Example: -[Your answer here] - -Sources: -- Python Documentation (/guides/python.md) -- /guides/python-basics.md - -Note: When using get_full_document, always use document_uri (not document_title). -""" - - -def load_message_history(context: list[Message]) -> list[ModelMessage]: - """Load pydantic-ai message history from A2A context. - - The context stores serialized pydantic-ai message history directly, - which we deserialize and return. - - Args: - context: A2A context messages - - Returns: - List of pydantic-ai ModelMessage objects - """ - if not context: - return [] - - # Context should contain a single "state" message with full history - for msg in context: - parts = msg.get("parts", []) - for part in parts: - if part.get("kind") == "data": - metadata = part.get("metadata", {}) - if metadata.get("type") == "conversation_state": - stored_history = part.get("data", {}).get("message_history", []) - if stored_history: - return ModelMessagesTypeAdapter.validate_python(stored_history) - - return [] - - -def save_message_history(message_history: list[ModelMessage]) -> Message: - """Save pydantic-ai message history to A2A context format. - - Args: - message_history: Full pydantic-ai message history - - Returns: - A2A Message containing the serialized state (stored as agent role) - """ - serialized = to_jsonable_python(message_history) - return Message( - role="agent", - parts=[ - DataPart( - kind="data", - data={"message_history": serialized}, - metadata={"type": "conversation_state"}, - ) - ], - kind="message", - message_id=str(uuid.uuid4()), - ) - - -class LRUMemoryStorage(Storage[list["Message"]]): # type: ignore - """Storage wrapper with LRU eviction for contexts. - - Enforces a maximum context limit using LRU (Least Recently Used) eviction. - """ - - def __init__(self, storage: InMemoryStorage, max_contexts: int): - self.storage = storage - self.max_contexts = max_contexts - # Track context access order (LRU cache) - self.context_order: OrderedDict[str, None] = OrderedDict() - - async def load_context(self, context_id: str) -> list["Message"] | None: - """Load context and update access order.""" - result = await self.storage.load_context(context_id) - if result is not None: - # Move to end (most recently used) - self.context_order.pop(context_id, None) - self.context_order[context_id] = None - return result - - async def update_context(self, context_id: str, context: list["Message"]) -> None: - """Update context and enforce LRU limit.""" - await self.storage.update_context(context_id, context) - # Move to end (most recently used) - self.context_order.pop(context_id, None) - self.context_order[context_id] = None - - # Enforce max contexts limit (LRU eviction) - while len(self.context_order) > self.max_contexts: - # Remove oldest (first item in OrderedDict) - oldest_context_id = next(iter(self.context_order)) - self.context_order.pop(oldest_context_id) - logger.debug( - f"Evicted context {oldest_context_id} (LRU, limit={self.max_contexts})" - ) - - async def load_task(self, task_id: str, history_length: int | None = None): - """Delegate to underlying storage.""" - return await self.storage.load_task(task_id, history_length) - - async def update_task( - self, - task_id: str, - state: TaskState, - new_artifacts: list["Artifact"] | None = None, - new_messages: list["Message"] | None = None, - ): - """Delegate to underlying storage.""" - return await self.storage.update_task( - task_id, state, new_artifacts, new_messages - ) - - async def submit_task(self, context_id: str, message: "Message"): - """Delegate to underlying storage.""" - return await self.storage.submit_task(context_id, message) - - -def get_agent_skills() -> list[Skill]: - """Define the skills exposed by the haiku.rag A2A agent. - - Returns: - List of skills describing the agent's capabilities - """ - return [ - Skill( - id="document-qa", - name="Document Question Answering", - description="Answer questions based on a knowledge base of documents using semantic search and retrieval", - tags=["question-answering", "search", "knowledge-base", "rag"], - input_modes=["application/json"], - output_modes=["application/json"], - examples=[ - "What does the documentation say about authentication?", - "Find information about Python best practices", - "Show me the full API documentation", - ], - ), - Skill( - id="deep-qa", - name="Deep Question Answering", - description="Multi-step question decomposition and research for complex queries (can take a long time)", - tags=["question-answering", "research", "multi-agent", "complex-queries"], - input_modes=["application/json"], - output_modes=["application/json"], - examples=[ - "What are the architectural patterns used in haiku.rag and how do they compare?", - "Analyze the trade-offs between the simple QA and research agents", - "What are all the configuration options and their effects?", - ], - ), - ] - - -def extract_skill_preference(task_history: list[Message]) -> str: - """Extract skill preference from task history metadata. - - Args: - task_history: Task history messages - - Returns: - Skill ID if found in metadata, otherwise "document-qa" (default) - """ - for msg in task_history: - if msg.get("role") == "user": - for part in msg.get("parts", []): - if part.get("kind") == "data": - metadata = part.get("metadata", {}) - if metadata.get("type") == "skill_preference": - skill = part.get("data", {}).get("skill") - if skill: - return skill - return "document-qa" - - -def extract_question_from_task(task_history: list[Message]) -> str | None: - """Extract the user's question from task history. - - Args: - task_history: Task history messages - - Returns: - The question text if found, None otherwise - """ - for msg in task_history: - if msg.get("role") == "user": - for part in msg.get("parts", []): - if part.get("kind") == "text": - text = part.get("text", "").strip() - if text: - return text - return None - - -def create_a2a_app(db_path: Path): - """Create an A2A app for the conversational QA agent. - - Args: - db_path: Path to the LanceDB database - - Returns: - A FastA2A ASGI application - """ - base_storage = InMemoryStorage() - storage = LRUMemoryStorage( - storage=base_storage, max_contexts=Config.A2A_MAX_CONTEXTS - ) - broker = InMemoryBroker() - - # Create the agent with native search tool - model = get_model(Config.QA_PROVIDER, Config.QA_MODEL) - agent = Agent( - model=model, - deps_type=AgentDependencies, - system_prompt=A2A_SYSTEM_PROMPT, - retries=3, - ) - - @agent.tool - async def search_documents( - ctx: RunContext[AgentDependencies], - query: str, - limit: int = 3, - ) -> list[SearchResult]: - """Search the knowledge base for relevant documents. - - Returns chunks of text with their relevance scores and document URIs. - Use get_full_document if you need to see the complete document content. - """ - search_results = await ctx.deps.client.search(query, limit=limit) - expanded_results = await ctx.deps.client.expand_context(search_results) - - return [ - SearchResult( - content=chunk.content, - score=score, - document_title=chunk.document_title, - document_uri=(chunk.document_uri or ""), - ) - for chunk, score in expanded_results - ] - - @agent.tool - async def get_full_document( - ctx: RunContext[AgentDependencies], - document_uri: str, - ) -> str: - """Retrieve the complete content of a document by its URI. - - Use this when you need more context than what's in a search result chunk. - The document_uri comes from search_documents results. - """ - document = await ctx.deps.client.get_document_by_uri(document_uri) - if document is None: - return f"Document not found: {document_uri}" - - return document.content - - @agent.tool - async def list_documents( - ctx: RunContext[AgentDependencies], - limit: int = 10, - ) -> list[str]: - """List documents in the knowledge base. - - Returns document URIs/titles. Use this to help users discover what's available. - """ - documents = await ctx.deps.client.list_documents(limit=limit) - return [doc.title or doc.uri or f"Document {doc.id}" for doc in documents] - - class ConversationalWorker(Worker[list[Message]]): - async def evaluate_answer_adequacy(self, question: str, answer: str) -> bool: - """Use LLM to evaluate if answer adequately addresses the question. - - Args: - question: The original question - answer: The answer to evaluate - - Returns: - True if answer is adequate, False if more research needed - """ - from pydantic import BaseModel, Field - - class AnswerEvaluation(BaseModel): - is_adequate: bool = Field( - description="True if the answer adequately addresses the question, False if more research is needed" - ) - reasoning: str = Field( - description="Brief explanation of the evaluation" - ) - - evaluation_agent = Agent( - model=get_model(Config.QA_PROVIDER, Config.QA_MODEL), - output_type=AnswerEvaluation, - system_prompt="""You evaluate whether an answer adequately addresses a question. - -Consider: -- Completeness: Does it answer all parts of the question? -- Specificity: Is it specific enough or too vague? -- Relevance: Does it directly address what was asked? -- Depth: For complex questions, does it provide sufficient depth? - -Return is_adequate=True if the answer satisfactorily addresses the question. -Return is_adequate=False if the answer is incomplete, too vague, or requires deeper research.""", - retries=1, - ) - - prompt = f"""Question: {question} - -Answer: {answer} - -Does this answer adequately address the question?""" - - result = await evaluation_agent.run(prompt) - logger.info( - f"Answer evaluation: is_adequate={result.output.is_adequate}, reasoning={result.output.reasoning}" - ) - return result.output.is_adequate - - async def run_task(self, params: TaskSendParams) -> None: - task = await self.storage.load_task(params["id"]) - if task is None: - raise ValueError(f"Task {params['id']} not found") - - if task["status"]["state"] != "submitted": - raise ValueError( - f"Task {params['id']} already processed: {task['status']['state']}" - ) - - await self.storage.update_task(task["id"], state="working") - - # Extract skill preference and question - task_history = task.get("history", []) - skill = extract_skill_preference(task_history) - question = extract_question_from_task(task_history) - - if not question: - await self.storage.update_task(task["id"], state="failed") - return - - logger.info(f"Task {task['id']} requested skill: {skill}") - - try: - async with HaikuRAG(db_path) as client: - if skill == "deep-qa": - # Explicitly requested deep QA - logger.info(f"Task {task['id']}: Running deep QA (explicit)") - deep_result, deep_state = await self.run_deep_qa( - client, question - ) - - response_message = Message( - role="agent", - parts=[TextPart(kind="text", text=deep_result.answer)], - kind="message", - message_id=str(uuid.uuid4()), - ) - - artifacts = self.build_deep_qa_artifacts( - deep_result, deep_state - ) - - await self.storage.update_task( - task["id"], - state="completed", - new_messages=[response_message], - new_artifacts=artifacts, - ) - else: - # Try simple QA first (default behavior or explicit document-qa) - logger.info(f"Task {task['id']}: Trying simple QA first") - - context = ( - await self.storage.load_context(task["context_id"]) or [] - ) - message_history = load_message_history(context) - - deps = AgentDependencies(client=client) - - result = await agent.run( - question, deps=deps, message_history=message_history - ) - - answer = str(result.output) - - # Evaluate answer adequacy - is_adequate = await self.evaluate_answer_adequacy( - question, answer - ) - - if not is_adequate: - # Escalate to deep QA - logger.info( - f"Task {task['id']}: Answer inadequate, escalating to deep QA" - ) - deep_result, deep_state = await self.run_deep_qa( - client, question - ) - - response_message = Message( - role="agent", - parts=[TextPart(kind="text", text=deep_result.answer)], - kind="message", - message_id=str(uuid.uuid4()), - ) - - artifacts = self.build_deep_qa_artifacts( - deep_result, deep_state - ) - - await self.storage.update_task( - task["id"], - state="completed", - new_messages=[response_message], - new_artifacts=artifacts, - ) - else: - # Simple QA answer is adequate - logger.info( - f"Task {task['id']}: Simple QA answer is adequate" - ) - - response_message = Message( - role="agent", - parts=[TextPart(kind="text", text=answer)], - kind="message", - message_id=str(uuid.uuid4()), - ) - - # Update context with complete conversation state - updated_history = message_history + result.new_messages() - state_message = save_message_history(updated_history) - - await self.storage.update_context( - task["context_id"], [state_message] - ) - - artifacts = self.build_artifacts(result) - - await self.storage.update_task( - task["id"], - state="completed", - new_messages=[response_message], - new_artifacts=artifacts, - ) - except Exception as e: - logger.error( - "Task execution failed: task_id=%s, question=%s, error=%s", - task["id"], - question, - str(e), - exc_info=True, - ) - await self.storage.update_task(task["id"], state="failed") - raise - - async def run_deep_qa(self, client: HaikuRAG, question: str): - """Run deep QA graph for complex questions. - - Args: - client: HaikuRAG client - question: User's question - - Returns: - Tuple of (DeepQAAnswer, DeepQAState) with answer and state - """ - graph = build_deep_qa_graph() - context = DeepQAContext(original_question=question, use_citations=False) - state = DeepQAState(context=context) - deps = DeepQADeps(client=client, console=None) - start_node = DeepQAPlanNode( - provider=Config.QA_PROVIDER, model=Config.QA_MODEL - ) - - result = await graph.run(start_node=start_node, state=state, deps=deps) - return result.output, state - - async def cancel_task(self, params: TaskIdParams) -> None: - """Cancel a task - not implemented for this worker.""" - pass - - def build_message_history(self, history: list[Message]) -> list[Message]: - """Required by Worker interface but unused - history stored in context.""" - return history - - def build_artifacts(self, result) -> list[Artifact]: - """Build artifacts from agent result. - - Note: Full conversation history (including tool calls) is stored in - context, so we only create a simple answer artifact here. - """ - return [ - Artifact( - artifact_id=str(uuid.uuid4()), - name="answer", - parts=[TextPart(kind="text", text=str(result.output))], - ) - ] - - def build_deep_qa_artifacts(self, result, state: DeepQAState) -> list[Artifact]: - """Build rich artifacts from deep QA result. - - Args: - result: DeepQAAnswer with final answer - state: DeepQAState with research process details - - Returns: - List of artifacts including answer and research breakdown - """ - artifacts = [ - # Final answer artifact - Artifact( - artifact_id=str(uuid.uuid4()), - name="answer", - parts=[TextPart(kind="text", text=result.answer)], - ) - ] - - # Add research process artifact with sub-questions and answers - if state.context.qa_responses: - research_data = { - "original_question": state.context.original_question, - "iterations": state.iterations, - "sub_questions_answered": [ - { - "question": qa.query, - "answer": qa.answer, - "sources": qa.sources, - } - for qa in state.context.qa_responses - ], - } - - artifacts.append( - Artifact( - artifact_id=str(uuid.uuid4()), - name="research_process", - parts=[ - DataPart( - kind="data", - data=research_data, - metadata={"type": "deep_qa_research"}, - ) - ], - ) - ) - - return artifacts - - worker = ConversationalWorker(storage=storage, broker=broker) - - # Create FastA2A app with custom worker lifecycle - @asynccontextmanager - async def lifespan(app): - logger.info(f"Started A2A server (max contexts: {Config.A2A_MAX_CONTEXTS})") - async with app.task_manager: - async with worker.run(): - yield - - return FastA2A( - storage=storage, - broker=broker, - name="haiku-rag", - description="Conversational question answering agent powered by haiku.rag RAG system", - skills=get_agent_skills(), - lifespan=lifespan, - ) diff --git a/src/haiku/rag/a2a/__init__.py b/src/haiku/rag/a2a/__init__.py new file mode 100644 index 00000000..7266cada --- /dev/null +++ b/src/haiku/rag/a2a/__init__.py @@ -0,0 +1,148 @@ +"""A2A (Agent-to-Agent) server integration for haiku.rag.""" + +import logging +from contextlib import asynccontextmanager +from pathlib import Path + +import logfire +from pydantic_ai import Agent, RunContext + +from haiku.rag.config import Config +from haiku.rag.graph.common import get_model + +from .context import load_message_history, save_message_history +from .models import AgentDependencies, SearchResult +from .prompts import A2A_SYSTEM_PROMPT +from .skills import ( + extract_question_from_task, + extract_skill_preference, + get_agent_skills, +) +from .storage import LRUMemoryStorage +from .worker import ConversationalWorker + +try: + from fasta2a import FastA2A # type: ignore + from fasta2a.broker import InMemoryBroker # type: ignore + from fasta2a.storage import InMemoryStorage # type: ignore +except ImportError as e: + raise ImportError( + "A2A support requires the 'a2a' extra. " + "Install with: uv pip install 'haiku.rag[a2a]'" + ) from e + +logfire.configure(send_to_logfire="if-token-present", service_name="a2a") +logfire.instrument_pydantic_ai() + +logger = logging.getLogger(__name__) + +__all__ = [ + "create_a2a_app", + "load_message_history", + "save_message_history", + "extract_question_from_task", + "extract_skill_preference", + "get_agent_skills", + "LRUMemoryStorage", +] + + +def create_a2a_app(db_path: Path): + """Create an A2A app for the conversational QA agent. + + Args: + db_path: Path to the LanceDB database + + Returns: + A FastA2A ASGI application + """ + base_storage = InMemoryStorage() + storage = LRUMemoryStorage( + storage=base_storage, max_contexts=Config.A2A_MAX_CONTEXTS + ) + broker = InMemoryBroker() + + # Create the agent with native search tool + model = get_model(Config.QA_PROVIDER, Config.QA_MODEL) + agent = Agent( + model=model, + deps_type=AgentDependencies, + system_prompt=A2A_SYSTEM_PROMPT, + retries=3, + ) + + @agent.tool + async def search_documents( + ctx: RunContext[AgentDependencies], + query: str, + limit: int = 3, + ) -> list[SearchResult]: + """Search the knowledge base for relevant documents. + + Returns chunks of text with their relevance scores and document URIs. + Use get_full_document if you need to see the complete document content. + """ + search_results = await ctx.deps.client.search(query, limit=limit) + expanded_results = await ctx.deps.client.expand_context(search_results) + + return [ + SearchResult( + content=chunk.content, + score=score, + document_title=chunk.document_title, + document_uri=(chunk.document_uri or ""), + ) + for chunk, score in expanded_results + ] + + @agent.tool + async def get_full_document( + ctx: RunContext[AgentDependencies], + document_uri: str, + ) -> str: + """Retrieve the complete content of a document by its URI. + + Use this when you need more context than what's in a search result chunk. + The document_uri comes from search_documents results. + """ + document = await ctx.deps.client.get_document_by_uri(document_uri) + if document is None: + return f"Document not found: {document_uri}" + + return document.content + + @agent.tool + async def list_documents( + ctx: RunContext[AgentDependencies], + limit: int = 10, + ) -> list[str]: + """List documents in the knowledge base. + + Returns document URIs/titles. Use this to help users discover what's available. + """ + documents = await ctx.deps.client.list_documents(limit=limit) + return [doc.title or doc.uri or f"Document {doc.id}" for doc in documents] + + worker = ConversationalWorker( + storage=storage, + broker=broker, + db_path=db_path, + agent=agent, # type: ignore + ) + + # Create FastA2A app with custom worker lifecycle + @asynccontextmanager + async def lifespan(app): + logger.info(f"Started A2A server (max contexts: {Config.A2A_MAX_CONTEXTS})") + async with app.task_manager: + async with worker.run(): + yield + + return FastA2A( + storage=storage, + broker=broker, + name="haiku-rag", + description="Conversational question answering agent powered by haiku.rag RAG system", + skills=get_agent_skills(), + lifespan=lifespan, + ) diff --git a/src/haiku/rag/a2a/context.py b/src/haiku/rag/a2a/context.py new file mode 100644 index 00000000..34f06823 --- /dev/null +++ b/src/haiku/rag/a2a/context.py @@ -0,0 +1,70 @@ +"""Context management for A2A conversations.""" + +import uuid + +from pydantic import TypeAdapter +from pydantic_ai.messages import ModelMessage +from pydantic_core import to_jsonable_python + +try: + from fasta2a.schema import DataPart, Message # type: ignore +except ImportError as e: + raise ImportError( + "A2A support requires the 'a2a' extra. " + "Install with: uv pip install 'haiku.rag[a2a]'" + ) from e + +ModelMessagesTypeAdapter = TypeAdapter(list[ModelMessage]) + + +def load_message_history(context: list[Message]) -> list[ModelMessage]: + """Load pydantic-ai message history from A2A context. + + The context stores serialized pydantic-ai message history directly, + which we deserialize and return. + + Args: + context: A2A context messages + + Returns: + List of pydantic-ai ModelMessage objects + """ + if not context: + return [] + + # Context should contain a single "state" message with full history + for msg in context: + parts = msg.get("parts", []) + for part in parts: + if part.get("kind") == "data": + metadata = part.get("metadata", {}) + if metadata.get("type") == "conversation_state": + stored_history = part.get("data", {}).get("message_history", []) + if stored_history: + return ModelMessagesTypeAdapter.validate_python(stored_history) + + return [] + + +def save_message_history(message_history: list[ModelMessage]) -> Message: + """Save pydantic-ai message history to A2A context format. + + Args: + message_history: Full pydantic-ai message history + + Returns: + A2A Message containing the serialized state (stored as agent role) + """ + serialized = to_jsonable_python(message_history) + return Message( + role="agent", + parts=[ + DataPart( + kind="data", + data={"message_history": serialized}, + metadata={"type": "conversation_state"}, + ) + ], + kind="message", + message_id=str(uuid.uuid4()), + ) diff --git a/src/haiku/rag/a2a/models.py b/src/haiku/rag/a2a/models.py new file mode 100644 index 00000000..8f4a8830 --- /dev/null +++ b/src/haiku/rag/a2a/models.py @@ -0,0 +1,23 @@ +"""Data models for A2A integration.""" + +from pydantic import BaseModel, Field + +from haiku.rag.client import HaikuRAG + + +class SearchResult(BaseModel): + """Search result with both title and URI for A2A agent.""" + + content: str = Field(description="The document text content") + score: float = Field(description="Relevance score (higher is more relevant)") + document_title: str | None = Field( + description="Human-readable document title", default=None + ) + document_uri: str = Field(description="Document URI/path for get_full_document") + + +class AgentDependencies(BaseModel): + """Dependencies for the A2A conversational agent.""" + + model_config = {"arbitrary_types_allowed": True} + client: HaikuRAG diff --git a/src/haiku/rag/a2a/prompts.py b/src/haiku/rag/a2a/prompts.py new file mode 100644 index 00000000..79f7386e --- /dev/null +++ b/src/haiku/rag/a2a/prompts.py @@ -0,0 +1,49 @@ +"""Prompts for A2A agents.""" + +A2A_SYSTEM_PROMPT = """You are Haiku.rag, an AI assistant that helps users find information from a document knowledge base. + +IMPORTANT: You are NOT any person mentioned in the documents. You retrieve and present information about them. + +Tools available: +- search_documents: Query for relevant text chunks +- get_full_document: Get complete document content by document_uri +- list_documents: Show available documents + +Your process: +1. Search phase: For straightforward questions use one search, for complex questions search multiple times with different queries +2. Synthesis phase: Combine the search results into a comprehensive answer +3. When user requests full document: use get_full_document with the exact document_uri from Sources + +Critical rules: +- ONLY answer based on information found via search_documents +- NEVER fabricate or assume information +- If not found, say: "I cannot find information about this in the knowledge base." +- For follow-ups, understand context (pronouns like "he", "it") but always search for facts +- ALWAYS include citations at the end showing document URIs used +- Be concise and direct + +Citation Format: +After your answer, include a "Sources:" section listing documents from search results. +Show both title and URI if available, otherwise just the URI. +Format: "Sources:\n- [document_title] ([document_uri])" or "Sources:\n- [document_uri]" + +Example: +[Your answer here] + +Sources: +- Python Documentation (/guides/python.md) +- /guides/python-basics.md + +Note: When using get_full_document, always use document_uri (not document_title). +""" + +ANSWER_EVALUATION_PROMPT = """You evaluate whether an answer adequately addresses a question. + +Consider: +- Completeness: Does it answer all parts of the question? +- Specificity: Is it specific enough or too vague? +- Relevance: Does it directly address what was asked? +- Depth: For complex questions, does it provide sufficient depth? + +Return is_adequate=True if the answer satisfactorily addresses the question. +Return is_adequate=False if the answer is incomplete, too vague, or requires deeper research.""" diff --git a/src/haiku/rag/a2a/skills.py b/src/haiku/rag/a2a/skills.py new file mode 100644 index 00000000..a181d74f --- /dev/null +++ b/src/haiku/rag/a2a/skills.py @@ -0,0 +1,85 @@ +"""A2A skill definitions and utilities.""" + +try: + from fasta2a.schema import Message, Skill # type: ignore +except ImportError as e: + raise ImportError( + "A2A support requires the 'a2a' extra. " + "Install with: uv pip install 'haiku.rag[a2a]'" + ) from e + + +def get_agent_skills() -> list[Skill]: + """Define the skills exposed by the haiku.rag A2A agent. + + Returns: + List of skills describing the agent's capabilities + """ + return [ + Skill( + id="document-qa", + name="Document Question Answering", + description="Answer questions based on a knowledge base of documents using semantic search and retrieval", + tags=["question-answering", "search", "knowledge-base", "rag"], + input_modes=["application/json"], + output_modes=["application/json"], + examples=[ + "What does the documentation say about authentication?", + "Find information about Python best practices", + "Show me the full API documentation", + ], + ), + Skill( + id="deep-qa", + name="Deep Question Answering", + description="Multi-step question decomposition and research for complex queries (can take a long time)", + tags=["question-answering", "research", "multi-agent", "complex-queries"], + input_modes=["application/json"], + output_modes=["application/json"], + examples=[ + "What are the architectural patterns used in haiku.rag and how do they compare?", + "Analyze the trade-offs between the simple QA and research agents", + "What are all the configuration options and their effects?", + ], + ), + ] + + +def extract_skill_preference(task_history: list[Message]) -> str: + """Extract skill preference from task history metadata. + + Args: + task_history: Task history messages + + Returns: + Skill ID if found in metadata, otherwise "document-qa" (default) + """ + for msg in task_history: + if msg.get("role") == "user": + for part in msg.get("parts", []): + if part.get("kind") == "data": + metadata = part.get("metadata", {}) + if metadata.get("type") == "skill_preference": + skill = part.get("data", {}).get("skill") + if skill: + return skill + return "document-qa" + + +def extract_question_from_task(task_history: list[Message]) -> str | None: + """Extract the user's question from task history. + + Args: + task_history: Task history messages + + Returns: + The question text if found, None otherwise + """ + for msg in task_history: + if msg.get("role") == "user": + for part in msg.get("parts", []): + if part.get("kind") == "text": + text = part.get("text", "").strip() + if text: + return text + return None diff --git a/src/haiku/rag/a2a/storage.py b/src/haiku/rag/a2a/storage.py new file mode 100644 index 00000000..56d776c4 --- /dev/null +++ b/src/haiku/rag/a2a/storage.py @@ -0,0 +1,73 @@ +"""Storage implementations for A2A contexts.""" + +import logging +from collections import OrderedDict + +try: + from fasta2a.schema import Artifact, Message, TaskState # type: ignore + from fasta2a.storage import InMemoryStorage, Storage # type: ignore +except ImportError as e: + raise ImportError( + "A2A support requires the 'a2a' extra. " + "Install with: uv pip install 'haiku.rag[a2a]'" + ) from e + +logger = logging.getLogger(__name__) + + +class LRUMemoryStorage(Storage[list[Message]]): # type: ignore + """Storage wrapper with LRU eviction for contexts. + + Enforces a maximum context limit using LRU (Least Recently Used) eviction. + """ + + def __init__(self, storage: InMemoryStorage, max_contexts: int): + self.storage = storage + self.max_contexts = max_contexts + # Track context access order (LRU cache) + self.context_order: OrderedDict[str, None] = OrderedDict() + + async def load_context(self, context_id: str) -> list[Message] | None: + """Load context and update access order.""" + result = await self.storage.load_context(context_id) + if result is not None: + # Move to end (most recently used) + self.context_order.pop(context_id, None) + self.context_order[context_id] = None + return result + + async def update_context(self, context_id: str, context: list[Message]) -> None: + """Update context and enforce LRU limit.""" + await self.storage.update_context(context_id, context) + # Move to end (most recently used) + self.context_order.pop(context_id, None) + self.context_order[context_id] = None + + # Enforce max contexts limit (LRU eviction) + while len(self.context_order) > self.max_contexts: + # Remove oldest (first item in OrderedDict) + oldest_context_id = next(iter(self.context_order)) + self.context_order.pop(oldest_context_id) + logger.debug( + f"Evicted context {oldest_context_id} (LRU, limit={self.max_contexts})" + ) + + async def load_task(self, task_id: str, history_length: int | None = None): + """Delegate to underlying storage.""" + return await self.storage.load_task(task_id, history_length) + + async def update_task( + self, + task_id: str, + state: TaskState, + new_artifacts: list[Artifact] | None = None, + new_messages: list[Message] | None = None, + ): + """Delegate to underlying storage.""" + return await self.storage.update_task( + task_id, state, new_artifacts, new_messages + ) + + async def submit_task(self, context_id: str, message: Message): + """Delegate to underlying storage.""" + return await self.storage.submit_task(context_id, message) diff --git a/src/haiku/rag/a2a/worker.py b/src/haiku/rag/a2a/worker.py new file mode 100644 index 00000000..36e382cf --- /dev/null +++ b/src/haiku/rag/a2a/worker.py @@ -0,0 +1,310 @@ +"""A2A worker implementation for conversational QA.""" + +import logging +import uuid +from pathlib import Path + +from pydantic import BaseModel, Field +from pydantic_ai import Agent + +from haiku.rag.a2a.context import load_message_history, save_message_history +from haiku.rag.a2a.models import AgentDependencies +from haiku.rag.a2a.skills import extract_question_from_task, extract_skill_preference +from haiku.rag.client import HaikuRAG +from haiku.rag.config import Config +from haiku.rag.graph.common import get_model +from haiku.rag.qa.deep.dependencies import DeepQAContext +from haiku.rag.qa.deep.graph import build_deep_qa_graph +from haiku.rag.qa.deep.nodes import DeepQAPlanNode +from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState + +try: + from fasta2a import Worker # type: ignore + from fasta2a.schema import ( # type: ignore + Artifact, + DataPart, + Message, + TaskIdParams, + TaskSendParams, + TextPart, + ) +except ImportError as e: + raise ImportError( + "A2A support requires the 'a2a' extra. " + "Install with: uv pip install 'haiku.rag[a2a]'" + ) from e + +logger = logging.getLogger(__name__) + + +class ConversationalWorker(Worker[list[Message]]): + """Worker that handles conversational QA tasks.""" + + def __init__( + self, + storage, + broker, + db_path: Path, + agent: "Agent[AgentDependencies, str]", + ): + super().__init__(storage=storage, broker=broker) + self.db_path = db_path + self.agent = agent + + async def evaluate_answer_adequacy(self, question: str, answer: str) -> bool: + """Use LLM to evaluate if answer adequately addresses the question. + + Args: + question: The original question + answer: The answer to evaluate + + Returns: + True if answer is adequate, False if more research needed + """ + + class AnswerEvaluation(BaseModel): + is_adequate: bool = Field( + description="True if the answer adequately addresses the question, False if more research is needed" + ) + reasoning: str = Field(description="Brief explanation of the evaluation") + + from .prompts import ANSWER_EVALUATION_PROMPT + + evaluation_agent = Agent( + model=get_model(Config.QA_PROVIDER, Config.QA_MODEL), + output_type=AnswerEvaluation, + system_prompt=ANSWER_EVALUATION_PROMPT, + retries=1, + ) + + prompt = f"""Question: {question} + +Answer: {answer} + +Does this answer adequately address the question?""" + + result = await evaluation_agent.run(prompt) + logger.info( + f"Answer evaluation: is_adequate={result.output.is_adequate}, reasoning={result.output.reasoning}" + ) + return result.output.is_adequate + + async def run_task(self, params: TaskSendParams) -> None: + task = await self.storage.load_task(params["id"]) + if task is None: + raise ValueError(f"Task {params['id']} not found") + + if task["status"]["state"] != "submitted": + raise ValueError( + f"Task {params['id']} already processed: {task['status']['state']}" + ) + + await self.storage.update_task(task["id"], state="working") + + # Extract skill preference and question + task_history = task.get("history", []) + skill = extract_skill_preference(task_history) + question = extract_question_from_task(task_history) + + if not question: + await self.storage.update_task(task["id"], state="failed") + return + + logger.info(f"Task {task['id']} requested skill: {skill}") + + try: + async with HaikuRAG(self.db_path) as client: + if skill == "deep-qa": + # Explicitly requested deep QA + logger.info(f"Task {task['id']}: Running deep QA (explicit)") + deep_result, deep_state = await self.run_deep_qa(client, question) + + response_message = Message( + role="agent", + parts=[TextPart(kind="text", text=deep_result.answer)], + kind="message", + message_id=str(uuid.uuid4()), + ) + + artifacts = self.build_deep_qa_artifacts(deep_result, deep_state) + + await self.storage.update_task( + task["id"], + state="completed", + new_messages=[response_message], + new_artifacts=artifacts, + ) + else: + # Try simple QA first (default behavior or explicit document-qa) + logger.info(f"Task {task['id']}: Trying simple QA first") + + context = await self.storage.load_context(task["context_id"]) or [] + message_history = load_message_history(context) + + from .models import AgentDependencies + + deps = AgentDependencies(client=client) + + result = await self.agent.run( + question, deps=deps, message_history=message_history + ) + + answer = str(result.output) + + # Evaluate answer adequacy + is_adequate = await self.evaluate_answer_adequacy(question, answer) + + if not is_adequate: + # Escalate to deep QA + logger.info( + f"Task {task['id']}: Answer inadequate, escalating to deep QA" + ) + deep_result, deep_state = await self.run_deep_qa( + client, question + ) + + response_message = Message( + role="agent", + parts=[TextPart(kind="text", text=deep_result.answer)], + kind="message", + message_id=str(uuid.uuid4()), + ) + + artifacts = self.build_deep_qa_artifacts( + deep_result, deep_state + ) + + await self.storage.update_task( + task["id"], + state="completed", + new_messages=[response_message], + new_artifacts=artifacts, + ) + else: + # Simple QA answer is adequate + logger.info(f"Task {task['id']}: Simple QA answer is adequate") + + response_message = Message( + role="agent", + parts=[TextPart(kind="text", text=answer)], + kind="message", + message_id=str(uuid.uuid4()), + ) + + # Update context with complete conversation state + updated_history = message_history + result.new_messages() + state_message = save_message_history(updated_history) + + await self.storage.update_context( + task["context_id"], [state_message] + ) + + artifacts = self.build_artifacts(result) + + await self.storage.update_task( + task["id"], + state="completed", + new_messages=[response_message], + new_artifacts=artifacts, + ) + except Exception as e: + logger.error( + "Task execution failed: task_id=%s, question=%s, error=%s", + task["id"], + question, + str(e), + exc_info=True, + ) + await self.storage.update_task(task["id"], state="failed") + raise + + async def run_deep_qa(self, client: HaikuRAG, question: str): + """Run deep QA graph for complex questions. + + Args: + client: HaikuRAG client + question: User's question + + Returns: + Tuple of (DeepQAAnswer, DeepQAState) with answer and state + """ + graph = build_deep_qa_graph() + context = DeepQAContext(original_question=question, use_citations=False) + state = DeepQAState(context=context) + deps = DeepQADeps(client=client, console=None) + start_node = DeepQAPlanNode(provider=Config.QA_PROVIDER, model=Config.QA_MODEL) + + result = await graph.run(start_node=start_node, state=state, deps=deps) + return result.output, state + + async def cancel_task(self, params: TaskIdParams) -> None: + """Cancel a task - not implemented for this worker.""" + pass + + def build_message_history(self, history: list[Message]) -> list[Message]: + """Required by Worker interface but unused - history stored in context.""" + return history + + def build_artifacts(self, result) -> list[Artifact]: + """Build artifacts from agent result. + + Note: Full conversation history (including tool calls) is stored in + context, so we only create a simple answer artifact here. + """ + return [ + Artifact( + artifact_id=str(uuid.uuid4()), + name="answer", + parts=[TextPart(kind="text", text=str(result.output))], + ) + ] + + def build_deep_qa_artifacts(self, result, state: DeepQAState) -> list[Artifact]: + """Build rich artifacts from deep QA result. + + Args: + result: DeepQAAnswer with final answer + state: DeepQAState with research process details + + Returns: + List of artifacts including answer and research breakdown + """ + artifacts = [ + # Final answer artifact + Artifact( + artifact_id=str(uuid.uuid4()), + name="answer", + parts=[TextPart(kind="text", text=result.answer)], + ) + ] + + # Add research process artifact with sub-questions and answers + if state.context.qa_responses: + research_data = { + "original_question": state.context.original_question, + "iterations": state.iterations, + "sub_questions_answered": [ + { + "question": qa.query, + "answer": qa.answer, + "sources": qa.sources, + } + for qa in state.context.qa_responses + ], + } + + artifacts.append( + Artifact( + artifact_id=str(uuid.uuid4()), + name="research_process", + parts=[ + DataPart( + kind="data", + data=research_data, + metadata={"type": "deep_qa_research"}, + ) + ], + ) + ) + + return artifacts diff --git a/tests/test_a2a.py b/tests/test_a2a.py index 6674e947..d0149478 100644 --- a/tests/test_a2a.py +++ b/tests/test_a2a.py @@ -9,6 +9,7 @@ from haiku.rag.a2a import ( load_message_history, save_message_history, ) +from haiku.rag.a2a.storage import LRUMemoryStorage from haiku.rag.client import HaikuRAG pytest.importorskip("fasta2a") @@ -145,8 +146,6 @@ async def test_lru_memory_storage_lru_eviction(): """Test that LRUMemoryStorage evicts least recently used contexts.""" from fasta2a.storage import InMemoryStorage - from haiku.rag.a2a import LRUMemoryStorage - base_storage = InMemoryStorage() storage = LRUMemoryStorage(storage=base_storage, max_contexts=3) @@ -186,8 +185,6 @@ async def test_lru_memory_storage_access_order(): """Test that accessing contexts updates their order.""" from fasta2a.storage import InMemoryStorage - from haiku.rag.a2a import LRUMemoryStorage - base_storage = InMemoryStorage() storage = LRUMemoryStorage(storage=base_storage, max_contexts=2)