Delete the standalone QA and analysis agents
This commit is contained in:
Yiorgis Gozadinos
parent
ceb645564c
commit
6f95e2bc27
47 changed files with 131 additions and 1482 deletions
26
CHANGELOG.md
26
CHANGELOG.md
|
|
@ -17,28 +17,28 @@
|
|||
- `AnalysisResult.program`. The per-execution programs are still tracked on `AnalysisState.executions` (the analysis skill's `execute_code` tool populates it); consumers that need the executed code should pull it from the skill state instead of the function return value.
|
||||
- `--cite` flag on `haiku-rag ask`. Citations always render after the answer now.
|
||||
- `system_prompt` kwarg on `client.ask`. No production caller used it; `config.prompts.domain_preamble` already covers the preamble use case.
|
||||
- `evaluations optimize` subcommand and the GEPA prompt-optimization module. Drops the `gepa` dependency. Hand-tuning SKILL.md / `prompts.qa` against `evaluations run` is the loop we actually iterate; the auto-mutation surface was unused and conflicted with the project's no-prompt-string-tests rule.
|
||||
- `--target qa` from `evaluations run`. The skill path supersedes the standalone QA agent — use `--target rag-skill` (now the default) or `--target analysis-skill`.
|
||||
|
||||
### Changed
|
||||
|
||||
- `search.limit` default lowered from `10` to `5`. Reduces text + binary noise in vision-tool returns (picture count tracks result count after expansion + dedup); the cite path still selects from all returned chunks.
|
||||
- Search result formatter surfaces picture captions on a labelled line when a chunk's expanded refs include pictures. The OpenAI vision API has no identifier field for binary parts, so the caption is the only signal a model can use to map a description to the figure it sees.
|
||||
- `AnalysisConfig.model` defaults to `None` (was `ollama:gpt-oss/no-thinking/temp=0`). Consumers resolve via `config.analysis.model or config.qa.model`, so the analysis skill inherits the QA driving model when no `analysis` block is in YAML. Set `analysis.model` explicitly to keep the analysis pipeline on a different model from QA. **Note**: the type is now `ModelConfig | None`; external code reading `cfg.analysis.model.X` directly will need to handle the `None` case.
|
||||
- `client.ask` and `client.analyze` now route through the rag and rag-analysis skills internally (via `haiku.skills.run_skill`). Return shapes preserved: `client.ask` still returns `tuple[str, list[Citation]]`; `client.analyze` returns `AnalysisResult(answer, citations)`. The standalone QA and analysis agents under `agents/qa/` and `agents/analysis/agent.py` remain for now but are scheduled for removal once GEPA optimisation and the `--target qa` eval CLI move to the skill path.
|
||||
|
||||
### Fixed
|
||||
|
||||
- Chat TUI's state-edit screen now syntax-highlights JSON instead of falling back to plain text. Adds `tree-sitter` + `tree-sitter-json` to the `[tui]` extra.
|
||||
- Standalone QA agent (`haiku.rag.agents.qa.*`) and analysis agent module (`haiku.rag.agents.analysis.agent`, `haiku.rag.agents.analysis.prompts`). Also drops `RawAnalysisResult`, `CodeExecution`, `AnalysisDeps`, and the dead `documents=` preload path in `Sandbox`.
|
||||
- `prompts.qa` config field.
|
||||
- `evaluations optimize` subcommand and GEPA prompt-optimization. Drops `gepa` dep.
|
||||
- `--target qa` from `evaluations run`. Default is now `rag-skill`.
|
||||
|
||||
### Changed
|
||||
|
||||
- `haiku.rag.agents.analysis` moved to `haiku.rag.sandbox`. Public surface: `from haiku.rag.sandbox import Sandbox, SandboxResult, AnalysisContext, AnalysisResult`.
|
||||
- `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.
|
||||
- `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`.
|
||||
- Bump `pydantic-ai-slim>=1.96.0`. Migrate off deprecated APIs: AG-UI imports use `pydantic_ai.ui.ag_ui`, docs/CLI examples use the explicit `openai-chat:` model prefix, and `Agent(retries=)` is split into `tool_retries=` + `output_retries=`.
|
||||
- Bump `pydantic-monty>=0.0.17`. Migrate off deprecated `pydantic_monty.run_repl_async(repl, ...)` to `repl.feed_run_async(...)`.
|
||||
- Cap `transformers<5.0.0` in the `mxbai` extra: `mxbai-rerank>=0.1.6` calls `tokenizer.prepare_for_model` which transformers 5 removed.
|
||||
- Refresh the rest of the lockfile to latest within current constraints (pydantic, pydantic-ai, rich, ruff, ty, pytest, torch, textual, textual-image, watchfiles, pre-commit, datasets, and transitives).
|
||||
|
||||
### Fixed
|
||||
|
||||
- Chat TUI's state-edit screen syntax-highlights JSON instead of falling back to plain text. Adds `tree-sitter` + `tree-sitter-json` to the `[tui]` extra.
|
||||
|
||||
## [0.47.0] - 2026-05-14
|
||||
|
||||
### Added
|
||||
|
|
|
|||
|
|
@ -11,11 +11,11 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p
|
|||
|
||||
- **Hybrid search** — Vector + full-text with Reciprocal Rank Fusion
|
||||
- **Multimodal & cross-modal search** — Multimodal embedders (vLLM) put picture vectors in the same space as text; supports text-as-query → figure hits and image-as-query
|
||||
- **Question answering** — QA agents with citations (page numbers, section headings)
|
||||
- **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 agents** — Multi-agent workflows via pydantic-graph: plan, search, evaluate, synthesize
|
||||
- **Analysis agent** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
|
||||
- **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
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
# Analysis Agent
|
||||
# Analysis
|
||||
|
||||
The analysis agent enables complex analytical tasks by writing and executing Python code in a sandboxed environment. It solves problems that traditional RAG struggles with:
|
||||
The analysis skill enables complex analytical tasks by writing and executing Python code in a sandboxed environment. It solves problems that traditional RAG struggles with:
|
||||
|
||||
- **Aggregation**: "How many documents mention security vulnerabilities?"
|
||||
- **Computation**: "What's the average revenue across all quarterly reports?"
|
||||
|
|
@ -9,10 +9,10 @@ The analysis agent enables complex analytical tasks by writing and executing Pyt
|
|||
|
||||
## How It Works
|
||||
|
||||
1. The agent receives a question
|
||||
1. The skill receives a question
|
||||
2. It writes Python code to explore the knowledge base
|
||||
3. Code executes in a sandboxed Python interpreter with access to search, LLM, and a virtual filesystem of documents
|
||||
4. The agent iterates: run code, examine results, refine approach
|
||||
3. Code executes in a sandboxed Python interpreter with access to search and a virtual filesystem of documents
|
||||
4. The skill iterates: run code, examine results, refine approach
|
||||
5. Final answer is synthesized from the gathered data
|
||||
|
||||
## CLI Usage
|
||||
|
|
@ -21,11 +21,8 @@ The analysis agent enables complex analytical tasks by writing and executing Pyt
|
|||
# Basic usage
|
||||
haiku-rag analyze "How many documents are in the database?"
|
||||
|
||||
# With document filter (restricts what the agent can access)
|
||||
# With document filter (restricts what the skill can access)
|
||||
haiku-rag analyze "Summarize the key points" --filter "uri LIKE '%report%'"
|
||||
|
||||
# Pre-load specific documents
|
||||
haiku-rag analyze "Compare these two reports" --document "Q1 Report" --document "Q2 Report"
|
||||
```
|
||||
|
||||
## Python Usage
|
||||
|
|
@ -34,24 +31,20 @@ haiku-rag analyze "Compare these two reports" --document "Q1 Report" --document
|
|||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
async with HaikuRAG(path_to_db) as client:
|
||||
# Basic question
|
||||
result = await client.analyze("How many documents mention 'security'?")
|
||||
print(result.answer) # The answer
|
||||
print(result.program) # The final consolidated program
|
||||
print(result.answer)
|
||||
for citation in result.citations:
|
||||
print(citation.uri, citation.title)
|
||||
|
||||
# With filter (agent can only see filtered documents)
|
||||
# With filter (skill can only see filtered documents)
|
||||
result = await client.analyze(
|
||||
"What is the total revenue?",
|
||||
filter="title LIKE '%Financial%'"
|
||||
)
|
||||
|
||||
# Pre-load specific documents
|
||||
result = await client.analyze(
|
||||
"Compare the conclusions",
|
||||
documents=["Report A", "Report B"]
|
||||
)
|
||||
```
|
||||
|
||||
The executed Python program(s) for each turn are not on `AnalysisResult` itself; they live on `AnalysisState.executions` while the skill runs.
|
||||
|
||||
## Sandbox Capabilities
|
||||
|
||||
The agent's code runs in a sandboxed Python interpreter ([pydantic-monty](https://github.com/pydantic/monty)) with:
|
||||
|
|
@ -60,9 +53,8 @@ The agent's code runs in a sandboxed Python interpreter ([pydantic-monty](https:
|
|||
|
||||
| Function | Description |
|
||||
|----------|-------------|
|
||||
| `search(query, limit)` | Hybrid search (vector + full-text) with automatic context expansion. Returns `doc_item_refs` for cross-referencing with `items.jsonl` |
|
||||
| `search(query, limit)` | Hybrid search (vector + full-text) with automatic context expansion. Returns `doc_item_refs` and `picture_refs` for cross-referencing with `items.jsonl` |
|
||||
| `list_documents()` | List all documents in the knowledge base |
|
||||
| `llm(prompt)` | Call an LLM for classification, summarization, or extraction |
|
||||
|
||||
### Document Filesystem
|
||||
|
||||
|
|
@ -81,13 +73,11 @@ All documents are mounted as a virtual filesystem at `/documents/`. The agent us
|
|||
|
||||
Search results include `doc_item_refs` (e.g. `["#/texts/5", "#/tables/0"]`) that match `self_ref` values in `items.jsonl`, enabling navigation from search hits to document structure.
|
||||
|
||||
When documents are pre-loaded via the `documents` parameter, they are also injected as a `documents` variable accessible in the sandbox code.
|
||||
|
||||
### Python Features
|
||||
|
||||
The interpreter supports a subset of Python: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `filter()`, `getattr()`, try/except, file I/O via `pathlib.Path`, and the `json`, `re`, `math` modules.
|
||||
|
||||
Not supported: most imports (only `json`, `re`, `math`, `pathlib` are available), class definitions, generators/yield, match statements, decorators, `with` statements. For pattern matching, the agent can use `import re`, string methods, or the `llm()` function.
|
||||
Not supported: most imports (only `json`, `re`, `math`, `pathlib` are available), class definitions, generators/yield, match statements, decorators, `with` statements. For pattern matching, use `import re` or string methods.
|
||||
|
||||
### Security
|
||||
|
||||
|
|
@ -101,10 +91,10 @@ Code executes in an isolated interpreter with:
|
|||
|
||||
## Context Filter
|
||||
|
||||
The `filter` parameter restricts what documents the agent can access. Unlike tool parameters, the filter is applied automatically and cannot be bypassed by the LLM — both the VFS and search results are scoped to the filter:
|
||||
The `filter` parameter restricts what documents the skill can access. Unlike tool parameters, the filter is applied automatically and cannot be bypassed by the LLM — both the VFS and search results are scoped to the filter:
|
||||
|
||||
```python
|
||||
# Agent can only see documents with "confidential" in the URI
|
||||
# Skill can only see documents with "confidential" in the URI
|
||||
result = await client.analyze(
|
||||
"Summarize all findings",
|
||||
filter="uri LIKE '%confidential%'"
|
||||
|
|
|
|||
|
|
@ -1,51 +1,30 @@
|
|||
# Agents
|
||||
|
||||
Three agentic flows are provided by haiku.rag:
|
||||
haiku.rag provides:
|
||||
|
||||
- **Simple QA Agent** — a focused question answering agent
|
||||
- **Research Graph** — a multi-step research workflow with question decomposition
|
||||
- **Analysis Agent** — complex analytical tasks via sandboxed Python code execution (see [Analysis Agent](analysis.md))
|
||||
- **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).
|
||||
|
||||
For multi-turn conversational RAG, haiku.rag provides [skills](../skills/index.md) built on [haiku.skills](https://github.com/ggozad/haiku.skills). The skills bundle search, Q&A, analysis, and research tools with session state management.
|
||||
`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.
|
||||
|
||||
## Simple QA Agent
|
||||
|
||||
The simple QA agent answers a single question using the knowledge base. It retrieves relevant chunks, optionally expands context around them, and asks the model to answer strictly based on that context.
|
||||
|
||||
Key points:
|
||||
|
||||
- Uses a single `search_documents` tool to fetch relevant chunks
|
||||
- Can be run with or without inline citations in the prompt
|
||||
- Returns a plain string answer
|
||||
|
||||
**CLI usage:**
|
||||
## Question Answering
|
||||
|
||||
```bash
|
||||
haiku-rag ask "What is climate change?"
|
||||
|
||||
# With citations
|
||||
haiku-rag ask "What is climate change?" --cite
|
||||
```
|
||||
|
||||
**Python usage:**
|
||||
|
||||
```python
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config.models import ModelConfig
|
||||
from haiku.rag.agents.qa.agent import QuestionAnswerAgent
|
||||
|
||||
async with HaikuRAG(path_to_db) as client:
|
||||
agent = QuestionAnswerAgent(
|
||||
client=client,
|
||||
model_config=ModelConfig(provider="openai", name="gpt-4o-mini"),
|
||||
)
|
||||
|
||||
answer, citations = await agent.answer("What is climate change?")
|
||||
print(answer)
|
||||
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.
|
||||
|
|
|
|||
|
|
@ -185,7 +185,7 @@ Search uses hybrid (vector + full-text) search across all chunks.
|
|||
|
||||
### Context Expansion
|
||||
|
||||
Press `c` while viewing a chunk to see the expanded context that would be provided to the QA agent:
|
||||
Press `c` while viewing a chunk to see the expanded context that would be provided to the rag skill:
|
||||
|
||||
- Section-aware expansion: expands to fill the current document section
|
||||
- Noise filtering: footnotes, page headers/footers excluded from structured documents
|
||||
|
|
|
|||
|
|
@ -61,9 +61,9 @@ evaluations run repliqa --config /path/to/haiku.rag.yaml --db /path/to/custom.la
|
|||
- `--skip-qa` - Skip QA benchmark
|
||||
- `--limit N` - Limit number of test cases
|
||||
- `--name NAME` - Override the evaluation name
|
||||
- `--judge-model PROVIDER:NAME` - Override the LLM judge model. Defaults to `ollama:qwen3.6` so the judge stays stable when the QA / skill model changes.
|
||||
- `--target {qa,rag-skill,analysis-skill}` - Choose what to benchmark (default: `qa`). `rag-skill` and `analysis-skill` run the corresponding [skill](skills/index.md) end-to-end against the same datasets and judge as the QA agent.
|
||||
- `--skill-model PROVIDER:NAME` - Override the skill model independently from the judge (default: `config.qa.model`). Only valid with skill targets.
|
||||
- `--judge-model PROVIDER:NAME` - Override the LLM judge model. Defaults to `ollama:qwen3.6` so the judge stays stable when the answering model changes.
|
||||
- `--target {rag-skill,analysis-skill}` - Choose which [skill](skills/index.md) to benchmark end-to-end against the same datasets and judge (default: `rag-skill`).
|
||||
- `--skill-model PROVIDER:NAME` - Override the skill model independently from the judge (default: `config.qa.model`, or `config.analysis.model` when set for `--target analysis-skill`).
|
||||
|
||||
If no config file is specified, the script searches standard locations: `./haiku.rag.yaml`, user config directory, then falls back to defaults.
|
||||
|
||||
|
|
@ -88,7 +88,7 @@ If no config file is specified, the script searches standard locations: `./haiku
|
|||
|
||||
### QA Accuracy
|
||||
|
||||
For question-answering evaluation, `pydantic-evals` coordinates an LLM judge to determine whether answers are correct. The default judge is `ollama:qwen3.6` — pinned so changes to the QA or skill model don't change the judge underneath. Override per run with `--judge-model provider:name`. Accuracy is the fraction of correctly answered questions.
|
||||
For question-answering evaluation, `pydantic-evals` coordinates an LLM judge to determine whether answers are correct. The default judge is `ollama:qwen3.6` — pinned so changes to the skill model don't change the judge underneath. Override per run with `--judge-model provider:name`. Accuracy is the fraction of correctly answered questions.
|
||||
|
||||
We picked `qwen3.6` over the previously-pinned `gpt-oss` after a 4-cell calibration (gpt-oss / qwen3.6 as both answerer and judge, with Claude Opus 4.7 as a reference). `qwen3.6` had κ ≥ 0.66 vs the reference on both same-family and cross-family answerers (vs ~0.39–0.55 for `gpt-oss`) and showed no measurable self-preference bias, while `gpt-oss` was ~10 pp more lenient on its own outputs.
|
||||
|
||||
|
|
|
|||
17
docs/cli.md
17
docs/cli.md
|
|
@ -171,21 +171,15 @@ Ask questions about your documents:
|
|||
haiku-rag ask "Who is the author of haiku.rag?"
|
||||
```
|
||||
|
||||
Ask questions with citations showing source documents:
|
||||
```bash
|
||||
haiku-rag ask "Who is the author of haiku.rag?" --cite
|
||||
```
|
||||
|
||||
Filter to specific documents:
|
||||
```bash
|
||||
haiku-rag ask "What are the main findings?" --filter "uri LIKE '%paper%'"
|
||||
```
|
||||
|
||||
The QA agent searches your documents for relevant information and provides a comprehensive answer. When available, citations use the document title; otherwise they fall back to the URI.
|
||||
`ask` runs the [rag skill](skills/index.md) and always renders citations under the answer. When available, citations use the document title; otherwise they fall back to the URI.
|
||||
|
||||
Flags:
|
||||
|
||||
- `--cite`: Include citations showing which documents were used
|
||||
- `--filter` / `-f`: Restrict searches to documents matching the filter (see [Filtering Search Results](python.md#filtering-search-results))
|
||||
|
||||
## Chat
|
||||
|
|
@ -271,18 +265,11 @@ Filter to specific documents:
|
|||
haiku-rag analyze "What is the total revenue?" --filter "title LIKE '%Financial%'"
|
||||
```
|
||||
|
||||
Pre-load specific documents for comparison:
|
||||
|
||||
```bash
|
||||
haiku-rag analyze "Compare the conclusions" --document "Report A" --document "Report B"
|
||||
```
|
||||
|
||||
Flags:
|
||||
|
||||
- `--filter` / `-f`: SQL WHERE clause to restrict document access
|
||||
- `--document` / `-d`: Pre-load a document by title or ID (can repeat)
|
||||
|
||||
See [Analysis Agent](agents/analysis.md) for details on capabilities and configuration.
|
||||
See [Analysis](agents/analysis.md) for details on capabilities and configuration.
|
||||
|
||||
## Create Skill
|
||||
|
||||
|
|
|
|||
|
|
@ -106,8 +106,7 @@ search:
|
|||
vector_refine_factor: 30
|
||||
|
||||
prompts:
|
||||
domain_preamble: "" # Prepended to all agent prompts
|
||||
qa: null # Custom QA agent prompt (null = use default)
|
||||
domain_preamble: "" # Prepended to skill instructions and research prompts
|
||||
synthesis: null # Custom research synthesis prompt (null = use default)
|
||||
|
||||
processing:
|
||||
|
|
|
|||
|
|
@ -1,20 +1,17 @@
|
|||
# Prompt Customization
|
||||
|
||||
Customize the prompts used by haiku.rag's AI agents to better match your domain and use case.
|
||||
Customize the prompts used by haiku.rag's skills and research workflow to better match your domain and use case.
|
||||
|
||||
## Configuration
|
||||
|
||||
```yaml
|
||||
prompts:
|
||||
# Domain context prepended to all agent prompts
|
||||
# Domain context prepended to skill instructions and research prompts
|
||||
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 QA agent prompt (optional)
|
||||
qa: null
|
||||
|
||||
# Full replacement for research synthesis prompt (optional)
|
||||
synthesis: null
|
||||
|
||||
|
|
@ -24,13 +21,13 @@ prompts:
|
|||
|
||||
## Domain Preamble
|
||||
|
||||
The `domain_preamble` field provides **domain context** that is prepended to all agent prompts — the main agent, skill subagents, and internal agents (QA, research planning, search, evaluation, and synthesis). Use this to:
|
||||
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:
|
||||
|
||||
- Describe what the knowledge base contains
|
||||
- Clarify domain-specific terminology
|
||||
- Provide context that helps agents interpret ambiguous queries
|
||||
- Provide context that helps the model interpret ambiguous queries
|
||||
|
||||
**Important:** `domain_preamble` is for domain context, not behavioral instructions. Descriptions of subject matter, terminology, and content scope belong here. Behavioral guidance (tone, response style, formatting rules) belongs in the agent's system prompt or custom `prompts.qa`.
|
||||
**Important:** `domain_preamble` is for domain context, not behavioral instructions. Descriptions of subject matter, terminology, and content scope belong here. Behavioral guidance (tone, response style, formatting rules) lives in the skill's SKILL.md — fork the skill via `haiku-rag create-skill` to customize behavior.
|
||||
|
||||
**Example:**
|
||||
|
||||
|
|
@ -42,34 +39,6 @@ prompts:
|
|||
"Deployment" refers to Acme's managed deployment service, not general CI/CD.
|
||||
```
|
||||
|
||||
## Custom QA Prompt
|
||||
|
||||
Replace the default QA agent prompt entirely by setting `prompts.qa`. The prompt should instruct the agent how to:
|
||||
|
||||
1. Use the `search_documents` tool to find relevant content
|
||||
2. Interpret search results with scores and metadata
|
||||
3. Cite sources using chunk IDs
|
||||
4. Handle insufficient information
|
||||
|
||||
**Example:**
|
||||
|
||||
```yaml
|
||||
prompts:
|
||||
qa: |
|
||||
You are a concise technical assistant. Answer questions using only the knowledge base.
|
||||
|
||||
Process:
|
||||
1. Search for relevant documents using the search_documents tool
|
||||
2. Review results ordered by relevance (rank 1 = most relevant)
|
||||
3. Provide a brief, direct answer based on retrieved content
|
||||
|
||||
Guidelines:
|
||||
- Use only information from search results
|
||||
- Include chunk IDs in cited_chunks for sources you use
|
||||
- If information is insufficient, say so clearly
|
||||
- Be concise - avoid unnecessary elaboration
|
||||
```
|
||||
|
||||
## 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.
|
||||
|
|
@ -130,7 +99,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.",
|
||||
qa=None, # Use default QA prompt
|
||||
synthesis=None, # Use default synthesis prompt
|
||||
picture_description="Describe this image for search indexing.",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ Context expansion is automatic and section-aware. For structured documents (with
|
|||
|
||||
## Question Answering Configuration
|
||||
|
||||
Configure the QA workflow:
|
||||
Configure the rag skill (used by `client.ask`, `haiku-rag ask`, and the MCP `ask_question` tool):
|
||||
|
||||
```yaml
|
||||
qa:
|
||||
|
|
@ -29,13 +29,13 @@ qa:
|
|||
name: gpt-oss
|
||||
enable_thinking: true
|
||||
temperature: 0.3 # Default: 0.3
|
||||
vision: false # Set true for vision-capable QA models
|
||||
vision: false # Set true for vision-capable models
|
||||
max_searches: 3 # Maximum search tool calls per question
|
||||
```
|
||||
|
||||
- **model**: LLM configuration (see [Providers](providers.md#model-settings))
|
||||
- **model.vision**: Set to `true` for vision-capable QA models (`qwen2.5vl`, `qwen3.6`, `gpt-4o`, `claude-sonnet`, …). The agent'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 QA agent can make per question (default: 3)
|
||||
- **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
|
||||
|
||||
|
|
@ -60,7 +60,7 @@ The research workflow uses an iterative feedback loop: the planner proposes one
|
|||
|
||||
## Analysis Configuration
|
||||
|
||||
Configure the analysis agent:
|
||||
Configure the analysis skill:
|
||||
|
||||
```yaml
|
||||
analysis:
|
||||
|
|
@ -72,8 +72,8 @@ analysis:
|
|||
max_output_chars: 50000 # Truncate output after this many chars
|
||||
```
|
||||
|
||||
- **model**: LLM configuration (see [Providers](providers.md#model-settings))
|
||||
- **model**: LLM configuration (see [Providers](providers.md#model-settings)). When unset, falls back to `qa.model`.
|
||||
- **code_timeout**: Maximum seconds for each code execution (default: 60)
|
||||
- **max_output_chars**: Truncate code output after this many characters (default: 50000)
|
||||
|
||||
See [Analysis Agent](../agents/analysis.md) for usage details.
|
||||
See [Analysis](../agents/analysis.md) for usage details.
|
||||
|
|
|
|||
|
|
@ -8,11 +8,11 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p
|
|||
|
||||
- **Hybrid search** — Vector + full-text with Reciprocal Rank Fusion
|
||||
- **Multimodal & cross-modal search** — Multimodal embedders (vLLM) put picture vectors in the same space as text; supports text-as-query → figure hits and image-as-query
|
||||
- **Question answering** — QA agents with citations (page numbers, section headings)
|
||||
- **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 agents** — Multi-agent workflows via pydantic-graph: plan, search, evaluate, synthesize
|
||||
- **Analysis agent** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
|
||||
- **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
|
||||
|
|
|
|||
|
|
@ -421,19 +421,6 @@ for cite in citations:
|
|||
print(f" [{cite.chunk_id}] {cite.document_title or cite.document_uri}")
|
||||
```
|
||||
|
||||
Customize the QA agent's behavior with a custom system prompt:
|
||||
|
||||
```python
|
||||
custom_prompt = """You are a technical support expert for WIX.
|
||||
Answer questions based on the knowledge base documents provided.
|
||||
Be concise and helpful."""
|
||||
|
||||
answer, citations = await client.ask(
|
||||
"How do I create a blog?",
|
||||
system_prompt=custom_prompt
|
||||
)
|
||||
```
|
||||
|
||||
Filter to specific documents:
|
||||
|
||||
```python
|
||||
|
|
@ -443,11 +430,11 @@ answer, citations = await client.ask(
|
|||
)
|
||||
```
|
||||
|
||||
The QA agent searches your documents for relevant information and uses the configured LLM to generate an answer. The method returns a tuple of `(answer_text, list[Citation])`. Citations include page numbers, section headings, and document references.
|
||||
`client.ask` runs the [rag skill](skills/index.md) under the hood and returns `(answer_text, list[Citation])`. Citations include page numbers, section headings, and document references.
|
||||
|
||||
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 the QA agent and the multi‑agent research workflow.
|
||||
See also: [Agents](agents/index.md) for details on question answering and the multi‑agent research workflow.
|
||||
|
||||
## Analysis
|
||||
|
||||
|
|
@ -456,25 +443,20 @@ Answer complex analytical questions via code execution:
|
|||
```python
|
||||
# Aggregation across documents
|
||||
result = await client.analyze("Which quarter had the highest revenue?")
|
||||
print(result.answer) # The answer
|
||||
print(result.program) # The final consolidated program
|
||||
print(result.answer)
|
||||
for citation in result.citations:
|
||||
print(citation.uri, citation.title)
|
||||
|
||||
# Computation within a document set
|
||||
result = await client.analyze(
|
||||
"What is the average deal size mentioned in these contracts?",
|
||||
filter="uri LIKE '%contracts%'"
|
||||
)
|
||||
|
||||
# Multi-document comparison
|
||||
result = await client.analyze(
|
||||
"What changed between these two versions of the policy?",
|
||||
documents=["Policy v1.0", "Policy v2.0"]
|
||||
)
|
||||
```
|
||||
|
||||
The analysis agent writes and executes Python code in a sandboxed environment to solve problems that traditional RAG struggles with: aggregation, computation, and multi-document analysis.
|
||||
`client.analyze` runs the [analysis skill](skills/index.md), which writes and executes Python code in a sandboxed environment to solve problems that traditional RAG struggles with: aggregation, computation, and multi-document analysis.
|
||||
|
||||
See [Analysis Agent](agents/analysis.md) for details on capabilities and configuration.
|
||||
See [Analysis](agents/analysis.md) for details on capabilities and configuration.
|
||||
|
||||
## Building Custom Agents
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ For lower-level access, `haiku.rag.tools` provides individual `FunctionToolset`
|
|||
|
||||
## Low-Level Toolsets
|
||||
|
||||
For advanced use cases, individual toolset factories are available in `haiku.rag.tools`. These are used internally by the QA agent and can be composed into custom agents.
|
||||
For advanced use cases, individual toolset factories are available in `haiku.rag.tools`. These are the same primitives the rag and rag-analysis skills compose, and can be reused to build custom agents.
|
||||
|
||||
### RAGDeps Protocol
|
||||
|
||||
|
|
|
|||
|
|
@ -1,21 +0,0 @@
|
|||
from haiku.rag.agents.analysis.agent import create_analysis_agent
|
||||
from haiku.rag.agents.analysis.dependencies import AnalysisContext, AnalysisDeps
|
||||
from haiku.rag.agents.analysis.models import (
|
||||
AnalysisResult,
|
||||
CodeExecution,
|
||||
RawAnalysisResult,
|
||||
)
|
||||
from haiku.rag.agents.analysis.prompts import ANALYSIS_SYSTEM_PROMPT
|
||||
from haiku.rag.agents.analysis.sandbox import Sandbox, SandboxResult
|
||||
|
||||
__all__ = [
|
||||
"ANALYSIS_SYSTEM_PROMPT",
|
||||
"AnalysisContext",
|
||||
"AnalysisDeps",
|
||||
"RawAnalysisResult",
|
||||
"AnalysisResult",
|
||||
"CodeExecution",
|
||||
"Sandbox",
|
||||
"SandboxResult",
|
||||
"create_analysis_agent",
|
||||
]
|
||||
|
|
@ -1,56 +0,0 @@
|
|||
from pydantic_ai import Agent, RunContext
|
||||
|
||||
from haiku.rag.agents.analysis.dependencies import AnalysisDeps
|
||||
from haiku.rag.agents.analysis.models import CodeExecution, RawAnalysisResult
|
||||
from haiku.rag.agents.analysis.prompts import ANALYSIS_SYSTEM_PROMPT
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.utils import get_model
|
||||
|
||||
|
||||
def create_analysis_agent(config: AppConfig) -> Agent[AnalysisDeps, RawAnalysisResult]:
|
||||
"""Create an analysis agent with code execution capability.
|
||||
|
||||
The analysis agent can write and execute Python code in a sandboxed
|
||||
environment to solve problems that require computation, aggregation,
|
||||
or complex traversal across documents.
|
||||
|
||||
Args:
|
||||
config: Application configuration.
|
||||
|
||||
Returns:
|
||||
A pydantic-ai Agent configured for analysis execution.
|
||||
"""
|
||||
model = get_model(config.analysis.model or config.qa.model, config)
|
||||
|
||||
agent: Agent[AnalysisDeps, RawAnalysisResult] = Agent( # type: ignore[assignment] # ty: ignore[invalid-assignment]
|
||||
model,
|
||||
deps_type=AnalysisDeps,
|
||||
output_type=RawAnalysisResult,
|
||||
instructions=ANALYSIS_SYSTEM_PROMPT,
|
||||
tool_retries=3,
|
||||
output_retries=3,
|
||||
)
|
||||
|
||||
@agent.tool
|
||||
async def execute_code(ctx: RunContext[AnalysisDeps], code: str) -> CodeExecution:
|
||||
"""Execute Python code in a sandboxed interpreter.
|
||||
|
||||
The code has access to search() and list_documents() external
|
||||
functions, plus a virtual filesystem at /documents/ with document
|
||||
content and structure. Use print() to output results.
|
||||
|
||||
Args:
|
||||
code: Python code to execute.
|
||||
|
||||
Returns:
|
||||
Structured result with success status, stdout, and stderr.
|
||||
"""
|
||||
result = await ctx.deps.sandbox.execute(code)
|
||||
return CodeExecution(
|
||||
code=code,
|
||||
stdout=result.stdout,
|
||||
stderr=result.stderr,
|
||||
success=result.success,
|
||||
)
|
||||
|
||||
return agent
|
||||
|
|
@ -1,23 +0,0 @@
|
|||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from haiku.rag.store.models import Document
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.agents.analysis.sandbox import Sandbox
|
||||
|
||||
|
||||
@dataclass
|
||||
class AnalysisContext:
|
||||
"""Mutable context accumulating data during analysis execution."""
|
||||
|
||||
documents: list[Document] | None = None
|
||||
filter: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class AnalysisDeps:
|
||||
"""Dependencies for analysis agent."""
|
||||
|
||||
sandbox: "Sandbox"
|
||||
context: AnalysisContext = field(default_factory=AnalysisContext)
|
||||
|
|
@ -1,31 +0,0 @@
|
|||
from pydantic import BaseModel, Field
|
||||
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
|
||||
|
||||
class CodeExecution(BaseModel):
|
||||
"""Result of executing a code block in the analysis sandbox."""
|
||||
|
||||
code: str = Field(description="The Python code that was executed")
|
||||
stdout: str = Field(description="Standard output captured during execution")
|
||||
stderr: str = Field(description="Standard error captured during execution")
|
||||
success: bool = Field(description="Whether execution completed without error")
|
||||
|
||||
|
||||
class RawAnalysisResult(BaseModel):
|
||||
"""Raw result from the analysis agent (LLM output)."""
|
||||
|
||||
answer: str = Field(description="The answer to the user's question")
|
||||
program: str = Field(description="The final consolidated program")
|
||||
|
||||
|
||||
class AnalysisResult(BaseModel):
|
||||
"""Result from analysis execution with resolved citations.
|
||||
|
||||
The per-execution program(s) are tracked on ``AnalysisState.executions``
|
||||
(populated by the analysis skill's ``execute_code`` tool), not on this
|
||||
return value. Consumers that need the executed code should pull it from
|
||||
the skill state."""
|
||||
|
||||
answer: str
|
||||
citations: list[Citation] = Field(default_factory=list)
|
||||
|
|
@ -1,167 +0,0 @@
|
|||
ANALYSIS_SYSTEM_PROMPT = """You are an analysis agent that solves complex research questions by writing and executing Python code.
|
||||
|
||||
You MUST use the `execute_code` tool to run Python code. The functions and filesystem described below are ONLY available inside execute_code. Always execute code to answer questions; do not just describe what code would do.
|
||||
|
||||
## Available Functions
|
||||
|
||||
Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
|
||||
- results = await search("query") ✓ CORRECT
|
||||
- import search ✗ WRONG - will fail
|
||||
- results = search("query") ✗ WRONG - must use await
|
||||
|
||||
### await search(query, limit=10) -> list[dict]
|
||||
Search the knowledge base using hybrid search (vector + full-text).
|
||||
Results are automatically expanded with surrounding context (adjacent paragraphs, complete tables, section content).
|
||||
Returns list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels, picture_refs.
|
||||
`picture_refs` is the subset of `doc_item_refs` whose label is `picture` — use it to spot results that contain figures.
|
||||
|
||||
### await list_documents() -> list[dict]
|
||||
List all documents in the knowledge base.
|
||||
Returns list of dicts with keys: id, title, uri, created_at
|
||||
|
||||
## Document Filesystem
|
||||
|
||||
All documents in the knowledge base are available as files under `/documents/`. Use `from pathlib import Path` and standard file I/O to access them.
|
||||
|
||||
### Directory structure
|
||||
```
|
||||
/documents/
|
||||
{document_id}/
|
||||
metadata.json # {"id", "title", "uri", "created_at"}
|
||||
content.txt # Full document text
|
||||
items.jsonl # Structured document items (one JSON object per line)
|
||||
toc.json # Section tree (nested or flat depending on source)
|
||||
```
|
||||
|
||||
### metadata.json
|
||||
Small file with document metadata. Use to discover and identify documents.
|
||||
```python
|
||||
from pathlib import Path
|
||||
import json
|
||||
for doc_dir in Path('/documents').iterdir():
|
||||
meta = json.loads((doc_dir / 'metadata.json').read_text())
|
||||
print(meta['title'], meta['uri'])
|
||||
```
|
||||
|
||||
### content.txt
|
||||
Full text content of the document. Use for regex, keyword search, or full-text analysis.
|
||||
```python
|
||||
content = Path(f'/documents/{doc_id}/content.txt').read_text()
|
||||
```
|
||||
|
||||
### items.jsonl
|
||||
Structured document items as JSONL. Each line is a JSON object with:
|
||||
- `position`: sequential position in the document
|
||||
- `self_ref`: item reference (e.g. "#/texts/5", "#/tables/0")
|
||||
- `label`: item type — "section_header", "text", "table", "list_item", "caption", "formula", "picture", "code", "footnote", etc.
|
||||
- `text`: rendered content (tables are markdown with `|` columns)
|
||||
- `page_numbers`: list of page numbers where the item appears
|
||||
- `heading_level`: H-level (1–6) on `section_header` items, `0` otherwise. Often `1` for everything when the source is a PDF (docling can't infer heading hierarchy from PDFs) — see `toc.json` for the derived tree.
|
||||
- `tree_depth`: DOM nesting depth from docling's structure. Useful for HTML where it varies meaningfully (sidebars, captions, nested lists); near-uniform on PDFs.
|
||||
|
||||
Use items.jsonl to find tables, section headers, or specific structural elements:
|
||||
```python
|
||||
import json
|
||||
items_text = Path(f'/documents/{doc_id}/items.jsonl').read_text()
|
||||
for line in items_text.strip().split(chr(10)):
|
||||
item = json.loads(line)
|
||||
if item['label'] == 'table':
|
||||
print(f"Table on page {item['page_numbers']}: {item['text'][:100]}")
|
||||
```
|
||||
|
||||
### toc.json
|
||||
Per-document section tree derived from `heading_level`. Shape:
|
||||
```json
|
||||
{"doc_id": "...", "title": "...", "tree": [
|
||||
{"self_ref": "#/texts/0", "level": 1, "title": "Intro",
|
||||
"position": 0, "page_numbers": [1], "item_range": [0, 18],
|
||||
"children": [
|
||||
{"self_ref": "#/texts/8", "level": 2, "title": "Background",
|
||||
"position": 8, "page_numbers": [2], "item_range": [8, 13],
|
||||
"children": []}
|
||||
]}
|
||||
]}
|
||||
```
|
||||
- `item_range = [start, end_exclusive]` over the same `position` ints used in items.jsonl. Slice items.jsonl by this range to read a whole section.
|
||||
- PDF-derived docs typically produce a flat list of level-1 siblings (docling collapses heading levels). HTML/markdown produce a real nested tree.
|
||||
- `tree: []` when the doc has no section_headers at all.
|
||||
|
||||
```python
|
||||
import json
|
||||
toc = json.loads(Path(f'/documents/{doc_id}/toc.json').read_text())
|
||||
items = [json.loads(line) for line in Path(f'/documents/{doc_id}/items.jsonl').read_text().strip().split(chr(10))]
|
||||
|
||||
# Read the contents of one section
|
||||
def items_in(node):
|
||||
start, end = node['item_range']
|
||||
return [it for it in items if start <= it['position'] < end]
|
||||
|
||||
# From a search hit's doc_item_refs, find the deepest TOC node containing it
|
||||
def find_containing_section(tree, position):
|
||||
best = None
|
||||
def walk(nodes):
|
||||
nonlocal best
|
||||
for n in nodes:
|
||||
s, e = n['item_range']
|
||||
if s <= position < e:
|
||||
best = n
|
||||
walk(n['children'])
|
||||
walk(tree)
|
||||
return best
|
||||
```
|
||||
|
||||
## Cross-referencing search results with items
|
||||
|
||||
Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) that correspond to `self_ref` values in items.jsonl. Use this to navigate from a search hit to the surrounding document structure:
|
||||
```python
|
||||
results = await search("revenue", limit=5)
|
||||
r = results[0]
|
||||
doc_id = r['document_id']
|
||||
refs = set(r['doc_item_refs'])
|
||||
|
||||
import json
|
||||
items_text = Path(f'/documents/{doc_id}/items.jsonl').read_text()
|
||||
for line in items_text.strip().split(chr(10)):
|
||||
item = json.loads(line)
|
||||
if item['self_ref'] in refs:
|
||||
print(f"Matched: {item['label']} on page {item['page_numbers']}")
|
||||
```
|
||||
|
||||
## Pre-loaded Documents Variable
|
||||
|
||||
If documents were pre-loaded for this session, a `documents` variable is available:
|
||||
```python
|
||||
# documents is a list of dicts with keys: id, title, uri, content
|
||||
for doc in documents:
|
||||
print(doc['title'], len(doc['content']))
|
||||
```
|
||||
Check if it exists with: `try: documents ... except NameError: ...`
|
||||
|
||||
## Available Python Features
|
||||
|
||||
The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `map()`, `filter()`, `getattr()`, `sorted()`/`.sort(key=...)`, try/except, and the `json`, `re`, `math` modules. File I/O via `pathlib.Path` is supported for the `/documents/` filesystem.
|
||||
|
||||
Not supported: most imports (only `json`, `re`, `math`, `pathlib` are available), class definitions, generators/yield, match statements, decorators, `with` statements.
|
||||
|
||||
## Strategy Guide
|
||||
|
||||
1. **Search First**: Start with `search()` to find relevant content. Results include expanded context and `doc_item_refs` for cross-referencing.
|
||||
2. **Discover Documents**: Use `list_documents()` to see what's in the knowledge base.
|
||||
3. **Navigate Structure**: Use `items.jsonl` to find tables, section headers, or specific elements by label and page number (tables are pre-rendered as markdown). When a question is scoped to a section, open `toc.json`, find the matching node, and slice `items.jsonl` by its `item_range` instead of streaming `content.txt`. For PDFs where the tree is flat, the sibling list is still useful as a TOC.
|
||||
4. **Use content.txt for Full Text**: When you need the complete document text (e.g., for regex across the whole document).
|
||||
5. **Iterate**: Run code, examine results, refine your approach. Don't try to solve everything in one execution.
|
||||
6. **Cite picture chunks for figure-driven questions**: When a question is about a figure or diagram, find the picture chunk (search results with non-empty `picture_refs`) and cite its chunk_id. The driving model already sees figures from search hits; the citation makes the picture visible in the user's UI as well.
|
||||
|
||||
## Output Format
|
||||
|
||||
Your final response MUST be valid JSON matching this exact schema:
|
||||
```json
|
||||
{"answer": "Your answer here", "program": "Your final program here"}
|
||||
```
|
||||
|
||||
- `answer`: A clear answer to the user's question with key findings and references to specific documents/chunks.
|
||||
- `program`: A single, self-contained Python program that produces the answer. Consolidate your exploratory code executions into one clean script.
|
||||
|
||||
Do NOT return arbitrary JSON structures. Always use the exact format above.
|
||||
|
||||
You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first."""
|
||||
|
|
@ -1,35 +0,0 @@
|
|||
from haiku.rag.agents.qa.agent import QuestionAnswerAgent
|
||||
from haiku.rag.agents.qa.prompts import QA_SYSTEM_PROMPT
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import AppConfig, Config
|
||||
from haiku.rag.utils import build_prompt
|
||||
|
||||
|
||||
def get_qa_agent(
|
||||
client: HaikuRAG,
|
||||
config: AppConfig = Config,
|
||||
system_prompt: str | None = None,
|
||||
) -> QuestionAnswerAgent:
|
||||
"""Factory function to get a QA agent based on the configuration.
|
||||
|
||||
Args:
|
||||
client: HaikuRAG client instance.
|
||||
config: Configuration to use. Defaults to global Config.
|
||||
system_prompt: Optional custom system prompt (overrides config).
|
||||
|
||||
Returns:
|
||||
A configured QuestionAnswerAgent instance.
|
||||
"""
|
||||
# Determine the base prompt: explicit > config > default
|
||||
if system_prompt is None:
|
||||
system_prompt = config.prompts.qa or QA_SYSTEM_PROMPT
|
||||
|
||||
# Prepend system_context if configured
|
||||
system_prompt = build_prompt(system_prompt, config)
|
||||
|
||||
return QuestionAnswerAgent(
|
||||
client=client,
|
||||
model_config=config.qa.model,
|
||||
config=config,
|
||||
system_prompt=system_prompt,
|
||||
)
|
||||
|
|
@ -1,81 +0,0 @@
|
|||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai import Agent
|
||||
|
||||
from haiku.rag.agents.qa.prompts import QA_SYSTEM_PROMPT
|
||||
from haiku.rag.agents.research.models import (
|
||||
Citation,
|
||||
RawSearchAnswer,
|
||||
resolve_citations,
|
||||
)
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.config.models import AppConfig, ModelConfig
|
||||
from haiku.rag.store.models import SearchResult
|
||||
from haiku.rag.tools.search import create_search_toolset
|
||||
from haiku.rag.utils import get_model
|
||||
|
||||
|
||||
@dataclass
|
||||
class _QARunDeps:
|
||||
client: HaikuRAG
|
||||
|
||||
|
||||
class QuestionAnswerAgent:
|
||||
def __init__(
|
||||
self,
|
||||
client: HaikuRAG,
|
||||
model_config: ModelConfig,
|
||||
config: AppConfig | None = None,
|
||||
system_prompt: str | None = None,
|
||||
):
|
||||
self._client = client
|
||||
self._config = config or Config
|
||||
self._model_config = model_config
|
||||
self._system_prompt = system_prompt or QA_SYSTEM_PROMPT
|
||||
|
||||
async def answer(
|
||||
self, question: str, filter: str | None = None
|
||||
) -> tuple[str, list[Citation]]:
|
||||
"""Answer a question using the RAG system.
|
||||
|
||||
Args:
|
||||
question: The question to answer
|
||||
filter: SQL WHERE clause to filter documents
|
||||
|
||||
Returns:
|
||||
Tuple of (answer text, list of resolved citations)
|
||||
"""
|
||||
accumulated_results: list[SearchResult] = []
|
||||
max_searches = self._config.qa.max_searches
|
||||
search_toolset = create_search_toolset(
|
||||
self._config,
|
||||
base_filter=filter,
|
||||
tool_name="search",
|
||||
on_results=accumulated_results.extend,
|
||||
max_searches=max_searches,
|
||||
)
|
||||
|
||||
# Agent created per-call: toolset varies with filter, and Agent
|
||||
# construction is pure Python (no IO).
|
||||
model = get_model(self._model_config, self._config)
|
||||
try:
|
||||
system_prompt = self._system_prompt.format(max_searches=max_searches)
|
||||
except KeyError:
|
||||
system_prompt = self._system_prompt
|
||||
agent: Agent[_QARunDeps, RawSearchAnswer] = Agent( # ty: ignore[invalid-assignment]
|
||||
model=model,
|
||||
deps_type=_QARunDeps,
|
||||
output_type=RawSearchAnswer,
|
||||
instructions=system_prompt,
|
||||
toolsets=[search_toolset],
|
||||
tool_retries=3,
|
||||
output_retries=3,
|
||||
)
|
||||
|
||||
deps = _QARunDeps(client=self._client)
|
||||
result = await agent.run(question, deps=deps)
|
||||
output = result.output
|
||||
|
||||
citations = resolve_citations(output.cited_chunks, accumulated_results)
|
||||
return output.answer, citations
|
||||
|
|
@ -1,41 +0,0 @@
|
|||
QA_SYSTEM_PROMPT = """You are a knowledgeable assistant that answers questions using a document knowledge base.
|
||||
|
||||
Process:
|
||||
1. Call search with relevant keywords from the question
|
||||
2. Review the results ordered by relevance
|
||||
3. If needed, perform follow-up searches with different keywords (max {max_searches} total)
|
||||
4. Provide a concise answer based strictly on the retrieved content
|
||||
|
||||
The search tool returns results like:
|
||||
[chunk_abc123] [rank 1 of 5]
|
||||
Source: "Document Title" > Section > Subsection
|
||||
Type: paragraph
|
||||
Content:
|
||||
The actual text content here...
|
||||
|
||||
[chunk_def456] [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
|
||||
|
||||
When a result is of `Type: picture` (a figure or diagram), the search tool may also attach the picture itself as image content alongside the text — use it for visual reasoning when answering. Reference it by its chunk_id like any other source.
|
||||
|
||||
IMPORTANT: You MUST include in cited_chunks the COMPLETE IDs of every chunk you reference. Copy the full ID string without brackets — e.g. "5ae52166-5329-42e9-b6a5-756fc0cb7200" not "[5ae52166]" or "5ae52166". Never truncate IDs. Never leave cited_chunks empty if you found relevant content.
|
||||
|
||||
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
|
||||
- Be concise and direct - avoid elaboration unless asked
|
||||
- Results are ordered by relevance, with rank 1 being most relevant
|
||||
- If the search tool tells you the search limit is reached, stop searching immediately and answer with what you have
|
||||
- If the retrieved documents do not directly address the question, say: "I cannot find enough information in the knowledge base to answer this question." Do not guess or infer an answer from tangentially related content.
|
||||
"""
|
||||
|
|
@ -30,11 +30,11 @@ if TYPE_CHECKING:
|
|||
from docling_core.types.doc.document import DoclingDocument
|
||||
from PIL import Image as PILImage
|
||||
|
||||
from haiku.rag.agents.analysis.models import AnalysisResult
|
||||
from haiku.rag.agents.research.models import (
|
||||
Citation,
|
||||
ResearchReport,
|
||||
)
|
||||
from haiku.rag.sandbox import AnalysisResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.agents.analysis.models import AnalysisResult
|
||||
from haiku.rag.agents.research.models import Citation, ResearchReport
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.sandbox import AnalysisResult
|
||||
|
||||
|
||||
async def ask(
|
||||
|
|
@ -89,7 +89,7 @@ async def analyze(
|
|||
Returns:
|
||||
AnalysisResult with the answer and resolved citations.
|
||||
"""
|
||||
from haiku.rag.agents.analysis.models import AnalysisResult
|
||||
from haiku.rag.sandbox import AnalysisResult
|
||||
from haiku.rag.skills.analysis import AnalysisState, create_skill
|
||||
from haiku.rag.utils import get_model
|
||||
from haiku.skills import run_skill
|
||||
|
|
|
|||
|
|
@ -241,7 +241,6 @@ class ProvidersConfig(BaseModel):
|
|||
|
||||
class PromptsConfig(BaseModel):
|
||||
domain_preamble: str = ""
|
||||
qa: str | None = None
|
||||
synthesis: str | None = None
|
||||
picture_description: str = (
|
||||
"Describe this image for a blind user. "
|
||||
|
|
|
|||
10
haiku_rag_slim/haiku/rag/sandbox/__init__.py
Normal file
10
haiku_rag_slim/haiku/rag/sandbox/__init__.py
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
from haiku.rag.sandbox.dependencies import AnalysisContext
|
||||
from haiku.rag.sandbox.models import AnalysisResult
|
||||
from haiku.rag.sandbox.sandbox import Sandbox, SandboxResult
|
||||
|
||||
__all__ = [
|
||||
"AnalysisContext",
|
||||
"AnalysisResult",
|
||||
"Sandbox",
|
||||
"SandboxResult",
|
||||
]
|
||||
8
haiku_rag_slim/haiku/rag/sandbox/dependencies.py
Normal file
8
haiku_rag_slim/haiku/rag/sandbox/dependencies.py
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class AnalysisContext:
|
||||
"""Mutable context accumulating data during analysis execution."""
|
||||
|
||||
filter: str | None = None
|
||||
14
haiku_rag_slim/haiku/rag/sandbox/models.py
Normal file
14
haiku_rag_slim/haiku/rag/sandbox/models.py
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
from pydantic import BaseModel, Field
|
||||
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
|
||||
|
||||
class AnalysisResult(BaseModel):
|
||||
"""Result from analysis execution with resolved citations.
|
||||
|
||||
Executed code is tracked on ``AnalysisState.executions`` (populated by the
|
||||
analysis skill's ``execute_code`` tool). Consumers that need the program
|
||||
should pull it from the skill state."""
|
||||
|
||||
answer: str
|
||||
citations: list[Citation] = Field(default_factory=list)
|
||||
|
|
@ -10,8 +10,8 @@ from typing import TYPE_CHECKING, Any, Literal
|
|||
import pydantic_monty
|
||||
from pydantic_monty import CallbackFile, MemoryFile, MontyRepl, OSAccess
|
||||
|
||||
from haiku.rag.agents.analysis.dependencies import AnalysisContext
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.sandbox.dependencies import AnalysisContext
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX, DocumentItem
|
||||
|
||||
|
|
@ -351,23 +351,6 @@ class Sandbox:
|
|||
"max_duration_secs": self._config.analysis.code_timeout,
|
||||
},
|
||||
)
|
||||
if self._context.documents:
|
||||
await self._repl.feed_run_async(
|
||||
"pass",
|
||||
inputs={
|
||||
"documents": [
|
||||
{
|
||||
"id": d.id,
|
||||
"title": d.title,
|
||||
"uri": d.uri,
|
||||
"content": d.content,
|
||||
}
|
||||
for d in self._context.documents
|
||||
]
|
||||
},
|
||||
external_functions=self._build_external_functions(),
|
||||
os=self._vfs,
|
||||
)
|
||||
assert self._repl is not None and self._vfs is not None
|
||||
return self._repl, self._vfs
|
||||
|
||||
|
|
@ -8,8 +8,8 @@ from haiku.rag.config.models import AppConfig
|
|||
from haiku.skills.state import SkillRunDeps
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.agents.analysis.sandbox import Sandbox
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.sandbox import Sandbox
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
@ -60,9 +60,8 @@ def make_rag_lifespan(db_path: Path, config: AppConfig):
|
|||
def make_analysis_lifespan(db_path: Path, config: AppConfig):
|
||||
@asynccontextmanager
|
||||
async def lifespan(deps: AnalysisRunDeps) -> AsyncIterator[None]:
|
||||
from haiku.rag.agents.analysis.dependencies import AnalysisContext
|
||||
from haiku.rag.agents.analysis.sandbox import Sandbox
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.sandbox import AnalysisContext, Sandbox
|
||||
|
||||
doc_filter = getattr(deps.state, "document_filter", None)
|
||||
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
||||
|
|
|
|||
|
|
@ -72,7 +72,7 @@ nav:
|
|||
- Tuning: tuning.md
|
||||
- Agents:
|
||||
- agents/index.md
|
||||
- Analysis Agent: agents/analysis.md
|
||||
- Analysis: agents/analysis.md
|
||||
- Skills:
|
||||
- skills/index.md
|
||||
- RAG: skills/rag.md
|
||||
|
|
|
|||
|
|
@ -1,32 +0,0 @@
|
|||
from haiku.rag.agents.analysis.models import AnalysisResult, CodeExecution
|
||||
|
||||
|
||||
class TestCodeExecution:
|
||||
def test_create_successful_execution(self):
|
||||
execution = CodeExecution(
|
||||
code="print('hello')",
|
||||
stdout="hello\n",
|
||||
stderr="",
|
||||
success=True,
|
||||
)
|
||||
assert execution.code == "print('hello')"
|
||||
assert execution.stdout == "hello\n"
|
||||
assert execution.stderr == ""
|
||||
assert execution.success is True
|
||||
|
||||
def test_create_failed_execution(self):
|
||||
execution = CodeExecution(
|
||||
code="1/0",
|
||||
stdout="",
|
||||
stderr="ZeroDivisionError: division by zero",
|
||||
success=False,
|
||||
)
|
||||
assert execution.success is False
|
||||
assert "ZeroDivisionError" in execution.stderr
|
||||
|
||||
|
||||
class TestAnalysisResult:
|
||||
def test_create_result(self):
|
||||
result = AnalysisResult(answer="The answer is 42")
|
||||
assert result.answer == "The answer is 42"
|
||||
assert result.citations == []
|
||||
|
|
@ -1,69 +0,0 @@
|
|||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from datasets import Dataset
|
||||
from evaluations.evaluators import LLMJudge
|
||||
|
||||
from haiku.rag.agents.qa.agent import QuestionAnswerAgent
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.config.models import ModelConfig
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def vcr_cassette_dir():
|
||||
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_qa")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_qa_agent_factory(temp_db_path):
|
||||
"""Test get_qa_agent factory function creates a properly configured agent."""
|
||||
from haiku.rag.agents.qa import get_qa_agent
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
agent = get_qa_agent(client, Config)
|
||||
|
||||
assert agent is not None
|
||||
assert isinstance(agent, QuestionAnswerAgent)
|
||||
# Verify internal client is set correctly
|
||||
assert agent._client is client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_qa_agent_with_custom_prompt(temp_db_path):
|
||||
"""Test get_qa_agent factory with custom system prompt."""
|
||||
from haiku.rag.agents.qa import get_qa_agent
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
custom_prompt = "You are a custom QA assistant."
|
||||
agent = get_qa_agent(client, Config, system_prompt=custom_prompt)
|
||||
|
||||
assert agent is not None
|
||||
assert isinstance(agent, QuestionAnswerAgent)
|
||||
assert agent._system_prompt == custom_prompt
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_qa_ollama(allow_model_requests, qa_corpus: Dataset, temp_db_path):
|
||||
"""Test Ollama QA with LLM judge (VCR recorded)."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
qa = QuestionAnswerAgent(
|
||||
client,
|
||||
ModelConfig(provider="ollama", name="gpt-oss", enable_thinking=True),
|
||||
)
|
||||
llm_judge = LLMJudge()
|
||||
|
||||
doc = qa_corpus[1]
|
||||
await client.create_document(
|
||||
content=doc["document_extracted"], uri=doc["document_id"]
|
||||
)
|
||||
|
||||
question = doc["question"]
|
||||
expected_answer = doc["answer"]
|
||||
|
||||
answer, _ = await qa.answer(question)
|
||||
is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer)
|
||||
|
||||
assert is_equivalent, (
|
||||
f"Generated answer not equivalent to expected answer.\nQuestion: {question}\nGenerated: {answer}\nExpected: {expected_answer}"
|
||||
)
|
||||
File diff suppressed because one or more lines are too long
|
|
@ -1,16 +1,8 @@
|
|||
import pytest
|
||||
|
||||
from haiku.rag.agents.analysis.dependencies import AnalysisContext
|
||||
from haiku.rag.agents.analysis.sandbox import Sandbox
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config.models import AppConfig
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def empty_client(temp_db_path):
|
||||
"""Create an empty HaikuRAG client without documents."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
yield client
|
||||
from haiku.rag.sandbox import AnalysisContext, Sandbox
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
8
tests/sandbox/test_models.py
Normal file
8
tests/sandbox/test_models.py
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
from haiku.rag.sandbox import AnalysisResult
|
||||
|
||||
|
||||
class TestAnalysisResult:
|
||||
def test_create_result(self):
|
||||
result = AnalysisResult(answer="The answer is 42")
|
||||
assert result.answer == "The answer is 42"
|
||||
assert result.citations == []
|
||||
|
|
@ -2,16 +2,14 @@ from pathlib import Path
|
|||
|
||||
import pytest
|
||||
|
||||
from haiku.rag.agents.analysis.dependencies import AnalysisContext
|
||||
from haiku.rag.agents.analysis.sandbox import Sandbox, SandboxResult
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.store.models import Document
|
||||
from haiku.rag.sandbox import AnalysisContext, Sandbox, SandboxResult
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def vcr_cassette_dir():
|
||||
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_sandbox")
|
||||
return str(Path(__file__).parent.parent / "cassettes" / "test_sandbox")
|
||||
|
||||
|
||||
class TestSandboxBasics:
|
||||
|
|
@ -414,35 +412,3 @@ class TestSandboxVFS:
|
|||
assert "1" in result.stdout
|
||||
assert "Public Doc" in result.stdout
|
||||
assert "Private Doc" not in result.stdout
|
||||
|
||||
|
||||
class TestSandboxPreloadedDocuments:
|
||||
"""Test pre-loaded documents context variable."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_documents_variable_not_available_without_preload(self, sandbox):
|
||||
"""documents variable is not available when context.documents is None."""
|
||||
result = await sandbox.execute("print(documents)")
|
||||
assert not result.success
|
||||
assert "NameError" in result.stderr
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_documents_variable_available_with_preload(self, temp_db_path):
|
||||
"""documents variable is available when context.documents is set."""
|
||||
async with HaikuRAG(temp_db_path, create=True):
|
||||
config = AppConfig()
|
||||
docs = [
|
||||
Document(id="1", content="Content A", title="Doc A", uri="a://1"),
|
||||
Document(id="2", content="Content B", title="Doc B", uri="b://2"),
|
||||
]
|
||||
context = AnalysisContext(documents=docs)
|
||||
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
|
||||
result = await sb.execute(
|
||||
"print(len(documents))\n"
|
||||
"print(documents[0]['title'])\n"
|
||||
"print(documents[1]['title'])"
|
||||
)
|
||||
assert result.success
|
||||
assert "2" in result.stdout
|
||||
assert "Doc A" in result.stdout
|
||||
assert "Doc B" in result.stdout
|
||||
|
|
@ -4,10 +4,9 @@ and the set of external functions exposed to the interpreter.
|
|||
|
||||
import pytest
|
||||
|
||||
from haiku.rag.agents.analysis.dependencies import AnalysisContext
|
||||
from haiku.rag.agents.analysis.sandbox import Sandbox
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.sandbox import AnalysisContext, Sandbox
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
|
||||
|
||||
|
|
@ -12,10 +12,9 @@ from pathlib import PurePosixPath
|
|||
|
||||
import pytest
|
||||
|
||||
from haiku.rag.agents.analysis.dependencies import AnalysisContext
|
||||
from haiku.rag.agents.analysis.sandbox import Sandbox
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.sandbox import AnalysisContext, Sandbox
|
||||
from haiku.rag.store.models.document import Document
|
||||
from haiku.rag.store.models.document_item import DocumentItem
|
||||
|
||||
|
|
@ -85,8 +85,7 @@ async def rag_client(rag_db):
|
|||
@pytest.fixture
|
||||
def sandbox_factory(rag_db, test_app_config):
|
||||
"""Build Sandbox instances bound to the sample db, optionally with a doc filter."""
|
||||
from haiku.rag.agents.analysis.dependencies import AnalysisContext
|
||||
from haiku.rag.agents.analysis.sandbox import Sandbox
|
||||
from haiku.rag.sandbox import AnalysisContext, Sandbox
|
||||
|
||||
def _make(filter: str | None = None) -> Sandbox:
|
||||
return Sandbox(
|
||||
|
|
|
|||
|
|
@ -285,7 +285,7 @@ class TestExecuteCodeTool:
|
|||
|
||||
class TestAnalysisLifespan:
|
||||
async def test_opens_client_and_sandbox_per_invocation(self, rag_db):
|
||||
from haiku.rag.agents.analysis.sandbox import Sandbox
|
||||
from haiku.rag.sandbox import Sandbox
|
||||
from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan
|
||||
|
||||
config = AppConfig()
|
||||
|
|
|
|||
|
|
@ -1,61 +1,22 @@
|
|||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from pydantic_ai import Agent
|
||||
|
||||
from haiku.rag.agents.analysis.agent import create_analysis_agent
|
||||
from haiku.rag.agents.analysis.dependencies import AnalysisDeps
|
||||
from haiku.rag.agents.analysis.models import CodeExecution, RawAnalysisResult
|
||||
from haiku.rag.config import AppConfig, Config
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import AppConfig
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def vcr_cassette_dir():
|
||||
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_analysis")
|
||||
|
||||
|
||||
class TestCreateAnalysisAgent:
|
||||
def test_creates_agent(self):
|
||||
agent = create_analysis_agent(Config)
|
||||
assert isinstance(agent, Agent)
|
||||
assert agent.deps_type is AnalysisDeps
|
||||
assert agent.output_type is RawAnalysisResult
|
||||
|
||||
def test_agent_has_execute_code_tool(self):
|
||||
agent = create_analysis_agent(Config)
|
||||
tool_names = list(agent._function_toolset.tools.keys())
|
||||
assert "execute_code" in tool_names
|
||||
|
||||
|
||||
class TestCodeExecutionModel:
|
||||
def test_code_execution_has_correct_fields(self):
|
||||
"""Test that CodeExecution has all expected fields."""
|
||||
execution = CodeExecution(
|
||||
code="print('hello')",
|
||||
stdout="hello\n",
|
||||
stderr="",
|
||||
success=True,
|
||||
)
|
||||
assert execution.code == "print('hello')"
|
||||
assert execution.stdout == "hello\n"
|
||||
assert execution.stderr == ""
|
||||
assert execution.success is True
|
||||
return str(Path(__file__).parent / "cassettes" / "test_client_analyze")
|
||||
|
||||
|
||||
class TestClientAnalysisIntegration:
|
||||
"""Integration tests for client.analyze() method."""
|
||||
"""Integration tests for client.analyze() through the rag-analysis skill."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.vcr()
|
||||
async def test_analyze_count_documents(self, allow_model_requests, temp_db_path):
|
||||
"""Test analysis agent can count documents.
|
||||
|
||||
Agent program:
|
||||
docs = list_documents(limit=1000)
|
||||
print(len(docs))
|
||||
"""
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
config = AppConfig()
|
||||
|
||||
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
||||
|
|
@ -70,26 +31,6 @@ class TestClientAnalysisIntegration:
|
|||
@pytest.mark.asyncio
|
||||
@pytest.mark.vcr()
|
||||
async def test_analyze_aggregation(self, allow_model_requests, temp_db_path):
|
||||
"""Test analysis agent can perform aggregation across documents.
|
||||
|
||||
Agent program:
|
||||
import re
|
||||
revs = {}
|
||||
for d in ['Q1 Report', 'Q2 Report', 'Q3 Report']:
|
||||
content = get_document(d)
|
||||
if content:
|
||||
vals = re.findall(r'\\$([\\d,]+)', content)
|
||||
if vals:
|
||||
rev = sum(int(v.replace(',', '')) for v in vals)
|
||||
else:
|
||||
rev = None
|
||||
else:
|
||||
rev = None
|
||||
revs[d] = rev
|
||||
print(revs)
|
||||
"""
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
config = AppConfig()
|
||||
|
||||
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
||||
|
|
@ -112,17 +53,6 @@ class TestClientAnalysisIntegration:
|
|||
@pytest.mark.asyncio
|
||||
@pytest.mark.vcr()
|
||||
async def test_analyze_with_filter(self, allow_model_requests, temp_db_path):
|
||||
"""Test analysis agent respects filter parameter.
|
||||
|
||||
Agent program:
|
||||
docs = list_documents(limit=1000)
|
||||
print(len(docs))
|
||||
print(docs[:5])
|
||||
|
||||
The filter is applied via context, so list_documents() only sees "Cats".
|
||||
"""
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
config = AppConfig()
|
||||
|
||||
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
||||
|
|
@ -142,9 +72,6 @@ class TestClientAnalysisIntegration:
|
|||
async def test_analyze_search_and_identify_source(
|
||||
self, allow_model_requests, temp_db_path
|
||||
):
|
||||
"""Test analysis agent can search and identify source documents."""
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
config = AppConfig()
|
||||
|
||||
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
||||
|
|
@ -163,20 +90,6 @@ class TestClientAnalysisIntegration:
|
|||
@pytest.mark.asyncio
|
||||
@pytest.mark.vcr()
|
||||
async def test_analyze_search_and_extract(self, allow_model_requests, temp_db_path):
|
||||
"""Test analysis agent can use search() to find content and extract information.
|
||||
|
||||
Agent program:
|
||||
results = search("document element types", limit=20)
|
||||
print(len(results))
|
||||
for r in results[:5]:
|
||||
print(r['document_title'], r['chunk_id'], r['score'])
|
||||
print(r['content'][:200])
|
||||
|
||||
results = search("DocBank element types", limit=10)
|
||||
...
|
||||
"""
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
pdf_path = Path("tests/data/doclaynet.pdf")
|
||||
config = AppConfig()
|
||||
config.processing.conversion_options.do_ocr = False
|
||||
|
|
@ -190,9 +103,7 @@ class TestClientAnalysisIntegration:
|
|||
"List them all."
|
||||
)
|
||||
|
||||
# The doclaynet.pdf defines exactly 11 class labels for document elements
|
||||
# Normalize Unicode hyphens (U+2011 non-breaking hyphen) to regular hyphens
|
||||
answer_lower = result.answer.lower().replace("\u2011", "-")
|
||||
answer_lower = result.answer.lower().replace("‑", "-")
|
||||
expected_labels = [
|
||||
"caption",
|
||||
"footnote",
|
||||
|
|
@ -206,8 +117,6 @@ class TestClientAnalysisIntegration:
|
|||
"text",
|
||||
"title",
|
||||
]
|
||||
# Check that the agent found at least 6 of the 11 labels
|
||||
# (LLM summaries may not always include all labels)
|
||||
found_labels = [
|
||||
label
|
||||
for label in expected_labels
|
||||
Loading…
Reference in a new issue