danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Remove list documents, make separate artifacts for q/a, search & get document

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-10-13 17:43:24 +03:00
parent 1f2552bd6a
commit c0691bc237
No known key found for this signature in database
9 changed files with 678 additions and 69 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
docs/a2a.md
View file

@ -9,7 +9,10 @@ The A2A server exposes `haiku.rag` as a conversational agent using the Agent-to-
- **Intelligent Search**: Performs single or multiple searches depending on question complexity
- **Source Citations**: Always includes sources with both titles and URIs
- **Full Document Retrieval**: Can fetch complete documents on request
- **Document Discovery**: Lists available documents to help users explore the knowledge base
- **Multiple Skills**: Exposes three distinct skills with appropriate artifacts:
- `document-qa`: Conversational question answering (default)
- `document-search`: Semantic search with structured results
- `document-retrieve`: Fetch complete documents by URI
## Starting A2A Server
@ -82,13 +85,37 @@ Each conversation is identified by a `context_id`. All messages within the same
- Track which documents were already found
- Provide contextual follow-up answers
### Skills
The agent exposes three skills:
- **document-qa** (default): Conversational question answering including follow-ups and multi-turn dialogue
- **document-search**: Direct semantic search returning formatted results
- **document-retrieve**: Fetch complete document content by URI
### Artifacts
All operations create artifacts for traceability:
- **search_results**: Created for each `search_documents` tool call
- Contains query and array of SearchResult objects (content, score, document_title, document_uri)
- **document**: Created for each `get_full_document` tool call
- Contains complete document text
- **qa_result**: Created for all document-qa operations
- Contains question, answer, and skill identifier
- Always created for Q&A, even when answering from conversation history without tools
### Memory Management
To prevent memory growth, the server uses LRU (Least Recently Used) eviction:
- Maximum 1000 contexts kept in memory (configurable via `A2A_MAX_CONTEXTS`)
- When limit exceeded, least recently used contexts are automatically evicted
- No periodic cleanup needed - eviction happens on-demand
Configure via environment variable:
```bash

14
src/haiku/rag/a2a/__init__.py
View file

@ -1,5 +1,3 @@
"""A2A (Agent-to-Agent) server integration for haiku.rag."""
import logging
from contextlib import asynccontextmanager
from pathlib import Path
@ -112,18 +110,6 @@ def create_a2a_app(
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,

2
src/haiku/rag/a2a/context.py
View file

@ -1,5 +1,3 @@
"""Context management for A2A conversations."""
import uuid
from pydantic import TypeAdapter

2
src/haiku/rag/a2a/models.py
View file

@ -1,5 +1,3 @@
"""Data models for A2A integration."""
from pydantic import BaseModel, Field
from haiku.rag.client import HaikuRAG

65
src/haiku/rag/a2a/prompts.py
View file

@ -1,38 +1,59 @@
"""Prompts for A2A agents."""
A2A_SYSTEM_PROMPT = """You are Haiku.rag, an AI assistant that helps users find information from a document knowledge base.
IMPORTANT: You are NOT any person mentioned in the documents. You retrieve and present information about them.
Tools available:
- search_documents: Query for relevant text chunks
- search_documents: Query for relevant text chunks (returns SearchResult objects with content, score, document_title, document_uri)
- get_full_document: Get complete document content by document_uri
- list_documents: Show available documents
Your process:
1. Search phase: For straightforward questions use one search, for complex questions search multiple times with different queries
2. Synthesis phase: Combine the search results into a comprehensive answer
3. When user requests full document: use get_full_document with the exact document_uri from Sources
Your behavior depends on the operation:
## For direct search requests:
When the user is explicitly searching (e.g., "search for X", "find documents about Y"):
- Use search_documents tool ONLY
- Format results as a numbered list using markdown formatting
- For each result show:
* First line: *Score in italic* | **source in bold** (title if available, otherwise URI)
* Second line: The FULL chunk content (do not summarize or truncate)
- Present results in order of relevance
- Be concise - just present the search results, do not synthesize or add commentary
Example format:
Found 3 relevant results:
1. *Score: 0.95* | **Python Documentation** (/guides/python.md)
Python is a high-level, general-purpose programming language. Its design philosophy emphasizes code readability with the use of significant indentation.
2. *Score: 0.87* | **/guides/python-basics.md**
Python supports multiple programming paradigms, including structured, object-oriented and functional programming.
## For question-answering:
When the user asks a question (e.g., "What is Python?", "How does X work?"):
- For complex questions, use search_documents MULTIPLE TIMES with DIFFERENT queries to gather comprehensive information
- Example: For "What are the benefits and drawbacks of Python?", search separately for:
* "Python benefits advantages"
* "Python drawbacks disadvantages limitations"
- Synthesize information from all searches into a comprehensive answer
- Include "Sources:" section at the end listing sources used
Sources Format:
List each source with its title/URI and the relevant chunk content (NOT the score).
Format: "- **[title or URI]**: [chunk content]"
Example:
[Your synthesized answer here]
Sources:
- **Python Documentation** (/guides/python.md): Python is a high-level, general-purpose programming language. Its design philosophy emphasizes code readability.
- **/guides/python-basics.md**: Python supports multiple programming paradigms, including structured, object-oriented and functional programming.
Critical rules:
- ONLY answer based on information found via search_documents
- For comprehensive questions, perform MULTIPLE searches with different query angles
- NEVER fabricate or assume information
- If not found, say: "I cannot find information about this in the knowledge base."
- For follow-ups, understand context (pronouns like "he", "it") but always search for facts
- ALWAYS include citations at the end showing document URIs used
- Be concise and direct
Citation Format:
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:
- Python Documentation (/guides/python.md)
- /guides/python-basics.md
- In Sources, include the actual chunk content from your search results, not summaries
Note: When using get_full_document, always use document_uri (not document_title).
"""

28
src/haiku/rag/a2a/skills.py
View file

@ -1,5 +1,3 @@
"""A2A skill definitions and utilities."""
try:
from fasta2a.schema import Message, Skill # type: ignore
except ImportError as e:
@ -29,6 +27,32 @@ def get_agent_skills() -> list[Skill]:
"Show me the full API documentation",
],
),
Skill(
id="document-search",
name="Document Search",
description="Search for relevant document chunks in the knowledge base using hybrid (semantic and BM25) search",
tags=["search", "retrieval", "semantic-search"],
input_modes=["application/json"],
output_modes=["application/json"],
examples=[
"Search for Python best practices",
"Find documents about authentication",
"Look for API documentation",
],
),
Skill(
id="document-retrieve",
name="Document Retrieval",
description="Retrieve the complete content of a specific document by its URI",
tags=["retrieval", "fetch", "document"],
input_modes=["application/json"],
output_modes=["application/json"],
examples=[
"Get the full content of document X",
"Retrieve document by URI",
"Show me the complete document",
],
),
]

2
src/haiku/rag/a2a/storage.py
View file

@ -1,5 +1,3 @@
"""Storage implementations for A2A contexts."""
import logging
from collections import OrderedDict

233
src/haiku/rag/a2a/worker.py
View file

@ -1,5 +1,3 @@
"""A2A worker implementation for conversational QA."""
import logging
import uuid
from pathlib import Path
@ -75,14 +73,11 @@ class ConversationalWorker(Worker[list[Message]]):
question, deps=deps, message_history=message_history
)
answer = str(result.output)
# Detect which skill was used
skill_type = self._detect_skill(result)
response_message = Message(
role="agent",
parts=[TextPart(kind="text", text=answer)],
kind="message",
message_id=str(uuid.uuid4()),
)
# Build messages based on skill type
response_messages = self._build_response_messages(result, skill_type)
# Update context with complete conversation state
updated_history = message_history + result.new_messages()
@ -90,12 +85,12 @@ class ConversationalWorker(Worker[list[Message]]):
await self.storage.update_context(task["context_id"], [state_message])
artifacts = self.build_artifacts(result)
artifacts = self.build_artifacts(result, skill_type, question)
await self.storage.update_task(
task["id"],
state="completed",
new_messages=[response_message],
new_messages=response_messages,
new_artifacts=artifacts,
)
except Exception as e:
@ -117,16 +112,212 @@ class ConversationalWorker(Worker[list[Message]]):
"""Required by Worker interface but unused - history stored in context."""
return history
def build_artifacts(self, result) -> list[Artifact]:
"""Build artifacts from agent result.
def _detect_skill(self, result) -> str:
"""Detect which skill was used based on tool calls and response pattern.
Note: Full conversation history (including tool calls) is stored in
context, so we only create a simple answer artifact here.
Returns:
"search", "retrieve", or "qa"
"""
return [
Artifact(
artifact_id=str(uuid.uuid4()),
name="answer",
parts=[TextPart(kind="text", text=str(result.output))],
from pydantic_ai.messages import ModelResponse, ToolCallPart
tool_calls = []
for msg in result.new_messages():
if isinstance(msg, ModelResponse):
for part in msg.parts:
if isinstance(part, ToolCallPart):
tool_calls.append(part.tool_name)
# Check if output looks like formatted search results
output_str = str(result.output).strip()
# Check for either format: "Found N relevant results" or "**Search results for"
is_search_format = (
output_str.startswith("Found ") and "relevant results" in output_str[:100]
) or output_str.startswith("**Search results for")
skill_type = "qa"
# If output is in search format and only search tools were used, it's a search
if is_search_format and all(tc == "search_documents" for tc in tool_calls):
skill_type = "search"
elif "get_full_document" in tool_calls and len(tool_calls) == 1:
skill_type = "retrieve"
return skill_type
def _build_response_messages(self, result, skill_type: str) -> list[Message]:
"""Build response messages based on skill type.
All skills return a single text message with LLM's response.
Structured data is provided via artifacts for search/retrieve.
"""
if skill_type == "search":
# Return LLM's formatted response
return [
Message(
role="agent",
parts=[TextPart(kind="text", text=str(result.output))],
kind="message",
message_id=str(uuid.uuid4()),
)
]
elif skill_type == "retrieve":
# Extract document content
from pydantic_ai.messages import ModelRequest, ToolReturnPart
document_content = ""
for msg in result.new_messages():
if isinstance(msg, ModelRequest):
for part in msg.parts:
if (
isinstance(part, ToolReturnPart)
and part.tool_name == "get_full_document"
):
document_content = part.content
break
return [
Message(
role="agent",
parts=[TextPart(kind="text", text=document_content)],
kind="message",
message_id=str(uuid.uuid4()),
)
]
else:
# Conversational Q&A - use agent's answer
return [
Message(
role="agent",
parts=[TextPart(kind="text", text=str(result.output))],
kind="message",
message_id=str(uuid.uuid4()),
)
]
def build_artifacts(
self, result, skill_type: str | None = None, question: str | None = None
) -> list[Artifact]:
"""Build artifacts from agent result based on tool calls.
Creates artifacts for:
- Each tool call (search_documents, get_full_document)
- Q&A operations: additional artifact with question and answer (only if tools were used)
"""
if skill_type is None:
skill_type = self._detect_skill(result)
artifacts = []
# Always create artifacts for all tool calls
tool_artifacts = self._build_all_tool_artifacts(result)
artifacts.extend(tool_artifacts)
# For Q&A, always add a Q&A artifact with question and answer
# This includes follow-up questions, clarifications, and conversational responses
if skill_type == "qa" and question:
from fasta2a.schema import DataPart
artifacts.append(
Artifact(
artifact_id=str(uuid.uuid4()),
name="qa_result",
parts=[
DataPart(
kind="data",
data={
"question": question,
"answer": str(result.output),
"skill": "document-qa",
},
metadata={"skill": "document-qa"},
)
],
)
)
]
return artifacts
def _build_all_tool_artifacts(self, result) -> list[Artifact]:
"""Build artifacts for all tool calls."""
from pydantic_ai.messages import (
ModelRequest,
ModelResponse,
ToolCallPart,
ToolReturnPart,
)
artifacts = []
# Track tool calls and their returns by call_id
tool_returns = {}
for msg in result.new_messages():
if isinstance(msg, ModelRequest):
for part in msg.parts:
if isinstance(part, ToolReturnPart):
result_count = (
len(part.content) if isinstance(part.content, list) else 1
)
logger.info(
"Tool return: tool_call_id=%s, tool_name=%s, result_count=%s",
part.tool_call_id,
part.tool_name,
result_count,
)
tool_returns[part.tool_call_id] = (part.tool_name, part.content)
# Create artifacts for each tool call
for msg in result.new_messages():
if isinstance(msg, ModelResponse):
for part in msg.parts:
if isinstance(part, ToolCallPart):
tool_name, content = tool_returns.get(
part.tool_call_id, (None, None)
)
if tool_name == "search_documents" and content:
import json
from fasta2a.schema import DataPart
# Extract query from tool call arguments
query = ""
if isinstance(part.args, dict):
query = part.args.get("query", "")
elif isinstance(part.args, str):
# Args is a JSON string - parse it
try:
args_dict = json.loads(part.args)
query = args_dict.get("query", "")
except (json.JSONDecodeError, AttributeError):
query = ""
elif hasattr(part.args, "get") and callable(
getattr(part.args, "get", None)
):
# ArgsDict or dict-like object
query = part.args.get("query", "") # type: ignore
elif hasattr(part.args, "query"):
# Object with query attribute
query = str(part.args.query) # type: ignore
artifacts.append(
Artifact(
artifact_id=str(uuid.uuid4()),
name="search_results",
parts=[
DataPart(
kind="data",
data={"results": content, "query": query},
metadata={"query": query},
)
],
)
)
elif tool_name == "get_full_document" and content:
artifacts.append(
Artifact(
artifact_id=str(uuid.uuid4()),
name="document",
parts=[TextPart(kind="text", text=content)],
)
)
return artifacts

370
tests/test_a2a.py
View file

@ -261,16 +261,382 @@ async def test_a2a_app_has_skills(temp_db_path):
def test_get_agent_skills():
"""Test that agent skills include document-qa."""
"""Test that agent skills include all three skills."""
skills = get_agent_skills()
assert len(skills) == 1
assert len(skills) == 3
skill_ids = [skill["id"] for skill in skills]
assert "document-qa" in skill_ids
assert "document-search" in skill_ids
assert "document-retrieve" in skill_ids
# Check document-qa skill
doc_qa = next(s for s in skills if s["id"] == "document-qa")
assert "Document Question Answering" in doc_qa["name"]
assert "semantic search" in doc_qa["description"]
assert "question-answering" in doc_qa["tags"]
# Check document-search skill
doc_search = next(s for s in skills if s["id"] == "document-search")
assert "Document Search" in doc_search["name"]
assert "search" in doc_search["tags"]
# Check document-retrieve skill
doc_retrieve = next(s for s in skills if s["id"] == "document-retrieve")
assert "Document Retrieval" in doc_retrieve["name"]
assert "retrieval" in doc_retrieve["tags"]
@pytest.mark.asyncio
async def test_build_artifacts_for_search():
"""Test that search operations produce structured search artifacts."""
from pydantic_ai.messages import (
ModelRequest,
ModelResponse,
ToolCallPart,
ToolReturnPart,
)
from haiku.rag.a2a.worker import ConversationalWorker
class MockResult:
output = "Found 1 relevant results:\n\n1. *Score: 0.9* | **test**\nresult"
def new_messages(self):
return [
ModelResponse(
parts=[
ToolCallPart(
tool_name="search_documents",
args={"query": "test", "limit": 3},
tool_call_id="call_1",
)
]
),
ModelRequest(
parts=[
ToolReturnPart(
tool_name="search_documents",
content=[{"content": "result", "score": 0.9}],
tool_call_id="call_1",
)
]
),
]
from pathlib import Path
from fasta2a.broker import InMemoryBroker
from fasta2a.storage import InMemoryStorage
worker = ConversationalWorker(
storage=InMemoryStorage(),
broker=InMemoryBroker(),
db_path=Path("/tmp/test.db"),
agent=None, # type: ignore
)
artifacts = worker.build_artifacts(MockResult(), "search", "test query")
assert len(artifacts) == 1
assert artifacts[0].get("name") == "search_results"
assert len(artifacts[0]["parts"]) == 1
assert artifacts[0]["parts"][0]["kind"] == "data"
assert "results" in artifacts[0]["parts"][0]["data"]
assert "query" in artifacts[0]["parts"][0]["data"]
@pytest.mark.asyncio
async def test_build_artifacts_for_retrieve():
"""Test that retrieve operations produce document artifacts."""
from pydantic_ai.messages import (
ModelRequest,
ModelResponse,
ToolCallPart,
ToolReturnPart,
)
from haiku.rag.a2a.worker import ConversationalWorker
class MockResult:
output = "Document content"
def new_messages(self):
return [
ModelResponse(
parts=[
ToolCallPart(
tool_name="get_full_document",
args={"document_uri": "test.txt"},
tool_call_id="call_1",
)
]
),
ModelRequest(
parts=[
ToolReturnPart(
tool_name="get_full_document",
content="Full document content here",
tool_call_id="call_1",
)
]
),
]
from pathlib import Path
from fasta2a.broker import InMemoryBroker
from fasta2a.storage import InMemoryStorage
worker = ConversationalWorker(
storage=InMemoryStorage(),
broker=InMemoryBroker(),
db_path=Path("/tmp/test.db"),
agent=None, # type: ignore
)
artifacts = worker.build_artifacts(MockResult(), "retrieve", "test query")
assert len(artifacts) == 1
assert artifacts[0].get("name") == "document"
assert artifacts[0]["parts"][0]["kind"] == "text"
assert artifacts[0]["parts"][0]["text"] == "Full document content here"
@pytest.mark.asyncio
async def test_build_artifacts_for_multiple_searches():
"""Test that multiple searches each get their own artifact with correct results."""
from pydantic_ai.messages import (
ModelRequest,
ModelResponse,
ToolCallPart,
ToolReturnPart,
)
from pydantic_ai.messages import TextPart as AITextPart
from haiku.rag.a2a.worker import ConversationalWorker
class MockResult:
output = "Answer based on multiple searches"
def new_messages(self):
return [
# First search
ModelResponse(
parts=[
ToolCallPart(
tool_name="search_documents",
args={"query": "first query", "limit": 2},
tool_call_id="call_1",
)
]
),
ModelRequest(
parts=[
ToolReturnPart(
tool_name="search_documents",
content=[
{"content": "result 1", "score": 0.9},
{"content": "result 2", "score": 0.8},
],
tool_call_id="call_1",
)
]
),
# Second search
ModelResponse(
parts=[
ToolCallPart(
tool_name="search_documents",
args={"query": "second query", "limit": 2},
tool_call_id="call_2",
)
]
),
ModelRequest(
parts=[
ToolReturnPart(
tool_name="search_documents",
content=[
{"content": "result 3", "score": 0.7},
{"content": "result 4", "score": 0.6},
],
tool_call_id="call_2",
)
]
),
ModelResponse(
parts=[AITextPart(content="Answer based on multiple searches")]
),
]
from pathlib import Path
from fasta2a.broker import InMemoryBroker
from fasta2a.storage import InMemoryStorage
worker = ConversationalWorker(
storage=InMemoryStorage(),
broker=InMemoryBroker(),
db_path=Path("/tmp/test.db"),
agent=None, # type: ignore
)
artifacts = worker.build_artifacts(MockResult(), "qa", "What is the answer?")
# Should have 2 search artifacts + 1 qa_result artifact
assert len(artifacts) == 3
# First search artifact
assert artifacts[0].get("name") == "search_results"
part_0 = artifacts[0]["parts"][0]
assert part_0.get("data", {}).get("query") == "first query"
results_1 = part_0.get("data", {}).get("results", [])
assert len(results_1) == 2
assert results_1[0]["content"] == "result 1"
assert results_1[1]["content"] == "result 2"
# Second search artifact
assert artifacts[1].get("name") == "search_results"
part_1 = artifacts[1]["parts"][0]
assert part_1.get("data", {}).get("query") == "second query"
results_2 = part_1.get("data", {}).get("results", [])
assert len(results_2) == 2
assert results_2[0]["content"] == "result 3"
assert results_2[1]["content"] == "result 4"
# Q&A artifact
assert artifacts[2].get("name") == "qa_result"
@pytest.mark.asyncio
async def test_qa_artifact_for_conversational_messages():
"""Test that conversational Q&A messages always create qa_result artifacts."""
from pydantic_ai.messages import ModelResponse
from pydantic_ai.messages import TextPart as AITextPart
from haiku.rag.a2a.worker import ConversationalWorker
class MockResult:
output = "Hello! How can I help you?"
def new_messages(self):
# No tool calls, just a conversational response
return [
ModelResponse(parts=[AITextPart(content="Hello! How can I help you?")]),
]
from pathlib import Path
from fasta2a.broker import InMemoryBroker
from fasta2a.storage import InMemoryStorage
worker = ConversationalWorker(
storage=InMemoryStorage(),
broker=InMemoryBroker(),
db_path=Path("/tmp/test.db"),
agent=None, # type: ignore
)
artifacts = worker.build_artifacts(MockResult(), "qa", "Hello")
# Should have qa_result artifact (even without tools, for A2A traceability)
assert len(artifacts) == 1
assert artifacts[0].get("name") == "qa_result"
part = artifacts[0]["parts"][0]
assert part.get("data", {}).get("question") == "Hello"
assert part.get("data", {}).get("answer") == "Hello! How can I help you?"
@pytest.mark.asyncio
async def test_build_artifacts_for_qa():
"""Test that Q&A operations produce artifacts for each tool call."""
from pydantic_ai.messages import (
ModelRequest,
ModelResponse,
ToolCallPart,
ToolReturnPart,
)
from pydantic_ai.messages import TextPart as AITextPart
from haiku.rag.a2a.worker import ConversationalWorker
class MockResult:
output = "This is the answer"
def new_messages(self):
# Multiple tool calls indicates Q&A workflow
return [
ModelResponse(
parts=[
ToolCallPart(
tool_name="search_documents",
args={"query": "test", "limit": 3},
tool_call_id="call_1",
)
]
),
ModelRequest(
parts=[
ToolReturnPart(
tool_name="search_documents",
content=[{"content": "result", "score": 0.9}],
tool_call_id="call_1",
)
]
),
ModelResponse(
parts=[
ToolCallPart(
tool_name="get_full_document",
args={"document_uri": "test.txt"},
tool_call_id="call_2",
)
]
),
ModelRequest(
parts=[
ToolReturnPart(
tool_name="get_full_document",
content="Full content",
tool_call_id="call_2",
)
]
),
ModelResponse(parts=[AITextPart(content="This is the answer")]),
]
from pathlib import Path
from fasta2a.broker import InMemoryBroker
from fasta2a.storage import InMemoryStorage
worker = ConversationalWorker(
storage=InMemoryStorage(),
broker=InMemoryBroker(),
db_path=Path("/tmp/test.db"),
agent=None, # type: ignore
)
artifacts = worker.build_artifacts(MockResult(), "qa", "What is Python?")
# Q&A should produce artifacts for each tool call (search + retrieve) + final Q&A artifact
assert len(artifacts) == 3
# First artifact is from search_documents
assert artifacts[0].get("name") == "search_results"
assert artifacts[0]["parts"][0]["kind"] == "data"
assert "results" in artifacts[0]["parts"][0]["data"]
assert artifacts[0]["parts"][0]["data"]["query"] == "test"
# Second artifact is from get_full_document
assert artifacts[1].get("name") == "document"
assert artifacts[1]["parts"][0]["kind"] == "text"
# Third artifact is the Q&A result
assert artifacts[2].get("name") == "qa_result"
assert artifacts[2]["parts"][0]["kind"] == "data"
assert artifacts[2]["parts"][0]["data"]["question"] == "What is Python?"
assert artifacts[2]["parts"][0]["data"]["answer"] == "This is the answer"
assert artifacts[2]["parts"][0]["data"]["skill"] == "document-qa"