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

Hook events carry batches and search parameters

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-07-15 15:38:00 +03:00
parent b4a3c678a0
commit 8083c57246
No known key found for this signature in database
6 changed files with 157 additions and 71 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
docs/hooks.md
View file

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

19
haiku_rag_slim/haiku/rag/client/__init__.py
View file

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

16
haiku_rag_slim/haiku/rag/client/documents.py
View file

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

20
haiku_rag_slim/haiku/rag/client/search.py
View file

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

75
haiku_rag_slim/haiku/rag/hooks.py
View file

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

77
tests/test_hooks.py
View file

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