danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Refactor into own module

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-10-13 13:22:27 +03:00
parent 4b2dfcd72f
commit 352637faa6
No known key found for this signature in database
9 changed files with 759 additions and 668 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

664
src/haiku/rag/a2a.py
View file

@ -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,
)

148
src/haiku/rag/a2a/__init__.py Normal file
View file

@ -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,
)

70
src/haiku/rag/a2a/context.py Normal file
View file

@ -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()),
)

23
src/haiku/rag/a2a/models.py Normal file
View file

@ -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

49
src/haiku/rag/a2a/prompts.py Normal file
View file

@ -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."""

85
src/haiku/rag/a2a/skills.py Normal file
View file

@ -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

73
src/haiku/rag/a2a/storage.py Normal file
View file

@ -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)

310
src/haiku/rag/a2a/worker.py Normal file
View file

@ -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

5
tests/test_a2a.py
View file

@ -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)