Remove list documents, make separate artifacts for q/a, search & get document
This commit is contained in:
parent
1f2552bd6a
commit
c0691bc237
9 changed files with 678 additions and 69 deletions
31
docs/a2a.md
31
docs/a2a.md
|
|
@ -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
|
- **Intelligent Search**: Performs single or multiple searches depending on question complexity
|
||||||
- **Source Citations**: Always includes sources with both titles and URIs
|
- **Source Citations**: Always includes sources with both titles and URIs
|
||||||
- **Full Document Retrieval**: Can fetch complete documents on request
|
- **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
|
## 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
|
- Track which documents were already found
|
||||||
- Provide contextual follow-up answers
|
- 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
|
### Memory Management
|
||||||
|
|
||||||
To prevent memory growth, the server uses LRU (Least Recently Used) eviction:
|
To prevent memory growth, the server uses LRU (Least Recently Used) eviction:
|
||||||
|
|
||||||
- Maximum 1000 contexts kept in memory (configurable via `A2A_MAX_CONTEXTS`)
|
- Maximum 1000 contexts kept in memory (configurable via `A2A_MAX_CONTEXTS`)
|
||||||
- When limit exceeded, least recently used contexts are automatically evicted
|
- When limit exceeded, least recently used contexts are automatically evicted
|
||||||
- No periodic cleanup needed - eviction happens on-demand
|
|
||||||
|
|
||||||
Configure via environment variable:
|
Configure via environment variable:
|
||||||
```bash
|
```bash
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,3 @@
|
||||||
"""A2A (Agent-to-Agent) server integration for haiku.rag."""
|
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
from contextlib import asynccontextmanager
|
from contextlib import asynccontextmanager
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
@ -112,18 +110,6 @@ def create_a2a_app(
|
||||||
|
|
||||||
return document.content
|
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(
|
worker = ConversationalWorker(
|
||||||
storage=storage,
|
storage=storage,
|
||||||
broker=broker,
|
broker=broker,
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,3 @@
|
||||||
"""Context management for A2A conversations."""
|
|
||||||
|
|
||||||
import uuid
|
import uuid
|
||||||
|
|
||||||
from pydantic import TypeAdapter
|
from pydantic import TypeAdapter
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,3 @@
|
||||||
"""Data models for A2A integration."""
|
|
||||||
|
|
||||||
from pydantic import BaseModel, Field
|
from pydantic import BaseModel, Field
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
|
|
|
||||||
|
|
@ -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.
|
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.
|
IMPORTANT: You are NOT any person mentioned in the documents. You retrieve and present information about them.
|
||||||
|
|
||||||
Tools available:
|
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
|
- get_full_document: Get complete document content by document_uri
|
||||||
- list_documents: Show available documents
|
|
||||||
|
|
||||||
Your process:
|
Your behavior depends on the operation:
|
||||||
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
|
## For direct search requests:
|
||||||
3. When user requests full document: use get_full_document with the exact document_uri from Sources
|
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:
|
Critical rules:
|
||||||
- ONLY answer based on information found via search_documents
|
- ONLY answer based on information found via search_documents
|
||||||
|
- For comprehensive questions, perform MULTIPLE searches with different query angles
|
||||||
- NEVER fabricate or assume information
|
- NEVER fabricate or assume information
|
||||||
- If not found, say: "I cannot find information about this in the knowledge base."
|
- 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
|
- For follow-ups, understand context (pronouns like "he", "it") but always search for facts
|
||||||
- ALWAYS include citations at the end showing document URIs used
|
- In Sources, include the actual chunk content from your search results, not summaries
|
||||||
- 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
|
|
||||||
|
|
||||||
Note: When using get_full_document, always use document_uri (not document_title).
|
Note: When using get_full_document, always use document_uri (not document_title).
|
||||||
"""
|
"""
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,3 @@
|
||||||
"""A2A skill definitions and utilities."""
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
from fasta2a.schema import Message, Skill # type: ignore
|
from fasta2a.schema import Message, Skill # type: ignore
|
||||||
except ImportError as e:
|
except ImportError as e:
|
||||||
|
|
@ -29,6 +27,32 @@ def get_agent_skills() -> list[Skill]:
|
||||||
"Show me the full API documentation",
|
"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",
|
||||||
|
],
|
||||||
|
),
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,3 @@
|
||||||
"""Storage implementations for A2A contexts."""
|
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
from collections import OrderedDict
|
from collections import OrderedDict
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,3 @@
|
||||||
"""A2A worker implementation for conversational QA."""
|
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
import uuid
|
import uuid
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
@ -75,14 +73,11 @@ class ConversationalWorker(Worker[list[Message]]):
|
||||||
question, deps=deps, message_history=message_history
|
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(
|
# Build messages based on skill type
|
||||||
role="agent",
|
response_messages = self._build_response_messages(result, skill_type)
|
||||||
parts=[TextPart(kind="text", text=answer)],
|
|
||||||
kind="message",
|
|
||||||
message_id=str(uuid.uuid4()),
|
|
||||||
)
|
|
||||||
|
|
||||||
# Update context with complete conversation state
|
# Update context with complete conversation state
|
||||||
updated_history = message_history + result.new_messages()
|
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])
|
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(
|
await self.storage.update_task(
|
||||||
task["id"],
|
task["id"],
|
||||||
state="completed",
|
state="completed",
|
||||||
new_messages=[response_message],
|
new_messages=response_messages,
|
||||||
new_artifacts=artifacts,
|
new_artifacts=artifacts,
|
||||||
)
|
)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|
@ -117,16 +112,212 @@ class ConversationalWorker(Worker[list[Message]]):
|
||||||
"""Required by Worker interface but unused - history stored in context."""
|
"""Required by Worker interface but unused - history stored in context."""
|
||||||
return history
|
return history
|
||||||
|
|
||||||
def build_artifacts(self, result) -> list[Artifact]:
|
def _detect_skill(self, result) -> str:
|
||||||
"""Build artifacts from agent result.
|
"""Detect which skill was used based on tool calls and response pattern.
|
||||||
|
|
||||||
Note: Full conversation history (including tool calls) is stored in
|
Returns:
|
||||||
context, so we only create a simple answer artifact here.
|
"search", "retrieve", or "qa"
|
||||||
"""
|
"""
|
||||||
return [
|
from pydantic_ai.messages import ModelResponse, ToolCallPart
|
||||||
Artifact(
|
|
||||||
artifact_id=str(uuid.uuid4()),
|
tool_calls = []
|
||||||
name="answer",
|
for msg in result.new_messages():
|
||||||
parts=[TextPart(kind="text", text=str(result.output))],
|
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
|
||||||
|
|
|
||||||
|
|
@ -261,16 +261,382 @@ async def test_a2a_app_has_skills(temp_db_path):
|
||||||
|
|
||||||
|
|
||||||
def test_get_agent_skills():
|
def test_get_agent_skills():
|
||||||
"""Test that agent skills include document-qa."""
|
"""Test that agent skills include all three skills."""
|
||||||
skills = get_agent_skills()
|
skills = get_agent_skills()
|
||||||
|
|
||||||
assert len(skills) == 1
|
assert len(skills) == 3
|
||||||
|
|
||||||
skill_ids = [skill["id"] for skill in skills]
|
skill_ids = [skill["id"] for skill in skills]
|
||||||
assert "document-qa" in skill_ids
|
assert "document-qa" in skill_ids
|
||||||
|
assert "document-search" in skill_ids
|
||||||
|
assert "document-retrieve" in skill_ids
|
||||||
|
|
||||||
# Check document-qa skill
|
# Check document-qa skill
|
||||||
doc_qa = next(s for s in skills if s["id"] == "document-qa")
|
doc_qa = next(s for s in skills if s["id"] == "document-qa")
|
||||||
assert "Document Question Answering" in doc_qa["name"]
|
assert "Document Question Answering" in doc_qa["name"]
|
||||||
assert "semantic search" in doc_qa["description"]
|
assert "semantic search" in doc_qa["description"]
|
||||||
assert "question-answering" in doc_qa["tags"]
|
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"
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue