diff --git a/haiku_rag_slim/haiku/rag/agents/chat/__init__.py b/haiku_rag_slim/haiku/rag/agents/chat/__init__.py index cfecef98..cde34e81 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/__init__.py +++ b/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", ] diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index c1081f57..140711ce 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py index d8ef026a..21ee0496 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/state.py +++ b/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, diff --git a/haiku_rag_slim/haiku/rag/tools/__init__.py b/haiku_rag_slim/haiku/rag/tools/__init__.py new file mode 100644 index 00000000..4dcfca6e --- /dev/null +++ b/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", +] diff --git a/haiku_rag_slim/haiku/rag/tools/context.py b/haiku_rag_slim/haiku/rag/tools/context.py new file mode 100644 index 00000000..450800e8 --- /dev/null +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/tools/filters.py b/haiku_rag_slim/haiku/rag/tools/filters.py new file mode 100644 index 00000000..b09d14db --- /dev/null +++ b/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]})" diff --git a/haiku_rag_slim/haiku/rag/tools/models.py b/haiku_rag_slim/haiku/rag/tools/models.py new file mode 100644 index 00000000..d911a338 --- /dev/null +++ b/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", + ) diff --git a/tests/agents/chat/test_chat_agent.py b/tests/agents/chat/test_chat_agent.py index 41f2f7d9..de1f4989 100644 --- a/tests/agents/chat/test_chat_agent.py +++ b/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) diff --git a/tests/agents/chat/test_state.py b/tests/agents/chat/test_state.py index 30ace7ab..788ead92 100644 --- a/tests/agents/chat/test_state.py +++ b/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, diff --git a/tests/tools/__init__.py b/tests/tools/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/tools/test_context.py b/tests/tools/test_context.py new file mode 100644 index 00000000..2adf2b24 --- /dev/null +++ b/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 diff --git a/tests/tools/test_filters.py b/tests/tools/test_filters.py new file mode 100644 index 00000000..de449dd8 --- /dev/null +++ b/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')" diff --git a/tests/tools/test_models.py b/tests/tools/test_models.py new file mode 100644 index 00000000..65a83a78 --- /dev/null +++ b/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