303 lines
10 KiB
Python
303 lines
10 KiB
Python
import asyncio
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from pydantic_ai import Agent, RunContext
|
|
from pydantic_ai.format_prompt import format_as_xml
|
|
from pydantic_ai.output import ToolOutput
|
|
from pydantic_graph import BaseNode, End, GraphRunContext
|
|
|
|
from haiku.rag.graph.common import get_model, log
|
|
from haiku.rag.graph.models import ResearchPlan, SearchAnswer
|
|
from haiku.rag.graph.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
|
|
from haiku.rag.qa.deep.dependencies import DeepQADependencies
|
|
from haiku.rag.qa.deep.models import DeepQAAnswer, DeepQAEvaluation
|
|
from haiku.rag.qa.deep.prompts import (
|
|
DECISION_PROMPT,
|
|
SYNTHESIS_PROMPT,
|
|
SYNTHESIS_PROMPT_WITH_CITATIONS,
|
|
)
|
|
from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
|
|
|
|
|
|
@dataclass
|
|
class DeepQAPlanNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
|
|
provider: str
|
|
model: str
|
|
|
|
async def run(
|
|
self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
|
|
) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]:
|
|
state = ctx.state
|
|
deps = ctx.deps
|
|
|
|
log(deps, state, "\n[bold cyan]📋 Planning approach...[/bold cyan]")
|
|
|
|
plan_agent = Agent(
|
|
model=get_model(self.provider, self.model),
|
|
output_type=ResearchPlan,
|
|
instructions=(
|
|
PLAN_PROMPT
|
|
+ "\n\nUse the gather_context tool once on the main question before planning."
|
|
),
|
|
retries=3,
|
|
deps_type=DeepQADependencies,
|
|
)
|
|
|
|
@plan_agent.tool
|
|
async def gather_context(
|
|
ctx2: RunContext[DeepQADependencies], query: str, limit: int = 6
|
|
) -> str:
|
|
results = await ctx2.deps.client.search(query, limit=limit)
|
|
expanded = await ctx2.deps.client.expand_context(results)
|
|
return "\n\n".join(chunk.content for chunk, _ in expanded)
|
|
|
|
prompt = (
|
|
"Plan a focused approach for answering the main question.\n\n"
|
|
f"Main question: {state.context.original_question}"
|
|
)
|
|
|
|
agent_deps = DeepQADependencies(
|
|
client=deps.client,
|
|
context=state.context,
|
|
console=deps.console,
|
|
)
|
|
plan_result = await plan_agent.run(prompt, deps=agent_deps)
|
|
state.context.sub_questions = list(plan_result.output.sub_questions)[
|
|
: state.max_sub_questions
|
|
]
|
|
|
|
log(deps, state, "\n[bold green]✅ Plan Created:[/bold green]")
|
|
log(
|
|
deps,
|
|
state,
|
|
f" [bold]Main Question:[/bold] {state.context.original_question}",
|
|
)
|
|
log(deps, state, " [bold]Sub-questions:[/bold]")
|
|
for i, sq in enumerate(state.context.sub_questions, 1):
|
|
log(deps, state, f" {i}. {sq}")
|
|
|
|
return DeepQASearchDispatchNode(self.provider, self.model)
|
|
|
|
|
|
@dataclass
|
|
class DeepQASearchDispatchNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
|
|
provider: str
|
|
model: str
|
|
|
|
async def run(
|
|
self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
|
|
) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]:
|
|
state = ctx.state
|
|
deps = ctx.deps
|
|
|
|
if not state.context.sub_questions:
|
|
return DeepQADecisionNode(self.provider, self.model)
|
|
|
|
# Take up to max_concurrency questions and answer them concurrently
|
|
take = max(1, state.max_concurrency)
|
|
batch: list[str] = []
|
|
while state.context.sub_questions and len(batch) < take:
|
|
batch.append(state.context.sub_questions.pop(0))
|
|
|
|
async def answer_one(sub_q: str) -> SearchAnswer | None:
|
|
log(
|
|
deps,
|
|
state,
|
|
f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}",
|
|
)
|
|
agent = Agent(
|
|
model=get_model(self.provider, self.model),
|
|
output_type=ToolOutput(SearchAnswer, max_retries=3),
|
|
instructions=SEARCH_AGENT_PROMPT,
|
|
retries=3,
|
|
deps_type=DeepQADependencies,
|
|
)
|
|
|
|
@agent.tool
|
|
async def search_and_answer(
|
|
ctx2: RunContext[DeepQADependencies], query: str, limit: int = 5
|
|
) -> str:
|
|
search_results = await ctx2.deps.client.search(query, limit=limit)
|
|
expanded = await ctx2.deps.client.expand_context(search_results)
|
|
|
|
entries: list[dict[str, Any]] = [
|
|
{
|
|
"text": chunk.content,
|
|
"score": score,
|
|
"document_uri": (
|
|
chunk.document_title or chunk.document_uri or ""
|
|
),
|
|
}
|
|
for chunk, score in expanded
|
|
]
|
|
if not entries:
|
|
return f"No relevant information found in the knowledge base for: {query}"
|
|
|
|
return format_as_xml(entries, root_tag="snippets")
|
|
|
|
agent_deps = DeepQADependencies(
|
|
client=deps.client,
|
|
context=state.context,
|
|
console=deps.console,
|
|
)
|
|
try:
|
|
result = await agent.run(sub_q, deps=agent_deps)
|
|
except Exception as e:
|
|
log(deps, state, f"[red]Search failed:[/red] {e}")
|
|
return None
|
|
|
|
return result.output
|
|
|
|
answers = await asyncio.gather(*(answer_one(q) for q in batch))
|
|
for ans in answers:
|
|
if ans is None:
|
|
continue
|
|
state.context.add_qa_response(ans)
|
|
preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "")
|
|
log(deps, state, f" [green]✓[/green] {preview}")
|
|
|
|
return DeepQASearchDispatchNode(self.provider, self.model)
|
|
|
|
|
|
@dataclass
|
|
class DeepQADecisionNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
|
|
provider: str
|
|
model: str
|
|
|
|
async def run(
|
|
self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
|
|
) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]:
|
|
state = ctx.state
|
|
deps = ctx.deps
|
|
|
|
log(
|
|
deps,
|
|
state,
|
|
"\n[bold cyan]📊 Evaluating information sufficiency...[/bold cyan]",
|
|
)
|
|
|
|
agent = Agent(
|
|
model=get_model(self.provider, self.model),
|
|
output_type=DeepQAEvaluation,
|
|
instructions=DECISION_PROMPT,
|
|
retries=3,
|
|
deps_type=DeepQADependencies,
|
|
)
|
|
|
|
context_data = {
|
|
"original_question": state.context.original_question,
|
|
"gathered_answers": [
|
|
{
|
|
"question": qa.query,
|
|
"answer": qa.answer,
|
|
"sources": qa.sources,
|
|
}
|
|
for qa in state.context.qa_responses
|
|
],
|
|
}
|
|
context_xml = format_as_xml(context_data, root_tag="gathered_information")
|
|
|
|
prompt = (
|
|
"Evaluate whether we have sufficient information to answer the question.\n\n"
|
|
f"{context_xml}"
|
|
)
|
|
|
|
agent_deps = DeepQADependencies(
|
|
client=deps.client,
|
|
context=state.context,
|
|
console=deps.console,
|
|
)
|
|
result = await agent.run(prompt, deps=agent_deps)
|
|
evaluation = result.output
|
|
|
|
state.iterations += 1
|
|
|
|
log(deps, state, f" [bold]Assessment:[/bold] {evaluation.reasoning}")
|
|
status = "[green]Yes[/green]" if evaluation.is_sufficient else "[red]No[/red]"
|
|
log(deps, state, f" Sufficient: {status}")
|
|
|
|
# Add new questions if not sufficient
|
|
for new_q in evaluation.new_questions:
|
|
if new_q not in state.context.sub_questions:
|
|
state.context.sub_questions.append(new_q)
|
|
|
|
if evaluation.new_questions:
|
|
log(deps, state, " [cyan]New questions:[/cyan]")
|
|
for question in evaluation.new_questions:
|
|
log(deps, state, f" • {question}")
|
|
|
|
# Decide next step
|
|
if evaluation.is_sufficient or state.iterations >= state.max_iterations:
|
|
if state.iterations >= state.max_iterations:
|
|
log(
|
|
deps,
|
|
state,
|
|
f"\n[bold yellow]⚠️ Reached max iterations ({state.max_iterations})[/bold yellow]",
|
|
)
|
|
log(deps, state, "\n[bold green]✅ Moving to synthesis.[/bold green]")
|
|
return DeepQASynthesizeNode(self.provider, self.model)
|
|
|
|
log(
|
|
deps,
|
|
state,
|
|
f"\n[bold cyan]🔄 Starting iteration {state.iterations + 1}...[/bold cyan]",
|
|
)
|
|
return DeepQASearchDispatchNode(self.provider, self.model)
|
|
|
|
|
|
@dataclass
|
|
class DeepQASynthesizeNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
|
|
provider: str
|
|
model: str
|
|
|
|
async def run(
|
|
self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
|
|
) -> End[DeepQAAnswer]:
|
|
state = ctx.state
|
|
deps = ctx.deps
|
|
|
|
log(
|
|
deps,
|
|
state,
|
|
"\n[bold cyan]📝 Synthesizing final answer...[/bold cyan]",
|
|
)
|
|
|
|
prompt_template = (
|
|
SYNTHESIS_PROMPT_WITH_CITATIONS
|
|
if state.context.use_citations
|
|
else SYNTHESIS_PROMPT
|
|
)
|
|
|
|
agent = Agent(
|
|
model=get_model(self.provider, self.model),
|
|
output_type=DeepQAAnswer,
|
|
instructions=prompt_template,
|
|
retries=3,
|
|
deps_type=DeepQADependencies,
|
|
)
|
|
|
|
context_data = {
|
|
"original_question": state.context.original_question,
|
|
"sub_answers": [
|
|
{
|
|
"question": qa.query,
|
|
"answer": qa.answer,
|
|
"sources": qa.sources,
|
|
}
|
|
for qa in state.context.qa_responses
|
|
],
|
|
}
|
|
context_xml = format_as_xml(context_data, root_tag="gathered_information")
|
|
|
|
prompt = f"Synthesize a comprehensive answer to the original question.\n\n{context_xml}"
|
|
|
|
agent_deps = DeepQADependencies(
|
|
client=deps.client,
|
|
context=state.context,
|
|
console=deps.console,
|
|
)
|
|
result = await agent.run(prompt, deps=agent_deps)
|
|
|
|
log(deps, state, "[bold green]✅ Answer complete![/bold green]")
|
|
return End(result.output)
|