Conversational a2a

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)