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

Merge pull request #197 from ggozad/feat/simplify-graphs

Browse source 
Siimplify and unifyresearch and deep QA into a single configurable graph
This commit is contained in:
Yiorgis Gozadinos 2025-12-18 11:42:45 +02:00 committed by GitHub
commit b60c5583aa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
37 changed files with 684 additions and 2300 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
CHANGELOG.md
View file

@ -15,6 +15,17 @@
### Changed
- **Chunker Sets Order**: Chunkers now set `chunk.order` directly
- **Unified Research Graph**: Simplified and unified research and deep QA into a single configurable graph
- Removed `analyze_insights` node - graph now flows directly from `collect_answers` to `decide`
- Simplified `EvaluationResult` to: `is_sufficient`, `confidence_score`, `reasoning`, `new_questions`
- Simplified `ResearchContext` - removed insight/gap tracking methods
- `ask --deep` now uses research graph with `max_iterations=2`, `confidence_threshold=0.0`
- `ask --deep` output now shows executive summary, key findings, and sources
- Added `include_plan` parameter to `build_research_graph()` for plan-less execution
- Added `max_iterations` and `confidence_threshold` overrides to `ResearchState.from_config()`
- **Improved Synthesis Prompt**: Updated synthesis agent prompt to produce direct answers
- Executive summary now directly answers the question instead of describing the report
- Added explicit examples of good vs bad output style
- **Evaluations Vacuum Strategy**: `populate_db` now uses periodic vacuum to prevent disk exhaustion with large datasets
- Disables auto_vacuum during population, vacuums every N documents with retention=0
- New `--vacuum-interval` CLI option (default: 100) to control vacuum frequency
@ -23,6 +34,16 @@
- Added dedicated Methodology section explaining MRR, MAP, and QA Accuracy metrics
- Organized results by dataset with retrieval and QA subsections
### Removed
- **Deep QA Graph**: Removed `haiku.rag.graph.deep_qa` module entirely
- Use `build_research_graph()` with appropriate parameters instead
- `ask --deep` CLI command now uses research graph internally
- **Insight/Gap Tracking**: Removed over-engineered insight and gap tracking from research graph
- Removed `InsightRecord`, `GapRecord`, `InsightAnalysis`, `InsightStatus`, `GapSeverity` models
- Removed `format_analysis_for_prompt()` helper
- Removed `INSIGHT_AGENT_PROMPT` from prompts
## [0.20.2] - 2025-12-12
### Fixed

267
docs/agents.md
View file

@ -1,156 +1,57 @@
# Agents
Three agentic flows are provided by haiku.rag:
Two agentic flows are provided by haiku.rag:
- Simple QA Agent — a focused question answering agent
- Deep QA Agent — multi-agent question decomposition for complex questions
- Research Multi‑Agent — a multi‑step, analyzable research workflow
- **Simple QA Agent** — a focused question answering agent
- **Research Graph** — a multi-step research workflow with question decomposition
For an interactive example using Pydantic AI and AG-UI, see the [Interactive Research Assistant](https://github.com/ggozad/haiku.rag/tree/main/examples/ag-ui-research) example ([demo video](https://vimeo.com/1128874386)). The demo uses a knowledge base containing haiku.rag's code and documentation.
For an interactive example using Pydantic AI and AG-UI, see the [Interactive Research Assistant](https://github.com/ggozad/haiku.rag/tree/main/examples/ag-ui-research) example ([demo video](https://vimeo.com/1128874386)).
See [QA and Research Configuration](configuration/qa-research.md) for configuring model, iterations, concurrency, and other settings.
### Simple QA Agent
## Simple QA Agent
The simple QA agent answers a single question using the knowledge base. It retrieves relevant chunks, optionally expands context around them, and asks the model to answer strictly based on that context.
Key points:
- Uses a single `search_documents` tool to fetch relevant chunks
- Can be run with or without inline citations in the prompt (citations prefer
document titles when present, otherwise URIs)
- Can be run with or without inline citations in the prompt
- Returns a plain string answer
Python usage:
**CLI usage:**
```bash
haiku-rag ask "What is climate change?"
# With citations
haiku-rag ask "What is climate change?" --cite
# Deep mode (uses research graph with optimized settings)
haiku-rag ask "What are the main features of haiku.rag?" --deep
```
**Python usage:**
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.qa.agent import QuestionAnswerAgent
async with HaikuRAG(path_to_db) as client:
# Choose a provider and model (see Configuration for env defaults)
agent = QuestionAnswerAgent(
client=client,
provider="openai", # or "ollama", "vllm", etc.
provider="openai",
model="gpt-4o-mini",
use_citations=False, # set True to bias prompt towards citing sources
use_citations=False,
)
answer = await agent.answer("What is climate change?")
print(answer)
```
### Deep QA Agent
## Research Graph
Deep QA is a multi-agent system that decomposes complex questions into sub-questions, answers them in batches, evaluates sufficiency, and iterates if needed before synthesizing a final answer. It's lighter than the full research workflow but more powerful than the simple QA agent.
```mermaid
---
title: Deep QA 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 QA
decide --> synthesize: Done with QA
synthesize --> [*]
```
Key nodes:
- **plan**: Decomposes the question into focused sub-questions using a presearch tool
- **get_batch**: Retrieves remaining sub-questions for the current iteration
- **search_one**: Answers a single sub-question using the knowledge base (mapped in parallel)
- **collect_answers**: Aggregates search results from parallel executions
- **decide**: Evaluates if sufficient information has been gathered or if more iterations are needed
- **synthesize**: Generates the final comprehensive answer from all gathered information
Key differences from Research:
- **Simpler evaluation**: Uses sufficiency check (not confidence + insight analysis)
- **Direct answers**: Returns just the answer (not a full research report)
- **Question-focused**: Optimized for answering specific questions, not open-ended research
- **Supports citations**: Can include inline source citations like `[document.md]`
- **Configurable iterations**: Control max_iterations (default: 2) and max_concurrency (default: 1)
Note on parallel execution:
- The `search_one` node is mapped over all questions in a batch
- Parallelism is controlled via `max_concurrency`
- All questions in an iteration are processed before evaluation
CLI usage:
```bash
# Deep QA without citations
haiku-rag ask "What are the main features of haiku.rag?" --deep
# Deep QA with citations
haiku-rag ask "What are the main features of haiku.rag?" --deep --cite
```
Python usage:
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
async with HaikuRAG(path_to_db) as client:
# Use global config (recommended)
graph = build_deep_qa_graph(config=Config)
context = DeepQAContext(
original_question="What are the main features of haiku.rag?",
use_citations=True
)
state = DeepQAState.from_config(context=context, config=Config)
deps = DeepQADeps(client=client)
result = await graph.run(
state=state,
deps=deps
)
print(result.answer)
print(result.sources)
```
Alternative usage with custom config:
```python
# Create a custom config with different settings
from haiku.rag.config.models import AppConfig, QAConfig
custom_config = AppConfig(
qa=QAConfig(
provider="openai",
model="gpt-4o-mini",
max_sub_questions=5,
max_iterations=3,
max_concurrency=2,
)
)
graph = build_deep_qa_graph(config=custom_config)
context = DeepQAContext(
original_question="What are the main features of haiku.rag?",
use_citations=True
)
state = DeepQAState.from_config(context=context, config=custom_config)
deps = DeepQADeps(client=client)
result = await graph.run(state=state, deps=deps)
```
### 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 — with clear stop conditions and shared state.
The research workflow is implemented as a typed pydantic-graph. It plans, searches (in parallel batches), evaluates, and synthesizes into a final report.
```mermaid
---
@ -162,47 +63,49 @@ stateDiagram-v2
get_batch --> search_one: Has questions (map)
get_batch --> synthesize: No questions
search_one --> collect_answers
collect_answers --> analyze_insights
analyze_insights --> decide
collect_answers --> decide
decide --> get_batch: Continue research
decide --> synthesize: Done researching
synthesize --> [*]
```
Key nodes:
**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 with minimal, verbatim context (mapped in parallel)
- **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
- **analyze_insights**: Synthesizes fresh insights, updates gaps, and suggests new sub-questions
- **decide**: Checks sufficiency/confidence thresholds and determines whether to continue research
- **decide**: Evaluates confidence and determines whether to continue or synthesize
- **synthesize**: Generates a final structured research report
Primary models:
**Primary models:**
- `SearchAnswer` — one per sub‑question (query, answer, context, sources)
- `InsightRecord` / `GapRecord` — structured tracking of findings and open issues
- `InsightAnalysis` — output of the analysis stage (insights, gaps, commentary)
- `EvaluationResult` — insights, new questions, sufficiency, confidence
- `SearchAnswer` — one per sub-question (query, answer, confidence, citations)
- `EvaluationResult` — confidence score, new questions, sufficiency assessment
- `ResearchReport` — final report (title, executive summary, findings, conclusions, …)
Note on parallel execution:
**Parallel execution:**
- The `search_one` node is mapped over all questions in a batch
- Parallelism is controlled via `max_concurrency`
- Analysis and decision nodes process results after each batch completes
- Decision nodes process results after each batch completes
CLI usage:
### CLI Usage
```bash
# Basic usage (uses config from file or defaults)
# Basic usage
haiku-rag research "How does haiku.rag organize and query documents?"
# With verbose output (shows progress)
haiku-rag research "How does haiku.rag organize and query documents?" --verbose
# With custom config file
haiku-rag --config my-research-config.yaml research "How does haiku.rag organize and query documents?" --verbose
# With document filter
haiku-rag research "What are the key findings?" --filter "uri LIKE '%report%'"
```
Python usage (blocking result):
### Python Usage
**Basic example:**
```python
from haiku.rag.client import HaikuRAG
@ -212,41 +115,25 @@ from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import ResearchDeps, ResearchState
async with HaikuRAG(path_to_db) as client:
# Use global config (recommended)
graph = build_research_graph(config=Config)
question = "What are the main drivers and trends of global temperature anomalies since 1990?"
context = ResearchContext(original_question=question)
context = ResearchContext(original_question="What are the main features?")
state = ResearchState.from_config(context=context, config=Config)
deps = ResearchDeps(client=client)
result = await graph.run(
state=state,
deps=deps,
)
report = await graph.run(state=state, deps=deps)
report = result
print(report.title)
print(report.executive_summary)
```
### Filtering Documents
Both Research and Deep QA graphs support restricting searches to specific documents via the `search_filter` parameter. Set it to a SQL WHERE clause before running:
```python
state = ResearchState.from_config(context=context, config=Config)
# Only search documents with these IDs
state.search_filter = "id IN ('doc-123', 'doc-456')"
result = await graph.run(state=state, deps=deps)
```
The filter applies to all search operations in the graph (context gathering and sub-question searches). See [Filtering Search Results](python.md#filtering-search-results) for available filter columns and syntax.
Alternative usage with custom config:
**With custom config:**
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig, ResearchConfig
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import ResearchDeps, ResearchState
custom_config = AppConfig(
research=ResearchConfig(
@ -258,28 +145,28 @@ custom_config = AppConfig(
)
)
graph = build_research_graph(config=custom_config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=custom_config)
deps = ResearchDeps(client=client)
async with HaikuRAG(path_to_db) as client:
graph = build_research_graph(config=custom_config)
context = ResearchContext(original_question="What are the main features?")
state = ResearchState.from_config(context=context, config=custom_config)
deps = ResearchDeps(client=client)
result = await graph.run(state=state, deps=deps)
report = await graph.run(state=state, deps=deps)
```
Python usage (streamed AG-UI events):
**Streaming AG-UI events:**
```python
from haiku.rag.graph.agui import stream_graph
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.graph.agui import stream_graph
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import ResearchDeps, ResearchState
async with HaikuRAG(path_to_db) as client:
graph = build_research_graph(config=Config)
question = "What are the main drivers and trends of global temperature anomalies since 1990?"
context = ResearchContext(original_question=question)
context = ResearchContext(original_question="What are the main features?")
state = ResearchState.from_config(context=context, config=Config)
deps = ResearchDeps(client=client)
@ -287,21 +174,25 @@ async with HaikuRAG(path_to_db) as client:
if event["type"] == "STEP_STARTED":
print(f"Starting step: {event['stepName']}")
elif event["type"] == "ACTIVITY_SNAPSHOT":
# Activity events include structured data alongside messages
content = event['content']
content = event["content"]
print(f" {content['message']}")
# Different activity types have different structured fields
if 'confidence' in content:
if "confidence" in content:
print(f" Confidence: {content['confidence']:.0%}")
if 'sub_questions' in content:
for q in content['sub_questions']:
print(f" - {q}")
if 'insights' in content:
print(f" New insights: {len(content['insights'])}")
elif event["type"] == "RUN_FINISHED":
print("\nResearch complete!\n")
result = event["result"]
print(result["title"])
print(result["executive_summary"])
report = event["result"]
print(report["executive_summary"])
```
### Filtering Documents
Restrict searches to specific documents via the `search_filter` parameter:
```python
# Set filter before running the graph
state = ResearchState.from_config(context=context, config=Config)
state.search_filter = "id IN ('doc-123', 'doc-456')"
report = await graph.run(state=state, deps=deps)
```
The filter applies to all search operations in the graph. See [Filtering Search Results](python.md#filtering-search-results) for available filter columns and syntax.

28
docs/server.md
View file

@ -119,7 +119,7 @@ URLs are also supported - the content is fetched and converted to markdown.
## AG-UI Server
The AG-UI server provides HTTP streaming of both research and deep ask graph execution using Server-Sent Events (SSE).
The AG-UI server provides HTTP streaming of research graph execution using Server-Sent Events (SSE).
### Starting the AG-UI Server
@ -131,7 +131,6 @@ This starts an HTTP server (default: http://0.0.0.0:8000) that exposes:
- `GET /health` - Health check endpoint
- `POST /v1/research/stream` - Research graph streaming endpoint
- `POST /v1/deep-ask/stream` - Deep ask graph streaming endpoint
### Configuration
@ -154,9 +153,9 @@ agui:
- **cors_methods**: Allowed HTTP methods (default: `["GET", "POST", "OPTIONS"]`)
- **cors_headers**: Allowed headers (default: `["*"]`)
### Using the Streaming Endpoints
### Using the Streaming Endpoint
Both endpoints accept POST requests with the same AG-UI RunAgentInput format and stream AG-UI events.
The endpoint accepts POST requests with AG-UI RunAgentInput format and streams AG-UI events.
**Request format:**
```json
@ -171,7 +170,7 @@ Both endpoints accept POST requests with the same AG-UI RunAgentInput format and
}
```
**Research endpoint example:**
**Example:**
```bash
curl -X POST http://localhost:8000/v1/research/stream \
-H "Content-Type: application/json" \
@ -183,24 +182,12 @@ curl -X POST http://localhost:8000/v1/research/stream \
--no-buffer
```
**Deep ask endpoint example:**
```bash
curl -X POST http://localhost:8000/v1/deep-ask/stream \
-H "Content-Type: application/json" \
-d '{
"state": {
"question": "How does haiku.rag handle document chunking?",
"use_citations": true
}
}' \
--no-buffer
```
The `--no-buffer` flag ensures curl displays events as they arrive instead of buffering them.
**Note:** The `state` object can include:
- `question`: The question to answer (required)
- `use_citations`: Enable citations in deep ask responses (optional, deep ask only)
- `max_iterations`: Maximum research iterations (optional, defaults to config)
- `confidence_threshold`: Confidence threshold for early termination (optional, defaults to config)
**Response:** Server-Sent Events stream with AG-UI protocol events:
- `RUN_STARTED` - Graph execution started
@ -239,8 +226,7 @@ data: {"type":"RUN_FINISHED","threadId":"abc123","runId":"xyz789","result":{"tit
- Additional structured fields depending on the activity type:
- **Planning**: `sub_questions` (list of strings)
- **Searching**: `query` (string), `confidence` (float, on completion), `error` (string, on failure)
- **Analyzing**: `insights` (list of insight objects), `gaps` (list of gap objects), `resolved_gaps` (list of strings)
- **Evaluating**: `confidence` (float), `is_sufficient` (boolean) for research; `is_sufficient` (boolean), `iterations` (int) for deep QA
- **Evaluating**: `confidence` (float), `is_sufficient` (boolean), `new_questions` (list of strings)
The `message` field is always present for simple rendering, while structured fields enable richer UI features like displaying lists, charts, and detailed status information.

10
examples/ag-ui-research/README.md
View file

@ -6,10 +6,9 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/),
## Features
- **Multi-iteration research graph**: Automated question decomposition, search, insight extraction, and gap analysis
- **Multi-iteration research graph**: Automated question decomposition and search
- **Intelligent evaluation**: Confidence-based decision making with automatic iteration until sufficient information is gathered
- **Live state synchronization**: Real-time delta updates of research progress via AG-UI protocol
- **Insight & gap tracking**: Structured insights with provenance and automatic gap identification
- **Rich reporting**: Generates comprehensive research reports with findings, conclusions, and sources
## Quick Start
@ -81,9 +80,8 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/),
- Gathers initial context about the topic
3. **Research iterations**: The graph autonomously:
- Searches the knowledge base for each sub-question in parallel
- Extracts structured insights with source provenance
- Identifies information gaps and assesses confidence
- Generates new follow-up questions for gaps
- Assesses confidence in gathered information
- Generates new follow-up questions if needed
- Iterates until confidence threshold is met or max iterations reached
4. **Synthesis**: Generates a comprehensive research report with:
- Executive summary
@ -126,7 +124,7 @@ This example demonstrates the **agent+graph** architecture pattern:
- CopilotKit for AG-UI protocol integration
- Split-pane UI: chat on left, live research state on right
- Real-time state synchronization via Server-Sent Events (SSE)
- `StateDisplay` component with collapsible sections for questions, insights, and gaps
- `StateDisplay` component with collapsible sections for questions and report
## Configuration

4
examples/ag-ui-research/backend/README.md
View file

@ -15,13 +15,11 @@ The server starts on `http://localhost:8000` and uses [haiku.rag configuration](
The backend uses `create_agui_server()` from `haiku.rag.graph.agui.server` which provides:
- **Research graph execution**: Multi-iteration research workflow with insight/gap tracking
- **Research graph execution**: Multi-iteration research workflow
- **AG-UI protocol**: Server-Sent Events (SSE) streaming for real-time state updates
- **Delta state updates**: Efficient incremental state synchronization using JSON Patch operations
- **Both research and deep_qa endpoints**: `/agent/research` and `/agent/deep_qa`
## Endpoints
- `GET /health` - Health check with configuration info
- `POST /agent/research/stream` - Research graph streaming endpoint (AG-UI protocol)
- `POST /agent/deep_qa/stream` - Deep QA graph streaming endpoint (AG-UI protocol)

3
examples/ag-ui-research/backend/agent.py
View file

@ -56,7 +56,7 @@ How to decide:
- "Tell me about Y" → Use run_research tool
When you use run_research, the graph will decompose questions, search the knowledge base,
extract insights, and generate a comprehensive report.
and generate a comprehensive report.
Be friendly and conversational in all responses.""",
)
@ -100,7 +100,6 @@ Main Findings:
Conclusions:
{chr(10).join(f"- {conclusion}" for conclusion in result.conclusions[:2])}
Total insights gathered: {len(state.context.insights)}
Confidence: {f"{state.last_eval.confidence_score:.0%}" if state.last_eval else "N/A"}
Iterations completed: {state.iterations}

31
examples/ag-ui-research/frontend/components/Agent.tsx
View file

@ -6,26 +6,6 @@ import "@copilotkit/react-ui/styles.css";
import DocumentSelector from "./DocumentSelector";
import StateDisplay from "./StateDisplay";
interface InsightRecord {
id: string;
summary: string;
status: string;
notes?: string;
supporting_sources: string[];
originating_questions: string[];
}
interface GapRecord {
id: string;
description: string;
severity: string;
blocking: boolean;
resolved: boolean;
notes?: string;
supporting_sources: string[];
resolved_by: string[];
}
interface Citation {
document_id: string;
chunk_id: string;
@ -48,14 +28,10 @@ interface ResearchContext {
original_question: string;
sub_questions: string[];
qa_responses: SearchAnswer[];
insights: InsightRecord[];
gaps: GapRecord[];
}
interface EvaluationResult {
key_insights: string[];
new_questions: string[];
gaps: string[];
confidence_score: number;
is_sufficient: boolean;
reasoning: string;
@ -78,10 +54,6 @@ interface ResearchState {
confidence_threshold: number;
max_concurrency: number;
last_eval: EvaluationResult | null;
last_analysis: {
insights_extracted: InsightRecord[];
gaps_identified: GapRecord[];
} | null;
result?: ResearchReport;
current_activity?: string;
current_activity_message?: string;
@ -96,15 +68,12 @@ function AgentContent() {
original_question: "",
sub_questions: [],
qa_responses: [],
insights: [],
gaps: [],
},
iterations: 0,
max_iterations: 3,
confidence_threshold: 0.8,
max_concurrency: 1,
last_eval: null,
last_analysis: null,
documentFilter: [],
},
});

329
examples/ag-ui-research/frontend/components/StateDisplay.tsx
View file

@ -11,26 +11,6 @@ interface VisualGroundingState {
error: string | null;
}
interface InsightRecord {
id: string;
summary: string;
status: string;
notes?: string;
supporting_sources: string[];
originating_questions: string[];
}
interface GapRecord {
id: string;
description: string;
severity: string;
blocking: boolean;
resolved: boolean;
notes?: string;
supporting_sources: string[];
resolved_by: string[];
}
interface Citation {
document_id: string;
chunk_id: string;
@ -53,14 +33,10 @@ interface ResearchContext {
original_question: string;
sub_questions: string[];
qa_responses: SearchAnswer[];
insights: InsightRecord[];
gaps: GapRecord[];
}
interface EvaluationResult {
key_insights: string[];
new_questions: string[];
gaps: string[];
confidence_score: number;
is_sufficient: boolean;
reasoning: string;
@ -83,10 +59,6 @@ interface ResearchState {
confidence_threshold: number;
max_concurrency: number;
last_eval: EvaluationResult | null;
last_analysis: {
insights_extracted: InsightRecord[];
gaps_identified: GapRecord[];
} | null;
result?: ResearchReport;
current_activity?: string;
current_activity_message?: string;
@ -101,8 +73,6 @@ export default function StateDisplay({ state }: StateDisplayProps) {
Record<string, boolean>
>({
questions: true,
insights: true,
gaps: true,
report: true,
});
@ -694,305 +664,6 @@ export default function StateDisplay({ state }: StateDisplayProps) {
</div>
)}
{/* Insights */}
{state.context.insights.length > 0 && (
<div
style={{
background: "white",
borderRadius: "8px",
boxShadow: "0 1px 3px rgba(0,0,0,0.1)",
overflow: "hidden",
}}
>
<button
type="button"
onClick={() => toggleSection("insights")}
style={{
width: "100%",
display: "flex",
justifyContent: "space-between",
alignItems: "center",
padding: "0.75rem",
background: "#edf2f7",
border: "1px solid #e2e8f0",
borderRadius: "4px",
cursor: "pointer",
fontSize: "1rem",
fontWeight: "600",
color: "#2d3748",
}}
>
<span>Key Insights ({state.context.insights.length})</span>
<span>{expandedSections.insights ? "▼" : "▶"}</span>
</button>
{expandedSections.insights && (
<div
style={{
padding: "1rem",
background: "#f7fafc",
border: "1px solid #e2e8f0",
borderTop: "none",
borderRadius: "0 0 4px 4px",
}}
>
{state.context.insights.map((insight) => (
<div
key={insight.id}
style={{
padding: "0.75rem",
background: "white",
borderRadius: "4px",
marginBottom: "0.5rem",
border: "1px solid #e2e8f0",
}}
>
<div
style={{
display: "flex",
justifyContent: "space-between",
marginBottom: "0.5rem",
}}
>
<span
style={{
fontSize: "0.75rem",
padding: "0.125rem 0.5rem",
background:
insight.status === "validated"
? "#c6f6d5"
: insight.status === "active"
? "#bee3f8"
: "#fed7d7",
color:
insight.status === "validated"
? "#22543d"
: insight.status === "active"
? "#2c5282"
: "#742a2a",
borderRadius: "4px",
}}
>
{insight.status}
</span>
<span
style={{
fontSize: "0.75rem",
color: "#718096",
}}
>
{insight.supporting_sources.length} sources
</span>
</div>
<div
style={{
fontSize: "0.875rem",
color: "#2d3748",
lineHeight: "1.5",
marginBottom: "0.5rem",
}}
>
<Markdown content={insight.summary} />
</div>
{insight.notes && (
<div
style={{
fontSize: "0.75rem",
color: "#718096",
marginTop: "0.5rem",
fontStyle: "italic",
}}
>
<Markdown content={insight.notes} />
</div>
)}
{insight.supporting_sources.length > 0 && (
<div
style={{
fontSize: "0.75rem",
color: "#718096",
marginTop: "0.5rem",
}}
>
<span style={{ fontWeight: "600" }}>Sources: </span>
{insight.supporting_sources.map((source, srcIdx) => (
<span key={`${insight.id}-src-${srcIdx}`}>
{srcIdx > 0 && ", "}
{source}
</span>
))}
</div>
)}
</div>
))}
</div>
)}
</div>
)}
{/* Knowledge Gaps */}
{state.context.gaps.length > 0 && (
<div
style={{
background: "white",
borderRadius: "8px",
boxShadow: "0 1px 3px rgba(0,0,0,0.1)",
overflow: "hidden",
}}
>
<button
type="button"
onClick={() => toggleSection("gaps")}
style={{
width: "100%",
display: "flex",
justifyContent: "space-between",
alignItems: "center",
padding: "0.75rem",
background: "#edf2f7",
border: "1px solid #e2e8f0",
borderRadius: "4px",
cursor: "pointer",
fontSize: "1rem",
fontWeight: "600",
color: "#2d3748",
}}
>
<span>Knowledge Gaps ({state.context.gaps.length})</span>
<span>{expandedSections.gaps ? "▼" : "▶"}</span>
</button>
{expandedSections.gaps && (
<div
style={{
padding: "1rem",
background: "#f7fafc",
border: "1px solid #e2e8f0",
borderTop: "none",
borderRadius: "0 0 4px 4px",
}}
>
{state.context.gaps.map((gap) => (
<div
key={gap.id}
style={{
padding: "0.75rem",
background: "white",
borderRadius: "4px",
marginBottom: "0.5rem",
border: "1px solid #e2e8f0",
}}
>
<div
style={{
display: "flex",
justifyContent: "space-between",
marginBottom: "0.5rem",
gap: "0.5rem",
flexWrap: "wrap",
}}
>
<div style={{ display: "flex", gap: "0.5rem" }}>
<span
style={{
fontSize: "0.75rem",
padding: "0.125rem 0.5rem",
background:
gap.severity === "critical"
? "#fed7d7"
: gap.severity === "high"
? "#feebc8"
: gap.severity === "medium"
? "#fef5e7"
: "#e6fffa",
color:
gap.severity === "critical"
? "#742a2a"
: gap.severity === "high"
? "#7c2d12"
: gap.severity === "medium"
? "#744210"
: "#234e52",
borderRadius: "4px",
fontWeight: "600",
}}
>
{gap.severity}
</span>
{gap.blocking && (
<span
style={{
fontSize: "0.75rem",
padding: "0.125rem 0.5rem",
background: "#fed7d7",
color: "#742a2a",
borderRadius: "4px",
fontWeight: "600",
}}
>
Blocking
</span>
)}
{gap.resolved && (
<span
style={{
fontSize: "0.75rem",
padding: "0.125rem 0.5rem",
background: "#c6f6d5",
color: "#22543d",
borderRadius: "4px",
fontWeight: "600",
}}
>
Resolved
</span>
)}
</div>
</div>
<div
style={{
fontSize: "0.875rem",
color: "#2d3748",
lineHeight: "1.5",
marginBottom: "0.5rem",
}}
>
<Markdown content={gap.description} />
</div>
{gap.notes && (
<div
style={{
fontSize: "0.75rem",
color: "#718096",
marginTop: "0.5rem",
fontStyle: "italic",
}}
>
<Markdown content={gap.notes} />
</div>
)}
{gap.resolved && gap.resolved_by.length > 0 && (
<div
style={{
fontSize: "0.75rem",
color: "#718096",
marginTop: "0.5rem",
}}
>
<span style={{ fontWeight: "600" }}>Resolved by: </span>
{gap.resolved_by.map((source, srcIdx) => (
<span key={`${gap.id}-resolved-${srcIdx}`}>
{srcIdx > 0 && ", "}
{source}
</span>
))}
</div>
)}
</div>
))}
</div>
)}
</div>
)}
{/* Final Report */}
{state.result && (
<div

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

@ -329,49 +329,58 @@ class HaikuRAGApp:
try:
citations = []
if deep:
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
from haiku.rag.graph.research.models import ResearchReport
graph = build_deep_qa_graph(config=self.config)
context = DeepQAContext(original_question=question)
state = DeepQAState.from_config(context=context, config=self.config)
graph = build_research_graph(config=self.config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(
context=context,
config=self.config,
max_iterations=2,
confidence_threshold=0.0,
)
state.search_filter = filter
deps = DeepQADeps(client=self.client)
deps = ResearchDeps(client=self.client)
if verbose:
# Use AG-UI renderer to process and display events
from haiku.rag.graph.common.models import Citation
renderer = AGUIConsoleRenderer(self.console)
result_dict = await renderer.render(
stream_graph(graph, state, deps)
)
# Result should be a dict with 'answer' and 'citations' keys
answer = result_dict.get("answer", "") if result_dict else ""
if cite and result_dict:
# Convert dicts to Citation objects
raw_citations = result_dict.get("citations", [])
citations = [
Citation(**c) if isinstance(c, dict) else c
for c in raw_citations
]
report = (
ResearchReport.model_validate(result_dict)
if result_dict
else None
)
else:
# Run without rendering events, just get the result
result = await graph.run(state=state, deps=deps)
answer = result.answer
if cite:
citations = result.citations
report = await graph.run(state=state, deps=deps)
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
if report:
self.console.print("[bold green]Answer:[/bold green]")
self.console.print(Markdown(report.executive_summary))
if report.main_findings:
self.console.print()
self.console.print("[bold cyan]Key Findings:[/bold cyan]")
for finding in report.main_findings:
self.console.print(f"• {finding}")
if report.sources_summary:
self.console.print()
self.console.print("[bold cyan]Sources:[/bold cyan]")
self.console.print(report.sources_summary)
else:
self.console.print("[yellow]No answer generated.[/yellow]")
else:
answer, citations = await self.client.ask(question, filter=filter)
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
self.console.print("[bold green]Answer:[/bold green]")
self.console.print(Markdown(answer))
if cite and citations:
for renderable in format_citations_rich(citations):
self.console.print(renderable)
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
self.console.print("[bold green]Answer:[/bold green]")
self.console.print(Markdown(answer))
if cite and citations:
for renderable in format_citations_rich(citations):
self.console.print(renderable)
except Exception as e:
self.console.print(f"[red]Error: {e}[/red]")

2
haiku_rag_slim/haiku/rag/client.py
View file

@ -27,7 +27,7 @@ from haiku.rag.store.repositories.settings import SettingsRepository
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.graph.common.models import Citation
from haiku.rag.graph.research.models import Citation
logger = logging.getLogger(__name__)

11
haiku_rag_slim/haiku/rag/graph/__init__.py
View file

@ -1,25 +1,14 @@
"""Graph module for haiku.rag.
This module contains all graph-related functionality including:
- AG-UI protocol for graph streaming
- Common graph utilities and models
- Research graph implementation
- Deep QA graph implementation
"""
from haiku.rag.graph.agui import (
AGUIConsoleRenderer,
AGUIEmitter,
create_agui_server,
stream_graph,
)
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.research.graph import build_research_graph
__all__ = [
"AGUIConsoleRenderer",
"AGUIEmitter",
"build_deep_qa_graph",
"build_research_graph",
"create_agui_server",
"stream_graph",

69
haiku_rag_slim/haiku/rag/graph/agui/server.py
View file

@ -154,19 +154,16 @@ def format_sse_event(event: AGUIEvent) -> str:
def create_agui_server( # pragma: no cover
config: "AppConfig", db_path: Path | None = None
) -> Starlette:
"""Create AG-UI server with both research and deep ask endpoints.
"""Create AG-UI server with research endpoint.
Args:
config: Application config with research and qa settings
config: Application config with research settings
db_path: Optional database path override
Returns:
Starlette app with research and deep ask endpoints
Starlette app with research endpoint
"""
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import ResearchDeps, ResearchState
@ -192,7 +189,14 @@ def create_agui_server( # pragma: no cover
if messages:
question = messages[0].get("content", "")
context = ResearchContext(original_question=question)
return ResearchState.from_config(context=context, config=config)
max_iterations = input_state.get("max_iterations")
confidence_threshold = input_state.get("confidence_threshold")
return ResearchState.from_config(
context=context,
config=config,
max_iterations=max_iterations,
confidence_threshold=confidence_threshold,
)
def research_deps_factory(input_config: dict[str, Any]) -> ResearchDeps:
effective_db_path = (
@ -202,28 +206,7 @@ def create_agui_server( # pragma: no cover
)
return ResearchDeps(client=get_client(effective_db_path))
# Deep ask graph factories
def deep_ask_graph_factory() -> Graph:
return build_deep_qa_graph(config)
def deep_ask_state_factory(input_state: dict[str, Any]) -> DeepQAState:
question = input_state.get("question", "")
if not question:
messages = input_state.get("messages", [])
if messages:
question = messages[0].get("content", "")
context = DeepQAContext(original_question=question)
return DeepQAState.from_config(context=context, config=config)
def deep_ask_deps_factory(input_config: dict[str, Any]) -> DeepQADeps:
effective_db_path = (
db_path
or input_config.get("db_path")
or config.storage.data_dir / "haiku.rag.lancedb"
)
return DeepQADeps(client=get_client(effective_db_path))
# Create event stream functions for each graph type
# Create event stream function
async def research_event_stream(
input_data: RunAgentInput,
) -> AsyncIterator[str]:
@ -236,18 +219,6 @@ def create_agui_server( # pragma: no cover
event_data = format_sse_event(event)
yield event_data
async def deep_ask_event_stream(
input_data: RunAgentInput,
) -> AsyncIterator[str]:
"""Generate SSE event stream from deep ask graph execution."""
graph = deep_ask_graph_factory()
initial_state = deep_ask_state_factory(input_data.state)
deps = deep_ask_deps_factory(input_data.config)
async for event in stream_graph(graph, initial_state, deps):
event_data = format_sse_event(event)
yield event_data
# Endpoint handlers
async def stream_research(request: Request) -> StreamingResponse:
"""Research graph streaming endpoint."""
@ -264,21 +235,6 @@ def create_agui_server( # pragma: no cover
},
)
async def stream_deep_ask(request: Request) -> StreamingResponse:
"""Deep ask graph streaming endpoint."""
body = await request.json()
input_data = RunAgentInput(**body)
return StreamingResponse(
deep_ask_event_stream(input_data),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
async def health_check(_: Request) -> JSONResponse:
"""Health check endpoint."""
return JSONResponse({"status": "healthy"})
@ -286,7 +242,6 @@ def create_agui_server( # pragma: no cover
# Define routes
routes = [
Route("/v1/research/stream", stream_research, methods=["POST"]),
Route("/v1/deep-ask/stream", stream_deep_ask, methods=["POST"]),
Route("/health", health_check, methods=["GET"]),
]

5
haiku_rag_slim/haiku/rag/graph/common/__init__.py
View file

@ -1,5 +0,0 @@
"""Common utilities for graph implementations."""
from haiku.rag.utils import get_model
__all__ = ["get_model"]

107
haiku_rag_slim/haiku/rag/graph/common/models.py
View file

@ -1,107 +0,0 @@
"""Common models used across different graph implementations."""
from typing import TYPE_CHECKING
from pydantic import BaseModel, Field, field_validator
if TYPE_CHECKING:
from haiku.rag.store.models import SearchResult
class ResearchPlan(BaseModel):
"""A structured research plan with sub-questions to explore."""
sub_questions: list[str] = Field(
...,
description="Specific questions to research, phrased as complete questions",
)
@field_validator("sub_questions")
@classmethod
def validate_sub_questions(cls, v: list[str]) -> list[str]:
if len(v) < 1:
raise ValueError("Must have at least 1 sub-question")
if len(v) > 12:
raise ValueError("Cannot have more than 12 sub-questions")
return v
class Citation(BaseModel):
"""Resolved citation with full metadata for display/visual grounding."""
document_id: str
chunk_id: str
document_uri: str
document_title: str | None = None
page_numbers: list[int] = Field(default_factory=list)
headings: list[str] | None = None
content: str
class RawSearchAnswer(BaseModel):
"""Answer to a search query with chunk references."""
query: str = Field(..., description="The question that was answered")
answer: str = Field(..., description="The answer to the question")
cited_chunks: list[str] = Field(
default_factory=list,
description="IDs of chunks used to form the answer",
)
confidence: float = Field(
default=1.0,
description="Confidence score for this answer (0-1)",
ge=0.0,
le=1.0,
)
class SearchAnswer(RawSearchAnswer):
"""Answer to a search query with resolved citations."""
citations: list[Citation] = Field(
default_factory=list,
description="Resolved citations with full metadata",
)
@classmethod
def from_raw(
cls,
raw: RawSearchAnswer,
search_results: "list[SearchResult]",
) -> "SearchAnswer":
"""Create SearchAnswer from RawSearchAnswer with resolved citations."""
citations = resolve_citations(raw.cited_chunks, search_results)
return cls(
query=raw.query,
answer=raw.answer,
cited_chunks=raw.cited_chunks,
confidence=raw.confidence,
citations=citations,
)
def resolve_citations(
cited_chunk_ids: list[str],
search_results: "list[SearchResult]",
) -> list[Citation]:
"""Resolve chunk IDs to full Citation objects with metadata."""
# Build lookup by chunk_id
by_id = {r.chunk_id: r for r in search_results if r.chunk_id}
citations = []
for chunk_id in cited_chunk_ids:
r = by_id.get(chunk_id)
if not r:
continue
citations.append(
Citation(
document_id=r.document_id or "",
chunk_id=chunk_id,
document_uri=r.document_uri or "",
document_title=r.document_title,
page_numbers=r.page_numbers,
headings=r.headings,
content=r.content,
)
)
return citations

315
haiku_rag_slim/haiku/rag/graph/common/nodes.py
View file

@ -1,315 +0,0 @@
"""Common node implementations for graph workflows."""
import asyncio
from collections.abc import Awaitable, Callable
from typing import Any, Protocol
from pydantic_ai import Agent, RunContext
from pydantic_ai.output import ToolOutput
from pydantic_graph.beta import StepContext
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig, ModelConfig
from haiku.rag.graph.agui.emitter import AGUIEmitter
from haiku.rag.graph.common import get_model
from haiku.rag.graph.common.models import RawSearchAnswer, ResearchPlan, SearchAnswer
from haiku.rag.graph.common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
from haiku.rag.store.models import SearchResult
class GraphContext(Protocol):
"""Protocol for graph context objects."""
original_question: str
sub_questions: list[str]
def add_qa_response(self, qa: SearchAnswer) -> None:
"""Add a QA response to context."""
...
class GraphState(Protocol):
"""Protocol for graph state objects."""
context: GraphContext
max_concurrency: int
search_filter: str | None
class GraphDeps(Protocol):
"""Protocol for graph dependencies."""
client: HaikuRAG
agui_emitter: AGUIEmitter[Any, Any] | None
semaphore: asyncio.Semaphore | None
class GraphAgentDeps(Protocol):
"""Protocol for agent dependencies."""
client: HaikuRAG
context: GraphContext
search_results: list[SearchResult]
def create_plan_node[AgentDepsT: GraphAgentDeps](
model_config: ModelConfig,
deps_type: type[AgentDepsT],
activity_message: str = "Creating plan",
output_retries: int | None = None,
config: AppConfig = Config,
) -> Callable[[StepContext[Any, Any, None]], Awaitable[None]]:
"""Create a plan node for any graph.
Args:
model_config: ModelConfig with provider, model, and settings
deps_type: Type of dependencies for the agent (e.g., ResearchDependencies, DeepQADependencies)
activity_message: Message to show during planning activity
output_retries: Number of output retries for the agent (optional)
config: AppConfig object (defaults to global Config)
Returns:
Async function that can be used as a graph step
"""
async def plan(ctx: StepContext[Any, Any, None], /) -> None:
state: GraphState = ctx.state # type: ignore[assignment]
deps: GraphDeps = ctx.deps # type: ignore[assignment]
if deps.agui_emitter:
deps.agui_emitter.start_step("plan")
deps.agui_emitter.update_activity(
"planning", {"stepName": "plan", "message": activity_message}
)
try:
# Build agent configuration
agent_config = {
"model": get_model(model_config, config),
"output_type": ResearchPlan,
"instructions": (
PLAN_PROMPT
+ "\n\nUse the gather_context tool once on the main question before planning."
),
"retries": 3,
"deps_type": deps_type,
}
if output_retries is not None:
agent_config["output_retries"] = output_retries
plan_agent = Agent(**agent_config)
# Capture search filter for use in tool
search_filter = state.search_filter
@plan_agent.tool
async def gather_context(
ctx2: RunContext[AgentDepsT], 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)
# Tool is registered via decorator above
_ = gather_context
prompt = (
"Plan a focused approach for the main question.\n\n"
f"Main question: {state.context.original_question}"
)
# Create agent dependencies
agent_deps = deps_type(client=deps.client, context=state.context) # type: ignore[call-arg]
plan_result = await plan_agent.run(prompt, deps=agent_deps)
state.context.sub_questions = list(plan_result.output.sub_questions)
# State now contains the plan - emit state update and narrate
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
count = len(state.context.sub_questions)
deps.agui_emitter.update_activity(
"planning",
{
"stepName": "plan",
"message": f"Created plan with {count} sub-questions",
"sub_questions": list(state.context.sub_questions),
},
)
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
return plan
def create_search_node[AgentDepsT: GraphAgentDeps](
model_config: ModelConfig,
deps_type: type[AgentDepsT],
with_step_wrapper: bool = True,
success_message_format: str = "Answered: {sub_q}",
handle_exceptions: bool = False,
config: AppConfig = Config,
) -> Callable[[StepContext[Any, Any, str]], Awaitable[SearchAnswer]]:
"""Create a search_one node for any graph.
Args:
model_config: ModelConfig with provider, model, and settings
deps_type: Type of dependencies for the agent
with_step_wrapper: Whether to wrap with agui_emitter start/finish step
success_message_format: Format string for success activity message
handle_exceptions: Whether to handle exceptions with fallback answer
config: AppConfig object (defaults to global Config)
Returns:
Async function that can be used as a graph step
"""
async def search_one(ctx: StepContext[Any, Any, str], /) -> SearchAnswer:
state: GraphState = ctx.state # type: ignore[assignment]
deps: GraphDeps = ctx.deps # type: ignore[assignment]
sub_q = ctx.inputs
# Create unique step name from question text
step_name = f"search: {sub_q}"
if deps.agui_emitter and with_step_wrapper:
deps.agui_emitter.start_step(step_name)
try:
# Create semaphore if not already provided
if deps.semaphore is None:
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
# Use semaphore to control concurrency
async with deps.semaphore:
return await _do_search(
state,
deps,
sub_q,
model_config,
deps_type,
success_message_format,
handle_exceptions,
config,
)
finally:
if deps.agui_emitter and with_step_wrapper:
deps.agui_emitter.finish_step()
return search_one
async def _do_search[AgentDepsT: GraphAgentDeps](
state: GraphState,
deps: GraphDeps,
sub_q: str,
model_config: ModelConfig,
deps_type: type[AgentDepsT],
success_message_format: str,
handle_exceptions: bool,
config: AppConfig,
) -> SearchAnswer:
"""Internal search implementation."""
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Searching: {sub_q}",
"query": sub_q,
},
)
agent = Agent(
model=get_model(model_config, config),
output_type=ToolOutput(RawSearchAnswer, max_retries=3),
instructions=SEARCH_AGENT_PROMPT,
retries=3,
deps_type=deps_type,
)
# Capture search filter for use in tool
search_filter = state.search_filter
@agent.tool
async def search_and_answer(
ctx2: RunContext[AgentDepsT], query: str, limit: int | None = None
) -> str:
"""Search the knowledge base for relevant documents.
Returns results with chunk IDs and relevance scores.
Reference results by their chunk_id in cited_chunks.
"""
results = await ctx2.deps.client.search(
query, limit=limit, filter=search_filter
)
results = await ctx2.deps.client.expand_context(results)
# Store results for citation resolution
ctx2.deps.search_results = results
# Format with metadata for agent context
parts = [r.format_for_agent() for r in results]
if not parts:
return f"No relevant information found in the knowledge base for: {query}"
return "\n\n".join(parts)
# Tool is registered via decorator above
_ = search_and_answer
agent_deps = deps_type(client=deps.client, context=state.context) # type: ignore[call-arg]
try:
result = await agent.run(sub_q, deps=agent_deps)
raw_answer = result.output
if raw_answer:
# Convert RawSearchAnswer to SearchAnswer with resolved citations
answer = SearchAnswer.from_raw(raw_answer, agent_deps.search_results)
state.context.add_qa_response(answer)
# State updated with new answer - emit state update and narrate
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
# Format the success message
if "{confidence" in success_message_format:
message = success_message_format.format(
sub_q=sub_q, confidence=answer.confidence
)
else:
message = success_message_format.format(sub_q=sub_q)
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": message,
"query": sub_q,
"confidence": answer.confidence,
},
)
return answer
# Return empty SearchAnswer if no result
return SearchAnswer(query=sub_q, answer="", confidence=0.0)
except Exception as e:
if handle_exceptions:
# Narrate the error
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Search failed: {e}",
"query": sub_q,
"error": str(e),
},
)
failure_answer = SearchAnswer(
query=sub_q,
answer=f"Search failed after retries: {str(e)}",
confidence=0.0,
)
return failure_answer
else:
raise

67
haiku_rag_slim/haiku/rag/graph/common/prompts.py
View file

@ -1,67 +0,0 @@
"""Common prompts used across different graph implementations."""
PLAN_PROMPT = """You are the research orchestrator for a focused, iterative workflow.
Responsibilities:
1. Understand and decompose the main question
2. Propose a minimal, high-leverage plan
3. Coordinate specialized agents to gather evidence
4. Iterate based on gaps and new findings
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).
Use the gather_context tool once on the main question before planning."""
SEARCH_AGENT_PROMPT = """You are a search and question-answering specialist.
Process:
1. Call search_and_answer with relevant keywords from the question.
2. Review the results and their relevance scores.
3. If needed, perform follow-up searches with different keywords (max 3 total).
4. Provide a concise answer based strictly on the retrieved content.
The search tool returns results like:
[9bde5847-44c9-400a-8997-0e6b65babf92] (score: 0.85)
Source: "Document Title" > Section > Subsection
Type: paragraph
Content:
The actual text content here...
[d5a63c82-cb40-439f-9b2e-de7d177829b7] (score: 0.72)
Source: "Another Document"
Type: table
Content:
| Column 1 | Column 2 |
...
Each result includes:
- chunk_id in brackets and relevance score
- Source: document title and section hierarchy (when available)
- Type: content type like paragraph, table, code, list_item (when available)
- Content: the actual text
Output format:
- query: Echo the question you are answering
- answer: Your concise answer based on the retrieved content
- cited_chunks: List of plain strings containing only the chunk UUIDs (not objects)
- confidence: A score from 0.0 to 1.0 indicating answer confidence
IMPORTANT: Use the EXACT, COMPLETE chunk ID (full UUID). Do NOT truncate IDs.
Guidelines:
- Base answers strictly on retrieved content - do not use external knowledge.
- Use the Source and Type metadata to understand context.
- If multiple results are relevant, synthesize them coherently.
- If information is insufficient, say so clearly.
- Be concise and direct; avoid meta commentary about the process.
- Higher scores indicate more relevant results."""

1
haiku_rag_slim/haiku/rag/graph/deep_qa/__init__.py
View file

@ -1 +0,0 @@
from haiku.rag.graph.deep_qa.models import DeepQAAnswer

29
haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py
View file

@ -1,29 +0,0 @@
from pydantic import BaseModel, Field
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.common.models import SearchAnswer
from haiku.rag.store.models import SearchResult
class DeepQAContext(BaseModel):
original_question: str = Field(description="The original question")
sub_questions: list[str] = Field(
default_factory=list, description="Decomposed sub-questions"
)
qa_responses: list[SearchAnswer] = Field(
default_factory=list, description="QA pairs collected during answering"
)
def add_qa_response(self, qa: SearchAnswer) -> None:
"""Add a QA response (citations already resolved)."""
self.qa_responses.append(qa)
class DeepQADependencies(BaseModel):
model_config = {"arbitrary_types_allowed": True}
client: HaikuRAG = Field(description="RAG client for document operations")
context: DeepQAContext = Field(description="Shared QA context")
search_results: list[SearchResult] = Field(
default_factory=list, description="Search results for citation resolution"
)

250
haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py
View file

@ -1,250 +0,0 @@
from pydantic_ai import Agent
from pydantic_ai.format_prompt import format_as_xml
from pydantic_graph.beta import Graph, GraphBuilder, StepContext
from pydantic_graph.beta.join import reduce_list_append
from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig
from haiku.rag.graph.common import get_model
from haiku.rag.graph.common.models import SearchAnswer, resolve_citations
from haiku.rag.graph.common.nodes import create_plan_node, create_search_node
from haiku.rag.graph.deep_qa.dependencies import DeepQADependencies
from haiku.rag.graph.deep_qa.models import DeepQAAnswer, DeepQAEvaluation
from haiku.rag.graph.deep_qa.prompts import DECISION_PROMPT, SYNTHESIS_PROMPT
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
from haiku.rag.store.models import SearchResult
def build_deep_qa_graph(
config: AppConfig = Config,
) -> Graph[DeepQAState, DeepQADeps, None, DeepQAAnswer]:
"""Build the Deep QA graph.
Args:
config: AppConfig object (uses config.qa for provider, model, and graph parameters)
Returns:
Configured Deep QA graph
"""
model_config = config.qa.model
g = GraphBuilder(
state_type=DeepQAState,
deps_type=DeepQADeps,
output_type=DeepQAAnswer,
)
# Create and register the plan node using the factory
plan = g.step(
create_plan_node(
model_config=model_config,
deps_type=DeepQADependencies, # type: ignore[arg-type]
activity_message="Planning approach",
output_retries=None, # Deep QA doesn't use output_retries
config=config,
)
) # type: ignore[arg-type]
# Create and register the search_one node using the factory
search_one = g.step(
create_search_node(
model_config=model_config,
deps_type=DeepQADependencies, # type: ignore[arg-type]
with_step_wrapper=False, # Deep QA doesn't wrap with agui_emitter step
success_message_format="Answered: {sub_q}",
handle_exceptions=True,
config=config,
)
) # type: ignore[arg-type]
@g.step
async def get_batch(
ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
) -> list[str] | None:
"""Get all remaining questions for this iteration."""
state = ctx.state
if not state.context.sub_questions:
return None
# Take ALL remaining questions - max_concurrency controls parallel execution within .map()
batch = list(state.context.sub_questions)
state.context.sub_questions.clear()
return batch
@g.step
async def decide(
ctx: StepContext[DeepQAState, DeepQADeps, list[SearchAnswer]],
) -> bool:
state = ctx.state
deps = ctx.deps
if deps.agui_emitter:
deps.agui_emitter.start_step("decide")
deps.agui_emitter.update_activity(
"evaluating", {"message": "Evaluating information sufficiency"}
)
try:
agent = Agent(
model=get_model(model_config, config),
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,
"confidence": qa.confidence,
}
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,
)
result = await agent.run(prompt, deps=agent_deps)
evaluation = result.output
state.iterations += 1
for new_q in evaluation.new_questions:
if new_q not in state.context.sub_questions:
state.context.sub_questions.append(new_q)
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
status = "sufficient" if evaluation.is_sufficient else "insufficient"
deps.agui_emitter.update_activity(
"evaluating",
{
"stepName": "decide",
"message": f"Information {status} after {state.iterations} iteration(s)",
"is_sufficient": evaluation.is_sufficient,
"iterations": state.iterations,
},
)
should_continue = (
not evaluation.is_sufficient and state.iterations < state.max_iterations
)
return should_continue
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
@g.step
async def synthesize(
ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
) -> DeepQAAnswer:
state = ctx.state
deps = ctx.deps
if deps.agui_emitter:
deps.agui_emitter.start_step("synthesize")
deps.agui_emitter.update_activity(
"synthesizing", {"message": "Synthesizing final answer"}
)
try:
agent = Agent(
model=get_model(model_config, config),
output_type=SearchAnswer,
instructions=SYNTHESIS_PROMPT,
retries=3,
deps_type=DeepQADependencies,
)
context_data = {
"original_question": state.context.original_question,
"sub_answers": [
{
"question": qa.query,
"answer": qa.answer,
"confidence": qa.confidence,
"cited_chunks": qa.cited_chunks,
}
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,
)
result = await agent.run(prompt, deps=agent_deps)
llm_answer = result.output
# Resolve citations by fetching chunks by ID
search_results = []
for chunk_id in llm_answer.cited_chunks:
chunk = await deps.client.chunk_repository.get_by_id(chunk_id)
if chunk:
search_results.append(SearchResult.from_chunk(chunk, score=1.0))
citations = resolve_citations(llm_answer.cited_chunks, search_results)
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"synthesizing", {"message": "Answer complete"}
)
return DeepQAAnswer(answer=llm_answer.answer, citations=citations)
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
# Build the graph structure
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),
)
# Branch based on whether we have questions
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),
)
# Branch based on decision
g.add(
g.edge_from(decide).to(
g.decision()
.branch(
g.match(bool, matches=lambda x: x).label("Continue QA").to(get_batch)
)
.branch(
g.match(bool, matches=lambda x: not x)
.label("Done with QA")
.to(synthesize)
)
),
g.edge_from(synthesize).to(g.end_node),
)
return g.build()

23
haiku_rag_slim/haiku/rag/graph/deep_qa/models.py
View file

@ -1,23 +0,0 @@
from pydantic import BaseModel, Field
from haiku.rag.graph.common.models import Citation
class DeepQAEvaluation(BaseModel):
is_sufficient: bool = Field(
description="Whether we have sufficient information to answer the question"
)
reasoning: str = Field(description="Explanation of the sufficiency assessment")
new_questions: list[str] = Field(
description="Additional sub-questions needed if insufficient",
default_factory=list,
)
class DeepQAAnswer(BaseModel):
"""Final deep QA answer with resolved citations."""
answer: str = Field(description="The comprehensive answer to the question")
citations: list[Citation] = Field(
default_factory=list, description="Resolved citations for the answer"
)

42
haiku_rag_slim/haiku/rag/graph/deep_qa/prompts.py
View file

@ -1,42 +0,0 @@
"""Deep QA specific prompts."""
SYNTHESIS_PROMPT = """You are an expert at synthesizing information into clear, concise answers.
Task:
- Combine the gathered information from sub-questions into a single comprehensive answer
- Answer the original question directly and completely
- Base your answer strictly on the provided evidence
- Be clear, accurate, and well-structured
Output format:
- query: Echo the original question being answered
- answer: The complete answer to the original question (2-4 paragraphs)
- cited_chunks: List of plain strings containing chunk IDs (UUIDs only, not objects)
- confidence: A score from 0.0 to 1.0 indicating answer confidence
Guidelines:
- Start directly with the answer - no preamble like "Based on the research..."
- Use a clear, professional tone
- Organize information logically
- If evidence is incomplete, state limitations clearly
- Do not include any claims not supported by the gathered information
- Each sub-answer includes cited_chunks IDs - include the relevant ones in your response"""
DECISION_PROMPT = """You are an expert at evaluating whether gathered information is sufficient to answer a question.
Task:
- Review the original question and all gathered sub-question answers
- Determine if we have enough information to provide a comprehensive answer
- If insufficient, suggest specific new sub-questions to fill the gaps
Output format:
- is_sufficient: Boolean indicating if we can answer the question comprehensively
- reasoning: Clear explanation of your assessment
- new_questions: List of plain strings, each a specific follow-up question (not objects)
Guidelines:
- Be strict but reasonable in your assessment
- Focus on whether core aspects of the question are addressed
- New questions should be specific and distinct from what's been asked
- Limit new questions to 2-3 maximum
- Consider whether additional searches would meaningfully improve the answer"""

59
haiku_rag_slim/haiku/rag/graph/deep_qa/state.py
View file

@ -1,59 +0,0 @@
import asyncio
from dataclasses import dataclass
from typing import TYPE_CHECKING
from pydantic import BaseModel, Field
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.models import DeepQAAnswer
if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig
from haiku.rag.graph.agui.emitter import AGUIEmitter
@dataclass
class DeepQADeps:
client: HaikuRAG
agui_emitter: "AGUIEmitter[DeepQAState, DeepQAAnswer] | None" = None
semaphore: asyncio.Semaphore | None = None
class DeepQAState(BaseModel):
"""Deep QA state for multi-agent question answering."""
model_config = {"arbitrary_types_allowed": True}
context: DeepQAContext = Field(description="Shared QA context")
max_sub_questions: int = Field(
default=3, description="Maximum number of sub-questions"
)
max_iterations: int = Field(
default=2, description="Maximum number of QA iterations"
)
max_concurrency: int = Field(
default=1, description="Maximum parallel sub-question searches"
)
iterations: int = Field(default=0, description="Current iteration number")
search_filter: str | None = Field(
default=None, description="SQL WHERE clause to filter search results"
)
@classmethod
def from_config(cls, context: DeepQAContext, config: "AppConfig") -> "DeepQAState":
"""Create a DeepQAState from an AppConfig.
Args:
context: The DeepQAContext containing the question and settings
config: The AppConfig object (uses config.qa for state parameters)
Returns:
A configured DeepQAState instance
"""
return cls(
context=context,
max_sub_questions=config.qa.max_sub_questions,
max_iterations=config.qa.max_iterations,
max_concurrency=config.qa.max_concurrency,
)

7
haiku_rag_slim/haiku/rag/graph/research/__init__.py
View file

@ -1,3 +1,6 @@
from haiku.rag.graph.common.models import SearchAnswer
from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.graph.research.models import EvaluationResult, ResearchReport
from haiku.rag.graph.research.models import (
EvaluationResult,
ResearchReport,
SearchAnswer,
)

95
haiku_rag_slim/haiku/rag/graph/research/common.py
View file

@ -1,95 +0,0 @@
from pydantic_ai import format_as_xml
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.models import InsightAnalysis
def format_context_for_prompt(context: ResearchContext) -> str:
"""Format the research context as XML for inclusion in prompts."""
context_data = {
"original_question": context.original_question,
"unanswered_questions": context.sub_questions,
"qa_responses": [
{
"question": qa.query,
"answer": qa.answer,
"confidence": qa.confidence,
"sources": [
{
"document_uri": c.document_uri,
"document_title": c.document_title,
"page_numbers": c.page_numbers,
"headings": c.headings,
}
for c in qa.citations
],
}
for qa in context.qa_responses
],
"insights": [
{
"id": insight.id,
"summary": insight.summary,
"status": insight.status.value,
"supporting_sources": insight.supporting_sources,
"originating_questions": insight.originating_questions,
"notes": insight.notes,
}
for insight in context.insights
],
"gaps": [
{
"id": gap.id,
"description": gap.description,
"severity": gap.severity.value,
"blocking": gap.blocking,
"resolved": gap.resolved,
"resolved_by": gap.resolved_by,
"supporting_sources": gap.supporting_sources,
"notes": gap.notes,
}
for gap in context.gaps
],
}
return format_as_xml(context_data, root_tag="research_context")
def format_analysis_for_prompt(
analysis: InsightAnalysis | None,
) -> str:
"""Format the latest insight analysis as XML for prompts."""
if analysis is None:
return "<latest_analysis />"
data = {
"commentary": analysis.commentary,
"highlights": [
{
"id": insight.id,
"summary": insight.summary,
"status": insight.status.value,
"supporting_sources": insight.supporting_sources,
"originating_questions": insight.originating_questions,
"notes": insight.notes,
}
for insight in analysis.highlights
],
"gap_assessments": [
{
"id": gap.id,
"description": gap.description,
"severity": gap.severity.value,
"blocking": gap.blocking,
"resolved": gap.resolved,
"resolved_by": gap.resolved_by,
"supporting_sources": gap.supporting_sources,
"notes": gap.notes,
}
for gap in analysis.gap_assessments
],
"resolved_gaps": analysis.resolved_gaps,
"new_questions": analysis.new_questions,
}
return format_as_xml(data, root_tag="latest_analysis")

134
haiku_rag_slim/haiku/rag/graph/research/dependencies.py
View file

@ -1,16 +1,13 @@
from collections.abc import Iterable
from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field, PrivateAttr
from pydantic import BaseModel, Field
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.common.models import SearchAnswer
from haiku.rag.graph.research.models import (
GapRecord,
InsightAnalysis,
InsightRecord,
)
from haiku.rag.store.models import SearchResult
if TYPE_CHECKING:
from haiku.rag.graph.research.models import SearchAnswer
class ResearchContext(BaseModel):
"""Context shared across research agents."""
@ -19,124 +16,14 @@ class ResearchContext(BaseModel):
sub_questions: list[str] = Field(
default_factory=list, description="Decomposed sub-questions"
)
qa_responses: list[SearchAnswer] = Field(
qa_responses: list[Any] = Field(
default_factory=list, description="Structured QA pairs used during research"
)
insights: list[InsightRecord] = Field(
default_factory=list, description="Key insights discovered"
)
gaps: list[GapRecord] = Field(
default_factory=list, description="Identified information gaps"
)
# Private dict indexes for O(1) lookups
_insights_by_id: dict[str, InsightRecord] = PrivateAttr(default_factory=dict)
_gaps_by_id: dict[str, GapRecord] = PrivateAttr(default_factory=dict)
def model_post_init(self, __context: object) -> None:
"""Build indexes after initialization."""
self._insights_by_id = {ins.id: ins for ins in self.insights}
self._gaps_by_id = {gap.id: gap for gap in self.gaps}
def add_qa_response(self, qa: SearchAnswer) -> None:
"""Add a structured QA response (citations already resolved)."""
def add_qa_response(self, qa: "SearchAnswer") -> None:
"""Add a structured QA response."""
self.qa_responses.append(qa)
def upsert_insights(self, records: Iterable[InsightRecord]) -> list[InsightRecord]:
"""Merge one or more insights into the shared context with deduplication."""
merged: list[InsightRecord] = []
for record in records:
candidate = InsightRecord.model_validate(record)
existing = self._insights_by_id.get(candidate.id)
if existing:
# Update existing insight
existing.summary = candidate.summary
existing.status = candidate.status
if candidate.notes:
existing.notes = candidate.notes
existing.supporting_sources = _merge_unique(
existing.supporting_sources, candidate.supporting_sources
)
existing.originating_questions = _merge_unique(
existing.originating_questions, candidate.originating_questions
)
merged.append(existing)
else:
# Add new insight
new_insight = candidate.model_copy(deep=True)
self.insights.append(new_insight)
self._insights_by_id[new_insight.id] = new_insight
merged.append(new_insight)
return merged
def upsert_gaps(self, records: Iterable[GapRecord]) -> list[GapRecord]:
"""Merge one or more gap records into the shared context with deduplication."""
merged: list[GapRecord] = []
for record in records:
candidate = GapRecord.model_validate(record)
existing = self._gaps_by_id.get(candidate.id)
if existing:
# Update existing gap
existing.description = candidate.description
existing.severity = candidate.severity
existing.blocking = candidate.blocking
existing.resolved = candidate.resolved
if candidate.notes:
existing.notes = candidate.notes
existing.supporting_sources = _merge_unique(
existing.supporting_sources, candidate.supporting_sources
)
existing.resolved_by = _merge_unique(
existing.resolved_by, candidate.resolved_by
)
merged.append(existing)
else:
# Add new gap
new_gap = candidate.model_copy(deep=True)
self.gaps.append(new_gap)
self._gaps_by_id[new_gap.id] = new_gap
merged.append(new_gap)
return merged
def mark_gap_resolved(
self, identifier: str, resolved_by: Iterable[str] | None = None
) -> GapRecord | None:
"""Mark a gap as resolved by identifier."""
gap = self._gaps_by_id.get(identifier)
if gap is None:
return None
gap.resolved = True
gap.blocking = False
if resolved_by:
gap.resolved_by = _merge_unique(gap.resolved_by, list(resolved_by))
return gap
def integrate_analysis(self, analysis: InsightAnalysis) -> None:
"""Apply an analysis result to the shared context."""
merged_insights: list[InsightRecord] = []
if analysis.highlights:
merged_insights = self.upsert_insights(analysis.highlights)
analysis.highlights = merged_insights
if analysis.gap_assessments:
merged_gaps = self.upsert_gaps(analysis.gap_assessments)
analysis.gap_assessments = merged_gaps
if analysis.resolved_gaps:
resolved_by_list = (
[ins.id for ins in merged_insights] if merged_insights else None
)
for resolved in analysis.resolved_gaps:
self.mark_gap_resolved(resolved, resolved_by=resolved_by_list)
for question in analysis.new_questions:
if question not in self.sub_questions:
self.sub_questions.append(question)
class ResearchDependencies(BaseModel):
"""Dependencies for research agents with multi-agent context."""
@ -148,8 +35,3 @@ class ResearchDependencies(BaseModel):
search_results: list[SearchResult] = Field(
default_factory=list, description="Search results for citation resolution"
)
def _merge_unique(existing: list[str], incoming: Iterable[str]) -> list[str]:
"""Merge two iterables preserving order while removing duplicates."""
return [k for k in dict.fromkeys([*existing, *incoming]) if k]

339
haiku_rag_slim/haiku/rag/graph/research/graph.py
View file

@ -1,37 +1,65 @@
from pydantic_ai import Agent
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.config import Config
from haiku.rag.config.models import AppConfig
from haiku.rag.graph.common import get_model
from haiku.rag.graph.common.models import SearchAnswer
from haiku.rag.graph.common.nodes import create_plan_node, create_search_node
from haiku.rag.graph.research.common import (
format_analysis_for_prompt,
format_context_for_prompt,
)
from haiku.rag.graph.research.dependencies import ResearchDependencies
from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.graph.research.models import (
EvaluationResult,
InsightAnalysis,
RawSearchAnswer,
ResearchPlan,
ResearchReport,
SearchAnswer,
)
from haiku.rag.graph.research.prompts import (
DECISION_AGENT_PROMPT,
INSIGHT_AGENT_PROMPT,
SYNTHESIS_AGENT_PROMPT,
DECISION_PROMPT,
PLAN_PROMPT,
SEARCH_PROMPT,
SYNTHESIS_PROMPT,
)
from haiku.rag.graph.research.state import ResearchDeps, ResearchState
from haiku.rag.utils import get_model
def format_context_for_prompt(context: ResearchContext) -> str:
"""Format the research context as XML for inclusion in prompts."""
context_data = {
"original_question": context.original_question,
"unanswered_questions": context.sub_questions,
"qa_responses": [
{
"question": qa.query,
"answer": qa.answer,
"confidence": qa.confidence,
"sources": [
{
"document_uri": c.document_uri,
"document_title": c.document_title,
"page_numbers": c.page_numbers,
"headings": c.headings,
}
for c in qa.citations
],
}
for qa in context.qa_responses
],
}
return format_as_xml(context_data, root_tag="research_context")
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
@ -43,28 +71,172 @@ def build_research_graph(
output_type=ResearchReport,
)
# Create and register the plan node using the factory
plan = g.step(
create_plan_node(
model_config=model_config,
deps_type=ResearchDependencies, # type: ignore[arg-type]
activity_message="Creating research plan",
output_retries=3,
config=config,
)
) # type: ignore[arg-type]
@g.step
async def plan(ctx: StepContext[ResearchState, ResearchDeps, None]) -> None:
"""Create research plan with sub-questions."""
state = ctx.state
deps = ctx.deps
# Create and register the search_one node using the factory
search_one = g.step(
create_search_node(
model_config=model_config,
deps_type=ResearchDependencies, # type: ignore[arg-type]
with_step_wrapper=True,
success_message_format="Found answer with {confidence:.0%} confidence",
handle_exceptions=True,
config=config,
)
) # type: ignore[arg-type]
if deps.agui_emitter:
deps.agui_emitter.start_step("plan")
deps.agui_emitter.update_activity(
"planning", {"stepName": "plan", "message": "Creating research plan"}
)
try:
plan_agent = Agent(
model=get_model(model_config, config),
output_type=ResearchPlan,
instructions=(
PLAN_PROMPT
+ "\n\nUse the gather_context tool once on the main question before planning."
),
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)
_ = gather_context
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)
state.context.sub_questions = list(plan_result.output.sub_questions)
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
count = len(state.context.sub_questions)
deps.agui_emitter.update_activity(
"planning",
{
"stepName": "plan",
"message": f"Created plan with {count} sub-questions",
"sub_questions": list(state.context.sub_questions),
},
)
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
@g.step
async def search_one(
ctx: StepContext[ResearchState, ResearchDeps, str],
) -> SearchAnswer:
"""Answer a single sub-question using the knowledge base."""
state = ctx.state
deps = ctx.deps
sub_q = ctx.inputs
step_name = f"search: {sub_q}"
if deps.agui_emitter:
deps.agui_emitter.start_step(step_name)
try:
if deps.semaphore is None:
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
async with deps.semaphore:
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Searching: {sub_q}",
"query": sub_q,
},
)
agent = 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
parts = [r.format_for_agent() for r in results]
if not parts:
return f"No relevant information found for: {query}"
return "\n\n".join(parts)
_ = search_and_answer
agent_deps = ResearchDependencies(
client=deps.client, context=state.context
)
try:
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)
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Found answer with {answer.confidence:.0%} confidence",
"query": sub_q,
"confidence": answer.confidence,
},
)
return answer
return SearchAnswer(query=sub_q, answer="", confidence=0.0)
except Exception as e:
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Search failed: {e}",
"query": sub_q,
"error": str(e),
},
)
return SearchAnswer(
query=sub_q,
answer=f"Search failed: {str(e)}",
confidence=0.0,
)
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
@g.step
async def get_batch(
@ -76,84 +248,15 @@ def build_research_graph(
if not state.context.sub_questions:
return None
# Take ALL remaining questions and process them in parallel
batch = list(state.context.sub_questions)
state.context.sub_questions.clear()
return batch
@g.step
async def analyze_insights(
async def decide(
ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]],
) -> None:
state = ctx.state
deps = ctx.deps
if deps.agui_emitter:
deps.agui_emitter.start_step("analyze_insights")
deps.agui_emitter.update_activity(
"analyzing", {"message": "Synthesizing insights and gaps"}
)
try:
agent = Agent(
model=get_model(model_config, config),
output_type=InsightAnalysis,
instructions=INSIGHT_AGENT_PROMPT,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Review the latest research context and update the shared ledger of insights, gaps,"
" and follow-up questions.\n\n"
f"{context_xml}"
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
analysis: InsightAnalysis = result.output
state.context.integrate_analysis(analysis)
state.last_analysis = analysis
# State updated with insights/gaps - emit state update and narrate
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
highlights = len(analysis.highlights)
gaps = len(analysis.gap_assessments)
resolved = len(analysis.resolved_gaps)
parts = []
if highlights:
parts.append(f"{highlights} insights")
if gaps:
parts.append(f"{gaps} gaps")
if resolved:
parts.append(f"{resolved} resolved")
summary = ", ".join(parts) if parts else "No updates"
deps.agui_emitter.update_activity(
"analyzing",
{
"stepName": "analyze_insights",
"message": f"Analysis: {summary}",
"insights": [
h.model_dump(mode="json") for h in analysis.highlights
],
"gaps": [
g.model_dump(mode="json") for g in analysis.gap_assessments
],
"resolved_gaps": list(analysis.resolved_gaps),
},
)
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
@g.step
async def decide(ctx: StepContext[ResearchState, ResearchDeps, None]) -> bool:
) -> bool:
"""Evaluate research sufficiency and decide whether to continue."""
state = ctx.state
deps = ctx.deps
@ -167,18 +270,16 @@ def build_research_graph(
agent = Agent(
model=get_model(model_config, config),
output_type=EvaluationResult,
instructions=DECISION_AGENT_PROMPT,
instructions=DECISION_PROMPT,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
analysis_xml = format_analysis_for_prompt(state.last_analysis)
prompt_parts = [
"Assess whether the research now answers the original question with adequate confidence.",
context_xml,
analysis_xml,
]
if state.last_eval is not None:
prev = state.last_eval
@ -205,7 +306,6 @@ def build_research_graph(
if new_q not in state.context.sub_questions:
state.context.sub_questions.append(new_q)
# State updated with evaluation - emit state update and narrate
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
sufficient = "Yes" if output.is_sufficient else "No"
@ -233,6 +333,7 @@ def build_research_graph(
async def synthesize(
ctx: StepContext[ResearchState, ResearchDeps, None | bool],
) -> ResearchReport:
"""Generate final research report."""
state = ctx.state
deps = ctx.deps
@ -246,7 +347,7 @@ def build_research_graph(
agent = Agent(
model=get_model(model_config, config),
output_type=ResearchReport,
instructions=SYNTHESIS_AGENT_PROMPT,
instructions=SYNTHESIS_PROMPT,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
@ -274,12 +375,14 @@ def build_research_graph(
initial_factory=list[SearchAnswer],
)
g.add(
g.edge_from(g.start_node).to(plan),
g.edge_from(plan).to(get_batch),
)
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))
# Branch based on whether we have questions
g.add(
g.edge_from(get_batch).to(
g.decision()
@ -287,11 +390,9 @@ def build_research_graph(
.branch(g.match(type(None)).label("No questions").to(synthesize))
),
g.edge_from(search_one).to(collect_answers),
g.edge_from(collect_answers).to(analyze_insights),
g.edge_from(analyze_insights).to(decide),
g.edge_from(collect_answers).to(decide),
)
# Branch based on decision
g.add(
g.edge_from(decide).to(
g.decision()

225
haiku_rag_slim/haiku/rag/graph/research/models.py
View file

@ -1,149 +1,128 @@
import uuid
from enum import Enum
from typing import TYPE_CHECKING
from pydantic import BaseModel, Field, field_validator
def _deduplicate_list(items: list[str]) -> list[str]:
"""Remove duplicates while preserving order."""
return list(dict.fromkeys(items))
if TYPE_CHECKING:
from haiku.rag.store.models import SearchResult
class InsightStatus(str, Enum):
OPEN = "open"
VALIDATED = "validated"
TENTATIVE = "tentative"
class ResearchPlan(BaseModel):
"""A structured research plan with sub-questions to explore."""
class GapSeverity(str, Enum):
LOW = "low"
MEDIUM = "medium"
HIGH = "high"
class TrackedRecord(BaseModel):
"""Base model for tracked entities with sources and metadata."""
model_config = {"validate_assignment": True}
id: str = Field(
default_factory=lambda: str(uuid.uuid4())[:8],
description="Unique identifier for the record",
)
supporting_sources: list[str] = Field(
default_factory=list,
description="Source identifiers backing this record",
)
notes: str | None = Field(
default=None,
description="Optional elaboration or caveats",
sub_questions: list[str] = Field(
...,
description="Specific questions to research, phrased as complete questions",
)
@field_validator("supporting_sources", mode="before")
@field_validator("sub_questions")
@classmethod
def deduplicate_sources(cls, v: list[str]) -> list[str]:
"""Ensure supporting_sources has no duplicates."""
return _deduplicate_list(v) if v else []
def validate_sub_questions(cls, v: list[str]) -> list[str]:
if len(v) < 1:
raise ValueError("Must have at least 1 sub-question")
if len(v) > 12:
raise ValueError("Cannot have more than 12 sub-questions")
return v
class InsightRecord(TrackedRecord):
"""Structured insight with provenance and lifecycle metadata."""
class Citation(BaseModel):
"""Resolved citation with full metadata for display/visual grounding."""
summary: str = Field(description="Concise description of the insight")
status: InsightStatus = Field(
default=InsightStatus.OPEN,
description="Lifecycle status for the insight",
)
originating_questions: list[str] = Field(
document_id: str
chunk_id: str
document_uri: str
document_title: str | None = None
page_numbers: list[int] = Field(default_factory=list)
headings: list[str] | None = None
content: str
class RawSearchAnswer(BaseModel):
"""Answer to a search query with chunk references."""
query: str = Field(..., description="The question that was answered")
answer: str = Field(..., description="The answer to the question")
cited_chunks: list[str] = Field(
default_factory=list,
description="Research sub-questions that produced this insight",
description="IDs of chunks used to form the answer",
)
@field_validator("originating_questions", mode="before")
@classmethod
def deduplicate_questions(cls, v: list[str]) -> list[str]:
"""Ensure originating_questions has no duplicates."""
return _deduplicate_list(v) if v else []
class GapRecord(TrackedRecord):
"""Structured representation of an identified research gap."""
description: str = Field(description="Concrete statement of what is missing")
severity: GapSeverity = Field(
default=GapSeverity.MEDIUM,
description="Severity of the gap for answering the main question",
)
blocking: bool = Field(
default=True,
description="Whether this gap blocks a confident answer",
)
resolved: bool = Field(
default=False,
description="Flag indicating if the gap has been resolved",
)
resolved_by: list[str] = Field(
default_factory=list,
description="Insight IDs or notes explaining how the gap was closed",
)
@field_validator("resolved_by", mode="before")
@classmethod
def deduplicate_resolved_by(cls, v: list[str]) -> list[str]:
"""Ensure resolved_by has no duplicates."""
return _deduplicate_list(v) if v else []
class InsightAnalysis(BaseModel):
"""Output of the insight aggregation agent."""
highlights: list[InsightRecord] = Field(
default_factory=list,
description="New or updated insights discovered this iteration",
)
gap_assessments: list[GapRecord] = Field(
default_factory=list,
description="New or updated gap records based on current evidence",
)
resolved_gaps: list[str] = Field(
default_factory=list,
description="Gap identifiers or descriptions considered resolved",
)
new_questions: list[str] = Field(
default_factory=list,
max_length=3,
description="Up to three follow-up sub-questions to pursue next",
)
commentary: str = Field(
description="Short narrative summary of the incremental findings",
)
class EvaluationResult(BaseModel):
"""Result of analysis and evaluation."""
key_insights: list[str] = Field(
description="Main insights extracted from the research so far"
)
new_questions: list[str] = Field(
description="New sub-questions to add to the research (max 3)",
max_length=3,
default=[],
)
gaps: list[str] = Field(
description="Concrete information gaps that remain", default_factory=list
)
confidence_score: float = Field(
description="Confidence level in the completeness of research (0-1)",
confidence: float = Field(
default=1.0,
description="Confidence score for this answer (0-1)",
ge=0.0,
le=1.0,
)
class SearchAnswer(RawSearchAnswer):
"""Answer to a search query with resolved citations."""
citations: list[Citation] = Field(
default_factory=list,
description="Resolved citations with full metadata",
)
@classmethod
def from_raw(
cls,
raw: RawSearchAnswer,
search_results: "list[SearchResult]",
) -> "SearchAnswer":
"""Create SearchAnswer from RawSearchAnswer with resolved citations."""
citations = resolve_citations(raw.cited_chunks, search_results)
return cls(
query=raw.query,
answer=raw.answer,
cited_chunks=raw.cited_chunks,
confidence=raw.confidence,
citations=citations,
)
def resolve_citations(
cited_chunk_ids: list[str],
search_results: "list[SearchResult]",
) -> list[Citation]:
"""Resolve chunk IDs to full Citation objects with metadata."""
by_id = {r.chunk_id: r for r in search_results if r.chunk_id}
citations = []
for chunk_id in cited_chunk_ids:
r = by_id.get(chunk_id)
if not r:
continue
citations.append(
Citation(
document_id=r.document_id or "",
chunk_id=chunk_id,
document_uri=r.document_uri or "",
document_title=r.document_title,
page_numbers=r.page_numbers,
headings=r.headings,
content=r.content,
)
)
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 ResearchReport(BaseModel):

184
haiku_rag_slim/haiku/rag/graph/research/prompts.py
View file

@ -1,92 +1,109 @@
INSIGHT_AGENT_PROMPT = """You are the insight aggregation specialist for the
research workflow.
PLAN_PROMPT = """You are the research orchestrator for a focused, iterative workflow.
Responsibilities:
1. Understand and decompose the main question
2. Propose a minimal, high-leverage plan
3. Coordinate specialized agents to gather evidence
4. Iterate based on gaps and new findings
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).
Use the gather_context tool once on the main question before planning."""
SEARCH_PROMPT = """You are a search and question-answering specialist.
Process:
1. Call search_and_answer with relevant keywords from the question.
2. Review the results and their relevance scores.
3. If needed, perform follow-up searches with different keywords (max 3 total).
4. Provide a concise answer based strictly on the retrieved content.
The search tool returns results like:
[9bde5847-44c9-400a-8997-0e6b65babf92] (score: 0.85)
Source: "Document Title" > Section > Subsection
Type: paragraph
Content:
The actual text content here...
[d5a63c82-cb40-439f-9b2e-de7d177829b7] (score: 0.72)
Source: "Another Document"
Type: table
Content:
| Column 1 | Column 2 |
...
Each result includes:
- chunk_id in brackets and relevance score
- Source: document title and section hierarchy (when available)
- Type: content type like paragraph, table, code, list_item (when available)
- Content: the actual text
Output format:
- query: Echo the question you are answering
- answer: Your concise answer based on the retrieved content
- cited_chunks: List of plain strings containing only the chunk UUIDs (not objects)
- confidence: A score from 0.0 to 1.0 indicating answer confidence
IMPORTANT: Use the EXACT, COMPLETE chunk ID (full UUID). Do NOT truncate IDs.
Guidelines:
- Base answers strictly on retrieved content - do not use external knowledge.
- Use the Source and Type metadata to understand context.
- If multiple results are relevant, synthesize them coherently.
- If information is insufficient, say so clearly.
- Be concise and direct; avoid meta commentary about the process.
- Higher scores indicate more relevant results."""
DECISION_PROMPT = """You are the research evaluator responsible for assessing
whether gathered evidence sufficiently answers the research question.
Inputs available:
- Original research question and sub-questions
- Question–answer pairs with supporting snippets and sources
- Existing insights and gaps (with status metadata)
- Original research question
- Question-answer pairs with supporting sources
- Previous evaluation (if any)
Tasks:
1. Extract new or refined insights that advance understanding of the question.
2. Update gap status, creating new gap entries when necessary and marking
resolved ones explicitly.
3. Suggest up to 3 high-impact follow-up sub_questions that would close the
most important remaining gaps.
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 format (map directly to fields):
- highlights: list of insights with fields {summary, status, supporting_sources,
originating_questions, notes}. Use status one of {validated, open, tentative}.
supporting_sources and originating_questions must be lists of plain strings.
- gap_assessments: list of gaps with fields {description, severity, blocking,
resolved, resolved_by, supporting_sources, notes}. Severity must be one of
{low, medium, high}. resolved_by and supporting_sources must be lists of plain strings.
- resolved_gaps: list of plain strings (identifiers or descriptions for gaps now closed).
- new_questions: list of plain strings, up to 3 standalone questions (no duplicates).
- commentary: 1–3 sentences summarizing what changed this round.
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
All list fields must contain plain strings only, not objects.
Be strict: only mark sufficient when key aspects are addressed with reliable evidence."""
Guidance:
- Be concise and avoid repeating previously recorded information unless it
changed materially.
- For supporting_sources, use only the document_uri strings from the sources.
- Only propose new sub_questions that directly address remaining gaps.
- When marking a gap as resolved, ensure the rationale is clear via
resolved_by or notes."""
DECISION_AGENT_PROMPT = """You are the research governor responsible for making
stop/go decisions.
Inputs available:
- Original research question and current plan
- Full insight ledger with status metadata
- Up-to-date gap tracker, including resolved indicators
- Latest insight analysis summary (highlights, gap changes, new questions)
- Previous evaluation decision (if any)
Tasks:
1. Determine whether the collected evidence now answers the original question.
2. Provide a confidence_score in [0,1] that reflects coverage, evidence quality,
and agreement across sources.
3. List the highest-priority gaps that still block a confident answer. Reference
existing gap descriptions rather than inventing new ones.
4. Optionally propose up to 3 new sub_questions only if they are not already in
the current backlog.
Strictness:
- Only mark research as sufficient when every critical aspect of the main
question is addressed with reliable, corroborated evidence.
- Treat unresolved high-severity or blocking gaps as a hard stop.
Output fields must line up with EvaluationResult:
- key_insights: list of plain strings, concise bullet-ready statements.
- new_questions: list of plain strings, follow-up sub-questions (max 3).
- gaps: list of plain strings, remaining blockers (reuse wording from tracked gaps).
- confidence_score: numeric in [0,1].
- is_sufficient: true only when no blocking gaps remain.
- reasoning: short narrative tying the decision to evidence coverage.
All list fields must contain plain strings only, not objects.
Remember: prefer maintaining continuity with the structured context over
introducing new terminology."""
SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist producing the final
research report.
SYNTHESIS_PROMPT = """You are a synthesis specialist producing the final
research report that directly answers the original question.
Goals:
1. Synthesize all gathered information into a coherent narrative.
1. Directly answer the research question using gathered evidence.
2. Present findings clearly and concisely.
3. Draw evidence‑based conclusions and recommendations.
3. Draw evidence-based conclusions and recommendations.
4. State limitations and uncertainties transparently.
Report guidelines (map to output fields):
- title: concise (5–12 words), informative.
- executive_summary: 3–5 sentences summarizing the overall answer.
- main_findings: list of plain strings, 4–8 one‑sentence bullets reflecting evidence.
- conclusions: list of plain strings, 2–4 bullets following logically from findings.
- recommendations: list of plain strings, 2–5 actionable bullets tied to findings.
- limitations: list of plain strings, 1–3 bullets describing constraints or uncertainties.
- title: concise (5-12 words), informative.
- executive_summary: 3-5 sentences that DIRECTLY ANSWER the original question.
Write the actual answer, not a description of what the report contains.
BAD: "This report examines the topic and presents findings..."
GOOD: "The system requires configuration X and supports features Y and Z..."
- main_findings: list of plain strings, 4-8 one-sentence bullets reflecting evidence.
- conclusions: list of plain strings, 2-4 bullets following logically from findings.
- recommendations: list of plain strings, 2-5 actionable bullets tied to findings.
- limitations: list of plain strings, 1-3 bullets describing constraints or uncertainties.
- sources_summary: single string listing sources with document paths and page numbers.
All list fields must contain plain strings only, not objects.
@ -94,16 +111,5 @@ All list fields must contain plain strings only, not objects.
Style:
- Base all content solely on the collected evidence.
- Be professional, objective, and specific.
- Avoid meta commentary and refrain from speculation beyond the evidence."""
PRESEARCH_AGENT_PROMPT = """You are a rapid research surveyor.
Task:
- Call gather_context once on the main question to obtain relevant text from
the knowledge base (KB).
- Read that context and produce a short natural‑language summary of what the
KB appears to contain relative to the question.
Rules:
- Base the summary strictly on the provided text; do not invent.
- Output only the summary as plain text (one short paragraph)."""
- NEVER use meta-commentary like "This report covers..." or "The findings show...".
Instead, state the actual information directly."""

46
haiku_rag_slim/haiku/rag/graph/research/state.py
View file

@ -6,11 +6,7 @@ from pydantic import BaseModel, Field
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.models import (
EvaluationResult,
InsightAnalysis,
ResearchReport,
)
from haiku.rag.graph.research.models import EvaluationResult, ResearchReport
if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig
@ -26,12 +22,7 @@ class ResearchDeps:
semaphore: asyncio.Semaphore | None = None
def emit_log(self, message: str, state: "ResearchState | None" = None) -> None:
"""Emit a log message through AG-UI events.
Args:
message: The message to log
state: Optional state to include in state update
"""
"""Emit a log message through AG-UI events."""
if self.agui_emitter:
self.agui_emitter.log(message)
if state:
@ -39,15 +30,12 @@ class ResearchDeps:
class ResearchState(BaseModel):
"""Research graph state model.
Fully JSON-serializable Pydantic model suitable for AG-UI state synchronization.
"""
"""Research graph state model."""
model_config = {"arbitrary_types_allowed": True}
context: ResearchContext = Field(
description="Shared research context with questions, insights, and gaps"
description="Shared research context with questions and QA responses"
)
iterations: int = Field(default=0, description="Current iteration number")
max_iterations: int = Field(default=3, description="Maximum allowed iterations")
@ -60,29 +48,33 @@ class ResearchState(BaseModel):
last_eval: EvaluationResult | None = Field(
default=None, description="Last evaluation result"
)
last_analysis: InsightAnalysis | None = Field(
default=None, description="Last insight analysis"
)
search_filter: str | None = Field(
default=None, description="SQL WHERE clause to filter search results"
)
@classmethod
def from_config(
cls, context: ResearchContext, config: "AppConfig"
cls,
context: ResearchContext,
config: "AppConfig",
max_iterations: int | None = None,
confidence_threshold: float | None = None,
) -> "ResearchState":
"""Create a ResearchState from an AppConfig.
Args:
context: The ResearchContext containing the question and settings
config: The AppConfig object (uses config.research for state parameters)
Returns:
A configured ResearchState instance
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=config.research.max_iterations,
confidence_threshold=config.research.confidence_threshold,
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,
)

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

@ -174,18 +174,26 @@ def create_mcp_server(db_path: Path, config: AppConfig = Config) -> FastMCP:
try:
async with HaikuRAG(db_path, config=config) as rag:
if deep:
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import (
ResearchDeps,
ResearchState,
)
graph = build_deep_qa_graph(config=config)
context = DeepQAContext(original_question=question)
state = DeepQAState.from_config(context=context, config=config)
deps = DeepQADeps(client=rag)
graph = build_research_graph(config=config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(
context=context,
config=config,
max_iterations=2,
confidence_threshold=0.0,
)
deps = ResearchDeps(client=rag)
result = await graph.run(state=state, deps=deps)
answer = result.answer
citations = result.citations
answer = result.executive_summary
citations = []
else:
answer, citations = await rag.ask(question)
if cite and citations:

4
haiku_rag_slim/haiku/rag/qa/agent.py
View file

@ -5,10 +5,10 @@ from pydantic_ai.output import ToolOutput
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.config.models import ModelConfig
from haiku.rag.graph.common import get_model
from haiku.rag.graph.common.models import Citation, RawSearchAnswer, resolve_citations
from haiku.rag.graph.research.models import Citation, RawSearchAnswer, resolve_citations
from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT
from haiku.rag.store.models import SearchResult
from haiku.rag.utils import get_model
class Dependencies(BaseModel):

2
haiku_rag_slim/haiku/rag/utils.py
View file

@ -8,7 +8,7 @@ from packaging.version import Version, parse
if TYPE_CHECKING:
from rich.console import RenderableType
from haiku.rag.graph.common.models import Citation
from haiku.rag.graph.research.models import Citation
def apply_common_settings(

41
tests/graph/test_deep_qa.py
View file

@ -1,41 +0,0 @@
import pytest
from pydantic_ai.models.test import TestModel
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
@pytest.mark.asyncio
async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path):
"""Test deep Q&A graph with mocked LLM using TestModel."""
# Mock get_model to return TestModel which generates valid schema-compliant data
def test_model_factory(provider, model, config=None):
return TestModel()
# Patch all locations where get_model is imported
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory)
graph = build_deep_qa_graph()
state = DeepQAState(
context=DeepQAContext(original_question="What is haiku.rag?"),
max_sub_questions=3,
)
# Use real client but with TestModel for LLM calls
client = HaikuRAG(temp_db_path, create=True)
deps = DeepQADeps(client=client)
result = await graph.run(state=state, deps=deps)
# TestModel will generate valid structured output based on schemas
assert result.answer is not None
assert isinstance(result.answer, str)
client.close()

2
tests/graph/test_research_graph.py
View file

@ -18,8 +18,6 @@ async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
# Patch all locations where get_model is imported
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
graph = build_research_graph()

53
tests/graph/test_search_filter.py
View file

@ -2,9 +2,6 @@ import pytest
from pydantic_ai.models.test import TestModel
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import ResearchDeps, ResearchState
@ -70,8 +67,6 @@ async def test_research_graph_uses_search_filter(monkeypatch, client_with_docs):
return TestModel()
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
graph = build_research_graph()
@ -97,52 +92,6 @@ async def test_research_graph_uses_search_filter(monkeypatch, client_with_docs):
)
@pytest.mark.asyncio
async def test_deep_qa_graph_uses_search_filter(monkeypatch, client_with_docs):
"""Test that deep QA graph passes search_filter to search operations."""
client, doc1_id, doc2_id = client_with_docs
# Track search calls to verify filter is passed
search_calls = []
original_search = client.search
async def tracking_search(query, limit=None, search_type="hybrid", filter=None):
search_calls.append({"query": query, "filter": filter})
return await original_search(query, limit, search_type, filter)
client.search = tracking_search
# Mock get_model to return TestModel
def test_model_factory(_provider, _model, _config=None):
return TestModel()
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory)
graph = build_deep_qa_graph()
# Create state with search_filter
filter_clause = f"id = '{doc2_id}'"
state = DeepQAState(
context=DeepQAContext(original_question="Tell me about animals"),
max_sub_questions=2,
search_filter=filter_clause,
)
deps = DeepQADeps(client=client)
await graph.run(state=state, deps=deps)
# Verify search was called with the filter
assert len(search_calls) > 0, "Expected search to be called"
for call in search_calls:
assert call["filter"] == filter_clause, (
f"Expected filter '{filter_clause}', got '{call['filter']}'"
)
@pytest.mark.asyncio
async def test_search_filter_none_searches_all(monkeypatch, client_with_docs):
"""Test that search_filter=None searches all documents."""
@ -163,8 +112,6 @@ async def test_search_filter_none_searches_all(monkeypatch, client_with_docs):
return TestModel()
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
graph = build_research_graph()

73
tests/test_app.py
View file

@ -310,7 +310,7 @@ async def test_ask_without_cite(app: HaikuRAGApp, monkeypatch):
@pytest.mark.asyncio
async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch):
"""Test asking a question with citations."""
from haiku.rag.graph.common.models import Citation
from haiku.rag.graph.research.models import Citation
mock_answer = "Test answer with citations"
mock_citations = [
@ -358,25 +358,36 @@ async def test_ask_with_verbose(app: HaikuRAGApp, monkeypatch):
@pytest.mark.asyncio
async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
"""Test asking a question with deep QA."""
from haiku.rag.graph.deep_qa.models import DeepQAAnswer
"""Test asking a question with deep mode uses research graph."""
import haiku.rag.app as app_module
from haiku.rag.graph.research.models import ResearchReport
mock_output = DeepQAAnswer(answer="Deep QA answer")
mock_output = ResearchReport(
title="Test",
executive_summary="Deep research answer",
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
sources_summary="Sources",
)
mock_graph = AsyncMock()
mock_graph.run.return_value = mock_output
mock_client = AsyncMock()
mock_client.__aenter__.return_value = mock_client
mock_print = MagicMock()
monkeypatch.setattr(app.console, "print", mock_print)
monkeypatch.setattr(app_module, "build_research_graph", lambda **kwargs: mock_graph)
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
with patch(
"haiku.rag.graph.deep_qa.graph.build_deep_qa_graph", return_value=mock_graph
):
await app.ask("test question", deep=True)
with patch("haiku.rag.app.HaikuRAG") as mock_rag_class:
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_client)
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
await app.ask("test question", deep=True)
# Check if there was an error printed
print_calls = [str(c) for c in mock_print.call_args_list]
error_calls = [c for c in print_calls if "Error" in c]
assert not error_calls, f"Error was printed: {error_calls}"
mock_graph.run.assert_called_once()
call_kwargs = mock_graph.run.call_args[1]
@ -385,25 +396,31 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
@pytest.mark.asyncio
async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch):
"""Test asking a question with deep QA and citations (cite ignored for deep)."""
from haiku.rag.graph.deep_qa.models import DeepQAAnswer
"""Test asking a question with deep mode (cite is ignored for research graph)."""
import haiku.rag.app as app_module
from haiku.rag.graph.research.models import ResearchReport
mock_output = DeepQAAnswer(answer="Deep QA answer")
mock_output = ResearchReport(
title="Test",
executive_summary="Deep research answer",
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
sources_summary="Sources",
)
mock_graph = AsyncMock()
mock_graph.run.return_value = mock_output
mock_client = AsyncMock()
mock_client.__aenter__.return_value = mock_client
mock_print = MagicMock()
monkeypatch.setattr(app.console, "print", mock_print)
monkeypatch.setattr(app_module, "build_research_graph", lambda **kwargs: mock_graph)
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
with patch(
"haiku.rag.graph.deep_qa.graph.build_deep_qa_graph", return_value=mock_graph
):
await app.ask("test question", deep=True, cite=True)
with patch("haiku.rag.app.HaikuRAG") as mock_rag_class:
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_client)
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
await app.ask("test question", deep=True, cite=True)
mock_graph.run.assert_called_once()
call_kwargs = mock_graph.run.call_args[1]
@ -412,9 +429,10 @@ async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch):
@pytest.mark.asyncio
async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch):
"""Test asking a question with deep QA and verbose output."""
"""Test asking a question with deep mode and verbose output."""
import haiku.rag.app as app_module
mock_output = {"answer": "Deep QA answer", "citations": []}
mock_output = {"executive_summary": "Deep research answer"}
mock_renderer = AsyncMock()
mock_renderer.render.return_value = mock_output
@ -422,17 +440,16 @@ async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch):
mock_graph = AsyncMock()
mock_client = AsyncMock()
mock_client.__aenter__.return_value = mock_client
mock_print = MagicMock()
monkeypatch.setattr(app.console, "print", mock_print)
monkeypatch.setattr(app_module, "build_research_graph", lambda **kwargs: mock_graph)
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
with patch(
"haiku.rag.graph.deep_qa.graph.build_deep_qa_graph", return_value=mock_graph
):
with patch("haiku.rag.app.AGUIConsoleRenderer", return_value=mock_renderer):
await app.ask("test question", deep=True, verbose=True)
with patch("haiku.rag.app.HaikuRAG") as mock_rag_class:
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_client)
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
with patch("haiku.rag.app.AGUIConsoleRenderer", return_value=mock_renderer):
await app.ask("test question", deep=True, verbose=True)
# With verbose, it should use AGUIConsoleRenderer.render, not graph.run
mock_renderer.render.assert_called_once()

9
tests/test_mcp.py
View file

@ -243,7 +243,7 @@ async def test_mcp_delete_document():
@pytest.mark.asyncio
async def test_mcp_ask_question_deep():
"""Test ask_question tool with deep=True is properly wired."""
"""Test ask_question tool with deep=True uses research graph."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test.lancedb"
mcp = create_mcp_server(db_path)
@ -251,7 +251,7 @@ async def test_mcp_ask_question_deep():
with (
patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class,
patch(
"haiku.rag.graph.deep_qa.graph.build_deep_qa_graph"
"haiku.rag.graph.research.graph.build_research_graph"
) as mock_graph_builder,
):
mock_rag = AsyncMock()
@ -260,8 +260,7 @@ async def test_mcp_ask_question_deep():
mock_graph = AsyncMock()
mock_result = AsyncMock()
mock_result.answer = "Deep answer"
mock_result.citations = []
mock_result.executive_summary = "Deep answer from research"
mock_graph.run = AsyncMock(return_value=mock_result)
mock_graph_builder.return_value = mock_graph
@ -273,7 +272,7 @@ async def test_mcp_ask_question_deep():
question="Deep question?", cite=False, deep=True
)
assert result == "Deep answer"
assert result == "Deep answer from research"
mock_graph.run.assert_called_once()