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

Remove deep q/a from a2a agent

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-10-13 13:34:15 +03:00
parent 352637faa6
commit 6ea7d73eed
No known key found for this signature in database
5 changed files with 29 additions and 331 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -13,11 +13,7 @@ from haiku.rag.graph.common import get_model
from .context import load_message_history, save_message_history
from .models import AgentDependencies, SearchResult
from .prompts import A2A_SYSTEM_PROMPT
from .skills import (
extract_question_from_task,
extract_skill_preference,
get_agent_skills,
)
from .skills import extract_question_from_task, get_agent_skills
from .storage import LRUMemoryStorage
from .worker import ConversationalWorker
@ -41,7 +37,6 @@ __all__ = [
"load_message_history",
"save_message_history",
"extract_question_from_task",
"extract_skill_preference",
"get_agent_skills",
"LRUMemoryStorage",
]

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

@ -36,14 +36,3 @@ Sources:
Note: When using get_full_document, always use document_uri (not document_title).
"""
ANSWER_EVALUATION_PROMPT = """You evaluate whether an answer adequately addresses a question.
Consider:
- Completeness: Does it answer all parts of the question?
- Specificity: Is it specific enough or too vague?
- Relevance: Does it directly address what was asked?
- Depth: For complex questions, does it provide sufficient depth?
Return is_adequate=True if the answer satisfactorily addresses the question.
Return is_adequate=False if the answer is incomplete, too vague, or requires deeper research."""

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

@ -29,43 +29,9 @@ def get_agent_skills() -> list[Skill]:
"Show me the full API documentation",
],
),
Skill(
id="deep-qa",
name="Deep Question Answering",
description="Multi-step question decomposition and research for complex queries (can take a long time)",
tags=["question-answering", "research", "multi-agent", "complex-queries"],
input_modes=["application/json"],
output_modes=["application/json"],
examples=[
"What are the architectural patterns used in haiku.rag and how do they compare?",
"Analyze the trade-offs between the simple QA and research agents",
"What are all the configuration options and their effects?",
],
),
]
def extract_skill_preference(task_history: list[Message]) -> str:
"""Extract skill preference from task history metadata.
Args:
task_history: Task history messages
Returns:
Skill ID if found in metadata, otherwise "document-qa" (default)
"""
for msg in task_history:
if msg.get("role") == "user":
for part in msg.get("parts", []):
if part.get("kind") == "data":
metadata = part.get("metadata", {})
if metadata.get("type") == "skill_preference":
skill = part.get("data", {}).get("skill")
if skill:
return skill
return "document-qa"
def extract_question_from_task(task_history: list[Message]) -> str | None:
"""Extract the user's question from task history.

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

@ -4,25 +4,17 @@ import logging
import uuid
from pathlib import Path
from pydantic import BaseModel, Field
from pydantic_ai import Agent
from haiku.rag.a2a.context import load_message_history, save_message_history
from haiku.rag.a2a.models import AgentDependencies
from haiku.rag.a2a.skills import extract_question_from_task, extract_skill_preference
from haiku.rag.a2a.skills import extract_question_from_task
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.graph.common import get_model
from haiku.rag.qa.deep.dependencies import DeepQAContext
from haiku.rag.qa.deep.graph import build_deep_qa_graph
from haiku.rag.qa.deep.nodes import DeepQAPlanNode
from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
try:
from fasta2a import Worker # type: ignore
from fasta2a.schema import ( # type: ignore
Artifact,
DataPart,
Message,
TaskIdParams,
TaskSendParams,
@ -51,44 +43,6 @@ class ConversationalWorker(Worker[list[Message]]):
self.db_path = db_path
self.agent = agent
async def evaluate_answer_adequacy(self, question: str, answer: str) -> bool:
"""Use LLM to evaluate if answer adequately addresses the question.
Args:
question: The original question
answer: The answer to evaluate
Returns:
True if answer is adequate, False if more research needed
"""
class AnswerEvaluation(BaseModel):
is_adequate: bool = Field(
description="True if the answer adequately addresses the question, False if more research is needed"
)
reasoning: str = Field(description="Brief explanation of the evaluation")
from .prompts import ANSWER_EVALUATION_PROMPT
evaluation_agent = Agent(
model=get_model(Config.QA_PROVIDER, Config.QA_MODEL),
output_type=AnswerEvaluation,
system_prompt=ANSWER_EVALUATION_PROMPT,
retries=1,
)
prompt = f"""Question: {question}
Answer: {answer}
Does this answer adequately address the question?"""
result = await evaluation_agent.run(prompt)
logger.info(
f"Answer evaluation: is_adequate={result.output.is_adequate}, reasoning={result.output.reasoning}"
)
return result.output.is_adequate
async def run_task(self, params: TaskSendParams) -> None:
task = await self.storage.load_task(params["id"])
if task is None:
@ -101,112 +55,49 @@ Does this answer adequately address the question?"""
await self.storage.update_task(task["id"], state="working")
# Extract skill preference and question
task_history = task.get("history", [])
skill = extract_skill_preference(task_history)
question = extract_question_from_task(task_history)
if not question:
await self.storage.update_task(task["id"], state="failed")
return
logger.info(f"Task {task['id']} requested skill: {skill}")
try:
async with HaikuRAG(self.db_path) as client:
if skill == "deep-qa":
# Explicitly requested deep QA
logger.info(f"Task {task['id']}: Running deep QA (explicit)")
deep_result, deep_state = await self.run_deep_qa(client, question)
context = await self.storage.load_context(task["context_id"]) or []
message_history = load_message_history(context)
response_message = Message(
role="agent",
parts=[TextPart(kind="text", text=deep_result.answer)],
kind="message",
message_id=str(uuid.uuid4()),
)
from haiku.rag.a2a.models import AgentDependencies
artifacts = self.build_deep_qa_artifacts(deep_result, deep_state)
deps = AgentDependencies(client=client)
await self.storage.update_task(
task["id"],
state="completed",
new_messages=[response_message],
new_artifacts=artifacts,
)
else:
# Try simple QA first (default behavior or explicit document-qa)
logger.info(f"Task {task['id']}: Trying simple QA first")
result = await self.agent.run(
question, deps=deps, message_history=message_history
)
context = await self.storage.load_context(task["context_id"]) or []
message_history = load_message_history(context)
answer = str(result.output)
from .models import AgentDependencies
response_message = Message(
role="agent",
parts=[TextPart(kind="text", text=answer)],
kind="message",
message_id=str(uuid.uuid4()),
)
deps = AgentDependencies(client=client)
# Update context with complete conversation state
updated_history = message_history + result.new_messages()
state_message = save_message_history(updated_history)
result = await self.agent.run(
question, deps=deps, message_history=message_history
)
await self.storage.update_context(task["context_id"], [state_message])
answer = str(result.output)
artifacts = self.build_artifacts(result)
# Evaluate answer adequacy
is_adequate = await self.evaluate_answer_adequacy(question, answer)
if not is_adequate:
# Escalate to deep QA
logger.info(
f"Task {task['id']}: Answer inadequate, escalating to deep QA"
)
deep_result, deep_state = await self.run_deep_qa(
client, question
)
response_message = Message(
role="agent",
parts=[TextPart(kind="text", text=deep_result.answer)],
kind="message",
message_id=str(uuid.uuid4()),
)
artifacts = self.build_deep_qa_artifacts(
deep_result, deep_state
)
await self.storage.update_task(
task["id"],
state="completed",
new_messages=[response_message],
new_artifacts=artifacts,
)
else:
# Simple QA answer is adequate
logger.info(f"Task {task['id']}: Simple QA answer is adequate")
response_message = Message(
role="agent",
parts=[TextPart(kind="text", text=answer)],
kind="message",
message_id=str(uuid.uuid4()),
)
# Update context with complete conversation state
updated_history = message_history + result.new_messages()
state_message = save_message_history(updated_history)
await self.storage.update_context(
task["context_id"], [state_message]
)
artifacts = self.build_artifacts(result)
await self.storage.update_task(
task["id"],
state="completed",
new_messages=[response_message],
new_artifacts=artifacts,
)
await self.storage.update_task(
task["id"],
state="completed",
new_messages=[response_message],
new_artifacts=artifacts,
)
except Exception as e:
logger.error(
"Task execution failed: task_id=%s, question=%s, error=%s",
@ -218,25 +109,6 @@ Does this answer adequately address the question?"""
await self.storage.update_task(task["id"], state="failed")
raise
async def run_deep_qa(self, client: HaikuRAG, question: str):
"""Run deep QA graph for complex questions.
Args:
client: HaikuRAG client
question: User's question
Returns:
Tuple of (DeepQAAnswer, DeepQAState) with answer and state
"""
graph = build_deep_qa_graph()
context = DeepQAContext(original_question=question, use_citations=False)
state = DeepQAState(context=context)
deps = DeepQADeps(client=client, console=None)
start_node = DeepQAPlanNode(provider=Config.QA_PROVIDER, model=Config.QA_MODEL)
result = await graph.run(start_node=start_node, state=state, deps=deps)
return result.output, state
async def cancel_task(self, params: TaskIdParams) -> None:
"""Cancel a task - not implemented for this worker."""
pass
@ -258,53 +130,3 @@ Does this answer adequately address the question?"""
parts=[TextPart(kind="text", text=str(result.output))],
)
]
def build_deep_qa_artifacts(self, result, state: DeepQAState) -> list[Artifact]:
"""Build rich artifacts from deep QA result.
Args:
result: DeepQAAnswer with final answer
state: DeepQAState with research process details
Returns:
List of artifacts including answer and research breakdown
"""
artifacts = [
# Final answer artifact
Artifact(
artifact_id=str(uuid.uuid4()),
name="answer",
parts=[TextPart(kind="text", text=result.answer)],
)
]
# Add research process artifact with sub-questions and answers
if state.context.qa_responses:
research_data = {
"original_question": state.context.original_question,
"iterations": state.iterations,
"sub_questions_answered": [
{
"question": qa.query,
"answer": qa.answer,
"sources": qa.sources,
}
for qa in state.context.qa_responses
],
}
artifacts.append(
Artifact(
artifact_id=str(uuid.uuid4()),
name="research_process",
parts=[
DataPart(
kind="data",
data=research_data,
metadata={"type": "deep_qa_research"},
)
],
)
)
return artifacts

78
tests/test_a2a.py
View file

@ -4,7 +4,6 @@ import pytest
from haiku.rag.a2a import (
extract_question_from_task,
extract_skill_preference,
get_agent_skills,
load_message_history,
save_message_history,
@ -262,89 +261,16 @@ async def test_a2a_app_has_skills(temp_db_path):
def test_get_agent_skills():
"""Test that agent skills include both document-qa and deep-qa."""
"""Test that agent skills include document-qa."""
skills = get_agent_skills()
assert len(skills) == 2
assert len(skills) == 1
skill_ids = [skill["id"] for skill in skills]
assert "document-qa" in skill_ids
assert "deep-qa" 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 deep-qa skill
deep_qa = next(s for s in skills if s["id"] == "deep-qa")
assert "Deep Question Answering" in deep_qa["name"]
assert "Multi-step" in deep_qa["description"]
assert "research" in deep_qa["tags"]
@pytest.mark.asyncio
async def test_extract_skill_preference_with_metadata():
"""Test extracting skill preference from message metadata."""
from fasta2a.schema import DataPart
task_history: list[Message] = [
Message(
role="user",
parts=[
TextPart(kind="text", text="Complex question"),
DataPart(
kind="data",
data={"skill": "deep-qa"},
metadata={"type": "skill_preference"},
),
],
kind="message",
message_id=str(uuid.uuid4()),
)
]
skill = extract_skill_preference(task_history)
assert skill == "deep-qa"
@pytest.mark.asyncio
async def test_extract_skill_preference_default():
"""Test that skill preference defaults to document-qa."""
task_history: list[Message] = [
Message(
role="user",
parts=[TextPart(kind="text", text="What is Python?")],
kind="message",
message_id=str(uuid.uuid4()),
)
]
skill = extract_skill_preference(task_history)
assert skill == "document-qa"
@pytest.mark.asyncio
async def test_extract_skill_preference_no_skill_in_data():
"""Test skill preference when DataPart exists but has no skill."""
from fasta2a.schema import DataPart
task_history: list[Message] = [
Message(
role="user",
parts=[
TextPart(kind="text", text="Question"),
DataPart(
kind="data",
data={"other": "value"},
metadata={"type": "skill_preference"},
),
],
kind="message",
message_id=str(uuid.uuid4()),
)
]
skill = extract_skill_preference(task_history)
assert skill == "document-qa"