548 lines
18 KiB
Python
548 lines
18 KiB
Python
import asyncio
|
|
|
|
from pydantic_ai import Agent, RunContext, format_as_xml
|
|
from pydantic_ai.output import ToolOutput
|
|
from pydantic_graph.beta import Graph, GraphBuilder, StepContext
|
|
from pydantic_graph.beta.join import reduce_list_append
|
|
|
|
from haiku.rag.agents.research.dependencies import ResearchContext, ResearchDependencies
|
|
from haiku.rag.agents.research.models import (
|
|
Citation,
|
|
ConversationalAnswer,
|
|
EvaluationResult,
|
|
RawSearchAnswer,
|
|
ResearchPlan,
|
|
ResearchReport,
|
|
SearchAnswer,
|
|
)
|
|
from haiku.rag.agents.research.prompts import (
|
|
CONVERSATIONAL_SYNTHESIS_PROMPT,
|
|
DECISION_PROMPT,
|
|
PLAN_PROMPT,
|
|
PLAN_PROMPT_WITH_CONTEXT,
|
|
SEARCH_PROMPT,
|
|
SYNTHESIS_PROMPT,
|
|
)
|
|
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
|
|
from haiku.rag.config import Config
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.utils import build_prompt, get_model
|
|
|
|
|
|
def format_context_for_prompt(context: ResearchContext) -> str:
|
|
"""Format the research context as XML for planning prompts."""
|
|
context_data: dict[str, object] = {}
|
|
|
|
if context.background_context:
|
|
context_data["background"] = context.background_context
|
|
|
|
context_data["question"] = context.original_question
|
|
|
|
if context.sub_questions:
|
|
context_data["pending_questions"] = context.sub_questions
|
|
|
|
if context.qa_responses:
|
|
context_data["prior_answers"] = [
|
|
{
|
|
"question": qa.query,
|
|
"answer": qa.answer,
|
|
"confidence": qa.confidence,
|
|
"source": qa.citations[0].document_title or qa.citations[0].document_uri
|
|
if qa.citations
|
|
else None,
|
|
}
|
|
for qa in context.qa_responses
|
|
]
|
|
|
|
return format_as_xml(context_data, root_tag="context")
|
|
|
|
|
|
def format_conversational_context_for_prompt(context: ResearchContext) -> str:
|
|
"""Format context for synthesis prompts."""
|
|
context_data: dict[str, object] = {}
|
|
|
|
if context.background_context:
|
|
context_data["background"] = context.background_context
|
|
|
|
context_data["question"] = context.original_question
|
|
|
|
if context.qa_responses:
|
|
context_data["prior_answers"] = [
|
|
{
|
|
"question": qa.query,
|
|
"answer": qa.answer,
|
|
"confidence": qa.confidence,
|
|
"source": qa.citations[0].document_title or qa.citations[0].document_uri
|
|
if qa.citations
|
|
else None,
|
|
}
|
|
for qa in context.qa_responses
|
|
]
|
|
|
|
return format_as_xml(context_data, root_tag="context")
|
|
|
|
|
|
# =============================================================================
|
|
# Shared step logic helpers
|
|
# =============================================================================
|
|
|
|
|
|
async def _plan_step_logic(
|
|
state: ResearchState,
|
|
deps: ResearchDeps,
|
|
config: AppConfig,
|
|
plan_prompt: str,
|
|
) -> None:
|
|
"""Shared logic for the plan step."""
|
|
model_config = config.research.model
|
|
|
|
# Use context-aware prompt if we have existing qa_responses
|
|
has_prior_answers = bool(state.context.qa_responses)
|
|
has_background = bool(state.context.background_context)
|
|
effective_plan_prompt = (
|
|
build_prompt(PLAN_PROMPT_WITH_CONTEXT, config)
|
|
if has_prior_answers
|
|
else plan_prompt
|
|
)
|
|
|
|
plan_agent: Agent[ResearchDependencies, ResearchPlan] = Agent(
|
|
model=get_model(model_config, config),
|
|
output_type=ResearchPlan,
|
|
instructions=effective_plan_prompt,
|
|
retries=3,
|
|
output_retries=3,
|
|
deps_type=ResearchDependencies,
|
|
)
|
|
|
|
search_filter = state.search_filter
|
|
|
|
@plan_agent.tool
|
|
async def gather_context(
|
|
ctx2: RunContext[ResearchDependencies],
|
|
query: str,
|
|
limit: int | None = None,
|
|
) -> str:
|
|
results = await ctx2.deps.client.search(
|
|
query, limit=limit, filter=search_filter
|
|
)
|
|
results = await ctx2.deps.client.expand_context(results)
|
|
return "\n\n".join(r.content for r in results)
|
|
|
|
# Build prompt with existing context if available
|
|
if has_prior_answers:
|
|
context_xml = format_context_for_prompt(state.context)
|
|
prompt = (
|
|
f"Review existing context and plan additional research if needed.\n\n"
|
|
f"{context_xml}\n\n"
|
|
f"Main question: {state.context.original_question}"
|
|
)
|
|
elif has_background:
|
|
context_xml = format_context_for_prompt(state.context)
|
|
prompt = (
|
|
f"Plan a focused approach for the main question.\n\n"
|
|
f"{context_xml}\n\n"
|
|
f"Main question: {state.context.original_question}"
|
|
)
|
|
else:
|
|
prompt = (
|
|
"Plan a focused approach for the main question.\n\n"
|
|
f"Main question: {state.context.original_question}"
|
|
)
|
|
|
|
agent_deps = ResearchDependencies(client=deps.client, context=state.context)
|
|
plan_result = await plan_agent.run(prompt, deps=agent_deps)
|
|
output = plan_result.output
|
|
state.context.sub_questions = list(output.sub_questions)
|
|
|
|
|
|
async def _search_one_step_logic(
|
|
state: ResearchState,
|
|
deps: ResearchDeps,
|
|
config: AppConfig,
|
|
search_prompt: str,
|
|
sub_q: str,
|
|
) -> SearchAnswer:
|
|
"""Shared logic for the search_one step."""
|
|
model_config = config.research.model
|
|
|
|
if deps.semaphore is None:
|
|
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
|
|
|
|
async with deps.semaphore:
|
|
agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent(
|
|
model=get_model(model_config, config),
|
|
output_type=ToolOutput(RawSearchAnswer, max_retries=3),
|
|
instructions=search_prompt,
|
|
retries=3,
|
|
deps_type=ResearchDependencies,
|
|
)
|
|
|
|
search_filter = state.search_filter
|
|
|
|
@agent.tool
|
|
async def search_and_answer(
|
|
ctx2: RunContext[ResearchDependencies],
|
|
query: str,
|
|
limit: int | None = None,
|
|
) -> str:
|
|
"""Search the knowledge base for relevant documents."""
|
|
results = await ctx2.deps.client.search(
|
|
query, limit=limit, filter=search_filter
|
|
)
|
|
results = await ctx2.deps.client.expand_context(results)
|
|
ctx2.deps.search_results = results
|
|
# Format with rank instead of raw score to avoid confusing LLMs
|
|
total = len(results)
|
|
parts = [
|
|
r.format_for_agent(rank=i + 1, total=total)
|
|
for i, r in enumerate(results)
|
|
]
|
|
if not parts:
|
|
return f"No relevant information found for: {query}"
|
|
return "\n\n".join(parts)
|
|
|
|
agent_deps = ResearchDependencies(client=deps.client, context=state.context)
|
|
|
|
result = await agent.run(sub_q, deps=agent_deps)
|
|
raw_answer = result.output
|
|
if raw_answer:
|
|
answer = SearchAnswer.from_raw(raw_answer, agent_deps.search_results)
|
|
state.context.add_qa_response(answer)
|
|
return answer
|
|
return SearchAnswer(query=sub_q, answer="", confidence=0.0)
|
|
|
|
|
|
def _get_batch_logic(state: ResearchState) -> list[str] | None:
|
|
"""Shared logic for the get_batch step."""
|
|
if not state.context.sub_questions:
|
|
return None
|
|
|
|
batch = list(state.context.sub_questions)
|
|
state.context.sub_questions.clear()
|
|
return batch
|
|
|
|
|
|
# =============================================================================
|
|
# Research graph (full version with decide loop)
|
|
# =============================================================================
|
|
|
|
|
|
def build_research_graph(
|
|
config: AppConfig = Config,
|
|
include_plan: bool = True,
|
|
) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]:
|
|
"""Build the Research graph.
|
|
|
|
Args:
|
|
config: AppConfig object (uses config.research for provider, model, and graph parameters)
|
|
include_plan: Whether to include the planning step (False for execute-only mode)
|
|
|
|
Returns:
|
|
Configured Research graph
|
|
"""
|
|
model_config = config.research.model
|
|
|
|
# Build prompts with system_context if configured
|
|
plan_prompt = build_prompt(
|
|
PLAN_PROMPT
|
|
+ "\n\nUse the gather_context tool once on the main question before planning.",
|
|
config,
|
|
)
|
|
search_prompt = build_prompt(SEARCH_PROMPT, config)
|
|
decision_prompt = build_prompt(DECISION_PROMPT, config)
|
|
synthesis_prompt = build_prompt(
|
|
config.prompts.synthesis or SYNTHESIS_PROMPT, config
|
|
)
|
|
g = GraphBuilder(
|
|
state_type=ResearchState,
|
|
deps_type=ResearchDeps,
|
|
output_type=ResearchReport,
|
|
)
|
|
|
|
@g.step
|
|
async def plan(ctx: StepContext[ResearchState, ResearchDeps, None]) -> None:
|
|
"""Create research plan with sub-questions."""
|
|
await _plan_step_logic(ctx.state, ctx.deps, config, plan_prompt)
|
|
|
|
@g.step
|
|
async def search_one(
|
|
ctx: StepContext[ResearchState, ResearchDeps, str],
|
|
) -> SearchAnswer:
|
|
"""Answer a single sub-question using the knowledge base."""
|
|
try:
|
|
return await _search_one_step_logic(
|
|
ctx.state, ctx.deps, config, search_prompt, ctx.inputs
|
|
)
|
|
except Exception as e:
|
|
return SearchAnswer(
|
|
query=ctx.inputs,
|
|
answer=f"Search failed: {str(e)}",
|
|
confidence=0.0,
|
|
)
|
|
|
|
@g.step
|
|
async def get_batch(
|
|
ctx: StepContext[ResearchState, ResearchDeps, None | bool | str],
|
|
) -> list[str] | None:
|
|
"""Get all remaining questions for this iteration."""
|
|
return _get_batch_logic(ctx.state)
|
|
|
|
@g.step
|
|
async def decide(
|
|
ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]],
|
|
) -> bool:
|
|
"""Evaluate research sufficiency and decide whether to continue."""
|
|
state = ctx.state
|
|
deps = ctx.deps
|
|
|
|
agent: Agent[ResearchDependencies, EvaluationResult] = Agent(
|
|
model=get_model(model_config, config),
|
|
output_type=EvaluationResult,
|
|
instructions=decision_prompt,
|
|
retries=3,
|
|
output_retries=3,
|
|
deps_type=ResearchDependencies,
|
|
)
|
|
|
|
context_xml = format_context_for_prompt(state.context)
|
|
prompt_parts = [
|
|
"Assess whether the research now answers the original question with adequate confidence.",
|
|
context_xml,
|
|
]
|
|
if state.last_eval is not None:
|
|
prev = state.last_eval
|
|
prompt_parts.append(
|
|
"<previous_evaluation>"
|
|
f"<confidence>{prev.confidence_score:.2f}</confidence>"
|
|
f"<is_sufficient>{str(prev.is_sufficient).lower()}</is_sufficient>"
|
|
f"<reasoning>{prev.reasoning}</reasoning>"
|
|
"</previous_evaluation>"
|
|
)
|
|
prompt = "\n\n".join(part for part in prompt_parts if part)
|
|
|
|
agent_deps = ResearchDependencies(
|
|
client=deps.client,
|
|
context=state.context,
|
|
)
|
|
decision_result = await agent.run(prompt, deps=agent_deps)
|
|
output = decision_result.output
|
|
|
|
state.last_eval = output
|
|
state.iterations += 1
|
|
|
|
# Get already-answered questions to avoid duplicates
|
|
answered_queries = {qa.query.lower() for qa in state.context.qa_responses}
|
|
|
|
for new_q in output.new_questions:
|
|
# Skip if already in pending or already answered
|
|
if new_q in state.context.sub_questions:
|
|
continue
|
|
if new_q.lower() in answered_queries:
|
|
continue
|
|
state.context.sub_questions.append(new_q)
|
|
|
|
should_continue = (
|
|
not output.is_sufficient
|
|
or output.confidence_score < state.confidence_threshold
|
|
) and state.iterations < state.max_iterations
|
|
|
|
return should_continue
|
|
|
|
@g.step
|
|
async def synthesize(
|
|
ctx: StepContext[ResearchState, ResearchDeps, None | bool | str],
|
|
) -> ResearchReport:
|
|
"""Generate final research report."""
|
|
state = ctx.state
|
|
deps = ctx.deps
|
|
|
|
agent: Agent[ResearchDependencies, ResearchReport] = Agent(
|
|
model=get_model(model_config, config),
|
|
output_type=ResearchReport,
|
|
instructions=synthesis_prompt,
|
|
retries=3,
|
|
output_retries=3,
|
|
deps_type=ResearchDependencies,
|
|
)
|
|
|
|
context_xml = format_context_for_prompt(state.context)
|
|
prompt = (
|
|
"Generate a comprehensive research report based on all gathered information.\n\n"
|
|
f"{context_xml}\n\n"
|
|
"Create a detailed report that synthesizes all findings into a coherent response."
|
|
)
|
|
agent_deps = ResearchDependencies(
|
|
client=deps.client,
|
|
context=state.context,
|
|
)
|
|
result = await agent.run(prompt, deps=agent_deps)
|
|
return result.output
|
|
|
|
# Build the graph structure
|
|
collect_answers = g.join(
|
|
reduce_list_append,
|
|
initial_factory=list[SearchAnswer],
|
|
)
|
|
|
|
if include_plan:
|
|
g.add(
|
|
g.edge_from(g.start_node).to(plan),
|
|
g.edge_from(plan).to(get_batch),
|
|
)
|
|
else:
|
|
g.add(g.edge_from(g.start_node).to(get_batch))
|
|
|
|
g.add(
|
|
g.edge_from(get_batch).to(
|
|
g.decision()
|
|
.branch(g.match(list).label("Has questions").map().to(search_one))
|
|
.branch(g.match(type(None)).label("No questions").to(synthesize))
|
|
),
|
|
g.edge_from(search_one).to(collect_answers),
|
|
g.edge_from(collect_answers).to(decide),
|
|
)
|
|
|
|
g.add(
|
|
g.edge_from(decide).to(
|
|
g.decision()
|
|
.branch(
|
|
g.match(bool, matches=lambda x: x)
|
|
.label("Continue research")
|
|
.to(get_batch)
|
|
)
|
|
.branch(
|
|
g.match(bool, matches=lambda x: not x)
|
|
.label("Done researching")
|
|
.to(synthesize)
|
|
)
|
|
),
|
|
g.edge_from(synthesize).to(g.end_node),
|
|
)
|
|
|
|
return g.build()
|
|
|
|
|
|
# =============================================================================
|
|
# Conversational graph (simplified, single iteration)
|
|
# =============================================================================
|
|
|
|
|
|
def build_conversational_graph(
|
|
config: AppConfig = Config,
|
|
) -> Graph[ResearchState, ResearchDeps, None, ConversationalAnswer]:
|
|
"""Build a simplified research graph for conversational chat.
|
|
|
|
This graph is optimized for single-iteration Q&A:
|
|
- Context-aware planning (generates fewer sub-questions when context exists)
|
|
- Single search iteration (no decide loop)
|
|
- Conversational output (direct answer, not formal report)
|
|
|
|
Args:
|
|
config: AppConfig object
|
|
|
|
Returns:
|
|
Graph that outputs ConversationalAnswer
|
|
"""
|
|
# Build prompts
|
|
plan_prompt = build_prompt(
|
|
PLAN_PROMPT
|
|
+ "\n\nUse the gather_context tool once on the main question before planning.",
|
|
config,
|
|
)
|
|
search_prompt = build_prompt(SEARCH_PROMPT, config)
|
|
conversational_prompt = build_prompt(CONVERSATIONAL_SYNTHESIS_PROMPT, config)
|
|
|
|
g = GraphBuilder(
|
|
state_type=ResearchState,
|
|
deps_type=ResearchDeps,
|
|
output_type=ConversationalAnswer,
|
|
)
|
|
|
|
@g.step
|
|
async def plan(ctx: StepContext[ResearchState, ResearchDeps, None]) -> None:
|
|
"""Create research plan with sub-questions."""
|
|
await _plan_step_logic(ctx.state, ctx.deps, config, plan_prompt)
|
|
|
|
@g.step
|
|
async def search_one(
|
|
ctx: StepContext[ResearchState, ResearchDeps, str],
|
|
) -> SearchAnswer:
|
|
"""Answer a single sub-question using the knowledge base."""
|
|
try:
|
|
return await _search_one_step_logic(
|
|
ctx.state, ctx.deps, config, search_prompt, ctx.inputs
|
|
)
|
|
except Exception as e:
|
|
return SearchAnswer(
|
|
query=ctx.inputs,
|
|
answer=f"Search failed: {str(e)}",
|
|
confidence=0.0,
|
|
)
|
|
|
|
@g.step
|
|
async def get_batch(
|
|
ctx: StepContext[ResearchState, ResearchDeps, None],
|
|
) -> list[str] | None:
|
|
"""Get all remaining questions for this iteration."""
|
|
return _get_batch_logic(ctx.state)
|
|
|
|
@g.step
|
|
async def synthesize(
|
|
ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer] | None],
|
|
) -> ConversationalAnswer:
|
|
"""Generate conversational answer from gathered evidence."""
|
|
state = ctx.state
|
|
deps = ctx.deps
|
|
|
|
agent: Agent[ResearchDependencies, ConversationalAnswer] = Agent(
|
|
model=get_model(config.research.model, config),
|
|
output_type=ConversationalAnswer,
|
|
instructions=conversational_prompt,
|
|
retries=3,
|
|
output_retries=3,
|
|
deps_type=ResearchDependencies,
|
|
)
|
|
|
|
context_xml = format_conversational_context_for_prompt(state.context)
|
|
prompt = f"Answer the question based on the gathered evidence.\n\n{context_xml}"
|
|
agent_deps = ResearchDependencies(
|
|
client=deps.client,
|
|
context=state.context,
|
|
)
|
|
result = await agent.run(prompt, deps=agent_deps)
|
|
|
|
# Collect unique citations from qa_responses (dedupe by chunk_id)
|
|
seen_chunks: set[str] = set()
|
|
unique_citations: list[Citation] = []
|
|
for qa in state.context.qa_responses:
|
|
for c in qa.citations:
|
|
if c.chunk_id not in seen_chunks:
|
|
seen_chunks.add(c.chunk_id)
|
|
unique_citations.append(c)
|
|
|
|
return ConversationalAnswer(
|
|
answer=result.output.answer,
|
|
citations=unique_citations,
|
|
confidence=result.output.confidence,
|
|
)
|
|
|
|
# Build the graph structure (simplified: plan → search → synthesize)
|
|
collect_answers = g.join(
|
|
reduce_list_append,
|
|
initial_factory=list[SearchAnswer],
|
|
)
|
|
|
|
g.add(
|
|
g.edge_from(g.start_node).to(plan),
|
|
g.edge_from(plan).to(get_batch),
|
|
g.edge_from(get_batch).to(
|
|
g.decision()
|
|
.branch(g.match(list).label("Has questions").map().to(search_one))
|
|
.branch(g.match(type(None)).label("No questions").to(synthesize))
|
|
),
|
|
g.edge_from(search_one).to(collect_answers),
|
|
g.edge_from(collect_answers).to(synthesize),
|
|
g.edge_from(synthesize).to(g.end_node),
|
|
)
|
|
|
|
return g.build()
|