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

Update docs & examples

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-20 11:08:50 +02:00
parent 09e6324add
commit bf724ad1ef
No known key found for this signature in database
7 changed files with 141 additions and 565 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

123
docs/agents/index.md
View file

@ -1,12 +1,13 @@
# Agents
Four agentic flows are provided by haiku.rag:
Three agentic flows are provided by haiku.rag:
- **Simple QA Agent** — a focused question answering agent
- **Chat Agent** — multi-turn conversational RAG with session memory
- **Research Graph** — a multi-step research workflow with question decomposition
- **RLM Agent** — complex analytical tasks via sandboxed Python code execution (see [RLM Agent](rlm.md))
For multi-turn conversational RAG, haiku.rag provides a [RAG skill](../tools.md#rag-skill) built on [haiku.skills](https://github.com/ggozad/haiku.skills). The skill bundles search, Q&A, analysis, and research tools with session state management.
See [QA and Research Configuration](../configuration/qa-research.md) for configuring model, iterations, concurrency, and other settings.
## Simple QA Agent
@ -26,9 +27,6 @@ haiku-rag ask "What is climate change?"
# With citations
haiku-rag ask "What is climate change?" --cite
# Deep mode (uses research graph with optimized settings)
haiku-rag ask "What are the main features of haiku.rag?" --deep
```
**Python usage:**
@ -48,121 +46,6 @@ async with HaikuRAG(path_to_db) as client:
print(answer)
```
## Chat Agent
The chat agent enables multi-turn conversational RAG. It is built from composable [toolsets](../tools.md) and maintains session state to improve follow-up answers.
Key features:
- **Composable toolsets**: Built from reusable `FunctionToolset` factories — see [Toolsets](../tools.md)
- **Semantic prior answer recall**: Similar prior Q/A pairs are retrieved and passed to the research planner, which can skip searching when they suffice
- **Background summarization**: After each `ask` call, the QA history is summarized into a compact session context for the next request
- **Document filtering**: Session-level or per-query document filtering
### CLI Usage
```bash
haiku-rag chat
haiku-rag chat --db /path/to/database.lancedb
```
See [Applications](../apps.md#chat-tui) for the full TUI interface guide.
### Python Usage
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.agents.chat import create_chat_agent, prepare_chat_context, ChatDeps
from haiku.rag.tools import ToolContext
agent = create_chat_agent(config)
async with HaikuRAG(path_to_db) as client:
context = ToolContext()
prepare_chat_context(context)
deps = ChatDeps(config=config, client=client, tool_context=context)
# First question
result = await agent.run("What is haiku.rag?", deps=deps)
print(result.output)
# Follow-up (uses session context)
result = await agent.run("How does it handle PDFs?", deps=deps)
print(result.output)
```
### Feature Selection
By default, `create_chat_agent` enables search, documents, and QA toolsets. You can customize which capabilities the agent has via the `features` parameter:
```python
from haiku.rag.agents.chat import (
create_chat_agent,
FEATURE_SEARCH,
FEATURE_DOCUMENTS,
FEATURE_QA,
FEATURE_ANALYSIS,
)
# Search-only agent
agent = create_chat_agent(config, features=[FEATURE_SEARCH])
# All features including code analysis
agent = create_chat_agent(
config,
features=[FEATURE_SEARCH, FEATURE_DOCUMENTS, FEATURE_QA, FEATURE_ANALYSIS],
)
```
Available features:
| Feature | Constant | Tools added |
|---------|----------|-------------|
| Search | `FEATURE_SEARCH` | `search` |
| Documents | `FEATURE_DOCUMENTS` | `list_documents`, `get_document`, `summarize_document` |
| QA | `FEATURE_QA` | `ask` |
| Analysis | `FEATURE_ANALYSIS` | `analyze` |
The system prompt is automatically composed to match the selected features. See [Toolsets](../tools.md) for details on each toolset's parameters and behavior.
### Session State
Session state is managed through `ToolContext` — a namespace-based state container shared across all toolsets. The chat agent uses two namespaces:
**`SessionState`** (session management):
- `document_filter` — List of document titles/URIs to restrict searches
- `citation_registry` — Stable mapping of chunk IDs to citation indices
- `citations` — Citations from the current query
**`QASessionState`** (QA history and context):
- `qa_history` — List of previous Q/A pairs with embeddings
- `session_context` — Automatically maintained session context summary
For multi-session applications (e.g., web backends), use `ToolContextCache` to cache `ToolContext` instances by external session/thread ID:
```python
from haiku.rag.tools import ToolContext, ToolContextCache
cache = ToolContextCache() # TTL-based, defaults to 1 hour
context, _is_new = cache.get_or_create(thread_id)
```
**Citation Registry**: Citation indices persist across tool calls within a session. The same `chunk_id` always returns the same citation index (first-occurrence-wins). This ensures consistent citation numbering in multi-turn conversations — `[1]` always refers to the same source.
### Conversational Memory
The chat agent maintains two layers of conversational memory:
**1. Semantic prior answer recall**
When the `ask` tool receives a question, it embeds the question and compares it against prior Q/A embeddings. Sufficiently similar prior answers are passed to the research planner, which can skip searching entirely if they already cover the question.
**2. Background session summarization**
After each `ask` call, a background task summarizes the full QA history into a compact session context. This summary is injected into the research planner on the next request, allowing it to resolve ambiguous references ("Tell me more about the authentication part") without having seen the full conversation.
## 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.

10
docs/apps.md
View file

@ -37,7 +37,7 @@ Press `Ctrl+P` to open the command palette:
| Command | Description |
|---------|-------------|
| Memory | Edit initial context (before first message) or view session context (after) |
| View state | View the current session state |
| Filter documents | Select documents to restrict searches |
| Show database info | View document/chunk counts and storage info |
| Visual grounding | View chunk source location in document |
@ -46,12 +46,10 @@ Press `Ctrl+P` to open the command palette:
### Session Management
- Conversation history is maintained in memory for the session
- Previous Q/A pairs are used as context for follow-up questions
- Previous Q/A pairs are used as context for follow-up questions via the `get_session_context` tool
- Citations are tracked per response and can be inspected
- Document filter restricts all searches to selected documents
- Initial context can be set via CLI (`--initial-context`) or command palette
- Initial context is editable until the first message is sent, then becomes read-only
- Clearing chat resets session state, restores CLI-provided context, and unlocks editing
- Clearing chat resets session state
## Web Application
@ -63,7 +61,7 @@ Browser-based conversational RAG with a CopilotKit frontend.
- Expandable citations with source documents, pages, and headings
- Visual grounding to view chunk source locations in documents
- Document filter to restrict searches to selected documents
- Memory panel: set initial context before first message, view session context after
- Session state view for inspecting accumulated Q&A history, citations, and documents
### Quick Start

20
docs/python.md
View file

@ -429,25 +429,23 @@ See [RLM Agent](agents/rlm.md) for details on capabilities and configuration.
## Building Custom Agents
haiku.rag provides composable toolset factories that can be mixed into any pydantic-ai agent. This lets you build custom agents with exactly the capabilities you need — search, document management, Q&A, or code analysis — sharing state across tools via `ToolContext`.
haiku.rag provides a RAG skill built on [haiku.skills](https://github.com/ggozad/haiku.skills) that bundles all capabilities into a composable agent:
```python
from pydantic_ai import Agent
from haiku.rag.tools import AgentDeps, build_toolkit
from haiku.rag.skills.rag import create_skill
from haiku.skills.agent import SkillToolset
toolkit = build_toolkit(config, features=["search", "qa"])
skill = create_skill(db_path=db_path, config=config)
toolset = SkillToolset(skills=[skill])
agent = Agent(
"openai:gpt-4o",
deps_type=AgentDeps,
instructions=f"You are a helpful assistant.\n{toolkit.prompt}",
toolsets=toolkit.toolsets,
instructions=toolset.system_prompt,
toolsets=[toolset],
)
async with HaikuRAG("path/to/db.lancedb") as client:
context = toolkit.create_context()
deps = AgentDeps(client=client, tool_context=context)
result = await agent.run("What are the main findings?", deps=deps)
result = await agent.run("What are the main findings?")
```
See [Toolsets](tools.md) for the full API reference and composition guide, and the [`examples/`](https://github.com/ggozad/haiku.rag/tree/main/examples) directory for runnable scripts.
See [Tools & Skills](tools.md) for the full API reference.

412
docs/tools.md
View file

@ -1,77 +1,107 @@
# Toolsets
# Tools & Skills
haiku.rag provides composable `FunctionToolset` factories in `haiku.rag.tools`. Each factory creates a pydantic-ai `FunctionToolset` that can be mixed into any agent. A shared `ToolContext` lets toolsets accumulate state (search results, citations, QA history) across invocations.
haiku.rag exposes its RAG capabilities through a [haiku.skills](https://github.com/ggozad/haiku.skills) skill. The skill provides tools for search, Q&A, analysis, and research that can be composed into any pydantic-ai agent via `SkillToolset`.
## ToolContext
For lower-level access, `haiku.rag.tools` provides individual `FunctionToolset` factories used internally by agents.
`ToolContext` is a namespace-based state container. Toolsets register Pydantic models under string namespaces, and any toolset sharing the same context can read or write the same state.
## RAG Skill
The RAG skill is the primary way to use haiku.rag tools. It bundles all capabilities into a single skill with managed state.
```python
from haiku.rag.tools import ToolContext
from haiku.rag.skills.rag import create_skill
from haiku.skills.agent import SkillToolset
from pydantic_ai import Agent
context = ToolContext()
skill = create_skill(db_path=db_path, config=config)
toolset = SkillToolset(skills=[skill])
agent = Agent(
"openai:gpt-4o",
instructions=toolset.system_prompt,
toolsets=[toolset],
)
result = await agent.run("What documents do we have?")
```
### Registering and accessing state
### `create_skill(db_path?, config?)`
Creates a RAG skill instance.
| Parameter | Default | Description |
|-----------|---------|-------------|
| `db_path` | `None` | Path to LanceDB database. Falls back to `HAIKU_RAG_DB` env var, then config default. |
| `config` | `None` | `AppConfig` instance. If None, uses `get_config()`. |
### Tools
| Tool | Purpose |
|------|---------|
| `search(query, limit?)` | Hybrid search (vector + full-text) with context expansion |
| `list_documents(limit?, offset?, filter?)` | Paginated document listing |
| `get_document(query)` | Retrieve a document by ID, title, or URI |
| `ask(question)` | Q&A with citations via the QA agent |
| `analyze(question, document?, filter?)` | Computational analysis via code execution (requires Docker) |
| `research(question)` | Deep multi-agent research producing comprehensive reports |
| `get_session_context(query)` | Retrieve relevant prior Q&A from the session |
### State
The skill manages a `RAGState` under the `"rag"` namespace:
```python
from pydantic import BaseModel
class MyState(BaseModel):
count: int = 0
context.register("my_namespace", MyState())
# Get state (returns None if not registered)
state = context.get("my_namespace")
# Get with type checking (returns None if wrong type)
state = context.get("my_namespace", MyState)
# Get or create (creates default if not registered)
state = context.get_or_create("my_namespace", MyState)
class RAGState(BaseModel):
citations: list[Any] = []
qa_history: list[QAHistoryEntry] = []
document_filter: str | None = None
searches: dict[str, list[SearchResult]] = {}
documents: list[DocumentInfo] = []
reports: list[ResearchEntry] = []
```
### Serialization
The entire context can be serialized and restored:
State is automatically synced via the AG-UI protocol when using `AGUIAdapter`. Access it programmatically:
```python
# Serialize all namespaces (keyed by namespace)
data = context.dump_namespaces()
# {"my_namespace": {"count": 0}}
# Restore a namespace from serialized data
context.load_namespace("my_namespace", MyState, data["my_namespace"])
rag_state = toolset.get_namespace("rag")
if rag_state:
print(f"Citations: {len(rag_state.citations)}")
print(f"Q&A history: {len(rag_state.qa_history)}")
```
For AG-UI state management, use flat snapshots:
### AG-UI Streaming
For web applications, use pydantic-ai's `AGUIAdapter` to stream tool calls, text, and state deltas:
```python
# Flat snapshot of all namespaces (for AG-UI state)
snapshot = context.build_state_snapshot()
# {"document_filter": [], "citations": [], "citation_registry": {}, "qa_history": []}
from pydantic_ai.ag_ui import AGUIAdapter
# Restore from flat snapshot (updates registered namespaces in place)
context.restore_state_snapshot(snapshot)
adapter = AGUIAdapter(agent=agent, run_input=run_input)
event_stream = adapter.run_stream()
sse_event_stream = adapter.encode_stream(event_stream)
```
### Preparing context for toolsets
See the [Web Application](apps.md#web-application) for a complete implementation.
`prepare_context()` registers the required namespaces for a given set of features:
## 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.
### RAGDeps Protocol
All toolsets use the `RAGDeps` protocol for dependency injection:
```python
from haiku.rag.tools import ToolContext, prepare_context
from haiku.rag.tools import RAGDeps
context = ToolContext()
prepare_context(context, features=["search", "qa"], state_key="my_app")
class MyDeps:
def __init__(self, client: HaikuRAG):
self.client = client
```
This is idempotent and registers `SessionState` (for search, QA, and analysis features) and `QASessionState` (for QA). The chat agent's `prepare_chat_context()` is a thin wrapper that defaults to chat features and sets the AG-UI state key.
### Search Toolset
## Search Toolset
`create_search_toolset()` provides hybrid search (vector + full-text) with context expansion and citation tracking.
`create_search_toolset()` provides hybrid search with context expansion.
```python
from haiku.rag.tools import create_search_toolset
@ -79,24 +109,17 @@ from haiku.rag.tools import create_search_toolset
search = create_search_toolset(config)
```
**Parameters:**
| Parameter | Default | Description |
|-----------|---------|-------------|
| `config` | required | AppConfig |
| `config` | required | `AppConfig` |
| `expand_context` | `True` | Expand results with surrounding chunks |
| `base_filter` | `None` | SQL WHERE clause applied to all searches |
| `tool_name` | `"search"` | Name of the tool exposed to the agent |
| `on_results` | `None` | Callback `(list[SearchResult]) -> None` invoked with results |
**Tool: `search(query, limit?, filter?)`**
### Document Toolset
Searches the knowledge base and returns formatted results. When a `ToolContext` with `SessionState` is registered, citations get stable indices via `citation_registry`.
**State:** Search results accumulate in `SearchState.results` under the `haiku.rag.search` namespace.
## Document Toolset
`create_document_toolset()` provides document browsing, retrieval, and summarization.
`create_document_toolset()` provides document browsing and retrieval.
```python
from haiku.rag.tools import create_document_toolset
@ -104,75 +127,20 @@ from haiku.rag.tools import create_document_toolset
docs = create_document_toolset(config)
```
**Parameters:**
| Parameter | Default | Description |
|-----------|---------|-------------|
| `config` | required | AppConfig (used for summarization LLM) |
| `config` | required | `AppConfig` |
| `base_filter` | `None` | SQL WHERE clause for list operations |
**Tools:**
- `list_documents(page?)` — Paginated document listing (50 per page). Returns `DocumentListResponse` with document titles, URIs, and pagination info.
- `get_document(query)` — Retrieve a document by title or URI. Uses `find_document()` which tries exact URI match, then partial URI match, then partial title match.
- `list_documents(page?)` — Paginated document listing (50 per page).
- `get_document(query)` — Retrieve a document by title or URI.
- `summarize_document(query)` — Generate an LLM summary of a document's content.
## QA Toolset
### Analysis Toolset
`create_qa_toolset()` provides question answering via the research graph, with prior answer recall and background summarization.
```python
from haiku.rag.tools import create_qa_toolset
qa = create_qa_toolset(config)
```
**Parameters:**
| Parameter | Default | Description |
|-----------|---------|-------------|
| `config` | required | AppConfig |
| `base_filter` | `None` | SQL WHERE clause applied to searches |
| `tool_name` | `"ask"` | Name of the tool exposed to the agent |
| `on_ask_complete` | `None` | Callback `(QASessionState, AppConfig) -> None` invoked after each QA cycle |
**Tool: `ask(question, document_name?)`**
Runs the research graph in conversational mode and returns a `QAResult`. When a `ToolContext` is provided:
- Prior answers from `QASessionState.qa_history` are matched via embedding similarity
- The answer is appended to `qa_history`
- `on_ask_complete` callback is invoked (if provided)
- Citations get stable indices via `SessionState.citation_registry`
**State:** QA history accumulates in `QASessionState` under the `haiku.rag.qa_session` namespace.
### Using `run_qa_core()` directly
For programmatic use without an agent, `run_qa_core()` provides the same QA flow:
```python
from haiku.rag.tools.qa import run_qa_core
result = await run_qa_core(
client=client,
config=config,
question="What are the main features?",
document_name="User Guide", # optional document filter
context=context, # optional ToolContext
session_context="User is building a web app", # optional
on_qa_complete=my_callback, # optional post-QA callback
)
print(result.answer)
print(result.confidence)
for citation in result.citations:
print(f" [{citation.index}] {citation.document_title}")
```
## Analysis Toolset
`create_analysis_toolset()` provides computational analysis via the RLM agent, which writes and executes Python code in a Docker sandbox.
`create_analysis_toolset()` provides computational analysis via the RLM agent (Docker sandbox).
```python
from haiku.rag.tools import create_analysis_toolset
@ -180,218 +148,16 @@ from haiku.rag.tools import create_analysis_toolset
analysis = create_analysis_toolset(config)
```
**Parameters:**
| Parameter | Default | Description |
|-----------|---------|-------------|
| `config` | required | AppConfig |
| `config` | required | `AppConfig` |
| `base_filter` | `None` | SQL WHERE clause applied to searches |
| `tool_name` | `"analyze"` | Name of the tool exposed to the agent |
**Tool: `analyze(task, document_name?)`**
Executes a computational task via code execution and returns an `AnalysisResult`. Requires Docker — see [RLM Agent](agents/rlm.md) for setup.
## Tool Prompts
`build_tools_prompt()` generates system prompt guidance for your toolsets — when to use each tool, the `document_name` parameter pattern, and usage examples. It's designed to be spliced into any agent's instructions alongside your own domain-specific guidance.
```python
from haiku.rag.tools import build_tools_prompt
# Generate guidance for the toolsets you're using
tools_prompt = build_tools_prompt(["search", "qa", "documents"])
```
Combine it with your own instructions:
```python
from pydantic_ai import Agent
from haiku.rag.tools import AgentDeps, build_tools_prompt
tools_prompt = build_tools_prompt(["search", "qa"])
agent = Agent(
"anthropic:claude-sonnet-4-5-20250929",
deps_type=AgentDeps,
instructions=f"""You are a medical research assistant.
{tools_prompt}
You also have access to:
- "check_interactions" - Use when the user asks about drug interactions.""",
toolsets=[search_toolset, qa_toolset, my_custom_toolset],
)
```
Available features: `"search"`, `"qa"`, `"documents"`, `"analysis"`.
## Composing Custom Agents
### Using `build_toolkit` (recommended)
`build_toolkit()` bundles toolsets, prompt, and context creation for a given feature set:
```python
from pydantic_ai import Agent
from haiku.rag.client import HaikuRAG
from haiku.rag.tools import AgentDeps, build_toolkit
toolkit = build_toolkit(config, features=["search", "documents", "qa"])
agent = Agent(
"openai:gpt-4o",
deps_type=AgentDeps,
instructions=f"You are a helpful research assistant.\n{toolkit.prompt}",
toolsets=toolkit.toolsets,
)
async with HaikuRAG("path/to/db.lancedb") as client:
context = toolkit.create_context()
deps = AgentDeps(client=client, tool_context=context)
result = await agent.run("What documents do we have about climate?", deps=deps)
print(result.output)
# Access accumulated state
from haiku.rag.tools.search import SearchState, SEARCH_NAMESPACE
search_state = context.get(SEARCH_NAMESPACE, SearchState)
if search_state:
print(f"Total search results: {len(search_state.results)}")
```
**Parameters:**
| Parameter | Default | Description |
|-----------|---------|-------------|
| `config` | required | AppConfig |
| `features` | `["search", "documents"]` | Features to enable |
| `base_filter` | `None` | SQL WHERE clause applied to all toolsets |
| `expand_context` | `True` | Expand search results with surrounding chunks |
| `on_qa_complete` | `None` | Callback invoked after each QA cycle |
**`Toolkit` properties:**
- `toolsets` — list of `FunctionToolset` instances to pass to the Agent
- `prompt` — tool guidance text for the system prompt
- `features` — the feature list this toolkit was built from
- `create_context(state_key=None)` — create a prepared `ToolContext` matching these features
- `prepare(context, state_key=None)` — register namespaces on an existing `ToolContext`
### Using individual factories
For full control, create toolsets individually with `create_*_toolset()`, `build_tools_prompt()`, and `prepare_context()`:
```python
from pydantic_ai import Agent
from haiku.rag.client import HaikuRAG
from haiku.rag.tools import (
AgentDeps,
ToolContext,
build_tools_prompt,
prepare_context,
create_search_toolset,
create_qa_toolset,
create_document_toolset,
)
search = create_search_toolset(config)
qa = create_qa_toolset(config)
docs = create_document_toolset(config)
features = ["search", "documents", "qa"]
tools_prompt = build_tools_prompt(features)
agent = Agent(
"openai:gpt-4o",
deps_type=AgentDeps,
instructions=f"You are a helpful research assistant.\n{tools_prompt}",
toolsets=[search, qa, docs],
)
async with HaikuRAG("path/to/db.lancedb") as client:
context = ToolContext()
prepare_context(context, features=features)
deps = AgentDeps(client=client, tool_context=context)
result = await agent.run("What documents do we have about climate?", deps=deps)
print(result.output)
```
`AgentDeps` satisfies the `RAGDeps` protocol and implements the AG-UI state protocol (`state` getter/setter). For AG-UI streaming, set `state_key` on the `ToolContext` (via `prepare_context` or `toolkit.create_context`):
```python
context = toolkit.create_context(state_key="my_app")
deps = AgentDeps(client=client, tool_context=context)
```
Tool functions access `client` and `tool_context` via pydantic-ai's `RunContext.deps`, so toolsets can be created once and reused across requests.
For complete runnable examples, see [`examples/custom_agent.py`](https://github.com/ggozad/haiku.rag/tree/main/examples/custom_agent.py) (standalone) and [`examples/custom_agent_agui.py`](https://github.com/ggozad/haiku.rag/tree/main/examples/custom_agent_agui.py) (AG-UI streaming server).
All toolsets respect session-level document filters when a `SessionState` is registered in the context. This means setting `SessionState.document_filter` restricts all tools simultaneously.
## AG-UI State Management
Both `AgentDeps` and `ChatDeps` implement the AG-UI `StateHandler` protocol. `ChatDeps` extends `AgentDeps` with chat-specific config and state handling. State is emitted under a namespaced key via `state_key` on the `ToolContext` — set it once via `prepare_context()`.
**Custom agents** use `AgentDeps` + `prepare_context`:
```python
from haiku.rag.tools import AgentDeps, ToolContext, ToolContextCache, prepare_context
context = ToolContext()
prepare_context(context, features=["search", "qa"], state_key="my_app")
deps = AgentDeps(client=client, tool_context=context)
```
**Chat agent** uses `ChatDeps` + `build_chat_toolkit` (adds chat-specific defaults like background summarization):
```python
from haiku.rag.agents.chat import (
AGUI_STATE_KEY, ChatDeps, build_chat_toolkit, create_chat_agent,
)
from haiku.rag.tools import ToolContextCache
chat_toolkit = build_chat_toolkit(config)
agent = create_chat_agent(config, toolkit=chat_toolkit)
# For multi-session apps, cache ToolContext per thread
cache = ToolContextCache()
context, is_new = cache.get_or_create(thread_id)
if is_new:
chat_toolkit.prepare(context, state_key=AGUI_STATE_KEY)
deps = ChatDeps(
config=config,
client=client,
tool_context=context,
)
```
The emitted state structure:
```json
{
"haiku.rag.chat": {
"citations": [],
"qa_history": [],
"session_context": null,
"document_filter": [],
"citation_registry": {}
}
}
```
State flows bidirectionally — the frontend sends its current state on each request, and the agent emits deltas (JSON Patch) reflecting server-side updates (new citations, QA history entries, session context). The server always prefers its own `session_context` over the client's value, since background summarization may have updated it between requests. See the [Web Application](apps.md#web-application) for a complete implementation.
## Filter Helpers
`haiku.rag.tools.filters` provides utilities for building SQL filters:
**`build_document_filter(document_name)`** — Builds a LIKE filter matching against both `uri` and `title`, case-insensitive. Also matches without spaces (e.g., "TB MED 593" matches "tbmed593").
**`build_multi_document_filter(document_names)`** — Combines multiple document name filters with OR logic.
**`combine_filters(filter1, filter2)`** — Combines two filters with AND logic. Returns `None` if both are `None`.
**`get_session_filter(context, base_filter?)`** — Extracts `document_filter` from `SessionState` in the `ToolContext`, builds a SQL filter from it, and combines with an optional `base_filter`.
- **`build_document_filter(document_name)`** — Builds a LIKE filter matching against both `uri` and `title`, case-insensitive. Also matches without spaces (e.g., "TB MED 593" matches "tbmed593").
- **`build_multi_document_filter(document_names)`** — Combines multiple document name filters with OR logic.
- **`combine_filters(filter1, filter2)`** — Combines two filters with AND logic. Returns `None` if both are `None`.

4
examples/README.md
View file

@ -16,7 +16,7 @@ See `docker/README.md` for setup instructions.
**Script:** `custom_agent.py`
Composes `search`, `qa`, and `document` toolsets into a pydantic-ai `Agent` using `AgentDeps` and `prepare_context`. Shows how to run queries and inspect accumulated state (citations, QA history).
Uses the RAG skill with `SkillToolset` to build a conversational agent.
```bash
uv run python examples/custom_agent.py /path/to/db.lancedb
@ -26,7 +26,7 @@ uv run python examples/custom_agent.py /path/to/db.lancedb
**Script:** `custom_agent_agui.py`
A Starlette app that serves an AG-UI streaming endpoint using composed toolsets, `AgentDeps`, and `ToolContextCache` for multi-session support.
A Starlette app that serves an AG-UI streaming endpoint using the RAG skill with `SkillToolset`.
```bash
DB_PATH=/path/to/db.lancedb uv run uvicorn examples.custom_agent_agui:app --reload --port 8000

71
examples/custom_agent.py
View file

@ -1,7 +1,7 @@
"""Custom agent using haiku.rag composable toolsets.
"""Custom agent using the haiku.rag RAG skill.
Demonstrates how to compose search, QA, and document toolsets into a
pydantic-ai Agent using AgentDeps and prepare_context.
Demonstrates how to use the RAG skill with haiku.skills SkillToolset
to build a conversational agent.
Requirements:
- An Ollama instance running locally (default embedder)
@ -14,61 +14,36 @@ Usage:
import asyncio
import sys
from pathlib import Path
from pydantic_ai import Agent
from haiku.rag.client import HaikuRAG
from haiku.rag.tools import (
AgentDeps,
ToolContext,
build_tools_prompt,
create_document_toolset,
create_qa_toolset,
create_search_toolset,
prepare_context,
)
from haiku.rag.skills.rag import create_skill
from haiku.skills.agent import SkillToolset
async def main(db_path: str) -> None:
async with HaikuRAG(db_path) as client:
# Compose toolsets into an agent
config = client.config
search_toolset = create_search_toolset(config)
qa_toolset = create_qa_toolset(config)
document_toolset = create_document_toolset(config)
skill = create_skill(db_path=Path(db_path))
toolset = SkillToolset(skills=[skill])
features = ["search", "documents", "qa"]
tools_prompt = build_tools_prompt(features)
agent = Agent(
"anthropic:claude-haiku-4-5-20251001",
instructions=toolset.system_prompt,
toolsets=[toolset],
)
agent = Agent(
"anthropic:claude-haiku-4-5-20251001",
deps_type=AgentDeps,
output_type=str,
instructions=(
"You are a helpful assistant with access to a knowledge base.\n"
f"{tools_prompt}"
),
toolsets=[search_toolset, qa_toolset, document_toolset],
)
print("Custom agent ready. Ctrl+C to exit.\n")
while True:
try:
user_input = input("You: ").strip()
except (EOFError, KeyboardInterrupt):
break
# Prepare a shared ToolContext
context = ToolContext()
prepare_context(context, features=["search", "documents", "qa"])
if not user_input:
continue
deps = AgentDeps(client=client, tool_context=context)
print("Custom agent ready. Ctrl+C to exit.\n")
while True:
try:
user_input = input("You: ").strip()
except (EOFError, KeyboardInterrupt):
break
if not user_input:
continue
result = await agent.run(user_input, deps=deps)
print(f"\nAgent: {result.output}\n")
result = await agent.run(user_input)
print(f"\nAgent: {result.output}\n")
if __name__ == "__main__":

66
examples/custom_agent_agui.py
View file

@ -1,7 +1,7 @@
"""Custom agent with AG-UI streaming.
A Starlette app that composes haiku.rag toolsets into an AG-UI compatible
agent. Multi-session support via ToolContextCache.
A Starlette app that serves an AG-UI streaming endpoint using the
haiku.rag RAG skill with haiku.skills SkillToolset.
Requirements:
- An Ollama instance running locally (default embedder)
@ -13,25 +13,18 @@ Usage:
import os
import sys
from pathlib import Path
from pydantic_ai import Agent
from pydantic_ai.ag_ui import AGUIAdapter
from pydantic_ai.ui import SSE_CONTENT_TYPE
from pydantic_ai.ui.ag_ui import AGUIAdapter
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import Route
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.tools import (
AgentDeps,
ToolContextCache,
build_tools_prompt,
create_qa_toolset,
create_search_toolset,
prepare_context,
)
from haiku.rag.skills.rag import create_skill
from haiku.skills.agent import SkillToolset
db_path = os.environ.get("DB_PATH")
if not db_path:
@ -40,39 +33,13 @@ if not db_path:
)
sys.exit(1)
AGUI_STATE_KEY = "my_app"
skill = create_skill(db_path=Path(db_path))
toolset = SkillToolset(skills=[skill])
config = AppConfig()
# ToolContextCache maintains per-thread state across requests
context_cache = ToolContextCache()
# Singleton client
_client: HaikuRAG | None = None
def get_client() -> HaikuRAG:
global _client
if _client is None:
_client = HaikuRAG(db_path=db_path)
return _client
features = ["search", "qa"]
tools_prompt = build_tools_prompt(features)
# Create the agent once at module level
agent = Agent(
"anthropic:claude-haiku-4-5-20251001",
deps_type=AgentDeps,
output_type=str,
instructions=(
f"You are a helpful assistant with access to a knowledge base.\n{tools_prompt}"
),
toolsets=[
create_search_toolset(config),
create_qa_toolset(config),
],
instructions=toolset.system_prompt,
toolsets=[toolset],
)
@ -81,19 +48,8 @@ async def stream_chat(request: Request) -> Response:
accept = request.headers.get("accept", SSE_CONTENT_TYPE)
run_input = AGUIAdapter.build_run_input(body)
thread_id = getattr(run_input, "thread_id", None) or "default"
context, is_new = context_cache.get_or_create(thread_id)
if is_new:
prepare_context(
context,
features=["search", "qa"],
state_key=AGUI_STATE_KEY,
)
deps = AgentDeps(client=get_client(), tool_context=context)
adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
event_stream = adapter.run_stream(deps=deps)
event_stream = adapter.run_stream()
sse_event_stream = adapter.encode_stream(event_stream)
return StreamingResponse(