haiku.rag/TOOLS_REFACTORING_PLAN.md

Tools Extraction Refactoring Plan

Goal

Extract tools from haiku.rag agents into a reusable tools/ module, enabling users to create pydantic-ai agents outside haiku.rag and compose toolsets as needed.

Target API

from pydantic_ai import Agent
from haiku.rag import HaikuRAG
from haiku.rag.tools import ToolContext, create_search_toolset, create_document_toolset

async with HaikuRAG(db_path) as client:
    context = ToolContext()
    search_tools = create_search_toolset(client, config, context)
    doc_tools = create_document_toolset(client, config, context)

    agent = Agent(
        'anthropic:claude-sonnet',
        toolsets=[search_tools, doc_tools]
    )
    result = await agent.run("Find documents about X")

    # Access accumulated state after run
    search_state = context.get("haiku.rag.search")
    for result in search_state.results:
        print(f"{result.document_title}")

Design Principles

1. ToolContext is a pure generic container - No special-cased fields. Toolsets register their own Pydantic model state under namespaces.

2. Shared state via same namespace - Multiple toolsets can share state (e.g., citations, filters) by registering under the same namespace.

3. App manages identity - ToolContext has no session/user identity. The app layer manages session_id -> ToolContext mapping.

4. Toolsets are stateless factories - create_*_toolset() returns a FunctionToolset. State lives in the context they're given.

ToolContext Design

class ToolContext(BaseModel):
    """Generic state container for toolsets.

    Toolsets register Pydantic model state under namespaces.
    Multiple toolsets can share state via the same namespace.
    """
    _namespaces: dict[str, BaseModel] = PrivateAttr(default_factory=dict)

    def register(self, namespace: str, state: BaseModel) -> None: ...
    def get(self, namespace: str) -> BaseModel | None: ...
    def get_or_create(self, namespace: str, factory: Callable[[], T]) -> T: ...
    def clear_namespace(self, namespace: str) -> None: ...
    def clear_all(self) -> None: ...
    def dump_namespaces(self) -> dict[str, dict[str, Any]]: ...
    def load_namespace(self, namespace: str, state_type: type[T], data: dict) -> T: ...

Toolset State Examples

Each toolset defines its own state model:

# Search toolset state
class SearchState(BaseModel):
    results: list[SearchResult] = []
    filter: str | None = None

SEARCH_NAMESPACE = "haiku.rag.search"

# QA toolset state
class QAState(BaseModel):
    history: list[QAResult] = []

QA_NAMESPACE = "haiku.rag.qa"

# Shared citation state (used by multiple toolsets)
class CitationState(BaseModel):
    registry: dict[str, int] = {}

    def get_or_assign_index(self, chunk_id: str) -> int:
        if chunk_id in self.registry:
            return self.registry[chunk_id]
        new_index = len(self.registry) + 1
        self.registry[chunk_id] = new_index
        return new_index

CITATION_NAMESPACE = "haiku.rag.citations"

Multi-User/Session Management

App layer manages context routing:

# App maintains context per session
contexts: dict[str, ToolContext] = {}

def get_context(session_id: str) -> ToolContext:
    if session_id not in contexts:
        contexts[session_id] = ToolContext()
    return contexts[session_id]

# When running agent
context = get_context(user_session_id)
toolsets = [create_search_toolset(client, config, context)]
await agent.run(prompt, toolsets=toolsets)

New Module Structure

haiku_rag_slim/haiku/rag/
├── tools/                      # NEW
│   ├── __init__.py             # Public exports
│   ├── context.py              # ToolContext (generic state container)
│   ├── models.py               # QAResult, AnalysisResult
│   ├── filters.py              # build_document_filter, combine_filters
│   ├── search.py               # create_search_toolset()
│   ├── document.py             # create_document_toolset()
│   ├── qa.py                   # create_qa_toolset()
│   └── analysis.py             # create_analysis_toolset()
├── agents/                     # REFACTORED to use tools/

Implementation Chunks

Chunk 1: Create tools module foundation ✅ DONE

• Created tools/__init__.py, tools/context.py, tools/models.py, tools/filters.py
• Created ToolContext as generic namespace-based Pydantic model
• Moved filter utilities from agents/chat/state.py to tools/filters.py
• Created result models (QAResult, AnalysisResult)
• Added tests for ToolContext and filters

Chunk 2: Create SearchToolset ✅ DONE

• Created tools/search.py with create_search_toolset()
• Defined SearchState model for accumulating search results
• Core search logic: client.search() → client.expand_context() → format_for_agent()
• Results accumulated in SearchState under SEARCH_NAMESPACE
• Added 13 tests for SearchToolset

Chunk 3: Refactor QA Agent to use SearchToolset ✅ DONE

• Updated agents/qa/agent.py to use create_search_toolset()
• Added base_filter and tool_name parameters to create_search_toolset()
• QA agent now uses ToolContext + SearchState for result accumulation
• Public interface (answer(question, filter)) unchanged
• All 5 QA tests pass

Chunk 4: Create DocumentToolset ✅ DONE

• Created tools/document.py with create_document_toolset()
• Defined DocumentState, DocumentInfo, DocumentListResponse models
• Extracted list_documents, get_document, summarize_document tools
• Moved find_document helper (now public)
• Added 13 tests

Chunk 5: Create QAToolset ✅ DONE

• Created tools/qa.py with create_qa_toolset()
• Defined QAState model (tracks QA history)
• Runs research graph, returns structured QAResult
• Supports base_filter, tool_name, session_context, prior_answers params
• Added 7 tests

Chunk 6: Create AnalysisToolset ✅ DONE

• Created tools/analysis.py with create_analysis_toolset()
• Defined AnalysisState model (tracks CodeExecution history)
• Extracted analyze tool (RLM delegation with filter support)
• Fixed circular import by using direct submodule imports
• Added 6 tests

Chunk 7: Refactor Chat Agent ✅ DONE

• Removed analyze tool from chat agent (kept hardcoded, not composing toolsets)
• Reverted system prompt to pre-analyze version
• Removed test_analyze_tool test and cassette file
• All 47 chat agent tests pass

Chunk 8: Refactor Research Graph

• Update _search_one_step_logic to use search toolset
• Verify research tests pass

Chunk 9: Public API and Documentation

• Export from haiku.rag.tools and haiku.rag
• Update CLAUDE.md
• Add usage examples

Verification

• Run pytest after each chunk
• Run ty check and ruff check
• Test with existing agents (QA, Chat, Research)
• Test with external agent using new toolsets

Critical Files

• haiku_rag_slim/haiku/rag/agents/chat/agent.py - largest tool collection
• haiku_rag_slim/haiku/rag/agents/qa/agent.py - simplest, good starting point
• haiku_rag_slim/haiku/rag/agents/chat/state.py - filter utilities (now moved)
• haiku_rag_slim/haiku/rag/agents/research/graph.py - search tool inside step
• haiku_rag_slim/haiku/rag/store/models/chunk.py - SearchResult.format_for_agent()