diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index 05163962..e06cb860 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/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", ] diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py index 01efb524..2a55bfff 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/state.py +++ b/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, - ) diff --git a/haiku_rag_slim/haiku/rag/tools/__init__.py b/haiku_rag_slim/haiku/rag/tools/__init__.py index d2a89bd2..fb30f93e 100644 --- a/haiku_rag_slim/haiku/rag/tools/__init__.py +++ b/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", diff --git a/haiku_rag_slim/haiku/rag/tools/analysis.py b/haiku_rag_slim/haiku/rag/tools/analysis.py index 8e987bdd..531d1c9c 100644 --- a/haiku_rag_slim/haiku/rag/tools/analysis.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/tools/document.py b/haiku_rag_slim/haiku/rag/tools/document.py index c96a6d0d..14fc52af 100644 --- a/haiku_rag_slim/haiku/rag/tools/document.py +++ b/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.session import SESSION_NAMESPACE, SessionState +from haiku.rag.tools.filters import get_session_filter 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 - 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) + effective_filter = get_session_filter(context, base_filter) docs = await client.list_documents( limit=page_size, offset=offset, filter=effective_filter diff --git a/haiku_rag_slim/haiku/rag/tools/filters.py b/haiku_rag_slim/haiku/rag/tools/filters.py index b09d14db..f7b1a729 100644 --- a/haiku_rag_slim/haiku/rag/tools/filters.py +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/tools/qa.py b/haiku_rag_slim/haiku/rag/tools/qa.py index 4385ed3c..bc0bce1f 100644 --- a/haiku_rag_slim/haiku/rag/tools/qa.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index 2e9ba7c2..17ff1f88 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/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 diff --git a/tests/agents/chat/test_chat_agent.py b/tests/agents/chat/test_chat_agent.py index 6cdf096a..656c52b6 100644 --- a/tests/agents/chat/test_chat_agent.py +++ b/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"]: diff --git a/tests/agents/chat/test_state.py b/tests/agents/chat/test_state.py index 91ee2952..2b633428 100644 --- a/tests/agents/chat/test_state.py +++ 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 ( - 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')" +from haiku.rag.tools.session import SessionState 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.get_or_assign_index("chunk-b") + original.citation_registry = {"chunk-a": 1, "chunk-b": 2} # 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) - - # 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 + assert restored.citation_registry == {"chunk-a": 1, "chunk-b": 2} 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 diff --git a/tests/tools/test_filters.py b/tests/tools/test_filters.py index de449dd8..e524e702 100644 --- a/tests/tools/test_filters.py +++ b/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