danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Extract get_session_filter helper, remove dead code and duplicate tests

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-06 14:53:26 +01:00
parent feee458d97
commit bd3b26f78b
No known key found for this signature in database
11 changed files with 93 additions and 280 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
haiku_rag_slim/haiku/rag/agents/chat/agent.py
View file

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

39
haiku_rag_slim/haiku/rag/agents/chat/state.py
View file

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

2
haiku_rag_slim/haiku/rag/tools/__init__.py
View file

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

14
haiku_rag_slim/haiku/rag/tools/analysis.py
View file

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

14
haiku_rag_slim/haiku/rag/tools/document.py
View file

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

35
haiku_rag_slim/haiku/rag/tools/filters.py
View file

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

14
haiku_rag_slim/haiku/rag/tools/qa.py
View file

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

9
haiku_rag_slim/haiku/rag/tools/search.py
View file

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

4
tests/agents/chat/test_chat_agent.py
View file

@ -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"]:

203
tests/agents/chat/test_state.py
View file

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

37
tests/tools/test_filters.py
View file

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