Introduce tool module & ToolContext to save state

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

@@ -13,7 +13,6 @@ from haiku.rag.agents.chat.state import (
     QAResponse,
     SearchDeps,
     SessionContext,
-    build_document_filter,
 )
 
 __all__ = [
@@ -27,7 +26,6 @@ __all__ = [
     "QAResponse",
     "SearchDeps",
     "SessionContext",
-    "build_document_filter",
     "summarize_session",
     "update_session_context",
 ]

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

@@ -17,9 +17,6 @@ from haiku.rag.agents.chat.state import (
     DocumentInfo,
     DocumentListResponse,
     QAResponse,
-    build_document_filter,
-    build_multi_document_filter,
-    combine_filters,
     emit_state_event,
 )
 from haiku.rag.agents.research.dependencies import ResearchContext
@@ -29,6 +26,11 @@ from haiku.rag.agents.research.state import ResearchDeps, ResearchState
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig
 from haiku.rag.embeddings import get_embedder
+from haiku.rag.tools.filters import (
+    build_document_filter,
+    build_multi_document_filter,
+    combine_filters,
+)
 from haiku.rag.utils import get_model
 
 # Similarity threshold for finding relevant prior answers

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

@@ -162,36 +162,6 @@ class SearchDeps:
     search_results: list[SearchResult] = field(default_factory=list)
 
 
-def build_document_filter(document_name: str) -> str:
-    """Build SQL filter for document name matching."""
-    escaped = document_name.replace("'", "''")
-    no_spaces = escaped.replace(" ", "")
-    return (
-        f"LOWER(uri) LIKE LOWER('%{escaped}%') OR LOWER(title) LIKE LOWER('%{escaped}%') "
-        f"OR LOWER(uri) LIKE LOWER('%{no_spaces}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')"
-    )
-
-
-def build_multi_document_filter(document_names: list[str]) -> str | None:
-    """Build SQL filter for multiple document names (OR combined)."""
-    if not document_names:
-        return None
-    filters = [build_document_filter(name) for name in document_names]
-    if len(filters) == 1:
-        return filters[0]
-    return " OR ".join(f"({f})" for f in filters)
-
-
-def combine_filters(filter1: str | None, filter2: str | None) -> str | None:
-    """Combine two SQL filters with AND logic."""
-    filters = [f for f in [filter1, filter2] if f]
-    if not filters:
-        return None
-    if len(filters) == 1:
-        return filters[0]
-    return f"({filters[0]}) AND ({filters[1]})"
-
-
 def emit_state_event(
     current_state: ChatSessionState,
     new_state: ChatSessionState,

haiku_rag_slim/haiku/rag/tools/__init__.py

@@ -0,0 +1,16 @@
+from haiku.rag.tools.context import ToolContext
+from haiku.rag.tools.filters import (
+    build_document_filter,
+    build_multi_document_filter,
+    combine_filters,
+)
+from haiku.rag.tools.models import AnalysisResult, QAResult
+
+__all__ = [
+    "ToolContext",
+    "QAResult",
+    "AnalysisResult",
+    "build_document_filter",
+    "build_multi_document_filter",
+    "combine_filters",
+]

haiku_rag_slim/haiku/rag/tools/context.py

@@ -0,0 +1,114 @@
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+from pydantic import BaseModel, PrivateAttr
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class ToolContext(BaseModel):
+    """Generic state container for haiku.rag toolsets.
+
+    Toolsets register their own Pydantic model state under namespaces.
+    Multiple toolsets can share state by registering under the same namespace.
+
+    All registered states must be Pydantic BaseModel subclasses, making
+    the entire context serializable via model_dump()/model_validate().
+
+    Example:
+        # Define toolset-specific state
+        class SearchState(BaseModel):
+            results: list[SearchResult] = []
+            filter: str | None = None
+
+        SEARCH_NAMESPACE = "haiku.rag.search"
+
+        # In toolset factory
+        def create_search_toolset(client, config, context=None):
+            if context:
+                state = context.get_or_create(SEARCH_NAMESPACE, SearchState)
+            ...
+
+        # Usage
+        context = ToolContext()
+        search_tools = create_search_toolset(client, config, context=context)
+
+        agent = Agent(..., toolsets=[search_tools])
+        await agent.run("...")
+
+        # Access accumulated state
+        search_state = context.get(SEARCH_NAMESPACE)
+        for result in search_state.results:
+            print(f"{result.document_title}")
+
+        # Serialize entire context
+        ns_data = context.dump_namespaces()
+    """
+
+    _namespaces: dict[str, BaseModel] = PrivateAttr(default_factory=dict)
+
+    def register(self, namespace: str, state: BaseModel) -> None:
+        """Register state for a namespace.
+
+        Args:
+            namespace: Unique identifier for the toolset (e.g., "haiku.rag.search")
+            state: A Pydantic BaseModel instance to store
+
+        Overwrites any existing state for the namespace.
+        """
+        self._namespaces[namespace] = state
+
+    def get(self, namespace: str) -> BaseModel | None:
+        """Get state for a namespace, or None if not registered."""
+        return self._namespaces.get(namespace)
+
+    def get_or_create(self, namespace: str, factory: Callable[[], T]) -> T:
+        """Get state for a namespace, creating it if not registered.
+
+        Args:
+            namespace: The namespace to get or create state for.
+            factory: A callable that returns a new Pydantic model instance.
+
+        Returns:
+            The state for the namespace.
+        """
+        if namespace not in self._namespaces:
+            self._namespaces[namespace] = factory()
+        return self._namespaces[namespace]  # type: ignore[return-value]
+
+    def clear_namespace(self, namespace: str) -> None:
+        """Clear state for a specific namespace."""
+        if namespace in self._namespaces:
+            del self._namespaces[namespace]
+
+    def clear_all(self) -> None:
+        """Clear all namespaces."""
+        self._namespaces.clear()
+
+    @property
+    def namespaces(self) -> list[str]:
+        """List all registered namespaces."""
+        return list(self._namespaces.keys())
+
+    def dump_namespaces(self) -> dict[str, dict[str, Any]]:
+        """Serialize all namespace states to a dictionary.
+
+        Returns:
+            Dict mapping namespace -> serialized state dict.
+        """
+        return {ns: state.model_dump() for ns, state in self._namespaces.items()}
+
+    def load_namespace(self, namespace: str, state_type: type[T], data: dict) -> T:
+        """Deserialize and register state for a namespace.
+
+        Args:
+            namespace: The namespace to register the state under.
+            state_type: The Pydantic model class to deserialize into.
+            data: The serialized state data.
+
+        Returns:
+            The deserialized and registered state.
+        """
+        state = state_type.model_validate(data)
+        self._namespaces[namespace] = state
+        return state

haiku_rag_slim/haiku/rag/tools/filters.py

@@ -0,0 +1,38 @@
+def build_document_filter(document_name: str) -> str:
+    """Build SQL filter for document name matching.
+
+    Matches against both uri and title fields, case-insensitive.
+    Also matches without spaces to handle cases like "TB MED 593" vs "TBMED593".
+    """
+    escaped = document_name.replace("'", "''")
+    no_spaces = escaped.replace(" ", "")
+    return (
+        f"LOWER(uri) LIKE LOWER('%{escaped}%') OR LOWER(title) LIKE LOWER('%{escaped}%') "
+        f"OR LOWER(uri) LIKE LOWER('%{no_spaces}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')"
+    )
+
+
+def build_multi_document_filter(document_names: list[str]) -> str | None:
+    """Build SQL filter for multiple document names (OR combined).
+
+    Returns None if the list is empty.
+    """
+    if not document_names:
+        return None
+    filters = [build_document_filter(name) for name in document_names]
+    if len(filters) == 1:
+        return filters[0]
+    return " OR ".join(f"({f})" for f in filters)
+
+
+def combine_filters(filter1: str | None, filter2: str | None) -> str | None:
+    """Combine two SQL filters with AND logic.
+
+    Returns None if both filters are None.
+    """
+    filters = [f for f in [filter1, filter2] if f]
+    if not filters:
+        return None
+    if len(filters) == 1:
+        return filters[0]
+    return f"({filters[0]}) AND ({filters[1]})"

haiku_rag_slim/haiku/rag/tools/models.py

@@ -0,0 +1,41 @@
+from pydantic import BaseModel, Field
+
+from haiku.rag.agents.research.models import Citation
+
+
+class QAResult(BaseModel):
+    """Result from the QA toolset."""
+
+    question: str = Field(description="The question that was answered")
+    answer: str = Field(description="The answer to the question")
+    confidence: float = Field(
+        default=1.0,
+        description="Confidence score for this answer (0-1)",
+        ge=0.0,
+        le=1.0,
+    )
+    citations: list[Citation] = Field(
+        default_factory=list,
+        description="Citations supporting the answer",
+    )
+
+    @property
+    def sources(self) -> list[str]:
+        """Source names for display."""
+        return list(
+            dict.fromkeys(c.document_title or c.document_uri for c in self.citations)
+        )
+
+
+class AnalysisResult(BaseModel):
+    """Result from the analysis toolset (RLM execution)."""
+
+    answer: str = Field(description="The answer produced by analysis")
+    code_executed: bool = Field(
+        default=True,
+        description="Whether code was executed to produce this answer",
+    )
+    execution_count: int = Field(
+        default=0,
+        description="Number of code executions performed",
+    )

tests/agents/chat/test_chat_agent.py

@@ -737,7 +737,7 @@ async def test_search_agent_with_session_filter(allow_model_requests, temp_db_pa
             title="DocLayNet Sources",
         )
 
-        from haiku.rag.agents.chat.state import build_multi_document_filter
+        from haiku.rag.tools.filters import build_multi_document_filter
 
         search_agent = SearchAgent(client, Config)
 

tests/agents/chat/test_state.py

@@ -5,6 +5,8 @@ from haiku.rag.agents.chat.state import (
     ChatSessionState,
     QAResponse,
     SessionContext,
+)
+from haiku.rag.tools.filters import (
     build_document_filter,
     build_multi_document_filter,
     combine_filters,

tests/tools/test_context.py

@@ -0,0 +1,173 @@
+from pydantic import BaseModel
+
+from haiku.rag.tools.context import ToolContext
+
+
+class TestState(BaseModel):
+    value: int = 0
+
+
+class TestStateWithList(BaseModel):
+    items: list[str] = []
+
+
+def test_tool_context_defaults():
+    """Test ToolContext has sensible defaults."""
+    ctx = ToolContext()
+    assert ctx._namespaces == {}
+
+
+def test_register_and_get():
+    """Test register and get state for a namespace."""
+    ctx = ToolContext()
+    state = TestState(value=42)
+
+    ctx.register("test.namespace", state)
+
+    retrieved = ctx.get("test.namespace")
+    assert retrieved is state
+    assert retrieved.value == 42
+
+
+def test_get_nonexistent_namespace():
+    """Test get returns None for unregistered namespace."""
+    ctx = ToolContext()
+    assert ctx.get("nonexistent") is None
+
+
+def test_get_or_create_creates_new():
+    """Test get_or_create creates state when namespace doesn't exist."""
+    ctx = ToolContext()
+
+    state = ctx.get_or_create("test.namespace", TestStateWithList)
+    assert isinstance(state, TestStateWithList)
+    assert state.items == []
+
+
+def test_get_or_create_returns_existing():
+    """Test get_or_create returns existing state."""
+    ctx = ToolContext()
+
+    state1 = ctx.get_or_create("test.namespace", TestStateWithList)
+    state1.items.append("item1")
+
+    state2 = ctx.get_or_create("test.namespace", TestStateWithList)
+    assert state2 is state1
+    assert state2.items == ["item1"]
+
+
+def test_clear_namespace():
+    """Test clear_namespace removes only the specified namespace."""
+    ctx = ToolContext()
+    ctx.register("ns1", TestState(value=1))
+    ctx.register("ns2", TestState(value=2))
+
+    ctx.clear_namespace("ns1")
+
+    assert ctx.get("ns1") is None
+    ns2 = ctx.get("ns2")
+    assert isinstance(ns2, TestState)
+    assert ns2.value == 2
+
+
+def test_clear_namespace_nonexistent():
+    """Test clear_namespace handles nonexistent namespace gracefully."""
+    ctx = ToolContext()
+    ctx.clear_namespace("nonexistent")  # Should not raise
+
+
+def test_clear_all():
+    """Test clear_all clears all namespaces."""
+    ctx = ToolContext()
+    ctx.register("ns1", TestState(value=1))
+    ctx.register("ns2", TestState(value=2))
+
+    ctx.clear_all()
+
+    assert ctx.get("ns1") is None
+    assert ctx.get("ns2") is None
+
+
+def test_namespaces_property():
+    """Test namespaces property lists all registered namespaces."""
+    ctx = ToolContext()
+    assert ctx.namespaces == []
+
+    ctx.register("ns1", TestState())
+    ctx.register("ns2", TestState())
+
+    assert set(ctx.namespaces) == {"ns1", "ns2"}
+
+
+def test_shared_namespace_between_toolsets():
+    """Test that toolsets can share state via the same namespace."""
+
+    class SharedState(BaseModel):
+        citations: dict[str, int] = {}
+
+    SHARED_NAMESPACE = "haiku.rag.citations"
+
+    ctx = ToolContext()
+
+    # First toolset registers the shared state
+    state1 = ctx.get_or_create(SHARED_NAMESPACE, SharedState)
+    state1.citations["chunk-a"] = 1
+
+    # Second toolset gets the same state
+    state2 = ctx.get_or_create(SHARED_NAMESPACE, SharedState)
+    assert state2 is state1
+    assert state2.citations == {"chunk-a": 1}
+
+    # Both see updates
+    state2.citations["chunk-b"] = 2
+    assert state1.citations == {"chunk-a": 1, "chunk-b": 2}
+
+
+def test_dump_namespaces():
+    """Test dump_namespaces serializes all registered states."""
+    ctx = ToolContext()
+    ctx.register("ns1", TestState(value=10))
+    ctx.register("ns2", TestStateWithList(items=["a", "b"]))
+
+    data = ctx.dump_namespaces()
+
+    assert data == {
+        "ns1": {"value": 10},
+        "ns2": {"items": ["a", "b"]},
+    }
+
+
+def test_load_namespace():
+    """Test load_namespace deserializes and registers state."""
+    ctx = ToolContext()
+
+    state = ctx.load_namespace("ns1", TestState, {"value": 42})
+
+    assert isinstance(state, TestState)
+    assert state.value == 42
+    assert ctx.get("ns1") is state
+
+
+def test_serialization_roundtrip():
+    """Test full serialization/deserialization roundtrip."""
+    # Create and populate context
+    original = ToolContext()
+    original.register("search", TestStateWithList(items=["result1", "result2"]))
+    original.register("qa", TestState(value=99))
+
+    # Serialize
+    ns_data = original.dump_namespaces()
+
+    # Deserialize
+    restored = ToolContext()
+    restored.load_namespace("search", TestStateWithList, ns_data["search"])
+    restored.load_namespace("qa", TestState, ns_data["qa"])
+
+    # Verify
+    search_state = restored.get("search")
+    assert isinstance(search_state, TestStateWithList)
+    assert search_state.items == ["result1", "result2"]
+
+    qa_state = restored.get("qa")
+    assert isinstance(qa_state, TestState)
+    assert qa_state.value == 99

tests/tools/test_filters.py

@@ -0,0 +1,79 @@
+from haiku.rag.tools.filters import (
+    build_document_filter,
+    build_multi_document_filter,
+    combine_filters,
+)
+
+
+def test_build_document_filter_simple():
+    """Test build_document_filter with simple name."""
+    result = build_document_filter("mytest")
+    assert "LOWER(uri) LIKE LOWER('%mytest%')" in result
+    assert "LOWER(title) LIKE LOWER('%mytest%')" in result
+
+
+def test_build_document_filter_with_spaces():
+    """Test build_document_filter handles spaces correctly."""
+    result = build_document_filter("TB MED 593")
+    # Should include both the original (with spaces) and without spaces
+    assert "LOWER(uri) LIKE LOWER('%TB MED 593%')" in result
+    assert "LOWER(uri) LIKE LOWER('%TBMED593%')" in result
+    assert "LOWER(title) LIKE LOWER('%TB MED 593%')" in result
+    assert "LOWER(title) LIKE LOWER('%TBMED593%')" in result
+
+
+def test_build_document_filter_escapes_quotes():
+    """Test build_document_filter escapes single quotes."""
+    result = build_document_filter("O'Reilly")
+    # Single quotes should be doubled for SQL escaping
+    assert "O''Reilly" in result
+
+
+def test_build_multi_document_filter_empty():
+    """Test build_multi_document_filter returns None for empty list."""
+    result = build_multi_document_filter([])
+    assert result is None
+
+
+def test_build_multi_document_filter_single():
+    """Test build_multi_document_filter with single document."""
+    result = build_multi_document_filter(["mytest"])
+    assert result is not None
+    assert "LOWER(uri) LIKE LOWER('%mytest%')" in result
+    assert "LOWER(title) LIKE LOWER('%mytest%')" in result
+    # Single document should not have extra wrapping parentheses
+    assert " OR (" not in result
+
+
+def test_build_multi_document_filter_multiple():
+    """Test build_multi_document_filter with multiple documents."""
+    result = build_multi_document_filter(["doc1", "doc2"])
+    assert result is not None
+    # Should have OR-combined filters
+    assert "doc1" in result
+    assert "doc2" in result
+    assert " OR (" in result
+
+
+def test_combine_filters_both_none():
+    """Test combine_filters with both None."""
+    result = combine_filters(None, None)
+    assert result is None
+
+
+def test_combine_filters_first_only():
+    """Test combine_filters with only first filter."""
+    result = combine_filters("uri = 'test'", None)
+    assert result == "uri = 'test'"
+
+
+def test_combine_filters_second_only():
+    """Test combine_filters with only second filter."""
+    result = combine_filters(None, "title = 'doc'")
+    assert result == "title = 'doc'"
+
+
+def test_combine_filters_both():
+    """Test combine_filters combines with AND."""
+    result = combine_filters("uri = 'test'", "title = 'doc'")
+    assert result == "(uri = 'test') AND (title = 'doc')"

tests/tools/test_models.py

@@ -0,0 +1,99 @@
+from haiku.rag.agents.research.models import Citation
+from haiku.rag.tools.models import AnalysisResult, QAResult
+
+
+def test_qa_result_defaults():
+    """Test QAResult has sensible defaults."""
+    result = QAResult(question="What is X?", answer="X is Y.")
+    assert result.confidence == 1.0
+    assert result.citations == []
+
+
+def test_qa_result_with_citations():
+    """Test QAResult with citations."""
+    citation = Citation(
+        index=1,
+        document_id="doc-1",
+        chunk_id="chunk-1",
+        document_uri="test.md",
+        document_title="Test Doc",
+        content="Citation content",
+    )
+    result = QAResult(
+        question="What is X?",
+        answer="X is Y.",
+        confidence=0.9,
+        citations=[citation],
+    )
+    assert result.confidence == 0.9
+    assert len(result.citations) == 1
+    assert result.citations[0].document_title == "Test Doc"
+
+
+def test_qa_result_sources_property():
+    """Test QAResult.sources returns unique source names."""
+    citations = [
+        Citation(
+            document_id="doc-1",
+            chunk_id="chunk-1",
+            document_uri="doc1.md",
+            document_title="Document One",
+            content="Content 1",
+        ),
+        Citation(
+            document_id="doc-1",
+            chunk_id="chunk-2",
+            document_uri="doc1.md",
+            document_title="Document One",
+            content="Content 2",
+        ),
+        Citation(
+            document_id="doc-2",
+            chunk_id="chunk-3",
+            document_uri="doc2.md",
+            document_title="Document Two",
+            content="Content 3",
+        ),
+    ]
+    result = QAResult(question="Q", answer="A", citations=citations)
+
+    sources = result.sources
+    assert len(sources) == 2
+    assert "Document One" in sources
+    assert "Document Two" in sources
+
+
+def test_qa_result_sources_uses_uri_as_fallback():
+    """Test QAResult.sources uses uri when title is None."""
+    citations = [
+        Citation(
+            document_id="doc-1",
+            chunk_id="chunk-1",
+            document_uri="test.md",
+            document_title=None,
+            content="Content",
+        ),
+    ]
+    result = QAResult(question="Q", answer="A", citations=citations)
+
+    sources = result.sources
+    assert sources == ["test.md"]
+
+
+def test_analysis_result_defaults():
+    """Test AnalysisResult has sensible defaults."""
+    result = AnalysisResult(answer="The result is 42")
+    assert result.code_executed is True
+    assert result.execution_count == 0
+
+
+def test_analysis_result_with_values():
+    """Test AnalysisResult with explicit values."""
+    result = AnalysisResult(
+        answer="The result is 42",
+        code_executed=True,
+        execution_count=3,
+    )
+    assert result.answer == "The result is 42"
+    assert result.code_executed is True
+    assert result.execution_count == 3