From cd2ff36684f214ae6d05f7602e920c5bde5d673f Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 9 Oct 2025 14:40:43 +0300
Subject: [PATCH] Always include URI so that we can get by uri later

---
 src/haiku/rag/a2a.py | 28 +++++++++++++++++++++-------
 1 file changed, 21 insertions(+), 7 deletions(-)

diff --git a/src/haiku/rag/a2a.py b/src/haiku/rag/a2a.py
index 0489d5b8..fff1d7c2 100644
--- a/src/haiku/rag/a2a.py
+++ b/src/haiku/rag/a2a.py
@@ -4,7 +4,7 @@ from contextlib import asynccontextmanager
 from pathlib import Path
 
 import logfire
-from pydantic import BaseModel, TypeAdapter
+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
@@ -12,7 +12,6 @@ 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
-from haiku.rag.qa.agent import SearchResult
 
 logger = logging.getLogger(__name__)
 
@@ -40,6 +39,17 @@ 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."""
 
@@ -70,15 +80,18 @@ Critical rules:
 - Be concise and direct
 
 Citation Format:
-After your answer, include a "Sources:" section listing document URIs from search results.
-Format: "Sources:\n- [document_uri]"
+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:
-- /path/to/document.pdf
-- /another/document.md
+- Python Documentation (/guides/python.md)
+- /guides/python-basics.md
+
+Note: When using get_full_document, always use document_uri (not document_title).
 """
 
 
@@ -195,7 +208,8 @@ def create_a2a_app(db_path: Path):
             SearchResult(
                 content=chunk.content,
                 score=score,
-                document_uri=(chunk.document_title or chunk.document_uri or ""),
+                document_title=chunk.document_title,
+                document_uri=(chunk.document_uri or ""),
             )
             for chunk, score in expanded_results
         ]