143 lines
4.3 KiB
Python
143 lines
4.3 KiB
Python
"""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, 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",
|
|
"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,
|
|
)
|