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

Merge pull request #267 from ggozad/feat/research-graph-improvements

Browse source 
Performance improvements for the research graph & chat agents
This commit is contained in:
Yiorgis Gozadinos 2026-01-31 11:11:40 +02:00 committed by GitHub
commit 5cdf079322
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
37 changed files with 5126 additions and 17626 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
CHANGELOG.md
View file

@ -1,6 +1,21 @@
# Changelog
## [Unreleased]
### Changed
- **Iterative Research Planning**: Research graph now uses an iterative feedback loop instead of batch question processing
- Planner proposes ONE question at a time, sees the answer, then decides whether to continue
- Removes `gather_context` tool — planner proposes questions directly
- Simpler flow: `plan_next` → `search_one` → loop back until complete → `synthesize`
- Consolidated `build_conversational_graph()` into `build_research_graph(output_mode="conversational")`
### Removed
- **Dead config options**: Removed vestigial fields from iterative planning refactor
- `confidence_threshold` from `ResearchConfig` and `ResearchState` (LLM decides completion via `is_complete`)
- `max_sub_questions` from `QAConfig` (iterative flow uses one question at a time)
- `sub_questions` field from `ResearchContext` (no longer populated)
## [0.27.2] - 2026-01-29
### Added

36
docs/agents.md
View file

@ -163,44 +163,39 @@ Frontend clients should extract state from under this key. See the [Web Applicat
## Research Graph
The research workflow is implemented as a typed pydantic-graph. It plans, searches (in parallel batches), evaluates, and synthesizes into a final report.
The research workflow is implemented as a typed pydantic-graph. It uses an iterative feedback loop where the planner proposes one question at a time, sees the answer, then decides whether to continue or synthesize.
```mermaid
---
title: Research graph
---
stateDiagram-v2
[*] --> plan
plan --> get_batch
get_batch --> search_one: Has questions (map)
get_batch --> synthesize: No questions
search_one --> collect_answers
collect_answers --> decide
decide --> get_batch: Continue research
decide --> synthesize: Done researching
[*] --> plan_next
plan_next --> search_one: Has next question
plan_next --> synthesize: Complete or max iterations
search_one --> plan_next
synthesize --> [*]
```
**Key nodes:**
- **plan**: Builds up to 3 standalone sub-questions (uses an internal presearch tool)
- **get_batch**: Retrieves remaining sub-questions for the current iteration
- **search_one**: Answers a single sub-question using the KB (mapped in parallel)
- **collect_answers**: Aggregates search results from parallel executions
- **decide**: Evaluates confidence and determines whether to continue or synthesize
- **plan_next**: Evaluates gathered evidence and either proposes the next question to investigate or marks research as complete
- **search_one**: Answers a single question using the knowledge base
- **synthesize**: Generates a final structured research report
**Primary models:**
- `SearchAnswer` — one per sub-question (query, answer, confidence, citations)
- `EvaluationResult` — confidence score, new questions, sufficiency assessment
- `IterativePlanResult` — planning decision (is_complete, next_question, reasoning)
- `SearchAnswer` — answer to a single question (query, answer, confidence, citations)
- `ResearchReport` — final report (title, executive summary, findings, conclusions, …)
- `ConversationalAnswer` — alternative output for chat integration (answer, citations, confidence)
**Parallel execution:**
**Iterative flow:**
- The `search_one` node is mapped over all questions in a batch
- Parallelism is controlled via `max_concurrency`
- Decision nodes process results after each batch completes
- Each iteration: planner evaluates context → proposes one question → search answers it → loop back
- Planner can decompose complex questions (e.g., "benefits and drawbacks" → start with "benefits")
- Session context is used to resolve ambiguous references and inform planning
- Loop terminates when planner marks `is_complete=True` or `max_iterations` is reached
### CLI Usage
@ -249,7 +244,6 @@ custom_config = AppConfig(
provider="openai",
model="gpt-4o-mini",
max_iterations=5,
confidence_threshold=0.85,
max_concurrency=3,
)
)

2
docs/cli.md
View file

@ -255,7 +255,7 @@ Flags:
- `--context`: Background context for the research
- `--context-file`: Path to a file containing background context
Research parameters like `max_iterations`, `confidence_threshold`, and `max_concurrency` are configured in your [configuration file](configuration/index.md) under the `research` section.
Research parameters like `max_iterations` and `max_concurrency` are configured in your [configuration file](configuration/index.md) under the `research` section.
## Server

2
docs/configuration/index.md
View file

@ -85,7 +85,6 @@ qa:
provider: ollama
name: gpt-oss
enable_thinking: false
max_sub_questions: 3
max_iterations: 2
max_concurrency: 1
@ -95,7 +94,6 @@ research:
name: ""
enable_thinking: false
max_iterations: 3
confidence_threshold: 0.8
max_concurrency: 1
search:

20
docs/configuration/qa-research.md
View file

@ -32,17 +32,15 @@ qa:
provider: ollama
name: gpt-oss
enable_thinking: false
max_sub_questions: 3 # Maximum sub-questions for deep QA
max_iterations: 2 # Maximum search iterations per sub-question
max_concurrency: 1 # Sub-questions processed in parallel
max_iterations: 2 # Maximum search iterations
max_concurrency: 1 # Concurrent search operations
```
- **model**: LLM configuration (see [Providers](providers.md#model-settings))
- **max_sub_questions**: For deep QA mode, maximum number of sub-questions to generate (default: 3)
- **max_iterations**: Maximum search/evaluate cycles per sub-question (default: 2)
- **max_concurrency**: Number of sub-questions to process in parallel (default: 1)
- **max_iterations**: Maximum search iterations (default: 2)
- **max_concurrency**: Number of concurrent search operations (default: 1)
Deep QA mode (`haiku-rag ask --deep`) decomposes complex questions into sub-questions, processes them in parallel batches, and synthesizes the results.
Deep QA mode (`haiku-rag ask --deep`) uses the research graph with a single iteration for quick, focused answers.
## Research Configuration
@ -55,13 +53,11 @@ research:
name: "" # Empty to use qa model
enable_thinking: false
max_iterations: 3
confidence_threshold: 0.8
max_concurrency: 1
```
- **model**: LLM configuration. Leave provider/model empty to inherit from `qa` (see [Providers](providers.md#model-settings))
- **max_iterations**: Maximum search/evaluate cycles (default: 3)
- **confidence_threshold**: Stop when confidence score meets/exceeds this (default: 0.8)
- **max_concurrency**: Sub-questions searched in parallel per iteration (default: 1)
- **max_iterations**: Maximum planning/search iterations (default: 3)
- **max_concurrency**: Concurrent search operations (default: 1)
The research workflow plans sub-questions, searches in parallel batches, evaluates findings, and iterates until reaching the confidence threshold or max iterations.
The research workflow uses an iterative feedback loop: the planner proposes one question at a time, sees the answer, then decides whether to continue or synthesize. This continues until the planner marks research as complete or `max_iterations` is reached.

7
evaluations/evaluations/benchmark.py
View file

@ -308,12 +308,7 @@ async def run_qa_benchmark(
async def answer_question(question: str) -> str:
context = ResearchContext(original_question=question)
state = ResearchState.from_config(
context=context,
config=config,
max_iterations=2,
confidence_threshold=0.0,
)
state = ResearchState.from_config(context=context, config=config)
deps = ResearchDeps(client=rag)
report = await graph.run(state=state, deps=deps)
return report.executive_summary if report else ""

10
haiku_rag_slim/haiku/rag/agents/__init__.py
View file

@ -9,16 +9,13 @@ from haiku.rag.agents.chat import (
from haiku.rag.agents.qa import QuestionAnswerAgent, get_qa_agent
from haiku.rag.agents.research import (
Citation,
EvaluationResult,
IterativePlanResult,
ResearchContext,
ResearchDependencies,
ResearchReport,
SearchAnswer,
)
from haiku.rag.agents.research.graph import (
build_conversational_graph,
build_research_graph,
)
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
__all__ = [
@ -27,7 +24,6 @@ __all__ = [
"QuestionAnswerAgent",
# Research
"build_research_graph",
"build_conversational_graph",
"ResearchContext",
"ResearchDependencies",
"ResearchDeps",
@ -35,7 +31,7 @@ __all__ = [
"ResearchReport",
"Citation",
"SearchAnswer",
"EvaluationResult",
"IterativePlanResult",
# Chat
"create_chat_agent",
"SearchAgent",

7
haiku_rag_slim/haiku/rag/agents/chat/agent.py
View file

@ -23,7 +23,7 @@ from haiku.rag.agents.chat.state import (
emit_state_event,
)
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_conversational_graph
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.models import Citation
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
from haiku.rag.client import HaikuRAG
@ -197,7 +197,9 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
doc_filter = combine_filters(session_filter, tool_filter)
# Build and run the conversational research graph
graph = build_conversational_graph(config=ctx.deps.config)
graph = build_research_graph(
config=ctx.deps.config, output_mode="conversational"
)
session_id = ctx.deps.session_state.session_id
# Get session context from server cache for planning, fallback to initial_context
@ -247,7 +249,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
state = ResearchState(
context=context,
max_iterations=1,
confidence_threshold=0.0,
search_filter=doc_filter,
max_concurrency=ctx.deps.config.research.max_concurrency,
)

2
haiku_rag_slim/haiku/rag/agents/research/__init__.py
View file

@ -1,7 +1,7 @@
from haiku.rag.agents.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.agents.research.models import (
Citation,
EvaluationResult,
IterativePlanResult,
ResearchReport,
SearchAnswer,
)

3
haiku_rag_slim/haiku/rag/agents/research/dependencies.py
View file

@ -13,9 +13,6 @@ class ResearchContext(BaseModel):
"""Context shared across research agents."""
original_question: str = Field(description="The original research question")
sub_questions: list[str] = Field(
default_factory=list, description="Decomposed sub-questions"
)
qa_responses: list[Any] = Field(
default_factory=list, description="Structured QA pairs used during research"
)

533
haiku_rag_slim/haiku/rag/agents/research/graph.py
View file

@ -1,25 +1,23 @@
import asyncio
from typing import Literal, overload
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,
IterativePlanResult,
RawSearchAnswer,
ResearchPlan,
ResearchReport,
SearchAnswer,
)
from haiku.rag.agents.research.prompts import (
CONVERSATIONAL_SYNTHESIS_PROMPT,
DECISION_PROMPT,
PLAN_PROMPT,
PLAN_PROMPT_WITH_CONTEXT,
ITERATIVE_PLAN_PROMPT,
ITERATIVE_PLAN_PROMPT_WITH_CONTEXT,
SEARCH_PROMPT,
SYNTHESIS_PROMPT,
)
@ -29,17 +27,8 @@ from haiku.rag.config.models import AppConfig
from haiku.rag.utils import build_prompt, get_model
def format_context_for_prompt(
context: ResearchContext,
include_pending_questions: bool = True,
) -> str:
"""Format the research context as XML for prompts.
Args:
context: The research context to format.
include_pending_questions: Whether to include pending sub-questions.
Set to False for synthesis prompts where pending questions aren't relevant.
"""
def format_context_for_prompt(context: ResearchContext) -> str:
"""Format the research context as XML for prompts."""
context_data: dict[str, object] = {}
if context.session_context:
@ -47,9 +36,6 @@ def format_context_for_prompt(
context_data["question"] = context.original_question
if include_pending_questions and context.sub_questions:
context_data["pending_questions"] = context.sub_questions
if context.qa_responses:
context_data["prior_answers"] = [
{
@ -64,80 +50,63 @@ def format_context_for_prompt(
return format_as_xml(context_data, root_tag="context")
# =============================================================================
# Shared step logic helpers
# =============================================================================
async def _plan_step_logic(
async def _iterative_plan_logic(
state: ResearchState,
deps: ResearchDeps,
config: AppConfig,
plan_prompt: str,
) -> None:
"""Shared logic for the plan step."""
) -> IterativePlanResult:
"""Evaluate context and decide next question or mark complete."""
has_prior_answers = bool(state.context.qa_responses)
# If max iterations reached, skip LLM and mark complete
if state.iterations >= state.max_iterations:
return IterativePlanResult(
is_complete=True,
next_question=None,
reasoning=f"Max iterations ({state.max_iterations}) reached.",
)
model_config = config.research.model
# Use context-aware prompt if we have existing qa_responses or session_context
has_prior_answers = bool(state.context.qa_responses)
has_session_context = bool(state.context.session_context)
effective_plan_prompt = (
build_prompt(PLAN_PROMPT_WITH_CONTEXT, config)
if has_prior_answers or has_session_context
else plan_prompt
)
if has_prior_answers:
effective_prompt = build_prompt(ITERATIVE_PLAN_PROMPT_WITH_CONTEXT, config)
else:
effective_prompt = build_prompt(ITERATIVE_PLAN_PROMPT, config)
plan_agent: Agent[ResearchDependencies, ResearchPlan] = Agent( # type: ignore[invalid-assignment]
plan_agent: Agent[ResearchDependencies, IterativePlanResult] = Agent( # type: ignore[assignment]
model=get_model(model_config, config),
output_type=ResearchPlan,
instructions=effective_plan_prompt,
output_type=IterativePlanResult,
instructions=effective_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
search_filter = state.search_filter
# Only register gather_context tool when we don't have existing context
if not has_prior_answers and not has_session_context:
@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
# Build prompt based on current state
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_session_context:
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}"
f"Review the gathered evidence and decide whether to continue or synthesize.\n\n"
f"{context_xml}"
)
else:
prompt = (
"Plan a focused approach for the main question.\n\n"
f"Main question: {state.context.original_question}"
)
context_xml = format_context_for_prompt(state.context)
prompt = f"Plan the research investigation.\n\n{context_xml}"
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)
result = await plan_agent.run(prompt, deps=agent_deps)
# Enforce: if no prior answers, must have a next_question to investigate
if not has_prior_answers:
if result.output.is_complete or not result.output.next_question:
return IterativePlanResult(
is_complete=False,
next_question=result.output.next_question
or state.context.original_question,
reasoning=result.output.reasoning,
)
return result.output
async def _search_one_step_logic(
@ -147,14 +116,14 @@ async def _search_one_step_logic(
search_prompt: str,
sub_q: str,
) -> SearchAnswer:
"""Shared logic for the search_one step."""
"""Answer a single question using the knowledge base."""
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( # type: ignore[invalid-assignment]
agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent( # type: ignore[assignment]
model=get_model(model_config, config),
output_type=ToolOutput(RawSearchAnswer, max_retries=3),
instructions=search_prompt,
@ -176,7 +145,6 @@ async def _search_one_step_logic(
)
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)
@ -190,6 +158,10 @@ async def _search_one_step_logic(
result = await agent.run(sub_q, deps=agent_deps)
raw_answer = result.output
# Increment iterations after each search completes
state.iterations += 1
if raw_answer:
answer = SearchAnswer.from_raw(raw_answer, agent_deps.search_results)
state.context.add_qa_response(answer)
@ -197,262 +169,62 @@ async def _search_one_step_logic(
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
@overload
def build_research_graph(
config: AppConfig = ...,
output_mode: Literal["report"] = ...,
) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]: ...
# =============================================================================
# Research graph (full version with decide loop)
# =============================================================================
@overload
def build_research_graph(
config: AppConfig = ...,
output_mode: Literal["conversational"] = ...,
) -> Graph[ResearchState, ResearchDeps, None, ConversationalAnswer]: ...
def build_research_graph(
config: AppConfig = Config,
include_plan: bool = True,
) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]:
"""Build the Research graph.
output_mode: Literal["report", "conversational"] = "report",
) -> Graph[ResearchState, ResearchDeps, None, ResearchReport | ConversationalAnswer]:
"""Build the iterative 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)
output_mode: Output format - "report" for ResearchReport, "conversational" for ConversationalAnswer
Returns:
Configured Research graph
Configured research graph with iterative planning
"""
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( # type: ignore[invalid-assignment]
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( # type: ignore[invalid-assignment]
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),
if output_mode == "report":
synthesis_prompt = build_prompt(
config.prompts.synthesis or SYNTHESIS_PROMPT, config
)
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)
synthesis_prompt = build_prompt(CONVERSATIONAL_SYNTHESIS_PROMPT, config)
g = GraphBuilder(
state_type=ResearchState,
deps_type=ResearchDeps,
output_type=ConversationalAnswer,
output_type=ResearchReport if output_mode == "report" else 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)
async def plan_next(
ctx: StepContext[ResearchState, ResearchDeps, None | SearchAnswer],
) -> IterativePlanResult:
"""Evaluate context and decide next question or complete."""
return await _iterative_plan_logic(ctx.state, ctx.deps, config)
@g.step
async def search_one(
ctx: StepContext[ResearchState, ResearchDeps, str],
) -> SearchAnswer:
"""Answer a single sub-question using the knowledge base."""
"""Answer a single question using the knowledge base."""
try:
return await _search_one_step_logic(
ctx.state, ctx.deps, config, search_prompt, ctx.inputs
@ -464,71 +236,114 @@ def build_conversational_graph(
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)
if output_mode == "report":
@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
@g.step
async def synthesize(
ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult],
) -> ResearchReport:
"""Generate final research report."""
state = ctx.state
deps = ctx.deps
agent: Agent[ResearchDependencies, ConversationalAnswer] = Agent( # type: ignore[invalid-assignment]
model=get_model(config.research.model, config),
output_type=ConversationalAnswer,
instructions=conversational_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
agent: Agent[ResearchDependencies, ResearchReport] = Agent( # type: ignore[assignment]
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, include_pending_questions=False
)
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)
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
# 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)
else:
return ConversationalAnswer(
answer=result.output.answer,
citations=unique_citations,
confidence=result.output.confidence,
)
@g.step
async def synthesize(
ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult],
) -> ConversationalAnswer:
"""Generate conversational answer from gathered evidence."""
state = ctx.state
deps = ctx.deps
# Build the graph structure (simplified: plan → search → synthesize)
collect_answers = g.join(
reduce_list_append,
initial_factory=list[SearchAnswer],
)
agent: Agent[ResearchDependencies, ConversationalAnswer] = Agent( # type: ignore[assignment]
model=get_model(model_config, config),
output_type=ConversationalAnswer,
instructions=synthesis_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_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 graph edges: iterative loop
#
# START -> plan_next -> [decision]
# |
# [is_complete or max_iterations] -> synthesize -> END
# |
# [has next_question] -> search_one -> plan_next (loop)
def extract_question(
ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult],
) -> str:
"""Extract next_question from IterativePlanResult."""
return ctx.inputs.next_question or ""
g.add(
g.edge_from(g.start_node).to(plan),
g.edge_from(plan).to(get_batch),
g.edge_from(get_batch).to(
g.edge_from(g.start_node).to(plan_next),
g.edge_from(plan_next).to(
g.decision()
.branch(g.match(list).label("Has questions").map().to(search_one))
.branch(g.match(type(None)).label("No questions").to(synthesize))
.branch(
g.match(
IterativePlanResult,
matches=lambda r: not r.is_complete and r.next_question is not None,
)
.label("Continue research")
.transform(extract_question)
.to(search_one)
)
.branch(
g.match(IterativePlanResult).label("Done researching").to(synthesize)
)
),
g.edge_from(search_one).to(collect_answers),
g.edge_from(collect_answers).to(synthesize),
g.edge_from(search_one).to(plan_next),
g.edge_from(synthesize).to(g.end_node),
)

43
haiku_rag_slim/haiku/rag/agents/research/models.py
View file

@ -1,25 +1,21 @@
from typing import TYPE_CHECKING
from pydantic import BaseModel, Field, field_validator
from pydantic import BaseModel, Field
if TYPE_CHECKING:
from haiku.rag.store.models import SearchResult
class ResearchPlan(BaseModel):
"""A structured research plan with sub-questions to explore."""
class IterativePlanResult(BaseModel):
"""Output from iterative planning step."""
sub_questions: list[str] = Field(
...,
description="Specific questions to research, phrased as complete questions",
is_complete: bool = Field(
description="Whether research is complete and can be synthesized"
)
@field_validator("sub_questions")
@classmethod
def validate_sub_questions(cls, v: list[str]) -> list[str]:
if len(v) > 12:
raise ValueError("Cannot have more than 12 sub-questions")
return v
next_question: str | None = Field(
default=None, description="Next question to investigate, if not complete"
)
reasoning: str = Field(description="Brief explanation of the decision")
class Citation(BaseModel):
@ -115,27 +111,6 @@ def resolve_citations(
return citations
class EvaluationResult(BaseModel):
"""Result of research sufficiency evaluation."""
is_sufficient: bool = Field(
description="Whether the research is sufficient to answer the original question"
)
confidence_score: float = Field(
ge=0.0,
le=1.0,
description="Confidence level in the completeness of research (0-1)",
)
reasoning: str = Field(
description="Explanation of why the research is or isn't complete"
)
new_questions: list[str] = Field(
default_factory=list,
max_length=3,
description="New sub-questions to add to the research (max 3)",
)
class ConversationalAnswer(BaseModel):
"""Conversational answer for chat context."""

85
haiku_rag_slim/haiku/rag/agents/research/prompts.py
View file

@ -1,47 +1,47 @@
PLAN_PROMPT = """You are the research orchestrator for a focused workflow.
ITERATIVE_PLAN_PROMPT = """You are the research orchestrator planning the investigation.
If a <background> section is provided, use it to understand the domain context.
If a <background> section is provided, use it to understand the conversation context.
Responsibilities:
1. Understand and decompose the main question
2. Propose a minimal, high-leverage plan
3. Coordinate specialized agents to gather evidence
Your task:
1. Analyze the original question
2. Propose the first question to investigate
Plan requirements:
- Produce at most 3 sub_questions that together cover the main question.
- sub_questions must be a list of plain strings, where each string is a complete
question. Do NOT use objects with nested fields like {question, details}.
- Each sub_question must be a standalone, self-contained query that can run
without extra context. Include concrete entities, scope, timeframe, and any
qualifiers. Avoid ambiguous pronouns (it/they/this/that).
- Prioritize the highest-value aspects first; avoid redundancy and overlap.
- Prefer questions that are likely answerable from the current knowledge base;
if coverage is uncertain, make scopes narrower and specific.
- Order sub_questions by execution priority (most valuable first).
For simple questions, investigate them directly. For composite or complex questions,
you may decompose into a focused sub-question. For example:
- "What are the benefits and drawbacks of X?" → Start with "What are the benefits of X?"
- Ambiguous references should be resolved using background context if available
Use the gather_context tool once on the main question before planning."""
Output requirements:
- Set is_complete=False (you are just starting the investigation)
- Set next_question to the question to investigate
- Provide brief reasoning explaining your choice
PLAN_PROMPT_WITH_CONTEXT = """You are the research orchestrator for a focused workflow.
The question must be standalone and self-contained:
- Include concrete entities, scope, and any qualifiers
- Avoid ambiguous pronouns (it/they/this/that)"""
ITERATIVE_PLAN_PROMPT_WITH_CONTEXT = """You are the research orchestrator evaluating gathered evidence.
You have access to context that may include:
- <background>: Domain context for the conversation
- <prior_answers>: Previous Q&A pairs with confidence scores
Review the provided context first. Use <background> to understand the domain.
If <prior_answers> exist and already answer the question completely,
return an empty sub_questions list. Only create sub-questions to fill gaps.
Your task:
1. Review the provided evidence carefully
2. Assess whether it sufficiently answers the original question
3. Decide whether to continue research or synthesize
Responsibilities:
1. Review provided context to understand what's already known
2. Identify gaps that need additional research
3. Propose minimal sub-questions only for missing information
Decision criteria:
- Set is_complete=True if the evidence adequately answers the question
- Set is_complete=False with a next_question if important gaps remain
Plan requirements:
- If existing context fully answers the question, return an empty sub_questions list.
- Only create new sub-questions for genuine gaps in existing knowledge.
- sub_questions must be a list of plain strings (max 3).
- Each sub_question must be standalone and self-contained.
- Prioritize the highest-value gaps first."""
If not complete, propose exactly ONE high-value follow-up question in next_question:
- Focus on the most critical gap not covered by prior_answers
- The question must be standalone and self-contained
- Avoid repeating questions that have already been answered
- Include concrete entities, scope, and any qualifiers
Provide brief reasoning explaining your decision."""
SEARCH_PROMPT = """You are a search and question-answering specialist.
@ -87,27 +87,6 @@ Guidelines:
- Be concise and direct; avoid meta commentary about the process.
- Results are ordered by relevance, with rank 1 being most relevant."""
DECISION_PROMPT = """You are the research evaluator responsible for assessing
whether gathered evidence sufficiently answers the research question.
Inputs available:
- Original research question
- Question-answer pairs with supporting sources
- Previous evaluation (if any)
Tasks:
1. Assess whether the collected evidence answers the original question.
2. Provide a confidence_score in [0,1] reflecting coverage and evidence quality.
3. Optionally propose up to 3 new sub-questions if important gaps remain.
Output fields:
- is_sufficient: true when the question is adequately answered
- confidence_score: numeric in [0,1]
- reasoning: brief explanation of the assessment
- new_questions: list of follow-up questions (max 3), only if needed
Be strict: only mark sufficient when key aspects are addressed with reliable evidence."""
SYNTHESIS_PROMPT = """You are a synthesis specialist producing the final
research report that directly answers the original question.

12
haiku_rag_slim/haiku/rag/agents/research/state.py
View file

@ -5,7 +5,6 @@ from typing import TYPE_CHECKING
from pydantic import BaseModel, Field
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.models import EvaluationResult
from haiku.rag.client import HaikuRAG
if TYPE_CHECKING:
@ -30,15 +29,9 @@ class ResearchState(BaseModel):
)
iterations: int = Field(default=0, description="Current iteration number")
max_iterations: int = Field(default=3, description="Maximum allowed iterations")
confidence_threshold: float = Field(
default=0.8, description="Confidence threshold for completion", ge=0.0, le=1.0
)
max_concurrency: int = Field(
default=1, description="Maximum concurrent search operations", ge=1
)
last_eval: EvaluationResult | None = Field(
default=None, description="Last evaluation result"
)
search_filter: str | None = Field(
default=None, description="SQL WHERE clause to filter search results"
)
@ -49,7 +42,6 @@ class ResearchState(BaseModel):
context: ResearchContext,
config: "AppConfig",
max_iterations: int | None = None,
confidence_threshold: float | None = None,
) -> "ResearchState":
"""Create a ResearchState from an AppConfig.
@ -57,15 +49,11 @@ class ResearchState(BaseModel):
context: The ResearchContext containing the question
config: The AppConfig object
max_iterations: Override max iterations (None uses config default)
confidence_threshold: Override threshold (None uses config, 0.0 disables check)
"""
return cls(
context=context,
max_iterations=max_iterations
if max_iterations is not None
else config.research.max_iterations,
confidence_threshold=confidence_threshold
if confidence_threshold is not None
else config.research.confidence_threshold,
max_concurrency=config.research.max_concurrency,
)

9
haiku_rag_slim/haiku/rag/app.py
View file

@ -398,8 +398,7 @@ class HaikuRAGApp:
state = ResearchState.from_config(
context=context,
config=self.config,
max_iterations=2,
confidence_threshold=0.0,
max_iterations=1,
)
state.search_filter = filter
deps = ResearchDeps(client=self.client)
@ -477,12 +476,6 @@ class HaikuRAGApp:
self.console.print(report.executive_summary)
self.console.print()
# Confidence (from last evaluation)
if state.last_eval:
conf = state.last_eval.confidence_score
self.console.print(f"[bold cyan]Confidence:[/bold cyan] {conf:.1%}")
self.console.print()
# Main Findings
if report.main_findings:
self.console.print("[bold cyan]Main Findings:[/bold cyan]")

2
haiku_rag_slim/haiku/rag/config/models.py
View file

@ -78,7 +78,6 @@ class QAConfig(BaseModel):
enable_thinking=False,
)
)
max_sub_questions: int = 3
max_iterations: int = 2
max_concurrency: int = 1
@ -92,7 +91,6 @@ class ResearchConfig(BaseModel):
)
)
max_iterations: int = 3
confidence_threshold: float = 0.8
max_concurrency: int = 1

1
haiku_rag_slim/haiku/rag/mcp.py
View file

@ -199,7 +199,6 @@ def create_mcp_server(
context=context,
config=config,
max_iterations=2,
confidence_threshold=0.0,
)
deps = ResearchDeps(client=rag)

59
tests/agents/research/test_models.py
View file

@ -120,62 +120,3 @@ class TestSearchAnswerPrimarySource:
citations=[],
)
assert answer.primary_source is None
class TestFormatContextMerged:
"""Tests for merged format_context_for_prompt function."""
def test_format_context_includes_pending_questions_by_default(self):
"""Test format_context_for_prompt includes pending_questions by default."""
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import format_context_for_prompt
context = ResearchContext(
original_question="What is X?",
sub_questions=["What is A?", "What is B?"],
)
result = format_context_for_prompt(context)
assert "<pending_questions>" in result
assert "What is A?" in result
assert "What is B?" in result
def test_format_context_excludes_pending_questions_when_flag_false(self):
"""Test format_context_for_prompt excludes pending_questions when flag is False."""
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import format_context_for_prompt
context = ResearchContext(
original_question="What is X?",
sub_questions=["What is A?", "What is B?"],
)
result = format_context_for_prompt(context, include_pending_questions=False)
assert "<pending_questions>" not in result
assert "What is A?" not in result
def test_format_context_uses_primary_source_helper(self):
"""Test format_context_for_prompt uses primary_source from SearchAnswer."""
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import format_context_for_prompt
context = ResearchContext(
original_question="What is X?",
)
# Add a QA response with citation
answer = SearchAnswer(
query="What is A?",
answer="A is...",
confidence=0.9,
citations=[
Citation(
document_id="doc-1",
chunk_id="chunk-1",
document_uri="test.md",
document_title="Test Document",
content="content",
),
],
)
context.add_qa_response(answer)
result = format_context_for_prompt(context)
assert "Test Document" in result

39
tests/agents/research/test_plan_prompt_selection.py
View file

@ -1,28 +1,31 @@
from haiku.rag.agents.research.prompts import PLAN_PROMPT, PLAN_PROMPT_WITH_CONTEXT
from haiku.rag.agents.research.prompts import (
ITERATIVE_PLAN_PROMPT,
ITERATIVE_PLAN_PROMPT_WITH_CONTEXT,
)
def test_plan_prompt_with_context_does_not_instruct_gather_context():
"""PLAN_PROMPT_WITH_CONTEXT should not instruct to use gather_context.
When session context already exists, we don't need to gather context again.
"""
assert "gather_context" not in PLAN_PROMPT_WITH_CONTEXT
def test_iterative_plan_prompt_proposes_first_question():
"""ITERATIVE_PLAN_PROMPT should instruct to propose the first question."""
assert "first question" in ITERATIVE_PLAN_PROMPT.lower()
assert "is_complete=False" in ITERATIVE_PLAN_PROMPT
def test_plan_prompt_instructs_gather_context():
"""PLAN_PROMPT should instruct to use gather_context for initial planning."""
assert "gather_context" in PLAN_PROMPT
def test_iterative_plan_prompt_with_context_evaluates_evidence():
"""ITERATIVE_PLAN_PROMPT_WITH_CONTEXT should evaluate prior answers."""
assert "prior_answers" in ITERATIVE_PLAN_PROMPT_WITH_CONTEXT
assert (
"evaluat" in ITERATIVE_PLAN_PROMPT_WITH_CONTEXT.lower()
) # matches evaluate/evaluating
def test_prompt_selection_uses_context_prompt_with_session_context():
"""When session_context exists, should use PLAN_PROMPT_WITH_CONTEXT."""
has_prior_answers = False
has_session_context = True
def test_prompt_selection_uses_context_prompt_with_prior_answers():
"""When prior_answers exist, should use ITERATIVE_PLAN_PROMPT_WITH_CONTEXT."""
has_prior_answers = True
effective_plan_prompt = (
PLAN_PROMPT_WITH_CONTEXT
if has_prior_answers or has_session_context
else PLAN_PROMPT
ITERATIVE_PLAN_PROMPT_WITH_CONTEXT
if has_prior_answers
else ITERATIVE_PLAN_PROMPT
)
assert effective_plan_prompt == PLAN_PROMPT_WITH_CONTEXT
assert effective_plan_prompt == ITERATIVE_PLAN_PROMPT_WITH_CONTEXT

70
tests/agents/research/test_research_graph.py
View file

@ -30,7 +30,6 @@ async def test_graph_end_to_end(allow_model_requests, temp_db_path, qa_corpus):
state = ResearchState(
context=ResearchContext(original_question=doc["question"]),
max_iterations=1,
confidence_threshold=0.5,
max_concurrency=1,
)
@ -46,22 +45,27 @@ async def test_graph_end_to_end(allow_model_requests, temp_db_path, qa_corpus):
client.close()
def test_research_plan_allows_empty_sub_questions():
"""Test ResearchPlan accepts empty sub_questions when context is sufficient."""
from haiku.rag.agents.research.models import ResearchPlan
def test_iterative_plan_result_model():
"""Test IterativePlanResult model validation."""
from haiku.rag.agents.research.models import IterativePlanResult
plan = ResearchPlan(sub_questions=[])
assert plan.sub_questions == []
# Test complete state
complete = IterativePlanResult(
is_complete=True,
next_question=None,
reasoning="All aspects covered.",
)
assert complete.is_complete is True
assert complete.next_question is None
def test_research_plan_rejects_too_many_sub_questions():
"""Test ResearchPlan rejects more than 12 sub_questions."""
from pydantic import ValidationError
from haiku.rag.agents.research.models import ResearchPlan
with pytest.raises(ValidationError, match="Cannot have more than 12"):
ResearchPlan(sub_questions=[f"q{i}" for i in range(13)])
# Test continue state
continue_result = IterativePlanResult(
is_complete=False,
next_question="What are the specific requirements?",
reasoning="Need more details.",
)
assert continue_result.is_complete is False
assert continue_result.next_question == "What are the specific requirements?"
# =============================================================================
@ -69,13 +73,20 @@ def test_research_plan_rejects_too_many_sub_questions():
# =============================================================================
def test_build_conversational_graph_returns_graph():
"""Test build_conversational_graph returns a valid Graph instance."""
def test_build_research_graph_conversational_mode_returns_graph():
"""Test build_research_graph with output_mode='conversational' returns a valid Graph instance."""
from pydantic_graph.beta import Graph
from haiku.rag.agents.research.graph import build_conversational_graph
graph = build_research_graph(output_mode="conversational")
assert graph is not None
assert isinstance(graph, Graph)
graph = build_conversational_graph()
def test_build_research_graph_report_mode_returns_graph():
"""Test build_research_graph with output_mode='report' returns a valid Graph instance."""
from pydantic_graph.beta import Graph
graph = build_research_graph(output_mode="report")
assert graph is not None
assert isinstance(graph, Graph)
@ -143,27 +154,6 @@ def test_format_context_for_prompt_with_session_context():
assert "What is Y?" in result
def test_format_context_for_prompt_excludes_pending_questions():
"""Test format_context_for_prompt can exclude pending questions."""
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import format_context_for_prompt
context = ResearchContext(
original_question="Main question?",
sub_questions=["Sub Q1?", "Sub Q2?"],
)
# With pending questions (default)
with_pending = format_context_for_prompt(context, include_pending_questions=True)
assert "Sub Q1?" in with_pending
# Without pending questions (for synthesis)
without_pending = format_context_for_prompt(
context, include_pending_questions=False
)
assert "Sub Q1?" not in without_pending
def test_format_context_for_prompt_with_prior_answers():
"""Test format_context_for_prompt includes prior_answers."""
from haiku.rag.agents.research.dependencies import ResearchContext

2
tests/agents/research/test_search_filter.py
View file

@ -78,7 +78,6 @@ async def test_research_graph_uses_search_filter(
state = ResearchState(
context=ResearchContext(original_question="Tell me about animals"),
max_iterations=1,
confidence_threshold=0.5,
search_filter=filter_clause,
)
@ -116,7 +115,6 @@ async def test_search_filter_none_searches_all(allow_model_requests, client_with
state = ResearchState(
context=ResearchContext(original_question="Tell me about animals"),
max_iterations=1,
confidence_threshold=0.5,
search_filter=None,
)

2310
tests/cassettes/test_chat_agent/test_chat_agent_ask_adds_citations.yaml
View file

File diff suppressed because one or more lines are too long

3498
tests/cassettes/test_chat_agent/test_chat_agent_ask_triggers_background_summarization.yaml
View file

File diff suppressed because one or more lines are too long

3773
tests/cassettes/test_chat_agent/test_chat_agent_ask_with_prior_answer_retrieval.yaml
View file

File diff suppressed because one or more lines are too long

214
tests/cassettes/test_chat_agent/test_chat_agent_get_document_not_found.yaml
View file

@ -8,7 +8,7 @@ interactions:
connection:
- keep-alive
content-length:
- '4073'
- '5211'
content-type:
- application/json
host:
@ -23,14 +23,16 @@ interactions:
CRITICAL RULES:
1. For greetings or casual chat: respond directly WITHOUT using any tools
2. For questions: Use the "ask" tool EXACTLY ONCE - it handles query expansion internally
2. For questions: Use the "ask" tool EXACTLY ONCE - it automatically uses prior conversation context
3. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally
4. NEVER call the same tool multiple times for a single user message
5. NEVER make up information - always use tools to get facts from the knowledge base
How to decide which tool to use:
- "get_document" - Use when the user references a SPECIFIC document by name, title, or URI (e.g., "summarize document X", "get the paper about Y", "fetch 2412.00566"). Retrieves the full document content.
- "ask" - Use for general questions about topics in the knowledge base when no specific document is named. It searches across all documents and returns answers with citations.
- "list_documents" - Use when the user wants to browse or see what documents are available (e.g., "what documents are available?", "show me the documents", "list available docs").
- "summarize_document" - Use when the user wants an overview or summary of a specific document (e.g., "summarize document X", "what does Y cover?", "give me an overview of Z").
- "get_document" - Use when the user wants the FULL content of a specific document (e.g., "get the paper about Y", "fetch 2412.00566", "show me the full document").
- "ask" - Use for questions about topics in the knowledge base. It automatically finds relevant prior answers from conversation history and searches across documents to return answers with citations.
- "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, copy the ENTIRE tool response to your output INCLUDING the content snippets. Do NOT shorten, summarize, or omit any part of the results.
IMPORTANT - When user mentions a document in search/ask:
@ -107,173 +109,17 @@ interactions:
type: function
- function:
description: |-
Retrieve a specific document by title or URI.
List available documents in the knowledge base.
Use this when the user wants to fetch/get/retrieve a specific document.
name: get_document
Use this when the user wants to browse or see what documents are available.
name: list_documents
parameters:
additionalProperties: false
properties:
query:
description: The document title or URI to look up
type: string
required:
- query
type: object
strict: true
type: function
uri: http://localhost:11434/v1/chat/completions
response:
headers:
content-length:
- '559'
content-type:
- application/json
parsed_body:
choices:
- finish_reason: tool_calls
index: 0
message:
content: ''
reasoning: User wants nonexistent document. We can use get_document but it may not exist. We'll try get_document.
role: assistant
tool_calls:
- function:
arguments: '{"query":"nonexistent document"}'
name: get_document
id: call_31uy8050
index: 0
type: function
created: 1768998264
id: chatcmpl-114
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 47
prompt_tokens: 842
total_tokens: 889
status:
code: 200
message: OK
- request:
headers:
accept:
- application/json
accept-encoding:
- gzip, deflate, zstd
connection:
- keep-alive
content-length:
- '4470'
content-type:
- application/json
host:
- localhost:11434
method: POST
parsed_body:
messages:
- content: |-
You are a helpful research assistant powered by haiku.rag, a knowledge base system.
You have access to a knowledge base of documents. Use your tools to search and answer questions.
CRITICAL RULES:
1. For greetings or casual chat: respond directly WITHOUT using any tools
2. For questions: Use the "ask" tool EXACTLY ONCE - it handles query expansion internally
3. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally
4. NEVER call the same tool multiple times for a single user message
5. NEVER make up information - always use tools to get facts from the knowledge base
How to decide which tool to use:
- "get_document" - Use when the user references a SPECIFIC document by name, title, or URI (e.g., "summarize document X", "get the paper about Y", "fetch 2412.00566"). Retrieves the full document content.
- "ask" - Use for general questions about topics in the knowledge base when no specific document is named. It searches across all documents and returns answers with citations.
- "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, copy the ENTIRE tool response to your output INCLUDING the content snippets. Do NOT shorten, summarize, or omit any part of the results.
IMPORTANT - When user mentions a document in search/ask:
- If user says "search in <doc>", "find in <doc>", "answer from <doc>", or "<topic> in <doc>":
- Extract the TOPIC as `query`/`question`
- Extract the DOCUMENT NAME as `document_name`
- Examples for search:
- "search for embeddings in the ML paper" → query="embeddings", document_name="ML paper"
- "find transformer architecture in 2412.00566" → query="transformer architecture", document_name="2412.00566"
- Examples for ask:
- "what does the ML paper say about embeddings?" → question="what are the embedding methods?", document_name="ML paper"
- "answer from 2412.00566 about model training" → question="how is the model trained?", document_name="2412.00566"
Be friendly and conversational. When you use the "ask" tool, summarize the key findings for the user.
role: system
- content: Get me the nonexistent document
role: user
- content: |-
<think>
User wants nonexistent document. We can use get_document but it may not exist. We'll try get_document.
</think>
role: assistant
tool_calls:
- function:
arguments: '{"query":"nonexistent document"}'
name: get_document
id: call_31uy8050
type: function
- content: 'Document not found: nonexistent document'
role: tool
tool_call_id: call_31uy8050
model: gpt-oss
reasoning_effort: low
stream: false
tool_choice: auto
tools:
- function:
description: |-
Search the knowledge base for relevant documents.
Use this when you need to find documents or explore the knowledge base.
Results are displayed to the user - just list the titles found.
name: search
parameters:
additionalProperties: false
properties:
document_name:
anyOf:
- type: string
- type: 'null'
default: null
description: Optional document name/title to search within
limit:
anyOf:
- type: integer
- type: 'null'
default: null
description: 'Number of results to return (default: 5)'
query:
description: The search query (what to search for)
type: string
required:
- query
type: object
type: function
- function:
description: |-
Answer a specific question using the knowledge base.
Use this for direct questions that need a focused answer with citations.
Uses a research graph for planning, searching, and synthesis.
name: ask
parameters:
additionalProperties: false
properties:
document_name:
anyOf:
- type: string
- type: 'null'
default: null
description: Optional document name/title to search within (e.g., "tbmed593", "army manual")
question:
description: The question to answer
type: string
required:
- question
page:
default: 1
description: 'Page number (default: 1, 50 documents per page)'
type: integer
type: object
type: function
- function:
@ -293,11 +139,28 @@ interactions:
type: object
strict: true
type: function
- function:
description: |-
Generate a summary of a specific document.
Use this when the user wants an overview or summary of a document's content.
name: summarize_document
parameters:
additionalProperties: false
properties:
query:
description: The document title or URI to summarize
type: string
required:
- query
type: object
strict: true
type: function
uri: http://localhost:11434/v1/chat/completions
response:
headers:
content-length:
- '459'
- '539'
content-type:
- application/json
parsed_body:
@ -305,18 +168,19 @@ interactions:
- finish_reason: stop
index: 0
message:
content: I’m sorry, but I couldn’t find a document titled “nonexistent document.” If you have another title or some
details to share, let me know and I’ll look it up for you!
content: I’m sorry, but that document isn’t available in the knowledge base. If there’s another topic or document
you’d like help with, just let me know!
reasoning: User asking for nonexistent document. Need to respond that none exists. No tool usage.
role: assistant
created: 1768998265
id: chatcmpl-968
created: 1769793913
id: chatcmpl-124
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 44
prompt_tokens: 912
total_tokens: 956
completion_tokens: 60
prompt_tokens: 1025
total_tokens: 1085
status:
code: 200
message: OK

136
tests/cassettes/test_chat_agent/test_chat_agent_get_document_tool.yaml
View file

@ -62,7 +62,7 @@ interactions:
connection:
- keep-alive
content-length:
- '4084'
- '5222'
content-type:
- application/json
host:
@ -77,14 +77,16 @@ interactions:
CRITICAL RULES:
1. For greetings or casual chat: respond directly WITHOUT using any tools
2. For questions: Use the "ask" tool EXACTLY ONCE - it handles query expansion internally
2. For questions: Use the "ask" tool EXACTLY ONCE - it automatically uses prior conversation context
3. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally
4. NEVER call the same tool multiple times for a single user message
5. NEVER make up information - always use tools to get facts from the knowledge base
How to decide which tool to use:
- "get_document" - Use when the user references a SPECIFIC document by name, title, or URI (e.g., "summarize document X", "get the paper about Y", "fetch 2412.00566"). Retrieves the full document content.
- "ask" - Use for general questions about topics in the knowledge base when no specific document is named. It searches across all documents and returns answers with citations.
- "list_documents" - Use when the user wants to browse or see what documents are available (e.g., "what documents are available?", "show me the documents", "list available docs").
- "summarize_document" - Use when the user wants an overview or summary of a specific document (e.g., "summarize document X", "what does Y cover?", "give me an overview of Z").
- "get_document" - Use when the user wants the FULL content of a specific document (e.g., "get the paper about Y", "fetch 2412.00566", "show me the full document").
- "ask" - Use for questions about topics in the knowledge base. It automatically finds relevant prior answers from conversation history and searches across documents to return answers with citations.
- "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, copy the ENTIRE tool response to your output INCLUDING the content snippets. Do NOT shorten, summarize, or omit any part of the results.
IMPORTANT - When user mentions a document in search/ask:
@ -159,6 +161,21 @@ interactions:
- question
type: object
type: function
- function:
description: |-
List available documents in the knowledge base.
Use this when the user wants to browse or see what documents are available.
name: list_documents
parameters:
additionalProperties: false
properties:
page:
default: 1
description: 'Page number (default: 1, 50 documents per page)'
type: integer
type: object
type: function
- function:
description: |-
Retrieve a specific document by title or URI.
@ -176,11 +193,28 @@ interactions:
type: object
strict: true
type: function
- function:
description: |-
Generate a summary of a specific document.
Use this when the user wants an overview or summary of a document's content.
name: summarize_document
parameters:
additionalProperties: false
properties:
query:
description: The document title or URI to summarize
type: string
required:
- query
type: object
strict: true
type: function
uri: http://localhost:11434/v1/chat/completions
response:
headers:
content-length:
- '477'
- '510'
content-type:
- application/json
parsed_body:
@ -189,24 +223,24 @@ interactions:
index: 0
message:
content: ''
reasoning: Need get_document.
reasoning: User wants a specific document. Use get_document.
role: assistant
tool_calls:
- function:
arguments: '{"query":"DocLayNet Class Labels"}'
name: get_document
id: call_m5izt73j
id: call_jxvkrbku
index: 0
type: function
created: 1768998258
id: chatcmpl-390
created: 1769793902
id: chatcmpl-566
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 32
prompt_tokens: 846
total_tokens: 878
completion_tokens: 38
prompt_tokens: 1029
total_tokens: 1067
status:
code: 200
message: OK
@ -219,7 +253,7 @@ interactions:
connection:
- keep-alive
content-length:
- '5169'
- '6338'
content-type:
- application/json
host:
@ -234,14 +268,16 @@ interactions:
CRITICAL RULES:
1. For greetings or casual chat: respond directly WITHOUT using any tools
2. For questions: Use the "ask" tool EXACTLY ONCE - it handles query expansion internally
2. For questions: Use the "ask" tool EXACTLY ONCE - it automatically uses prior conversation context
3. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally
4. NEVER call the same tool multiple times for a single user message
5. NEVER make up information - always use tools to get facts from the knowledge base
How to decide which tool to use:
- "get_document" - Use when the user references a SPECIFIC document by name, title, or URI (e.g., "summarize document X", "get the paper about Y", "fetch 2412.00566"). Retrieves the full document content.
- "ask" - Use for general questions about topics in the knowledge base when no specific document is named. It searches across all documents and returns answers with citations.
- "list_documents" - Use when the user wants to browse or see what documents are available (e.g., "what documents are available?", "show me the documents", "list available docs").
- "summarize_document" - Use when the user wants an overview or summary of a specific document (e.g., "summarize document X", "what does Y cover?", "give me an overview of Z").
- "get_document" - Use when the user wants the FULL content of a specific document (e.g., "get the paper about Y", "fetch 2412.00566", "show me the full document").
- "ask" - Use for questions about topics in the knowledge base. It automatically finds relevant prior answers from conversation history and searches across documents to return answers with citations.
- "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, copy the ENTIRE tool response to your output INCLUDING the content snippets. Do NOT shorten, summarize, or omit any part of the results.
IMPORTANT - When user mentions a document in search/ask:
@ -261,21 +297,21 @@ interactions:
role: user
- content: |-
<think>
Need get_document.
User wants a specific document. Use get_document.
</think>
role: assistant
tool_calls:
- function:
arguments: '{"query":"DocLayNet Class Labels"}'
name: get_document
id: call_m5izt73j
id: call_jxvkrbku
type: function
- content: |-
**DocLayNet Class Labels**
- ID: cd0fc5b8-ec16-4f1c-99fe-20979ecaab1c
- ID: 99c3503e-a8ad-4116-a8be-63fb8048dceb
- URI: doclaynet-labels
- Created: 2026-01-21 14:24
- Created: 2026-01-30 19:25
**Content:**
DocLayNet Dataset - Class Labels
@ -296,7 +332,7 @@ interactions:
The Text class has the highest count with 510,377 instances in the dataset.
role: tool
tool_call_id: call_m5izt73j
tool_call_id: call_jxvkrbku
model: gpt-oss
reasoning_effort: low
stream: false
@ -354,6 +390,21 @@ interactions:
- question
type: object
type: function
- function:
description: |-
List available documents in the knowledge base.
Use this when the user wants to browse or see what documents are available.
name: list_documents
parameters:
additionalProperties: false
properties:
page:
default: 1
description: 'Page number (default: 1, 50 documents per page)'
type: integer
type: object
type: function
- function:
description: |-
Retrieve a specific document by title or URI.
@ -371,11 +422,28 @@ interactions:
type: object
strict: true
type: function
- function:
description: |-
Generate a summary of a specific document.
Use this when the user wants an overview or summary of a document's content.
name: summarize_document
parameters:
additionalProperties: false
properties:
query:
description: The document title or URI to summarize
type: string
required:
- query
type: object
strict: true
type: function
uri: http://localhost:11434/v1/chat/completions
response:
headers:
content-length:
- '1131'
- '1137'
content-type:
- application/json
parsed_body:
@ -383,24 +451,24 @@ interactions:
- finish_reason: stop
index: 0
message:
content: "Here’s the full **DocLayNet Class Labels** document:\n\n---\n\n## DocLayNet Dataset – Class Labels\n\nDocLayNet
defines **11 distinct class labels** for document layout analysis:\n\n1. **Caption** – Text describing figures
or tables \n2. **Footnote** – Notes at the bottom of pages \n3. **Formula** – Mathematical expressions \n4.
**List‑item** – Items in bulleted or numbered lists \n5. **Page‑footer** – Footer content on pages \n6. **Page‑header**
– Header content on pages \n7. **Picture** – Images and diagrams \n8. **Section‑header** – Headings for document
content: "Here’s the **DocLayNet Class Labels** document:\n\n**DocLayNet Dataset – Class Labels**\n\nDocLayNet defines
**11** distinct class labels for document layout analysis:\n\n1. **Caption** – Text describing figures or tables
\ \n2. **Footnote** – Notes at the bottom of pages \n3. **Formula** – Mathematical expressions \n4. **List‑item**
– Items in bulleted or numbered lists \n5. **Page‑footer** – Footer content on pages \n6. **Page‑header** –
Header content on pages \n7. **Picture** – Images and diagrams \n8. **Section‑header** – Headings for document
sections \n9. **Table** – Tabular data \n10. **Text** – Regular paragraph text (highest count: **510,377** instances)
\ \n11. **Title** – Document titles \n\nThe **Text** class has the highest count with 510,377 instances in the
dataset."
\ \n11. **Title** – Document titles \n\nThe **Text** class has the largest number of instances, with **510,377**
entries in the dataset."
role: assistant
created: 1768998262
id: chatcmpl-746
created: 1769793910
id: chatcmpl-948
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 202
prompt_tokens: 1110
total_tokens: 1312
completion_tokens: 204
prompt_tokens: 1297
total_tokens: 1501
status:
code: 200
message: OK

196
tests/cassettes/test_chat_agent/test_chat_agent_search_tool.yaml
View file

File diff suppressed because one or more lines are too long

359
tests/cassettes/test_chat_agent/test_chat_agent_search_tool_with_filter.yaml
View file

File diff suppressed because one or more lines are too long

377
tests/cassettes/test_chat_agent/test_chat_agent_search_with_session_filter.yaml
View file

File diff suppressed because one or more lines are too long

233
tests/cassettes/test_chat_agent/test_search_agent_deduplication.yaml
View file

File diff suppressed because one or more lines are too long

99
tests/cassettes/test_chat_agent/test_search_agent_no_results.yaml
View file

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_chat_agent/test_search_agent_with_context.yaml
View file

File diff suppressed because one or more lines are too long

511
tests/cassettes/test_chat_agent/test_search_agent_with_filter.yaml
View file

File diff suppressed because one or more lines are too long

252
tests/cassettes/test_chat_agent/test_search_agent_with_session_filter.yaml
View file

File diff suppressed because one or more lines are too long

2787
tests/cassettes/test_research_graph/test_graph_end_to_end.yaml
View file

File diff suppressed because one or more lines are too long

3037
tests/cassettes/test_search_filter/test_research_graph_uses_search_filter.yaml
View file

File diff suppressed because one or more lines are too long

3929
tests/cassettes/test_search_filter/test_search_filter_none_searches_all.yaml
View file

File diff suppressed because one or more lines are too long