Drop the multi-agent research workflow
This commit is contained in:
Yiorgis Gozadinos
parent
a317a951d9
commit
8c57a3ca99
57 changed files with 130 additions and 4373 deletions
14
CHANGELOG.md
14
CHANGELOG.md
|
|
@ -4,13 +4,19 @@
|
|||
### Added
|
||||
|
||||
- `heading_level` and `tree_depth` on `DocumentItem`, populated by `extract_items` and persisted on `document_items`. 0.48.0 migration backfills existing rows from each doc's docling structure blob.
|
||||
- `toc.json` in the analysis sandbox VFS at `/documents/{id}/toc.json`. Nested tree on HTML/markdown sources, flat sibling list on PDFs. `items.jsonl` rows now include `heading_level` and `tree_depth`.
|
||||
- `toc.json` in the analysis sandbox VFS at `/documents/{id}/toc.json`. Nested tree on HTML/markdown sources, flat sibling list on PDFs. Each node carries `{self_ref, level, title, page_numbers, item_range, chunk_ids, children}`. `chunk_ids` aggregates the citable chunks across the section's `item_range`, so the analysis skill can `cite()` a section from one VFS read instead of falling back to a corpus-wide `search()` that risks cross-document hits.
|
||||
- `chunk_ids` on every `items.jsonl` row (the citable chunks that contain that item).
|
||||
- `picture_refs` on sandbox `search()` result dicts and on `Citation` (subset of `doc_item_refs` starting with `#/pictures/`).
|
||||
- `picture_captions: dict[str, str]` on `SearchResult`, populated alongside `image_data` and rendered as a labelled line in `format_for_agent` for picture-bearing chunks.
|
||||
- Chat TUI renders picture citations inline via `textual_image.widget.Image` inside the existing `CitationWidget`.
|
||||
- CLI citation panel renders `picture_refs` inline via `textual_image.renderable.Image` next to the text preview. `format_citations_rich` is async and takes an optional `HaikuRAG` client; without one, figures fall back to `[Figure: <ref>]` markers.
|
||||
- BTree scalar indexes on `document_items.{document_id, position, self_ref}`. The 0.48.0 migration creates them on existing DBs. Per-doc lookups go from full-table scans (~100–300 ms) to point queries (~3–21 ms) on small/medium corpora.
|
||||
- Per-doc lazy cache for `items.jsonl` and `toc.json` in the analysis sandbox. First read fetches; subsequent reads of either file in the same `execute_code` session hit a serialized cache. One DB fetch per doc per session.
|
||||
- `cite` tool now accepts chunk_ids that resolve via the chunks table, not only chunk_ids from a prior `search()` result. Lets the model cite directly from `items.jsonl` / `toc.json` rows. The hallucination guard (`ModelRetry` on chunk_ids that don't exist in the DB) is preserved.
|
||||
|
||||
### Removed
|
||||
|
||||
- **Multi-agent research workflow.** Removes `agents/research/` (graph, state, deps, models, prompts), `client.research`, the CLI `research` command, the MCP `research_question` tool, `ResearchConfig`, `AppConfig.research`, `PromptsConfig.synthesis`, and the corresponding wiring in `chat/__init__.py` and `client/downloads.py`. The pydantic-graph workflow was a three-node loop whose differentiators vs the rag skill (planner step, structured `ResearchReport`, iteration bound) were either redundant with `qa.max_searches` or sat on pre-cite-tool legacy. Multi-step questions go through `client.ask` (rag skill).
|
||||
- `llm()` from the analysis sandbox. Sandbox externals are now `search` and `list_documents` only.
|
||||
- `list_documents` top-level tool from the analysis skill (still available as `await list_documents()` inside `execute_code`).
|
||||
- `documents=` kwarg on `client.analyze` (and the `--document` flag on `haiku-rag analyze` / MCP `analyze` tool). The pre-loaded `documents` Python variable inside the sandbox is no longer populated. Use `filter=` (SQL WHERE clause) to scope analysis to specific documents.
|
||||
|
|
@ -21,12 +27,18 @@
|
|||
- `prompts.qa` config field.
|
||||
- `evaluations optimize` subcommand and GEPA prompt-optimization. Drops `gepa` dep.
|
||||
- `--target qa` from `evaluations run`. Default is now `rag-skill`.
|
||||
- `position` field on `toc.json` nodes (redundant with `item_range[0]`).
|
||||
- `position` and `tree_depth` from `items.jsonl` row serialization. Both fields are still persisted on `DocumentItem`; they are no longer surfaced to the sandbox.
|
||||
|
||||
### Changed
|
||||
|
||||
- `haiku.rag.agents.analysis` moved to `haiku.rag.sandbox`. Public surface: `from haiku.rag.sandbox import Sandbox, SandboxResult, AnalysisContext, AnalysisResult`.
|
||||
- `Citation` and `resolve_citations` moved to `haiku.rag.store.models.citation` (was `haiku.rag.agents.research.models`), peer to the other output domain models.
|
||||
- `search.limit` default lowered from `10` to `5`.
|
||||
- Search result formatter surfaces picture captions on a labelled line when a chunk's expanded refs include pictures.
|
||||
- Picture bytes attached to `search()` results are bounded to the pre-expansion chunk's `doc_item_refs`. Section expansion that sweeps in adjacent picture-bearing items no longer pulls their bytes into the response. Observed ~16× reduction on tool-response payload sizes; eliminates a class of cross-figure contamination in the agent's view.
|
||||
- rag-analysis SKILL.md steers structural lookups ("which section X", "list sections of Y", "summarize section Z") to read `/documents/{id}/toc.json` first and cite the matching node's `chunk_ids` directly, instead of calling `search()`. Empirically validated on the ORB multimodal and Wix corpora: locator and orientation questions now cite the correct document instead of falling back to cross-document search hits.
|
||||
- CLI citation panel compacted: 300-char text preview (no full chunk dump), `[N] Title (URI) — pp. — §Section` header, dimmed `doc: <id> chunk: <id>` footer. Green "Citations" label matches the green "Answer:" label.
|
||||
- `AnalysisConfig.model` defaults to `None` (was `ollama:gpt-oss/no-thinking/temp=0`). Resolves via `config.analysis.model or config.qa.model`.
|
||||
- `client.ask` and `client.analyze` route through the rag and rag-analysis skills internally.
|
||||
- Bump `docling>=2.93.0` and `docling-core>=2.75.0`.
|
||||
|
|
|
|||
16
README.md
16
README.md
|
|
@ -14,11 +14,10 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p
|
|||
- **Question answering** — RAG skill with citations (page numbers, section headings)
|
||||
- **Vision QA** — Vision-capable models receive figure bytes alongside chunk text
|
||||
- **Reranking** — MxBAI, Cohere, Zero Entropy, or vLLM
|
||||
- **Research workflow** — Multi-agent pydantic-graph: plan, search, evaluate, synthesize
|
||||
- **Analysis skill** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
|
||||
- **Conversational RAG** — Chat TUI and web application for multi-turn conversations with session memory
|
||||
- **Document structure** — Stores full [DoclingDocument](https://docling-project.github.io/docling/concepts/docling_document/), enabling structure-aware context expansion
|
||||
- **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, LM Studio, vLLM (multimodal). QA/Research: any model supported by Pydantic AI
|
||||
- **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, LM Studio, vLLM (multimodal). QA: any model supported by Pydantic AI
|
||||
- **Local-first** — Embedded LanceDB, no servers required. Also supports S3, GCS, Azure, and LanceDB Cloud
|
||||
- **CLI & Python API** — Full functionality from command line or code
|
||||
- **MCP server** — Expose as tools for AI assistants (Claude Desktop, etc.)
|
||||
|
|
@ -63,9 +62,6 @@ haiku-rag search "attention mechanism"
|
|||
# Ask questions with citations
|
||||
haiku-rag ask "What datasets were used for evaluation?" --cite
|
||||
|
||||
# Research mode — iterative planning and search
|
||||
haiku-rag research "What are the limitations of the approach?"
|
||||
|
||||
# Analyze — complex analytical tasks via code execution
|
||||
haiku-rag analyze "How many documents mention transformers?"
|
||||
|
||||
|
|
@ -83,7 +79,7 @@ See [Configuration](https://ggozad.github.io/haiku.rag/configuration/) for custo
|
|||
```python
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
async with HaikuRAG("research.lancedb", create=True) as rag:
|
||||
async with HaikuRAG("knowledge.lancedb", create=True) as rag:
|
||||
# Index documents
|
||||
await rag.create_document_from_source("paper.pdf")
|
||||
await rag.create_document_from_source("https://arxiv.org/pdf/1706.03762")
|
||||
|
|
@ -100,7 +96,7 @@ async with HaikuRAG("research.lancedb", create=True) as rag:
|
|||
print(f" [{cite.chunk_id}] p.{cite.page_numbers}: {cite.content[:80]}")
|
||||
```
|
||||
|
||||
For research agents and chat, see the [Agents docs](https://ggozad.github.io/haiku.rag/agents/).
|
||||
For details on the skills the client wraps, see the [Skills docs](https://ggozad.github.io/haiku.rag/skills/).
|
||||
|
||||
## MCP Server
|
||||
|
||||
|
|
@ -123,7 +119,7 @@ Add to your Claude Desktop configuration:
|
|||
}
|
||||
```
|
||||
|
||||
Provides tools for document management, search, QA, and research directly in your AI assistant.
|
||||
Provides tools for document management, search, QA, and analysis directly in your AI assistant.
|
||||
|
||||
## Examples
|
||||
|
||||
|
|
@ -141,8 +137,8 @@ Full documentation at: https://ggozad.github.io/haiku.rag/
|
|||
- [Configuration](https://ggozad.github.io/haiku.rag/configuration/) - YAML configuration
|
||||
- [CLI](https://ggozad.github.io/haiku.rag/cli/) - Command reference
|
||||
- [Python API](https://ggozad.github.io/haiku.rag/python/) - Complete API docs
|
||||
- [Agents](https://ggozad.github.io/haiku.rag/agents/) - QA and research agents
|
||||
- [Analysis Agent](https://ggozad.github.io/haiku.rag/agents/analysis/) - Complex analytical tasks via code execution
|
||||
- [Skills](https://ggozad.github.io/haiku.rag/skills/) - The RAG and analysis skills the client wraps
|
||||
- [Analysis](https://ggozad.github.io/haiku.rag/agents/analysis/) - Complex analytical tasks via code execution
|
||||
- [Applications](https://ggozad.github.io/haiku.rag/apps/) - Chat TUI, web app, and inspector
|
||||
- [Server](https://ggozad.github.io/haiku.rag/server/) - File monitoring and MCP
|
||||
- [MCP](https://ggozad.github.io/haiku.rag/mcp/) - Model Context Protocol integration
|
||||
|
|
|
|||
|
|
@ -1,142 +0,0 @@
|
|||
# Agents
|
||||
|
||||
haiku.rag provides:
|
||||
|
||||
- **Question Answering** via `client.ask` and the [RAG skill](../skills/index.md) — search + cite over the knowledge base.
|
||||
- **Analysis** via `client.analyze` and the [analysis skill](../skills/index.md) — sandboxed Python code execution (see [Analysis](analysis.md)).
|
||||
- **Research Graph** — a multi-step research workflow with question decomposition (this page).
|
||||
|
||||
`client.ask` and `client.analyze` are thin wrappers over the rag and rag-analysis skills built on [haiku.skills](https://github.com/ggozad/haiku.skills). For multi-turn conversational RAG with the same primitives, use the skills directly via `SkillToolset`.
|
||||
|
||||
See [QA and Research Configuration](../configuration/qa-research.md) for configuring model, iterations, concurrency, and other settings.
|
||||
|
||||
## Question Answering
|
||||
|
||||
```bash
|
||||
haiku-rag ask "What is climate change?"
|
||||
```
|
||||
|
||||
```python
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
async with HaikuRAG(path_to_db) as client:
|
||||
answer, citations = await client.ask("What is climate change?")
|
||||
```
|
||||
|
||||
Citations are always returned in the second element of the tuple and rendered after the answer on the CLI.
|
||||
|
||||
## Research Graph
|
||||
|
||||
The research workflow is implemented as a typed pydantic-graph. It uses an iterative feedback loop where the planner proposes one question at a time, sees the answer, then decides whether to continue or synthesize.
|
||||
|
||||
```mermaid
|
||||
---
|
||||
title: Research graph
|
||||
---
|
||||
stateDiagram-v2
|
||||
state plan_next_decision <<choice>>
|
||||
[*] --> plan_next
|
||||
plan_next --> plan_next_decision
|
||||
plan_next_decision --> search_one: Has next question
|
||||
plan_next_decision --> synthesize: Complete or max iterations
|
||||
search_one --> plan_next: Answer added to context
|
||||
synthesize --> [*]
|
||||
|
||||
note right of plan_next
|
||||
Uses prior_answers from previous iterations.
|
||||
Uses a different prompt when prior answers exist.
|
||||
end note
|
||||
```
|
||||
|
||||
The graph receives a `ResearchContext` containing:
|
||||
|
||||
- `original_question` — the user's question
|
||||
- `qa_responses` — prior answers from previous iterations (injected as `<prior_answers>` XML)
|
||||
|
||||
When prior answers are provided, the planner uses a context-aware prompt that evaluates whether existing evidence is sufficient. If it is, the planner marks `is_complete=True` and the graph skips directly to synthesis without any searches.
|
||||
|
||||
**Key nodes:**
|
||||
|
||||
- **plan_next**: Evaluates gathered evidence and either proposes the next question to investigate or marks research as complete. Uses a context-aware prompt when prior answers exist, allowing it to skip research entirely.
|
||||
- **search_one**: Answers a single question using the knowledge base (up to 3 search calls per question). Each answer is added to `ResearchContext.qa_responses` for the next planning iteration.
|
||||
- **synthesize**: Generates the final output from all gathered evidence.
|
||||
|
||||
**Iterative flow:**
|
||||
|
||||
- Each iteration: planner evaluates context → proposes one question → search answers it → loop back
|
||||
- Planner can decompose complex questions (e.g., "benefits and drawbacks" → start with "benefits")
|
||||
- Prior answers let the planner skip redundant searches
|
||||
- Loop terminates when planner marks `is_complete=True` or `max_iterations` is reached
|
||||
|
||||
### CLI Usage
|
||||
|
||||
```bash
|
||||
# Basic usage
|
||||
haiku-rag research "How does haiku.rag organize and query documents?"
|
||||
|
||||
# With document filter
|
||||
haiku-rag research "What are the key findings?" --filter "uri LIKE '%report%'"
|
||||
```
|
||||
|
||||
### Python Usage
|
||||
|
||||
**Basic example:**
|
||||
|
||||
```python
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.agents.research.dependencies import ResearchContext
|
||||
from haiku.rag.agents.research.graph import build_research_graph
|
||||
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
|
||||
|
||||
async with HaikuRAG(path_to_db) as client:
|
||||
graph = build_research_graph(config=Config)
|
||||
context = ResearchContext(original_question="What are the main features?")
|
||||
state = ResearchState.from_config(context=context, config=Config)
|
||||
deps = ResearchDeps(client=client)
|
||||
|
||||
report = await graph.run(state=state, deps=deps)
|
||||
|
||||
print(report.title)
|
||||
print(report.executive_summary)
|
||||
```
|
||||
|
||||
**With custom config:**
|
||||
|
||||
```python
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config.models import AppConfig, ModelConfig, ResearchConfig
|
||||
from haiku.rag.agents.research.dependencies import ResearchContext
|
||||
from haiku.rag.agents.research.graph import build_research_graph
|
||||
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
|
||||
|
||||
custom_config = AppConfig(
|
||||
research=ResearchConfig(
|
||||
model=ModelConfig(provider="openai", name="gpt-4o-mini"),
|
||||
max_iterations=5,
|
||||
max_concurrency=3,
|
||||
)
|
||||
)
|
||||
|
||||
async with HaikuRAG(path_to_db) as client:
|
||||
graph = build_research_graph(config=custom_config)
|
||||
context = ResearchContext(original_question="What are the main features?")
|
||||
state = ResearchState.from_config(context=context, config=custom_config)
|
||||
deps = ResearchDeps(client=client)
|
||||
|
||||
report = await graph.run(state=state, deps=deps)
|
||||
```
|
||||
|
||||
### 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.
|
||||
22
docs/cli.md
22
docs/cli.md
|
|
@ -231,26 +231,6 @@ The inspector provides:
|
|||
|
||||
See [Applications](apps.md#inspector) for details.
|
||||
|
||||
## Research
|
||||
|
||||
Run the multi-step research graph:
|
||||
|
||||
```bash
|
||||
haiku-rag research "How does haiku.rag organize and query documents?"
|
||||
```
|
||||
|
||||
Filter to specific documents:
|
||||
|
||||
```bash
|
||||
haiku-rag research "What are the key findings?" --filter "uri LIKE '%paper%'"
|
||||
```
|
||||
|
||||
Flags:
|
||||
|
||||
- `--filter` / `-f`: SQL WHERE clause to filter documents (see [Filtering Search Results](python.md#filtering-search-results))
|
||||
|
||||
Research parameters like `max_iterations` and `max_concurrency` are configured in your [configuration file](configuration/index.md) under the `research` section.
|
||||
|
||||
## Analyze
|
||||
|
||||
Answer complex analytical questions via code execution:
|
||||
|
|
@ -510,7 +490,7 @@ This command downloads:
|
|||
|
||||
- Docling OCR/conversion models
|
||||
- HuggingFace tokenizer (for chunking)
|
||||
- Ollama models referenced in your configuration (embeddings, QA, research, rerank)
|
||||
- Ollama models referenced in your configuration (embeddings, QA, rerank)
|
||||
|
||||
Progress is displayed in real-time with download status and progress bars for Ollama model pulls.
|
||||
|
||||
|
|
|
|||
|
|
@ -90,15 +90,6 @@ qa:
|
|||
temperature: 0.3
|
||||
max_searches: 3
|
||||
|
||||
research:
|
||||
model:
|
||||
provider: "" # Empty to use qa settings
|
||||
name: ""
|
||||
enable_thinking: false
|
||||
temperature: 0.3
|
||||
max_iterations: 3
|
||||
max_concurrency: 1
|
||||
|
||||
search:
|
||||
limit: 10 # Default number of results to return
|
||||
max_context_chars: 10000 # Maximum characters in expanded context
|
||||
|
|
@ -106,8 +97,7 @@ search:
|
|||
vector_refine_factor: 30
|
||||
|
||||
prompts:
|
||||
domain_preamble: "" # Prepended to skill instructions and research prompts
|
||||
synthesis: null # Custom research synthesis prompt (null = use default)
|
||||
domain_preamble: "" # Prepended to skill instructions
|
||||
|
||||
processing:
|
||||
converter: docling-local # docling-local or docling-serve
|
||||
|
|
@ -189,7 +179,7 @@ This is useful for:
|
|||
For detailed configuration of specific topics, see:
|
||||
|
||||
- **[Providers](providers.md)** - Model settings and provider-specific configuration (embeddings, reranking)
|
||||
- **[Search and Question Answering](qa-research.md)** - Search settings, question answering, and research workflows
|
||||
- **[Search and Question Answering](qa.md)** - Search settings and question answering
|
||||
- **[Document Processing](processing.md)** - Document conversion, chunking, and file monitoring
|
||||
- **[Storage](storage.md)** - Database, remote storage, and vector indexing
|
||||
- **[Prompts](prompts.md)** - Customize agent prompts for your domain
|
||||
|
|
|
|||
|
|
@ -289,7 +289,7 @@ processing:
|
|||
chunk_size: 256 # Maximum tokens per chunk
|
||||
```
|
||||
|
||||
Context expansion settings (for enriching search results with surrounding content) are configured in the `search` section. See [Search Settings](qa-research.md#search-settings).
|
||||
Context expansion settings (for enriching search results with surrounding content) are configured in the `search` section. See [Search Settings](qa.md#search-settings).
|
||||
|
||||
## File Monitoring
|
||||
|
||||
|
|
|
|||
|
|
@ -1,27 +1,24 @@
|
|||
# Prompt Customization
|
||||
|
||||
Customize the prompts used by haiku.rag's skills and research workflow to better match your domain and use case.
|
||||
Customize the prompts used by haiku.rag's skills to better match your domain and use case.
|
||||
|
||||
## Configuration
|
||||
|
||||
```yaml
|
||||
prompts:
|
||||
# Domain context prepended to skill instructions and research prompts
|
||||
# Domain context prepended to skill instructions
|
||||
domain_preamble: |
|
||||
This knowledge base contains technical documentation for the Helios solar panel
|
||||
system, including installation manuals, maintenance procedures, and safety guidelines.
|
||||
Questions about "the system" or unqualified specs refer to the Helios panel.
|
||||
|
||||
# Full replacement for research synthesis prompt (optional)
|
||||
synthesis: null
|
||||
|
||||
# VLM prompt for image description during conversion (optional)
|
||||
picture_description: null # Uses default prompt
|
||||
```
|
||||
|
||||
## Domain Preamble
|
||||
|
||||
The `domain_preamble` field provides **domain context** prepended to the rag and rag-analysis skill instructions and to the research planner/search/synthesis prompts. Use this to:
|
||||
The `domain_preamble` field provides **domain context** prepended to the rag and rag-analysis skill instructions. Use this to:
|
||||
|
||||
- Describe what the knowledge base contains
|
||||
- Clarify domain-specific terminology
|
||||
|
|
@ -39,34 +36,6 @@ prompts:
|
|||
"Deployment" refers to Acme's managed deployment service, not general CI/CD.
|
||||
```
|
||||
|
||||
## Custom Synthesis Prompt
|
||||
|
||||
Replace the research report synthesis prompt by setting `prompts.synthesis`. This controls how the multi-agent research workflow generates its final report.
|
||||
|
||||
The prompt should produce a `ResearchReport` with: `title`, `executive_summary`, `main_findings`, `conclusions`, `recommendations`, `limitations`, and `sources_summary`.
|
||||
|
||||
**Example:**
|
||||
|
||||
```yaml
|
||||
prompts:
|
||||
synthesis: |
|
||||
Generate a research report based on the gathered evidence.
|
||||
|
||||
Output format:
|
||||
- title: 5-12 word title
|
||||
- executive_summary: 3-5 sentence overview
|
||||
- main_findings: 4-8 bullet points of key findings
|
||||
- conclusions: 2-4 bullet points
|
||||
- recommendations: 2-5 actionable recommendations
|
||||
- limitations: 1-3 limitations or gaps
|
||||
- sources_summary: Brief description of sources used
|
||||
|
||||
Guidelines:
|
||||
- Base all content strictly on collected evidence
|
||||
- Be specific and objective
|
||||
- Avoid meta-commentary like "This report covers..."
|
||||
```
|
||||
|
||||
## Picture Description Prompt
|
||||
|
||||
Customize the prompt used when generating VLM descriptions for embedded images during document conversion. This prompt is sent to the configured Vision Language Model for each image.
|
||||
|
|
@ -99,7 +68,6 @@ from haiku.rag.config.models import PromptsConfig
|
|||
config = AppConfig(
|
||||
prompts=PromptsConfig(
|
||||
domain_preamble="This knowledge base contains Acme Corp product documentation and API references.",
|
||||
synthesis=None, # Use default synthesis prompt
|
||||
picture_description="Describe this image for search indexing.",
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ haiku.rag supports multiple AI providers for embeddings, question answering, and
|
|||
|
||||
## Model Settings
|
||||
|
||||
Configure model behavior for `qa` and `research` workflows. These settings apply to any provider that supports them.
|
||||
Configure model behavior for the `qa` and `analysis` skills. These settings apply to any provider that supports them.
|
||||
|
||||
### Basic Settings
|
||||
|
||||
|
|
@ -22,7 +22,7 @@ qa:
|
|||
|
||||
**Available options:**
|
||||
|
||||
- **temperature**: Sampling temperature (0.0-1.0+). Defaults vary by task: 0.3 for QA, research, and title generation; 0.0 for analysis and picture description.
|
||||
- **temperature**: Sampling temperature (0.0-1.0+). Defaults vary by task: 0.3 for QA and title generation; 0.0 for analysis and picture description.
|
||||
- Lower (0.0-0.3): Deterministic, focused responses
|
||||
- Medium (0.4-0.7): Balanced
|
||||
- Higher (0.8-1.0+): Creative, varied responses
|
||||
|
|
@ -39,10 +39,6 @@ The `enable_thinking` setting controls whether models use explicit reasoning ste
|
|||
qa:
|
||||
model:
|
||||
enable_thinking: true # Better grounded answers
|
||||
|
||||
research:
|
||||
model:
|
||||
enable_thinking: true # Deeper reasoning
|
||||
```
|
||||
|
||||
**Values:**
|
||||
|
|
@ -64,7 +60,7 @@ See the [Pydantic AI thinking documentation](https://ai.pydantic.dev/thinking/)
|
|||
- **LM Studio**: Models supporting reasoning (gpt-oss, etc.)
|
||||
|
||||
**When to use:**
|
||||
- Enable for QA, research, complex reasoning, and mathematical problems
|
||||
- Enable for QA, complex reasoning, and mathematical problems
|
||||
- Disable for speed-critical applications, title generation, and simple tasks
|
||||
|
||||
### Raw Provider Pass-through
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ search:
|
|||
max_context_chars: 10000 # Maximum characters in expanded context
|
||||
```
|
||||
|
||||
- **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, QA, and research workflows. Default: 10
|
||||
- **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, and QA. Default: 10
|
||||
- **max_context_chars**: Hard limit on total characters in expanded content. Default: 10000.
|
||||
|
||||
Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers) — this naturally crosses into adjacent sections until the budget is filled. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
|
||||
|
|
@ -37,27 +37,6 @@ qa:
|
|||
- **model.vision**: Set to `true` for vision-capable models (`qwen2.5vl`, `qwen3.6`, `gpt-4o`, `claude-sonnet`, …). The skill's `search` tool only attaches picture bytes (`BinaryContent`) to its `ToolReturn` when this is `true`; otherwise picture bytes are withheld. See [Pictures × embedder × QA model](processing.md#pictures--embedder--qa-model-how-the-pieces-compose) for the full matrix.
|
||||
- **max_searches**: Maximum number of search tool calls the rag skill can make per question (default: 3)
|
||||
|
||||
## Research Configuration
|
||||
|
||||
Configure the multi-agent research workflow:
|
||||
|
||||
```yaml
|
||||
research:
|
||||
model:
|
||||
provider: "" # Empty to use qa settings
|
||||
name: "" # Empty to use qa model
|
||||
enable_thinking: false
|
||||
temperature: 0.3 # Default: 0.3
|
||||
max_iterations: 3
|
||||
max_concurrency: 1
|
||||
```
|
||||
|
||||
- **model**: LLM configuration. Leave provider/model empty to inherit from `qa` (see [Providers](providers.md#model-settings))
|
||||
- **max_iterations**: Maximum planning/search iterations (default: 3)
|
||||
- **max_concurrency**: Concurrent search operations (default: 1)
|
||||
|
||||
The research workflow uses an iterative feedback loop: the planner proposes one question at a time, sees the answer, then decides whether to continue or synthesize. This continues until the planner marks research as complete or `max_iterations` is reached.
|
||||
|
||||
## Analysis Configuration
|
||||
|
||||
Configure the analysis skill:
|
||||
|
|
@ -123,7 +123,7 @@ search:
|
|||
vector_refine_factor: 30 # Re-ranking factor for accuracy
|
||||
```
|
||||
|
||||
For search behavior settings (`limit`, `max_context_chars`), see [QA and Research](qa-research.md#search-settings).
|
||||
For search behavior settings (`limit`, `max_context_chars`), see [Search and Question Answering](qa.md#search-settings).
|
||||
|
||||
- **vector_index_metric**: Distance metric for vector similarity:
|
||||
- `cosine`: Cosine similarity (default, best for most embeddings)
|
||||
|
|
|
|||
|
|
@ -11,11 +11,10 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p
|
|||
- **Question answering** — RAG skill with citations (page numbers, section headings)
|
||||
- **Vision QA** — Vision-capable models receive figure bytes alongside chunk text via pydantic-ai `BinaryContent` when `qa.model.vision = true`
|
||||
- **Reranking** — MxBAI, Cohere, Zero Entropy, or vLLM
|
||||
- **Research workflow** — Multi-agent pydantic-graph: plan, search, evaluate, synthesize
|
||||
- **Analysis skill** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
|
||||
- **Conversational RAG** — Chat TUI and web application for multi-turn conversations with session memory
|
||||
- **Document structure** — Stores full [DoclingDocument](https://docling-project.github.io/docling/concepts/docling_document/), enabling structure-aware context expansion
|
||||
- **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, LM Studio, vLLM (multimodal). QA/Research: any model supported by Pydantic AI
|
||||
- **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, LM Studio, vLLM (multimodal). QA: any model supported by Pydantic AI
|
||||
- **Local-first** — Embedded LanceDB, no servers required. Also supports S3, GCS, Azure, and LanceDB Cloud
|
||||
- **CLI & Python API** — Full functionality from command line or code
|
||||
- **MCP server** — Expose as tools for AI assistants (Claude Desktop, etc.)
|
||||
|
|
@ -67,8 +66,8 @@ haiku-rag chat # Interactive conversation mode
|
|||
- [CLI](cli.md) - Command line interface usage
|
||||
- [Python](python.md) - Python API reference
|
||||
- [Custom Pipelines](custom-pipelines.md) - Build custom processing workflows
|
||||
- [Agents](agents/index.md) - QA, chat, and research agents
|
||||
- [Analysis Agent](agents/analysis.md) - Complex analytical tasks via code execution
|
||||
- [Skills](skills/index.md) - The RAG and analysis skills the client wraps
|
||||
- [Analysis](agents/analysis.md) - Complex analytical tasks via code execution
|
||||
- [Applications](apps.md) - Chat TUI, web app, and inspector
|
||||
- [Server](server.md) - File monitoring and server mode
|
||||
- [MCP](mcp.md) - Model Context Protocol integration
|
||||
|
|
|
|||
|
|
@ -52,10 +52,6 @@ The MCP server exposes `haiku.rag` as MCP tools for compatible MCP clients like
|
|||
- `cite` (optional): Include source citations (default: false)
|
||||
- `deep` (optional): Use multi-agent deep QA for complex questions (default: false)
|
||||
|
||||
- **`research_question`** - Run multi-agent research on complex topics
|
||||
- `question` (required): The research question
|
||||
- Returns a structured research report with findings, conclusions, and sources
|
||||
|
||||
- **`analyze`** - Answer complex analytical questions via code execution
|
||||
- `question` (required): The question to answer
|
||||
- `filter` (optional): SQL WHERE clause to restrict document access
|
||||
|
|
|
|||
|
|
@ -434,7 +434,7 @@ answer, citations = await client.ask(
|
|||
|
||||
The QA provider and model are configured in `haiku.rag.yaml` or can be passed directly to the client (see [Configuration](configuration/index.md)).
|
||||
|
||||
See also: [Agents](agents/index.md) for details on question answering and the multi‑agent research workflow.
|
||||
See also: [Skills](skills/index.md) for details on the skills the client wraps.
|
||||
|
||||
## Analysis
|
||||
|
||||
|
|
|
|||
|
|
@ -20,11 +20,11 @@ Larger embedding models produce better representations at the cost of slower ind
|
|||
|
||||
### Reranking
|
||||
|
||||
When configured, a cross-encoder reranker re-scores 10x the requested candidates and returns the top results. This adds latency but improves precision — on the Wix benchmark, adding `mxbai-rerank-base-v2` raised MAP from 0.34 to 0.39 on HTML content. See [Search Settings](configuration/qa-research.md#search-settings) for how reranking integrates with search.
|
||||
When configured, a cross-encoder reranker re-scores 10x the requested candidates and returns the top results. This adds latency but improves precision — on the Wix benchmark, adding `mxbai-rerank-base-v2` raised MAP from 0.34 to 0.39 on HTML content. See [Search Settings](configuration/qa.md#search-settings) for how reranking integrates with search.
|
||||
|
||||
### Search Settings
|
||||
|
||||
`limit` controls how many results reach the LLM. More candidates improve recall but increase token usage. See [Search Settings](configuration/qa-research.md#search-settings).
|
||||
`limit` controls how many results reach the LLM. More candidates improve recall but increase token usage. See [Search Settings](configuration/qa.md#search-settings).
|
||||
|
||||
Context expansion is automatic and section-aware — search results are expanded to include surrounding content from the same document section. For structured documents, expansion stays within section boundaries and filters noise (footnotes, page headers). For unstructured documents, expansion grows outward until the character budget is filled. `max_context_chars` caps expansion to prevent context bloat.
|
||||
|
||||
|
|
@ -32,7 +32,7 @@ Context expansion is automatic and section-aware — search results are expanded
|
|||
|
||||
Model and temperature selection affect answer quality directly — see [Providers](configuration/providers.md#model-settings) for options.
|
||||
|
||||
`domain_preamble` prepends domain context to all agent prompts — including the main agent, skill subagents, and internal agents (QA, research). Use it to describe what the knowledge base contains and clarify domain-specific terminology. For full prompt replacement, set `prompts.qa` directly. See [Prompt Customization](configuration/prompts.md).
|
||||
`domain_preamble` prepends domain context to the rag and rag-analysis skill instructions. Use it to describe what the knowledge base contains and clarify domain-specific terminology. See [Prompt Customization](configuration/prompts.md).
|
||||
|
||||
## What Requires a Rebuild
|
||||
|
||||
|
|
|
|||
|
|
@ -206,6 +206,6 @@ The following people are presenting talks at PyCon Finland 2025:
|
|||
- **[Chat](apps.md#chat-tui)** - Interactive conversations with `haiku-rag chat`
|
||||
- **[CLI Reference](cli.md)** - All available commands and options
|
||||
- **[Python API](python.md)** - Use haiku.rag in your Python applications
|
||||
- **[Agents](agents/index.md)** - Deep QA and multi-agent research workflows
|
||||
- **[Skills](skills/index.md)** - The RAG and analysis skills the client wraps
|
||||
- **[Configuration](configuration/index.md)** - Complete YAML configuration reference
|
||||
- **[Server Mode](server.md)** - File monitoring and MCP server
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ from typing import Protocol, cast
|
|||
|
||||
from pydantic_ai.models import Model
|
||||
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
from haiku.skills import run_skill
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import pytest
|
|||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
from evaluations.skill_runner import SkillRunResult, run_skill_question
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.embeddings import EmbedderWrapper
|
||||
|
|
|
|||
|
|
@ -1,7 +0,0 @@
|
|||
from haiku.rag.agents.research.dependencies import ResearchContext, ResearchDependencies
|
||||
from haiku.rag.agents.research.models import (
|
||||
Citation,
|
||||
IterativePlanResult,
|
||||
ResearchReport,
|
||||
SearchAnswer,
|
||||
)
|
||||
|
|
@ -1,34 +0,0 @@
|
|||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.store.models import SearchResult
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.agents.research.models import SearchAnswer
|
||||
|
||||
|
||||
class ResearchContext(BaseModel):
|
||||
"""Context shared across research agents."""
|
||||
|
||||
original_question: str = Field(description="The original research question")
|
||||
qa_responses: list[Any] = Field(
|
||||
default_factory=list, description="Structured QA pairs used during research"
|
||||
)
|
||||
|
||||
def add_qa_response(self, qa: "SearchAnswer") -> None:
|
||||
"""Add a structured QA response."""
|
||||
self.qa_responses.append(qa)
|
||||
|
||||
|
||||
class ResearchDependencies(BaseModel):
|
||||
"""Dependencies for research agents with multi-agent context."""
|
||||
|
||||
model_config = {"arbitrary_types_allowed": True}
|
||||
|
||||
client: HaikuRAG = Field(description="RAG client for document operations")
|
||||
context: ResearchContext = Field(description="Shared research context")
|
||||
search_results: list[SearchResult] = Field(
|
||||
default_factory=list, description="Search results for citation resolution"
|
||||
)
|
||||
|
|
@ -1,280 +0,0 @@
|
|||
import asyncio
|
||||
|
||||
from pydantic_ai import Agent, RunContext, format_as_xml
|
||||
from pydantic_graph.beta import Graph, GraphBuilder, StepContext
|
||||
|
||||
from haiku.rag.agents.research.dependencies import ResearchContext, ResearchDependencies
|
||||
from haiku.rag.agents.research.models import (
|
||||
IterativePlanResult,
|
||||
RawSearchAnswer,
|
||||
ResearchReport,
|
||||
SearchAnswer,
|
||||
)
|
||||
from haiku.rag.agents.research.prompts import (
|
||||
ITERATIVE_PLAN_PROMPT,
|
||||
ITERATIVE_PLAN_PROMPT_WITH_CONTEXT,
|
||||
SEARCH_PROMPT,
|
||||
SYNTHESIS_PROMPT,
|
||||
)
|
||||
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.utils import build_prompt, get_model
|
||||
|
||||
|
||||
def format_context_for_prompt(context: ResearchContext) -> str:
|
||||
"""Format the research context as XML for prompts."""
|
||||
context_data: dict[str, object] = {}
|
||||
|
||||
context_data["question"] = context.original_question
|
||||
|
||||
if context.qa_responses:
|
||||
context_data["prior_answers"] = [
|
||||
{
|
||||
"question": qa.query,
|
||||
"answer": qa.answer,
|
||||
"confidence": qa.confidence,
|
||||
"source": qa.primary_source,
|
||||
}
|
||||
for qa in context.qa_responses
|
||||
]
|
||||
|
||||
return format_as_xml(context_data, root_tag="context")
|
||||
|
||||
|
||||
async def _iterative_plan_logic(
|
||||
state: ResearchState,
|
||||
deps: ResearchDeps,
|
||||
config: AppConfig,
|
||||
) -> IterativePlanResult:
|
||||
"""Evaluate context and decide next question or mark complete."""
|
||||
has_prior_answers = bool(state.context.qa_responses)
|
||||
|
||||
# If max iterations reached, skip LLM and mark complete
|
||||
if state.iterations >= state.max_iterations:
|
||||
return IterativePlanResult(
|
||||
is_complete=True,
|
||||
next_question=None,
|
||||
reasoning=f"Max iterations ({state.max_iterations}) reached.",
|
||||
)
|
||||
|
||||
model_config = config.research.model
|
||||
|
||||
if has_prior_answers:
|
||||
effective_prompt = build_prompt(ITERATIVE_PLAN_PROMPT_WITH_CONTEXT, config)
|
||||
else:
|
||||
effective_prompt = build_prompt(ITERATIVE_PLAN_PROMPT, config)
|
||||
|
||||
model = get_model(model_config, config)
|
||||
plan_agent: Agent[ResearchDependencies, IterativePlanResult] = Agent( # type: ignore[assignment] # ty: ignore[invalid-assignment]
|
||||
model=model,
|
||||
output_type=IterativePlanResult,
|
||||
instructions=effective_prompt,
|
||||
tool_retries=3,
|
||||
output_retries=3,
|
||||
deps_type=ResearchDependencies,
|
||||
)
|
||||
|
||||
# Build prompt based on current state
|
||||
if has_prior_answers:
|
||||
context_xml = format_context_for_prompt(state.context)
|
||||
prompt = (
|
||||
f"Review the gathered evidence and decide whether to continue or synthesize.\n\n"
|
||||
f"{context_xml}"
|
||||
)
|
||||
else:
|
||||
context_xml = format_context_for_prompt(state.context)
|
||||
prompt = f"Plan the research investigation.\n\n{context_xml}"
|
||||
|
||||
agent_deps = ResearchDependencies(client=deps.client, context=state.context)
|
||||
result = await plan_agent.run(prompt, deps=agent_deps)
|
||||
|
||||
# Enforce: if no prior answers, must have a next_question to investigate
|
||||
if not has_prior_answers:
|
||||
if result.output.is_complete or not result.output.next_question:
|
||||
return IterativePlanResult(
|
||||
is_complete=False,
|
||||
next_question=result.output.next_question
|
||||
or state.context.original_question,
|
||||
reasoning=result.output.reasoning,
|
||||
)
|
||||
|
||||
return result.output
|
||||
|
||||
|
||||
async def _search_one_step_logic(
|
||||
state: ResearchState,
|
||||
deps: ResearchDeps,
|
||||
config: AppConfig,
|
||||
search_prompt: str,
|
||||
sub_q: str,
|
||||
) -> SearchAnswer:
|
||||
"""Answer a single question using the knowledge base."""
|
||||
model_config = config.research.model
|
||||
|
||||
if deps.semaphore is None:
|
||||
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
|
||||
|
||||
async with deps.semaphore:
|
||||
model = get_model(model_config, config)
|
||||
agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent( # type: ignore[assignment] # ty: ignore[invalid-assignment]
|
||||
model=model,
|
||||
output_type=RawSearchAnswer,
|
||||
instructions=search_prompt,
|
||||
tool_retries=3,
|
||||
output_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
|
||||
total = len(results)
|
||||
parts = [
|
||||
r.format_for_agent(rank=i + 1, total=total)
|
||||
for i, r in enumerate(results)
|
||||
]
|
||||
if not parts:
|
||||
return f"No relevant information found for: {query}"
|
||||
return "\n\n".join(parts)
|
||||
|
||||
agent_deps = ResearchDependencies(client=deps.client, context=state.context)
|
||||
|
||||
result = await agent.run(sub_q, deps=agent_deps)
|
||||
raw_answer = result.output
|
||||
|
||||
# Increment iterations after each search completes
|
||||
state.iterations += 1
|
||||
|
||||
if raw_answer:
|
||||
answer = SearchAnswer.from_raw(raw_answer, agent_deps.search_results)
|
||||
state.context.add_qa_response(answer)
|
||||
return answer
|
||||
return SearchAnswer(query=sub_q, answer="", confidence=0.0)
|
||||
|
||||
|
||||
def build_research_graph(
|
||||
config: AppConfig = Config,
|
||||
) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]:
|
||||
"""Build the iterative research graph.
|
||||
|
||||
Args:
|
||||
config: AppConfig object (uses config.research for provider, model, and graph parameters)
|
||||
|
||||
Returns:
|
||||
Configured research graph with iterative planning
|
||||
"""
|
||||
model_config = config.research.model
|
||||
|
||||
search_prompt = build_prompt(SEARCH_PROMPT, config)
|
||||
synthesis_prompt = build_prompt(
|
||||
config.prompts.synthesis or SYNTHESIS_PROMPT, config
|
||||
)
|
||||
|
||||
g = GraphBuilder(
|
||||
state_type=ResearchState,
|
||||
deps_type=ResearchDeps,
|
||||
output_type=ResearchReport,
|
||||
)
|
||||
|
||||
@g.step
|
||||
async def plan_next(
|
||||
ctx: StepContext[ResearchState, ResearchDeps, None | SearchAnswer],
|
||||
) -> IterativePlanResult:
|
||||
"""Evaluate context and decide next question or complete."""
|
||||
return await _iterative_plan_logic(ctx.state, ctx.deps, config)
|
||||
|
||||
@g.step
|
||||
async def search_one(
|
||||
ctx: StepContext[ResearchState, ResearchDeps, str],
|
||||
) -> SearchAnswer:
|
||||
"""Answer a single question using the knowledge base."""
|
||||
try:
|
||||
return await _search_one_step_logic(
|
||||
ctx.state, ctx.deps, config, search_prompt, ctx.inputs
|
||||
)
|
||||
except Exception as e:
|
||||
return SearchAnswer(
|
||||
query=ctx.inputs,
|
||||
answer=f"Search failed: {str(e)}",
|
||||
confidence=0.0,
|
||||
)
|
||||
|
||||
@g.step
|
||||
async def synthesize(
|
||||
ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult],
|
||||
) -> ResearchReport:
|
||||
"""Generate final research report."""
|
||||
state = ctx.state
|
||||
deps = ctx.deps
|
||||
|
||||
model = get_model(model_config, config)
|
||||
agent: Agent[ResearchDependencies, ResearchReport] = Agent( # type: ignore[assignment] # ty: ignore[invalid-assignment]
|
||||
model=model,
|
||||
output_type=ResearchReport,
|
||||
instructions=synthesis_prompt,
|
||||
tool_retries=3,
|
||||
output_retries=3,
|
||||
deps_type=ResearchDependencies,
|
||||
)
|
||||
|
||||
context_xml = format_context_for_prompt(state.context)
|
||||
prompt = (
|
||||
"Generate a comprehensive research report based on all gathered information.\n\n"
|
||||
f"{context_xml}\n\n"
|
||||
"Create a detailed report that synthesizes all findings into a coherent response."
|
||||
)
|
||||
agent_deps = ResearchDependencies(
|
||||
client=deps.client,
|
||||
context=state.context,
|
||||
)
|
||||
result = await agent.run(prompt, deps=agent_deps)
|
||||
return result.output
|
||||
|
||||
# Build graph edges: iterative loop
|
||||
#
|
||||
# START -> plan_next -> [decision]
|
||||
# |
|
||||
# [is_complete or max_iterations] -> synthesize -> END
|
||||
# |
|
||||
# [has next_question] -> search_one -> plan_next (loop)
|
||||
|
||||
def extract_question(
|
||||
ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult],
|
||||
) -> str:
|
||||
"""Extract next_question from IterativePlanResult."""
|
||||
return ctx.inputs.next_question or ""
|
||||
|
||||
g.add(
|
||||
g.edge_from(g.start_node).to(plan_next),
|
||||
g.edge_from(plan_next).to(
|
||||
g.decision()
|
||||
.branch(
|
||||
g.match(
|
||||
IterativePlanResult,
|
||||
matches=lambda r: not r.is_complete and r.next_question is not None,
|
||||
)
|
||||
.label("Continue research")
|
||||
.transform(extract_question)
|
||||
.to(search_one)
|
||||
)
|
||||
.branch(
|
||||
g.match(IterativePlanResult).label("Done researching").to(synthesize)
|
||||
)
|
||||
),
|
||||
g.edge_from(search_one).to(plan_next),
|
||||
g.edge_from(synthesize).to(g.end_node),
|
||||
)
|
||||
|
||||
return g.build()
|
||||
|
|
@ -1,144 +0,0 @@
|
|||
from typing import TYPE_CHECKING
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.store.models import SearchResult
|
||||
|
||||
|
||||
class IterativePlanResult(BaseModel):
|
||||
"""Output from iterative planning step."""
|
||||
|
||||
is_complete: bool = Field(
|
||||
description="Whether research is complete and can be synthesized"
|
||||
)
|
||||
next_question: str | None = Field(
|
||||
default=None, description="Next question to investigate, if not complete"
|
||||
)
|
||||
reasoning: str = Field(description="Brief explanation of the decision")
|
||||
|
||||
|
||||
class Citation(BaseModel):
|
||||
"""Resolved citation with full metadata for display/visual grounding.
|
||||
|
||||
Used by research graph and chat applications. The optional index field
|
||||
supports UI display ordering in chat contexts.
|
||||
|
||||
``picture_refs`` lists the ``self_ref`` values of picture items in the
|
||||
cited chunk. Empty for text-only citations. UIs can fetch the picture
|
||||
bytes via ``DocumentItemRepository.get_picture_bytes(document_id, ref)``
|
||||
and render them alongside the text content.
|
||||
"""
|
||||
|
||||
index: int | None = None
|
||||
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
|
||||
picture_refs: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
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="Complete chunk IDs from search results (e.g. '5ae52166-5329-42e9-b6a5-756fc0cb7200'). Copy the full UUID without brackets. Must not be empty when providing an 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",
|
||||
)
|
||||
|
||||
@property
|
||||
def primary_source(self) -> str | None:
|
||||
"""Get primary source title from citations."""
|
||||
if not self.citations:
|
||||
return None
|
||||
first = self.citations[0]
|
||||
return first.document_title or first.document_uri
|
||||
|
||||
@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 raw_id in cited_chunk_ids:
|
||||
chunk_id = raw_id.strip("[]")
|
||||
r = by_id.get(chunk_id)
|
||||
if not r:
|
||||
continue
|
||||
picture_refs = [
|
||||
ref for ref in r.doc_item_refs if ref.startswith(PICTURE_REF_PREFIX)
|
||||
]
|
||||
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,
|
||||
picture_refs=picture_refs,
|
||||
)
|
||||
)
|
||||
return citations
|
||||
|
||||
|
||||
class ResearchReport(BaseModel):
|
||||
"""Final research report structure."""
|
||||
|
||||
title: str = Field(description="Concise title for the research")
|
||||
executive_summary: str = Field(description="Brief overview of key findings")
|
||||
main_findings: list[str] = Field(
|
||||
description="Primary research findings with supporting evidence"
|
||||
)
|
||||
conclusions: list[str] = Field(description="Evidence-based conclusions")
|
||||
limitations: list[str] = Field(
|
||||
description="Limitations of the current research", default=[]
|
||||
)
|
||||
recommendations: list[str] = Field(
|
||||
description="Actionable recommendations based on findings", default=[]
|
||||
)
|
||||
sources_summary: str = Field(
|
||||
description="Summary of sources used and their reliability"
|
||||
)
|
||||
|
|
@ -1,114 +0,0 @@
|
|||
ITERATIVE_PLAN_PROMPT = """You are the research orchestrator planning the investigation.
|
||||
|
||||
Your task:
|
||||
1. Analyze the original question
|
||||
2. Propose the first question to investigate
|
||||
|
||||
For simple questions, investigate them directly. For composite or complex questions,
|
||||
you may decompose into a focused sub-question. For example:
|
||||
- "What are the benefits and drawbacks of X?" → Start with "What are the benefits of X?"
|
||||
- Ambiguous references should be resolved using background context if available
|
||||
|
||||
Output requirements:
|
||||
- Set is_complete=False (you are just starting the investigation)
|
||||
- Set next_question to the question to investigate
|
||||
- Provide brief reasoning explaining your choice
|
||||
|
||||
The question must be standalone and self-contained:
|
||||
- Include concrete entities, scope, and any qualifiers
|
||||
- Avoid ambiguous pronouns (it/they/this/that)"""
|
||||
|
||||
ITERATIVE_PLAN_PROMPT_WITH_CONTEXT = """You are the research orchestrator evaluating gathered evidence.
|
||||
|
||||
You have access to context that may include:
|
||||
- <prior_answers>: Previous Q&A pairs with confidence scores
|
||||
|
||||
Your task:
|
||||
1. Review the provided evidence carefully
|
||||
2. Assess whether it sufficiently answers the original question
|
||||
3. Decide whether to continue research or synthesize
|
||||
|
||||
Decision criteria:
|
||||
- Set is_complete=True if the evidence adequately answers the question
|
||||
- Set is_complete=False with a next_question if important gaps remain
|
||||
|
||||
If not complete, propose exactly ONE high-value follow-up question in next_question:
|
||||
- Focus on the most critical gap not covered by prior_answers
|
||||
- The question must be standalone and self-contained
|
||||
- Avoid repeating questions that have already been answered
|
||||
- Include concrete entities, scope, and any qualifiers
|
||||
|
||||
Provide brief reasoning explaining your decision."""
|
||||
|
||||
SEARCH_PROMPT = """You are a search and question-answering specialist.
|
||||
|
||||
Process:
|
||||
1. Call search_and_answer with relevant keywords from the question.
|
||||
2. Review the results ordered by relevance.
|
||||
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] [rank 1 of 5]
|
||||
Source: "Document Title" > Section > Subsection
|
||||
Type: paragraph
|
||||
Content:
|
||||
The actual text content here...
|
||||
|
||||
[d5a63c82-cb40-439f-9b2e-de7d177829b7] [rank 2 of 5]
|
||||
Source: "Another Document"
|
||||
Type: table
|
||||
Content:
|
||||
| Column 1 | Column 2 |
|
||||
...
|
||||
|
||||
Each result includes:
|
||||
- chunk_id in brackets and rank position (rank 1 = most relevant)
|
||||
- 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.
|
||||
- Results are ordered by relevance, with rank 1 being most relevant."""
|
||||
|
||||
SYNTHESIS_PROMPT = """You are a synthesis specialist producing the final
|
||||
research report that directly answers the original question.
|
||||
|
||||
Goals:
|
||||
1. Directly answer the research question using gathered evidence.
|
||||
2. Present findings clearly and concisely.
|
||||
3. Draw evidence-based conclusions and recommendations.
|
||||
4. State limitations and uncertainties transparently.
|
||||
|
||||
Report guidelines (map to output fields):
|
||||
- title: concise (5-12 words), informative.
|
||||
- executive_summary: 3-5 sentences that DIRECTLY ANSWER the original question.
|
||||
Write the actual answer, not a description of what the report contains.
|
||||
BAD: "This report examines the topic and presents findings..."
|
||||
GOOD: "The system requires configuration X and supports features Y and Z..."
|
||||
- main_findings: list of plain strings, 4-8 one-sentence bullets reflecting evidence.
|
||||
- conclusions: list of plain strings, 2-4 bullets following logically from findings.
|
||||
- recommendations: list of plain strings, 2-5 actionable bullets tied to findings.
|
||||
- limitations: list of plain strings, 1-3 bullets describing constraints or uncertainties.
|
||||
- sources_summary: single string listing sources with document paths and page numbers.
|
||||
|
||||
All list fields must contain plain strings only, not objects.
|
||||
|
||||
Style:
|
||||
- Base all content solely on the collected evidence.
|
||||
- Be professional, objective, and specific.
|
||||
- NEVER use meta-commentary like "This report covers..." or "The findings show...".
|
||||
Instead, state the actual information directly."""
|
||||
|
|
@ -1,59 +0,0 @@
|
|||
import asyncio
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from haiku.rag.agents.research.dependencies import ResearchContext
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.config.models import AppConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResearchDeps:
|
||||
"""Dependencies for research graph execution."""
|
||||
|
||||
client: HaikuRAG
|
||||
semaphore: asyncio.Semaphore | None = None
|
||||
|
||||
|
||||
class ResearchState(BaseModel):
|
||||
"""Research graph state model."""
|
||||
|
||||
model_config = {"arbitrary_types_allowed": True}
|
||||
|
||||
context: ResearchContext = Field(
|
||||
description="Shared research context with questions and QA responses"
|
||||
)
|
||||
iterations: int = Field(default=0, description="Current iteration number")
|
||||
max_iterations: int = Field(default=3, description="Maximum allowed iterations")
|
||||
max_concurrency: int = Field(
|
||||
default=1, description="Maximum concurrent search operations", ge=1
|
||||
)
|
||||
search_filter: str | None = Field(
|
||||
default=None, description="SQL WHERE clause to filter search results"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_config(
|
||||
cls,
|
||||
context: ResearchContext,
|
||||
config: "AppConfig",
|
||||
max_iterations: int | None = None,
|
||||
) -> "ResearchState":
|
||||
"""Create a ResearchState from an AppConfig.
|
||||
|
||||
Args:
|
||||
context: The ResearchContext containing the question
|
||||
config: The AppConfig object
|
||||
max_iterations: Override max iterations (None uses config default)
|
||||
"""
|
||||
return cls(
|
||||
context=context,
|
||||
max_iterations=max_iterations
|
||||
if max_iterations is not None
|
||||
else config.research.max_iterations,
|
||||
max_concurrency=config.research.max_concurrency,
|
||||
)
|
||||
|
|
@ -500,77 +500,6 @@ class HaikuRAGApp: # pragma: no cover
|
|||
):
|
||||
self.console.print(renderable)
|
||||
|
||||
async def research(
|
||||
self,
|
||||
question: str,
|
||||
filter: str | None = None,
|
||||
):
|
||||
"""Run research via the pydantic-graph pipeline.
|
||||
|
||||
Args:
|
||||
question: The research question
|
||||
filter: SQL WHERE clause to filter documents
|
||||
"""
|
||||
async with HaikuRAG(
|
||||
db_path=self.db_path,
|
||||
config=self.config,
|
||||
read_only=self.read_only,
|
||||
before=self.before,
|
||||
) as client:
|
||||
self.console.print("[bold cyan]Starting research[/bold cyan]")
|
||||
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
|
||||
self.console.print()
|
||||
|
||||
report = await client.research(question=question, filter=filter)
|
||||
|
||||
if report is None:
|
||||
self.console.print("[red]Research did not produce a report.[/red]")
|
||||
return
|
||||
|
||||
# Display the report
|
||||
self.console.print("[bold green]Research Report[/bold green]")
|
||||
self.console.rule()
|
||||
|
||||
# Title and Executive Summary
|
||||
self.console.print(f"[bold]{report.title}[/bold]")
|
||||
self.console.print()
|
||||
self.console.print("[bold cyan]Executive Summary:[/bold cyan]")
|
||||
self.console.print(report.executive_summary)
|
||||
self.console.print()
|
||||
|
||||
# Main Findings
|
||||
if report.main_findings:
|
||||
self.console.print("[bold cyan]Main Findings:[/bold cyan]")
|
||||
for finding in report.main_findings:
|
||||
self.console.print(f"• {finding}")
|
||||
self.console.print()
|
||||
|
||||
# Conclusions
|
||||
if report.conclusions:
|
||||
self.console.print("[bold cyan]Conclusions:[/bold cyan]")
|
||||
for conclusion in report.conclusions:
|
||||
self.console.print(f"• {conclusion}")
|
||||
self.console.print()
|
||||
|
||||
# Recommendations
|
||||
if report.recommendations:
|
||||
self.console.print("[bold cyan]Recommendations:[/bold cyan]")
|
||||
for rec in report.recommendations:
|
||||
self.console.print(f"• {rec}")
|
||||
self.console.print()
|
||||
|
||||
# Limitations
|
||||
if report.limitations:
|
||||
self.console.print("[bold yellow]Limitations:[/bold yellow]")
|
||||
for limitation in report.limitations:
|
||||
self.console.print(f"• {limitation}")
|
||||
self.console.print()
|
||||
|
||||
# Sources Summary
|
||||
if report.sources_summary:
|
||||
self.console.print("[bold cyan]Sources:[/bold cyan]")
|
||||
self.console.print(report.sources_summary)
|
||||
|
||||
async def rebuild(self, mode: RebuildMode = RebuildMode.FULL):
|
||||
async with HaikuRAG(
|
||||
db_path=self.db_path,
|
||||
|
|
|
|||
|
|
@ -36,7 +36,6 @@ def run_chat(
|
|||
if model:
|
||||
model_config = parse_model_option(model)
|
||||
config.qa.model = model_config
|
||||
config.research.model = model_config
|
||||
config.analysis.model = model_config
|
||||
|
||||
enabled = skills or ["rag"]
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from textual.widgets import Collapsible, LoadingIndicator, Markdown, Static
|
|||
from textual.widgets.markdown import MarkdownStream
|
||||
from textual_image.widget import Image as TextualImage
|
||||
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from textual.app import ComposeResult
|
||||
|
|
|
|||
|
|
@ -411,25 +411,6 @@ def analyze( # pragma: no cover
|
|||
)
|
||||
|
||||
|
||||
@_cli.command("research", help="Run multi-agent research and output a concise report")
|
||||
def research( # pragma: no cover
|
||||
question: str = typer.Argument(..., help="The research question to investigate"),
|
||||
db: Path | None = typer.Option(
|
||||
None,
|
||||
"--db",
|
||||
help="Path to the LanceDB database file",
|
||||
),
|
||||
filter: str | None = typer.Option(
|
||||
None,
|
||||
"--filter",
|
||||
"-f",
|
||||
help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
|
||||
),
|
||||
):
|
||||
app = create_app(db)
|
||||
asyncio.run(app.research(question=question, filter=filter))
|
||||
|
||||
|
||||
@_cli.command("settings", help="Display current configuration settings")
|
||||
def settings(): # pragma: no cover
|
||||
config = get_config()
|
||||
|
|
|
|||
|
|
@ -30,11 +30,8 @@ if TYPE_CHECKING:
|
|||
from docling_core.types.doc.document import DoclingDocument
|
||||
from PIL import Image as PILImage
|
||||
|
||||
from haiku.rag.agents.research.models import (
|
||||
Citation,
|
||||
ResearchReport,
|
||||
)
|
||||
from haiku.rag.sandbox import AnalysisResult
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -374,19 +371,6 @@ class HaikuRAG:
|
|||
|
||||
return await ask(self, question, filter)
|
||||
|
||||
async def research(
|
||||
self,
|
||||
question: str,
|
||||
*,
|
||||
filter: str | None = None,
|
||||
max_iterations: int | None = None,
|
||||
) -> "ResearchReport":
|
||||
from haiku.rag.client.agents import research
|
||||
|
||||
return await research(
|
||||
self, question, filter=filter, max_iterations=max_iterations
|
||||
)
|
||||
|
||||
async def analyze(
|
||||
self,
|
||||
question: str,
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.agents.research.models import Citation, ResearchReport
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.sandbox import AnalysisResult
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
|
||||
|
||||
async def ask(
|
||||
|
|
@ -37,39 +37,6 @@ async def ask(
|
|||
return answer, citations
|
||||
|
||||
|
||||
async def research(
|
||||
client: "HaikuRAG",
|
||||
question: str,
|
||||
*,
|
||||
filter: str | None = None,
|
||||
max_iterations: int | None = None,
|
||||
) -> "ResearchReport":
|
||||
"""Run multi-agent research to investigate a question.
|
||||
|
||||
Args:
|
||||
client: The HaikuRAG client.
|
||||
question: The research question to investigate.
|
||||
filter: SQL WHERE clause to filter documents.
|
||||
max_iterations: Override max iterations (None uses config default).
|
||||
|
||||
Returns:
|
||||
ResearchReport with structured findings.
|
||||
"""
|
||||
from haiku.rag.agents.research.dependencies import ResearchContext
|
||||
from haiku.rag.agents.research.graph import build_research_graph
|
||||
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
|
||||
|
||||
graph = build_research_graph(config=client._config)
|
||||
context = ResearchContext(original_question=question)
|
||||
state = ResearchState.from_config(
|
||||
context=context, config=client._config, max_iterations=max_iterations
|
||||
)
|
||||
state.search_filter = filter
|
||||
deps = ResearchDeps(client=client)
|
||||
|
||||
return await graph.run(state=state, deps=deps)
|
||||
|
||||
|
||||
async def analyze(
|
||||
client: "HaikuRAG",
|
||||
question: str,
|
||||
|
|
|
|||
|
|
@ -101,8 +101,6 @@ async def download_models(
|
|||
required_models.add(config.embeddings.model.name)
|
||||
if config.qa.model.provider == "ollama":
|
||||
required_models.add(config.qa.model.name)
|
||||
if config.research.model.provider == "ollama":
|
||||
required_models.add(config.research.model.name)
|
||||
if config.reranking.model and config.reranking.model.provider == "ollama":
|
||||
required_models.add(config.reranking.model.name)
|
||||
pic_desc = config.processing.conversion_options.picture_description
|
||||
|
|
|
|||
|
|
@ -17,7 +17,6 @@ from haiku.rag.config.models import (
|
|||
ProvidersConfig,
|
||||
QAConfig,
|
||||
RerankingConfig,
|
||||
ResearchConfig,
|
||||
S3MonitorEntry,
|
||||
StorageConfig,
|
||||
)
|
||||
|
|
@ -37,7 +36,6 @@ __all__ = [
|
|||
"ProvidersConfig",
|
||||
"QAConfig",
|
||||
"RerankingConfig",
|
||||
"ResearchConfig",
|
||||
"S3MonitorEntry",
|
||||
"StorageConfig",
|
||||
"find_config_file",
|
||||
|
|
|
|||
|
|
@ -102,19 +102,6 @@ class QAConfig(BaseModel):
|
|||
max_searches: int = 5
|
||||
|
||||
|
||||
class ResearchConfig(BaseModel):
|
||||
model: ModelConfig = Field(
|
||||
default_factory=lambda: ModelConfig(
|
||||
provider="ollama",
|
||||
name="gpt-oss",
|
||||
enable_thinking=False,
|
||||
temperature=0.3,
|
||||
)
|
||||
)
|
||||
max_iterations: int = 3
|
||||
max_concurrency: int = 1
|
||||
|
||||
|
||||
class AnalysisConfig(BaseModel):
|
||||
"""Driving model + sandbox limits for the analysis skill.
|
||||
|
||||
|
|
@ -241,7 +228,6 @@ class ProvidersConfig(BaseModel):
|
|||
|
||||
class PromptsConfig(BaseModel):
|
||||
domain_preamble: str = ""
|
||||
synthesis: str | None = None
|
||||
picture_description: str = (
|
||||
"Describe this image for a blind user. "
|
||||
"State the image type (screenshot, chart, photo, etc.), "
|
||||
|
|
@ -258,7 +244,6 @@ class AppConfig(BaseModel):
|
|||
embeddings: EmbeddingsConfig = Field(default_factory=EmbeddingsConfig)
|
||||
reranking: RerankingConfig = Field(default_factory=RerankingConfig)
|
||||
qa: QAConfig = Field(default_factory=QAConfig)
|
||||
research: ResearchConfig = Field(default_factory=ResearchConfig)
|
||||
analysis: AnalysisConfig = Field(default_factory=AnalysisConfig)
|
||||
processing: ProcessingConfig = Field(default_factory=ProcessingConfig)
|
||||
search: SearchConfig = Field(default_factory=SearchConfig)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ from typing import Any
|
|||
|
||||
from fastmcp import FastMCP
|
||||
|
||||
from haiku.rag.agents.research.models import ResearchReport
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import AppConfig, Config
|
||||
from haiku.rag.store.models import Document, SearchResult
|
||||
|
|
@ -202,27 +201,6 @@ def create_mcp_server(
|
|||
except Exception as e:
|
||||
return f"Error answering question: {e!s}"
|
||||
|
||||
@mcp.tool()
|
||||
async def research_question(
|
||||
question: str,
|
||||
) -> ResearchReport | None:
|
||||
"""Run multi-agent research to investigate a complex question.
|
||||
|
||||
The research process uses multiple agents to plan, search, evaluate, and synthesize
|
||||
information iteratively until confidence threshold is met or max iterations reached.
|
||||
|
||||
Args:
|
||||
question: The research question to investigate.
|
||||
|
||||
Returns:
|
||||
A research report with findings, or None if an error occurred.
|
||||
"""
|
||||
try:
|
||||
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
||||
return await rag.research(question=question)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@mcp.tool()
|
||||
async def analyze(
|
||||
question: str,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from pydantic import BaseModel, Field
|
||||
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
|
||||
|
||||
class AnalysisResult(BaseModel):
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from haiku.rag.config.models import AppConfig
|
|||
from haiku.skills.models import Skill
|
||||
from haiku.skills.parser import parse_skill_md
|
||||
{% if "cite" in tool_names %}
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
{% endif %}
|
||||
{% if "search" in tool_names %}
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
|
|
|
|||
|
|
@ -5,11 +5,11 @@ from pydantic import BaseModel
|
|||
from pydantic_ai import ModelRetry, RunContext
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config.models import AppConfig, ModelConfig
|
||||
from haiku.rag.skills._deps import AnalysisRunDeps, RAGRunDeps
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.tools.search import build_binary_parts_from_results
|
||||
|
||||
|
||||
|
|
@ -302,8 +302,8 @@ def create_skill_tools(
|
|||
Args:
|
||||
chunk_ids: List of chunk_id values from search results or VFS reads.
|
||||
"""
|
||||
from haiku.rag.agents.research.models import resolve_citations
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
from haiku.rag.store.models.citation import resolve_citations
|
||||
|
||||
state = _get_state(ctx, state_type)
|
||||
if not state:
|
||||
|
|
|
|||
|
|
@ -4,10 +4,10 @@ from pathlib import Path
|
|||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.skills._tools import CodeExecutionEntry
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.skills.models import Skill, SkillMetadata, SkillSource, StateMetadata
|
||||
from haiku.skills.parser import parse_skill_md
|
||||
|
||||
|
|
|
|||
|
|
@ -4,9 +4,9 @@ from pathlib import Path
|
|||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.skills.models import Skill, SkillMetadata, SkillSource, StateMetadata
|
||||
from haiku.skills.parser import parse_skill_md
|
||||
|
||||
|
|
|
|||
62
haiku_rag_slim/haiku/rag/store/models/citation.py
Normal file
62
haiku_rag_slim/haiku/rag/store/models/citation.py
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
from typing import TYPE_CHECKING
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.store.models import SearchResult
|
||||
|
||||
|
||||
class Citation(BaseModel):
|
||||
"""Resolved citation with full metadata for display/visual grounding.
|
||||
|
||||
Used by the rag and analysis skills and rendered by the CLI / chat
|
||||
application. The optional index field supports UI display ordering.
|
||||
|
||||
``picture_refs`` lists the ``self_ref`` values of picture items in the
|
||||
cited chunk. Empty for text-only citations. UIs can fetch the picture
|
||||
bytes via ``DocumentItemRepository.get_picture_bytes(document_id, ref)``
|
||||
and render them alongside the text content.
|
||||
"""
|
||||
|
||||
index: int | None = None
|
||||
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
|
||||
picture_refs: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
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 raw_id in cited_chunk_ids:
|
||||
chunk_id = raw_id.strip("[]")
|
||||
r = by_id.get(chunk_id)
|
||||
if not r:
|
||||
continue
|
||||
picture_refs = [
|
||||
ref for ref in r.doc_item_refs if ref.startswith(PICTURE_REF_PREFIX)
|
||||
]
|
||||
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,
|
||||
picture_refs=picture_refs,
|
||||
)
|
||||
)
|
||||
return citations
|
||||
|
|
@ -11,9 +11,9 @@ from packaging.version import Version, parse
|
|||
if TYPE_CHECKING:
|
||||
from rich.console import RenderableType
|
||||
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config.models import AppConfig, ModelConfig
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
|
||||
|
||||
def parse_model_option(value: str) -> "ModelConfig":
|
||||
|
|
|
|||
|
|
@ -62,7 +62,7 @@ nav:
|
|||
- Configuration:
|
||||
- configuration/index.md
|
||||
- Providers: configuration/providers.md
|
||||
- Search and Question Answering: configuration/qa-research.md
|
||||
- Search and Question Answering: configuration/qa.md
|
||||
- Document Processing: configuration/processing.md
|
||||
- Storage: configuration/storage.md
|
||||
- Prompts: configuration/prompts.md
|
||||
|
|
@ -70,9 +70,7 @@ nav:
|
|||
- Python: python.md
|
||||
- Custom Pipelines: custom-pipelines.md
|
||||
- Tuning: tuning.md
|
||||
- Agents:
|
||||
- agents/index.md
|
||||
- Analysis: agents/analysis.md
|
||||
- Analysis: agents/analysis.md
|
||||
- Skills:
|
||||
- skills/index.md
|
||||
- RAG: skills/rag.md
|
||||
|
|
|
|||
|
|
@ -1,159 +0,0 @@
|
|||
from haiku.rag.agents.research.models import Citation, SearchAnswer, resolve_citations
|
||||
from haiku.rag.store.models import SearchResult
|
||||
|
||||
|
||||
class TestCitation:
|
||||
"""Tests for unified Citation class."""
|
||||
|
||||
def test_citation_without_index(self):
|
||||
"""Test Citation can be created without index (research graph use case)."""
|
||||
citation = Citation(
|
||||
document_id="doc-1",
|
||||
chunk_id="chunk-1",
|
||||
document_uri="test.md",
|
||||
document_title="Test Document",
|
||||
page_numbers=[1, 2],
|
||||
headings=["Introduction"],
|
||||
content="Test content",
|
||||
)
|
||||
assert citation.document_id == "doc-1"
|
||||
assert citation.chunk_id == "chunk-1"
|
||||
assert citation.document_uri == "test.md"
|
||||
assert citation.document_title == "Test Document"
|
||||
assert citation.page_numbers == [1, 2]
|
||||
assert citation.headings == ["Introduction"]
|
||||
assert citation.content == "Test content"
|
||||
assert citation.index is None
|
||||
|
||||
def test_citation_with_index(self):
|
||||
"""Test Citation can be created with index (chat use case)."""
|
||||
citation = Citation(
|
||||
index=1,
|
||||
document_id="doc-1",
|
||||
chunk_id="chunk-1",
|
||||
document_uri="test.md",
|
||||
content="Test content",
|
||||
)
|
||||
assert citation.index == 1
|
||||
assert citation.document_id == "doc-1"
|
||||
|
||||
def test_citation_index_defaults_to_none(self):
|
||||
"""Test Citation index defaults to None."""
|
||||
citation = Citation(
|
||||
document_id="doc-1",
|
||||
chunk_id="chunk-1",
|
||||
document_uri="test.md",
|
||||
content="Test content",
|
||||
)
|
||||
assert citation.index is None
|
||||
|
||||
def test_citation_serialization_includes_index_when_set(self):
|
||||
"""Test Citation serialization includes index when set."""
|
||||
citation = Citation(
|
||||
index=2,
|
||||
document_id="doc-1",
|
||||
chunk_id="chunk-1",
|
||||
document_uri="test.md",
|
||||
content="Test content",
|
||||
)
|
||||
data = citation.model_dump()
|
||||
assert data["index"] == 2
|
||||
|
||||
def test_citation_deserialization_from_dict_with_index(self):
|
||||
"""Test Citation can be deserialized from dict with index (AG-UI state sync)."""
|
||||
data = {
|
||||
"index": 1,
|
||||
"document_id": "doc-1",
|
||||
"chunk_id": "chunk-1",
|
||||
"document_uri": "test.md",
|
||||
"document_title": "Test Doc",
|
||||
"page_numbers": [1, 2],
|
||||
"headings": ["Intro"],
|
||||
"content": "Test content",
|
||||
}
|
||||
citation = Citation.model_validate(data)
|
||||
assert citation.index == 1
|
||||
assert citation.document_id == "doc-1"
|
||||
|
||||
|
||||
class TestSearchAnswerPrimarySource:
|
||||
"""Tests for SearchAnswer.primary_source property."""
|
||||
|
||||
def test_primary_source_returns_title_when_available(self):
|
||||
"""Test primary_source returns first citation's title."""
|
||||
answer = SearchAnswer(
|
||||
query="test query",
|
||||
answer="test answer",
|
||||
citations=[
|
||||
Citation(
|
||||
document_id="doc-1",
|
||||
chunk_id="chunk-1",
|
||||
document_uri="test.md",
|
||||
document_title="Test Document",
|
||||
content="content",
|
||||
),
|
||||
],
|
||||
)
|
||||
assert answer.primary_source == "Test Document"
|
||||
|
||||
def test_primary_source_returns_uri_when_no_title(self):
|
||||
"""Test primary_source returns URI when title is None."""
|
||||
answer = SearchAnswer(
|
||||
query="test query",
|
||||
answer="test answer",
|
||||
citations=[
|
||||
Citation(
|
||||
document_id="doc-1",
|
||||
chunk_id="chunk-1",
|
||||
document_uri="test.md",
|
||||
document_title=None,
|
||||
content="content",
|
||||
),
|
||||
],
|
||||
)
|
||||
assert answer.primary_source == "test.md"
|
||||
|
||||
def test_primary_source_returns_none_when_no_citations(self):
|
||||
"""Test primary_source returns None when no citations."""
|
||||
answer = SearchAnswer(
|
||||
query="test query",
|
||||
answer="test answer",
|
||||
citations=[],
|
||||
)
|
||||
assert answer.primary_source is None
|
||||
|
||||
|
||||
class TestResolveCitations:
|
||||
"""Tests for resolve_citations function."""
|
||||
|
||||
def _make_result(self, chunk_id: str) -> SearchResult:
|
||||
return SearchResult(
|
||||
content="test content",
|
||||
score=1.0,
|
||||
chunk_id=chunk_id,
|
||||
document_id="doc-1",
|
||||
document_uri="test.md",
|
||||
document_title="Test Doc",
|
||||
)
|
||||
|
||||
def test_resolves_exact_ids(self):
|
||||
results = [self._make_result("abc123")]
|
||||
citations = resolve_citations(["abc123"], results)
|
||||
assert len(citations) == 1
|
||||
assert citations[0].chunk_id == "abc123"
|
||||
|
||||
def test_strips_brackets_from_ids(self):
|
||||
results = [self._make_result("abc123")]
|
||||
citations = resolve_citations(["[abc123]"], results)
|
||||
assert len(citations) == 1
|
||||
assert citations[0].chunk_id == "abc123"
|
||||
|
||||
def test_skips_unmatched_ids(self):
|
||||
results = [self._make_result("abc123")]
|
||||
citations = resolve_citations(["nonexistent"], results)
|
||||
assert len(citations) == 0
|
||||
|
||||
def test_empty_cited_chunks(self):
|
||||
results = [self._make_result("abc123")]
|
||||
citations = resolve_citations([], results)
|
||||
assert len(citations) == 0
|
||||
|
|
@ -1,31 +0,0 @@
|
|||
from haiku.rag.agents.research.prompts import (
|
||||
ITERATIVE_PLAN_PROMPT,
|
||||
ITERATIVE_PLAN_PROMPT_WITH_CONTEXT,
|
||||
)
|
||||
|
||||
|
||||
def test_iterative_plan_prompt_proposes_first_question():
|
||||
"""ITERATIVE_PLAN_PROMPT should instruct to propose the first question."""
|
||||
assert "first question" in ITERATIVE_PLAN_PROMPT.lower()
|
||||
assert "is_complete=False" in ITERATIVE_PLAN_PROMPT
|
||||
|
||||
|
||||
def test_iterative_plan_prompt_with_context_evaluates_evidence():
|
||||
"""ITERATIVE_PLAN_PROMPT_WITH_CONTEXT should evaluate prior answers."""
|
||||
assert "prior_answers" in ITERATIVE_PLAN_PROMPT_WITH_CONTEXT
|
||||
assert (
|
||||
"evaluat" in ITERATIVE_PLAN_PROMPT_WITH_CONTEXT.lower()
|
||||
) # matches evaluate/evaluating
|
||||
|
||||
|
||||
def test_prompt_selection_uses_context_prompt_with_prior_answers():
|
||||
"""When prior_answers exist, should use ITERATIVE_PLAN_PROMPT_WITH_CONTEXT."""
|
||||
has_prior_answers = True
|
||||
|
||||
effective_plan_prompt = (
|
||||
ITERATIVE_PLAN_PROMPT_WITH_CONTEXT
|
||||
if has_prior_answers
|
||||
else ITERATIVE_PLAN_PROMPT
|
||||
)
|
||||
|
||||
assert effective_plan_prompt == ITERATIVE_PLAN_PROMPT_WITH_CONTEXT
|
||||
|
|
@ -1,109 +0,0 @@
|
|||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from haiku.rag.agents.research.dependencies import ResearchContext
|
||||
from haiku.rag.agents.research.graph import build_research_graph
|
||||
from haiku.rag.agents.research.models import ResearchReport
|
||||
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def vcr_cassette_dir():
|
||||
return str(
|
||||
Path(__file__).parent.parent.parent / "cassettes" / "test_research_graph"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_graph_end_to_end(allow_model_requests, temp_db_path, qa_corpus):
|
||||
"""Test research graph with real LLM calls recorded via VCR."""
|
||||
graph = build_research_graph()
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = qa_corpus[0]
|
||||
await client.create_document(
|
||||
content=doc["document_extracted"], uri=doc["document_id"]
|
||||
)
|
||||
|
||||
state = ResearchState(
|
||||
context=ResearchContext(original_question=doc["question"]),
|
||||
max_iterations=1,
|
||||
max_concurrency=1,
|
||||
)
|
||||
|
||||
deps = ResearchDeps(client=client)
|
||||
|
||||
result = await graph.run(state=state, deps=deps)
|
||||
|
||||
assert result is not None
|
||||
assert isinstance(result, ResearchReport)
|
||||
assert result.title
|
||||
assert result.executive_summary
|
||||
|
||||
|
||||
def test_iterative_plan_result_model():
|
||||
"""Test IterativePlanResult model validation."""
|
||||
from haiku.rag.agents.research.models import IterativePlanResult
|
||||
|
||||
# Test complete state
|
||||
complete = IterativePlanResult(
|
||||
is_complete=True,
|
||||
next_question=None,
|
||||
reasoning="All aspects covered.",
|
||||
)
|
||||
assert complete.is_complete is True
|
||||
assert complete.next_question is None
|
||||
|
||||
# Test continue state
|
||||
continue_result = IterativePlanResult(
|
||||
is_complete=False,
|
||||
next_question="What are the specific requirements?",
|
||||
reasoning="Need more details.",
|
||||
)
|
||||
assert continue_result.is_complete is False
|
||||
assert continue_result.next_question == "What are the specific requirements?"
|
||||
|
||||
|
||||
def test_build_research_graph_returns_graph():
|
||||
"""Test build_research_graph returns a valid Graph instance."""
|
||||
from pydantic_graph.beta import Graph
|
||||
|
||||
graph = build_research_graph()
|
||||
assert graph is not None
|
||||
assert isinstance(graph, Graph)
|
||||
|
||||
|
||||
def test_format_context_for_prompt_basic():
|
||||
"""Test format_context_for_prompt with basic context."""
|
||||
from haiku.rag.agents.research.dependencies import ResearchContext
|
||||
from haiku.rag.agents.research.graph import format_context_for_prompt
|
||||
|
||||
context = ResearchContext(original_question="What is X?")
|
||||
result = format_context_for_prompt(context)
|
||||
|
||||
assert "<context>" in result
|
||||
assert "What is X?" in result
|
||||
|
||||
|
||||
def test_format_context_for_prompt_with_prior_answers():
|
||||
"""Test format_context_for_prompt includes prior_answers."""
|
||||
from haiku.rag.agents.research.dependencies import ResearchContext
|
||||
from haiku.rag.agents.research.graph import format_context_for_prompt
|
||||
from haiku.rag.agents.research.models import SearchAnswer
|
||||
|
||||
context = ResearchContext(original_question="Main question?")
|
||||
context.add_qa_response(
|
||||
SearchAnswer(
|
||||
query="Sub question?",
|
||||
answer="The answer is here.",
|
||||
confidence=0.9,
|
||||
)
|
||||
)
|
||||
|
||||
result = format_context_for_prompt(context)
|
||||
|
||||
assert "<prior_answers>" in result
|
||||
assert "Sub question?" in result
|
||||
assert "The answer is here." in result
|
||||
|
|
@ -1,125 +0,0 @@
|
|||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from haiku.rag.agents.research.dependencies import ResearchContext
|
||||
from haiku.rag.agents.research.graph import build_research_graph
|
||||
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def vcr_cassette_dir():
|
||||
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_search_filter")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def client_with_docs(temp_db_path):
|
||||
"""Create a client with two distinct documents."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
# Add two documents with distinct content
|
||||
doc1 = await client.create_document(
|
||||
"Document about cats: Cats are small furry mammals that purr.",
|
||||
title="Cat Facts",
|
||||
)
|
||||
doc2 = await client.create_document(
|
||||
"Document about dogs: Dogs are loyal companions that bark.",
|
||||
title="Dog Facts",
|
||||
)
|
||||
|
||||
yield client, doc1.id, doc2.id
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_search_filter_restricts_results(client_with_docs):
|
||||
"""Test that search_filter restricts search to specified documents."""
|
||||
client, doc1_id, doc2_id = client_with_docs
|
||||
|
||||
# Search without filter - should find both
|
||||
results_all = await client.search("animals mammals companions")
|
||||
assert len(results_all) >= 1
|
||||
|
||||
# Search with filter for doc1 only
|
||||
filter_doc1 = f"id = '{doc1_id}'"
|
||||
results_filtered = await client.search(
|
||||
"animals mammals companions", filter=filter_doc1
|
||||
)
|
||||
|
||||
# All results should be from doc1
|
||||
for result in results_filtered:
|
||||
assert result.document_id == doc1_id
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
@pytest.mark.asyncio
|
||||
async def test_research_graph_uses_search_filter(
|
||||
allow_model_requests, client_with_docs
|
||||
):
|
||||
"""Test that research 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
|
||||
|
||||
graph = build_research_graph()
|
||||
|
||||
# Create state with search_filter
|
||||
filter_clause = f"id = '{doc1_id}'"
|
||||
state = ResearchState(
|
||||
context=ResearchContext(original_question="Tell me about animals"),
|
||||
max_iterations=1,
|
||||
search_filter=filter_clause,
|
||||
)
|
||||
|
||||
deps = ResearchDeps(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.vcr()
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_filter_none_searches_all(allow_model_requests, client_with_docs):
|
||||
"""Test that search_filter=None searches all documents."""
|
||||
client, doc1_id, doc2_id = client_with_docs
|
||||
|
||||
# Track search calls
|
||||
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
|
||||
|
||||
graph = build_research_graph()
|
||||
|
||||
# Create state without search_filter (None)
|
||||
state = ResearchState(
|
||||
context=ResearchContext(original_question="Tell me about animals"),
|
||||
max_iterations=1,
|
||||
search_filter=None,
|
||||
)
|
||||
|
||||
deps = ResearchDeps(client=client)
|
||||
|
||||
await graph.run(state=state, deps=deps)
|
||||
|
||||
# Verify search was called with None filter
|
||||
assert len(search_calls) > 0, "Expected search to be called"
|
||||
for call in search_calls:
|
||||
assert call["filter"] is None, f"Expected filter None, got '{call['filter']}'"
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
|
@ -158,8 +158,8 @@ async def test_chat_history_can_add_tool_calls(temp_db_path: Path):
|
|||
@pytest.mark.asyncio
|
||||
async def test_chat_history_can_add_citations(temp_db_path: Path):
|
||||
"""Test that ChatHistory can display inline citations."""
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
|
||||
app, mock_client = _make_app(temp_db_path)
|
||||
|
||||
|
|
@ -241,8 +241,8 @@ async def test_clear_chat_resets_state(temp_db_path: Path):
|
|||
@pytest.mark.asyncio
|
||||
async def test_citation_expand_collapse_with_enter(temp_db_path: Path):
|
||||
"""Test that pressing Enter on a focused citation toggles expand/collapse."""
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
|
||||
app, mock_client = _make_app(temp_db_path)
|
||||
|
||||
|
|
@ -279,9 +279,9 @@ async def test_citation_expand_collapse_with_enter(temp_db_path: Path):
|
|||
@pytest.mark.asyncio
|
||||
async def test_show_citations_renders_from_flat_state(temp_db_path: Path):
|
||||
"""Citations in state (flat list[str]) render into the chat history."""
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.chat.app import RAG_STATE_NAMESPACE
|
||||
from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
|
||||
app, mock_client = _make_app_with_state(temp_db_path)
|
||||
|
||||
|
|
|
|||
|
|
@ -335,10 +335,10 @@ class TestAnalysisLifespan:
|
|||
assert result
|
||||
|
||||
async def test_lifespan_clears_executions_citations_searches(self, rag_db):
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan
|
||||
from haiku.rag.skills._tools import CodeExecutionEntry
|
||||
from haiku.rag.skills.analysis import AnalysisState
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
|
||||
config = AppConfig()
|
||||
lifespan = make_analysis_lifespan(rag_db, config)
|
||||
|
|
|
|||
|
|
@ -458,9 +458,9 @@ class TestLifespan:
|
|||
assert result
|
||||
|
||||
async def test_lifespan_clears_citations_and_searches_but_keeps_index(self, rag_db):
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.skills._deps import RAGRunDeps, make_rag_lifespan
|
||||
from haiku.rag.skills.rag import RAGState
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
|
||||
config = AppConfig()
|
||||
lifespan = make_rag_lifespan(rag_db, config)
|
||||
|
|
|
|||
|
|
@ -1,84 +0,0 @@
|
|||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from haiku.rag.agents.research.models import ResearchReport
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
|
||||
async def test_client_research_report(temp_db_path):
|
||||
"""Test client.research() delegates to research graph in report mode."""
|
||||
mock_report = ResearchReport(
|
||||
title="Test Report",
|
||||
executive_summary="Summary",
|
||||
main_findings=["Finding 1"],
|
||||
conclusions=["Conclusion 1"],
|
||||
sources_summary="Sources",
|
||||
)
|
||||
|
||||
with patch("haiku.rag.agents.research.graph.build_research_graph") as mock_build:
|
||||
mock_graph = AsyncMock()
|
||||
mock_graph.run = AsyncMock(return_value=mock_report)
|
||||
mock_build.return_value = mock_graph
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
result = await client.research(question="What is X?")
|
||||
|
||||
assert result is mock_report
|
||||
mock_build.assert_called_once()
|
||||
|
||||
# Verify graph.run was called with correct state/deps
|
||||
mock_graph.run.assert_called_once()
|
||||
call_kwargs = mock_graph.run.call_args[1]
|
||||
assert call_kwargs["state"].context.original_question == "What is X?"
|
||||
assert isinstance(call_kwargs["deps"].client, HaikuRAG)
|
||||
|
||||
|
||||
async def test_client_research_passes_filter(temp_db_path):
|
||||
"""Test client.research() passes filter to state."""
|
||||
mock_report = ResearchReport(
|
||||
title="Test",
|
||||
executive_summary="Summary",
|
||||
main_findings=[],
|
||||
conclusions=[],
|
||||
sources_summary="",
|
||||
)
|
||||
|
||||
with patch("haiku.rag.agents.research.graph.build_research_graph") as mock_build:
|
||||
mock_graph = AsyncMock()
|
||||
mock_graph.run = AsyncMock(return_value=mock_report)
|
||||
mock_build.return_value = mock_graph
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
await client.research(
|
||||
question="What is X?",
|
||||
filter="uri LIKE '%test%'",
|
||||
)
|
||||
|
||||
call_kwargs = mock_graph.run.call_args[1]
|
||||
assert call_kwargs["state"].search_filter == "uri LIKE '%test%'"
|
||||
|
||||
|
||||
async def test_client_research_uses_config(temp_db_path):
|
||||
"""Test client.research() passes config to graph builder and state."""
|
||||
mock_report = ResearchReport(
|
||||
title="Test",
|
||||
executive_summary="Summary",
|
||||
main_findings=[],
|
||||
conclusions=[],
|
||||
sources_summary="",
|
||||
)
|
||||
|
||||
with patch("haiku.rag.agents.research.graph.build_research_graph") as mock_build:
|
||||
mock_graph = AsyncMock()
|
||||
mock_graph.run = AsyncMock(return_value=mock_report)
|
||||
mock_build.return_value = mock_graph
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
await client.research(question="What is X?")
|
||||
|
||||
_, kwargs = mock_build.call_args
|
||||
assert kwargs["config"] is client._config
|
||||
|
||||
call_kwargs = mock_graph.run.call_args[1]
|
||||
state = call_kwargs["state"]
|
||||
assert state.max_iterations == client._config.research.max_iterations
|
||||
assert state.max_concurrency == client._config.research.max_concurrency
|
||||
|
|
@ -200,7 +200,6 @@ def test_generate_default_config_completeness():
|
|||
assert config.environment == "production"
|
||||
assert config.embeddings.model.provider == "ollama"
|
||||
assert config.qa.model.provider == "ollama"
|
||||
assert config.research.model.provider == "ollama"
|
||||
assert config.reranking.model is None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -76,7 +76,7 @@ async def test_download_models_ollama_pulls_models(mock_to_thread):
|
|||
async for progress in download_models(Config):
|
||||
events.append(progress)
|
||||
|
||||
# Default config has embeddings=qwen3-embedding:4b, qa/research=gpt-oss
|
||||
# Default config has embeddings=qwen3-embedding:4b, qa=gpt-oss
|
||||
ollama_models = {"gpt-oss", "qwen3-embedding:4b"}
|
||||
ollama_events = [e for e in events if e.model in ollama_models]
|
||||
pulling_events = [e for e in ollama_events if e.status == "pulling"]
|
||||
|
|
@ -100,7 +100,6 @@ async def test_download_models_no_ollama_models(mock_to_thread):
|
|||
config = AppConfig()
|
||||
config.embeddings.model.provider = "openai"
|
||||
config.qa.model.provider = "openai"
|
||||
config.research.model.provider = "openai"
|
||||
|
||||
events = []
|
||||
async for progress in download_models(config):
|
||||
|
|
|
|||
|
|
@ -524,7 +524,7 @@ def test_format_citations_empty():
|
|||
|
||||
|
||||
def test_format_citations_with_citation():
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.utils import format_citations
|
||||
|
||||
citation = Citation(
|
||||
|
|
@ -547,7 +547,7 @@ def test_format_citations_with_citation():
|
|||
|
||||
|
||||
def test_format_citations_multiple_pages():
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.utils import format_citations
|
||||
|
||||
citation = Citation(
|
||||
|
|
@ -563,7 +563,7 @@ def test_format_citations_multiple_pages():
|
|||
|
||||
|
||||
def test_format_citations_no_title():
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.utils import format_citations
|
||||
|
||||
citation = Citation(
|
||||
|
|
@ -578,7 +578,7 @@ def test_format_citations_no_title():
|
|||
|
||||
|
||||
def test_format_citations_with_index():
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.utils import format_citations
|
||||
|
||||
citation = Citation(
|
||||
|
|
@ -594,7 +594,7 @@ def test_format_citations_with_index():
|
|||
|
||||
|
||||
def test_format_citations_sequential_indices():
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.utils import format_citations
|
||||
|
||||
citations = [
|
||||
|
|
@ -622,7 +622,7 @@ def test_format_citations_sequential_indices():
|
|||
|
||||
|
||||
def test_format_citations_picture_refs_render_as_markers():
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.utils import format_citations
|
||||
|
||||
citation = Citation(
|
||||
|
|
@ -657,7 +657,7 @@ async def test_format_citations_rich_empty():
|
|||
|
||||
|
||||
async def test_format_citations_rich_header_and_footer():
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.utils import format_citations_rich
|
||||
|
||||
citation = Citation(
|
||||
|
|
@ -679,7 +679,7 @@ async def test_format_citations_rich_header_and_footer():
|
|||
|
||||
|
||||
async def test_format_citations_rich_truncates_long_content():
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.utils import CITATION_PREVIEW_CHARS, format_citations_rich
|
||||
|
||||
citation = Citation(
|
||||
|
|
@ -694,7 +694,7 @@ async def test_format_citations_rich_truncates_long_content():
|
|||
|
||||
|
||||
async def test_format_citations_rich_picture_marker_without_client():
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.utils import format_citations_rich
|
||||
|
||||
citation = Citation(
|
||||
|
|
|
|||
Loading…
Reference in a new issue