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 ### Changed
- **Chunker Sets Order**: Chunkers now set `chunk.order` directly - **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 - **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 - Disables auto_vacuum during population, vacuums every N documents with retention=0
- New `--vacuum-interval` CLI option (default: 100) to control vacuum frequency - 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 - Added dedicated Methodology section explaining MRR, MAP, and QA Accuracy metrics
- Organized results by dataset with retrieval and QA subsections - 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 ## [0.20.2] - 2025-12-12
### Fixed ### Fixed

267
docs/agents.md
View file

@ -1,156 +1,57 @@
# Agents # 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 - **Simple QA Agent** — a focused question answering agent
- Deep QA Agent — multi-agent question decomposition for complex questions - **Research Graph** — a multi-step research workflow with question decomposition
- Research Multi‑Agent — a multi‑step, analyzable research workflow
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. 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. 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: Key points:
- Uses a single `search_documents` tool to fetch relevant chunks - Uses a single `search_documents` tool to fetch relevant chunks
- Can be run with or without inline citations in the prompt (citations prefer - Can be run with or without inline citations in the prompt
document titles when present, otherwise URIs)
- Returns a plain string answer - 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 ```python
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.qa.agent import QuestionAnswerAgent from haiku.rag.qa.agent import QuestionAnswerAgent
async with HaikuRAG(path_to_db) as client: async with HaikuRAG(path_to_db) as client:
# Choose a provider and model (see Configuration for env defaults)
agent = QuestionAnswerAgent( agent = QuestionAnswerAgent(
client=client, client=client,
provider="openai", # or "ollama", "vllm", etc. provider="openai",
model="gpt-4o-mini", 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?") answer = await agent.answer("What is climate change?")
print(answer) 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. The research workflow is implemented as a typed pydantic-graph. It plans, searches (in parallel batches), evaluates, and synthesizes into a final report.
```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.
```mermaid ```mermaid
--- ---
@ -162,47 +63,49 @@ stateDiagram-v2
get_batch --> search_one: Has questions (map) get_batch --> search_one: Has questions (map)
get_batch --> synthesize: No questions get_batch --> synthesize: No questions
search_one --> collect_answers search_one --> collect_answers
collect_answers --> analyze_insights collect_answers --> decide
analyze_insights --> decide
decide --> get_batch: Continue research decide --> get_batch: Continue research
decide --> synthesize: Done researching decide --> synthesize: Done researching
synthesize --> [*] synthesize --> [*]
``` ```
Key nodes: **Key nodes:**
- **plan**: Builds up to 3 standalone sub‑questions (uses an internal presearch tool) - **plan**: Builds up to 3 standalone sub-questions (uses an internal presearch tool)
- **get_batch**: Retrieves remaining sub‑questions for the current iteration - **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) - **search_one**: Answers a single sub-question using the KB (mapped in parallel)
- **collect_answers**: Aggregates search results from parallel executions - **collect_answers**: Aggregates search results from parallel executions
- **analyze_insights**: Synthesizes fresh insights, updates gaps, and suggests new sub-questions - **decide**: Evaluates confidence and determines whether to continue or synthesize
- **decide**: Checks sufficiency/confidence thresholds and determines whether to continue research
- **synthesize**: Generates a final structured research report - **synthesize**: Generates a final structured research report
Primary models: **Primary models:**
- `SearchAnswer` — one per sub‑question (query, answer, context, sources) - `SearchAnswer` — one per sub-question (query, answer, confidence, citations)
- `InsightRecord` / `GapRecord` — structured tracking of findings and open issues - `EvaluationResult` — confidence score, new questions, sufficiency assessment
- `InsightAnalysis` — output of the analysis stage (insights, gaps, commentary)
- `EvaluationResult` — insights, new questions, sufficiency, confidence
- `ResearchReport` — final report (title, executive summary, findings, conclusions, …) - `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 - The `search_one` node is mapped over all questions in a batch
- Parallelism is controlled via `max_concurrency` - 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 ```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 haiku-rag research "How does haiku.rag organize and query documents?" --verbose
# With custom config file # With document filter
haiku-rag --config my-research-config.yaml research "How does haiku.rag organize and query documents?" --verbose haiku-rag research "What are the key findings?" --filter "uri LIKE '%report%'"
``` ```
Python usage (blocking result): ### Python Usage
**Basic example:**
```python ```python
from haiku.rag.client import HaikuRAG 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 from haiku.rag.graph.research.state import ResearchDeps, ResearchState
async with HaikuRAG(path_to_db) as client: async with HaikuRAG(path_to_db) as client:
# Use global config (recommended)
graph = build_research_graph(config=Config) graph = build_research_graph(config=Config)
question = "What are the main drivers and trends of global temperature anomalies since 1990?" context = ResearchContext(original_question="What are the main features?")
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=Config) state = ResearchState.from_config(context=context, config=Config)
deps = ResearchDeps(client=client) deps = ResearchDeps(client=client)
result = await graph.run( report = await graph.run(state=state, deps=deps)
state=state,
deps=deps,
)
report = result
print(report.title) print(report.title)
print(report.executive_summary) print(report.executive_summary)
``` ```
### Filtering Documents **With custom config:**
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:
```python ```python
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig, ResearchConfig 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( custom_config = AppConfig(
research=ResearchConfig( research=ResearchConfig(
@ -258,28 +145,28 @@ custom_config = AppConfig(
) )
) )
graph = build_research_graph(config=custom_config) async with HaikuRAG(path_to_db) as client:
context = ResearchContext(original_question=question) graph = build_research_graph(config=custom_config)
state = ResearchState.from_config(context=context, config=custom_config) context = ResearchContext(original_question="What are the main features?")
deps = ResearchDeps(client=client) 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 ```python
from haiku.rag.graph.agui import stream_graph
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config from haiku.rag.config import Config
from haiku.rag.graph.agui import stream_graph
from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import ResearchDeps, ResearchState from haiku.rag.graph.research.state import ResearchDeps, ResearchState
async with HaikuRAG(path_to_db) as client: async with HaikuRAG(path_to_db) as client:
graph = build_research_graph(config=Config) graph = build_research_graph(config=Config)
question = "What are the main drivers and trends of global temperature anomalies since 1990?" context = ResearchContext(original_question="What are the main features?")
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=Config) state = ResearchState.from_config(context=context, config=Config)
deps = ResearchDeps(client=client) deps = ResearchDeps(client=client)
@ -287,21 +174,25 @@ async with HaikuRAG(path_to_db) as client:
if event["type"] == "STEP_STARTED": if event["type"] == "STEP_STARTED":
print(f"Starting step: {event['stepName']}") print(f"Starting step: {event['stepName']}")
elif event["type"] == "ACTIVITY_SNAPSHOT": elif event["type"] == "ACTIVITY_SNAPSHOT":
# Activity events include structured data alongside messages content = event["content"]
content = event['content']
print(f" {content['message']}") print(f" {content['message']}")
if "confidence" in content:
# Different activity types have different structured fields
if 'confidence' in content:
print(f" Confidence: {content['confidence']:.0%}") 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": elif event["type"] == "RUN_FINISHED":
print("\nResearch complete!\n") report = event["result"]
result = event["result"] print(report["executive_summary"])
print(result["title"])
print(result["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 ## 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 ### 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 - `GET /health` - Health check endpoint
- `POST /v1/research/stream` - Research graph streaming endpoint - `POST /v1/research/stream` - Research graph streaming endpoint
- `POST /v1/deep-ask/stream` - Deep ask graph streaming endpoint
### Configuration ### Configuration
@ -154,9 +153,9 @@ agui:
- **cors_methods**: Allowed HTTP methods (default: `["GET", "POST", "OPTIONS"]`) - **cors_methods**: Allowed HTTP methods (default: `["GET", "POST", "OPTIONS"]`)
- **cors_headers**: Allowed headers (default: `["*"]`) - **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:** **Request format:**
```json ```json
@ -171,7 +170,7 @@ Both endpoints accept POST requests with the same AG-UI RunAgentInput format and
} }
``` ```
**Research endpoint example:** **Example:**
```bash ```bash
curl -X POST http://localhost:8000/v1/research/stream \ curl -X POST http://localhost:8000/v1/research/stream \
-H "Content-Type: application/json" \ -H "Content-Type: application/json" \
@ -183,24 +182,12 @@ curl -X POST http://localhost:8000/v1/research/stream \
--no-buffer --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. The `--no-buffer` flag ensures curl displays events as they arrive instead of buffering them.
**Note:** The `state` object can include: **Note:** The `state` object can include:
- `question`: The question to answer (required) - `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: **Response:** Server-Sent Events stream with AG-UI protocol events:
- `RUN_STARTED` - Graph execution started - `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: - Additional structured fields depending on the activity type:
- **Planning**: `sub_questions` (list of strings) - **Planning**: `sub_questions` (list of strings)
- **Searching**: `query` (string), `confidence` (float, on completion), `error` (string, on failure) - **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), `new_questions` (list of strings)
- **Evaluating**: `confidence` (float), `is_sufficient` (boolean) for research; `is_sufficient` (boolean), `iterations` (int) for deep QA
The `message` field is always present for simple rendering, while structured fields enable richer UI features like displaying lists, charts, and detailed status information. 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 ## 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 - **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 - **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 - **Rich reporting**: Generates comprehensive research reports with findings, conclusions, and sources
## Quick Start ## Quick Start
@ -81,9 +80,8 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/),
- Gathers initial context about the topic - Gathers initial context about the topic
3. **Research iterations**: The graph autonomously: 3. **Research iterations**: The graph autonomously:
- Searches the knowledge base for each sub-question in parallel - Searches the knowledge base for each sub-question in parallel
- Extracts structured insights with source provenance - Assesses confidence in gathered information
- Identifies information gaps and assesses confidence - Generates new follow-up questions if needed
- Generates new follow-up questions for gaps
- Iterates until confidence threshold is met or max iterations reached - Iterates until confidence threshold is met or max iterations reached
4. **Synthesis**: Generates a comprehensive research report with: 4. **Synthesis**: Generates a comprehensive research report with:
- Executive summary - Executive summary
@ -126,7 +124,7 @@ This example demonstrates the **agent+graph** architecture pattern:
- CopilotKit for AG-UI protocol integration - CopilotKit for AG-UI protocol integration
- Split-pane UI: chat on left, live research state on right - Split-pane UI: chat on left, live research state on right
- Real-time state synchronization via Server-Sent Events (SSE) - 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 ## 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: 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 - **AG-UI protocol**: Server-Sent Events (SSE) streaming for real-time state updates
- **Delta state updates**: Efficient incremental state synchronization using JSON Patch operations - **Delta state updates**: Efficient incremental state synchronization using JSON Patch operations
- **Both research and deep_qa endpoints**: `/agent/research` and `/agent/deep_qa`
## Endpoints ## Endpoints
- `GET /health` - Health check with configuration info - `GET /health` - Health check with configuration info
- `POST /agent/research/stream` - Research graph streaming endpoint (AG-UI protocol) - `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 - "Tell me about Y" → Use run_research tool
When you use run_research, the graph will decompose questions, search the knowledge base, 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.""", Be friendly and conversational in all responses.""",
) )
@ -100,7 +100,6 @@ Main Findings:
Conclusions: Conclusions:
{chr(10).join(f"- {conclusion}" for conclusion in result.conclusions[:2])} {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"} Confidence: {f"{state.last_eval.confidence_score:.0%}" if state.last_eval else "N/A"}
Iterations completed: {state.iterations} 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 DocumentSelector from "./DocumentSelector";
import StateDisplay from "./StateDisplay"; 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 { interface Citation {
document_id: string; document_id: string;
chunk_id: string; chunk_id: string;
@ -48,14 +28,10 @@ interface ResearchContext {
original_question: string; original_question: string;
sub_questions: string[]; sub_questions: string[];
qa_responses: SearchAnswer[]; qa_responses: SearchAnswer[];
insights: InsightRecord[];
gaps: GapRecord[];
} }
interface EvaluationResult { interface EvaluationResult {
key_insights: string[];
new_questions: string[]; new_questions: string[];
gaps: string[];
confidence_score: number; confidence_score: number;
is_sufficient: boolean; is_sufficient: boolean;
reasoning: string; reasoning: string;
@ -78,10 +54,6 @@ interface ResearchState {
confidence_threshold: number; confidence_threshold: number;
max_concurrency: number; max_concurrency: number;
last_eval: EvaluationResult | null; last_eval: EvaluationResult | null;
last_analysis: {
insights_extracted: InsightRecord[];
gaps_identified: GapRecord[];
} | null;
result?: ResearchReport; result?: ResearchReport;
current_activity?: string; current_activity?: string;
current_activity_message?: string; current_activity_message?: string;
@ -96,15 +68,12 @@ function AgentContent() {
original_question: "", original_question: "",
sub_questions: [], sub_questions: [],
qa_responses: [], qa_responses: [],
insights: [],
gaps: [],
}, },
iterations: 0, iterations: 0,
max_iterations: 3, max_iterations: 3,
confidence_threshold: 0.8, confidence_threshold: 0.8,
max_concurrency: 1, max_concurrency: 1,
last_eval: null, last_eval: null,
last_analysis: null,
documentFilter: [], documentFilter: [],
}, },
}); });

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

@ -11,26 +11,6 @@ interface VisualGroundingState {
error: string | null; 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 { interface Citation {
document_id: string; document_id: string;
chunk_id: string; chunk_id: string;
@ -53,14 +33,10 @@ interface ResearchContext {
original_question: string; original_question: string;
sub_questions: string[]; sub_questions: string[];
qa_responses: SearchAnswer[]; qa_responses: SearchAnswer[];
insights: InsightRecord[];
gaps: GapRecord[];
} }
interface EvaluationResult { interface EvaluationResult {
key_insights: string[];
new_questions: string[]; new_questions: string[];
gaps: string[];
confidence_score: number; confidence_score: number;
is_sufficient: boolean; is_sufficient: boolean;
reasoning: string; reasoning: string;
@ -83,10 +59,6 @@ interface ResearchState {
confidence_threshold: number; confidence_threshold: number;
max_concurrency: number; max_concurrency: number;
last_eval: EvaluationResult | null; last_eval: EvaluationResult | null;
last_analysis: {
insights_extracted: InsightRecord[];
gaps_identified: GapRecord[];
} | null;
result?: ResearchReport; result?: ResearchReport;
current_activity?: string; current_activity?: string;
current_activity_message?: string; current_activity_message?: string;
@ -101,8 +73,6 @@ export default function StateDisplay({ state }: StateDisplayProps) {
Record<string, boolean> Record<string, boolean>
>({ >({
questions: true, questions: true,
insights: true,
gaps: true,
report: true, report: true,
}); });
@ -694,305 +664,6 @@ export default function StateDisplay({ state }: StateDisplayProps) {
</div> </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 */} {/* Final Report */}
{state.result && ( {state.result && (
<div <div

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

@ -329,49 +329,58 @@ class HaikuRAGApp:
try: try:
citations = [] citations = []
if deep: if deep:
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext from haiku.rag.graph.research.models import ResearchReport
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
graph = build_deep_qa_graph(config=self.config) graph = build_research_graph(config=self.config)
context = DeepQAContext(original_question=question) context = ResearchContext(original_question=question)
state = DeepQAState.from_config(context=context, config=self.config) state = ResearchState.from_config(
context=context,
config=self.config,
max_iterations=2,
confidence_threshold=0.0,
)
state.search_filter = filter state.search_filter = filter
deps = DeepQADeps(client=self.client) deps = ResearchDeps(client=self.client)
if verbose: if verbose:
# Use AG-UI renderer to process and display events
from haiku.rag.graph.common.models import Citation
renderer = AGUIConsoleRenderer(self.console) renderer = AGUIConsoleRenderer(self.console)
result_dict = await renderer.render( result_dict = await renderer.render(
stream_graph(graph, state, deps) stream_graph(graph, state, deps)
) )
# Result should be a dict with 'answer' and 'citations' keys report = (
answer = result_dict.get("answer", "") if result_dict else "" ResearchReport.model_validate(result_dict)
if cite and result_dict: if result_dict
# Convert dicts to Citation objects else None
raw_citations = result_dict.get("citations", []) )
citations = [
Citation(**c) if isinstance(c, dict) else c
for c in raw_citations
]
else: else:
# Run without rendering events, just get the result report = await graph.run(state=state, deps=deps)
result = await graph.run(state=state, deps=deps)
answer = result.answer self.console.print(f"[bold blue]Question:[/bold blue] {question}")
if cite: self.console.print()
citations = result.citations 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: else:
answer, citations = await self.client.ask(question, filter=filter) answer, citations = await self.client.ask(question, filter=filter)
self.console.print(f"[bold blue]Question:[/bold blue] {question}") self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print() self.console.print()
self.console.print("[bold green]Answer:[/bold green]") self.console.print("[bold green]Answer:[/bold green]")
self.console.print(Markdown(answer)) self.console.print(Markdown(answer))
if cite and citations: if cite and citations:
for renderable in format_citations_rich(citations): for renderable in format_citations_rich(citations):
self.console.print(renderable) self.console.print(renderable)
except Exception as e: except Exception as e:
self.console.print(f"[red]Error: {e}[/red]") 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: if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument 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__) 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 ( from haiku.rag.graph.agui import (
AGUIConsoleRenderer, AGUIConsoleRenderer,
AGUIEmitter, AGUIEmitter,
create_agui_server, create_agui_server,
stream_graph, stream_graph,
) )
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.research.graph import build_research_graph from haiku.rag.graph.research.graph import build_research_graph
__all__ = [ __all__ = [
"AGUIConsoleRenderer", "AGUIConsoleRenderer",
"AGUIEmitter", "AGUIEmitter",
"build_deep_qa_graph",
"build_research_graph", "build_research_graph",
"create_agui_server", "create_agui_server",
"stream_graph", "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 def create_agui_server( # pragma: no cover
config: "AppConfig", db_path: Path | None = None config: "AppConfig", db_path: Path | None = None
) -> Starlette: ) -> Starlette:
"""Create AG-UI server with both research and deep ask endpoints. """Create AG-UI server with research endpoint.
Args: Args:
config: Application config with research and qa settings config: Application config with research settings
db_path: Optional database path override db_path: Optional database path override
Returns: Returns:
Starlette app with research and deep ask endpoints Starlette app with research endpoint
""" """
from haiku.rag.client import HaikuRAG 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.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import ResearchDeps, ResearchState from haiku.rag.graph.research.state import ResearchDeps, ResearchState
@ -192,7 +189,14 @@ def create_agui_server( # pragma: no cover
if messages: if messages:
question = messages[0].get("content", "") question = messages[0].get("content", "")
context = ResearchContext(original_question=question) 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: def research_deps_factory(input_config: dict[str, Any]) -> ResearchDeps:
effective_db_path = ( effective_db_path = (
@ -202,28 +206,7 @@ def create_agui_server( # pragma: no cover
) )
return ResearchDeps(client=get_client(effective_db_path)) return ResearchDeps(client=get_client(effective_db_path))
# Deep ask graph factories # Create event stream function
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
async def research_event_stream( async def research_event_stream(
input_data: RunAgentInput, input_data: RunAgentInput,
) -> AsyncIterator[str]: ) -> AsyncIterator[str]:
@ -236,18 +219,6 @@ def create_agui_server( # pragma: no cover
event_data = format_sse_event(event) event_data = format_sse_event(event)
yield event_data 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 # Endpoint handlers
async def stream_research(request: Request) -> StreamingResponse: async def stream_research(request: Request) -> StreamingResponse:
"""Research graph streaming endpoint.""" """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: async def health_check(_: Request) -> JSONResponse:
"""Health check endpoint.""" """Health check endpoint."""
return JSONResponse({"status": "healthy"}) return JSONResponse({"status": "healthy"})
@ -286,7 +242,6 @@ def create_agui_server( # pragma: no cover
# Define routes # Define routes
routes = [ routes = [
Route("/v1/research/stream", stream_research, methods=["POST"]), Route("/v1/research/stream", stream_research, methods=["POST"]),
Route("/v1/deep-ask/stream", stream_deep_ask, methods=["POST"]),
Route("/health", health_check, methods=["GET"]), 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.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.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 from haiku.rag.store.models import SearchResult
if TYPE_CHECKING:
from haiku.rag.graph.research.models import SearchAnswer
class ResearchContext(BaseModel): class ResearchContext(BaseModel):
"""Context shared across research agents.""" """Context shared across research agents."""
@ -19,124 +16,14 @@ class ResearchContext(BaseModel):
sub_questions: list[str] = Field( sub_questions: list[str] = Field(
default_factory=list, description="Decomposed sub-questions" 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" 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 def add_qa_response(self, qa: "SearchAnswer") -> None:
_insights_by_id: dict[str, InsightRecord] = PrivateAttr(default_factory=dict) """Add a structured QA response."""
_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)."""
self.qa_responses.append(qa) 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): class ResearchDependencies(BaseModel):
"""Dependencies for research agents with multi-agent context.""" """Dependencies for research agents with multi-agent context."""
@ -148,8 +35,3 @@ class ResearchDependencies(BaseModel):
search_results: list[SearchResult] = Field( search_results: list[SearchResult] = Field(
default_factory=list, description="Search results for citation resolution" 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 import Graph, GraphBuilder, StepContext
from pydantic_graph.beta.join import reduce_list_append from pydantic_graph.beta.join import reduce_list_append
from haiku.rag.config import Config from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
from haiku.rag.graph.common import get_model from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies
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.models import ( from haiku.rag.graph.research.models import (
EvaluationResult, EvaluationResult,
InsightAnalysis, RawSearchAnswer,
ResearchPlan,
ResearchReport, ResearchReport,
SearchAnswer,
) )
from haiku.rag.graph.research.prompts import ( from haiku.rag.graph.research.prompts import (
DECISION_AGENT_PROMPT, DECISION_PROMPT,
INSIGHT_AGENT_PROMPT, PLAN_PROMPT,
SYNTHESIS_AGENT_PROMPT, SEARCH_PROMPT,
SYNTHESIS_PROMPT,
) )
from haiku.rag.graph.research.state import ResearchDeps, ResearchState 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( def build_research_graph(
config: AppConfig = Config, config: AppConfig = Config,
include_plan: bool = True,
) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]: ) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]:
"""Build the Research graph. """Build the Research graph.
Args: Args:
config: AppConfig object (uses config.research for provider, model, and graph parameters) 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: Returns:
Configured Research graph Configured Research graph
@ -43,28 +71,172 @@ def build_research_graph(
output_type=ResearchReport, output_type=ResearchReport,
) )
# Create and register the plan node using the factory @g.step
plan = g.step( async def plan(ctx: StepContext[ResearchState, ResearchDeps, None]) -> None:
create_plan_node( """Create research plan with sub-questions."""
model_config=model_config, state = ctx.state
deps_type=ResearchDependencies, # type: ignore[arg-type] deps = ctx.deps
activity_message="Creating research plan",
output_retries=3,
config=config,
)
) # type: ignore[arg-type]
# Create and register the search_one node using the factory if deps.agui_emitter:
search_one = g.step( deps.agui_emitter.start_step("plan")
create_search_node( deps.agui_emitter.update_activity(
model_config=model_config, "planning", {"stepName": "plan", "message": "Creating research plan"}
deps_type=ResearchDependencies, # type: ignore[arg-type] )
with_step_wrapper=True,
success_message_format="Found answer with {confidence:.0%} confidence", try:
handle_exceptions=True, plan_agent = Agent(
config=config, model=get_model(model_config, config),
) output_type=ResearchPlan,
) # type: ignore[arg-type] 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 @g.step
async def get_batch( async def get_batch(
@ -76,84 +248,15 @@ def build_research_graph(
if not state.context.sub_questions: if not state.context.sub_questions:
return None return None
# Take ALL remaining questions and process them in parallel
batch = list(state.context.sub_questions) batch = list(state.context.sub_questions)
state.context.sub_questions.clear() state.context.sub_questions.clear()
return batch return batch
@g.step @g.step
async def analyze_insights( async def decide(
ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]], ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]],
) -> None: ) -> bool:
state = ctx.state """Evaluate research sufficiency and decide whether to continue."""
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:
state = ctx.state state = ctx.state
deps = ctx.deps deps = ctx.deps
@ -167,18 +270,16 @@ def build_research_graph(
agent = Agent( agent = Agent(
model=get_model(model_config, config), model=get_model(model_config, config),
output_type=EvaluationResult, output_type=EvaluationResult,
instructions=DECISION_AGENT_PROMPT, instructions=DECISION_PROMPT,
retries=3, retries=3,
output_retries=3, output_retries=3,
deps_type=ResearchDependencies, deps_type=ResearchDependencies,
) )
context_xml = format_context_for_prompt(state.context) context_xml = format_context_for_prompt(state.context)
analysis_xml = format_analysis_for_prompt(state.last_analysis)
prompt_parts = [ prompt_parts = [
"Assess whether the research now answers the original question with adequate confidence.", "Assess whether the research now answers the original question with adequate confidence.",
context_xml, context_xml,
analysis_xml,
] ]
if state.last_eval is not None: if state.last_eval is not None:
prev = state.last_eval prev = state.last_eval
@ -205,7 +306,6 @@ def build_research_graph(
if new_q not in state.context.sub_questions: if new_q not in state.context.sub_questions:
state.context.sub_questions.append(new_q) state.context.sub_questions.append(new_q)
# State updated with evaluation - emit state update and narrate
if deps.agui_emitter: if deps.agui_emitter:
deps.agui_emitter.update_state(state) deps.agui_emitter.update_state(state)
sufficient = "Yes" if output.is_sufficient else "No" sufficient = "Yes" if output.is_sufficient else "No"
@ -233,6 +333,7 @@ def build_research_graph(
async def synthesize( async def synthesize(
ctx: StepContext[ResearchState, ResearchDeps, None | bool], ctx: StepContext[ResearchState, ResearchDeps, None | bool],
) -> ResearchReport: ) -> ResearchReport:
"""Generate final research report."""
state = ctx.state state = ctx.state
deps = ctx.deps deps = ctx.deps
@ -246,7 +347,7 @@ def build_research_graph(
agent = Agent( agent = Agent(
model=get_model(model_config, config), model=get_model(model_config, config),
output_type=ResearchReport, output_type=ResearchReport,
instructions=SYNTHESIS_AGENT_PROMPT, instructions=SYNTHESIS_PROMPT,
retries=3, retries=3,
output_retries=3, output_retries=3,
deps_type=ResearchDependencies, deps_type=ResearchDependencies,
@ -274,12 +375,14 @@ def build_research_graph(
initial_factory=list[SearchAnswer], initial_factory=list[SearchAnswer],
) )
g.add( if include_plan:
g.edge_from(g.start_node).to(plan), g.add(
g.edge_from(plan).to(get_batch), 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.add(
g.edge_from(get_batch).to( g.edge_from(get_batch).to(
g.decision() g.decision()
@ -287,11 +390,9 @@ def build_research_graph(
.branch(g.match(type(None)).label("No questions").to(synthesize)) .branch(g.match(type(None)).label("No questions").to(synthesize))
), ),
g.edge_from(search_one).to(collect_answers), g.edge_from(search_one).to(collect_answers),
g.edge_from(collect_answers).to(analyze_insights), g.edge_from(collect_answers).to(decide),
g.edge_from(analyze_insights).to(decide),
) )
# Branch based on decision
g.add( g.add(
g.edge_from(decide).to( g.edge_from(decide).to(
g.decision() g.decision()

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

@ -1,149 +1,128 @@
import uuid from typing import TYPE_CHECKING
from enum import Enum
from pydantic import BaseModel, Field, field_validator from pydantic import BaseModel, Field, field_validator
if TYPE_CHECKING:
def _deduplicate_list(items: list[str]) -> list[str]: from haiku.rag.store.models import SearchResult
"""Remove duplicates while preserving order."""
return list(dict.fromkeys(items))
class InsightStatus(str, Enum): class ResearchPlan(BaseModel):
OPEN = "open" """A structured research plan with sub-questions to explore."""
VALIDATED = "validated"
TENTATIVE = "tentative"
sub_questions: list[str] = Field(
class GapSeverity(str, Enum): ...,
LOW = "low" description="Specific questions to research, phrased as complete questions",
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",
) )
@field_validator("supporting_sources", mode="before") @field_validator("sub_questions")
@classmethod @classmethod
def deduplicate_sources(cls, v: list[str]) -> list[str]: def validate_sub_questions(cls, v: list[str]) -> list[str]:
"""Ensure supporting_sources has no duplicates.""" if len(v) < 1:
return _deduplicate_list(v) if v else [] 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): class Citation(BaseModel):
"""Structured insight with provenance and lifecycle metadata.""" """Resolved citation with full metadata for display/visual grounding."""
summary: str = Field(description="Concise description of the insight") document_id: str
status: InsightStatus = Field( chunk_id: str
default=InsightStatus.OPEN, document_uri: str
description="Lifecycle status for the insight", document_title: str | None = None
) page_numbers: list[int] = Field(default_factory=list)
originating_questions: list[str] = Field( 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, default_factory=list,
description="Research sub-questions that produced this insight", description="IDs of chunks used to form the answer",
) )
confidence: float = Field(
@field_validator("originating_questions", mode="before") default=1.0,
@classmethod description="Confidence score for this answer (0-1)",
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)",
ge=0.0, ge=0.0,
le=1.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( is_sufficient: bool = Field(
description="Whether the research is sufficient to answer the original question" 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( reasoning: str = Field(
description="Explanation of why the research is or isn't complete" 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): 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 PLAN_PROMPT = """You are the research orchestrator for a focused, iterative workflow.
research 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: Inputs available:
- Original research question and sub-questions - Original research question
- Question–answer pairs with supporting snippets and sources - Question-answer pairs with supporting sources
- Existing insights and gaps (with status metadata) - Previous evaluation (if any)
Tasks: Tasks:
1. Extract new or refined insights that advance understanding of the question. 1. Assess whether the collected evidence answers the original question.
2. Update gap status, creating new gap entries when necessary and marking 2. Provide a confidence_score in [0,1] reflecting coverage and evidence quality.
resolved ones explicitly. 3. Optionally propose up to 3 new sub-questions if important gaps remain.
3. Suggest up to 3 high-impact follow-up sub_questions that would close the
most important remaining gaps.
Output format (map directly to fields): Output fields:
- highlights: list of insights with fields {summary, status, supporting_sources, - is_sufficient: true when the question is adequately answered
originating_questions, notes}. Use status one of {validated, open, tentative}. - confidence_score: numeric in [0,1]
supporting_sources and originating_questions must be lists of plain strings. - reasoning: brief explanation of the assessment
- gap_assessments: list of gaps with fields {description, severity, blocking, - new_questions: list of follow-up questions (max 3), only if needed
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.
All list fields must contain plain strings only, not objects. Be strict: only mark sufficient when key aspects are addressed with reliable evidence."""
Guidance: SYNTHESIS_PROMPT = """You are a synthesis specialist producing the final
- Be concise and avoid repeating previously recorded information unless it research report that directly answers the original question.
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.
Goals: 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. 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. 4. State limitations and uncertainties transparently.
Report guidelines (map to output fields): Report guidelines (map to output fields):
- title: concise (5–12 words), informative. - title: concise (5-12 words), informative.
- executive_summary: 3–5 sentences summarizing the overall answer. - executive_summary: 3-5 sentences that DIRECTLY ANSWER the original question.
- main_findings: list of plain strings, 4–8 one‑sentence bullets reflecting evidence. Write the actual answer, not a description of what the report contains.
- conclusions: list of plain strings, 2–4 bullets following logically from findings. BAD: "This report examines the topic and presents findings..."
- recommendations: list of plain strings, 2–5 actionable bullets tied to findings. GOOD: "The system requires configuration X and supports features Y and Z..."
- limitations: list of plain strings, 1–3 bullets describing constraints or uncertainties. - 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. - sources_summary: single string listing sources with document paths and page numbers.
All list fields must contain plain strings only, not objects. 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: Style:
- Base all content solely on the collected evidence. - Base all content solely on the collected evidence.
- Be professional, objective, and specific. - Be professional, objective, and specific.
- Avoid meta commentary and refrain from speculation beyond the evidence.""" - NEVER use meta-commentary like "This report covers..." or "The findings show...".
Instead, state the actual information directly."""
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)."""

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.client import HaikuRAG
from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.models import ( from haiku.rag.graph.research.models import EvaluationResult, ResearchReport
EvaluationResult,
InsightAnalysis,
ResearchReport,
)
if TYPE_CHECKING: if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
@ -26,12 +22,7 @@ class ResearchDeps:
semaphore: asyncio.Semaphore | None = None semaphore: asyncio.Semaphore | None = None
def emit_log(self, message: str, state: "ResearchState | None" = None) -> None: def emit_log(self, message: str, state: "ResearchState | None" = None) -> None:
"""Emit a log message through AG-UI events. """Emit a log message through AG-UI events."""
Args:
message: The message to log
state: Optional state to include in state update
"""
if self.agui_emitter: if self.agui_emitter:
self.agui_emitter.log(message) self.agui_emitter.log(message)
if state: if state:
@ -39,15 +30,12 @@ class ResearchDeps:
class ResearchState(BaseModel): class ResearchState(BaseModel):
"""Research graph state model. """Research graph state model."""
Fully JSON-serializable Pydantic model suitable for AG-UI state synchronization.
"""
model_config = {"arbitrary_types_allowed": True} model_config = {"arbitrary_types_allowed": True}
context: ResearchContext = Field( 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") iterations: int = Field(default=0, description="Current iteration number")
max_iterations: int = Field(default=3, description="Maximum allowed iterations") max_iterations: int = Field(default=3, description="Maximum allowed iterations")
@ -60,29 +48,33 @@ class ResearchState(BaseModel):
last_eval: EvaluationResult | None = Field( last_eval: EvaluationResult | None = Field(
default=None, description="Last evaluation result" default=None, description="Last evaluation result"
) )
last_analysis: InsightAnalysis | None = Field(
default=None, description="Last insight analysis"
)
search_filter: str | None = Field( search_filter: str | None = Field(
default=None, description="SQL WHERE clause to filter search results" default=None, description="SQL WHERE clause to filter search results"
) )
@classmethod @classmethod
def from_config( def from_config(
cls, context: ResearchContext, config: "AppConfig" cls,
context: ResearchContext,
config: "AppConfig",
max_iterations: int | None = None,
confidence_threshold: float | None = None,
) -> "ResearchState": ) -> "ResearchState":
"""Create a ResearchState from an AppConfig. """Create a ResearchState from an AppConfig.
Args: Args:
context: The ResearchContext containing the question and settings context: The ResearchContext containing the question
config: The AppConfig object (uses config.research for state parameters) config: The AppConfig object
max_iterations: Override max iterations (None uses config default)
Returns: confidence_threshold: Override threshold (None uses config, 0.0 disables check)
A configured ResearchState instance
""" """
return cls( return cls(
context=context, context=context,
max_iterations=config.research.max_iterations, max_iterations=max_iterations
confidence_threshold=config.research.confidence_threshold, 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, 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: try:
async with HaikuRAG(db_path, config=config) as rag: async with HaikuRAG(db_path, config=config) as rag:
if deep: if deep:
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState from haiku.rag.graph.research.state import (
ResearchDeps,
ResearchState,
)
graph = build_deep_qa_graph(config=config) graph = build_research_graph(config=config)
context = DeepQAContext(original_question=question) context = ResearchContext(original_question=question)
state = DeepQAState.from_config(context=context, config=config) state = ResearchState.from_config(
deps = DeepQADeps(client=rag) context=context,
config=config,
max_iterations=2,
confidence_threshold=0.0,
)
deps = ResearchDeps(client=rag)
result = await graph.run(state=state, deps=deps) result = await graph.run(state=state, deps=deps)
answer = result.answer answer = result.executive_summary
citations = result.citations citations = []
else: else:
answer, citations = await rag.ask(question) answer, citations = await rag.ask(question)
if cite and citations: 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.client import HaikuRAG
from haiku.rag.config import Config from haiku.rag.config import Config
from haiku.rag.config.models import ModelConfig from haiku.rag.config.models import ModelConfig
from haiku.rag.graph.common import get_model from haiku.rag.graph.research.models import Citation, RawSearchAnswer, resolve_citations
from haiku.rag.graph.common.models import Citation, RawSearchAnswer, resolve_citations
from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT
from haiku.rag.store.models import SearchResult from haiku.rag.store.models import SearchResult
from haiku.rag.utils import get_model
class Dependencies(BaseModel): 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: if TYPE_CHECKING:
from rich.console import RenderableType 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( 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 # Patch all locations where get_model is imported
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory) 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) monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
graph = build_research_graph() 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 pydantic_ai.models.test import TestModel
from haiku.rag.client import HaikuRAG 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.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import ResearchDeps, ResearchState 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() return TestModel()
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory) 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) monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
graph = build_research_graph() 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 @pytest.mark.asyncio
async def test_search_filter_none_searches_all(monkeypatch, client_with_docs): async def test_search_filter_none_searches_all(monkeypatch, client_with_docs):
"""Test that search_filter=None searches all documents.""" """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() return TestModel()
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory) 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) monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
graph = build_research_graph() 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 @pytest.mark.asyncio
async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch): async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch):
"""Test asking a question with citations.""" """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_answer = "Test answer with citations"
mock_citations = [ mock_citations = [
@ -358,25 +358,36 @@ async def test_ask_with_verbose(app: HaikuRAGApp, monkeypatch):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch): async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
"""Test asking a question with deep QA.""" """Test asking a question with deep mode uses research graph."""
from haiku.rag.graph.deep_qa.models import DeepQAAnswer 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 = AsyncMock()
mock_graph.run.return_value = mock_output mock_graph.run.return_value = mock_output
mock_client = AsyncMock() mock_client = AsyncMock()
mock_client.__aenter__.return_value = mock_client
mock_print = MagicMock() mock_print = MagicMock()
monkeypatch.setattr(app.console, "print", mock_print) 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.app.HaikuRAG") as mock_rag_class:
with patch( mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_client)
"haiku.rag.graph.deep_qa.graph.build_deep_qa_graph", return_value=mock_graph mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
): await app.ask("test question", deep=True)
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() mock_graph.run.assert_called_once()
call_kwargs = mock_graph.run.call_args[1] call_kwargs = mock_graph.run.call_args[1]
@ -385,25 +396,31 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch): async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch):
"""Test asking a question with deep QA and citations (cite ignored for deep).""" """Test asking a question with deep mode (cite is ignored for research graph)."""
from haiku.rag.graph.deep_qa.models import DeepQAAnswer 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 = AsyncMock()
mock_graph.run.return_value = mock_output mock_graph.run.return_value = mock_output
mock_client = AsyncMock() mock_client = AsyncMock()
mock_client.__aenter__.return_value = mock_client
mock_print = MagicMock() mock_print = MagicMock()
monkeypatch.setattr(app.console, "print", mock_print) 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.app.HaikuRAG") as mock_rag_class:
with patch( mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_client)
"haiku.rag.graph.deep_qa.graph.build_deep_qa_graph", return_value=mock_graph mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
): await app.ask("test question", deep=True, cite=True)
await app.ask("test question", deep=True, cite=True)
mock_graph.run.assert_called_once() mock_graph.run.assert_called_once()
call_kwargs = mock_graph.run.call_args[1] 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 @pytest.mark.asyncio
async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch): 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 = AsyncMock()
mock_renderer.render.return_value = mock_output 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_graph = AsyncMock()
mock_client = AsyncMock() mock_client = AsyncMock()
mock_client.__aenter__.return_value = mock_client
mock_print = MagicMock() mock_print = MagicMock()
monkeypatch.setattr(app.console, "print", mock_print) 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.app.HaikuRAG") as mock_rag_class:
with patch( mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_client)
"haiku.rag.graph.deep_qa.graph.build_deep_qa_graph", return_value=mock_graph mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
): with patch("haiku.rag.app.AGUIConsoleRenderer", return_value=mock_renderer):
with patch("haiku.rag.app.AGUIConsoleRenderer", return_value=mock_renderer): await app.ask("test question", deep=True, verbose=True)
await app.ask("test question", deep=True, verbose=True)
# With verbose, it should use AGUIConsoleRenderer.render, not graph.run # With verbose, it should use AGUIConsoleRenderer.render, not graph.run
mock_renderer.render.assert_called_once() 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 @pytest.mark.asyncio
async def test_mcp_ask_question_deep(): 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: with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test.lancedb" db_path = Path(temp_dir) / "test.lancedb"
mcp = create_mcp_server(db_path) mcp = create_mcp_server(db_path)
@ -251,7 +251,7 @@ async def test_mcp_ask_question_deep():
with ( with (
patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class, patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class,
patch( patch(
"haiku.rag.graph.deep_qa.graph.build_deep_qa_graph" "haiku.rag.graph.research.graph.build_research_graph"
) as mock_graph_builder, ) as mock_graph_builder,
): ):
mock_rag = AsyncMock() mock_rag = AsyncMock()
@ -260,8 +260,7 @@ async def test_mcp_ask_question_deep():
mock_graph = AsyncMock() mock_graph = AsyncMock()
mock_result = AsyncMock() mock_result = AsyncMock()
mock_result.answer = "Deep answer" mock_result.executive_summary = "Deep answer from research"
mock_result.citations = []
mock_graph.run = AsyncMock(return_value=mock_result) mock_graph.run = AsyncMock(return_value=mock_result)
mock_graph_builder.return_value = mock_graph 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 question="Deep question?", cite=False, deep=True
) )
assert result == "Deep answer" assert result == "Deep answer from research"
mock_graph.run.assert_called_once() mock_graph.run.assert_called_once()