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

Add deep QA to a2a agent

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-10-10 14:31:59 +03:00
parent 1a88b0355a
commit 9aad6b65b8
No known key found for this signature in database
1 changed files with 86 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

120
src/haiku/rag/a2a.py
View file

@ -13,6 +13,10 @@ from pydantic_core import to_jsonable_python
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config from haiku.rag.config import Config
from haiku.rag.graph.common import get_model 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
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -386,47 +390,74 @@ def create_a2a_app(db_path: Path):
logger.info(f"Task {task['id']} using skill: {skill}") logger.info(f"Task {task['id']} using skill: {skill}")
try: try:
# Load conversation context
context = await self.storage.load_context(task["context_id"]) or []
# Load conversation history
message_history = load_message_history(context)
# Create fresh client for this task and run agent
async with HaikuRAG(db_path) as client: async with HaikuRAG(db_path) as client:
deps = AgentDependencies(client=client) if skill == "deep-qa":
# Run deep QA graph
deep_result = await self.run_deep_qa(client, question)
# Run agent with full conversation history including tool calls # Build response message
result = await agent.run( response_message = Message(
question, deps=deps, message_history=message_history role="agent",
) parts=[TextPart(kind="text", text=deep_result.answer)],
kind="message",
message_id=str(uuid.uuid4()),
)
# Build response message for A2A protocol # Build artifacts (basic for now, will be enhanced in commit 3)
response_message = Message( artifacts = [
role="agent", Artifact(
parts=[TextPart(kind="text", text=str(result.output))], artifact_id=str(uuid.uuid4()),
kind="message", name="answer",
message_id=str(uuid.uuid4()), parts=[TextPart(kind="text", text=deep_result.answer)],
) )
]
# Update context with complete conversation state await self.storage.update_task(
# Store all messages from this run (includes tool calls & results) task["id"],
updated_history = message_history + result.new_messages() state="completed",
state_message = save_message_history(updated_history) new_messages=[response_message],
new_artifacts=artifacts,
)
else:
# Load conversation context for simple QA
context = (
await self.storage.load_context(task["context_id"]) or []
)
message_history = load_message_history(context)
# Replace old state with new complete state deps = AgentDependencies(client=client)
await self.storage.update_context(
task["context_id"], [state_message]
)
# Build rich artifacts with search results and answer # Run agent with full conversation history including tool calls
artifacts = self.build_artifacts(result) result = await agent.run(
question, deps=deps, message_history=message_history
)
await self.storage.update_task( # Build response message for A2A protocol
task["id"], response_message = Message(
state="completed", role="agent",
new_messages=[response_message], parts=[TextPart(kind="text", text=str(result.output))],
new_artifacts=artifacts, 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)
# Replace old state with new complete state
await self.storage.update_context(
task["context_id"], [state_message]
)
# Build rich artifacts with search results and answer
artifacts = self.build_artifacts(result)
await self.storage.update_task(
task["id"],
state="completed",
new_messages=[response_message],
new_artifacts=artifacts,
)
except Exception as e: except Exception as e:
logger.error( logger.error(
"Task execution failed: task_id=%s, question=%s, error=%s", "Task execution failed: task_id=%s, question=%s, error=%s",
@ -438,6 +469,27 @@ def create_a2a_app(db_path: Path):
await self.storage.update_task(task["id"], state="failed") await self.storage.update_task(task["id"], state="failed")
raise 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:
DeepQAAnswer with answer and sources
"""
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
async def cancel_task(self, params: TaskIdParams) -> None: async def cancel_task(self, params: TaskIdParams) -> None:
"""Cancel a task - not implemented for this worker.""" """Cancel a task - not implemented for this worker."""
pass pass