Add client lifecycle hooks (after_ingest, after_delete, before_search, after_search)

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

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(

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
 
 

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
 
 

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)

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

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 == []