diff --git a/src/haiku/rag/a2a.py b/src/haiku/rag/a2a.py index e75781d3..d0a3437c 100644 --- a/src/haiku/rag/a2a.py +++ b/src/haiku/rag/a2a.py @@ -3,6 +3,8 @@ 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 @@ -12,6 +14,7 @@ try: from fasta2a.broker import InMemoryBroker # type: ignore from fasta2a.schema import ( # type: ignore Artifact, + DataPart, Message, TaskIdParams, TaskSendParams, @@ -27,6 +30,56 @@ except ImportError as 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, @@ -71,13 +124,13 @@ def create_qa_a2a_app( await self.storage.update_task(task["id"], state="working") - # Load context and build simple message for agent + # Load full conversation context from previous tasks context = await self.storage.load_context(task["context_id"]) or [] - context.extend(task.get("history", [])) + current_task_history = task.get("history", []) # Extract the user's question from the latest message user_messages = [ - msg for msg in task.get("history", []) if msg["role"] == "user" + msg for msg in current_task_history if msg["role"] == "user" ] if not user_messages: await self.storage.update_task(task["id"], state="failed") @@ -94,7 +147,14 @@ def create_qa_a2a_app( # Create fresh client for this task and run QA agent async with HaikuRAG(db_path) as client: deps = Dependencies(client=client) - result = await qa_agent._agent.run(question, deps=deps) + + # 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( @@ -104,12 +164,14 @@ def create_qa_a2a_app( message_id=str(uuid.uuid4()), ) - # Update context with new message + # 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 artifacts (optional) - artifacts = self.build_artifacts(result.output) + # Build rich artifacts with search results and answer + artifacts = self.build_artifacts(result) await self.storage.update_task( task["id"], @@ -127,15 +189,53 @@ def create_qa_a2a_app( def build_message_history(self, history: list[Message]) -> list[Message]: return history - def build_artifacts(self, result: str) -> list[Artifact]: - # Simple artifact with the result text - return [ + 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="result", - parts=[TextPart(kind="text", text=result)], + 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)