diff --git a/docs/hooks.md b/docs/hooks.md index 8f894983..a757d956 100644 --- a/docs/hooks.md +++ b/docs/hooks.md @@ -10,12 +10,12 @@ Subclass `haiku.rag.hooks.Hook` and override any subset: | Method | Fires | Use for | |--------|-------|---------| -| `after_ingest(client, document)` | A document's content was written (create, import, batch import, update) | Deriving state from documents | -| `after_delete(client, document_id)` | A document was deleted, once per document in a cascade | Cleaning up derived state | -| `before_search(client, query, filter)` | Before retrieval, text queries only | Query expansion, filter injection | -| `after_search(client, query, results)` | After retrieval, reranking, and deduplication | Annotating, reordering, or filtering results | +| `after_ingest(client, event)` | Document content was written (create, import, batch import, update) | Deriving state from documents | +| `after_delete(client, event)` | Documents were deleted | Cleaning up derived state | +| `before_search(client, request)` | Before retrieval, text queries only | Query expansion, filter injection | +| `after_search(client, request, results)` | After retrieval, reranking, and deduplication | Annotating, reordering, or filtering results | -`before_search` returns the `(query, filter)` pair to search with. The returned query feeds both the vector and the full-text side. `after_search` returns the result list. Hooks run in the order listed in config, each receiving the previous hook's output. +Events are batch shaped. `IngestEvent` carries `documents` (a batch import arrives as one event with the whole batch) and `operation` (`"create"` or `"update"`). `DeleteEvent` carries the deleted `documents` in their last-known state, since the rows are already gone; a cascade arrives as one event. `SearchRequest` carries `query`, `filter`, `search_type`, and `limit`. `before_search` returns the request to search with, and may modify any of its fields. The query feeds both the vector and the full-text side. `after_search` returns the result list. Hooks run in the order listed in config, each receiving the previous hook's output. Hooks receive the `HaikuRAG` client, so they can search, read repositories, and store their own state. @@ -25,11 +25,11 @@ Hooks receive the `HaikuRAG` client, so they can search, read repositories, and from haiku.rag.hooks import Hook class AbbreviationHook(Hook): - async def before_search(self, client, query, filter): - expanded = my_glossary.expand(query) - return expanded, filter + async def before_search(self, client, request): + request.query = my_glossary.expand(request.query) + return request - async def after_search(self, client, query, results): + async def after_search(self, client, request, results): for result in results: result.annotations = [ f"{term}: {definition}" @@ -60,7 +60,8 @@ An unknown name in `hooks:` raises `ValueError` when the client is constructed, ## Semantics -- **Update equals ingest.** `after_ingest` fires for both creation and content updates. Treat it as "replace any state you derived from this document". Metadata-only and title-only updates do not fire. +- **Update equals ingest.** `after_ingest` fires for both creation and content updates, with `event.operation` set to `"create"` or `"update"` (creation against an already-stored URI reports `"update"`). Treat both as "replace any state you derived from these documents". The operation is informational, for notification or sync hooks. Metadata-only and title-only updates do not fire. +- **Batch your writes.** A batch import delivers all its documents in one event. A hook keeping LanceDB state should write once per event, not once per document, to avoid creating a table version per document. - **Hooks run after the write commits.** They execute outside the store's write lock, so a hook may itself write to the database, and a hook failure never rolls back the document write. - **Rebuild does not fire hooks.** `rebuild` re-chunks and re-embeds but never changes document content, so content-derived state is unaffected. - **Backfill is your loop.** A hook enabled on an existing database can backfill by iterating `client.list_documents()` and calling its own `after_ingest`. diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 88063fbc..0b42e974 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -17,7 +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.hooks import DeleteEvent, 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 @@ -400,7 +400,7 @@ class HaikuRAG: # can't appear or move between collection and deletion. parent_uri # links a child to its parent's uri; walk transitively, guarding # against cycles. - ids_to_delete: list[str] = [] + docs_to_delete: list[Document] = [] seen: set[str] = set() queue = [await self.get_document_by_id(document_id)] while queue: @@ -408,23 +408,24 @@ class HaikuRAG: if doc is None or doc.id is None or doc.id in seen: continue seen.add(doc.id) - ids_to_delete.append(doc.id) + docs_to_delete.append(doc) if doc.uri: queue.extend( await self.list_documents(filter=parent_uri_filter(doc.uri)) ) - if not ids_to_delete: + if not docs_to_delete: return False - for doc_id in ids_to_delete: - await self.document_repository.delete(doc_id) + for doc in docs_to_delete: + assert doc.id is not None + await self.document_repository.delete(doc.id) 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) + event = DeleteEvent(documents=docs_to_delete) + for hook in self._hooks: + await hook.after_delete(self, event) 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 d6c55947..b3fa1db6 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -16,6 +16,7 @@ from haiku.rag.client.processing import ( ) from haiku.rag.client.titles import resolve_title from haiku.rag.converters import get_converter +from haiku.rag.hooks import IngestEvent from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document from haiku.rag.store.models.document_item import DocumentItem, extract_items @@ -162,8 +163,12 @@ async def _store_document_with_chunks( client._schedule_vacuum() + event = IngestEvent( + documents=[stored_doc], + operation="create" if existing is None else "update", + ) for hook in client._hooks: - await hook.after_ingest(client, stored_doc) + await hook.after_ingest(client, event) return stored_doc @@ -216,8 +221,9 @@ async def _update_document_with_chunks( client._schedule_vacuum() + event = IngestEvent(documents=[updated_doc], operation="update") for hook in client._hooks: - await hook.after_ingest(client, updated_doc) + await hook.after_ingest(client, event) return updated_doc @@ -321,9 +327,9 @@ async def _store_documents_with_chunks( client._schedule_vacuum() - for doc in created: - for hook in client._hooks: - await hook.after_ingest(client, doc) + event = IngestEvent(documents=created, operation="create") + for hook in client._hooks: + await hook.after_ingest(client, event) return created diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 13c363d1..f05f23bc 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -35,15 +35,27 @@ async def search( Returns: List of SearchResult objects ordered by relevance. """ + from haiku.rag.hooks import SearchRequest + if limit is None: limit = client._config.search.limit + request = SearchRequest( + query=query, filter=filter, search_type=search_type, limit=limit + ) + if isinstance(query, str): - if search_type is None: - search_type = "hybrid" + if request.search_type is None: + request.search_type = "hybrid" for hook in client._hooks: - query, filter = await hook.before_search(client, query, filter) + request = await hook.before_search(client, request) + + query = request.query + assert isinstance(query, str), "before_search must keep text queries text" + filter = request.filter + limit = request.limit + search_type = request.search_type or "hybrid" reranker = client.reranker @@ -83,7 +95,7 @@ async def search( await _populate_image_data(client, results) for hook in client._hooks: - results = await hook.after_search(client, query, results) + results = await hook.after_search(client, request, results) return results diff --git a/haiku_rag_slim/haiku/rag/hooks.py b/haiku_rag_slim/haiku/rag/hooks.py index a8eeb5f7..7268494a 100644 --- a/haiku_rag_slim/haiku/rag/hooks.py +++ b/haiku_rag_slim/haiku/rag/hooks.py @@ -1,16 +1,49 @@ from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from importlib.metadata import entry_points -from typing import TYPE_CHECKING, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable + +from haiku.rag.store.models.chunk import SearchResult, SearchType +from haiku.rag.store.models.document import Document 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" +IngestOperation = Literal["create", "update"] + + +@dataclass +class IngestEvent: + """Documents whose content was written in one operation. Batch imports + carry the whole batch in a single event.""" + + documents: list[Document] + operation: IngestOperation + + +@dataclass +class DeleteEvent: + """Documents removed in one operation. A cascade delete carries the root + and all its children in a single event. The documents no longer exist in + the database; the models are the last-known state.""" + + documents: list[Document] + + +@dataclass +class SearchRequest: + """The parameters a search will run with. ``before_search`` hooks may + modify ``query``, ``filter``, ``search_type``, and ``limit``.""" + + query: "str | bytes | PILImage.Image" + filter: str | None + search_type: SearchType | None + limit: int + class Hook: """Base class for client lifecycle hooks. Subclasses override any subset. @@ -23,29 +56,33 @@ class 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_ingest(self, client: "HaikuRAG", event: IngestEvent) -> None: + """Content was written for ``event.documents``. ``event.operation`` + is ``"create"`` for new documents and ``"update"`` when an existing + document's content was rewritten (including creation against an + already-stored URI). Replace any state derived from the documents + regardless of the operation: even a creation may be a retry. + 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 after_delete(self, client: "HaikuRAG", event: DeleteEvent) -> None: + """``event.documents`` were deleted; cascades arrive as one event.""" 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 + self, client: "HaikuRAG", request: SearchRequest + ) -> SearchRequest: + """Transform the search parameters before retrieval. Text queries + only; the returned request's query feeds both the vector and FTS + sides.""" + return request 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.""" + request: SearchRequest, + results: list[SearchResult], + ) -> list[SearchResult]: + """Transform or annotate search results before they are returned. + ``request`` reflects any ``before_search`` transformations.""" return results diff --git a/tests/test_hooks.py b/tests/test_hooks.py index 51d74062..c79e1b5a 100644 --- a/tests/test_hooks.py +++ b/tests/test_hooks.py @@ -11,29 +11,33 @@ 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_ingest(self, client, event): + self.events.append( + ("ingest", event.operation, tuple((d.id, d.uri) for d in event.documents)) + ) - async def after_delete(self, client, document_id): - self.events.append(("delete", document_id)) + async def after_delete(self, client, event): + self.events.append(("delete", tuple((d.id, d.uri) for d in event.documents))) 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 + async def before_search(self, client, request): + request.query = f"{request.query} {self.token}" + return request class FilterHook(Hook): - async def before_search(self, client, query, filter): - return query, "uri = 'mem://hooked'" + async def before_search(self, client, request): + request.filter = "uri = 'mem://hooked'" + return request class ReverseResultsHook(Hook): - async def after_search(self, client, query, results): - self.seen_query = query + async def after_search(self, client, request, results): + self.seen_query = request.query return list(reversed(results)) @@ -111,22 +115,30 @@ async def _capture_repo_search(client): @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()] + spy = SpyBeforeSearchHook() + client._hooks = [ + AppendTokenHook("one"), + AppendTokenHook("two"), + FilterHook(), + spy, + ] captured = await _capture_repo_search(client) await client.search("alpha") assert captured["query"] == "alpha one two" assert captured["filter"] == "uri = 'mem://hooked'" + # The request carries the resolved search parameters. + assert spy.requests == [("alpha one two", "hybrid", Config.search.limit)] class SpyBeforeSearchHook(Hook): def __init__(self): - self.called: list[str] = [] + self.requests: list[tuple] = [] - async def before_search(self, client, query, filter): - self.called.append(query) - return query, filter + async def before_search(self, client, request): + self.requests.append((request.query, request.search_type, request.limit)) + return request @pytest.mark.asyncio @@ -151,7 +163,7 @@ async def test_before_search_skips_non_text_queries(temp_db_path): await client.search(b"image-bytes") - assert hook.called == [] + assert hook.requests == [] @pytest.mark.asyncio @@ -190,8 +202,9 @@ async def test_after_ingest_fires_on_import_batch_update(temp_db_path): uri="mem://a", title="Alpha", ) - assert spy.events == [("ingest", doc.id, "mem://a")] + assert spy.events == [("ingest", "create", ((doc.id, "mem://a"),))] + # A batch import arrives as one event carrying all documents. spy.events.clear() batch = await client.import_documents( [ @@ -200,8 +213,11 @@ async def test_after_ingest_fires_on_import_batch_update(temp_db_path): ] ) assert spy.events == [ - ("ingest", batch[0].id, "mem://b"), - ("ingest", batch[1].id, "mem://c"), + ( + "ingest", + "create", + ((batch[0].id, "mem://b"), (batch[1].id, "mem://c")), + ) ] spy.events.clear() @@ -211,7 +227,17 @@ async def test_after_ingest_fires_on_import_batch_update(temp_db_path): 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")] + assert spy.events == [("ingest", "update", ((doc.id, "mem://a"),))] + + # Creation against an already-stored URI updates in place. + spy.events.clear() + await client.import_document( + _docling_doc("a3", "Alpha again"), + [Chunk(content="Alpha again", embedding=[0.3] * dim, order=0)], + uri="mem://a", + title="Alpha", + ) + assert spy.events == [("ingest", "update", ((doc.id, "mem://a"),))] @pytest.mark.asyncio @@ -260,13 +286,16 @@ async def test_after_delete_fires_for_cascade(temp_db_path): 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) + # One event for the whole cascade, carrying the deleted documents' + # last-known state (uri still resolvable). + assert len(spy.events) == 1 + kind, deleted = spy.events[0] + assert kind == "delete" + assert set(deleted) == {(parent.id, "mem://parent"), (child.id, "mem://child")} class AnnotateHook(Hook): - async def after_search(self, client, query, results): + async def after_search(self, client, request, results): for result in results: result.annotations = ["XMT: transmit"] return results