416 lines
15 KiB
Python
416 lines
15 KiB
Python
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
|
|
|
|
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,
|
|
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 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 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 the user's question
|
|
question = extract_question_from_task(task.get("history", []))
|
|
if not question:
|
|
await self.storage.update_task(task["id"], state="failed")
|
|
return
|
|
|
|
try:
|
|
# Load conversation context
|
|
context = await self.storage.load_context(task["context_id"]) or []
|
|
# Load conversation history
|
|
message_history = load_message_history(context)
|
|
|
|
# Create fresh client for this task and run agent
|
|
async with HaikuRAG(db_path) as client:
|
|
deps = AgentDependencies(client=client)
|
|
|
|
# Run agent with full conversation history including tool calls
|
|
result = await agent.run(
|
|
question, deps=deps, message_history=message_history
|
|
)
|
|
|
|
# Build response message for A2A protocol
|
|
response_message = Message(
|
|
role="agent",
|
|
parts=[TextPart(kind="text", text=str(result.output))],
|
|
kind="message",
|
|
message_id=str(uuid.uuid4()),
|
|
)
|
|
|
|
# Update context with complete conversation state
|
|
# Store all messages from this run (includes tool calls & results)
|
|
updated_history = message_history + result.new_messages()
|
|
state_message = save_message_history(updated_history)
|
|
|
|
# Replace old state with new complete state
|
|
await self.storage.update_context(
|
|
task["context_id"], [state_message]
|
|
)
|
|
|
|
# Build rich artifacts with search results and answer
|
|
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 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))],
|
|
)
|
|
]
|
|
|
|
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",
|
|
lifespan=lifespan,
|
|
)
|