diff --git a/src/haiku/rag/a2a/__init__.py b/src/haiku/rag/a2a/__init__.py index 7266cada..06a5467c 100644 --- a/src/haiku/rag/a2a/__init__.py +++ b/src/haiku/rag/a2a/__init__.py @@ -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", ] diff --git a/src/haiku/rag/a2a/prompts.py b/src/haiku/rag/a2a/prompts.py index 79f7386e..1e46f85d 100644 --- a/src/haiku/rag/a2a/prompts.py +++ b/src/haiku/rag/a2a/prompts.py @@ -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.""" diff --git a/src/haiku/rag/a2a/skills.py b/src/haiku/rag/a2a/skills.py index a181d74f..10750174 100644 --- a/src/haiku/rag/a2a/skills.py +++ b/src/haiku/rag/a2a/skills.py @@ -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. diff --git a/src/haiku/rag/a2a/worker.py b/src/haiku/rag/a2a/worker.py index 36e382cf..6029dd83 100644 --- a/src/haiku/rag/a2a/worker.py +++ b/src/haiku/rag/a2a/worker.py @@ -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 diff --git a/tests/test_a2a.py b/tests/test_a2a.py index d0149478..58d77ab9 100644 --- a/tests/test_a2a.py +++ b/tests/test_a2a.py @@ -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"