Update tui and app to use the toolkit

CHANGELOG.md

@@ -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`

app/backend/main.py

@@ -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)
+# Toolkit and agent are created once at module level
-agent = create_chat_agent(Config)
+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,

docs/agents/index.md

@@ -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_config=ModelConfig(provider="openai", name="gpt-4o-mini"),
-        model="gpt-4o-mini",
-        use_citations=False,
     )
 
-    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=ModelConfig(provider="openai", name="gpt-4o-mini"),
-        model="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

docs/python.md

@@ -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 (
+from haiku.rag.tools import AgentDeps, build_toolkit
-    AgentDeps, ToolContext, prepare_context,
-    create_search_toolset, create_qa_toolset,
-)
 
-search = create_search_toolset(config)
+toolkit = build_toolkit(config, features=["search", "qa"])
-qa = create_qa_toolset(config)
 
 agent = Agent(
     "openai:gpt-4o",
     deps_type=AgentDeps,
-    instructions="You are a helpful assistant.",
+    instructions=f"You are a helpful assistant.\n{toolkit.prompt}",
-    toolsets=[search, qa],
+    toolsets=toolkit.toolsets,
 )
 
 async with HaikuRAG("path/to/db.lancedb") as client:
-    context = ToolContext()
+    context = toolkit.create_context()
-    prepare_context(context, features=["search", "qa"])
     deps = AgentDeps(client=client, tool_context=context)
     result = await agent.run("What are the main findings?", deps=deps)
 ```

docs/tools.md

@@ -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()
+context = toolkit.create_context(state_key="my_app")
-prepare_context(context, features=["search", "qa"], 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)
+context, is_new = cache.get_or_create(thread_id)
-prepare_chat_context(context)  # idempotent; sets state_key="haiku.rag.chat"
+if is_new:
+    chat_toolkit.prepare(context, state_key=AGUI_STATE_KEY)
 
 deps = ChatDeps(
     config=config,

haiku_rag_slim/haiku/rag/agents/chat/__init__.py

@@ -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",

haiku_rag_slim/haiku/rag/agents/chat/agent.py

@@ -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
-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.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()
+            toolkit = build_chat_toolkit(config)
-            prepare_chat_context(context)
+            context = toolkit.create_context(state_key=AGUI_STATE_KEY)
-            agent = create_chat_agent(config)
+            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 toolkit is None:
-    if FEATURE_SEARCH in features:
+        toolkit = build_chat_toolkit(config, features=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))
 
     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",

haiku_rag_slim/haiku/rag/agents/qa/agent.py

@@ -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

haiku_rag_slim/haiku/rag/chat/app.py

@@ -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
+        # Create toolkit, context, and agent
-        self.tool_context = ToolContext()
+        self.toolkit = build_chat_toolkit(self.config)
-        prepare_chat_context(self.tool_context)
+        self.tool_context = self.toolkit.create_context(state_key=AGUI_STATE_KEY)
-        self.agent = create_chat_agent(self.config)
+        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)

haiku_rag_slim/haiku/rag/tools/analysis.py

@@ -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

haiku_rag_slim/haiku/rag/tools/document.py

@@ -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)

haiku_rag_slim/haiku/rag/tools/qa.py

@@ -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

haiku_rag_slim/haiku/rag/tools/search.py

@@ -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

haiku_rag_slim/haiku/rag/tools/toolkit.py

@@ -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