Extract get_session_filter helper, remove dead code and duplicate tests

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

@@ -15,7 +15,6 @@ from haiku.rag.agents.chat.state import (
     AGUI_STATE_KEY,
     ChatSessionState,
     SessionContext,
-    emit_state_event,
 )
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig
@@ -273,6 +272,5 @@ __all__ = [
     "ChatDeps",
     "ChatSessionState",
     "SessionContext",
-    "emit_state_event",
     "AGUI_STATE_KEY",
 ]

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

@@ -1,8 +1,6 @@
 from datetime import datetime
 from typing import TYPE_CHECKING
 
-import jsonpatch
-from ag_ui.core import EventType, StateDeltaEvent
 from pydantic import BaseModel, Field
 
 from haiku.rag.agents.research.models import Citation
@@ -35,20 +33,6 @@ class ChatSessionState(BaseModel):
     document_filter: list[str] = []
     citation_registry: dict[str, int] = {}
 
-    def get_or_assign_index(self, chunk_id: str) -> int:
-        """Get or assign a stable citation index for a chunk_id.
-
-        Citation indices persist across tool calls within a session.
-        The first chunk gets index 1, subsequent new chunks get incrementing indices.
-        Same chunk_id always returns the same index.
-        """
-        if chunk_id in self.citation_registry:
-            return self.citation_registry[chunk_id]
-
-        new_index = len(self.citation_registry) + 1
-        self.citation_registry[chunk_id] = new_index
-        return new_index
-
 
 def _rebuild_models(qa_history_entry_cls: type) -> None:
     """Resolve ChatSessionState forward reference to QAHistoryEntry.
@@ -58,26 +42,3 @@ def _rebuild_models(qa_history_entry_cls: type) -> None:
     ChatSessionState.model_rebuild(
         _types_namespace={"QAHistoryEntry": qa_history_entry_cls}
     )
-
-
-def emit_state_event(
-    current_state: ChatSessionState,
-    new_state: ChatSessionState,
-    state_key: str | None = None,
-) -> StateDeltaEvent | None:
-    """Emit state delta against current state, or None if no changes."""
-    new_snapshot = new_state.model_dump(mode="json")
-    wrapped_new = {state_key: new_snapshot} if state_key else new_snapshot
-
-    current_snapshot = current_state.model_dump(mode="json")
-    wrapped_current = {state_key: current_snapshot} if state_key else current_snapshot
-
-    patch = jsonpatch.make_patch(wrapped_current, wrapped_new)
-
-    if not patch.patch:
-        return None
-
-    return StateDeltaEvent(
-        type=EventType.STATE_DELTA,
-        delta=patch.patch,
-    )

haiku_rag_slim/haiku/rag/tools/__init__.py

@@ -16,6 +16,7 @@ from haiku.rag.tools.filters import (
     build_document_filter,
     build_multi_document_filter,
     combine_filters,
+    get_session_filter,
 )
 from haiku.rag.tools.models import AnalysisResult, QAResult
 from haiku.rag.tools.qa import (
@@ -40,6 +41,7 @@ __all__ = [
     "build_document_filter",
     "build_multi_document_filter",
     "combine_filters",
+    "get_session_filter",
     "SEARCH_NAMESPACE",
     "SearchState",
     "create_search_toolset",

haiku_rag_slim/haiku/rag/tools/analysis.py

@@ -9,11 +9,10 @@ from haiku.rag.config.models import AppConfig
 from haiku.rag.tools.context import ToolContext
 from haiku.rag.tools.filters import (
     build_document_filter,
-    build_multi_document_filter,
     combine_filters,
+    get_session_filter,
 )
 from haiku.rag.tools.models import AnalysisResult
-from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
 
 ANALYSIS_NAMESPACE = "haiku.rag.analysis"
 
@@ -70,19 +69,10 @@ def create_analysis_toolset(
         Returns:
             AnalysisResult with answer and execution metadata.
         """
-        # Get session filter from session state
-        session_filter = None
-        if context is not None:
-            session_state = context.get_typed(SESSION_NAMESPACE, SessionState)
-            if session_state is not None and session_state.document_filter:
-                session_filter = build_multi_document_filter(
-                    session_state.document_filter
-                )
-
         # Build filter from base_filter, session_filter, and document_name
         doc_filter = build_document_filter(document_name) if document_name else None
         effective_filter = combine_filters(
-            combine_filters(base_filter, session_filter), doc_filter
+            get_session_filter(context, base_filter), doc_filter
         )
 
         # Create RLM context

haiku_rag_slim/haiku/rag/tools/document.py

@@ -4,8 +4,7 @@ from pydantic_ai import Agent, FunctionToolset
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig
 from haiku.rag.tools.context import ToolContext
-from haiku.rag.tools.filters import build_multi_document_filter, combine_filters
+from haiku.rag.tools.filters import get_session_filter
-from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
 from haiku.rag.utils import get_model
 
 DOCUMENT_NAMESPACE = "haiku.rag.document"
@@ -117,16 +116,7 @@ def create_document_toolset(
         page_size = 50
         offset = (page - 1) * page_size
 
-        # Get session filter from session state
+        effective_filter = get_session_filter(context, base_filter)
-        session_filter = None
-        if context is not None:
-            session_state = context.get_typed(SESSION_NAMESPACE, SessionState)
-            if session_state is not None and session_state.document_filter:
-                session_filter = build_multi_document_filter(
-                    session_state.document_filter
-                )
-
-        effective_filter = combine_filters(base_filter, session_filter)
 
         docs = await client.list_documents(
             limit=page_size, offset=offset, filter=effective_filter

haiku_rag_slim/haiku/rag/tools/filters.py

@@ -1,3 +1,9 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from haiku.rag.tools.context import ToolContext
+
+
 def build_document_filter(document_name: str) -> str:
     """Build SQL filter for document name matching.
 
@@ -25,6 +31,35 @@ def build_multi_document_filter(document_names: list[str]) -> str | None:
     return " OR ".join(f"({f})" for f in filters)
 
 
+def get_session_filter(
+    context: "ToolContext | None",
+    base_filter: str | None = None,
+) -> str | None:
+    """Build effective filter from session state document filter and base filter.
+
+    Checks the ToolContext for a registered SessionState. If it has a
+    document_filter, builds a SQL filter from it and combines with base_filter.
+
+    Args:
+        context: Optional ToolContext that may contain a SessionState.
+        base_filter: Optional base SQL WHERE clause to combine with.
+
+    Returns:
+        Combined filter string, or None if no filters apply.
+    """
+    if context is None:
+        return base_filter
+
+    from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
+
+    session_state = context.get_typed(SESSION_NAMESPACE, SessionState)
+    if session_state is None or not session_state.document_filter:
+        return base_filter
+
+    session_filter = build_multi_document_filter(session_state.document_filter)
+    return combine_filters(base_filter, session_filter)
+
+
 def combine_filters(filter1: str | None, filter2: str | None) -> str | None:
     """Combine two SQL filters with AND logic.
 

haiku_rag_slim/haiku/rag/tools/qa.py

@@ -17,7 +17,11 @@ from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig
 from haiku.rag.embeddings import get_embedder
 from haiku.rag.tools.context import ToolContext
-from haiku.rag.tools.filters import build_document_filter, combine_filters
+from haiku.rag.tools.filters import (
+    build_document_filter,
+    combine_filters,
+    get_session_filter,
+)
 from haiku.rag.tools.models import QAResult
 from haiku.rag.tools.session import (
     SESSION_NAMESPACE,
@@ -177,15 +181,9 @@ def create_qa_toolset(
                         old_state_snapshot["session_context"] = None
 
         # Build filter from session state, base_filter, and document_name
-        session_filter = None
-        if session_state is not None and session_state.document_filter:
-            from haiku.rag.tools.filters import build_multi_document_filter
-
-            session_filter = build_multi_document_filter(session_state.document_filter)
-
         doc_filter = build_document_filter(document_name) if document_name else None
         effective_filter = combine_filters(
-            combine_filters(base_filter, session_filter), doc_filter
+            get_session_filter(context, base_filter), doc_filter
         )
 
         # Determine session context

haiku_rag_slim/haiku/rag/tools/search.py

@@ -6,7 +6,7 @@ from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig
 from haiku.rag.store.models import SearchResult
 from haiku.rag.tools.context import ToolContext
-from haiku.rag.tools.filters import build_multi_document_filter, combine_filters
+from haiku.rag.tools.filters import combine_filters, get_session_filter
 from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState, compute_state_delta
 
 SEARCH_NAMESPACE = "haiku.rag.search"
@@ -75,14 +75,9 @@ def create_search_toolset(
             if session_state is not None:
                 old_session_state = session_state.model_copy(deep=True)
 
-        # Build session filter from session state's document_filter
-        session_filter = None
-        if session_state is not None and session_state.document_filter:
-            session_filter = build_multi_document_filter(session_state.document_filter)
-
         # Combine all filters: base_filter AND session_filter AND tool filter
         effective_filter = combine_filters(
-            combine_filters(base_filter, session_filter), filter
+            get_session_filter(context, base_filter), filter
         )
 
         effective_limit = limit or config.search.limit

tests/agents/chat/test_chat_agent.py

@@ -717,7 +717,7 @@ def test_ask_tool_citation_registry_logic():
     3. Same chunk_id always gets same index
     4. Indices don't reset between calls
     """
-    session_state = ChatSessionState(session_id="test-registry")
+    session_state = SessionState()
 
     # Simulate first ask tool building citations
     first_ask_chunks = ["chunk-a", "chunk-b"]
@@ -768,7 +768,7 @@ def test_search_tool_citation_registry_logic():
     Verifies that search and ask tools share the same registry,
     maintaining stable indices across different tool calls.
     """
-    session_state = ChatSessionState(session_id="test-registry")
+    session_state = SessionState()
 
     # Simulate ask tool first (assigns indices 1, 2)
     for chunk_id in ["chunk-a", "chunk-b"]:

tests/agents/chat/test_state.py

@@ -1,89 +1,8 @@
-from ag_ui.core import StateDeltaEvent
-
 from haiku.rag.agents.chat.state import (
     ChatSessionState,
     SessionContext,
 )
-from haiku.rag.tools.filters import (
+from haiku.rag.tools.session import SessionState
-    build_document_filter,
-    build_multi_document_filter,
-    combine_filters,
-)
-from haiku.rag.tools.qa import QAHistoryEntry
-
-
-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')"
 
 
 def test_max_qa_history_constant():
@@ -101,7 +20,7 @@ def test_citation_registry_index_assignment():
     - Second unique chunk gets index 2
     - Same chunk_id always returns same index
     """
-    session_state = ChatSessionState(session_id="test")
+    session_state = SessionState()
 
     # First chunk gets index 1
     index1 = session_state.get_or_assign_index("chunk-abc")
@@ -118,7 +37,7 @@ def test_citation_registry_index_assignment():
 
 def test_citation_registry_stability():
     """Test citation indices are stable across multiple calls in any order."""
-    session_state = ChatSessionState(session_id="test")
+    session_state = SessionState()
 
     # First round assigns indices 1, 2, 3
     idx_a = session_state.get_or_assign_index("chunk-a")
@@ -139,8 +58,7 @@ def test_citation_registry_serialization_roundtrip():
     """Test citation_registry serializes and deserializes correctly for AG-UI state."""
     # Create state and assign indices
     original = ChatSessionState(session_id="test")
-    original.get_or_assign_index("chunk-a")
+    original.citation_registry = {"chunk-a": 1, "chunk-b": 2}
-    original.get_or_assign_index("chunk-b")
 
     # Serialize
     state_dict = original.model_dump()
@@ -149,12 +67,7 @@ def test_citation_registry_serialization_roundtrip():
 
     # Deserialize (simulating AG-UI state restoration)
     restored = ChatSessionState.model_validate(state_dict)
-
+    assert restored.citation_registry == {"chunk-a": 1, "chunk-b": 2}
-    # Existing chunks should return their persisted indices
-    assert restored.get_or_assign_index("chunk-a") == 1
-    assert restored.get_or_assign_index("chunk-b") == 2
-    # New chunk should get next index
-    assert restored.get_or_assign_index("chunk-c") == 3
 
 
 def test_chat_session_state_defaults_to_empty_session_id():
@@ -221,109 +134,3 @@ def test_chat_session_state_model_dump_json_serializes_datetime():
     # datetime should be serialized as ISO string, not datetime object
     assert isinstance(snapshot["session_context"]["last_updated"], str)
     assert snapshot["session_context"]["last_updated"] == "2025-01-27T12:00:00"
-
-
-def test_emit_state_event_includes_session_id_when_assigned():
-    """emit_state_event detects session_id change from empty to UUID.
-
-    When session_id defaults to "" and the tool assigns a UUID,
-    the delta must include session_id so clients can persist it.
-    """
-    from haiku.rag.agents.chat.state import emit_state_event
-
-    current_state = ChatSessionState()  # session_id=""
-    new_state = ChatSessionState(session_id="assigned-uuid-123")
-
-    event = emit_state_event(current_state, new_state)
-
-    assert event is not None
-    session_id_op = next(
-        (op for op in event.delta if op["path"] == "/session_id"), None
-    )
-    assert session_id_op is not None
-    assert session_id_op["value"] == "assigned-uuid-123"
-
-
-def test_emit_state_event_returns_none_when_no_changes():
-    """emit_state_event returns None when states are identical."""
-    from haiku.rag.agents.chat.state import emit_state_event
-
-    state = ChatSessionState(session_id="test-123", qa_history=[], citations=[])
-
-    event = emit_state_event(state, state)
-
-    assert event is None
-
-
-def test_emit_state_event_returns_delta_with_changes():
-    """emit_state_event returns StateDeltaEvent with JSON Patch ops for changes."""
-    from ag_ui.core import EventType, StateDeltaEvent
-
-    from haiku.rag.agents.chat.state import emit_state_event
-
-    current_state = ChatSessionState(session_id="test-123", qa_history=[], citations=[])
-    new_state = ChatSessionState(
-        session_id="test-123",
-        qa_history=[QAHistoryEntry(question="Q1", answer="A1", confidence=0.9)],
-        citations=[],
-    )
-
-    event = emit_state_event(current_state, new_state)
-
-    assert isinstance(event, StateDeltaEvent)
-    assert event.type == EventType.STATE_DELTA
-    assert len(event.delta) > 0
-    # Delta should contain an "add" operation for the new qa_history entry
-    ops = event.delta
-    qa_history_op = next((op for op in ops if "/qa_history" in op["path"]), None)
-    assert qa_history_op is not None
-
-
-def test_emit_state_event_delta_with_state_key():
-    """emit_state_event wraps delta paths with state_key namespace."""
-    from ag_ui.core import StateDeltaEvent
-
-    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, emit_state_event
-
-    current_state = ChatSessionState(session_id="test-123", qa_history=[])
-    new_state = ChatSessionState(
-        session_id="test-123",
-        qa_history=[QAHistoryEntry(question="Q1", answer="A1", confidence=0.9)],
-    )
-
-    event = emit_state_event(current_state, new_state, state_key=AGUI_STATE_KEY)
-
-    assert isinstance(event, StateDeltaEvent)
-    # Paths should be namespaced under state_key
-    for op in event.delta:
-        assert op["path"].startswith(f"/{AGUI_STATE_KEY}")
-
-
-def test_emit_state_event_delta_produces_valid_patch():
-    """emit_state_event delta can be applied to reproduce new state."""
-    import jsonpatch
-
-    from haiku.rag.agents.chat.state import emit_state_event
-
-    current_state = ChatSessionState(
-        session_id="test-123",
-        qa_history=[QAHistoryEntry(question="Q1", answer="A1", confidence=0.9)],
-        citations=[],
-    )
-    new_state = ChatSessionState(
-        session_id="test-123",
-        qa_history=[
-            QAHistoryEntry(question="Q1", answer="A1", confidence=0.9),
-            QAHistoryEntry(question="Q2", answer="A2", confidence=0.8),
-        ],
-        citations=[],
-    )
-
-    event = emit_state_event(current_state, new_state)
-    assert isinstance(event, StateDeltaEvent)
-
-    # Apply patch to current state and verify it produces new state
-    current_snapshot = current_state.model_dump(mode="json")
-    patched = jsonpatch.apply_patch(current_snapshot, event.delta)
-    new_snapshot = new_state.model_dump(mode="json")
-    assert patched == new_snapshot

tests/tools/test_filters.py

@@ -1,8 +1,11 @@
+from haiku.rag.tools.context import ToolContext
 from haiku.rag.tools.filters import (
     build_document_filter,
     build_multi_document_filter,
     combine_filters,
+    get_session_filter,
 )
+from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
 
 
 def test_build_document_filter_simple():
@@ -77,3 +80,37 @@ 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')"
+
+
+def test_get_session_filter_no_context():
+    """Returns base_filter as-is when context is None."""
+    assert get_session_filter(None) is None
+    assert get_session_filter(None, "uri = 'test'") == "uri = 'test'"
+
+
+def test_get_session_filter_no_document_filter():
+    """Returns base_filter when SessionState has no document_filter."""
+    context = ToolContext()
+    context.register(SESSION_NAMESPACE, SessionState())
+    assert get_session_filter(context) is None
+    assert get_session_filter(context, "uri = 'test'") == "uri = 'test'"
+
+
+def test_get_session_filter_with_document_filter():
+    """Builds filter from SessionState.document_filter."""
+    context = ToolContext()
+    context.register(SESSION_NAMESPACE, SessionState(document_filter=["mytest"]))
+    result = get_session_filter(context)
+    assert result is not None
+    assert "mytest" in result
+
+
+def test_get_session_filter_combines_with_base_filter():
+    """Combines session filter with base_filter using AND."""
+    context = ToolContext()
+    context.register(SESSION_NAMESPACE, SessionState(document_filter=["mytest"]))
+    result = get_session_filter(context, "uri = 'base'")
+    assert result is not None
+    assert "uri = 'base'" in result
+    assert "mytest" in result
+    assert "AND" in result