Add deep QA to a2a agent
This commit is contained in:
parent
1a88b0355a
commit
9aad6b65b8
1 changed files with 86 additions and 34 deletions
|
|
@ -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,13 +390,41 @@ 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
|
async with HaikuRAG(db_path) as client:
|
||||||
context = await self.storage.load_context(task["context_id"]) or []
|
if skill == "deep-qa":
|
||||||
# Load conversation history
|
# Run deep QA graph
|
||||||
|
deep_result = await self.run_deep_qa(client, question)
|
||||||
|
|
||||||
|
# Build response message
|
||||||
|
response_message = Message(
|
||||||
|
role="agent",
|
||||||
|
parts=[TextPart(kind="text", text=deep_result.answer)],
|
||||||
|
kind="message",
|
||||||
|
message_id=str(uuid.uuid4()),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Build artifacts (basic for now, will be enhanced in commit 3)
|
||||||
|
artifacts = [
|
||||||
|
Artifact(
|
||||||
|
artifact_id=str(uuid.uuid4()),
|
||||||
|
name="answer",
|
||||||
|
parts=[TextPart(kind="text", text=deep_result.answer)],
|
||||||
|
)
|
||||||
|
]
|
||||||
|
|
||||||
|
await self.storage.update_task(
|
||||||
|
task["id"],
|
||||||
|
state="completed",
|
||||||
|
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)
|
message_history = load_message_history(context)
|
||||||
|
|
||||||
# Create fresh client for this task and run agent
|
|
||||||
async with HaikuRAG(db_path) as client:
|
|
||||||
deps = AgentDependencies(client=client)
|
deps = AgentDependencies(client=client)
|
||||||
|
|
||||||
# Run agent with full conversation history including tool calls
|
# Run agent with full conversation history including tool calls
|
||||||
|
|
@ -409,7 +441,6 @@ def create_a2a_app(db_path: Path):
|
||||||
)
|
)
|
||||||
|
|
||||||
# Update context with complete conversation state
|
# Update context with complete conversation state
|
||||||
# Store all messages from this run (includes tool calls & results)
|
|
||||||
updated_history = message_history + result.new_messages()
|
updated_history = message_history + result.new_messages()
|
||||||
state_message = save_message_history(updated_history)
|
state_message = save_message_history(updated_history)
|
||||||
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue