267 lines
9 KiB
Python
267 lines
9 KiB
Python
import uuid
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
|
|
import logfire
|
|
from pydantic import TypeAdapter
|
|
from pydantic_ai.messages import ModelMessage, ModelResponse, ToolCallPart
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import Config
|
|
|
|
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,
|
|
TextPart,
|
|
)
|
|
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()
|
|
|
|
ModelMessagesTypeAdapter = TypeAdapter(list[ModelMessage])
|
|
|
|
|
|
def a2a_to_pydantic_messages(a2a_messages: list[Message]) -> list[ModelMessage]:
|
|
"""Convert A2A messages to pydantic-ai ModelMessage format.
|
|
|
|
Args:
|
|
a2a_messages: List of A2A Message objects
|
|
|
|
Returns:
|
|
List of pydantic-ai ModelMessage objects suitable for agent.run()
|
|
"""
|
|
pydantic_messages = []
|
|
|
|
for msg in a2a_messages:
|
|
role = msg.get("role", "user")
|
|
parts = msg.get("parts", [])
|
|
|
|
# Extract text content from all text parts
|
|
text_content = " ".join(
|
|
part.get("text", "") for part in parts if part.get("kind") == "text"
|
|
)
|
|
|
|
if not text_content:
|
|
continue
|
|
|
|
# Build message dict with proper part_kind discriminators
|
|
if role == "user":
|
|
pydantic_messages.append(
|
|
{
|
|
"parts": [{"content": text_content, "part_kind": "user-prompt"}],
|
|
"kind": "request",
|
|
}
|
|
)
|
|
elif role == "agent":
|
|
# Agent responses become ModelResponse with TextPart
|
|
pydantic_messages.append(
|
|
{
|
|
"parts": [{"content": text_content, "part_kind": "text"}],
|
|
"kind": "response",
|
|
"model_name": "unknown",
|
|
}
|
|
)
|
|
|
|
# Validate and convert to proper ModelMessage objects
|
|
if pydantic_messages:
|
|
return ModelMessagesTypeAdapter.validate_python(pydantic_messages)
|
|
|
|
return []
|
|
|
|
|
|
def create_qa_a2a_app(
|
|
db_path: Path,
|
|
deep: bool = False,
|
|
):
|
|
"""Create an A2A app for the QA agent.
|
|
|
|
Args:
|
|
db_path: Path to the LanceDB database
|
|
deep: Use deep multi-agent QA for complex questions
|
|
|
|
Returns:
|
|
A FastA2A ASGI application
|
|
"""
|
|
if deep:
|
|
raise NotImplementedError("Deep QA agent not yet implemented for A2A")
|
|
|
|
from haiku.rag.qa.agent import Dependencies, QuestionAnswerAgent
|
|
|
|
# Create the agent (client will be provided per-task in custom worker)
|
|
temp_client = HaikuRAG(db_path)
|
|
qa_agent = QuestionAnswerAgent(
|
|
client=temp_client,
|
|
provider=Config.QA_PROVIDER,
|
|
model=Config.QA_MODEL,
|
|
)
|
|
|
|
# Create custom worker using base Worker class
|
|
storage = InMemoryStorage()
|
|
broker = InMemoryBroker()
|
|
|
|
class QAWorker(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")
|
|
|
|
# Load full conversation context from previous tasks
|
|
context = await self.storage.load_context(task["context_id"]) or []
|
|
current_task_history = task.get("history", [])
|
|
|
|
# Extract the user's question from the latest message
|
|
user_messages = [
|
|
msg for msg in current_task_history if msg["role"] == "user"
|
|
]
|
|
if not user_messages:
|
|
await self.storage.update_task(task["id"], state="failed")
|
|
return
|
|
|
|
last_user_msg = user_messages[-1]
|
|
question = ""
|
|
for part in last_user_msg.get("parts", []):
|
|
if part.get("kind") == "text":
|
|
question = part.get("text", "")
|
|
break
|
|
|
|
try:
|
|
# Create fresh client for this task and run QA agent
|
|
async with HaikuRAG(db_path) as client:
|
|
deps = Dependencies(client=client)
|
|
|
|
# Convert conversation history to pydantic-ai format
|
|
message_history = a2a_to_pydantic_messages(context)
|
|
|
|
# Run agent with full conversation history
|
|
result = await qa_agent._agent.run(
|
|
question, deps=deps, message_history=message_history
|
|
)
|
|
|
|
# Build response message
|
|
response_message = Message(
|
|
role="agent",
|
|
parts=[TextPart(kind="text", text=str(result.output))],
|
|
kind="message",
|
|
message_id=str(uuid.uuid4()),
|
|
)
|
|
|
|
# Store complete agent state (all messages including tool calls)
|
|
# Add both the user question and agent response to context
|
|
context.extend(current_task_history)
|
|
context.append(response_message)
|
|
await self.storage.update_context(task["context_id"], context)
|
|
|
|
# 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:
|
|
await self.storage.update_task(task["id"], state="failed")
|
|
raise
|
|
|
|
async def cancel_task(self, params: TaskIdParams) -> None:
|
|
pass
|
|
|
|
def build_message_history(self, history: list[Message]) -> list[Message]:
|
|
return history
|
|
|
|
def build_artifacts(self, result) -> list[Artifact]:
|
|
"""Build rich artifacts from agent result including search details."""
|
|
artifacts: list[Artifact] = []
|
|
|
|
# Main answer artifact
|
|
artifacts.append(
|
|
Artifact(
|
|
artifact_id=str(uuid.uuid4()),
|
|
name="answer",
|
|
parts=[TextPart(kind="text", text=str(result.output))],
|
|
)
|
|
)
|
|
|
|
# Extract search tool calls and results from message history
|
|
search_results = []
|
|
for msg in result.all_messages():
|
|
if isinstance(msg, ModelResponse):
|
|
for part in msg.parts:
|
|
if isinstance(part, ToolCallPart):
|
|
if part.tool_name == "search_documents":
|
|
search_results.append(
|
|
{
|
|
"tool_call": part.tool_name,
|
|
"args": part.args,
|
|
}
|
|
)
|
|
|
|
# Create search results artifact if we found any searches
|
|
if search_results:
|
|
artifacts.append(
|
|
Artifact(
|
|
artifact_id=str(uuid.uuid4()),
|
|
name="search_activity",
|
|
parts=[
|
|
DataPart(
|
|
kind="data",
|
|
data={
|
|
"searches": search_results,
|
|
"count": len(search_results),
|
|
},
|
|
metadata={"type": "search_history"},
|
|
)
|
|
],
|
|
)
|
|
)
|
|
|
|
return artifacts
|
|
|
|
worker = QAWorker(storage=storage, broker=broker)
|
|
|
|
# Create FastA2A app with custom worker lifecycle
|
|
@asynccontextmanager
|
|
async def lifespan(app):
|
|
async with app.task_manager:
|
|
async with worker.run():
|
|
yield
|
|
|
|
return FastA2A(
|
|
storage=storage,
|
|
broker=broker,
|
|
name="haiku-rag-qa",
|
|
description="Question answering agent powered by haiku.rag RAG system",
|
|
lifespan=lifespan,
|
|
)
|
|
|
|
|
|
def create_research_a2a_app(db_path: Path):
|
|
"""Create an A2A app for the research agent.
|
|
|
|
Args:
|
|
db_path: Path to the LanceDB database
|
|
|
|
Returns:
|
|
A FastA2A ASGI application
|
|
"""
|
|
raise NotImplementedError("Research agent not yet implemented for A2A")
|