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

Update tui and app to use the toolkit

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-16 12:15:03 +02:00
parent 7ff2123806
commit 8c76897af2
No known key found for this signature in database
14 changed files with 150 additions and 81 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
CHANGELOG.md
View file

@ -8,6 +8,7 @@
- `create_document_toolset()` — document listing, retrieval, and summarization
- `create_qa_toolset()` — question answering via research graph with prior answer recall
- `create_analysis_toolset()` — computational analysis via RLM agent (Docker sandbox)
- **`Toolkit` and `build_toolkit()`**: High-level factory that bundles toolsets, prompt, and context creation for a given feature set. Reduces agent composition from ~15 lines to ~5. `build_chat_toolkit()` adds chat-specific defaults (background summarization callback)
- **`ToolContext`**: Namespace-based state container shared across toolsets. Toolsets register Pydantic models under string namespaces, enabling state accumulation (search results, citations, QA history) across invocations
- **`ToolContextCache`**: In-memory TTL-based cache for `ToolContext` instances, keyed by external session/thread ID. Replaces module-level caches for embeddings and summaries
- **`run_qa_core()`**: Extracted core QA function for direct programmatic use without an agent
@ -17,6 +18,8 @@
### Changed
- **Toolset factories decoupled from runtime dependencies**: `create_search_toolset()`, `create_qa_toolset()`, `create_document_toolset()`, `create_analysis_toolset()`, and `create_chat_agent()` no longer take `client` or `context` parameters. Instead, tool functions receive these via pydantic-ai's `RunContext.deps`. This enables toolset and agent creation at configuration time (cacheable, created once), with only lightweight deps created per-request. Deps must satisfy the `RAGDeps` protocol (`client: HaikuRAG`, `tool_context: ToolContext | None`)
- **Toolset factory return types narrowed to `FunctionToolset[RAGDeps]`**: All four toolset factories now declare their return type as `FunctionToolset[RAGDeps]` instead of bare `FunctionToolset`
- **`create_chat_agent()` accepts optional `toolkit` parameter**: Pass a pre-built `Toolkit` to share toolsets between agent and context creation, avoiding duplicate construction
- **`ChatDeps` now includes `client`**: `ChatDeps(config=..., client=..., tool_context=...)` — the `client` field was added since it's no longer captured by the agent factory
- **`prepare_chat_context()` helper**: Extracted from `create_chat_agent()` for idempotent namespace registration, since the agent factory no longer has access to the context
- **Chat agent architecture**: Rebuilt on composable toolsets instead of monolithic tool definitions. Chat agent is now a thin wrapper around `create_search_toolset`, `create_document_toolset`, `create_qa_toolset`, and `create_analysis_toolset`

10
app/backend/main.py
View file

@ -13,9 +13,10 @@ from starlette.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import Route
from haiku.rag.agents.chat import (
AGUI_STATE_KEY,
ChatDeps,
build_chat_toolkit,
create_chat_agent,
prepare_chat_context,
)
from haiku.rag.client import HaikuRAG
from haiku.rag.config import load_yaml_config
@ -68,8 +69,9 @@ def get_client() -> HaikuRAG:
return _client
# Agent is created once at module level (no runtime deps needed)
agent = create_chat_agent(Config)
# Toolkit and agent are created once at module level
chat_toolkit = build_chat_toolkit(Config)
agent = create_chat_agent(Config, toolkit=chat_toolkit)
async def stream_chat(request: Request) -> Response:
@ -85,7 +87,7 @@ async def stream_chat(request: Request) -> Response:
thread_id = getattr(run_input, "thread_id", None) or "default"
context, is_new = context_cache.get_or_create(thread_id)
if is_new:
prepare_chat_context(context)
chat_toolkit.prepare(context, state_key=AGUI_STATE_KEY)
deps = ChatDeps(
config=Config,

13
docs/agents/index.md
View file

@ -35,17 +35,16 @@ haiku-rag ask "What are the main features of haiku.rag?" --deep
```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,
provider="openai",
model="gpt-4o-mini",
use_citations=False,
model_config=ModelConfig(provider="openai", name="gpt-4o-mini"),
)
answer = await agent.answer("What is climate change?")
answer, citations = await agent.answer("What is climate change?")
print(answer)
```
@ -256,15 +255,14 @@ async with HaikuRAG(path_to_db) as client:
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig, ResearchConfig
from haiku.rag.config.models import AppConfig, ModelConfig, ResearchConfig
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
custom_config = AppConfig(
research=ResearchConfig(
provider="openai",
model="gpt-4o-mini",
model=ModelConfig(provider="openai", name="gpt-4o-mini"),
max_iterations=5,
max_concurrency=3,
)
@ -282,6 +280,7 @@ async with HaikuRAG(path_to_db) as client:
**Conversational mode with prior answers:**
```python
from haiku.rag.config import Config
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.models import SearchAnswer

15
docs/python.md
View file

@ -433,24 +433,19 @@ haiku.rag provides composable toolset factories that can be mixed into any pydan
```python
from pydantic_ai import Agent
from haiku.rag.tools import (
AgentDeps, ToolContext, prepare_context,
create_search_toolset, create_qa_toolset,
)
from haiku.rag.tools import AgentDeps, build_toolkit
search = create_search_toolset(config)
qa = create_qa_toolset(config)
toolkit = build_toolkit(config, features=["search", "qa"])
agent = Agent(
"openai:gpt-4o",
deps_type=AgentDeps,
instructions="You are a helpful assistant.",
toolsets=[search, qa],
instructions=f"You are a helpful assistant.\n{toolkit.prompt}",
toolsets=toolkit.toolsets,
)
async with HaikuRAG("path/to/db.lancedb") as client:
context = ToolContext()
prepare_context(context, features=["search", "qa"])
context = toolkit.create_context()
deps = AgentDeps(client=client, tool_context=context)
result = await agent.run("What are the main findings?", deps=deps)
```

80
docs/tools.md
View file

@ -227,7 +227,59 @@ Available features: `"search"`, `"qa"`, `"documents"`, `"analysis"`.
## Composing Custom Agents
Toolsets are designed to be composed into custom pydantic-ai agents. Use `AgentDeps`, `prepare_context`, and `build_tools_prompt` for minimal boilerplate:
### 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
@ -242,7 +294,6 @@ from haiku.rag.tools import (
create_document_toolset,
)
# Toolsets are created once at configuration time
search = create_search_toolset(config)
qa = create_qa_toolset(config)
docs = create_document_toolset(config)
@ -264,19 +315,12 @@ async with HaikuRAG("path/to/db.lancedb") as client:
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)}")
```
`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`):
`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 = ToolContext()
prepare_context(context, features=["search", "qa"], state_key="my_app")
context = toolkit.create_context(state_key="my_app")
deps = AgentDeps(client=client, tool_context=context)
```
@ -300,20 +344,22 @@ prepare_context(context, features=["search", "qa"], state_key="my_app")
deps = AgentDeps(client=client, tool_context=context)
```
**Chat agent** uses `ChatDeps` + `prepare_chat_context` (adds chat-specific overrides like background summarization and initial context handling):
**Chat agent** uses `ChatDeps` + `build_chat_toolkit` (adds chat-specific defaults like background summarization):
```python
from haiku.rag.agents.chat import (
ChatDeps, create_chat_agent, prepare_chat_context,
AGUI_STATE_KEY, ChatDeps, build_chat_toolkit, create_chat_agent,
)
from haiku.rag.tools import ToolContext, ToolContextCache
from haiku.rag.tools import ToolContextCache
agent = create_chat_agent(config)
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)
prepare_chat_context(context) # idempotent; sets state_key="haiku.rag.chat"
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,

2
haiku_rag_slim/haiku/rag/agents/chat/__init__.py
View file

@ -5,6 +5,7 @@ from haiku.rag.agents.chat.agent import (
FEATURE_QA,
FEATURE_SEARCH,
ChatDeps,
build_chat_toolkit,
create_chat_agent,
prepare_chat_context,
run_chat_agent,
@ -28,6 +29,7 @@ __all__ = [
"FEATURE_QA",
"FEATURE_SEARCH",
"build_chat_prompt",
"build_chat_toolkit",
"create_chat_agent",
"prepare_chat_context",
"run_chat_agent",

69
haiku_rag_slim/haiku/rag/agents/chat/agent.py
View file

@ -9,19 +9,20 @@ from haiku.rag.agents.chat.context import (
from haiku.rag.agents.chat.prompts import build_chat_prompt
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
from haiku.rag.config.models import AppConfig
from haiku.rag.tools.context import ToolContext, prepare_context
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.deps import AgentDeps
from haiku.rag.tools.document import create_document_toolset
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState, create_qa_toolset
from haiku.rag.tools.search import create_search_toolset
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
from haiku.rag.tools.session import SessionContext
from haiku.rag.tools.toolkit import (
FEATURE_ANALYSIS,
FEATURE_DOCUMENTS,
FEATURE_QA,
FEATURE_SEARCH,
Toolkit,
build_toolkit,
)
from haiku.rag.utils import get_model
FEATURE_SEARCH = "search"
FEATURE_DOCUMENTS = "documents"
FEATURE_QA = "qa"
FEATURE_ANALYSIS = "analysis"
DEFAULT_FEATURES = [FEATURE_SEARCH, FEATURE_DOCUMENTS, FEATURE_QA]
@ -71,6 +72,28 @@ class ChatDeps(AgentDeps):
)
def build_chat_toolkit(
config: AppConfig,
features: list[str] | None = None,
) -> Toolkit:
"""Build a Toolkit configured for the chat agent.
Includes the on_qa_complete callback that triggers background
session summarization.
Args:
config: Application configuration.
features: List of features to enable. Defaults to DEFAULT_FEATURES.
Returns:
A Toolkit ready for chat agent composition and context creation.
"""
if features is None:
features = DEFAULT_FEATURES
return build_toolkit(config, features=features, on_qa_complete=_on_qa_complete)
def prepare_chat_context(
context: ToolContext,
features: list[str] | None = None,
@ -83,6 +106,8 @@ def prepare_chat_context(
context: ToolContext to prepare.
features: List of enabled features. Defaults to DEFAULT_FEATURES.
"""
from haiku.rag.tools.context import prepare_context
if features is None:
features = DEFAULT_FEATURES
@ -93,6 +118,7 @@ def create_chat_agent(
config: AppConfig,
features: list[str] | None = None,
preamble: str | None = None,
toolkit: Toolkit | None = None,
) -> Agent[ChatDeps, str]:
"""Create the chat agent with composed toolsets.
@ -104,32 +130,26 @@ def create_chat_agent(
preamble: Optional custom identity/rules section for the system prompt.
When provided, replaces the default identity prompt. Tool guidance,
feature rules, and closing are still appended by the builder.
toolkit: Optional pre-built Toolkit. When provided, its toolsets are
used directly. When omitted, a toolkit is built from config and
features.
Returns:
The configured chat agent.
Example:
async with HaikuRAG(db_path, create=True) as client:
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(config)
toolkit = build_chat_toolkit(config)
context = toolkit.create_context(state_key=AGUI_STATE_KEY)
agent = create_chat_agent(config, toolkit=toolkit)
deps = ChatDeps(config=config, client=client, tool_context=context)
result = await agent.run("Search for X", deps=deps)
"""
if features is None:
features = DEFAULT_FEATURES
toolsets = []
if FEATURE_SEARCH in features:
toolsets.append(create_search_toolset(config))
if FEATURE_DOCUMENTS in features:
toolsets.append(create_document_toolset(config))
if FEATURE_QA in features:
toolsets.append(create_qa_toolset(config, on_ask_complete=_on_qa_complete))
if FEATURE_ANALYSIS in features:
from haiku.rag.tools.analysis import create_analysis_toolset
toolsets.append(create_analysis_toolset(config))
if toolkit is None:
toolkit = build_chat_toolkit(config, features=features)
model = get_model(config.qa.model, config)
@ -138,7 +158,7 @@ def create_chat_agent(
deps_type=ChatDeps,
output_type=str,
instructions=build_chat_prompt(features, preamble=preamble),
toolsets=toolsets,
toolsets=toolkit.toolsets,
retries=3,
)
@ -182,6 +202,7 @@ async def run_chat_agent(
__all__ = [
"build_chat_toolkit",
"create_chat_agent",
"prepare_chat_context",
"run_chat_agent",

6
haiku_rag_slim/haiku/rag/agents/qa/agent.py
View file

@ -56,18 +56,18 @@ class QuestionAnswerAgent:
# Agent created per-call: toolset varies with filter, and Agent
# construction is pure Python (no IO).
agent = Agent(
agent: Agent[_QARunDeps, RawSearchAnswer] = Agent( # ty: ignore[invalid-assignment]
model=get_model(self._model_config, self._config),
deps_type=_QARunDeps,
output_type=RawSearchAnswer,
output_retries=3,
instructions=self._system_prompt,
toolsets=[search_toolset], # ty: ignore[invalid-argument-type]
toolsets=[search_toolset],
retries=3,
)
deps = _QARunDeps(client=self._client, tool_context=context)
result = await agent.run(question, deps=deps) # ty: ignore[invalid-argument-type]
result = await agent.run(question, deps=deps)
output = result.output
# Get search results from context for citation resolution

12
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -17,13 +17,13 @@ from pydantic_ai.messages import ModelMessage
from haiku.rag.agents.chat.agent import (
ChatDeps,
build_chat_toolkit,
create_chat_agent,
prepare_chat_context,
trigger_background_summarization,
)
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
from haiku.rag.client import HaikuRAG
from haiku.rag.config import get_config
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
if TYPE_CHECKING:
@ -151,10 +151,10 @@ class ChatApp(App):
)
await self.client.__aenter__()
# Create tool context and agent
self.tool_context = ToolContext()
prepare_chat_context(self.tool_context)
self.agent = create_chat_agent(self.config)
# Create toolkit, context, and agent
self.toolkit = build_chat_toolkit(self.config)
self.tool_context = self.toolkit.create_context(state_key=AGUI_STATE_KEY)
self.agent = create_chat_agent(self.config, toolkit=self.toolkit)
# Sync document filter to tool context
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)

4
haiku_rag_slim/haiku/rag/tools/analysis.py
View file

@ -17,7 +17,7 @@ def create_analysis_toolset(
config: AppConfig,
base_filter: str | None = None,
tool_name: str = "analyze",
) -> FunctionToolset:
) -> FunctionToolset[RAGDeps]:
"""Create a toolset with code analysis capabilities via RLM agent.
Args:
@ -77,6 +77,6 @@ def create_analysis_toolset(
code_executed=bool(program),
)
toolset = FunctionToolset()
toolset: FunctionToolset[RAGDeps] = FunctionToolset()
toolset.add_function(analyze, name=tool_name)
return toolset

4
haiku_rag_slim/haiku/rag/tools/document.py
View file

@ -68,7 +68,7 @@ async def find_document(client: HaikuRAG, query: str):
def create_document_toolset(
config: AppConfig,
base_filter: str | None = None,
) -> FunctionToolset:
) -> FunctionToolset[RAGDeps]:
"""Create a toolset with document management capabilities.
Args:
@ -169,7 +169,7 @@ def create_document_toolset(
return f"**Summary of {doc.title or doc.uri}:**\n\n{result.output}"
toolset = FunctionToolset()
toolset: FunctionToolset[RAGDeps] = FunctionToolset()
toolset.add_function(list_documents)
toolset.add_function(get_document)
toolset.add_function(summarize_document)

4
haiku_rag_slim/haiku/rag/tools/qa.py
View file

@ -210,7 +210,7 @@ def create_qa_toolset(
base_filter: str | None = None,
tool_name: str = "ask",
on_ask_complete: Callable[[QASessionState, AppConfig], None] | None = None,
) -> FunctionToolset:
) -> FunctionToolset[RAGDeps]:
"""Create a toolset with Q&A capabilities using research graph.
Args:
@ -277,6 +277,6 @@ def create_qa_toolset(
return qa_result
toolset = FunctionToolset()
toolset: FunctionToolset[RAGDeps] = FunctionToolset()
toolset.add_function(ask, name=tool_name)
return toolset

4
haiku_rag_slim/haiku/rag/tools/search.py
View file

@ -25,7 +25,7 @@ def create_search_toolset(
expand_context: bool = True,
base_filter: str | None = None,
tool_name: str = "search",
) -> FunctionToolset:
) -> FunctionToolset[RAGDeps]:
"""Create a toolset with search capabilities.
Args:
@ -150,6 +150,6 @@ def create_search_toolset(
]
return "\n\n".join(formatted)
toolset = FunctionToolset()
toolset: FunctionToolset[RAGDeps] = FunctionToolset()
toolset.add_function(search, name=tool_name)
return toolset

5
haiku_rag_slim/haiku/rag/tools/toolkit.py
View file

@ -1,5 +1,6 @@
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any
from pydantic_ai import FunctionToolset
@ -21,7 +22,7 @@ class Toolkit:
an agent with haiku.rag toolsets and create matching ToolContexts.
"""
toolsets: list[FunctionToolset] = field(default_factory=list)
toolsets: list[FunctionToolset[Any]] = field(default_factory=list)
prompt: str = ""
features: list[str] = field(default_factory=list)
@ -72,7 +73,7 @@ def build_toolkit(
if features is None:
features = [FEATURE_SEARCH, FEATURE_DOCUMENTS]
toolsets: list[FunctionToolset] = []
toolsets: list[FunctionToolset[Any]] = []
if FEATURE_SEARCH in features:
from haiku.rag.tools.search import create_search_toolset