270 lines
9.3 KiB
Markdown
270 lines
9.3 KiB
Markdown
# Toolsets
|
|
|
|
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.
|
|
|
|
## ToolContext
|
|
|
|
`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.
|
|
|
|
```python
|
|
from haiku.rag.tools import ToolContext
|
|
|
|
context = ToolContext()
|
|
```
|
|
|
|
### Registering and accessing state
|
|
|
|
```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)
|
|
```
|
|
|
|
### Serialization
|
|
|
|
The entire context can be serialized and restored:
|
|
|
|
```python
|
|
# Serialize all namespaces
|
|
data = context.dump_namespaces()
|
|
# {"my_namespace": {"count": 0}}
|
|
|
|
# Restore a namespace from serialized data
|
|
context.load_namespace("my_namespace", MyState, data["my_namespace"])
|
|
```
|
|
|
|
## Search Toolset
|
|
|
|
`create_search_toolset()` provides hybrid search (vector + full-text) with context expansion and citation tracking.
|
|
|
|
```python
|
|
from haiku.rag.tools import ToolContext, create_search_toolset
|
|
|
|
context = ToolContext()
|
|
search = create_search_toolset(client, config, context=context)
|
|
```
|
|
|
|
**Parameters:**
|
|
|
|
| Parameter | Default | Description |
|
|
|-----------|---------|-------------|
|
|
| `client` | required | HaikuRAG client |
|
|
| `config` | required | AppConfig |
|
|
| `context` | `None` | ToolContext for state accumulation |
|
|
| `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 |
|
|
|
|
**Tool: `search(query, limit?, filter?)`**
|
|
|
|
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.
|
|
|
|
```python
|
|
from haiku.rag.tools import ToolContext, create_document_toolset
|
|
|
|
context = ToolContext()
|
|
docs = create_document_toolset(client, config, context=context)
|
|
```
|
|
|
|
**Parameters:**
|
|
|
|
| Parameter | Default | Description |
|
|
|-----------|---------|-------------|
|
|
| `client` | required | HaikuRAG client |
|
|
| `config` | required | AppConfig (used for summarization LLM) |
|
|
| `context` | `None` | ToolContext for session filtering |
|
|
| `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.
|
|
- `summarize_document(query)` — Generate an LLM summary of a document's content.
|
|
|
|
## QA Toolset
|
|
|
|
`create_qa_toolset()` provides question answering via the research graph, with prior answer recall and background summarization.
|
|
|
|
```python
|
|
from haiku.rag.tools import ToolContext, create_qa_toolset
|
|
|
|
context = ToolContext()
|
|
qa = create_qa_toolset(client, config, context=context)
|
|
```
|
|
|
|
**Parameters:**
|
|
|
|
| Parameter | Default | Description |
|
|
|-----------|---------|-------------|
|
|
| `client` | required | HaikuRAG client |
|
|
| `config` | required | AppConfig |
|
|
| `context` | `None` | ToolContext for state accumulation |
|
|
| `base_filter` | `None` | SQL WHERE clause applied to searches |
|
|
| `tool_name` | `"ask"` | Name of the tool exposed to the agent |
|
|
| `session_context` | `None` | Session context for the research graph |
|
|
| `prior_answers` | `None` | Prior answers for context |
|
|
|
|
**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`
|
|
- Background summarization is triggered
|
|
- 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
|
|
)
|
|
|
|
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.
|
|
|
|
```python
|
|
from haiku.rag.tools import create_analysis_toolset
|
|
|
|
analysis = create_analysis_toolset(client, config, context=context)
|
|
```
|
|
|
|
**Parameters:**
|
|
|
|
| Parameter | Default | Description |
|
|
|-----------|---------|-------------|
|
|
| `client` | required | HaikuRAG client |
|
|
| `config` | required | AppConfig |
|
|
| `context` | `None` | ToolContext for session filtering |
|
|
| `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.
|
|
|
|
## Composing Custom Agents
|
|
|
|
Toolsets are designed to be composed into custom pydantic-ai agents:
|
|
|
|
```python
|
|
from pydantic_ai import Agent
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import Config
|
|
from haiku.rag.tools import (
|
|
ToolContext,
|
|
create_search_toolset,
|
|
create_qa_toolset,
|
|
create_document_toolset,
|
|
)
|
|
|
|
async with HaikuRAG("path/to/db.lancedb") as client:
|
|
# Shared context across all toolsets
|
|
context = ToolContext()
|
|
|
|
# Pick the toolsets you need
|
|
search = create_search_toolset(client, Config, context=context)
|
|
qa = create_qa_toolset(client, Config, context=context)
|
|
docs = create_document_toolset(client, Config, context=context)
|
|
|
|
agent = Agent(
|
|
"openai:gpt-4o",
|
|
instructions="You are a helpful research assistant.",
|
|
toolsets=[search, qa, docs],
|
|
)
|
|
|
|
result = await agent.run("What documents do we have about climate?")
|
|
print(result.output)
|
|
|
|
# Access accumulated state
|
|
from haiku.rag.tools import SearchState, SEARCH_NAMESPACE
|
|
search_state = context.get(SEARCH_NAMESPACE, SearchState)
|
|
if search_state:
|
|
print(f"Total search results: {len(search_state.results)}")
|
|
```
|
|
|
|
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
|
|
|
|
When using the chat agent with [AG-UI](https://docs.ag-ui.com) streaming, `ChatDeps` implements the `StateHandler` protocol. State is emitted under a namespaced key via `state_key`:
|
|
|
|
```python
|
|
from haiku.rag.agents.chat import AGUI_STATE_KEY, ChatDeps, create_chat_agent
|
|
from haiku.rag.tools import ToolContext, ToolContextCache
|
|
|
|
# For multi-session apps, cache ToolContext per thread
|
|
cache = ToolContextCache()
|
|
context, _is_new = cache.get_or_create(thread_id)
|
|
|
|
agent = create_chat_agent(config, client, context)
|
|
deps = ChatDeps(
|
|
config=config,
|
|
tool_context=context,
|
|
state_key=AGUI_STATE_KEY, # "haiku.rag.chat"
|
|
)
|
|
```
|
|
|
|
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`.
|