From 330468f2ca206dfc137c7244460e27fd32023d48 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 13 Jul 2026 15:56:26 +0300 Subject: [PATCH] Add client lifecycle hooks (after_ingest, after_delete, before_search, after_search) --- CHANGELOG.md | 4 + haiku_rag_slim/haiku/rag/client/__init__.py | 5 + haiku_rag_slim/haiku/rag/client/documents.py | 10 + haiku_rag_slim/haiku/rag/client/search.py | 6 + haiku_rag_slim/haiku/rag/config/models.py | 1 + haiku_rag_slim/haiku/rag/hooks.py | 89 ++++++ tests/test_hooks.py | 274 +++++++++++++++++++ 7 files changed, 389 insertions(+) create mode 100644 haiku_rag_slim/haiku/rag/hooks.py create mode 100644 tests/test_hooks.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a43776c..9ab6d844 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## [Unreleased] +### Added + +- Client lifecycle hooks (`after_ingest`, `after_delete`, `before_search`, `after_search`) registered under the `haiku.rag.hooks` entry-point group and activated via the `hooks:` config list. + ## [0.78.0] - 2026-08-24 ### Added diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index e97dd695..88063fbc 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -17,6 +17,7 @@ import httpx from haiku.rag.client.documents import DocumentImport from haiku.rag.config import AppConfig, get_config from haiku.rag.converters import get_converter +from haiku.rag.hooks import build_hooks, load_hooks from haiku.rag.reranking import get_reranker from haiku.rag.store.engine import Store from haiku.rag.store.models.chunk import Chunk, SearchResult, SearchType @@ -93,6 +94,7 @@ class HaikuRAG: self._vacuum_tasks: set[asyncio.Task] = set() self._last_vacuum_at: float | None = None self._vacuum_dirty = False + self._hooks = build_hooks(config.hooks, load_hooks()) if config.hooks else [] @property def is_read_only(self) -> bool: @@ -420,6 +422,9 @@ class HaikuRAG: if self._config.storage.auto_vacuum: self._schedule_vacuum() + for doc_id in ids_to_delete: + for hook in self._hooks: + await hook.after_delete(self, doc_id) return True async def list_documents( diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index d049a723..d6c55947 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -161,6 +161,9 @@ async def _store_document_with_chunks( if client._config.storage.auto_vacuum: client._schedule_vacuum() + + for hook in client._hooks: + await hook.after_ingest(client, stored_doc) return stored_doc @@ -212,6 +215,9 @@ async def _update_document_with_chunks( if client._config.storage.auto_vacuum: client._schedule_vacuum() + + for hook in client._hooks: + await hook.after_ingest(client, updated_doc) return updated_doc @@ -314,6 +320,10 @@ async def _store_documents_with_chunks( if client._config.storage.auto_vacuum: client._schedule_vacuum() + + for doc in created: + for hook in client._hooks: + await hook.after_ingest(client, doc) return created diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index e2f169da..13c363d1 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -42,6 +42,9 @@ async def search( if search_type is None: search_type = "hybrid" + for hook in client._hooks: + query, filter = await hook.before_search(client, query, filter) + reranker = client.reranker if reranker is None: @@ -79,6 +82,9 @@ async def search( if include_images: await _populate_image_data(client, results) + for hook in client._hooks: + results = await hook.after_search(client, query, results) + return results diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 181f70b9..289ab2df 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -615,6 +615,7 @@ class IngesterConfig(ConfigModel): class AppConfig(ConfigModel): environment: str = "production" + hooks: list[str] = Field(default_factory=list) storage: StorageConfig = Field(default_factory=StorageConfig) lancedb: LanceDBConfig = Field(default_factory=LanceDBConfig) embeddings: EmbeddingsConfig = Field(default_factory=EmbeddingsConfig) diff --git a/haiku_rag_slim/haiku/rag/hooks.py b/haiku_rag_slim/haiku/rag/hooks.py new file mode 100644 index 00000000..a8eeb5f7 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/hooks.py @@ -0,0 +1,89 @@ +from collections.abc import Callable, Mapping, Sequence +from importlib.metadata import entry_points +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +if TYPE_CHECKING: + from PIL import Image as PILImage + + from haiku.rag.client import HaikuRAG + from haiku.rag.store.models.chunk import SearchResult + from haiku.rag.store.models.document import Document + +ENTRY_POINT_GROUP = "haiku.rag.hooks" + + +class Hook: + """Base class for client lifecycle hooks. Subclasses override any subset. + + A package registers a zero-arg factory under the ``haiku.rag.hooks`` + entry-point group; ``config.hooks`` lists the hooks to activate, and they + run in the listed order at every hook point. Hooks receive the ``HaikuRAG`` + client, so they may search, read repositories, or keep their own state in + the database via ``client.store`` (table names must use the ``hook_`` + prefix to stay clear of core tables and migrations). + """ + + async def after_ingest(self, client: "HaikuRAG", document: "Document") -> None: + """A document's content was written (create, import, batch import, + update). Replace any state derived from this document: create and + update are deliberately the same event. Metadata/title-only updates + do not fire.""" + + async def after_delete(self, client: "HaikuRAG", document_id: str) -> None: + """A document was deleted; fires once per document in a cascade.""" + + async def before_search( + self, client: "HaikuRAG", query: str, filter: str | None + ) -> tuple[str, str | None]: + """Transform the query and/or filter before retrieval. Text queries + only; the returned query feeds both the vector and FTS sides.""" + return query, filter + + async def after_search( + self, + client: "HaikuRAG", + query: "str | bytes | PILImage.Image", + results: "list[SearchResult]", + ) -> "list[SearchResult]": + """Transform or annotate search results before they are returned.""" + return results + + +HookFactory = Callable[[], Hook] + + +@runtime_checkable +class LoadableEntryPoint(Protocol): + """The slice of ``importlib.metadata.EntryPoint`` ``build_hooks`` needs: + a deferred ``load()`` returning the hook factory.""" + + def load(self) -> HookFactory: ... + + +def load_hooks() -> dict[str, LoadableEntryPoint]: + """Discover registered hook entry points, keyed by name. The entry points + are not imported here; ``build_hooks`` loads only the ones the config + references, so an unused hook with a broken import does not fail client + construction.""" + return {ep.name: ep for ep in entry_points(group=ENTRY_POINT_GROUP)} + + +def build_hooks( + names: Sequence[str], + discovered: Mapping[str, LoadableEntryPoint], +) -> list[Hook]: + """Load and instantiate the named hooks in configured order. Raises + ValueError for a name with no registered entry point so a misconfigured + client fails at construction rather than silently skipping a hook.""" + hooks: list[Hook] = [] + for name in names: + try: + entry_point = discovered[name] + except KeyError: + raise ValueError( + f"Config references unknown hook {name!r}; no entry point " + f"registered under {ENTRY_POINT_GROUP!r}." + ) from None + factory: HookFactory = entry_point.load() + hooks.append(factory()) + return hooks diff --git a/tests/test_hooks.py b/tests/test_hooks.py new file mode 100644 index 00000000..45696357 --- /dev/null +++ b/tests/test_hooks.py @@ -0,0 +1,274 @@ +import pytest + +from haiku.rag.client import HaikuRAG +from haiku.rag.config import Config +from haiku.rag.hooks import ENTRY_POINT_GROUP, Hook, build_hooks +from haiku.rag.store.models.chunk import Chunk +from tests.test_client import _docling_doc, _import + + +class RecordingHook(Hook): + def __init__(self): + self.events: list[tuple] = [] + + async def after_ingest(self, client, document): + self.events.append(("ingest", document.id, document.uri)) + + async def after_delete(self, client, document_id): + self.events.append(("delete", document_id)) + + +class AppendTokenHook(Hook): + def __init__(self, token: str = "expanded"): + self.token = token + + async def before_search(self, client, query, filter): + return f"{query} {self.token}", filter + + +class FilterHook(Hook): + async def before_search(self, client, query, filter): + return query, "uri = 'mem://hooked'" + + +class ReverseResultsHook(Hook): + async def after_search(self, client, query, results): + self.seen_query = query + return list(reversed(results)) + + +class _EntryPointStub: + def __init__(self, factory): + self._factory = factory + + def load(self): + return self._factory + + +class _BrokenEntryPoint: + def load(self): + raise AssertionError("unreferenced entry point must not be loaded") + + +def test_build_hooks_unknown_name_raises(): + with pytest.raises(ValueError, match=ENTRY_POINT_GROUP): + build_hooks(["missing"], {}) + + +def test_build_hooks_loads_lazily_in_configured_order(): + discovered = { + "recording": _EntryPointStub(RecordingHook), + "append": _EntryPointStub(AppendTokenHook), + "broken": _BrokenEntryPoint(), + } + hooks = build_hooks(["append", "recording"], discovered) + assert [type(h) for h in hooks] == [AppendTokenHook, RecordingHook] + + +def test_client_init_unknown_hook_raises(temp_db_path): + config = Config.model_copy(deep=True) + config.hooks = ["missing"] + with pytest.raises(ValueError, match="missing"): + HaikuRAG(temp_db_path, config=config, create=True) + + +def test_client_builds_hooks_from_entry_points(temp_db_path, monkeypatch): + class _NamedEntryPoint: + name = "recording" + + def load(self): + return RecordingHook + + def fake_entry_points(group): + assert group == ENTRY_POINT_GROUP + return [_NamedEntryPoint()] + + monkeypatch.setattr("haiku.rag.hooks.entry_points", fake_entry_points) + config = Config.model_copy(deep=True) + config.hooks = ["recording"] + client = HaikuRAG(temp_db_path, config=config, create=True) + assert len(client._hooks) == 1 + assert isinstance(client._hooks[0], RecordingHook) + + +def test_client_without_hooks_builds_none(temp_db_path): + client = HaikuRAG(temp_db_path, create=True) + assert client._hooks == [] + + +async def _capture_repo_search(client): + captured = {} + + async def fake_search(query, limit, search_type=None, filter=None, **kwargs): + captured["query"] = query + captured["filter"] = filter + return [] + + client.chunk_repository.search = fake_search + return captured + + +@pytest.mark.asyncio +async def test_before_search_hooks_chain_in_order(temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as client: + client._hooks = [AppendTokenHook("one"), AppendTokenHook("two"), FilterHook()] + captured = await _capture_repo_search(client) + + await client.search("alpha") + + assert captured["query"] == "alpha one two" + assert captured["filter"] == "uri = 'mem://hooked'" + + +class SpyBeforeSearchHook(Hook): + def __init__(self): + self.called: list[str] = [] + + async def before_search(self, client, query, filter): + self.called.append(query) + return query, filter + + +@pytest.mark.asyncio +async def test_before_search_skips_non_text_queries(temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as client: + hook = SpyBeforeSearchHook() + client._hooks = [hook] + + captured = {} + + async def fake_search(query, limit, search_type=None, filter=None, **kwargs): + captured["query"] = query + return [] + + client.chunk_repository.search = fake_search + + async def fake_embed_image(image): + return [0.1] * Config.embeddings.model.vector_dim + + client.store.embedder.embed_image = fake_embed_image + client.store.embedder.supports_images = True + + await client.search(b"image-bytes") + + assert hook.called == [] + + +@pytest.mark.asyncio +async def test_after_search_transforms_results(temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as client: + hook = ReverseResultsHook() + client._hooks = [hook] + + chunks = [ + (Chunk(id="c1", content="first", document_id="d1", order=0), 0.9), + (Chunk(id="c2", content="second", document_id="d1", order=1), 0.5), + ] + + async def fake_search(query, limit, search_type=None, filter=None, **kwargs): + return chunks + + client.chunk_repository.search = fake_search + + results = await client.search("alpha", include_images=False) + + assert [r.content for r in results] == ["second", "first"] + assert hook.seen_query == "alpha" + + +@pytest.mark.asyncio +async def test_after_ingest_fires_on_import_batch_update(temp_db_path): + spy = RecordingHook() + dim = Config.embeddings.model.vector_dim + + async with HaikuRAG(temp_db_path, create=True) as client: + client._hooks = [spy] + + doc = await client.import_document( + _docling_doc("a", "Alpha body"), + [Chunk(content="Alpha body", embedding=[0.1] * dim, order=0)], + uri="mem://a", + title="Alpha", + ) + assert spy.events == [("ingest", doc.id, "mem://a")] + + spy.events.clear() + batch = await client.import_documents( + [ + _import("b", "Beta body", uri="mem://b", title="Beta"), + _import("c", "Gamma body", uri="mem://c", title="Gamma"), + ] + ) + assert spy.events == [ + ("ingest", batch[0].id, "mem://b"), + ("ingest", batch[1].id, "mem://c"), + ] + + spy.events.clear() + assert doc.id is not None + await client.update_document( + doc.id, + docling_document=_docling_doc("a2", "Alpha updated"), + chunks=[Chunk(content="Alpha updated", embedding=[0.2] * dim, order=0)], + ) + assert spy.events == [("ingest", doc.id, "mem://a")] + + +@pytest.mark.asyncio +async def test_metadata_only_update_does_not_fire_after_ingest(temp_db_path): + spy = RecordingHook() + dim = Config.embeddings.model.vector_dim + + async with HaikuRAG(temp_db_path, create=True) as client: + client._hooks = [spy] + doc = await client.import_document( + _docling_doc("a", "Alpha body"), + [Chunk(content="Alpha body", embedding=[0.1] * dim, order=0)], + uri="mem://a", + title="Alpha", + ) + assert doc.id is not None + spy.events.clear() + + await client.update_document(doc.id, title="Renamed") + + assert spy.events == [] + + +@pytest.mark.asyncio +async def test_after_delete_fires_for_cascade(temp_db_path): + spy = RecordingHook() + dim = Config.embeddings.model.vector_dim + + async with HaikuRAG(temp_db_path, create=True) as client: + client._hooks = [spy] + parent = await client.import_document( + _docling_doc("p", "Parent body"), + [Chunk(content="Parent body", embedding=[0.1] * dim, order=0)], + uri="mem://parent", + title="Parent", + ) + child = await client.import_document( + _docling_doc("k", "Child body"), + [Chunk(content="Child body", embedding=[0.1] * dim, order=0)], + uri="mem://child", + title="Child", + metadata={"parent_uri": "mem://parent"}, + ) + assert parent.id is not None and child.id is not None + spy.events.clear() + + assert await client.delete_document(parent.id) is True + + deleted = {event[1] for event in spy.events} + assert deleted == {parent.id, child.id} + assert all(event[0] == "delete" for event in spy.events) + + +@pytest.mark.asyncio +async def test_delete_missing_document_fires_nothing(temp_db_path): + spy = RecordingHook() + async with HaikuRAG(temp_db_path, create=True) as client: + client._hooks = [spy] + assert await client.delete_document("does-not-exist") is False + assert spy.events == []